holetron-lab/fleet-memory
1
0
Fork 0
Code Releases Packages Activity Actions

think feat

Browse source
This commit is contained in:
Nicolò Boschi 2025-11-03 18:43:21 +01:00
parent d8d48d6f80
commit 42260c29f7
43 changed files with 42192 additions and 24303 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

202
README.md
View file

@ -14,6 +14,29 @@ The combination of these three networks enables powerful memory retrieval that g
## Architecture
### Triple Network Design
The system maintains three separate but interconnected memory networks:
1. **World Network** (`fact_type='world'`)
- General knowledge and facts about the world
- Information not specific to the agent's actions
- Example: "Alice works at Google", "Yosemite is in California"
2. **Agent Network** (`fact_type='agent'`)
- Facts about what the AI agent specifically did
- Agent's own actions and experiences
- Example: "The agent helped debug a Python script", "The agent recommended Yosemite"
3. **Opinion Network** (`fact_type='opinion'`)
- Agent's formed opinions and perspectives
- Automatically extracted during think operations
- Includes reasons and confidence scores (0.0-1.0)
- Immutable once formed (event_date = when opinion was formed)
- Example: "Python is better for data science than JavaScript (Reasons: has better libraries like pandas and numpy) [confidence: 0.85]"
All three networks share the same infrastructure (temporal/semantic/entity links) but can be searched independently or together. The **think** operation combines all three networks to formulate consistent, contextual answers while forming new opinions.
### Core Concepts
**Memory Units**: Individual sentence-level memories that are:
@ -22,6 +45,7 @@ The combination of these three networks enables powerful memory retrieval that g
- Embedded as vectors for semantic similarity
- Timestamped for temporal relationships
- Linked to extracted entities
- Classified as either 'world' or 'agent' fact type
**Entity Resolution**: Named entities (PERSON, ORG, GPE, etc.) are:
- Extracted using spaCy NER
@ -229,7 +253,6 @@ Raw content is processed through an LLM to extract meaningful facts before stora
- `langchain-text-splitters` - Intelligent text chunking
- `networkx` - Graph operations
- `pyvis` - Interactive HTML graph visualization
- `matplotlib` - Static graph visualization
- `rich` - Terminal UI
**Models**:
@ -299,12 +322,120 @@ This will:
Open `memory_graph_interactive.html` in your browser to explore the memory graph!
## Using as a Library (Local Import)
You can import this project from another Poetry project using a local path dependency:
### 1. Add to your project's `pyproject.toml`:
```toml
[tool.poetry.dependencies]
memory-poc = {path = "../memory-poc", develop = true}
```
Or using poetry CLI:
```bash
poetry add ../memory-poc --editable
```
### 2. Import the memory system:
```python
from memory import TemporalSemanticMemory
# Initialize memory
memory = TemporalSemanticMemory()
await memory.initialize()
# Use the memory system
await memory.put_batch_async(
agent_id="my_agent",
contents=["Alice works at Google", "Bob loves hiking"],
event_date=datetime.now(timezone.utc)
)
results, trace = await memory.search_async(
agent_id="my_agent",
query="Who works at Google?"
)
```
### 3. Import the FastAPI app:
```python
from web import app, memory
# Use the FastAPI app in your own project
# You can mount it as a sub-application or run it directly
import uvicorn
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)
```
### 4. Example: Extending the FastAPI app
```python
from fastapi import FastAPI
from web import app as memory_app, memory
# Create your own app
my_app = FastAPI()
# Mount the memory app as a sub-application
my_app.mount("/memory", memory_app)
# Add your own endpoints that use the memory system
@my_app.post("/my-custom-endpoint")
async def my_endpoint():
# Use the shared memory instance
results, _ = await memory.search_async(
agent_id="my_agent",
query="some query"
)
return {"results": results}
if __name__ == "__main__":
import uvicorn
uvicorn.run(my_app, host="0.0.0.0", port=8000)
```
### Available Exports
**From `memory` package:**
- `TemporalSemanticMemory` - Main memory system class
- `SearchTrace`, `SearchTracer` - Search tracing utilities
- `QueryInfo`, `EntryPoint`, `NodeVisit`, etc. - Trace data structures
**From `web` package:**
- `app` - FastAPI application instance
- `memory` - Shared TemporalSemanticMemory instance
### Web Server
To run the web interface:
```bash
# Development mode with auto-reload
uvicorn web.server:app --reload --port 8000
# Production mode
uvicorn web.server:app --host 0.0.0.0 --port 8000 --workers 4
```
Then open http://localhost:8000 in your browser to access the visualization interface.
## Project Structure
```
memory-poc/
├── memory/ # Core memory system package
│ ├── temporal_semantic_memory.py # Main memory system class
│ ├── operations/ # Modular operation mixins
│ │ ├── embedding_operations.py # Embedding generation with process pool
│ │ ├── link_operations.py # Entity, temporal, semantic links
│ │ ├── batch_operations.py # Placeholder for future extraction
│ │ └── search_operations.py # Placeholder for future extraction
│ ├── entity_resolver.py # Entity extraction and disambiguation
│ ├── llm_client.py # LLM-based fact extraction
│ └── utils.py # Utility functions
@ -320,15 +451,25 @@ memory-poc/
└── README.md # This file
```
The memory system uses a **mixin pattern** for code organization:
- `TemporalSemanticMemory` inherits from `EmbeddingOperationsMixin` and `LinkOperationsMixin`
- This reduced the main file from 1,720 lines to 1,420 lines (17% reduction)
- See `memory/operations/README.md` for detailed refactoring documentation
## Key Features
✅ **Triple network architecture**: Separate networks for world knowledge, agent actions, and opinions
✅ **Opinion formation**: Automatically extracts and stores opinions with confidence scores during thinking
✅ **Three-layered linking**: Temporal + Semantic + Entity
✅ **Entity disambiguation**: Resolves "Alice" across different contexts
✅ **Self-contained units**: Pronouns resolved to actual referents
✅ **Spreading activation**: Graph-aware search beyond vector similarity
✅ **Think operation**: Combines all three networks for consistent, contextual answers
✅ **Confidence scores**: Opinions include confidence ratings (0.0-1.0) based on supporting evidence
✅ **Interactive visualization**: Explore memory graph in browser
✅ **Recency & frequency weighting**: Recent and important memories boosted
✅ **Linguistic validation**: Memory units verified to have subject + verb
✅ **Modular architecture**: Mixin pattern with 17% code size reduction
## API Usage
@ -361,6 +502,27 @@ results, trace = memory.search(
for result in results:
print(f"{result['text']} (weight: {result['weight']:.3f})")
# Search only world facts
results, trace = memory.search(
agent_id="agent_1",
query="What does Alice do?",
fact_type="world" # Only search world network
)
# Search only agent facts
results, trace = memory.search(
agent_id="agent_1",
query="What have I done?",
fact_type="agent" # Only search agent network
)
# Search only opinions
results, trace = memory.search(
agent_id="agent_1",
query="What do I think about Python?",
fact_type="opinion" # Only search opinion network
)
# Search with tracing for debugging
results, trace = memory.search(
agent_id="agent_1",
@ -388,6 +550,44 @@ results, trace = memory.search(
)
```
### Think and Formulate Answers
The `think` operation combines all three networks to formulate consistent, contextual answers:
```python
result = await memory.think_async(
agent_id="agent_1",
query="What do you think about Python?",
thinking_budget=50,
top_k=10
)
print(result["text"]) # Plain text answer from LLM
# Access facts used to formulate the answer
for fact in result["based_on"]["world"]:
print(f"World: {fact['text']}")
for fact in result["based_on"]["agent"]:
print(f"Agent: {fact['text']}")
for fact in result["based_on"]["opinion"]:
print(f"Opinion: {fact['text']} (confidence: {fact.get('confidence_score', 'N/A')})")
# Check for newly formed opinions
for opinion in result["new_opinions"]:
print(f"New opinion formed: {opinion['text']} (confidence: {opinion['confidence']})")
```
The think operation:
1. Searches the agent network to understand the agent's identity and actions
2. Searches the world network for relevant general knowledge
3. Searches the opinion network for existing perspectives
4. Uses an LLM (Groq by default) to formulate a coherent answer, being consistent with existing opinions
5. Extracts any new opinions formed during thinking with confidence scores
6. Stores new opinions with the current timestamp and query context
7. Returns plain text response with supporting facts from all networks and new opinions
## How It Works: Example
**Input memories**:

147
alembic.ini Normal file
View file

@ -0,0 +1,147 @@
# A generic, single database configuration.
[alembic]
# path to migration scripts.
# this is typically a path given in POSIX (e.g. forward slashes)
# format, relative to the token %(here)s which refers to the location of this
# ini file
script_location = %(here)s/alembic
# template used to generate migration file names; The default value is %%(rev)s_%%(slug)s
# Uncomment the line below if you want the files to be prepended with date and time
# see https://alembic.sqlalchemy.org/en/latest/tutorial.html#editing-the-ini-file
# for all available tokens
# file_template = %%(year)d_%%(month).2d_%%(day).2d_%%(hour).2d%%(minute).2d-%%(rev)s_%%(slug)s
# sys.path path, will be prepended to sys.path if present.
# defaults to the current working directory. for multiple paths, the path separator
# is defined by "path_separator" below.
prepend_sys_path = .
# timezone to use when rendering the date within the migration file
# as well as the filename.
# If specified, requires the tzdata library which can be installed by adding
# `alembic[tz]` to the pip requirements.
# string value is passed to ZoneInfo()
# leave blank for localtime
# timezone =
# max length of characters to apply to the "slug" field
# truncate_slug_length = 40
# set to 'true' to run the environment during
# the 'revision' command, regardless of autogenerate
# revision_environment = false
# set to 'true' to allow .pyc and .pyo files without
# a source .py file to be detected as revisions in the
# versions/ directory
# sourceless = false
# version location specification; This defaults
# to <script_location>/versions. When using multiple version
# directories, initial revisions must be specified with --version-path.
# The path separator used here should be the separator specified by "path_separator"
# below.
# version_locations = %(here)s/bar:%(here)s/bat:%(here)s/alembic/versions
# path_separator; This indicates what character is used to split lists of file
# paths, including version_locations and prepend_sys_path within configparser
# files such as alembic.ini.
# The default rendered in new alembic.ini files is "os", which uses os.pathsep
# to provide os-dependent path splitting.
#
# Note that in order to support legacy alembic.ini files, this default does NOT
# take place if path_separator is not present in alembic.ini. If this
# option is omitted entirely, fallback logic is as follows:
#
# 1. Parsing of the version_locations option falls back to using the legacy
# "version_path_separator" key, which if absent then falls back to the legacy
# behavior of splitting on spaces and/or commas.
# 2. Parsing of the prepend_sys_path option falls back to the legacy
# behavior of splitting on spaces, commas, or colons.
#
# Valid values for path_separator are:
#
# path_separator = :
# path_separator = ;
# path_separator = space
# path_separator = newline
#
# Use os.pathsep. Default configuration used for new projects.
path_separator = os
# set to 'true' to search source files recursively
# in each "version_locations" directory
# new in Alembic version 1.10
# recursive_version_locations = false
# the output encoding used when revision files
# are written from script.py.mako
# output_encoding = utf-8
# database URL. This is consumed by the user-maintained env.py script only.
# other means of configuring database URLs may be customized within the env.py
# file.
# sqlalchemy.url = driver://user:pass@localhost/dbname # Disabled, using DATABASE_URL from .env
[post_write_hooks]
# post_write_hooks defines scripts or Python functions that are run
# on newly generated revision scripts. See the documentation for further
# detail and examples
# format using "black" - use the console_scripts runner, against the "black" entrypoint
# hooks = black
# black.type = console_scripts
# black.entrypoint = black
# black.options = -l 79 REVISION_SCRIPT_FILENAME
# lint with attempts to fix using "ruff" - use the module runner, against the "ruff" module
# hooks = ruff
# ruff.type = module
# ruff.module = ruff
# ruff.options = check --fix REVISION_SCRIPT_FILENAME
# Alternatively, use the exec runner to execute a binary found on your PATH
# hooks = ruff
# ruff.type = exec
# ruff.executable = ruff
# ruff.options = check --fix REVISION_SCRIPT_FILENAME
# Logging configuration. This is also consumed by the user-maintained
# env.py script only.
[loggers]
keys = root,sqlalchemy,alembic
[handlers]
keys = console
[formatters]
keys = generic
[logger_root]
level = WARNING
handlers = console
qualname =
[logger_sqlalchemy]
level = WARNING
handlers =
qualname = sqlalchemy.engine
[logger_alembic]
level = INFO
handlers =
qualname = alembic
[handler_console]
class = StreamHandler
args = (sys.stderr,)
level = NOTSET
formatter = generic
[formatter_generic]
format = %(levelname)-5.5s [%(name)s] %(message)s
datefmt = %H:%M:%S

1
alembic/README Normal file
View file

@ -0,0 +1 @@
Generic single-database configuration.

99
alembic/env.py Normal file
View file

@ -0,0 +1,99 @@
"""
Alembic environment configuration for SQLAlchemy with pgvector.
Uses synchronous psycopg2 driver for migrations to avoid pgbouncer issues.
"""
import os
from logging.config import fileConfig
from sqlalchemy import pool, engine_from_config
from sqlalchemy.engine import Connection
from alembic import context
from dotenv import load_dotenv
# Import your models here
from memory.models import Base
# Load environment variables
load_dotenv()
# this is the Alembic Config object, which provides
# access to the values within the .ini file in use.
config = context.config
# Interpret the config file for Python logging.
# This line sets up loggers basically.
if config.config_file_name is not None:
fileConfig(config.config_file_name)
# Get database URL from environment
database_url = os.getenv("DATABASE_URL")
if not database_url:
raise ValueError("DATABASE_URL environment variable is not set")
# For migrations, use psycopg2 (sync driver) to avoid pgbouncer prepared statement issues
# The application uses asyncpg, but migrations work better with psycopg2
if database_url.startswith("postgresql+asyncpg://"):
database_url = database_url.replace("postgresql+asyncpg://", "postgresql://", 1)
elif database_url.startswith("postgres+asyncpg://"):
database_url = database_url.replace("postgres+asyncpg://", "postgresql://", 1)
# Override the sqlalchemy.url in alembic.ini
config.set_main_option("sqlalchemy.url", database_url)
# add your model's MetaData object here
# for 'autogenerate' support
target_metadata = Base.metadata
# other values from the config, defined by the needs of env.py,
# can be acquired:
# my_important_option = config.get_main_option("my_important_option")
# ... etc.
def run_migrations_offline() -> None:
"""Run migrations in 'offline' mode.
This configures the context with just a URL
and not an Engine, though an Engine is acceptable
here as well. By skipping the Engine creation
we don't even need a DBAPI to be available.
Calls to context.execute() here emit the given string to the
script output.
"""
url = config.get_main_option("sqlalchemy.url")
context.configure(
url=url,
target_metadata=target_metadata,
literal_binds=True,
dialect_opts={"paramstyle": "named"},
)
with context.begin_transaction():
context.run_migrations()
def run_migrations_online() -> None:
"""Run migrations in 'online' mode with synchronous engine."""
connectable = engine_from_config(
config.get_section(config.config_ini_section, {}),
prefix="sqlalchemy.",
poolclass=pool.NullPool,
)
with connectable.connect() as connection:
context.configure(
connection=connection,
target_metadata=target_metadata
)
with context.begin_transaction():
context.run_migrations()
if context.is_offline_mode():
run_migrations_offline()
else:
run_migrations_online()

28
alembic/script.py.mako Normal file
View file

@ -0,0 +1,28 @@
"""${message}
Revision ID: ${up_revision}
Revises: ${down_revision | comma,n}
Create Date: ${create_date}
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
${imports if imports else ""}
# revision identifiers, used by Alembic.
revision: str = ${repr(up_revision)}
down_revision: Union[str, Sequence[str], None] = ${repr(down_revision)}
branch_labels: Union[str, Sequence[str], None] = ${repr(branch_labels)}
depends_on: Union[str, Sequence[str], None] = ${repr(depends_on)}
def upgrade() -> None:
"""Upgrade schema."""
${upgrades if upgrades else "pass"}
def downgrade() -> None:
"""Downgrade schema."""
${downgrades if downgrades else "pass"}

40
alembic/versions/01f989db9079_fix_memory_links_entity_id_to_be_.py Normal file
View file

@ -0,0 +1,40 @@
"""Fix memory_links entity_id to be nullable
Revision ID: 01f989db9079
Revises: af0413383b3e
Create Date: 2025-11-03 14:43:18.721430
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
# revision identifiers, used by Alembic.
revision: str = '01f989db9079'
down_revision: Union[str, Sequence[str], None] = 'af0413383b3e'
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Drop the existing primary key
op.execute('ALTER TABLE memory_links DROP CONSTRAINT memory_links_pkey')
# Change entity_id to nullable
op.alter_column('memory_links', 'entity_id',
existing_type=sa.UUID(),
nullable=True)
# Create a unique index with COALESCE expression to handle NULL entity_id
op.execute("""
CREATE UNIQUE INDEX idx_memory_links_unique
ON memory_links (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid))
""")
def downgrade() -> None:
"""Downgrade schema."""
pass

196
alembic/versions/af0413383b3e_initial_schema.py Normal file
View file

@ -0,0 +1,196 @@
"""Initial schema
Revision ID: af0413383b3e
Revises:
Create Date: 2025-11-03 14:31:53.245542
"""
from typing import Sequence, Union
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects import postgresql
import pgvector.sqlalchemy
# revision identifiers, used by Alembic.
revision: str = 'af0413383b3e'
down_revision: Union[str, Sequence[str], None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None
def upgrade() -> None:
"""Upgrade schema."""
# Create pgvector extension
op.execute('CREATE EXTENSION IF NOT EXISTS vector')
op.execute('CREATE EXTENSION IF NOT EXISTS "uuid-ossp"')
# ### commands auto generated by Alembic - please adjust! ###
op.create_table('documents',
sa.Column('id', sa.Text(), nullable=False),
sa.Column('agent_id', sa.Text(), nullable=False),
sa.Column('original_text', sa.Text(), nullable=True),
sa.Column('content_hash', sa.Text(), nullable=True),
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()),
server_default=sa.text("'{}'::jsonb"), nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.PrimaryKeyConstraint('id', 'agent_id')
)
op.create_index('idx_documents_agent_id', 'documents', ['agent_id'], unique=False)
op.create_index('idx_documents_content_hash', 'documents', ['content_hash'], unique=False)
op.create_table('entities',
sa.Column('id', sa.UUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
sa.Column('canonical_name', sa.Text(), nullable=False),
sa.Column('entity_type', sa.Text(), nullable=False),
sa.Column('agent_id', sa.Text(), nullable=False),
sa.Column('metadata', postgresql.JSONB(astext_type=sa.Text()),
server_default=sa.text("'{}'::jsonb"), nullable=False),
sa.Column('first_seen', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.Column('last_seen', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.Column('mention_count', sa.Integer(), server_default='1', nullable=False),
sa.PrimaryKeyConstraint('id')
)
op.create_index('idx_entities_agent_id', 'entities', ['agent_id'], unique=False)
op.create_index('idx_entities_agent_name_type', 'entities', ['agent_id', 'canonical_name', 'entity_type'],
unique=False)
op.create_index('idx_entities_canonical_name', 'entities', ['canonical_name'], unique=False)
op.create_index('idx_entities_type', 'entities', ['entity_type'], unique=False)
op.create_table('entity_cooccurrences',
sa.Column('entity_id_1', sa.UUID(), nullable=False),
sa.Column('entity_id_2', sa.UUID(), nullable=False),
sa.Column('cooccurrence_count', sa.Integer(), server_default='1', nullable=False),
sa.Column('last_cooccurred', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.CheckConstraint('entity_id_1 < entity_id_2', name='entity_cooccurrence_order_check'),
sa.ForeignKeyConstraint(['entity_id_1'], ['entities.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['entity_id_2'], ['entities.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('entity_id_1', 'entity_id_2')
)
op.create_index('idx_entity_cooccurrences_count', 'entity_cooccurrences', ['cooccurrence_count'], unique=False,
postgresql_ops={'cooccurrence_count': 'DESC'})
op.create_index('idx_entity_cooccurrences_entity1', 'entity_cooccurrences', ['entity_id_1'], unique=False)
op.create_index('idx_entity_cooccurrences_entity2', 'entity_cooccurrences', ['entity_id_2'], unique=False)
op.create_table('memory_units',
sa.Column('id', sa.UUID(), server_default=sa.text('uuid_generate_v4()'), nullable=False),
sa.Column('agent_id', sa.Text(), nullable=False),
sa.Column('document_id', sa.Text(), nullable=True),
sa.Column('text', sa.Text(), nullable=False),
sa.Column('embedding', pgvector.sqlalchemy.vector.VECTOR(dim=384), nullable=True),
sa.Column('context', sa.Text(), nullable=True),
sa.Column('event_date', postgresql.TIMESTAMP(timezone=True), nullable=False),
sa.Column('fact_type', sa.Text(), server_default='world', nullable=False),
sa.Column('confidence_score', sa.Float(), nullable=True),
sa.Column('access_count', sa.Integer(), server_default='0', nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.Column('updated_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.CheckConstraint(
"(fact_type = 'opinion' AND confidence_score IS NOT NULL) OR (fact_type != 'opinion' AND confidence_score IS NULL)",
name='confidence_score_fact_type_check'),
sa.CheckConstraint("fact_type IN ('world', 'agent', 'opinion')"),
sa.CheckConstraint(
'confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0)'),
sa.ForeignKeyConstraint(['document_id', 'agent_id'], ['documents.id', 'documents.agent_id'],
name='memory_units_document_fkey', ondelete='CASCADE'),
sa.PrimaryKeyConstraint('id')
)
op.create_index('idx_memory_units_access_count', 'memory_units', ['access_count'], unique=False,
postgresql_ops={'access_count': 'DESC'})
op.create_index('idx_memory_units_agent_date', 'memory_units', ['agent_id', 'event_date'], unique=False,
postgresql_ops={'event_date': 'DESC'})
op.create_index('idx_memory_units_agent_fact_type', 'memory_units', ['agent_id', 'fact_type'], unique=False)
op.create_index('idx_memory_units_agent_id', 'memory_units', ['agent_id'], unique=False)
op.create_index('idx_memory_units_agent_type_date', 'memory_units', ['agent_id', 'fact_type', 'event_date'],
unique=False, postgresql_ops={'event_date': 'DESC'})
op.create_index('idx_memory_units_document_id', 'memory_units', ['document_id'], unique=False)
op.create_index('idx_memory_units_embedding', 'memory_units', ['embedding'], unique=False, postgresql_using='hnsw',
postgresql_ops={'embedding': 'vector_cosine_ops'})
op.create_index('idx_memory_units_event_date', 'memory_units', ['event_date'], unique=False,
postgresql_ops={'event_date': 'DESC'})
op.create_index('idx_memory_units_fact_type', 'memory_units', ['fact_type'], unique=False)
op.create_index('idx_memory_units_opinion_confidence', 'memory_units', ['agent_id', 'confidence_score'],
unique=False, postgresql_where=sa.text("fact_type = 'opinion'"),
postgresql_ops={'confidence_score': 'DESC'})
op.create_index('idx_memory_units_opinion_date', 'memory_units', ['agent_id', 'event_date'], unique=False,
postgresql_where=sa.text("fact_type = 'opinion'"), postgresql_ops={'event_date': 'DESC'})
op.create_table('memory_links',
sa.Column('from_unit_id', sa.UUID(), nullable=False),
sa.Column('to_unit_id', sa.UUID(), nullable=False),
sa.Column('link_type', sa.Text(), nullable=False),
sa.Column('entity_id', sa.UUID(), nullable=False),
sa.Column('weight', sa.Float(), server_default='1.0', nullable=False),
sa.Column('created_at', postgresql.TIMESTAMP(timezone=True), server_default=sa.text('now()'),
nullable=False),
sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['from_unit_id'], ['memory_units.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['to_unit_id'], ['memory_units.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('from_unit_id', 'to_unit_id', 'link_type', 'entity_id')
)
op.create_index('idx_memory_links_entity', 'memory_links', ['entity_id'], unique=False,
postgresql_where=sa.text('entity_id IS NOT NULL'))
op.create_index('idx_memory_links_from', 'memory_links', ['from_unit_id'], unique=False)
op.create_index('idx_memory_links_from_weight', 'memory_links', ['from_unit_id', 'weight'], unique=False,
postgresql_where=sa.text('weight >= 0.1'), postgresql_ops={'weight': 'DESC'})
op.create_index('idx_memory_links_to', 'memory_links', ['to_unit_id'], unique=False)
op.create_index('idx_memory_links_type', 'memory_links', ['link_type'], unique=False)
op.create_table('unit_entities',
sa.Column('unit_id', sa.UUID(), nullable=False),
sa.Column('entity_id', sa.UUID(), nullable=False),
sa.ForeignKeyConstraint(['entity_id'], ['entities.id'], ondelete='CASCADE'),
sa.ForeignKeyConstraint(['unit_id'], ['memory_units.id'], ondelete='CASCADE'),
sa.PrimaryKeyConstraint('unit_id', 'entity_id')
)
op.create_index('idx_unit_entities_entity', 'unit_entities', ['entity_id'], unique=False)
op.create_index('idx_unit_entities_unit', 'unit_entities', ['unit_id'], unique=False)
# ### end Alembic commands ###
def downgrade() -> None:
"""Downgrade schema."""
# ### commands auto generated by Alembic - please adjust! ###
op.drop_index('idx_unit_entities_unit', table_name='unit_entities')
op.drop_index('idx_unit_entities_entity', table_name='unit_entities')
op.drop_table('unit_entities')
op.drop_index('idx_memory_links_type', table_name='memory_links')
op.drop_index('idx_memory_links_to', table_name='memory_links')
op.drop_index('idx_memory_links_from_weight', table_name='memory_links', postgresql_where=sa.text('weight >= 0.1'),
postgresql_ops={'weight': 'DESC'})
op.drop_index('idx_memory_links_from', table_name='memory_links')
op.drop_index('idx_memory_links_entity', table_name='memory_links',
postgresql_where=sa.text('entity_id IS NOT NULL'))
op.drop_table('memory_links')
op.drop_index('idx_memory_units_opinion_date', table_name='memory_units',
postgresql_where=sa.text("fact_type = 'opinion'"), postgresql_ops={'event_date': 'DESC'})
op.drop_index('idx_memory_units_opinion_confidence', table_name='memory_units',
postgresql_where=sa.text("fact_type = 'opinion'"), postgresql_ops={'confidence_score': 'DESC'})
op.drop_index('idx_memory_units_fact_type', table_name='memory_units')
op.drop_index('idx_memory_units_event_date', table_name='memory_units', postgresql_ops={'event_date': 'DESC'})
op.drop_index('idx_memory_units_embedding', table_name='memory_units', postgresql_using='hnsw',
postgresql_ops={'embedding': 'vector_cosine_ops'})
op.drop_index('idx_memory_units_document_id', table_name='memory_units')
op.drop_index('idx_memory_units_agent_type_date', table_name='memory_units', postgresql_ops={'event_date': 'DESC'})
op.drop_index('idx_memory_units_agent_id', table_name='memory_units')
op.drop_index('idx_memory_units_agent_fact_type', table_name='memory_units')
op.drop_index('idx_memory_units_agent_date', table_name='memory_units', postgresql_ops={'event_date': 'DESC'})
op.drop_index('idx_memory_units_access_count', table_name='memory_units', postgresql_ops={'access_count': 'DESC'})
op.drop_table('memory_units')
op.drop_index('idx_entity_cooccurrences_entity2', table_name='entity_cooccurrences')
op.drop_index('idx_entity_cooccurrences_entity1', table_name='entity_cooccurrences')
op.drop_index('idx_entity_cooccurrences_count', table_name='entity_cooccurrences',
postgresql_ops={'cooccurrence_count': 'DESC'})
op.drop_table('entity_cooccurrences')
op.drop_index('idx_entities_type', table_name='entities')
op.drop_index('idx_entities_canonical_name', table_name='entities')
op.drop_index('idx_entities_agent_name_type', table_name='entities')
op.drop_index('idx_entities_agent_id', table_name='entities')
op.drop_table('entities')
op.drop_index('idx_documents_content_hash', table_name='documents')
op.drop_index('idx_documents_agent_id', table_name='documents')
op.drop_table('documents')
# ### end Alembic commands ###

54340
benchmarks/locomo/benchmark_results.json
View file

File diff suppressed because it is too large Load diff

6
benchmarks/locomo/results_table.md
View file

@ -1,8 +1,8 @@
# LoComo Benchmark Results
**Overall Accuracy**: 66.00% (66/100)
**Overall Accuracy**: 62.00% (62/100)
| Sample ID | Sessions | Questions | Correct | Accuracy | Multi-hop | Single-hop | Temporal | Open-domain |
|-----------|----------|-----------|---------|----------|-----------|------------|----------|-------------|
| conv-26 | 19 | 50 | 35 | 70.00% | N/A | N/A | N/A | N/A |
| conv-30 | 19 | 50 | 31 | 62.00% | N/A | N/A | N/A | N/A |
| conv-26 | 19 | 50 | 30 | 60.00% | N/A | N/A | N/A | N/A |
| conv-30 | 19 | 50 | 32 | 64.00% | N/A | N/A | N/A | N/A |

3
benchmarks/locomo/run_benchmark.py
View file

@ -37,6 +37,7 @@ async def run_benchmark(
answer_generator = LoComoAnswerGenerator()
answer_evaluator = LoComoAnswerEvaluator()
memory = TemporalSemanticMemory()
await memory.initialize()
# Create benchmark runner
runner = BenchmarkRunner(
@ -130,7 +131,7 @@ def generate_markdown_table(results: dict):
if __name__ == "__main__":
import logging
logging.basicConfig(level=logging.INFO)
logging.basicConfig(level=logging.INFO, format='%(asctime)s %(levelname)s %(message)s')
parser = argparse.ArgumentParser(description='Run LoComo benchmark')
parser.add_argument('--max-conversations', type=int, default=None, help='Maximum conversations to evaluate')

94
examples/library_usage_example.py Normal file
View file

@ -0,0 +1,94 @@
"""
Example of using memory-poc as a library in another project.
This demonstrates how to:
1. Import and use the memory system directly
2. Import and extend the FastAPI app
3. Mount the memory app as a sub-application
"""
import asyncio
from web import app, memory
async def example_memory_usage():
"""Example of using the memory system directly."""
# Initialize memory system
await memory.initialize()
try:
# Store some memories
await memory.put_async(
agent_id="example_agent",
content="Example memory content",
context="test context"
)
# Search for memories
results, trace = await memory.search_async(
agent_id="example_agent",
query="example",
top_k=5
)
print(f"Found {len(results)} results")
for result in results:
print(f" - {result['text']} (score: {result['score']:.4f})")
# Use think functionality
think_result = await memory.think_async(
agent_id="example_agent",
query="What do you know?",
thinking_budget=50
)
print(f"\nThink result: {think_result['text']}")
if think_result.get('new_opinions'):
print(f"New opinions formed: {len(think_result['new_opinions'])}")
finally:
# Clean up
await memory.close()
def example_fastapi_extension():
"""Example of extending the FastAPI app with custom endpoints."""
from fastapi import FastAPI
# Option 1: Add endpoints directly to the imported app
@app.get("/api/custom")
async def custom_endpoint():
return {"message": "Custom endpoint added to memory-poc app"}
# Option 2: Mount as sub-application
main_app = FastAPI(title="My Application")
@main_app.get("/")
async def root():
return {"message": "My main application"}
# Mount the memory app at /memory
main_app.mount("/memory", app)
# Now you can access:
# - / -> your main app
# - /memory/ -> memory visualization
# - /memory/api/graph -> memory graph API
# - /memory/api/search -> memory search API
return main_app
if __name__ == "__main__":
# Example 1: Use memory system directly
print("=" * 80)
print("Example 1: Direct memory system usage")
print("=" * 80)
asyncio.run(example_memory_usage())
# Example 2: FastAPI extension
print("\n" + "=" * 80)
print("Example 2: FastAPI app extension")
print("=" * 80)
extended_app = example_fastapi_extension()
print("FastAPI app extended successfully")
print("To run: uvicorn library_usage_example:extended_app --reload")

180
examples/trace_example.py
View file

@ -1,180 +0,0 @@
"""
Example demonstrating search tracing functionality.
This script shows how to:
1. Enable search tracing
2. Retrieve the trace object
3. Export trace to JSON for visualization
"""
import asyncio
import json
from datetime import datetime, timezone
from memory import TemporalSemanticMemory
async def main():
"""Run the trace example."""
# Initialize memory system
memory = TemporalSemanticMemory()
try:
# Create a test agent
agent_id = f"trace_demo_{datetime.now(timezone.utc).timestamp()}"
print("=" * 70)
print("SEARCH TRACE EXAMPLE")
print("=" * 70)
# Store some test memories
print("\n1. Storing test memories...")
await memory.put_async(
agent_id=agent_id,
content="Alice works at Google as a software engineer in Mountain View",
context="conversation",
)
await memory.put_async(
agent_id=agent_id,
content="Bob also works at Google but in the New York office",
context="conversation",
)
await memory.put_async(
agent_id=agent_id,
content="Charlie founded TechCorp, a startup in San Francisco",
context="conversation",
)
await memory.put_async(
agent_id=agent_id,
content="Alice and Bob met at a Google conference last year",
context="conversation",
)
print(" ✓ 4 memories stored")
# Perform search with tracing enabled
print("\n2. Searching with trace enabled...")
query = "Who works at Google?"
results, trace = await memory.search_async(
agent_id=agent_id,
query=query,
thinking_budget=30,
top_k=5,
enable_trace=True,
)
print(f" ✓ Search completed")
# Display trace summary
print("\n3. Trace Summary:")
print(f" - Query: {trace.query.query_text}")
print(f" - Thinking budget: {trace.query.thinking_budget}")
print(f" - Entry points found: {len(trace.entry_points)}")
print(f" - Total nodes visited: {trace.summary.total_nodes_visited}")
print(f" - Total nodes pruned: {trace.summary.total_nodes_pruned}")
print(f" - Budget used: {trace.summary.budget_used}")
print(f" - Budget remaining: {trace.summary.budget_remaining}")
print(f" - Results returned: {trace.summary.results_returned}")
print(f" - Total duration: {trace.summary.total_duration_seconds:.3f}s")
print(f" - Temporal links followed: {trace.summary.temporal_links_followed}")
print(f" - Semantic links followed: {trace.summary.semantic_links_followed}")
print(f" - Entity links followed: {trace.summary.entity_links_followed}")
# Show entry points
print("\n4. Entry Points:")
for ep in trace.entry_points:
print(f" [{ep.rank}] {ep.text[:60]}... (similarity: {ep.similarity_score:.3f})")
# Show visited nodes with their paths
print("\n5. Search Path (First 5 visits):")
for i, visit in enumerate(trace.visits[:5], 1):
indent = " "
if visit.is_entry_point:
print(f"{indent}[{i}] ENTRY POINT: {visit.text[:60]}...")
else:
parent = f"from {visit.parent_node_id[:8]}" if visit.parent_node_id else "?"
link_info = f"via {visit.link_type}" if visit.link_type else ""
print(f"{indent}[{i}] {parent} {link_info}: {visit.text[:60]}...")
print(f"{indent} - Activation: {visit.weights.activation:.3f}")
print(f"{indent} - Semantic sim: {visit.weights.semantic_similarity:.3f}")
print(f"{indent} - Recency: {visit.weights.recency:.3f}")
print(f"{indent} - Final weight: {visit.weights.final_weight:.3f}")
if visit.neighbors_explored:
followed = sum(1 for n in visit.neighbors_explored if n.followed)
pruned = len(visit.neighbors_explored) - followed
print(f"{indent} - Neighbors: {followed} followed, {pruned} pruned")
# Show pruning decisions
if trace.pruned:
print(f"\n6. Pruning Decisions (showing first 5 of {len(trace.pruned)}):")
for prune in trace.pruned[:5]:
print(f" - Node {prune.node_id[:8]}: {prune.reason} (activation: {prune.activation:.3f})")
# Show phase metrics
print("\n7. Phase Metrics:")
for pm in trace.summary.phase_metrics:
print(f" - {pm.phase_name}: {pm.duration_seconds:.3f}s")
if pm.details:
for key, value in pm.details.items():
if isinstance(value, float):
print(f" • {key}: {value:.3f}")
else:
print(f" • {key}: {value}")
# Export to JSON
print("\n8. Exporting trace to JSON...")
trace_json = trace.to_json()
output_file = f"trace_{agent_id}.json"
with open(output_file, "w") as f:
f.write(trace_json)
print(f" ✓ Trace saved to: {output_file}")
print(f" ✓ JSON size: {len(trace_json):,} bytes")
# Show search results
print("\n9. Search Results:")
for i, result in enumerate(results, 1):
print(f" [{i}] {result['text'][:70]}...")
print(f" Weight: {result['weight']:.3f} "
f"(act: {result['activation']:.2f}, "
f"sem: {result['semantic_similarity']:.2f}, "
f"rec: {result['recency']:.2f})")
# Test helper methods
print("\n10. Testing Helper Methods:")
# Get path to first result
if results:
first_result_id = results[0]['id']
path = trace.get_search_path_to_node(first_result_id)
print(f" - Path to top result has {len(path)} steps")
# Count nodes by link type
temporal_nodes = trace.get_nodes_by_link_type("temporal")
semantic_nodes = trace.get_nodes_by_link_type("semantic")
entity_nodes = trace.get_nodes_by_link_type("entity")
print(f" - Nodes reached via temporal links: {len(temporal_nodes)}")
print(f" - Nodes reached via semantic links: {len(semantic_nodes)}")
print(f" - Nodes reached via entity links: {len(entity_nodes)}")
print("\n" + "=" * 70)
print("TRACE EXAMPLE COMPLETE!")
print("=" * 70)
print(f"\nYou can now build a visualization using the trace data in:")
print(f" {output_file}")
print("\nThe trace contains:")
print(f" - Complete search path with all nodes visited")
print(f" - Weight calculations for each node")
print(f" - Link information (type, weight, whether followed)")
print(f" - Pruning decisions with reasons")
print(f" - Performance metrics for each phase")
# Cleanup
print("\nCleaning up test agent...")
await memory.delete_agent(agent_id)
finally:
await memory.close()
if __name__ == "__main__":
asyncio.run(main())

2
memory/__init__.py
View file

@ -4,7 +4,6 @@ Memory System for AI Agents.
Temporal + Semantic Memory Architecture using PostgreSQL with pgvector.
"""
from .temporal_semantic_memory import TemporalSemanticMemory
from .visualizer import MemoryVisualizer
from .search_trace import (
SearchTrace,
QueryInfo,
@ -20,7 +19,6 @@ from .search_tracer import SearchTracer
__all__ = [
"TemporalSemanticMemory",
"MemoryVisualizer",
"SearchTrace",
"SearchTracer",
"QueryInfo",

115
memory/entity_resolver.py
View file

@ -4,7 +4,6 @@ Entity extraction and resolution for memory system.
Uses spaCy for entity extraction and implements resolution logic
to disambiguate entities across memory units.
"""
import spacy
import asyncpg
from typing import List, Dict, Optional, Set
from difflib import SequenceMatcher
@ -15,78 +14,6 @@ from datetime import datetime, timezone
_nlp = None
def get_nlp():
"""Get or load spaCy model."""
global _nlp
if _nlp is None:
_nlp = spacy.load("en_core_web_sm")
return _nlp
def extract_entities(text: str) -> List[Dict[str, any]]:
"""
Extract entities from text using spaCy.
Args:
text: Input text
Returns:
List of entities with text, type, and span info
"""
nlp = get_nlp()
doc = nlp(text)
entities = []
for ent in doc.ents:
# Filter to important entity types
if ent.label_ in ['PERSON', 'ORG', 'GPE', 'LOC', 'PRODUCT', 'EVENT']:
entities.append({
'text': ent.text,
'type': ent.label_,
'start': ent.start_char,
'end': ent.end_char,
})
return entities
def extract_entities_batch(texts: List[str]) -> List[List[Dict[str, any]]]:
"""
Extract entities from multiple texts in batch (MUCH faster than sequential).
Uses spaCy's nlp.pipe() for efficient batch processing.
Args:
texts: List of input texts
Returns:
List of entity lists, one per input text
"""
if not texts:
return []
nlp = get_nlp()
# Process all texts in batch using nlp.pipe (significantly faster!)
docs = list(nlp.pipe(texts, batch_size=50))
all_entities = []
for doc in docs:
entities = []
for ent in doc.ents:
# Filter to important entity types
if ent.label_ in ['PERSON', 'ORG', 'GPE', 'LOC', 'PRODUCT', 'EVENT']:
entities.append({
'text': ent.text,
'type': ent.label_,
'start': ent.start_char,
'end': ent.end_char,
})
all_entities.append(entities)
return all_entities
class EntityResolver:
"""
Resolves entities to canonical IDs with disambiguation.
@ -248,14 +175,44 @@ class EntityResolver:
entities_to_update
)
# Batch create new entities
# Batch create new entities using multi-row VALUES
if entities_to_create:
import logging
logger = logging.getLogger(__name__)
create_start = time.time()
# 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:
entity_id = await self._create_entity(
conn, agent_id, entity_data['text'],
entity_data['type'], unit_event_date
)
entity_ids[idx] = entity_id
values_clauses.append(f"(${param_idx}, ${param_idx+1}, ${param_idx+2}, ${param_idx+3}, ${param_idx+4}, ${param_idx+5})")
params.extend([
agent_id,
entity_data['text'],
entity_data['type'],
unit_event_date,
unit_event_date,
1
])
param_idx += 6
# Single INSERT with multiple VALUES rows
query = f"""
INSERT INTO entities (agent_id, canonical_name, entity_type, 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']
logger.info(f" [6.2.2.X] Batch created {len(entities_to_create)} new entities in {time.time() - create_start:.3f}s")
return entity_ids

62
memory/llm_client.py
View file

@ -31,6 +31,9 @@ class ExtractedFact(BaseModel):
date: str = Field(
description="Absolute date/time when this fact occurred in ISO format (YYYY-MM-DDTHH:MM:SSZ). If text mentions relative time (yesterday, last week, this morning), calculate absolute date from the provided context date."
)
fact_type: Literal["world", "agent", "opinion"] = Field(
description="Type of fact: 'world' for general facts about the world (events, people, things that happen), 'agent' for facts about what the AI agent specifically did or actions the agent took (conversations with the user, tasks performed by the agent), 'opinion' for the agent's formed opinions and perspectives"
)
entities: List[Entity] = Field(
default_factory=list,
description="List of important entities mentioned in this fact with their types"
@ -212,6 +215,21 @@ Examples of date extraction:
- Pure reactions without content ("wow", "cool", "nice")
- Incomplete thoughts or sentence fragments with no meaning
## FACT TYPE CLASSIFICATION (CRITICAL):
For EACH fact, classify it as either 'world' or 'agent':
- **'world'**: General facts about the world, events, people, things that happen
- Examples: "Alice works at Google", "Bob went hiking in Yosemite", "The meeting is scheduled for Monday"
- Most facts will be 'world' type
- **'agent'**: Facts specifically about what the AI agent did or actions the agent took
- Examples: "The AI agent helped the user debug their code", "The agent answered a question about Python", "The agent created a new file"
- ONLY use 'agent' if the fact is explicitly about the AI agent's actions
- Conversations with the user where the agent participated are 'agent' type
- Tasks performed BY the agent are 'agent' type
When in doubt, classify as 'world'.
## ENTITY EXTRACTION (CRITICAL):
For EACH fact, extract ALL important entities mentioned with their types:
- **PERSON**: Names of individuals (Alice, Bob, Dr. Smith)
@ -232,6 +250,7 @@ Entity extraction rules:
Input: "Alice mentioned she works at Google in Mountain View. She joined the AI team last year."
GOOD fact: "Alice works at Google in Mountain View on the AI team, which she joined last year"
GOOD fact_type: "world"
GOOD date: Calculate based on reference date (if reference is 2024-03-20, "last year" = 2023-03-20)
GOOD entities: [
{{"text": "Alice", "type": "PERSON"}},
@ -242,6 +261,7 @@ GOOD entities: [
Input: "Yesterday Bob went hiking in Yosemite because it helps him clear his mind."
GOOD fact: "Bob went hiking in Yosemite because it helps him clear his mind"
GOOD fact_type: "world"
GOOD date: Reference date minus 1 day
GOOD entities: [
{{"text": "Bob", "type": "PERSON"}},
@ -256,11 +276,21 @@ NOTE: Extract the event (photo taken/shared with friends at beach), NOT just tha
Input: "I sent you that article about AI last Tuesday."
GOOD fact: "Someone sent an article about AI"
GOOD fact_type: "world"
GOOD date: Calculate last Tuesday from reference date
GOOD entities: [
{{"text": "AI", "type": "CONCEPT"}}
]
Input: "The AI agent helped me write a Python script to analyze my data."
GOOD fact: "The AI agent helped someone write a Python script to analyze their data"
GOOD fact_type: "agent"
GOOD date: Reference date (no specific time mentioned)
GOOD entities: [
{{"text": "Python", "type": "PRODUCT"}},
{{"text": "data analysis", "type": "CONCEPT"}}
]
Input: "I bought an Apple laptop and some apples from the store."
GOOD fact: "Someone bought an Apple laptop and some apples from the store"
GOOD entities: [
@ -307,10 +337,17 @@ Remember:
5. **Extract biographical details as SEPARATE facts** - "my home country Sweden" should create a fact "Person is from Sweden"
6. Include ALL details, names, numbers, reasons, and context in the fact text
7. Extract the absolute date for EACH fact by calculating relative times from the reference date
8. Extract ALL entities with their types (PERSON, ORG, PLACE, PRODUCT, CONCEPT) for each fact
9. Use types to disambiguate entities (Apple the company = ORG, apple the fruit = PRODUCT)
10. When in doubt, EXTRACT IT - better to have too many facts than miss important events"""
8. **CLASSIFY EACH FACT**: 'world' for general facts, 'agent' for AI agent actions
9. Extract ALL entities with their types (PERSON, ORG, PLACE, PRODUCT, CONCEPT) for each fact
10. Use types to disambiguate entities (Apple the company = ORG, apple the fruit = PRODUCT)
11. When in doubt, EXTRACT IT - better to have too many facts than miss important events"""
import time
import logging
logger = logging.getLogger(__name__)
llm_call_start = time.time()
response = await client.beta.chat.completions.parse(
model=model,
messages=[
@ -328,6 +365,7 @@ Remember:
response_format=FactExtractionResponse,
extra_body={"service_tier": "auto"},
)
llm_call_time = time.time() - llm_call_start
# Extract the parsed response
extraction_response = response.choices[0].message.parsed
@ -335,6 +373,8 @@ Remember:
# Convert to dict format
chunk_facts = [fact.model_dump() for fact in extraction_response.facts]
logger.info(f" [1.3.{chunk_index + 1}] Chunk {chunk_index + 1}/{total_chunks} LLM call: {len(chunk_facts)} facts from {len(chunk)} chars in {llm_call_time:.3f}s")
return chunk_facts
@ -365,12 +405,21 @@ async def extract_facts_from_text(
Returns:
List of fact dictionaries with 'fact' and 'date' keys
"""
import time
import logging
logger = logging.getLogger(__name__)
client = get_llm_client()
# Chunk text if necessary
chunk_start = time.time()
chunks = chunk_text(text, max_chars=chunk_size)
chunk_time = time.time() - chunk_start
logger.info(f" [1.1] Text chunking: {len(chunks)} chunks from {len(text)} chars in {chunk_time:.3f}s")
# Process all chunks in parallel using asyncio.gather
task_creation_start = time.time()
tasks = [
_extract_facts_from_chunk(
chunk=chunk,
@ -385,13 +434,20 @@ async def extract_facts_from_text(
)
for i, chunk in enumerate(chunks)
]
logger.info(f" [1.2] Task creation: {len(tasks)} tasks in {time.time() - task_creation_start:.3f}s")
# Wait for all chunks to complete in parallel
llm_start = time.time()
chunk_results = await asyncio.gather(*tasks)
llm_time = time.time() - llm_start
logger.info(f" [1.3] LLM extraction (parallel): {len(chunks)} chunks in {llm_time:.3f}s")
# Flatten results from all chunks
flatten_start = time.time()
all_facts = []
for chunk_facts in chunk_results:
all_facts.extend(chunk_facts)
flatten_time = time.time() - flatten_start
logger.info(f" [1.4] Result flattening: {len(all_facts)} facts in {flatten_time:.3f}s")
return all_facts

268
memory/models.py Normal file
View file

@ -0,0 +1,268 @@
"""
SQLAlchemy models for the memory system.
"""
from datetime import datetime
from typing import Optional
from uuid import UUID as PyUUID, uuid4
from sqlalchemy import (
CheckConstraint,
Column,
Float,
ForeignKey,
ForeignKeyConstraint,
Index,
Integer,
PrimaryKeyConstraint,
Text,
func,
text as sql_text,
)
from sqlalchemy.dialects.postgresql import JSONB, TIMESTAMP, UUID
from sqlalchemy.ext.asyncio import AsyncAttrs
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column, relationship
from pgvector.sqlalchemy import Vector
class Base(AsyncAttrs, DeclarativeBase):
"""Base class for all models."""
pass
class Document(Base):
"""Source documents for memory units."""
__tablename__ = "documents"
id: Mapped[str] = mapped_column(Text, primary_key=True)
agent_id: Mapped[str] = mapped_column(Text, primary_key=True)
original_text: Mapped[Optional[str]] = mapped_column(Text)
content_hash: Mapped[Optional[str]] = mapped_column(Text)
doc_metadata: Mapped[dict] = mapped_column("metadata", JSONB, server_default=sql_text("'{}'::jsonb"))
created_at: Mapped[datetime] = mapped_column(
TIMESTAMP(timezone=True), server_default=func.now()
)
updated_at: Mapped[datetime] = mapped_column(
TIMESTAMP(timezone=True), server_default=func.now()
)
# Relationships
memory_units = relationship("MemoryUnit", back_populates="document", cascade="all, delete-orphan")
__table_args__ = (
Index("idx_documents_agent_id", "agent_id"),
Index("idx_documents_content_hash", "content_hash"),
)
class MemoryUnit(Base):
"""Individual sentence-level memories."""
__tablename__ = "memory_units"
id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), primary_key=True, server_default=sql_text("uuid_generate_v4()")
)
agent_id: Mapped[str] = mapped_column(Text, nullable=False)
document_id: Mapped[Optional[str]] = mapped_column(Text)
text: Mapped[str] = mapped_column(Text, nullable=False)
embedding = mapped_column(Vector(384)) # pgvector type
context: Mapped[Optional[str]] = mapped_column(Text)
event_date: Mapped[datetime] = mapped_column(TIMESTAMP(timezone=True), nullable=False)
fact_type: Mapped[str] = mapped_column(Text, nullable=False, server_default="world")
confidence_score: Mapped[Optional[float]] = mapped_column(Float)
access_count: Mapped[int] = mapped_column(Integer, server_default="0")
created_at: Mapped[datetime] = mapped_column(
TIMESTAMP(timezone=True), server_default=func.now()
)
updated_at: Mapped[datetime] = mapped_column(
TIMESTAMP(timezone=True), server_default=func.now()
)
# Relationships
document = relationship("Document", back_populates="memory_units")
unit_entities = relationship("UnitEntity", back_populates="memory_unit", cascade="all, delete-orphan")
outgoing_links = relationship(
"MemoryLink",
foreign_keys="MemoryLink.from_unit_id",
back_populates="from_unit",
cascade="all, delete-orphan"
)
incoming_links = relationship(
"MemoryLink",
foreign_keys="MemoryLink.to_unit_id",
back_populates="to_unit",
cascade="all, delete-orphan"
)
__table_args__ = (
ForeignKeyConstraint(
["document_id", "agent_id"],
["documents.id", "documents.agent_id"],
name="memory_units_document_fkey",
ondelete="CASCADE",
),
CheckConstraint("fact_type IN ('world', 'agent', 'opinion')"),
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)",
name="confidence_score_fact_type_check"
),
Index("idx_memory_units_agent_id", "agent_id"),
Index("idx_memory_units_document_id", "document_id"),
Index("idx_memory_units_event_date", "event_date", postgresql_ops={"event_date": "DESC"}),
Index("idx_memory_units_agent_date", "agent_id", "event_date", postgresql_ops={"event_date": "DESC"}),
Index("idx_memory_units_access_count", "access_count", postgresql_ops={"access_count": "DESC"}),
Index("idx_memory_units_fact_type", "fact_type"),
Index("idx_memory_units_agent_fact_type", "agent_id", "fact_type"),
Index("idx_memory_units_agent_type_date", "agent_id", "fact_type", "event_date", postgresql_ops={"event_date": "DESC"}),
Index(
"idx_memory_units_opinion_confidence",
"agent_id",
"confidence_score",
postgresql_where=sql_text("fact_type = 'opinion'"),
postgresql_ops={"confidence_score": "DESC"}
),
Index(
"idx_memory_units_opinion_date",
"agent_id",
"event_date",
postgresql_where=sql_text("fact_type = 'opinion'"),
postgresql_ops={"event_date": "DESC"}
),
Index(
"idx_memory_units_embedding",
"embedding",
postgresql_using="hnsw",
postgresql_ops={"embedding": "vector_cosine_ops"}
),
)
class Entity(Base):
"""Resolved entities (people, organizations, locations, etc.)."""
__tablename__ = "entities"
id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), primary_key=True, server_default=sql_text("uuid_generate_v4()")
)
canonical_name: Mapped[str] = mapped_column(Text, nullable=False)
entity_type: Mapped[str] = mapped_column(Text, nullable=False)
agent_id: Mapped[str] = mapped_column(Text, nullable=False)
entity_metadata: Mapped[dict] = mapped_column("metadata", JSONB, server_default=sql_text("'{}'::jsonb"))
first_seen: Mapped[datetime] = mapped_column(
TIMESTAMP(timezone=True), server_default=func.now()
)
last_seen: Mapped[datetime] = mapped_column(
TIMESTAMP(timezone=True), server_default=func.now()
)
mention_count: Mapped[int] = mapped_column(Integer, server_default="1")
# Relationships
unit_entities = relationship("UnitEntity", back_populates="entity", cascade="all, delete-orphan")
memory_links = relationship("MemoryLink", back_populates="entity", cascade="all, delete-orphan")
cooccurrences_1 = relationship(
"EntityCooccurrence",
foreign_keys="EntityCooccurrence.entity_id_1",
back_populates="entity_1",
cascade="all, delete-orphan"
)
cooccurrences_2 = relationship(
"EntityCooccurrence",
foreign_keys="EntityCooccurrence.entity_id_2",
back_populates="entity_2",
cascade="all, delete-orphan"
)
__table_args__ = (
Index("idx_entities_agent_id", "agent_id"),
Index("idx_entities_canonical_name", "canonical_name"),
Index("idx_entities_type", "entity_type"),
Index("idx_entities_agent_name_type", "agent_id", "canonical_name", "entity_type"),
)
class UnitEntity(Base):
"""Association between memory units and entities."""
__tablename__ = "unit_entities"
unit_id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("memory_units.id", ondelete="CASCADE"), primary_key=True
)
entity_id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
)
# Relationships
memory_unit = relationship("MemoryUnit", back_populates="unit_entities")
entity = relationship("Entity", back_populates="unit_entities")
__table_args__ = (
Index("idx_unit_entities_unit", "unit_id"),
Index("idx_unit_entities_entity", "entity_id"),
)
class EntityCooccurrence(Base):
"""Materialized cache of entity co-occurrences."""
__tablename__ = "entity_cooccurrences"
entity_id_1: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
)
entity_id_2: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
)
cooccurrence_count: Mapped[int] = mapped_column(Integer, server_default="1")
last_cooccurred: Mapped[datetime] = mapped_column(
TIMESTAMP(timezone=True), server_default=func.now()
)
# Relationships
entity_1 = relationship("Entity", foreign_keys=[entity_id_1], back_populates="cooccurrences_1")
entity_2 = relationship("Entity", foreign_keys=[entity_id_2], back_populates="cooccurrences_2")
__table_args__ = (
CheckConstraint("entity_id_1 < entity_id_2", name="entity_cooccurrence_order_check"),
Index("idx_entity_cooccurrences_entity1", "entity_id_1"),
Index("idx_entity_cooccurrences_entity2", "entity_id_2"),
Index("idx_entity_cooccurrences_count", "cooccurrence_count", postgresql_ops={"cooccurrence_count": "DESC"}),
)
class MemoryLink(Base):
"""Links between memory units (temporal, semantic, entity)."""
__tablename__ = "memory_links"
from_unit_id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("memory_units.id", ondelete="CASCADE"), primary_key=True
)
to_unit_id: Mapped[PyUUID] = mapped_column(
UUID(as_uuid=True), ForeignKey("memory_units.id", ondelete="CASCADE"), primary_key=True
)
link_type: Mapped[str] = mapped_column(Text, primary_key=True)
entity_id: Mapped[Optional[PyUUID]] = mapped_column(
UUID(as_uuid=True), ForeignKey("entities.id", ondelete="CASCADE"), primary_key=True
)
weight: Mapped[float] = mapped_column(Float, nullable=False, server_default="1.0")
created_at: Mapped[datetime] = mapped_column(
TIMESTAMP(timezone=True), server_default=func.now()
)
# Relationships
from_unit = relationship("MemoryUnit", foreign_keys=[from_unit_id], back_populates="outgoing_links")
to_unit = relationship("MemoryUnit", foreign_keys=[to_unit_id], back_populates="incoming_links")
entity = relationship("Entity", back_populates="memory_links")
__table_args__ = (
Index("idx_memory_links_from", "from_unit_id"),
Index("idx_memory_links_to", "to_unit_id"),
Index("idx_memory_links_type", "link_type"),
Index("idx_memory_links_entity", "entity_id", postgresql_where=sql_text("entity_id IS NOT NULL")),
Index(
"idx_memory_links_from_weight",
"from_unit_id",
"weight",
postgresql_where=sql_text("weight >= 0.1"),
postgresql_ops={"weight": "DESC"}
),
)

103
memory/operations/README.md Normal file
View file

@ -0,0 +1,103 @@
# Memory Operations Modules
This directory contains specialized operation modules for the TemporalSemanticMemory class.
## Refactoring Results
✅ **Successfully Completed!**
**File Size Reduction:**
- Before: 1,720 lines (temporal_semantic_memory.py)
- After: 1,420 lines (temporal_semantic_memory.py)
- **Removed: 300 lines (17% reduction)**
**Modules Created:**
- `embedding_operations.py` - Embedding generation with process pool parallelism
- `link_operations.py` - Entity, temporal, and semantic link creation (300+ lines)
- `batch_operations.py` - Placeholder for future extraction
- `search_operations.py` - Placeholder for future extraction
## Architecture
The memory system now uses a **mixin pattern** for better code organization:
```python
class TemporalSemanticMemory(
EmbeddingOperationsMixin,
LinkOperationsMixin,
):
"""
Advanced memory system using temporal and semantic linking.
Mixins provide:
- EmbeddingOperationsMixin: _generate_embedding, _generate_embeddings_batch
- LinkOperationsMixin: Entity, temporal, semantic link operations
"""
# Core infrastructure and batch operations
pass
```
## What Was Extracted
### EmbeddingOperationsMixin (embedding_operations.py)
- `_generate_embedding()` - Single embedding generation
- `_generate_embeddings_batch()` - Parallel batch embedding generation
- Process pool worker functions for CPU parallelism
### LinkOperationsMixin (link_operations.py)
- `_extract_entities_batch_optimized()` - Entity resolution and linking
- `_create_temporal_links_batch_per_fact()` - Time-based connections
- `_create_semantic_links_batch()` - Meaning-based connections
- `_insert_entity_links_batch()` - Batch link insertion
### Remaining in Main Class
- Database connection management (`__init__`, `_get_pool`, `close`)
- Batch storage operations (`put`, `put_async`, `put_batch_async`)
- Search operations (`search`, `search_async`, `_apply_mmr`)
- Document management (`get_document`, `delete_document`, `delete_agent`)
- Think operations (`think_async`)
- Deduplication (`_find_duplicate_facts_batch`)
## Benefits Achieved
1. ✅ **Better Organization** - Related methods grouped in focused modules
2. ✅ **Reduced Complexity** - Main file is 17% smaller
3. ✅ **Reusability** - Mixins can be composed and tested independently
4. ✅ **Maintainability** - Easier to find and modify specific operations
5. ✅ **All Tests Pass** - No breaking changes to public API
## Usage
The public API remains unchanged:
```python
from memory import TemporalSemanticMemory
memory = TemporalSemanticMemory()
# All operations work exactly as before
result = await memory.think_async(
agent_id="agent_1",
query="What have I done?"
)
results, trace = await memory.search_async(
agent_id="agent_1",
query="example query",
fact_type="world"
)
```
## Future Work (Optional)
The foundation is now in place for further extraction:
- Extract batch operations to `batch_operations.py`
- Extract search operations to `search_operations.py`
- Split large methods into smaller, focused functions
## Design Principles Followed
1. ✅ **Preserved batch mechanisms** - Performance maintained
2. ✅ **No breaking changes** - All tests pass
3. ✅ **Clear separation** - Each mixin has focused responsibility
4. ✅ **Gradual refactoring** - Can continue incrementally

13
memory/operations/__init__.py Normal file
View file

@ -0,0 +1,13 @@
"""
Memory operations modules.
This package contains specialized operation modules for the TemporalSemanticMemory class.
"""
from .embedding_operations import EmbeddingOperationsMixin
from .link_operations import LinkOperationsMixin
__all__ = [
'EmbeddingOperationsMixin',
'LinkOperationsMixin',
]

19
memory/operations/batch_operations.py Normal file
View file

@ -0,0 +1,19 @@
"""
Batch operations for storing memories.
NOTE: This is a placeholder for future refactoring.
The actual implementation is currently in temporal_semantic_memory.py
"""
class BatchOperationsMixin:
"""
Mixin class for batch operations.
Methods to be extracted:
- put
- put_async
- put_batch_async
- _find_duplicate_facts_batch
"""
pass

88
memory/operations/embedding_operations.py Normal file
View file

@ -0,0 +1,88 @@
"""
Embedding generation operations for memory units.
"""
import asyncio
import logging
from typing import List
from concurrent.futures import ProcessPoolExecutor
logger = logging.getLogger(__name__)
# Global process pool for parallel embedding generation
_PROCESS_POOL = None
def _get_worker_model():
"""Get or load the embedding model in worker process."""
from sentence_transformers import SentenceTransformer
global _worker_model
if '_worker_model' not in globals():
globals()['_worker_model'] = SentenceTransformer("BAAI/bge-small-en-v1.5")
return globals()['_worker_model']
def _encode_batch_worker(texts: List[str]) -> List[List[float]]:
"""
Worker function for process pool - encodes texts to embeddings.
This function runs in a separate process and loads its own model.
"""
model = _get_worker_model()
embeddings = model.encode(texts, convert_to_numpy=True, show_progress_bar=False)
return [emb.tolist() for emb in embeddings]
def _get_process_pool():
"""Get or create the global process pool."""
global _PROCESS_POOL
if _PROCESS_POOL is None:
# Use 4 worker processes for true parallelism
_PROCESS_POOL = ProcessPoolExecutor(max_workers=4)
return _PROCESS_POOL
class EmbeddingOperationsMixin:
"""Mixin class for embedding operations."""
def _generate_embedding(self, text: str) -> List[float]:
"""
Generate embedding for text using local SentenceTransformer model.
Args:
text: Text to embed
Returns:
384-dimensional embedding vector (bge-small-en-v1.5)
"""
try:
embedding = self.embedding_model.encode(text, convert_to_numpy=True, show_progress_bar=False)
return embedding.tolist()
except Exception as e:
raise Exception(f"Failed to generate embedding: {str(e)}")
async def _generate_embeddings_batch(self, texts: List[str]) -> List[List[float]]:
"""
Generate embeddings for multiple texts using local model in parallel.
Uses a ProcessPoolExecutor to achieve TRUE parallelism for CPU-bound
embedding generation. Each worker process loads its own model copy.
Args:
texts: List of texts to embed
Returns:
List of 384-dimensional embeddings in same order as input texts
"""
try:
# Run in process pool for true parallelism
loop = asyncio.get_event_loop()
pool = _get_process_pool()
embeddings = await loop.run_in_executor(
pool,
_encode_batch_worker,
texts
)
return embeddings
except Exception as e:
raise Exception(f"Failed to generate batch embeddings: {str(e)}")

400
memory/operations/link_operations.py Normal file
View file

@ -0,0 +1,400 @@
"""
Link creation operations for temporal, semantic, and entity links.
"""
import time
import logging
from typing import List
from datetime import timedelta
logger = logging.getLogger(__name__)
class LinkOperationsMixin:
"""Mixin class for link creation operations."""
async def _extract_entities_batch_optimized(
self,
conn,
agent_id: str,
unit_ids: List[str],
sentences: List[str],
context: str,
fact_dates: List,
llm_entities: List[List[dict]],
) -> List[tuple]:
"""
Process LLM-extracted entities for ALL facts in batch.
Uses entities provided by the LLM (no spaCy needed), then resolves
and links them in bulk.
Returns list of tuples for batch insertion: (from_unit_id, to_unit_id, link_type, weight, entity_id)
"""
try:
# Step 1: Convert LLM entities to the format expected by entity resolver
substep_start = time.time()
all_entities = []
for entity_list in llm_entities:
# Convert List[Entity] or List[dict] to List[Dict] format
formatted_entities = []
for ent in entity_list:
# Handle both Entity objects and dicts
if hasattr(ent, 'text'):
formatted_entities.append({'text': ent.text, 'type': ent.type})
elif isinstance(ent, dict):
formatted_entities.append({'text': ent.get('text', ''), 'type': ent.get('type', 'CONCEPT')})
all_entities.append(formatted_entities)
total_entities = sum(len(ents) for ents in all_entities)
logger.info(f" [6.1] Process LLM entities: {total_entities} entities from {len(sentences)} facts in {time.time() - substep_start:.3f}s")
# Step 2: Resolve entities in BATCH (much faster!)
substep_start = time.time()
step_6_2_start = time.time()
# [6.2.1] Prepare all entities for batch resolution
substep_6_2_1_start = time.time()
all_entities_flat = []
entity_to_unit = [] # Maps flat index to (unit_id, local_index)
for unit_id, entities, fact_date in zip(unit_ids, all_entities, fact_dates):
if not entities:
continue
for local_idx, entity in enumerate(entities):
all_entities_flat.append({
'text': entity['text'],
'type': entity['type'],
'nearby_entities': entities,
})
entity_to_unit.append((unit_id, local_idx, fact_date))
logger.info(f" [6.2.1] Prepare entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_1_start:.3f}s")
# Resolve ALL entities in one batch call
if all_entities_flat:
# [6.2.2] Batch resolve entities
substep_6_2_2_start = time.time()
# Group by date for batch resolution (most will have same date)
entities_by_date = {}
for idx, (unit_id, local_idx, fact_date) in enumerate(entity_to_unit):
date_key = fact_date
if date_key not in entities_by_date:
entities_by_date[date_key] = []
entities_by_date[date_key].append((idx, all_entities_flat[idx]))
logger.info(f" [6.2.2] Grouped into {len(entities_by_date)} date buckets, resolving...")
# Resolve each date group in batch
resolved_entity_ids = [None] * len(all_entities_flat)
for date_idx, (fact_date, entities_group) in enumerate(entities_by_date.items(), 1):
date_bucket_start = time.time()
indices = [idx for idx, _ in entities_group]
entities_data = [entity_data for _, entity_data in entities_group]
batch_resolved = await self.entity_resolver.resolve_entities_batch(
agent_id=agent_id,
entities_data=entities_data,
context=context,
unit_event_date=fact_date,
conn=conn
)
for idx, entity_id in zip(indices, batch_resolved):
resolved_entity_ids[idx] = entity_id
logger.info(f" [6.2.2.{date_idx}] Resolved {len(entities_data)} entities in {time.time() - date_bucket_start:.3f}s")
logger.info(f" [6.2.2] Resolve entities: {len(all_entities_flat)} entities in {time.time() - substep_6_2_2_start:.3f}s")
# [6.2.3] Create unit-entity links in BATCH
substep_6_2_3_start = time.time()
# Map resolved entities back to units and collect all (unit, entity) pairs
unit_to_entity_ids = {}
unit_entity_pairs = []
for idx, (unit_id, local_idx, fact_date) in enumerate(entity_to_unit):
if unit_id not in unit_to_entity_ids:
unit_to_entity_ids[unit_id] = []
entity_id = resolved_entity_ids[idx]
unit_to_entity_ids[unit_id].append(entity_id)
unit_entity_pairs.append((unit_id, entity_id))
# Batch insert all unit-entity links (MUCH faster!)
await self.entity_resolver.link_units_to_entities_batch(unit_entity_pairs, conn=conn)
logger.info(f" [6.2.3] Create unit-entity links (batched): {len(unit_entity_pairs)} links in {time.time() - substep_6_2_3_start:.3f}s")
logger.info(f" [6.2] Entity resolution (batched): {len(all_entities_flat)} entities resolved in {time.time() - step_6_2_start:.3f}s")
else:
unit_to_entity_ids = {}
logger.info(f" [6.2] Entity resolution (batched): 0 entities in {time.time() - step_6_2_start:.3f}s")
# Step 3: Create entity links between units that share entities
substep_start = time.time()
# Collect all unique entity IDs
all_entity_ids = set()
for entity_ids in unit_to_entity_ids.values():
all_entity_ids.update(entity_ids)
logger.info(f" [6.3] Creating entity links for {len(all_entity_ids)} unique entities...")
# Find all units that reference these entities (ONE batched query)
entity_to_units = {}
if all_entity_ids:
query_start = time.time()
import uuid
entity_id_list = [uuid.UUID(eid) if isinstance(eid, str) else eid for eid in all_entity_ids]
rows = await conn.fetch(
"""
SELECT entity_id, unit_id
FROM unit_entities
WHERE entity_id = ANY($1::uuid[])
""",
entity_id_list
)
logger.info(f" [6.3.1] Query unit_entities: {len(rows)} rows in {time.time() - query_start:.3f}s")
# Group by entity_id
group_start = time.time()
for row in rows:
entity_id = row['entity_id']
if entity_id not in entity_to_units:
entity_to_units[entity_id] = []
entity_to_units[entity_id].append(row['unit_id'])
logger.info(f" [6.3.2] Group by entity_id: {time.time() - group_start:.3f}s")
# Create bidirectional links between units that share entities
link_gen_start = time.time()
links = []
for entity_id, units_with_entity in entity_to_units.items():
# For each pair of units with this entity, create bidirectional links
for i, unit_id_1 in enumerate(units_with_entity):
for unit_id_2 in units_with_entity[i+1:]:
# Bidirectional links
links.append((unit_id_1, unit_id_2, 'entity', 1.0, entity_id))
links.append((unit_id_2, unit_id_1, 'entity', 1.0, entity_id))
logger.info(f" [6.3.3] Generate {len(links)} links: {time.time() - link_gen_start:.3f}s")
logger.info(f" [6.3] Entity link creation: {len(links)} links for {len(all_entity_ids)} unique entities in {time.time() - substep_start:.3f}s")
return links
except Exception as e:
logger.error(f"Failed to extract entities in batch: {str(e)}")
import traceback
traceback.print_exc()
raise
async def _create_temporal_links_batch_per_fact(
self,
conn,
agent_id: str,
unit_ids: List[str],
time_window_hours: int = 24,
):
"""
Create temporal links for multiple units, each with their own event_date.
Queries the event_date for each unit from the database and creates temporal
links based on individual dates (supports per-fact dating).
"""
if not unit_ids:
return
try:
import time as time_mod
# Get the event_date for each new unit
fetch_dates_start = time_mod.time()
rows = await conn.fetch(
"""
SELECT id, event_date
FROM memory_units
WHERE id::text = ANY($1)
""",
unit_ids
)
new_units = {str(row['id']): row['event_date'] for row in rows}
logger.info(f" [7.1] Fetch event_dates for {len(unit_ids)} units: {time_mod.time() - fetch_dates_start:.3f}s")
# Fetch ALL potential temporal neighbors in ONE query (much faster!)
# Get time range across all units
all_dates = list(new_units.values())
min_date = min(all_dates) - timedelta(hours=time_window_hours)
max_date = max(all_dates) + timedelta(hours=time_window_hours)
fetch_neighbors_start = time_mod.time()
all_candidates = await conn.fetch(
"""
SELECT id, event_date
FROM memory_units
WHERE agent_id = $1
AND event_date BETWEEN $2 AND $3
AND id::text != ALL($4)
ORDER BY event_date DESC
""",
agent_id,
min_date,
max_date,
unit_ids
)
logger.info(f" [7.2] Fetch {len(all_candidates)} candidate neighbors (1 query): {time_mod.time() - fetch_neighbors_start:.3f}s")
# Filter and create links in memory (much faster than N queries)
link_gen_start = time_mod.time()
links = []
for unit_id, unit_event_date in new_units.items():
# Filter candidates within this unit's time window
time_lower = unit_event_date - timedelta(hours=time_window_hours)
time_upper = unit_event_date + timedelta(hours=time_window_hours)
matching_neighbors = [
(row['id'], row['event_date'])
for row in all_candidates
if time_lower <= row['event_date'] <= time_upper
][:10] # Limit to top 10
for recent_id, recent_event_date in matching_neighbors:
# Calculate temporal proximity weight
time_diff_hours = abs((unit_event_date - recent_event_date).total_seconds() / 3600)
weight = max(0.3, 1.0 - (time_diff_hours / time_window_hours))
links.append((unit_id, str(recent_id), 'temporal', weight, None))
logger.info(f" [7.3] Generate {len(links)} temporal links: {time_mod.time() - link_gen_start:.3f}s")
if links:
insert_start = time_mod.time()
await conn.executemany(
"""
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
""",
links
)
logger.info(f" [7.4] Insert {len(links)} temporal links: {time_mod.time() - insert_start:.3f}s")
except Exception as e:
logger.error(f"Failed to create temporal links: {str(e)}")
import traceback
traceback.print_exc()
raise
async def _create_semantic_links_batch(
self,
conn,
agent_id: str,
unit_ids: List[str],
embeddings: List[List[float]],
top_k: int = 5,
threshold: float = 0.7,
):
"""
Create semantic links for multiple units efficiently.
For each unit, finds similar units and creates links.
"""
if not unit_ids or not embeddings:
return
try:
import time as time_mod
import numpy as np
# Fetch ALL existing units with embeddings in ONE query
fetch_start = time_mod.time()
all_existing = await conn.fetch(
"""
SELECT id, embedding
FROM memory_units
WHERE agent_id = $1
AND embedding IS NOT NULL
AND id::text != ALL($2)
""",
agent_id,
unit_ids
)
logger.info(f" [8.1] Fetch {len(all_existing)} existing embeddings (1 query): {time_mod.time() - fetch_start:.3f}s")
# Convert to numpy for vectorized similarity computation
compute_start = time_mod.time()
all_links = []
if all_existing:
# Convert existing embeddings to numpy array
existing_ids = [str(row['id']) for row in all_existing]
# Stack embeddings as 2D array: (num_embeddings, embedding_dim)
embedding_arrays = []
for row in all_existing:
raw_emb = row['embedding']
# Handle different pgvector formats
if isinstance(raw_emb, str):
# Parse string format: "[1.0, 2.0, ...]"
import json
emb = np.array(json.loads(raw_emb), dtype=np.float32)
elif isinstance(raw_emb, (list, tuple)):
emb = np.array(raw_emb, dtype=np.float32)
else:
# Try direct conversion (works for numpy arrays, pgvector objects, etc.)
emb = np.array(raw_emb, dtype=np.float32)
embedding_arrays.append(emb)
existing_embeddings = np.vstack(embedding_arrays) if embedding_arrays else np.array([])
# For each new unit, compute similarities with ALL existing units
for unit_id, new_embedding in zip(unit_ids, embeddings):
new_emb_array = np.array(new_embedding)
# Compute cosine similarities (dot product for normalized vectors)
similarities = np.dot(existing_embeddings, new_emb_array)
# Find top-k above threshold
# Get indices of similarities above threshold
above_threshold = np.where(similarities >= threshold)[0]
if len(above_threshold) > 0:
# Sort by similarity (descending) and take top-k
sorted_indices = above_threshold[np.argsort(-similarities[above_threshold])][:top_k]
for idx in sorted_indices:
similar_id = existing_ids[idx]
similarity = float(similarities[idx])
all_links.append((unit_id, similar_id, 'semantic', similarity, None))
logger.info(f" [8.2] Compute similarities & generate {len(all_links)} semantic links: {time_mod.time() - compute_start:.3f}s")
if all_links:
insert_start = time_mod.time()
await conn.executemany(
"""
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
""",
all_links
)
logger.info(f" [8.3] Insert {len(all_links)} semantic links: {time_mod.time() - insert_start:.3f}s")
except Exception as e:
logger.error(f"Failed to create semantic links: {str(e)}")
import traceback
traceback.print_exc()
raise
async def _insert_entity_links_batch(self, conn, links: List[tuple]):
"""Insert all entity links in a single batch."""
if not links:
return
try:
await conn.executemany(
"""
INSERT INTO memory_links (from_unit_id, to_unit_id, link_type, weight, entity_id)
VALUES ($1, $2, $3, $4, $5)
ON CONFLICT (from_unit_id, to_unit_id, link_type, COALESCE(entity_id, '00000000-0000-0000-0000-000000000000'::uuid)) DO NOTHING
""",
links
)
except Exception as e:
logger.warning(f"Failed to insert entity links: {str(e)}")

18
memory/operations/search_operations.py Normal file
View file

@ -0,0 +1,18 @@
"""
Search operations for memory retrieval using spreading activation.
NOTE: This is a placeholder for future refactoring.
The actual implementation is currently in temporal_semantic_memory.py
"""
class SearchOperationsMixin:
"""
Mixin class for search operations.
Methods to be extracted:
- search
- search_async
- _apply_mmr
"""
pass

112
memory/operations/think_operations.py Normal file
View file

@ -0,0 +1,112 @@
"""
Think operations for formulating answers based on agent and world facts.
"""
import os
from typing import Dict, List, Any
from openai import AsyncOpenAI
class ThinkOperationsMixin:
"""Mixin class for think operations."""
async def think_async(
self,
agent_id: str,
query: str,
thinking_budget: int = 50,
top_k: int = 10,
model: str = "llama-3.3-70b-versatile",
temperature: float = 0.7,
max_tokens: int = 1000,
) -> Dict[str, Any]:
"""
Think and formulate an answer using agent identity and world facts.
This method:
1. Retrieves agent facts (agent's identity and past actions)
2. Retrieves world facts (general knowledge)
3. Uses Groq LLM to formulate an answer
4. Returns plain text answer and the facts used
Args:
agent_id: Agent identifier
query: Question to answer
thinking_budget: Number of memory units to explore
top_k: Maximum facts to retrieve
model: LLM model to use (default: llama-3.3-70b-versatile)
temperature: Sampling temperature
max_tokens: Maximum tokens in response
Returns:
Dict with:
- text: Plain text answer (no markdown)
- based_on: Dict with 'world' and 'agent' fact lists
"""
# Initialize Groq client
groq_api_key = os.getenv("GROQ_API_KEY")
if not groq_api_key:
raise ValueError("GROQ_API_KEY environment variable not set")
client = AsyncOpenAI(
api_key=groq_api_key,
base_url="https://api.groq.com/openai/v1"
)
# Step 1: Get agent facts (identity)
agent_results, _ = await self.search_async(
agent_id=agent_id,
query=query,
thinking_budget=thinking_budget,
top_k=top_k,
enable_trace=False,
fact_type='agent'
)
# Step 2: Get world facts
world_results, _ = await self.search_async(
agent_id=agent_id,
query=query,
thinking_budget=thinking_budget,
top_k=top_k,
enable_trace=False,
fact_type='world'
)
# Step 3: Format facts for LLM
agent_facts_text = "\n".join([f"- {fact['text']}" for fact in agent_results]) if agent_results else "None"
world_facts_text = "\n".join([f"- {fact['text']}" for fact in world_results]) if world_results else "None"
# Step 4: Call Groq to formulate answer
prompt = f"""You are an AI assistant answering a question based on retrieved facts.
AGENT IDENTITY (what the agent has done):
{agent_facts_text}
WORLD FACTS (general knowledge):
{world_facts_text}
QUESTION: {query}
Provide a helpful, accurate answer based on the facts above. If the facts don't contain enough information to answer the question, say so clearly. Do not use markdown formatting - respond in plain text only."""
response = await client.chat.completions.create(
model=model,
messages=[
{"role": "system", "content": "You are a helpful AI assistant. Always respond in plain text without markdown formatting."},
{"role": "user", "content": prompt}
],
temperature=temperature,
max_tokens=max_tokens
)
answer_text = response.choices[0].message.content.strip()
# Step 5: Return response with facts split by type
return {
"text": answer_text,
"based_on": {
"world": world_results,
"agent": agent_results
}
}

1228
memory/temporal_semantic_memory.py
View file

File diff suppressed because it is too large Load diff

165
memory/visualizer.py
View file

@ -1,165 +0,0 @@
"""
Memory visualization module.
Provides visual representations of memory networks and search paths.
"""
import time
from typing import List, Dict, Any, Optional, Tuple
import networkx as nx
import matplotlib.pyplot as plt
from matplotlib.patches import FancyBboxPatch
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.layout import Layout
from rich.live import Live
from rich.text import Text
from rich import box
class MemoryVisualizer:
"""
Visualizes memory networks and search paths.
"""
def __init__(self):
"""Initialize the visualizer."""
self.console = Console()
def visualize_memory_graph(
self,
units: List[Dict[str, Any]],
links: List[Dict[str, Any]],
output_file: str = "memory_graph.png",
highlight_nodes: Optional[List[str]] = None,
):
"""
Create a visual representation of the memory graph.
Args:
units: List of memory units (id, text, context, etc.)
links: List of links (from_unit_id, to_unit_id, link_type, weight)
output_file: Output file path for the visualization
highlight_nodes: Optional list of node IDs to highlight
"""
# Create directed graph
G = nx.DiGraph()
# Add nodes
node_labels = {}
for unit in units:
unit_id = str(unit['id'])
# Truncate text for display
label = unit['text'][:40] + "..." if len(unit['text']) > 40 else unit['text']
G.add_node(unit_id)
node_labels[unit_id] = label
# Add edges
temporal_edges = []
semantic_edges = []
for link in links:
from_id = str(link['from_unit_id'])
to_id = str(link['to_unit_id'])
weight = link['weight']
link_type = link['link_type']
if link_type == 'temporal':
temporal_edges.append((from_id, to_id, weight))
else: # semantic
semantic_edges.append((from_id, to_id, weight))
G.add_edge(from_id, to_id, weight=weight, type=link_type)
# Create figure
fig, ax = plt.subplots(figsize=(20, 14))
ax.set_facecolor('#1a1a2e')
fig.patch.set_facecolor('#0f0f1e')
# Use spring layout for better visualization
pos = nx.spring_layout(G, k=2, iterations=50, seed=42)
# Draw temporal edges (blue)
if temporal_edges:
nx.draw_networkx_edges(
G, pos,
edgelist=[(e[0], e[1]) for e in temporal_edges],
edge_color='#4ecdc4',
alpha=0.6,
width=2,
arrows=True,
arrowsize=15,
arrowstyle='->',
connectionstyle='arc3,rad=0.1',
ax=ax
)
# Draw semantic edges (purple)
if semantic_edges:
nx.draw_networkx_edges(
G, pos,
edgelist=[(e[0], e[1]) for e in semantic_edges],
edge_color='#ff6b9d',
alpha=0.6,
width=2,
arrows=True,
arrowsize=15,
arrowstyle='->',
connectionstyle='arc3,rad=0.1',
ax=ax
)
# Determine node colors
node_colors = []
for node in G.nodes():
if highlight_nodes and node in highlight_nodes:
node_colors.append('#ffd93d') # Yellow for highlighted
else:
node_colors.append('#6c63ff') # Purple for normal
# Draw nodes
nx.draw_networkx_nodes(
G, pos,
node_color=node_colors,
node_size=3000,
alpha=0.9,
ax=ax
)
# Draw labels
nx.draw_networkx_labels(
G, pos,
node_labels,
font_size=8,
font_color='white',
font_weight='bold',
ax=ax
)
# Add legend
legend_elements = [
plt.Line2D([0], [0], color='#4ecdc4', lw=2, label='Temporal Links'),
plt.Line2D([0], [0], color='#ff6b9d', lw=2, label='Semantic Links'),
plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='#6c63ff',
markersize=10, label='Memory Unit', linestyle=''),
]
if highlight_nodes:
legend_elements.append(
plt.Line2D([0], [0], marker='o', color='w', markerfacecolor='#ffd93d',
markersize=10, label='Highlighted', linestyle='')
)
ax.legend(handles=legend_elements, loc='upper left', facecolor='#2d2d44',
edgecolor='white', fontsize=10, labelcolor='white')
# Title
ax.set_title('Memory Network Graph\nTemporal + Semantic Architecture',
color='white', fontsize=16, fontweight='bold', pad=20)
ax.axis('off')
plt.tight_layout()
plt.savefig(output_file, dpi=150, facecolor='#0f0f1e')
plt.close()
self.console.print(f"[green]✓[/green] Memory graph saved to [cyan]{output_file}[/cyan]")

8
migrations/001_add_links_composite_index.sql
View file

@ -1,8 +0,0 @@
-- Migration: Add composite index for spreading activation neighbor queries
-- This index optimizes the WHERE ml.from_unit_id::text = ANY($1) AND ml.weight >= 0.1 query
-- which is used during spreading activation search.
-- Composite index for spreading activation neighbor queries (from_unit_id + weight filter)
CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_memory_links_from_weight
ON memory_links(from_unit_id, weight DESC)
WHERE weight >= 0.1;

22
migrations/002_add_documents_table.sql
View file

@ -1,22 +0,0 @@
-- Migration: Add documents table and document_id to memory_units
-- This enables document tracking, upsert, and cascade deletion
-- Create documents table
CREATE TABLE IF NOT EXISTS documents (
id TEXT NOT NULL,
agent_id TEXT NOT NULL,
PRIMARY KEY (id, agent_id),
original_text TEXT,
content_hash TEXT,
metadata JSONB DEFAULT '{}'::jsonb,
created_at TIMESTAMPTZ DEFAULT NOW(),
updated_at TIMESTAMPTZ DEFAULT NOW()
);
-- Add document_id column to memory_units (nullable for backward compatibility)
ALTER TABLE memory_units ADD COLUMN IF NOT EXISTS document_id TEXT REFERENCES documents(id) ON DELETE CASCADE;
-- Create indexes
CREATE INDEX IF NOT EXISTS idx_documents_agent_id ON documents(agent_id);
CREATE INDEX IF NOT EXISTS idx_documents_content_hash ON documents(content_hash);
CREATE INDEX IF NOT EXISTS idx_memory_units_document_id ON memory_units(document_id);

19
migrations/003_fix_documents_primary_key_v2.sql
View file

@ -1,19 +0,0 @@
-- Migration: Fix documents table primary key to be composite (id, agent_id)
-- This is a safer approach that doesn't drop the table
-- Step 1: Drop the foreign key constraint from memory_units
ALTER TABLE memory_units DROP CONSTRAINT IF EXISTS memory_units_document_id_fkey;
ALTER TABLE memory_units DROP CONSTRAINT IF EXISTS memory_units_document_fkey;
-- Step 2: Drop the old primary key on documents
ALTER TABLE documents DROP CONSTRAINT IF EXISTS documents_pkey;
-- Step 3: Add the new composite primary key
ALTER TABLE documents ADD PRIMARY KEY (id, agent_id);
-- Step 4: Add back the foreign key constraint with the composite key
ALTER TABLE memory_units
ADD CONSTRAINT memory_units_document_fkey
FOREIGN KEY (document_id, agent_id)
REFERENCES documents(id, agent_id)
ON DELETE CASCADE;

29
pyproject.toml
View file

@ -1,5 +1,9 @@
[build-system]
requires = ["hatchling"]
build-backend = "hatchling.build"
[project]
name = "memory-poc"
name = "agent_memory"
version = "0.1.0"
description = "Temporal + Semantic + Entity Memory System for AI agents using PostgreSQL"
readme = "README.md"
@ -9,17 +13,28 @@ dependencies = [
"python-dotenv>=1.0.0",
"openai>=1.0.0",
"pydantic>=2.0.0",
"nltk>=3.8.0",
"networkx>=3.0",
"matplotlib>=3.7.0",
"rich>=13.0.0",
"spacy>=3.7.0",
"sentence-transformers>=2.2.0",
"torch>=2.0.0",
"pytest>=7.0.0",
"pytest-asyncio>=0.21.0",
"langchain-text-splitters>=0.3.0",
"flask>=3.1.2",
"fastapi[standard]>=0.120.3",
"uvicorn>=0.38.0",
"sqlalchemy>=2.0.44",
"alembic>=1.17.1",
"pgvector>=0.4.1",
"greenlet>=3.2.4",
"psycopg2-binary>=2.9.11",
"pytest-timeout>=2.4.0",
]
[tool.hatch.build.targets.wheel]
packages = ["memory", "web"]
[tool.pytest.ini_options]
log_cli = true
log_cli_level = "INFO"
log_cli_format = "%(asctime)s %(levelname)s %(message)s"
log_cli_date_format = "%Y-%m-%d %H:%M:%S"
addopts = "--timeout 60 -p no:warnings"
asyncio_default_fixture_loop_scope = "session"

8
schema.sql
View file

@ -27,9 +27,13 @@ CREATE TABLE IF NOT EXISTS memory_units (
embedding vector(384), -- bge-small-en-v1.5 dimension
context TEXT, -- What was happening when this memory was formed
event_date TIMESTAMPTZ NOT NULL, -- When the event occurred
fact_type TEXT NOT NULL DEFAULT 'world', -- 'world' (general facts), 'agent' (agent actions), or 'opinion' (agent opinions)
confidence_score FLOAT, -- Confidence score for opinions (0.0 to 1.0, only used for fact_type='opinion')
access_count INTEGER DEFAULT 0, -- For recency/frequency weighting
created_at TIMESTAMPTZ DEFAULT NOW(),
updated_at TIMESTAMPTZ DEFAULT NOW()
updated_at TIMESTAMPTZ DEFAULT NOW(),
CHECK (fact_type IN ('world', 'agent', 'opinion')),
CHECK (confidence_score IS NULL OR (confidence_score >= 0.0 AND confidence_score <= 1.0))
);
-- Entities: Resolved entities (people, organizations, locations, etc.)
@ -108,6 +112,8 @@ CREATE INDEX IF NOT EXISTS idx_memory_units_document_id ON memory_units(document
CREATE INDEX IF NOT EXISTS idx_memory_units_event_date ON memory_units(event_date DESC);
CREATE INDEX IF NOT EXISTS idx_memory_units_agent_date ON memory_units(agent_id, event_date DESC);
CREATE INDEX IF NOT EXISTS idx_memory_units_access_count ON memory_units(access_count DESC);
CREATE INDEX IF NOT EXISTS idx_memory_units_fact_type ON memory_units(fact_type);
CREATE INDEX IF NOT EXISTS idx_memory_units_agent_fact_type ON memory_units(agent_id, fact_type);
-- Vector similarity index (HNSW for fast approximate nearest neighbor)
CREATE INDEX IF NOT EXISTS idx_memory_units_embedding ON memory_units

34
tests/conftest.py
View file

@ -2,6 +2,7 @@
Pytest configuration and shared fixtures.
"""
import pytest
import pytest_asyncio
import os
import asyncio
from dotenv import load_dotenv
@ -11,18 +12,23 @@ import asyncpg
load_dotenv()
@pytest.fixture(scope="function")
def memory():
@pytest_asyncio.fixture(scope="session")
async def memory():
"""
Provide a clean memory system instance for each test.
Provide a shared memory system instance for all tests in the session.
This avoids reloading the embedding model for every test (saves 3+ seconds per test).
"""
mem = TemporalSemanticMemory()
yield mem
# Cleanup is handled by individual tests
# Ensure cleanup happens at end of session
try:
await mem.close()
except Exception as e:
print(f"Warning: Error during memory cleanup: {e}")
@pytest.fixture(scope="function")
def clean_agent(memory):
@pytest_asyncio.fixture(scope="function")
async def clean_agent(memory):
"""
Provide a clean agent ID and clean up data after test.
Uses agent_id='test' for all tests (multi-tenant isolation).
@ -30,19 +36,25 @@ def clean_agent(memory):
agent_id = "test"
# Clean up before test
asyncio.run(memory.delete_agent(agent_id))
await memory.delete_agent(agent_id)
yield agent_id
# Clean up after test
asyncio.run(memory.delete_agent(agent_id))
try:
await memory.delete_agent(agent_id)
except Exception as e:
print(f"Warning: Error during agent cleanup: {e}")
@pytest.fixture
@pytest_asyncio.fixture
async def db_connection():
"""
Provide a database connection for direct DB queries in tests.
"""
conn = await asyncpg.connect(os.getenv('DATABASE_URL'))
conn = await asyncpg.connect(os.getenv('DATABASE_URL'), statement_cache_size=0)
yield conn
await conn.close()
try:
await conn.close()
except Exception as e:
print(f"Warning: Error closing connection: {e}")

19
tests/fixtures/README.md vendored Normal file
View file

@ -0,0 +1,19 @@
# Test Fixtures
## locomo_conversation_sample.json
Sample conversation from the LoComo benchmark (conv-26) used for performance tuning tests.
**Stats:**
- Sample ID: conv-26
- Sessions: 19
- Total dialogues: 419
- Questions: 199
**Usage:**
Used by `test_performance_tuning.py` to measure:
- Batch ingestion performance
- Search performance
- Entity resolution performance
This is a realistic long-form conversation for stress testing the memory system.

5271
tests/fixtures/locomo_conversation_sample.json vendored Normal file
View file

File diff suppressed because it is too large Load diff

5
tests/test_document_tracking.py
View file

@ -1,6 +1,7 @@
"""
Tests for document tracking and upsert functionality.
"""
import logging
import os
import pytest
from datetime import datetime, timezone
@ -15,6 +16,7 @@ async def test_document_creation_and_retrieval():
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
try:
agent_id = f"test_doc_{datetime.now(timezone.utc).timestamp()}"
@ -51,6 +53,7 @@ async def test_document_upsert():
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
try:
agent_id = f"test_upsert_{datetime.now(timezone.utc).timestamp()}"
@ -100,6 +103,7 @@ async def test_document_deletion():
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
try:
agent_id = f"test_delete_{datetime.now(timezone.utc).timestamp()}"
@ -139,6 +143,7 @@ async def test_memory_without_document():
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
try:
agent_id = f"test_no_doc_{datetime.now(timezone.utc).timestamp()}"

131
tests/test_performance_tuning.py Normal file
View file

@ -0,0 +1,131 @@
"""
Performance tuning test using real LoComo conversation.
This test loads a long conversation (419 dialogues across 19 sessions),
ingests it into memory, and runs searches to measure performance.
"""
import logging
import os
import json
import pytest
from datetime import datetime, timezone
from pathlib import Path
from memory import TemporalSemanticMemory
# Configure logging to show performance metrics
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s %(levelname)s:%(name)s: %(message)s'
)
@pytest.mark.asyncio
@pytest.mark.timeout(300) # 5 minute timeout for performance test
async def test_batch_ingestion_single_call():
"""
Test ingesting entire conversation in ONE batch call.
This is the most efficient way - all sessions in one put_batch_async.
"""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
# Load conversation fixture
fixture_path = Path(__file__).parent / "fixtures" / "locomo_conversation_sample.json"
with open(fixture_path) as f:
conversation_data = json.load(f)
sample_id = conversation_data['sample_id']
logging.info(f"\n{'='*80}")
logging.info(f"BATCH INGESTION TEST: {sample_id}")
logging.info(f"{'='*80}")
# Initialize memory
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
agent_id = f"batch_test_{sample_id}_{datetime.now(timezone.utc).timestamp()}"
try:
# Parse all sessions into batch format
# LIMIT to first 5 sessions for faster iteration during perf tuning
MAX_SESSIONS = 5
logging.info(f"\nPreparing batch contents (limiting to {MAX_SESSIONS} sessions for perf tuning)...")
conversation = conversation_data['conversation']
batch_contents = []
for i in range(1, MAX_SESSIONS + 1):
session_key = f'session_{i}'
session_date_key = f'session_{i}_date_time'
if session_key not in conversation or not conversation[session_key]:
break
session_dialogues = conversation[session_key]
session_date = conversation.get(session_date_key, datetime.now(timezone.utc).isoformat())
# Combine dialogues
session_text = "\n".join([
f"{d['speaker']}: {d['text']}"
for d in session_dialogues
])
# Parse date
try:
from dateutil import parser as date_parser
event_date = date_parser.isoparse(session_date)
except:
event_date = datetime.now(timezone.utc)
batch_contents.append({
'content': session_text,
'context': f'session_{i}',
'event_date': event_date
})
logging.info(f"Prepared {len(batch_contents)} sessions for batch ingestion")
# Single batch call
logging.info(f"\nIngesting all {len(batch_contents)} sessions in ONE batch call...")
result_ids = await memory.put_batch_async(
agent_id=agent_id,
contents=batch_contents,
document_id=f"{agent_id}_full_conversation"
)
total_units = sum(len(ids) for ids in result_ids)
logging.info(f"\n{'='*80}")
logging.info(f"BATCH INGESTION COMPLETE: {total_units} memory units created")
logging.info(f"{'='*80}")
# Run one sample search
logging.info(f"\nRunning sample search...")
question = conversation_data['qa'][0]['question']
logging.info(f"Question: {question}")
results, _ = await memory.search_async(
agent_id=agent_id,
query=question,
thinking_budget=100,
top_k=5,
enable_trace=False
)
logging.info(f"Found {len(results)} results")
if results:
logging.info(f"Top result: {results[0]['text'][:100]}...")
finally:
# Cleanup
logging.info("\nCleaning up...")
await memory.delete_agent(agent_id)
await memory.close()
if __name__ == "__main__":
# Allow running directly for quick perf checks
import asyncio
logging.info("Running performance test...")
asyncio.run(test_batch_ingestion_single_call())

179
tests/test_think.py Normal file
View file

@ -0,0 +1,179 @@
"""
Test think function for opinion generation and consistency.
"""
import pytest
import os
from datetime import datetime, timezone
from memory import TemporalSemanticMemory
@pytest.mark.asyncio
async def test_think_opinion_consistency():
"""
Test that think function:
1. Generates an opinion
2. Stores the opinion in the database
3. Returns consistent response on subsequent calls with the same query
"""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
try:
agent_id = f"test_think_{datetime.now(timezone.utc).timestamp()}"
# Store some initial facts to give context for opinion formation
await memory.put_async(
agent_id=agent_id,
content="Alice is a software engineer who has worked on 5 major projects. She always delivers on time and writes clean, well-documented code.",
context="performance review",
event_date=datetime(2024, 1, 15, tzinfo=timezone.utc)
)
await memory.put_async(
agent_id=agent_id,
content="Bob recently joined the team. He missed his first deadline and his code had many bugs.",
context="performance review",
event_date=datetime(2024, 2, 1, tzinfo=timezone.utc)
)
# First think call - should generate opinions
query = "Who is a more reliable engineer?"
result1 = await memory.think_async(
agent_id=agent_id,
query=query,
thinking_budget=30,
top_k=10
)
print(f"\n=== First Think Call ===")
print(f"Answer: {result1['text']}")
print(f"New opinions formed: {len(result1.get('new_opinions', []))}")
for opinion in result1.get('new_opinions', []):
print(f" - {opinion['text']} (confidence: {opinion['confidence']:.2f})")
# Verify we got an answer
assert result1['text'], "First think call should return an answer"
assert 'based_on' in result1, "Should return based_on facts"
# Verify opinions were formed
new_opinions_count = len(result1.get('new_opinions', []))
print(f"\nNew opinions formed: {new_opinions_count}")
# Wait a moment to ensure opinions are stored and any background tasks complete
import asyncio
await asyncio.sleep(2.0)
# Search for stored opinions to verify they were actually saved
pool = await memory._get_pool()
async with pool.acquire() as conn:
stored_opinions = await conn.fetch(
"""
SELECT id, text, confidence_score, fact_type
FROM memory_units
WHERE agent_id = $1 AND fact_type = 'opinion'
ORDER BY created_at DESC
""",
agent_id
)
print(f"\n=== Stored Opinions in Database ===")
print(f"Total opinions stored: {len(stored_opinions)}")
for op in stored_opinions:
print(f" - {op['text']} (confidence: {op['confidence_score']:.2f})")
# Verify opinions were actually written to database
assert len(stored_opinions) > 0, "Opinions should be stored in the database"
assert all(op['fact_type'] == 'opinion' for op in stored_opinions), "All stored items should have fact_type='opinion'"
# Second think call - should use the stored opinions
result2 = await memory.think_async(
agent_id=agent_id,
query=query,
thinking_budget=30,
top_k=10
)
print(f"\n=== Second Think Call ===")
print(f"Answer: {result2['text']}")
print(f"Existing opinions used: {len(result2['based_on'].get('opinion', []))}")
for opinion in result2['based_on'].get('opinion', []):
print(f" - {opinion['text']}")
print(f"New opinions formed: {len(result2.get('new_opinions', []))}")
# Verify second call also got an answer
assert result2['text'], "Second think call should return an answer"
# Verify second call used the stored opinions
assert len(result2['based_on'].get('opinion', [])) > 0, "Second call should retrieve stored opinions"
# The responses should be consistent (both should mention the same person as more reliable)
# We'll do a basic check that they're not contradictory
text1_lower = result1['text'].lower()
text2_lower = result2['text'].lower()
print(f"\n=== Consistency Check ===")
# Check if Alice is mentioned as more reliable in first response
if 'alice' in text1_lower and ('reliable' in text1_lower or 'better' in text1_lower):
print("First response favors Alice")
# Second response should also favor Alice (consistency)
assert 'alice' in text2_lower, "Second response should also mention Alice"
print("Second response also mentions Alice - CONSISTENT ✓")
# Check if Bob is mentioned
if 'bob' in text1_lower:
print("First response mentions Bob")
if 'bob' in text2_lower:
print("Second response also mentions Bob - CONSISTENT ✓")
print(f"\n✅ Test passed - opinions were formed, stored, and used consistently")
finally:
# Clean up agent data
try:
await memory.delete_agent(agent_id)
except Exception as e:
print(f"Warning: Error during cleanup: {e}")
await memory.close()
@pytest.mark.asyncio
async def test_think_without_prior_context():
"""
Test that think function handles queries when there's no relevant context.
"""
db_url = os.getenv("DATABASE_URL")
if not db_url:
pytest.skip("DATABASE_URL not set")
memory = TemporalSemanticMemory(db_url=db_url)
await memory.initialize()
try:
agent_id = f"test_think_no_context_{datetime.now(timezone.utc).timestamp()}"
# Call think without storing any prior facts
result = await memory.think_async(
agent_id=agent_id,
query="What is the capital of France?",
thinking_budget=20,
top_k=5
)
print(f"\n=== Think Without Context ===")
print(f"Answer: {result['text']}")
# Should still return an answer (even if it says it doesn't have enough info)
assert result['text'], "Should return some answer"
assert 'based_on' in result, "Should return based_on structure"
finally:
await memory.close()
if __name__ == "__main__":
pytest.main([__file__, "-v", "-s"])

53
update_fact_types.py Executable file
View file

@ -0,0 +1,53 @@
#!/usr/bin/env python3
"""Update existing memory_units to have fact_type='world' if NULL."""
import asyncio
import asyncpg
import os
from dotenv import load_dotenv
load_dotenv()
async def main():
conn = await asyncpg.connect(os.getenv('DATABASE_URL'))
# Check if fact_type column exists
result = await conn.fetchrow("""
SELECT column_name
FROM information_schema.columns
WHERE table_name = 'memory_units' AND column_name = 'fact_type'
""")
if not result:
print("fact_type column does not exist! Run migrations first:")
print(" uv run alembic upgrade head")
await conn.close()
return
print("fact_type column exists ✓")
# Update existing rows
count = await conn.fetchval("""
UPDATE memory_units
SET fact_type = 'world'
WHERE fact_type IS NULL
RETURNING (SELECT COUNT(*) FROM memory_units WHERE fact_type IS NULL)
""")
print(f"Updated {count} rows with fact_type='world'")
# Show distribution
distribution = await conn.fetch("""
SELECT fact_type, COUNT(*) as count
FROM memory_units
GROUP BY fact_type
ORDER BY count DESC
""")
print("\nFact type distribution:")
for row in distribution:
print(f" {row['fact_type']}: {row['count']}")
await conn.close()
if __name__ == "__main__":
asyncio.run(main())

1113
uv.lock
View file

File diff suppressed because it is too large Load diff

8
web/__init__.py Normal file
View file

@ -0,0 +1,8 @@
"""
Web interface for memory system.
Provides FastAPI app and visualization interface.
"""
from .server import app, memory
__all__ = ["app", "memory"]

306
web/server.py
View file

@ -4,7 +4,6 @@ FastAPI server for memory graph visualization and API.
Provides REST API endpoints for memory operations and serves
the interactive visualization interface.
"""
import asyncpg
import asyncio
from fastapi import FastAPI, HTTPException
from fastapi.staticfiles import StaticFiles
@ -58,147 +57,35 @@ class BatchPutRequest(BaseModel):
upsert: bool = False
async def get_graph_data():
"""Fetch graph data from database."""
conn = await asyncpg.connect(
os.getenv('DATABASE_URL'),
statement_cache_size=0 # Disable statement caching for pgbouncer compatibility
)
class ThinkRequest(BaseModel):
"""Request model for think endpoint."""
query: str
agent_id: str = "default"
thinking_budget: int = 50
top_k: int = 10
# Get all memory units
units = await conn.fetch("""
SELECT id, text, event_date, context
FROM memory_units
ORDER BY event_date
""")
# Get all links with weights
links = await conn.fetch("""
SELECT
ml.from_unit_id,
ml.to_unit_id,
ml.link_type,
ml.weight,
e.canonical_name as entity_name
FROM memory_links ml
LEFT JOIN entities e ON ml.entity_id = e.id
ORDER BY ml.link_type, ml.weight DESC
""")
class ThinkResponse(BaseModel):
"""Response model for think endpoint."""
text: str
based_on: Dict[str, List[Dict[str, Any]]] # {"world": [...], "agent": [...], "opinion": [...]}
new_opinions: List[str] = [] # List of newly formed opinions
# Get entity information
unit_entities = await conn.fetch("""
SELECT ue.unit_id, e.canonical_name, e.entity_type
FROM unit_entities ue
JOIN entities e ON ue.entity_id = e.id
ORDER BY ue.unit_id
""")
await conn.close()
# Build entity mapping
entity_map = {}
for row in unit_entities:
unit_id = row['unit_id']
entity_name = row['canonical_name']
entity_type = row['entity_type']
if unit_id not in entity_map:
entity_map[unit_id] = []
entity_map[unit_id].append(f"{entity_name} ({entity_type})")
# Build nodes
nodes = []
for row in units:
unit_id = row['id']
text = row['text']
event_date = row['event_date']
context = row['context']
entities = entity_map.get(unit_id, [])
entity_count = len(entities)
# Color by entity count
if entity_count == 0:
color = "#e0e0e0"
elif entity_count == 1:
color = "#90caf9"
else:
color = "#42a5f5"
nodes.append({
"data": {
"id": str(unit_id),
"label": text[:50] + "..." if len(text) > 50 else text,
"text": text,
"context": context,
"date": str(event_date.date()),
"entities": ", ".join(entities) if entities else "None",
"color": color
}
})
# Build edges
edges = []
for row in links:
from_id = row['from_unit_id']
to_id = row['to_unit_id']
link_type = row['link_type']
weight = row['weight']
entity_name = row['entity_name']
# Set color based on link type
if link_type == 'temporal':
color = "#00bcd4"
line_style = "dashed"
elif link_type == 'semantic':
color = "#ff69b4"
line_style = "solid"
elif link_type == 'entity':
color = "#ffd700"
line_style = "solid"
else:
color = "#999999"
line_style = "solid"
edges.append({
"data": {
"id": f"{from_id}-{to_id}-{link_type}",
"source": str(from_id),
"target": str(to_id),
"weight": weight,
"linkType": link_type,
"entityName": entity_name or "",
"color": color,
"lineStyle": line_style
}
})
# Build table rows
table_rows = []
for row in units:
unit_id = row['id']
text = row['text']
event_date = row['event_date']
context = row['context']
entities = entity_map.get(unit_id, [])
entity_str = ", ".join(entities) if entities else "None"
table_rows.append({
"id": str(unit_id)[:8] + "...",
"text": text,
"context": context,
"date": str(event_date.date()),
"entities": entity_str
})
return {
"nodes": nodes,
"edges": edges,
"table_rows": table_rows,
"total_units": len(units)
}
memory = TemporalSemanticMemory()
@app.on_event("startup")
async def startup_event():
"""Initialize memory system on startup."""
await memory.initialize()
logging.info("Memory system initialized")
@app.on_event("shutdown")
async def shutdown_event():
"""Cleanup memory system on shutdown."""
await memory.close()
logging.info("Memory system closed")
@app.get("/")
async def index():
"""Serve the visualization page."""
@ -206,10 +93,10 @@ async def index():
@app.get("/api/graph")
async def api_graph():
"""Get graph data from database."""
async def api_graph(agent_id: Optional[str] = None, fact_type: Optional[str] = None):
"""Get graph data from database, optionally filtered by agent_id and fact_type."""
try:
data = await get_graph_data()
data = await memory.get_graph_data(agent_id, fact_type)
return data
except Exception as e:
import traceback
@ -247,26 +134,133 @@ async def api_search(request: SearchRequest):
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/world_search")
async def api_world_search(request: SearchRequest):
"""Search only world facts (general knowledge about the world)."""
try:
# Run search with fact_type filter for 'world'
results, trace = await memory.search_async(
agent_id=request.agent_id,
query=request.query,
thinking_budget=request.thinking_budget,
top_k=request.top_k,
enable_trace=request.trace,
mmr_lambda=request.mmr_lambda,
fact_type='world'
)
# Convert trace to dict
trace_dict = trace.to_dict() if trace else None
return {
'results': results,
'trace': trace_dict
}
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
print(f"Error in /api/world_search: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/agent_search")
async def api_agent_search(request: SearchRequest):
"""Search only agent facts (facts about what the agent did)."""
try:
# Run search with fact_type filter for 'agent'
results, trace = await memory.search_async(
agent_id=request.agent_id,
query=request.query,
thinking_budget=request.thinking_budget,
top_k=request.top_k,
enable_trace=request.trace,
mmr_lambda=request.mmr_lambda,
fact_type='agent'
)
# Convert trace to dict
trace_dict = trace.to_dict() if trace else None
return {
'results': results,
'trace': trace_dict
}
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
print(f"Error in /api/agent_search: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/opinion_search")
async def api_opinion_search(request: SearchRequest):
"""Search only opinion facts (agent's formed opinions and perspectives)."""
try:
# Run search with fact_type filter for 'opinion'
results, trace = await memory.search_async(
agent_id=request.agent_id,
query=request.query,
thinking_budget=request.thinking_budget,
top_k=request.top_k,
enable_trace=request.trace,
mmr_lambda=request.mmr_lambda,
fact_type='opinion'
)
# Convert trace to dict
trace_dict = trace.to_dict() if trace else None
return {
'results': results,
'trace': trace_dict
}
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
print(f"Error in /api/opinion_search: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.post("/api/think")
async def api_think(request: ThinkRequest):
"""
Think and formulate an answer using agent identity, world facts, and opinions.
This endpoint:
1. Retrieves agent facts (agent's identity)
2. Retrieves world facts relevant to the query
3. Retrieves existing opinions (agent's perspectives)
4. Uses Groq LLM to formulate an answer
5. Extracts and stores any new opinions formed
6. Returns plain text answer, the facts used, and new opinions
"""
try:
# Use the memory system's think_async method
result = await memory.think_async(
agent_id=request.agent_id,
query=request.query,
thinking_budget=request.thinking_budget,
top_k=request.top_k
)
return ThinkResponse(
text=result["text"],
based_on=result["based_on"],
new_opinions=result.get("new_opinions", [])
)
except Exception as e:
import traceback
error_detail = f"{str(e)}\n\nTraceback:\n{traceback.format_exc()}"
print(f"Error in /api/think: {error_detail}")
raise HTTPException(status_code=500, detail=str(e))
@app.get("/api/agents")
async def api_agents():
"""Get list of available agents from database."""
try:
conn = await asyncpg.connect(
os.getenv('DATABASE_URL'),
statement_cache_size=0
)
# Get distinct agent IDs from memory_units
agents = await conn.fetch("""
SELECT DISTINCT agent_id
FROM memory_units
WHERE agent_id IS NOT NULL
ORDER BY agent_id
""")
await conn.close()
agent_list = [row['agent_id'] for row in agents]
agent_list = await memory.list_agents()
return {"agents": agent_list}
except Exception as e:
import traceback
@ -325,8 +319,6 @@ async def api_batch_put(request: BatchPutRequest):
upsert=request.upsert
)
await memory.close()
return {
"success": True,
"message": f"Successfully stored {len(contents)} memory items",

308
web/static/css/styles.css
View file

@ -5,6 +5,52 @@ body {
background: #f5f5f5;
}
/* Breadcrumb */
.breadcrumb-container {
background: #333;
color: white;
padding: 12px 20px;
border-bottom: 3px solid #42a5f5;
}
.breadcrumb {
display: flex;
align-items: center;
gap: 10px;
font-size: 14px;
}
.breadcrumb-item {
font-weight: 500;
}
.breadcrumb-separator {
color: #999;
font-weight: bold;
}
.agent-selector {
padding: 6px 12px;
border: 2px solid #42a5f5;
border-radius: 4px;
background: white;
color: #333;
font-size: 14px;
font-weight: bold;
cursor: pointer;
transition: all 0.2s;
}
.agent-selector:hover {
background: #e3f2fd;
border-color: #1e88e5;
}
.agent-selector:focus {
outline: none;
box-shadow: 0 0 0 3px rgba(66, 165, 245, 0.3);
}
.tab-container {
background: white;
}
@ -51,6 +97,192 @@ body {
display: block;
}
/* Data Tab Styles */
.data-sub-tabs {
background: #e3f2fd;
padding: 10px 20px;
border-bottom: 2px solid #333;
display: flex;
gap: 10px;
}
.data-sub-tab-button {
background: #fff;
border: 2px solid #42a5f5;
padding: 8px 20px;
cursor: pointer;
font-size: 14px;
font-weight: bold;
border-radius: 4px;
transition: all 0.2s;
}
.data-sub-tab-button:hover {
background: #e3f2fd;
}
.data-sub-tab-button.active {
background: #42a5f5;
color: white;
}
.data-subtab-content {
display: none;
}
.data-subtab-content.active {
display: block;
}
.view-toggle {
background: #f9f9f9;
padding: 10px 20px;
border-bottom: 2px solid #ddd;
display: flex;
gap: 10px;
}
.view-toggle-button {
background: #fff;
border: 2px solid #ccc;
padding: 6px 16px;
cursor: pointer;
font-size: 14px;
font-weight: bold;
border-radius: 4px;
transition: all 0.2s;
}
.view-toggle-button:hover {
background: #e0e0e0;
}
.view-toggle-button.active {
background: #333;
color: white;
border-color: #333;
}
.no-agent-message {
padding: 40px;
text-align: center;
color: #666;
background: #f9f9f9;
}
.data-view {
position: relative;
}
.data-controls {
padding: 15px;
background: #f9f9f9;
border-bottom: 2px solid #333;
display: flex;
gap: 15px;
align-items: center;
flex-wrap: wrap;
}
.data-controls h2 {
margin: 0;
flex: 1 0 100%;
}
.load-button {
padding: 8px 20px;
background: #66bb6a;
color: white;
border: none;
border-radius: 4px;
cursor: pointer;
font-weight: bold;
font-size: 14px;
}
.load-button:hover {
background: #43a047;
}
.control-input {
width: 80px;
padding: 5px;
border: 1px solid #ccc;
border-radius: 4px;
}
.control-select {
padding: 5px;
border: 1px solid #ccc;
border-radius: 4px;
}
.apply-button {
padding: 6px 15px;
background: #42a5f5;
color: white;
border: none;
border-radius: 4px;
cursor: pointer;
font-weight: bold;
}
.apply-button:hover {
background: #1e88e5;
}
.node-count {
color: #666;
font-size: 14px;
}
.graph-canvas {
width: 100%;
height: 800px;
background: #ffffff;
}
.table-filter {
width: 100%;
max-width: 600px;
padding: 10px;
margin: 0 20px 15px 20px;
border: 2px solid #ccc;
border-radius: 4px;
font-size: 14px;
box-sizing: border-box;
}
.table-container {
overflow-x: auto;
padding: 0 20px 20px 20px;
}
.memory-table {
width: 100%;
border-collapse: collapse;
font-size: 13px;
max-width: 1400px;
}
.memory-table th {
padding: 10px;
text-align: left;
border: 1px solid #ddd;
background: #f0f0f0;
}
.memory-table td {
padding: 8px;
border: 1px solid #ddd;
}
.empty-message {
padding: 40px;
text-align: center;
color: #666;
}
#cy {
width: 100%;
height: 800px;
@ -84,6 +316,10 @@ body {
padding-bottom: 5px;
}
.legend h4 {
margin: 10px 0 5px 0;
}
.legend-item {
margin: 8px 0;
display: flex;
@ -96,6 +332,19 @@ body {
margin-right: 10px;
}
.legend-line.temporal {
background: #00bcd4;
border-top: 1px dashed #00bcd4;
}
.legend-line.semantic {
background: #ff69b4;
}
.legend-line.entity {
background: #ffd700;
}
.legend-node {
width: 20px;
height: 20px;
@ -104,35 +353,18 @@ body {
border-radius: 3px;
}
#table-filter {
width: 100%;
max-width: 600px;
padding: 10px;
margin-bottom: 15px;
border: 2px solid #ccc;
border-radius: 4px;
font-size: 14px;
box-sizing: border-box;
.legend-node.no-entities {
background: #e0e0e0;
}
#memory-table {
width: 100%;
border-collapse: collapse;
font-size: 13px;
max-width: 1400px;
.legend-node.one-entity {
background: #90caf9;
}
#memory-table th {
padding: 10px;
text-align: left;
border: 1px solid #ddd;
background: #f0f0f0;
.legend-node.multi-entities {
background: #42a5f5;
}
#memory-table td {
padding: 8px;
border: 1px solid #ddd;
}
.tooltip {
position: absolute;
@ -496,3 +728,33 @@ body {
max-width: 320px;
font-size: 12px;
}
/* Think Tab Styles */
#think-tab {
padding: 20px;
max-width: 1200px;
margin: 0 auto;
}
.think-controls {
background: #f9f9f9;
padding: 20px;
border-radius: 8px;
border: 2px solid #333;
}
.think-answer {
background: white;
padding: 20px;
border-radius: 8px;
border: 2px solid #333;
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}
.think-sources {
background: white;
padding: 20px;
border-radius: 8px;
border: 2px solid #333;
box-shadow: 0 2px 4px rgba(0,0,0,0.1);
}

694
web/static/js/app.js
View file

@ -3,11 +3,367 @@ let allGraphData = null;
let cy = null;
let debugPanes = [];
let debugPaneCounter = 0;
let currentAgentId = null; // Global agent context
let dataGraphs = {
world: null,
agent: null,
opinions: null
};
let dataCache = {
world: null,
agent: null,
opinions: null
};
let currentDataSubTab = 'world';
// Load data from API
// Main tab switching (Data, Debug, Think, Benchmark)
window.switchMainTab = function(tabName) {
// Remove active class from all main tabs and buttons
document.querySelectorAll('.tab-content').forEach(tab => {
tab.classList.remove('active');
});
document.querySelectorAll('.tab-button').forEach(btn => {
btn.classList.remove('active');
});
// Activate the selected tab
const tabElement = document.getElementById(`${tabName}-tab`);
if (tabElement) {
tabElement.classList.add('active');
}
// Find and activate the corresponding button
const buttons = document.querySelectorAll('.tab-button');
buttons.forEach(btn => {
const btnText = btn.textContent.toLowerCase();
if (
(tabName === 'data' && btnText.includes('data')) ||
(tabName === 'debug' && btnText.includes('debug')) ||
(tabName === 'think' && btnText.includes('think')) ||
(tabName === 'benchmark' && btnText.includes('benchmark'))
) {
btn.classList.add('active');
}
});
// Tab-specific logic
if (tabName === 'data') {
// Resize current data graph if exists
const factType = currentDataSubTab;
if (dataGraphs[factType]) {
setTimeout(() => dataGraphs[factType].resize(), 10);
}
} else if (tabName === 'debug') {
if (debugPanes.length === 0) {
addDebugPane();
}
debugPanes.forEach(pane => {
if (pane.cy) {
pane.cy.resize();
}
});
} else if (tabName === 'think') {
// Think tab uses global agent selector
}
}
// Data subtab switching (World, Agent, Opinions)
window.switchDataSubTab = function(subTab) {
currentDataSubTab = subTab;
// Remove active class from all subtab buttons and content
document.querySelectorAll('.data-sub-tab-button').forEach(btn => {
btn.classList.remove('active');
});
document.querySelectorAll('.data-subtab-content').forEach(content => {
content.classList.remove('active');
});
// Activate selected subtab
const buttons = document.querySelectorAll('.data-sub-tab-button');
buttons.forEach(btn => {
if (btn.textContent.toLowerCase().includes(subTab.toLowerCase())) {
btn.classList.add('active');
}
});
const subtabElement = document.getElementById(`${subTab}-subtab`);
if (subtabElement) {
subtabElement.classList.add('active');
}
// Resize graph if exists
if (dataGraphs[subTab]) {
setTimeout(() => dataGraphs[subTab].resize(), 10);
}
}
// Switch between graph and table view for a fact type
window.switchDataView = function(factType, viewType) {
const graphView = document.getElementById(`${factType}-graph-view`);
const tableView = document.getElementById(`${factType}-table-view`);
const buttons = document.querySelectorAll(`#${factType}-subtab .view-toggle-button`);
// Update button states
buttons.forEach(btn => {
btn.classList.remove('active');
if ((viewType === 'graph' && btn.textContent.includes('Graph')) ||
(viewType === 'table' && btn.textContent.includes('Table'))) {
btn.classList.add('active');
}
});
// Show/hide views
if (viewType === 'graph') {
graphView.style.display = 'block';
tableView.style.display = 'none';
if (dataGraphs[factType]) {
setTimeout(() => dataGraphs[factType].resize(), 10);
}
} else {
graphView.style.display = 'none';
tableView.style.display = 'block';
}
}
// Load data for a specific fact type
window.loadDataView = async function(factType) {
if (!currentAgentId) {
alert('Please select an agent first');
return;
}
try {
// Build URL with agent filter and fact_type filter
let url = `/api/graph?agent_id=${encodeURIComponent(currentAgentId)}`;
if (factType !== 'all') {
url += `&fact_type=${factType}`;
}
const response = await fetch(url);
if (!response.ok) {
const error = await response.json();
throw new Error(error.detail || `HTTP ${response.status}`);
}
const data = await response.json();
// Validate response structure
if (!data || !data.nodes || !data.edges) {
throw new Error('Invalid response format from server');
}
// Cache the data
dataCache[factType] = data;
// Update table
updateDataTable(factType, data);
// Update graph if in graph view
const graphView = document.getElementById(`${factType}-graph-view`);
if (graphView && graphView.style.display !== 'none') {
reloadDataGraph(factType);
}
return data;
} catch (e) {
console.error(`Error loading ${factType} data:`, e);
alert(`Error loading ${factType} data: ` + e.message);
}
}
// Reload graph for a specific fact type
window.reloadDataGraph = function(factType) {
const data = dataCache[factType];
if (!data) return;
const nodeLimit = parseInt(document.getElementById(`${factType}-node-limit`).value) || 50;
const layoutName = document.getElementById(`${factType}-layout-select`).value;
// Filter nodes to limit
const limitedNodes = data.nodes.slice(0, nodeLimit);
const nodeIds = new Set(limitedNodes.map(n => n.data.id));
// Filter edges to only include those between visible nodes
const limitedEdges = data.edges.filter(e =>
nodeIds.has(e.data.source) && nodeIds.has(e.data.target)
);
// Update count display
document.getElementById(`${factType}-node-count`).textContent =
`Showing ${limitedNodes.length} of ${data.nodes.length} nodes`;
// Destroy existing graph if any
if (dataGraphs[factType]) {
dataGraphs[factType].destroy();
}
// Layout configurations
const layouts = {
'circle': {
name: 'circle',
animate: false,
radius: 300,
spacingFactor: 1.5
},
'grid': {
name: 'grid',
animate: false,
rows: Math.ceil(Math.sqrt(limitedNodes.length)),
cols: Math.ceil(Math.sqrt(limitedNodes.length)),
spacingFactor: 2
},
'cose': {
name: 'cose',
animate: false,
nodeRepulsion: 15000,
idealEdgeLength: 150,
edgeElasticity: 100,
nestingFactor: 1.2,
gravity: 1,
numIter: 1000,
initialTemp: 200,
coolingFactor: 0.95,
minTemp: 1.0
}
};
// Initialize Cytoscape
dataGraphs[factType] = cytoscape({
container: document.getElementById(`${factType}-cy`),
elements: [
...limitedNodes.map(n => ({ data: n.data })),
...limitedEdges.map(e => ({ data: e.data }))
],
style: [
{
selector: 'node',
style: {
'background-color': 'data(color)',
'label': 'data(label)',
'text-valign': 'center',
'text-halign': 'center',
'font-size': '10px',
'font-weight': 'bold',
'text-wrap': 'wrap',
'text-max-width': '100px',
'width': 40,
'height': 40,
'border-width': 2,
'border-color': '#333'
}
},
{
selector: 'edge',
style: {
'width': 1,
'line-color': 'data(color)',
'line-style': 'data(lineStyle)',
'target-arrow-shape': 'triangle',
'target-arrow-color': 'data(color)',
'curve-style': 'bezier',
'opacity': 0.7
}
},
{
selector: 'node:selected',
style: {
'border-width': 4,
'border-color': '#000'
}
}
],
layout: layouts[layoutName] || layouts['circle']
});
// Add tooltip on hover
let tooltip = null;
dataGraphs[factType].on('mouseover', 'node', function(evt) {
const node = evt.target;
const data = node.data();
const renderedPosition = node.renderedPosition();
tooltip = document.createElement('div');
tooltip.className = 'tooltip';
tooltip.innerHTML = `
<b>Text:</b> ${data.text}<br>
<b>Context:</b> ${data.context}<br>
<b>Date:</b> ${data.date}<br>
<b>Entities:</b> ${data.entities}
`;
tooltip.style.left = renderedPosition.x + 20 + 'px';
tooltip.style.top = renderedPosition.y + 'px';
document.body.appendChild(tooltip);
});
dataGraphs[factType].on('mouseout', 'node', function(evt) {
if (tooltip) {
tooltip.remove();
tooltip = null;
}
});
}
// Update table for a specific fact type
function updateDataTable(factType, data) {
if (!data) return;
const tbody = document.getElementById(`${factType}-table-body`);
const countSpan = document.getElementById(`${factType}-table-count`);
if (countSpan) {
countSpan.textContent = `(${data.total_units})`;
}
if (tbody) {
tbody.innerHTML = data.table_rows.map(row => `
<tr>
<td>${row.id}</td>
<td>${row.text}</td>
<td>${row.context}</td>
<td>${row.date}</td>
<td>${row.entities}</td>
</tr>
`).join('');
}
// Setup table filter
const filterInput = document.getElementById(`${factType}-table-filter`);
if (filterInput) {
filterInput.removeEventListener('input', filterInput._filterHandler);
filterInput._filterHandler = function() {
const filterValue = this.value.toLowerCase();
const rows = document.querySelectorAll(`#${factType}-table-body tr`);
rows.forEach(row => {
const text = row.textContent.toLowerCase();
if (text.includes(filterValue)) {
row.style.display = '';
} else {
row.style.display = 'none';
}
});
};
filterInput.addEventListener('input', filterInput._filterHandler);
}
}
// Load data from API (old function - kept for backward compatibility)
async function loadGraphData() {
try {
const response = await fetch('/api/graph');
// Require agent selection
if (!currentAgentId) {
alert('Please select an agent first');
return;
}
// Build URL with agent filter
let url = `/api/graph?agent_id=${encodeURIComponent(currentAgentId)}`;
const response = await fetch(url);
if (!response.ok) {
const error = await response.json();
@ -245,41 +601,6 @@ async function loadAgents() {
}
}
// Tab switching
function switchTab(tabName) {
document.querySelectorAll('.tab-content').forEach(tab => {
tab.classList.remove('active');
});
document.querySelectorAll('.tab-button').forEach(btn => {
btn.classList.remove('active');
});
if (tabName === 'graph') {
document.getElementById('graph-tab').classList.add('active');
document.querySelectorAll('.tab-button')[0].classList.add('active');
if (cy) cy.resize();
} else if (tabName === 'table') {
document.getElementById('table-tab').classList.add('active');
document.querySelectorAll('.tab-button')[1].classList.add('active');
} else if (tabName === 'debug') {
document.getElementById('debug-tab').classList.add('active');
document.querySelectorAll('.tab-button')[2].classList.add('active');
// Initialize with one pane if empty
if (debugPanes.length === 0) {
addDebugPane();
}
// Resize all debug graphs
debugPanes.forEach(pane => {
if (pane.cy) {
pane.cy.resize();
}
});
} else if (tabName === 'locomo') {
document.getElementById('locomo-tab').classList.add('active');
document.querySelectorAll('.tab-button')[3].classList.add('active');
}
}
// Debug pane management
function addDebugPane() {
const paneId = debugPaneCounter++;
@ -299,6 +620,14 @@ function addDebugPane() {
<label style="font-weight: bold; display: block; margin-bottom: 3px; font-size: 12px;">Query:</label>
<input type="text" id="search-query-${paneId}" placeholder="Enter search query..." style="width: 250px; padding: 6px; border: 1px solid #ccc; border-radius: 4px; font-size: 12px;">
</div>
<div>
<label style="font-weight: bold; display: block; margin-bottom: 3px; font-size: 12px;">Search Type:</label>
<select id="search-type-${paneId}" style="width: 120px; padding: 6px; border: 1px solid #ccc; border-radius: 4px; font-size: 12px;">
<option value="all">All Facts</option>
<option value="world">World Facts</option>
<option value="agent">Agent Facts</option>
</select>
</div>
<div>
<label style="font-weight: bold; display: block; margin-bottom: 3px; font-size: 12px;">Agent:</label>
<select id="search-agent-${paneId}" style="width: 120px; padding: 6px; border: 1px solid #ccc; border-radius: 4px; font-size: 12px;">
@ -459,6 +788,7 @@ window.runSearchInPane = async function(paneId) {
if (!pane) return;
const query = document.getElementById(`search-query-${paneId}`).value;
const searchType = document.getElementById(`search-type-${paneId}`).value;
const agentId = document.getElementById(`search-agent-${paneId}`).value;
const thinkingBudget = parseInt(document.getElementById(`search-budget-${paneId}`).value);
const topK = parseInt(document.getElementById(`search-top-k-${paneId}`).value);
@ -471,9 +801,17 @@ window.runSearchInPane = async function(paneId) {
}
try {
// Determine endpoint based on search type
let endpoint = '/api/search';
if (searchType === 'world') {
endpoint = '/api/world_search';
} else if (searchType === 'agent') {
endpoint = '/api/agent_search';
}
statusBar.innerHTML = '<span style="color: #ff9800;">🔄 Searching...</span>';
const response = await fetch('/api/search', {
const response = await fetch(endpoint, {
method: 'POST',
headers: {
'Content-Type': 'application/json'
@ -1139,19 +1477,281 @@ function highlightMatchingNodes(paneId, searchText) {
}
}
// Table filtering
document.getElementById('table-filter').addEventListener('input', function() {
const filterValue = this.value.toLowerCase();
const rows = document.querySelectorAll('#memory-table tbody tr');
// Load agents into global selector
async function loadGlobalAgents() {
try {
console.log('Loading global agents...'); // Debug
const select = document.getElementById('global-agent-selector');
rows.forEach(row => {
const text = row.textContent.toLowerCase();
if (text.includes(filterValue)) {
row.style.display = '';
if (!select) {
console.error('global-agent-selector element not found!');
return;
}
console.log('Fetching /api/agents...'); // Debug
const response = await fetch('/api/agents');
if (!response.ok) {
throw new Error(`HTTP ${response.status}: ${response.statusText}`);
}
const data = await response.json();
console.log('Agents data received:', data); // Debug
// Start with placeholder
select.innerHTML = '<option value="">Select an agent...</option>';
if (data.agents && data.agents.length > 0) {
console.log(`Adding ${data.agents.length} agents to dropdown`); // Debug
data.agents.forEach(agent => {
const option = document.createElement('option');
option.value = agent;
option.textContent = agent;
select.appendChild(option);
});
// Auto-select first agent
select.value = data.agents[0];
currentAgentId = data.agents[0];
console.log('Auto-selected agent:', currentAgentId); // Debug
// Update UI to show agent is selected
updateUIForAgentSelection();
} else {
row.style.display = 'none';
console.warn('No agents found in response');
}
} catch (e) {
console.error('Error loading global agents:', e);
alert('Failed to load agents: ' + e.message);
}
}
// Handle global agent selection change
function onGlobalAgentChange() {
const select = document.getElementById('global-agent-selector');
currentAgentId = select.value || null;
// Show/hide UI elements based on agent selection
updateUIForAgentSelection();
// Refresh all tabs if agent is selected
if (currentAgentId) {
refreshAllTabs();
}
}
// Update UI visibility based on agent selection
function updateUIForAgentSelection() {
const hasAgent = !!currentAgentId;
// Update each data subtab
['world', 'agent', 'opinions'].forEach(factType => {
const noAgentMsg = document.getElementById(`${factType}-no-agent-message`);
const graphView = document.getElementById(`${factType}-graph-view`);
const tableView = document.getElementById(`${factType}-table-view`);
if (noAgentMsg) noAgentMsg.style.display = hasAgent ? 'none' : 'block';
if (graphView) graphView.style.display = hasAgent ? 'none' : 'none'; // Start hidden, load on demand
if (tableView) tableView.style.display = hasAgent ? 'none' : 'none'; // Start hidden, load on demand
});
}
// Refresh all tabs with new agent context
async function refreshAllTabs() {
// Clear existing data
dataCache = {
world: null,
agent: null,
opinions: null
};
// Destroy existing graphs
['world', 'agent', 'opinions'].forEach(factType => {
if (dataGraphs[factType]) {
dataGraphs[factType].destroy();
dataGraphs[factType] = null;
}
});
// Update active debug panes with new agent
debugPanes.forEach(pane => {
const agentSelect = document.getElementById(`search-agent-${pane.id}`);
if (agentSelect && currentAgentId) {
agentSelect.value = currentAgentId;
}
});
}
// Run Think query
window.runThink = async function() {
console.log('runThink called'); // Debug log
const query = document.getElementById('think-query').value;
const agentSelect = document.getElementById('global-agent-selector');
const agentId = agentSelect ? agentSelect.value : null;
const thinkingBudget = parseInt(document.getElementById('think-budget').value);
const topK = parseInt(document.getElementById('think-top-k').value);
console.log('Query:', query, 'Agent:', agentId); // Debug log
if (!query || query.trim() === '') {
alert('Please enter a question');
return;
}
if (!agentId) {
alert('Please select an agent from the breadcrumb');
return;
}
const resultDiv = document.getElementById('think-result');
const loadingDiv = document.getElementById('think-loading');
if (!resultDiv || !loadingDiv) {
console.error('Think result divs not found');
return;
}
try {
// Show loading
resultDiv.style.display = 'none';
loadingDiv.style.display = 'block';
console.log('Calling /api/think with', { query, agentId, thinkingBudget, topK }); // Debug log
const response = await fetch('/api/think', {
method: 'POST',
headers: {
'Content-Type': 'application/json'
},
body: JSON.stringify({
query: query,
agent_id: agentId,
thinking_budget: thinkingBudget,
top_k: topK
})
});
console.log('Response status:', response.status); // Debug log
const data = await response.json();
console.log('Response data:', data); // Debug log
if (data.detail) {
alert('Error: ' + data.detail);
loadingDiv.style.display = 'none';
return;
}
// Display answer
document.getElementById('think-answer-text').textContent = data.text;
// Display world facts
const worldFactsDiv = document.getElementById('think-world-facts');
if (data.based_on.world && data.based_on.world.length > 0) {
worldFactsDiv.innerHTML = data.based_on.world.map((fact, idx) => `
<div style="margin-bottom: 10px; padding: 10px; background: white; border-radius: 4px; border-left: 3px solid #1976d2;">
<div style="font-size: 13px; color: #333; margin-bottom: 5px;">${fact.text}</div>
<div style="font-size: 11px; color: #666;">
Score: ${fact.score ? fact.score.toFixed(4) : 'N/A'} |
${fact.context ? 'Context: ' + fact.context : ''}
</div>
</div>
`).join('');
} else {
worldFactsDiv.innerHTML = '<div style="color: #666; font-style: italic;">No world facts used</div>';
}
// Display agent facts
const agentFactsDiv = document.getElementById('think-agent-facts');
if (data.based_on.agent && data.based_on.agent.length > 0) {
agentFactsDiv.innerHTML = data.based_on.agent.map((fact, idx) => `
<div style="margin-bottom: 10px; padding: 10px; background: white; border-radius: 4px; border-left: 3px solid #f57c00;">
<div style="font-size: 13px; color: #333; margin-bottom: 5px;">${fact.text}</div>
<div style="font-size: 11px; color: #666;">
Score: ${fact.score ? fact.score.toFixed(4) : 'N/A'} |
${fact.context ? 'Context: ' + fact.context : ''}
</div>
</div>
`).join('');
} else {
agentFactsDiv.innerHTML = '<div style="color: #666; font-style: italic;">No agent facts used</div>';
}
// Display opinions
const opinionsDiv = document.getElementById('think-opinions');
if (data.based_on.opinion && data.based_on.opinion.length > 0) {
opinionsDiv.innerHTML = data.based_on.opinion.map((fact, idx) => `
<div style="margin-bottom: 10px; padding: 10px; background: white; border-radius: 4px; border-left: 3px solid #7b1fa2;">
<div style="font-size: 13px; color: #333; margin-bottom: 5px;">${fact.text}</div>
<div style="font-size: 11px; color: #666;">
Score: ${fact.score ? fact.score.toFixed(4) : 'N/A'} |
Confidence: ${fact.confidence_score ? (fact.confidence_score * 100).toFixed(1) + '%' : 'N/A'} |
${fact.context ? 'Context: ' + fact.context : ''}
</div>
</div>
`).join('');
} else {
opinionsDiv.innerHTML = '<div style="color: #666; font-style: italic;">No opinions used</div>';
}
// Display new opinions
const newOpinionsDiv = document.getElementById('think-new-opinions');
const newOpinionsListDiv = document.getElementById('think-new-opinions-list');
if (data.new_opinions && data.new_opinions.length > 0) {
newOpinionsListDiv.innerHTML = data.new_opinions.map((opinion, idx) => `
<div style="margin-bottom: 15px; padding: 15px; background: white; border-radius: 6px; border-left: 4px solid #4caf50; box-shadow: 0 2px 4px rgba(0,0,0,0.1);">
<div style="display: flex; align-items: center; margin-bottom: 8px;">
<span style="background: #4caf50; color: white; padding: 4px 8px; border-radius: 12px; font-size: 11px; font-weight: bold; margin-right: 10px;">NEW</span>
<span style="color: #666; font-size: 12px;">#${idx + 1}</span>
</div>
<div style="font-size: 14px; color: #333; line-height: 1.5;">${opinion}</div>
</div>
`).join('');
newOpinionsDiv.style.display = 'block';
} else {
newOpinionsDiv.style.display = 'none';
}
// Show result
loadingDiv.style.display = 'none';
resultDiv.style.display = 'block';
} catch (e) {
console.error('Error running Think:', e);
alert('Error: ' + e.message);
loadingDiv.style.display = 'none';
}
};
// Initialize global agent selector on page load
document.addEventListener('DOMContentLoaded', function() {
console.log('DOM Content Loaded - initializing...'); // Debug
// Add change listener to global agent selector
const agentSelector = document.getElementById('global-agent-selector');
if (agentSelector) {
console.log('Found global-agent-selector, adding change listener'); // Debug
agentSelector.addEventListener('change', onGlobalAgentChange);
} else {
console.error('global-agent-selector not found in DOM!'); // Debug
}
// Add Enter key listener for Think query input
const thinkQuery = document.getElementById('think-query');
if (thinkQuery) {
thinkQuery.addEventListener('keypress', function(e) {
if (e.key === 'Enter') {
runThink();
}
});
}
// Initialize UI visibility
updateUIForAgentSelection();
// Load agents
loadGlobalAgents();
});
// Don't auto-load data on page load - wait for user to click load button

359
web/templates/index.html
View file

@ -7,99 +7,211 @@
<link rel="stylesheet" href="/static/css/styles.css">
</head>
<body>
<div class="breadcrumb-container">
<div class="breadcrumb">
<span class="breadcrumb-item">Memory Graph</span>
<span class="breadcrumb-separator">/</span>
<span class="breadcrumb-item">Agent:</span>
<select id="global-agent-selector" class="agent-selector">
<option value="">Select an agent...</option>
</select>
</div>
</div>
<div class="tab-container">
<div class="tab-buttons">
<button class="tab-button active" onclick="switchTab('graph')">Graph View</button>
<button class="tab-button" onclick="switchTab('table')">Table View</button>
<button class="tab-button" onclick="switchTab('debug')">Search Debug</button>
<button class="tab-button" onclick="switchTab('locomo')">Locomo Benchmark</button>
<button class="tab-button active" onclick="switchMainTab('data')">Data</button>
<button class="tab-button" onclick="switchMainTab('debug')">Search Debug</button>
<button class="tab-button" onclick="switchMainTab('think')">Think</button>
<button class="tab-button" onclick="switchMainTab('benchmark')">Benchmark</button>
</div>
<div id="graph-tab" class="tab-content active">
<div style="padding: 15px; background: #f9f9f9; border-bottom: 2px solid #333;">
<div style="display: flex; gap: 15px; align-items: center; flex-wrap: wrap;">
<button onclick="loadGraphData()" style="padding: 8px 20px; background: #66bb6a; color: white; border: none; border-radius: 4px; cursor: pointer; font-weight: bold; font-size: 14px;">
📊 Load Graph Data
</button>
<div>
<label style="font-weight: bold; margin-right: 5px;">Limit nodes:</label>
<input type="number" id="node-limit" value="50" min="10" max="1000" step="10"
style="width: 80px; padding: 5px; border: 1px solid #ccc; border-radius: 4px;">
</div>
<div>
<label style="font-weight: bold; margin-right: 5px;">Layout:</label>
<select id="layout-select" style="padding: 5px; border: 1px solid #ccc; border-radius: 4px;">
<option value="circle">Circle (fast)</option>
<option value="grid">Grid (fast)</option>
<option value="cose">Force-directed (slow)</option>
</select>
</div>
<button onclick="reloadGraph()" style="padding: 6px 15px; background: #42a5f5; color: white; border: none; border-radius: 4px; cursor: pointer; font-weight: bold;">
Apply
</button>
<button onclick="loadGraphData()" class="refresh-button">
🔄 Refresh
</button>
<span id="node-count" style="color: #666; font-size: 14px;"></span>
</div>
<!-- Data Tab -->
<div id="data-tab" class="tab-content active">
<div class="data-sub-tabs">
<button class="data-sub-tab-button active" onclick="switchDataSubTab('world')">World</button>
<button class="data-sub-tab-button" onclick="switchDataSubTab('agent')">Agent</button>
<button class="data-sub-tab-button" onclick="switchDataSubTab('opinions')">Opinions</button>
</div>
<div id="cy"><div style="padding: 40px; text-align: center; color: #666;">
<p>Click "Load Graph Data" to visualize the memory graph</p>
</div></div>
<div class="legend">
<h3>Legend</h3>
<h4 style="margin: 10px 0 5px 0;">Link Types:</h4>
<div class="legend-item">
<div class="legend-line" style="background: #00bcd4; border-top: 1px dashed #00bcd4;"></div>
<span><b>Temporal</b></span>
</div>
<div class="legend-item">
<div class="legend-line" style="background: #ff69b4;"></div>
<span><b>Semantic</b></span>
</div>
<div class="legend-item">
<div class="legend-line" style="background: #ffd700;"></div>
<span><b>Entity</b></span>
</div>
<h4 style="margin: 15px 0 5px 0;">Nodes:</h4>
<div class="legend-item">
<div class="legend-node" style="background: #e0e0e0;"></div>
<span>No entities</span>
</div>
<div class="legend-item">
<div class="legend-node" style="background: #90caf9;"></div>
<span>1 entity</span>
</div>
<div class="legend-item">
<div class="legend-node" style="background: #42a5f5;"></div>
<span>2+ entities</span>
</div>
</div>
</div>
<div id="table-tab" class="tab-content">
<h2>Memory Units <span id="table-count"></span></h2>
<div style="margin-bottom: 15px;">
<button onclick="loadGraphData()" style="padding: 8px 20px; background: #66bb6a; color: white; border: none; border-radius: 4px; cursor: pointer; font-weight: bold; font-size: 14px;">
📊 Load Table Data
</button>
<!-- World, Agent, and Opinions subtabs share the same structure -->
<div id="world-subtab" class="data-subtab-content active">
<div class="view-toggle">
<button class="view-toggle-button active" onclick="switchDataView('world', 'graph')">Graph</button>
<button class="view-toggle-button" onclick="switchDataView('world', 'table')">Table</button>
</div>
<div id="world-no-agent-message" class="no-agent-message">
<h3>No Agent Selected</h3>
<p>Please select an agent from the dropdown above to view world facts.</p>
</div>
<div id="world-graph-view" class="data-view" style="display: none;">
<div class="data-controls">
<button onclick="loadDataView('world')" class="load-button">📊 Load World Facts</button>
<div>
<label>Limit nodes:</label>
<input type="number" id="world-node-limit" value="50" min="10" max="1000" step="10" class="control-input">
</div>
<div>
<label>Layout:</label>
<select id="world-layout-select" class="control-select">
<option value="circle">Circle (fast)</option>
<option value="grid">Grid (fast)</option>
<option value="cose">Force-directed (slow)</option>
</select>
</div>
<button onclick="reloadDataGraph('world')" class="apply-button">Apply</button>
<button onclick="loadDataView('world')" class="refresh-button">🔄 Refresh</button>
<span id="world-node-count" class="node-count"></span>
</div>
<div id="world-cy" class="graph-canvas"></div>
<div class="legend">
<h3>Legend</h3>
<h4>Link Types:</h4>
<div class="legend-item"><div class="legend-line temporal"></div><span><b>Temporal</b></span></div>
<div class="legend-item"><div class="legend-line semantic"></div><span><b>Semantic</b></span></div>
<div class="legend-item"><div class="legend-line entity"></div><span><b>Entity</b></span></div>
<h4>Nodes:</h4>
<div class="legend-item"><div class="legend-node no-entities"></div><span>No entities</span></div>
<div class="legend-item"><div class="legend-node one-entity"></div><span>1 entity</span></div>
<div class="legend-item"><div class="legend-node multi-entities"></div><span>2+ entities</span></div>
</div>
</div>
<div id="world-table-view" class="data-view" style="display: none;">
<div class="data-controls">
<h2>World Facts <span id="world-table-count"></span></h2>
<button onclick="loadDataView('world')" class="load-button">📊 Load World Facts</button>
</div>
<input type="text" id="world-table-filter" placeholder="Filter by text, context, or entities..." class="table-filter">
<div class="table-container">
<table class="memory-table">
<thead>
<tr><th>ID</th><th>Text</th><th>Context</th><th>Date</th><th>Entities</th></tr>
</thead>
<tbody id="world-table-body">
<tr><td colspan="5" class="empty-message">Click "Load World Facts" to view data</td></tr>
</tbody>
</table>
</div>
</div>
</div>
<input type="text" id="table-filter" placeholder="Filter by text, context, or entities...">
<div style="overflow-x: auto;">
<table id="memory-table">
<thead>
<tr>
<th>ID</th>
<th>Text</th>
<th>Context</th>
<th>Date</th>
<th>Entities</th>
</tr>
</thead>
<tbody id="table-body">
<tr><td colspan="5" style="padding: 40px; text-align: center; color: #666;">Click "Load Table Data" to view memory units</td></tr>
</tbody>
</table>
<div id="agent-subtab" class="data-subtab-content">
<div class="view-toggle">
<button class="view-toggle-button active" onclick="switchDataView('agent', 'graph')">Graph</button>
<button class="view-toggle-button" onclick="switchDataView('agent', 'table')">Table</button>
</div>
<div id="agent-no-agent-message" class="no-agent-message">
<h3>No Agent Selected</h3>
<p>Please select an agent from the dropdown above to view agent facts.</p>
</div>
<div id="agent-graph-view" class="data-view" style="display: none;">
<div class="data-controls">
<button onclick="loadDataView('agent')" class="load-button">📊 Load Agent Facts</button>
<div>
<label>Limit nodes:</label>
<input type="number" id="agent-node-limit" value="50" min="10" max="1000" step="10" class="control-input">
</div>
<div>
<label>Layout:</label>
<select id="agent-layout-select" class="control-select">
<option value="circle">Circle (fast)</option>
<option value="grid">Grid (fast)</option>
<option value="cose">Force-directed (slow)</option>
</select>
</div>
<button onclick="reloadDataGraph('agent')" class="apply-button">Apply</button>
<button onclick="loadDataView('agent')" class="refresh-button">🔄 Refresh</button>
<span id="agent-node-count" class="node-count"></span>
</div>
<div id="agent-cy" class="graph-canvas"></div>
<div class="legend">
<h3>Legend</h3>
<h4>Link Types:</h4>
<div class="legend-item"><div class="legend-line temporal"></div><span><b>Temporal</b></span></div>
<div class="legend-item"><div class="legend-line semantic"></div><span><b>Semantic</b></span></div>
<div class="legend-item"><div class="legend-line entity"></div><span><b>Entity</b></span></div>
<h4>Nodes:</h4>
<div class="legend-item"><div class="legend-node no-entities"></div><span>No entities</span></div>
<div class="legend-item"><div class="legend-node one-entity"></div><span>1 entity</span></div>
<div class="legend-item"><div class="legend-node multi-entities"></div><span>2+ entities</span></div>
</div>
</div>
<div id="agent-table-view" class="data-view" style="display: none;">
<div class="data-controls">
<h2>Agent Facts <span id="agent-table-count"></span></h2>
<button onclick="loadDataView('agent')" class="load-button">📊 Load Agent Facts</button>
</div>
<input type="text" id="agent-table-filter" placeholder="Filter by text, context, or entities..." class="table-filter">
<div class="table-container">
<table class="memory-table">
<thead>
<tr><th>ID</th><th>Text</th><th>Context</th><th>Date</th><th>Entities</th></tr>
</thead>
<tbody id="agent-table-body">
<tr><td colspan="5" class="empty-message">Click "Load Agent Facts" to view data</td></tr>
</tbody>
</table>
</div>
</div>
</div>
<div id="opinions-subtab" class="data-subtab-content">
<div class="view-toggle">
<button class="view-toggle-button active" onclick="switchDataView('opinions', 'graph')">Graph</button>
<button class="view-toggle-button" onclick="switchDataView('opinions', 'table')">Table</button>
</div>
<div id="opinions-no-agent-message" class="no-agent-message">
<h3>No Agent Selected</h3>
<p>Please select an agent from the dropdown above to view opinions.</p>
</div>
<div id="opinions-graph-view" class="data-view" style="display: none;">
<div class="data-controls">
<button onclick="loadDataView('opinions')" class="load-button">📊 Load Opinions</button>
<div>
<label>Limit nodes:</label>
<input type="number" id="opinions-node-limit" value="50" min="10" max="1000" step="10" class="control-input">
</div>
<div>
<label>Layout:</label>
<select id="opinions-layout-select" class="control-select">
<option value="circle">Circle (fast)</option>
<option value="grid">Grid (fast)</option>
<option value="cose">Force-directed (slow)</option>
</select>
</div>
<button onclick="reloadDataGraph('opinions')" class="apply-button">Apply</button>
<button onclick="loadDataView('opinions')" class="refresh-button">🔄 Refresh</button>
<span id="opinions-node-count" class="node-count"></span>
</div>
<div id="opinions-cy" class="graph-canvas"></div>
<div class="legend">
<h3>Legend</h3>
<h4>Link Types:</h4>
<div class="legend-item"><div class="legend-line temporal"></div><span><b>Temporal</b></span></div>
<div class="legend-item"><div class="legend-line semantic"></div><span><b>Semantic</b></span></div>
<div class="legend-item"><div class="legend-line entity"></div><span><b>Entity</b></span></div>
<h4>Nodes:</h4>
<div class="legend-item"><div class="legend-node no-entities"></div><span>No entities</span></div>
<div class="legend-item"><div class="legend-node one-entity"></div><span>1 entity</span></div>
<div class="legend-item"><div class="legend-node multi-entities"></div><span>2+ entities</span></div>
</div>
</div>
<div id="opinions-table-view" class="data-view" style="display: none;">
<div class="data-controls">
<h2>Opinions <span id="opinions-table-count"></span></h2>
<button onclick="loadDataView('opinions')" class="load-button">📊 Load Opinions</button>
</div>
<input type="text" id="opinions-table-filter" placeholder="Filter by text, context, or entities..." class="table-filter">
<div class="table-container">
<table class="memory-table">
<thead>
<tr><th>ID</th><th>Text</th><th>Context</th><th>Date</th><th>Entities</th></tr>
</thead>
<tbody id="opinions-table-body">
<tr><td colspan="5" class="empty-message">Click "Load Opinions" to view data</td></tr>
</tbody>
</table>
</div>
</div>
</div>
</div>
@ -113,18 +225,77 @@
</div>
</div>
<div id="locomo-tab" class="tab-content">
<h2>Locomo Benchmark Results</h2>
<div id="think-tab" class="tab-content">
<h2>Think - AI-Powered Answers</h2>
<p style="color: #666; margin-bottom: 15px;">
Analyze benchmark results and debug incorrect answers.
Ask questions and get AI-generated answers based on agent identity and world facts.
</p>
<div style="margin-bottom: 15px;">
<button onclick="loadLocomoResults()" style="padding: 8px 20px; background: #66bb6a; color: white; border: none; border-radius: 4px; cursor: pointer; font-weight: bold; font-size: 14px;">
📊 Load Benchmark Results
<div class="think-controls">
<div style="display: flex; gap: 15px; align-items: flex-end; flex-wrap: wrap; margin-bottom: 15px;">
<div style="flex: 1; min-width: 300px;">
<label style="font-weight: bold; display: block; margin-bottom: 5px;">Question:</label>
<input type="text" id="think-query" placeholder="Enter your question..." style="width: 100%; padding: 10px; border: 2px solid #ccc; border-radius: 4px; font-size: 14px;">
</div>
<div>
<label style="font-weight: bold; display: block; margin-bottom: 5px;">Budget:</label>
<input type="number" id="think-budget" value="50" min="10" max="1000" style="width: 80px; padding: 10px; border: 2px solid #ccc; border-radius: 4px; font-size: 14px;">
</div>
<div>
<label style="font-weight: bold; display: block; margin-bottom: 5px;">Top K:</label>
<input type="number" id="think-top-k" value="10" min="1" max="50" style="width: 70px; padding: 10px; border: 2px solid #ccc; border-radius: 4px; font-size: 14px;">
</div>
<button id="think-button" onclick="runThink()" style="padding: 10px 24px; background: #66bb6a; color: white; border: none; border-radius: 4px; cursor: pointer; font-weight: bold; font-size: 14px;">
💭 Think
</button>
</div>
</div>
<div id="think-result" style="display: none; margin-top: 20px;">
<div class="think-answer">
<h3 style="margin-top: 0; color: #333; border-bottom: 2px solid #333; padding-bottom: 10px;">Answer</h3>
<div id="think-answer-text" style="padding: 15px; background: #f9f9f9; border-left: 4px solid #66bb6a; font-size: 15px; line-height: 1.6; white-space: pre-wrap;"></div>
</div>
<div class="think-sources" style="margin-top: 30px;">
<h3 style="margin-top: 0; color: #333; border-bottom: 2px solid #333; padding-bottom: 10px;">Based On</h3>
<div style="display: grid; grid-template-columns: 1fr 1fr 1fr; gap: 20px; margin-top: 15px;">
<div>
<h4 style="margin: 0 0 10px 0; color: #1976d2;">World Facts (General Knowledge)</h4>
<div id="think-world-facts" style="background: #e3f2fd; padding: 15px; border-radius: 4px; border: 2px solid #1976d2; min-height: 100px;"></div>
</div>
<div>
<h4 style="margin: 0 0 10px 0; color: #f57c00;">Agent Facts (Identity)</h4>
<div id="think-agent-facts" style="background: #fff3e0; padding: 15px; border-radius: 4px; border: 2px solid #f57c00; min-height: 100px;"></div>
</div>
<div>
<h4 style="margin: 0 0 10px 0; color: #7b1fa2;">Opinions (Agent Beliefs)</h4>
<div id="think-opinions" style="background: #f3e5f5; padding: 15px; border-radius: 4px; border: 2px solid #7b1fa2; min-height: 100px;"></div>
</div>
</div>
</div>
<div id="think-new-opinions" style="display: none; margin-top: 30px;">
<div style="background: #e8f5e9; padding: 20px; border-radius: 8px; border: 2px solid #4caf50;">
<h3 style="margin-top: 0; color: #2e7d32; border-bottom: 2px solid #4caf50; padding-bottom: 10px;">✨ New Opinions Formed</h3>
<div id="think-new-opinions-list" style="margin-top: 15px;"></div>
</div>
</div>
</div>
<div id="think-loading" style="display: none; text-align: center; padding: 40px; color: #666;">
<div style="font-size: 48px; margin-bottom: 10px;">💭</div>
<div style="font-size: 18px;">Thinking...</div>
</div>
</div>
<div id="benchmark-tab" class="tab-content">
<h2>Benchmark</h2>
<p style="color: #666; margin-bottom: 15px; padding: 0 20px;">
Run and analyze benchmark results.
</p>
<div style="margin-bottom: 15px; padding: 0 20px;">
<button onclick="loadLocomoResults()" class="load-button">
📊 Load Locomo Benchmark
</button>
</div>
<div id="locomo-content">
<p style="padding: 20px; text-align: center; color: #666;">Click "Load Benchmark Results" to view the data</p>
<div id="locomo-content" style="padding: 20px;">
<p style="text-align: center; color: #666;">Click "Load Locomo Benchmark" to view results</p>
</div>
</div>
</div>