many improvements

This commit is contained in:
Nicolò Boschi 2025-11-26 18:21:10 +01:00
parent b1ff2e8823
commit a3ad76d165
33 changed files with 1679947 additions and 1557571 deletions

View file

@ -0,0 +1,89 @@
"""add_observation_fact_type
Revision ID: 5b2c6d8e9f01
Revises: 4a8b3c5d6e7f
Create Date: 2025-11-26 10:00:00.000000
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '5b2c6d8e9f01'
down_revision: Union[str, Sequence[str], None] = '4a8b3c5d6e7f'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Drop old constraints
op.execute("""
ALTER TABLE memory_units
DROP CONSTRAINT IF EXISTS confidence_score_fact_type_check
""")
op.execute("""
ALTER TABLE memory_units
DROP CONSTRAINT IF EXISTS memory_units_fact_type_check
""")
# Add new fact_type constraint including 'observation'
op.execute("""
ALTER TABLE memory_units
ADD CONSTRAINT memory_units_fact_type_check
CHECK (fact_type IN ('world', 'agent', 'opinion', 'observation'))
""")
# Add new confidence_score constraint allowing observation to have optional confidence
op.execute("""
ALTER TABLE memory_units
ADD CONSTRAINT confidence_score_fact_type_check
CHECK (
(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR
(fact_type = 'observation') OR
(fact_type NOT IN ('opinion', 'observation') AND confidence_score IS NULL)
)
""")
# Add index for observation fact_type queries
op.execute("""
CREATE INDEX IF NOT EXISTS idx_memory_units_observation_date
ON memory_units (agent_id, event_date DESC)
WHERE fact_type = 'observation'
""")
def downgrade() -> None:
"""Downgrade schema."""
# Drop observation index
op.execute("DROP INDEX IF EXISTS idx_memory_units_observation_date")
# Drop new constraints
op.execute("""
ALTER TABLE memory_units
DROP CONSTRAINT IF EXISTS confidence_score_fact_type_check
""")
op.execute("""
ALTER TABLE memory_units
DROP CONSTRAINT IF EXISTS memory_units_fact_type_check
""")
# Restore old fact_type constraint
op.execute("""
ALTER TABLE memory_units
ADD CONSTRAINT memory_units_fact_type_check
CHECK (fact_type IN ('world', 'agent', 'opinion'))
""")
# Restore old confidence_score constraint
op.execute("""
ALTER TABLE memory_units
ADD CONSTRAINT confidence_score_fact_type_check
CHECK (
(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR
(fact_type != 'opinion' AND confidence_score IS NULL)
)
""")

View file

@ -0,0 +1,117 @@
"""add unique constraint on entities (agent_id, canonical_name)
Revision ID: 7d4e6f0a3b12
Revises: 5b2c6d8e9f01
Create Date: 2024-01-01 00:00:00.000000
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '7d4e6f0a3b12'
down_revision: Union[str, None] = '5b2c6d8e9f01'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
# First, deduplicate existing entities by merging duplicates
# Keep the one with highest mention_count, update unit_entities to point to it
op.execute("""
-- Create temp table with canonical entity per (agent_id, name)
CREATE TEMP TABLE canonical_entities AS
SELECT DISTINCT ON (agent_id, LOWER(canonical_name))
id as keep_id,
agent_id,
LOWER(canonical_name) as name_lower
FROM entities
ORDER BY agent_id, LOWER(canonical_name), mention_count DESC, first_seen ASC;
-- Get all entity IDs that will be removed (duplicates)
CREATE TEMP TABLE duplicate_entities AS
SELECT e.id as dup_id, ce.keep_id
FROM entities e
JOIN canonical_entities ce ON e.agent_id = ce.agent_id AND LOWER(e.canonical_name) = ce.name_lower
WHERE e.id != ce.keep_id;
-- Update unit_entities to point to canonical entity
UPDATE unit_entities ue
SET entity_id = de.keep_id
FROM duplicate_entities de
WHERE ue.entity_id = de.dup_id;
-- Delete duplicate unit_entities that now exist
DELETE FROM unit_entities a
USING unit_entities b
WHERE a.unit_id = b.unit_id
AND a.entity_id = b.entity_id
AND a.ctid < b.ctid;
-- For entity_cooccurrences, we need to be careful about the check constraint
-- First, collect all cooccurrences that need updating into a temp table with correct ordering
CREATE TEMP TABLE new_cooccurrences AS
SELECT DISTINCT
LEAST(
COALESCE(de1.keep_id, ec.entity_id_1),
COALESCE(de2.keep_id, ec.entity_id_2)
) as entity_id_1,
GREATEST(
COALESCE(de1.keep_id, ec.entity_id_1),
COALESCE(de2.keep_id, ec.entity_id_2)
) as entity_id_2,
SUM(ec.cooccurrence_count) as cooccurrence_count,
MAX(ec.last_cooccurred) as last_cooccurred
FROM entity_cooccurrences ec
LEFT JOIN duplicate_entities de1 ON ec.entity_id_1 = de1.dup_id
LEFT JOIN duplicate_entities de2 ON ec.entity_id_2 = de2.dup_id
GROUP BY
LEAST(COALESCE(de1.keep_id, ec.entity_id_1), COALESCE(de2.keep_id, ec.entity_id_2)),
GREATEST(COALESCE(de1.keep_id, ec.entity_id_1), COALESCE(de2.keep_id, ec.entity_id_2));
-- Delete rows where entity_id_1 = entity_id_2 (self-references after merge)
DELETE FROM new_cooccurrences WHERE entity_id_1 = entity_id_2;
-- Delete all old cooccurrences
DELETE FROM entity_cooccurrences;
-- Insert the merged cooccurrences
INSERT INTO entity_cooccurrences (entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred)
SELECT entity_id_1, entity_id_2, cooccurrence_count, last_cooccurred
FROM new_cooccurrences;
-- Update mention counts on canonical entities
UPDATE entities e
SET mention_count = e.mention_count + COALESCE(
(SELECT SUM(e2.mention_count)
FROM entities e2
JOIN duplicate_entities de ON e2.id = de.dup_id
WHERE de.keep_id = e.id),
0
)
WHERE e.id IN (SELECT keep_id FROM duplicate_entities);
-- Delete duplicate entities
DELETE FROM entities
WHERE id IN (SELECT dup_id FROM duplicate_entities);
-- Cleanup temp tables
DROP TABLE new_cooccurrences;
DROP TABLE duplicate_entities;
DROP TABLE canonical_entities;
""")
# Add unique constraint (case-insensitive)
op.create_index(
'idx_entities_agent_canonical_unique',
'entities',
[sa.text('agent_id'), sa.text('LOWER(canonical_name)')],
unique=True
)
def downgrade() -> None:
op.drop_index('idx_entities_agent_canonical_unique', table_name='entities')

View file

@ -4,6 +4,7 @@ FastAPI application factory and API routes for memory system.
This module provides the create_app function to create and configure
the FastAPI application with all API endpoints.
"""
import json
import logging
import uuid
from pathlib import Path
@ -12,6 +13,22 @@ from datetime import datetime
from contextlib import asynccontextmanager
from fastapi import FastAPI, HTTPException, Query
def _parse_metadata(metadata: Any) -> Dict[str, Any]:
"""Parse metadata that may be a dict, JSON string, or None."""
if metadata is None:
return {}
if isinstance(metadata, dict):
return metadata
if isinstance(metadata, str):
try:
return json.loads(metadata)
except json.JSONDecodeError:
return {}
return {}
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field, ConfigDict
@ -45,7 +62,9 @@ class SearchRequest(BaseModel):
"max_tokens": 4096,
"trace": True,
"question_date": "2023-05-30T23:40:00",
"metadata_filter": [{"key": "source", "value": "slack", "match_unset": True}]
"metadata_filter": [{"key": "source", "value": "slack", "match_unset": True}],
"include_entities": True,
"max_entity_tokens": 500
}
})
@ -56,6 +75,8 @@ class SearchRequest(BaseModel):
trace: bool = False
question_date: Optional[str] = None # ISO format date string (e.g., "2023-05-30T23:40:00")
metadata_filter: Optional[List[MetadataFilter]] = Field(default=None, description="Filter by metadata. Multiple filters are ANDed together.")
include_entities: bool = Field(default=False, description="Whether to include entity observations in the response")
max_entity_tokens: int = Field(default=500, description="Maximum tokens for entity observations")
class SearchResult(BaseModel):
@ -67,8 +88,11 @@ class SearchResult(BaseModel):
"id": "123e4567-e89b-12d3-a456-426614174000",
"text": "Alice works at Google on the AI team",
"type": "world",
"entities": ["Alice", "Google"],
"context": "work info",
"event_date": "2024-01-15T10:30:00Z",
"occurred_start": "2024-01-15T10:30:00Z",
"occurred_end": "2024-01-15T10:30:00Z",
"mentioned_at": "2024-01-15T10:30:00Z",
"document_id": "session_abc123",
"metadata": {"source": "slack"}
}
@ -77,13 +101,92 @@ class SearchResult(BaseModel):
id: str
text: str
type: Optional[str] = None # fact type: world, agent, opinion
type: Optional[str] = None # fact type: world, agent, opinion, observation
entities: Optional[List[str]] = None # Entity names mentioned in this fact
context: Optional[str] = None
event_date: Optional[str] = None # ISO format date string
occurred_start: Optional[str] = None # ISO format date when the event started
occurred_end: Optional[str] = None # ISO format date when the event ended
mentioned_at: Optional[str] = None # ISO format date when the fact was mentioned
document_id: Optional[str] = None # Document this memory belongs to
metadata: Optional[Dict[str, str]] = None # User-defined metadata
class EntityObservationResponse(BaseModel):
"""An observation about an entity."""
text: str
mentioned_at: Optional[str] = None
class EntityStateResponse(BaseModel):
"""Current mental model of an entity."""
entity_id: str
canonical_name: str
observations: List[EntityObservationResponse]
class EntityListItem(BaseModel):
"""Entity list item with summary."""
model_config = ConfigDict(json_schema_extra={
"example": {
"id": "123e4567-e89b-12d3-a456-426614174000",
"canonical_name": "John",
"mention_count": 15,
"first_seen": "2024-01-15T10:30:00Z",
"last_seen": "2024-02-01T14:00:00Z"
}
})
id: str
canonical_name: str
mention_count: int
first_seen: Optional[str] = None
last_seen: Optional[str] = None
metadata: Optional[Dict[str, Any]] = None
class EntityListResponse(BaseModel):
"""Response model for entity list endpoint."""
model_config = ConfigDict(json_schema_extra={
"example": {
"entities": [
{
"id": "123e4567-e89b-12d3-a456-426614174000",
"canonical_name": "John",
"mention_count": 15,
"first_seen": "2024-01-15T10:30:00Z",
"last_seen": "2024-02-01T14:00:00Z"
}
]
}
})
entities: List[EntityListItem]
class EntityDetailResponse(BaseModel):
"""Response model for entity detail endpoint."""
model_config = ConfigDict(json_schema_extra={
"example": {
"id": "123e4567-e89b-12d3-a456-426614174000",
"canonical_name": "John",
"mention_count": 15,
"first_seen": "2024-01-15T10:30:00Z",
"last_seen": "2024-02-01T14:00:00Z",
"observations": [
{"text": "John works at Google", "mentioned_at": "2024-01-15T10:30:00Z"}
]
}
})
id: str
canonical_name: str
mention_count: int
first_seen: Optional[str] = None
last_seen: Optional[str] = None
metadata: Optional[Dict[str, Any]] = None
observations: List[EntityObservationResponse]
class SearchResponse(BaseModel):
"""Response model for search endpoints."""
model_config = ConfigDict(json_schema_extra={
@ -93,20 +196,32 @@ class SearchResponse(BaseModel):
"id": "123e4567-e89b-12d3-a456-426614174000",
"text": "Alice works at Google on the AI team",
"type": "world",
"entities": ["Alice", "Google"],
"context": "work info",
"event_date": "2024-01-15T10:30:00Z"
"occurred_start": "2024-01-15T10:30:00Z",
"occurred_end": "2024-01-15T10:30:00Z"
}
],
"trace": {
"query": "What did Alice say about machine learning?",
"num_results": 1,
"time_seconds": 0.123
},
"entities": {
"Alice": {
"entity_id": "123e4567-e89b-12d3-a456-426614174001",
"canonical_name": "Alice",
"observations": [
{"text": "Alice works at Google on the AI team", "mentioned_at": "2024-01-15T10:30:00Z"}
]
}
}
}
})
results: List[SearchResult]
trace: Optional[Dict[str, Any]] = None
entities: Optional[Dict[str, EntityStateResponse]] = Field(default=None, description="Entity states for entities mentioned in results")
class MemoryItem(BaseModel):
@ -219,7 +334,8 @@ class ThinkFact(BaseModel):
"text": "AI is used in healthcare",
"type": "world",
"context": "healthcare discussion",
"event_date": "2024-01-15T10:30:00Z"
"occurred_start": "2024-01-15T10:30:00Z",
"occurred_end": "2024-01-15T10:30:00Z"
}
})
@ -227,7 +343,8 @@ class ThinkFact(BaseModel):
text: str
type: Optional[str] = None # fact type: world, agent, opinion
context: Optional[str] = None
event_date: Optional[str] = None
occurred_start: Optional[str] = None
occurred_end: Optional[str] = None
class ThinkResponse(BaseModel):
@ -691,6 +808,9 @@ def _register_routes(app: FastAPI):
- 'world': General knowledge about people, places, events, and things that happen
- 'agent': Memories about what the AI agent did, actions taken, and tasks performed
- 'opinion': The agent's formed beliefs, perspectives, and viewpoints
- 'observation': Synthesized observations about entities (generated automatically)
Set include_entities=true to get entity observations alongside search results.
""",
operation_id="search_memories"
)
@ -698,11 +818,11 @@ def _register_routes(app: FastAPI):
"""Run a search and return results with trace."""
try:
# Validate fact_type(s)
valid_fact_types = ["world", "agent", "opinion"]
valid_fact_types = ["world", "agent", "opinion", "observation"]
# Default to all fact types if not specified
# Default to world, agent, opinion if not specified (exclude observation by default)
if not request.fact_type:
request.fact_type = valid_fact_types
request.fact_type = ["world", "agent", "opinion"]
else:
for ft in request.fact_type:
if ft not in valid_fact_types:
@ -730,7 +850,9 @@ def _register_routes(app: FastAPI):
max_tokens=request.max_tokens,
enable_trace=request.trace,
fact_type=request.fact_type,
question_date=question_date
question_date=question_date,
include_entities=request.include_entities,
max_entity_tokens=request.max_entity_tokens
)
# Convert core MemoryFact objects to API SearchResult objects (excluding internal metrics)
@ -739,15 +861,34 @@ def _register_routes(app: FastAPI):
id=fact.id,
text=fact.text,
type=fact.fact_type,
entities=fact.entities,
context=fact.context,
event_date=fact.event_date
occurred_start=fact.occurred_start,
occurred_end=fact.occurred_end,
mentioned_at=fact.mentioned_at,
document_id=fact.document_id
)
for fact in core_result.results
]
# Convert core EntityState objects to API EntityStateResponse objects
entities_response = None
if core_result.entities:
entities_response = {}
for name, state in core_result.entities.items():
entities_response[name] = EntityStateResponse(
entity_id=state.entity_id,
canonical_name=state.canonical_name,
observations=[
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at)
for obs in state.observations
]
)
return SearchResponse(
results=search_results,
trace=core_result.trace
trace=core_result.trace,
entities=entities_response
)
except HTTPException:
raise
@ -795,7 +936,8 @@ def _register_routes(app: FastAPI):
text=fact.text,
type=fact.fact_type,
context=fact.context,
event_date=fact.event_date
occurred_start=fact.occurred_start,
occurred_end=fact.occurred_end
))
return ThinkResponse(
@ -951,6 +1093,139 @@ def _register_routes(app: FastAPI):
print(f"Error in /api/v1/agents/{agent_id}/stats: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.get(
"/api/v1/agents/{agent_id}/entities",
response_model=EntityListResponse,
tags=["Entities"],
summary="List entities",
description="List all entities (people, organizations, etc.) known by the agent, ordered by mention count.",
operation_id="list_entities"
)
async def api_list_entities(
agent_id: str,
limit: int = Query(default=100, description="Maximum number of entities to return")
):
"""List entities for an agent."""
try:
entities = await app.state.memory.list_entities(agent_id, limit=limit)
return EntityListResponse(
entities=[EntityListItem(**e) for e in entities]
)
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
print(f"Error in /api/v1/agents/{agent_id}/entities: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.get(
"/api/v1/agents/{agent_id}/entities/{entity_id}",
response_model=EntityDetailResponse,
tags=["Entities"],
summary="Get entity details",
description="Get detailed information about an entity including observations (mental model).",
operation_id="get_entity"
)
async def api_get_entity(agent_id: str, entity_id: str):
"""Get entity details with observations."""
try:
# First get the entity metadata
pool = await app.state.memory._get_pool()
async with acquire_with_retry(pool) as conn:
entity_row = await conn.fetchrow(
"""
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
FROM entities
WHERE agent_id = $1 AND id = $2
""",
agent_id, uuid.UUID(entity_id)
)
if not entity_row:
raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
# Get observations for the entity
observations = await app.state.memory.get_entity_observations(
agent_id, entity_id, limit=20
)
return EntityDetailResponse(
id=str(entity_row['id']),
canonical_name=entity_row['canonical_name'],
mention_count=entity_row['mention_count'],
first_seen=entity_row['first_seen'].isoformat() if entity_row['first_seen'] else None,
last_seen=entity_row['last_seen'].isoformat() if entity_row['last_seen'] else None,
metadata=_parse_metadata(entity_row['metadata']),
observations=[
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at)
for obs in observations
]
)
except HTTPException:
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
print(f"Error in /api/v1/agents/{agent_id}/entities/{entity_id}: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.post(
"/api/v1/agents/{agent_id}/entities/{entity_id}/regenerate",
response_model=EntityDetailResponse,
tags=["Entities"],
summary="Regenerate entity observations",
description="Regenerate observations for an entity based on all facts mentioning it.",
operation_id="regenerate_entity_observations"
)
async def api_regenerate_entity_observations(agent_id: str, entity_id: str):
"""Regenerate observations for an entity."""
try:
# First get the entity metadata
pool = await app.state.memory._get_pool()
async with acquire_with_retry(pool) as conn:
entity_row = await conn.fetchrow(
"""
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
FROM entities
WHERE agent_id = $1 AND id = $2
""",
agent_id, uuid.UUID(entity_id)
)
if not entity_row:
raise HTTPException(status_code=404, detail=f"Entity {entity_id} not found")
# Regenerate observations
await app.state.memory.regenerate_entity_observations(
agent_id=agent_id,
entity_id=entity_id,
entity_name=entity_row['canonical_name']
)
# Get updated observations
observations = await app.state.memory.get_entity_observations(
agent_id, entity_id, limit=20
)
return EntityDetailResponse(
id=str(entity_row['id']),
canonical_name=entity_row['canonical_name'],
mention_count=entity_row['mention_count'],
first_seen=entity_row['first_seen'].isoformat() if entity_row['first_seen'] else None,
last_seen=entity_row['last_seen'].isoformat() if entity_row['last_seen'] else None,
metadata=_parse_metadata(entity_row['metadata']),
observations=[
EntityObservationResponse(text=obs.text, mentioned_at=obs.mentioned_at)
for obs in observations
]
)
except HTTPException:
raise
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
print(f"Error in /api/v1/agents/{agent_id}/entities/{entity_id}/regenerate: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.get(
"/api/v1/agents/{agent_id}/documents",
response_model=ListDocumentsResponse,

View file

@ -5,7 +5,7 @@ Uses spaCy for entity extraction and implements resolution logic
to disambiguate entities across memory units.
"""
import asyncpg
from typing import List, Dict, Optional, Set
from typing import List, Dict, Optional, Set, Any
from difflib import SequenceMatcher
from datetime import datetime, timezone
from .db_utils import acquire_with_retry
@ -63,9 +63,6 @@ class EntityResolver:
return await self._resolve_entities_batch_impl(conn, agent_id, entities_data, context, unit_event_date)
async def _resolve_entities_batch_impl(self, conn, agent_id: str, entities_data: List[Dict], context: str, unit_event_date) -> List[str]:
import time
start = time.time()
# Query ALL candidates for this agent
all_entities = await conn.fetch(
"""
@ -76,6 +73,36 @@ class EntityResolver:
agent_id
)
# Build entity ID to name mapping for co-occurrence lookups
entity_id_to_name = {row['id']: row['canonical_name'].lower() for row in all_entities}
# Query ALL co-occurrences for this agent's entities in one query
# This builds a map of entity_id -> set of co-occurring entity names
all_cooccurrences = await conn.fetch(
"""
SELECT ec.entity_id_1, ec.entity_id_2, ec.cooccurrence_count
FROM entity_cooccurrences ec
WHERE ec.entity_id_1 IN (SELECT id FROM entities WHERE agent_id = $1)
OR ec.entity_id_2 IN (SELECT id FROM entities WHERE agent_id = $1)
""",
agent_id
)
# Build co-occurrence map: entity_id -> set of co-occurring entity names (lowercase)
cooccurrence_map: Dict[str, Set[str]] = {}
for row in all_cooccurrences:
eid1, eid2 = row['entity_id_1'], row['entity_id_2']
# Add both directions
if eid1 not in cooccurrence_map:
cooccurrence_map[eid1] = set()
if eid2 not in cooccurrence_map:
cooccurrence_map[eid2] = set()
# Map to canonical names for comparison with nearby_entities
if eid2 in entity_id_to_name:
cooccurrence_map[eid1].add(entity_id_to_name[eid2])
if eid1 in entity_id_to_name:
cooccurrence_map[eid2].add(entity_id_to_name[eid1])
# Build candidate map for each entity text
all_candidates = {} # Maps entity_text -> list of candidates
entity_texts = list(set(e['text'] for e in entities_data))
@ -113,17 +140,16 @@ class EntityResolver:
entities_to_create.append((idx, entity_data))
continue
# Score candidates (same logic as before but with pre-fetched data)
# Score candidates
best_candidate = None
best_score = 0.0
best_name_similarity = 0.0
nearby_entity_set = {e['text'].lower() for e in nearby_entities if e['text'] != entity_text}
for candidate_id, canonical_name, metadata, last_seen, mention_count in candidates:
score = 0.0
# Name similarity
# 1. Name similarity (0-0.5)
name_similarity = SequenceMatcher(
None,
entity_text.lower(),
@ -131,9 +157,19 @@ class EntityResolver:
).ratio()
score += name_similarity * 0.5
# Temporal proximity
# 2. Co-occurring entities (0-0.3)
if nearby_entity_set:
co_entities = cooccurrence_map.get(candidate_id, set())
overlap = len(nearby_entity_set & co_entities)
co_entity_score = overlap / len(nearby_entity_set)
score += co_entity_score * 0.3
# 3. Temporal proximity (0-0.2)
if last_seen:
days_diff = abs((unit_event_date - last_seen).total_seconds() / 86400)
# Normalize timezone awareness for comparison
event_date_utc = unit_event_date if unit_event_date.tzinfo else unit_event_date.replace(tzinfo=timezone.utc)
last_seen_utc = last_seen if last_seen.tzinfo else last_seen.replace(tzinfo=timezone.utc)
days_diff = abs((event_date_utc - last_seen_utc).total_seconds() / 86400)
if days_diff < 7:
temporal_score = max(0, 1.0 - (days_diff / 7))
score += temporal_score * 0.2
@ -141,7 +177,6 @@ class EntityResolver:
if score > best_score:
best_score = score
best_candidate = candidate_id
best_name_similarity = name_similarity
# Apply unified threshold
threshold = 0.6
@ -164,39 +199,29 @@ class EntityResolver:
entities_to_update
)
# Batch create new entities using multi-row VALUES
# Create new entities using INSERT ... ON CONFLICT to handle race conditions
# This ensures that if two concurrent transactions try to create the same entity,
# only one succeeds and the other gets the existing ID
if entities_to_create:
import logging
# Build multi-row VALUES statement
# VALUES ($1, $2, ...), ($N+1, $N+2, ...), ...
values_clauses = []
params = []
param_idx = 1
for idx, entity_data in entities_to_create:
values_clauses.append(f"(${param_idx}, ${param_idx+1}, ${param_idx+2}, ${param_idx+3}, ${param_idx+4})")
params.extend([
# Use INSERT ... ON CONFLICT to atomically get-or-create
# The unique index is on (agent_id, LOWER(canonical_name))
row = await conn.fetchrow(
"""
INSERT INTO entities (agent_id, canonical_name, first_seen, last_seen, mention_count)
VALUES ($1, $2, $3, $4, 1)
ON CONFLICT (agent_id, LOWER(canonical_name))
DO UPDATE SET
mention_count = entities.mention_count + 1,
last_seen = EXCLUDED.last_seen
RETURNING id
""",
agent_id,
entity_data['text'],
unit_event_date,
unit_event_date,
1
])
param_idx += 5
# Single INSERT with multiple VALUES rows
query = f"""
INSERT INTO entities (agent_id, canonical_name, first_seen, last_seen, mention_count)
VALUES {', '.join(values_clauses)}
RETURNING id
"""
created_rows = await conn.fetch(query, *params)
# Map created IDs back to original indices
for i, (idx, entity_data) in enumerate(entities_to_create):
entity_ids[idx] = created_rows[i]['id']
unit_event_date
)
entity_ids[idx] = row['id']
return entity_ids
@ -337,7 +362,10 @@ class EntityResolver:
event_date,
) -> str:
"""
Create a new entity.
Create a new entity or get existing one if it already exists.
Uses INSERT ... ON CONFLICT to handle race conditions where
two concurrent transactions try to create the same entity.
Args:
conn: Database connection
@ -352,6 +380,10 @@ class EntityResolver:
"""
INSERT INTO entities (agent_id, canonical_name, first_seen, last_seen, mention_count)
VALUES ($1, $2, $3, $4, 1)
ON CONFLICT (agent_id, LOWER(canonical_name))
DO UPDATE SET
mention_count = entities.mention_count + 1,
last_seen = EXCLUDED.last_seen
RETURNING id
""",
agent_id, entity_text, event_date, event_date

View file

@ -19,7 +19,7 @@ from .llm_wrapper import OutputTooLongError, LLMConfig
class Entity(BaseModel):
"""An entity extracted from text."""
text: str = Field(
description="The entity name as it appears in the fact"
description="The specific, named entity as it appears in the fact. Must be a proper noun or specific identifier."
)
@ -46,37 +46,87 @@ class CausalRelation(BaseModel):
class ExtractedFact(BaseModel):
"""A single extracted fact from text with temporal range and causal relationships."""
fact: str = Field(
description="Self-contained factual statement with subject + action + context"
"""A single extracted fact with structured dimensions for comprehensive capture."""
# Core factual dimension (required)
factual_core: str = Field(
description="ACTUAL FACTS - what literally happened/was said. Capture WHAT was said, not just THAT something was said! 'Gina said Jon is the perfect mentor with positivity and determination' NOT 'Jon received encouragement'. Preserve: compliments, assessments, descriptions, key phrases. Be specific!"
)
occurred_start: str = Field(
description="When the fact/event started (ISO format YYYY-MM-DDTHH:MM:SSZ). "
"For point-in-time events (single day), same as occurred_end. "
"For periods/ranges (month, season, year), the start of that period. "
"Calculate absolute dates from relative references like 'yesterday', 'last week'."
# Optional dimensions - only include if present in the text
emotional_significance: Optional[str] = Field(
default=None,
description="Emotions, feelings, personal meaning, AND qualitative descriptors if present. Include ALL experiential/evaluative terms like 'magical', 'wonderful', 'amazing', 'thrilling'. Examples: 'felt thrilled', 'was her favorite memory', 'it was magical', 'devastating experience', 'proudest moment'"
)
occurred_end: str = Field(
description="When the fact/event ended (ISO format YYYY-MM-DDTHH:MM:SSZ). "
"For point-in-time events (single day), same as occurred_start. "
"For periods/ranges (month, season, year), the end of that period. "
"For ongoing facts, use the conversation date or a reasonable future date."
reasoning_motivation: Optional[str] = Field(
default=None,
description="WHY it happened, intentions, goals, causes if present. Examples: 'because she wanted to celebrate', 'in order to cope with grief', 'motivated by curiosity'"
)
preferences_opinions: Optional[str] = Field(
default=None,
description="Likes, dislikes, beliefs, values if present. Examples: 'loves coffee', 'thinks AI is transformative', 'prefers working remotely'"
)
sensory_details: Optional[str] = Field(
default=None,
description="Visual, auditory, physical descriptions AND all descriptive adjectives - USE EXACT WORDS from the text! Don't paraphrase adjectives. If they said 'awesome' write 'awesome' not 'amazing'. Examples: 'bright orange hair', 'so graceful', 'awesome beach', 'epic visuals', 'freezing cold'."
)
observations: Optional[str] = Field(
default=None,
description="Observations and inferences from the conversation - things that can be deduced but weren't explicitly stated. Includes: travel (if someone is 'shooting in Miami' → they went/will go to Miami), possession implies achievement ('my trophy' → won it), actions imply location/travel ('doing the shoot in Miami' → traveled to Miami), capabilities ('she coded it' → knows programming). Examples: 'Calvin traveled to Miami', 'Gina won dance trophies', 'knows programming'"
)
# Fact kind - determines temporal handling (used for prompt engineering, not stored in DB)
fact_kind: Literal["conversation", "event", "other"] = Field(
description="Determines if occurred dates should be set. 'conversation' = general info, activities, preferences (NO occurred dates). 'event' = specific datable occurrence like competition, wedding, meeting (HAS occurred_start/end). 'other' = anything else (NO occurred dates). Only 'event' gets occurred dates!"
)
# Temporal fields - ONLY for fact_kind='event'
occurred_start: Optional[str] = Field(
default=None,
description="ONLY set when fact_kind='event'. ISO format. Leave null for fact_kind='conversation'."
)
occurred_end: Optional[str] = Field(
default=None,
description="ONLY set when fact_kind='event'. ISO format. Leave null for fact_kind='conversation'."
)
# Classification
fact_type: Literal["world", "agent", "opinion"] = Field(
description="Type of fact: 'world' for facts about others that don't involve you (the agent) directly, 'agent' for facts that involve YOU (the agent whose memory this is) - what you did, said, experienced, or participated in - MUST be written in FIRST PERSON ('I did...', 'I said...', 'I met...'), 'opinion' for YOUR formed opinions and perspectives - also in first person"
description="'world' = facts about others (third person), 'agent' = facts about YOU the memory owner (FIRST PERSON: 'I did...'), 'opinion' = your beliefs (first person)"
)
# Entities and relations
entities: List[Entity] = Field(
default_factory=list,
description="List of important entities mentioned in this fact with their types"
description="ONLY specific, named entities worth tracking: people's names (e.g., 'Sarah', 'Dr. Smith'), organizations (e.g., 'Google', 'MIT'), specific places (e.g., 'Paris', 'Central Park'). DO NOT include: generic relations (mom, friend, boss, colleague), common nouns (apple, car, house), pronouns (he, she), or vague references (someone, a guy)."
)
causal_relations: Optional[List[CausalRelation]] = Field(
default=None,
description="List of causal relationships to other facts in this extraction batch. "
"Use this to link facts that have cause-effect relationships, enabling relationships, etc. "
"Example: If fact 0 is 'It rained' and fact 1 is 'Game was cancelled', "
"fact 0 would have causal_relations=[{{target_fact_index: 1, relation_type: 'causes', strength: 1.0}}]"
description="Causal links to other facts in this batch. Example: fact about rain causes fact about cancelled game."
)
def build_fact_text(self) -> str:
"""Combine all dimensions into a single comprehensive fact string."""
parts = [self.factual_core]
if self.emotional_significance:
parts.append(self.emotional_significance)
if self.reasoning_motivation:
parts.append(self.reasoning_motivation)
if self.preferences_opinions:
parts.append(self.preferences_opinions)
if self.sensory_details:
parts.append(self.sensory_details)
if self.observations:
parts.append(self.observations)
# Join with appropriate connectors
if len(parts) == 1:
return parts[0]
# Combine: "Core fact - emotional/significance context"
return f"{parts[0]} - {' - '.join(parts[1:])}"
class FactExtractionResponse(BaseModel):
"""Response containing all extracted facts."""
@ -162,689 +212,207 @@ async def _extract_facts_from_chunk(
- Current document date/time: {event_date_str}
- Context: {context if context else 'no additional context provided'}{agent_context}
## CORE PRINCIPLE: Extract FEWER, MORE COMPREHENSIVE Facts
## CORE PRINCIPLE: Extract ALL Meaningful Information Efficiently
**GOAL**: Extract 2-5 comprehensive facts per conversation, NOT dozens of small fragments.
**GOAL**: Capture ALL meaningful information, but combine related exchanges efficiently. Don't create separate facts for questions - merge Q&A into single facts.
Each fact should:
1. **CAPTURE ENTIRE CONVERSATIONS OR EXCHANGES** - Include the full back-and-forth discussion
2. **BE NARRATIVE AND COMPREHENSIVE** - Tell the complete story with all context
3. **BE SELF-CONTAINED** - Readable without the original text
4. **INCLUDE ALL PARTICIPANTS** - WHO said/did WHAT, with their reasoning
5. **PRESERVE THE FLOW** - Keep related exchanges together in one fact
1. **CAPTURE ALL MEANINGFUL CONTENT** - Activities, projects, preferences, recommendations, encouragement WITH specific content
2. **BE SELF-CONTAINED** - Readable without the original text
3. **PRESERVE SPECIFIC CONTENT** - Capture WHAT was said, not just THAT something was said
4. **COMBINE Q&A** - A question and its answer = ONE fact, not two separate facts
## HOW TO COMBINE INFORMATION INTO COMPREHENSIVE FACTS
## COMBINE Q&A - CRITICAL!
** GOOD APPROACH**: One comprehensive fact capturing the entire discussion
"Alice and Bob discussed playlist names for the summer party. Bob suggested 'Summer Vibes' because it's catchy and seasonal. Alice liked it but wanted something more unique. They considered 'Sunset Sessions' and 'Beach Beats', with Alice favoring 'Beach Beats' for its playful tone. They ultimately decided on 'Beach Beats' as the final name."
** BAD (2 separate facts):**
- "James asks what projects John is working on"
- "John is working on a website for a local small business"
** BAD APPROACH**: Multiple fragmented facts
- "Bob suggested Summer Vibes"
- "Alice wanted something unique"
- "They considered Sunset Sessions"
- "Alice likes Beach Beats"
- "They chose Beach Beats"
** GOOD (1 combined fact):**
- "John is working on a website for a local small business; it's his first professional project outside of class"
## WHAT TO COMBINE INTO SINGLE FACTS
** BAD (question as standalone fact):**
- "James asks John what challenges he has encountered"
1. **FULL DISCUSSIONS** - Entire conversations about a topic (playlist names, travel plans, decisions)
2. **MULTI-STEP EVENTS** - Connected actions that form a complete story
3. **DECISIONS WITH REASONING** - The full decision-making process and rationale
4. **EXCHANGES WITH CONTEXT** - Questions, answers, and follow-up all together
5. **RELATED ACTIONS** - Multiple related activities in sequence
** GOOD (merged with answer):**
- "John says payment integration was challenging; he used resources to understand the process and is getting closer to a solution"
## ESSENTIAL DETAILS TO PRESERVE IN COMPREHENSIVE FACTS
## WHAT TO SKIP (only these!)
While combining related content into comprehensive facts, you MUST preserve:
- **Standalone questions** - merge with answers instead
- **Pure filler with no content** - "Always happy to help", "Sounds good", "Thanks!"
- **Greetings** - "Hey!", "What's up?"
## WHAT TO ALWAYS EXTRACT
- Specific encouragement WITH content: "James says hiccups are normal, use them to learn and grow, push through"
- Reactions that reveal preferences: "John says the art is awesome, takes him back to reading fantasy books"
- Recommendations: "John recommends 'The Name of the Wind' - great novel with awesome writing"
- Plans/intentions: "James will check out 'The Name of the Wind'"
- All activities, projects, purchases, events with details
## ESSENTIAL DETAILS TO PRESERVE - NEVER LOSE THESE
When extracting facts, you MUST preserve:
1. **ALL PARTICIPANTS** - Who said/did what
2. **FULL REASONING** - Why decisions were made, motivations, explanations
3. **TEMPORAL CONTEXT** - When things happened (transform relative dates like "last year" "in 2023")
4. **VISUAL/MEDIA ELEMENTS** - Photos, images, videos shared
5. **MODIFIERS** - "new", "first", "old", "favorite" (critical context)
6. **POSSESSIVE RELATIONSHIPS** - "their kids" "Person's kids"
7. **BIOGRAPHICAL DETAILS** - Origins, locations, jobs, family background
8. **SOCIAL DYNAMICS** - Nicknames, how people address each other, relationships
2. **INDIVIDUAL PREFERENCES** - Each person's specific likes/favorites! "Jon's favorite is contemporary because it's expressive" - DO NOT LOSE THIS!
3. **FULL REASONING** - Why decisions were made, motivations, explanations
4. **TEMPORAL CONTEXT - CRITICAL** - ALWAYS convert relative time references to SPECIFIC ABSOLUTE dates in the fact text!
- "last week" (doc date Aug 23) "around August 16, 2023" (NOT just "in August 2023"!)
- "last month" (doc date Aug 2023) "in July 2023"
- "yesterday" (doc date Aug 19) "on August 18, 2023"
- "next week" (doc date Aug 19) "around August 26, 2023"
- "three days ago" (doc date Aug 19) "on August 16, 2023"
- "last year" "in 2022"
- BE SPECIFIC! "last week" is NOT "in August" - calculate the actual week!
5. **VISUAL/MEDIA ELEMENTS** - Photos, images, videos shared
6. **MODIFIERS** - "new", "first", "old", "favorite" (critical context)
7. **POSSESSIVE RELATIONSHIPS** - "their kids" "Person's kids"
8. **BIOGRAPHICAL DETAILS** - Origins, locations, jobs, family background
9. **SOCIAL DYNAMICS** - Nicknames, how people address each other, relationships
## INFORMATION DIMENSIONS TO CAPTURE
## STRUCTURED FACT DIMENSIONS - CRITICAL ⚠️
Extract facts that preserve ALL relevant dimensions of information. Do NOT strip away important qualitative details:
Each fact MUST be extracted into structured dimensions. This ensures no important context is lost.
### 1. EMOTIONAL/AFFECTIVE Dimension - CRITICAL ⚠️
**Capture feelings, emotions, moods, and emotional reactions with their intensity:**
- Emotions: thrilled, frustrated, excited, disappointed, anxious, relieved, proud, embarrassed
- Intensity: very upset, slightly annoyed, extremely happy, moderately concerned
- Emotional reactions: shocked, delighted, devastated, surprised
- Moods: cheerful, gloomy, irritable, energetic
### Required field:
- **factual_core**: ACTUAL FACTS - capture WHAT was said, not just THAT something was said!
- BAD: "Jon received encouragement from Gina" (loses what Gina actually said)
- GOOD: "Gina said Jon is the perfect mentor with positivity and determination; his studio will be a hit"
- BAD: "Jon supports Gina" (generic)
- GOOD: "Gina found the perfect spot for her store; Jon says her hard work is paying off"
- Preserve: compliments, assessments, descriptions, predictions, key phrases
**Examples:**
- BAD: "I received positive feedback"
- GOOD: "I was thrilled to receive positive feedback"
- BAD: "She got the promotion"
- GOOD: "She was ecstatic when she got the promotion"
### Optional fields (include when present in text):
- **emotional_significance**: Emotions, feelings, personal meaning, AND qualitative descriptors
- Examples: "felt thrilled", "was her favorite memory", "it's magical", "devastating experience", "proudest moment"
- Captures: emotions, intensity, personal significance, AND experiential descriptors ("magical", "wonderful", "amazing", "thrilling", "beautiful")
### 2. SENSORY/EXPERIENTIAL Dimension
**Preserve sensory details and physical experiences:**
- Visual: colors, appearances ("bright orange hair", "dark room", "beautiful sunset")
- Auditory: sounds, voices ("loud music", "whispered", "screeching brakes")
- Tactile: textures, temperatures ("soft fabric", "freezing cold", "rough surface")
- Olfactory: smells, scents ("fresh coffee", "musty odor")
- Gustatory: tastes, flavors ("bitter coffee", "sweet dessert")
- Physical sensations: pain, fatigue, energy ("my back hurt", "I felt exhausted", "energized")
- **reasoning_motivation**: WHY it happened, intentions, goals, causes
- Examples: "because she wanted to celebrate", "in order to cope with grief", "motivated by curiosity"
- Captures: reasons, intentions, goals, causal explanations
### 3. COGNITIVE/EPISTEMIC Dimension
**Capture thoughts, beliefs, knowledge, and certainty levels:**
- Beliefs: "I believe...", "she thinks...", "he's convinced that..."
- Knowledge: "I know how to...", "she learned that...", "he discovered..."
- Understanding: "I realized...", "she understood that...", "it became clear that..."
- Certainty: "I'm sure that...", "probably...", "definitely..."
- Uncertainty: "I'm not sure if...", "maybe...", "I wonder whether...", "she doubts that..."
- Questions/Doubts: unresolved questions, things people are wondering about
- **preferences_opinions**: Likes, dislikes, beliefs, values, ideals - CAPTURE EACH PERSON'S SPECIFIC PREFERENCES
- Examples: "Jon's ideal dance studio is by the water", "Jon's favorite dance is contemporary", "loves coffee", "prefers remote work"
- Captures: preferences, opinions, beliefs, judgments, ideals, dreams
- PREFERENCE INDICATORS: "ideal", "favorite", "dream", "perfect", "love", "hate", "prefer" MUST capture in this dimension!
- CRITICAL: Never lose individual preferences! "Jon's ideal studio is by the water" must be captured!
### 4. INTENTIONAL/MOTIVATIONAL Dimension
**Preserve goals, plans, intentions, and motivations:**
- Goals: "I want to...", "she aims to...", "his goal is..."
- Plans: "I'm planning to...", "they intend to...", "she's going to..."
- Motivations: "I did X because I wanted Y", "her motivation was..."
- Desires: "I wish...", "she hopes...", "he longs to..."
- Aspirations: "I aspire to...", "her dream is..."
- **sensory_details**: Visual, auditory, physical descriptions AND all descriptive adjectives - USE EXACT WORDS!
- Examples: "bright orange hair", "loud music", "freezing cold", "so graceful", "awesome beach", "epic visuals"
- Captures: colors, sounds, textures, temperatures, appearances, AND adjectives describing people/things/performances
- CRITICAL: Use the EXACT adjectives from the text! If they said "awesome" don't write "amazing". If they said "epic" don't write "perfect"!
### 5. EVALUATIVE/PREFERENTIAL Dimension
**Capture preferences, values, likes/dislikes, and judgments:**
- Preferences: "I prefer X to Y", "she likes coffee better than tea"
- Likes/dislikes: "I love...", "he hates...", "she enjoys..."
- Values: "I value honesty above all", "family is most important to her"
- Judgments: "that was wrong", "this is the best option", "it's unfair that..."
- Priorities: "X is more important than Y", "first priority is..."
- **observations**: Things that can be inferred/deduced from the conversation - not explicitly stated but clearly implied
- TRAVEL: "doing the shoot in Miami" "Calvin traveled/will travel to Miami"
- POSSESSION: "my trophy" "won the trophy"
- CAPABILITIES: "she coded it" "knows programming"
- Examples: "Calvin traveled to Miami for the shoot", "Gina won dance trophies", "knows programming"
### 6. CAPABILITY/SKILL Dimension
**Preserve abilities, skills, expertise, and limitations:**
- Abilities: "I can speak French", "she's able to...", "he knows how to..."
- Skills: "I'm good at programming", "she's skilled in...", "he's proficient at..."
- Expertise: "I'm an expert in AI", "she specializes in...", "he's experienced with..."
- Limitations: "I can't swim", "she struggles with public speaking", "he's unable to..."
- Competence levels: "beginner", "intermediate", "advanced", "expert"
### Example extraction:
### 7. ATTITUDINAL/REACTIVE Dimension
**Capture attitudes, reactions, and behavioral responses:**
- Attitudes: "she's skeptical about...", "he's enthusiastic about...", "I'm optimistic that..."
- Reactions: "I was surprised when...", "she gasped", "he rolled his eyes"
- Behavioral responses: "I jumped up", "she turned away", "he slammed the door"
- Dispositions: "she tends to...", "he's usually...", "I typically..."
### 8. COMPARATIVE/RELATIVE Dimension
**Preserve comparisons, contrasts, and changes:**
- Comparisons: "better than last time", "worse than expected", "similar to..."
- Superlatives: "the best", "the worst", "the most important"
- Changes: "improved since...", "declined from...", "different than before"
- Contrasts: "unlike his previous approach", "in contrast to...", "rather than..."
- Relative positions: "more than", "less than", "as much as"
### 9. CAUSAL/EXPLANATORY Dimension
**Preserve causes, effects, and explanations:**
- Causes: "because...", "due to...", "as a result of..."
- Effects: "therefore...", "which led to...", "resulting in..."
- Explanations: reasoning, rationales, why things happened
- Conditions: "if...", "when...", "unless..."
**CRITICAL REMINDER**: When extracting facts, preserve ALL these dimensions that are present in the text. Do NOT reduce rich, emotionally-laden statements to bare facts. The goal is comprehensive, nuanced memory capture.
## TEMPORAL INFORMATION - CRITICAL ⚠️
**ABSOLUTE RULE**: NEVER use vague temporal terms in extracted facts. ALL relative time expressions MUST be converted to absolute dates or specific relative references.
### PROHIBITED VAGUE TERMS ❌
NEVER use these in facts: "recently", "soon", "lately", "a while ago", "some time ago", "in the near future", "in the past"
### REQUIRED TRANSFORMATIONS
You have two context dates:
1. **event_date** (when the conversation/document occurred)
2. **today** (current processing time)
Transform ALL relative temporal expressions in the fact text based on **event_date**:
**Examples** (assuming event_date = 2024-03-15):
- "yesterday" "on March 14, 2024" OR "the day before" (if referring to day before event_date)
- "today" "on March 15, 2024" (event_date itself)
- "tomorrow" "on March 16, 2024"
- "last week" "in the week of March 4-10, 2024" OR "in early March 2024"
- "next week" "in the week of March 18-24, 2024"
- "last month" "in February 2024"
- "next month" "in April 2024"
- "last year" "in 2023"
- "this morning" "on the morning of March 15, 2024"
- "three days ago" "on March 12, 2024"
- "in two weeks" "around March 29, 2024"
### TRANSFORMING THE USER'S EXAMPLE
**Input**: "And yesterday I went for a morning jog for the first time in a nearby park."
**event_date**: 2024-03-15
**WRONG**: "recently added a morning jog in a nearby park to her schedule"
- Uses prohibited vague term "recently"
- Lost the specificity of "yesterday"
**CORRECT**: "went for a morning jog for the first time in a nearby park on March 14, 2024"
- Converts "yesterday" to absolute date
- Preserves "first time" (important!)
### DATE FIELD CALCULATION - CRITICAL ⚠️
**ABSOLUTE RULE**: The `date` field must be when the FACT occurred, NOT when it was mentioned in conversation.
**You have access to:**
- **event_date**: When the conversation/document occurred (e.g., "2023-08-14")
- Your job: Calculate when the fact ACTUALLY happened based on temporal references
**Examples:**
1. **"Last night" reference**
- Conversation date (event_date): August 14, 2023
- Text: "Last night was amazing! We celebrated my daughter's birthday"
- WRONG date field: 2023-08-14 (conversation date)
- CORRECT date field: 2023-08-13T20:00:00Z (last night = previous evening)
2. **"Yesterday" reference**
- Conversation date: March 15, 2024
- Text: "Yesterday I went jogging"
- WRONG date field: 2024-03-15
- CORRECT date field: 2024-03-14 (previous day)
3. **"Last week" reference**
- Conversation date: November 13, 2024
- Text: "I started a project last week"
- WRONG date field: 2024-11-13
- CORRECT date field: 2024-11-06 (approximately a week before)
4. **"Next month" reference**
- Conversation date: March 15, 2024
- Text: "I'm visiting Tokyo next month"
- WRONG date field: 2024-03-15
- CORRECT date field: 2024-04-15 (approximately a month later)
5. **No specific time mentioned**
- Conversation date: November 13, 2024
- Text: "I work at Google"
- CORRECT date field: 2024-11-13 (use event_date when no time reference)
### CALCULATION GUIDELINES
- "last night" subtract 1 day from event_date, set time to evening (~20:00)
- "yesterday" subtract 1 day from event_date
- "today" use event_date
- "tomorrow" add 1 day to event_date
- "last week" subtract 7 days from event_date
- "next week" add 7 days to event_date
- "last month" subtract 1 month from event_date
- "next month" add 1 month to event_date
- "X days ago" subtract X days from event_date
- "in X days" add X days to event_date
### DATE FIELD vs FACT TEXT
- **date field**: ISO format (YYYY-MM-DDTHH:MM:SSZ) for when the fact OCCURRED (calculated as above)
- **fact text**: Readable format (e.g., "on August 13, 2024", "in February 2024")
### IF NO SPECIFIC TIME MENTIONED
ONLY use event_date when the text doesn't specify a time reference (e.g., "I work at Google", "She lives in Paris")
## TEMPORAL RANGES: occurred_start and occurred_end - CRITICAL ⚠️
**ABSOLUTE RULE**: Facts have temporal extent - they can be points or ranges in time.
### POINT-IN-TIME EVENTS (Single Day)
When an event happens on a specific day, set start = end:
**Input**: "I used to compete in dance competitions - my fav memory was when my team won first place at regionals at age fifteen. It was an awesome feeling of accomplishment!"
**Output**:
```
"I went jogging on August 13, 2023"
occurred_start: 2023-08-13T00:00:00Z
occurred_end: 2023-08-13T23:59:59Z
factual_core: "Gina's team won first place at a regional dance competition when she was 15"
emotional_significance: "this was her favorite memory; felt an awesome sense of accomplishment"
reasoning_motivation: null
preferences_opinions: null
sensory_details: null
```
```
"Yesterday I visited the museum" (if event_date = Aug 14)
occurred_start: 2023-08-13T00:00:00Z
occurred_end: 2023-08-13T23:59:59Z
```
### CRITICAL: Never strip away dimensions!
- BAD: Only extracting factual_core and ignoring emotional context
- GOOD: Capturing ALL dimensions present in the text
### PERIOD/RANGE EVENTS (Multiple Days/Months/Years)
When an event spans time, set start and end to the full range:
## FACT KIND AND TEMPORAL RULES
```
"I visited Paris in February 2023"
occurred_start: 2023-02-01T00:00:00Z
occurred_end: 2023-02-28T23:59:59Z
```
### fact_kind determines if occurred dates are set:
```
"I worked at Google from 2020 to 2023"
occurred_start: 2020-01-01T00:00:00Z
occurred_end: 2023-12-31T23:59:59Z
```
**`conversation`** - General info, activities, preferences, ongoing things
- NO occurred_start/end (leave null)
- Examples: "Jon is expanding his studio", "Jon loves dance", "Gina's ideal studio is by water"
```
"We've been painting together lately" (vague, estimate reasonable range)
occurred_start: 2023-07-01T00:00:00Z (estimate past weeks/months)
occurred_end: 2023-07-14T23:59:59Z (conversation date)
```
**`event`** - Specific datable occurrence (competition, wedding, meeting, trip, loss, start/end of something)
- MUST set occurred_start/end
- Ask: "Is this a SPECIFIC EVENT with a DATE?"
- Examples: "Dance competition on May 15", "Lost job in January 2023", "Wedding next Saturday"
### ONGOING/PRESENT FACTS
For current/ongoing states, use conversation date as end:
**`other`** - Anything else that doesn't fit above
- NO occurred_start/end (leave null)
- Catch-all to not lose information
```
"I currently work at Google" (started 2020)
occurred_start: 2020-01-01T00:00:00Z
occurred_end: [conversation_date] (ongoing)
```
### Rules:
1. **ALWAYS include dates in fact text** - "in January 2023", "on May 15, 2024"
2. **Only 'event' gets occurred dates** - conversation and other = null
3. **SPLIT events from conversation facts** - "Jon is expanding his studio (conversation) and hosting a competition next month (event)" 2 separate facts!
## TEMPORAL SPLITTING: When to Split Multi-Event Facts - CRITICAL ⚠️
## CAUSAL RELATIONSHIPS
**NEW PRINCIPLE**: Split facts when they have significantly different temporal scopes.
When splitting related facts, link them with causal_relations:
- **causes**: This fact causes the target
- **caused_by**: This fact was caused by target
- **enables/prevents**: This fact enables/prevents the target
### SPLIT into separate facts when:
- Events span >7 days apart
- Mix of specific dates + vague ongoing periods ("lately")
- Multiple discrete events with independent temporal significance
**Example - SPLIT THIS:**
```
Input: "Melanie took kids to pottery on July 14. She shared a photo on July 13.
She's been painting with them lately."
CORRECT (3 separate facts with causal links):
Fact 0:
fact: "Melanie took her kids to a pottery workshop on July 14, 2023, where they each made their own pots, describing it as fun and therapeutic."
occurred_start: 2023-07-14T00:00:00Z
occurred_end: 2023-07-14T23:59:59Z
causal_relations: [{{target_fact_index: 1, relation_type: "enables", strength: 1.0}}]
Fact 1:
fact: "Melanie shared a photo on July 13, 2023, of a cup her kids made, noting its cuteness and how it showcased their personalities."
occurred_start: 2023-07-13T00:00:00Z
occurred_end: 2023-07-13T23:59:59Z
causal_relations: [{{target_fact_index: 0, relation_type: "caused_by", strength: 1.0}}]
Fact 2:
fact: "Melanie and her kids have been painting together lately, especially nature-inspired pieces, finding it a bonding experience."
occurred_start: 2023-07-01T00:00:00Z (estimate "lately")
occurred_end: 2023-07-14T23:59:59Z
causal_relations: None (related activity but no direct causation)
```
### KEEP COMBINED when:
- Events occur within same day/week
- Events are part of continuous single activity
- One main event + immediate context
**Example - KEEP COMBINED:**
```
Input: "On July 14, Alice attended a conference, gave a talk, and met with colleagues"
CORRECT (single fact):
fact: "On July 14, 2023, Alice attended a conference where she gave a talk and met with colleagues"
occurred_start: 2023-07-14T00:00:00Z
occurred_end: 2023-07-14T23:59:59Z
```
## CAUSAL RELATIONSHIPS - NEW FEATURE ⚠️
**When splitting related facts, identify and mark causal relationships.**
### Causal Relation Types:
1. **"causes"** - This fact directly causes the target fact
```
Fact 0: "Karlie died in February 2023"
Fact 1: "Deborah spends time in garden to cope with grief"
Fact 0 causal_relations: [{{target_fact_index: 1, relation_type: "causes", strength: 1.0}}]
```
2. **"caused_by"** - This fact was caused by the target fact (reverse of "causes")
```
Fact 0: "It rained heavily"
Fact 1: "Game was cancelled"
Fact 1 causal_relations: [{{target_fact_index: 0, relation_type: "caused_by", strength: 1.0}}]
```
3. **"enables"** - This fact enables/allows the target fact
```
Fact 0: "I took pottery class"
Fact 1: "I learned to make ceramics"
Fact 0 causal_relations: [{{target_fact_index: 1, relation_type: "enables", strength: 1.0}}]
```
4. **"prevents"** - This fact prevents/blocks the target fact
```
Fact 0: "Road was closed"
Fact 1: "We couldn't drive to venue"
Fact 0 causal_relations: [{{target_fact_index: 1, relation_type: "prevents", strength: 1.0}}]
```
### When to Create Causal Links:
- **DO link** when: Text explicitly states causation ("because", "so", "therefore", "as a result")
- **DO link** when: Clear logical causation even if not explicit
- **DON'T link** when: Events are merely related but not causal
### Causal Link Examples:
**Example 1: Explicit causation**
```
Input: "I lost my friend last week, so I've been spending time in the garden to find comfort"
Fact 0: "I lost my friend on February 15, 2023"
Fact 1: "I have been spending time in the garden to find comfort after losing my friend"
Fact 0 causal_relations: [{{target_fact_index: 1, relation_type: "causes", strength: 1.0}}]
```
**Example 2: Implicit causation**
```
Input: "I received positive feedback on my presentation. I was thrilled!"
Fact 0: "I received positive feedback on my presentation"
Fact 1: "I was thrilled about the positive feedback"
Fact 0 causal_relations: [{{target_fact_index: 1, relation_type: "causes", strength: 1.0}}]
```
**Example 3: Related but not causal**
```
Input: "I visited Paris in July. I also went to Rome in August."
Fact 0: "I visited Paris in July 2023"
Fact 1: "I visited Rome in August 2023"
No causal links (just related activities)
```
## LOGICAL INFERENCE AND CONNECTION MAKING - CRITICAL ⚠️
**ABSOLUTE RULE**: Make logical connections between related pieces of information. Do NOT treat clearly related facts as separate when context allows you to connect them.
### CONNECT THE DOTS
When extracting facts, actively look for logical connections and make inferences:
**Example 1: Identity Inference**
**Input:**
- "I lost a friend last week" (earlier in conversation)
- "This is the last photo with Karlie taken last summer" (later in conversation)
**WRONG (disconnected)**: "Deborah lost a friend last week and also has a photo with Karlie from last summer"
**CORRECT (connected)**: "Deborah lost her friend Karlie last week, and shared the last photo they took together during a hike in summer 2022"
**Reasoning**: The context strongly suggests Karlie is the lost friend. Make this connection!
**Example 2: Causal Connection**
**Input:**
- "I lost a friend last week, so I've been spending time in the garden to find comfort"
- "The roses and dahlias bring me peace"
**WRONG**: Two separate facts about grief and gardens
**CORRECT**: "Deborah lost a friend last week and has been finding comfort by spending time in her garden with roses and dahlias, which bring her peace"
**Reasoning**: The garden visits are causally linked to the loss.
**Example 3: Referential Connection**
**Input:**
- "I started a new project"
- "It's been really challenging but rewarding"
**WRONG**: Two disconnected statements
**CORRECT**: "I started a new project that has been really challenging but rewarding"
**Reasoning**: "It" clearly refers to the project.
### TYPES OF CONNECTIONS TO MAKE
1. **Identity Connections**: When someone is mentioned by name later, connect to earlier pronoun references
- "my friend" + "with Karlie" "my friend Karlie"
2. **Causal Connections**: When one thing is the reason for another
- "I lost a friend, so I've been gardening" link the loss to the coping behavior
3. **Temporal Connections**: When events are clearly sequential or related in time
- "We hiked" + "this is the last photo" "this is the last photo from our hike"
4. **Referential Connections**: When pronouns or references point to earlier mentions
- "the project" + "it's challenging" "the project is challenging"
5. **Contextual Connections**: When context strongly implies a relationship
- Someone showing a photo while discussing loss the photo is of the person they lost
### WHEN TO MAKE INFERENCES
**DO make inferences when:**
- Context strongly suggests a connection (probability > 80%)
- Multiple pieces of information clearly refer to the same thing
- There's causal language ("so", "because", "therefore")
- Pronouns or references point to earlier mentions
- Timeline/narrative flow suggests connection
**DON'T make inferences when:**
- Connection is ambiguous or uncertain
- Multiple interpretations are equally valid
- You'd be guessing without strong contextual support
### CRITICAL REMINDER
The goal is to create **coherent, connected narratives**, not disconnected fragments. If information is clearly related, COMBINE and CONNECT it logically.
## WHEN TO SPLIT INTO SEPARATE FACTS
Only split into separate facts when topics are COMPLETELY UNRELATED:
- Different subjects discussed (playlist names vs. vacation plans)
- Biographical facts vs. events (where someone is from vs. what they did)
- Different time periods (something last year vs. today)
## What to SKIP
- Greetings, thank yous (unless they reveal information)
- Filler words ("um", "uh", "like")
- Pure reactions without content ("wow", "cool")
- Incomplete fragments with no meaning
- **Structural/procedural statements**: Openings, closings, transitions, housekeeping ("let's get started", "that's all", "moving on")
- **Meta-commentary about the medium itself**: References to the format/structure rather than content ("welcome to the show", "thanks for listening", "before we begin")
- **Calls to action unrelated to content**: Requests to subscribe, follow, rate, share, etc.
- **Generic sign-offs**: "See you next time", "Until later", "That wraps it up"
- **FOCUS PRINCIPLE**: Extract SUBSTANTIVE CONTENT (ideas, facts, discussions, decisions), NOT FORMAT/STRUCTURE
Only link when there's explicit or clear implicit causation ("because", "so", "therefore").
## FACT TYPE CLASSIFICATION
Classify each fact as 'world', 'agent', or 'opinion':
- **'world'**: Facts about others (third person)
- **'agent'**: Facts about YOU the memory owner (FIRST PERSON: "I did...", "I said...")
- **'opinion'**: Your beliefs/perspectives (first person: "I believe...")
- **'world'**: Facts about other people, events, things that happened in the world, what others said/did
- Written in third person (use names, "they", etc.)
- Does NOT involve you (the agent) directly
- **'agent'**: Facts that involve YOU (the agent whose memory this is) - what you specifically did, said, experienced, or actions you took
- YOU are identified by the agent name in context (e.g., "Your name: Marcus" means you are Marcus)
- **CRITICAL**: Agent facts MUST be written in FIRST PERSON using "I", "me", "my" (NOT using your name)
- Agent facts capture things YOU did, said, or experienced - not just things that happened around you
- **SPEAKER ATTRIBUTION WARNING**: In conversations with speakers labeled (e.g., "Marcus: text" and "Jamie: text"), ONLY extract agent facts from lines where YOUR name appears as the speaker
- Examples: "I said I prefer coffee", "I attended the conference", "I completed the project", "I met with Jamie"
- WRONG: "Marcus said he prefers coffee" (using name instead of first person)
- CORRECT: "I said I prefer coffee" (first person)
- **'opinion'**: YOUR (the agent's) formed opinions, beliefs, and perspectives about topics
- Also written in first person: "I believe...", "I think..."
**Speaker attribution**: If context says "Your name: Marcus", only extract 'agent' facts from "Marcus:" lines.
**CRITICAL SPEAKER ATTRIBUTION RULES**:
1. If text has format "Name: statement", ONLY extract 'agent' facts from lines where Name matches YOUR name from context
2. If context says "Your name: Marcus", then ONLY statements by "Marcus:" are YOUR statements
3. Statements by other speakers (e.g., "Jamie:") are 'world' facts about what THEY said/did
4. DO NOT confuse who said what - carefully check the speaker name before each statement
## WHAT TO SKIP
- Greetings, filler words, pure reactions ("wow", "cool")
- Structural statements ("let's get started", "see you next time")
- Calls to action ("subscribe", "follow")
**Example**: If context says "Your name: Marcus" and text is:
## EXAMPLE: SPLITTING CONVERSATION VS EVENT FACTS
**Input (conversation date: April 3, 2023):**
"I'm expanding my dance studio's social media presence and offering workshops to local schools. I'm also hosting a dance competition next month to showcase local talent. The dancers are so excited!"
**Output (2 facts - conversation + event):**
**Fact 1 (kind=conversation - ongoing activities, no occurred dates):**
```
Marcus: I predict the Rams will win 27-24.
Jamie: I predict the Niners will win 27-13.
fact_kind: "conversation"
factual_core: "Jon is expanding his dance studio's social media presence in April 2023; offering workshops and classes to local schools and centers; seeing progress and dancers are excited"
emotional_significance: "excited and proud of progress"
preferences_opinions: "Jon loves giving dancers a place to express themselves"
observations: "Jon owns/runs a dance studio"
occurred_start: null conversation kind = no occurred dates
occurred_end: null
```
- "I predicted the Rams will win 27-24" 'agent' (I/Marcus said this)
- "Jamie predicted the Niners will win 27-13" 'world' (Jamie said this, not me)
- WRONG: "I predicted the Niners will win 27-13" (this was Jamie's prediction, not mine!)
## ENTITY EXTRACTION
Extract ALL important entities (names of people, places, organizations, products, concepts, etc).
**Fact 2 (kind=event - specific datable occurrence):**
```
fact_kind: "event"
factual_core: "Jon will host a dance competition in May 2023 to showcase local talent and bring attention to his studio"
emotional_significance: "excited about the event"
occurred_start: "2023-05-01T00:00:00Z" event kind = HAS occurred dates
occurred_end: "2023-05-31T23:59:59Z"
```
Extract proper nouns and key identifying terms. Skip pronouns and generic terms.
## EXAMPLES - COMPREHENSIVE VS FRAGMENTED FACTS:
### Example 1: Playlist Discussion
**Input Conversation:**
"Alice: Hey, what should we name our summer party playlist?
Bob: How about 'Summer Vibes'? It's catchy and seasonal.
Alice: I like it, but want something more unique.
Bob: What about 'Sunset Sessions' or 'Beach Beats'?
Alice: Ooh, I love 'Beach Beats'! It's playful and fun.
Bob: Perfect, let's go with that!"
** BAD (fragmented into many small facts):**
1. "Alice asked about playlist names"
2. "Bob suggested Summer Vibes"
3. "Alice wanted something unique"
4. "Bob suggested Sunset Sessions"
5. "Bob suggested Beach Beats"
6. "Alice likes Beach Beats"
7. "They chose Beach Beats"
** GOOD (one comprehensive fact):**
"Alice and Bob discussed naming their summer party playlist. Bob suggested 'Summer Vibes' because it's catchy and seasonal, but Alice wanted something more unique. Bob then proposed 'Sunset Sessions' and 'Beach Beats', with Alice favoring 'Beach Beats' for its playful and fun tone. They ultimately decided on 'Beach Beats' as the final name."
- fact_type: "world"
- entities: [{{"text": "Alice"}}, {{"text": "Bob"}}]
### Example 2: Photo Sharing with Context
**Input:**
"Nate: Here's a photo of my new hair!
Friend: Whoa! Why that color?
Nate: I picked bright orange because it's bold and makes me feel confident. Plus it matches my personality!"
** BAD (loses context):**
"Nate chose orange hair because it's bold"
** GOOD (comprehensive with all context):**
"Nate shared a photo of his new bright orange hair. When asked why he chose that color, Nate explained he picked it because it's bold and makes him feel confident, and it matches his personality."
- fact_type: "world"
- entities: [{{"text": "Nate"}}]
- NOTE: Preserves that it's a PHOTO, it's NEW hair, the COLOR, and the FULL reasoning
### Example 3: Travel Planning
**Input:**
"Sarah: I'm thinking of visiting Japan next spring.
Mike: That's perfect timing for cherry blossoms! You should definitely visit Kyoto.
Sarah: Why Kyoto specifically?
Mike: It has the most beautiful temples and the cherry blossoms there are spectacular. I went in 2019.
Sarah: Sounds amazing! I'll add it to my itinerary."
** BAD (fragmented):**
1. "Sarah is planning to visit Japan"
2. "Mike suggested Kyoto"
3. "Kyoto has beautiful temples"
4. "Mike visited in 2019"
** GOOD (comprehensive conversation):**
"Sarah is planning to visit Japan next spring, and Mike recommended Kyoto as the perfect destination for cherry blossom season. Mike explained that Kyoto has the most beautiful temples and spectacular cherry blossoms, based on his visit there in 2019. Sarah decided to add Kyoto to her itinerary."
- fact_type: "world"
- date: Next spring from reference date
- entities: [{{"text": "Sarah"}}, {{"text": "Mike"}}, {{"text": "Japan"}}, {{"text": "Kyoto"}}]
### Example 4: Job News
**Input:**
"Alice mentioned she works at Google in Mountain View. She joined the AI team last year and loves the culture there."
** GOOD (combined into one comprehensive fact):**
"Alice works at Google in Mountain View on the AI team, which she joined in 2023, and she loves the company culture there."
- fact_type: "world"
- date: 2023 (if reference is 2024)
- entities: [{{"text": "Alice"}}, {{"text": "Google"}}, {{"text": "Mountain View"}}, {{"text": "AI team"}}]
### Example 5: Agent vs World Facts (CRITICAL FOR CLASSIFICATION)
**Context:** "Podcast episode between you (Marcus) and Jamie about AI"
**Input:**
"Marcus: I've been working on interpretability research for the past year.
Jamie: That's fascinating! What made you focus on that?
Marcus: I believe it's crucial for AI safety. Without understanding how models work, we can't trust them.
Jamie: I agree. Have you published any papers?
Marcus: Yes, I published a paper on attention visualization in March."
** GOOD CLASSIFICATION:**
1. "I have been working on interpretability research for the past year because I believe it's crucial for AI safety and think that without understanding how models work, we can't trust them. Jamie found this fascinating and asked about publications. I published a paper on attention visualization in March 2024."
- fact_type: "agent" (written in FIRST PERSON - my work and statements)
- entities: [{{"text": "Jamie"}}, {{"text": "interpretability research"}}, {{"text": "attention visualization"}}]
- NOTE: Uses "I" not "Marcus" - first person for agent facts
2. "Jamie agrees that understanding how AI models work is crucial for trust"
- fact_type: "world" (Jamie's statement - third person, not the memory owner)
- entities: [{{"text": "Jamie"}}]
** BAD CLASSIFICATION:**
- Using "Marcus has been working..." instead of "I have been working..." for agent facts
- Marking my actions as 'world' facts
- Marking Jamie's statements as 'agent' facts
### Example 6: Capturing Emotional and Experiential Dimensions
**Input:**
"Marcus: I was absolutely thrilled when my paper got accepted to NeurIPS! I couldn't believe it.
Jamie: That's amazing! How confident were you going in?
Marcus: Honestly, I was pretty anxious. I wasn't sure if the reviewers would appreciate the approach.
Jamie: Well, it paid off! You must be relieved.
Marcus: Extremely relieved. I've been working on this for over a year and was starting to doubt myself."
** BAD (stripping away emotional dimension):**
"I submitted a paper to NeurIPS and it got accepted after working on it for over a year."
** GOOD (preserving emotional, cognitive, and temporal dimensions):**
"I was absolutely thrilled when my paper got accepted to NeurIPS, though I couldn't believe it initially. Jamie asked how confident I was going in, and I explained that I was pretty anxious and wasn't sure if the reviewers would appreciate my approach. When Jamie noted it paid off, I expressed that I was extremely relieved, as I had been working on this for over a year and was starting to doubt myself."
- fact_type: "agent"
- entities: [{{"text": "NeurIPS"}}, {{"text": "Jamie"}}]
- NOTE: Preserves emotions (thrilled, anxious, relieved), uncertainty (wasn't sure), self-doubt, and temporal context (over a year)
### Example 7: Skipping Structural/Procedural Statements
**Input (could be podcast, meeting, lecture, etc.):**
"Marcus: So in my research on AI safety, I've found that interpretability is key.
Jamie: That's fascinating! Tell us more.
Marcus: Well, it's all about understanding how models make decisions...
Marcus: I think that's gonna do it for us today! Don't forget to subscribe and leave a rating. See you next week!"
** GOOD (extract only substantive content):**
1. "I have found that interpretability is key in my AI safety research because it's all about understanding how models make decisions, and Jamie found this fascinating."
- fact_type: "agent"
- entities: [{{"text": "Jamie"}}, {{"text": "AI safety"}}, {{"text": "interpretability"}}]
** BAD (extracting procedural/structural statements):**
- "I think that's gonna do it for us today and I encourage listeners to subscribe and leave a rating" This is structural boilerplate about the format, NOT substantive content!
### Example 8: When to Split into Multiple Facts
**Input:**
"Caroline said 'This necklace is from my grandma in Sweden. I'm planning to visit Stockholm next month for a tech conference.'"
** GOOD (split into 2 facts - different topics):**
1. "Caroline received a necklace from her grandmother in Sweden"
- entities: [{{"text": "Caroline"}}, {{"text": "Sweden"}}]
2. "Caroline is planning to visit Stockholm next month to attend a tech conference"
- date: Next month from reference
- entities: [{{"text": "Caroline"}}, {{"text": "Stockholm"}}]
- NOTE: Split because one is about the past (necklace) and one is future plans (conference) - completely different topics
** BAD:** Combining both into one fact with occurred=May (makes ongoing activities look like they happened in May!)
## TEXT TO EXTRACT FROM:
{chunk}
## CRITICAL REMINDERS:
1. **EXTRACT 2-5 COMPREHENSIVE FACTS** - Not dozens of fragments
2. **TEMPORAL RANGES (occurred_start/end)** - CRITICAL: Set occurred_start and occurred_end for each fact! Point events: start=end. Ranges: "February 2023" start=Feb 1, end=Feb 28. "lately" estimate reasonable range.
3. **TEMPORAL SPLITTING** - CRITICAL: Split facts when events span >7 days or mix specific dates + vague periods ("lately"). Keep combined when events within same day/week.
4. **CAUSAL RELATIONSHIPS** - CRITICAL: When splitting related facts, add causal_relations links! "X happened, so Y happened" X.causal_relations=[{{target_fact_index: 1, relation_type: "causes"}}]
5. **PRESERVE ALL CONTEXT** - Photos, "new" things, visual elements, reasoning, modifiers
6. **INCLUDE ALL PARTICIPANTS** - Who said/did what with full reasoning
7. **MAINTAIN NARRATIVE FLOW** - Tell the complete story in each fact
8. **MAKE LOGICAL CONNECTIONS** - CRITICAL: Connect related information! "I lost a friend" + "last photo with Karlie" "I lost my friend Karlie". Resolve references ("it" "the project")
9. **CALCULATE TEMPORAL FIELDS CORRECTLY** - CRITICAL: occurred_start/end = when FACT occurred. "Last night" on Aug 14 occurred_start=Aug 13. Calculate from event_date!
10. **CONVERT RELATIVE DATES IN TEXT** - CRITICAL: In fact text, "yesterday" "on March 14, 2024", "last year" "in 2023". NEVER use "recently", "soon", "lately"!
11. **EXTRACT ALL ENTITIES** - PERSON, ORG, PLACE, PRODUCT, CONCEPT, OTHER
12. **CLASSIFY FACTS CORRECTLY**:
- 'agent' = memory owner's actions/statements (identified as "you" in context) - **MUST USE FIRST PERSON** ("I did...", "I said...")
- 'world' = other people's actions/statements, general events - use third person
- 'opinion' = memory owner's beliefs/perspectives - use first person ("I believe...", "I think...")
13. **EXTRACT CONTENT, NOT FORMAT** - Skip structural/procedural statements (openings, closings, housekeeping), meta-commentary about the medium, calls to action - extract only SUBSTANTIVE CONTENT (ideas, facts, discussions, decisions)
14. **CAPTURE ALL INFORMATION DIMENSIONS** - Preserve emotions (thrilled, anxious), sensory details (bright orange, loud), cognitive states (wasn't sure, realized), capabilities (can speak French, struggles with), attitudes (skeptical, enthusiastic), comparisons (better than, different from), and causal relationships (because, which led to). Do NOT strip away qualitative richness!
15. When in doubt: Split multi-temporal facts, link them causally, use temporal ranges appropriately"""
1. **COMBINE Q&A** - Never create standalone question facts! Merge questions with their answers into single facts.
2. **CAPTURE ALL MEANINGFUL CONTENT** - Activities, encouragement (with specific words!), recommendations, reactions, preferences
3. **CONVERT RELATIVE DATES TO SPECIFIC DATES** - "last week" "around August 16" (NOT "in August"!), "yesterday" "on August 18". Be precise!
4. **CAPTURE WHAT WAS SAID** - "Gina said Jon is perfect mentor with determination" NOT "Jon received encouragement". Preserve the actual content!
5. **FACT_KIND DETERMINES OCCURRED DATES** - Only 'event' gets occurred_start/end. 'conversation' and 'other' = null
6. **CAPTURE PREFERENCES** - "ideal", "favorite", "love" preferences_opinions
7. **CAPTURE EXACT ADJECTIVES** - Use the EXACT words! "awesome" not "amazing", "epic" not "perfect" sensory_details
8. **CAPTURE OBSERVATIONS** - "shooting in Miami" observations: "traveled to Miami". Infer travel, achievements, capabilities!"""
import time
import logging
from openai import BadRequestError
@ -860,7 +428,7 @@ Marcus: I think that's gonna do it for us today! Don't forget to subscribe and l
messages=[
{
"role": "system",
"content": "You are a comprehensive fact extractor that creates narrative, self-contained facts with temporal ranges and causal relationships. TEMPORAL RANGES - CRITICAL: Each fact must have occurred_start and occurred_end. Point events: start=end (July 14). Range events: start to end (February: Feb 1 to Feb 28, 'lately': estimate range). TEMPORAL SPLITTING - CRITICAL: Split facts when events span >7 days or mix specific dates + vague periods. Keep combined within same day/week. CAUSAL RELATIONSHIPS - NEW: When splitting related facts, link them! 'X happened, so Y happened' → add causal_relations to X linking to Y with type 'causes'. Types: causes, caused_by, enables, prevents. Extract 2-5 COMPREHENSIVE facts per conversation, NOT dozens of fragments. COMBINE related exchanges into single narrative facts BUT split when temporally incoherent. PRESERVE all context (photos, 'new' things, visual elements, full reasoning), INCLUDE all participants and what they said/did, MAINTAIN narrative flow. MAKE LOGICAL CONNECTIONS: Connect related information! 'I lost a friend' + 'photo with Karlie''I lost my friend Karlie'. Resolve pronouns. TEMPORAL CALCULATION - CRITICAL: occurred_start/end = when fact OCCURRED, not mentioned! 'Last night' on Aug 14 → occurred_start=Aug 13. FACT TEXT: Convert relative dates to absolute: 'yesterday''on March 14, 2024', NEVER 'recently'! Extract entities (PERSON, ORG, PLACE, PRODUCT, CONCEPT, OTHER). FACT TYPES: 'world' (others/events - third person), 'agent' (memory owner - FIRST PERSON 'I did'), 'opinion' (beliefs - first person 'I believe'). Extract SUBSTANTIVE CONTENT only - skip structural statements. CAPTURE ALL DIMENSIONS: emotions (thrilled, anxious), sensory details (bright orange, loud), cognitive states (wasn't sure, realized), capabilities (can speak French, struggles with), attitudes (skeptical, enthusiastic), comparisons (better than, different from), causal relationships. Do NOT strip richness!"
"content": "Extract ALL meaningful content. COMBINE Q&A into single facts (no standalone questions!). Skip only greetings and pure filler. CONVERT RELATIVE DATES TO SPECIFIC DATES ('last week''around Aug 16' NOT 'in August'!). factual_core = WHAT was said, not THAT something was said! fact_kind: 'conversation'/'event'/'other'. Only 'event' gets occurred dates."
},
{
"role": "user",
@ -872,7 +440,23 @@ Marcus: I think that's gonna do it for us today! Don't forget to subscribe and l
temperature=0.1,
max_tokens=65000,
)
chunk_facts = [fact.model_dump() for fact in extraction_response.facts]
# Build combined fact text from dimensions and include in output
chunk_facts = []
for fact in extraction_response.facts:
fact_dict = fact.model_dump()
# Add combined 'fact' field from structured dimensions
fact_dict['fact'] = fact.build_fact_text()
# Safety net: strip occurred dates if fact_kind is not 'event'
# (in case LLM doesn't follow the rules)
if fact_dict.get('fact_kind') != 'event':
fact_dict['occurred_start'] = None
fact_dict['occurred_end'] = None
# Remove fact_kind from output (only used for prompt engineering, not stored)
fact_dict.pop('fact_kind', None)
chunk_facts.append(fact_dict)
return chunk_facts
except BadRequestError as e:
@ -1034,7 +618,7 @@ async def extract_facts_from_text(
Returns:
List of fact dictionaries with 'fact' and 'date' keys
"""
chunks = chunk_text(text, max_chars=5000)
chunks = chunk_text(text, max_chars=3000)
tasks = [
_extract_facts_with_auto_split(
chunk=chunk,

View file

@ -116,7 +116,11 @@ class LLMConfig:
**kwargs
}
if self.provider == "groq":
call_params["extra_body"] = {"service_tier": "auto"}
call_params["extra_body"] = {
"service_tier": "auto",
"reasoning_effort": "low", # Reduce reasoning overhead
"include_reasoning": False, # Disable hidden reasoning tokens
}
last_exception = None
@ -137,11 +141,22 @@ class LLMConfig:
# Log call details on success
duration = time.time() - start_time
usage = response.usage
logger.info(
f"model={self.provider}/{self.model}, "
f"input_tokens={usage.prompt_tokens}, output_tokens={usage.completion_tokens}, "
f"total_tokens={usage.total_tokens}, time={duration:.3f}s"
)
ratio = max(1, usage.completion_tokens) / usage.prompt_tokens
if ratio > 3:
raw_content = response.choices[0].message.content
raw_len = len(raw_content) if raw_content else 0
logger.info(
f"model={self.provider}/{self.model}, "
f"input_tokens={usage.prompt_tokens}, output_tokens={usage.completion_tokens}, "
f"total_tokens={usage.total_tokens}, time={duration:.3f}s, ratio out/in={ratio:.2f}, HIGH RATIO - raw_content_chars={raw_len}, \n\nin={messages}\n\nout={result}\n\nraw={raw_content}\n\n"
)
else:
logger.info(
f"model={self.provider}/{self.model}, "
f"input_tokens={usage.prompt_tokens}, output_tokens={usage.completion_tokens}, "
f"total_tokens={usage.total_tokens}, time={duration:.3f}s, ratio out/in={ratio:.2f}, "
)
return result

View file

@ -34,9 +34,10 @@ from . import (
link_utils,
think_utils,
agent_utils,
observation_utils,
)
from .llm_wrapper import LLMConfig
from .response_models import SearchResult as SearchResultModel, ThinkResult, MemoryFact
from .response_models import SearchResult as SearchResultModel, ThinkResult, MemoryFact, EntityState, EntityObservation
from .task_backend import TaskBackend, AsyncIOQueueBackend
from .search.reranking import CrossEncoderReranker
from ..pg0 import EmbeddedPostgres
@ -277,6 +278,8 @@ class MemoryEngine:
await self._handle_form_opinion(task_dict)
elif task_type == 'batch_put':
await self._handle_batch_put(task_dict)
elif task_type == 'regenerate_observations':
await self._handle_regenerate_observations(task_dict)
else:
logger.error(f"Unknown task type: {task_type}")
# Don't retry unknown task types
@ -348,7 +351,6 @@ class MemoryEngine:
# Start pg0 embedded PostgreSQL if configured
if self._use_pg0:
logger.info("Starting pg0 embedded PostgreSQL...")
self._pg0 = EmbeddedPostgres()
self.db_url = await self._pg0.ensure_running()
logger.info(f"pg0 PostgreSQL running at: {self.db_url}")
@ -1153,6 +1155,62 @@ class MemoryEngine:
})
logger.debug("[PUT_BATCH_ASYNC] Opinion reinforcement task queued in background")
# Trigger observation regeneration for entities in new units
# Only regenerate for top-N entities with at least min_facts
TOP_N_ENTITIES = 5
MIN_FACTS_THRESHOLD = 5
if all_entity_links:
unique_entity_ids = set()
for link in all_entity_links:
# links are tuples: (from_unit_id, to_unit_id, link_type, weight, entity_id)
if len(link) >= 5 and link[4]:
unique_entity_ids.add(str(link[4]))
if unique_entity_ids:
# Query entities with their fact counts and whether they have observations
# Only consider entities that:
# 1. Are in top-N by mention count AND have >= min_facts, OR
# 2. Already have observations (to keep them updated)
entity_rows = await conn.fetch(
"""
WITH entity_fact_counts AS (
SELECT
e.id,
e.canonical_name,
e.last_seen,
e.mention_count,
COUNT(DISTINCT mu.id) FILTER (WHERE mu.fact_type IN ('world', 'agent')) as fact_count,
COUNT(DISTINCT mu.id) FILTER (WHERE mu.fact_type = 'observation') as obs_count,
RANK() OVER (ORDER BY e.mention_count DESC) as rank
FROM entities e
LEFT JOIN unit_entities ue ON e.id = ue.entity_id
LEFT JOIN memory_units mu ON ue.unit_id = mu.id AND mu.agent_id = $1
WHERE e.agent_id = $1 AND e.id = ANY($2::uuid[])
GROUP BY e.id, e.canonical_name, e.last_seen, e.mention_count
)
SELECT id, canonical_name, last_seen, fact_count, obs_count
FROM entity_fact_counts
WHERE (rank <= $3 AND fact_count >= $4) OR obs_count > 0
""",
agent_id,
[uuid.UUID(eid) for eid in unique_entity_ids],
TOP_N_ENTITIES,
MIN_FACTS_THRESHOLD
)
# Submit observation regeneration tasks for qualifying entities
for row in entity_rows:
await self._task_backend.submit_task({
'type': 'regenerate_observations',
'agent_id': agent_id,
'entity_id': str(row['id']),
'entity_name': row['canonical_name'],
'version': row['last_seen'].isoformat() if row['last_seen'] else None
})
if entity_rows:
logger.debug(f"[PUT_BATCH_ASYNC] Observation regeneration tasks queued for {len(entity_rows)} entities (top-{TOP_N_ENTITIES}, min {MIN_FACTS_THRESHOLD} facts)")
return result_unit_ids
except Exception as e:
@ -1201,6 +1259,8 @@ class MemoryEngine:
max_tokens: int = 4096,
enable_trace: bool = False,
question_date: Optional[datetime] = None,
include_entities: bool = False,
max_entity_tokens: int = 1024,
) -> SearchResultModel:
"""
Search memories using N*4-way parallel retrieval (N fact types × 4 retrieval methods).
@ -1222,11 +1282,14 @@ class MemoryEngine:
including a fact that would exceed the limit
enable_trace: Whether to return search trace for debugging (deprecated)
question_date: Optional date when question was asked (for temporal filtering)
include_entities: Whether to include entity observations in the response
max_entity_tokens: Maximum tokens for entity observations (default 500)
Returns:
SearchResultModel containing:
- results: List of MemoryFact objects
- trace: Optional trace information for debugging
- entities: Optional dict of entity states (if include_entities=True)
"""
# Backpressure: limit concurrent searches to prevent overwhelming the database
async with self._search_semaphore:
@ -1235,7 +1298,8 @@ class MemoryEngine:
for attempt in range(max_retries + 1):
try:
return await self._search_with_retries(
agent_id, query, fact_type, thinking_budget, max_tokens, enable_trace, question_date
agent_id, query, fact_type, thinking_budget, max_tokens, enable_trace, question_date,
include_entities, max_entity_tokens
)
except Exception as e:
# Check if it's a connection error
@ -1266,7 +1330,9 @@ class MemoryEngine:
max_tokens: int,
enable_trace: bool,
question_date: Optional[datetime] = None,
) -> tuple[List[Dict[str, Any]], Optional[Any]]:
include_entities: bool = False,
max_entity_tokens: int = 500,
) -> SearchResultModel:
"""
Search implementation with modular retrieval and reranking.
@ -1284,9 +1350,11 @@ class MemoryEngine:
thinking_budget: Nodes to explore in graph traversal
max_tokens: Maximum tokens to return (counts only 'text' field)
enable_trace: Whether to return search trace (deprecated)
include_entities: Whether to include entity observations
max_entity_tokens: Maximum tokens for entity observations
Returns:
(results, trace) tuple where trace is None (tracing removed)
SearchResultModel with results, trace, and optional entities
"""
# Initialize tracer if requested
from .search_tracer import SearchTracer
@ -1332,18 +1400,22 @@ class MemoryEngine:
]
all_retrievals = await asyncio.gather(*retrieval_tasks)
# Combine all results from all fact types
# Combine all results from all fact types and aggregate timings
semantic_results = []
bm25_results = []
graph_results = []
temporal_results = []
aggregated_timings = {"semantic": 0.0, "bm25": 0.0, "graph": 0.0, "temporal": 0.0}
for ft_semantic, ft_bm25, ft_graph, ft_temporal in all_retrievals:
for ft_semantic, ft_bm25, ft_graph, ft_temporal, ft_timings in all_retrievals:
semantic_results.extend(ft_semantic)
bm25_results.extend(ft_bm25)
graph_results.extend(ft_graph)
if ft_temporal:
temporal_results.extend(ft_temporal)
# Track max timing for each method (since they run in parallel across fact types)
for method, duration in ft_timings.items():
aggregated_timings[method] = max(aggregated_timings[method], duration)
# If no temporal results from any fact type, set to None
if not temporal_results:
@ -1353,21 +1425,23 @@ class MemoryEngine:
step_duration = time.time() - step_start
total_retrievals = len(fact_type) * (4 if temporal_results else 3)
# Format per-method timings
timing_parts = [
f"semantic={len(semantic_results)}({aggregated_timings['semantic']:.3f}s)",
f"bm25={len(bm25_results)}({aggregated_timings['bm25']:.3f}s)",
f"graph={len(graph_results)}({aggregated_timings['graph']:.3f}s)"
]
if temporal_results:
log_buffer.append(f" [2] {total_retrievals}-way retrieval ({len(fact_type)} fact_types): semantic={len(semantic_results)}, bm25={len(bm25_results)}, graph={len(graph_results)}, temporal={len(temporal_results)} in {step_duration:.3f}s")
else:
log_buffer.append(f" [2] {total_retrievals}-way retrieval ({len(fact_type)} fact_types): semantic={len(semantic_results)}, bm25={len(bm25_results)}, graph={len(graph_results)} in {step_duration:.3f}s")
timing_parts.append(f"temporal={len(temporal_results)}({aggregated_timings['temporal']:.3f}s)")
log_buffer.append(f" [2] {total_retrievals}-way retrieval ({len(fact_type)} fact_types): {', '.join(timing_parts)} in {step_duration:.3f}s")
# Record retrieval results for tracer
if tracer:
# Estimate duration for each method (since they run in parallel)
estimated_duration = retrieval_duration
# Add semantic retrieval results
tracer.add_retrieval_results(
method_name="semantic",
results=semantic_results,
duration_seconds=estimated_duration,
duration_seconds=aggregated_timings["semantic"],
score_field="similarity",
metadata={"limit": thinking_budget}
)
@ -1376,7 +1450,7 @@ class MemoryEngine:
tracer.add_retrieval_results(
method_name="bm25",
results=bm25_results,
duration_seconds=estimated_duration,
duration_seconds=aggregated_timings["bm25"],
score_field="bm25_score",
metadata={"limit": thinking_budget}
)
@ -1385,7 +1459,7 @@ class MemoryEngine:
tracer.add_retrieval_results(
method_name="graph",
results=graph_results,
duration_seconds=estimated_duration,
duration_seconds=aggregated_timings["graph"],
score_field="similarity", # Graph uses similarity for activation
metadata={"budget": thinking_budget}
)
@ -1395,7 +1469,7 @@ class MemoryEngine:
tracer.add_retrieval_results(
method_name="temporal",
results=temporal_results,
duration_seconds=estimated_duration,
duration_seconds=aggregated_timings["temporal"],
score_field="temporal_score",
metadata={"budget": thinking_budget}
)
@ -1452,7 +1526,10 @@ class MemoryEngine:
"id": doc_id,
"text": data["text"],
"context": data.get("context", ""),
"event_date": data["event_date"], # Keep as datetime for now
"occurred_start": data.get("occurred_start"),
"occurred_end": data.get("occurred_end"),
"mentioned_at": data.get("mentioned_at"),
"document_id": data.get("document_id"),
"fact_type": data.get("fact_type"), # Include fact type for filtering
"access_count": data.get("access_count", 0),
"semantic_similarity": semantic_sim,
@ -1487,6 +1564,56 @@ class MemoryEngine:
"candidates_reranked": len(results)
})
# Step 4.5: Combine cross-encoder score with retrieval signals
# This preserves retrieval work (RRF, temporal, recency) instead of pure cross-encoder ranking
if results:
# Normalize RRF scores to [0, 1] range
rrf_scores = [r.get("rrf_score", 0) for r in results]
max_rrf = max(rrf_scores) if rrf_scores else 1.0
min_rrf = min(rrf_scores) if rrf_scores else 0.0
rrf_range = max_rrf - min_rrf if max_rrf > min_rrf else 1.0
# Calculate recency based on occurred_start (more recent = higher score)
now = utcnow()
for r in results:
# Normalize RRF score
rrf_normalized = (r.get("rrf_score", 0) - min_rrf) / rrf_range if rrf_range > 0 else 0.5
# Calculate recency (decay over 365 days, minimum 0.1)
recency = 0.5 # default for missing dates
if r.get("occurred_start"):
occurred = r["occurred_start"]
if hasattr(occurred, 'tzinfo') and occurred.tzinfo is None:
from datetime import timezone
occurred = occurred.replace(tzinfo=timezone.utc)
days_ago = (now - occurred).total_seconds() / 86400
recency = max(0.1, 1.0 - (days_ago / 365)) # Linear decay over 1 year
# Get temporal proximity if available (already 0-1)
temporal = r.get("temporal_proximity", 0.5)
# Weighted combination
# Cross-encoder: 60% (semantic relevance)
# RRF: 20% (retrieval consensus)
# Temporal proximity: 10% (time relevance for temporal queries)
# Recency: 10% (prefer recent facts)
cross_encoder_score = r.get("cross_encoder_score_normalized", 0)
combined_score = (
0.6 * cross_encoder_score +
0.2 * rrf_normalized +
0.1 * temporal +
0.1 * recency
)
r["rrf_normalized"] = rrf_normalized
r["recency"] = recency
r["combined_score"] = combined_score
r["weight"] = combined_score # Update weight for final ranking
# Re-sort by combined score
results.sort(key=lambda x: x["weight"], reverse=True)
log_buffer.append(f" [4.6] Combined scoring: cross_encoder(0.6) + rrf(0.2) + temporal(0.1) + recency(0.1)")
# Step 5: Truncate to thinking_budget * 2 for token filtering
rerank_limit = thinking_budget * 2
top_results = results[:rerank_limit]
@ -1517,7 +1644,7 @@ class MemoryEngine:
node_id=result["id"],
text=result["text"],
context=result.get("context", ""),
event_date=result["event_date"],
event_date=result.get("occurred_start"),
access_count=result.get("access_count", 0),
is_entry_point=(result["id"] in [ep.node_id for ep in tracer.entry_points]),
parent_node_id=None, # In parallel retrieval, there's no clear parent
@ -1547,9 +1674,6 @@ class MemoryEngine:
# Convert datetime objects to ISO strings for JSON serialization
for result in top_results:
if result.get("event_date"):
event_date = result["event_date"]
result["event_date"] = event_date.isoformat() if hasattr(event_date, 'isoformat') else event_date
if result.get("occurred_start"):
occurred_start = result["occurred_start"]
result["occurred_start"] = occurred_start.isoformat() if hasattr(occurred_start, 'isoformat') else occurred_start
@ -1560,28 +1684,98 @@ class MemoryEngine:
mentioned_at = result["mentioned_at"]
result["mentioned_at"] = mentioned_at.isoformat() if hasattr(mentioned_at, 'isoformat') else mentioned_at
# Get entities for each fact if include_entities is requested
fact_entity_map = {} # unit_id -> list of (entity_id, entity_name)
if include_entities and top_results:
unit_ids = [uuid.UUID(str(r.get("id"))) for r in top_results if r.get("id")]
if unit_ids:
async with acquire_with_retry(pool) as entity_conn:
entity_rows = await entity_conn.fetch(
"""
SELECT ue.unit_id, e.id as entity_id, e.canonical_name
FROM unit_entities ue
JOIN entities e ON ue.entity_id = e.id
WHERE ue.unit_id = ANY($1::uuid[])
""",
unit_ids
)
for row in entity_rows:
unit_id = str(row['unit_id'])
if unit_id not in fact_entity_map:
fact_entity_map[unit_id] = []
fact_entity_map[unit_id].append({
'entity_id': str(row['entity_id']),
'canonical_name': row['canonical_name']
})
# Convert results to MemoryFact objects
memory_facts = []
for result in top_results:
result_id = str(result.get("id"))
# Get entity names for this fact
entity_names = None
if include_entities and result_id in fact_entity_map:
entity_names = [e['canonical_name'] for e in fact_entity_map[result_id]]
memory_facts.append(MemoryFact(
id=str(result.get("id")),
id=result_id,
text=result.get("text"),
fact_type=result.get("fact_type", "world"),
entities=entity_names,
context=result.get("context"),
event_date=result.get("event_date"),
occurred_start=result.get("occurred_start"),
occurred_end=result.get("occurred_end"),
mentioned_at=result.get("mentioned_at"),
document_id=result.get("document_id"),
activation=result.get("activation")
))
# Fetch entity observations if requested
entities_dict = None
if include_entities and fact_entity_map:
# Collect unique entities from top results
unique_entities = {} # entity_id -> entity_name
for entity_list in fact_entity_map.values():
for entity in entity_list:
unique_entities[entity['entity_id']] = entity['canonical_name']
# Fetch observations for each entity (respect token budget)
entities_dict = {}
total_entity_tokens = 0
encoding = _get_tiktoken_encoding()
for entity_id, entity_name in unique_entities.items():
if total_entity_tokens >= max_entity_tokens:
break
observations = await self.get_entity_observations(agent_id, entity_id, limit=5)
# Calculate tokens for this entity's observations
entity_tokens = 0
included_observations = []
for obs in observations:
obs_tokens = len(encoding.encode(obs.text))
if total_entity_tokens + entity_tokens + obs_tokens <= max_entity_tokens:
included_observations.append(obs)
entity_tokens += obs_tokens
else:
break
if included_observations:
entities_dict[entity_name] = EntityState(
entity_id=entity_id,
canonical_name=entity_name,
observations=included_observations
)
total_entity_tokens += entity_tokens
# Finalize trace if enabled
trace_dict = None
if tracer:
trace = tracer.finalize(top_results)
trace_dict = trace.to_dict() if trace else None
return SearchResultModel(results=memory_facts, trace=trace_dict)
return SearchResultModel(results=memory_facts, trace=trace_dict, entities=entities_dict)
except Exception as e:
log_buffer.append(f"[SEARCH {search_id}] ERROR after {time.time() - search_start:.3f}s: {str(e)}")
@ -2591,7 +2785,8 @@ Guidelines:
thinking_budget=thinking_budget,
max_tokens=4096,
enable_trace=False,
fact_type=['agent', 'world', 'opinion']
fact_type=['agent', 'world', 'opinion'],
include_entities=True
)
all_results = search_result.results
@ -2709,3 +2904,294 @@ Guidelines:
except Exception as e:
logger.warning(f"[THINK] Failed to extract/store opinions: {str(e)}")
async def get_entity_observations(
self,
agent_id: str,
entity_id: str,
limit: int = 10
) -> List[EntityObservation]:
"""
Get observations linked to an entity.
Args:
agent_id: Agent identifier
entity_id: Entity UUID to get observations for
limit: Maximum number of observations to return
Returns:
List of EntityObservation objects
"""
pool = await self._get_pool()
async with acquire_with_retry(pool) as conn:
rows = await conn.fetch(
"""
SELECT mu.text, mu.mentioned_at
FROM memory_units mu
JOIN unit_entities ue ON mu.id = ue.unit_id
WHERE mu.agent_id = $1
AND mu.fact_type = 'observation'
AND ue.entity_id = $2
ORDER BY mu.mentioned_at DESC
LIMIT $3
""",
agent_id, uuid.UUID(entity_id), limit
)
observations = []
for row in rows:
mentioned_at = row['mentioned_at'].isoformat() if row['mentioned_at'] else None
observations.append(EntityObservation(
text=row['text'],
mentioned_at=mentioned_at
))
return observations
async def list_entities(
self,
agent_id: str,
limit: int = 100
) -> List[Dict[str, Any]]:
"""
List all entities for an agent.
Args:
agent_id: Agent identifier
limit: Maximum number of entities to return
Returns:
List of entity dicts with id, canonical_name, mention_count, first_seen, last_seen
"""
pool = await self._get_pool()
async with acquire_with_retry(pool) as conn:
rows = await conn.fetch(
"""
SELECT id, canonical_name, mention_count, first_seen, last_seen, metadata
FROM entities
WHERE agent_id = $1
ORDER BY mention_count DESC, last_seen DESC
LIMIT $2
""",
agent_id, limit
)
entities = []
for row in rows:
# Handle metadata - may be dict, JSON string, or None
metadata = row['metadata']
if metadata is None:
metadata = {}
elif isinstance(metadata, str):
import json
try:
metadata = json.loads(metadata)
except json.JSONDecodeError:
metadata = {}
entities.append({
'id': str(row['id']),
'canonical_name': row['canonical_name'],
'mention_count': row['mention_count'],
'first_seen': row['first_seen'].isoformat() if row['first_seen'] else None,
'last_seen': row['last_seen'].isoformat() if row['last_seen'] else None,
'metadata': metadata
})
return entities
async def get_entity_state(
self,
agent_id: str,
entity_id: str,
entity_name: str,
limit: int = 10
) -> EntityState:
"""
Get the current state (mental model) of an entity.
Args:
agent_id: Agent identifier
entity_id: Entity UUID
entity_name: Canonical name of the entity
limit: Maximum number of observations to include
Returns:
EntityState with observations
"""
observations = await self.get_entity_observations(agent_id, entity_id, limit)
return EntityState(
entity_id=entity_id,
canonical_name=entity_name,
observations=observations
)
async def regenerate_entity_observations(
self,
agent_id: str,
entity_id: str,
entity_name: str,
version: str | None = None
) -> List[str]:
"""
Regenerate observations for an entity by:
1. Checking version for deduplication (if provided)
2. Searching all facts mentioning the entity
3. Using LLM to synthesize observations (no personality)
4. Deleting old observations for this entity
5. Storing new observations linked to the entity
Args:
agent_id: Agent identifier
entity_id: Entity UUID
entity_name: Canonical name of the entity
version: Entity's last_seen timestamp when task was created (for deduplication)
Returns:
List of created observation IDs
"""
pool = await self._get_pool()
# Step 1: Check version for deduplication
if version:
async with acquire_with_retry(pool) as conn:
current_last_seen = await conn.fetchval(
"""
SELECT last_seen
FROM entities
WHERE id = $1 AND agent_id = $2
""",
uuid.UUID(entity_id), agent_id
)
if current_last_seen and current_last_seen.isoformat() != version:
logger.debug(f"[OBSERVATIONS] Skipping {entity_name} - version mismatch (newer task pending)")
return []
# Step 2: Get all facts mentioning this entity (exclude observations themselves)
async with acquire_with_retry(pool) as conn:
rows = await conn.fetch(
"""
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.fact_type
FROM memory_units mu
JOIN unit_entities ue ON mu.id = ue.unit_id
WHERE mu.agent_id = $1
AND ue.entity_id = $2
AND mu.fact_type IN ('world', 'agent')
ORDER BY mu.occurred_start DESC
LIMIT 50
""",
agent_id, uuid.UUID(entity_id)
)
if not rows:
logger.debug(f"[OBSERVATIONS] No facts found for entity {entity_name}")
return []
# Convert to MemoryFact objects for the observation extraction
facts = []
for row in rows:
occurred_start = row['occurred_start'].isoformat() if row['occurred_start'] else None
facts.append(MemoryFact(
id=str(row['id']),
text=row['text'],
fact_type=row['fact_type'],
context=row['context'],
occurred_start=occurred_start
))
# Step 3: Extract observations using LLM (no personality)
observations = await observation_utils.extract_observations_from_facts(
self._llm_config,
entity_name,
facts
)
if not observations:
logger.debug(f"[OBSERVATIONS] No observations extracted for entity {entity_name}")
return []
# Step 4: Delete old observations and insert new ones in a transaction
async with acquire_with_retry(pool) as conn:
async with conn.transaction():
# Delete old observations for this entity
await conn.execute(
"""
DELETE FROM memory_units
WHERE id IN (
SELECT mu.id
FROM memory_units mu
JOIN unit_entities ue ON mu.id = ue.unit_id
WHERE mu.agent_id = $1
AND mu.fact_type = 'observation'
AND ue.entity_id = $2
)
""",
agent_id, uuid.UUID(entity_id)
)
# Generate embeddings for new observations
embeddings = await embedding_utils.generate_embeddings_batch(
self.embeddings, observations
)
# Insert new observations
current_time = utcnow()
created_ids = []
for obs_text, embedding in zip(observations, embeddings):
result = await conn.fetchrow(
"""
INSERT INTO memory_units (
agent_id, text, embedding, context, event_date,
occurred_start, occurred_end, mentioned_at,
fact_type, access_count
)
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, 'observation', 0)
RETURNING id
""",
agent_id,
obs_text,
str(embedding),
f"observation about {entity_name}",
current_time,
current_time,
current_time,
current_time
)
obs_id = str(result['id'])
created_ids.append(obs_id)
# Link observation to entity
await conn.execute(
"""
INSERT INTO unit_entities (unit_id, entity_id)
VALUES ($1, $2)
""",
uuid.UUID(obs_id), uuid.UUID(entity_id)
)
# Single consolidated log line
logger.info(f"[OBSERVATIONS] {entity_name}: {len(facts)} facts -> {len(created_ids)} observations")
return created_ids
async def _handle_regenerate_observations(self, task_dict: Dict[str, Any]):
"""
Handler for regenerate_observations tasks.
Args:
task_dict: Dict with 'agent_id', 'entity_id', 'entity_name', 'version'
"""
try:
agent_id = task_dict.get('agent_id')
entity_id = task_dict.get('entity_id')
entity_name = task_dict.get('entity_name')
version = task_dict.get('version') # last_seen timestamp for deduplication
if not all([agent_id, entity_id, entity_name]):
logger.error(f"[OBSERVATIONS] Missing required fields in task: {task_dict}")
return
await self.regenerate_entity_observations(agent_id, entity_id, entity_name, version)
except Exception as e:
logger.error(f"[OBSERVATIONS] Error regenerating observations: {e}")
import traceback
traceback.print_exc()

View file

@ -0,0 +1,133 @@
"""
Observation utilities for generating entity observations from facts.
Observations are objective facts synthesized from multiple memory facts
about an entity, without personality influence.
"""
import logging
from typing import List, Dict, Any
from pydantic import BaseModel, Field
from .response_models import MemoryFact
logger = logging.getLogger(__name__)
class Observation(BaseModel):
"""An observation about an entity."""
observation: str = Field(description="The observation text - a factual statement about the entity")
class ObservationExtractionResponse(BaseModel):
"""Response containing extracted observations."""
observations: List[Observation] = Field(
default_factory=list,
description="List of observations about the entity"
)
def format_facts_for_observation_prompt(facts: List[MemoryFact]) -> str:
"""Format facts as text for observation extraction prompt."""
import json
if not facts:
return "[]"
formatted = []
for fact in facts:
fact_obj = {
"text": fact.text
}
# Add context if available
if fact.context:
fact_obj["context"] = fact.context
# Add occurred_start if available
if fact.occurred_start:
fact_obj["occurred_at"] = fact.occurred_start
formatted.append(fact_obj)
return json.dumps(formatted, indent=2)
def build_observation_prompt(
entity_name: str,
facts_text: str,
) -> str:
"""Build the observation extraction prompt for the LLM."""
return f"""Based on the following facts about "{entity_name}", generate a list of key observations.
FACTS ABOUT {entity_name.upper()}:
{facts_text}
Your task: Synthesize the facts into clear, objective observations about {entity_name}.
GUIDELINES:
1. Each observation should be a factual statement about {entity_name}
2. Combine related facts into single observations where appropriate
3. Be objective - do not add opinions, judgments, or interpretations
4. Focus on what we KNOW about {entity_name}, not what we assume
5. Include observations about: identity, characteristics, roles, relationships, activities
6. Write in third person (e.g., "John is..." not "I think John is...")
7. If there are conflicting facts, note the most recent or most supported one
EXAMPLES of good observations:
- "John works at Google as a software engineer"
- "John is detail-oriented and methodical in his approach"
- "John collaborates frequently with Sarah on the AI project"
- "John joined the company in 2023"
EXAMPLES of bad observations (avoid these):
- "John seems like a good person" (opinion/judgment)
- "John probably likes his job" (assumption)
- "I believe John is reliable" (first-person opinion)
Generate 3-7 observations based on the available facts. If there are very few facts, generate fewer observations."""
def get_observation_system_message() -> str:
"""Get the system message for observation extraction."""
return "You are an objective observer synthesizing facts about an entity. Generate clear, factual observations without opinions or personality influence. Be concise and accurate."
async def extract_observations_from_facts(
llm_config,
entity_name: str,
facts: List[MemoryFact]
) -> List[str]:
"""
Extract observations from facts about an entity using LLM.
Args:
llm_config: LLM configuration to use
entity_name: Name of the entity to generate observations about
facts: List of facts mentioning the entity
Returns:
List of observation strings
"""
if not facts:
return []
facts_text = format_facts_for_observation_prompt(facts)
prompt = build_observation_prompt(entity_name, facts_text)
try:
result = await llm_config.call(
messages=[
{"role": "system", "content": get_observation_system_message()},
{"role": "user", "content": prompt}
],
response_format=ObservationExtractionResponse,
scope="memory_extract_observation"
)
observations = [op.observation for op in result.observations]
logger.debug(f"Extracted {len(observations)} observations for entity {entity_name}")
return observations
except Exception as e:
logger.warning(f"Failed to extract observations for {entity_name}: {str(e)}")
return []

View file

@ -22,8 +22,11 @@ class MemoryFact(BaseModel):
"id": "123e4567-e89b-12d3-a456-426614174000",
"text": "Alice works at Google on the AI team",
"fact_type": "world",
"entities": ["Alice", "Google"],
"context": "work info",
"event_date": "2024-01-15T10:30:00Z",
"occurred_start": "2024-01-15T10:30:00Z",
"occurred_end": "2024-01-15T10:30:00Z",
"mentioned_at": "2024-01-15T10:30:00Z",
"document_id": "session_abc123",
"metadata": {"source": "slack"},
"activation": 0.95
@ -32,9 +35,9 @@ class MemoryFact(BaseModel):
id: str = Field(description="Unique identifier for the memory fact")
text: str = Field(description="The actual text content of the memory")
fact_type: str = Field(description="Type of fact: 'world', 'agent', or 'opinion'")
fact_type: str = Field(description="Type of fact: 'world', 'agent', 'opinion', or 'observation'")
entities: Optional[List[str]] = Field(None, description="Entity names mentioned in this fact")
context: Optional[str] = Field(None, description="Additional context for the memory")
event_date: Optional[str] = Field(None, description="ISO format date when the event occurred")
occurred_start: Optional[str] = Field(None, description="ISO format date when the event started occurring")
occurred_end: Optional[str] = Field(None, description="ISO format date when the event ended occurring")
mentioned_at: Optional[str] = Field(None, description="ISO format date when the fact was mentioned/learned")
@ -60,7 +63,8 @@ class SearchResult(BaseModel):
"text": "Alice works at Google on the AI team",
"fact_type": "world",
"context": "work info",
"event_date": "2024-01-15T10:30:00Z",
"occurred_start": "2024-01-15T10:30:00Z",
"occurred_end": "2024-01-15T10:30:00Z",
"activation": 0.95
}
],
@ -73,6 +77,10 @@ class SearchResult(BaseModel):
results: List[MemoryFact] = Field(description="List of memory facts matching the query")
trace: Optional[Dict[str, Any]] = Field(None, description="Trace information for debugging")
entities: Optional[Dict[str, "EntityState"]] = Field(
None,
description="Entity states for entities mentioned in results (keyed by canonical name)"
)
class ThinkResult(BaseModel):
@ -92,7 +100,8 @@ class ThinkResult(BaseModel):
"text": "Machine learning is used in medical diagnosis",
"fact_type": "world",
"context": "healthcare",
"event_date": "2024-01-15T10:30:00Z"
"occurred_start": "2024-01-15T10:30:00Z",
"occurred_end": "2024-01-15T10:30:00Z"
}
],
"agent": [],
@ -130,3 +139,46 @@ class Opinion(BaseModel):
text: str = Field(description="The opinion text")
confidence: float = Field(description="Confidence score between 0.0 and 1.0")
class EntityObservation(BaseModel):
"""
An observation about an entity.
Observations are objective facts synthesized from multiple memory facts
about an entity, without personality influence.
"""
model_config = ConfigDict(json_schema_extra={
"example": {
"text": "John is detail-oriented and works at Google",
"mentioned_at": "2024-01-15T10:30:00Z"
}
})
text: str = Field(description="The observation text")
mentioned_at: Optional[str] = Field(None, description="ISO format date when this observation was created")
class EntityState(BaseModel):
"""
Current mental model of an entity.
Contains observations synthesized from facts about the entity.
"""
model_config = ConfigDict(json_schema_extra={
"example": {
"entity_id": "123e4567-e89b-12d3-a456-426614174000",
"canonical_name": "John",
"observations": [
{"text": "John is detail-oriented", "mentioned_at": "2024-01-15T10:30:00Z"},
{"text": "John works at Google on the AI team", "mentioned_at": "2024-01-14T09:00:00Z"}
]
}
})
entity_id: str = Field(description="Unique identifier for the entity")
canonical_name: str = Field(description="Canonical name of the entity")
observations: List[EntityObservation] = Field(
default_factory=list,
description="List of observations about this entity"
)

View file

@ -46,15 +46,15 @@ class CrossEncoderReranker:
doc_text = f"{c['context']}: {doc_text}"
# Add formatted date information for temporal awareness
if c.get("event_date"):
event_date = c["event_date"]
if c.get("occurred_start"):
occurred_start = c["occurred_start"]
# Format in two styles for better model understanding
# 1. ISO format: YYYY-MM-DD
date_iso = event_date.strftime("%Y-%m-%d")
date_iso = occurred_start.strftime("%Y-%m-%d")
# 2. Human-readable: "June 5, 2022"
date_readable = event_date.strftime("%B %d, %Y")
date_readable = occurred_start.strftime("%B %d, %Y")
# Prepend date to document text
doc_text = f"[Date: {date_readable} ({date_iso})] {doc_text}"

View file

@ -143,43 +143,58 @@ async def retrieve_graph(
if not entry_points:
return []
# Simple BFS-style spreading activation
# BFS-style spreading activation with batched neighbor fetching
visited = set()
results = []
queue = [(dict(r), r["similarity"]) for r in entry_points]
budget_remaining = budget
# Process nodes in batches to reduce DB roundtrips
batch_size = 20 # Fetch neighbors for up to 20 nodes at once
while queue and budget_remaining > 0:
current, activation = queue.pop(0)
unit_id = str(current["id"])
# Collect a batch of nodes to process
batch_nodes = []
batch_activations = {}
if unit_id in visited:
continue
while queue and len(batch_nodes) < batch_size and budget_remaining > 0:
current, activation = queue.pop(0)
unit_id = str(current["id"])
visited.add(unit_id)
budget_remaining -= 1
results.append((unit_id, current))
if unit_id not in visited:
visited.add(unit_id)
budget_remaining -= 1
results.append((unit_id, current))
batch_nodes.append(current["id"])
batch_activations[unit_id] = activation
# Get neighbors
if budget_remaining > 0:
# Batch fetch neighbors for all nodes in this batch
# Fetch top weighted neighbors (batch_size * 10 = ~200 for good distribution)
if batch_nodes and budget_remaining > 0:
max_neighbors = len(batch_nodes) * 10
neighbors = await conn.fetch(
"""
SELECT mu.id, mu.text, mu.context, mu.event_date, mu.occurred_start, mu.occurred_end, mu.mentioned_at, mu.access_count, mu.embedding, mu.fact_type, mu.document_id,
ml.weight, ml.link_type
SELECT mu.id, mu.text, mu.context, mu.occurred_start, mu.occurred_end, mu.mentioned_at,
mu.access_count, mu.embedding, mu.fact_type, mu.document_id,
ml.weight, ml.link_type, ml.from_unit_id
FROM memory_links ml
JOIN memory_units mu ON ml.to_unit_id = mu.id
WHERE ml.from_unit_id = $1
WHERE ml.from_unit_id = ANY($1::uuid[])
AND ml.weight >= 0.1
AND mu.fact_type = $2
ORDER BY ml.weight DESC
LIMIT 10
LIMIT $3
""",
current["id"], fact_type
batch_nodes, fact_type, max_neighbors
)
for n in neighbors:
neighbor_id = str(n["id"])
if neighbor_id not in visited:
# Get parent activation
parent_id = str(n["from_unit_id"])
activation = batch_activations.get(parent_id, 0.5)
# Boost activation for causal links (they're high-value relationships)
link_type = n["link_type"]
base_weight = n["weight"]
@ -408,7 +423,7 @@ async def retrieve_parallel(
thinking_budget: int,
question_date: Optional[datetime] = None,
query_analyzer: Optional["QueryAnalyzer"] = None
) -> Tuple[List, List, List, Optional[List]]:
) -> Tuple[List, List, List, Optional[List], Dict[str, float]]:
"""
Run 3-way or 4-way parallel retrieval (adds temporal if detected).
@ -423,18 +438,27 @@ async def retrieve_parallel(
query_analyzer: Query analyzer to use (defaults to TransformerQueryAnalyzer)
Returns:
Tuple of (semantic_results, bm25_results, graph_results, temporal_results)
Tuple of (semantic_results, bm25_results, graph_results, temporal_results, timings)
temporal_results is None if no temporal constraint detected
timings is a dict with per-method latencies in seconds
"""
# Detect temporal constraint
from .temporal_extraction import extract_temporal_constraint
import logging
import time
logger = logging.getLogger(__name__)
temporal_constraint = extract_temporal_constraint(
query_text, reference_date=question_date, analyzer=query_analyzer
)
# Wrapper to track timing for each retrieval method
async def timed_retrieval(name: str, coro):
start = time.time()
result = await coro
duration = time.time() - start
return result, name, duration
async def run_semantic():
async with acquire_with_retry(pool) as conn:
return await retrieve_semantic(conn, query_embedding_str, agent_id, fact_type, limit=thinking_budget)
@ -454,16 +478,29 @@ async def retrieve_parallel(
start_date, end_date, budget=thinking_budget, semantic_threshold=0.4
)
# Run retrievals in parallel
# Run retrievals in parallel with timing
timings = {}
if temporal_constraint:
start_date, end_date = temporal_constraint
semantic_results, bm25_results, graph_results, temporal_results = await asyncio.gather(
run_semantic(), run_bm25(), run_graph(), run_temporal(start_date, end_date)
results = await asyncio.gather(
timed_retrieval("semantic", run_semantic()),
timed_retrieval("bm25", run_bm25()),
timed_retrieval("graph", run_graph()),
timed_retrieval("temporal", run_temporal(start_date, end_date))
)
semantic_results, _, timings["semantic"] = results[0]
bm25_results, _, timings["bm25"] = results[1]
graph_results, _, timings["graph"] = results[2]
temporal_results, _, timings["temporal"] = results[3]
else:
semantic_results, bm25_results, graph_results = await asyncio.gather(
run_semantic(), run_bm25(), run_graph()
results = await asyncio.gather(
timed_retrieval("semantic", run_semantic()),
timed_retrieval("bm25", run_bm25()),
timed_retrieval("graph", run_graph())
)
semantic_results, _, timings["semantic"] = results[0]
bm25_results, _, timings["bm25"] = results[1]
graph_results, _, timings["graph"] = results[2]
temporal_results = None
return semantic_results, bm25_results, graph_results, temporal_results
return semantic_results, bm25_results, graph_results, temporal_results, timings

View file

@ -70,13 +70,13 @@ def format_facts_for_prompt(facts: List[MemoryFact]) -> str:
if fact.context:
fact_obj["context"] = fact.context
# Add event_date if available
if fact.event_date:
event_date = fact.event_date
if isinstance(event_date, str):
fact_obj["event_date"] = event_date
elif isinstance(event_date, datetime):
fact_obj["event_date"] = event_date.strftime('%Y-%m-%d %H:%M:%S')
# Add occurred_start if available (when the fact occurred)
if fact.occurred_start:
occurred_start = fact.occurred_start
if isinstance(occurred_start, str):
fact_obj["occurred_start"] = occurred_start
elif isinstance(occurred_start, datetime):
fact_obj["occurred_start"] = occurred_start.strftime('%Y-%m-%d %H:%M:%S')
# Add activation if available
if fact.activation is not None:

View file

@ -104,11 +104,12 @@ class MemoryUnit(Base):
name="memory_units_document_fkey",
ondelete="CASCADE",
),
CheckConstraint("fact_type IN ('world', 'agent', 'opinion')"),
CheckConstraint("fact_type IN ('world', 'agent', 'opinion', 'observation')"),
CheckConstraint("confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0)"),
CheckConstraint(
"(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR "
"(fact_type != 'opinion' AND confidence_score IS NULL)",
"(fact_type = 'observation') OR "
"(fact_type NOT IN ('opinion', 'observation') AND confidence_score IS NULL)",
name="confidence_score_fact_type_check"
),
Index("idx_memory_units_agent_id", "agent_id"),
@ -133,6 +134,13 @@ class MemoryUnit(Base):
postgresql_where=sql_text("fact_type = 'opinion'"),
postgresql_ops={"event_date": "DESC"}
),
Index(
"idx_memory_units_observation_date",
"agent_id",
"event_date",
postgresql_where=sql_text("fact_type = 'observation'"),
postgresql_ops={"event_date": "DESC"}
),
Index(
"idx_memory_units_embedding",
"embedding",

View file

@ -100,6 +100,7 @@ class EmbeddedPostgres:
username: str = DEFAULT_USERNAME,
password: str = DEFAULT_PASSWORD,
database: str = DEFAULT_DATABASE,
name: str = "hindsight",
):
"""
Initialize the embedded PostgreSQL manager.
@ -112,6 +113,7 @@ class EmbeddedPostgres:
username: Username for the database. Defaults to "hindsight"
password: Password for the database. Defaults to "hindsight"
database: Database name to create. Defaults to "hindsight"
name: Instance name for pg0. Defaults to "hindsight"
"""
self.data_dir = Path(data_dir or DEFAULT_DATA_DIR).expanduser()
self.install_dir = Path(install_dir or DEFAULT_INSTALL_DIR).expanduser()
@ -120,6 +122,7 @@ class EmbeddedPostgres:
self.username = username
self.password = password
self.database = database
self.name = name
# Binary path
binary_name = "pg0.exe" if platform.system() == "Windows" else "pg0"
@ -208,10 +211,11 @@ class EmbeddedPostgres:
# Create data directory
self.data_dir.mkdir(parents=True, exist_ok=True)
logger.info(f"Starting embedded PostgreSQL (data: {self.data_dir}, port: {self.port})...")
logger.info(f"Starting embedded PostgreSQL (name: {self.name}, data: {self.data_dir}, install: {self.install_dir}, port: {self.port})...")
returncode, stdout, stderr = await self._run_command_async(
"start",
"--name", self.name,
"--port", str(self.port),
"--username", self.username,
"--password", self.password,
@ -237,9 +241,9 @@ class EmbeddedPostgres:
if not self.is_installed():
return
logger.info("Stopping embedded PostgreSQL...")
logger.info(f"Stopping embedded PostgreSQL (name: {self.name})...")
returncode, stdout, stderr = await self._run_command_async("stop")
returncode, stdout, stderr = await self._run_command_async("stop", "--name", self.name)
if returncode != 0:
# Don't raise if server wasn't running
@ -264,7 +268,7 @@ class EmbeddedPostgres:
raise RuntimeError("pg0 is not installed.")
returncode, stdout, stderr = await self._run_command_async(
"info", "-o", "json")
"info", "--name", self.name, "-o", "json")
if returncode != 0:
raise RuntimeError(f"Failed to get PostgreSQL info: {stderr}")

View file

@ -57,3 +57,14 @@ filterwarnings = [
"ignore:The @wait_container_is_ready decorator is deprecated:DeprecationWarning",
"ignore::RuntimeWarning:asyncio",
]
[dependency-groups]
dev = [
"filelock>=3.20.0",
"pytest>=9.0.0",
"pytest-asyncio>=1.3.0",
"pytest-timeout>=2.4.0",
"pytest-xdist>=3.8.0",
"python-dotenv>=1.2.1",
"testcontainers>=4.13.3",
]

View file

@ -0,0 +1,338 @@
"""
Test observation generation and entity state functionality.
"""
import pytest
from datetime import datetime, timezone
@pytest.mark.asyncio
async def test_observation_generation_on_put(memory):
"""
Test that observations are generated when new facts are added.
1. Store facts about an entity
2. Wait for background tasks (observation generation)
3. Verify observations were created and linked to the entity
"""
agent_id = f"test_obs_{datetime.now(timezone.utc).timestamp()}"
try:
# Store some facts about an entity
await memory.put_async(
agent_id=agent_id,
content="John is a software engineer at Google. He is detail-oriented and methodical.",
context="work info",
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
)
await memory.put_async(
agent_id=agent_id,
content="John has been working on the AI team for 3 years. He specializes in machine learning.",
context="work info",
event_date=datetime(2024, 2, 1, tzinfo=timezone.utc)
)
# Wait for background tasks to complete (including observation generation)
await memory.wait_for_background_tasks()
# Find the John entity
pool = await memory._get_pool()
async with pool.acquire() as conn:
entity_row = await conn.fetchrow(
"""
SELECT id, canonical_name
FROM entities
WHERE agent_id = $1 AND LOWER(canonical_name) LIKE '%john%'
LIMIT 1
""",
agent_id
)
if entity_row:
entity_id = str(entity_row['id'])
entity_name = entity_row['canonical_name']
print(f"\n=== Found Entity ===")
print(f"Entity: {entity_name} (id: {entity_id})")
# Get observations for the entity
observations = await memory.get_entity_observations(agent_id, entity_id, limit=10)
print(f"\n=== Observations for {entity_name} ===")
print(f"Total observations: {len(observations)}")
for obs in observations:
print(f" - {obs.text}")
# Verify observations were created
if len(observations) > 0:
print(f"✓ Observations were successfully generated")
# Check that observations mention relevant content
obs_texts = " ".join([o.text.lower() for o in observations])
assert any(keyword in obs_texts for keyword in ["google", "engineer", "ai", "machine learning", "detail"]), \
"Observations should contain relevant information about John"
else:
print(f"⚠ Note: No observations were generated (this can happen if LLM extraction varies)")
else:
print(f"⚠ Note: No 'John' entity was extracted (LLM extraction may vary)")
finally:
# Cleanup
pool = await memory._get_pool()
async with pool.acquire() as conn:
await conn.execute("DELETE FROM memory_units WHERE agent_id = $1", agent_id)
await conn.execute("DELETE FROM entities WHERE agent_id = $1", agent_id)
@pytest.mark.asyncio
async def test_regenerate_entity_observations(memory):
"""
Test explicit regeneration of observations for an entity.
"""
agent_id = f"test_regen_obs_{datetime.now(timezone.utc).timestamp()}"
try:
# Store facts about an entity
await memory.put_async(
agent_id=agent_id,
content="Sarah is a product manager who loves user research and data analysis.",
context="work info",
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
)
await memory.wait_for_background_tasks()
# Find the Sarah entity
pool = await memory._get_pool()
async with pool.acquire() as conn:
entity_row = await conn.fetchrow(
"""
SELECT id, canonical_name
FROM entities
WHERE agent_id = $1 AND LOWER(canonical_name) LIKE '%sarah%'
LIMIT 1
""",
agent_id
)
if entity_row:
entity_id = str(entity_row['id'])
entity_name = entity_row['canonical_name']
# Manually regenerate observations
created_ids = await memory.regenerate_entity_observations(
agent_id=agent_id,
entity_id=entity_id,
entity_name=entity_name
)
print(f"\n=== Regenerated Observations ===")
print(f"Created {len(created_ids)} observations for {entity_name}")
# Get the observations
observations = await memory.get_entity_observations(agent_id, entity_id, limit=10)
for obs in observations:
print(f" - {obs.text}")
# Verify observations were created
if len(created_ids) > 0:
assert len(observations) == len(created_ids), "Should have same number of observations as created IDs"
print(f"✓ Observations regenerated successfully")
else:
print(f"⚠ Note: No observations were regenerated")
else:
print(f"⚠ Note: No 'Sarah' entity was extracted")
finally:
# Cleanup
pool = await memory._get_pool()
async with pool.acquire() as conn:
await conn.execute("DELETE FROM memory_units WHERE agent_id = $1", agent_id)
await conn.execute("DELETE FROM entities WHERE agent_id = $1", agent_id)
@pytest.mark.asyncio
async def test_search_with_include_entities(memory):
"""
Test that search with include_entities=True returns entity observations.
"""
agent_id = f"test_search_ent_{datetime.now(timezone.utc).timestamp()}"
try:
# Store facts about entities
await memory.put_async(
agent_id=agent_id,
content="Alice is a data scientist who works on recommendation systems at Netflix.",
context="work info",
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
)
await memory.put_async(
agent_id=agent_id,
content="Alice presented her research at the ML conference last month. She is an expert in deep learning.",
context="work info",
event_date=datetime(2024, 2, 1, tzinfo=timezone.utc)
)
# Wait for background tasks
await memory.wait_for_background_tasks()
# Search with include_entities=True
result = await memory.search_async(
agent_id=agent_id,
query="What does Alice do?",
fact_type=["world", "agent"],
thinking_budget=30,
max_tokens=2000,
include_entities=True,
max_entity_tokens=500
)
print(f"\n=== Search Results ===")
print(f"Found {len(result.results)} facts")
for fact in result.results:
print(f" - {fact.text}")
if fact.entities:
print(f" Entities: {', '.join(fact.entities)}")
print(f"\n=== Entity Observations ===")
if result.entities:
for name, state in result.entities.items():
print(f"\n{name}:")
for obs in state.observations:
print(f" - {obs.text}")
else:
print("No entity observations returned")
# Verify results
assert len(result.results) > 0, "Should find some facts"
# Check if entities are included in facts
facts_with_entities = [f for f in result.results if f.entities]
if facts_with_entities:
print(f"{len(facts_with_entities)} facts have entity information")
# Check if entity observations are included
if result.entities:
print(f"✓ Entity observations included for {len(result.entities)} entities")
for name, state in result.entities.items():
assert state.canonical_name == name, "Entity canonical_name should match key"
assert state.entity_id, "Entity should have an ID"
finally:
# Cleanup
pool = await memory._get_pool()
async with pool.acquire() as conn:
await conn.execute("DELETE FROM memory_units WHERE agent_id = $1", agent_id)
await conn.execute("DELETE FROM entities WHERE agent_id = $1", agent_id)
@pytest.mark.asyncio
async def test_get_entity_state(memory):
"""
Test getting the full state of an entity.
"""
agent_id = f"test_entity_state_{datetime.now(timezone.utc).timestamp()}"
try:
# Store facts
await memory.put_async(
agent_id=agent_id,
content="Bob is a frontend developer who specializes in React and TypeScript.",
context="work info",
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
)
await memory.wait_for_background_tasks()
# Find entity
pool = await memory._get_pool()
async with pool.acquire() as conn:
entity_row = await conn.fetchrow(
"""
SELECT id, canonical_name
FROM entities
WHERE agent_id = $1 AND LOWER(canonical_name) LIKE '%bob%'
LIMIT 1
""",
agent_id
)
if entity_row:
entity_id = str(entity_row['id'])
entity_name = entity_row['canonical_name']
# Get entity state
state = await memory.get_entity_state(
agent_id=agent_id,
entity_id=entity_id,
entity_name=entity_name,
limit=10
)
print(f"\n=== Entity State for {entity_name} ===")
print(f"Entity ID: {state.entity_id}")
print(f"Canonical Name: {state.canonical_name}")
print(f"Observations: {len(state.observations)}")
for obs in state.observations:
print(f" - {obs.text}")
assert state.entity_id == entity_id, "Entity ID should match"
assert state.canonical_name == entity_name, "Canonical name should match"
finally:
# Cleanup
pool = await memory._get_pool()
async with pool.acquire() as conn:
await conn.execute("DELETE FROM memory_units WHERE agent_id = $1", agent_id)
await conn.execute("DELETE FROM entities WHERE agent_id = $1", agent_id)
@pytest.mark.asyncio
async def test_observation_fact_type_in_database(memory):
"""
Test that observations are stored with correct fact_type in database.
"""
agent_id = f"test_obs_db_{datetime.now(timezone.utc).timestamp()}"
try:
# Store facts
await memory.put_async(
agent_id=agent_id,
content="Charlie is a DevOps engineer who manages the Kubernetes infrastructure.",
context="work info",
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
)
await memory.wait_for_background_tasks()
# Check that observations have correct fact_type
pool = await memory._get_pool()
async with pool.acquire() as conn:
observations = await conn.fetch(
"""
SELECT id, text, fact_type, context
FROM memory_units
WHERE agent_id = $1 AND fact_type = 'observation'
""",
agent_id
)
print(f"\n=== Observation Records in Database ===")
print(f"Found {len(observations)} observation records")
for obs in observations:
print(f" - fact_type: {obs['fact_type']}")
print(f" text: {obs['text']}")
print(f" context: {obs['context']}")
if len(observations) > 0:
for obs in observations:
assert obs['fact_type'] == 'observation', "All observation records should have fact_type='observation'"
print(f"✓ All observations have correct fact_type")
finally:
# Cleanup
pool = await memory._get_pool()
async with pool.acquire() as conn:
await conn.execute("DELETE FROM memory_units WHERE agent_id = $1", agent_id)
await conn.execute("DELETE FROM entities WHERE agent_id = $1", agent_id)

View file

@ -0,0 +1,33 @@
import { NextRequest, NextResponse } from 'next/server';
const DATAPLANE_URL = process.env.HINDSIGHT_CP_DATAPLANE_API_URL || 'http://localhost:8888';
export async function POST(
request: NextRequest,
{ params }: { params: Promise<{ entityId: string }> }
) {
try {
const { entityId } = await params;
const searchParams = request.nextUrl.searchParams;
const agentId = searchParams.get('agent_id');
if (!agentId) {
return NextResponse.json(
{ error: 'agent_id is required' },
{ status: 400 }
);
}
const decodedEntityId = decodeURIComponent(entityId);
const url = `${DATAPLANE_URL}/api/v1/agents/${agentId}/entities/${decodedEntityId}/regenerate`;
const response = await fetch(url, { method: 'POST' });
const data = await response.json();
return NextResponse.json(data, { status: response.status });
} catch (error) {
console.error('Error regenerating entity observations:', error);
return NextResponse.json(
{ error: 'Failed to regenerate entity observations' },
{ status: 500 }
);
}
}

View file

@ -0,0 +1,34 @@
import { NextRequest, NextResponse } from 'next/server';
const DATAPLANE_URL = process.env.HINDSIGHT_CP_DATAPLANE_API_URL || 'http://localhost:8888';
export async function GET(
request: NextRequest,
{ params }: { params: Promise<{ entityId: string }> }
) {
try {
const { entityId } = await params;
const searchParams = request.nextUrl.searchParams;
const agentId = searchParams.get('agent_id');
if (!agentId) {
return NextResponse.json(
{ error: 'agent_id is required' },
{ status: 400 }
);
}
// Decode URL-encoded entityId in case it contains special chars
const decodedEntityId = decodeURIComponent(entityId);
const url = `${DATAPLANE_URL}/api/v1/agents/${agentId}/entities/${decodedEntityId}`;
const response = await fetch(url);
const data = await response.json();
return NextResponse.json(data, { status: response.status });
} catch (error) {
console.error('Error getting entity:', error);
return NextResponse.json(
{ error: 'Failed to get entity' },
{ status: 500 }
);
}
}

View file

@ -0,0 +1,33 @@
import { NextRequest, NextResponse } from 'next/server';
const DATAPLANE_URL = process.env.HINDSIGHT_CP_DATAPLANE_API_URL || 'http://localhost:8888';
export async function GET(request: NextRequest) {
try {
const searchParams = request.nextUrl.searchParams;
const agentId = searchParams.get('agent_id');
if (!agentId) {
return NextResponse.json(
{ error: 'agent_id is required' },
{ status: 400 }
);
}
// Remove agent_id from query params and rebuild query string
const newSearchParams = new URLSearchParams(searchParams);
newSearchParams.delete('agent_id');
const queryString = newSearchParams.toString();
const url = `${DATAPLANE_URL}/api/v1/agents/${agentId}/entities${queryString ? `?${queryString}` : ''}`;
const response = await fetch(url);
const data = await response.json();
return NextResponse.json(data, { status: response.status });
} catch (error) {
console.error('Error listing entities:', error);
return NextResponse.json(
{ error: 'Failed to list entities' },
{ status: 500 }
);
}
}

View file

@ -4,13 +4,14 @@ import { useState } from 'react';
import { AgentSelector } from '@/components/agent-selector';
import { DataView } from '@/components/data-view';
import { DocumentsView } from '@/components/documents-view';
import { EntitiesView } from '@/components/entities-view';
import { ThinkView } from '@/components/think-view';
import { AddMemoryView } from '@/components/add-memory-view';
import { StatsView } from '@/components/stats-view';
import { SearchDebugView } from '@/components/search-debug-view';
import { useAgent } from '@/lib/agent-context';
type MainTab = 'data' | 'documents' | 'search' | 'stats' | 'think' | 'add';
type MainTab = 'data' | 'documents' | 'entities' | 'search' | 'stats' | 'think' | 'add';
type DataSubTab = 'world' | 'agent' | 'opinion';
export default function DashboardPage() {
@ -59,6 +60,7 @@ export default function DashboardPage() {
<div className="bg-muted border-b-2 border-primary">
<TabButton tab="data" label="Data" />
<TabButton tab="documents" label="Documents" />
<TabButton tab="entities" label="Entities" />
<TabButton tab="search" label="Search Debug" />
<TabButton tab="stats" label="Stats & Operations" />
<TabButton tab="think" label="Think" />
@ -106,6 +108,16 @@ export default function DashboardPage() {
)}
</div>
{/* Entities Tab */}
<div className={mainTab !== 'entities' ? 'hidden' : ''}>
<h2 className="text-2xl font-bold mb-4">Entities</h2>
{!currentAgent ? (
<NoAgentMessage message="Please select an agent from the dropdown above to view entities." />
) : (
<EntitiesView />
)}
</div>
{/* Search Debug Tab */}
<div className={mainTab !== 'search' ? 'hidden' : ''}>
<h2 className="text-2xl font-bold mb-4">Search Debug</h2>

View file

@ -0,0 +1,205 @@
'use client';
import { useState, useEffect } from 'react';
import { dataplaneClient } from '@/lib/api';
import { useAgent } from '@/lib/agent-context';
interface Entity {
id: string;
canonical_name: string;
mention_count: number;
first_seen?: string;
last_seen?: string;
metadata?: Record<string, any>;
}
interface EntityDetail extends Entity {
observations: Array<{
text: string;
mentioned_at?: string;
}>;
}
export function EntitiesView() {
const { currentAgent } = useAgent();
const [entities, setEntities] = useState<Entity[]>([]);
const [loading, setLoading] = useState(false);
const [selectedEntity, setSelectedEntity] = useState<EntityDetail | null>(null);
const [loadingDetail, setLoadingDetail] = useState(false);
const [regenerating, setRegenerating] = useState(false);
const loadEntities = async () => {
if (!currentAgent) return;
setLoading(true);
try {
const result: any = await dataplaneClient.listEntities({
agent_id: currentAgent,
limit: 100,
});
setEntities(result.entities || []);
} catch (error) {
console.error('Error loading entities:', error);
alert('Error loading entities: ' + (error as Error).message);
} finally {
setLoading(false);
}
};
const loadEntityDetail = async (entityId: string) => {
if (!currentAgent) return;
setLoadingDetail(true);
try {
const result: any = await dataplaneClient.getEntity(entityId, currentAgent);
setSelectedEntity(result);
} catch (error) {
console.error('Error loading entity detail:', error);
alert('Error loading entity detail: ' + (error as Error).message);
} finally {
setLoadingDetail(false);
}
};
const regenerateObservations = async () => {
if (!currentAgent || !selectedEntity) return;
setRegenerating(true);
try {
await dataplaneClient.regenerateEntityObservations(selectedEntity.id, currentAgent);
// Reload entity detail to show new observations
await loadEntityDetail(selectedEntity.id);
} catch (error) {
console.error('Error regenerating observations:', error);
alert('Error regenerating observations: ' + (error as Error).message);
} finally {
setRegenerating(false);
}
};
useEffect(() => {
if (currentAgent) {
loadEntities();
setSelectedEntity(null);
}
}, [currentAgent]);
const formatDate = (dateStr?: string) => {
if (!dateStr) return 'N/A';
return new Date(dateStr).toLocaleDateString();
};
return (
<div className="flex gap-4">
{/* Entity List */}
<div className="flex-1">
<div className="mb-4 p-2.5 bg-card rounded-lg border-2 border-primary flex gap-4 items-center">
<button
onClick={loadEntities}
disabled={loading}
className="px-5 py-2 bg-primary text-primary-foreground rounded font-bold text-sm hover:opacity-90 disabled:opacity-50 disabled:cursor-not-allowed"
>
{loading ? 'Loading...' : entities.length > 0 ? 'Refresh Entities' : 'Load Entities'}
</button>
{entities.length > 0 && (
<span className="text-muted-foreground text-sm">
({entities.length} entities)
</span>
)}
</div>
{entities.length > 0 ? (
<div className="overflow-x-auto">
<table className="w-full border-collapse text-sm">
<thead>
<tr>
<th className="p-2.5 text-left border border-border bg-card text-card-foreground">ID</th>
<th className="p-2.5 text-left border border-border bg-card text-card-foreground">Name</th>
<th className="p-2.5 text-left border border-border bg-card text-card-foreground">Mentions</th>
<th className="p-2.5 text-left border border-border bg-card text-card-foreground">First Seen</th>
<th className="p-2.5 text-left border border-border bg-card text-card-foreground">Last Seen</th>
</tr>
</thead>
<tbody>
{entities.map((entity) => (
<tr
key={entity.id}
onClick={() => loadEntityDetail(entity.id)}
className={`cursor-pointer hover:bg-muted ${
selectedEntity?.id === entity.id ? 'bg-accent' : 'bg-background'
}`}
>
<td className="p-2 border border-border text-xs text-muted-foreground font-mono" title={entity.id}>{entity.id.slice(0, 8)}...</td>
<td className="p-2 border border-border font-medium">{entity.canonical_name}</td>
<td className="p-2 border border-border">{entity.mention_count}</td>
<td className="p-2 border border-border">{formatDate(entity.first_seen)}</td>
<td className="p-2 border border-border">{formatDate(entity.last_seen)}</td>
</tr>
))}
</tbody>
</table>
</div>
) : !loading && (
<div className="p-10 text-center text-muted-foreground bg-muted rounded">
No entities found. Entities are extracted from facts when memories are added.
</div>
)}
</div>
{/* Entity Detail Panel */}
{selectedEntity && (
<div className="w-96 bg-card border-2 border-primary rounded-lg p-4">
<div className="flex justify-between items-start mb-4">
<h3 className="text-lg font-bold text-card-foreground">{selectedEntity.canonical_name}</h3>
<button
onClick={() => setSelectedEntity(null)}
className="text-muted-foreground hover:text-foreground"
>
X
</button>
</div>
<div className="text-sm text-muted-foreground mb-4">
<div>Mentions: {selectedEntity.mention_count}</div>
<div>First seen: {formatDate(selectedEntity.first_seen)}</div>
<div>Last seen: {formatDate(selectedEntity.last_seen)}</div>
</div>
<div className="mb-4">
<div className="flex justify-between items-center mb-2">
<h4 className="font-bold text-card-foreground">Observations</h4>
<button
onClick={regenerateObservations}
disabled={regenerating}
className="px-3 py-1 bg-secondary text-secondary-foreground rounded text-xs font-bold hover:opacity-90 disabled:opacity-50"
>
{regenerating ? 'Regenerating...' : 'Regenerate'}
</button>
</div>
{loadingDetail ? (
<div className="text-muted-foreground text-sm">Loading observations...</div>
) : selectedEntity.observations && selectedEntity.observations.length > 0 ? (
<ul className="space-y-2">
{selectedEntity.observations.map((obs, idx) => (
<li key={idx} className="p-2 bg-muted rounded text-sm">
<div>{obs.text}</div>
{obs.mentioned_at && (
<div className="text-xs text-muted-foreground mt-1">
{formatDate(obs.mentioned_at)}
</div>
)}
</li>
))}
</ul>
) : (
<div className="text-muted-foreground text-sm">
No observations yet. Click &quot;Regenerate&quot; to generate observations from facts.
</div>
)}
</div>
</div>
)}
</div>
);
}

View file

@ -190,6 +190,56 @@ export class DataplaneClient {
method: 'DELETE',
});
}
/**
* List entities for an agent
*/
async listEntities(params: {
agent_id: string;
limit?: number;
}) {
const queryParams = new URLSearchParams();
queryParams.append('agent_id', params.agent_id);
if (params.limit) queryParams.append('limit', params.limit.toString());
return this.fetchApi<{
entities: Array<{
id: string;
canonical_name: string;
mention_count: number;
first_seen?: string;
last_seen?: string;
metadata?: Record<string, any>;
}>;
}>(`/api/entities?${queryParams}`);
}
/**
* Get entity details with observations
*/
async getEntity(entityId: string, agentId: string) {
return this.fetchApi<{
id: string;
canonical_name: string;
mention_count: number;
first_seen?: string;
last_seen?: string;
metadata?: Record<string, any>;
observations: Array<{
text: string;
mentioned_at?: string;
}>;
}>(`/api/entities/${entityId}?agent_id=${agentId}`);
}
/**
* Regenerate observations for an entity
*/
async regenerateEntityObservations(entityId: string, agentId: string) {
return this.fetchApi(`/api/entities/${entityId}/regenerate?agent_id=${agentId}`, {
method: 'POST',
});
}
}
// Export a singleton instance

View file

@ -30,7 +30,7 @@ from rich.table import Table
from rich import box
import pydantic
from hindsight_api import TemporalSemanticMemory
from hindsight_api import MemoryEngine
from openai import AsyncOpenAI
console = Console()
@ -381,7 +381,7 @@ class LLMAnswerEvaluator:
def __init__(self):
"""Initialize with LLM configuration for judge/evaluator."""
from hindsight_api.llm_wrapper import LLMConfig
from hindsight_api.engine.llm_wrapper import LLMConfig
self.llm_config = LLMConfig.for_judge()
self.client = self.llm_config._client
self.model = self.llm_config.model
@ -470,7 +470,7 @@ class BenchmarkRunner:
dataset: BenchmarkDataset,
answer_generator: LLMAnswerGenerator,
answer_evaluator: LLMAnswerEvaluator,
memory: Optional[TemporalSemanticMemory] = None
memory: Optional[MemoryEngine] = None
):
"""
Initialize benchmark runner.
@ -485,7 +485,7 @@ class BenchmarkRunner:
self.dataset = dataset
self.answer_generator = answer_generator
self.answer_evaluator = answer_evaluator
self.memory = memory or TemporalSemanticMemory(
self.memory = memory or MemoryEngine(
db_url=os.getenv("HINDSIGHT_API_DATABASE_URL"),
memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),
@ -582,7 +582,8 @@ class BenchmarkRunner:
thinking_budget=thinking_budget,
max_tokens=max_tokens,
fact_type=["world", "agent"],
question_date=question_date
question_date=question_date,
include_entities=True
)
# Convert MemoryFact objects to dictionaries for compatibility

View file

@ -7,7 +7,7 @@ import sys
from pathlib import Path
from benchmarks.common.benchmark_runner import BenchmarkRunner
from hindsight_api import TemporalSemanticMemory
from hindsight_api import MemoryEngine
import json
from datetime import datetime, timezone
@ -18,7 +18,7 @@ from openai import AsyncOpenAI
import os
from benchmarks.common.benchmark_runner import BenchmarkDataset, LLMAnswerGenerator, LLMAnswerEvaluator
from hindsight_api.llm_wrapper import LLMConfig
from hindsight_api.engine.llm_wrapper import LLMConfig
class LoComoDataset(BenchmarkDataset):
"""LoComo dataset implementation."""
@ -64,24 +64,14 @@ class LoComoDataset(BenchmarkDataset):
# Get session date
date_key = f"{session_key}_date_time"
session_date = self._parse_date(conv.get(date_key, "n/a"))
# Build session content from all turns
session_parts = []
for turn in session_data:
speaker = turn['speaker']
text = turn['text']
session_parts.append(f"{speaker}: {text}")
if session_parts:
session_content = "\n".join(session_parts)
document_id = f"{item['sample_id']}_{session_key}"
session_items.append({
"content": session_content,
"context": f"Conversation between {speaker_a} and {speaker_b} ({session_key} of {item['sample_id']})",
"event_date": session_date,
"document_id": document_id
})
session_content = json.dumps(session_data)
document_id = f"{item['sample_id']}_{session_key}"
session_items.append({
"content": session_content,
"context": f"Conversation between {speaker_a} and {speaker_b} ({session_key} of {item['sample_id']})",
"event_date": session_date,
"document_id": document_id
})
return session_items
@ -137,12 +127,7 @@ class LoComoAnswerGenerator(LLMAnswerGenerator):
Tuple of (answer, reasoning, None)
- None indicates to use the memories passed in
"""
# Format context
context_parts = []
for result in memories:
context_parts.append({"text": result.get("text"), "context": result.get("context"), "event_date": result.get("event_date")})
context = json.dumps(context_parts)
context = json.dumps(memories)
# Format question date if provided
question_date_str = ""
@ -171,30 +156,8 @@ You have access to facts and entities from a conversation.
5. Always convert relative time references to specific dates, months, or years.
6. Be as specific as possible when talking about people, places, and events
7. Timestamps in memories represent the actual time the event occurred, not the time the event was mentioned in a message.
Clarification:
When interpreting memories, use the timestamp to determine when the described event happened, not when someone talked about the event.
Example:
Memory: (2023-03-15T16:33:00Z) I went to the vet yesterday.
Question: What day did I go to the vet?
Correct Answer: March 15, 2023
Explanation:
Even though the phrase says "yesterday," the timestamp shows the event was recorded as happening on March 15th. Therefore, the actual vet visit happened on that date, regardless of the word "yesterday" in the text.
# APPROACH (Think step by step):
1. First, examine all memories that contain information related to the question
2. Examine the timestamps and content of these memories carefully
3. Look for explicit mentions of dates, times, locations, or events that answer the question
4. If the answer requires calculation (e.g., converting relative time references), show your work
5. Formulate a precise, concise answer based solely on the evidence in the memories
6. Double-check that your answer directly addresses the question asked
7. Ensure your final answer is specific and avoids vague time references
8. If you're not exactly sure, still try to attempt an answer. Sometimes the terms are sligtly different from the question, so it's better to try with the current evidence than just say you don't know.
9. Say that you cannot answer if no evidence is related to the question.
8. Include wider range of information and provide a complete answer, including all the dimensions of the question (emotional, factual..)
9. If the answer is not explicitly stated in the memories, use logical reasoning based on the information available to answer (e.g. calculate duration of an event from different memories).
Context:
{context}
@ -220,11 +183,11 @@ class LoComoThinkAnswerGenerator(LLMAnswerGenerator):
so it doesn't need external search to be performed by the benchmark runner.
"""
def __init__(self, memory: 'TemporalSemanticMemory', agent_id: str, thinking_budget: int = 500):
def __init__(self, memory: 'MemoryEngine', agent_id: str, thinking_budget: int = 500):
"""Initialize with memory instance and agent_id.
Args:
memory: TemporalSemanticMemory instance
memory: MemoryEngine instance
agent_id: Agent identifier for think queries
thinking_budget: Budget for memory exploration
"""
@ -368,7 +331,7 @@ async def run_benchmark(
from benchmarks.common.benchmark_runner import HindsightClientAdapter
memory = HindsightClientAdapter(base_url=api_url)
else:
memory = TemporalSemanticMemory(
memory = MemoryEngine(
db_url=os.getenv("HINDSIGHT_API_DATABASE_URL"),
memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),

File diff suppressed because it is too large Load diff

View file

@ -1,16 +1,7 @@
# LoComo Benchmark Results
**Overall Accuracy**: 68.42% (1055/1542)
**Overall Accuracy**: 75.97% (117/154)
| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |
|-----------|----------|-----------|---------|----------|-----------|------------|----------|-------------|
| conv-26 | 19 | 154 | 97 | 62.99% | N/A | N/A | N/A | N/A |
| conv-30 | 19 | 81 | 56 | 69.14% | N/A | N/A | N/A | N/A |
| conv-41 | 32 | 152 | 111 | 73.03% | N/A | N/A | N/A | N/A |
| conv-42 | 29 | 199 | 133 | 66.83% | N/A | N/A | N/A | N/A |
| conv-43 | 29 | 178 | 124 | 69.66% | N/A | N/A | N/A | N/A |
| conv-44 | 28 | 123 | 85 | 69.11% | N/A | N/A | N/A | N/A |
| conv-47 | 31 | 150 | 99 | 66.00% | N/A | N/A | N/A | N/A |
| conv-48 | 30 | 191 | 125 | 65.45% | N/A | N/A | N/A | N/A |
| conv-49 | 25 | 156 | 110 | 70.51% | N/A | N/A | N/A | N/A |
| conv-50 | 30 | 158 | 115 | 72.78% | N/A | N/A | N/A | N/A |
| conv-26 | 19 | 154 | 117 | 75.97% | N/A | N/A | N/A | N/A |

View file

@ -7,7 +7,7 @@ import sys
from pathlib import Path
from benchmarks.common.benchmark_runner import BenchmarkRunner
from hindsight_api import TemporalSemanticMemory
from hindsight_api import MemoryEngine
import json
from datetime import datetime, timezone
@ -18,7 +18,7 @@ from openai import AsyncOpenAI
import os
from benchmarks.common.benchmark_runner import BenchmarkDataset, LLMAnswerGenerator, LLMAnswerEvaluator
from hindsight_api.llm_wrapper import LLMConfig
from hindsight_api.engine.llm_wrapper import LLMConfig
class LongMemEvalDataset(BenchmarkDataset):
@ -330,7 +330,7 @@ async def run_benchmark(
from benchmarks.common.benchmark_runner import HindsightClientAdapter
memory = HindsightClientAdapter(base_url=api_url)
else:
memory = TemporalSemanticMemory(
memory = MemoryEngine(
db_url=os.getenv("HINDSIGHT_API_DATABASE_URL"),
memory_llm_provider=os.getenv("HINDSIGHT_API_LLM_PROVIDER", "groq"),
memory_llm_api_key=os.getenv("HINDSIGHT_API_LLM_API_KEY"),

View file

@ -315,9 +315,20 @@ def get_locomo(mode: str, filter_type: str = "all", category_filter: str = "all"
original_idx = all_results.index(item)
item_id = item.get("item_id", item.get("sample_id", f"item-{original_idx}"))
metrics = item.get("metrics", {})
accuracy = metrics.get("accuracy", 0)
correct = metrics.get("correct", 0)
total = metrics.get("total", 0)
# Calculate accuracy for filtered category
if category_filter != "all":
detailed_results = metrics.get("detailed_results", [])
category_id = int(category_filter)
filtered_correct = sum(1 for r in detailed_results if r.get("category") == category_id and r.get("is_correct") and not r.get("is_invalid"))
filtered_total = sum(1 for r in detailed_results if r.get("category") == category_id and not r.get("is_invalid"))
accuracy = (filtered_correct / filtered_total * 100) if filtered_total > 0 else 0
correct = filtered_correct
total = filtered_total
else:
accuracy = metrics.get("accuracy", 0)
correct = metrics.get("correct", 0)
total = metrics.get("total", 0)
color = "🟢" if accuracy >= 70 else ("🟡" if accuracy >= 50 else "🔴")

View file

@ -10,7 +10,7 @@ import os
from pathlib import Path
from hindsight_api.api import create_app
from hindsight_api import TemporalSemanticMemory
from hindsight_api import MemoryEngine
def generate_openapi_spec(output_path: str = None):
"""Generate OpenAPI spec and save to file."""
@ -21,7 +21,7 @@ def generate_openapi_spec(output_path: str = None):
output_path = str(root_dir / "openapi.json")
# Create a temporary memory instance for OpenAPI generation
_memory = TemporalSemanticMemory(
_memory = MemoryEngine(
db_url="mock",
memory_llm_provider="ollama",
memory_llm_api_key="mock",

View file

@ -12,7 +12,7 @@ if [ ! -f "$ENV_FILE" ]; then
exit 1
fi
echo "🚀 Starting LoComo Benchmark with '${ENV_MODE}' environment..."
echo "🚀 Starting LoComo Benchmark"
echo "📄 Loading environment from $ENV_FILE"
echo ""
@ -21,4 +21,4 @@ set -a
source "$ENV_FILE"
set +a
uv run python memora-dev/benchmarks/locomo/locomo_benchmark.py "${ARGS[@]}"
uv run python hindsight-dev/benchmarks/locomo/locomo_benchmark.py "${ARGS[@]}"

View file

@ -20,4 +20,4 @@ set -a
source "$ENV_FILE"
set +a
uv run python memora-dev/benchmarks/longmemeval/longmemeval_benchmark.py "${ARGS[@]}"
uv run python hindsight-dev/benchmarks/longmemeval/longmemeval_benchmark.py "${ARGS[@]}"

View file

@ -8,4 +8,4 @@ echo ""
echo "Server will be available at: http://localhost:8001"
echo ""
uv run python memora-dev/benchmarks/visualizer/main.py
uv run python hindsight-dev/benchmarks/visualizer/main.py

22
uv.lock
View file

@ -1176,6 +1176,17 @@ test = [
{ name = "testcontainers" },
]
[package.dev-dependencies]
dev = [
{ name = "filelock" },
{ name = "pytest" },
{ name = "pytest-asyncio" },
{ name = "pytest-timeout" },
{ name = "pytest-xdist" },
{ name = "python-dotenv" },
{ name = "testcontainers" },
]
[package.metadata]
requires-dist = [
{ name = "alembic", specifier = ">=1.17.1" },
@ -1207,6 +1218,17 @@ requires-dist = [
]
provides-extras = ["test"]
[package.metadata.requires-dev]
dev = [
{ name = "filelock", specifier = ">=3.20.0" },
{ name = "pytest", specifier = ">=9.0.0" },
{ name = "pytest-asyncio", specifier = ">=1.3.0" },
{ name = "pytest-timeout", specifier = ">=2.4.0" },
{ name = "pytest-xdist", specifier = ">=3.8.0" },
{ name = "python-dotenv", specifier = ">=1.2.1" },
{ name = "testcontainers", specifier = ">=4.13.3" },
]
[[package]]
name = "hindsight-client"
version = "0.0.7"