fixes
This commit is contained in:
parent
27bb335d7c
commit
7a9f4ee33c
2 changed files with 380 additions and 244 deletions
|
|
@ -154,6 +154,11 @@ class TemporalSemanticMemory(
|
|||
# we use ~20-40 connections max, staying well within pool limits
|
||||
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]):
|
||||
"""
|
||||
Handler for access count update tasks.
|
||||
|
|
@ -177,6 +182,36 @@ class TemporalSemanticMemory(
|
|||
except Exception as 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]):
|
||||
"""
|
||||
Execute a task by routing it to the appropriate handler.
|
||||
|
|
@ -196,6 +231,8 @@ class TemporalSemanticMemory(
|
|||
await self._handle_reinforce_opinion(task_dict)
|
||||
elif task_type == 'form_opinion':
|
||||
await self._handle_form_opinion(task_dict)
|
||||
elif task_type == 'batch_put':
|
||||
await self._handle_batch_put(task_dict)
|
||||
else:
|
||||
logger.error(f"Unknown task type: {task_type}")
|
||||
|
||||
|
|
@ -591,7 +628,11 @@ class TemporalSemanticMemory(
|
|||
|
||||
Assumes contents are already appropriately sized (< 50k chars).
|
||||
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()
|
||||
total_chars = sum(len(item.get("content", "")) for item in contents)
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
"""Request model for think endpoint."""
|
||||
query: str
|
||||
|
|
@ -637,6 +659,79 @@ def _register_routes(app: FastAPI):
|
|||
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(
|
||||
"/api/memory/{unit_id}",
|
||||
tags=["Memory Storage"],
|
||||
|
|
|
|||
Loading…
Reference in a new issue