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
|
# 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)
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -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"],
|
||||||
|
|
|
||||||
Loading…
Reference in a new issue