Merge branch 'main' of http://gitea.bubuit.net:3000/oib/aitbc
Some checks failed
Cross-Node Transaction Testing / transaction-test (push) Has been cancelled
Deploy to Testnet / deploy-testnet (push) Has been cancelled
Documentation Validation / validate-docs (push) Has been cancelled
Documentation Validation / validate-policies-strict (push) Has been cancelled
Multi-Node Stress Testing / stress-test (push) Has been cancelled
Node Failover Simulation / failover-test (push) Has been cancelled

This commit is contained in:
aitbc
2026-05-07 21:09:16 +02:00
2 changed files with 77 additions and 34 deletions

View File

@@ -1,5 +1,6 @@
from datetime import datetime, timezone
from typing import Any, Dict, List, Optional
import json
from aitbc import get_logger
from fastapi import APIRouter, BackgroundTasks, Depends, HTTPException, Query, Response
@@ -60,34 +61,36 @@ async def send_message(request: MessageRequest):
payload=request.payload
)
# Send message with specified protocol
success = await state.communication_manager.send_message(protocol, message)
# Send message
# Store message in Redis first (always)
message_data = {
"sender_id": message.sender_id,
"receiver_id": message.receiver_id,
"message_type": message.message_type.value,
"priority": message.priority.value,
"payload": json.dumps(message.payload),
"protocol": protocol,
"timestamp": datetime.now(timezone.utc).isoformat()
}
if success:
# Store message in Redis for history
if state.message_storage:
message_data = {
"message_id": message.id,
"sender_id": message.sender_id,
"receiver_id": message.receiver_id,
"message_type": message.message_type.value,
"priority": message.priority.value,
"payload": json.dumps(message.payload),
"protocol": protocol,
"timestamp": datetime.now(timezone.utc).isoformat()
}
await state.message_storage.store_message(message.id, message_data)
if state.message_storage:
await state.message_storage.store_message(message.id, message_data)
return {
"status": "success",
"message": "Message sent successfully",
"message_id": message.id,
"receiver_id": request.receiver_id,
"protocol": protocol,
"sent_at": datetime.now(timezone.utc).isoformat()
}
else:
raise HTTPException(status_code=500, detail="Failed to send message")
# Try to send via protocol (optional, for real-time notification)
if state.communication_manager:
try:
await state.communication_manager.send_message(protocol, message)
except:
pass # Protocol send is optional
return {
"status": "success",
"message": "Message sent successfully",
"message_id": message.id,
"receiver_id": request.receiver_id,
"protocol": protocol,
"sent_at": datetime.now(timezone.utc).isoformat()
}
except HTTPException:
raise
@@ -150,10 +153,28 @@ async def broadcast_message(request: BroadcastRequest):
payload=request.payload
)
success = await state.communication_manager.send_message("broadcast", message)
if success:
# Store in Redis first (always)
message_data = {
"sender_id": message.sender_id,
"receiver_id": message.receiver_id,
"message_type": message.message_type.value,
"priority": message.priority.value,
"payload": json.dumps(message.payload),
"protocol": "broadcast",
"timestamp": datetime.now(timezone.utc).isoformat()
}
if state.message_storage:
await state.message_storage.store_message(message.id, message_data)
recipients.append(agent.agent_id)
# Optionally try to send via protocol (for real-time notification)
if state.communication_manager:
try:
await state.communication_manager.send_message("broadcast", message)
except:
pass # Protocol send is optional
return {
"status": "success",
"message": f"Broadcast sent to {len(recipients)} agents",
@@ -188,10 +209,16 @@ async def get_message_history(
else:
messages = await state.message_storage.get_all_messages(limit, offset)
# Get total count
total = 0
if state.message_storage:
total = await state.message_storage.get_message_count()
return {
"status": "success",
"messages": messages,
"count": len(messages),
"total": total,
"limit": limit,
"offset": offset,
"timestamp": datetime.now(timezone.utc).isoformat()

View File

@@ -49,8 +49,16 @@ class MessageStorage:
await self.redis.sadd(f"messages:receiver:{receiver_id}", message_id)
# Index by timestamp (for time-based queries)
timestamp = message_data.get("timestamp", datetime.now(timezone.utc).isoformat())
await self.redis.zadd(f"messages:timestamp", {message_id: timestamp})
timestamp_str = message_data.get("timestamp", datetime.now(timezone.utc).isoformat())
# Convert to float for sorted set
try:
# Try to parse ISO format
dt = datetime.fromisoformat(timestamp_str.replace("Z", "+00:00"))
timestamp_float = dt.timestamp()
except:
# Already a float or int
timestamp_float = float(timestamp_str)
await self.redis.zadd(f"messages:timestamp", {message_id: timestamp_float})
logger.debug(f"Stored message {message_id} in Redis")
return True
@@ -59,6 +67,14 @@ class MessageStorage:
logger.error(f"Error storing message {message_id}: {e}")
return False
async def get_message_count(self) -> int:
"""Get total count of messages"""
try:
return await self.redis.zcard("messages:timestamp")
except Exception as e:
logger.error(f"Error getting message count: {e}")
return 0
async def get_message(self, message_id: str) -> Optional[Dict[str, Any]]:
"""Retrieve a specific message by ID"""
try: