This commit is contained in:
Nicolò Boschi 2025-11-07 13:22:36 +01:00
parent 27bb335d7c
commit 7a9f4ee33c
2 changed files with 380 additions and 244 deletions

View file

@ -154,6 +154,11 @@ class TemporalSemanticMemory(
# we use ~20-40 connections max, staying well within pool limits # we use ~20-40 connections max, staying well within pool limits
self._search_semaphore = asyncio.Semaphore(10) self._search_semaphore = asyncio.Semaphore(10)
# Backpressure for put operations: limit concurrent puts to prevent database contention
# Each put_batch holds a connection for the entire transaction, so we limit to 5
# concurrent puts to avoid connection pool exhaustion and reduce write contention
self._put_semaphore = asyncio.Semaphore(5)
async def _handle_access_count_update(self, task_dict: Dict[str, Any]): async def _handle_access_count_update(self, task_dict: Dict[str, Any]):
""" """
Handler for access count update tasks. Handler for access count update tasks.
@ -177,6 +182,36 @@ class TemporalSemanticMemory(
except Exception as e: except Exception as e:
logger.error(f"Access count handler: Error updating access counts: {e}") logger.error(f"Access count handler: Error updating access counts: {e}")
async def _handle_batch_put(self, task_dict: Dict[str, Any]):
"""
Handler for batch put tasks.
Args:
task_dict: Dict with 'agent_id', 'contents', 'document_id', 'document_metadata', 'upsert'
"""
try:
agent_id = task_dict.get('agent_id')
contents = task_dict.get('contents', [])
document_id = task_dict.get('document_id')
document_metadata = task_dict.get('document_metadata')
upsert = task_dict.get('upsert', False)
logger.info(f"[BATCH_PUT_TASK] Starting background batch put for agent_id={agent_id}, {len(contents)} items")
await self.put_batch_async(
agent_id=agent_id,
contents=contents,
document_id=document_id,
document_metadata=document_metadata,
upsert=upsert
)
logger.info(f"[BATCH_PUT_TASK] Completed background batch put for agent_id={agent_id}")
except Exception as e:
logger.error(f"Batch put handler: Error processing batch put: {e}")
import traceback
traceback.print_exc()
async def execute_task(self, task_dict: Dict[str, Any]): async def execute_task(self, task_dict: Dict[str, Any]):
""" """
Execute a task by routing it to the appropriate handler. Execute a task by routing it to the appropriate handler.
@ -196,6 +231,8 @@ class TemporalSemanticMemory(
await self._handle_reinforce_opinion(task_dict) await self._handle_reinforce_opinion(task_dict)
elif task_type == 'form_opinion': elif task_type == 'form_opinion':
await self._handle_form_opinion(task_dict) await self._handle_form_opinion(task_dict)
elif task_type == 'batch_put':
await self._handle_batch_put(task_dict)
else: else:
logger.error(f"Unknown task type: {task_type}") logger.error(f"Unknown task type: {task_type}")
@ -591,7 +628,11 @@ class TemporalSemanticMemory(
Assumes contents are already appropriately sized (< 50k chars). Assumes contents are already appropriately sized (< 50k chars).
Called by put_batch_async after chunking large batches. Called by put_batch_async after chunking large batches.
Uses semaphore for backpressure to limit concurrent puts.
""" """
# Backpressure: limit concurrent puts to prevent database contention
async with self._put_semaphore:
start_time = time.time() start_time = time.time()
total_chars = sum(len(item.get("content", "")) for item in contents) total_chars = sum(len(item.get("content", "")) for item in contents)

View file

@ -206,6 +206,28 @@ class BatchPutResponse(BaseModel):
} }
class BatchPutAsyncResponse(BaseModel):
"""Response model for async batch put endpoint."""
success: bool
message: str
agent_id: str
document_id: Optional[str] = None
items_count: int
queued: bool
class Config:
json_schema_extra = {
"example": {
"success": True,
"message": "Batch put task queued for background processing",
"agent_id": "user123",
"document_id": "conversation_123",
"items_count": 2,
"queued": True
}
}
class ThinkRequest(BaseModel): class ThinkRequest(BaseModel):
"""Request model for think endpoint.""" """Request model for think endpoint."""
query: str query: str
@ -637,6 +659,79 @@ def _register_routes(app: FastAPI):
raise HTTPException(status_code=500, detail=str(e)) raise HTTPException(status_code=500, detail=str(e))
@app.post(
"/api/memories/batch_async",
response_model=BatchPutAsyncResponse,
tags=["Memory Storage"],
summary="Store multiple memories asynchronously",
description="""
Store multiple memory items in batch asynchronously using the task backend.
This endpoint returns immediately after queuing the task, without waiting for completion.
The actual processing happens in the background.
Features:
- Immediate response (non-blocking)
- Background processing via task queue
- Efficient batch processing
- Automatic fact extraction from natural language
- Entity recognition and linking
- Document tracking with optional upsert
- Temporal and semantic linking
The system automatically:
1. Queues the batch put task
2. Returns immediately with success=True, queued=True
3. Processes in background: extracts facts, generates embeddings, creates links
"""
)
async def api_batch_put_async(request: BatchPutRequest):
try:
# Validate agent_id - prevent writing to reserved agents
RESERVED_AGENT_IDS = {"locomo"}
if request.agent_id in RESERVED_AGENT_IDS:
raise HTTPException(
status_code=403,
detail=f"Cannot write to reserved agent_id '{request.agent_id}'. Reserved agents: {', '.join(RESERVED_AGENT_IDS)}"
)
# Prepare contents for put_batch_async
contents = []
for item in request.items:
content_dict = {"content": item.content}
if item.event_date:
content_dict["event_date"] = item.event_date
if item.context:
content_dict["context"] = item.context
contents.append(content_dict)
# Submit task to background queue
await app.state.memory._task_backend.submit_task({
'type': 'batch_put',
'agent_id': request.agent_id,
'contents': contents,
'document_id': request.document_id,
'document_metadata': request.document_metadata,
'upsert': request.upsert
})
logging.info(f"Batch put task queued for agent_id={request.agent_id}, {len(contents)} items")
return BatchPutAsyncResponse(
success=True,
message=f"Batch put task queued for background processing ({len(contents)} items)",
agent_id=request.agent_id,
document_id=request.document_id,
items_count=len(contents),
queued=True
)
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
print(f"Error in /api/memories/batch_async: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.delete( @app.delete(
"/api/memory/{unit_id}", "/api/memory/{unit_id}",
tags=["Memory Storage"], tags=["Memory Storage"],