This commit is contained in:
Nicolò Boschi 2025-11-07 13:02:52 +01:00
parent 18c7e0845b
commit 27bb335d7c
6 changed files with 126 additions and 60 deletions

View file

@ -718,25 +718,23 @@ class BenchmarkRunner:
await self.memory.delete_agent(agent_id)
console.print(f" [green]✓[/green] Cleared agent data")
with Progress(
SpinnerColumn(),
TextColumn("[progress.description]{task.description}"),
BarColumn(),
TextColumn("[progress.percentage]{task.percentage:>3.0f}%"),
console=console
) as progress:
task = progress.add_task(
f"[cyan]Ingesting {len(items)} items...",
total=len(items)
)
# Collect all sessions from all items into one large batch
console.print(f" [yellow]Collecting sessions from all items...[/yellow]")
all_sessions = []
for item in items:
item_sessions = self.dataset.prepare_sessions_for_ingestion(item)
all_sessions.extend(item_sessions)
total_sessions = 0
for item in items:
num_sessions = await self.ingest_conversation(item, agent_id)
total_sessions += num_sessions
progress.update(task, advance=1)
console.print(f" [cyan]Collected {len(all_sessions)} sessions from {len(items)} items[/cyan]")
console.print(f" [yellow]Ingesting in one batch (auto-chunks if needed)...[/yellow]")
console.print(f" [green]✓[/green] Ingested {total_sessions} sessions from {len(items)} items")
# Ingest all sessions in one batch call (will auto-chunk if too large)
await self.memory.put_batch_async(
agent_id=agent_id,
contents=all_sessions
)
console.print(f" [green]✓[/green] Ingested {len(all_sessions)} sessions from {len(items)} items")
else:
console.print(f"\n[3] Skipping ingestion (using existing data)")

View file

@ -423,8 +423,6 @@ Remember:
# Convert to dict format
chunk_facts = [fact.model_dump() for fact in extraction_response.facts]
logger.info(f" [1.3.{chunk_index + 1}] Chunk {chunk_index + 1}/{total_chunks} LLM call: {len(chunk_facts)} facts from {len(chunk)} chars in {llm_call_time:.3f}s")
return chunk_facts
except BadRequestError as e:

View file

@ -476,6 +476,7 @@ class TemporalSemanticMemory(
- Extracts facts from all contents in parallel
- Generates ALL embeddings in ONE batch
- Does ALL database operations in ONE transaction
- Automatically chunks large batches to prevent timeouts
Args:
agent_id: Unique identifier for the agent
@ -505,15 +506,102 @@ class TemporalSemanticMemory(
# Returns: [["unit-id-1"], ["unit-id-2"]]
"""
start_time = time.time()
log_buffer = [] # Buffer all logs to avoid interleaving
log_buffer.append(f"{'='*60}")
log_buffer.append(f"PUT_BATCH_ASYNC START: {agent_id}")
log_buffer.append(f"Batch size: {len(contents)} content items")
log_buffer.append(f"{'='*60}")
if not contents:
return []
# Auto-chunk large batches by character count to avoid timeouts and memory issues
# Calculate total character count
total_chars = sum(len(item.get("content", "")) for item in contents)
# Threshold: 50,000 characters per sub-batch (roughly 12k tokens with average 4 chars/token)
CHARS_PER_BATCH = 50_000
if total_chars > CHARS_PER_BATCH:
# Split into smaller batches based on character count
logger.info(f"Large batch detected ({total_chars:,} chars from {len(contents)} items). Splitting into sub-batches of ~{CHARS_PER_BATCH:,} chars each...")
sub_batches = []
current_batch = []
current_batch_chars = 0
for item in contents:
item_chars = len(item.get("content", ""))
# If adding this item would exceed the limit, start a new batch
# (unless current batch is empty - then we must include it even if it's large)
if current_batch and current_batch_chars + item_chars > CHARS_PER_BATCH:
sub_batches.append(current_batch)
current_batch = [item]
current_batch_chars = item_chars
else:
current_batch.append(item)
current_batch_chars += item_chars
# Add the last batch
if current_batch:
sub_batches.append(current_batch)
logger.info(f"Split into {len(sub_batches)} sub-batches: {[len(b) for b in sub_batches]} items each")
# Process each sub-batch using internal method (skip chunking check)
all_results = []
for i, sub_batch in enumerate(sub_batches, 1):
sub_batch_chars = sum(len(item.get("content", "")) for item in sub_batch)
logger.info(f"Processing sub-batch {i}/{len(sub_batches)}: {len(sub_batch)} items, {sub_batch_chars:,} chars")
sub_results = await self._put_batch_async_internal(
agent_id=agent_id,
contents=sub_batch,
document_id=document_id,
document_metadata=document_metadata,
upsert=upsert and i == 1, # Only upsert on first batch
fact_type_override=fact_type_override,
confidence_score=confidence_score
)
all_results.extend(sub_results)
total_time = time.time() - start_time
logger.info(f"PUT_BATCH_ASYNC (chunked) COMPLETE: {len(all_results)} results from {len(contents)} contents in {total_time:.3f}s")
return all_results
# Small batch - use internal method directly
return await self._put_batch_async_internal(
agent_id=agent_id,
contents=contents,
document_id=document_id,
document_metadata=document_metadata,
upsert=upsert,
fact_type_override=fact_type_override,
confidence_score=confidence_score
)
async def _put_batch_async_internal(
self,
agent_id: str,
contents: List[Dict[str, Any]],
document_id: Optional[str] = None,
document_metadata: Optional[Dict[str, Any]] = None,
upsert: bool = False,
fact_type_override: Optional[str] = None,
confidence_score: Optional[float] = None,
) -> List[List[str]]:
"""
Internal method for batch processing without chunking logic.
Assumes contents are already appropriately sized (< 50k chars).
Called by put_batch_async after chunking large batches.
"""
start_time = time.time()
total_chars = sum(len(item.get("content", "")) for item in contents)
# Buffer all logs to avoid interleaving
log_buffer = []
log_buffer.append(f"{'='*60}")
log_buffer.append(f"PUT_BATCH_ASYNC START: {agent_id}")
log_buffer.append(f"Batch size: {len(contents)} content items, {total_chars:,} chars")
log_buffer.append(f"{'='*60}")
# Step 1: Extract facts from ALL contents in parallel
step_start = time.time()

View file

@ -3,13 +3,6 @@ Web interface for memory system.
Provides FastAPI app and visualization interface.
"""
from .server import create_app
# Lazy import of app to avoid initialization issues
def __getattr__(name):
if name == "app":
from .server import _get_app
return _get_app()
raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
from .server import app, create_app
__all__ = ["app", "create_app"]

View file

@ -22,7 +22,7 @@ from memora.embeddings import Embeddings
import logging
logging.basicConfig(level=logging.INFO)
# Environment variables are loaded by the shell script that calls this module
# No need to load .env files here as they're sourced by start-server.sh
@ -621,6 +621,7 @@ def _register_routes(app: FastAPI):
document_metadata=request.document_metadata,
upsert=request.upsert
)
logging.info(f"Batch put result: {result}")
return BatchPutResponse(
success=True,
@ -662,33 +663,21 @@ def _register_routes(app: FastAPI):
def _create_default_app():
"""Create app instance with default environment configuration."""
_memory = TemporalSemanticMemory(
db_url=os.getenv("DATABASE_URL"),
memory_llm_provider=os.getenv("MEMORY_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("MEMORY_LLM_API_KEY"),
memory_llm_model=os.getenv("MEMORY_LLM_MODEL", "openai/gpt-oss-120b"),
memory_llm_base_url=os.getenv("MEMORY_LLM_BASE_URL") or None,
)
return create_app(_memory)
# Module-level app instance (lazy initialization)
_app_instance = None
def _get_app():
"""Get or create the app instance."""
global _app_instance
if _app_instance is None:
_app_instance = _create_default_app()
return _app_instance
# Create app at module level (required for uvicorn import string)
_memory = TemporalSemanticMemory(
db_url=os.getenv("DATABASE_URL"),
memory_llm_provider=os.getenv("MEMORY_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("MEMORY_LLM_API_KEY"),
memory_llm_model=os.getenv("MEMORY_LLM_MODEL", "openai/gpt-oss-120b"),
memory_llm_base_url=os.getenv("MEMORY_LLM_BASE_URL") or None,
)
app = create_app(_memory)
if __name__ == "__main__":
import uvicorn
import argparse
logging.basicConfig(level=logging.INFO)
# Parse CLI arguments
parser = argparse.ArgumentParser(description="Memory Graph API Server")
@ -708,9 +697,6 @@ if __name__ == "__main__":
args = parser.parse_args()
# Create app after parsing args (so --help works without DB connection)
app = _create_default_app()
print("\n" + "=" * 80)
print("Memory Graph API Server")
print("=" * 80)
@ -721,9 +707,12 @@ if __name__ == "__main__":
print(f"Log Level: {args.log_level}")
print("=" * 80 + "\n")
# Always use import string for uvicorn (required for reload and workers)
app_ref = "memora.web.server:app"
# Prepare uvicorn config
uvicorn_config = {
"app": app,
"app": app_ref,
"host": args.host,
"port": args.port,
"reload": args.reload,

View file

@ -1127,7 +1127,7 @@ window.runSearchInPane = async function(paneId) {
return;
}
try:
try {
// Prepare request body with optional fact_type
const requestBody = {
query: query,