fixes
This commit is contained in:
Nicolò Boschi
parent
7a9f4ee33c
commit
5d32e79955
11 changed files with 16558 additions and 760 deletions
File diff suppressed because it is too large
Load diff
|
|
@ -1,10 +1,11 @@
|
|||
"""Benchmark Visualizer Web Service.
|
||||
|
||||
A standalone web service for visualizing benchmark results.
|
||||
Currently supports LoComo benchmark visualization.
|
||||
Supports LoComo and LongMemEval benchmark visualization.
|
||||
"""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
|
@ -63,10 +64,41 @@ async def get_locomo_results(mode: str = "search") -> dict[str, Any]:
|
|||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.get("/api/longmemeval")
|
||||
async def get_longmemeval_results() -> dict[str, Any]:
|
||||
"""Get LongMemEval benchmark results.
|
||||
|
||||
Returns pre-computed benchmark results from the longmemeval directory.
|
||||
"""
|
||||
try:
|
||||
results_path = BENCHMARKS_DIR / "longmemeval" / "results" / "benchmark_results.json"
|
||||
print(f"chcking path {results_path}")
|
||||
logging.info(f"chcking path {results_path}")
|
||||
|
||||
if not results_path.exists():
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail="Benchmark results not found. Please run the benchmark first."
|
||||
)
|
||||
|
||||
with open(results_path) as f:
|
||||
results = json.load(f)
|
||||
|
||||
return results
|
||||
except json.JSONDecodeError as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail=f"Failed to parse benchmark results: {str(e)}"
|
||||
)
|
||||
except Exception as e:
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
# Mount static files
|
||||
app.mount("/static", StaticFiles(directory=Path(__file__).parent / "static"), name="static")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
uvicorn.run(app, host="127.0.0.1", port=8001)
|
||||
|
|
|
|||
|
|
@ -17,6 +17,7 @@
|
|||
<option value="">-- Select a benchmark --</option>
|
||||
<option value="locomo-search">LoComo (search)</option>
|
||||
<option value="locomo-think">LoComo (think)</option>
|
||||
<option value="longmemeval">LongMemEval</option>
|
||||
</select>
|
||||
</div>
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,8 @@ function selectBenchmark() {
|
|||
loadLocomoResults('search');
|
||||
} else if (currentBenchmark === 'locomo-think') {
|
||||
loadLocomoResults('think');
|
||||
} else if (currentBenchmark === 'longmemeval') {
|
||||
loadLongMemEvalResults();
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -369,3 +371,259 @@ function filterAnswers() {
|
|||
}
|
||||
});
|
||||
}
|
||||
|
||||
// LongMemEval functions
|
||||
async function loadLongMemEvalResults() {
|
||||
try {
|
||||
const response = await fetch('/api/longmemeval');
|
||||
|
||||
if (!response.ok) {
|
||||
const errorData = await response.json();
|
||||
document.getElementById('benchmark-content').innerHTML = `
|
||||
<div class="error-message">
|
||||
<h3>⚠️ Benchmark Results Not Found</h3>
|
||||
<p>${errorData.detail || 'The requested benchmark results are not available.'}</p>
|
||||
<p><strong>To generate results:</strong></p>
|
||||
<pre style="background: #f5f5f5; padding: 10px; border-radius: 4px; overflow-x: auto;">cd benchmarks/longmemeval
|
||||
uv run python longmemeval_benchmark.py</pre>
|
||||
<p style="margin-top: 15px; font-size: 14px; color: #666;">
|
||||
Once the benchmark completes, refresh this page and select "LongMemEval" again.
|
||||
</p>
|
||||
</div>
|
||||
`;
|
||||
return;
|
||||
}
|
||||
|
||||
benchmarkData = await response.json();
|
||||
console.log('Loaded longmemeval data:', benchmarkData);
|
||||
renderLongMemEvalResults();
|
||||
} catch (e) {
|
||||
console.error('Error loading longmemeval results:', e);
|
||||
document.getElementById('benchmark-content').innerHTML = `
|
||||
<div class="error-message">
|
||||
<h3>❌ Error Loading Results</h3>
|
||||
<p>${e.message}</p>
|
||||
<p style="font-size: 12px; color: #666; margin-top: 10px;">Check the browser console for more details.</p>
|
||||
</div>
|
||||
`;
|
||||
}
|
||||
}
|
||||
|
||||
function renderLongMemEvalResults() {
|
||||
if (!benchmarkData) return;
|
||||
|
||||
const content = document.getElementById('benchmark-content');
|
||||
|
||||
try {
|
||||
const results = benchmarkData.item_results || [];
|
||||
const numItems = benchmarkData.num_items || results.length;
|
||||
|
||||
console.log('Rendering longmemeval results:', { resultsCount: results.length, numItems });
|
||||
|
||||
// Calculate per-category statistics
|
||||
const categoryStats = {};
|
||||
|
||||
// Aggregate across all items
|
||||
let totalInvalid = 0;
|
||||
results.forEach(item => {
|
||||
if (item.metrics && item.metrics.category_stats) {
|
||||
Object.entries(item.metrics.category_stats).forEach(([category, stats]) => {
|
||||
if (!categoryStats[category]) {
|
||||
categoryStats[category] = { name: category, correct: 0, total: 0, invalid: 0 };
|
||||
}
|
||||
categoryStats[category].correct += stats.correct || 0;
|
||||
categoryStats[category].total += stats.total || 0;
|
||||
categoryStats[category].invalid += stats.invalid || 0;
|
||||
});
|
||||
}
|
||||
if (item.metrics && item.metrics.detailed_results) {
|
||||
item.metrics.detailed_results.forEach(result => {
|
||||
if (result.is_invalid) {
|
||||
totalInvalid++;
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// Overall stats
|
||||
const totalInvalidDisplay = totalInvalid > 0
|
||||
? `<div class="stat-item">
|
||||
<div class="stat-label">Invalid Questions</div>
|
||||
<div class="stat-value" style="color: #ff9800;">${totalInvalid}</div>
|
||||
</div>`
|
||||
: '';
|
||||
|
||||
const overallHtml = `
|
||||
<div style="background: #f9f9f9; padding: 20px; border: 2px solid #333; border-radius: 8px; margin-bottom: 20px;">
|
||||
<h3 style="margin-top: 0;">LongMemEval Benchmark - Overall Performance</h3>
|
||||
${totalInvalid > 0 ? `<div style="background: #fff3cd; border: 1px solid #ffc107; padding: 10px; border-radius: 4px; margin-bottom: 15px;">
|
||||
<strong>⚠️ Note:</strong> ${totalInvalid} question(s) marked as invalid due to errors (excluded from accuracy calculation)
|
||||
</div>` : ''}
|
||||
<div class="stats-grid">
|
||||
<div class="stat-item">
|
||||
<div class="stat-label">Overall Accuracy</div>
|
||||
<div class="stat-value">${benchmarkData.overall_accuracy.toFixed(2)}%</div>
|
||||
${totalInvalid > 0 ? `<div style="font-size: 11px; color: #666; margin-top: 4px;">(${benchmarkData.total_correct} / ${benchmarkData.total_valid || (benchmarkData.total_questions - totalInvalid)})</div>` : ''}
|
||||
</div>
|
||||
<div class="stat-item">
|
||||
<div class="stat-label">Correct Answers</div>
|
||||
<div class="stat-value">${benchmarkData.total_correct} / ${benchmarkData.total_questions}</div>
|
||||
</div>
|
||||
${totalInvalidDisplay}
|
||||
<div class="stat-item">
|
||||
<div class="stat-label">Items</div>
|
||||
<div class="stat-value">${numItems}</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<h4 style="margin: 20px 0 10px 0; padding-top: 15px; border-top: 1px solid #ddd;">Accuracy by Category</h4>
|
||||
<div class="stats-grid" style="grid-template-columns: repeat(auto-fit, minmax(250px, 1fr));">
|
||||
${Object.values(categoryStats).map(cat => {
|
||||
const invalidCount = cat.invalid || 0;
|
||||
const validTotal = cat.total - invalidCount;
|
||||
const accuracy = validTotal > 0 ? ((cat.correct / validTotal) * 100).toFixed(1) : 0;
|
||||
const color = accuracy >= 70 ? '#43a047' : accuracy >= 50 ? '#ff9800' : '#e53935';
|
||||
const invalidNote = invalidCount > 0 ? ` <span style="color: #ff9800; font-size: 10px;">(${invalidCount} invalid)</span>` : '';
|
||||
return `
|
||||
<div class="stat-item">
|
||||
<div class="stat-label">${cat.name}</div>
|
||||
<div class="stat-value" style="color: ${color};">${accuracy}%</div>
|
||||
<div style="font-size: 11px; color: #666; margin-top: 4px;">${cat.correct} / ${cat.total}${invalidNote}</div>
|
||||
</div>
|
||||
`;
|
||||
}).join('')}
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
|
||||
// Filter controls
|
||||
const filterHtml = `
|
||||
<div style="margin-bottom: 20px; display: flex; gap: 10px; align-items: center;">
|
||||
<label style="font-weight: bold;">Show:</label>
|
||||
<label><input type="radio" name="answer-filter" value="all" checked onchange="filterAnswers()"> All Answers</label>
|
||||
<label><input type="radio" name="answer-filter" value="incorrect" onchange="filterAnswers()"> ❌ Incorrect Only</label>
|
||||
<label><input type="radio" name="answer-filter" value="correct" onchange="filterAnswers()"> ✅ Correct Only</label>
|
||||
${totalInvalid > 0 ? '<label><input type="radio" name="answer-filter" value="invalid" onchange="filterAnswers()"> ⚠️ Invalid Only</label>' : ''}
|
||||
</div>
|
||||
`;
|
||||
|
||||
// Build item sections
|
||||
let itemsHtml = '';
|
||||
results.forEach((item, idx) => {
|
||||
const itemId = item.item_id || `item-${idx}`;
|
||||
const accuracy = item.metrics.accuracy.toFixed(2);
|
||||
const correctCount = item.metrics.correct;
|
||||
const totalCount = item.metrics.total;
|
||||
|
||||
itemsHtml += `
|
||||
<div style="margin-bottom: 30px; border: 2px solid #333; border-radius: 8px; overflow: hidden;">
|
||||
<div style="background: #f0f0f0; padding: 15px; border-bottom: 2px solid #333; cursor: pointer;" onclick="toggleConversation(${idx})">
|
||||
<h3 style="margin: 0; display: flex; justify-content: space-between; align-items: center;">
|
||||
<span>📊 ${itemId}</span>
|
||||
<span style="font-size: 18px; color: ${accuracy >= 70 ? '#43a047' : accuracy >= 50 ? '#ff9800' : '#e53935'};">
|
||||
${accuracy}% (${correctCount}/${totalCount})
|
||||
</span>
|
||||
</h3>
|
||||
</div>
|
||||
<div id="conv-${idx}" style="display: none; padding: 20px;">
|
||||
${renderLongMemEvalItemDetails(item)}
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
});
|
||||
|
||||
content.innerHTML = overallHtml + filterHtml + itemsHtml;
|
||||
} catch (e) {
|
||||
console.error('Error rendering LongMemEval results:', e);
|
||||
content.innerHTML = `
|
||||
<div class="error-message">
|
||||
<strong>Error rendering results:</strong> ${e.message}<br>
|
||||
<pre style="margin-top: 10px; font-size: 11px; overflow: auto;">${e.stack}</pre>
|
||||
</div>
|
||||
`;
|
||||
}
|
||||
}
|
||||
|
||||
function renderLongMemEvalItemDetails(item) {
|
||||
if (!item || !item.metrics) {
|
||||
return '<div style="padding: 20px; color: #666;">No metrics available</div>';
|
||||
}
|
||||
|
||||
const results = item.metrics.detailed_results;
|
||||
if (!results || !Array.isArray(results) || results.length === 0) {
|
||||
return '<div style="padding: 20px; color: #666;">No detailed results available</div>';
|
||||
}
|
||||
|
||||
let html = '<div class="qa-results">';
|
||||
|
||||
results.forEach((result, idx) => {
|
||||
const isInvalid = result.is_invalid || false;
|
||||
const isCorrect = result.is_correct;
|
||||
const bgColor = isInvalid ? '#fff3cd' : (isCorrect ? '#e8f5e9' : '#ffebee');
|
||||
const icon = isInvalid ? '⚠️' : (isCorrect ? '✅' : '❌');
|
||||
const category = result.category || 'Unknown';
|
||||
|
||||
html += `
|
||||
<div class="qa-item" data-correct="${isCorrect}" data-invalid="${isInvalid}" style="background: ${bgColor}; padding: 15px; margin-bottom: 15px; border: 1px solid #ddd; border-radius: 8px;">
|
||||
<div style="display: flex; justify-content: space-between; align-items: flex-start; margin-bottom: 10px;">
|
||||
<div style="flex: 1;">
|
||||
<div style="font-weight: bold; font-size: 16px; margin-bottom: 8px;">
|
||||
${icon} Question ${idx + 1} ${isInvalid ? '<span style="font-size: 12px; background: #ff9800; color: white; padding: 2px 8px; border-radius: 4px; margin-left: 8px;">INVALID</span>' : ''} <span style="font-size: 12px; background: #666; color: white; padding: 2px 8px; border-radius: 4px; margin-left: 8px;">${category}</span>
|
||||
</div>
|
||||
<div style="margin-bottom: 8px;">
|
||||
<b>Q:</b> ${result.question}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div style="display: grid; grid-template-columns: 1fr 1fr; gap: 15px; margin-bottom: 10px;">
|
||||
<div>
|
||||
<div style="font-weight: bold; color: #43a047; margin-bottom: 4px;">✓ Correct Answer:</div>
|
||||
<div style="background: white; padding: 8px; border-radius: 4px; border: 1px solid #ccc;">
|
||||
${result.correct_answer}
|
||||
</div>
|
||||
</div>
|
||||
<div>
|
||||
<div style="font-weight: bold; color: ${isCorrect ? '#43a047' : '#e53935'}; margin-bottom: 4px;">
|
||||
${isCorrect ? '✓' : '✗'} Predicted Answer:
|
||||
</div>
|
||||
<div style="background: white; padding: 8px; border-radius: 4px; border: 1px solid #ccc;">
|
||||
${result.predicted_answer}
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<details style="margin-top: 10px;" ${isInvalid ? 'open' : ''}>
|
||||
<summary style="cursor: pointer; font-weight: bold; padding: 5px; background: rgba(255,255,255,0.5); border-radius: 4px;">
|
||||
📝 Show Reasoning & Retrieved Memories
|
||||
</summary>
|
||||
<div style="margin-top: 10px; padding: 10px; background: white; border-radius: 4px;">
|
||||
${isInvalid ? `<div style="margin-bottom: 10px; padding: 10px; background: #ffebee; border-left: 4px solid #e53935; border-radius: 4px;">
|
||||
<b style="color: #c62828;">⚠️ Error:</b>
|
||||
<div style="margin-top: 4px; color: #333;">${result.error || 'Question marked as invalid'}</div>
|
||||
</div>` : ''}
|
||||
<div style="margin-bottom: 10px;">
|
||||
<b>System Reasoning:</b>
|
||||
<div style="padding: 8px; background: #f5f5f5; border-radius: 4px; margin-top: 4px;">
|
||||
${result.reasoning || 'N/A'}
|
||||
</div>
|
||||
</div>
|
||||
<div style="margin-bottom: 10px;">
|
||||
<b>Judge Reasoning:</b>
|
||||
<div style="padding: 8px; background: #f5f5f5; border-radius: 4px; margin-top: 4px;">
|
||||
${result.correctness_reasoning || 'N/A'}
|
||||
</div>
|
||||
</div>
|
||||
<div>
|
||||
<b>Retrieved Memories (${result.retrieved_memories ? result.retrieved_memories.length : 0}):</b>
|
||||
${renderRetrievedMemories(result.retrieved_memories)}
|
||||
</div>
|
||||
</div>
|
||||
</details>
|
||||
</div>
|
||||
`;
|
||||
});
|
||||
|
||||
html += '</div>';
|
||||
return html;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from datetime import datetime
|
|||
from typing import List, Dict, Optional, Literal
|
||||
from openai import AsyncOpenAI
|
||||
from pydantic import BaseModel, Field
|
||||
from .llm_wrapper import OutputTooLongError
|
||||
|
||||
|
||||
class Entity(BaseModel):
|
||||
|
|
@ -439,6 +440,112 @@ Remember:
|
|||
raise last_error
|
||||
|
||||
|
||||
async def _extract_facts_with_auto_split(
|
||||
chunk: str,
|
||||
chunk_index: int,
|
||||
total_chunks: int,
|
||||
event_date: datetime,
|
||||
context: str,
|
||||
llm_config: 'LLMConfig'
|
||||
) -> List[Dict[str, str]]:
|
||||
"""
|
||||
Extract facts from a chunk with automatic splitting if output exceeds token limits.
|
||||
|
||||
If the LLM output is too long (OutputTooLongError), this function automatically
|
||||
splits the chunk in half and processes each half recursively.
|
||||
|
||||
Args:
|
||||
chunk: Text chunk to process
|
||||
chunk_index: Index of this chunk in the original list
|
||||
total_chunks: Total number of original chunks
|
||||
event_date: Reference date for temporal information
|
||||
context: Context about the conversation/document
|
||||
llm_config: LLM configuration to use
|
||||
|
||||
Returns:
|
||||
List of fact dictionaries extracted from the chunk (possibly from sub-chunks)
|
||||
"""
|
||||
import logging
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
try:
|
||||
# Try to extract facts from the full chunk
|
||||
return await _extract_facts_from_chunk(
|
||||
chunk=chunk,
|
||||
chunk_index=chunk_index,
|
||||
total_chunks=total_chunks,
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config
|
||||
)
|
||||
except OutputTooLongError as e:
|
||||
# Output exceeded token limits - split the chunk in half and retry
|
||||
logger.warning(
|
||||
f"Output too long for chunk {chunk_index + 1}/{total_chunks} "
|
||||
f"({len(chunk)} chars). Splitting in half and retrying..."
|
||||
)
|
||||
|
||||
# Split at the midpoint, preferring sentence boundaries
|
||||
mid_point = len(chunk) // 2
|
||||
|
||||
# Try to find a sentence boundary near the midpoint
|
||||
# Look for ". ", "! ", "? " within 20% of midpoint
|
||||
search_range = int(len(chunk) * 0.2)
|
||||
search_start = max(0, mid_point - search_range)
|
||||
search_end = min(len(chunk), mid_point + search_range)
|
||||
|
||||
sentence_endings = ['. ', '! ', '? ', '\n\n']
|
||||
best_split = mid_point
|
||||
|
||||
for ending in sentence_endings:
|
||||
pos = chunk.rfind(ending, search_start, search_end)
|
||||
if pos != -1:
|
||||
best_split = pos + len(ending)
|
||||
break
|
||||
|
||||
# Split the chunk
|
||||
first_half = chunk[:best_split].strip()
|
||||
second_half = chunk[best_split:].strip()
|
||||
|
||||
logger.info(
|
||||
f"Split chunk {chunk_index + 1} into two sub-chunks: "
|
||||
f"{len(first_half)} chars and {len(second_half)} chars"
|
||||
)
|
||||
|
||||
# Process both halves recursively (in parallel)
|
||||
sub_tasks = [
|
||||
_extract_facts_with_auto_split(
|
||||
chunk=first_half,
|
||||
chunk_index=chunk_index,
|
||||
total_chunks=total_chunks,
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config
|
||||
),
|
||||
_extract_facts_with_auto_split(
|
||||
chunk=second_half,
|
||||
chunk_index=chunk_index,
|
||||
total_chunks=total_chunks,
|
||||
event_date=event_date,
|
||||
context=context,
|
||||
llm_config=llm_config
|
||||
)
|
||||
]
|
||||
|
||||
sub_results = await asyncio.gather(*sub_tasks)
|
||||
|
||||
# Combine results from both halves
|
||||
all_facts = []
|
||||
for sub_result in sub_results:
|
||||
all_facts.extend(sub_result)
|
||||
|
||||
logger.info(
|
||||
f"Successfully extracted {len(all_facts)} facts from split chunk {chunk_index + 1}"
|
||||
)
|
||||
|
||||
return all_facts
|
||||
|
||||
|
||||
async def extract_facts_from_text(
|
||||
text: str,
|
||||
event_date: datetime,
|
||||
|
|
@ -452,6 +559,9 @@ async def extract_facts_from_text(
|
|||
For large texts (>chunk_size chars), automatically chunks at sentence boundaries
|
||||
to avoid hitting output token limits. Processes ALL chunks in PARALLEL for speed.
|
||||
|
||||
If a chunk produces output that exceeds token limits (OutputTooLongError), it is
|
||||
automatically split in half and retried recursively until successful.
|
||||
|
||||
Args:
|
||||
text: Input text (conversation, article, etc.)
|
||||
event_date: Reference date for resolving relative times
|
||||
|
|
@ -473,7 +583,7 @@ async def extract_facts_from_text(
|
|||
|
||||
chunks = chunk_text(text, max_chars=chunk_size)
|
||||
tasks = [
|
||||
_extract_facts_from_chunk(
|
||||
_extract_facts_with_auto_split(
|
||||
chunk=chunk,
|
||||
chunk_index=i,
|
||||
total_chunks=len(chunks),
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import os
|
|||
import time
|
||||
import asyncio
|
||||
from typing import Optional, Any, Dict, List
|
||||
from openai import AsyncOpenAI, RateLimitError, APIError, APIStatusError
|
||||
from openai import AsyncOpenAI, RateLimitError, APIError, APIStatusError, LengthFinishReasonError
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
|
@ -14,6 +14,17 @@ logger = logging.getLogger(__name__)
|
|||
logging.getLogger("httpx").setLevel(logging.WARNING)
|
||||
|
||||
|
||||
class OutputTooLongError(Exception):
|
||||
"""
|
||||
Bridge exception raised when LLM output exceeds token limits.
|
||||
|
||||
This wraps provider-specific errors (e.g., OpenAI's LengthFinishReasonError)
|
||||
to allow callers to handle output length issues without depending on
|
||||
provider-specific implementations.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
class LLMConfig:
|
||||
"""Configuration for an LLM provider."""
|
||||
|
||||
|
|
@ -134,6 +145,13 @@ class LLMConfig:
|
|||
|
||||
return result
|
||||
|
||||
except LengthFinishReasonError as e:
|
||||
# Output exceeded token limits - raise bridge exception for caller to handle
|
||||
logger.warning(f"LLM output exceeded token limits: {str(e)}")
|
||||
raise OutputTooLongError(
|
||||
f"LLM output exceeded token limits. Input may need to be split into smaller chunks."
|
||||
) from e
|
||||
|
||||
except APIStatusError as e:
|
||||
last_exception = e
|
||||
if attempt < max_retries:
|
||||
|
|
|
|||
|
|
@ -551,8 +551,7 @@ class TemporalSemanticMemory(
|
|||
# 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
|
||||
CHARS_PER_BATCH = 500_000
|
||||
|
||||
if total_chars > CHARS_PER_BATCH:
|
||||
# Split into smaller batches based on character count
|
||||
|
|
|
|||
|
|
@ -3,6 +3,7 @@ Web interface for memory system.
|
|||
|
||||
Provides FastAPI app and visualization interface.
|
||||
"""
|
||||
from .server import app, create_app
|
||||
from memora.api import create_app
|
||||
from .server import app
|
||||
|
||||
__all__ = ["app", "create_app"]
|
||||
|
|
|
|||
|
|
@ -4,758 +4,12 @@ FastAPI server for memory graph visualization and API.
|
|||
Provides REST API endpoints for memory operations and serves
|
||||
the interactive visualization interface.
|
||||
"""
|
||||
import asyncio
|
||||
from fastapi import FastAPI, HTTPException
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.responses import FileResponse
|
||||
from pydantic import BaseModel
|
||||
from dotenv import load_dotenv
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Optional, List, Dict, Any
|
||||
from datetime import datetime
|
||||
|
||||
# Import from parent memora package
|
||||
from memora import TemporalSemanticMemory
|
||||
from memora.embeddings import Embeddings
|
||||
|
||||
import logging
|
||||
import os
|
||||
import argparse
|
||||
|
||||
|
||||
|
||||
# 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
|
||||
|
||||
|
||||
def create_app(memory: TemporalSemanticMemory) -> FastAPI:
|
||||
"""
|
||||
Create and configure the FastAPI application.
|
||||
|
||||
Args:
|
||||
memory: TemporalSemanticMemory instance (already initialized with required parameters)
|
||||
|
||||
Returns:
|
||||
Configured FastAPI application
|
||||
"""
|
||||
app = FastAPI(
|
||||
title="Agent Memory API",
|
||||
version="1.0.0",
|
||||
description="""
|
||||
A temporal-semantic memory system for AI agents that stores, retrieves, and reasons over memories.
|
||||
|
||||
## Features
|
||||
|
||||
* **Batch Memory Storage**: Store multiple memories efficiently with automatic fact extraction
|
||||
* **Semantic Search**: Find relevant memories using natural language queries
|
||||
* **Fact Type Filtering**: Search across world facts, agent actions, and opinions separately
|
||||
* **Think Endpoint**: Generate contextual answers based on agent identity and memories
|
||||
* **Graph Visualization**: Interactive memory graph visualization
|
||||
* **Document Tracking**: Track and manage memory documents with upsert support
|
||||
|
||||
## Architecture
|
||||
|
||||
The system uses:
|
||||
- **Temporal Links**: Connect memories that are close in time
|
||||
- **Semantic Links**: Connect semantically similar memories
|
||||
- **Entity Links**: Connect memories that mention the same entities
|
||||
- **Spreading Activation**: Intelligent traversal for memory retrieval
|
||||
""",
|
||||
contact={
|
||||
"name": "Memory System",
|
||||
},
|
||||
license_info={
|
||||
"name": "Apache 2.0",
|
||||
"url": "https://www.apache.org/licenses/LICENSE-2.0.html",
|
||||
}
|
||||
)
|
||||
|
||||
# Mount static files
|
||||
app.mount("/static", StaticFiles(directory=str(Path(__file__).parent / "static")), name="static")
|
||||
|
||||
@app.on_event("startup")
|
||||
async def startup_event():
|
||||
"""Initialize memory system on startup."""
|
||||
await memory.initialize()
|
||||
logging.info("Memory system initialized")
|
||||
|
||||
@app.on_event("shutdown")
|
||||
async def shutdown_event():
|
||||
"""Cleanup memory system on shutdown."""
|
||||
await memory.close()
|
||||
logging.info("Memory system closed")
|
||||
|
||||
# Store memory instance on app for route handlers to access
|
||||
app.state.memory = memory
|
||||
|
||||
# Register all routes
|
||||
_register_routes(app)
|
||||
|
||||
return app
|
||||
|
||||
|
||||
class SearchRequest(BaseModel):
|
||||
"""Request model for search endpoint."""
|
||||
query: str
|
||||
agent_id: str = "default"
|
||||
thinking_budget: int = 100
|
||||
max_tokens: int = 4096
|
||||
reranker: str = "heuristic"
|
||||
trace: bool = False
|
||||
fact_type: Optional[str] = None
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"query": "What did Alice say about machine learning?",
|
||||
"agent_id": "user123",
|
||||
"thinking_budget": 100,
|
||||
"max_tokens": 4096,
|
||||
"reranker": "heuristic",
|
||||
"trace": True,
|
||||
"fact_type": "world"
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class SearchResponse(BaseModel):
|
||||
"""Response model for search endpoints."""
|
||||
results: List[Dict[str, Any]]
|
||||
trace: Optional[Dict[str, Any]] = None
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"results": [
|
||||
{
|
||||
"text": "Alice works at Google on the AI team",
|
||||
"score": 0.95,
|
||||
"id": "123e4567-e89b-12d3-a456-426614174000"
|
||||
}
|
||||
],
|
||||
"trace": {
|
||||
"query": "What did Alice say about machine learning?",
|
||||
"num_results": 1,
|
||||
"time_seconds": 0.123
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class MemoryItem(BaseModel):
|
||||
"""Single memory item for batch put."""
|
||||
content: str
|
||||
event_date: Optional[datetime] = None
|
||||
context: Optional[str] = None
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"content": "Alice mentioned she's working on a new ML model",
|
||||
"event_date": "2024-01-15T10:30:00Z",
|
||||
"context": "team meeting"
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class BatchPutRequest(BaseModel):
|
||||
"""Request model for batch put endpoint."""
|
||||
agent_id: str
|
||||
items: List[MemoryItem]
|
||||
document_id: Optional[str] = None
|
||||
document_metadata: Optional[Dict[str, Any]] = None
|
||||
upsert: bool = False
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"agent_id": "user123",
|
||||
"items": [
|
||||
{
|
||||
"content": "Alice works at Google",
|
||||
"context": "work"
|
||||
},
|
||||
{
|
||||
"content": "Bob went hiking yesterday",
|
||||
"event_date": "2024-01-15T10:00:00Z"
|
||||
}
|
||||
],
|
||||
"document_id": "conversation_123",
|
||||
"upsert": False
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class BatchPutResponse(BaseModel):
|
||||
"""Response model for batch put endpoint."""
|
||||
success: bool
|
||||
message: str
|
||||
agent_id: str
|
||||
document_id: Optional[str] = None
|
||||
items_count: int
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"success": True,
|
||||
"message": "Successfully stored 2 memory items",
|
||||
"agent_id": "user123",
|
||||
"document_id": "conversation_123",
|
||||
"items_count": 2
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
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
|
||||
agent_id: str = "default"
|
||||
thinking_budget: int = 50
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"query": "What do you think about artificial intelligence?",
|
||||
"agent_id": "user123",
|
||||
"thinking_budget": 50
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class OpinionItem(BaseModel):
|
||||
"""Model for an opinion with confidence score."""
|
||||
text: str
|
||||
confidence: float
|
||||
|
||||
|
||||
class ThinkResponse(BaseModel):
|
||||
"""Response model for think endpoint."""
|
||||
text: str
|
||||
based_on: Dict[str, List[Dict[str, Any]]] # {"world": [...], "agent": [...], "opinion": [...]}
|
||||
new_opinions: List[OpinionItem] = [] # List of newly formed opinions with confidence
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"text": "Based on my understanding, AI is a transformative technology...",
|
||||
"based_on": {
|
||||
"world": [{"text": "AI is used in healthcare", "score": 0.9}],
|
||||
"agent": [{"text": "I discussed AI applications last week", "score": 0.85}],
|
||||
"opinion": [{"text": "I believe AI should be used ethically", "score": 0.8}]
|
||||
},
|
||||
"new_opinions": [
|
||||
{"text": "AI has great potential when used responsibly", "confidence": 0.95}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class AgentsResponse(BaseModel):
|
||||
"""Response model for agents list endpoint."""
|
||||
agents: List[str]
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"agents": ["user123", "agent_alice", "agent_bob"]
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class GraphDataResponse(BaseModel):
|
||||
"""Response model for graph data endpoint."""
|
||||
nodes: List[Dict[str, Any]]
|
||||
edges: List[Dict[str, Any]]
|
||||
table_rows: List[Dict[str, Any]]
|
||||
total_units: int
|
||||
|
||||
class Config:
|
||||
json_schema_extra = {
|
||||
"example": {
|
||||
"nodes": [
|
||||
{"id": "1", "label": "Alice works at Google", "type": "world"},
|
||||
{"id": "2", "label": "Bob went hiking", "type": "world"}
|
||||
],
|
||||
"edges": [
|
||||
{"from": "1", "to": "2", "type": "semantic", "weight": 0.8}
|
||||
],
|
||||
"table_rows": [
|
||||
{"id": "abc12345...", "text": "Alice works at Google", "context": "Work info", "date": "2024-01-15 10:30", "entities": "Alice (PERSON), Google (ORGANIZATION)"}
|
||||
],
|
||||
"total_units": 2
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
def _register_routes(app: FastAPI):
|
||||
"""Register all API routes on the given app instance."""
|
||||
|
||||
@app.get("/", include_in_schema=False)
|
||||
async def index():
|
||||
"""Serve the visualization page."""
|
||||
return FileResponse(str(Path(__file__).parent / "templates" / "index.html"))
|
||||
|
||||
|
||||
@app.get(
|
||||
"/api/graph",
|
||||
response_model=GraphDataResponse,
|
||||
tags=["Visualization"],
|
||||
summary="Get memory graph data",
|
||||
description="Retrieve graph data for visualization, optionally filtered by agent_id and fact_type (world/agent/opinion)"
|
||||
)
|
||||
async def api_graph(
|
||||
agent_id: Optional[str] = None,
|
||||
fact_type: Optional[str] = None
|
||||
):
|
||||
"""Get graph data from database, optionally filtered by agent_id and fact_type."""
|
||||
try:
|
||||
data = await app.state.memory.get_graph_data(agent_id, fact_type)
|
||||
return data
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
print(f"Error in /api/graph: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.post(
|
||||
"/api/search",
|
||||
response_model=SearchResponse,
|
||||
tags=["Search"],
|
||||
summary="Search memory",
|
||||
description="Search memory using semantic similarity and spreading activation. Optionally filter by fact_type (world, agent, opinion)"
|
||||
)
|
||||
async def api_search(request: SearchRequest):
|
||||
"""Run a search and return results with trace."""
|
||||
try:
|
||||
# Run search with tracing
|
||||
results, trace = await app.state.memory.search_async(
|
||||
agent_id=request.agent_id,
|
||||
query=request.query,
|
||||
thinking_budget=request.thinking_budget,
|
||||
max_tokens=request.max_tokens,
|
||||
enable_trace=request.trace,
|
||||
reranker=request.reranker,
|
||||
fact_type=request.fact_type
|
||||
)
|
||||
|
||||
# Convert trace to dict
|
||||
trace_dict = trace.to_dict() if trace else None
|
||||
|
||||
return SearchResponse(
|
||||
results=results,
|
||||
trace=trace_dict
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
print(f"Error in /api/search: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.post(
|
||||
"/api/world_search",
|
||||
response_model=SearchResponse,
|
||||
tags=["Search"],
|
||||
summary="Search world facts",
|
||||
description="Search only world facts - general knowledge about people, places, events, and things that happen"
|
||||
)
|
||||
async def api_world_search(request: SearchRequest):
|
||||
"""Search only world facts (general knowledge about the world)."""
|
||||
try:
|
||||
# Run search with fact_type filter for 'world'
|
||||
results, trace = await app.state.memory.search_async(
|
||||
agent_id=request.agent_id,
|
||||
query=request.query,
|
||||
thinking_budget=request.thinking_budget,
|
||||
max_tokens=request.max_tokens,
|
||||
enable_trace=request.trace,
|
||||
reranker=request.reranker,
|
||||
fact_type='world'
|
||||
)
|
||||
|
||||
# Convert trace to dict
|
||||
trace_dict = trace.to_dict() if trace else None
|
||||
|
||||
return SearchResponse(
|
||||
results=results,
|
||||
trace=trace_dict
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
print(f"Error in /api/world_search: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.post(
|
||||
"/api/agent_search",
|
||||
response_model=SearchResponse,
|
||||
tags=["Search"],
|
||||
summary="Search agent action facts",
|
||||
description="Search only agent facts - memories about what the AI agent did, actions taken, and tasks performed"
|
||||
)
|
||||
async def api_agent_search(request: SearchRequest):
|
||||
"""Search only agent facts (facts about what the agent did)."""
|
||||
try:
|
||||
# Run search with fact_type filter for 'agent'
|
||||
results, trace = await app.state.memory.search_async(
|
||||
agent_id=request.agent_id,
|
||||
query=request.query,
|
||||
thinking_budget=request.thinking_budget,
|
||||
max_tokens=request.max_tokens,
|
||||
enable_trace=request.trace,
|
||||
reranker=request.reranker,
|
||||
fact_type='agent'
|
||||
)
|
||||
|
||||
# Convert trace to dict
|
||||
trace_dict = trace.to_dict() if trace else None
|
||||
|
||||
return SearchResponse(
|
||||
results=results,
|
||||
trace=trace_dict
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
print(f"Error in /api/agent_search: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.post(
|
||||
"/api/opinion_search",
|
||||
response_model=SearchResponse,
|
||||
tags=["Search"],
|
||||
summary="Search agent opinions",
|
||||
description="Search only opinion facts - the agent's formed beliefs, perspectives, and viewpoints"
|
||||
)
|
||||
async def api_opinion_search(request: SearchRequest):
|
||||
"""Search only opinion facts (agent's formed opinions and perspectives)."""
|
||||
try:
|
||||
# Run search with fact_type filter for 'opinion'
|
||||
results, trace = await app.state.memory.search_async(
|
||||
agent_id=request.agent_id,
|
||||
query=request.query,
|
||||
thinking_budget=request.thinking_budget,
|
||||
max_tokens=request.max_tokens,
|
||||
enable_trace=request.trace,
|
||||
reranker=request.reranker,
|
||||
fact_type='opinion'
|
||||
)
|
||||
|
||||
# Convert trace to dict
|
||||
trace_dict = trace.to_dict() if trace else None
|
||||
|
||||
return SearchResponse(
|
||||
results=results,
|
||||
trace=trace_dict
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
print(f"Error in /api/opinion_search: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.post(
|
||||
"/api/think",
|
||||
response_model=ThinkResponse,
|
||||
tags=["Reasoning"],
|
||||
summary="Think and generate answer",
|
||||
description="""
|
||||
Think and formulate an answer using agent identity, world facts, and opinions.
|
||||
|
||||
This endpoint:
|
||||
1. Retrieves agent facts (agent's identity)
|
||||
2. Retrieves world facts relevant to the query
|
||||
3. Retrieves existing opinions (agent's perspectives)
|
||||
4. Uses LLM to formulate a contextual answer
|
||||
5. Extracts and stores any new opinions formed
|
||||
6. Returns plain text answer, the facts used, and new opinions
|
||||
"""
|
||||
)
|
||||
async def api_think(request: ThinkRequest):
|
||||
try:
|
||||
# Use the memory system's think_async method
|
||||
result = await app.state.memory.think_async(
|
||||
agent_id=request.agent_id,
|
||||
query=request.query,
|
||||
thinking_budget=request.thinking_budget
|
||||
)
|
||||
|
||||
return ThinkResponse(
|
||||
text=result["text"],
|
||||
based_on=result["based_on"],
|
||||
new_opinions=result.get("new_opinions", [])
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
print(f"Error in /api/think: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
@app.get(
|
||||
"/api/agents",
|
||||
response_model=AgentsResponse,
|
||||
tags=["Management"],
|
||||
summary="List all agents",
|
||||
description="Get a list of all agent IDs that have stored memories in the system"
|
||||
)
|
||||
async def api_agents():
|
||||
"""Get list of available agents from database."""
|
||||
try:
|
||||
agent_list = await app.state.memory.list_agents()
|
||||
return AgentsResponse(agents=agent_list)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
print(f"Error in /api/agents: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.get(
|
||||
"/api/stats/{agent_id}",
|
||||
tags=["Memory Statistics"],
|
||||
summary="Get memory statistics for an agent",
|
||||
description="Get statistics about nodes and links for a specific agent"
|
||||
)
|
||||
async def api_stats(agent_id: str):
|
||||
"""Get statistics about memory nodes and links for an agent."""
|
||||
try:
|
||||
pool = await app.state.memory._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
# Get node counts by fact_type
|
||||
node_stats = await conn.fetch(
|
||||
"""
|
||||
SELECT fact_type, COUNT(*) as count
|
||||
FROM memory_units
|
||||
WHERE agent_id = $1
|
||||
GROUP BY fact_type
|
||||
""",
|
||||
agent_id
|
||||
)
|
||||
|
||||
# Get link counts by link_type
|
||||
link_stats = await conn.fetch(
|
||||
"""
|
||||
SELECT ml.link_type, COUNT(*) as count
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.from_unit_id = mu.id
|
||||
WHERE mu.agent_id = $1
|
||||
GROUP BY ml.link_type
|
||||
""",
|
||||
agent_id
|
||||
)
|
||||
|
||||
# Format results
|
||||
nodes_by_type = {row['fact_type']: row['count'] for row in node_stats}
|
||||
links_by_type = {row['link_type']: row['count'] for row in link_stats}
|
||||
|
||||
total_nodes = sum(nodes_by_type.values())
|
||||
total_links = sum(links_by_type.values())
|
||||
|
||||
return {
|
||||
"agent_id": agent_id,
|
||||
"total_nodes": total_nodes,
|
||||
"total_links": total_links,
|
||||
"nodes_by_type": nodes_by_type,
|
||||
"links_by_type": links_by_type
|
||||
}
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
print(f"Error in /api/stats/{agent_id}: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
@app.post(
|
||||
"/api/memories/batch",
|
||||
response_model=BatchPutResponse,
|
||||
tags=["Memory Storage"],
|
||||
summary="Store multiple memories",
|
||||
description="""
|
||||
Store multiple memory items in batch with automatic fact extraction.
|
||||
|
||||
Features:
|
||||
- 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. Extracts semantic facts from the content
|
||||
2. Generates embeddings
|
||||
3. Deduplicates similar facts
|
||||
4. Creates temporal, semantic, and entity links
|
||||
5. Tracks document metadata
|
||||
"""
|
||||
)
|
||||
async def api_batch_put(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)
|
||||
|
||||
# Call put_batch_async
|
||||
result = await app.state.memory.put_batch_async(
|
||||
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 result: {result}")
|
||||
|
||||
return BatchPutResponse(
|
||||
success=True,
|
||||
message=f"Successfully stored {len(contents)} memory items",
|
||||
agent_id=request.agent_id,
|
||||
document_id=request.document_id,
|
||||
items_count=len(contents)
|
||||
)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
print(f"Error in /api/memories/batch: {error_detail}")
|
||||
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"],
|
||||
summary="Delete a memory unit",
|
||||
description="Delete a single memory unit and all its associated links (temporal, semantic, and entity links)"
|
||||
)
|
||||
async def api_delete_memory_unit(unit_id: str):
|
||||
"""Delete a memory unit and all its links."""
|
||||
try:
|
||||
result = await app.state.memory.delete_memory_unit(unit_id)
|
||||
|
||||
if not result["success"]:
|
||||
raise HTTPException(status_code=404, detail=result["message"])
|
||||
|
||||
return result
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
import traceback
|
||||
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
|
||||
print(f"Error in /api/memory/{unit_id}: {error_detail}")
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
|
||||
|
||||
from memora import TemporalSemanticMemory
|
||||
from memora.api import create_app
|
||||
|
||||
|
||||
# Create app at module level (required for uvicorn import string)
|
||||
|
|
@ -771,7 +25,7 @@ app = create_app(_memory)
|
|||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
import argparse
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
# Parse CLI arguments
|
||||
|
|
|
|||
|
|
@ -9,5 +9,4 @@ echo "Server will be available at: http://localhost:8001"
|
|||
echo ""
|
||||
|
||||
cd benchmarks/visualizer
|
||||
open http://localhost:8001
|
||||
uv run uvicorn server:app --reload --host 0.0.0.0 --port 8001
|
||||
|
|
|
|||
|
|
@ -82,5 +82,4 @@ if [[ ${#SERVER_ARGS[@]} -eq 0 ]]; then
|
|||
SERVER_ARGS=(--reload --host 0.0.0.0 --port 8080)
|
||||
fi
|
||||
|
||||
open "http://localhost:${PORT}"
|
||||
uv run python -m memora.web.server "${SERVER_ARGS[@]}"
|
||||
|
|
|
|||
Loading…
Reference in a new issue