mrr and links reinforcement
This commit is contained in:
Nicolò Boschi
parent
75f40fcc3e
commit
b3259d2907
17 changed files with 23537 additions and 10291 deletions
File diff suppressed because it is too large
Load diff
|
|
@ -2328,7 +2328,7 @@
|
|||
"img_url": [
|
||||
"https://www.speakers.co.uk/microsites/tom-oliver/wp-content/uploads/2014/11/Book-Cover-3D1.jpg"
|
||||
],
|
||||
"blip_caption": "a photography of a book cover with a gold coin on it",
|
||||
"blip_caption": "a photography of a book cover with a gold coin on it called 'Nothing is Impossible'",
|
||||
"query": "painted canvas follow your dreams",
|
||||
"dia_id": "D7:8",
|
||||
"re-download": true,
|
||||
|
|
|
|||
|
|
@ -113,7 +113,19 @@ class QuestionAnswer(pydantic.BaseModel):
|
|||
|
||||
|
||||
class LoComoAnswerGenerator(LLMAnswerGenerator):
|
||||
"""LoComo-specific answer generator using OpenAI."""
|
||||
"""LoComo-specific answer generator using Groq."""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize with Groq client."""
|
||||
groq_api_key = os.getenv('GROQ_API_KEY')
|
||||
if not groq_api_key:
|
||||
raise ValueError("GROQ_API_KEY environment variable not set")
|
||||
|
||||
base_url = os.getenv('GROQ_BASE_URL', 'https://api.groq.com/openai/v1')
|
||||
self.client = AsyncOpenAI(
|
||||
api_key=groq_api_key,
|
||||
base_url=base_url
|
||||
)
|
||||
|
||||
async def generate_answer(
|
||||
self,
|
||||
|
|
@ -121,23 +133,22 @@ class LoComoAnswerGenerator(LLMAnswerGenerator):
|
|||
memories: List[Dict[str, Any]]
|
||||
) -> Tuple[str, str]:
|
||||
"""
|
||||
Generate answer from retrieved memories using OpenAI.
|
||||
Generate answer from retrieved memories using Groq.
|
||||
|
||||
Returns:
|
||||
Tuple of (answer, reasoning)
|
||||
"""
|
||||
# Format context
|
||||
context_parts = []
|
||||
for i, result in enumerate(memories):
|
||||
context_parts.append(f"{i}. {result['text']}")
|
||||
for result in memories:
|
||||
context_parts.append({"text": result.get("text"), "context": result.get("context"), "event_date": result.get("event_date")})
|
||||
|
||||
context = "\n".join(context_parts)
|
||||
context = json.dumps(context_parts)
|
||||
|
||||
# Use OpenAI to generate answer
|
||||
# Use Groq to generate answer
|
||||
try:
|
||||
client = AsyncOpenAI()
|
||||
response = await client.beta.chat.completions.parse(
|
||||
model="gpt-5",
|
||||
response = await self.client.beta.chat.completions.parse(
|
||||
model="openai/gpt-oss-120b",
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
|
|
@ -238,15 +249,37 @@ class LoComoAnswerEvaluator(LLMAnswerEvaluator):
|
|||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are an objective judge. Determine if the predicted answer contains the correct answer or they are the same content (with different form is fine)."
|
||||
"content": "You are an expert grader that determines if answers to questions match a gold standard answer"
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": f"Question: {question}\nCorrect answer: {correct_answer}\nPredicted answer: {predicted_answer}\n\nAre they equivalent?"
|
||||
"content": f"""
|
||||
Your task is to label an answer to a question as ’CORRECT’ or ’WRONG’. You williolw23 be given the following data:
|
||||
(1) a question (posed by one user to another user),
|
||||
(2) a ’gold’ (ground truth) answer,
|
||||
(3) a generated answer
|
||||
which you will score as CORRECT/WRONG.
|
||||
|
||||
The point of the question is to ask about something one user should know about the other user based on their prior conversations.
|
||||
The gold answer will usually be a concise and short answer that includes the referenced topic, for example:
|
||||
Question: Do you remember what I got the last time I went to Hawaii?
|
||||
Gold answer: A shell necklace
|
||||
The generated answer might be much longer, but you should be generous with your grading - as long as it touches on the same topic as the gold answer, it should be counted as CORRECT.
|
||||
|
||||
For time related questions, the gold answer will be a specific date, month, year, etc. The generated answer might be much longer or use relative time references (like "last Tuesday" or "next month"), but you should be generous with your grading - as long as it refers to the same date or time period as the gold answer, it should be counted as CORRECT. Even if the format differs (e.g., "May 7th" vs "7 May"), consider it CORRECT if it's the same date.
|
||||
|
||||
Now it’s time for the real question:
|
||||
Question: {question}
|
||||
Gold answer: {correct_answer}
|
||||
Generated answer: {predicted_answer}
|
||||
|
||||
First, provide a short (one sentence) explanation of your reasoning. Short reasoning is preferred.
|
||||
If it's correct, set correct=true.
|
||||
"""
|
||||
}
|
||||
],
|
||||
temperature=0,
|
||||
max_tokens=512,
|
||||
max_tokens=4096,
|
||||
response_format=JudgeResponse
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -1,8 +1,8 @@
|
|||
# LoComo Benchmark Results
|
||||
|
||||
**Overall Accuracy**: 41.70% (98/235)
|
||||
**Overall Accuracy**: 66.00% (66/100)
|
||||
|
||||
| Sample ID | Turns | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |
|
||||
|-----------|-------|-----------|---------|----------|-----------|------------|----------|-------------|
|
||||
| conv-26 | 419 | 154 | 69 | 44.81% | N/A | N/A | N/A | N/A |
|
||||
| conv-30 | 369 | 81 | 29 | 35.80% | N/A | N/A | N/A | N/A |
|
||||
| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |
|
||||
|-----------|----------|-----------|---------|----------|-----------|------------|----------|-------------|
|
||||
| conv-26 | 19 | 50 | 35 | 70.00% | N/A | N/A | N/A | N/A |
|
||||
| conv-30 | 19 | 50 | 31 | 62.00% | N/A | N/A | N/A | N/A |
|
||||
|
|
@ -129,6 +129,9 @@ def generate_markdown_table(results: dict):
|
|||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
parser = argparse.ArgumentParser(description='Run LoComo benchmark')
|
||||
parser.add_argument('--max-conversations', type=int, default=None, help='Maximum conversations to evaluate')
|
||||
parser.add_argument('--max-questions', type=int, default=None, help='Maximum questions per conversation')
|
||||
|
|
|
|||
|
|
@ -13,6 +13,16 @@ from openai import AsyncOpenAI
|
|||
from pydantic import BaseModel, Field
|
||||
|
||||
|
||||
class Entity(BaseModel):
|
||||
"""An entity extracted from text."""
|
||||
text: str = Field(
|
||||
description="The entity name as it appears in the fact"
|
||||
)
|
||||
type: Literal["PERSON", "ORG", "PLACE", "PRODUCT", "CONCEPT"] = Field(
|
||||
description="Entity type: PERSON, ORG, PLACE, PRODUCT, or CONCEPT"
|
||||
)
|
||||
|
||||
|
||||
class ExtractedFact(BaseModel):
|
||||
"""A single extracted fact from text."""
|
||||
fact: str = Field(
|
||||
|
|
@ -21,6 +31,10 @@ class ExtractedFact(BaseModel):
|
|||
date: str = Field(
|
||||
description="Absolute date/time when this fact occurred in ISO format (YYYY-MM-DDTHH:MM:SSZ). If text mentions relative time (yesterday, last week, this morning), calculate absolute date from the provided context date."
|
||||
)
|
||||
entities: List[Entity] = Field(
|
||||
default_factory=list,
|
||||
description="List of important entities mentioned in this fact with their types"
|
||||
)
|
||||
|
||||
|
||||
class FactExtractionResponse(BaseModel):
|
||||
|
|
@ -132,10 +146,20 @@ async def _extract_facts_from_chunk(
|
|||
Each fact should:
|
||||
1. Be SELF-CONTAINED - readable without the original context
|
||||
2. Include ALL relevant details: WHO, WHAT, WHERE, WHEN, WHY, HOW
|
||||
3. Preserve specific names, dates, numbers, locations, relationships
|
||||
4. Resolve pronouns to actual names/entities
|
||||
5. Include surrounding context that makes the fact meaningful
|
||||
6. Capture nuances, reasons, causes, and implications
|
||||
3. **CRITICAL: ALWAYS include the SUBJECT (who is doing/saying/experiencing)**
|
||||
4. Preserve specific names, dates, numbers, locations, relationships
|
||||
5. Resolve pronouns to actual names/entities (I → speaker name, their → possessor name)
|
||||
6. **CRITICAL: Preserve possessive relationships** (their kids → whose kids, his car → whose car)
|
||||
7. Include surrounding context that makes the fact meaningful
|
||||
8. Capture nuances, reasons, causes, and implications
|
||||
|
||||
**COMMON MISTAKES TO AVOID:**
|
||||
- ❌ "The kids were excited" → Missing WHO the kids belong to
|
||||
- ✅ "Melanie's kids were excited" or "Melanie took her kids who were excited"
|
||||
- ❌ "Someone went hiking" → Missing WHO
|
||||
- ✅ "Bob went hiking"
|
||||
- ❌ "The car broke down" → Missing whose car
|
||||
- ✅ "Alice's car broke down"
|
||||
|
||||
## TEMPORAL INFORMATION (VERY IMPORTANT)
|
||||
For each fact, extract the ABSOLUTE date/time when it occurred:
|
||||
|
|
@ -153,7 +177,12 @@ Examples of date extraction:
|
|||
- "I work at Google" (no time mentioned) → date: 2024-03-20T10:00:00Z (use reference)
|
||||
|
||||
## What to EXTRACT (BE EXHAUSTIVE - DO NOT SKIP ANYTHING):
|
||||
- **Biographical information**: jobs, roles, backgrounds, experiences, skills
|
||||
- **Biographical information (CRITICAL - NEVER MISS)**:
|
||||
- Origins: home country, birthplace, where someone is from ("my home country Sweden" = Caroline is from Sweden)
|
||||
- Current location: where they live now
|
||||
- Jobs, roles, backgrounds, experiences, skills
|
||||
- Family background, heritage, cultural identity
|
||||
- Education, training, certifications
|
||||
- **Events (NEVER MISS THESE)**:
|
||||
- ANY action that happened (went, did, attended, joined, started, finished, etc.)
|
||||
- Photos, images, videos shared or taken ("here's a photo", "took a picture", "captured")
|
||||
|
|
@ -161,6 +190,10 @@ Examples of date extraction:
|
|||
- Achievements, milestones, accomplishments
|
||||
- Travels, visits, locations visited
|
||||
- Purchases, acquisitions, creations
|
||||
- **Identity and personal details**:
|
||||
- Origins, nationality, home country, roots
|
||||
- Cultural background, heritage
|
||||
- Family connections (grandmother from X, parents in Y)
|
||||
- **Opinions and beliefs**: who believes what and why
|
||||
- **Recommendations and advice**: specific suggestions with reasoning
|
||||
- **Descriptions**: detailed explanations of how things work
|
||||
|
|
@ -179,24 +212,89 @@ Examples of date extraction:
|
|||
- Pure reactions without content ("wow", "cool", "nice")
|
||||
- Incomplete thoughts or sentence fragments with no meaning
|
||||
|
||||
## ENTITY EXTRACTION (CRITICAL):
|
||||
For EACH fact, extract ALL important entities mentioned with their types:
|
||||
- **PERSON**: Names of individuals (Alice, Bob, Dr. Smith)
|
||||
- **ORG**: Companies, institutions, teams (Google, MIT, AI Team)
|
||||
- **PLACE**: Cities, countries, locations, venues (Mountain View, Yosemite, The Coffee Shop)
|
||||
- **PRODUCT**: Specific products, tools, technologies (iPhone, Python, TensorFlow)
|
||||
- **CONCEPT**: Important topics, projects, subjects (AI, machine learning, Project Phoenix)
|
||||
|
||||
Entity extraction rules:
|
||||
- Use the EXACT form as it appears in the fact (preserve capitalization)
|
||||
- Assign the correct type to distinguish ambiguous entities (Apple the company = ORG, apple the fruit = PRODUCT/CONCEPT)
|
||||
- Include both full names and commonly used short forms if both appear
|
||||
- Extract proper nouns and key identifying terms
|
||||
- Skip generic terms (the, a, some) and pronouns (he, she, they)
|
||||
- Each entity must have both text and type
|
||||
|
||||
## EXAMPLES of GOOD facts (detailed, comprehensive):
|
||||
|
||||
Input: "Alice mentioned she works at Google in Mountain View. She joined the AI team last year."
|
||||
GOOD fact: "Alice works at Google in Mountain View on the AI team, which she joined last year"
|
||||
GOOD date: Calculate based on reference date (if reference is 2024-03-20, "last year" = 2023-03-20)
|
||||
GOOD entities: [
|
||||
{{"text": "Alice", "type": "PERSON"}},
|
||||
{{"text": "Google", "type": "ORG"}},
|
||||
{{"text": "Mountain View", "type": "PLACE"}},
|
||||
{{"text": "AI team", "type": "ORG"}}
|
||||
]
|
||||
|
||||
Input: "Yesterday Bob went hiking in Yosemite because it helps him clear his mind."
|
||||
GOOD fact: "Bob went hiking in Yosemite because it helps him clear his mind"
|
||||
GOOD date: Reference date minus 1 day
|
||||
GOOD entities: [
|
||||
{{"text": "Bob", "type": "PERSON"}},
|
||||
{{"text": "Yosemite", "type": "PLACE"}}
|
||||
]
|
||||
|
||||
Input: "Here's a photo of me with my friends taken last week at the beach."
|
||||
GOOD fact: "Someone shared/took a photo with their friends at the beach"
|
||||
GOOD date: Reference date minus 7 days (last week)
|
||||
GOOD entities: []
|
||||
NOTE: Extract the event (photo taken/shared with friends at beach), NOT just that a photo exists
|
||||
|
||||
Input: "I sent you that article about AI last Tuesday."
|
||||
GOOD fact: "Someone sent an article about AI"
|
||||
GOOD date: Calculate last Tuesday from reference date
|
||||
GOOD entities: [
|
||||
{{"text": "AI", "type": "CONCEPT"}}
|
||||
]
|
||||
|
||||
Input: "I bought an Apple laptop and some apples from the store."
|
||||
GOOD fact: "Someone bought an Apple laptop and some apples from the store"
|
||||
GOOD entities: [
|
||||
{{"text": "Apple", "type": "ORG"}},
|
||||
{{"text": "apples", "type": "PRODUCT"}}
|
||||
]
|
||||
NOTE: Use type to distinguish "Apple" the company from "apples" the fruit
|
||||
|
||||
Input: "Melanie said 'Yesterday I took the kids to the museum - it was so cool seeing their eyes light up!'"
|
||||
BAD fact: "The kids were excited about the museum"
|
||||
BAD entities: [{{"text": "museum", "type": "PLACE"}}]
|
||||
PROBLEM: Missing WHO (Melanie) and whose kids!
|
||||
|
||||
GOOD fact: "Melanie took her kids to the museum yesterday and they were excited, with their eyes lighting up"
|
||||
GOOD date: Reference date minus 1 day
|
||||
GOOD entities: [
|
||||
{{"text": "Melanie", "type": "PERSON"}},
|
||||
{{"text": "museum", "type": "PLACE"}}
|
||||
]
|
||||
NOTE: Preserved the subject (Melanie) and possessive relationship (her kids)
|
||||
|
||||
Input: "Caroline said 'This necklace is from my grandma in my home country, Sweden. She gave it to me when I was young.'"
|
||||
BAD fact: "Caroline received a necklace as a gift from her grandmother when she was young"
|
||||
BAD entities: [{{"text": "Caroline", "type": "PERSON"}}, {{"text": "necklace", "type": "PRODUCT"}}]
|
||||
PROBLEM: Missing the CRITICAL biographical info that Caroline is from Sweden!
|
||||
|
||||
GOOD facts (extract MULTIPLE facts):
|
||||
1. "Caroline is from Sweden, which is her home country"
|
||||
entities: [{{"text": "Caroline", "type": "PERSON"}}, {{"text": "Sweden", "type": "PLACE"}}]
|
||||
2. "Caroline's grandmother is from Sweden"
|
||||
entities: [{{"text": "Caroline", "type": "PERSON"}}, {{"text": "Sweden", "type": "PLACE"}}]
|
||||
3. "Caroline received a necklace as a gift from her grandmother in Sweden when she was young"
|
||||
entities: [{{"text": "Caroline", "type": "PERSON"}}, {{"text": "Sweden", "type": "PLACE"}}, {{"text": "necklace", "type": "PRODUCT"}}]
|
||||
NOTE: Extract SEPARATE facts for biographical details (home country) AND events (gift received)
|
||||
|
||||
## TEXT TO EXTRACT FROM:
|
||||
{chunk}
|
||||
|
|
@ -204,16 +302,21 @@ GOOD date: Calculate last Tuesday from reference date
|
|||
Remember:
|
||||
1. BE EXHAUSTIVE - Extract EVERY event, action, and fact mentioned
|
||||
2. DO NOT skip casual mentions like "here's a photo", "I was with X", "sent you Y"
|
||||
3. Include ALL details, names, numbers, reasons, and context in the fact text
|
||||
4. Extract the absolute date for EACH fact by calculating relative times from the reference date
|
||||
5. When in doubt, EXTRACT IT - better to have too many facts than miss important events"""
|
||||
3. **ALWAYS include the SUBJECT** - never say "the kids" without saying whose kids
|
||||
4. **Preserve possessive relationships** - "their kids" must become "Person's kids"
|
||||
5. **Extract biographical details as SEPARATE facts** - "my home country Sweden" should create a fact "Person is from Sweden"
|
||||
6. Include ALL details, names, numbers, reasons, and context in the fact text
|
||||
7. Extract the absolute date for EACH fact by calculating relative times from the reference date
|
||||
8. Extract ALL entities with their types (PERSON, ORG, PLACE, PRODUCT, CONCEPT) for each fact
|
||||
9. Use types to disambiguate entities (Apple the company = ORG, apple the fruit = PRODUCT)
|
||||
10. When in doubt, EXTRACT IT - better to have too many facts than miss important events"""
|
||||
|
||||
response = await client.beta.chat.completions.parse(
|
||||
model=model,
|
||||
messages=[
|
||||
{
|
||||
"role": "system",
|
||||
"content": "You are an EXHAUSTIVE fact extractor. Extract EVERY event, action, and fact mentioned - never skip anything. This includes casual mentions like photos shared, things sent, meetups, gatherings, or any action. Preserve all context, details, and nuances. Calculate absolute dates from relative time expressions. When in doubt, extract it - missing facts is worse than extracting too many."
|
||||
"content": "You are an EXHAUSTIVE fact and entity extractor. CRITICAL RULES: 1) ALWAYS include the SUBJECT (never 'the kids' without whose kids), 2) Extract biographical details as SEPARATE facts (if someone mentions 'my home country Sweden', extract 'Person is from Sweden' as its own fact), 3) Extract EVERY event, action, and fact - never skip anything. For each fact, extract ALL important entities with their types: PERSON, ORG, PLACE, PRODUCT, CONCEPT. Use types to disambiguate (Apple=ORG vs apples=PRODUCT). Preserve possessive relationships (their→whose). Include casual mentions (photos, meetups). Calculate absolute dates from relative times. When in doubt, extract it - better too many facts than missing critical biographical/identity information."
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
|
|
|
|||
|
|
@ -47,9 +47,10 @@ class LinkInfo(BaseModel):
|
|||
link_type: Literal["temporal", "semantic", "entity"] = Field(description="Type of link")
|
||||
link_weight: float = Field(description="Weight of the link", ge=0.0, le=1.0)
|
||||
entity_id: Optional[str] = Field(default=None, description="Entity ID if link_type is 'entity'")
|
||||
new_activation: float = Field(description="Activation that would be passed to neighbor")
|
||||
new_activation: Optional[float] = Field(default=None, description="Activation that would be passed to neighbor (None for supplementary links)")
|
||||
followed: bool = Field(description="Whether this link was followed (or pruned)")
|
||||
prune_reason: Optional[str] = Field(default=None, description="Why link was not followed (if not followed)")
|
||||
is_supplementary: bool = Field(default=False, description="Whether this is a supplementary link (multiple connections to same node)")
|
||||
|
||||
|
||||
class NodeVisit(BaseModel):
|
||||
|
|
|
|||
|
|
@ -182,9 +182,10 @@ class SearchTracer:
|
|||
link_type: Literal["temporal", "semantic", "entity"],
|
||||
link_weight: float,
|
||||
entity_id: Optional[str],
|
||||
new_activation: float,
|
||||
new_activation: Optional[float],
|
||||
followed: bool,
|
||||
prune_reason: Optional[str] = None,
|
||||
is_supplementary: bool = False,
|
||||
):
|
||||
"""
|
||||
Record a link to a neighbor (whether followed or not).
|
||||
|
|
@ -195,9 +196,10 @@ class SearchTracer:
|
|||
link_type: Type of link
|
||||
link_weight: Weight of link
|
||||
entity_id: Entity ID if link is entity-based
|
||||
new_activation: Activation passed to neighbor
|
||||
new_activation: Activation passed to neighbor (None for supplementary links)
|
||||
followed: Whether link was followed
|
||||
prune_reason: Why link was not followed (if not followed)
|
||||
is_supplementary: Whether this is a supplementary link (multiple connections)
|
||||
"""
|
||||
# Find the visit for the source node
|
||||
visit = None
|
||||
|
|
@ -218,6 +220,7 @@ class SearchTracer:
|
|||
new_activation=new_activation,
|
||||
followed=followed,
|
||||
prune_reason=prune_reason,
|
||||
is_supplementary=is_supplementary,
|
||||
)
|
||||
|
||||
visit.neighbors_explored.append(link_info)
|
||||
|
|
|
|||
|
|
@ -18,6 +18,8 @@ import asyncio
|
|||
import time
|
||||
from concurrent.futures import ProcessPoolExecutor
|
||||
import numpy as np
|
||||
import uuid
|
||||
import logging
|
||||
|
||||
from .utils import (
|
||||
extract_facts,
|
||||
|
|
@ -32,6 +34,9 @@ def utcnow():
|
|||
return datetime.now(timezone.utc)
|
||||
|
||||
|
||||
# Logger for memory system
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
# Global process pool for parallel embedding generation
|
||||
# Each process loads its own copy of the embedding model
|
||||
# This provides TRUE parallelism for CPU-bound embedding operations
|
||||
|
|
@ -106,9 +111,9 @@ class TemporalSemanticMemory:
|
|||
self.entity_resolver = None
|
||||
|
||||
# Initialize local embedding model (384 dimensions)
|
||||
print(f"Loading embedding model: {embedding_model}...")
|
||||
logger.info(f"Loading embedding model: {embedding_model}...")
|
||||
self.embedding_model = SentenceTransformer(embedding_model)
|
||||
print(f"✓ Model loaded (embedding dim: {self.embedding_model.get_sentence_embedding_dimension()})")
|
||||
logger.info(f"Model loaded (embedding dim: {self.embedding_model.get_sentence_embedding_dimension()})")
|
||||
|
||||
# Background queue for access count updates (to avoid blocking searches)
|
||||
self._access_count_queue = asyncio.Queue()
|
||||
|
|
@ -143,18 +148,20 @@ class TemporalSemanticMemory:
|
|||
if updates:
|
||||
node_id_list = list(updates.keys())
|
||||
try:
|
||||
# Convert string UUIDs to UUID type for faster matching
|
||||
uuid_list = [uuid.UUID(nid) for nid in node_id_list]
|
||||
async with pool.acquire() as conn:
|
||||
await conn.execute(
|
||||
"UPDATE memory_units SET access_count = access_count + 1 WHERE id::text = ANY($1)",
|
||||
node_id_list
|
||||
"UPDATE memory_units SET access_count = access_count + 1 WHERE id = ANY($1::uuid[])",
|
||||
uuid_list
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"[ACCESS_COUNT_WORKER] Error updating access counts: {e}")
|
||||
logger.error(f"Access count worker: Error updating access counts: {e}")
|
||||
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
print(f"[ACCESS_COUNT_WORKER] Unexpected error: {e}")
|
||||
logger.error(f"Access count worker: Unexpected error: {e}")
|
||||
await asyncio.sleep(1) # Backoff on error
|
||||
|
||||
async def _get_pool(self) -> asyncpg.Pool:
|
||||
|
|
@ -330,6 +337,9 @@ class TemporalSemanticMemory:
|
|||
content: str,
|
||||
context: str = "",
|
||||
event_date: Optional[datetime] = None,
|
||||
document_id: Optional[str] = None,
|
||||
document_metadata: Optional[Dict[str, Any]] = None,
|
||||
upsert: bool = False,
|
||||
) -> List[str]:
|
||||
"""
|
||||
Store content as memory units with temporal and semantic links (ASYNC version).
|
||||
|
|
@ -341,6 +351,9 @@ class TemporalSemanticMemory:
|
|||
content: Text content to store
|
||||
context: Context about when/why this memory was formed
|
||||
event_date: When the event occurred (defaults to now)
|
||||
document_id: Optional document ID for tracking and upsert
|
||||
document_metadata: Optional metadata about the document
|
||||
upsert: If True and document_id exists, delete old units and create new ones
|
||||
|
||||
Returns:
|
||||
List of created unit IDs
|
||||
|
|
@ -352,7 +365,10 @@ class TemporalSemanticMemory:
|
|||
"content": content,
|
||||
"context": context,
|
||||
"event_date": event_date
|
||||
}]
|
||||
}],
|
||||
document_id=document_id,
|
||||
document_metadata=document_metadata,
|
||||
upsert=upsert
|
||||
)
|
||||
|
||||
# Return the first (and only) list of unit IDs
|
||||
|
|
@ -362,6 +378,9 @@ class TemporalSemanticMemory:
|
|||
self,
|
||||
agent_id: str,
|
||||
contents: List[Dict[str, Any]],
|
||||
document_id: Optional[str] = None,
|
||||
document_metadata: Optional[Dict[str, Any]] = None,
|
||||
upsert: bool = False,
|
||||
) -> List[List[str]]:
|
||||
"""
|
||||
Store multiple content items as memory units in ONE batch operation.
|
||||
|
|
@ -377,6 +396,9 @@ class TemporalSemanticMemory:
|
|||
- "content" (required): Text content to store
|
||||
- "context" (optional): Context about the memory
|
||||
- "event_date" (optional): When the event occurred
|
||||
document_id: Optional document ID for tracking and upsert
|
||||
document_metadata: Optional metadata about the document
|
||||
upsert: If True and document_id exists, delete old units and create new ones
|
||||
|
||||
Returns:
|
||||
List of lists of unit IDs (one list per content item)
|
||||
|
|
@ -387,15 +409,17 @@ class TemporalSemanticMemory:
|
|||
contents=[
|
||||
{"content": "Alice works at Google", "context": "conversation"},
|
||||
{"content": "Bob loves Python", "context": "conversation"},
|
||||
]
|
||||
],
|
||||
document_id="meeting-2024-01-15",
|
||||
upsert=True
|
||||
)
|
||||
# Returns: [["unit-id-1"], ["unit-id-2"]]
|
||||
"""
|
||||
start_time = time.time()
|
||||
print(f"\n{'='*60}")
|
||||
print(f"PUT_BATCH_ASYNC START: {agent_id}")
|
||||
print(f"Batch size: {len(contents)} content items")
|
||||
print(f"{'='*60}")
|
||||
logger.debug(f"\n{'='*60}")
|
||||
logger.debug(f"PUT_BATCH_ASYNC START: {agent_id}")
|
||||
logger.debug(f"Batch size: {len(contents)} content items")
|
||||
logger.debug(f"{'='*60}")
|
||||
|
||||
if not contents:
|
||||
return []
|
||||
|
|
@ -420,6 +444,7 @@ class TemporalSemanticMemory:
|
|||
all_fact_texts = []
|
||||
all_fact_dates = []
|
||||
all_contexts = []
|
||||
all_fact_entities = [] # NEW: Store LLM-extracted entities per fact
|
||||
content_boundaries = [] # [(start_idx, end_idx), ...]
|
||||
|
||||
current_idx = 0
|
||||
|
|
@ -435,6 +460,8 @@ class TemporalSemanticMemory:
|
|||
except Exception:
|
||||
all_fact_dates.append(event_date)
|
||||
all_contexts.append(context)
|
||||
# Extract entities from fact (default to empty list if not present)
|
||||
all_fact_entities.append(fact_dict.get('entities', []))
|
||||
|
||||
end_idx = current_idx + len(fact_dicts)
|
||||
content_boundaries.append((start_idx, end_idx))
|
||||
|
|
@ -448,13 +475,51 @@ class TemporalSemanticMemory:
|
|||
# Step 2: Generate ALL embeddings in ONE batch (HUGE speedup!)
|
||||
step_start = time.time()
|
||||
all_embeddings = await self._generate_embeddings_batch(all_fact_texts)
|
||||
print(f"[2] Generate embeddings (parallel): {len(all_embeddings)} embeddings in {time.time() - step_start:.3f}s")
|
||||
logger.debug(f"[2] Generate embeddings (parallel): {len(all_embeddings)} embeddings in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Step 3: Process everything in ONE database transaction
|
||||
pool = await self._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
async with conn.transaction():
|
||||
try:
|
||||
# Handle document tracking and upsert
|
||||
if document_id:
|
||||
import hashlib
|
||||
import json
|
||||
|
||||
# Calculate content hash from all content items
|
||||
combined_content = "\n".join([c.get("content", "") for c in contents])
|
||||
content_hash = hashlib.sha256(combined_content.encode()).hexdigest()
|
||||
|
||||
# If upsert, delete old document first (cascades to units and links)
|
||||
if upsert:
|
||||
deleted = await conn.fetchval(
|
||||
"DELETE FROM documents WHERE id = $1 AND agent_id = $2 RETURNING id",
|
||||
document_id, agent_id
|
||||
)
|
||||
if deleted:
|
||||
logger.debug(f"[3.1] Upsert: Deleted existing document '{document_id}' and all its units")
|
||||
|
||||
# Insert or update document
|
||||
# Always use ON CONFLICT for idempotent behavior
|
||||
await conn.execute(
|
||||
"""
|
||||
INSERT INTO documents (id, agent_id, original_text, content_hash, metadata)
|
||||
VALUES ($1, $2, $3, $4, $5)
|
||||
ON CONFLICT (id, agent_id) DO UPDATE
|
||||
SET original_text = EXCLUDED.original_text,
|
||||
content_hash = EXCLUDED.content_hash,
|
||||
metadata = EXCLUDED.metadata,
|
||||
updated_at = NOW()
|
||||
""",
|
||||
document_id,
|
||||
agent_id,
|
||||
combined_content,
|
||||
content_hash,
|
||||
json.dumps(document_metadata or {})
|
||||
)
|
||||
logger.debug(f"[3.2] Document '{document_id}' stored/updated")
|
||||
|
||||
# Deduplication check for all facts
|
||||
step_start = time.time()
|
||||
all_is_duplicate = []
|
||||
|
|
@ -466,16 +531,17 @@ class TemporalSemanticMemory:
|
|||
|
||||
duplicates_filtered = sum(all_is_duplicate)
|
||||
new_facts = total_facts - duplicates_filtered
|
||||
print(f"[3] Deduplication check: {duplicates_filtered} duplicates filtered, {new_facts} new facts in {time.time() - step_start:.3f}s")
|
||||
logger.debug(f"[3] Deduplication check: {duplicates_filtered} duplicates filtered, {new_facts} new facts in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Filter out duplicates
|
||||
filtered_sentences = [s for s, is_dup in zip(all_fact_texts, all_is_duplicate) if not is_dup]
|
||||
filtered_embeddings = [e for e, is_dup in zip(all_embeddings, all_is_duplicate) if not is_dup]
|
||||
filtered_dates = [d for d, is_dup in zip(all_fact_dates, all_is_duplicate) if not is_dup]
|
||||
filtered_contexts = [c for c, is_dup in zip(all_contexts, all_is_duplicate) if not is_dup]
|
||||
filtered_entities = [ents for ents, is_dup in zip(all_fact_entities, all_is_duplicate) if not is_dup]
|
||||
|
||||
if not filtered_sentences:
|
||||
print(f"[PUT_BATCH_ASYNC] All facts were duplicates, returning empty")
|
||||
logger.debug(f"[PUT_BATCH_ASYNC] All facts were duplicates, returning empty")
|
||||
return [[] for _ in contents]
|
||||
|
||||
# Batch insert ALL units
|
||||
|
|
@ -484,11 +550,12 @@ class TemporalSemanticMemory:
|
|||
filtered_embeddings_str = [str(emb) for emb in filtered_embeddings]
|
||||
results = await conn.fetch(
|
||||
"""
|
||||
INSERT INTO memory_units (agent_id, text, context, embedding, event_date, access_count)
|
||||
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::vector[], $5::timestamptz[], $6::integer[])
|
||||
INSERT INTO memory_units (agent_id, document_id, text, context, embedding, event_date, access_count)
|
||||
SELECT * FROM unnest($1::text[], $2::text[], $3::text[], $4::text[], $5::vector[], $6::timestamptz[], $7::integer[])
|
||||
RETURNING id
|
||||
""",
|
||||
[agent_id] * len(filtered_sentences),
|
||||
[document_id] * len(filtered_sentences) if document_id else [None] * len(filtered_sentences),
|
||||
filtered_sentences,
|
||||
filtered_contexts,
|
||||
filtered_embeddings_str,
|
||||
|
|
@ -497,34 +564,34 @@ class TemporalSemanticMemory:
|
|||
)
|
||||
|
||||
created_unit_ids = [str(row['id']) for row in results]
|
||||
print(f"[5] Batch insert units: {len(created_unit_ids)} units in {time.time() - step_start:.3f}s")
|
||||
logger.debug(f"[5] Batch insert units: {len(created_unit_ids)} units in {time.time() - step_start:.3f}s")
|
||||
|
||||
# Process entities for ALL units
|
||||
step_start = time.time()
|
||||
all_entity_links = await self._extract_entities_batch_optimized(
|
||||
conn, agent_id, created_unit_ids, filtered_sentences, "", filtered_dates
|
||||
conn, agent_id, created_unit_ids, filtered_sentences, "", filtered_dates, filtered_entities
|
||||
)
|
||||
print(f"[6] Extract entities (batched): {time.time() - step_start:.3f}s")
|
||||
logger.debug(f"[6] Process entities (batched): {time.time() - step_start:.3f}s")
|
||||
|
||||
# Create temporal links
|
||||
step_start = time.time()
|
||||
await self._create_temporal_links_batch_per_fact(conn, agent_id, created_unit_ids)
|
||||
print(f"[7] Batch create temporal links: {time.time() - step_start:.3f}s")
|
||||
logger.debug(f"[7] Batch create temporal links: {time.time() - step_start:.3f}s")
|
||||
|
||||
# Create semantic links
|
||||
step_start = time.time()
|
||||
await self._create_semantic_links_batch(conn, agent_id, created_unit_ids, filtered_embeddings)
|
||||
print(f"[8] Batch create semantic links: {time.time() - step_start:.3f}s")
|
||||
logger.debug(f"[8] Batch create semantic links: {time.time() - step_start:.3f}s")
|
||||
|
||||
# Insert entity links
|
||||
step_start = time.time()
|
||||
if all_entity_links:
|
||||
await self._insert_entity_links_batch(conn, all_entity_links)
|
||||
print(f"[9] Batch insert entity links: {time.time() - step_start:.3f}s")
|
||||
logger.debug(f"[9] Batch insert entity links: {time.time() - step_start:.3f}s")
|
||||
|
||||
# Transaction auto-commits on success
|
||||
commit_start = time.time()
|
||||
print(f"[10] Commit: {time.time() - commit_start:.3f}s")
|
||||
logger.debug(f"[10] Commit: {time.time() - commit_start:.3f}s")
|
||||
|
||||
# Map created unit IDs back to original content items
|
||||
# Account for duplicates when mapping back
|
||||
|
|
@ -540,9 +607,9 @@ class TemporalSemanticMemory:
|
|||
result_unit_ids.append(content_unit_ids)
|
||||
|
||||
total_time = time.time() - start_time
|
||||
print(f"\n{'='*60}")
|
||||
print(f"PUT_BATCH_ASYNC COMPLETE: {len(created_unit_ids)} units from {len(contents)} contents in {total_time:.3f}s")
|
||||
print(f"{'='*60}\n")
|
||||
logger.debug(f"\n{'='*60}")
|
||||
logger.debug(f"PUT_BATCH_ASYNC COMPLETE: {len(created_unit_ids)} units from {len(contents)} contents in {total_time:.3f}s")
|
||||
logger.debug(f"{'='*60}\n")
|
||||
|
||||
return result_unit_ids
|
||||
|
||||
|
|
@ -563,6 +630,7 @@ class TemporalSemanticMemory:
|
|||
weight_semantic: float = 0.30,
|
||||
weight_recency: float = 0.25,
|
||||
weight_frequency: float = 0.15,
|
||||
mmr_lambda: float = 0.5,
|
||||
) -> tuple[List[Dict[str, Any]], Optional[Any]]:
|
||||
"""
|
||||
Search memories using spreading activation (synchronous wrapper).
|
||||
|
|
@ -580,6 +648,7 @@ class TemporalSemanticMemory:
|
|||
weight_semantic: Weight for semantic similarity component (default: 0.30)
|
||||
weight_recency: Weight for recency component (default: 0.25)
|
||||
weight_frequency: Weight for frequency component (default: 0.15)
|
||||
mmr_lambda: Lambda for MMR diversification (0=max diversity, 1=no diversity, default: 0.5)
|
||||
|
||||
Returns:
|
||||
Tuple of (results, trace)
|
||||
|
|
@ -587,7 +656,7 @@ class TemporalSemanticMemory:
|
|||
# Run async version synchronously
|
||||
return asyncio.run(self.search_async(
|
||||
agent_id, query, thinking_budget, top_k, enable_trace,
|
||||
weight_activation, weight_semantic, weight_recency, weight_frequency
|
||||
weight_activation, weight_semantic, weight_recency, weight_frequency, mmr_lambda
|
||||
))
|
||||
|
||||
async def search_async(
|
||||
|
|
@ -601,6 +670,7 @@ class TemporalSemanticMemory:
|
|||
weight_semantic: float = 0.30,
|
||||
weight_recency: float = 0.25,
|
||||
weight_frequency: float = 0.15,
|
||||
mmr_lambda: float = 0.5,
|
||||
) -> tuple[List[Dict[str, Any]], Optional[Any]]:
|
||||
"""
|
||||
Search memories using spreading activation (ASYNC version).
|
||||
|
|
@ -629,14 +699,18 @@ class TemporalSemanticMemory:
|
|||
|
||||
pool = await self._get_pool()
|
||||
search_start = time.time()
|
||||
print(f"\n[SEARCH] Starting search for query: '{query[:50]}...' (thinking_budget={thinking_budget}, top_k={top_k})")
|
||||
|
||||
# Buffer logs for clean output in concurrent scenarios
|
||||
search_id = f"{agent_id[:8]}-{int(time.time() * 1000) % 100000}"
|
||||
log_buffer = []
|
||||
log_buffer.append(f"[SEARCH {search_id}] Query: '{query[:50]}...' (budget={thinking_budget}, top_k={top_k})")
|
||||
|
||||
try:
|
||||
# Step 1: Generate query embedding (CPU-bound, no DB needed)
|
||||
step_start = time.time()
|
||||
query_embedding = self._generate_embedding(query)
|
||||
step_duration = time.time() - step_start
|
||||
print(f" [1] Generate query embedding: {step_duration:.3f}s")
|
||||
log_buffer.append(f" [1] Generate query embedding: {step_duration:.3f}s")
|
||||
|
||||
if tracer:
|
||||
tracer.record_query_embedding(query_embedding)
|
||||
|
|
@ -651,7 +725,7 @@ class TemporalSemanticMemory:
|
|||
async with pool.acquire() as conn:
|
||||
conn_acquire_time = time.time() - conn_acquire_start
|
||||
if conn_acquire_time > 0.1: # Log if waiting > 100ms
|
||||
print(f" [2.1] Waited {conn_acquire_time:.3f}s for connection (pool busy)")
|
||||
log_buffer.append(f" [2.1] Waited {conn_acquire_time:.3f}s for connection (pool busy)")
|
||||
|
||||
entry_points = await conn.fetch(
|
||||
"""
|
||||
|
|
@ -668,7 +742,7 @@ class TemporalSemanticMemory:
|
|||
)
|
||||
|
||||
step_duration = time.time() - step_start
|
||||
print(f" [2] Find entry points: {len(entry_points)} found in {step_duration:.3f}s")
|
||||
log_buffer.append(f" [2] Find entry points: {len(entry_points)} found in {step_duration:.3f}s")
|
||||
|
||||
if tracer:
|
||||
tracer.add_phase_metric("find_entry_points", step_duration, {"count": len(entry_points)})
|
||||
|
|
@ -681,7 +755,7 @@ class TemporalSemanticMemory:
|
|||
)
|
||||
|
||||
if not entry_points:
|
||||
print(f"[SEARCH] Complete: 0 results in {time.time() - search_start:.3f}s")
|
||||
logger.debug(f"[SEARCH] Complete: 0 results in {time.time() - search_start:.3f}s")
|
||||
if tracer:
|
||||
trace = tracer.finalize([])
|
||||
return [], trace
|
||||
|
|
@ -734,10 +808,12 @@ class TemporalSemanticMemory:
|
|||
async with pool.acquire() as conn:
|
||||
batch_conn_acquire = time.time() - batch_conn_start
|
||||
if batch_conn_acquire > 0.1: # Log if waiting > 100ms
|
||||
print(f" [3.3.1] Waited {batch_conn_acquire:.3f}s for connection (pool busy) - batch size: {len(node_ids)}")
|
||||
log_buffer.append(f" [3.3.1] Waited {batch_conn_acquire:.3f}s for connection (pool busy) - batch size: {len(node_ids)}")
|
||||
|
||||
# Query neighbors for ALL nodes in batch at once (without embeddings for speed)
|
||||
# Convert string UUIDs to UUID type for faster matching
|
||||
substep_start = time.time()
|
||||
uuid_array = [uuid.UUID(nid) for nid in node_ids]
|
||||
all_neighbors = await conn.fetch(
|
||||
"""
|
||||
SELECT ml.from_unit_id, ml.to_unit_id, ml.weight, ml.link_type, ml.entity_id,
|
||||
|
|
@ -745,27 +821,27 @@ class TemporalSemanticMemory:
|
|||
mu.id as neighbor_id
|
||||
FROM memory_links ml
|
||||
JOIN memory_units mu ON ml.to_unit_id = mu.id
|
||||
WHERE ml.from_unit_id::text = ANY($1)
|
||||
WHERE ml.from_unit_id = ANY($1::uuid[])
|
||||
AND ml.weight >= 0.1
|
||||
ORDER BY ml.from_unit_id, ml.weight DESC
|
||||
""",
|
||||
node_ids
|
||||
uuid_array
|
||||
)
|
||||
neighbor_query_time = time.time() - substep_start
|
||||
if neighbor_query_time > 1.0: # Log slow neighbor queries
|
||||
print(f" [3.3.3] Slow NEIGHBOR query: {neighbor_query_time:.3f}s for {len(node_ids)} nodes → {len(all_neighbors)} neighbors")
|
||||
log_buffer.append(f" [3.3.3] Slow NEIGHBOR query: {neighbor_query_time:.3f}s for {len(node_ids)} nodes → {len(all_neighbors)} neighbors")
|
||||
query_neighbors_time += neighbor_query_time
|
||||
|
||||
# Fetch embeddings for current batch nodes (needed for weight calculation)
|
||||
substep_start = time.time()
|
||||
embeddings = await conn.fetch(
|
||||
"SELECT id, embedding FROM memory_units WHERE id::text = ANY($1)",
|
||||
node_ids
|
||||
"SELECT id, embedding FROM memory_units WHERE id = ANY($1::uuid[])",
|
||||
uuid_array
|
||||
)
|
||||
embedding_map = {str(row["id"]): row["embedding"] for row in embeddings}
|
||||
fetch_embeddings_time = time.time() - substep_start
|
||||
if fetch_embeddings_time > 0.5:
|
||||
print(f" [3.3.4] Slow EMBEDDING fetch: {fetch_embeddings_time:.3f}s for {len(node_ids)} nodes")
|
||||
log_buffer.append(f" [3.3.4] Slow EMBEDDING fetch: {fetch_embeddings_time:.3f}s for {len(node_ids)} nodes")
|
||||
query_neighbors_time += fetch_embeddings_time
|
||||
|
||||
# Group neighbors by from_unit_id (in-memory, no DB)
|
||||
|
|
@ -851,49 +927,90 @@ class TemporalSemanticMemory:
|
|||
"semantic_similarity": semantic_similarity,
|
||||
"recency": recency_weight,
|
||||
"frequency": frequency_weight,
|
||||
"embedding": memory_embedding, # Store for MMR
|
||||
})
|
||||
|
||||
# Spread to neighbors (from batch query results)
|
||||
neighbors = neighbors_by_node.get(unit_id, [])
|
||||
|
||||
# Group neighbors by to_unit_id to handle multiple connections
|
||||
neighbors_grouped = {}
|
||||
for neighbor in neighbors:
|
||||
neighbor_id = str(neighbor["to_unit_id"])
|
||||
link_weight = neighbor["weight"]
|
||||
link_type = neighbor["link_type"]
|
||||
entity_id = str(neighbor["entity_id"]) if neighbor["entity_id"] else None
|
||||
new_activation = activation * link_weight * 0.8 # 0.8 = decay factor
|
||||
if neighbor_id not in neighbors_grouped:
|
||||
neighbors_grouped[neighbor_id] = []
|
||||
neighbors_grouped[neighbor_id].append(neighbor)
|
||||
|
||||
if neighbor_id not in visited:
|
||||
if new_activation > 0.1:
|
||||
queue.append(({
|
||||
"id": neighbor["to_unit_id"],
|
||||
"text": neighbor["text"],
|
||||
"context": neighbor.get("context", ""),
|
||||
"event_date": neighbor["event_date"],
|
||||
"access_count": neighbor["access_count"],
|
||||
}, new_activation, False, unit_id, link_type, link_weight)) # parent_id, link_type, link_weight
|
||||
# Process each unique neighbor (aggregating multiple links)
|
||||
for neighbor_id, neighbor_links in neighbors_grouped.items():
|
||||
if neighbor_id in visited:
|
||||
continue
|
||||
|
||||
if tracer:
|
||||
tracer.add_neighbor_link(
|
||||
from_node_id=unit_id,
|
||||
to_node_id=neighbor_id,
|
||||
link_type=link_type,
|
||||
link_weight=link_weight,
|
||||
entity_id=entity_id,
|
||||
new_activation=new_activation,
|
||||
followed=True
|
||||
)
|
||||
elif tracer:
|
||||
# Sort links by weight descending to identify primary link
|
||||
neighbor_links_sorted = sorted(neighbor_links, key=lambda x: x["weight"], reverse=True)
|
||||
primary_link = neighbor_links_sorted[0]
|
||||
|
||||
# Aggregate link weights: max + 30% bonus for additional links
|
||||
max_weight = primary_link["weight"]
|
||||
bonus_weight = sum(link["weight"] for link in neighbor_links_sorted[1:]) * 0.3
|
||||
combined_weight = max_weight + bonus_weight
|
||||
|
||||
# Calculate new activation using combined weight
|
||||
new_activation = activation * combined_weight * 0.8 # 0.8 = decay factor
|
||||
|
||||
# Use primary link metadata for queue and trace
|
||||
primary_link_type = primary_link["link_type"]
|
||||
primary_entity_id = str(primary_link["entity_id"]) if primary_link["entity_id"] else None
|
||||
|
||||
if new_activation > 0.1:
|
||||
queue.append(({
|
||||
"id": primary_link["to_unit_id"],
|
||||
"text": primary_link["text"],
|
||||
"context": primary_link.get("context", ""),
|
||||
"event_date": primary_link["event_date"],
|
||||
"access_count": primary_link["access_count"],
|
||||
}, new_activation, False, unit_id, primary_link_type, combined_weight)) # parent_id, link_type, combined_weight
|
||||
|
||||
# Record all links in trace (primary + additional)
|
||||
if tracer:
|
||||
# Add primary link with combined activation
|
||||
tracer.add_neighbor_link(
|
||||
from_node_id=unit_id,
|
||||
to_node_id=neighbor_id,
|
||||
link_type=link_type,
|
||||
link_weight=link_weight,
|
||||
entity_id=entity_id,
|
||||
link_type=primary_link_type,
|
||||
link_weight=combined_weight,
|
||||
entity_id=primary_entity_id,
|
||||
new_activation=new_activation,
|
||||
followed=False,
|
||||
prune_reason="activation_too_low"
|
||||
followed=True
|
||||
)
|
||||
|
||||
# Add additional links as supplementary (if multiple connections exist)
|
||||
for additional_link in neighbor_links_sorted[1:]:
|
||||
additional_link_type = additional_link["link_type"]
|
||||
additional_entity_id = str(additional_link["entity_id"]) if additional_link["entity_id"] else None
|
||||
tracer.add_neighbor_link(
|
||||
from_node_id=unit_id,
|
||||
to_node_id=neighbor_id,
|
||||
link_type=additional_link_type,
|
||||
link_weight=additional_link["weight"],
|
||||
entity_id=additional_entity_id,
|
||||
new_activation=None, # Don't show activation for supplementary links
|
||||
followed=True,
|
||||
is_supplementary=True # Mark as supplementary link
|
||||
)
|
||||
elif tracer:
|
||||
# Record pruned link
|
||||
tracer.add_neighbor_link(
|
||||
from_node_id=unit_id,
|
||||
to_node_id=neighbor_id,
|
||||
link_type=primary_link_type,
|
||||
link_weight=combined_weight,
|
||||
entity_id=primary_entity_id,
|
||||
new_activation=new_activation,
|
||||
followed=False,
|
||||
prune_reason="activation_too_low"
|
||||
)
|
||||
|
||||
calculate_weight_time += time.time() - substep_start
|
||||
process_neighbors_time += time.time() - substep_start
|
||||
|
||||
|
|
@ -902,10 +1019,10 @@ class TemporalSemanticMemory:
|
|||
|
||||
spreading_activation_time = time.time() - step_start
|
||||
num_batches = (len(visited) + BATCH_SIZE - 1) // BATCH_SIZE # Ceiling division
|
||||
print(f" [3] Spreading activation: {len(visited)} nodes visited in {spreading_activation_time:.3f}s")
|
||||
print(f" [3.1] Calculate weights: {calculate_weight_time:.3f}s")
|
||||
print(f" [3.2] Query neighbors: {query_neighbors_time:.3f}s ({num_batches} batched queries)")
|
||||
print(f" [3.3] Process neighbors: {process_neighbors_time:.3f}s")
|
||||
log_buffer.append(f" [3] Spreading activation: {len(visited)} nodes visited in {spreading_activation_time:.3f}s")
|
||||
log_buffer.append(f" [3.1] Calculate weights: {calculate_weight_time:.3f}s")
|
||||
log_buffer.append(f" [3.2] Query neighbors: {query_neighbors_time:.3f}s ({num_batches} batched queries)")
|
||||
log_buffer.append(f" [3.3] Process neighbors: {process_neighbors_time:.3f}s")
|
||||
|
||||
if tracer:
|
||||
tracer.add_phase_metric("spreading_activation", spreading_activation_time, {
|
||||
|
|
@ -916,15 +1033,33 @@ class TemporalSemanticMemory:
|
|||
# Step 4: Queue access count updates (background worker will process them)
|
||||
if visited_node_ids:
|
||||
await self._access_count_queue.put(visited_node_ids)
|
||||
print(f" [4] Queued access count updates for {len(visited_node_ids)} nodes")
|
||||
log_buffer.append(f" [4] Queued access count updates for {len(visited_node_ids)} nodes")
|
||||
|
||||
# Step 5: Sort by final weight and return top results
|
||||
# Step 5: Sort by final weight and apply MMR for diversity
|
||||
step_start = time.time()
|
||||
results.sort(key=lambda x: x["weight"], reverse=True)
|
||||
top_results = results[:top_k]
|
||||
print(f" [5] Sort and return top {top_k}: {time.time() - step_start:.3f}s")
|
||||
|
||||
print(f"[SEARCH] Complete: {len(top_results)} results in {time.time() - search_start:.3f}s\n")
|
||||
# Apply MMR (Maximal Marginal Relevance) for diversity if lambda < 1.0
|
||||
if mmr_lambda < 1.0 and len(results) > top_k:
|
||||
top_results = self._apply_mmr(results, top_k, mmr_lambda, log_buffer)
|
||||
log_buffer.append(f" [5] MMR diversification (λ={mmr_lambda}): {time.time() - step_start:.3f}s")
|
||||
else:
|
||||
top_results = results[:top_k]
|
||||
# Add original rank and remove embeddings from results
|
||||
for idx, result in enumerate(top_results):
|
||||
result["original_rank"] = idx + 1
|
||||
result["mmr_score"] = None
|
||||
result["mmr_relevance"] = None
|
||||
result["mmr_max_similarity"] = None
|
||||
result["mmr_diversified"] = False
|
||||
result.pop("embedding", None)
|
||||
log_buffer.append(f" [5] Sort and return top {top_k} (no MMR): {time.time() - step_start:.3f}s")
|
||||
|
||||
total_time = time.time() - search_start
|
||||
log_buffer.append(f"[SEARCH {search_id}] Complete: {len(top_results)} results in {total_time:.3f}s")
|
||||
|
||||
# Log all buffered logs at once
|
||||
logger.info("\n" + "\n".join(log_buffer))
|
||||
|
||||
# Finalize trace if enabled
|
||||
if tracer:
|
||||
|
|
@ -933,9 +1068,207 @@ class TemporalSemanticMemory:
|
|||
return top_results, None
|
||||
|
||||
except Exception as e:
|
||||
print(f"[SEARCH] ERROR after {time.time() - search_start:.3f}s: {str(e)}")
|
||||
log_buffer.append(f"[SEARCH {search_id}] ERROR after {time.time() - search_start:.3f}s: {str(e)}")
|
||||
logger.error("\n" + "\n".join(log_buffer))
|
||||
raise Exception(f"Failed to search memories: {str(e)}")
|
||||
|
||||
def _apply_mmr(
|
||||
self,
|
||||
results: List[Dict[str, Any]],
|
||||
top_k: int,
|
||||
mmr_lambda: float,
|
||||
log_buffer: List[str]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Apply Maximal Marginal Relevance (MMR) to diversify search results.
|
||||
|
||||
MMR balances relevance with diversity by selecting results that are:
|
||||
1. Relevant to the query (high score)
|
||||
2. Different from already selected results (low similarity)
|
||||
|
||||
Formula: MMR = λ * relevance - (1-λ) * max_similarity_to_selected
|
||||
|
||||
Args:
|
||||
results: Sorted list of all results with embeddings
|
||||
top_k: Number of results to select
|
||||
mmr_lambda: Balance parameter (0=max diversity, 1=max relevance)
|
||||
log_buffer: Logging buffer
|
||||
|
||||
Returns:
|
||||
Diversified list of top_k results
|
||||
"""
|
||||
if not results or top_k <= 0:
|
||||
return []
|
||||
|
||||
# Normalize weights to [0, 1] for fair comparison with similarity
|
||||
max_weight = max(r["weight"] for r in results)
|
||||
min_weight = min(r["weight"] for r in results)
|
||||
weight_range = max_weight - min_weight if max_weight > min_weight else 1.0
|
||||
|
||||
# Pre-compute normalized relevance scores for all results
|
||||
for idx, result in enumerate(results):
|
||||
result["original_rank"] = idx + 1
|
||||
result["normalized_relevance"] = (result["weight"] - min_weight) / weight_range
|
||||
|
||||
# Extract embeddings as a numpy array for vectorized operations
|
||||
# Shape: (num_results, embedding_dim)
|
||||
embeddings_list = []
|
||||
valid_indices = []
|
||||
for idx, result in enumerate(results):
|
||||
if result.get("embedding") is not None:
|
||||
embeddings_list.append(result["embedding"])
|
||||
valid_indices.append(idx)
|
||||
|
||||
if not embeddings_list:
|
||||
# No embeddings available, just return top-k by relevance
|
||||
return results[:top_k]
|
||||
|
||||
# Stack embeddings into a matrix (num_results, embedding_dim)
|
||||
embeddings_matrix = np.array(embeddings_list, dtype=np.float32)
|
||||
|
||||
# Normalize embeddings for faster cosine similarity (just dot product after normalization)
|
||||
norms = np.linalg.norm(embeddings_matrix, axis=1, keepdims=True)
|
||||
norms[norms == 0] = 1.0 # Avoid division by zero
|
||||
embeddings_matrix = embeddings_matrix / norms
|
||||
|
||||
selected_indices = []
|
||||
remaining_indices = list(range(len(results)))
|
||||
diversified_count = 0
|
||||
|
||||
for selection_round in range(min(top_k, len(results))):
|
||||
if not remaining_indices:
|
||||
break
|
||||
|
||||
best_mmr_score = float('-inf')
|
||||
best_remaining_idx = 0
|
||||
|
||||
# Vectorized computation for all remaining candidates
|
||||
for remaining_idx, candidate_idx in enumerate(remaining_indices):
|
||||
candidate = results[candidate_idx]
|
||||
normalized_relevance = candidate["normalized_relevance"]
|
||||
|
||||
# Calculate max similarity to selected results
|
||||
max_similarity = 0.0
|
||||
if selected_indices and candidate_idx in valid_indices:
|
||||
# Find position in embeddings_matrix
|
||||
embedding_idx = valid_indices.index(candidate_idx)
|
||||
candidate_embedding = embeddings_matrix[embedding_idx]
|
||||
|
||||
# Vectorized similarity calculation with all selected embeddings
|
||||
if selected_indices:
|
||||
selected_embedding_indices = [valid_indices.index(idx) for idx in selected_indices if idx in valid_indices]
|
||||
if selected_embedding_indices:
|
||||
selected_embeddings = embeddings_matrix[selected_embedding_indices]
|
||||
# Compute cosine similarities in one operation (already normalized, so just dot product)
|
||||
similarities = np.dot(selected_embeddings, candidate_embedding)
|
||||
max_similarity = float(np.max(similarities))
|
||||
|
||||
# MMR score: balance relevance and diversity
|
||||
mmr_score = mmr_lambda * normalized_relevance - (1 - mmr_lambda) * max_similarity
|
||||
|
||||
if mmr_score > best_mmr_score:
|
||||
best_mmr_score = mmr_score
|
||||
best_remaining_idx = remaining_idx
|
||||
best_max_similarity = max_similarity
|
||||
|
||||
# Select the best candidate
|
||||
best_candidate_idx = remaining_indices.pop(best_remaining_idx)
|
||||
best_candidate = results[best_candidate_idx]
|
||||
|
||||
# Store MMR metadata
|
||||
best_candidate["mmr_score"] = best_mmr_score
|
||||
best_candidate["mmr_relevance"] = best_candidate["normalized_relevance"]
|
||||
best_candidate["mmr_max_similarity"] = best_max_similarity
|
||||
best_candidate["mmr_diversified"] = best_remaining_idx > 0
|
||||
|
||||
selected_indices.append(best_candidate_idx)
|
||||
|
||||
if best_remaining_idx > 0:
|
||||
diversified_count += 1
|
||||
|
||||
log_buffer.append(f" MMR: Selected {len(selected_indices)} results, {diversified_count} diversified picks")
|
||||
|
||||
# Return selected results in order
|
||||
selected_results = [results[idx] for idx in selected_indices]
|
||||
|
||||
# Remove embeddings from final results (not needed in response)
|
||||
for result in selected_results:
|
||||
result.pop("embedding", None)
|
||||
result.pop("normalized_relevance", None) # Clean up temp field
|
||||
|
||||
return selected_results
|
||||
|
||||
async def get_document(self, document_id: str, agent_id: str) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Retrieve document metadata and statistics.
|
||||
|
||||
Args:
|
||||
document_id: Document ID to retrieve
|
||||
agent_id: Agent ID that owns the document
|
||||
|
||||
Returns:
|
||||
Dictionary with document info or None if not found
|
||||
"""
|
||||
pool = await self._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
doc = await conn.fetchrow(
|
||||
"""
|
||||
SELECT d.id, d.agent_id, d.original_text, d.content_hash, d.metadata,
|
||||
d.created_at, d.updated_at, COUNT(mu.id) as unit_count
|
||||
FROM documents d
|
||||
LEFT JOIN memory_units mu ON mu.document_id = d.id
|
||||
WHERE d.id = $1 AND d.agent_id = $2
|
||||
GROUP BY d.id, d.agent_id, d.original_text, d.content_hash, d.metadata, d.created_at, d.updated_at
|
||||
""",
|
||||
document_id, agent_id
|
||||
)
|
||||
|
||||
if not doc:
|
||||
return None
|
||||
|
||||
import json
|
||||
return {
|
||||
"id": doc["id"],
|
||||
"agent_id": doc["agent_id"],
|
||||
"original_text": doc["original_text"],
|
||||
"content_hash": doc["content_hash"],
|
||||
"metadata": json.loads(doc["metadata"]) if doc["metadata"] else {},
|
||||
"unit_count": doc["unit_count"],
|
||||
"created_at": doc["created_at"],
|
||||
"updated_at": doc["updated_at"]
|
||||
}
|
||||
|
||||
async def delete_document(self, document_id: str, agent_id: str) -> Dict[str, int]:
|
||||
"""
|
||||
Delete a document and all its associated memory units and links.
|
||||
|
||||
Args:
|
||||
document_id: Document ID to delete
|
||||
agent_id: Agent ID that owns the document
|
||||
|
||||
Returns:
|
||||
Dictionary with counts of deleted items
|
||||
"""
|
||||
pool = await self._get_pool()
|
||||
async with pool.acquire() as conn:
|
||||
async with conn.transaction():
|
||||
# Count units before deletion
|
||||
units_count = await conn.fetchval(
|
||||
"SELECT COUNT(*) FROM memory_units WHERE document_id = $1",
|
||||
document_id
|
||||
)
|
||||
|
||||
# Delete document (cascades to memory_units and all their links)
|
||||
deleted = await conn.fetchval(
|
||||
"DELETE FROM documents WHERE id = $1 AND agent_id = $2 RETURNING id",
|
||||
document_id, agent_id
|
||||
)
|
||||
|
||||
return {
|
||||
"document_deleted": 1 if deleted else 0,
|
||||
"memory_units_deleted": units_count if deleted else 0
|
||||
}
|
||||
|
||||
async def delete_agent(self, agent_id: str) -> Dict[str, int]:
|
||||
"""
|
||||
Delete all data for a specific agent (multi-tenant cleanup).
|
||||
|
|
@ -984,23 +1317,33 @@ class TemporalSemanticMemory:
|
|||
sentences: List[str],
|
||||
context: str,
|
||||
fact_dates: List,
|
||||
llm_entities: List[List[Dict]], # NEW: Entities from LLM
|
||||
) -> List[tuple]:
|
||||
"""
|
||||
Extract entities from ALL sentences in one batch (MUCH faster than sequential).
|
||||
Process LLM-extracted entities for ALL facts in batch.
|
||||
|
||||
Uses spaCy's batch processing to extract entities from all texts at once,
|
||||
then resolves and links them in bulk.
|
||||
Uses entities provided by the LLM (no spaCy needed), then resolves
|
||||
and links them in bulk.
|
||||
|
||||
Returns list of tuples for batch insertion: (from_unit_id, to_unit_id, link_type, weight, entity_id)
|
||||
"""
|
||||
from .entity_resolver import extract_entities_batch
|
||||
|
||||
try:
|
||||
# Step 1: Extract entities from ALL sentences in one batch (fast!)
|
||||
# Step 1: Convert LLM entities to the format expected by entity resolver
|
||||
substep_start = time.time()
|
||||
all_entities = extract_entities_batch(sentences)
|
||||
all_entities = []
|
||||
for entity_list in llm_entities:
|
||||
# Convert List[Entity] or List[dict] to List[Dict] format
|
||||
formatted_entities = []
|
||||
for ent in entity_list:
|
||||
# Handle both Entity objects and dicts
|
||||
if hasattr(ent, 'text'):
|
||||
formatted_entities.append({'text': ent.text, 'type': ent.type})
|
||||
elif isinstance(ent, dict):
|
||||
formatted_entities.append({'text': ent.get('text', ''), 'type': ent.get('type', 'CONCEPT')})
|
||||
all_entities.append(formatted_entities)
|
||||
|
||||
total_entities = sum(len(ents) for ents in all_entities)
|
||||
print(f" [6.1] spaCy NER (batch): {total_entities} entities from {len(sentences)} sentences in {time.time() - substep_start:.3f}s")
|
||||
logger.debug(f" [6.1] Process LLM entities: {total_entities} entities from {len(sentences)} facts in {time.time() - substep_start:.3f}s")
|
||||
|
||||
# Step 2: Resolve entities in BATCH (much faster!)
|
||||
substep_start = time.time()
|
||||
|
|
@ -1022,7 +1365,7 @@ class TemporalSemanticMemory:
|
|||
'nearby_entities': entities,
|
||||
})
|
||||
entity_to_unit.append((unit_id, local_idx, fact_date))
|
||||
print(f" [6.2.1] Prepare entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_1_start:.3f}s")
|
||||
logger.debug(f" [6.2.1] Prepare entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_1_start:.3f}s")
|
||||
|
||||
# Resolve ALL entities in one batch call
|
||||
if all_entities_flat:
|
||||
|
|
@ -1052,7 +1395,7 @@ class TemporalSemanticMemory:
|
|||
|
||||
for idx, entity_id in zip(indices, batch_resolved):
|
||||
resolved_entity_ids[idx] = entity_id
|
||||
print(f" [6.2.2] Resolve entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_2_start:.3f}s")
|
||||
logger.debug(f" [6.2.2] Resolve entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_2_start:.3f}s")
|
||||
|
||||
# [6.2.3] Create unit-entity links in BATCH
|
||||
substep_6_2_3_start = time.time()
|
||||
|
|
@ -1069,12 +1412,12 @@ class TemporalSemanticMemory:
|
|||
|
||||
# Batch insert all unit-entity links (MUCH faster!)
|
||||
await self.entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
|
||||
print(f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_6_2_3_start:.3f}s")
|
||||
logger.debug(f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_6_2_3_start:.3f}s")
|
||||
|
||||
print(f" [6.2] Entity resolution (batched): {len(all_entities_flat)} entities resolved in {time.time() - step_6_2_start:.3f}s")
|
||||
logger.debug(f" [6.2] Entity resolution (batched): {len(all_entities_flat)} entities resolved in {time.time() - step_6_2_start:.3f}s")
|
||||
else:
|
||||
unit_to_entity_ids = {}
|
||||
print(f" [6.2] Entity resolution (batched): 0 entities in {time.time() - step_6_2_start:.3f}s")
|
||||
logger.debug(f" [6.2] Entity resolution (batched): 0 entities in {time.time() - step_6_2_start:.3f}s")
|
||||
|
||||
# Step 3: Create entity links between units that share entities
|
||||
substep_start = time.time()
|
||||
|
|
@ -1106,12 +1449,12 @@ class TemporalSemanticMemory:
|
|||
links.append((unit_id_1, unit_id_2, 'entity', 1.0, entity_id))
|
||||
links.append((unit_id_2, unit_id_1, 'entity', 1.0, entity_id))
|
||||
|
||||
print(f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s")
|
||||
logger.debug(f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s")
|
||||
|
||||
return links
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to extract entities in batch: {str(e)}")
|
||||
logger.error(f" Failed to extract entities in batch: {str(e)}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
# Re-raise to trigger rollback at put_async level
|
||||
|
|
@ -1184,7 +1527,7 @@ class TemporalSemanticMemory:
|
|||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to create temporal links: {str(e)}")
|
||||
logger.error(f" Failed to create temporal links: {str(e)}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
# Re-raise to trigger rollback at put_async level
|
||||
|
|
@ -1244,7 +1587,7 @@ class TemporalSemanticMemory:
|
|||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"ERROR: Failed to create semantic links: {str(e)}")
|
||||
logger.error(f" Failed to create semantic links: {str(e)}")
|
||||
import traceback
|
||||
traceback.print_exc()
|
||||
# Re-raise to trigger rollback at put_async level
|
||||
|
|
@ -1265,4 +1608,4 @@ class TemporalSemanticMemory:
|
|||
links
|
||||
)
|
||||
except Exception as e:
|
||||
print(f"Warning: Failed to insert entity links: {str(e)}")
|
||||
logger.warning(f" Failed to insert entity links: {str(e)}")
|
||||
|
|
|
|||
|
|
@ -1,423 +0,0 @@
|
|||
|
||||
<!DOCTYPE html>
|
||||
<html>
|
||||
<head>
|
||||
<title>Memory Graph - Interactive Visualization</title>
|
||||
<meta charset="utf-8">
|
||||
<script src="https://cdnjs.cloudflare.com/ajax/libs/cytoscape/3.28.1/cytoscape.min.js"></script>
|
||||
<style>
|
||||
body {
|
||||
font-family: Tahoma, sans-serif;
|
||||
margin: 0;
|
||||
padding: 0;
|
||||
background: #f5f5f5;
|
||||
}
|
||||
|
||||
.tab-container {
|
||||
background: white;
|
||||
}
|
||||
|
||||
.tab-buttons {
|
||||
background: #f0f0f0;
|
||||
border-bottom: 2px solid #333;
|
||||
padding: 0;
|
||||
margin: 0;
|
||||
}
|
||||
|
||||
.tab-button {
|
||||
background: #e0e0e0;
|
||||
border: none;
|
||||
padding: 12px 24px;
|
||||
cursor: pointer;
|
||||
font-size: 16px;
|
||||
font-weight: bold;
|
||||
border-top: 2px solid transparent;
|
||||
border-left: 2px solid transparent;
|
||||
border-right: 2px solid transparent;
|
||||
transition: background 0.2s;
|
||||
}
|
||||
|
||||
.tab-button:hover {
|
||||
background: #d0d0d0;
|
||||
}
|
||||
|
||||
.tab-button.active {
|
||||
background: white;
|
||||
border-top: 2px solid #333;
|
||||
border-left: 2px solid #333;
|
||||
border-right: 2px solid #333;
|
||||
border-bottom: 2px solid white;
|
||||
margin-bottom: -2px;
|
||||
}
|
||||
|
||||
.tab-content {
|
||||
display: none;
|
||||
background: white;
|
||||
}
|
||||
|
||||
.tab-content.active {
|
||||
display: block;
|
||||
}
|
||||
|
||||
#cy {
|
||||
width: 100%;
|
||||
height: 800px;
|
||||
background: #ffffff;
|
||||
}
|
||||
|
||||
#graph-tab {
|
||||
position: relative;
|
||||
}
|
||||
|
||||
#table-tab {
|
||||
padding: 20px;
|
||||
}
|
||||
|
||||
.legend {
|
||||
position: absolute;
|
||||
top: 20px;
|
||||
left: 20px;
|
||||
background: white;
|
||||
padding: 15px;
|
||||
border: 2px solid #333;
|
||||
border-radius: 8px;
|
||||
box-shadow: 2px 2px 8px rgba(0,0,0,0.3);
|
||||
z-index: 1000;
|
||||
max-width: 250px;
|
||||
}
|
||||
|
||||
.legend h3 {
|
||||
margin-top: 0;
|
||||
border-bottom: 2px solid #333;
|
||||
padding-bottom: 5px;
|
||||
}
|
||||
|
||||
.legend-item {
|
||||
margin: 8px 0;
|
||||
display: flex;
|
||||
align-items: center;
|
||||
}
|
||||
|
||||
.legend-line {
|
||||
width: 30px;
|
||||
height: 2px;
|
||||
margin-right: 10px;
|
||||
}
|
||||
|
||||
.legend-node {
|
||||
width: 20px;
|
||||
height: 20px;
|
||||
margin-right: 10px;
|
||||
border: 1px solid #999;
|
||||
border-radius: 3px;
|
||||
}
|
||||
|
||||
#table-filter {
|
||||
width: 100%;
|
||||
max-width: 600px;
|
||||
padding: 10px;
|
||||
margin-bottom: 15px;
|
||||
border: 2px solid #ccc;
|
||||
border-radius: 4px;
|
||||
font-size: 14px;
|
||||
box-sizing: border-box;
|
||||
}
|
||||
|
||||
#memory-table {
|
||||
width: 100%;
|
||||
border-collapse: collapse;
|
||||
font-size: 13px;
|
||||
max-width: 1400px;
|
||||
}
|
||||
|
||||
#memory-table th {
|
||||
padding: 10px;
|
||||
text-align: left;
|
||||
border: 1px solid #ddd;
|
||||
background: #f0f0f0;
|
||||
}
|
||||
|
||||
#memory-table td {
|
||||
padding: 8px;
|
||||
border: 1px solid #ddd;
|
||||
}
|
||||
|
||||
.tooltip {
|
||||
position: absolute;
|
||||
background: white;
|
||||
border: 2px solid #333;
|
||||
border-radius: 4px;
|
||||
padding: 10px;
|
||||
box-shadow: 2px 2px 8px rgba(0,0,0,0.3);
|
||||
max-width: 300px;
|
||||
font-size: 12px;
|
||||
pointer-events: none;
|
||||
z-index: 9999;
|
||||
}
|
||||
</style>
|
||||
</head>
|
||||
<body>
|
||||
<div class="tab-container">
|
||||
<div class="tab-buttons">
|
||||
<button class="tab-button active" onclick="switchTab('graph')">Graph View</button>
|
||||
<button class="tab-button" onclick="switchTab('table')">Table View</button>
|
||||
</div>
|
||||
|
||||
<div id="graph-tab" class="tab-content active">
|
||||
<div style="padding: 15px; background: #f9f9f9; border-bottom: 2px solid #333;">
|
||||
<div style="display: flex; gap: 15px; align-items: center; flex-wrap: wrap;">
|
||||
<div>
|
||||
<label style="font-weight: bold; margin-right: 5px;">Limit nodes:</label>
|
||||
<input type="number" id="node-limit" value="50" min="10" max="1000" step="10"
|
||||
style="width: 80px; padding: 5px; border: 1px solid #ccc; border-radius: 4px;">
|
||||
</div>
|
||||
<div>
|
||||
<label style="font-weight: bold; margin-right: 5px;">Layout:</label>
|
||||
<select id="layout-select" style="padding: 5px; border: 1px solid #ccc; border-radius: 4px;">
|
||||
<option value="circle">Circle (fast)</option>
|
||||
<option value="grid">Grid (fast)</option>
|
||||
<option value="cose">Force-directed (slow)</option>
|
||||
</select>
|
||||
</div>
|
||||
<button onclick="reloadGraph()" style="padding: 6px 15px; background: #42a5f5; color: white; border: none; border-radius: 4px; cursor: pointer; font-weight: bold;">
|
||||
Apply
|
||||
</button>
|
||||
<span id="node-count" style="color: #666; font-size: 14px;"></span>
|
||||
</div>
|
||||
</div>
|
||||
<div id="cy"></div>
|
||||
<div class="legend">
|
||||
<h3>Legend</h3>
|
||||
<h4 style="margin: 10px 0 5px 0;">Link Types:</h4>
|
||||
<div class="legend-item">
|
||||
<div class="legend-line" style="background: #00bcd4; border-top: 1px dashed #00bcd4;"></div>
|
||||
<span><b>Temporal</b></span>
|
||||
</div>
|
||||
<div class="legend-item">
|
||||
<div class="legend-line" style="background: #ff69b4;"></div>
|
||||
<span><b>Semantic</b></span>
|
||||
</div>
|
||||
<div class="legend-item">
|
||||
<div class="legend-line" style="background: #ffd700;"></div>
|
||||
<span><b>Entity</b></span>
|
||||
</div>
|
||||
<h4 style="margin: 15px 0 5px 0;">Nodes:</h4>
|
||||
<div class="legend-item">
|
||||
<div class="legend-node" style="background: #e0e0e0;"></div>
|
||||
<span>No entities</span>
|
||||
</div>
|
||||
<div class="legend-item">
|
||||
<div class="legend-node" style="background: #90caf9;"></div>
|
||||
<span>1 entity</span>
|
||||
</div>
|
||||
<div class="legend-item">
|
||||
<div class="legend-node" style="background: #42a5f5;"></div>
|
||||
<span>2+ entities</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<div id="table-tab" class="tab-content">
|
||||
<h2>Memory Units (0)</h2>
|
||||
<input type="text" id="table-filter" placeholder="Filter by text, context, or entities...">
|
||||
<div style="overflow-x: auto;">
|
||||
<table id="memory-table">
|
||||
<thead>
|
||||
<tr>
|
||||
<th>ID</th>
|
||||
<th>Text</th>
|
||||
<th>Context</th>
|
||||
<th>Date</th>
|
||||
<th>Entities</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
|
||||
</tbody>
|
||||
</table>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<script>
|
||||
// Graph data
|
||||
const allGraphData = {"nodes": [], "edges": []};
|
||||
let cy = null;
|
||||
|
||||
// Initialize graph with filtering
|
||||
function initGraph(nodeLimit, layoutName) {
|
||||
// Filter nodes to limit
|
||||
const limitedNodes = allGraphData.nodes.slice(0, nodeLimit);
|
||||
const nodeIds = new Set(limitedNodes.map(n => n.data.id));
|
||||
|
||||
// Filter edges to only include those between visible nodes
|
||||
const limitedEdges = allGraphData.edges.filter(e =>
|
||||
nodeIds.has(e.data.source) && nodeIds.has(e.data.target)
|
||||
);
|
||||
|
||||
// Update count display
|
||||
document.getElementById('node-count').textContent =
|
||||
`Showing ${limitedNodes.length} of ${allGraphData.nodes.length} nodes`;
|
||||
|
||||
// Destroy existing graph if any
|
||||
if (cy) {
|
||||
cy.destroy();
|
||||
}
|
||||
|
||||
// Layout configurations
|
||||
const layouts = {
|
||||
'circle': {
|
||||
name: 'circle',
|
||||
animate: false,
|
||||
radius: 300,
|
||||
spacingFactor: 1.5
|
||||
},
|
||||
'grid': {
|
||||
name: 'grid',
|
||||
animate: false,
|
||||
rows: Math.ceil(Math.sqrt(limitedNodes.length)),
|
||||
cols: Math.ceil(Math.sqrt(limitedNodes.length)),
|
||||
spacingFactor: 2
|
||||
},
|
||||
'cose': {
|
||||
name: 'cose',
|
||||
animate: false,
|
||||
nodeRepulsion: 15000,
|
||||
idealEdgeLength: 150,
|
||||
edgeElasticity: 100,
|
||||
nestingFactor: 1.2,
|
||||
gravity: 1,
|
||||
numIter: 1000,
|
||||
initialTemp: 200,
|
||||
coolingFactor: 0.95,
|
||||
minTemp: 1.0
|
||||
}
|
||||
};
|
||||
|
||||
// Initialize Cytoscape
|
||||
cy = cytoscape({
|
||||
container: document.getElementById('cy'),
|
||||
|
||||
elements: [
|
||||
...limitedNodes.map(n => ({ data: n.data })),
|
||||
...limitedEdges.map(e => ({ data: e.data }))
|
||||
],
|
||||
|
||||
style: [
|
||||
{
|
||||
selector: 'node',
|
||||
style: {
|
||||
'background-color': 'data(color)',
|
||||
'label': 'data(label)',
|
||||
'text-valign': 'center',
|
||||
'text-halign': 'center',
|
||||
'font-size': '10px',
|
||||
'font-weight': 'bold',
|
||||
'text-wrap': 'wrap',
|
||||
'text-max-width': '100px',
|
||||
'width': 40,
|
||||
'height': 40,
|
||||
'border-width': 2,
|
||||
'border-color': '#333'
|
||||
}
|
||||
},
|
||||
{
|
||||
selector: 'edge',
|
||||
style: {
|
||||
'width': 1,
|
||||
'line-color': 'data(color)',
|
||||
'line-style': 'data(lineStyle)',
|
||||
'target-arrow-shape': 'triangle',
|
||||
'target-arrow-color': 'data(color)',
|
||||
'curve-style': 'bezier',
|
||||
'opacity': 0.7
|
||||
}
|
||||
},
|
||||
{
|
||||
selector: 'node:selected',
|
||||
style: {
|
||||
'border-width': 4,
|
||||
'border-color': '#000'
|
||||
}
|
||||
}
|
||||
],
|
||||
|
||||
layout: layouts[layoutName] || layouts['circle']
|
||||
});
|
||||
|
||||
// Simple tooltip on hover
|
||||
let tooltip = null;
|
||||
|
||||
cy.on('mouseover', 'node', function(evt) {
|
||||
const node = evt.target;
|
||||
const data = node.data();
|
||||
const renderedPosition = node.renderedPosition();
|
||||
|
||||
// Create tooltip
|
||||
tooltip = document.createElement('div');
|
||||
tooltip.className = 'tooltip';
|
||||
tooltip.innerHTML = `
|
||||
<b>Text:</b> ${data.text}<br>
|
||||
<b>Context:</b> ${data.context}<br>
|
||||
<b>Date:</b> ${data.date}<br>
|
||||
<b>Entities:</b> ${data.entities}
|
||||
`;
|
||||
tooltip.style.left = renderedPosition.x + 20 + 'px';
|
||||
tooltip.style.top = renderedPosition.y + 'px';
|
||||
document.body.appendChild(tooltip);
|
||||
});
|
||||
|
||||
cy.on('mouseout', 'node', function(evt) {
|
||||
if (tooltip) {
|
||||
tooltip.remove();
|
||||
tooltip = null;
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
// Reload graph with current settings
|
||||
function reloadGraph() {
|
||||
const nodeLimit = parseInt(document.getElementById('node-limit').value) || 50;
|
||||
const layoutName = document.getElementById('layout-select').value;
|
||||
initGraph(nodeLimit, layoutName);
|
||||
}
|
||||
|
||||
// Initialize with default settings (50 nodes, circle layout)
|
||||
initGraph(50, 'circle');
|
||||
|
||||
// Tab switching
|
||||
function switchTab(tabName) {
|
||||
document.querySelectorAll('.tab-content').forEach(tab => {
|
||||
tab.classList.remove('active');
|
||||
});
|
||||
document.querySelectorAll('.tab-button').forEach(btn => {
|
||||
btn.classList.remove('active');
|
||||
});
|
||||
|
||||
if (tabName === 'graph') {
|
||||
document.getElementById('graph-tab').classList.add('active');
|
||||
document.querySelectorAll('.tab-button')[0].classList.add('active');
|
||||
cy.resize(); // Resize graph when switching to it
|
||||
} else if (tabName === 'table') {
|
||||
document.getElementById('table-tab').classList.add('active');
|
||||
document.querySelectorAll('.tab-button')[1].classList.add('active');
|
||||
}
|
||||
}
|
||||
|
||||
// Table filtering
|
||||
document.getElementById('table-filter').addEventListener('input', function() {
|
||||
const filterValue = this.value.toLowerCase();
|
||||
const rows = document.querySelectorAll('#memory-table tbody tr');
|
||||
|
||||
rows.forEach(row => {
|
||||
const text = row.textContent.toLowerCase();
|
||||
if (text.includes(filterValue)) {
|
||||
row.style.display = '';
|
||||
} else {
|
||||
row.style.display = 'none';
|
||||
}
|
||||
});
|
||||
});
|
||||
</script>
|
||||
</body>
|
||||
</html>
|
||||
22
migrations/002_add_documents_table.sql
Normal file
22
migrations/002_add_documents_table.sql
Normal file
|
|
@ -0,0 +1,22 @@
|
|||
-- Migration: Add documents table and document_id to memory_units
|
||||
-- This enables document tracking, upsert, and cascade deletion
|
||||
|
||||
-- Create documents table
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id TEXT NOT NULL,
|
||||
agent_id TEXT NOT NULL,
|
||||
PRIMARY KEY (id, agent_id),
|
||||
original_text TEXT,
|
||||
content_hash TEXT,
|
||||
metadata JSONB DEFAULT '{}'::jsonb,
|
||||
created_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ DEFAULT NOW()
|
||||
);
|
||||
|
||||
-- Add document_id column to memory_units (nullable for backward compatibility)
|
||||
ALTER TABLE memory_units ADD COLUMN IF NOT EXISTS document_id TEXT REFERENCES documents(id) ON DELETE CASCADE;
|
||||
|
||||
-- Create indexes
|
||||
CREATE INDEX IF NOT EXISTS idx_documents_agent_id ON documents(agent_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_documents_content_hash ON documents(content_hash);
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_document_id ON memory_units(document_id);
|
||||
19
migrations/003_fix_documents_primary_key_v2.sql
Normal file
19
migrations/003_fix_documents_primary_key_v2.sql
Normal file
|
|
@ -0,0 +1,19 @@
|
|||
-- Migration: Fix documents table primary key to be composite (id, agent_id)
|
||||
-- This is a safer approach that doesn't drop the table
|
||||
|
||||
-- Step 1: Drop the foreign key constraint from memory_units
|
||||
ALTER TABLE memory_units DROP CONSTRAINT IF EXISTS memory_units_document_id_fkey;
|
||||
ALTER TABLE memory_units DROP CONSTRAINT IF EXISTS memory_units_document_fkey;
|
||||
|
||||
-- Step 2: Drop the old primary key on documents
|
||||
ALTER TABLE documents DROP CONSTRAINT IF EXISTS documents_pkey;
|
||||
|
||||
-- Step 3: Add the new composite primary key
|
||||
ALTER TABLE documents ADD PRIMARY KEY (id, agent_id);
|
||||
|
||||
-- Step 4: Add back the foreign key constraint with the composite key
|
||||
ALTER TABLE memory_units
|
||||
ADD CONSTRAINT memory_units_document_fkey
|
||||
FOREIGN KEY (document_id, agent_id)
|
||||
REFERENCES documents(id, agent_id)
|
||||
ON DELETE CASCADE;
|
||||
36
schema.sql
36
schema.sql
|
|
@ -6,10 +6,23 @@ CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
|
|||
-- TEMPORAL + SEMANTIC + ENTITY MEMORY ARCHITECTURE
|
||||
-- ============================================================================
|
||||
|
||||
-- Documents: Source of memory units (for tracking, updates, and deletion)
|
||||
CREATE TABLE IF NOT EXISTS documents (
|
||||
id TEXT NOT NULL, -- User-provided document ID
|
||||
agent_id TEXT NOT NULL,
|
||||
PRIMARY KEY (id, agent_id),
|
||||
original_text TEXT, -- Full original content (for context expansion)
|
||||
content_hash TEXT, -- SHA256 hash for deduplication
|
||||
metadata JSONB DEFAULT '{}'::jsonb, -- User-provided metadata
|
||||
created_at TIMESTAMPTZ DEFAULT NOW(),
|
||||
updated_at TIMESTAMPTZ DEFAULT NOW()
|
||||
);
|
||||
|
||||
-- Memory Units: Individual sentence-level memories
|
||||
CREATE TABLE IF NOT EXISTS memory_units (
|
||||
id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
|
||||
agent_id TEXT NOT NULL,
|
||||
document_id TEXT, -- Link to source document
|
||||
text TEXT NOT NULL,
|
||||
embedding vector(384), -- bge-small-en-v1.5 dimension
|
||||
context TEXT, -- What was happening when this memory was formed
|
||||
|
|
@ -63,12 +76,35 @@ CREATE TABLE IF NOT EXISTS memory_links (
|
|||
CREATE UNIQUE INDEX IF NOT EXISTS idx_memory_links_unique
|
||||
ON memory_links (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid));
|
||||
|
||||
-- ============================================================================
|
||||
-- FOREIGN KEY CONSTRAINTS
|
||||
-- ============================================================================
|
||||
|
||||
-- Add foreign key from memory_units to documents (composite key)
|
||||
DO $$
|
||||
BEGIN
|
||||
IF NOT EXISTS (
|
||||
SELECT 1 FROM pg_constraint WHERE conname = 'memory_units_document_fkey'
|
||||
) THEN
|
||||
ALTER TABLE memory_units
|
||||
ADD CONSTRAINT memory_units_document_fkey
|
||||
FOREIGN KEY (document_id, agent_id)
|
||||
REFERENCES documents(id, agent_id)
|
||||
ON DELETE CASCADE;
|
||||
END IF;
|
||||
END $$;
|
||||
|
||||
-- ============================================================================
|
||||
-- INDEXES
|
||||
-- ============================================================================
|
||||
|
||||
-- Document indexes
|
||||
CREATE INDEX IF NOT EXISTS idx_documents_agent_id ON documents(agent_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_documents_content_hash ON documents(content_hash);
|
||||
|
||||
-- Memory unit indexes
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_agent_id ON memory_units(agent_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_document_id ON memory_units(document_id);
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_event_date ON memory_units(event_date DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_agent_date ON memory_units(agent_id, event_date DESC);
|
||||
CREATE INDEX IF NOT EXISTS idx_memory_units_access_count ON memory_units(access_count DESC);
|
||||
|
|
|
|||
156
tests/test_document_tracking.py
Normal file
156
tests/test_document_tracking.py
Normal file
|
|
@ -0,0 +1,156 @@
|
|||
"""
|
||||
Tests for document tracking and upsert functionality.
|
||||
"""
|
||||
import os
|
||||
import pytest
|
||||
from datetime import datetime, timezone
|
||||
from memory import TemporalSemanticMemory
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_creation_and_retrieval():
|
||||
"""Test that documents are created and can be retrieved."""
|
||||
db_url = os.getenv("DATABASE_URL")
|
||||
if not db_url:
|
||||
pytest.skip("DATABASE_URL not set")
|
||||
|
||||
memory = TemporalSemanticMemory(db_url=db_url)
|
||||
|
||||
try:
|
||||
agent_id = f"test_doc_{datetime.now(timezone.utc).timestamp()}"
|
||||
document_id = "meeting-001"
|
||||
|
||||
# Store memory with document tracking
|
||||
await memory.put_async(
|
||||
agent_id=agent_id,
|
||||
content="Alice works at Google. Bob works at Microsoft.",
|
||||
context="Team meeting",
|
||||
document_id=document_id,
|
||||
document_metadata={"source": "meeting", "participants": ["Alice", "Bob"]}
|
||||
)
|
||||
|
||||
# Retrieve document
|
||||
doc = await memory.get_document(document_id, agent_id)
|
||||
|
||||
assert doc is not None
|
||||
assert doc["id"] == document_id
|
||||
assert doc["agent_id"] == agent_id
|
||||
assert "Alice works at Google" in doc["original_text"]
|
||||
assert doc["metadata"]["source"] == "meeting"
|
||||
assert doc["unit_count"] > 0
|
||||
|
||||
finally:
|
||||
await memory.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_upsert():
|
||||
"""Test that upsert deletes old units and creates new ones."""
|
||||
db_url = os.getenv("DATABASE_URL")
|
||||
if not db_url:
|
||||
pytest.skip("DATABASE_URL not set")
|
||||
|
||||
memory = TemporalSemanticMemory(db_url=db_url)
|
||||
|
||||
try:
|
||||
agent_id = f"test_upsert_{datetime.now(timezone.utc).timestamp()}"
|
||||
document_id = "meeting-002"
|
||||
|
||||
# First version
|
||||
units_v1 = await memory.put_async(
|
||||
agent_id=agent_id,
|
||||
content="Alice works at Google.",
|
||||
context="Initial",
|
||||
document_id=document_id
|
||||
)
|
||||
|
||||
# Get document stats
|
||||
doc_v1 = await memory.get_document(document_id, agent_id)
|
||||
count_v1 = doc_v1["unit_count"]
|
||||
|
||||
# Upsert with different content
|
||||
units_v2 = await memory.put_async(
|
||||
agent_id=agent_id,
|
||||
content="Alice works at Microsoft. Bob works at Apple.",
|
||||
context="Updated",
|
||||
document_id=document_id,
|
||||
upsert=True
|
||||
)
|
||||
|
||||
# Get updated document stats
|
||||
doc_v2 = await memory.get_document(document_id, agent_id)
|
||||
count_v2 = doc_v2["unit_count"]
|
||||
|
||||
# Verify old units were replaced
|
||||
assert "Microsoft" in doc_v2["original_text"]
|
||||
assert doc_v2["updated_at"] > doc_v1["created_at"]
|
||||
|
||||
# Different unit IDs (old ones deleted, new ones created)
|
||||
assert set(units_v1).isdisjoint(set(units_v2))
|
||||
|
||||
finally:
|
||||
await memory.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_document_deletion():
|
||||
"""Test that deleting a document cascades to memory units."""
|
||||
db_url = os.getenv("DATABASE_URL")
|
||||
if not db_url:
|
||||
pytest.skip("DATABASE_URL not set")
|
||||
|
||||
memory = TemporalSemanticMemory(db_url=db_url)
|
||||
|
||||
try:
|
||||
agent_id = f"test_delete_{datetime.now(timezone.utc).timestamp()}"
|
||||
document_id = "meeting-003"
|
||||
|
||||
# Create document
|
||||
await memory.put_async(
|
||||
agent_id=agent_id,
|
||||
content="Alice works at Google.",
|
||||
context="Test",
|
||||
document_id=document_id
|
||||
)
|
||||
|
||||
# Verify it exists
|
||||
doc = await memory.get_document(document_id, agent_id)
|
||||
assert doc is not None
|
||||
assert doc["unit_count"] > 0
|
||||
|
||||
# Delete document
|
||||
result = await memory.delete_document(document_id, agent_id)
|
||||
assert result["document_deleted"] == 1
|
||||
assert result["memory_units_deleted"] > 0
|
||||
|
||||
# Verify it's gone
|
||||
doc_after = await memory.get_document(document_id, agent_id)
|
||||
assert doc_after is None
|
||||
|
||||
finally:
|
||||
await memory.close()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_memory_without_document():
|
||||
"""Test that memories can still be created without document tracking."""
|
||||
db_url = os.getenv("DATABASE_URL")
|
||||
if not db_url:
|
||||
pytest.skip("DATABASE_URL not set")
|
||||
|
||||
memory = TemporalSemanticMemory(db_url=db_url)
|
||||
|
||||
try:
|
||||
agent_id = f"test_no_doc_{datetime.now(timezone.utc).timestamp()}"
|
||||
|
||||
# Create memory without document_id (backward compatibility)
|
||||
units = await memory.put_async(
|
||||
agent_id=agent_id,
|
||||
content="Alice works at Google.",
|
||||
context="Test"
|
||||
)
|
||||
|
||||
assert len(units) > 0
|
||||
|
||||
finally:
|
||||
await memory.close()
|
||||
|
|
@ -20,7 +20,10 @@ from typing import Optional
|
|||
sys.path.insert(0, str(Path(__file__).parent.parent))
|
||||
from memory import TemporalSemanticMemory
|
||||
|
||||
import logging
|
||||
|
||||
load_dotenv()
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
app = FastAPI(title="Memory Graph API", version="1.0.0")
|
||||
|
||||
|
|
@ -34,6 +37,7 @@ class SearchRequest(BaseModel):
|
|||
agent_id: str = "default"
|
||||
thinking_budget: int = 100
|
||||
top_k: int = 10
|
||||
mmr_lambda: float = 0.5
|
||||
|
||||
|
||||
async def get_graph_data():
|
||||
|
|
@ -208,7 +212,8 @@ async def api_search(request: SearchRequest):
|
|||
query=request.query,
|
||||
thinking_budget=request.thinking_budget,
|
||||
top_k=request.top_k,
|
||||
enable_trace=True
|
||||
enable_trace=True,
|
||||
mmr_lambda=request.mmr_lambda
|
||||
)
|
||||
|
||||
# Convert trace to dict
|
||||
|
|
|
|||
|
|
@ -313,6 +313,10 @@ function addDebugPane() {
|
|||
<label style="font-weight: bold; display: block; margin-bottom: 3px; font-size: 12px;">Top K:</label>
|
||||
<input type="number" id="search-top-k-${paneId}" value="10" min="1" max="50" style="width: 60px; padding: 6px; border: 1px solid #ccc; border-radius: 4px; font-size: 12px;">
|
||||
</div>
|
||||
<div>
|
||||
<label style="font-weight: bold; display: block; margin-bottom: 3px; font-size: 12px;" title="MMR Lambda: 0=max diversity, 1=no diversity">MMR λ:</label>
|
||||
<input type="number" id="search-mmr-lambda-${paneId}" value="0.5" min="0" max="1" step="0.1" style="width: 60px; padding: 6px; border: 1px solid #ccc; border-radius: 4px; font-size: 12px;">
|
||||
</div>
|
||||
<button onclick="runSearchInPane(${paneId})" style="padding: 6px 16px; background: #42a5f5; color: white; border: none; border-radius: 4px; cursor: pointer; font-weight: bold; font-size: 12px;">🔍 Search</button>
|
||||
</div>
|
||||
</div>
|
||||
|
|
@ -336,6 +340,10 @@ function addDebugPane() {
|
|||
<label>
|
||||
<input type="checkbox" id="debug-highlight-path-${paneId}"> Highlight top result path
|
||||
</label>
|
||||
<span style="margin-left: 15px;">
|
||||
<input type="text" id="graph-search-${paneId}" placeholder="Find nodes..." style="width: 150px; padding: 4px 8px; border: 1px solid #ccc; border-radius: 4px; font-size: 12px;">
|
||||
<span id="graph-search-count-${paneId}" style="margin-left: 5px; font-size: 12px; color: #666;"></span>
|
||||
</span>
|
||||
</span>
|
||||
</div>
|
||||
<div class="debug-viz-container">
|
||||
|
|
@ -394,6 +402,14 @@ function addDebugPane() {
|
|||
}
|
||||
});
|
||||
|
||||
// Add search input listener for finding nodes
|
||||
document.getElementById(`graph-search-${paneId}`).addEventListener('input', function(e) {
|
||||
const pane = debugPanes.find(p => p.id === paneId);
|
||||
if (pane && pane.cy) {
|
||||
highlightMatchingNodes(paneId, e.target.value);
|
||||
}
|
||||
});
|
||||
|
||||
debugPanes.push({
|
||||
id: paneId,
|
||||
element: paneDiv,
|
||||
|
|
@ -446,6 +462,7 @@ window.runSearchInPane = async function(paneId) {
|
|||
const agentId = document.getElementById(`search-agent-${paneId}`).value;
|
||||
const thinkingBudget = parseInt(document.getElementById(`search-budget-${paneId}`).value);
|
||||
const topK = parseInt(document.getElementById(`search-top-k-${paneId}`).value);
|
||||
const mmrLambda = parseFloat(document.getElementById(`search-mmr-lambda-${paneId}`).value);
|
||||
const statusBar = document.getElementById(`debug-status-${paneId}`);
|
||||
|
||||
if (!query) {
|
||||
|
|
@ -465,7 +482,8 @@ window.runSearchInPane = async function(paneId) {
|
|||
query: query,
|
||||
agent_id: agentId,
|
||||
thinking_budget: thinkingBudget,
|
||||
top_k: topK
|
||||
top_k: topK,
|
||||
mmr_lambda: mmrLambda
|
||||
})
|
||||
});
|
||||
|
||||
|
|
@ -536,16 +554,28 @@ function renderResultsTable(paneId, results, trace) {
|
|||
return;
|
||||
}
|
||||
|
||||
// Check if MMR was used
|
||||
const mmrUsed = results.some(r => r.mmr_score !== null && r.mmr_score !== undefined);
|
||||
|
||||
let html = `
|
||||
<div style="padding: 20px; overflow: auto; height: 100%;">
|
||||
<h3>Search Results (${results.length} memories)</h3>
|
||||
<p style="color: #666; font-size: 13px; margin-bottom: 15px;">
|
||||
<p style="color: #666; font-size: 13px; margin-bottom: 10px;">
|
||||
Query: "${trace.query.query_text}"
|
||||
</p>
|
||||
${mmrUsed ? `
|
||||
<div style="background: #e3f2fd; padding: 10px; border-left: 4px solid #2196f3; margin-bottom: 15px; font-size: 12px;">
|
||||
<strong>MMR Diversification Active:</strong>
|
||||
🎯 = Diversified pick (selected for variety) |
|
||||
<span style="background: #fff3e0; padding: 2px 6px; border-radius: 3px;">Orange background</span> = Rank changed by MMR |
|
||||
<strong>Orig Rank</strong> shows position before MMR reranking
|
||||
</div>
|
||||
` : ''}
|
||||
<table style="width: 100%; border-collapse: collapse; font-size: 12px;">
|
||||
<thead>
|
||||
<tr style="background: #f0f0f0; border: 2px solid #333;">
|
||||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;">Rank</th>
|
||||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;" title="Original rank before MMR">Orig Rank</th>
|
||||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;">Text</th>
|
||||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;">Context</th>
|
||||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;">Date</th>
|
||||
|
|
@ -554,6 +584,9 @@ function renderResultsTable(paneId, results, trace) {
|
|||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;" title="Semantic similarity to query">Similarity</th>
|
||||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;" title="Recency boost">Recency</th>
|
||||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;" title="Frequency boost">Frequency</th>
|
||||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;" title="MMR score (relevance - diversity penalty)">MMR Score</th>
|
||||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;" title="Normalized relevance component">MMR Rel</th>
|
||||
<th style="padding: 8px; text-align: left; border: 1px solid #ddd;" title="Max similarity to already selected">MMR Sim</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
|
|
@ -570,9 +603,21 @@ function renderResultsTable(paneId, results, trace) {
|
|||
const recency = visit ? (visit.weights.recency || 0) : 0;
|
||||
const frequency = visit ? (visit.weights.frequency || 0) : 0;
|
||||
|
||||
// Get MMR information
|
||||
const originalRank = result.original_rank || (idx + 1);
|
||||
const mmrScore = result.mmr_score;
|
||||
const mmrRelevance = result.mmr_relevance;
|
||||
const mmrMaxSim = result.mmr_max_similarity;
|
||||
const isDiversified = result.mmr_diversified || false;
|
||||
|
||||
// Highlight diversified results
|
||||
const rowBg = isDiversified ? '#fff3e0' : 'white';
|
||||
const rankDisplay = originalRank !== (idx + 1) ? `<span style="color: #ff9800;" title="Rank changed by MMR">↑${originalRank}</span>` : originalRank;
|
||||
|
||||
html += `
|
||||
<tr style="border: 1px solid #ddd;">
|
||||
<td style="padding: 8px; border: 1px solid #ddd; font-weight: bold;">#${idx + 1}</td>
|
||||
<tr style="border: 1px solid #ddd; background: ${rowBg};">
|
||||
<td style="padding: 8px; border: 1px solid #ddd; font-weight: bold;">#${idx + 1}${isDiversified ? ' 🎯' : ''}</td>
|
||||
<td style="padding: 8px; border: 1px solid #ddd;">${rankDisplay}</td>
|
||||
<td style="padding: 8px; border: 1px solid #ddd; max-width: 300px;">${result.text}</td>
|
||||
<td style="padding: 8px; border: 1px solid #ddd; max-width: 150px;">${result.context || 'N/A'}</td>
|
||||
<td style="padding: 8px; border: 1px solid #ddd; white-space: nowrap;">${result.event_date ? new Date(result.event_date).toLocaleDateString() : 'N/A'}</td>
|
||||
|
|
@ -581,6 +626,9 @@ function renderResultsTable(paneId, results, trace) {
|
|||
<td style="padding: 8px; border: 1px solid #ddd;">${similarity.toFixed(4)}</td>
|
||||
<td style="padding: 8px; border: 1px solid #ddd;">${recency.toFixed(4)}</td>
|
||||
<td style="padding: 8px; border: 1px solid #ddd;">${frequency.toFixed(4)}</td>
|
||||
<td style="padding: 8px; border: 1px solid #ddd;">${mmrScore !== null && mmrScore !== undefined ? mmrScore.toFixed(4) : '-'}</td>
|
||||
<td style="padding: 8px; border: 1px solid #ddd;">${mmrRelevance !== null && mmrRelevance !== undefined ? mmrRelevance.toFixed(4) : '-'}</td>
|
||||
<td style="padding: 8px; border: 1px solid #ddd;">${mmrMaxSim !== null && mmrMaxSim !== undefined ? mmrMaxSim.toFixed(4) : '-'}</td>
|
||||
</tr>
|
||||
`;
|
||||
});
|
||||
|
|
@ -1023,6 +1071,73 @@ function visualizeTrace(paneId, trace) {
|
|||
}
|
||||
}
|
||||
|
||||
function highlightMatchingNodes(paneId, searchText) {
|
||||
const pane = debugPanes.find(p => p.id === paneId);
|
||||
if (!pane || !pane.cy) return;
|
||||
|
||||
const countSpan = document.getElementById(`graph-search-count-${paneId}`);
|
||||
|
||||
// If search is empty, reset all styles
|
||||
if (!searchText || searchText.trim() === '') {
|
||||
pane.cy.nodes().style({
|
||||
'opacity': 1,
|
||||
'border-width': 2,
|
||||
'border-color': '#333'
|
||||
});
|
||||
pane.cy.edges().style({
|
||||
'opacity': 0.8
|
||||
});
|
||||
if (countSpan) countSpan.textContent = '';
|
||||
return;
|
||||
}
|
||||
|
||||
const searchLower = searchText.toLowerCase();
|
||||
let matchCount = 0;
|
||||
const totalNodes = pane.cy.nodes().length;
|
||||
|
||||
// Check each node for matches
|
||||
pane.cy.nodes().forEach(node => {
|
||||
const data = node.data();
|
||||
const text = (data.text || '').toLowerCase();
|
||||
const label = (data.label || '').toLowerCase();
|
||||
|
||||
const matches = text.includes(searchLower) || label.includes(searchLower);
|
||||
|
||||
if (matches) {
|
||||
matchCount++;
|
||||
// Highlight matching nodes - full opacity with thicker orange border
|
||||
node.style({
|
||||
'opacity': 1,
|
||||
'border-width': 4,
|
||||
'border-color': '#ff6f00'
|
||||
});
|
||||
} else {
|
||||
// Dim non-matching nodes
|
||||
node.style({
|
||||
'opacity': 0.2,
|
||||
'border-width': 2,
|
||||
'border-color': '#333'
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// Dim all edges
|
||||
pane.cy.edges().style({
|
||||
'opacity': 0.2
|
||||
});
|
||||
|
||||
// Update counter
|
||||
if (countSpan) {
|
||||
if (matchCount === 0) {
|
||||
countSpan.textContent = '(no matches)';
|
||||
countSpan.style.color = '#d32f2f';
|
||||
} else {
|
||||
countSpan.textContent = `(${matchCount} of ${totalNodes})`;
|
||||
countSpan.style.color = '#43a047';
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Table filtering
|
||||
document.getElementById('table-filter').addEventListener('input', function() {
|
||||
const filterValue = this.value.toLowerCase();
|
||||
|
|
|
|||
|
|
@ -6,10 +6,13 @@ window.loadLocomoResults = async function() {
|
|||
try {
|
||||
const response = await fetch('/api/locomo');
|
||||
locomoData = await response.json();
|
||||
console.log('Loaded locomo data:', locomoData);
|
||||
renderLocomoResults();
|
||||
} catch (e) {
|
||||
console.error('Error loading benchmark results:', e);
|
||||
document.getElementById('locomo-content').innerHTML = `
|
||||
<div class="error-message">Error loading benchmark results: ${e.message}</div>
|
||||
<div class="error-message">Error loading benchmark results: ${e.message}<br>
|
||||
Check console for details.</div>
|
||||
`;
|
||||
}
|
||||
}
|
||||
|
|
@ -19,6 +22,36 @@ function renderLocomoResults() {
|
|||
|
||||
const content = document.getElementById('locomo-content');
|
||||
|
||||
try {
|
||||
// Handle both old and new structure
|
||||
const results = locomoData.item_results || locomoData.conversation_results || [];
|
||||
const numItems = locomoData.num_items || results.length;
|
||||
|
||||
console.log('Rendering results:', { resultsCount: results.length, numItems });
|
||||
|
||||
// Calculate per-category statistics
|
||||
const categoryStats = {
|
||||
1: { name: 'Multi-hop', correct: 0, total: 0 }, // category 1
|
||||
2: { name: 'Single-hop', correct: 0, total: 0 }, // category 2
|
||||
3: { name: 'Temporal', correct: 0, total: 0 }, // category 3
|
||||
4: { name: 'Open-domain', correct: 0, total: 0 } // category 4
|
||||
};
|
||||
|
||||
// Aggregate across all items
|
||||
results.forEach(item => {
|
||||
if (item.metrics && item.metrics.detailed_results) {
|
||||
item.metrics.detailed_results.forEach(result => {
|
||||
const category = result.category;
|
||||
if (categoryStats[category]) {
|
||||
categoryStats[category].total++;
|
||||
if (result.is_correct) {
|
||||
categoryStats[category].correct++;
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
});
|
||||
|
||||
// Overall stats
|
||||
const overallHtml = `
|
||||
<div style="background: #f9f9f9; padding: 20px; border: 2px solid #333; border-radius: 8px; margin-bottom: 20px;">
|
||||
|
|
@ -33,10 +66,25 @@ function renderLocomoResults() {
|
|||
<div class="stat-value">${locomoData.total_correct} / ${locomoData.total_questions}</div>
|
||||
</div>
|
||||
<div class="stat-item">
|
||||
<div class="stat-label">Conversations</div>
|
||||
<div class="stat-value">${locomoData.conversation_results.length}</div>
|
||||
<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(200px, 1fr));">
|
||||
${Object.values(categoryStats).map(cat => {
|
||||
const accuracy = cat.total > 0 ? ((cat.correct / cat.total) * 100).toFixed(1) : 0;
|
||||
const color = accuracy >= 70 ? '#43a047' : accuracy >= 50 ? '#ff9800' : '#e53935';
|
||||
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}</div>
|
||||
</div>
|
||||
`;
|
||||
}).join('')}
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
|
||||
|
|
@ -50,35 +98,52 @@ function renderLocomoResults() {
|
|||
</div>
|
||||
`;
|
||||
|
||||
// Build conversation sections
|
||||
let conversationsHtml = '';
|
||||
locomoData.conversation_results.forEach((conv, idx) => {
|
||||
const accuracy = conv.metrics.accuracy.toFixed(2);
|
||||
const correctCount = conv.metrics.correct;
|
||||
const totalCount = conv.metrics.total;
|
||||
// Build item sections
|
||||
let itemsHtml = '';
|
||||
results.forEach((item, idx) => {
|
||||
const itemId = item.item_id || item.sample_id || `item-${idx}`;
|
||||
const accuracy = item.metrics.accuracy.toFixed(2);
|
||||
const correctCount = item.metrics.correct;
|
||||
const totalCount = item.metrics.total;
|
||||
|
||||
conversationsHtml += `
|
||||
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>📊 ${conv.sample_id}</span>
|
||||
<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;">
|
||||
${renderConversationDetails(conv)}
|
||||
${renderConversationDetails(item)}
|
||||
</div>
|
||||
</div>
|
||||
`;
|
||||
});
|
||||
|
||||
content.innerHTML = overallHtml + filterHtml + conversationsHtml;
|
||||
content.innerHTML = overallHtml + filterHtml + itemsHtml;
|
||||
} catch (e) {
|
||||
console.error('Error rendering Locomo 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 renderConversationDetails(conv) {
|
||||
if (!conv || !conv.metrics) {
|
||||
return '<div style="padding: 20px; color: #666;">No metrics available</div>';
|
||||
}
|
||||
|
||||
const results = conv.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">';
|
||||
|
||||
|
|
@ -150,12 +215,13 @@ function renderConversationDetails(conv) {
|
|||
}
|
||||
|
||||
function renderRetrievedMemories(memories) {
|
||||
if (!memories || memories.length === 0) {
|
||||
if (!memories || !Array.isArray(memories) || memories.length === 0) {
|
||||
return '<div style="padding: 8px; color: #999;">No memories retrieved</div>';
|
||||
}
|
||||
|
||||
let html = '<div style="margin-top: 8px;">';
|
||||
memories.forEach((mem, idx) => {
|
||||
if (!mem) return;
|
||||
html += `
|
||||
<div style="padding: 8px; background: #f5f5f5; border-left: 3px solid #42a5f5; margin-bottom: 8px;">
|
||||
<div style="font-size: 11px; color: #666; margin-bottom: 4px;">
|
||||
|
|
|
|||
Loading…
Reference in a new issue