From f988aae9962e69f2ed07a373193e91e8f19add66 Mon Sep 17 00:00:00 2001
From: doria <93405247+dr-frmr@users.noreply.github.com>
Date: Tue, 7 Oct 2025 15:28:44 -0400
Subject: [PATCH] create Representation class and use it to unify all
formatting (#214)
* feat: add optional JWT and webhook secrets to honcho instance creation
* chore: ignore spurious warnings
* feat: add response format if using gpt-5 model family
* feat: add response models to all apis except anthropic
* fix: raise NotImplementedError for response models in AsyncAnthropic client
* chore: address review
* [WIP] representation structure + deriver cleanup
* chore: add tests, cleanup
* feat: [WIP: semi-working] representation object
* fix: alignment
* fix: make observations hashable for dedup
* fix: datetime formatting, observation counting
* fix: switch to int for message id, clean up representation
* feat: remove need for metadata working rep
* chore: cleanup
* fix: use tenacity instead of custom fns
* feat: add representation and card to context if desired
* feat: add semantically relevant observations
* fix: pass all params to streaming, nonblocking streaming
* feat: consolidate document saving, make working representation fetching much smarter
* chore: add 100% test coverage of representation util
* feat: basic dream infra
* feat: dream queue item first pass
* chore: fixes & cleanup from coderabbit
* fix: dreams scheduled when new document count reaches a certain threshold
* feat: wip: timed dreams (not working)
* fix: test
* fix: remove useless pyright ignore
* fix: executing dreams
* feat: dreaming
* feat: [WIP] longmemeval bench
* feat: add USE_PEER_CARD setting, fix longmem test driver
* feat: get full working rep for dialectic in one swoop -- fix representation_from_documents to use the proper timestamp!
* fix: timestamps for real, handle assistant qs in longmem
* fix: remove old client, add batching to longmem
* perf: remove duplicate detection, will move to background task
* feat: track perf metrics on evals
* feat: adjust deriver prompt to use peer_id, add question date to question, clean up deriver
* fix: label metrics by task for better perf trace
* chore: code review
* feat: add efficiency score to longmem bench
* chore: tuning and cleaning up eval
* chore: bring in the big prompts
* feat: add support for vllm client
* feat: perf: bundle db calls in deriver and dialectic, increase max conns in docker db
* feat: add merge-sessions flag to longmemeval, add SUMMARY_ENABLED flag
* fix: COLLECT_METRICS default false
* chore: display start/end message ids, don't include in metrics
* fix: break large messages apart for eval
* fix: only get/create collection when needed
* feat: properly attribute documents with message id ranges and add session name column to documents
* fix: revert move of get_or_create_collection (need for fkey)
* fix: always get collection with peer name even if it's none
* chore: coderabbit
* fix: give peer card its own config, expand document schema, refactor get_context to be parallel, various cleanup chores and bugfixes
* chore: refactor: reify observer/observed system across entire codebase, including db migration
* refactor: cleanup code organization, make singletons where desired
* refactor: replace embeddings store with representation manager
* chore: coderabbit cleanup
* chore: update migration to non-null session param in documents, general review and cleanup
* chore: merge branch 'main' into ben/deriver-tidy
* chore: review fixes
---
.env.template | 12 +-
.gitignore | 2 +
CLAUDE.md | 30 +-
config.toml.example | 12 +-
docker-compose.yml.example | 1 +
..._replace_collection_name_with_observer_.py | 564 ++++++++
...c5_add_session_name_column_to_documents.py | 104 ++
pyproject.toml | 1 +
scripts/generate_message_embeddings.py | 12 +-
sdks/python/CHANGELOG.md | 2 +-
sdks/typescript/CHANGELOG.md | 2 +-
sdks/typescript/src/session_context.ts | 24 +-
src/config.py | 52 +-
src/crud/__init__.py | 24 +-
src/crud/collection.py | 36 +-
src/crud/deriver.py | 34 +-
src/crud/document.py | 183 ++-
src/crud/peer_card.py | 96 ++
src/crud/representation.py | 886 ++++++++-----
src/dependencies.py | 5 +-
src/deriver/consumer.py | 59 +-
src/deriver/deriver.py | 687 +++-------
src/deriver/enqueue.py | 74 +-
src/deriver/prompts.py | 72 +-
src/deriver/queue_manager.py | 77 +-
src/deriver/utils.py | 82 --
src/dialectic/chat.py | 253 ++--
src/dialectic/prompts.py | 88 +-
src/dialectic/utils.py | 246 ----
src/dreamer/__init__.py | 13 +
src/dreamer/dream_scheduler.py | 354 +++++
src/dreamer/dreamer.py | 195 +++
src/dreamer/prompts.py | 26 +
src/embedding_client.py | 207 ++-
src/models.py | 50 +-
src/routers/peers.py | 28 +-
src/routers/sessions.py | 167 ++-
src/routers/workspaces.py | 4 +-
src/schemas.py | 48 +-
src/utils/clients.py | 170 ++-
src/utils/embedding_store.py | 458 -------
src/utils/formatting.py | 189 +--
src/utils/json_parser.py | 378 ++++++
src/utils/langfuse_client.py | 40 +
src/utils/logging.py | 195 +--
src/utils/metrics_collector.py | 364 ++++++
src/utils/peer_card.py | 25 +
src/{deriver => utils}/queue_payload.py | 50 +-
src/utils/representation.py | 337 +++++
src/utils/shared_models.py | 205 ---
src/utils/summarizer.py | 4 +-
src/utils/tokens.py | 15 +
src/utils/types.py | 2 +-
src/utils/work_unit.py | 110 ++
src/webhooks/events.py | 13 +-
src/webhooks/webhook_delivery.py | 2 +-
tests/bench/.gitignore | 3 +
tests/bench/longmem.py | 1151 +++++++++++++++++
tests/bench/peer_card_bench.py | 30 +-
tests/bench/run_tests.py | 34 +-
tests/conftest.py | 88 +-
tests/crud/test_peer_card.py | 451 +++++++
tests/deriver/README.md | 12 +-
tests/deriver/conftest.py | 36 +-
tests/deriver/test_deriver_processing.py | 62 +-
tests/deriver/test_queue_processing.py | 62 +-
tests/deriver/test_representation_crud.py | 664 ++--------
tests/dreamer/__init__.py | 0
tests/dreamer/test_dream_scheduler.py | 64 +
tests/dreamer/test_queue_manager_singleton.py | 56 +
tests/integration/test_enqueue.py | 126 +-
tests/integration/test_representation.py | 533 ++++++++
tests/routes/test_peers.py | 8 +-
tests/routes/test_queue_status.py | 20 +-
tests/test_datetime_parsing.py | 34 +-
tests/test_llm_mock.py | 25 +-
tests/test_schema_validations.py | 40 +-
tests/utils/test_clients.py | 18 +-
uv.lock | 11 +
79 files changed, 7254 insertions(+), 3643 deletions(-)
create mode 100644 migrations/versions/08894082221a_replace_collection_name_with_observer_.py
create mode 100644 migrations/versions/564ba40505c5_add_session_name_column_to_documents.py
create mode 100644 src/crud/peer_card.py
delete mode 100644 src/deriver/utils.py
delete mode 100644 src/dialectic/utils.py
create mode 100644 src/dreamer/__init__.py
create mode 100644 src/dreamer/dream_scheduler.py
create mode 100644 src/dreamer/dreamer.py
create mode 100644 src/dreamer/prompts.py
delete mode 100644 src/utils/embedding_store.py
create mode 100644 src/utils/json_parser.py
create mode 100644 src/utils/langfuse_client.py
create mode 100644 src/utils/metrics_collector.py
create mode 100644 src/utils/peer_card.py
rename src/{deriver => utils}/queue_payload.py (78%)
create mode 100644 src/utils/representation.py
delete mode 100644 src/utils/shared_models.py
create mode 100644 src/utils/tokens.py
create mode 100644 src/utils/work_unit.py
create mode 100644 tests/bench/.gitignore
create mode 100644 tests/bench/longmem.py
create mode 100644 tests/crud/test_peer_card.py
create mode 100644 tests/dreamer/__init__.py
create mode 100644 tests/dreamer/test_dream_scheduler.py
create mode 100644 tests/dreamer/test_queue_manager_singleton.py
create mode 100644 tests/integration/test_representation.py
diff --git a/.env.template b/.env.template
index 2ddb6802..e3be457d 100644
--- a/.env.template
+++ b/.env.template
@@ -86,13 +86,18 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here
# DERIVER_MAX_OUTPUT_TOKENS=2500
# only applied when using Anthropic as provider
# DERIVER_THINKING_BUDGET_TOKENS=1024
-# DERIVER_PEER_CARD_PROVIDER=openai
-# DERIVER_PEER_CARD_MODEL=gpt-5-nano-2025-08-07
-# DERIVER_PEER_CARD_MAX_OUTPUT_TOKENS=2000
# DERIVER_WORKING_REPRESENTATION_MAX_OBSERVATIONS=100
# DERIVER_REPRESENTATION_BATCH_MAX_TOKENS=4096
# DERIVER_MAX_INPUT_TOKENS=23000
+# =============================================================================
+# Peer Card Configuration
+# =============================================================================
+# ENABLED=true
+# PROVIDER=openai
+# MODEL=gpt-5-nano-2025-08-07
+# MAX_OUTPUT_TOKENS=4000
+
# =============================================================================
# Dialectic Settings
# =============================================================================
@@ -110,6 +115,7 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here
# =============================================================================
# Summary Settings
# =============================================================================
+# SUMMARY_ENABLED=true
# SUMMARY_MESSAGES_PER_SHORT_SUMMARY=20
# SUMMARY_MESSAGES_PER_LONG_SUMMARY=60
# SUMMARY_PROVIDER=google
diff --git a/.gitignore b/.gitignore
index 8e15cb72..f75448c4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -185,4 +185,6 @@ config.toml
CRUSH.md
.crush/
+
+metrics.jsonl
AGENTS.md
diff --git a/CLAUDE.md b/CLAUDE.md
index 4adce6ac..0ffa2017 100644
--- a/CLAUDE.md
+++ b/CLAUDE.md
@@ -115,21 +115,22 @@ src/
├── embedding_client.py # Embedding service client
├── crud/ # Database operations
│ ├── __init__.py
-│ ├── collection.py # Collection CRUD operations
-│ ├── deriver.py # Deriver-related CRUD operations
-│ ├── document.py # Document CRUD operations
-│ ├── message.py # Message CRUD operations
-│ ├── peer.py # Peer CRUD operations
-│ ├── representation.py # Representation CRUD operations
-│ ├── session.py # Session CRUD operations
-│ ├── webhook.py # Webhook CRUD operations
-│ └── workspace.py # Workspace CRUD operations
-├── dialectic/ # Dialectic API implementation
+│ ├── collection.py # Collection CRUD operations
+│ ├── deriver.py # Deriver-related CRUD operations
+│ ├── document.py # Document CRUD operations
+│ ├── message.py # Message CRUD operations
+│ ├── peer.py # Peer CRUD operations
+│ ├── peer_card.py # Peer Card CRUD operations
+│ ├── representation.py # RepresentationManager and representation operations
+│ ├── session.py # Session CRUD operations
+│ ├── webhook.py # Webhook CRUD operations
+│ └── workspace.py # Workspace CRUD operations
+├── dialectic/ # Dialectic API implementation
│ ├── __init__.py
-│ ├── chat.py # Chat functionality
-│ ├── prompts.py # Prompt templates
-│ └── utils.py # Dialectic utilities
-├── routers/ # API endpoints
+│ ├── chat.py # Chat functionality
+│ ├── prompts.py # Prompt templates
+│ └── utils.py # Dialectic utilities
+├── routers/ # API endpoints
│ ├── workspaces.py
│ ├── peers.py
│ ├── sessions.py
@@ -149,7 +150,6 @@ src/
├── utils/ # Utilities
│ ├── __init__.py
│ ├── clients.py # LLM client abstraction
-│ ├── embedding_store.py # Vector storage management
│ ├── files.py # File handling utilities
│ ├── filter.py # Query filtering utilities
│ ├── formatting.py # Message formatting utilities
diff --git a/config.toml.example b/config.toml.example
index 6209f255..a26d4672 100644
--- a/config.toml.example
+++ b/config.toml.example
@@ -66,13 +66,18 @@ PROVIDER = "google"
MODEL = "gemini-2.0-flash-lite"
MAX_OUTPUT_TOKENS = 2500
THINKING_BUDGET_TOKENS = 1024 # only applied when using Anthropic
-PEER_CARD_PROVIDER = "openai"
-PEER_CARD_MODEL = "gpt-5-nano-2025-08-07"
-PEER_CARD_MAX_OUTPUT_TOKENS = 2000
+
WORKING_REPRESENTATION_MAX_OBSERVATIONS = 100
REPRESENTATION_BATCH_MAX_TOKENS = 4096
MAX_INPUT_TOKENS = 23000
+# Peer card settings
+[peer_card]
+ENABLED = true
+PROVIDER = "openai"
+MODEL = "gpt-5-nano-2025-08-07"
+MAX_OUTPUT_TOKENS = 4000
+
# Dialectic settings
[dialectic]
PROVIDER = "anthropic"
@@ -88,6 +93,7 @@ CONTEXT_WINDOW_SIZE = 100000
# Summary settings
[summary]
+ENABLED = true
MESSAGES_PER_SHORT_SUMMARY = 20
MESSAGES_PER_LONG_SUMMARY = 60
PROVIDER = "google"
diff --git a/docker-compose.yml.example b/docker-compose.yml.example
index 452c9697..e11e93ad 100644
--- a/docker-compose.yml.example
+++ b/docker-compose.yml.example
@@ -30,6 +30,7 @@ services:
restart: always
ports:
- 5432:5432
+ command: ["postgres", "-c", "max_connections=800"]
environment:
- POSTGRES_DB=honcho
- POSTGRES_USER=testuser
diff --git a/migrations/versions/08894082221a_replace_collection_name_with_observer_.py b/migrations/versions/08894082221a_replace_collection_name_with_observer_.py
new file mode 100644
index 00000000..9b387574
--- /dev/null
+++ b/migrations/versions/08894082221a_replace_collection_name_with_observer_.py
@@ -0,0 +1,564 @@
+"""replace collection name with observer_observed
+
+Revision ID: 08894082221a
+Revises: 564ba40505c5
+Create Date: 2025-10-03 16:48:13.270834
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+from sqlalchemy import text
+
+from migrations.utils import column_exists, constraint_exists, fk_exists, index_exists
+from src.config import settings
+
+# revision identifiers, used by Alembic.
+revision: str = "08894082221a"
+down_revision: str | None = "564ba40505c5"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+ """Replace collections.name and documents.collection_name with observer and observed fields."""
+ schema = settings.DB.SCHEMA
+ inspector = sa.inspect(op.get_bind())
+ connection = op.get_bind()
+
+ # SESSION_NAME MIGRATION
+ # Replace NULL session_name values with empty strings and make column non-nullable
+ # This only applies to documents table
+
+ # Update documents table
+ connection.execute(
+ text(
+ """
+ UPDATE documents
+ SET session_name = ''
+ WHERE session_name IS NULL
+ """
+ )
+ )
+ if column_exists("documents", "session_name", inspector):
+ op.alter_column("documents", "session_name", nullable=False, schema=schema)
+
+ # COLLECTIONS TABLE
+ # Step 1: Add new observer and observed columns to collections
+ if not column_exists("collections", "observer", inspector):
+ op.add_column(
+ "collections",
+ sa.Column("observer", sa.TEXT(), nullable=True),
+ schema=schema,
+ )
+
+ if not column_exists("collections", "observed", inspector):
+ op.add_column(
+ "collections",
+ sa.Column("observed", sa.TEXT(), nullable=True),
+ schema=schema,
+ )
+
+ # Step 2: Populate collections observer and observed from existing name field
+ # The logic is:
+ # - observer = peer_name (the exact peer ID)
+ # - If name is "global_representation", observed = peer_name
+ # - Otherwise, observed = name with the "observer_" prefix stripped
+ connection.execute(
+ text(
+ """
+ UPDATE collections
+ SET
+ observer = peer_name,
+ observed = CASE
+ WHEN name = 'global_representation' THEN peer_name
+ ELSE substring(name from length(peer_name) + 2)
+ END
+ WHERE observer IS NULL OR observed IS NULL
+ """
+ )
+ )
+
+ # Step 3: Make collections observer and observed NOT NULL
+ op.alter_column("collections", "observer", nullable=False, schema=schema)
+ op.alter_column("collections", "observed", nullable=False, schema=schema)
+
+ # DOCUMENTS TABLE
+ # Step 4: Add new observer and observed columns to documents
+ if not column_exists("documents", "observer", inspector):
+ op.add_column(
+ "documents",
+ sa.Column("observer", sa.TEXT(), nullable=True),
+ schema=schema,
+ )
+
+ if not column_exists("documents", "observed", inspector):
+ op.add_column(
+ "documents",
+ sa.Column("observed", sa.TEXT(), nullable=True),
+ schema=schema,
+ )
+
+ # Step 5: Populate documents observer and observed from collections table
+ # Join to the already-populated collections table to get authoritative values
+ # Process in batches of 1000 to reduce query size
+ batch_size = 1000
+ while True:
+ result = connection.execute(
+ text(
+ """
+ WITH batch AS (
+ SELECT d.ctid
+ FROM documents d
+ WHERE d.observer IS NULL OR d.observed IS NULL
+ LIMIT :batch_size
+ )
+ UPDATE documents d
+ SET
+ observer = c.observer,
+ observed = c.observed
+ FROM collections c, batch
+ WHERE d.ctid = batch.ctid
+ AND d.collection_name = c.name
+ AND d.peer_name = c.peer_name
+ AND d.workspace_name = c.workspace_name
+ """
+ ),
+ {"batch_size": batch_size},
+ )
+ if result.rowcount == 0:
+ break
+
+ # Step 6: Make documents observer and observed NOT NULL
+ op.alter_column("documents", "observer", nullable=False, schema=schema)
+ op.alter_column("documents", "observed", nullable=False, schema=schema)
+
+ # CONSTRAINTS AND INDEXES
+ # Step 7: Drop ALL foreign key constraints from documents that might reference collections.peer_name
+ # Get all foreign keys on documents table
+ documents_fks = inspector.get_foreign_keys("documents", schema=schema)
+ for fk in documents_fks:
+ fk_name = fk.get("name")
+ # Drop any FK that references collections or peers and includes peer_name or collection_name
+ if fk_name and any(
+ pattern in fk_name
+ for pattern in [
+ "collection_name",
+ "peer_name",
+ ]
+ ):
+ op.drop_constraint(
+ fk_name,
+ "documents",
+ type_="foreignkey",
+ schema=schema,
+ )
+
+ # Step 7a: Drop ALL foreign key constraints from collections that reference peer_name
+ collections_fks = inspector.get_foreign_keys("collections", schema=schema)
+ for fk in collections_fks:
+ fk_name = fk.get("name")
+ if fk_name and "peer_name" in fk_name:
+ op.drop_constraint(
+ fk_name,
+ "collections",
+ type_="foreignkey",
+ schema=schema,
+ )
+
+ # Step 7c: Drop the peer_name column from collections (observer replaces it)
+ if column_exists("collections", "peer_name", inspector):
+ op.drop_column("collections", "peer_name", schema=schema)
+
+ # Step 7d: Drop the peer_name column from documents (observed replaces it)
+ if column_exists("documents", "peer_name", inspector):
+ op.drop_column("documents", "peer_name", schema=schema)
+
+ # Step 8: Drop the old unique constraint on collections that includes name
+ if constraint_exists(
+ "collections", "unique_name_collection_peer", "unique", inspector
+ ):
+ op.drop_constraint(
+ "unique_name_collection_peer", "collections", type_="unique", schema=schema
+ )
+
+ # Step 9: Create new unique constraint on collections (without peer_name)
+ if not constraint_exists(
+ "collections", "unique_observer_observed_collection", "unique", inspector
+ ):
+ op.create_unique_constraint(
+ "unique_observer_observed_collection",
+ "collections",
+ ["observer", "observed", "workspace_name"],
+ schema=schema,
+ )
+
+ # Step 10: Add composite foreign key constraint for observer peer on collections
+ if not fk_exists(
+ "collections", "collections_observer_workspace_name_fkey", inspector
+ ):
+ op.create_foreign_key(
+ "collections_observer_workspace_name_fkey",
+ "collections",
+ "peers",
+ ["observer", "workspace_name"],
+ ["name", "workspace_name"],
+ source_schema=schema,
+ referent_schema=schema,
+ )
+
+ # Step 11: Add composite foreign key constraint for observed peer on collections
+ if not fk_exists(
+ "collections", "collections_observed_workspace_name_fkey", inspector
+ ):
+ op.create_foreign_key(
+ "collections_observed_workspace_name_fkey",
+ "collections",
+ "peers",
+ ["observed", "workspace_name"],
+ ["name", "workspace_name"],
+ source_schema=schema,
+ referent_schema=schema,
+ )
+
+ # Step 12: Add composite foreign key constraint from documents to collections using observer/observed
+ if not fk_exists(
+ "documents", "documents_observer_observed_workspace_name_fkey", inspector
+ ):
+ op.create_foreign_key(
+ "documents_observer_observed_workspace_name_fkey",
+ "documents",
+ "collections",
+ ["observer", "observed", "workspace_name"],
+ ["observer", "observed", "workspace_name"],
+ source_schema=schema,
+ referent_schema=schema,
+ )
+
+ # Step 13: Add composite foreign key constraint for observer peer on documents
+ if not fk_exists("documents", "documents_observer_workspace_name_fkey", inspector):
+ op.create_foreign_key(
+ "documents_observer_workspace_name_fkey",
+ "documents",
+ "peers",
+ ["observer", "workspace_name"],
+ ["name", "workspace_name"],
+ source_schema=schema,
+ referent_schema=schema,
+ )
+
+ # Step 14: Add composite foreign key constraint for observed peer on documents
+ if not fk_exists("documents", "documents_observed_workspace_name_fkey", inspector):
+ op.create_foreign_key(
+ "documents_observed_workspace_name_fkey",
+ "documents",
+ "peers",
+ ["observed", "workspace_name"],
+ ["name", "workspace_name"],
+ source_schema=schema,
+ referent_schema=schema,
+ )
+
+ # Step 15: Create indexes for observer and observed on collections
+ if not index_exists("collections", "idx_collections_observer", inspector):
+ op.create_index(
+ "idx_collections_observer",
+ "collections",
+ ["observer"],
+ schema=schema,
+ )
+
+ if not index_exists("collections", "idx_collections_observed", inspector):
+ op.create_index(
+ "idx_collections_observed",
+ "collections",
+ ["observed"],
+ schema=schema,
+ )
+
+ # Step 16: Create indexes for observer and observed on documents
+ if not index_exists("documents", "idx_documents_observer", inspector):
+ op.create_index(
+ "idx_documents_observer",
+ "documents",
+ ["observer"],
+ schema=schema,
+ )
+
+ if not index_exists("documents", "idx_documents_observed", inspector):
+ op.create_index(
+ "idx_documents_observed",
+ "documents",
+ ["observed"],
+ schema=schema,
+ )
+
+ # Step 17: Drop the name column from collections
+ if column_exists("collections", "name", inspector):
+ op.drop_column("collections", "name", schema=schema)
+
+ # Step 18: Drop the collection_name column from documents
+ if column_exists("documents", "collection_name", inspector):
+ op.drop_column("documents", "collection_name", schema=schema)
+
+
+def downgrade() -> None:
+ """Restore collections.name and documents.collection_name from observer and observed fields."""
+ schema = settings.DB.SCHEMA
+ inspector = sa.inspect(op.get_bind())
+ connection = op.get_bind()
+
+ # SESSION_NAME MIGRATION ROLLBACK
+ # Make session_name nullable again for documents table
+
+ # Revert documents table
+ if column_exists("documents", "session_name", inspector):
+ op.alter_column("documents", "session_name", nullable=True, schema=schema)
+
+ # COLLECTIONS TABLE
+ # Step 1: Add back the name column to collections
+ if not column_exists("collections", "name", inspector):
+ op.add_column(
+ "collections",
+ sa.Column("name", sa.TEXT(), nullable=True),
+ schema=schema,
+ )
+
+ # Step 2: Populate collections name from observer and observed
+ connection.execute(
+ text(
+ """
+ UPDATE collections
+ SET name = CASE
+ WHEN observer = observed THEN 'global_representation'
+ ELSE observer || '_' || observed
+ END
+ WHERE name IS NULL
+ """
+ )
+ )
+
+ # Step 3: Make collections name NOT NULL
+ op.alter_column("collections", "name", nullable=False, schema=schema)
+
+ # Step 3a: Restore peer_name column to collections (set to observer value)
+ if not column_exists("collections", "peer_name", inspector):
+ op.add_column(
+ "collections",
+ sa.Column("peer_name", sa.TEXT(), nullable=True),
+ schema=schema,
+ )
+
+ # Populate peer_name with observer value
+ connection.execute(
+ text(
+ """
+ UPDATE collections
+ SET peer_name = observer
+ WHERE peer_name IS NULL
+ """
+ )
+ )
+
+ # Make peer_name NOT NULL
+ op.alter_column("collections", "peer_name", nullable=False, schema=schema)
+
+ # Recreate the foreign key constraint for peer_name
+ if not fk_exists(
+ "collections", "collections_peer_name_workspace_name_fkey", inspector
+ ):
+ op.create_foreign_key(
+ "collections_peer_name_workspace_name_fkey",
+ "collections",
+ "peers",
+ ["peer_name", "workspace_name"],
+ ["name", "workspace_name"],
+ source_schema=schema,
+ referent_schema=schema,
+ )
+
+ # DOCUMENTS TABLE
+ # Step 4: Add back the collection_name column to documents
+ if not column_exists("documents", "collection_name", inspector):
+ op.add_column(
+ "documents",
+ sa.Column("collection_name", sa.TEXT(), nullable=True),
+ schema=schema,
+ )
+
+ # Step 5: Populate documents collection_name from observer and observed
+ connection.execute(
+ text(
+ """
+ UPDATE documents
+ SET collection_name = CASE
+ WHEN observer = observed THEN 'global_representation'
+ ELSE observer || '_' || observed
+ END
+ WHERE collection_name IS NULL
+ """
+ )
+ )
+
+ # Step 6: Make documents collection_name NOT NULL
+ op.alter_column("documents", "collection_name", nullable=False, schema=schema)
+
+ # Step 6a: Restore peer_name column to documents (set to observed value)
+ if not column_exists("documents", "peer_name", inspector):
+ op.add_column(
+ "documents",
+ sa.Column("peer_name", sa.TEXT(), nullable=True),
+ schema=schema,
+ )
+
+ # Populate peer_name with observed value
+ connection.execute(
+ text(
+ """
+ UPDATE documents
+ SET peer_name = observed
+ WHERE peer_name IS NULL
+ """
+ )
+ )
+
+ # Make peer_name NOT NULL
+ op.alter_column("documents", "peer_name", nullable=False, schema=schema)
+
+ # Recreate the foreign key constraint for peer_name on documents
+ if not fk_exists("documents", "documents_peer_name_workspace_name_fkey", inspector):
+ op.create_foreign_key(
+ "documents_peer_name_workspace_name_fkey",
+ "documents",
+ "peers",
+ ["peer_name", "workspace_name"],
+ ["name", "workspace_name"],
+ source_schema=schema,
+ referent_schema=schema,
+ )
+
+ # Step 7: Add check constraint for name length on collections
+ if not constraint_exists("collections", "name_length", "check", inspector):
+ op.create_check_constraint(
+ "name_length",
+ "collections",
+ "length(name) <= 1025",
+ schema=schema,
+ )
+
+ # CONSTRAINTS AND INDEXES
+ # Step 8: Drop new foreign key constraints from documents
+ if fk_exists(
+ "documents", "documents_observer_observed_workspace_name_fkey", inspector
+ ):
+ op.drop_constraint(
+ "documents_observer_observed_workspace_name_fkey",
+ "documents",
+ type_="foreignkey",
+ schema=schema,
+ )
+
+ if fk_exists("documents", "documents_observer_workspace_name_fkey", inspector):
+ op.drop_constraint(
+ "documents_observer_workspace_name_fkey",
+ "documents",
+ type_="foreignkey",
+ schema=schema,
+ )
+
+ if fk_exists("documents", "documents_observed_workspace_name_fkey", inspector):
+ op.drop_constraint(
+ "documents_observed_workspace_name_fkey",
+ "documents",
+ type_="foreignkey",
+ schema=schema,
+ )
+
+ # Step 9: Drop the new unique constraint on collections
+ if constraint_exists(
+ "collections", "unique_observer_observed_collection", "unique", inspector
+ ):
+ op.drop_constraint(
+ "unique_observer_observed_collection",
+ "collections",
+ type_="unique",
+ schema=schema,
+ )
+
+ # Step 10: Recreate the old unique constraint on collections
+ if not constraint_exists(
+ "collections", "unique_name_collection_peer", "unique", inspector
+ ):
+ op.create_unique_constraint(
+ "unique_name_collection_peer",
+ "collections",
+ ["name", "peer_name", "workspace_name"],
+ schema=schema,
+ )
+
+ # Step 11: Recreate the old foreign key constraint from documents to collections
+ if not fk_exists(
+ "documents",
+ "documents_collection_name_peer_name_workspace_name_fkey",
+ inspector,
+ ):
+ op.create_foreign_key(
+ "documents_collection_name_peer_name_workspace_name_fkey",
+ "documents",
+ "collections",
+ ["collection_name", "peer_name", "workspace_name"],
+ ["name", "peer_name", "workspace_name"],
+ source_schema=schema,
+ referent_schema=schema,
+ )
+
+ # Step 12: Drop foreign key constraints from collections
+ if fk_exists("collections", "collections_observer_workspace_name_fkey", inspector):
+ op.drop_constraint(
+ "collections_observer_workspace_name_fkey",
+ "collections",
+ type_="foreignkey",
+ schema=schema,
+ )
+
+ if fk_exists("collections", "collections_observed_workspace_name_fkey", inspector):
+ op.drop_constraint(
+ "collections_observed_workspace_name_fkey",
+ "collections",
+ type_="foreignkey",
+ schema=schema,
+ )
+
+ # Step 13: Drop indexes from collections
+ if index_exists("collections", "idx_collections_observer", inspector):
+ op.drop_index(
+ "idx_collections_observer", table_name="collections", schema=schema
+ )
+
+ if index_exists("collections", "idx_collections_observed", inspector):
+ op.drop_index(
+ "idx_collections_observed", table_name="collections", schema=schema
+ )
+
+ # Step 14: Drop indexes from documents
+ if index_exists("documents", "idx_documents_observer", inspector):
+ op.drop_index("idx_documents_observer", table_name="documents", schema=schema)
+
+ if index_exists("documents", "idx_documents_observed", inspector):
+ op.drop_index("idx_documents_observed", table_name="documents", schema=schema)
+
+ # Step 15: Drop observer and observed columns from collections
+ if column_exists("collections", "observer", inspector):
+ op.drop_column("collections", "observer", schema=schema)
+
+ if column_exists("collections", "observed", inspector):
+ op.drop_column("collections", "observed", schema=schema)
+
+ # Step 16: Drop observer and observed columns from documents
+ if column_exists("documents", "observer", inspector):
+ op.drop_column("documents", "observer", schema=schema)
+
+ if column_exists("documents", "observed", inspector):
+ op.drop_column("documents", "observed", schema=schema)
diff --git a/migrations/versions/564ba40505c5_add_session_name_column_to_documents.py b/migrations/versions/564ba40505c5_add_session_name_column_to_documents.py
new file mode 100644
index 00000000..d6a5fde4
--- /dev/null
+++ b/migrations/versions/564ba40505c5_add_session_name_column_to_documents.py
@@ -0,0 +1,104 @@
+"""add_session_name_column_to_documents
+
+Revision ID: 564ba40505c5
+Revises: 88b0fb10906f
+Create Date: 2025-10-01 15:32:13.210971
+
+"""
+
+from collections.abc import Sequence
+
+import sqlalchemy as sa
+from alembic import op
+
+from migrations.utils import column_exists, fk_exists, index_exists
+from src.config import settings
+
+# revision identifiers, used by Alembic.
+revision: str = "564ba40505c5"
+down_revision: str | None = "88b0fb10906f"
+branch_labels: str | Sequence[str] | None = None
+depends_on: str | Sequence[str] | None = None
+
+
+def upgrade() -> None:
+ """Add session_name column to documents table and migrate data from internal_metadata."""
+ schema = settings.DB.SCHEMA
+ inspector = sa.inspect(op.get_bind())
+
+ # Step 1: Add session_name column as nullable
+ if not column_exists("documents", "session_name", inspector):
+ op.add_column(
+ "documents",
+ sa.Column("session_name", sa.TEXT(), nullable=True),
+ schema=schema,
+ )
+
+ # Step 2: Migrate data from internal_metadata to session_name column
+ op.execute(
+ sa.text(
+ f"""
+ UPDATE {schema}.documents
+ SET session_name = internal_metadata->>'session_name'
+ WHERE internal_metadata ? 'session_name'
+ """
+ )
+ )
+
+ # Step 3: Create index on session_name for efficient querying
+ if not index_exists("documents", "idx_documents_session_name", inspector):
+ op.create_index(
+ "idx_documents_session_name",
+ "documents",
+ ["session_name"],
+ schema=schema,
+ )
+
+ # Step 4: Add foreign key constraint to sessions table
+ # Documents can have NULL session_name (for global observations not tied to a session)
+ if not fk_exists("documents", "fk_documents_session_workspace", inspector):
+ op.create_foreign_key(
+ "fk_documents_session_workspace",
+ "documents",
+ "sessions",
+ ["session_name", "workspace_name"],
+ ["name", "workspace_name"],
+ source_schema=schema,
+ referent_schema=schema,
+ )
+
+ # Note: We keep session_name nullable since some documents may not have a session
+
+
+def downgrade() -> None:
+ """Remove session_name column and restore data to internal_metadata."""
+ schema = settings.DB.SCHEMA
+ inspector = sa.inspect(op.get_bind())
+
+ # Step 1: Copy session_name back to internal_metadata
+ op.execute(
+ sa.text(
+ f"""
+ UPDATE {schema}.documents
+ SET internal_metadata = internal_metadata || jsonb_build_object('session_name', session_name)
+ WHERE session_name IS NOT NULL
+ """
+ )
+ )
+
+ # Step 2: Drop foreign key constraint
+ if fk_exists("documents", "fk_documents_session_workspace", inspector):
+ op.drop_constraint(
+ "fk_documents_session_workspace",
+ "documents",
+ type_="foreignkey",
+ schema=schema,
+ )
+
+ # Step 3: Drop index
+ if index_exists("documents", "idx_documents_session_name", inspector):
+ op.drop_index("idx_documents_session_name", "documents", schema=schema)
+
+ # Step 4: Drop session_name column
+ if column_exists("documents", "session_name", inspector):
+ op.drop_column("documents", "session_name", schema=schema)
diff --git a/pyproject.toml b/pyproject.toml
index 40a13384..329e5604 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -31,6 +31,7 @@ dependencies = [
"google-genai>=1.32.0",
"pdfplumber>=0.11.7",
"typing-extensions>=4.11.0",
+ "json-repair>=0.49.0",
]
[tool.uv]
dev-dependencies = [
diff --git a/scripts/generate_message_embeddings.py b/scripts/generate_message_embeddings.py
index 60ac5782..fc19c3fe 100644
--- a/scripts/generate_message_embeddings.py
+++ b/scripts/generate_message_embeddings.py
@@ -25,9 +25,8 @@ from sqlalchemy import select # noqa: E402
from sqlalchemy.ext.asyncio import AsyncSession # noqa: E402
from src import models # noqa: E402
-from src.config import settings # noqa: E402
from src.dependencies import tracked_db # noqa: E402
-from src.embedding_client import EmbeddingClient # noqa: E402
+from src.embedding_client import embedding_client # noqa: E402
async def get_messages_without_embeddings(
@@ -76,7 +75,6 @@ async def get_messages_without_embeddings(
async def create_embeddings_for_messages(
db: AsyncSession,
messages: list[models.Message],
- embedding_client: EmbeddingClient,
) -> int:
"""
Create embeddings for a batch of messages.
@@ -84,7 +82,6 @@ async def create_embeddings_for_messages(
Args:
db: Database session
messages: List of messages to create embeddings for
- embedding_client: Embedding client instance
Returns:
Number of embeddings created
@@ -159,9 +156,6 @@ async def main() -> None:
args = parser.parse_args()
- # Initialize embedding client
- embedding_client = EmbeddingClient(settings.LLM.OPENAI_API_KEY)
-
print("Generating embeddings for messages...")
if args.workspace_name:
print(f" Filtering by workspace: {args.workspace_name}")
@@ -200,9 +194,7 @@ async def main() -> None:
f"Processing batch {batch_num}/{total_batches} ({len(batch)} messages)..."
)
- embeddings_created = await create_embeddings_for_messages(
- db, batch, embedding_client
- )
+ embeddings_created = await create_embeddings_for_messages(db, batch)
total_embeddings += embeddings_created
print(
diff --git a/sdks/python/CHANGELOG.md b/sdks/python/CHANGELOG.md
index 8a06e418..2ce0448c 100644
--- a/sdks/python/CHANGELOG.md
+++ b/sdks/python/CHANGELOG.md
@@ -5,7 +5,7 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](http://keepachangelog.com/)
and this project adheres to [Semantic Versioning](http://semver.org/).
-## [1.4.1] — 2025-01-01
+## [1.4.1] — 2025-10-09
### Added
diff --git a/sdks/typescript/CHANGELOG.md b/sdks/typescript/CHANGELOG.md
index e4366558..1cbc337b 100644
--- a/sdks/typescript/CHANGELOG.md
+++ b/sdks/typescript/CHANGELOG.md
@@ -5,7 +5,7 @@ All notable changes to this project will be documented in this file.
The format is based on [Keep a Changelog](http://keepachangelog.com/)
and this project adheres to [Semantic Versioning](http://semver.org/).
-## [1.4.1] — 2025-01-01
+## [1.4.1] — 2025-10-09
### Added
diff --git a/sdks/typescript/src/session_context.ts b/sdks/typescript/src/session_context.ts
index 3dd4c136..538844c2 100644
--- a/sdks/typescript/src/session_context.ts
+++ b/sdks/typescript/src/session_context.ts
@@ -138,9 +138,9 @@ export class SessionContext {
const assistantId = typeof assistant === 'string' ? assistant : assistant.id
const summaryMessage = this.summary
? {
- role: 'system',
- content: `${this.summary.content}`,
- }
+ role: 'system',
+ content: `${this.summary.content}`,
+ }
: null
const messages = this.messages.map((message) => ({
role: message.peer_id === assistantId ? 'assistant' : 'user',
@@ -172,20 +172,20 @@ export class SessionContext {
const assistantId = typeof assistant === 'string' ? assistant : assistant.id
const summaryMessage = this.summary
? {
- role: 'user',
- content: `${this.summary.content}`,
- }
+ role: 'user',
+ content: `${this.summary.content}`,
+ }
: null
const messages = this.messages.map((message) =>
message.peer_id === assistantId
? {
- role: 'assistant',
- content: message.content,
- }
+ role: 'assistant',
+ content: message.content,
+ }
: {
- role: 'user',
- content: `${message.peer_id}: ${message.content}`,
- }
+ role: 'user',
+ content: `${message.peer_id}: ${message.content}`,
+ }
)
return summaryMessage ? [summaryMessage, ...messages] : messages
}
diff --git a/src/config.py b/src/config.py
index a87ac9b5..170bccd5 100644
--- a/src/config.py
+++ b/src/config.py
@@ -1,6 +1,6 @@
import logging
from pathlib import Path
-from typing import Annotated, Any, ClassVar
+from typing import Annotated, Any, ClassVar, Literal
import tomllib
from dotenv import load_dotenv
@@ -52,9 +52,11 @@ class TomlConfigSettingsSource(PydanticBaseSettingsSource):
"SENTRY": "sentry",
"LLM": "llm",
"DERIVER": "deriver",
+ "PEER_CARD": "peer_card",
"DIALECTIC": "dialectic",
"SUMMARY": "summary",
"WEBHOOK": "webhook",
+ "DREAM": "dream",
"": "app", # For AppSettings with no prefix
}
@@ -174,6 +176,8 @@ class LLMSettings(HonchoSettings):
GROQ_API_KEY: str | None = None
OPENAI_COMPATIBLE_BASE_URL: str | None = None
+ EMBEDDING_PROVIDER: Literal["openai", "gemini"] = "openai"
+
# General LLM settings
DEFAULT_MAX_TOKENS: Annotated[int, Field(default=1000, gt=0, le=100_000)] = 2500
@@ -190,22 +194,15 @@ class DeriverSettings(HonchoSettings):
PROVIDER: SupportedProviders = "google"
MODEL: str = "gemini-2.5-flash-lite"
- MAX_OUTPUT_TOKENS: Annotated[int, Field(default=2500, gt=0, le=100_000)] = 2500
+ MAX_OUTPUT_TOKENS: Annotated[int, Field(default=10_000, gt=0, le=100_000)] = 10_000
# Thinking budget tokens are only applied when using Anthropic as provider
THINKING_BUDGET_TOKENS: Annotated[int, Field(default=1024, gt=0, le=5000)] = 1024
- PEER_CARD_PROVIDER: SupportedProviders = "openai"
- PEER_CARD_MODEL: str = "gpt-5-nano-2025-08-07"
- # Note: peer cards should be very short, but GPT-5 models need output tokens for thinking which cannot be turned off...
- PEER_CARD_MAX_OUTPUT_TOKENS: Annotated[
- int, Field(default=4000, gt=1000, le=10_000)
- ] = 4000
-
# Maximum number of observations to store in working representation
# This is applied to both explicit and deductive observations
WORKING_REPRESENTATION_MAX_OBSERVATIONS: Annotated[
- int, Field(default=100, gt=0, le=500)
- ] = 100
+ int, Field(default=50, gt=0, le=500)
+ ] = 50
REPRESENTATION_BATCH_MAX_TOKENS: Annotated[
int,
@@ -226,6 +223,17 @@ class DeriverSettings(HonchoSettings):
return self
+class PeerCardSettings(HonchoSettings):
+ model_config = SettingsConfigDict(env_prefix="PEER_CARD_", extra="ignore") # pyright: ignore
+
+ ENABLED: bool = True
+
+ PROVIDER: SupportedProviders = "openai"
+ MODEL: str = "gpt-5-nano-2025-08-07"
+ # Note: peer cards should be very short, but GPT-5 models need output tokens for thinking which cannot be turned off...
+ MAX_OUTPUT_TOKENS: Annotated[int, Field(default=4000, gt=1000, le=10_000)] = 4000
+
+
class DialecticSettings(HonchoSettings):
model_config = SettingsConfigDict(env_prefix="DIALECTIC_", extra="ignore") # pyright: ignore
@@ -253,6 +261,8 @@ class DialecticSettings(HonchoSettings):
class SummarySettings(HonchoSettings):
model_config = SettingsConfigDict(env_prefix="SUMMARY_", extra="ignore") # pyright: ignore
+ ENABLED: bool = True
+
MESSAGES_PER_SHORT_SUMMARY: Annotated[int, Field(default=20, gt=0, le=100)] = 20
MESSAGES_PER_LONG_SUMMARY: Annotated[int, Field(default=60, gt=0, le=500)] = 60
@@ -271,6 +281,21 @@ class WebhookSettings(HonchoSettings):
MAX_WORKSPACE_LIMIT: int = 10
+class DreamSettings(HonchoSettings):
+ model_config = SettingsConfigDict(env_prefix="DREAM_", extra="ignore") # pyright: ignore
+
+ ENABLED: bool = True
+ DOCUMENT_THRESHOLD: Annotated[int, Field(default=50, gt=0, le=1000)] = 50
+ IDLE_TIMEOUT_MINUTES: Annotated[int, Field(default=60, gt=0, le=1440)] = 60
+ MIN_HOURS_BETWEEN_DREAMS: Annotated[int, Field(default=8, gt=0, le=72)] = 8
+ ENABLED_TYPES: list[str] = ["consolidate"]
+
+ # LLM settings for dream processing
+ PROVIDER: SupportedProviders = "openai"
+ MODEL: str = "gpt-4o-mini-2024-07-18"
+ MAX_OUTPUT_TOKENS: Annotated[int, Field(default=2000, gt=0, le=10_000)] = 2000
+
+
class AppSettings(HonchoSettings):
# No env_prefix for app-level settings
model_config = SettingsConfigDict( # pyright: ignore
@@ -294,6 +319,9 @@ class AppSettings(HonchoSettings):
LANGFUSE_HOST: str | None = None
LANGFUSE_PUBLIC_KEY: str | None = None
+ COLLECT_METRICS_LOCAL: bool = False
+ LOCAL_METRICS_FILE: str = "metrics.jsonl"
+
# Nested settings models
DB: DBSettings = Field(default_factory=DBSettings)
AUTH: AuthSettings = Field(default_factory=AuthSettings)
@@ -301,8 +329,10 @@ class AppSettings(HonchoSettings):
LLM: LLMSettings = Field(default_factory=LLMSettings)
DERIVER: DeriverSettings = Field(default_factory=DeriverSettings)
DIALECTIC: DialecticSettings = Field(default_factory=DialecticSettings)
+ PEER_CARD: PeerCardSettings = Field(default_factory=PeerCardSettings)
SUMMARY: SummarySettings = Field(default_factory=SummarySettings)
WEBHOOK: WebhookSettings = Field(default_factory=WebhookSettings)
+ DREAM: DreamSettings = Field(default_factory=DreamSettings)
@field_validator("LOG_LEVEL")
def validate_log_level(cls, v: str) -> str:
diff --git a/src/crud/__init__.py b/src/crud/__init__.py
index 45f2133f..f00a6ae1 100644
--- a/src/crud/__init__.py
+++ b/src/crud/__init__.py
@@ -1,6 +1,10 @@
from .collection import get_collection, get_or_create_collection
from .deriver import get_deriver_status
-from .document import create_document, get_duplicate_documents, query_documents
+from .document import (
+ create_documents,
+ get_all_documents,
+ query_documents,
+)
from .message import (
create_messages,
get_message,
@@ -17,13 +21,9 @@ from .peer import (
get_sessions_for_peer,
update_peer,
)
+from .peer_card import get_peer_card, set_peer_card
from .representation import (
- construct_collection_name,
- get_peer_card,
get_working_representation,
- get_working_representation_data,
- set_peer_card,
- set_working_representation,
)
from .session import (
clone_session,
@@ -53,9 +53,9 @@ __all__ = [
# Deriver
"get_deriver_status",
# Document
+ "create_documents",
+ "get_all_documents",
"query_documents",
- "create_document",
- "get_duplicate_documents",
# Message
"create_messages",
"get_messages",
@@ -70,13 +70,11 @@ __all__ = [
"get_peers",
"update_peer",
"get_sessions_for_peer",
- # Representation
- "construct_collection_name",
+ # Peer Card
"get_peer_card",
- "get_working_representation",
- "get_working_representation_data",
"set_peer_card",
- "set_working_representation",
+ # Representation
+ "get_working_representation",
# Session
"get_sessions",
"get_or_create_session",
diff --git a/src/crud/collection.py b/src/crud/collection.py
index 49fe9d88..c1759b9c 100644
--- a/src/crud/collection.py
+++ b/src/crud/collection.py
@@ -13,17 +13,18 @@ logger = getLogger(__name__)
async def get_collection(
db: AsyncSession,
workspace_name: str,
- collection_name: str,
- peer_name: str | None = None,
+ *,
+ observer: str,
+ observed: str,
) -> models.Collection:
"""
- Get a collection by name for a specific peer and workspace.
+ Get a collection by observer/observed for a workspace.
Args:
db: Database session
workspace_name: Name of the workspace
- peer_name: Name of the peer
- collection_name: Name of the collection
+ observer: Name of the observing peer (owns the collection)
+ observed: Name of the observed peer
Returns:
The collection if found
@@ -34,35 +35,34 @@ async def get_collection(
stmt = (
select(models.Collection)
.where(models.Collection.workspace_name == workspace_name)
- .where(models.Collection.name == collection_name)
+ .where(models.Collection.observer == observer)
+ .where(models.Collection.observed == observed)
)
- if peer_name:
- stmt = stmt.where(models.Collection.peer_name == peer_name)
result = await db.execute(stmt)
collection = result.scalar_one_or_none()
if collection is None:
- raise ResourceNotFoundException(
- "Collection not found or does not belong to peer"
- )
+ raise ResourceNotFoundException("Collection not found")
return collection
async def get_or_create_collection(
db: AsyncSession,
workspace_name: str,
- collection_name: str,
- peer_name: str | None = None,
*,
+ observer: str,
+ observed: str,
_retry: bool = False,
) -> models.Collection:
try:
- return await get_collection(db, workspace_name, collection_name, peer_name)
+ return await get_collection(
+ db, workspace_name, observer=observer, observed=observed
+ )
except ResourceNotFoundException:
try:
honcho_collection = models.Collection(
workspace_name=workspace_name,
- peer_name=peer_name,
- name=collection_name,
+ observer=observer,
+ observed=observed,
)
db.add(honcho_collection)
await db.commit()
@@ -71,8 +71,8 @@ async def get_or_create_collection(
await db.rollback()
if _retry:
raise ConflictException(
- f"Unable to create or get collection: {collection_name}"
+ f"Unable to create or get collection: {observer}/{observed}"
) from None
return await get_or_create_collection(
- db, workspace_name, collection_name, peer_name, _retry=True
+ db, workspace_name, observer=observer, observed=observed, _retry=True
)
diff --git a/src/crud/deriver.py b/src/crud/deriver.py
index 6066b878..326b71a1 100644
--- a/src/crud/deriver.py
+++ b/src/crud/deriver.py
@@ -14,9 +14,10 @@ logger = getLogger(__name__)
async def get_deriver_status(
db: AsyncSession,
workspace_name: str,
- observer_name: str | None = None,
- sender_name: str | None = None,
session_name: str | None = None,
+ *,
+ observer: str | None = None,
+ observed: str | None = None,
) -> schemas.DeriverStatus:
"""
Get the deriver processing status, optionally filtered by observer, sender, and/or session.
@@ -24,20 +25,20 @@ async def get_deriver_status(
Args:
db: Database session
workspace_name: Name of the workspace
- observer_name: Optional name of the observer (target) to filter by
- sender_name: Optional name of the sender to filter by
session_name: Optional session name to filter by
+ observer: Optional name of the observer to filter by
+ observed: Optional name of the observed (message sender) to filter by
"""
# Normalize empty strings to None for consistent handling
- normalized_observer_name = observer_name if observer_name else None
- normalized_sender_name = sender_name if sender_name else None
+ normalized_observer = observer if observer else None
+ normalized_observed = observed if observed else None
normalized_session_name = session_name if session_name else None
stmt = _build_queue_status_query(
workspace_name,
- normalized_observer_name,
- normalized_sender_name,
normalized_session_name,
+ observer=normalized_observer,
+ observed=normalized_observed,
)
result = await db.execute(stmt)
rows = result.fetchall()
@@ -51,13 +52,14 @@ async def get_deriver_status(
def _build_queue_status_query(
workspace_name: str,
- observer_name: str | None,
- sender_name: str | None,
session_name: str | None,
+ *,
+ observer: str | None = None,
+ observed: str | None = None,
) -> Select[Any]:
"""Build SQL query for queue status with validation and aggregation."""
- sender_name_expr = models.QueueItem.payload["sender_name"].astext
- target_name_expr = models.QueueItem.payload["target_name"].astext
+ observer_name_expr = models.QueueItem.payload["observer"].astext
+ observed_name_expr = models.QueueItem.payload["observed"].astext
# Define conditions for cleaner window functions
is_completed = models.QueueItem.processed
@@ -103,10 +105,10 @@ def _build_queue_status_query(
stmt = stmt.where(models.Session.name == session_name)
peer_conditions = []
- if observer_name is not None:
- peer_conditions.append(target_name_expr == observer_name) # pyright: ignore
- if sender_name is not None:
- peer_conditions.append(sender_name_expr == sender_name) # pyright: ignore
+ if observer is not None:
+ peer_conditions.append(observer_name_expr == observer) # pyright: ignore
+ if observed is not None:
+ peer_conditions.append(observed_name_expr == observed) # pyright: ignore
if peer_conditions:
stmt = stmt.where(or_(*peer_conditions)) # pyright: ignore
diff --git a/src/crud/document.py b/src/crud/document.py
index 3a2e240a..127b5d7b 100644
--- a/src/crud/document.py
+++ b/src/crud/document.py
@@ -3,6 +3,7 @@ from logging import getLogger
from typing import Any
from sqlalchemy import select
+from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
from src import models, schemas
@@ -11,17 +12,41 @@ from src.embedding_client import embedding_client
from src.exceptions import ValidationException
from src.utils.filter import apply_filter
-from .collection import get_collection
-
logger = getLogger(__name__)
+async def get_all_documents(
+ db: AsyncSession,
+ workspace_name: str,
+ *,
+ observer: str,
+ observed: str,
+ limit: int = 1000,
+) -> Sequence[models.Document]:
+ """
+ Get all documents in a collection.
+
+ NOTE: Order is nondeterministic. Also this may return a massive amount of documents. Don't use this on large collections.
+ TODO: add pagination and update dreaming logic to deduplicate more effectively
+ """
+ stmt = (
+ select(models.Document)
+ .limit(limit)
+ .where(models.Document.workspace_name == workspace_name)
+ .where(models.Document.observer == observer)
+ .where(models.Document.observed == observed)
+ )
+ result = await db.execute(stmt)
+ return result.scalars().all()
+
+
async def query_documents(
db: AsyncSession,
workspace_name: str,
- peer_name: str,
- collection_name: str,
query: str,
+ *,
+ observer: str,
+ observed: str,
filters: dict[str, Any] | None = None,
max_distance: float | None = None,
top_k: int = 5,
@@ -33,20 +58,19 @@ async def query_documents(
Args:
db: Database session
workspace_name: Name of the workspace
- peer_name: Name of the peer
- collection_name: Name of the collection
query: Search query text
+ observer: Name of the observing peer
+ observed: Name of the observed peer
filters: Optional filters to apply
max_distance: Maximum cosine distance for results
top_k: Number of results to return
- embedding: Optional pre-computed embedding for the query (avoids API call)
+ embedding: Optional pre-computed embedding for the query (avoids extra API call if possible)
Returns:
Sequence of matching documents
"""
# Use provided embedding or generate one
if embedding is None:
- # Using ModelClient for embeddings
try:
embedding = await embedding_client.embed(query)
except ValueError as e:
@@ -57,9 +81,8 @@ async def query_documents(
stmt = (
select(models.Document)
.where(models.Document.workspace_name == workspace_name)
- .where(models.Document.peer_name == peer_name)
- .where(models.Document.collection_name == collection_name)
- # .limit(top_k)
+ .where(models.Document.observer == observer)
+ .where(models.Document.observed == observed)
)
if max_distance is not None:
stmt = stmt.where(
@@ -73,114 +96,54 @@ async def query_documents(
return result.scalars().all()
-async def create_document(
+async def create_documents(
db: AsyncSession,
- document: schemas.DocumentCreate,
+ documents: list[schemas.DocumentCreate],
workspace_name: str,
- peer_name: str,
- collection_name: str,
- duplicate_threshold: float | None = None,
-) -> models.Document:
+ *,
+ observer: str,
+ observed: str,
+) -> int:
"""
- Embed text as a vector and create a document.
+ Create multiple documents with NO duplicate detection.
Args:
db: Database session
- document: Document creation schema
+ documents: List of document creation schemas
workspace_name: Name of the workspace
- peer_name: Name of the peer
- collection_name: Name of the collection
+ observer: Name of the observing peer
+ observed: Name of the observed peer
Returns:
- The created document
-
- Raises:
- ResourceNotFoundException: If the collection does not exist
- ValidationException: If the document data is invalid
+ Count of new documents
"""
-
- # This will raise ResourceNotFoundException if collection not found
- await get_collection(
- db,
- workspace_name=workspace_name,
- collection_name=collection_name,
- peer_name=peer_name,
- )
-
- # Using ModelClient for embeddings
- embedding = await embedding_client.embed(document.content)
-
- if duplicate_threshold is not None:
- # Check if there are duplicates within the threshold
- stmt = (
- select(models.Document)
- .where(models.Document.workspace_name == workspace_name)
- .where(models.Document.peer_name == peer_name)
- .where(models.Document.collection_name == collection_name)
- .where(
- models.Document.embedding.cosine_distance(embedding)
- < duplicate_threshold
+ honcho_documents: list[models.Document] = []
+ for doc in documents:
+ try:
+ metadata_dict = doc.metadata.model_dump(exclude_none=True)
+ honcho_documents.append(
+ models.Document(
+ workspace_name=workspace_name,
+ observer=observer,
+ observed=observed,
+ content=doc.content,
+ internal_metadata=metadata_dict,
+ embedding=doc.embedding,
+ session_name=doc.session_name,
+ )
)
- .order_by(models.Document.embedding.cosine_distance(embedding))
- .limit(1)
- )
- result = await db.execute(stmt)
- duplicate = result.scalar_one_or_none() # Get the closest match if any exist
- if duplicate is not None:
- logger.info(f"Duplicate found: {duplicate.content}. Ignoring new document.")
- return duplicate
+ except Exception as e:
+ logger.error(
+ f"Error adding new document to {workspace_name}/{doc.session_name}/{observer}/{observed}: {e}"
+ )
+ continue
+ try:
+ db.add_all(honcho_documents)
+ await db.commit()
+ except IntegrityError as e:
+ await db.rollback()
+ raise ValidationException(
+ "Failed to create documents due to integrity constraint violation"
+ ) from e
- honcho_document = models.Document(
- workspace_name=workspace_name,
- peer_name=peer_name,
- collection_name=collection_name,
- content=document.content,
- internal_metadata=document.metadata,
- embedding=embedding,
- )
- db.add(honcho_document)
- await db.commit()
- await db.refresh(honcho_document)
- return honcho_document
-
-
-async def get_duplicate_documents(
- db: AsyncSession,
- workspace_name: str,
- peer_name: str,
- collection_name: str,
- content: str,
- similarity_threshold: float = 0.85,
-) -> list[models.Document]:
- """Check if a document with similar content already exists in the collection.
-
- Args:
- db: Database session
- workspace_name: Name of the workspace
- peer_name: Name of the peer
- collection_name: Name of the collection
- content: Document content to check for duplicates
- similarity_threshold: Similarity threshold (0-1) for considering documents as duplicates
-
- Returns:
- List of documents that are similar to the provided content
- """
- # Get embedding for the content
- # Using ModelClient for embeddings
- embedding = await embedding_client.embed(content)
-
- # Find documents with similar embeddings
- stmt = (
- select(models.Document)
- .where(models.Document.workspace_name == workspace_name)
- .where(models.Document.peer_name == peer_name)
- .where(models.Document.collection_name == collection_name)
- .where(
- models.Document.embedding.cosine_distance(embedding)
- < (1 - similarity_threshold)
- ) # Convert similarity to distance
- .order_by(models.Document.embedding.cosine_distance(embedding))
- )
-
- result = await db.execute(stmt)
- return list(result.scalars().all()) # Convert to list to match the return type
+ return len(honcho_documents)
diff --git a/src/crud/peer_card.py b/src/crud/peer_card.py
new file mode 100644
index 00000000..295d0de8
--- /dev/null
+++ b/src/crud/peer_card.py
@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import logging
+from typing import cast
+
+from sqlalchemy import update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from src import exceptions, models, schemas
+from src.crud.peer import get_peer
+
+logger = logging.getLogger(__name__)
+
+
+async def get_peer_card(
+ db: AsyncSession,
+ workspace_name: str,
+ *,
+ observer: str,
+ observed: str,
+) -> list[str] | None:
+ """
+ Get peer card from internal_metadata.
+
+ The peer card is returned for the observer/observed relationship.
+
+ Args:
+ db: Database session
+ workspace_name: Name of the workspace
+ observed: Peer name of the peer described in the peer card
+ observer: Peer name of the observer
+
+ Returns:
+ The peer's card text if present, otherwise None (also None if peer not found).
+ """
+ try:
+ peer = await get_peer(db, workspace_name, schemas.PeerCreate(name=observer))
+ return cast(
+ list[str] | None,
+ peer.internal_metadata.get(
+ construct_peer_card_label(observer=observer, observed=observed)
+ ),
+ )
+ except exceptions.ResourceNotFoundException:
+ return None
+
+
+async def set_peer_card(
+ db: AsyncSession,
+ workspace_name: str,
+ peer_card: list[str],
+ *,
+ observer: str,
+ observed: str,
+) -> None:
+ """
+ Set peer card for a peer.
+
+ If observer_name is provided, the peer card is set for the observer/observed relationship.
+
+ Args:
+ db: Database session
+ workspace_name: Name of the workspace
+ peer_card: List of strings to set as the peer card
+ observed: Peer name of the peer described in the peer card
+ observer: Peer name of the observer
+
+ Raises:
+ ResourceNotFoundException: If the peer does not exist
+ """
+ stmt = (
+ update(models.Peer)
+ .where(models.Peer.workspace_name == workspace_name)
+ .where(models.Peer.name == observer)
+ .values(
+ internal_metadata=models.Peer.internal_metadata.op("||")(
+ {
+ construct_peer_card_label(
+ observer=observer, observed=observed
+ ): peer_card
+ }
+ )
+ )
+ )
+ result = await db.execute(stmt)
+ if result.rowcount == 0:
+ raise exceptions.ResourceNotFoundException(
+ f"Peer {observer} not found in workspace {workspace_name}"
+ )
+ await db.commit()
+
+
+def construct_peer_card_label(*, observer: str, observed: str) -> str:
+ if observer == observed:
+ return "peer_card"
+ return f"{observed}_peer_card"
diff --git a/src/crud/representation.py b/src/crud/representation.py
index 787dcbdd..ed88bf4b 100644
--- a/src/crud/representation.py
+++ b/src/crud/representation.py
@@ -1,371 +1,569 @@
-from logging import getLogger
-from typing import Any, Final, cast
+from __future__ import annotations
-from sqlalchemy import select, update
+import datetime
+import logging
+import time
+from typing import Any
+
+from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
-from src import exceptions, models, schemas
+from src import crud, exceptions, models, schemas
from src.config import settings
-from src.crud.peer import get_peer
-from src.utils.shared_models import ObservationDict
+from src.dependencies import tracked_db
+from src.dreamer.dream_scheduler import check_and_schedule_dream
+from src.embedding_client import embedding_client
+from src.utils.formatting import format_datetime_utc
+from src.utils.logging import accumulate_metric, conditional_observe
+from src.utils.representation import (
+ DeductiveObservation,
+ ExplicitObservation,
+ Representation,
+)
-logger = getLogger(__name__)
+logger = logging.getLogger(__name__)
-# The collection name for documents that make up a peer's global representation
-GLOBAL_REPRESENTATION_COLLECTION_NAME: Final[str] = "global_representation"
-
-# The key for the working representation in the session peer's internal_metadata
-WORKING_REPRESENTATION_METADATA_KEY = "working_representation"
-
-# Old working representation key--remove in 2.3.0?
-WORKING_REPRESENTATION_LEGACY_METADATA_KEY = "global_representation"
+# Fetch extra documents to ensure we have enough after filtering
+FILTER_OVERSAMPLING_FACTOR = 3
-Observation = str | ObservationDict
+class RepresentationManager:
+ """Unified manager for representation and document queries."""
+ def __init__(
+ self,
+ workspace_name: str,
+ *,
+ observer: str,
+ observed: str,
+ ) -> None:
+ self.workspace_name: str = workspace_name
+ self.observer: str = observer
+ self.observed: str = observed
-async def get_peer_card(
- db: AsyncSession,
- workspace_name: str,
- observed_name: str,
- observer_name: str,
-) -> list[str] | None:
- """
- Get peer card from internal_metadata.
+ @conditional_observe
+ async def save_representation(
+ self,
+ representation: Representation,
+ message_id_range: tuple[int, int],
+ session_name: str,
+ message_created_at: datetime.datetime,
+ ) -> int:
+ """
+ Save Representation objects to the collection as a set of documents.
- The peer card is returned for the observer/observed relationship.
+ Args:
+ representation: Representation object
+ message_id_range: Message ID range to link with observations
+ session_name: Session name to link with existing summary context
+ message_created_at: Timestamp when the message was created
- Args:
- db: Database session
- workspace_name: Name of the workspace
- observed_name: Peer name of the peer described in the peer card
- observer_name: Peer name of the observer
+ Returns:
+ The number of *new documents saved*
+ """
- Returns:
- The peer's card text if present, otherwise None (also None if peer not found).
- """
- try:
- peer = await get_peer(
- db, workspace_name, schemas.PeerCreate(name=observer_name)
+ new_documents = 0
+
+ if not representation.deductive and not representation.explicit:
+ logger.debug("No observations to save")
+ return new_documents
+
+ all_observations = representation.deductive + representation.explicit
+
+ # Batch embed all observations
+ batch_embed_start = time.perf_counter()
+
+ observation_texts = [
+ obs.conclusion if isinstance(obs, DeductiveObservation) else obs.content
+ for obs in all_observations
+ ]
+ try:
+ embeddings = await embedding_client.simple_batch_embed(observation_texts)
+ except ValueError as e:
+ raise exceptions.ValidationException(
+ f"Observation content exceeds maximum token limit of {settings.MAX_EMBEDDING_TOKENS}."
+ ) from e
+
+ batch_embed_duration = (time.perf_counter() - batch_embed_start) * 1000
+ accumulate_metric(
+ f"deriver_{message_id_range[1]}_{self.observer}",
+ "embed_new_observations",
+ batch_embed_duration,
+ "ms",
)
- return cast(
- list[str] | None,
- peer.internal_metadata.get(
- construct_peer_card_label(
- observer=observer_name, observed=observed_name
- )
- ),
- )
- except exceptions.ResourceNotFoundException:
- return None
-
-async def set_peer_card(
- db: AsyncSession,
- workspace_name: str,
- observed_name: str,
- observer_name: str,
- peer_card: list[str] | None,
-) -> None:
- """
- Set peer card for a peer.
-
- If observer_name is provided, the peer card is set for the observer/observed relationship.
-
- Args:
- db: Database session
- workspace_name: Name of the workspace
- observed_name: Peer name of the peer described in the peer card
- observer_name: Peer name of the observer
- peer_card: List of strings to set as the peer card
-
- Raises:
- ResourceNotFoundException: If the peer does not exist
- """
- stmt = (
- update(models.Peer)
- .where(models.Peer.workspace_name == workspace_name)
- .where(models.Peer.name == observer_name)
- .values(
- internal_metadata=models.Peer.internal_metadata.op("||")(
- {
- construct_peer_card_label(
- observer=observer_name, observed=observed_name
- ): peer_card
- }
+ # Batch create document objects
+ create_document_start = time.perf_counter()
+ async with tracked_db("representation_manager.save_representation") as db:
+ new_documents = await self._save_representation_internal(
+ db,
+ all_observations,
+ embeddings,
+ message_id_range,
+ session_name,
+ message_created_at,
)
+
+ create_document_duration = (time.perf_counter() - create_document_start) * 1000
+ accumulate_metric(
+ f"deriver_{message_id_range[1]}_{self.observer}",
+ "save_new_observations",
+ create_document_duration,
+ "ms",
)
- )
- result = await db.execute(stmt)
- if result.rowcount == 0:
- raise exceptions.ResourceNotFoundException(
- f"Peer {observer_name} not found in workspace {workspace_name}"
+
+ return new_documents
+
+ async def _save_representation_internal(
+ self,
+ db: AsyncSession,
+ all_observations: list[ExplicitObservation | DeductiveObservation],
+ embeddings: list[list[float]],
+ message_id_range: tuple[int, int],
+ session_name: str,
+ message_created_at: datetime.datetime,
+ ) -> int:
+ # get_or_create_collection already handles IntegrityError with rollback and a retry
+ collection = await crud.get_or_create_collection(
+ db,
+ self.workspace_name,
+ observer=self.observer,
+ observed=self.observed,
)
- await db.commit()
+
+ # Prepare all documents for bulk creation
+ documents_to_create: list[schemas.DocumentCreate] = []
+ for obs, embedding in zip(all_observations, embeddings, strict=True):
+ # NOTE: will add additional levels of reasoning in the future
+ if isinstance(obs, DeductiveObservation):
+ obs_level = "deductive"
+ obs_content = obs.conclusion
+ obs_premises = obs.premises
+ else:
+ obs_level = "explicit"
+ obs_content = obs.content
+ obs_premises = None
+
+ metadata: schemas.DocumentMetadata = schemas.DocumentMetadata(
+ message_ids=[message_id_range],
+ level=obs_level,
+ premises=obs_premises,
+ message_created_at=format_datetime_utc(message_created_at),
+ )
+
+ documents_to_create.append(
+ schemas.DocumentCreate(
+ content=obs_content,
+ session_name=session_name,
+ metadata=metadata,
+ embedding=embedding,
+ )
+ )
+
+ # Use bulk creation with NO duplicate detection
+ new_documents = await crud.create_documents(
+ db,
+ documents_to_create,
+ self.workspace_name,
+ observer=self.observer,
+ observed=self.observed,
+ )
+
+ try:
+ await check_and_schedule_dream(db, collection)
+ except Exception as e:
+ logger.warning(f"Failed to check dream scheduling: {e}")
+
+ return new_documents
+
+ async def get_relevant_observations(
+ self,
+ query: str,
+ *,
+ top_k: int = 5,
+ max_distance: float = 0.3,
+ level: str | None = None,
+ conversation_context: str = "",
+ ) -> Representation:
+ """
+ Unified method to get relevant observations with flexible options.
+
+ Args:
+ query: The search query
+ top_k: Number of results to return
+ max_distance: Maximum distance for semantic similarity
+ level: Optional reasoning level to filter by
+ conversation_context: Additional conversation context
+
+ Returns:
+ Representation
+ """
+ async with tracked_db("representation_manager.get_relevant_observations") as db:
+ documents = await self._get_observations_internal(
+ db,
+ query,
+ top_k,
+ max_distance,
+ level,
+ conversation_context,
+ )
+
+ # convert documents to representation
+ return Representation.from_documents(documents)
+
+ async def get_working_representation(
+ self,
+ *,
+ session_name: str | None = None,
+ include_semantic_query: str | None = None,
+ semantic_search_top_k: int | None = None,
+ semantic_search_max_distance: float | None = None,
+ include_most_derived: bool = False,
+ max_observations: int = settings.DERIVER.WORKING_REPRESENTATION_MAX_OBSERVATIONS,
+ ) -> Representation:
+ """
+ Get working representation with flexible query options.
+
+ Args:
+ session_name: Optional session to filter by
+ include_semantic_query: Query for semantic search
+ semantic_search_top_k: Number of semantic results
+ semantic_search_max_distance: Maximum distance for semantic search
+ include_most_derived: Include most derived observations
+ max_observations: Maximum total observations to return
+
+ Returns:
+ Representation combining various query strategies
+ """
+ async with tracked_db(
+ "representation_manager.get_working_representation"
+ ) as db:
+ return await self._get_working_representation_internal(
+ db,
+ session_name=session_name,
+ include_semantic_query=include_semantic_query,
+ semantic_search_top_k=semantic_search_top_k,
+ semantic_search_max_distance=semantic_search_max_distance,
+ include_most_derived=include_most_derived,
+ max_observations=max_observations,
+ )
+
+ # Private helper methods
+
+ async def _get_working_representation_internal(
+ self,
+ db: AsyncSession,
+ *,
+ session_name: str | None = None,
+ include_semantic_query: str | None = None,
+ semantic_search_top_k: int | None = None,
+ semantic_search_max_distance: float | None = None,
+ include_most_derived: bool = False,
+ max_observations: int = settings.DERIVER.WORKING_REPRESENTATION_MAX_OBSERVATIONS,
+ ) -> Representation:
+ """Internal implementation of get_working_representation."""
+ total = max_observations
+
+ # Calculate how many observations to get from each source
+ semantic_observations = (
+ min(
+ max(
+ 0,
+ semantic_search_top_k
+ if semantic_search_top_k is not None
+ else total // 3,
+ ),
+ total,
+ )
+ if include_semantic_query
+ else 0
+ )
+
+ if include_semantic_query and include_most_derived:
+ # three-way blend: both semantic and derived requested
+ top_observations = min(max(0, total // 3), total - semantic_observations)
+ elif include_most_derived:
+ # two-way blend: only derived requested
+ top_observations = min(max(0, total // 2), total - semantic_observations)
+ else:
+ # no derived observations requested
+ top_observations = 0
+
+ # remaining observations are recent
+ recent_observations = total - semantic_observations - top_observations
+
+ representation = Representation()
+
+ # Get semantic observations if requested
+ if include_semantic_query:
+ semantic_docs = await self._query_documents_semantic(
+ db,
+ query=include_semantic_query,
+ top_k=semantic_observations,
+ max_distance=semantic_search_max_distance
+ if semantic_search_max_distance is not None
+ else 0.3,
+ )
+ representation.merge_representation(
+ Representation.from_documents(semantic_docs)
+ )
+
+ # Get most derived observations if requested
+ if include_most_derived:
+ derived_docs = await self._query_documents_most_derived(
+ db, top_k=top_observations
+ )
+ representation.merge_representation(
+ Representation.from_documents(derived_docs)
+ )
+
+ # Get recent observations
+ recent_docs = await self._query_documents_recent(
+ db, top_k=recent_observations, session_name=session_name
+ )
+
+ if not recent_docs:
+ logger.warning(
+ f"No observations for {self.observed} (observer: {self.observer}) found. Normal if brand-new peer."
+ )
+
+ representation.merge_representation(Representation.from_documents(recent_docs))
+
+ return representation
+
+ async def _query_documents_semantic(
+ self,
+ db: AsyncSession,
+ query: str,
+ top_k: int,
+ max_distance: float,
+ level: str | None = None,
+ conversation_context: str = "",
+ ) -> list[models.Document]:
+ """Query documents by semantic similarity."""
+ try:
+ if level:
+ return await self._query_documents_for_level(
+ db,
+ query,
+ level,
+ conversation_context,
+ max_distance,
+ top_k,
+ )
+ else:
+ documents = await crud.query_documents(
+ db,
+ workspace_name=self.workspace_name,
+ observer=self.observer,
+ observed=self.observed,
+ query=self._build_truncated_query(query, conversation_context),
+ max_distance=max_distance,
+ top_k=top_k,
+ )
+ db.expunge_all()
+ return list(documents)
+
+ except Exception as e:
+ logger.error(f"Error getting relevant observations: {e}")
+ return []
+
+ async def _query_documents_recent(
+ self, db: AsyncSession, top_k: int, session_name: str | None = None
+ ) -> list[models.Document]:
+ """Query most recent documents."""
+ stmt = (
+ select(models.Document)
+ .limit(top_k)
+ .where(
+ models.Document.workspace_name == self.workspace_name,
+ models.Document.observer == self.observer,
+ models.Document.observed == self.observed,
+ *(
+ [models.Document.session_name == session_name]
+ if session_name is not None
+ else []
+ ),
+ )
+ .order_by(models.Document.created_at.desc())
+ )
+
+ result = await db.execute(stmt)
+ documents = result.scalars().all()
+ db.expunge_all()
+ return list(documents)
+
+ async def _query_documents_most_derived(
+ self, db: AsyncSession, top_k: int
+ ) -> list[models.Document]:
+ """Query most derived documents."""
+ stmt = (
+ select(models.Document)
+ .limit(top_k)
+ .where(
+ models.Document.workspace_name == self.workspace_name,
+ models.Document.observer == self.observer,
+ models.Document.observed == self.observed,
+ )
+ .order_by(models.Document.internal_metadata["times_derived"].desc())
+ )
+
+ result = await db.execute(stmt)
+ documents = result.scalars().all()
+ db.expunge_all()
+ return list(documents)
+
+ async def _get_observations_internal(
+ self,
+ db: AsyncSession,
+ query: str,
+ top_k: int,
+ max_distance: float,
+ level: str | None,
+ conversation_context: str,
+ ) -> list[models.Document]:
+ """Internal method that does the actual observation retrieval."""
+ return await self._query_documents_semantic(
+ db, query, top_k, max_distance, level, conversation_context
+ )
+
+ async def _query_documents_for_level(
+ self,
+ db: AsyncSession,
+ query: str,
+ level: str,
+ conversation_context: str,
+ max_distance: float,
+ count: int,
+ ) -> list[models.Document]:
+ """Query documents for a specific level."""
+ documents = await crud.query_documents(
+ db,
+ workspace_name=self.workspace_name,
+ observer=self.observer,
+ observed=self.observed,
+ query=self._build_truncated_query(query, conversation_context),
+ max_distance=max_distance,
+ top_k=count * FILTER_OVERSAMPLING_FACTOR,
+ filters=self._build_filter_conditions(level),
+ )
+
+ # Sort by creation time and return top count
+ docs_sorted: list[models.Document] = sorted(
+ list(documents), key=lambda x: x.created_at, reverse=True
+ )
+ return docs_sorted[:count]
+
+ def _build_filter_conditions(
+ self,
+ level: str | None = None,
+ ) -> dict[str, Any]:
+ """Build complete filter conditions for document queries."""
+ conditions: list[dict[str, Any]] = []
+
+ if level:
+ conditions.append({"internal_metadata": {"level": level}})
+
+ if not conditions:
+ return {}
+
+ return conditions[0] if len(conditions) == 1 else {"AND": conditions}
+
+ def _build_truncated_query(
+ self,
+ query: str,
+ conversation_context: str = "",
+ max_tokens: int | None = None,
+ ) -> str:
+ """Build a query that fits within token limits with clear priorities.
+
+ Args:
+ query: The search query
+ conversation_context: Optional conversation context to include
+ max_tokens: Maximum tokens allowed (defaults to setting with buffer)
+
+ Returns:
+ Truncated query string that fits within token limits
+ """
+ max_tokens = max_tokens or (settings.MAX_EMBEDDING_TOKENS - 100)
+ encoding = embedding_client.encoding
+
+ # Pre-calculate all token counts once
+ query_prefix = "Current message: "
+ context_prefix = "\nContext: "
+
+ prefix_tokens = len(encoding.encode(query_prefix))
+ context_prefix_tokens = len(encoding.encode(context_prefix))
+ query_tokens = encoding.encode(query)
+
+ # Simple case: query alone fits
+ if prefix_tokens + len(query_tokens) <= max_tokens:
+ if not conversation_context:
+ return f"{query_prefix}{query}"
+
+ # Try to add context
+ context_tokens = encoding.encode(conversation_context)
+ total_without_context = (
+ prefix_tokens + len(query_tokens) + context_prefix_tokens
+ )
+
+ if total_without_context + len(context_tokens) <= max_tokens:
+ return f"{query_prefix}{query}{context_prefix}{conversation_context}"
+
+ # Truncate context to fit
+ available_context_tokens = max_tokens - total_without_context
+ if available_context_tokens > 0:
+ truncated_context = encoding.decode(
+ context_tokens[-available_context_tokens:]
+ )
+ return f"{query_prefix}{query}{context_prefix}{truncated_context}"
+ else:
+ # No room left for context; keep full query intact
+ return f"{query_prefix}{query}"
+
+ # Query itself is too long - truncate it
+ available_query_tokens = max_tokens - prefix_tokens
+ if available_query_tokens > 0:
+ # Keep the end (recency) of the query
+ truncated_query = encoding.decode(query_tokens[-available_query_tokens:])
+ return f"{query_prefix}{truncated_query}"
+
+ # Pathological case - just return what we can
+ logger.warning("Token limit too restrictive: %s", max_tokens)
+ return encoding.decode(query_tokens[:max_tokens])
+
+
+# Module-level functions for backward compatibility and convenience
async def get_working_representation(
- db: AsyncSession,
workspace_name: str,
- observer_name: str,
- observed_name: str,
- session_name: str,
-) -> str:
+ *,
+ observer: str,
+ observed: str,
+ session_name: str | None = None,
+ include_semantic_query: str | None = None,
+ semantic_search_top_k: int | None = None,
+ semantic_search_max_distance: float | None = None,
+ include_most_derived: bool = False,
+ max_observations: int = settings.DERIVER.WORKING_REPRESENTATION_MAX_OBSERVATIONS,
+) -> Representation:
"""
- Get working representation for observer/observed relationship.
+ Get raw working representation data from the relevant document collection.
- Args:
- db: Database session
- workspace_name: Name of the workspace
- observer_name: Name of the peer doing the observing
- observed_name: Name of the peer being observed
- session_name: Name of the session
-
- Returns:
- Formatted working representation string
+ This is a convenience function that creates a RepresentationManager and calls
+ get_working_representation on it.
"""
- working_rep_data = await get_working_representation_data(
- db, workspace_name, observer_name, observed_name, session_name
+ manager = RepresentationManager(
+ workspace_name=workspace_name,
+ observer=observer,
+ observed=observed,
)
-
- if not working_rep_data:
- logger.warning(
- f"No working representation found for observer: {observer_name}, observed: {observed_name}"
- )
- return ""
-
- # Handle both old format (string) and new format (structured data)
- if isinstance(working_rep_data, str):
- return working_rep_data
-
- # New structured format - extract and format final_observations
- try:
- final_observations = working_rep_data.get("final_observations", {})
- if not final_observations:
- logger.warning("No final_observations found in working representation data")
- return ""
-
- return _format_observations_by_level(final_observations)
- except Exception:
- logger.exception("Error processing working representation")
- return ""
-
-
-async def get_working_representation_data(
- db: AsyncSession,
- workspace_name: str,
- observer_name: str,
- observed_name: str,
- session_name: str,
-) -> dict[str, Any] | str | None:
- """
- Get raw working representation data from internal_metadata.
-
- Returns either structured data (new format) or string (legacy format).
- """
- # Determine metadata key based on observer/observed relationship
- if observer_name == observed_name:
- metadata_key = WORKING_REPRESENTATION_METADATA_KEY
- else:
- metadata_key = construct_collection_name(
- observer=observer_name, observed=observed_name
- )
-
- stmt = select(models.SessionPeer.internal_metadata).where(
- models.SessionPeer.peer_name == observer_name,
- models.SessionPeer.workspace_name == workspace_name,
- models.SessionPeer.session_name == session_name,
+ return await manager.get_working_representation(
+ session_name=session_name,
+ include_semantic_query=include_semantic_query,
+ semantic_search_top_k=semantic_search_top_k,
+ semantic_search_max_distance=semantic_search_max_distance,
+ include_most_derived=include_most_derived,
+ max_observations=max_observations,
)
-
- result = await db.execute(stmt)
- peer_metadata = result.scalar_one_or_none()
-
- if not peer_metadata:
- return None
-
- working_rep_data = peer_metadata.get(metadata_key)
- if working_rep_data:
- return cast(dict[str, Any] | str, working_rep_data)
-
- # Try legacy key--remove in 2.3.0?
- if observer_name == observed_name:
- working_rep_data = peer_metadata.get(WORKING_REPRESENTATION_LEGACY_METADATA_KEY)
- if working_rep_data:
- return cast(dict[str, Any] | str, working_rep_data)
-
- return None
-
-
-def _format_observations_by_level(final_observations: dict[str, Any]) -> str:
- """Format final observations into structured text by level."""
- formatted_sections: list[str] = []
-
- for level in ["explicit", "deductive"]:
- observations_raw: Any = final_observations.get(level, [])
- observations: list[Any] = cast(list[Any], observations_raw or [])
- if observations:
- formatted_sections.append(f"{level.upper()} OBSERVATIONS:")
- formatted_sections.extend(_format_observation_list(observations))
- formatted_sections.append("")
-
- return "\n".join(formatted_sections) if formatted_sections else ""
-
-
-def _format_observation_list(observations: list[Observation]) -> list[str]:
- """Format a list of observations into consistent string format."""
- formatted: list[str] = []
- for obs in observations:
- if isinstance(obs, dict):
- # Determine core content and premises
- if "conclusion" in obs:
- conclusion_text: str = obs["conclusion"]
- premises: list[str] = obs.get("premises", [])
- if premises:
- premises_text = "; ".join(premises)
- formatted_obs = f"{conclusion_text} (based on: {premises_text})"
- else:
- formatted_obs = conclusion_text
- else:
- content_text: str = obs.get("content", str(obs))
- formatted_obs = content_text
-
- formatted.append(f"- {formatted_obs}")
- else:
- # Handle string fallback
- formatted.append(f"- {str(obs)}")
- return formatted
-
-
-def _merge_working_representation(
- existing: dict[str, Any] | str | None, new: dict[str, Any]
-) -> dict[str, Any]:
- """Merge a new working representation into an existing one.
-
- - Appends `explicit` and `deductive` observations in that order
- - Trims each list to the most recent `WORKING_REPRESENTATION_MAX_OBSERVATIONS` entries (FIFO)
- - Uses the latest `thinking`, `message_id`, and `created_at`
- """
- max_observations = settings.DERIVER.WORKING_REPRESENTATION_MAX_OBSERVATIONS
-
- new_final_raw: Any = new.get("final_observations") or {}
- new_explicit: list[Observation] = cast(
- list[Observation], (new_final_raw.get("explicit") or [])
- )
- new_deductive: list[ObservationDict] = cast(
- list[ObservationDict], (new_final_raw.get("deductive") or [])
- )
-
- existing_explicit: list[Observation] = []
- existing_deductive: list[ObservationDict] = []
- if isinstance(existing, dict):
- existing_final_raw: Any = existing.get("final_observations") or {}
- existing_explicit = cast(
- list[Observation], (existing_final_raw.get("explicit") or [])
- )
- existing_deductive = cast(
- list[ObservationDict], (existing_final_raw.get("deductive") or [])
- )
-
- merged_explicit: list[Observation] = existing_explicit + new_explicit
- merged_deductive: list[ObservationDict] = existing_deductive + new_deductive
-
- if len(merged_explicit) > max_observations:
- merged_explicit = merged_explicit[-max_observations:]
- if len(merged_deductive) > max_observations:
- merged_deductive = merged_deductive[-max_observations:]
-
- return {
- "final_observations": {
- "explicit": merged_explicit,
- "thinking": cast(str | None, new_final_raw.get("thinking")),
- "deductive": merged_deductive,
- },
- "message_id": cast(str | None, new.get("message_id")),
- "created_at": cast(str | None, new.get("created_at")),
- }
-
-
-async def set_working_representation(
- db: AsyncSession,
- representation: str | dict[str, Any],
- workspace_name: str,
- observer_name: str,
- observed_name: str,
- session_name: str,
-) -> None:
- """
- Set working representation for observer/observed relationship.
-
- If the provided representation is structured (dict with `final_observations`),
- append new observations to the existing ones for both `explicit` and `deductive`
- kinds, update `message_id` and `created_at`, and cap each observations list to
- the most recent `WORKING_REPRESENTATION_MAX_OBSERVATIONS` items (FIFO trimming
- of oldest entries).
-
- Args:
- db: Database session
- representation: Working representation data (string or structured dict)
- workspace_name: Name of the workspace
- observer_name: Name of the peer doing the observing
- observed_name: Name of the peer being observed (required for explicit global/local)
- session_name: Name of the session
- """
- # Determine metadata key based on observer/observed relationship
- if observer_name == observed_name:
- metadata_key = WORKING_REPRESENTATION_METADATA_KEY
- else:
- metadata_key = construct_collection_name(
- observer=observer_name, observed=observed_name
- )
-
- merged_value: str | dict[str, Any] = representation
- if isinstance(representation, dict):
- try:
- existing = await get_working_representation_data(
- db=db,
- workspace_name=workspace_name,
- observer_name=observer_name,
- observed_name=observed_name,
- session_name=session_name,
- )
- merged_value = _merge_working_representation(
- existing,
- representation,
- )
- except Exception:
- logger.exception(
- "Failed to merge working representation; storing as provided"
- )
- merged_value = representation
-
- stmt = (
- update(models.SessionPeer)
- .where(models.SessionPeer.workspace_name == workspace_name)
- .where(models.SessionPeer.peer_name == observer_name)
- .where(models.SessionPeer.session_name == session_name)
- .values(
- internal_metadata=models.SessionPeer.internal_metadata.op("||")(
- {metadata_key: merged_value}
- )
- )
- )
-
- await db.execute(stmt)
- await db.commit()
-
- logger.info(
- "Saved working representation to session peer %s - %s with key %s",
- session_name,
- observer_name,
- metadata_key,
- )
-
-
-def construct_collection_name(*, observer: str, observed: str) -> str:
- return f"{observer}_{observed}"
-
-
-def construct_peer_card_label(*, observer: str, observed: str) -> str:
- if observer == observed:
- return "peer_card"
- return f"{observed}_peer_card"
diff --git a/src/dependencies.py b/src/dependencies.py
index b4664760..77699150 100644
--- a/src/dependencies.py
+++ b/src/dependencies.py
@@ -49,13 +49,12 @@ async def tracked_db(operation_name: str | None = None):
)
yield db
- # Explicitly end transaction if still open
- if db.in_transaction():
- await db.rollback() # Or commit if needed for write operations
except Exception:
await db.rollback()
raise
finally:
+ if db.in_transaction():
+ await db.rollback()
await db.close()
if token: # Only reset if we set it
request_context.reset(token)
diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py
index 95e363d7..59aab370 100644
--- a/src/deriver/consumer.py
+++ b/src/deriver/consumer.py
@@ -2,7 +2,6 @@ import logging
from typing import Any
import sentry_sdk
-from langfuse import get_client
from pydantic import ValidationError
from rich.console import Console
from sqlalchemy import select
@@ -11,22 +10,24 @@ from src import models
from src.config import settings
from src.dependencies import tracked_db
from src.deriver.deriver import process_representation_tasks_batch
+from src.dreamer.dreamer import process_dream
from src.models import Message
from src.utils import summarizer
+from src.utils.langfuse_client import get_langfuse_client
from src.utils.logging import log_performance_metrics
-from src.webhooks import webhook_delivery
-
-from .queue_payload import (
+from src.utils.queue_payload import (
+ DreamPayload,
SummaryPayload,
WebhookPayload,
)
+from src.webhooks import webhook_delivery
logger = logging.getLogger(__name__)
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
console = Console(markup=True)
-lf = get_client()
+lf = get_langfuse_client() if settings.LANGFUSE_PUBLIC_KEY else None
async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None:
@@ -79,7 +80,7 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None:
message_public_id = message.public_id
with sentry_sdk.start_transaction(name="process_summary_task", op="deriver"):
- if settings.LANGFUSE_PUBLIC_KEY:
+ if lf:
with lf.start_as_current_span(
name="summary_processing",
input={
@@ -99,7 +100,7 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None:
message_public_id,
)
log_performance_metrics(
- f"summary_{validated.workspace_name}_{validated.message_id}"
+ "summary", f"{validated.workspace_name}_{validated.message_id}"
)
else:
await summarizer.summarize_if_needed(
@@ -110,16 +111,30 @@ async def process_item(task_type: str, queue_payload: dict[str, Any]) -> None:
message_public_id,
)
log_performance_metrics(
- f"summary_{validated.workspace_name}_{validated.message_id}"
+ "summary", f"{validated.workspace_name}_{validated.message_id}"
)
+
+ elif task_type == "dream":
+ with sentry_sdk.start_transaction(name="process_dream_task", op="deriver"):
+ try:
+ validated = DreamPayload(**queue_payload)
+ except ValidationError as e:
+ logger.error(
+ "Invalid dream payload received: %s. Payload: %s",
+ str(e),
+ queue_payload,
+ )
+ raise ValueError(f"Invalid payload structure: {str(e)}") from e
+ await process_dream(validated)
else:
raise ValueError(f"Invalid task type: {task_type}")
async def process_representation_batch(
messages: list[Message],
- sender_name: str | None,
- target_name: str | None,
+ *,
+ observer: str | None,
+ observed: str | None,
) -> None:
"""Validate incoming queue payloads and dispatch to the appropriate handler.
@@ -130,34 +145,32 @@ async def process_representation_batch(
Args:
task_type: The type of task to process
queue_payloads: List of payload dictionaries to process
- sender_name (optional): For representation tasks, the sender_name from work_unit_key
+ observed (optional): For representation tasks, the observed from work_unit_key
to identify which messages should be focused on
- target_name (optional): For representation tasks, the target_name from work_unit_key
+ observer (optional): For representation tasks, the observer from work_unit_key
to identify which messages should be focused on
"""
if not messages or not messages[0]:
logger.debug("process_representation_batch received no payloads")
return
- if sender_name is None or target_name is None:
- raise ValueError(
- "sender_name and target_name are required for representation tasks"
- )
+ if observed is None or observer is None:
+ raise ValueError("observed and observer are required for representation tasks")
logger.debug(
"process_representation_batch received %s payloads",
len(messages),
)
- if settings.LANGFUSE_PUBLIC_KEY:
+ if lf:
with lf.start_as_current_span(
name="representation_processing",
input={
"payloads": [
{
"message_id": msg.id,
- "sender_name": sender_name,
- "target_name": target_name,
+ "observer": observer,
+ "observed": observed,
"session_name": msg.session_name,
}
for msg in messages
@@ -167,6 +180,10 @@ async def process_representation_batch(
"critical_analysis_model": settings.DERIVER.MODEL,
},
):
- await process_representation_tasks_batch(sender_name, target_name, messages)
+ await process_representation_tasks_batch(
+ messages, observer=observer, observed=observed
+ )
else:
- await process_representation_tasks_batch(sender_name, target_name, messages)
+ await process_representation_tasks_batch(
+ messages, observer=observer, observed=observed
+ )
diff --git a/src/deriver/deriver.py b/src/deriver/deriver.py
index e94ccd19..365586f2 100644
--- a/src/deriver/deriver.py
+++ b/src/deriver/deriver.py
@@ -1,45 +1,27 @@
import datetime
-import json
import logging
import time
-from typing import Any
import sentry_sdk
-from langfuse import get_client
from src import crud, exceptions
from src.config import settings
-from src.crud.representation import GLOBAL_REPRESENTATION_COLLECTION_NAME
+from src.crud.representation import RepresentationManager
from src.dependencies import tracked_db
-from src.deriver.utils import estimate_tokens
from src.models import Message
from src.utils import summarizer
from src.utils.clients import honcho_llm_call
-from src.utils.embedding_store import EmbeddingStore
-from src.utils.formatting import (
- REASONING_LEVELS,
- extract_observation_content,
- find_new_observations,
- format_context_for_prompt,
- format_new_turn_with_timestamp,
- utc_now_iso,
-)
+from src.utils.formatting import format_new_turn_with_timestamp
+from src.utils.langfuse_client import get_langfuse_client
from src.utils.logging import (
accumulate_metric,
conditional_observe,
- format_reasoning_response_as_markdown,
- log_observations_tree,
log_performance_metrics,
- log_thinking_panel,
-)
-from src.utils.shared_models import (
- DeductiveObservation,
- ObservationContext,
- PeerCardQuery,
- ReasoningResponse,
- ReasoningResponseWithThinking,
- UnifiedObservation,
+ log_representation,
)
+from src.utils.peer_card import PeerCardQuery
+from src.utils.representation import PromptRepresentation, Representation
+from src.utils.tokens import estimate_tokens
from src.utils.tracing import with_sentry_transaction
from .prompts import (
@@ -51,17 +33,17 @@ from .prompts import (
logger = logging.getLogger(__name__)
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
-lf = get_client()
+lf = get_langfuse_client() if settings.LANGFUSE_PUBLIC_KEY else None
async def critical_analysis_call(
peer_id: str,
peer_card: list[str] | None,
message_created_at: datetime.datetime,
- working_representation: str | None,
+ working_representation: Representation,
history: str,
new_turns: list[str],
-) -> ReasoningResponse | None:
+) -> PromptRepresentation:
prompt = critical_analysis_prompt(
peer_id=peer_id,
peer_card=peer_card,
@@ -78,8 +60,9 @@ async def critical_analysis_call(
max_tokens=settings.DERIVER.MAX_OUTPUT_TOKENS
or settings.LLM.DEFAULT_MAX_TOKENS,
track_name="Critical Analysis Call",
- response_model=ReasoningResponse,
+ response_model=PromptRepresentation,
json_mode=True,
+ stop_seqs=[" \n", "\n\n\n\n"],
thinking_budget_tokens=settings.DERIVER.THINKING_BUDGET_TOKENS,
enable_retry=True,
retry_attempts=3,
@@ -90,21 +73,21 @@ async def critical_analysis_call(
async def peer_card_call(
old_peer_card: list[str] | None,
- new_observations: list[str],
+ new_observations: Representation,
) -> PeerCardQuery:
"""
Generate peer card prompt, call LLM with response model.
"""
prompt = peer_card_prompt(
old_peer_card=old_peer_card,
- new_observations=new_observations,
+ new_observations=new_observations.str_no_timestamps(),
)
response = await honcho_llm_call(
- provider=settings.DERIVER.PEER_CARD_PROVIDER,
- model=settings.DERIVER.PEER_CARD_MODEL,
+ provider=settings.PEER_CARD.PROVIDER,
+ model=settings.PEER_CARD.MODEL,
prompt=prompt,
- max_tokens=settings.DERIVER.PEER_CARD_MAX_OUTPUT_TOKENS
+ max_tokens=settings.PEER_CARD.MAX_OUTPUT_TOKENS
or settings.LLM.DEFAULT_MAX_TOKENS,
track_name="Peer Card Call",
response_model=PeerCardQuery,
@@ -119,9 +102,10 @@ async def peer_card_call(
@with_sentry_transaction("process_representation_tasks_batch", op="deriver")
async def process_representation_tasks_batch(
- sender_name: str,
- target_name: str,
messages: list[Message],
+ *,
+ observer: str,
+ observed: str,
) -> None:
"""
Process a batch of representation tasks by extracting insights and updating working representations.
@@ -134,41 +118,59 @@ async def process_representation_tasks_batch(
latest_message = messages[-1]
earliest_message = messages[0]
+ accumulate_metric(
+ f"deriver_{latest_message.id}_{observer}",
+ "starting_message_id",
+ earliest_message.id,
+ "id",
+ )
+
+ accumulate_metric(
+ f"deriver_{latest_message.id}_{observer}",
+ "ending_message_id",
+ latest_message.id,
+ "id",
+ )
+
# Start overall timing
overall_start = time.perf_counter()
- logger.debug(
- "Starting insight extraction for message batch starting with: %s",
- earliest_message.id,
+ # Time context preparation
+ context_prep_start = time.perf_counter()
+
+ # Use get_session_context_formatted with configurable token limit
+
+ working_representation = await crud.get_working_representation(
+ latest_message.workspace_name,
+ observer=observer,
+ observed=observed,
+ # include_semantic_query=latest_message.content,
+ # include_most_derived=False,
)
- async with tracked_db("deriver.get_peer_card") as db:
- speaker_peer_card: list[str] | None = await crud.get_peer_card(
- db,
- latest_message.workspace_name,
- sender_name,
- target_name,
- )
- if speaker_peer_card is None:
- logger.warning("No peer card found for %s", sender_name)
+ if settings.PEER_CARD.ENABLED:
+ async with tracked_db("deriver.get_peer_card") as db:
+ speaker_peer_card: list[str] | None = await crud.get_peer_card(
+ db,
+ latest_message.workspace_name,
+ observer=observer,
+ observed=observed,
+ )
+ if speaker_peer_card is None:
+ logger.warning(
+ "No peer card found for %s. Normal if brand-new peer.",
+ observed,
+ )
+ else:
+ logger.info("Using peer card: %s", speaker_peer_card)
else:
- logger.debug("Using peer card for %s", sender_name)
-
- # Get working representation data early for token estimation
- async with tracked_db("deriver.get_working_representation_data") as db:
- working_rep_data: (
- dict[str, Any] | str | None
- ) = await crud.get_working_representation_data(
- db,
- latest_message.workspace_name,
- target_name,
- sender_name,
- latest_message.session_name,
- )
+ speaker_peer_card = None
# Estimate tokens for deriver input
peer_card_tokens = estimate_tokens(speaker_peer_card)
- working_rep_tokens = _estimate_working_representation_tokens(working_rep_data)
+ working_rep_tokens = estimate_tokens(
+ str(working_representation) if not working_representation.is_empty() else None
+ )
base_prompt_tokens = estimate_base_prompt_tokens()
# Estimate tokens for new conversation turns
@@ -200,8 +202,7 @@ async def process_representation_tasks_batch(
available_context_tokens,
)
- # Use get_session_context_formatted with dynamic token limit
- async with tracked_db("deriver.get_session_context") as db:
+ async with tracked_db("deriver.get_session_context_formatted") as db:
formatted_history = await summarizer.get_session_context_formatted(
db,
latest_message.workspace_name,
@@ -211,138 +212,39 @@ async def process_representation_tasks_batch(
include_summary=True,
)
- # instantiate embedding store from collection
- # if the sender is also the target, we're handling a global representation task.
- # otherwise, we're handling a directional representation task where the sender is
- # being observed by the target.
- collection_name = (
- crud.construct_collection_name(observer=target_name, observed=sender_name)
- if sender_name != target_name
- else GLOBAL_REPRESENTATION_COLLECTION_NAME
- )
-
- # get_or_create_collection already handles IntegrityError with rollback and a retry
- async with tracked_db("deriver.get_or_create_collection") as db:
- collection = await crud.get_or_create_collection(
- db,
- latest_message.workspace_name,
- collection_name,
- sender_name,
- )
- collection_name_loaded = collection.name
-
- # Use the embedding store directly
- embedding_store = EmbeddingStore(
- workspace_name=latest_message.workspace_name,
- peer_name=sender_name,
- collection_name=collection_name_loaded,
- )
-
- # Create reasoner instance
- reasoner = CertaintyReasoner(
- embedding_store=embedding_store,
- ctx=messages,
- sender_name=sender_name,
- target_name=target_name,
- )
-
- # Time context preparation
- context_prep_start = time.perf_counter()
- if (
- working_rep_data
- and isinstance(working_rep_data, dict)
- and working_rep_data.get("final_observations")
- ):
- # Reconstruct ReasoningResponse from stored peer data
- final_obs: dict[str, Any] = working_rep_data["final_observations"]
- deductive_observations: list[DeductiveObservation] = []
- for deductive_data in final_obs.get("deductive", []):
- deductive_observations.append(
- DeductiveObservation(
- conclusion=deductive_data["conclusion"],
- premises=deductive_data.get("premises", []),
- )
- )
-
- working_representation = ReasoningResponseWithThinking(
- thinking=final_obs.get("thinking"),
- explicit=final_obs.get("explicit", []),
- deductive=deductive_observations,
- )
- logger.info(
- "Using existing working representation with %s explicit, %s deductive observations",
- len(working_representation.explicit),
- len(working_representation.deductive),
- )
- else:
- # No existing working representation, use global search
- query_text = [m.content for m in messages]
- query_text = "\n".join(
- query_text
- ) # TODO: consider a smarter strategy than concatenation
- working_representation = await embedding_store.get_relevant_observations(
- query=query_text,
- conversation_context=formatted_history,
- for_reasoning=True,
- )
-
- working_representation = observation_context_to_reasoning_response(
- working_representation
- )
- logger.info("No working representation found, using global semantic search")
-
- # Recalculate tokens now that we have a new working representation
- new_working_rep_tokens = _estimate_working_representation_tokens(
- {
- "final_observations": {
- "thinking": working_representation.thinking,
- "explicit": working_representation.explicit,
- "deductive": [
- {
- "conclusion": obs.conclusion,
- "premises": obs.premises,
- }
- for obs in working_representation.deductive
- ],
- }
- }
- )
-
- # Update estimated input tokens with the new working representation
- estimated_input_tokens = (
- peer_card_tokens
- + new_working_rep_tokens
- + base_prompt_tokens
- + new_turns_tokens
- )
-
- # Recalculate available tokens for context
- available_context_tokens = max(
- 0,
- settings.DERIVER.MAX_INPUT_TOKENS - estimated_input_tokens - safety_buffer,
- )
-
- # Recalculate formatted_history with updated token limit
- async with tracked_db("deriver.recalc_session_context") as db:
- formatted_history = await summarizer.get_session_context_formatted(
- db,
- latest_message.workspace_name,
- latest_message.session_name,
- token_limit=available_context_tokens,
- cutoff=earliest_message.id,
- include_summary=True,
- )
+ # got working representation and peer card, log timing
context_prep_duration = (time.perf_counter() - context_prep_start) * 1000
accumulate_metric(
- f"deriver_representation_{latest_message.id}_{target_name}",
+ f"deriver_{latest_message.id}_{observer}",
"context_preparation",
context_prep_duration,
"ms",
)
- # Run consolidated reasoning that handles explicit and deductive levels
- logger.debug(
- "REASONING: Running unified insight derivation across explicit and deductive reasoning levels"
+ logger.info(
+ "Using working representation with %s explicit, %s deductive observations",
+ len(working_representation.explicit),
+ len(working_representation.deductive),
+ )
+
+ # instantiate representation manager from collection
+ # if the sender is also the target, we're handling a global representation task.
+ # otherwise, we're handling a directional representation task where the sender is
+ # being observed by the target.
+
+ # Use the representation manager directly
+ representation_manager = RepresentationManager(
+ workspace_name=latest_message.workspace_name,
+ observer=observer,
+ observed=observed,
+ )
+
+ # Create reasoner instance
+ reasoner = CertaintyReasoner(
+ representation_manager=representation_manager,
+ ctx=messages,
+ observed=observed,
+ observer=observer,
)
# Run single-pass reasoning
@@ -352,96 +254,81 @@ async def process_representation_tasks_batch(
speaker_peer_card,
)
- logger.debug("REASONING COMPLETION: Unified reasoning completed across all levels.")
-
# Display final observations in a beautiful tree
- final_obs_dict = {
- level: getattr(final_observations, level, []) for level in REASONING_LEVELS
- }
- log_observations_tree(final_obs_dict)
+ log_representation(final_observations)
- # Always save working representation to peer for dialectic access
- await save_working_representation_to_peer(
- latest_message, final_observations, sender_name, target_name
- )
# Calculate and log overall timing
overall_duration = (time.perf_counter() - overall_start) * 1000
accumulate_metric(
- f"deriver_representation_{latest_message.id}_{target_name}",
+ f"deriver_{latest_message.id}_{observer}",
"total_processing_time",
overall_duration,
"ms",
)
- total_observations = sum(len(obs_list) for obs_list in final_obs_dict.values())
+ total_observations = len(final_observations.explicit) + len(
+ final_observations.deductive
+ )
accumulate_metric(
- f"deriver_representation_{latest_message.id}_{target_name}",
- "final_observation_count",
+ f"deriver_{latest_message.id}_{observer}",
+ "observation_count",
total_observations,
"count",
)
- log_performance_metrics(f"deriver_representation_{latest_message.id}_{target_name}")
- if settings.LANGFUSE_PUBLIC_KEY:
- lf.update_current_trace(
- output=format_reasoning_response_as_markdown(final_observations)
- )
+ log_performance_metrics("deriver", f"{latest_message.id}_{observer}")
+
+ if lf:
+ lf.update_current_trace(output=final_observations.format_as_markdown())
class CertaintyReasoner:
"""Certainty reasoner for analyzing and deriving insights."""
- embedding_store: EmbeddingStore
+ representation_manager: RepresentationManager
ctx: list[Message]
- sender_name: str
- target_name: str
+ observer: str
+ observed: str
def __init__(
self,
- embedding_store: EmbeddingStore,
+ representation_manager: RepresentationManager,
ctx: list[Message],
- sender_name: str,
- target_name: str,
+ *,
+ observed: str,
+ observer: str,
) -> None:
- self.embedding_store = embedding_store
+ self.representation_manager = representation_manager
self.ctx = ctx
- self.sender_name = sender_name
- self.target_name = target_name
+ self.observed = observed
+ self.observer = observer
@conditional_observe
@sentry_sdk.trace
- async def derive_new_insights(
+ async def reason(
self,
- working_representation: ReasoningResponseWithThinking,
+ working_representation: Representation,
history: str,
speaker_peer_card: list[str] | None,
- ) -> ReasoningResponseWithThinking:
+ ) -> Representation:
"""
- Critically analyzes and revises understanding, returning structured observations.
- """
- # For logging, we can just show the content of the last message
- latest_message = self.ctx[-1]
+ Single-pass reasoning function that critically analyzes and derives insights.
+ Performs one analysis pass and returns the final observations.
- # if settings.LANGFUSE_PUBLIC_KEY:
- # lf.update_current_generation(
- # input=format_reasoning_inputs_as_markdown(
- # working_representation,
- # history,
- # latest_message.content,
- # latest_message.created_at,
- # )
- # )
+ Returns:
+ Representation: Final observations
+ """
+ analysis_start = time.perf_counter()
+
+ earliest_message = self.ctx[0]
+ latest_message = self.ctx[-1]
new_turns = [
format_new_turn_with_timestamp(m.content, m.created_at, m.peer_name)
for m in self.ctx
]
- formatted_working_representation = format_context_for_prompt(
- working_representation
- )
-
logger.debug(
"CRITICAL ANALYSIS: message_created_at='%s', new_turns_count=%s",
latest_message.created_at,
@@ -449,222 +336,84 @@ class CertaintyReasoner:
)
try:
- response_obj = await critical_analysis_call(
- peer_id=self.sender_name,
+ reasoning_response = await critical_analysis_call(
+ peer_id=self.observed,
peer_card=speaker_peer_card,
message_created_at=latest_message.created_at,
- working_representation=formatted_working_representation,
+ working_representation=working_representation,
history=history,
new_turns=new_turns,
)
except Exception as e:
raise exceptions.LLMError(
speaker_peer_card=speaker_peer_card,
- working_representation=formatted_working_representation,
+ working_representation=working_representation,
history=history,
new_turns=new_turns,
) from e
- # Handle None response from LLM
- if response_obj is None:
- logger.warning("LLM returned None response, using empty insights")
- new_insights = ReasoningResponse(explicit=[], deductive=[])
- # If response is a string, try to parse as JSON
- elif isinstance(response_obj, str):
- try:
- response_data = json.loads(response_obj)
- new_insights = ReasoningResponse(
- explicit=response_data.get("explicit", []),
- deductive=[
- DeductiveObservation(**item)
- for item in response_data.get("deductive", [])
- ],
- )
- except (json.JSONDecodeError, KeyError, TypeError) as e:
- if settings.SENTRY.ENABLED:
- sentry_sdk.capture_exception(e)
- logger.warning("Failed to parse string response as JSON: %s", e)
- new_insights = ReasoningResponse(explicit=[], deductive=[])
- else:
- # If response is already a ReasoningResponse object
- new_insights = response_obj
-
- # Extract thinking content from the response
- thinking: str | None = None
- try:
- # Try to get thinking from the response object using getattr for safety
- response_attr = getattr(response_obj, "_response", None)
- if response_attr:
- thinking = getattr(response_attr, "thinking", None)
- else:
- thinking = getattr(response_obj, "thinking", None)
-
- if thinking is None:
- logger.debug("No thinking content found in response")
- except (AttributeError, TypeError) as e:
- logger.warning("Error accessing thinking content: %s, setting to None", e)
- thinking = None
-
- response = ReasoningResponseWithThinking(
- thinking=thinking,
- explicit=new_insights.explicit,
- deductive=new_insights.deductive,
+ reasoning_response = Representation.from_prompt_representation(
+ reasoning_response,
+ (earliest_message.id, latest_message.id),
+ latest_message.session_name,
+ latest_message.created_at,
)
- logger.debug(
- "🚀 DEBUG: new_insights=%s, thinking_length=%s",
- new_insights,
- len(thinking) if thinking else 0,
- )
-
- # if settings.LANGFUSE_PUBLIC_KEY:
- # lf.update_current_generation(
- # output=format_reasoning_response_as_markdown(response),
- # )
-
- return response
-
- @conditional_observe
- @sentry_sdk.trace
- async def reason(
- self,
- working_representation: ReasoningResponseWithThinking,
- history: str,
- speaker_peer_card: list[str] | None,
- ) -> ReasoningResponseWithThinking:
- """
- Single-pass reasoning function that critically analyzes and derives insights.
- Performs one analysis pass and returns the final observations.
- """
- latest_message = self.ctx[-1]
- analysis_start = time.perf_counter()
-
- # Perform critical analysis to get observation lists
- reasoning_response = await self.derive_new_insights(
- working_representation,
- history,
- speaker_peer_card,
- )
-
- # Output the thinking content for this analysis
- log_thinking_panel(reasoning_response.thinking)
+ if lf:
+ lf.update_current_generation(
+ output=reasoning_response.format_as_markdown(),
+ )
analysis_duration_ms = (time.perf_counter() - analysis_start) * 1000
accumulate_metric(
- f"deriver_representation_{latest_message.id}_{self.target_name}",
+ f"deriver_{latest_message.id}_{self.observer}",
"critical_analysis_duration",
analysis_duration_ms,
"ms",
)
- save_observations_start = time.perf_counter()
- # Save only the NEW observations that weren't in the original context
- new_observations_by_level: dict[
- str, list[str]
- ] = await self._save_new_observations(
- working_representation, reasoning_response, latest_message
+ # Save only the new observations that weren't in the original context
+ new_observations = working_representation.diff_representation(
+ reasoning_response
)
- save_observations_duration = (
- time.perf_counter() - save_observations_start
- ) * 1000
- accumulate_metric(
- f"deriver_representation_{latest_message.id}_{self.target_name}",
- "save_new_observations",
- save_observations_duration,
- "ms",
- )
-
- update_peer_card_start = time.perf_counter()
- # flatten new observations by level into a list
- new_observations = [
- extract_observation_content(observation)
- for level in new_observations_by_level.values()
- for observation in level
- ]
- if new_observations:
- await self._update_peer_card(speaker_peer_card, new_observations)
- update_peer_card_duration = (
- time.perf_counter() - update_peer_card_start
- ) * 1000
- accumulate_metric(
- f"deriver_representation_{latest_message.id}_{self.target_name}",
- "update_peer_card",
- update_peer_card_duration,
- "ms",
- )
-
- return reasoning_response
-
- @conditional_observe
- @sentry_sdk.trace
- async def _save_new_observations(
- self,
- original_working_representation: ReasoningResponse
- | ReasoningResponseWithThinking,
- revised_observations: ReasoningResponse | ReasoningResponseWithThinking,
- latest_message: Message,
- ) -> dict[str, list[str]]:
- """Save only the observations that are new compared to the original context."""
- # Use the utility function to find new observations
- new_observations_by_level: dict[str, list[str]] = find_new_observations(
- original_working_representation, revised_observations
- )
-
- all_unified_observations: list[UnifiedObservation] = []
- total_observations_count: int = 0
-
- for level, new_observations in new_observations_by_level.items():
- if not new_observations:
- logger.debug("No new observations to save for %s level", level)
- continue
-
- logger.debug("Found %s new %s observations", len(new_observations), level)
-
- # Convert each observation to UnifiedObservation with proper premises and level
- for observation in new_observations:
- if isinstance(observation, DeductiveObservation):
- # Create UnifiedObservation with premises from DeductiveObservation
- unified_obs = UnifiedObservation(
- conclusion=observation.conclusion,
- premises=observation.premises,
- level=level,
- )
- all_unified_observations.append(unified_obs)
- logger.debug(
- "Added %s observation: %s... with %s premises",
- level,
- observation.conclusion[:50],
- len(observation.premises),
- )
-
- else:
- # String observations (explicit) have no premises
- unified_obs = UnifiedObservation.from_string(
- observation, level=level
- )
- all_unified_observations.append(unified_obs)
- logger.debug("Added %s observation: %s...", level, observation[:50])
-
- total_observations_count += 1
-
- if all_unified_observations:
- await self.embedding_store.save_unified_observations(
- all_unified_observations,
- latest_message.id,
+ if not new_observations.is_empty():
+ await self.representation_manager.save_representation(
+ new_observations,
+ (earliest_message.id, latest_message.id),
latest_message.session_name,
latest_message.created_at,
)
- else:
- logger.debug("No new observations to save")
- return new_observations_by_level
+ # not currently deduplicating at the save_representation step, so this isn't useful
+ # accumulate_metric(
+ # f"deriver_{latest_payload.message_id}_{latest_payload.observer}",
+ # "new_observation_count",
+ # new_observations_saved,
+ # "count",
+ # )
+
+ if settings.PEER_CARD.ENABLED:
+ update_peer_card_start = time.perf_counter()
+ if not new_observations.is_empty():
+ await self._update_peer_card(speaker_peer_card, new_observations)
+ update_peer_card_duration = (
+ time.perf_counter() - update_peer_card_start
+ ) * 1000
+ accumulate_metric(
+ f"deriver_{latest_message.id}_{self.observer}",
+ "update_peer_card",
+ update_peer_card_duration,
+ "ms",
+ )
+
+ return reasoning_response
@conditional_observe
@sentry_sdk.trace
async def _update_peer_card(
self,
old_peer_card: list[str] | None,
- new_observations: list[str],
+ new_observations: Representation,
) -> None:
"""
Update the peer card by calling LLM with the old peer card and new observations.
@@ -672,6 +421,7 @@ class CertaintyReasoner:
"""
try:
response = await peer_card_call(old_peer_card, new_observations)
+ logger.info("Jettisoned notes from peer card: %s", response.notes)
new_peer_card = response.card
if not new_peer_card:
logger.info("No changes to peer card")
@@ -680,109 +430,18 @@ class CertaintyReasoner:
new_peer_card = [
observation
for observation in new_peer_card
- if not observation.lower().startswith("notes")
+ if not observation.lower().startswith(("note", "notes"))
]
logger.info("New peer card: %s", new_peer_card)
async with tracked_db("deriver.update_peer_card") as db:
await crud.set_peer_card(
db,
self.ctx[0].workspace_name,
- self.sender_name,
- self.target_name,
new_peer_card,
+ observer=self.observer,
+ observed=self.observed,
)
except Exception as e:
if settings.SENTRY.ENABLED:
sentry_sdk.capture_exception(e)
logger.error("Error updating peer card! Skipping... %s", e)
-
-
-def observation_context_to_reasoning_response(
- context: ObservationContext,
-) -> ReasoningResponseWithThinking:
- """Convert ObservationContext to ReasoningResponse for compatibility."""
- thinking = context.thinking
-
- # Convert explicit observations to new structure
- explicit: list[str] = []
- for obs in context.explicit:
- explicit.append(obs.content)
-
- # Convert deductive observations
- deductive: list[DeductiveObservation] = []
- for obs in context.deductive:
- deductive_obs = DeductiveObservation(
- conclusion=obs.content,
- premises=obs.metadata.premises if obs.metadata else [],
- )
- deductive.append(deductive_obs)
-
- return ReasoningResponseWithThinking(
- thinking=thinking,
- explicit=explicit,
- deductive=deductive,
- )
-
-
-@sentry_sdk.trace
-async def save_working_representation_to_peer(
- latest_message: Message,
- final_observations: ReasoningResponseWithThinking,
- sender_name: str,
- target_name: str,
-) -> None:
- """Save working representation to peer internal_metadata for dialectic access."""
-
- # Convert ReasoningResponse to serializable dict
- final_obs_dict = {
- "thinking": final_observations.thinking,
- "explicit": final_observations.explicit,
- "deductive": [
- {
- "conclusion": obs.conclusion,
- "premises": obs.premises,
- }
- for obs in final_observations.deductive
- ],
- }
-
- working_rep_data = {
- "final_observations": final_obs_dict,
- "message_id": latest_message.id,
- "created_at": utc_now_iso(),
- }
-
- async with tracked_db("deriver.save_working_representation") as db:
- await crud.set_working_representation(
- db,
- working_rep_data,
- latest_message.workspace_name,
- target_name,
- sender_name,
- latest_message.session_name,
- )
-
-
-def _estimate_working_representation_tokens(
- working_rep_data: dict[str, Any] | str | None,
-) -> int:
- """Estimate tokens for working representation data."""
- if isinstance(working_rep_data, str):
- return estimate_tokens(working_rep_data)
-
- if (
- not isinstance(working_rep_data, dict)
- or "final_observations" not in working_rep_data
- ):
- return 0
-
- final_obs = working_rep_data["final_observations"]
-
- explicit_tokens = estimate_tokens(final_obs.get("explicit"))
- thinking_tokens = estimate_tokens(final_obs.get("thinking"))
- deductive_tokens = sum(
- estimate_tokens(d.get("conclusion")) + estimate_tokens(d.get("premises"))
- for d in final_obs.get("deductive", [])
- )
-
- return explicit_tokens + thinking_tokens + deductive_tokens
diff --git a/src/deriver/enqueue.py b/src/deriver/enqueue.py
index c097e99c..ee66a6be 100644
--- a/src/deriver/enqueue.py
+++ b/src/deriver/enqueue.py
@@ -7,11 +7,11 @@ from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, schemas
from src.config import settings
from src.dependencies import tracked_db
-from src.deriver.utils import get_work_unit_key
+from src.dreamer.dream_scheduler import get_affected_dream_keys, get_dream_scheduler
from src.exceptions import ValidationException
from src.models import QueueItem
-
-from .queue_payload import create_payload
+from src.utils.queue_payload import create_payload
+from src.utils.work_unit import get_work_unit_key
logger = logging.getLogger(__name__)
@@ -24,7 +24,22 @@ async def enqueue(payload: list[dict[str, Any]]) -> None:
payload: List of message payload dictionaries
"""
- # Use the get_db dependency to ensure proper transaction handling
+ # Cancel any pending dreams for affected collections since user is active again
+ dream_scheduler = get_dream_scheduler()
+ if dream_scheduler and payload:
+ cancelled_dreams: set[str] = set()
+ for message in payload:
+ # Generate work unit keys for dreams that might be affected by this message
+ dream_keys: list[str] = get_affected_dream_keys(message)
+ for dream_key in dream_keys:
+ if dream_scheduler.cancel_dream(dream_key):
+ cancelled_dreams.add(dream_key)
+
+ if cancelled_dreams:
+ logger.info(
+ f"Cancelled {len(cancelled_dreams)} pending dreams due to new activity"
+ )
+
async with tracked_db("message_enqueue") as db_session:
try:
# Determine if batch or single processing
@@ -136,17 +151,18 @@ async def get_peers_with_configuration(
def create_representation_record(
message: dict[str, Any],
- sender_name: str,
- target_name: str,
session_id: str | None = None,
+ *,
+ observer: str,
+ observed: str,
) -> dict[str, Any]:
"""
Create a queue record for representation task.
Args:
message: The message payload
- sender_name: Name of the sender
- target_name: Name of the target
+ observed: Name of the sender
+ observer: Name of the target
session_id: Optional session ID
Returns:
@@ -154,14 +170,12 @@ def create_representation_record(
"""
processed_payload = create_payload(
message=message,
- sender_name=sender_name,
- target_name=target_name,
task_type="representation",
+ observer=observer,
+ observed=observed,
)
return {
- "work_unit_key": get_work_unit_key(
- task_type="representation", payload=processed_payload
- ),
+ "work_unit_key": get_work_unit_key(processed_payload),
"payload": processed_payload,
"session_id": session_id,
"task_type": "representation",
@@ -178,8 +192,8 @@ def create_summary_record(
Args:
message: The message payload
- sender_name: Name of the sender
- target_name: Name of the target
+ observed: Name of the sender
+ observer: Name of the target
session_id: Session ID
Returns:
@@ -191,9 +205,7 @@ def create_summary_record(
message_seq_in_session=message_seq_in_session,
)
return {
- "work_unit_key": get_work_unit_key(
- task_type="summary", payload=processed_payload
- ),
+ "work_unit_key": get_work_unit_key(processed_payload),
"payload": processed_payload,
"session_id": session_id,
"task_type": "summary",
@@ -201,13 +213,13 @@ def create_summary_record(
def get_effective_observe_me(
- sender_name: str, peers_with_configuration: dict[str, list[dict[str, Any]]]
+ observed: str, peers_with_configuration: dict[str, list[dict[str, Any]]]
) -> bool:
"""
Determine the effective observe_me setting for a sender, considering session and peer configurations.
Args:
- sender_name: Name of the sender
+ observed: Name of the sender
peers_with_configuration: Dictionary of peer configurations
Returns:
@@ -216,7 +228,7 @@ def get_effective_observe_me(
# If the sender is not in peers_with_configuration, they left after sending a message.
# We'll use the default behavior of observing the sender by instantiating the default
# peer-level and session-level configs.
- configuration: list[Any] = peers_with_configuration.get(sender_name, [{}, {}])
+ configuration: list[Any] = peers_with_configuration.get(observed, [{}, {}])
sender_session_peer_config = (
schemas.SessionPeerConfig(**configuration[1]) if configuration[1] else None
)
@@ -257,7 +269,7 @@ async def generate_queue_records(
Returns:
List of queue records for this message
"""
- sender_name = message["peer_name"]
+ observed = message["peer_name"]
message_id: int = message["message_id"]
# Use pre-fetched sequence if available, otherwise fall back to individual query
@@ -273,7 +285,7 @@ async def generate_queue_records(
records: list[dict[str, Any]] = []
- if (
+ if settings.SUMMARY.ENABLED and (
message_seq_in_session % settings.SUMMARY.MESSAGES_PER_SHORT_SUMMARY == 0
or message_seq_in_session % settings.SUMMARY.MESSAGES_PER_LONG_SUMMARY == 0
):
@@ -288,19 +300,19 @@ async def generate_queue_records(
if deriver_disabled:
return records
- if get_effective_observe_me(sender_name, peers_with_configuration):
+ if get_effective_observe_me(observed, peers_with_configuration):
# global representation task
records.append(
create_representation_record(
message,
- sender_name=sender_name,
- target_name=sender_name,
+ observed=observed,
+ observer=observed,
session_id=session_id,
)
)
for peer_name, configuration in peers_with_configuration.items():
- if peer_name == sender_name:
+ if peer_name == observed:
continue
# If the observer peer has left the session, we don't need to enqueue a representation task for them.
@@ -320,21 +332,21 @@ async def generate_queue_records(
# peer representation task
create_representation_record(
message,
- sender_name=sender_name,
- target_name=peer_name,
+ observed=observed,
+ observer=peer_name,
session_id=session_id,
)
)
logger.debug(
"enqueued representation task for %s's representation of %s",
peer_name,
- sender_name,
+ observed,
)
logger.info(
"message %s from %s created %s queue items",
message_id,
- sender_name,
+ observed,
len(records),
)
diff --git a/src/deriver/prompts.py b/src/deriver/prompts.py
index 1371fb65..2ef199eb 100644
--- a/src/deriver/prompts.py
+++ b/src/deriver/prompts.py
@@ -9,14 +9,15 @@ import datetime
from functools import cache
from inspect import cleandoc as c
-from src.deriver.utils import estimate_tokens
+from src.utils.representation import Representation
+from src.utils.tokens import estimate_tokens
def critical_analysis_prompt(
peer_id: str,
peer_card: list[str] | None,
message_created_at: datetime.datetime,
- working_representation: str | None,
+ working_representation: Representation,
history: str,
new_turns: list[str],
) -> str:
@@ -27,7 +28,7 @@ def critical_analysis_prompt(
peer_id (str): The ID of the user being analyzed.
peer_card (list[str] | None): The bio card of the user being analyzed.
message_created_at (datetime.datetime): Timestamp of the message.
- working_representation (str | None): Current user understanding context.
+ working_representation (Representation): Current user understanding context.
history (str): Recent conversation history.
new_turns (list[str]): New conversation turns to analyze.
@@ -50,10 +51,10 @@ def critical_analysis_prompt(
f"""
Current understanding of {peer_id}:
-{working_representation}
+{str(working_representation)}
"""
- if working_representation is not None
+ if not working_representation.is_empty()
else ""
)
@@ -61,7 +62,7 @@ Current understanding of {peer_id}:
return c(
f"""
-You are an agent who critically analyzes user messages through rigorous logical reasoning to produce only conclusions about the user that are CERTAIN.
+You are an agent who critically analyzes messages from {peer_id} through rigorous logical reasoning to produce only conclusions about them that are CERTAIN.
TARGET USER TO ANALYZE
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
@@ -80,18 +81,55 @@ Your goal is to IMPROVE understanding of {peer_id} through careful analysis. You
Here are strict definitions for the reasoning modes you are to employ:
1. **EXPLICIT REASONING**:
- - Conclusions about the user that MUST be true given premises ONLY of the following types:
- - Most recent user message
+ - Conclusions about {peer_id} that MUST be true given premises ONLY of the following types:
+ - Recent messages
- Knowledge about the conversation history
- Current date and time (which is: {message_created_at})
- Timestamps from conversation history
+ - Follow strict literal necessity--if stated directly in message, extract a conclusion
+ - Latest message MUST be a premise, previous messages and timestamps may be used to contextualize
+ - Transforms a single message (premise) into ONE OR MULTIPLE conclusions
+ - Derive EVERYTHING that can be explicitly concluded
+ - Make sure EVERY conclusion is sufficiently contextualized, i.e. ensure each conclusion contains enough specific information about subjects and objects to make it self-contained and useful (e.g. instead of "Ann is nervous about the interview", use "Ann is nervous about the job interview at the pharmacy")
+ - When possible, always use absolute dates and times, and avoid relative dates and times (e.g. instead of 'Mary went to the store yesterday', use 'Mary went to the store on June 26, 2025')
2. **DEDUCTIVE REASONING**:
- - Conclusions about the user that MUST be true given premises ONLY of the following types:
+ - Conclusions about {peer_id} that MUST be true given premises ONLY of the following types:
- Explicit conclusions
- Previous deductive conclusions
- General, open domain knowledge known to be true
- Current date and time (which is: {message_created_at})
- - Timestamps for user messages, and previous premises and conclusions
+ - Timestamps for {peer_id}'s messages, and previous premises and conclusions
+ - Follow strict logical necessity--if premises are true, conclusion MUST be true
+ - Multiple premises may be used in a deduction, but only one conclusion may be drawn
+ - Complete ONLY as many deductions as needed to form useful and additive knowledge about {peer_id}
+ - May scaffold previous conclusions and known facts to do further deduction
+ - But MAY NOT use previous **probabilistic** deductive conclusions (including qualifiers like probably, likely, typically, may, etc) as premises in further deductions
+ - Use current timestamp as needed to provide absolute dates
+
+Here are examples of the reasoning modes in action:
+
+- **EXPLICIT REASONING EXAMPLES**
+ 1. PREMISE(S): "I just had my 25th birthday last Saturday" (latest message), Current date is June 26, 2025 (timestamp) → CONCLUSION(S): "Maria is 25 years old", "Maria's birthday is June 21st"
+ 2. PREMISE(S): "I took my dog for a walk in a park near my house in NYC—it was such a beautiful day" (latest message) → CONCLUSION(S): "Liam has a dog", "Liam took his dog for a walk", "Liam has a house in NYC", "Liam lives near a park", "Liam prefers to take advantage of nice weather to walk his dog"
+ 3. PREMISE(S): "Whenever I think about my college experience I feel nostalgic" (latest message) → CONCLUSION(S): "Aisha attended college", "Aisha feels nostalgic about her college experience"
+ 4. PREMISE(S): "That's so cool!" (latest message), The speaker is reacting to learning the definition of Kant's categorical imperative (conversation knowledge) → CONCLUSION(S): "Carlos thinks Kant's categorical imperative is cool"
+- **DEDUCTIVE REASONING EXAMPLES**
+ 1. PREMISE(S): "Maria attended college" (explicit), All people who attended college have completed high school or equivalent (general) → CONCLUSION: "Maria completed high school or equivalent education"
+ 2. PREMISE(S): "Liam is 25 years old" (explicit), Current date is June 26, 2025 (timestamp), "Liam's birthday was last Saturday" (explicit) → CONCLUSION: "Liam was born on June 21, 1998"
+ 3. PREMISE(S): "Aisha has a dog" (explicit), "Aisha took her dog for a walk" (explicit), All dogs require regular walks for health (general) → CONCLUSION: "Aisha provides care for her dog"
+ 4. PREMISE(S): "Carlos prefers to take advantage of nice weather to walk his dog" (explicit), Message timestamp shows afternoon hours (timestamp), Nice weather is typically during daylight (general) → CONCLUSION: "Carlos has flexibility in his schedule during typical work hours"
+
+Based on our definitions and examples, here's a summary of the logical reasoning task:
+
+**REASONING INTERACTIONS:**
+
+- Message (required)/Conversation History (optional)/Temporal (optional) → Explicit: Derive certain conclusions only from literal statements
+- Explicit/Deductive/Temporal/General → Deductive: When logical necessity allows certain conclusion
+- Explicit/Deductive/Temporal/General → Further Deductive: Can use certain conclusions and known facts to deduce additional certain conclusions
+- Probabilistic Deductive ↛ Further Deductive: If a deductive conclusion includes probabilistic qualifiers (likely, potentially, typically, might, etc) it may NOT be used as a premise for further deductions
+
+**INSTRUCTIONS:** Given the above, first think critically about what it means to do explicit and deductive reasoning, then consider how to apply that to the latest message, finally do explicit and deductive reasoning about the user to reach useful, contextually-rich conclusions.
+
{peer_card_section}
@@ -112,11 +150,18 @@ New conversation turns to analyze:
def peer_card_prompt(
old_peer_card: list[str] | None,
- new_observations: list[str],
+ new_observations: str,
) -> str:
"""
Generate the peer card prompt for the deriver.
Currently optimized for GPT-5 mini/nano.
+
+ Args:
+ old_peer_card: Existing biographical card lines, if any.
+ new_observations: Pre-formatted observations block (multiple lines).
+
+ Returns:
+ Formatted prompt string for (re)generating the peer card JSON.
"""
old_peer_card_section = (
f"""
@@ -161,7 +206,8 @@ Example 2:
{old_peer_card_section}
New observations:
-{chr(10).join(new_observations)}
+
+{new_observations}
If there's no new key info, set "card" to null (or omit it) to signal no update. **NEVER** include notes or temporary information in the card itself, instead use the notes field. There are no mandatory fields -- if you can't find a value, just leave it out. **ONLY** include information that is **GIVEN**.
""" # nosec B608 <-- this is a really dumb false positive
@@ -180,7 +226,7 @@ def estimate_base_prompt_tokens() -> int:
peer_id="",
peer_card=None,
message_created_at=datetime.datetime.now(datetime.timezone.utc),
- working_representation=None,
+ working_representation=Representation(),
history="",
new_turns=[],
)
diff --git a/src/deriver/queue_manager.py b/src/deriver/queue_manager.py
index 539e685f..b5706a76 100644
--- a/src/deriver/queue_manager.py
+++ b/src/deriver/queue_manager.py
@@ -22,8 +22,17 @@ from src.deriver.consumer import (
process_item,
process_representation_batch,
)
-from src.deriver.utils import parse_work_unit_key
+from src.dreamer.dream_scheduler import (
+ DreamScheduler,
+ get_dream_scheduler,
+ set_dream_scheduler,
+)
from src.models import QueueItem
+from src.utils.work_unit import parse_work_unit_key
+from src.webhooks.events import (
+ QueueEmptyEvent,
+ publish_webhook_event,
+)
logger = getLogger(__name__)
@@ -48,6 +57,14 @@ class QueueManager:
self.workers: int = settings.DERIVER.WORKERS
self.semaphore: asyncio.Semaphore = asyncio.Semaphore(self.workers)
+ # Get or create the singleton dream scheduler
+ existing_scheduler = get_dream_scheduler()
+ if existing_scheduler is None:
+ self.dream_scheduler: DreamScheduler = DreamScheduler()
+ set_dream_scheduler(self.dream_scheduler)
+ else:
+ self.dream_scheduler = existing_scheduler
+
# Initialize Sentry if enabled, using settings
if settings.SENTRY.ENABLED:
sentry_sdk.init(
@@ -110,6 +127,9 @@ class QueueManager:
logger.info(f"Received exit signal {sig.name}...")
self.shutdown_event.set()
+ # Cancel all pending dreams
+ await self.dream_scheduler.shutdown()
+
if self.active_tasks:
logger.info(
f"Waiting for {len(self.active_tasks)} active tasks to complete..."
@@ -188,6 +208,7 @@ class QueueManager:
) as db: # Get number of available workers
query = (
select(models.QueueItem.work_unit_key)
+ .limit(limit)
.outerjoin(
models.ActiveQueueSession,
models.QueueItem.work_unit_key
@@ -197,7 +218,6 @@ class QueueManager:
.where(models.QueueItem.work_unit_key.isnot(None))
.where(models.ActiveQueueSession.work_unit_key.is_(None))
.distinct()
- .limit(limit)
)
result = await db.execute(query)
@@ -290,16 +310,12 @@ class QueueManager:
######################
async def process_work_unit(self, work_unit_key: str, worker_id: str) -> None:
"""Process all messages for a specific work unit by routing to the correct handler."""
- logger.debug(
- f"Worker {worker_id} starting to process work unit {work_unit_key}"
- )
+ logger.debug(f"Starting to process work unit {work_unit_key}")
+ work_unit = parse_work_unit_key(work_unit_key)
async with self.semaphore:
message_count = 0
messages_to_process: list[QueueItem] = []
try:
- parsed_key = parse_work_unit_key(work_unit_key)
- task_type = parsed_key["task_type"]
-
while not self.shutdown_event.is_set():
# Get worker ownership info for verification
ownership = self.worker_ownership.get(worker_id)
@@ -309,12 +325,12 @@ class QueueManager:
)
break
try:
- if task_type == "representation":
+ if work_unit.task_type == "representation":
(
messages_context,
items_to_process,
) = await self.get_message_batch(
- task_type, work_unit_key, ownership.aqs_id
+ work_unit.task_type, work_unit_key, ownership.aqs_id
)
logger.debug(
f"Worker {worker_id} retrieved {len(messages_context)} messages and {len(items_to_process)} queue items for work unit {work_unit_key} (AQS ID: {ownership.aqs_id})"
@@ -326,13 +342,10 @@ class QueueManager:
break
# Build payloads from the unique messages context window
- sender_name = parsed_key["sender_name"]
- target_name = parsed_key["target_name"]
-
await process_representation_batch(
messages_context,
- sender_name=sender_name,
- target_name=target_name,
+ observer=work_unit.observer,
+ observed=work_unit.observed,
)
await self.mark_messages_as_processed(
@@ -342,7 +355,7 @@ class QueueManager:
else:
messages_to_process = await self.get_next_message(
- task_type, work_unit_key, ownership.aqs_id
+ work_unit.task_type, work_unit_key, ownership.aqs_id
)
if not messages_to_process:
logger.debug(
@@ -350,7 +363,7 @@ class QueueManager:
)
break
await process_item(
- task_type, messages_to_process[0].payload
+ work_unit.task_type, messages_to_process[0].payload
)
await self.mark_messages_as_processed(
messages_to_process, work_unit_key
@@ -387,23 +400,17 @@ class QueueManager:
if removed and message_count > 0:
# Only publish webhook if we actually removed an active session
try:
- from src.webhooks.events import (
- QueueEmptyEvent,
- publish_webhook_event,
- )
-
- parsed_key = parse_work_unit_key(work_unit_key)
- if parsed_key["task_type"] in ["representation", "summary"]:
+ if work_unit.task_type in ["representation", "summary"]:
logger.debug(
f"Publishing queue.empty event for {work_unit_key}"
)
await publish_webhook_event(
QueueEmptyEvent(
- workspace_id=parsed_key["workspace_name"],
- queue_type=parsed_key["task_type"],
- session_id=parsed_key["session_name"],
- sender_name=parsed_key["sender_name"],
- observer_name=parsed_key["target_name"],
+ workspace_id=work_unit.workspace_name,
+ queue_type=work_unit.task_type,
+ session_id=work_unit.session_name,
+ observer=work_unit.observer,
+ observed=work_unit.observed,
)
)
else:
@@ -475,8 +482,6 @@ class QueueManager:
# For representation tasks, get a batch based on token limit.
# Step 1: Parse work_unit_key to get session context and focused sender
parsed_key = parse_work_unit_key(work_unit_key)
- session_name = parsed_key["session_name"]
- workspace_name = parsed_key["workspace_name"]
# Verify worker still owns the work_unit_key
ownership_check = await db.execute(
@@ -506,8 +511,8 @@ class QueueManager:
== models.Message.id,
)
.where(~models.QueueItem.processed)
- .where(models.Message.session_name == session_name)
- .where(models.Message.workspace_name == workspace_name)
+ .where(models.Message.session_name == parsed_key.session_name)
+ .where(models.Message.workspace_name == parsed_key.workspace_name)
.where(models.QueueItem.work_unit_key == work_unit_key)
.scalar_subquery()
)
@@ -518,13 +523,13 @@ class QueueManager:
select(
models.Message.id.label("message_id"),
models.Message.token_count.label("token_count"),
- models.Message.peer_name.label("sender_name"),
+ models.Message.peer_name.label("peer_name"),
func.sum(models.Message.token_count)
.over(order_by=models.Message.id)
.label("cumulative_token_count"),
)
- .where(models.Message.session_name == session_name)
- .where(models.Message.workspace_name == workspace_name)
+ .where(models.Message.session_name == parsed_key.session_name)
+ .where(models.Message.workspace_name == parsed_key.workspace_name)
.where(models.Message.id >= min_unprocessed_message_id_subq)
.order_by(models.Message.id)
.cte()
diff --git a/src/deriver/utils.py b/src/deriver/utils.py
deleted file mode 100644
index 244dbd33..00000000
--- a/src/deriver/utils.py
+++ /dev/null
@@ -1,82 +0,0 @@
-import tiktoken
-from typing_extensions import Any, TypedDict
-
-tokenizer = tiktoken.get_encoding("cl100k_base")
-
-
-class ParsedWorkUnit(TypedDict):
- task_type: str
- workspace_name: str
- session_name: str | None
- sender_name: str | None
- target_name: str | None
-
-
-def get_work_unit_key(task_type: str, payload: dict[str, Any]) -> str:
- """
- Generate a work unit key for a given task type, workspace name, and event type.
- """
- workspace_name = payload.get("workspace_name")
- if not workspace_name:
- raise ValueError("workspace_name is required to generate a work_unit_key")
-
- if task_type in ["representation", "summary"]:
- sender_name = payload.get("sender_name", "None")
- target_name = payload.get("target_name", "None")
- session_name = payload.get("session_name", "None")
- return (
- f"{task_type}:{workspace_name}:{session_name}:{sender_name}:{target_name}"
- )
-
- if task_type == "webhook":
- return f"webhook:{workspace_name}"
-
- raise ValueError(f"Invalid task type: {task_type}")
-
-
-def parse_work_unit_key(work_unit_key: str) -> ParsedWorkUnit:
- """
- Parse a work unit key to extract its components.
- """
- parts = work_unit_key.split(":")
- task_type = parts[0]
-
- if task_type in ["representation", "summary"]:
- if len(parts) != 5:
- raise ValueError(
- f"Invalid work_unit_key format for task_type {task_type}: {work_unit_key}"
- )
- return {
- "task_type": task_type,
- "workspace_name": parts[1],
- "session_name": parts[2],
- "sender_name": parts[3],
- "target_name": parts[4],
- }
-
- if task_type == "webhook":
- if len(parts) != 2:
- raise ValueError(
- f"Invalid work_unit_key format for task_type {task_type}: {work_unit_key}"
- )
- return {
- "task_type": task_type,
- "workspace_name": parts[1],
- "session_name": None,
- "sender_name": None,
- "target_name": None,
- }
-
- raise ValueError(f"Invalid task type in work_unit_key: {task_type}")
-
-
-def estimate_tokens(text: str | list[str] | None) -> int:
- """Estimate token count using tiktoken for text or list of strings."""
- if not text:
- return 0
- if isinstance(text, list):
- text = "\n".join(text)
- try:
- return len(tokenizer.encode(text))
- except Exception:
- return len(text) // 4
diff --git a/src/dialectic/chat.py b/src/dialectic/chat.py
index 99e7a468..cfca960b 100644
--- a/src/dialectic/chat.py
+++ b/src/dialectic/chat.py
@@ -6,29 +6,27 @@ and understand users through context synthesis of working representations and
historical observations.
"""
-import asyncio
import logging
+import time
import uuid
from collections.abc import AsyncIterator
-import tiktoken
from dotenv import load_dotenv
-from langfuse import get_client
from src import crud
from src.config import settings
-from src.crud.representation import GLOBAL_REPRESENTATION_COLLECTION_NAME
from src.dependencies import tracked_db
from src.utils import summarizer
from src.utils.clients import HonchoLLMCallStreamChunk, honcho_llm_call
-from src.utils.embedding_store import EmbeddingStore
+from src.utils.langfuse_client import get_langfuse_client
from src.utils.logging import (
accumulate_metric,
log_performance_metrics,
)
+from src.utils.representation import Representation
+from src.utils.tokens import estimate_tokens
from .prompts import dialectic_prompt
-from .utils import get_observations
# Configure logging
logger = logging.getLogger(__name__)
@@ -37,27 +35,30 @@ logger = logging.getLogger(__name__)
load_dotenv()
# Create langfuse client
-lf = get_client()
+lf = get_langfuse_client() if settings.LANGFUSE_PUBLIC_KEY else None
async def dialectic_call(
query: str,
- working_representation: str | None,
+ working_representation: str,
recent_conversation_history: str | None,
- additional_context: str | None,
- peer_name: str,
peer_card: list[str] | None,
- target_name: str | None = None,
- target_peer_card: list[str] | None = None,
+ observed_peer_card: list[str] | None = None,
+ *,
+ observer: str,
+ observed: str,
):
"""
Make a direct call to the dialectic model for context synthesis.
Args:
query: The user query
- working_representation: Current session conclusions
- additional_context: Historical context from semantic search
+ working_representation: Current session conclusions AND historical conclusions from the user's global representation
+ recent_conversation_history: Recent conversation history
peer_name: Name of the user/peer
+ peer_card: Known biographical information about the user
+ observed: Name of the user/peer being queried about
+ observed_peer_card: Known biographical information about the target, if applicable
Returns:
Model response
@@ -67,11 +68,10 @@ async def dialectic_call(
query,
working_representation,
recent_conversation_history,
- additional_context,
- peer_name,
peer_card,
- target_name,
- target_peer_card,
+ observed_peer_card,
+ observer=observer,
+ observed=observed,
)
response = await honcho_llm_call(
@@ -96,22 +96,25 @@ async def dialectic_call(
async def dialectic_stream(
query: str,
- working_representation: str | None,
+ working_representation: str,
recent_conversation_history: str | None,
- additional_context: str | None,
- peer_name: str,
peer_card: list[str] | None,
- target_name: str | None = None,
- target_peer_card: list[str] | None = None,
+ observed_peer_card: list[str] | None = None,
+ *,
+ observer: str,
+ observed: str,
):
"""
Make a streaming call to the dialectic model for context synthesis.
Args:
query: The user query
- working_representation: Current session conclusions
- additional_context: Historical context from semantic search
+ working_representation: Current session conclusions AND historical conclusions from the user's global representation
+ recent_conversation_history: Recent conversation history
peer_name: Name of the user/peer
+ peer_card: Known biographical information about the user
+ observed: Name of the user/peer being queried about
+ observed_peer_card: Known biographical information about the target, if applicable
Returns:
Streaming model response
@@ -121,11 +124,10 @@ async def dialectic_stream(
query,
working_representation,
recent_conversation_history,
- additional_context,
- peer_name,
peer_card,
- target_name,
- target_peer_card,
+ observed_peer_card,
+ observer=observer,
+ observed=observed,
)
response = await honcho_llm_call(
@@ -151,11 +153,11 @@ async def dialectic_stream(
async def chat(
workspace_name: str,
- peer_name: str,
- target_name: str | None,
session_name: str | None,
query: str,
*,
+ observer: str,
+ observed: str,
stream: bool = False,
) -> str | AsyncIterator[HonchoLLMCallStreamChunk]:
"""
@@ -170,7 +172,7 @@ async def chat(
Args:
workspace_name: Name of the workspace
peer_name: Name of the peer making the query
- target_name: Optional name of the peer being queried about
+ observed: Optional name of the peer being queried about
session_name: Optional session name for scoping
query: Input Dialectic Query
stream: Whether to stream the response
@@ -181,15 +183,13 @@ async def chat(
dialectic_chat_uuid = str(uuid.uuid4())
- tokenizer = tiktoken.get_encoding("cl100k_base")
-
context_window_size = (
settings.DIALECTIC.CONTEXT_WINDOW_SIZE - 750
) # this is a hardcoded (accurate, slightly conservative) estimate of system prompt
- context_window_size -= len(tokenizer.encode(query))
+ context_window_size -= estimate_tokens(query)
- if settings.LANGFUSE_PUBLIC_KEY:
+ if lf:
lf.update_current_trace(
metadata={
"query_generation_model": settings.DIALECTIC.QUERY_GENERATION_MODEL,
@@ -198,87 +198,68 @@ async def chat(
}
)
logger.info(
- "Received query:\n'%s'\nobserver: %s%s%s\n",
+ "Received query:\n'%s'\nobserver: %s, observed: %s%s\n",
query,
- peer_name,
- f", target: {target_name}" if target_name else "",
+ observer,
+ observed,
f", session: {session_name}" if session_name else "",
)
- start_time = asyncio.get_event_loop().time()
+ start_time = time.perf_counter()
# 1. Working representation (short-term) -----------------------------------
- # Only useful for session-scoped queries, not global queries
- if session_name:
- working_rep_start_time = asyncio.get_event_loop().time()
- async with tracked_db("chat.get_working_representation") as db:
- # If no target specified, get global representation (peer observing themselves)
- target_peer = target_name if target_name is not None else peer_name
- working_representation = await crud.get_working_representation(
- db, workspace_name, peer_name, target_peer, session_name
- )
- working_rep_duration = asyncio.get_event_loop().time() - working_rep_start_time
- accumulate_metric(
- f"dialectic_chat_{dialectic_chat_uuid}",
- "retrieve_working_rep",
- working_rep_duration,
- "s",
- )
- logger.info("Retrieved working representation:\n%s\n", working_representation)
- context_window_size -= len(tokenizer.encode(working_representation))
- else:
- # For global queries, working representation isn't useful - use historical context instead
- working_representation = None
- logger.info("Query is not session-scoped, skipping working representation")
-
- # 2. Additional context (long-term semantic search) ------------------------
- # If the query is not targeted, get global_representation facts from other sessions
- # If the query is targeted, get facts from other sessions for our target
- additional_context_start_time = asyncio.get_event_loop().time()
- embedding_store = EmbeddingStore(
- workspace_name=workspace_name,
- peer_name=target_name if target_name else peer_name,
- collection_name=GLOBAL_REPRESENTATION_COLLECTION_NAME
- if not target_name
- else crud.construct_collection_name(observer=peer_name, observed=target_name),
- )
- additional_context: str = await get_observations(
- query,
- target_name if target_name else peer_name,
- embedding_store,
- include_premises=True,
- )
- additional_context_duration = (
- asyncio.get_event_loop().time() - additional_context_start_time
+ working_rep_start_time = time.perf_counter()
+ # If no target specified, get global representation (peer observing themselves)
+ working_representation: Representation = await crud.get_working_representation(
+ workspace_name,
+ observer=observer,
+ observed=observed,
+ session_name=session_name,
+ include_semantic_query=query,
+ semantic_search_top_k=settings.DIALECTIC.SEMANTIC_SEARCH_TOP_K,
+ semantic_search_max_distance=settings.DIALECTIC.SEMANTIC_SEARCH_MAX_DISTANCE,
+ include_most_derived=True,
)
+ working_rep_duration = (time.perf_counter() - working_rep_start_time) * 1000
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
- "retrieve_additional_context",
- additional_context_duration,
- "s",
+ "retrieve_working_rep",
+ working_rep_duration,
+ "ms",
+ )
+ logger.info(
+ "Retrieved working representation with %s explicit, %s deductive observations",
+ len(working_representation.explicit),
+ len(working_representation.deductive),
)
- logger.info("Retrieved additional context:\n%s", additional_context)
- context_window_size -= len(tokenizer.encode(additional_context))
+ working_representation_str = str(working_representation)
- # 3. Recent conversation history --------------------------------------------
+ context_window_size -= max(0, estimate_tokens(working_representation_str))
+
+ logger.info(
+ "Constructed working representation:\n%s\n",
+ working_representation_str,
+ )
+
+ # 2. Recent conversation history --------------------------------------------
# If query is session-scoped, get recent conversation history from that session
- if session_name:
- async with tracked_db("chat.get_session_context") as db:
- recent_conversation_history = (
- await summarizer.get_session_context_formatted(
- db,
- workspace_name=workspace_name,
- session_name=session_name,
- token_limit=context_window_size,
- include_summary=True,
- )
+ async with tracked_db("chat.get_context") as db:
+ if session_name:
+ recent_history = await summarizer.get_session_context_formatted(
+ db,
+ workspace_name=workspace_name,
+ session_name=session_name,
+ token_limit=context_window_size,
+ include_summary=True,
+ )
+ logger.info("Retrieved recent conversation history")
+ else:
+ recent_history = None
+ logger.info(
+ "Query is not session-scoped, skipping recent conversation history"
)
- logger.info("Retrieved recent conversation history")
- else:
- recent_conversation_history = None
- logger.info("Query is not session-scoped, skipping recent conversation history")
- context_window_size -= len(tokenizer.encode(recent_conversation_history or ""))
+ context_window_size -= max(0, estimate_tokens(recent_history or ""))
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
@@ -287,61 +268,63 @@ async def chat(
"tokens",
)
- # 4. Peer card(s) ----------------------------------------------------------
- async with tracked_db("chat.get_peer_card") as db:
- peer_card = await crud.get_peer_card(db, workspace_name, peer_name, peer_name)
- if target_name:
- target_peer_card = await crud.get_peer_card(
- db, workspace_name, target_name, peer_name
+ # 3. Peer card(s) ----------------------------------------------------------
+ if settings.PEER_CARD.ENABLED:
+ async with tracked_db("chat.get_peer_card") as db:
+ peer_card = await crud.get_peer_card(
+ db, workspace_name, observer=observer, observed=observed
)
+ if observer != observed:
+ observed_peer_card = await crud.get_peer_card(
+ db, workspace_name, observer=observer, observed=observed
+ )
+ else:
+ observed_peer_card = None
+
+ if observed_peer_card:
+ logger.info("Retrieved peer cards:\n%s\n%s", peer_card, observed_peer_card)
else:
- target_peer_card = None
-
- if target_peer_card:
- logger.info("Retrieved peer cards:\n%s\n%s", peer_card, target_peer_card)
+ logger.info("Retrieved peer card:\n%s", peer_card)
else:
- logger.info("Retrieved peer card:\n%s", peer_card)
+ peer_card = None
+ observed_peer_card = None
- # 5. Dialectic call --------------------------------------------------------
- dialectic_call_start_time = asyncio.get_event_loop().time()
+ # 4. Dialectic call --------------------------------------------------------
+ dialectic_call_start_time = time.perf_counter()
if stream:
return await dialectic_stream(
query,
- working_representation,
- recent_conversation_history,
- additional_context,
- peer_name,
+ working_representation_str,
+ recent_history,
peer_card,
- target_name,
- target_peer_card,
+ observed_peer_card,
+ observer=observer,
+ observed=observed,
)
response = await dialectic_call(
query,
- working_representation,
- recent_conversation_history,
- additional_context,
- peer_name,
+ working_representation_str,
+ recent_history,
peer_card,
- target_name,
- target_peer_card,
- )
- dialectic_call_duration = (
- asyncio.get_event_loop().time() - dialectic_call_start_time
+ observed_peer_card,
+ observer=observer,
+ observed=observed,
)
+ dialectic_call_duration = (time.perf_counter() - dialectic_call_start_time) * 1000
accumulate_metric(
f"dialectic_chat_{dialectic_chat_uuid}",
"dialectic_call",
dialectic_call_duration,
- "s",
+ "ms",
)
- elapsed = asyncio.get_event_loop().time() - start_time
+ elapsed = (time.perf_counter() - start_time) * 1000
accumulate_metric(
- f"dialectic_chat_{dialectic_chat_uuid}", "total_duration", elapsed, "s"
+ f"dialectic_chat_{dialectic_chat_uuid}", "total_duration", elapsed, "ms"
)
- log_performance_metrics(f"dialectic_chat_{dialectic_chat_uuid}")
+ log_performance_metrics("dialectic_chat", dialectic_chat_uuid)
# Convert AnthropicCallResponse to string for compatibility
return str(response)
diff --git a/src/dialectic/prompts.py b/src/dialectic/prompts.py
index 99cd1141..db4a9392 100644
--- a/src/dialectic/prompts.py
+++ b/src/dialectic/prompts.py
@@ -3,44 +3,47 @@ from inspect import cleandoc as c
def dialectic_prompt(
query: str,
- working_representation: str | None,
+ working_representation: str,
recent_conversation_history: str | None,
- additional_context: str | None,
- peer_name: str,
- peer_card: list[str] | None,
- target_name: str | None = None,
- target_peer_card: list[str] | None = None,
+ observer_peer_card: list[str] | None,
+ observed_peer_card: list[str] | None = None,
+ *,
+ observer: str,
+ observed: str,
) -> str:
"""
Generate the main dialectic prompt for context synthesis.
Args:
query: The specific question or request from the application about the user
- working_representation: Current session conclusions from recent conversation analysis
- additional_context: Historical conclusions from the user's global representation
- peer_name: Name of the user/peer being queried about
+ working_representation: Conclusions from recent conversation analysis AND historical conclusions from the user's global representation
+ recent_conversation_history: Recent conversation history
+ peer_card: Known biographical information about the user
+ observed_peer_card: Known biographical information about the target, if applicable
Returns:
Formatted prompt string for the dialectic model
"""
- if target_name:
- query_target = f"""The query is about user {peer_name}'s understanding of {target_name}.
+ if observer != observed:
+ # this is a directional query from the observer's view of the observed
+ query_target = f"""The query is about user {observer}'s understanding of {observed}.
The user's known biographical information:
-{chr(10).join(peer_card) if peer_card else "(none)"}
+{chr(10).join(observer_peer_card) if observer_peer_card else "(none)"}
The target's known biographical information:
-{chr(10).join(target_peer_card) if target_peer_card else "(none)"}
+{chr(10).join(observed_peer_card) if observed_peer_card else "(none)"}
If the user's name or nickname is known, exclusively refer to them by that name.
If the target's name or nickname is known, exclusively refer to them by that name.
"""
else:
- query_target = f"""The query is about user {peer_name}.
+ # this is a global query: honcho's omniscient view of the observed
+ query_target = f"""The query is about user {observed}.
The user's known biographical information:
-{chr(10).join(peer_card) if peer_card else "(none)"}
+{chr(10).join(observer_peer_card) if observer_peer_card else "(none)"}
If the user's name or nickname is known, exclusively refer to them by that name.
"""
@@ -62,6 +65,52 @@ Each conclusion contains:
- **Type**: Either Explicit or Deductive
- **Temporal Data**: When conclusions were made
+## CONCLUSION TYPE DEFINITIONS
+
+**Explicit Conclusions** (Direct Facts)
+- Direct, literal conclusions which were extracted from statements by the user in their messages
+- No interpretation - only derived from what was explicitly written
+
+**Deductive Conclusions** (Logical Certainties)
+- Conclusions that MUST be true given the premises
+- Built from premises that may include explicit conclusions, deductive conclusions, temporal premises, and/or general knowledge known to be true
+
+## SYNTHESIS PROCESS
+
+1. **Query Analysis**: Identify what specific information the application needs
+2. **Conclusion Gathering**: Collect all conclusions relevant to the query
+3. **Evidence Evaluation**: Assess conclusions quality based on:
+ - Reasoning type (explicit > deductive in certainty)
+ - Recency (newer = more current state)
+ - Premise strength (more supporting evidence = stronger)
+ - Qualifiers (likely, probably, typically, etc)
+1. **Synthesis**: Build a coherent answer that:
+ - Directly addresses the query
+ - Provides additional useful context
+ - Connects related conclusions logically
+ - Acknowledges gaps or uncertainties
+
+## SYNTHESIS PRINCIPLES
+
+**Logical Chaining**:
+- Connect conclusions across time to build deeper understanding
+- Use general knowledge to bridge gaps between user observations
+- Apply established user patterns from one domain to predict behavior in another
+
+**Temporal Awareness**:
+- Recent conclusions reflect current state
+- Historical patterns show consistent traits
+- Note when conclusions may be outdated
+
+**Evidence Integration**:
+- Multiple converging conclusions strengthen synthesis
+- Contradictions require resolution (prioritize: recency > explicit > deductive)
+- Build from certainties toward useful query answers
+
+**Response Requirements**:
+- Answer the specific question asked
+- Ground responses in actual conclusions
+
## OUTPUT FORMAT
Provide a natural language response that:
@@ -79,25 +128,24 @@ Provide a natural language response that:
{query}
{working_representation}
-{f"{additional_context}" if additional_context else ""}"""
+"""
)
-def query_generation_prompt(query: str, target_peer_name: str) -> str:
+def query_generation_prompt(query: str, observed: str) -> str:
"""
Generate the prompt for semantic query expansion.
Args:
query: The original user query
- peer_name: Name of the user/peer
- target_name: Name of the target/peer if dialectic query is targeted
+ observed: Name of the target peer
Returns:
Formatted prompt string for query generation
"""
return c(
f"""
-You are a query expansion agent helping AI applications understand their users. The user's name is {target_peer_name}. Your job is to take application queries about this user and generate targeted search queries that will retrieve the most relevant observations using semantic search over an embedding store containing observations about the user.
+You are a query expansion agent helping AI applications understand their users. The user's name is {observed}. Your job is to take application queries about this user and generate targeted search queries that will retrieve the most relevant observations using semantic search over an embedding store containing observations about the user.
## QUERY EXPANSION STRATEGY FOR SEMANTIC SIMILARITY
diff --git a/src/dialectic/utils.py b/src/dialectic/utils.py
deleted file mode 100644
index bc834266..00000000
--- a/src/dialectic/utils.py
+++ /dev/null
@@ -1,246 +0,0 @@
-import asyncio
-import json
-import logging
-from typing import Any
-
-from langfuse import get_client
-
-from src.config import settings
-from src.models import Document
-from src.utils.clients import honcho_llm_call
-from src.utils.embedding_store import EmbeddingStore
-from src.utils.formatting import (
- format_premises_for_display,
- parse_datetime_iso,
-)
-from src.utils.logging import conditional_observe
-from src.utils.shared_models import SemanticQueries
-
-from .prompts import query_generation_prompt
-
-# Configure logging
-logger = logging.getLogger(__name__)
-
-lf = get_client()
-
-
-@conditional_observe
-async def get_observations(
- query: str,
- target_peer_name: str,
- embedding_store: EmbeddingStore,
- *,
- include_premises: bool = False,
-) -> str:
- """
- Generate queries based on the dialectic query and retrieve relevant observations.
-
- Uses semantic search to find additional relevant historical context beyond
- what's already in the working representation.
-
- Args:
- query: The user query
- embedding_store: The embedding store to search
- include_premises: Whether to include premises from document metadata
-
- Returns:
- String containing additional relevant observations from semantic search
- """
- logger.info("Starting observation retrieval for query: %s", query)
-
- if settings.DIALECTIC.PERFORM_QUERY_GENERATION:
- logger.debug(
- "Attempting to generate semantic queries using %s",
- settings.DIALECTIC.QUERY_GENERATION_PROVIDER,
- )
- search_queries_result = await generate_semantic_queries(query, target_peer_name)
- logger.debug(
- "Successfully generated queries via %s: %s",
- settings.DIALECTIC.QUERY_GENERATION_PROVIDER,
- search_queries_result,
- )
-
- # search_queries_result should never be None based on function return types
-
- search_queries = search_queries_result.queries
- # Include the original query in the search queries
- search_queries.append(query)
- logger.info(
- "Generated %s search queries: \n%s",
- len(search_queries),
- json.dumps(search_queries, indent=2),
- )
-
- # Execute all queries in parallel
- tasks = [_execute_single_query(q, embedding_store) for q in search_queries]
- all_results = await asyncio.gather(*tasks)
- unique_observations = _deduplicate_observations(all_results)
-
- if settings.LANGFUSE_PUBLIC_KEY:
- lf.update_current_generation(
- input={
- "query": query,
- "include_premises": include_premises,
- },
- output={
- "search_queries": search_queries_result,
- "all_results": all_results,
- "unique_observations": unique_observations,
- },
- )
-
- lf.update_current_trace(
- metadata={
- "search_queries": search_queries,
- "observations_retrieved": unique_observations,
- }
- )
- else:
- all_results = [await _execute_single_query(query, embedding_store)]
- unique_observations = _deduplicate_observations(all_results)
-
- # Format observations
- if not unique_observations:
- logger.info("No unique historical observations found after filtering")
- return "No additional relevant context found."
-
- return _format_observations(unique_observations, include_premises=include_premises)
-
-
-async def _execute_single_query(
- query: str,
- embedding_store: EmbeddingStore,
-) -> list[tuple[str, str, dict[str, Any]]]:
- """
- Execute a single semantic search query and return formatted results.
-
- Args:
- query: The query to search for
- embedding_store: The embedding store to use.
-
- Returns:
- A list of tuples containing the content, timestamp, and metadata of the retrieved observations.
- """
- documents: list[Document] = await embedding_store.get_relevant_observations(
- query,
- top_k=settings.DIALECTIC.SEMANTIC_SEARCH_TOP_K,
- max_distance=settings.DIALECTIC.SEMANTIC_SEARCH_MAX_DISTANCE,
- for_reasoning=False,
- )
-
- # Extract data to avoid DetachedInstanceError
- return [
- (
- doc.content,
- doc.created_at.strftime("%Y-%m-%d-%H:%M:%S"),
- doc.internal_metadata or {},
- )
- for doc in documents
- ]
-
-
-def _deduplicate_observations(
- all_results: list[list[tuple[str, str, dict[str, Any]]]],
-) -> list[tuple[str, str, dict[str, Any]]]:
- """Deduplicate observations based on content."""
- unique_observations: list[tuple[str, str, dict[str, Any]]] = []
- seen_content: set[str] = set()
-
- for results in all_results:
- for content, timestamp, metadata in results:
- if content not in seen_content:
- unique_observations.append((content, timestamp, metadata))
- seen_content.add(content)
-
- return unique_observations
-
-
-def _format_observations(
- observations: list[tuple[str, str, dict[str, Any]]], *, include_premises: bool
-) -> str:
- """Format observations grouped by level and date, including access metadata."""
- grouped: dict[str, dict[str, list[str]]] = {}
-
- for content, timestamp, metadata in observations:
- level: str = metadata.get("level", "unknown")
- date_str: str = timestamp[:10] # Extract YYYY-MM-DD
-
- if level not in grouped:
- grouped[level] = {}
- if date_str not in grouped[level]:
- grouped[level][date_str] = []
-
- # Build formatted content with premises and access metadata
- formatted_content: str = content
-
- # Add premises if requested and available
- if include_premises and metadata.get("premises"):
- premises_text: str = format_premises_for_display(metadata["premises"])
- formatted_content = f"{content}{premises_text}"
-
- # Prefix with full timestamp for clarity
- if timestamp:
- formatted_content = f"{timestamp}: {formatted_content}"
-
- # Add access metadata if available
- access_parts: list[str] = []
- access_count: int = metadata.get("access_count", 0)
- last_accessed: Any = metadata.get("last_accessed")
-
- if access_count > 0:
- access_parts.append(f"accessed {access_count}x")
-
- if last_accessed:
- # Format the last_accessed datetime for display
- try:
- if isinstance(last_accessed, str):
- # Parse ISO format datetime string
- dt = parse_datetime_iso(last_accessed)
- formatted_last_accessed: str = dt.strftime("%Y-%m-%d %H:%M")
- access_parts.append(f"last accessed {formatted_last_accessed}")
- except (ValueError, AttributeError):
- # If parsing fails, just show the raw value
- access_parts.append(f"last accessed {last_accessed}")
-
- # Append access metadata to the formatted content
- if access_parts:
- access_info: str = ", ".join(access_parts)
- formatted_content = f"{formatted_content} [{access_info}]"
-
- grouped[level][date_str].append(formatted_content)
-
- # Build output
- parts: list[str] = []
- for level in sorted(grouped.keys()):
- header: str = (
- f"\n{level.upper()} OBSERVATIONS:"
- if level != "unknown"
- else "\nOBSERVATIONS:"
- )
- parts.append(header)
-
- for date_str in sorted(
- grouped[level].keys(), reverse=True
- ): # Most recent first
- parts.append(f"\n{date_str}:")
- for obs in grouped[level][date_str]:
- parts.append(f" • {obs}")
-
- return "\n".join(parts).strip()
-
-
-async def generate_semantic_queries(
- query: str, target_peer_name: str
-) -> SemanticQueries:
- """Generate semantic search queries for observation retrieval."""
- prompt = query_generation_prompt(query, target_peer_name)
- response = await honcho_llm_call(
- provider=settings.DIALECTIC.QUERY_GENERATION_PROVIDER,
- model=settings.DIALECTIC.QUERY_GENERATION_MODEL,
- prompt=prompt,
- max_tokens=settings.LLM.DEFAULT_MAX_TOKENS,
- response_model=SemanticQueries,
- enable_retry=True,
- retry_attempts=3,
- )
- return response.content
diff --git a/src/dreamer/__init__.py b/src/dreamer/__init__.py
new file mode 100644
index 00000000..bd9bfee7
--- /dev/null
+++ b/src/dreamer/__init__.py
@@ -0,0 +1,13 @@
+from .dream_scheduler import (
+ check_and_schedule_dream,
+ get_affected_dream_keys,
+ get_dream_scheduler,
+)
+from .dreamer import process_dream
+
+__all__ = [
+ "get_affected_dream_keys",
+ "get_dream_scheduler",
+ "check_and_schedule_dream",
+ "process_dream",
+]
diff --git a/src/dreamer/dream_scheduler.py b/src/dreamer/dream_scheduler.py
new file mode 100644
index 00000000..8c572e1c
--- /dev/null
+++ b/src/dreamer/dream_scheduler.py
@@ -0,0 +1,354 @@
+import asyncio
+from datetime import datetime, timezone
+from logging import getLogger
+from typing import Any
+
+import sentry_sdk
+from sqlalchemy import func, insert, select, update
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from src import models
+from src.config import settings
+from src.dependencies import tracked_db
+from src.utils.queue_payload import create_dream_payload
+from src.utils.work_unit import get_work_unit_key, parse_work_unit_key
+
+logger = getLogger(__name__)
+
+
+_dream_scheduler: "DreamScheduler | None" = None
+
+
+def set_dream_scheduler(dream_scheduler: "DreamScheduler") -> None:
+ """Set the global dream scheduler reference."""
+ global _dream_scheduler
+ _dream_scheduler = dream_scheduler
+
+
+def get_dream_scheduler() -> "DreamScheduler | None":
+ """Get the global dream scheduler reference."""
+ return _dream_scheduler
+
+
+def get_affected_dream_keys(message: dict[str, Any]) -> list[str]:
+ """
+ Get all work unit keys for dreams that might be affected by this message.
+
+ Args:
+ message: The message payload
+
+ Returns:
+ List of work unit keys that should have their dreams cancelled
+ """
+ workspace_name = message.get("workspace_name")
+ peer_name = message.get("peer_name")
+
+ if not workspace_name or not peer_name:
+ return []
+
+ # Generate dream work unit key for this peer's collection
+ dream_key = get_work_unit_key(
+ {
+ "task_type": "dream",
+ "workspace_name": workspace_name,
+ "observer": peer_name,
+ "observed": peer_name,
+ }
+ )
+
+ return [dream_key]
+
+
+class DreamScheduler:
+ _instance: "DreamScheduler | None" = None
+ _initialized: bool = False
+
+ def __new__(cls):
+ if cls._instance is None:
+ cls._instance = super().__new__(cls)
+ return cls._instance
+
+ def __init__(self):
+ # Only initialize once
+ if not DreamScheduler._initialized:
+ self.pending_dreams: dict[str, asyncio.Task[None]] = {}
+ DreamScheduler._initialized = True
+
+ @classmethod
+ def reset_singleton(cls) -> None:
+ """Reset the singleton instance. Only use this in tests."""
+ cls._instance = None
+ cls._initialized = False
+
+ def schedule_dream(
+ self,
+ work_unit_key: str,
+ workspace_name: str,
+ document_count: int,
+ delay_minutes: int,
+ *,
+ observer: str,
+ observed: str,
+ ) -> None:
+ """Schedule a dream for a collection after a delay."""
+ if not settings.DREAM.ENABLED:
+ return
+
+ # Cancel any existing dream for this collection
+ self.cancel_dream(work_unit_key)
+
+ task = asyncio.create_task(
+ self._delayed_dream(
+ work_unit_key,
+ workspace_name,
+ document_count,
+ delay_minutes,
+ observer=observer,
+ observed=observed,
+ )
+ )
+ self.pending_dreams[work_unit_key] = task
+ task.add_done_callback(lambda t: self.pending_dreams.pop(work_unit_key, None))
+
+ def cancel_dream(self, work_unit_key: str) -> bool:
+ """Cancel a pending dream. Returns True if a dream was cancelled."""
+ if work_unit_key in self.pending_dreams:
+ task = self.pending_dreams.pop(work_unit_key)
+ task.cancel()
+ logger.debug(f"Cancelled pending dream for {work_unit_key}")
+ return True
+ return False
+
+ async def _delayed_dream(
+ self,
+ work_unit_key: str,
+ workspace_name: str,
+ document_count: int,
+ delay_minutes: int,
+ *,
+ observer: str,
+ observed: str,
+ ) -> None:
+ try:
+ await asyncio.sleep(delay_minutes * 60)
+
+ # Check if collection is still inactive before executing dream
+ if await self._should_execute_dream(
+ workspace_name, observer=observer, observed=observed
+ ):
+ await self._execute_dream(
+ work_unit_key,
+ workspace_name,
+ document_count,
+ observer=observer,
+ observed=observed,
+ )
+ logger.info(f"Executed dream for {work_unit_key}")
+ else:
+ logger.info(
+ f"Skipping dream for {work_unit_key} - collection is active"
+ )
+
+ except asyncio.CancelledError:
+ logger.info(f"Dream task cancelled for {work_unit_key}")
+ except Exception as e:
+ logger.error(f"Error in delayed dream for {work_unit_key}: {str(e)}")
+ if settings.SENTRY.ENABLED:
+ sentry_sdk.capture_exception(e)
+
+ async def _should_execute_dream(
+ self, workspace_name: str, *, observer: str, observed: str
+ ) -> bool:
+ """Check if the collection is inactive and should be dreamed upon."""
+ async with tracked_db("dream_activity_check") as db:
+ # Check for active queue sessions related to this collection
+ query = select(models.ActiveQueueSession)
+ result = await db.execute(query)
+ active_sessions = result.scalars().all()
+
+ # Look for any active work units that match this collection
+ for active_session in active_sessions:
+ parsed_key = parse_work_unit_key(active_session.work_unit_key)
+ if (
+ parsed_key.workspace_name == workspace_name
+ and parsed_key.observer == observer
+ and parsed_key.observed == observed
+ ):
+ logger.debug("Collection is active, skipping dream")
+ return False
+
+ return True
+
+ async def _execute_dream(
+ self,
+ work_unit_key: str,
+ workspace_name: str,
+ document_count: int,
+ *,
+ observer: str,
+ observed: str,
+ ) -> None:
+ """Execute the dream by enqueueing it and updating collection metadata."""
+ dream_payload = create_dream_payload(
+ workspace_name=workspace_name,
+ dream_type="consolidate",
+ observer=observer,
+ observed=observed,
+ )
+
+ async with tracked_db("dream_execute") as db:
+ dream_record = {
+ "work_unit_key": work_unit_key,
+ "payload": dream_payload,
+ "session_id": None,
+ "task_type": "dream",
+ }
+
+ await db.execute(insert(models.QueueItem), [dream_record])
+
+ now_iso = datetime.now(timezone.utc).isoformat()
+ stmt = (
+ update(models.Collection)
+ .where(
+ models.Collection.workspace_name == workspace_name,
+ models.Collection.observer == observer,
+ models.Collection.observed == observed,
+ )
+ .values(
+ internal_metadata=models.Collection.internal_metadata.op("||")(
+ {
+ "dream": {
+ "last_dream_document_count": document_count,
+ "last_dream_at": now_iso,
+ }
+ }
+ )
+ )
+ )
+ await db.execute(stmt)
+ await db.commit()
+
+ logger.info(
+ "Enqueued dream task",
+ extra={
+ "workspace_name": workspace_name,
+ "observer": observer,
+ "observed": observed,
+ },
+ )
+
+ async def shutdown(self) -> None:
+ """Cancel all pending dreams during shutdown."""
+ if self.pending_dreams:
+ logger.info(f"Cancelling {len(self.pending_dreams)} pending dreams...")
+ for task in self.pending_dreams.values():
+ task.cancel()
+ await asyncio.gather(*self.pending_dreams.values(), return_exceptions=True)
+ self.pending_dreams.clear()
+
+
+async def check_and_schedule_dream(
+ db: AsyncSession,
+ collection: models.Collection,
+) -> bool:
+ """
+ Check if a collection has reached the document threshold and schedule a timer-based dream.
+
+ This function only schedules a timer-based dream if:
+ 1. Dreams are enabled
+ 2. Document threshold is reached
+ 3. Minimum hours between dreams have passed
+ 4. No dream is already scheduled for this collection
+
+ Args:
+ db: Database session
+ collection: Collection model to check
+
+ Returns:
+ True if a dream timer was scheduled, False otherwise
+ """
+ if not settings.DREAM.ENABLED:
+ return False
+
+ # Get dream metadata from internal_metadata
+ dream_metadata = collection.internal_metadata.get("dream", {})
+ last_dream_document_count = dream_metadata.get("last_dream_document_count", 0)
+ last_dream_at = dream_metadata.get("last_dream_at")
+
+ # Count current documents in the collection
+ count_stmt = select(func.count(models.Document.id)).where(
+ models.Document.workspace_name == collection.workspace_name,
+ models.Document.observer == collection.observer,
+ models.Document.observed == collection.observed,
+ )
+ current_document_count = int(await db.scalar(count_stmt) or 0)
+
+ # Calculate documents added since last dream
+ documents_since_last_dream = current_document_count - last_dream_document_count
+
+ logger.info(
+ "Dream check",
+ extra={
+ "workspace_name": collection.workspace_name,
+ "observer": collection.observer,
+ "observed": collection.observed,
+ "current_document_count": current_document_count,
+ "last_dream_document_count": last_dream_document_count,
+ "documents_since_last_dream": documents_since_last_dream,
+ "document_threshold": settings.DREAM.DOCUMENT_THRESHOLD,
+ },
+ )
+
+ # Only schedule timer if document threshold is reached
+ if documents_since_last_dream >= settings.DREAM.DOCUMENT_THRESHOLD:
+ # Check if we're within minimum hours between dreams
+ if last_dream_at:
+ try:
+ last_dream_time = datetime.fromisoformat(last_dream_at)
+ hours_since_last_dream = (
+ datetime.now(timezone.utc) - last_dream_time
+ ).total_seconds() / 3600
+
+ if hours_since_last_dream < settings.DREAM.MIN_HOURS_BETWEEN_DREAMS:
+ logger.info(
+ f"Skipping dream for {collection.observer}/{collection.observed}: only {hours_since_last_dream:.1f} hours "
+ + f"since last dream (minimum: {settings.DREAM.MIN_HOURS_BETWEEN_DREAMS})"
+ )
+ return False
+ except (ValueError, TypeError) as e:
+ logger.warning(
+ f"Invalid last_dream_at timestamp: {last_dream_at}, error: {e}"
+ )
+
+ dream_scheduler = get_dream_scheduler()
+ if dream_scheduler:
+ collection_work_unit_key = get_work_unit_key(
+ {
+ "task_type": "dream",
+ "workspace_name": collection.workspace_name,
+ "observer": collection.observer,
+ "observed": collection.observed,
+ }
+ )
+
+ dream_scheduler.schedule_dream(
+ collection_work_unit_key,
+ collection.workspace_name,
+ current_document_count,
+ settings.DREAM.IDLE_TIMEOUT_MINUTES,
+ observer=collection.observer,
+ observed=collection.observed,
+ )
+ logger.info(
+ "Scheduled dream",
+ extra={
+ "workspace_name": collection.workspace_name,
+ "observer": collection.observer,
+ "observed": collection.observed,
+ "documents_since_last_dream": documents_since_last_dream,
+ "document_threshold": settings.DREAM.DOCUMENT_THRESHOLD,
+ },
+ )
+ return True
+
+ return False
diff --git a/src/dreamer/dreamer.py b/src/dreamer/dreamer.py
new file mode 100644
index 00000000..bac32db9
--- /dev/null
+++ b/src/dreamer/dreamer.py
@@ -0,0 +1,195 @@
+import logging
+from collections.abc import Sequence
+
+import sentry_sdk
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from src import crud, models, schemas
+from src.config import settings
+from src.dependencies import tracked_db
+from src.dreamer.prompts import consolidation_prompt
+from src.embedding_client import embedding_client
+from src.utils.clients import honcho_llm_call
+from src.utils.formatting import format_datetime_utc
+from src.utils.queue_payload import DreamPayload
+from src.utils.representation import (
+ ExplicitObservation,
+ Representation,
+)
+
+logger = logging.getLogger(__name__)
+
+
+@sentry_sdk.trace
+async def process_dream(
+ payload: DreamPayload,
+) -> None:
+ """
+ Process a dream task by performing collection maintenance operations.
+
+ Args:
+ payload: The dream task payload containing workspace, peer, and dream type information
+ """
+ logger.info(
+ f"Processing dream task: {payload.dream_type} for {payload.workspace_name}/{payload.observer}/{payload.observed}"
+ )
+
+ try:
+ if payload.dream_type == "consolidate":
+ await _process_consolidate_dream(payload)
+ ## TODO other dream types
+
+ except Exception as e:
+ logger.error(
+ f"Error processing dream task {payload.dream_type} for {payload.observer}/{payload.observed}: {str(e)}",
+ exc_info=True,
+ )
+ if settings.SENTRY.ENABLED:
+ sentry_sdk.capture_exception(e)
+ # Don't re-raise - we want to mark the dream task as processed even if it fails
+
+
+async def _process_consolidate_dream(payload: DreamPayload) -> None:
+ """
+ Process a consolidation dream task.
+
+ Consolidation means taking all the documents in a collection and merging
+ similar observations into a single, best-quality observation document.
+
+ TODO: need to determine a way to do this on a subset of documents since
+ collections will grow very large.
+ """
+ logger.info(
+ f"""
+(っ- ‸ - ς)ᶻ z 𐰁 ᶻ z 𐰁 ᶻ z 𐰁\n
+DREAM: consolidating documents for {payload.workspace_name}/{payload.observer}/{payload.observed}\n
+𐰁 z ᶻ 𐰁 z ᶻ 𐰁 z ᶻ(っ- ‸ - ς)"""
+ )
+
+ # get all documents in the collection
+ async with tracked_db("dream_consolidate") as db:
+ documents = await crud.get_all_documents(
+ db,
+ payload.workspace_name,
+ observer=payload.observer,
+ observed=payload.observed,
+ )
+
+ logger.info("found %d documents to consolidate", len(documents))
+
+ # TODO: create clusters of documents based on cosine similarity
+ # clusters = await create_document_clusters(documents)
+
+ # logger.info("created %d clusters", len(clusters))
+ clusters = [documents]
+
+ # for each cluster, call llm to consolidate the representation if possible
+ for cluster in clusters:
+ await _consolidate_cluster(
+ cluster,
+ payload.workspace_name,
+ db,
+ observer=payload.observer,
+ observed=payload.observed,
+ )
+
+
+async def _consolidate_cluster(
+ cluster: Sequence[models.Document],
+ workspace_name: str,
+ db: AsyncSession,
+ *,
+ observer: str,
+ observed: str,
+) -> None:
+ """
+ Consolidate a cluster of documents, treated as a Representation, into a smaller one.
+ Removes old documents and replaces them with consolidated versions while preserving metadata.
+ """
+ if len(cluster) <= 1:
+ logger.info("Cluster has %d documents, skipping consolidation", len(cluster))
+ return
+
+ cluster_representation = Representation.from_documents(cluster)
+ logger.info("unconsolidated representation:\n%s", cluster_representation)
+
+ consolidated_representation = await consolidate_call(cluster_representation)
+ logger.info("consolidated representation:\n%s", consolidated_representation)
+
+ # TODO: less hacky preservation of times_derived
+ total_times_derived = sum(
+ doc.internal_metadata.get("times_derived", 1) for doc in cluster
+ )
+
+ new_documents = [
+ *consolidated_representation.explicit,
+ *consolidated_representation.deductive,
+ ]
+
+ documents_to_create: list[schemas.DocumentCreate] = []
+
+ for obs in new_documents:
+ if isinstance(obs, ExplicitObservation):
+ content = obs.content
+ level = "explicit"
+ premises = None
+ else:
+ content = obs.conclusion
+ level = "deductive"
+ premises = obs.premises
+ # NOTE: other kinds of observations here in the future
+
+ metadata = schemas.DocumentMetadata(
+ times_derived=total_times_derived,
+ message_ids=obs.message_ids,
+ message_created_at=format_datetime_utc(obs.created_at),
+ level=level,
+ premises=premises,
+ )
+
+ embedding = await embedding_client.embed(content)
+
+ documents_to_create.append(
+ schemas.DocumentCreate(
+ content=content,
+ session_name=obs.session_name,
+ metadata=metadata,
+ embedding=embedding,
+ )
+ )
+
+ # bulk create documents
+ await crud.create_documents(
+ db, documents_to_create, workspace_name, observer=observer, observed=observed
+ )
+
+ # delete old documents
+ for doc in cluster:
+ await db.delete(doc)
+
+ await db.commit()
+
+ logger.info(
+ "consolidated %d documents into %d new documents",
+ len(cluster),
+ len(new_documents),
+ )
+
+
+async def consolidate_call(
+ representation: Representation,
+) -> Representation:
+ prompt = consolidation_prompt(representation)
+
+ response = await honcho_llm_call(
+ provider=settings.DREAM.PROVIDER,
+ model=settings.DREAM.MODEL,
+ prompt=prompt,
+ max_tokens=settings.DREAM.MAX_OUTPUT_TOKENS,
+ track_name="Dream Call",
+ response_model=Representation,
+ enable_retry=True,
+ retry_attempts=3,
+ )
+
+ return response.content
diff --git a/src/dreamer/prompts.py b/src/dreamer/prompts.py
new file mode 100644
index 00000000..f6ed1713
--- /dev/null
+++ b/src/dreamer/prompts.py
@@ -0,0 +1,26 @@
+from inspect import cleandoc as c
+
+from src.utils.representation import Representation
+
+
+def consolidation_prompt(
+ representation: Representation,
+) -> str:
+ """
+ Generate the prompt for user representation consolidation.
+
+ Args:
+ representation: The user representation to consolidate
+
+ Returns:
+ A consolidated user representation
+ """
+ representation_as_json = representation.model_dump_json(indent=2)
+
+ return c(
+ f"""
+You are an agent that consolidates observations about an entity. You will be presented with a list of EXPLICIT and DEDUCTIVE observations. **Reduce** the number of observations, if possible, by combining similar observations. **ONLY** include information that is **GIVEN**. Create the highest-quality observations with the given information. Observations must always be maximally concise.
+
+{representation_as_json}
+"""
+ )
diff --git a/src/embedding_client.py b/src/embedding_client.py
index bdb58140..6e0a1195 100644
--- a/src/embedding_client.py
+++ b/src/embedding_client.py
@@ -1,9 +1,11 @@
import asyncio
import logging
+import threading
from collections import defaultdict
from typing import NamedTuple
import tiktoken
+from google import genai
from openai import AsyncOpenAI
from .config import settings
@@ -19,19 +21,36 @@ class BatchItem(NamedTuple):
chunk_index: int
-class EmbeddingClient:
+class _EmbeddingClient:
"""
- Embedding client for OpenAI with chunking and batching support.
+ Embedding client supporting OpenAI and Gemini with chunking and batching support.
"""
- def __init__(self, api_key: str | None = None):
- if api_key is None:
- api_key = settings.LLM.OPENAI_API_KEY
- if not api_key:
- raise ValueError("API key is required")
- self.client: AsyncOpenAI = AsyncOpenAI(api_key=api_key)
+ def __init__(self, api_key: str | None = None, provider: str | None = None):
+ self.provider: str = provider or settings.LLM.EMBEDDING_PROVIDER
+
+ if self.provider == "gemini":
+ if api_key is None:
+ api_key = settings.LLM.GEMINI_API_KEY
+ if not api_key:
+ raise ValueError("Gemini API key is required")
+ self.client: genai.Client | AsyncOpenAI = genai.Client(api_key=api_key)
+ self.model: str = "gemini-embedding-001"
+ # Gemini has a 2048 token limit
+ self.max_embedding_tokens: int = min(settings.MAX_EMBEDDING_TOKENS, 2048)
+ # Gemini batch size is not documented, using conservative estimate
+ self.max_batch_size: int = 100
+ else: # openai
+ if api_key is None:
+ api_key = settings.LLM.OPENAI_API_KEY
+ if not api_key:
+ raise ValueError("OpenAI API key is required")
+ self.client = AsyncOpenAI(api_key=api_key)
+ self.model = "text-embedding-3-small"
+ self.max_embedding_tokens = settings.MAX_EMBEDDING_TOKENS
+ self.max_batch_size = 2048 # OpenAI batch limit
+
self.encoding: tiktoken.Encoding = tiktoken.get_encoding("cl100k_base")
- self.max_embedding_tokens: int = settings.MAX_EMBEDDING_TOKENS
self.max_embedding_tokens_per_request: int = (
settings.MAX_EMBEDDING_TOKENS_PER_REQUEST
)
@@ -44,10 +63,65 @@ class EmbeddingClient:
f"Query exceeds maximum token limit of {self.max_embedding_tokens} tokens (got {token_count} tokens)"
)
- response = await self.client.embeddings.create(
- model="text-embedding-3-small", input=query
- )
- return response.data[0].embedding
+ if isinstance(self.client, genai.Client):
+ response = await self.client.aio.models.embed_content(
+ model=self.model,
+ contents=query,
+ config={"output_dimensionality": 1536},
+ )
+ if not response.embeddings or not response.embeddings[0].values:
+ raise ValueError("No embedding returned from Gemini API")
+ return response.embeddings[0].values
+ else: # openai
+ response = await self.client.embeddings.create(
+ model=self.model, input=query
+ )
+ return response.data[0].embedding
+
+ async def simple_batch_embed(self, texts: list[str]) -> list[list[float]]:
+ """
+ Simple batch embedding for a list of text strings.
+
+ Args:
+ texts: List of text strings to embed
+
+ Returns:
+ List of embedding vectors corresponding to input texts
+
+ Raises:
+ ValueError: If any text exceeds token limits
+ """
+ embeddings: list[list[float]] = []
+
+ for i in range(0, len(texts), self.max_batch_size):
+ batch = texts[i : i + self.max_batch_size]
+ try:
+ if isinstance(self.client, genai.Client):
+ # Type cast needed due to genai type signature complexity
+ response = await self.client.aio.models.embed_content(
+ model=self.model,
+ contents=batch, # pyright: ignore[reportArgumentType]
+ config={"output_dimensionality": 1536},
+ )
+ if response.embeddings:
+ for emb in response.embeddings:
+ if emb.values:
+ embeddings.append(emb.values)
+ else: # openai
+ response = await self.client.embeddings.create(
+ input=batch,
+ model=self.model,
+ )
+ embeddings.extend([data.embedding for data in response.data])
+ except Exception as e:
+ # Check if it's a token limit error and re-raise as ValueError for consistency
+ if "token" in str(e).lower():
+ raise ValueError(
+ f"Text content exceeds maximum token limit of {self.max_embedding_tokens}."
+ ) from e
+ raise
+
+ return embeddings
async def batch_embed(
self, id_resource_dict: dict[str, tuple[str, list[int]]]
@@ -124,7 +198,7 @@ class EmbeddingClient:
current_tokens + chunk_tokens
> self.max_embedding_tokens_per_request
)
- would_exceed_count = len(current_batch) >= 2048 # OpenAI's input limit
+ would_exceed_count = len(current_batch) >= self.max_batch_size
if current_batch and (would_exceed_tokens or would_exceed_count):
batches.append(current_batch)
@@ -152,14 +226,25 @@ class EmbeddingClient:
Maps text IDs to {chunk_index: embedding_vector} dictionaries
"""
try:
- response = await self.client.embeddings.create(
- model="text-embedding-3-small", input=[item.text for item in batch]
- )
-
# Organize embeddings by text_id and chunk_index
result: dict[str, dict[int, list[float]]] = defaultdict(dict)
- for item, embedding_data in zip(batch, response.data, strict=True):
- result[item.text_id][item.chunk_index] = embedding_data.embedding
+
+ if isinstance(self.client, genai.Client):
+ response = await self.client.aio.models.embed_content(
+ model=self.model,
+ contents=[item.text for item in batch],
+ config={"output_dimensionality": 1536},
+ )
+ if response.embeddings:
+ for item, embedding in zip(batch, response.embeddings, strict=True):
+ if embedding.values:
+ result[item.text_id][item.chunk_index] = embedding.values
+ else: # openai
+ response = await self.client.embeddings.create(
+ model=self.model, input=[item.text for item in batch]
+ )
+ for item, embedding_data in zip(batch, response.data, strict=True):
+ result[item.text_id][item.chunk_index] = embedding_data.embedding
return dict(result)
@@ -228,5 +313,83 @@ def _chunk_text_with_tokens(
]
-# Shared embedding client instance
-embedding_client = EmbeddingClient(settings.LLM.OPENAI_API_KEY)
+class EmbeddingClient:
+ """
+ Singleton wrapper for the embedding client with deferred loading.
+
+ The actual client is only initialized on first use, improving startup time
+ and allowing the application to start even if API keys are not yet configured.
+ """
+
+ _instance: "_EmbeddingClient | None" = None
+ _lock: threading.Lock = threading.Lock()
+ _wrapper_instance: "EmbeddingClient | None" = None
+
+ def __new__(cls):
+ """Ensure only one instance of EmbeddingClient exists."""
+ # We always return the same wrapper instance
+ if cls._wrapper_instance is None:
+ cls._wrapper_instance = super().__new__(cls)
+ return cls._wrapper_instance
+
+ def _get_client(self) -> _EmbeddingClient:
+ """
+ Get or create the underlying embedding client instance.
+
+ Uses double-checked locking for thread-safe lazy initialization.
+ """
+ if self._instance is None:
+ with self._lock:
+ if self._instance is None:
+ provider = settings.LLM.EMBEDDING_PROVIDER
+ if provider == "gemini":
+ api_key = settings.LLM.GEMINI_API_KEY
+ else:
+ api_key = settings.LLM.OPENAI_API_KEY
+
+ self._instance = _EmbeddingClient(
+ api_key=api_key, provider=provider
+ )
+ logger.info(
+ f"Initialized embedding client with provider: {provider}"
+ )
+
+ return self._instance
+
+ async def embed(self, query: str) -> list[float]:
+ """Embed a single query string."""
+ return await self._get_client().embed(query)
+
+ async def simple_batch_embed(self, texts: list[str]) -> list[list[float]]:
+ """Simple batch embedding for a list of text strings."""
+ return await self._get_client().simple_batch_embed(texts)
+
+ async def batch_embed(
+ self, id_resource_dict: dict[str, tuple[str, list[int]]]
+ ) -> dict[str, list[list[float]]]:
+ """Embed multiple texts, chunking long ones and batching API calls."""
+ return await self._get_client().batch_embed(id_resource_dict)
+
+ @property
+ def provider(self) -> str:
+ """Get the provider name."""
+ return self._get_client().provider
+
+ @property
+ def model(self) -> str:
+ """Get the model name."""
+ return self._get_client().model
+
+ @property
+ def max_embedding_tokens(self) -> int:
+ """Get the maximum embedding tokens."""
+ return self._get_client().max_embedding_tokens
+
+ @property
+ def encoding(self) -> tiktoken.Encoding:
+ """Get the tiktoken encoding."""
+ return self._get_client().encoding
+
+
+# Shared singleton embedding client instance
+embedding_client = EmbeddingClient()
diff --git a/src/models.py b/src/models.py
index defe2110..a50df97c 100644
--- a/src/models.py
+++ b/src/models.py
@@ -121,7 +121,6 @@ class Peer(Base):
sessions = relationship(
"Session", secondary=session_peers_table, back_populates="peers"
)
- collections = relationship("Collection", back_populates="peer")
__table_args__ = (
UniqueConstraint("name", "workspace_name", name="unique_name_workspace_peer"),
@@ -280,7 +279,8 @@ class Collection(Base):
__tablename__: str = "collections"
id: Mapped[str] = mapped_column(TEXT, default=generate_nanoid, primary_key=True)
- name: Mapped[str] = mapped_column(TEXT, index=True)
+ observer: Mapped[str] = mapped_column(TEXT, index=True)
+ observed: Mapped[str] = mapped_column(TEXT, index=True)
created_at: Mapped[datetime.datetime] = mapped_column(
DateTime(timezone=True), index=True, default=func.now()
)
@@ -291,22 +291,27 @@ class Collection(Base):
documents = relationship(
"Document", back_populates="collection", cascade="all, delete, delete-orphan"
)
- peer = relationship("Peer", back_populates="collections")
- peer_name: Mapped[str] = mapped_column(TEXT, index=True)
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"), index=True
)
__table_args__ = (
UniqueConstraint(
- "name", "peer_name", "workspace_name", name="unique_name_collection_peer"
+ "observer",
+ "observed",
+ "workspace_name",
+ name="unique_observer_observed_collection",
),
CheckConstraint("length(id) = 21", name="id_length"),
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
- CheckConstraint("length(name) <= 1025", name="name_length"),
- # Composite foreign key constraint for peers
+ # Composite foreign key constraint for observer peer
ForeignKeyConstraint(
- ["peer_name", "workspace_name"],
+ ["observer", "workspace_name"],
+ ["peers.name", "peers.workspace_name"],
+ ),
+ # Composite foreign key constraint for observed peer
+ ForeignKeyConstraint(
+ ["observed", "workspace_name"],
["peers.name", "peers.workspace_name"],
),
)
@@ -325,11 +330,12 @@ class Document(Base):
DateTime(timezone=True), index=True, default=func.now()
)
- collection_name: Mapped[str] = mapped_column(TEXT, index=True)
- peer_name: Mapped[str] = mapped_column(index=True)
+ observer: Mapped[str] = mapped_column(TEXT, index=True)
+ observed: Mapped[str] = mapped_column(TEXT, index=True)
workspace_name: Mapped[str] = mapped_column(
ForeignKey("workspaces.name"), index=True
)
+ session_name: Mapped[str] = mapped_column(TEXT, index=True)
collection = relationship("Collection", back_populates="documents")
__table_args__ = (
@@ -338,14 +344,28 @@ class Document(Base):
CheckConstraint("id ~ '^[A-Za-z0-9_-]+$'", name="id_format"),
# Composite foreign key constraint for collections
ForeignKeyConstraint(
- ["collection_name", "peer_name", "workspace_name"],
- ["collections.name", "collections.peer_name", "collections.workspace_name"],
+ ["observer", "observed", "workspace_name"],
+ [
+ "collections.observer",
+ "collections.observed",
+ "collections.workspace_name",
+ ],
),
- # Composite foreign key constraint for peers
+ # Composite foreign key constraint for observer peer
ForeignKeyConstraint(
- ["peer_name", "workspace_name"],
+ ["observer", "workspace_name"],
["peers.name", "peers.workspace_name"],
),
+ # Composite foreign key constraint for observed peer
+ ForeignKeyConstraint(
+ ["observed", "workspace_name"],
+ ["peers.name", "peers.workspace_name"],
+ ),
+ # Composite foreign key constraint for sessions
+ ForeignKeyConstraint(
+ ["session_name", "workspace_name"],
+ ["sessions.name", "sessions.workspace_name"],
+ ),
# HNSW index on embedding column
Index(
"idx_documents_embedding_hnsw",
@@ -359,7 +379,7 @@ class Document(Base):
)
-TaskType = Literal["webhook", "summary", "representation"]
+TaskType = Literal["webhook", "summary", "representation", "dream"]
@final
diff --git a/src/routers/peers.py b/src/routers/peers.py
index 423aeb3d..a6ecb10e 100644
--- a/src/routers/peers.py
+++ b/src/routers/peers.py
@@ -170,11 +170,13 @@ async def chat(
if not options.stream:
response = await dialectic_chat(
workspace_name=workspace_id,
- peer_name=peer_id,
- target_name=options.target,
session_name=options.session_id,
query=options.query,
stream=options.stream,
+ observer=peer_id,
+ # if target is given, that's the observed peer. otherwise, observer==observed
+ # and it's answered from the omniscient Honcho perspective
+ observed=options.target if options.target is not None else peer_id,
)
return schemas.DialecticResponse(content=str(response))
@@ -182,11 +184,11 @@ async def chat(
try:
stream = await dialectic_chat(
workspace_name=workspace_id,
- peer_name=peer_id,
- target_name=options.target,
session_name=options.session_id,
query=options.query,
stream=options.stream,
+ observer=peer_id,
+ observed=options.target if options.target is not None else peer_id,
)
if isinstance(stream, AsyncIterator):
async for chunk in stream:
@@ -225,20 +227,20 @@ async def get_working_representation(
options: schemas.PeerRepresentationGet = Body(
..., description="Options for getting the peer representation"
),
- db: AsyncSession = db,
):
"""Get a peer's working representation for a session.
If a session_id is provided in the body, we get the working representation of the peer in that session.
If a target is provided, we get the representation of the target from the perspective of the peer.
- If no target is provided, we get the global representation of the peer.
+ If no target is provided, we get the omniscient Honcho representation of the peer.
"""
try:
- # If no target specified, get global representation (peer observing themselves)
- target_peer = options.target if options.target is not None else peer_id
-
+ # If no target specified, get global representation (omniscient Honcho perspective)
representation = await crud.get_working_representation(
- db, workspace_id, peer_id, target_peer, options.session_id
+ workspace_id,
+ observer=peer_id,
+ observed=options.target if options.target is not None else peer_id,
+ session_name=options.session_id,
)
return {"representation": representation}
except ValueError as e:
@@ -270,9 +272,11 @@ async def get_peer_card(
If no target is specified, returns the observer's own peer card.
"""
# If no target specified, get the observer's own card
- target_peer = target if target is not None else peer_id
+ observed = target if target is not None else peer_id
- peer_card = await crud.get_peer_card(db, workspace_id, target_peer, peer_id)
+ peer_card = await crud.get_peer_card(
+ db, workspace_id, observer=peer_id, observed=observed
+ )
return schemas.PeerCardResponse(peer_card=peer_card)
diff --git a/src/routers/sessions.py b/src/routers/sessions.py
index b86a8481..fafb310e 100644
--- a/src/routers/sessions.py
+++ b/src/routers/sessions.py
@@ -1,12 +1,14 @@
+import asyncio
import logging
+from typing import cast
from fastapi import APIRouter, Body, Depends, Path, Query, Response
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import apaginate
from sqlalchemy.ext.asyncio import AsyncSession
-from src import config, crud, schemas
-from src.dependencies import db
+from src import config, crud, models, schemas
+from src.dependencies import db, tracked_db
from src.exceptions import (
AuthenticationException,
ResourceNotFoundException,
@@ -14,7 +16,9 @@ from src.exceptions import (
)
from src.security import JWTParams, require_auth
from src.utils import summarizer
+from src.utils.representation import Representation
from src.utils.search import search
+from src.utils.tokens import estimate_tokens
logger = logging.getLogger(__name__)
@@ -24,6 +28,88 @@ router = APIRouter(
)
+async def _get_working_representation_task(
+ workspace_id: str,
+ last_message: str | None,
+ *,
+ observer: str,
+ observed: str,
+) -> Representation:
+ """
+ Atomic task to get working representation using tracked_db.
+
+ Args:
+ workspace_id: The workspace identifier
+ last_message: Optional last message for semantic query
+ observer: Name of the observer peer
+ observed: Name of the observed peer
+
+ Returns:
+ The working representation
+ """
+ return await crud.get_working_representation(
+ workspace_name=workspace_id,
+ include_semantic_query=last_message,
+ include_most_derived=True,
+ observer=observer,
+ observed=observed,
+ )
+
+
+async def _get_peer_card_task(
+ workspace_id: str,
+ *,
+ observer: str,
+ observed: str,
+) -> list[str] | None:
+ """
+ Atomic task to get peer card using tracked_db.
+
+ Args:
+ workspace_id: The workspace identifier
+ observer: Name of the observer peer
+ observed: Name of the observed peer
+
+ Returns:
+ The peer card or None if not found
+ """
+ async with tracked_db("get_peer_card") as db:
+ return await crud.get_peer_card(
+ db,
+ workspace_name=workspace_id,
+ observer=observer,
+ observed=observed,
+ )
+
+
+async def _get_session_context_task(
+ workspace_id: str,
+ session_id: str,
+ token_limit: int,
+ include_summary: bool,
+) -> tuple[schemas.Summary | None, list[models.Message]]:
+ """
+ Atomic task to get session context using tracked_db.
+
+ Args:
+ workspace_id: The workspace identifier
+ session_id: The session identifier
+ token_limit: Maximum tokens for the context
+ include_summary: Whether to include summary if available
+
+ Returns:
+ Tuple of (summary, messages)
+ """
+ async with tracked_db("get_session_context") as db:
+ return await summarizer.get_session_context(
+ db,
+ workspace_name=workspace_id,
+ session_name=session_id,
+ token_limit=token_limit,
+ include_summary=include_summary,
+ )
+
+
@router.post(
"",
response_model=schemas.Session,
@@ -372,15 +458,26 @@ async def get_session_context(
tokens: int | None = Query(
None,
le=config.settings.GET_CONTEXT_MAX_TOKENS,
- description=f"Number of tokens to use for the context. Includes summary if set to true. If not provided, the context will be exhaustive (within {config.settings.GET_CONTEXT_MAX_TOKENS} tokens)",
+ description=f"Number of tokens to use for the context. Includes summary if set to true. Includes representation and peer card if they are included in the response. If not provided, the context will be exhaustive (within {config.settings.GET_CONTEXT_MAX_TOKENS} tokens)",
),
*,
+ last_message: str | None = Query(
+ None,
+ description="The most recent message, used to fetch semantically relevant observations",
+ ),
include_summary: bool = Query(
default=True,
description="Whether or not to include a summary *if* one is available for the session",
alias="summary",
),
- db: AsyncSession = db,
+ peer_target: str | None = Query(
+ None,
+ description="The target of the perspective. If given without `peer_perspective`, will get the Honcho-level representation and peer card for this peer. If given with `peer_perspective`, will get the representation and card for this peer *from the perspective of that peer*.",
+ ),
+ peer_perspective: str | None = Query(
+ None,
+ description="A peer to get context for. If given, response will attempt to include representation and card from the perspective of that peer. Must be provided with `peer_target`.",
+ ),
):
"""
Produce a context object from the session. The caller provides an optional token limit which the entire context must fit into.
@@ -388,21 +485,65 @@ async def get_session_context(
to the summary, and 60% to recent messages -- as many as can fit. Note that the summary will usually take up less space than
this. If the caller does not want a summary, we allocate all the tokens to recent messages.
"""
- token_limit = tokens or config.settings.GET_CONTEXT_MAX_TOKENS
+ token_limit = (
+ tokens if tokens is not None else config.settings.GET_CONTEXT_MAX_TOKENS
+ )
- # Use the shared get_session_context function from summarizer
- summary_obj, messages = await summarizer.get_session_context(
- db,
- workspace_name=workspace_id,
- session_name=session_id,
- token_limit=token_limit,
- include_summary=include_summary,
+ if peer_perspective and not peer_target:
+ raise ValidationException(
+ "peer_target must be provided if peer_perspective is provided"
+ )
+
+ if not peer_target:
+ # No representation or card needed
+ summary, messages = await _get_session_context_task(
+ workspace_id, session_id, token_limit, include_summary
+ )
+ return schemas.SessionContext(
+ name=session_id,
+ messages=messages, # pyright: ignore -- db message type and schema message type are different, but excess gets removed by schema
+ summary=summary,
+ )
+
+ observer = peer_perspective or peer_target
+ observed = peer_target
+
+ # Run representation and card tasks in parallel
+ representation, card = await asyncio.gather(
+ _get_working_representation_task(
+ workspace_id, last_message, observer=observer, observed=observed
+ ),
+ _get_peer_card_task(workspace_id, observer=observer, observed=observed),
+ return_exceptions=True,
+ )
+
+ # Handle any exceptions from the parallel tasks
+ if isinstance(representation, Exception):
+ raise representation
+ if isinstance(card, Exception):
+ raise card
+
+ # At this point, we know the types are correct - cast to help type checker
+ representation = cast(Representation, representation)
+ card = cast(list[str] | None, card)
+
+ # adjust token limit downward to account for approximate token count of representation and card
+ # TODO determine if this impacts performance too much
+ adjusted_token_limit = (
+ token_limit - estimate_tokens(str(representation)) - estimate_tokens(card)
+ )
+
+ # Get the session context with the adjusted limit
+ summary, messages = await _get_session_context_task(
+ workspace_id, session_id, adjusted_token_limit, include_summary
)
return schemas.SessionContext(
name=session_id,
messages=messages, # pyright: ignore -- db message type and schema message type are different, but excess gets removed by schema
- summary=summary_obj,
+ summary=summary,
+ peer_representation=representation,
+ peer_card=card,
)
diff --git a/src/routers/workspaces.py b/src/routers/workspaces.py
index 40bebbb2..2c10b5d0 100644
--- a/src/routers/workspaces.py
+++ b/src/routers/workspaces.py
@@ -132,9 +132,9 @@ async def get_deriver_status(
return await crud.get_deriver_status(
db,
workspace_name=workspace_id,
- observer_name=observer_id,
- sender_name=sender_id,
session_name=session_id,
+ observer=observer_id,
+ observed=sender_id,
)
except ValueError as e:
logger.warning(f"Invalid request parameters: {str(e)}")
diff --git a/src/schemas.py b/src/schemas.py
index cf4f382c..a1eac978 100644
--- a/src/schemas.py
+++ b/src/schemas.py
@@ -1,7 +1,6 @@
-# pyright: reportUnannotatedClassAttribute=false # pyright: ignore
import datetime
import ipaddress
-from typing import Annotated, Any, Self
+from typing import Annotated, Any, Literal, Self
from urllib.parse import urlparse
import tiktoken
@@ -15,6 +14,7 @@ from pydantic import (
)
from src.config import settings
+from src.utils.representation import Representation
RESOURCE_NAME_PATTERN = r"^[a-zA-Z0-9_-]+$"
@@ -95,8 +95,8 @@ class Peer(PeerBase):
class PeerRepresentationGet(BaseModel):
- session_id: str = Field(
- ..., description="Get the working representation within this session"
+ session_id: str | None = Field(
+ None, description="Get the working representation within this session"
)
target: str | None = Field(
None,
@@ -255,6 +255,14 @@ class SessionContext(SessionBase):
summary: Summary | None = Field(
default=None, description="The summary if available"
)
+ peer_representation: Representation | None = Field(
+ default=None,
+ description="The peer representation, if context is requested from a specific perspective",
+ )
+ peer_card: list[str] | None = Field(
+ default=None,
+ description="The peer card, if context is requested from a specific perspective",
+ )
model_config = ConfigDict( # pyright: ignore
from_attributes=True, populate_by_name=True
@@ -279,14 +287,34 @@ class DocumentBase(BaseModel):
pass
+class DocumentMetadata(BaseModel):
+ times_derived: int | None = Field(
+ default=None,
+ ge=1,
+ description="The number of times that a semantic duplicate document to this one has been derived",
+ )
+ message_ids: list[tuple[int, int]] = Field(
+ description="The ID range(s) of the messages that this document was derived from. Acts as a link to the primary source of the document. Note that as a document gets deduplicated, additional ranges will be added, because the same document could be derived from completely separate message ranges."
+ )
+ message_created_at: str = Field(
+ description="The timestamp of the message that this document was derived from. Note that this is not the same as the created_at timestamp of the document. This timestamp is usually only saved with second-level precision."
+ )
+ level: Literal["explicit", "deductive"] = Field(
+ description="The level of the document (explicit or deductive)"
+ )
+ premises: list[str] | None = Field(
+ default=None,
+ description="The premises of the deduction -- only applicable for deductive observations",
+ )
+
+
class DocumentCreate(DocumentBase):
content: Annotated[str, Field(min_length=1, max_length=100000)]
- metadata: dict[str, Any] = {}
-
-
-class DocumentUpdate(DocumentBase):
- content: Annotated[str, Field(min_length=1, max_length=100000)]
- metadata: dict[str, Any] | None = None
+ session_name: str = Field(
+ description="The session from which the document was derived"
+ )
+ metadata: DocumentMetadata = Field()
+ embedding: list[float] = Field()
class MessageSearchOptions(BaseModel):
diff --git a/src/utils/clients.py b/src/utils/clients.py
index 58691440..d5f1e235 100644
--- a/src/utils/clients.py
+++ b/src/utils/clients.py
@@ -1,3 +1,5 @@
+import json
+import logging
from collections.abc import AsyncIterator, Callable
from functools import wraps
from typing import Any, Generic, Literal, TypeVar, cast, overload
@@ -8,20 +10,24 @@ from anthropic.types.message import Message as AnthropicMessage
from google import genai
from google.genai.types import GenerateContentResponse
from groq import AsyncGroq
-from langfuse import get_client
from openai import AsyncOpenAI
from openai.types.chat import ChatCompletion, ChatCompletionChunk
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, ValidationError
from sentry_sdk.ai.monitoring import ai_track
from tenacity import retry, stop_after_attempt, wait_exponential
from src.config import settings
+from src.utils.json_parser import validate_and_repair_json
+from src.utils.langfuse_client import get_langfuse_client
+from src.utils.representation import PromptRepresentation
from src.utils.types import SupportedProviders
+logger = logging.getLogger(__name__)
+
T = TypeVar("T")
M = TypeVar("M", bound=BaseModel)
-lf = get_client()
+lf = get_langfuse_client() if settings.LANGFUSE_PUBLIC_KEY else None
CLIENTS: dict[
SupportedProviders,
@@ -38,12 +44,19 @@ if settings.LLM.OPENAI_API_KEY:
)
CLIENTS["openai"] = openai_client
-if settings.LLM.OPENAI_COMPATIBLE_BASE_URL:
+if settings.LLM.OPENAI_COMPATIBLE_API_KEY and settings.LLM.OPENAI_COMPATIBLE_BASE_URL:
CLIENTS["custom"] = AsyncOpenAI(
api_key=settings.LLM.OPENAI_COMPATIBLE_API_KEY,
base_url=settings.LLM.OPENAI_COMPATIBLE_BASE_URL,
)
+# NOTE: user must know whether they want to use 'custom' or 'vllm'
+if settings.LLM.OPENAI_COMPATIBLE_API_KEY and settings.LLM.OPENAI_COMPATIBLE_BASE_URL:
+ CLIENTS["vllm"] = AsyncOpenAI(
+ api_key=settings.LLM.OPENAI_COMPATIBLE_API_KEY,
+ base_url=settings.LLM.OPENAI_COMPATIBLE_BASE_URL,
+ )
+
if settings.LLM.GEMINI_API_KEY:
google = genai.client.Client(api_key=settings.LLM.GEMINI_API_KEY)
CLIENTS["google"] = google
@@ -56,7 +69,6 @@ SELECTED_PROVIDERS = [
("Dialectic", settings.DIALECTIC.PROVIDER),
("Summary", settings.SUMMARY.PROVIDER),
("Deriver", settings.DERIVER.PROVIDER),
- ("Query Generation Provider", settings.DIALECTIC.QUERY_GENERATION_PROVIDER),
]
for provider_name, provider_value in SELECTED_PROVIDERS:
@@ -105,6 +117,7 @@ async def honcho_llm_call(
*,
response_model: type[M],
json_mode: bool = False,
+ stop_seqs: list[str] | None = None,
reasoning_effort: Literal["low", "medium", "high", "minimal"]
| None = None, # OpenAI only
verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only
@@ -124,6 +137,7 @@ async def honcho_llm_call(
track_name: str | None = None,
response_model: None = None,
json_mode: bool = False,
+ stop_seqs: list[str] | None = None,
reasoning_effort: Literal["low", "medium", "high", "minimal"]
| None = None, # OpenAI only
verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only
@@ -143,6 +157,7 @@ async def honcho_llm_call(
track_name: str | None = None,
response_model: type[BaseModel] | None = None,
json_mode: bool = False,
+ stop_seqs: list[str] | None = None,
reasoning_effort: Literal["low", "medium", "high", "minimal"]
| None = None, # OpenAI only
verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only
@@ -161,6 +176,7 @@ async def honcho_llm_call(
track_name: str | None = None,
response_model: type[BaseModel] | None = None,
json_mode: bool = False,
+ stop_seqs: list[str] | None = None,
reasoning_effort: Literal["low", "medium", "high", "minimal"]
| None = None, # OpenAI only
verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only
@@ -198,6 +214,7 @@ async def honcho_llm_call(
max_tokens,
response_model,
json_mode,
+ stop_seqs,
reasoning_effort,
verbosity,
thinking_budget_tokens,
@@ -211,6 +228,7 @@ async def honcho_llm_call(
max_tokens,
response_model,
json_mode,
+ stop_seqs,
reasoning_effort,
verbosity,
thinking_budget_tokens,
@@ -226,6 +244,7 @@ async def honcho_llm_call_inner(
max_tokens: int,
response_model: type[M],
json_mode: bool = False,
+ stop_seqs: list[str] | None = None,
reasoning_effort: Literal["low", "medium", "high", "minimal"]
| None = None, # OpenAI only
verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only
@@ -242,6 +261,7 @@ async def honcho_llm_call_inner(
max_tokens: int,
response_model: None = None,
json_mode: bool = False,
+ stop_seqs: list[str] | None = None,
reasoning_effort: Literal["low", "medium", "high", "minimal"]
| None = None, # OpenAI only
verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only
@@ -258,6 +278,7 @@ async def honcho_llm_call_inner(
max_tokens: int,
response_model: type[BaseModel] | None = None,
json_mode: bool = False,
+ stop_seqs: list[str] | None = None,
reasoning_effort: Literal["low", "medium", "high", "minimal"]
| None = None, # OpenAI only
verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only
@@ -273,6 +294,7 @@ async def honcho_llm_call_inner(
max_tokens: int,
response_model: type[BaseModel] | None = None,
json_mode: bool = False,
+ stop_seqs: list[str] | None = None,
reasoning_effort: Literal["low", "medium", "high", "minimal"]
| None = None, # OpenAI only
verbosity: Literal["low", "medium", "high"] | None = None, # OpenAI only
@@ -342,11 +364,14 @@ async def honcho_llm_call_inner(
output_tokens=usage.output_tokens if usage else 0, # pyright: ignore
finish_reasons=[stop_reason] if stop_reason else [],
)
+
case AsyncOpenAI():
openai_params: dict[str, Any] = {
"model": params["model"],
"messages": params["messages"],
}
+ if stop_seqs:
+ openai_params["stop"] = stop_seqs
if "gpt-5" in model:
openai_params["max_completion_tokens"] = params["max_tokens"]
if reasoning_effort:
@@ -355,9 +380,91 @@ async def honcho_llm_call_inner(
openai_params["verbosity"] = verbosity
else:
openai_params["max_tokens"] = params["max_tokens"]
- if json_mode:
+
+ if json_mode and provider != "vllm":
openai_params["response_format"] = {"type": "json_object"}
- if response_model:
+
+ # custom shim for vLLM response model formatting
+ # NOTE: this is all specific to the Representation model.
+ # Do not call with any other response model.
+ if provider == "vllm" and response_model:
+ if response_model is not PromptRepresentation:
+ raise NotImplementedError(
+ "vLLM structured output currently supports only PromptRepresentation"
+ )
+ openai_params["response_format"] = {
+ "type": "json_schema",
+ "json_schema": {
+ "name": response_model.__name__,
+ "schema": response_model.model_json_schema(),
+ },
+ }
+ response: ChatCompletion = await client.chat.completions.create( # pyright: ignore
+ **openai_params
+ )
+
+ usage = response.usage # pyright: ignore
+ finish_reason = response.choices[0].finish_reason # pyright: ignore
+
+ try:
+ test_rep = ""
+ if response.choices[0].message.content is not None: # pyright: ignore
+ test_rep = response.choices[0].message.content # pyright: ignore
+
+ final = validate_and_repair_json(test_rep) # pyright: ignore
+
+ # Schema-aware repair: ensure deductive observations have required fields
+
+ repaired_data = json.loads(final)
+
+ # Fix deductive observations that might be missing conclusion
+ if "deductive" in repaired_data and isinstance(
+ repaired_data["deductive"], list
+ ):
+ for i, item in enumerate(repaired_data["deductive"]):
+ if isinstance(item, dict):
+ # If conclusion is missing but premises exist, create a placeholder
+ if "conclusion" not in item and "premises" in item:
+ logger.warning(
+ f"Deductive observation {i} missing conclusion, adding placeholder"
+ )
+ # Try to generate a conclusion from premises if possible
+ if item["premises"]:
+ item["conclusion"] = (
+ f"[Incomplete reasoning from premises: {item['premises'][0][:100]}...]"
+ )
+ else:
+ item["conclusion"] = (
+ "[Incomplete reasoning - conclusion missing]"
+ )
+ # If premises is missing, add empty list (it's optional with default)
+ if "premises" not in item:
+ item["premises"] = []
+
+ final = json.dumps(repaired_data)
+ except (json.JSONDecodeError, KeyError, TypeError) as e:
+ final = ""
+ logger.warning(f"Could not perform schema-aware repair: {e}")
+ # Continue with original final value if repair fails
+
+ try:
+ response_obj = PromptRepresentation.model_validate_json(final)
+ except ValidationError as e:
+ logger.error(f"Validation error after repair: {e}")
+ logger.debug(f"Problematic JSON: {final}")
+
+ # Fallback: return empty response rather than failing
+ logger.warning(
+ "Using fallback empty Representation due to validation error"
+ )
+ response_obj = PromptRepresentation(explicit=[], deductive=[])
+
+ return HonchoLLMCallResponse(
+ content=response_obj,
+ output_tokens=usage.completion_tokens if usage else 0, # pyright: ignore
+ finish_reasons=[finish_reason] if finish_reason else [],
+ )
+ elif response_model:
openai_params["response_format"] = response_model
response: ChatCompletion = await client.chat.completions.parse( # pyright: ignore
**openai_params
@@ -367,10 +474,15 @@ async def honcho_llm_call_inner(
if parsed_content is None:
raise ValueError("No parsed content in structured response")
- # Safely extract usage and finish_reason
usage = response.usage
finish_reason = response.choices[0].finish_reason
+ # Validate that parsed content matches the response model
+ if not isinstance(parsed_content, response_model):
+ raise ValueError(
+ f"Parsed content does not match the response model: {parsed_content} != {response_model}"
+ )
+
return HonchoLLMCallResponse(
content=parsed_content,
output_tokens=usage.completion_tokens if usage else 0,
@@ -381,7 +493,6 @@ async def honcho_llm_call_inner(
**openai_params
)
- # Safely extract usage and finish_reason
usage = response.usage # pyright: ignore
finish_reason = response.choices[0].finish_reason # pyright: ignore
@@ -390,6 +501,7 @@ async def honcho_llm_call_inner(
output_tokens=usage.completion_tokens if usage else 0, # pyright: ignore
finish_reasons=[finish_reason] if finish_reason else [],
)
+
case genai.Client():
if response_model is None:
gemini_response: GenerateContentResponse = (
@@ -446,6 +558,12 @@ async def honcho_llm_call_inner(
else "stop"
)
+ # Validate that parsed content matches the response model
+ if not isinstance(gemini_response.parsed, response_model):
+ raise ValueError(
+ f"Parsed content does not match the response model: {gemini_response.parsed} != {response_model}"
+ )
+
return HonchoLLMCallResponse(
content=gemini_response.parsed,
output_tokens=token_count,
@@ -464,6 +582,7 @@ async def honcho_llm_call_inner(
elif json_mode:
groq_params["response_format"] = {"type": "json_object"}
+ # TODO: figure out why groq returns unknown type and fix it
response: ChatCompletion = await client.chat.completions.create( # pyright: ignore
**groq_params
)
@@ -474,11 +593,27 @@ async def honcho_llm_call_inner(
usage = response.usage # pyright: ignore
finish_reason = response.choices[0].finish_reason # pyright: ignore
- return HonchoLLMCallResponse(
- content=response.choices[0].message.content, # pyright: ignore
- output_tokens=usage.completion_tokens if usage else 0, # pyright: ignore
- finish_reasons=[finish_reason] if finish_reason else [],
- )
+ # Handle response model parsing for Groq
+ if response_model:
+ try:
+ json_content = json.loads(response.choices[0].message.content) # pyright: ignore
+ parsed_content = response_model.model_validate(json_content)
+
+ return HonchoLLMCallResponse(
+ content=parsed_content,
+ output_tokens=usage.completion_tokens if usage else 0, # pyright: ignore
+ finish_reasons=[finish_reason] if finish_reason else [],
+ )
+ except (json.JSONDecodeError, ValidationError, ValueError) as e:
+ raise ValueError(
+ f"Failed to parse Groq response as {response_model}: {e}. Raw content: {response.choices[0].message.content}" # pyright: ignore
+ ) from e
+ else:
+ return HonchoLLMCallResponse(
+ content=response.choices[0].message.content, # pyright: ignore
+ output_tokens=usage.completion_tokens if usage else 0, # pyright: ignore
+ finish_reasons=[finish_reason] if finish_reason else [],
+ )
async def handle_streaming_response(
@@ -651,7 +786,10 @@ async def handle_streaming_response(
def with_langfuse(func: Callable[..., Any]) -> Callable[..., Any]:
@wraps(func)
async def wrapper(*args: Any, **kwargs: Any) -> Any:
- lf.start_as_current_generation(name="LLM Call")
- return await func(*args, **kwargs)
+ if lf:
+ with lf.start_as_current_generation(name="LLM Call"):
+ return await func(*args, **kwargs)
+ else:
+ return await func(*args, **kwargs)
return wrapper
diff --git a/src/utils/embedding_store.py b/src/utils/embedding_store.py
deleted file mode 100644
index 8144b070..00000000
--- a/src/utils/embedding_store.py
+++ /dev/null
@@ -1,458 +0,0 @@
-from __future__ import annotations
-
-import datetime
-import logging
-from typing import Any, Literal, overload
-
-from langfuse import get_client
-from openai.types import CreateEmbeddingResponse
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from src import crud, models
-from src.config import settings
-from src.dependencies import tracked_db
-from src.embedding_client import embedding_client
-from src.utils.formatting import format_datetime_utc
-from src.utils.logging import conditional_observe
-from src.utils.shared_models import (
- Observation,
- ObservationContext,
- ObservationMetadata,
- ReasoningLevel,
- UnifiedObservation,
-)
-
-logger = logging.getLogger(__name__)
-
-lf = get_client()
-
-
-class EmbeddingStore:
- """Embedding store specialized for observation-based reasoning with structured metadata."""
-
- def __init__(
- self, workspace_name: str, peer_name: str, collection_name: str
- ) -> None:
- self.workspace_name: str = workspace_name
- self.peer_name: str = peer_name
- self.collection_name: str = collection_name
-
- @conditional_observe
- async def save_unified_observations(
- self,
- observations: list[UnifiedObservation],
- message_id: int,
- session_name: str,
- message_created_at: datetime.datetime,
- fallback_level: str = "explicit",
- similarity_threshold: float = 0.85,
- ) -> None:
- """Save UnifiedObservation objects to the collection.
-
- This method handles UnifiedObservation objects by:
- 1. Generating embeddings only from conclusions
- 2. Storing premises in metadata for reference
-
- Args:
- observations: List of UnifiedObservation objects or strings
- message_id: Message ID to link with observations
- session_name: Session name to link with existing summary context
- message_created_at: Timestamp when the message was created
- fallback_level: Reasoning level for the observations if not provided
- similarity_threshold: Threshold for considering observations similar
- """
- # Extract conclusions for deduplication and embedding
- conclusions: list[str] = [obs.conclusion for obs in observations]
-
- # Remove duplicates before saving
- unique_conclusions: list[str] = await self.remove_duplicates(
- conclusions, similarity_threshold=similarity_threshold
- )
- if settings.LANGFUSE_PUBLIC_KEY:
- lf.update_current_trace(
- input={"observations": [obs.model_dump() for obs in observations]},
- output={"unique_conclusions": unique_conclusions},
- )
-
- if not unique_conclusions:
- logger.debug("No unique observations to save after deduplication")
- return
-
- # Create mapping from conclusion back to original observation
- conclusion_to_observation: dict[str, UnifiedObservation] = {
- obs.conclusion: obs for obs in observations
- }
-
- # Filter unified observations to only unique ones
- unique_observations: list[UnifiedObservation] = [
- conclusion_to_observation[conclusion] for conclusion in unique_conclusions
- ]
-
- # Batch embed all unique conclusions (not premises)
- embeddings: list[list[float]] = []
- batch_size: int = 2048 # OpenAI batch limit
-
- for i in range(0, len(unique_conclusions), batch_size):
- batch = unique_conclusions[i : i + batch_size]
- response: CreateEmbeddingResponse = (
- await embedding_client.client.embeddings.create(
- input=batch, model="text-embedding-3-small"
- )
- )
- embeddings.extend([data.embedding for data in response.data])
-
- # Batch create document objects
- document_objects: list[models.Document] = []
- for obs, embedding in zip(unique_observations, embeddings, strict=True):
- # Use the observation's own level or fall back to parameter level
- obs_level = obs.level or fallback_level
-
- # Build metadata including premises
- metadata: dict[str, Any] = {
- "level": obs_level,
- "message_id": message_id,
- "session_name": session_name,
- "premises": obs.premises, # Store premises in metadata
- "created_at": format_datetime_utc(message_created_at),
- }
-
- doc = models.Document(
- workspace_name=self.workspace_name,
- peer_name=self.peer_name,
- collection_name=self.collection_name,
- content=obs.conclusion, # Store only conclusion as content
- internal_metadata=metadata,
- embedding=embedding, # Embedding generated from conclusion only
- created_at=message_created_at,
- )
- document_objects.append(doc)
- async with tracked_db("ed_embedding_store.save_unified_observations") as db:
- # Batch insert all documents
- db.add_all(document_objects)
- await db.commit()
- logger.debug("Batch created %s unified observations", len(document_objects))
-
- @overload
- async def get_relevant_observations(
- self,
- query: str,
- *,
- top_k: int = 5,
- max_distance: float = 0.3,
- level: str | None = None,
- conversation_context: str = "",
- for_reasoning: Literal[True],
- ) -> ObservationContext: ...
-
- @overload
- async def get_relevant_observations(
- self,
- query: str,
- *,
- top_k: int = 5,
- max_distance: float = 0.3,
- level: str | None = None,
- conversation_context: str = "",
- for_reasoning: Literal[False],
- ) -> list[models.Document]: ...
-
- async def get_relevant_observations(
- self,
- query: str,
- *,
- top_k: int = 5,
- max_distance: float = 0.3,
- level: str | None = None,
- conversation_context: str = "",
- for_reasoning: bool = False,
- ) -> list[models.Document] | ObservationContext:
- """
- Unified method to get relevant observations with flexible options.
-
- Args:
- query: The search query
- top_k: Number of results to return
- max_distance: Maximum distance for semantic similarity
- level: Optional reasoning level to filter by
- conversation_context: Additional conversation context
- for_reasoning: If True, returns ObservationContext for ed reasoning
-
- Returns:
- List of documents or ObservationContext (if for_reasoning=True)
- """
- async with tracked_db("embedding_store.get_relevant_observations") as db:
- return await self._get_observations_internal(
- db,
- query,
- top_k,
- max_distance,
- level,
- conversation_context,
- for_reasoning,
- )
-
- def _build_filter_conditions(
- self,
- level: str | None = None,
- ) -> dict[str, Any]:
- """Build complete filter conditions for document queries."""
- conditions: list[dict[str, Any]] = []
-
- if level:
- conditions.append({"internal_metadata": {"level": level}})
-
- if not conditions:
- return {}
-
- return conditions[0] if len(conditions) == 1 else {"AND": conditions}
-
- async def _get_observations_internal(
- self,
- db: AsyncSession,
- query: str,
- top_k: int,
- max_distance: float,
- level: str | None,
- conversation_context: str,
- for_reasoning: bool,
- ) -> Any:
- """Internal method that does the actual observation retrieval."""
- try:
- if for_reasoning:
- return await self._get_observations_for_reasoning(
- db,
- query,
- max_distance,
- conversation_context,
- )
- else:
- # Regular document list return
- if level:
- return await self._query_documents_for_level(
- db,
- query,
- level,
- conversation_context,
- max_distance,
- top_k,
- )
- else:
- documents = await crud.query_documents(
- db,
- workspace_name=self.workspace_name,
- peer_name=self.peer_name,
- collection_name=self.collection_name,
- query=self._build_truncated_query(query, ""),
- max_distance=max_distance,
- top_k=top_k,
- )
- db.expunge_all()
- return list(documents)
-
- except Exception as e:
- logger.error(f"Error getting relevant observations: {e}")
- if for_reasoning:
- return ObservationContext()
- return []
-
- async def _get_observations_for_reasoning(
- self,
- db: AsyncSession,
- query: str,
- max_distance: float,
- conversation_context: str,
- ) -> ObservationContext:
- """Get observations formatted for reasoning with ObservationContext."""
- context = ObservationContext()
-
- for level_name in ["explicit", "deductive"]:
- count: int = getattr(self, f"{level_name}_observations_count", 5)
- level_enum = ReasoningLevel(level_name)
-
- docs = await self._query_documents_for_level(
- db,
- query,
- level_name,
- conversation_context,
- max_distance,
- count,
- )
-
- seen_observations: set[str] = set()
- for doc in docs:
- normalized_content: str = doc.content.strip().lower()
- if normalized_content not in seen_observations:
- metadata = self._extract_observation_metadata(doc)
- observation = Observation(
- content=doc.content,
- metadata=metadata,
- created_at=doc.created_at,
- )
- context.add_observation(observation, level_enum)
- seen_observations.add(normalized_content)
-
- return context
-
- def _build_truncated_query(
- self,
- query: str,
- conversation_context: str = "",
- max_tokens: int | None = None,
- ) -> str:
- """Build a query that fits within token limits with clear priorities.
-
- Args:
- query: The search query
- conversation_context: Optional conversation context to include
- max_tokens: Maximum tokens allowed (defaults to setting with buffer)
-
- Returns:
- Truncated query string that fits within token limits
- """
- max_tokens = max_tokens or (settings.MAX_EMBEDDING_TOKENS - 100)
- encoding = embedding_client.encoding
-
- # Pre-calculate all token counts once
- query_prefix = "Current message: "
- context_prefix = "\nContext: "
-
- prefix_tokens = len(encoding.encode(query_prefix))
- context_prefix_tokens = len(encoding.encode(context_prefix))
- query_tokens = encoding.encode(query)
-
- # Simple case: query alone fits
- if prefix_tokens + len(query_tokens) <= max_tokens:
- if not conversation_context:
- return f"{query_prefix}{query}"
-
- # Try to add context
- context_tokens = encoding.encode(conversation_context)
- total_without_context = (
- prefix_tokens + len(query_tokens) + context_prefix_tokens
- )
-
- if total_without_context + len(context_tokens) <= max_tokens:
- return f"{query_prefix}{query}{context_prefix}{conversation_context}"
-
- # Truncate context to fit
- available_context_tokens = max_tokens - total_without_context
- if available_context_tokens > 0:
- truncated_context = encoding.decode(
- context_tokens[-available_context_tokens:]
- )
- return f"{query_prefix}{query}{context_prefix}{truncated_context}"
- else:
- # No room left for context; keep full query intact
- return f"{query_prefix}{query}"
-
- # Query itself is too long - truncate it
- available_query_tokens = max_tokens - prefix_tokens
- if available_query_tokens > 0:
- # Keep the end (recency) of the query
- truncated_query = encoding.decode(query_tokens[-available_query_tokens:])
- return f"{query_prefix}{truncated_query}"
-
- # Pathological case - just return what we can
- logger.warning("Token limit too restrictive: %s", max_tokens)
- return encoding.decode(query_tokens[:max_tokens])
-
- async def _query_documents_for_level(
- self,
- db: AsyncSession,
- query: str,
- level: str,
- conversation_context: str,
- max_distance: float,
- count: int,
- ) -> list[models.Document]:
- """Query documents for a specific level."""
- # Construct the combined query with truncation to prevent token limit errors
- combined_query = self._build_truncated_query(query, conversation_context)
-
- documents = await crud.query_documents(
- db,
- workspace_name=self.workspace_name,
- peer_name=self.peer_name,
- collection_name=self.collection_name,
- query=combined_query,
- max_distance=max_distance,
- top_k=count * 3,
- filters=self._build_filter_conditions(level),
- )
-
- # Sort by creation time and return top count
- docs_sorted: list[models.Document] = sorted(
- list(documents), key=lambda x: x.created_at, reverse=True
- )
- return docs_sorted[:count]
-
- def _extract_observation_metadata(self, doc: models.Document) -> Any:
- """Extract metadata from a document for ObservationMetadata."""
- metadata = ObservationMetadata()
- if doc.internal_metadata:
- metadata.session_context = doc.internal_metadata.get("session_context", "")
- metadata.summary_id = doc.internal_metadata.get("summary_id", "")
- metadata.message_id = doc.internal_metadata.get("message_id")
- metadata.level = doc.internal_metadata.get("level")
- metadata.session_name = doc.internal_metadata.get("session_name")
- metadata.premises = doc.internal_metadata.get("premises", [])
- return metadata
-
- async def remove_duplicates(
- self,
- facts: list[str],
- *,
- similarity_threshold: float = 0.85,
- ) -> list[str]:
- """Remove duplicate observations based on similarity threshold.
-
- Args:
- facts: List of observation strings
- similarity_threshold: Threshold for considering observations similar
-
- Returns:
- List of unique observations
- """
- if not facts:
- return []
-
- # Batch generate embeddings for all facts at once
- embeddings: list[list[float]] = []
- batch_size: int = 2048 # OpenAI batch limit
-
- for i in range(0, len(facts), batch_size):
- batch = facts[i : i + batch_size]
- response: CreateEmbeddingResponse = (
- await embedding_client.client.embeddings.create(
- input=batch, model="text-embedding-3-small"
- )
- )
- embeddings.extend([data.embedding for data in response.data])
-
- # Now check each fact for duplicates using query_documents with pre-computed embeddings
- unique_observations: list[str] = []
-
- async with tracked_db("embedding_store.remove_duplicates") as db:
- for fact, embedding in zip(facts, embeddings, strict=True):
- documents = await crud.query_documents(
- db,
- workspace_name=self.workspace_name,
- peer_name=self.peer_name,
- collection_name=self.collection_name,
- query=fact,
- max_distance=1.0 - similarity_threshold,
- top_k=1,
- embedding=embedding, # Pass pre-computed embedding
- )
-
- docs_list: list[models.Document] = list(documents)
- if not docs_list:
- unique_observations.append(fact)
-
- logger.debug(
- "Batch remove duplicates: %s input facts, %s unique after deduplication",
- len(facts),
- len(unique_observations),
- )
-
- return unique_observations
diff --git a/src/utils/formatting.py b/src/utils/formatting.py
index cb205e37..a7276e5c 100644
--- a/src/utils/formatting.py
+++ b/src/utils/formatting.py
@@ -6,10 +6,6 @@ and handling temporal metadata for the reasoning system.
"""
from datetime import datetime, timezone
-from typing import Any, Protocol, cast, runtime_checkable
-
-from src.utils.logging import conditional_observe
-from src.utils.shared_models import ReasoningResponse
def format_datetime_utc(dt: datetime) -> str:
@@ -39,6 +35,9 @@ def format_datetime_utc(dt: datetime) -> str:
if dt.tzinfo != timezone.utc:
dt = dt.astimezone(timezone.utc)
+ # Remove subsecond precision
+ dt = dt.replace(microsecond=0)
+
# Format and replace +00:00 with Z
return dt.isoformat().replace("+00:00", "Z")
@@ -46,13 +45,14 @@ def format_datetime_utc(dt: datetime) -> str:
def utc_now_iso() -> str:
"""
Get current UTC time as ISO 8601 string with Z suffix.
+ Removes subsecond precision.
Returns:
Current UTC time in ISO 8601 format with Z suffix
Example:
>>> utc_now_iso()
- '2023-01-01T12:34:56.789123Z'
+ '2023-01-01T12:34:56Z'
"""
return format_datetime_utc(datetime.now(timezone.utc))
@@ -115,85 +115,6 @@ def parse_datetime_iso(iso_string: str) -> datetime:
raise ValueError(f"Invalid ISO 8601 datetime format: {e}") from e
-@runtime_checkable
-class StructuredObservation(Protocol):
- """Protocol for observations that have conclusion and premises attributes."""
-
- conclusion: str
- premises: list[str]
-
-
-REASONING_LEVELS: list[str] = ["explicit", "deductive"]
-LEVEL_LABELS: dict[str, str] = {
- "explicit": "Explicit (Literal facts directly stated by the user)",
- "deductive": "Deductive (Logically necessary conclusions from explicit facts)",
-}
-
-
-def format_premises_for_display(premises: list[str]) -> str:
- """
- Format premises as a clean bulleted list for display.
-
- Args:
- premises: List of premise strings
-
- Returns:
- Formatted premises text with newlines and bullets, or empty string if no premises
- """
- if not premises:
- return ""
-
- premises_formatted: list[str] = []
- for premise in premises:
- premises_formatted.append(f" - {premise}")
- return "\n" + "\n".join(premises_formatted)
-
-
-def format_structured_observation(conclusion: str, premises: list[str]) -> str:
- """
- Format a structured observation with conclusion and premises for display.
-
- Args:
- conclusion: The main conclusion
- premises: List of supporting premises
-
- Returns:
- Formatted observation string
- """
- premises_text = format_premises_for_display(premises)
- return f"{conclusion}{premises_text}"
-
-
-def extract_observation_content(observation: str | dict[str, Any] | Any) -> str:
- """Extract content string from an observation (dict or string)."""
- # Handle StructuredObservation objects (Pydantic models)
- if isinstance(observation, StructuredObservation):
- return format_structured_observation(
- observation.conclusion, observation.premises
- )
-
- # Handle explicit observations as simple strings
- if isinstance(observation, str):
- return observation
-
- if isinstance(observation, dict):
- # For explicit observations with conclusions
- if "conclusions" in observation:
- conclusions_value: str = cast(str, observation["conclusions"])
- if isinstance(conclusions_value, list):
- return "; ".join(cast(list[str], conclusions_value))
- return conclusions_value
- # For structured observations, return conclusion with premises formatted
- if "conclusion" in observation:
- conclusion: str = cast(str, observation["conclusion"])
- premises: list[str] = cast(list[str], observation.get("premises", [])) # pyright: ignore
- return format_structured_observation(conclusion, premises)
- # Fallback to content field or string representation
- content_value: str | None = observation.get("content") # pyright: ignore
- return content_value if content_value is not None else str(observation) # pyright: ignore
- return str(observation)
-
-
def format_new_turn_with_timestamp(
new_turn: str, current_time: datetime, speaker: str
) -> str:
@@ -210,103 +131,3 @@ def format_new_turn_with_timestamp(
"""
current_time_str = current_time.strftime("%Y-%m-%d %H:%M:%S")
return f"{current_time_str} {speaker}: {new_turn}"
-
-
-def format_context_for_prompt(
- context: ReasoningResponse | dict[str, Any] | None,
-) -> str:
- """
- Format context into a clean, readable string for LLM prompts.
-
- Args:
- context: ReasoningResponse object or dict with reasoning levels as keys and observation lists as values
- Observations can be strings or dicts - will be normalized
-
- Returns:
- Formatted string with clear sections and bullet points including temporal metadata
- """
- if not context:
- return "No context available."
-
- formatted_sections: list[str] = []
-
- # Handle both ReasoningResponse objects and dicts
- if isinstance(context, ReasoningResponse):
- # It's a ReasoningResponse object
- observations_by_level = {
- "explicit": context.explicit,
- "deductive": context.deductive,
- }
- else:
- # It's a dict
- observations_by_level = context
- # Process each level in a consistent order
- for level in REASONING_LEVELS:
- observations = observations_by_level.get(level, [])
- if not observations:
- continue
-
- label = LEVEL_LABELS.get(level, level.title())
- formatted_sections.append(f"{label}:")
-
- # Format observations with temporal metadata when available
- for observation in observations:
- observation_content = extract_observation_content(observation)
- formatted_sections.append(f" • {observation_content}")
-
- formatted_sections.append("") # Blank line between sections
-
- # Remove trailing blank line if exists
- if formatted_sections and formatted_sections[-1] == "":
- formatted_sections.pop()
-
- return (
- "\n".join(formatted_sections)
- if formatted_sections
- else "No relevant context available."
- )
-
-
-def normalize_observations_for_comparison(observations: list[Any]) -> set[str]:
- """Convert observations to normalized strings for comparison."""
- normalized: set[str] = set()
- for observation in observations:
- observation_content = extract_observation_content(observation)
- normalized.add(observation_content.strip().lower())
- return normalized
-
-
-@conditional_observe
-def find_new_observations(
- original_context: ReasoningResponse, revised_observations: ReasoningResponse
-) -> dict[str, list[str]]:
- """
- Find observations that are new in revised_observations compared to original_context.
-
- Args:
- original_context: Original observation context
- revised_observations: Revised observation context
-
- Returns:
- Dictionary with new observations by level
- """
- new_observations_by_level: dict[str, list[str]] = {}
-
- for level in REASONING_LEVELS:
- original_observations = normalize_observations_for_comparison(
- getattr(original_context, level, [])
- )
- revised_list = getattr(revised_observations, level, [])
-
- # Find genuinely new observations
- new_observations: list[str] = []
- for observation in revised_list:
- normalized_observation = (
- extract_observation_content(observation).strip().lower()
- )
- if normalized_observation not in original_observations:
- new_observations.append(observation)
-
- new_observations_by_level[level] = new_observations
-
- return new_observations_by_level
diff --git a/src/utils/json_parser.py b/src/utils/json_parser.py
new file mode 100644
index 00000000..65673b46
--- /dev/null
+++ b/src/utils/json_parser.py
@@ -0,0 +1,378 @@
+import json
+import logging
+import re
+from typing import Any
+
+from json_repair import repair_json # pyright: ignore
+
+logger = logging.getLogger(__name__)
+# logging.getLogger("sqlalchemy.engine.Engine").disabled = True
+
+
+def comprehensive_json_repair(json_str: str) -> str:
+ """Comprehensively repair malformed JSON with multiple strategies"""
+
+ # Strategy 1: Handle truncated JSON by parsing what we can
+ repaired = try_partial_parse_repair(json_str)
+ if repaired:
+ return repaired
+
+ # Strategy 2: Smart bracket/brace matching with context awareness
+ repaired = try_contextual_closure_repair(json_str)
+ if repaired:
+ return repaired
+
+ # Strategy 3: Line-by-line reconstruction
+ repaired = try_line_reconstruction_repair(json_str)
+ if repaired:
+ return repaired
+
+ # Strategy 4: Regex-based common pattern fixes
+ repaired = try_regex_pattern_repair(json_str)
+ if repaired:
+ return repaired
+
+ # Fallback: Original simple method
+ return simple_bracket_repair(json_str)
+
+
+def try_partial_parse_repair(json_str: str) -> str | None:
+ """Try to parse JSON incrementally and reconstruct from valid parts"""
+ try:
+ # First, try to find the last complete object/array
+ lines = json_str.split("\n")
+
+ for i in range(len(lines), 0, -1):
+ partial = "\n".join(lines[:i])
+
+ # Try different closure strategies
+ for closure_attempt in generate_closure_attempts(partial):
+ try:
+ json.loads(closure_attempt)
+ return closure_attempt
+ except json.JSONDecodeError:
+ continue
+
+ return None
+ except Exception:
+ return None
+
+
+def generate_closure_attempts(partial_json: str) -> list[str]:
+ """Generate different ways to close the JSON structure"""
+ attempts: list[str] = []
+
+ # Analyze the structure to understand what's open
+ stack: list[tuple[str, int]] = []
+ in_string = False
+ escape_next = False
+
+ for i, char in enumerate(partial_json):
+ if escape_next:
+ escape_next = False
+ continue
+
+ if char == "\\":
+ escape_next = True
+ continue
+
+ if char == '"' and not escape_next:
+ in_string = not in_string
+ continue
+
+ if in_string:
+ continue
+
+ if char in "({[":
+ stack.append((char, i))
+ elif char in ")}]" and stack:
+ opener, _ = stack.pop()
+ # Verify matching pairs
+ if not (
+ (char == ")" and opener == "(")
+ or (char == "}" and opener == "{")
+ or (char == "]" and opener == "[")
+ ):
+ # Mismatched - this is likely where corruption started
+ break
+
+ # Generate closure attempts based on what's still open
+ base = partial_json.rstrip()
+
+ # Remove trailing comma if present
+ if base.rstrip().endswith(","):
+ base = base[:-1]
+ attempts.append(base)
+
+ # Close based on stack
+ closures: list[str] = []
+ for opener, _ in reversed(stack):
+ if opener == "{":
+ closures.append("}")
+ elif opener == "[":
+ closures.append("]")
+ elif opener == "(":
+ closures.append(")")
+
+ # Try different combinations
+ attempts.append(base + "".join(closures))
+
+ # Try closing just objects/arrays (ignore parentheses)
+ obj_closures = [c for c in closures if c in "]}"]
+ attempts.append(base + "".join(obj_closures))
+
+ # Try adding missing quotes if we're in a string
+ if in_string:
+ attempts.append(base + '"' + "".join(closures))
+
+ return attempts
+
+
+def try_contextual_closure_repair(json_str: str) -> str | None:
+ """Smart closure repair based on JSON context"""
+ try:
+ # Find the last valid JSON token
+ tokens: list[dict[str, Any]] = tokenize_json(json_str)
+
+ # Look for patterns that indicate what should come next
+ if not tokens:
+ return None
+
+ last_token: dict[str, Any] = tokens[-1]
+
+ # If last token is a value, we might need to close objects/arrays
+ if last_token["type"] in ["string", "number", "boolean", "null"]:
+ return try_close_after_value(json_str, tokens)
+
+ # If last token is a structural element, handle appropriately
+ elif last_token["type"] in ["comma", "colon"]:
+ return try_complete_structure(json_str, tokens)
+
+ return None
+ except Exception:
+ return None
+
+
+def tokenize_json(json_str: str) -> list[dict[str, Any]]:
+ """Tokenize JSON string into meaningful components"""
+ tokens: list[dict[str, Any]] = []
+ i = 0
+
+ while i < len(json_str):
+ char = json_str[i]
+
+ # Skip whitespace
+ if char.isspace():
+ i += 1
+ continue
+
+ # String literals
+ if char == '"':
+ start = i
+ i += 1
+ while i < len(json_str):
+ if json_str[i] == '"' and json_str[i - 1] != "\\":
+ break
+ i += 1
+ tokens.append(
+ {
+ "type": "string",
+ "value": json_str[start : i + 1],
+ "start": start,
+ "end": i,
+ }
+ )
+
+ # Numbers
+ elif char.isdigit() or char == "-":
+ start = i
+ while i < len(json_str) and (
+ json_str[i].isdigit() or json_str[i] in ".-eE"
+ ):
+ i += 1
+ tokens.append(
+ {
+ "type": "number",
+ "value": json_str[start:i],
+ "start": start,
+ "end": i - 1,
+ }
+ )
+ continue # Don't increment i again
+
+ # Structural characters
+ elif char in "{}[],:":
+ token_type = {
+ "{": "object_start",
+ "}": "object_end",
+ "[": "array_start",
+ "]": "array_end",
+ ",": "comma",
+ ":": "colon",
+ }[char]
+
+ tokens.append({"type": token_type, "value": char, "start": i, "end": i})
+
+ # Boolean/null literals
+ elif char in "tfn":
+ if json_str[i : i + 4] == "true":
+ tokens.append(
+ {"type": "boolean", "value": "true", "start": i, "end": i + 3}
+ )
+ i += 3
+ elif json_str[i : i + 5] == "false":
+ tokens.append(
+ {"type": "boolean", "value": "false", "start": i, "end": i + 4}
+ )
+ i += 4
+ elif json_str[i : i + 4] == "null":
+ tokens.append(
+ {"type": "null", "value": "null", "start": i, "end": i + 3}
+ )
+ i += 3
+
+ i += 1
+
+ return tokens
+
+
+def try_close_after_value(json_str: str, tokens: list[dict[str, Any]]) -> str | None:
+ """Try to close JSON after a value token"""
+ # Analyze nesting to determine what needs to be closed
+ nesting_stack: list[str] = []
+
+ for token in tokens[:-1]: # Exclude the last token (which is the value)
+ if token["type"] == "object_start":
+ nesting_stack.append("}")
+ elif token["type"] == "array_start":
+ nesting_stack.append("]")
+ elif (
+ token["type"] in ["object_end", "array_end"]
+ and nesting_stack
+ and nesting_stack[-1] == token["value"]
+ ):
+ nesting_stack.pop()
+
+ # Close remaining open structures
+ closure = "".join(reversed(nesting_stack))
+ candidate = json_str + closure
+
+ try:
+ json.loads(candidate)
+ return candidate
+ except json.JSONDecodeError:
+ return None
+
+
+def try_complete_structure(json_str: str, tokens: list[dict[str, Any]]) -> str | None:
+ """Try to complete JSON ending with structural tokens like comma or colon"""
+ last_token = tokens[-1]
+
+ if last_token["type"] == "comma":
+ # After comma, we might be missing a key-value pair or array element
+ # Try removing the trailing comma first
+ trimmed = json_str.rstrip().rstrip(",")
+ return try_contextual_closure_repair(trimmed)
+
+ elif last_token["type"] == "colon":
+ # After colon, we're missing a value - try adding a placeholder
+ candidates = [
+ json_str + "null",
+ json_str + '""',
+ json_str + "[]",
+ json_str + "{}",
+ ]
+
+ for candidate in candidates:
+ try:
+ repaired = try_contextual_closure_repair(candidate)
+ if repaired:
+ return repaired
+ except (json.JSONDecodeError, TypeError, ValueError):
+ continue
+
+ return None
+
+
+def try_line_reconstruction_repair(json_str: str) -> str | None:
+ """Try to reconstruct JSON line by line"""
+ lines = json_str.split("\n")
+
+ # Find the last line that makes the JSON valid when truncated there
+ for i in range(len(lines), 0, -1):
+ partial_lines = lines[:i]
+ partial_json = "\n".join(partial_lines)
+
+ # Try to repair this partial JSON
+ repaired = try_contextual_closure_repair(partial_json)
+ if repaired:
+ return repaired
+
+ return None
+
+
+def try_regex_pattern_repair(json_str: str) -> str | None:
+ """Use regex to fix common JSON formatting issues"""
+ fixed = json_str
+
+ # Remove trailing commas before closing braces/brackets
+ fixed = re.sub(r",(\s*[}\]])", r"\1", fixed)
+
+ # Fix unescaped quotes in strings (basic attempt)
+ fixed = re.sub(r'(? str:
+ """Fallback: Original simple bracket counting method"""
+ open_braces = json_str.count("{")
+ close_braces = json_str.count("}")
+ open_brackets = json_str.count("[")
+ close_brackets = json_str.count("]")
+
+ missing_brackets = open_brackets - close_brackets
+ missing_braces = open_braces - close_braces
+
+ repaired = json_str
+ repaired += "]" * max(0, missing_brackets)
+ repaired += "}" * max(0, missing_braces)
+
+ return repaired
+
+
+def validate_and_repair_json(json_str: str) -> str:
+ """Main function with comprehensive repair strategies"""
+ json_str = json_str.strip()
+
+ # Try parsing with repair library
+ good_json = repair_json(json_str)
+ if good_json:
+ return good_json
+
+ # Try comprehensive repair
+ try:
+ repaired = comprehensive_json_repair(json_str)
+
+ # Validate the repair
+ json.loads(repaired)
+ logger.info("✅ JSON successfully repaired!")
+ return repaired
+
+ except json.JSONDecodeError as repair_error:
+ logger.error(f"❌ Repair failed: {repair_error}")
+ raise ValueError(
+ f"Could not repair JSON. Original error: {repair_error.msg}, "
+ + f"Repair error: {repair_error.msg}"
+ ) from repair_error
diff --git a/src/utils/langfuse_client.py b/src/utils/langfuse_client.py
new file mode 100644
index 00000000..c6e68c0b
--- /dev/null
+++ b/src/utils/langfuse_client.py
@@ -0,0 +1,40 @@
+"""
+Centralized Langfuse client management.
+
+This module provides a singleton Langfuse client to avoid multiple initialization
+errors when modules import get_client() at the module level.
+"""
+
+from typing import Any
+
+from langfuse import get_client
+
+_langfuse_client: Any = None
+
+
+def get_langfuse_client() -> Any:
+ """
+ Get the singleton Langfuse client instance.
+
+ This function ensures that get_client() is only called once, regardless of
+ how many modules import this function. This prevents multiple authentication
+ error messages when LANGFUSE_PUBLIC_KEY is not configured.
+
+ Returns:
+ Any: The singleton Langfuse client instance
+ """
+ global _langfuse_client
+
+ if _langfuse_client is None:
+ _langfuse_client = get_client()
+
+ return _langfuse_client
+
+
+# For backward compatibility, provide the client as a module-level variable
+# but only initialize it when first accessed
+def __getattr__(name: str):
+ """Lazy initialization of module-level 'lf' attribute."""
+ if name == "lf":
+ return get_langfuse_client()
+ raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
diff --git a/src/utils/logging.py b/src/utils/logging.py
index b9ceae32..dc9d6ad6 100644
--- a/src/utils/logging.py
+++ b/src/utils/logging.py
@@ -5,21 +5,25 @@ and a conditional observe decorator that only applies when Langfuse is configure
"""
import datetime
-from collections.abc import Callable, Sequence
-from typing import Any, Protocol
+from collections.abc import Callable
+from typing import Any
from rich import box
from rich.console import Console
-from rich.panel import Panel
from rich.table import Table
from rich.tree import Tree
from src.config import settings
-from src.utils.shared_models import ObservationDict, ReasoningResponseWithThinking
+from src.utils.metrics_collector import append_metrics_to_file
+from src.utils.representation import (
+ Representation,
+)
# Global console instance for consistent formatting
console = Console(markup=True)
+COLLECT_METRICS_LOCAL = settings.COLLECT_METRICS_LOCAL
+
def conditional_observe(func: Callable[..., Any]) -> Callable[..., Any]:
"""
@@ -41,80 +45,12 @@ def conditional_observe(func: Callable[..., Any]) -> Callable[..., Any]:
return func
-class ObservationWithContent(Protocol):
- """Protocol for objects with content attribute."""
-
- content: str
-
-
-class ObservationWithConclusion(Protocol):
- """Protocol for objects with conclusion and optional premises."""
-
- conclusion: str
- premises: Sequence[str] | None
-
-
-# Union type for all possible observation types
-ObservationType = (
- str
- | ObservationDict
- | ObservationWithContent
- | ObservationWithConclusion
- | dict[str, Any]
-)
-
# dict[task_name, list[tuple[metric_name, metric_value, metric_unit]]]
accumulated_metrics: dict[str, list[tuple[str, str | int | float, str]]] = {}
-def format_reasoning_response_as_markdown(
- response: ReasoningResponseWithThinking | None,
-) -> str:
- """
- Format a ReasoningResponse object as markdown.
- Args:
- response: ReasoningResponse object or similar structure
- Returns:
- Formatted markdown string
- """
- if not response:
- return "No reasoning response available"
-
- parts: list[str] = []
-
- # Add thinking section if available
- if hasattr(response, "thinking") and response.thinking:
- parts.append("## Thinking\n")
- parts.append(response.thinking.strip())
- parts.append("")
-
- # Add explicit observations
- if hasattr(response, "explicit") and response.explicit:
- parts.append("## Explicit Observations\n")
- for i, obs in enumerate(response.explicit, 1):
- parts.append(f"{i}. {obs}")
- parts.append("")
-
- # Add deductive observations
- if hasattr(response, "deductive") and response.deductive:
- parts.append("## Deductive Observations\n")
- for i, obs in enumerate(response.deductive, 1):
- if hasattr(obs, "conclusion"):
- parts.append(f"{i}. **Conclusion**: {obs.conclusion}")
- if hasattr(obs, "premises") and obs.premises:
- parts.append(" **Premises**:")
- for premise in obs.premises:
- parts.append(f" - {premise}")
- parts.append("")
- else:
- parts.append(f"{i}. {obs}")
- parts.append("")
-
- return "\n".join(parts)
-
-
def format_reasoning_inputs_as_markdown(
- context: ReasoningResponseWithThinking | None,
+ representation: Representation,
history: str,
new_turn: str,
message_created_at: datetime.datetime,
@@ -122,7 +58,7 @@ def format_reasoning_inputs_as_markdown(
"""
Format reasoning inputs as markdown for logging.
Args:
- context: Current context/observations
+ representation: Current/working representation
history: Conversation history
new_turn: New user message
message_created_at: Message timestamp
@@ -137,23 +73,9 @@ def format_reasoning_inputs_as_markdown(
)
parts.append("")
- # Add context if available
- if context:
- parts.append("### Current Context\n")
- if hasattr(context, "explicit") and context.explicit:
- parts.append("**Explicit Observations**:")
- for obs in context.explicit:
- parts.append(f"- {obs}")
- parts.append("")
-
- if hasattr(context, "deductive") and context.deductive:
- parts.append("**Deductive Observations**:")
- for obs in context.deductive:
- if hasattr(obs, "conclusion"):
- parts.append(f"- {obs.conclusion}")
- else:
- parts.append(f"- {obs}")
- parts.append("")
+ # Add representation
+ parts.append("### Current Representation\n")
+ parts.append(representation.format_as_markdown())
# Add history
if history:
@@ -170,50 +92,23 @@ def format_reasoning_inputs_as_markdown(
return "\n".join(parts)
-def log_thinking_panel(
- thinking: str | None,
+def log_representation(
+ representation: Representation,
) -> None:
"""
- Log thinking content in a beautiful panel.
+ Log representation in a tree structure.
Args:
- thinking: Thinking content to display (can be None)
+ representation: Representation to log
"""
- if not thinking:
- return
+ tree = Tree("📊 REPRESENTATION")
- panel = Panel(
- thinking.strip(),
- title="🧠 THINKING",
- title_align="left",
- border_style="blue",
- padding=(1, 2),
- )
+ type_branch = tree.add(f"[bold cyan]EXPLICIT[/] ({len(representation.explicit)})")
+ for i, obs in enumerate(representation.explicit, 1):
+ type_branch.add(f"[dim]{i}.[/] {obs}")
- # Use console.print for immediate output only
- console.print(panel)
- console.print()
-
-
-def log_observations_tree(
- observations: dict[str, list[Any]],
-) -> None:
- """
- Log observations in a tree structure.
- Args:
- observations: Dictionary of observation types and their lists
- """
- tree = Tree("📊 OBSERVATIONS")
-
- for obs_type, obs_list in observations.items():
- if obs_list:
- type_branch = tree.add(
- f"[bold cyan]{obs_type.title()}[/] ({len(obs_list)})"
- )
-
- for i, obs in enumerate(obs_list): # Show all observations
- content = _extract_observation_text(obs)
- truncated = content[:120] + "..." if len(content) > 120 else content
- type_branch.add(f"[dim]{i + 1}.[/] {truncated}")
+ type_branch = tree.add(f"[bold cyan]DEDUCTIVE[/] ({len(representation.deductive)})")
+ for i, obs in enumerate(representation.deductive, 1):
+ type_branch.add(f"[dim]{i}.[/] {obs}")
console.print(tree)
console.print()
@@ -236,16 +131,21 @@ def accumulate_metric(
def log_performance_metrics(
+ task_slug: str,
task_name: str,
metrics: list[tuple[str, str | int | float, str]] | None = None,
title: str = "⚡ PERFORMANCE",
) -> None:
"""
- Log performance metrics in a clean table.
+ Log performance metrics in a clean table and optionally send to global collector.
+
Args:
+ task_slug: Slug of the task that generated these metrics
+ task_name: Name of the task that generated these metrics
metrics: Dictionary of metric names and (value, unit) tuples
title: Table title
"""
+ task_name = f"{task_slug}_{task_name}"
if not accumulated_metrics.get(task_name) and not metrics:
return
if metrics is None:
@@ -253,6 +153,9 @@ def log_performance_metrics(
metrics = accumulated_metrics.get(task_name, []) + metrics
accumulated_metrics[task_name].clear()
+ if COLLECT_METRICS_LOCAL:
+ append_metrics_to_file(task_slug, task_name, metrics)
+
table = Table(
title=f"{title} - {task_name}",
show_header=True,
@@ -276,35 +179,3 @@ def log_performance_metrics(
if metrics:
console.print(table)
console.print()
-
-
-def _extract_observation_text(obs: ObservationType) -> str:
- """Extract text content from various observation types, including premises."""
- if isinstance(obs, str):
- return obs
- elif isinstance(obs, dict):
- # Handle dict-based structured observations first
- if "conclusion" in obs:
- conclusion: str = str(obs["conclusion"])
- premises: list[Any] = list(obs.get("premises", []))
- if premises:
- premises_text = "\n" + "\n".join(f" - {str(p)}" for p in premises)
- return f"{conclusion}{premises_text}"
- return conclusion
- return str(obs.get("content", obs))
- else:
- # Handle object-based observations
- # Use Any type for this branch since we're doing dynamic attribute checking
- obj: Any = obs
- if hasattr(obj, "conclusion"):
- conclusion = str(obj.conclusion)
- if hasattr(obj, "premises") and obj.premises:
- premises_text = "\n" + "\n".join(
- f" - {str(p)}" for p in obj.premises
- )
- return f"{conclusion}{premises_text}"
- return conclusion
- elif hasattr(obj, "content"):
- return str(obj.content)
- else:
- return str(obj)
diff --git a/src/utils/metrics_collector.py b/src/utils/metrics_collector.py
new file mode 100644
index 00000000..8c7cb35e
--- /dev/null
+++ b/src/utils/metrics_collector.py
@@ -0,0 +1,364 @@
+"""
+Global metrics collector for aggregating performance metrics across benchmark runs.
+
+This module provides functionality to collect, aggregate, and export performance
+metrics from deriver and dialectic operations during benchmarking.
+"""
+
+import json
+import statistics
+from datetime import datetime
+from pathlib import Path
+
+from typing_extensions import TypedDict
+
+from src.config import settings
+
+
+class MetricStats(TypedDict):
+ """Statistics for a single metric type."""
+
+ count: int
+ mean: float
+ median: float
+ min: float
+ max: float
+ std_dev: float
+ unit: str
+ raw_values: list[float]
+
+
+class MetricsExport(TypedDict):
+ """Complete metrics export structure."""
+
+ run_id: str
+ start_time: str
+ end_time: str
+ total_tasks: int
+ metrics_by_type: dict[str, list[float]]
+ aggregated_stats: dict[str, MetricStats]
+
+
+class MetricsCollector:
+ """
+ Collects and aggregates performance metrics across multiple deriver runs.
+
+ This collector is designed to work alongside the existing logging system,
+ capturing metrics from individual tasks and providing aggregated statistics
+ for benchmarking analysis.
+ """
+
+ def __init__(self) -> None:
+ """Initialize the metrics collector."""
+ self.run_id: str | None = None
+ self.start_time: datetime | None = None
+ self.end_time: datetime | None = None
+ self.metrics_by_type: dict[str, list[float]] = {}
+ self.task_count: int = 0
+ self.is_collecting: bool = False
+
+ def start_collection(self, run_id: str) -> None:
+ """
+ Initialize metrics collection for a new benchmark run.
+
+ Args:
+ run_id: Unique identifier for this benchmark run
+ """
+ self.run_id = run_id
+ self.start_time = datetime.now()
+ self.end_time = None
+ self.metrics_by_type.clear()
+ self.task_count = 0
+ self.is_collecting = True
+ print(f"📊 Started metrics collection for run: {run_id}")
+
+ def collect_metrics(
+ self, metrics_list: list[tuple[str, str | int | float, str]]
+ ) -> None:
+ """
+ Collect metrics from a completed task.
+
+ Args:
+ metrics_list: List of (metric_name, value, unit) tuples
+ """
+ if not self.is_collecting:
+ return
+
+ self.task_count += 1
+
+ for metric_name, value, unit in metrics_list:
+ # Normalize metric names to be consistent
+ normalized_name = metric_name.lower().replace(" ", "_")
+
+ # skip metrics whose unit is "id" (more in future possibly)
+ if unit in ["id"]:
+ continue
+
+ # Convert value to float for aggregation
+ try:
+ numeric_value = float(value)
+ except (ValueError, TypeError):
+ continue # Skip non-numeric metrics
+
+ # Store the metric with its unit
+ metric_key = f"{normalized_name}_{unit}"
+
+ if metric_key not in self.metrics_by_type:
+ self.metrics_by_type[metric_key] = []
+
+ self.metrics_by_type[metric_key].append(numeric_value)
+
+ def load_from_file(self, filepath: Path) -> None:
+ """
+ Load metrics from a file into this collector.
+ Creates the file if it doesn't exist.
+
+ Args:
+ filepath: Path to the metrics file to load
+ """
+ if not filepath.exists():
+ filepath.parent.mkdir(parents=True, exist_ok=True)
+ filepath.touch()
+ with open(filepath, "w") as f:
+ f.write("")
+
+ file_metrics = load_metrics_from_file(filepath)
+ for _task_name, metrics_list in file_metrics:
+ self.collect_metrics(metrics_list)
+
+ def finalize_collection(self) -> None:
+ """
+ Finalize metrics collection and calculate aggregated statistics.
+ """
+ if not self.is_collecting:
+ return
+
+ # Load any metrics from the file before finalizing
+ metrics_file = get_metrics_file_path()
+ if metrics_file:
+ self.load_from_file(metrics_file)
+
+ self.end_time = datetime.now()
+ self.is_collecting = False
+
+ duration = (
+ (self.end_time - self.start_time).total_seconds()
+ if self.end_time and self.start_time
+ else 0
+ )
+ print(f"📊 Finalized metrics collection for run: {self.run_id}")
+ print(f" Tasks processed: {self.task_count}")
+ print(f" Collection duration: {duration:.2f}s")
+ print(f" Metric types collected: {len(self.metrics_by_type)}")
+
+ def get_aggregated_stats(self) -> dict[str, MetricStats]:
+ """
+ Calculate aggregated statistics for all collected metrics.
+
+ Returns:
+ Dictionary mapping metric names to their statistics
+ """
+ stats: dict[str, MetricStats] = {}
+
+ for metric_key, values in self.metrics_by_type.items():
+ if not values:
+ continue
+
+ # Extract unit from metric key (format: metric_name_unit)
+ parts = metric_key.rsplit("_", 1)
+ if len(parts) == 2:
+ metric_name, unit = parts
+ else:
+ metric_name, unit = metric_key, "unknown"
+
+ # Calculate statistics
+ stats[metric_name] = MetricStats(
+ count=len(values),
+ mean=statistics.mean(values),
+ median=statistics.median(values),
+ min=min(values),
+ max=max(values),
+ std_dev=statistics.stdev(values) if len(values) > 1 else 0.0,
+ unit=unit,
+ raw_values=values.copy(),
+ )
+
+ return stats
+
+ def export_to_json(self, filepath: Path) -> None:
+ """
+ Export collected metrics to a JSON file.
+
+ Args:
+ filepath: Path where the JSON file should be written
+ """
+ if not self.run_id or not self.start_time:
+ raise ValueError("Cannot export metrics - collection was never started")
+
+ # Prepare raw metrics data (without units in keys for cleaner JSON)
+ raw_metrics: dict[str, list[float]] = {}
+ for metric_key, values in self.metrics_by_type.items():
+ # Remove unit suffix for cleaner JSON keys
+ parts = metric_key.rsplit("_", 1)
+ clean_name = parts[0] if len(parts) == 2 else metric_key
+ raw_metrics[clean_name] = values
+
+ export_data: MetricsExport = {
+ "run_id": self.run_id,
+ "start_time": self.start_time.isoformat(),
+ "end_time": self.end_time.isoformat() if self.end_time else "",
+ "total_tasks": self.task_count,
+ "metrics_by_type": raw_metrics,
+ "aggregated_stats": self.get_aggregated_stats(),
+ }
+
+ # Ensure parent directory exists
+ filepath.parent.mkdir(parents=True, exist_ok=True)
+
+ with open(filepath, "w") as f:
+ json.dump(export_data, f, indent=2, default=str)
+
+ print(f"📊 Exported metrics to: {filepath}")
+
+ def print_summary(self) -> None:
+ """
+ Print a summary of collected metrics to the console.
+ """
+ if not self.metrics_by_type:
+ print("📊 No metrics collected")
+ return
+
+ stats = self.get_aggregated_stats()
+
+ print(f"\n{'=' * 80}")
+ print(f"📊 PERFORMANCE METRICS SUMMARY - {self.run_id}")
+ print(f"{'=' * 80}")
+ print(f"Tasks processed: {self.task_count}")
+
+ if self.start_time and self.end_time:
+ duration = (self.end_time - self.start_time).total_seconds()
+ print(f"Collection duration: {duration:.2f}s")
+
+ print("\nAggregated Performance Metrics:")
+ print(
+ f"{'Metric':<40} {'Count':<8} {'Mean':<12} {'Median':<12} {'Min':<12} {'Max':<12} {'Unit'}"
+ )
+ print(
+ f"{'-' * 40} {'-' * 8} {'-' * 12} {'-' * 12} {'-' * 12} {'-' * 12} {'-' * 8}"
+ )
+
+ # Sort metrics by name for consistent display
+ for metric_name in sorted(stats.keys()):
+ stat = stats[metric_name]
+ unit_display = stat["unit"]
+
+ # Format values based on unit
+ if unit_display == "ms":
+ mean_str = f"{stat['mean']:.1f}"
+ median_str = f"{stat['median']:.1f}"
+ min_str = f"{stat['min']:.1f}"
+ max_str = f"{stat['max']:.1f}"
+ elif unit_display == "s":
+ mean_str = f"{stat['mean']:.3f}"
+ median_str = f"{stat['median']:.3f}"
+ min_str = f"{stat['min']:.3f}"
+ max_str = f"{stat['max']:.3f}"
+ else:
+ mean_str = f"{stat['mean']:.2f}"
+ median_str = f"{stat['median']:.2f}"
+ min_str = f"{stat['min']:.2f}"
+ max_str = f"{stat['max']:.2f}"
+
+ print(
+ f"{metric_name:<40} {stat['count']:<8} {mean_str:<12} {median_str:<12} {min_str:<12} {max_str:<12} {unit_display}"
+ )
+
+ print(f"{'=' * 80}")
+
+ def cleanup_collection(self) -> None:
+ """
+ Cleanup metrics collection.
+ """
+ self.is_collecting = False
+ self.end_time = datetime.now()
+ # delete the metrics file
+ metrics_file = get_metrics_file_path()
+ if metrics_file:
+ metrics_file.unlink()
+
+
+def get_metrics_file_path() -> Path | None:
+ """Get the current metrics file path."""
+ env_path = settings.LOCAL_METRICS_FILE
+ if env_path:
+ _metrics_file_path = Path(env_path)
+ return _metrics_file_path
+
+ return None
+
+
+def append_metrics_to_file(
+ task_slug: str,
+ task_name: str,
+ metrics_list: list[tuple[str, str | int | float, str]],
+) -> None:
+ """
+ Append metrics to the shared metrics file for cross-process collection.
+
+ Args:
+ task_slug: Slug of the task that generated these metrics
+ task_name: Name of the task that generated these metrics
+ metrics_list: List of (metric_name, value, unit) tuples
+ """
+ metrics_file = get_metrics_file_path()
+ if not metrics_file:
+ return
+
+ import fcntl
+ import time
+
+ # Prepare metrics data
+ timestamp = time.time()
+ metrics_entry = {
+ "timestamp": timestamp,
+ "task_name": f"{task_slug}_{task_name}",
+ "metrics": [
+ {"name": f"{task_slug}_{name}", "value": value, "unit": unit}
+ for name, value, unit in metrics_list
+ ],
+ }
+
+ # Use file locking to handle concurrent writes from multiple processes
+ with open(metrics_file, "a") as f:
+ fcntl.flock(f.fileno(), fcntl.LOCK_EX)
+ f.write(json.dumps(metrics_entry) + "\n")
+ fcntl.flock(f.fileno(), fcntl.LOCK_UN)
+
+
+def load_metrics_from_file(
+ filepath: Path,
+) -> list[tuple[str, list[tuple[str, str | int | float, str]]]]:
+ """
+ Load all metrics from the shared metrics file.
+
+ Args:
+ filepath: Path to the metrics file
+
+ Returns:
+ List of (task_name, metrics_list) tuples
+ """
+ if not filepath.exists():
+ return []
+
+ all_metrics: list[tuple[str, list[tuple[str, str | int | float, str]]]] = []
+ with open(filepath) as f:
+ for line in f:
+ if line.strip():
+ entry = json.loads(line.strip())
+ task_name = entry["task_name"]
+ metrics_list = [
+ (m["name"], m["value"], m["unit"]) for m in entry["metrics"]
+ ]
+ all_metrics.append((task_name, metrics_list))
+
+ return all_metrics
diff --git a/src/utils/peer_card.py b/src/utils/peer_card.py
new file mode 100644
index 00000000..b528cee4
--- /dev/null
+++ b/src/utils/peer_card.py
@@ -0,0 +1,25 @@
+"""
+Shared Pydantic models used by both dialectic and deriver modules.
+"""
+
+from __future__ import annotations
+
+from pydantic import BaseModel, Field
+
+
+class PeerCardQuery(BaseModel):
+ """
+ Model for peer card query generation responses.
+
+ Contains the new peer card, or None if there are no new key observations.
+ The notes field is just a place for stupid models to dump useless info.
+ """
+
+ card: list[str] | None = Field(
+ default=None,
+ description="Generated peer card as list of strings. None if no new useful biographical observations.",
+ )
+ notes: str | None = Field(
+ default=None,
+ description="Optional additional notes from the model; may include non-actionable info.",
+ )
diff --git a/src/deriver/queue_payload.py b/src/utils/queue_payload.py
similarity index 78%
rename from src/deriver/queue_payload.py
rename to src/utils/queue_payload.py
index c3c111e0..e4d4a964 100644
--- a/src/deriver/queue_payload.py
+++ b/src/utils/queue_payload.py
@@ -18,8 +18,8 @@ class RepresentationPayload(BasePayload):
session_name: str
message_id: int
content: str
- sender_name: str
- target_name: str
+ observer: str
+ observed: str
created_at: datetime
@@ -50,6 +50,16 @@ class WebhookPayload(BasePayload):
data: dict[str, Any]
+class DreamPayload(BasePayload):
+ """Payload for dream tasks."""
+
+ task_type: Literal["dream"] = "dream"
+ workspace_name: str
+ dream_type: Literal["consolidate"] = "consolidate"
+ observer: str
+ observed: str
+
+
def create_webhook_payload(
workspace_name: str,
event_type: str,
@@ -60,12 +70,28 @@ def create_webhook_payload(
).model_dump(mode="json")
+def create_dream_payload(
+ workspace_name: str,
+ dream_type: Literal["consolidate"] = "consolidate",
+ *,
+ observer: str,
+ observed: str,
+) -> dict[str, Any]:
+ return DreamPayload(
+ workspace_name=workspace_name,
+ dream_type=dream_type,
+ observer=observer,
+ observed=observed,
+ ).model_dump(mode="json")
+
+
def create_payload(
message: dict[str, Any],
task_type: Literal["representation", "summary"],
- sender_name: str | None = None,
- target_name: str | None = None,
message_seq_in_session: int | None = None,
+ *,
+ observer: str | None = None,
+ observed: str | None = None,
) -> dict[str, Any]:
"""
Create a processed payload from a message for queue processing.
@@ -73,8 +99,8 @@ def create_payload(
Args:
message: The original message dictionary
task_type: Type of task ('representation' or 'summary')
- sender_name: Name of the message sender (required for representation tasks)
- target_name: Name of the observer peer (required for representation tasks)
+ observer: Name of the observer peer (required for representation tasks)
+ observed: Name of the observed peer (*always* the peer who sent the message) (required for representation tasks)
message_seq_in_session: Required for summary tasks, must be None for representation
Returns:
@@ -108,20 +134,20 @@ def create_payload(
if not isinstance(created_at, datetime):
raise TypeError("created_at must be a datetime object")
- if sender_name is None:
- raise ValueError("sender_name is required for representation tasks")
+ if observer is None:
+ raise ValueError("observer is required for representation tasks")
- if target_name is None:
- raise ValueError("target_name is required for representation tasks")
+ if observed is None:
+ raise ValueError("observed is required for representation tasks")
validated_payload = RepresentationPayload(
content=content,
workspace_name=workspace_name,
- sender_name=sender_name,
- target_name=target_name,
session_name=session_name,
message_id=message_id,
created_at=created_at,
+ observer=observer,
+ observed=observed,
)
elif task_type == "summary":
if message_seq_in_session is None:
diff --git a/src/utils/representation.py b/src/utils/representation.py
new file mode 100644
index 00000000..30ebe71d
--- /dev/null
+++ b/src/utils/representation.py
@@ -0,0 +1,337 @@
+from collections.abc import Sequence
+from datetime import datetime
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+from src import models
+from src.utils.formatting import parse_datetime_iso
+
+
+class ObservationMetadata(BaseModel):
+ created_at: datetime
+ message_ids: list[tuple[int, int]]
+ session_name: str
+
+
+class ExplicitObservationBase(BaseModel):
+ content: str = Field(description="The explicit observation")
+
+
+class DeductiveObservationBase(BaseModel):
+ premises: list[str] = Field(
+ description="Supporting premises or evidence for this conclusion",
+ default_factory=list,
+ )
+ conclusion: str = Field(description="The deductive conclusion")
+
+
+class PromptRepresentation(BaseModel):
+ """
+ The representation format that is used when getting structured output from an LLM.
+ """
+
+ explicit: list[ExplicitObservationBase] = Field(
+ description="Facts LITERALLY stated by the user - direct quotes or clear paraphrases only, no interpretation or inference. Example: ['The user is 25 years old', 'The user has a dog named Rover']",
+ default_factory=list,
+ )
+ deductive: list[DeductiveObservationBase] = Field(
+ description="Conclusions that MUST be true given explicit facts and premises - strict logical necessities. Each deduction should have premises and a single conclusion.",
+ default_factory=list,
+ )
+
+
+class ExplicitObservation(ExplicitObservationBase, ObservationMetadata):
+ """Explicit observation with content and metadata."""
+
+ def __str__(self) -> str:
+ return f"[{self.created_at.replace(microsecond=0)}] {self.content}"
+
+ def __hash__(self) -> int:
+ """
+ Make ExplicitObservation hashable for use in sets.
+ """
+ return hash((self.content, self.created_at, self.session_name))
+
+ def __eq__(self, other: object) -> bool:
+ """
+ Define equality for ExplicitObservation objects.
+ Two observations are equal if all their fields match.
+ """
+ if not isinstance(other, ExplicitObservation):
+ return False
+ return (
+ self.content == other.content
+ and self.created_at == other.created_at
+ and self.session_name == other.session_name
+ )
+
+
+class DeductiveObservation(DeductiveObservationBase, ObservationMetadata):
+ """Deductive observation with multiple premises and one conclusion, plus metadata."""
+
+ def __str__(self) -> str:
+ premises_text = "\n".join(f" - {premise}" for premise in self.premises)
+ return f"[{self.created_at.replace(microsecond=0)}] {self.conclusion}\n{premises_text}"
+
+ def str_no_timestamps(self) -> str:
+ premises_text = "\n".join(f" - {premise}" for premise in self.premises)
+ return f"{self.conclusion}\n{premises_text}"
+
+ def __hash__(self) -> int:
+ """
+ Make DeductiveObservation hashable for use in sets. NOTE: premises are not included in the hash.
+ """
+ return hash((self.conclusion, self.created_at, self.session_name))
+
+ def __eq__(self, other: object) -> bool:
+ """
+ Define equality for DeductiveObservation objects.
+ Two observations are equal if all their fields match -- NOTE: premises are not included in the equality check.
+ """
+ if not isinstance(other, DeductiveObservation):
+ return False
+ return (
+ self.conclusion == other.conclusion
+ and self.created_at == other.created_at
+ and self.session_name == other.session_name
+ )
+
+
+class Representation(BaseModel):
+ """
+ A Representation is a traversable and diffable map of observations.
+ At the base, we have a list of explicit observations, derived from a peer's messages.
+
+ From there, deductive observations can be made by establishing logical relationships between explicit observations.
+
+ In the future, we can add more levels of reasoning on top of these.
+
+ All of a peer's observations are stored as documents in a collection. These documents can be queried in various ways
+ to produce this Representation object.
+
+ Additionally, a "working representation" is a version of this data structure representing the most recent observations
+ within a single session.
+
+ A representation can have a maximum number of observations, which is applied individually to each level of reasoning.
+ If a maximum is set, observations are added and removed in FIFO order.
+ """
+
+ explicit: list[ExplicitObservation] = Field(
+ description="Facts LITERALLY stated by the user - direct quotes or clear paraphrases only, no interpretation or inference. Example: ['The user is 25 years old', 'The user has a dog']",
+ default_factory=list,
+ )
+ deductive: list[DeductiveObservation] = Field(
+ description="Conclusions that MUST be true given explicit facts and premises - strict logical necessities. Each deduction should have premises and a single conclusion.",
+ default_factory=list,
+ )
+
+ def is_empty(self) -> bool:
+ """
+ Check if the representation is empty.
+ """
+ return len(self.explicit) == 0 and len(self.deductive) == 0
+
+ def diff_representation(self, other: "Representation") -> "Representation":
+ """
+ Given this and another representation, return a new representation with only observations that are unique to the other.
+ Note that this only removes literal duplicates, not semantically equivalent ones.
+ """
+ diff = Representation()
+ diff.explicit = [o for o in other.explicit if o not in self.explicit]
+ diff.deductive = [o for o in other.deductive if o not in self.deductive]
+ return diff
+
+ def merge_representation(
+ self, other: "Representation", max_observations: int | None = None
+ ):
+ """
+ Merge another representation object into this one.
+ This will automatically deduplicate explicit and deductive observations.
+ This *preserves order* of observations so that they retain FIFO order.
+
+ NOTE: observations with the *same* timestamp will not have order preserved.
+ That's fine though, because they are from the same timestamp...
+ """
+ # removing duplicates by going list->set->list
+ self.explicit = list(set(self.explicit + other.explicit))
+ self.deductive = list(set(self.deductive + other.deductive))
+ # sort by created_at
+ self.explicit.sort(key=lambda x: x.created_at)
+ self.deductive.sort(key=lambda x: x.created_at)
+
+ if max_observations:
+ self.explicit = self.explicit[-max_observations:]
+ self.deductive = self.deductive[-max_observations:]
+
+ def __str__(self) -> str:
+ """
+ Format representation into a clean, readable string for LLM prompts.
+ NOTE: we always strip subsecond precision from the timestamps.
+
+ Returns:
+ Formatted string with clear sections and bullet points including temporal metadata
+ Example:
+ EXPLICIT:
+ 1. [2025-01-01 12:00:00] The user has a dog named Rover
+ 2. [2025-01-01 12:01:00] The user's dog is 5 years old
+ 3. [2025-01-01 12:05:00] The user is 25 years old
+ DEDUCTIVE:
+ 1. [2025-01-01 12:01:00] Rover is 5 years old
+ - The user has a dog named Rover
+ - The user's dog is 5 years old
+
+ """
+
+ parts: list[str] = []
+
+ parts.append("EXPLICIT:\n")
+ for i, observation in enumerate(self.explicit, 1):
+ parts.append(f"{i}. {observation}")
+ parts.append("")
+
+ parts.append("DEDUCTIVE:\n")
+ for i, observation in enumerate(self.deductive, 1):
+ parts.append(f"{i}. {observation}")
+ parts.append("")
+
+ return "\n".join(parts)
+
+ def str_no_timestamps(self) -> str:
+ """
+ Format representation into a clean, readable string for LLM prompts... but without timestamps.
+
+ Returns:
+ Formatted string with clear sections and bullet points including temporal metadata
+ Example:
+ EXPLICIT:
+ 1. The user has a dog named Rover
+ 2. The user's dog is 5 years old
+ 3. The user is 25 years old
+ DEDUCTIVE:
+ 1. Rover is 5 years old
+ - The user has a dog named Rover
+ - The user's dog is 5 years old
+
+ """
+ parts: list[str] = []
+
+ parts.append("EXPLICIT:\n")
+ for i, observation in enumerate(self.explicit, 1):
+ parts.append(f"{i}. {observation.content}")
+ parts.append("")
+
+ parts.append("DEDUCTIVE:\n")
+ for i, observation in enumerate(self.deductive, 1):
+ parts.append(f"{i}. {observation.str_no_timestamps()}")
+ parts.append("")
+
+ return "\n".join(parts)
+
+ def format_as_markdown(self) -> str:
+ """
+ Format a Representation object as markdown.
+ NOTE: we always strip subsecond precision from the timestamps.
+
+ Returns:
+ Formatted markdown string
+ """
+
+ parts: list[str] = []
+
+ # Add explicit observations
+ parts.append("## Explicit Observations\n")
+ for i, obs in enumerate(self.explicit, 1):
+ parts.append(f"{i}. {obs}")
+ parts.append("")
+
+ # Add deductive observations
+ parts.append("## Deductive Observations\n")
+ for i, obs in enumerate(self.deductive, 1):
+ parts.append(f"{i}. **Conclusion**: {obs.conclusion}")
+ if obs.premises:
+ parts.append(" **Premises**:")
+ for premise in obs.premises:
+ parts.append(f" - {premise}")
+ parts.append("")
+ parts.append("")
+
+ return "\n".join(parts)
+
+ @classmethod
+ def from_documents(cls, documents: Sequence[models.Document]) -> "Representation":
+ return cls(
+ explicit=[
+ ExplicitObservation(
+ created_at=_safe_datetime_from_metadata(
+ doc.internal_metadata, doc.created_at
+ ),
+ content=doc.content,
+ message_ids=doc.internal_metadata.get("message_ids", [(0, 0)]),
+ session_name=doc.session_name,
+ )
+ for doc in documents
+ if doc.internal_metadata.get("level") == "explicit"
+ ],
+ deductive=[
+ DeductiveObservation(
+ created_at=_safe_datetime_from_metadata(
+ doc.internal_metadata, doc.created_at
+ ),
+ conclusion=doc.content,
+ message_ids=doc.internal_metadata.get("message_ids", [(0, 0)]),
+ session_name=doc.session_name,
+ premises=doc.internal_metadata.get("premises", []),
+ )
+ for doc in documents
+ if doc.internal_metadata.get("level") == "deductive"
+ ],
+ )
+
+ @classmethod
+ def from_prompt_representation(
+ cls,
+ prompt_representation: "PromptRepresentation",
+ message_ids: tuple[int, int],
+ session_name: str,
+ created_at: datetime,
+ ) -> "Representation":
+ return cls(
+ explicit=[
+ ExplicitObservation(
+ content=e.content,
+ created_at=created_at,
+ message_ids=[message_ids],
+ session_name=session_name,
+ )
+ for e in prompt_representation.explicit
+ ],
+ deductive=[
+ DeductiveObservation(
+ conclusion=d.conclusion,
+ created_at=created_at,
+ message_ids=[message_ids],
+ session_name=session_name,
+ premises=d.premises,
+ )
+ for d in prompt_representation.deductive
+ ],
+ )
+
+
+def _safe_datetime_from_metadata(
+ internal_metadata: dict[str, Any], fallback_datetime: datetime
+) -> datetime:
+ message_created_at = internal_metadata.get("message_created_at")
+ if message_created_at is None:
+ return fallback_datetime.replace(microsecond=0)
+
+ if isinstance(message_created_at, str):
+ try:
+ return parse_datetime_iso(message_created_at)
+ except ValueError:
+ return fallback_datetime.replace(microsecond=0)
+
+ if isinstance(message_created_at, datetime):
+ return message_created_at.replace(microsecond=0)
+ return fallback_datetime.replace(microsecond=0)
diff --git a/src/utils/shared_models.py b/src/utils/shared_models.py
deleted file mode 100644
index 70f83a76..00000000
--- a/src/utils/shared_models.py
+++ /dev/null
@@ -1,205 +0,0 @@
-"""
-Shared Pydantic models used by both dialectic and deriver modules.
-"""
-
-from __future__ import annotations
-
-from datetime import datetime, timezone
-from enum import Enum
-
-from pydantic import BaseModel, Field
-from typing_extensions import TypedDict
-
-
-class ReasoningLevel(str, Enum):
- EXPLICIT = "explicit"
- DEDUCTIVE = "deductive"
-
-
-class ObservationMetadata(BaseModel):
- """Actual metadata structure from the database."""
-
- session_context: str = ""
- summary_id: str = ""
- message_id: str | None = None
- level: str | None = None
- session_name: str | None = None
- premises: list[str] = Field(default_factory=list)
-
-
-class Observation(BaseModel):
- """Observation matching the actual document structure."""
-
- content: str
- metadata: ObservationMetadata = Field(default_factory=ObservationMetadata)
- created_at: datetime
-
- def __str__(self) -> str:
- return self.content
-
-
-class DeductiveObservation(BaseModel):
- """Deductive observation with multiple premises and one conclusion."""
-
- premises: list[str] = Field(
- description="Supporting premises or evidence for this conclusion",
- default_factory=list,
- )
- conclusion: str = Field(description="The deductive conclusion")
-
-
-class UnifiedObservation(BaseModel):
- """Unified observation model with conclusion and optional premises.
-
- This model separates the core observation (conclusion) from its supporting
- evidence (premises), enabling proper embedding generation from conclusions
- while preserving premise information in metadata.
- """
-
- conclusion: str = Field(description="The actual observation content")
- premises: list[str] = Field(
- description="Optional supporting premises or evidence", default_factory=list
- )
- level: str | None = Field(
- description="Reasoning level (explicit, deductive)", default=None
- )
-
- @property
- def has_premises(self) -> bool:
- """Check if this observation has premises."""
- return len(self.premises) > 0
-
- def to_deductive_observation(self) -> DeductiveObservation:
- """Convert to DeductiveObservation for backward compatibility."""
- return DeductiveObservation(conclusion=self.conclusion, premises=self.premises)
-
- @classmethod
- def from_deductive_observation(
- cls, deductive_obs: DeductiveObservation
- ) -> UnifiedObservation:
- """Create from DeductiveObservation."""
- return cls(conclusion=deductive_obs.conclusion, premises=deductive_obs.premises)
-
- @classmethod
- def from_string(
- cls, observation: str, level: str | None = None
- ) -> UnifiedObservation:
- """Create from simple string observation (no premises)."""
- return cls(conclusion=observation, level=level)
-
-
-class ReasoningResponse(BaseModel):
- """Reasoning response with explicit and deductive observation types."""
-
- explicit: list[str] = Field(
- description="Facts LITERALLY stated by the user - direct quotes or clear paraphrases only, no interpretation or inference. Example: ['The user is 25 years old', 'The user has a dog']",
- default_factory=list,
- )
- deductive: list[DeductiveObservation] = Field(
- description="Conclusions that MUST be true given explicit facts and premises - strict logical necessities. Each deduction should have premises and a single conclusion.",
- default_factory=list,
- )
-
-
-class ReasoningResponseWithThinking(ReasoningResponse):
- thinking: str | None = Field(
- description="Critical thinking about what it means to do explicit and deductive reasoning and how to apply it here",
- default=None,
- )
-
-
-class ObservationContext(BaseModel):
- """Type-safe context container."""
-
- thinking: str | None = Field(default=None)
- explicit: list[Observation] = Field(default_factory=list)
- deductive: list[Observation] = Field(default_factory=list)
-
- @property
- def all_observations(self) -> list[Observation]:
- return self.explicit + self.deductive
-
- def get_by_level(self, level: ReasoningLevel) -> list[Observation]:
- return getattr(self, level.value)
-
- def add_observation(self, observation: Observation, level: ReasoningLevel) -> None:
- getattr(self, level.value).append(observation)
-
- @classmethod
- def from_reasoning_response(
- cls,
- response: ReasoningResponse,
- base_metadata: ObservationMetadata | None = None,
- ) -> ObservationContext:
- """Create ObservationContext from ReasoningResponse."""
- context = cls()
-
- # Add thinking trace if available
- context.thinking = getattr(response, "thinking", None)
-
- # Add explicit observations
- for conclusion in response.explicit:
- explicit_metadata: ObservationMetadata = (
- base_metadata.model_copy() if base_metadata else ObservationMetadata()
- )
- explicit_metadata.level = "explicit"
-
- obs = Observation(
- content=conclusion,
- metadata=explicit_metadata,
- created_at=datetime.now(timezone.utc),
- )
- context.add_observation(obs, ReasoningLevel.EXPLICIT)
-
- # Add deductive observations
- for level_name in ["deductive"]:
- level = ReasoningLevel(level_name)
- structured_obs_list: list[DeductiveObservation] = getattr(
- response, level_name
- )
-
- for structured_obs in structured_obs_list:
- deductive_metadata: ObservationMetadata = (
- base_metadata.model_copy()
- if base_metadata
- else ObservationMetadata()
- )
- deductive_metadata.level = level_name
- deductive_metadata.premises = structured_obs.premises
- obs = Observation(
- content=structured_obs.conclusion,
- metadata=deductive_metadata,
- created_at=datetime.now(timezone.utc),
- )
- context.add_observation(obs, level)
-
- return context
-
-
-class SemanticQueries(BaseModel):
- """Model for semantic query generation responses."""
-
- queries: list[str] = Field(
- description="List of semantic search queries to retrieve relevant observations"
- )
-
-
-class ObservationDict(TypedDict, total=False):
- """Type definition for observation dictionary structures."""
-
- conclusion: str
- content: str
- premises: list[str]
- created_at: str
-
-
-class PeerCardQuery(BaseModel):
- """
- Model for peer card query generation responses.
-
- Contains the new peer card, or None if there are no new key observations.
- The notes field is just a place for stupid models to dump useless info.
- """
-
- card: list[str] | None
- notes: str | None
diff --git a/src/utils/summarizer.py b/src/utils/summarizer.py
index 6b7b2cf3..e99bd14b 100644
--- a/src/utils/summarizer.py
+++ b/src/utils/summarizer.py
@@ -3,10 +3,10 @@ import logging
import time
from enum import Enum
from inspect import cleandoc as c
+from typing import TypedDict
from sqlalchemy import update
from sqlalchemy.ext.asyncio import AsyncSession
-from typing_extensions import TypedDict
from src import schemas
from src.config import settings
@@ -640,7 +640,7 @@ async def get_session_context(
messages_start_id = latest_short_summary["message_id"]
else:
logger.warning(
- "No summary available for get_context call with token limit %s, returning empty string. long_summary_len: %s, short_summary_len: %s",
+ "No summary available for get_context call with token limit %s, returning empty string. Normal if brand-new session. long_summary_len: %s, short_summary_len: %s",
token_limit,
long_len,
short_len,
diff --git a/src/utils/tokens.py b/src/utils/tokens.py
new file mode 100644
index 00000000..2af439c9
--- /dev/null
+++ b/src/utils/tokens.py
@@ -0,0 +1,15 @@
+import tiktoken
+
+tokenizer = tiktoken.get_encoding("cl100k_base")
+
+
+def estimate_tokens(text: str | list[str] | None) -> int:
+ """Estimate token count using tiktoken for text or list of strings."""
+ if not text:
+ return 0
+ if isinstance(text, list):
+ text = "\n".join(text)
+ try:
+ return len(tokenizer.encode(text))
+ except Exception:
+ return len(text) // 4
diff --git a/src/utils/types.py b/src/utils/types.py
index 494f2225..317e545a 100644
--- a/src/utils/types.py
+++ b/src/utils/types.py
@@ -1,3 +1,3 @@
from typing import Literal
-SupportedProviders = Literal["anthropic", "openai", "google", "groq", "custom"]
+SupportedProviders = Literal["anthropic", "openai", "google", "groq", "custom", "vllm"]
diff --git a/src/utils/work_unit.py b/src/utils/work_unit.py
new file mode 100644
index 00000000..d7fd12bb
--- /dev/null
+++ b/src/utils/work_unit.py
@@ -0,0 +1,110 @@
+"""Work unit utility functions for generating and parsing work unit keys."""
+
+from typing import Any
+
+from pydantic import BaseModel
+
+
+class ParsedWorkUnit(BaseModel):
+ """Parsed work unit components."""
+
+ task_type: str
+ workspace_name: str
+ session_name: str | None
+ observer: str | None
+ observed: str | None
+
+
+def get_work_unit_key(payload: dict[str, Any] | ParsedWorkUnit) -> str:
+ """
+ Generate a work unit key for a given task type, workspace name, and event type.
+
+ Args:
+ payload: Dictionary containing work unit information
+
+ Returns:
+ Formatted work unit key string
+
+ Raises:
+ ValueError: If required fields are missing or task type is invalid
+ """
+ if isinstance(payload, ParsedWorkUnit):
+ payload = payload.model_dump()
+
+ workspace_name: str | None = payload.get("workspace_name")
+ task_type: str | None = payload.get("task_type")
+ if not workspace_name or not task_type:
+ raise ValueError(
+ "workspace_name and task_type are required to generate a work_unit_key"
+ )
+
+ if task_type in ["representation", "summary", "dream"]:
+ observer = payload.get("observer", "None")
+ observed = payload.get("observed", "None")
+ session_name = payload.get("session_name", "None")
+ if task_type == "dream":
+ return f"{task_type}:{workspace_name}:{observer}:{observed}"
+ return f"{task_type}:{workspace_name}:{session_name}:{observer}:{observed}"
+
+ if task_type == "webhook":
+ return f"webhook:{workspace_name}"
+
+ raise ValueError(f"Invalid task type: {task_type}")
+
+
+def parse_work_unit_key(work_unit_key: str) -> ParsedWorkUnit:
+ """
+ Parse a work unit key to extract its components.
+
+ Args:
+ work_unit_key: The work unit key string to parse
+
+ Returns:
+ ParsedWorkUnit with extracted components
+
+ Raises:
+ ValueError: If the work unit key format is invalid
+ """
+ parts = work_unit_key.split(":")
+ task_type = parts[0]
+
+ if task_type in ["representation", "summary"]:
+ if len(parts) != 5:
+ raise ValueError(
+ f"Invalid work_unit_key format for task_type {task_type}: {work_unit_key}"
+ )
+ return ParsedWorkUnit(
+ task_type=task_type,
+ workspace_name=parts[1],
+ session_name=parts[2],
+ observer=parts[3],
+ observed=parts[4],
+ )
+
+ if task_type == "dream":
+ if len(parts) != 4:
+ raise ValueError(
+ f"Invalid work_unit_key format for task_type {task_type}: {work_unit_key}"
+ )
+ return ParsedWorkUnit(
+ task_type=task_type,
+ workspace_name=parts[1],
+ session_name=None,
+ observer=parts[2],
+ observed=parts[3],
+ )
+
+ if task_type == "webhook":
+ if len(parts) != 2:
+ raise ValueError(
+ f"Invalid work_unit_key format for task_type {task_type}: {work_unit_key}"
+ )
+ return ParsedWorkUnit(
+ task_type=task_type,
+ workspace_name=parts[1],
+ session_name=None,
+ observer=None,
+ observed=None,
+ )
+
+ raise ValueError(f"Invalid task type in work_unit_key: {task_type}")
diff --git a/src/webhooks/events.py b/src/webhooks/events.py
index 236521ae..f942bb5d 100644
--- a/src/webhooks/events.py
+++ b/src/webhooks/events.py
@@ -5,9 +5,9 @@ from typing import Literal
from pydantic import BaseModel
from src.dependencies import tracked_db
-from src.deriver.queue_payload import create_webhook_payload
-from src.deriver.utils import get_work_unit_key
from src.models import QueueItem
+from src.utils.queue_payload import create_webhook_payload
+from src.utils.work_unit import get_work_unit_key
logger = logging.getLogger(__name__)
@@ -29,8 +29,8 @@ class QueueEmptyEvent(BaseWebhookEvent):
type: Literal[WebhookEventType.QUEUE_EMPTY] = WebhookEventType.QUEUE_EMPTY
queue_type: str
session_id: str | None = None
- sender_name: str | None = None
- observer_name: str | None = None
+ observer: str | None = None
+ observed: str | None = None
class TestEvent(BaseWebhookEvent):
@@ -60,7 +60,10 @@ async def publish_webhook_event(event: WebhookEvent) -> None:
async with tracked_db("publish_webhook_event") as db:
queue_item = QueueItem(
work_unit_key=get_work_unit_key(
- "webhook", {"workspace_name": event.workspace_id}
+ {
+ "task_type": "webhook",
+ "workspace_name": event.workspace_id,
+ }
),
payload=payload,
session_id=None,
diff --git a/src/webhooks/webhook_delivery.py b/src/webhooks/webhook_delivery.py
index 96f9311a..4573c79c 100644
--- a/src/webhooks/webhook_delivery.py
+++ b/src/webhooks/webhook_delivery.py
@@ -9,8 +9,8 @@ from sqlalchemy.ext.asyncio import AsyncSession
from src.config import settings
from src.crud.webhook import list_webhook_endpoints
-from src.deriver.queue_payload import WebhookPayload
from src.utils.formatting import utc_now_iso
+from src.utils.queue_payload import WebhookPayload
logger = logging.getLogger(__name__)
diff --git a/tests/bench/.gitignore b/tests/bench/.gitignore
new file mode 100644
index 00000000..e8d34dc6
--- /dev/null
+++ b/tests/bench/.gitignore
@@ -0,0 +1,3 @@
+longmemeval_data
+eval_results
+perf_metrics
diff --git a/tests/bench/longmem.py b/tests/bench/longmem.py
new file mode 100644
index 00000000..05d2b63e
--- /dev/null
+++ b/tests/bench/longmem.py
@@ -0,0 +1,1151 @@
+"""
+Honcho LongMemEval Test Runner
+
+A script that executes longmemeval tests against a running Honcho instance.
+This script:
+1. Loads longmemeval test definitions from JSON files
+2. Creates a workspace for each question (using question_id and question_type)
+3. Creates sessions with haystack conversations
+4. Adds the answer session if present
+5. Waits for the deriver queue to be empty
+6. Executes the question and judges the response using an LLM
+
+## To use
+
+0. Set up env:
+```
+uv sync
+source .venv/bin/activate
+```
+NOTE: you may create a .env file in this directory to customize honcho config. The harness will print the config it is using.
+
+1. Run the test harness:
+```
+python -m tests.bench.harness
+```
+
+2. Choose a test file:
+ should be formatted as longmemeval_data.
+ see: https://github.com/xiaowu0162/LongMemEval?tab=readme-ov-file
+ or: https://huggingface.co/datasets/xiaowu0162/longmemeval
+
+3. Run this file with a selected test file:
+```
+python -m tests.bench.longmem --test-file tests/bench/longmemeval_data/longmemeval_oracle.json
+```
+
+Optional arguments:
+```
+--anthropic-api-key: Anthropic API key for response judging (can be set in .env as LLM_ANTHROPIC_API_KEY or provided as an argument)
+--timeout: Timeout for deriver queue to empty in seconds (default: 10 minutes)
+--honcho-url: URL of the running Honcho instance (default: http://localhost:8000)
+--batch-size: Number of questions to run concurrently in each batch (default: 10)
+--json-output: Path to write JSON summary results for analytics (if not provided, creates timestamped file in tests/bench/eval_results)
+--merge-sessions: Merge all sessions within a question into a single session (default: False)
+```
+
+## Other notes
+- Judge is Claude Sonnet 4
+- If processing lots of data, set timeout very high or all will be lost
+"""
+
+import argparse
+import asyncio
+import json
+import logging
+import os
+import time
+from datetime import datetime
+from pathlib import Path
+from typing import Any
+
+import tiktoken
+from anthropic import AsyncAnthropic
+from dotenv import load_dotenv
+from honcho import AsyncHoncho
+from honcho.async_client.session import SessionPeerConfig
+from honcho_core.types.workspaces.sessions.message_create_param import (
+ MessageCreateParam,
+)
+from typing_extensions import TypedDict
+
+from src.config import settings
+from src.utils.metrics_collector import MetricsCollector
+
+load_dotenv()
+
+
+class SessionResult(TypedDict):
+ """Type definition for session creation results."""
+
+ name: str
+ message_count: int
+
+
+class QueryResult(TypedDict):
+ """Type definition for query execution results."""
+
+ question: str
+ expected_answer: str
+ actual_response: str
+ judgment: dict[str, Any]
+ token_efficiency: dict[str, Any] | None
+
+
+class TestResult(TypedDict):
+ """Type definition for test execution results."""
+
+ question_id: str
+ question_type: str
+ workspace_id: str
+ sessions_created: list[SessionResult]
+ query_executed: QueryResult | None
+ passed: bool
+ error: str | None
+ start_time: float
+ end_time: float
+ duration_seconds: float
+ output_lines: list[str]
+
+
+class LongMemEvalRunner:
+ """
+ Executes longmemeval JSON tests against a Honcho instance.
+ """
+
+ def __init__(
+ self,
+ honcho_url: str = "http://localhost:8000",
+ anthropic_api_key: str | None = None,
+ timeout_seconds: int | None = None,
+ merge_sessions: bool = False,
+ ):
+ """
+ Initialize the test runner.
+
+ Args:
+ honcho_url: URL of the running Honcho instance
+ anthropic_api_key: Anthropic API key for judging responses
+ timeout_seconds: Timeout for deriver queue in seconds
+ merge_sessions: If True, merge all sessions within a question into one session
+ """
+ self.honcho_url: str = honcho_url
+ self.anthropic_api_key: str | None = anthropic_api_key
+ self.timeout_seconds: int = (
+ timeout_seconds if timeout_seconds is not None else 10000
+ )
+ self.merge_sessions: bool = merge_sessions
+
+ # Initialize metrics collector
+ self.metrics_collector: MetricsCollector = MetricsCollector()
+ self.metrics_collector.start_collection(
+ f"longmem_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+ )
+
+ # Configure logging
+ logging.basicConfig(
+ level=logging.WARNING, format="%(asctime)s - %(levelname)s - %(message)s"
+ )
+ self.logger: logging.Logger = logging.getLogger(__name__)
+
+ # Suppress HTTP request logs from the Honcho SDK
+ logging.getLogger("httpx").setLevel(logging.ERROR)
+ logging.getLogger("httpcore").setLevel(logging.ERROR)
+
+ if self.anthropic_api_key:
+ self.anthropic_client: AsyncAnthropic = AsyncAnthropic(
+ api_key=self.anthropic_api_key
+ )
+ else:
+ api_key = os.getenv("LLM_ANTHROPIC_API_KEY")
+ if not api_key:
+ raise ValueError("LLM_ANTHROPIC_API_KEY is not set")
+ self.anthropic_client = AsyncAnthropic(api_key=api_key)
+
+ def _format_duration(self, total_seconds: float) -> str:
+ """Format a duration in seconds into a human-readable string.
+
+ If the duration is at least one minute, this returns a string in the
+ form "XmYYs" with zero-padded seconds. Otherwise, it returns the
+ duration in seconds with two decimal places, e.g., "12.34s".
+
+ Args:
+ total_seconds: The duration in seconds.
+
+ Returns:
+ A formatted duration string.
+ """
+ minutes = int(total_seconds // 60)
+ if minutes > 0:
+ seconds_rounded = int(round(total_seconds - minutes * 60))
+ if seconds_rounded == 60:
+ minutes += 1
+ seconds_rounded = 0
+ return f"{minutes}m{seconds_rounded:02d}s"
+ return f"{total_seconds:.2f}s"
+
+ def _calculate_total_tokens(
+ self, haystack_sessions: list[list[dict[str, str]]]
+ ) -> int:
+ """Calculate total tokens from all messages in all sessions.
+
+ Args:
+ haystack_sessions: List of sessions, each containing messages
+
+ Returns:
+ Total number of tokens across all messages
+ """
+ tokenizer = tiktoken.get_encoding("cl100k_base")
+ total_tokens = 0
+
+ for session_messages in haystack_sessions:
+ for msg in session_messages:
+ content = msg.get("content", "")
+ total_tokens += len(tokenizer.encode(content))
+
+ return total_tokens
+
+ def _get_latest_tokens_used(self) -> int | None:
+ """Get the tokens_used_estimate from the most recent dialectic_chat metric.
+
+ Returns:
+ Number of tokens used, or None if not found
+ """
+ metrics_file = Path(settings.LOCAL_METRICS_FILE)
+ if not metrics_file.exists():
+ return None
+
+ # Read the file and find the most recent dialectic_chat metric
+ try:
+ with open(metrics_file) as f:
+ lines = f.readlines()
+
+ # Search backwards through the file for the most recent dialectic_chat
+ for line in reversed(lines):
+ if not line.strip():
+ continue
+ try:
+ data = json.loads(line)
+ task_name = data.get("task_name", "")
+ if task_name.startswith("dialectic_chat_"):
+ for metric in data.get("metrics", []):
+ metric_name = metric.get("name", "")
+ if metric_name.endswith("tokens_used_estimate"):
+ return int(metric.get("value", 0))
+ except (json.JSONDecodeError, KeyError, ValueError):
+ continue
+
+ except Exception as e:
+ self.logger.warning(f"Error reading metrics file: {e}")
+
+ return None
+
+ def _parse_date(self, date_str: str) -> datetime:
+ """Parse longmemeval date format to datetime.
+
+ Args:
+ date_str: Date string in format "YYYY/MM/DD (Day) HH:MM"
+
+ Returns:
+ Parsed datetime object
+
+ Raises:
+ ValueError: If date format is invalid
+ """
+ try:
+ # Extract the date and time parts, ignoring the day name in parentheses
+ # Format: "2023/05/20 (Sat) 02:21"
+ parts = date_str.split(") ")
+ if len(parts) != 2:
+ raise ValueError(f"Invalid date format: {date_str}")
+
+ date_part = parts[0].split(" (")[0] # "2023/05/20"
+ time_part = parts[1] # "02:21"
+
+ # Combine and parse
+ datetime_str = f"{date_part} {time_part}"
+ return datetime.strptime(datetime_str, "%Y/%m/%d %H:%M")
+ except (ValueError, IndexError) as e:
+ raise ValueError(f"Failed to parse date '{date_str}': {e}") from e
+
+ def load_test_file(self, test_file: Path) -> list[dict[str, Any]]:
+ """
+ Load longmemeval test definitions from a JSON file.
+
+ Args:
+ test_file: Path to the JSON test file
+
+ Returns:
+ List of test question dictionaries
+ """
+ with open(test_file) as f:
+ return json.load(f)
+
+ async def create_honcho_client(self, workspace_id: str) -> AsyncHoncho:
+ """
+ Create a Honcho client for a specific workspace.
+
+ Args:
+ workspace_id: Workspace ID for the test
+
+ Returns:
+ AsyncHoncho client instance
+ """
+ return AsyncHoncho(
+ environment="local",
+ workspace_id=workspace_id,
+ base_url=self.honcho_url,
+ )
+
+ async def wait_for_deriver_queue_empty(
+ self, honcho_client: AsyncHoncho, session_id: str | None = None
+ ) -> bool:
+ start_time = time.time()
+ while True:
+ try:
+ status = await honcho_client.get_deriver_status(session_id=session_id)
+ except Exception as _e:
+ await asyncio.sleep(1)
+ elapsed_time = time.time() - start_time
+ if elapsed_time >= self.timeout_seconds:
+ return False
+ continue
+
+ if status.pending_work_units == 0 and status.in_progress_work_units == 0:
+ return True
+
+ elapsed_time = time.time() - start_time
+ if elapsed_time >= self.timeout_seconds:
+ return False
+ await asyncio.sleep(1)
+
+ async def judge_response(
+ self, question: str, expected_answer: str, actual_response: str
+ ) -> dict[str, Any]:
+ """
+ Use an LLM to judge if the actual response matches the expected answer.
+
+ Args:
+ question: The question asked
+ expected_answer: Expected answer from the test
+ actual_response: Actual response from Honcho
+
+ Returns:
+ Judgment result with pass/fail and reasoning
+ """
+ try:
+ system_prompt = """
+You are an expert judge evaluating AI responses to memory questions. Your task is to determine if an actual response contains the correct answer from long-term memory.
+
+CRITICAL JUDGING PRINCIPLES:
+1. SEMANTIC UNDERSTANDING: Focus on whether the actual response conveys the same core factual information as expected, even if expressed differently
+2. FLEXIBLE INTERPRETATION: Accept responses that are longer, more detailed, or use different phrasing as long as they contain the correct answer
+3. MEMORY ACCURACY: The key is whether the AI correctly recalled and stated the factual information from memory
+4. PARTIAL CREDIT: If the response shows the AI accessed relevant memories but made minor errors in details, consider partial credit
+5. IMPLICIT vs EXPLICIT: Accept responses that clearly imply the correct answer through context
+
+ONLY FAIL when:
+- The core factual answer is demonstrably wrong
+- The response shows no evidence of accessing the relevant memory
+- The AI explicitly states incorrect information that contradicts the expected answer
+
+Always respond with valid JSON: {"passed": boolean, "reasoning": "short (1-3 sentences) explanation of why the response is correct or incorrect"}"""
+
+ user_prompt = f"""Question: "{question}"
+Expected answer: "{expected_answer}"
+Actual response: "{actual_response}"
+
+Evaluate whether the actual response correctly answers the question based on the expected answer. Focus on factual accuracy and evidence that the AI accessed the correct memory."""
+
+ response = await self.anthropic_client.messages.create(
+ model="claude-sonnet-4-20250514",
+ max_tokens=300,
+ temperature=0.0,
+ system=system_prompt,
+ messages=[
+ {
+ "role": "user",
+ "content": user_prompt,
+ }
+ ],
+ )
+
+ if not response.content:
+ raise ValueError("Anthropic returned empty response")
+
+ content_block = response.content[0]
+ judgment_text = getattr(content_block, "text", None)
+ if judgment_text is None:
+ raise ValueError(
+ f"No text content in response block: {type(content_block)}"
+ )
+
+ # Extract JSON from the response if it's wrapped in markdown
+ if "```json" in judgment_text:
+ json_start = judgment_text.find("```json") + 7
+ json_end = judgment_text.find("```", json_start)
+ judgment_text = judgment_text[json_start:json_end].strip()
+ elif "```" in judgment_text:
+ json_start = judgment_text.find("```") + 3
+ json_end = judgment_text.find("```", json_start)
+ judgment_text = judgment_text[json_start:json_end].strip()
+
+ judgment = json.loads(judgment_text)
+ return judgment
+
+ except Exception as e:
+ self.logger.error(f"Error judging response: {e}")
+ # Fallback to simple string matching
+ is_correct = expected_answer.lower() in actual_response.lower()
+ return {
+ "passed": is_correct,
+ "reasoning": f"Fallback string matching due to error: {'Match found' if is_correct else 'No match found'}",
+ }
+
+ async def execute_question(self, question_data: dict[str, Any]) -> TestResult:
+ """
+ Execute a single longmemeval question.
+
+ Args:
+ question_data: Dictionary containing question data
+
+ Returns:
+ Test execution results
+ """
+ question_id = question_data["question_id"]
+ question_type = question_data["question_type"]
+ question = question_data["question"]
+ expected_answer = question_data["answer"]
+ question_date = question_data.get("question_date", "")
+
+ question_with_date = (
+ f"[{question_date}] {question}" if question_date else question
+ )
+
+ output_lines: list[str] = []
+ output_lines.append(
+ f"\033[1mExecuting question {question_id} ({question_type})\033[0m"
+ )
+ output_lines.append(f"Question: {question_with_date}")
+ output_lines.append(f"Expected: {expected_answer}")
+
+ # Create workspace for this question
+ workspace_id = f"{question_id}_{question_type}"
+ honcho_client = await self.create_honcho_client(workspace_id)
+
+ results: TestResult = {
+ "question_id": question_id,
+ "question_type": question_type,
+ "workspace_id": workspace_id,
+ "sessions_created": [],
+ "query_executed": None,
+ "passed": False,
+ "error": None,
+ "start_time": time.time(),
+ "end_time": 0.0,
+ "duration_seconds": 0.0,
+ "output_lines": output_lines,
+ }
+
+ try:
+ user_peer = await honcho_client.peer(id="user")
+ assistant_peer = await honcho_client.peer(id="assistant")
+
+ # Process haystack sessions
+ haystack_dates = question_data.get("haystack_dates", [])
+ haystack_sessions = question_data.get("haystack_sessions", [])
+ haystack_session_ids = question_data.get("haystack_session_ids", [])
+
+ # Validate alignment of dates, session IDs, and sessions
+ if len(haystack_dates) != len(haystack_sessions):
+ raise ValueError(
+ f"Misaligned data: {len(haystack_dates)} dates but {len(haystack_sessions)} sessions"
+ )
+ if len(haystack_session_ids) != len(haystack_sessions):
+ raise ValueError(
+ f"Misaligned data: {len(haystack_session_ids)} session IDs but {len(haystack_sessions)} sessions"
+ )
+
+ # Parse all dates upfront to catch parsing errors early
+ parsed_dates: list[datetime] = []
+ for date_str in haystack_dates:
+ try:
+ parsed_dates.append(self._parse_date(date_str))
+ except ValueError as e:
+ raise ValueError(f"Error parsing date '{date_str}': {e}") from e
+
+ haystack_total_messages = sum(len(session) for session in haystack_sessions)
+
+ # Calculate total tokens available in the sessions for this question
+ total_available_tokens = self._calculate_total_tokens(haystack_sessions)
+
+ print(
+ f"[{workspace_id}] processing {len(haystack_sessions)} sessions with {haystack_total_messages} total messages ({total_available_tokens} total tokens)"
+ )
+
+ # Determine which peer should be observed based on question type
+ is_assistant_type = question_type == "single-session-assistant"
+
+ if self.merge_sessions:
+ # Create a single merged session for all messages
+ merged_session_id = f"{workspace_id}_merged"
+ session = await honcho_client.session(id=merged_session_id)
+
+ # Configure peer observation based on question type
+ if is_assistant_type:
+ await session.add_peers(
+ [
+ (
+ user_peer,
+ SessionPeerConfig(
+ observe_me=False, observe_others=False
+ ),
+ ),
+ (
+ assistant_peer,
+ SessionPeerConfig(
+ observe_me=True, observe_others=False
+ ),
+ ),
+ ]
+ )
+ else:
+ await session.add_peers(
+ [
+ (
+ user_peer,
+ SessionPeerConfig(
+ observe_me=True, observe_others=False
+ ),
+ ),
+ (
+ assistant_peer,
+ SessionPeerConfig(
+ observe_me=False, observe_others=False
+ ),
+ ),
+ ]
+ )
+
+ # Collect all messages from all sessions in chronological order
+ all_messages: list[MessageCreateParam] = []
+ for session_date, session_messages in zip(
+ parsed_dates, haystack_sessions, strict=True
+ ):
+ for msg in session_messages:
+ role = msg["role"]
+ content = msg["content"]
+
+ # Split message if it exceeds 25000 characters
+ if len(content) > 25000:
+ chunks = [
+ content[i : i + 25000]
+ for i in range(0, len(content), 25000)
+ ]
+ for chunk in chunks:
+ if role == "user":
+ all_messages.append(
+ user_peer.message(
+ chunk, created_at=session_date
+ )
+ )
+ elif role == "assistant":
+ all_messages.append(
+ assistant_peer.message(
+ chunk, created_at=session_date
+ )
+ )
+ else:
+ if role == "user":
+ all_messages.append(
+ user_peer.message(content, created_at=session_date)
+ )
+ elif role == "assistant":
+ all_messages.append(
+ assistant_peer.message(
+ content, created_at=session_date
+ )
+ )
+
+ # Add messages in batches of 100 (max supported by add_messages)
+ if all_messages:
+ for i in range(0, len(all_messages), 100):
+ batch = all_messages[i : i + 100]
+ await session.add_messages(batch)
+
+ results["sessions_created"].append(
+ SessionResult(
+ name=merged_session_id, message_count=len(all_messages)
+ )
+ )
+ else:
+ # create separate sessions
+ # Zip together dates, session IDs, and session content
+ for session_date, session_id, session_messages in zip(
+ parsed_dates, haystack_session_ids, haystack_sessions, strict=True
+ ):
+ session = await honcho_client.session(id=session_id)
+
+ # Configure peer observation based on question type
+ if is_assistant_type:
+ # For assistant questions, observe the assistant peer
+ await session.add_peers(
+ [
+ (
+ user_peer,
+ SessionPeerConfig(
+ observe_me=False, observe_others=False
+ ),
+ ),
+ (
+ assistant_peer,
+ SessionPeerConfig(
+ observe_me=True, observe_others=False
+ ),
+ ),
+ ]
+ )
+ else:
+ # For user questions, observe the user peer (default behavior)
+ await session.add_peers(
+ [
+ (
+ user_peer,
+ SessionPeerConfig(
+ observe_me=True, observe_others=False
+ ),
+ ),
+ (
+ assistant_peer,
+ SessionPeerConfig(
+ observe_me=False, observe_others=False
+ ),
+ ),
+ ]
+ )
+
+ honcho_messages: list[MessageCreateParam] = []
+ for msg in session_messages:
+ role = msg["role"]
+ content = msg["content"]
+
+ # Split message if it exceeds 25000 characters
+ if len(content) > 25000:
+ chunks = [
+ content[i : i + 25000]
+ for i in range(0, len(content), 25000)
+ ]
+ for chunk in chunks:
+ # Use the session date as the timestamp for all messages in this session
+ if role == "user":
+ honcho_messages.append(
+ user_peer.message(
+ chunk, created_at=session_date
+ )
+ )
+ elif role == "assistant":
+ honcho_messages.append(
+ assistant_peer.message(
+ chunk, created_at=session_date
+ )
+ )
+ else:
+ # Use the session date as the timestamp for all messages in this session
+ if role == "user":
+ honcho_messages.append(
+ user_peer.message(content, created_at=session_date)
+ )
+ elif role == "assistant":
+ honcho_messages.append(
+ assistant_peer.message(
+ content, created_at=session_date
+ )
+ )
+
+ if honcho_messages:
+ await session.add_messages(honcho_messages)
+
+ results["sessions_created"].append(
+ SessionResult(
+ name=session_id, message_count=len(honcho_messages)
+ )
+ )
+
+ print(
+ f"[{workspace_id}] fired all messages.\nwaiting for deriver queue to be empty... will time out in {self.timeout_seconds} seconds"
+ )
+ await asyncio.sleep(
+ 1
+ ) # Give time for at least some tasks to be queued, so deriver queue size check doesn't immediately return 0
+
+ queue_empty = await self.wait_for_deriver_queue_empty(honcho_client)
+ if not queue_empty:
+ output_lines.append("Deriver queue never emptied!!!")
+ results["error"] = "Deriver queue timeout"
+ return results
+
+ # Execute the question
+ output_lines.append(f"\nAsking question: {question_with_date}")
+
+ try:
+ # Use the appropriate peer based on question type
+ if is_assistant_type:
+ # For assistant questions, use the assistant peer
+ actual_response = await assistant_peer.chat(question_with_date)
+ else:
+ # For user questions, use the user peer (default behavior)
+ actual_response = await user_peer.chat(question_with_date)
+
+ actual_response = (
+ actual_response if isinstance(actual_response, str) else ""
+ )
+
+ tokens_used = self._get_latest_tokens_used()
+
+ token_efficiency = None
+ if tokens_used is not None and total_available_tokens > 0:
+ efficiency_ratio = tokens_used / total_available_tokens
+ token_efficiency = {
+ "total_available_tokens": total_available_tokens,
+ "tokens_used": tokens_used,
+ "efficiency_ratio": efficiency_ratio,
+ }
+ output_lines.append(
+ f" token efficiency: {efficiency_ratio:.4f} ({tokens_used}/{total_available_tokens} tokens, {efficiency_ratio * 100:.2f}%)"
+ )
+
+ judgment = await self.judge_response(
+ question_with_date, expected_answer, actual_response
+ )
+
+ query_result: QueryResult = {
+ "question": question_with_date,
+ "expected_answer": expected_answer,
+ "actual_response": actual_response,
+ "judgment": judgment,
+ "token_efficiency": token_efficiency,
+ }
+
+ results["query_executed"] = query_result
+ results["passed"] = judgment["passed"]
+
+ output_lines.append(
+ " judgment: \033[1m\033[32mPASS\033[0m"
+ if judgment["passed"]
+ else " judgment: \033[1m\033[31mFAIL\033[0m"
+ )
+ if not judgment["passed"]:
+ output_lines.append(
+ f" got response: \033[3m{actual_response}\033[0m"
+ )
+ output_lines.append(f" expected: {expected_answer}")
+ output_lines.append(f" reasoning: {judgment['reasoning']}")
+
+ except Exception as e:
+ self.logger.error(f"Error executing question: {e}")
+ query_result = QueryResult(
+ question=question_with_date,
+ expected_answer=expected_answer,
+ actual_response=f"ERROR: {e}",
+ judgment={
+ "passed": False,
+ "reasoning": f"Question execution failed: {e}",
+ },
+ token_efficiency=None,
+ )
+ results["query_executed"] = query_result
+ results["passed"] = False
+
+ results["end_time"] = time.time()
+ results["duration_seconds"] = results["end_time"] - results["start_time"]
+
+ output_lines.append(
+ f"\nQuestion {question_id} completed. Status: {'PASS' if results['passed'] else 'FAIL'} (Duration: {self._format_duration(results['duration_seconds'])})"
+ )
+
+ except Exception as e:
+ self.logger.error(f"Error executing question {question_id}: {e}")
+ results["error"] = str(e)
+ results["passed"] = False
+ results["end_time"] = time.time()
+ results["duration_seconds"] = results["end_time"] - results["start_time"]
+ output_lines.append(f"Error executing question {question_id}: {e}")
+
+ return results
+
+ async def run_all_questions(
+ self, test_file: Path, batch_size: int = 10
+ ) -> tuple[list[TestResult], float]:
+ """
+ Run all questions in a longmemeval test file.
+
+ Args:
+ test_file: Path to the longmemeval JSON file
+ batch_size: Number of questions to run concurrently in each batch
+
+ Returns:
+ Tuple of (list of test results, total duration)
+ """
+ questions = self.load_test_file(test_file)
+ print(
+ f"found {len(questions)} {'question' if len(questions) == 1 else 'questions'} in {test_file}"
+ )
+
+ overall_start = time.time()
+
+ # Process questions in batches
+ all_results: list[TestResult] = []
+
+ for i in range(0, len(questions), batch_size):
+ batch = questions[i : i + batch_size]
+ batch_num = (i // batch_size) + 1
+ total_batches = (len(questions) + batch_size - 1) // batch_size
+
+ print(f"\n{'=' * 60}")
+ print(
+ f"Processing batch {batch_num}/{total_batches} ({len(batch)} questions)"
+ )
+ print(f"{'=' * 60}")
+
+ # Run questions in current batch concurrently
+ batch_results: list[TestResult] = await asyncio.gather(
+ *[self.execute_question(q) for q in batch]
+ )
+
+ # Print detailed per-question outputs for this batch
+ for result in batch_results:
+ print(f"\n{'=' * 60}")
+ print("\n".join(result.get("output_lines", [])))
+ print(f"{'=' * 60}\n")
+
+ all_results.extend(batch_results)
+
+ overall_end = time.time()
+ overall_duration = overall_end - overall_start
+
+ # Finalize metrics collection
+ self.metrics_collector.finalize_collection()
+
+ return all_results, overall_duration
+
+ def print_summary(
+ self, results: list[TestResult], total_elapsed_seconds: float | None = None
+ ) -> None:
+ """
+ Print a summary of all test results.
+
+ Args:
+ results: List of test results
+ total_elapsed_seconds: Total elapsed time
+ """
+ print(f"\n{'=' * 80}")
+ print("LONGMEMEVAL TEST EXECUTION SUMMARY")
+ print(f"{'=' * 80}")
+
+ total_questions = len(results)
+ passed_questions = sum(1 for r in results if r.get("passed", False))
+ failed_questions = total_questions - passed_questions
+ total_test_time = (
+ total_elapsed_seconds
+ if total_elapsed_seconds is not None
+ else sum(r["duration_seconds"] for r in results)
+ )
+
+ print(f"Total Questions: {total_questions}")
+ print(f"Passed: {passed_questions}")
+ print(f"Failed: {failed_questions}")
+ print(f"Success Rate: {(passed_questions / total_questions) * 100:.1f}%")
+ print(f"Total Test Time: {self._format_duration(total_test_time)}")
+
+ efficiency_ratios: list[float] = []
+ for result in results:
+ query = result.get("query_executed")
+ if query:
+ token_eff = query.get("token_efficiency")
+ if token_eff:
+ efficiency_ratios.append(token_eff["efficiency_ratio"])
+
+ if efficiency_ratios:
+ avg_efficiency = sum(efficiency_ratios) / len(efficiency_ratios)
+ min_efficiency = min(efficiency_ratios)
+ max_efficiency = max(efficiency_ratios)
+ print("\nToken Efficiency:")
+ print(
+ f" Average: {avg_efficiency:.4f} ({avg_efficiency * 100:.2f}% of available tokens used)"
+ )
+ print(f" Min: {min_efficiency:.4f} ({min_efficiency * 100:.2f}%)")
+ print(f" Max: {max_efficiency:.4f} ({max_efficiency * 100:.2f}%)")
+
+ print("\nDetailed Results:")
+ print(
+ f"{'Question ID':<15} {'Type':<20} {'Status':<8} {'Duration':<10} {'Workspace ID':<30}"
+ )
+ print(f"{'-' * 15} {'-' * 20} {'-' * 8} {'-' * 10} {'-' * 30}")
+
+ for result in results:
+ question_id = result["question_id"]
+ question_type = result["question_type"]
+ status = "PASS" if result.get("passed", False) else "FAIL"
+ duration = self._format_duration(result["duration_seconds"])
+ workspace = result["workspace_id"]
+
+ print(
+ f"{question_id:<15} {question_type:<20} {status:<8} {duration:<10} {workspace:<30}"
+ )
+
+ print(f"{'=' * 80}")
+
+ def generate_json_summary(
+ self,
+ results: list[TestResult],
+ test_file: Path,
+ total_elapsed_seconds: float,
+ output_file: Path | None = None,
+ ) -> None:
+ """
+ Generate a comprehensive JSON summary of test results for analytics.
+
+ Args:
+ results: List of test results
+ test_file: Path to the test file that was executed
+ total_elapsed_seconds: Total elapsed time for all tests
+ output_file: Optional path to write JSON output to
+ """
+ total_questions = len(results)
+ passed_questions = sum(1 for r in results if r.get("passed", False))
+ failed_questions = total_questions - passed_questions
+
+ # Calculate statistics by question type
+ type_stats: dict[str, dict[str, int | float]] = {}
+ for result in results:
+ q_type = result["question_type"]
+ if q_type not in type_stats:
+ type_stats[q_type] = {"total": 0, "passed": 0, "failed": 0}
+ type_stats[q_type]["total"] += 1
+ if result.get("passed", False):
+ type_stats[q_type]["passed"] += 1
+ else:
+ type_stats[q_type]["failed"] += 1
+
+ # Add success rates to type stats
+ for q_type in type_stats:
+ stats = type_stats[q_type]
+ stats["success_rate"] = (
+ (stats["passed"] / stats["total"]) * 100 if stats["total"] > 0 else 0
+ )
+
+ # Calculate timing statistics
+ durations = [r["duration_seconds"] for r in results]
+ timing_stats = {
+ "total_duration_seconds": total_elapsed_seconds,
+ "individual_test_durations": {
+ "min_seconds": min(durations) if durations else 0,
+ "max_seconds": max(durations) if durations else 0,
+ "mean_seconds": sum(durations) / len(durations) if durations else 0,
+ "median_seconds": sorted(durations)[len(durations) // 2]
+ if durations
+ else 0,
+ },
+ }
+
+ # Calculate token efficiency statistics
+ efficiency_ratios: list[float] = []
+ total_available_tokens_list: list[int] = []
+ tokens_used_list: list[int] = []
+ for result in results:
+ query = result.get("query_executed")
+ if query:
+ eff = query.get("token_efficiency")
+ if eff:
+ efficiency_ratios.append(eff["efficiency_ratio"])
+ total_available_tokens_list.append(eff["total_available_tokens"])
+ tokens_used_list.append(eff["tokens_used"])
+
+ token_efficiency_stats = None
+ if efficiency_ratios:
+ token_efficiency_stats = {
+ "mean_efficiency_ratio": sum(efficiency_ratios)
+ / len(efficiency_ratios),
+ "min_efficiency_ratio": min(efficiency_ratios),
+ "max_efficiency_ratio": max(efficiency_ratios),
+ "median_efficiency_ratio": sorted(efficiency_ratios)[
+ len(efficiency_ratios) // 2
+ ],
+ "mean_tokens_available": sum(total_available_tokens_list)
+ / len(total_available_tokens_list),
+ "mean_tokens_used": sum(tokens_used_list) / len(tokens_used_list),
+ "total_questions_with_metrics": len(efficiency_ratios),
+ }
+
+ # Create the full summary
+ summary = {
+ "metadata": {
+ "test_file": str(test_file),
+ "execution_timestamp": datetime.now().isoformat(),
+ "runner_version": "1.0.0",
+ "honcho_url": self.honcho_url,
+ "timeout_seconds": self.timeout_seconds,
+ "deriver_settings": settings.DERIVER.model_dump(),
+ "dialectic_settings": settings.DIALECTIC.model_dump(),
+ "dream_settings": settings.DREAM.model_dump(),
+ "summary_settings": settings.SUMMARY.model_dump(),
+ },
+ "summary_statistics": {
+ "total_questions": total_questions,
+ "passed": passed_questions,
+ "failed": failed_questions,
+ "success_rate_percent": (passed_questions / total_questions) * 100
+ if total_questions > 0
+ else 0,
+ "statistics_by_type": type_stats,
+ },
+ "timing": timing_stats,
+ "token_efficiency": token_efficiency_stats,
+ "detailed_results": [
+ {
+ "question_id": result["question_id"],
+ "question_type": result["question_type"],
+ "workspace_id": result["workspace_id"],
+ "passed": result.get("passed", False),
+ "duration_seconds": result["duration_seconds"],
+ "start_time": result["start_time"],
+ "end_time": result["end_time"],
+ "error": result.get("error"),
+ "query_executed": result.get("query_executed"),
+ }
+ for result in results
+ ],
+ }
+
+ if output_file:
+ output_file.parent.mkdir(parents=True, exist_ok=True)
+ with open(output_file, "w") as f:
+ json.dump(summary, f, indent=2, default=str)
+ print(f"\nJSON summary written to: {output_file}")
+
+
+async def main() -> int:
+ """
+ Main entry point for the longmemeval test runner.
+ """
+ parser = argparse.ArgumentParser(
+ description="Run longmemeval tests against a Honcho instance",
+ formatter_class=argparse.RawDescriptionHelpFormatter,
+ epilog="""
+Examples:
+ %(prog)s --test-file tests/bench/longmemeval_data/longmemeval_s.json # Run longmemeval tests
+ %(prog)s --honcho-url http://localhost:8000 # Custom Honcho URL
+ """,
+ )
+
+ parser.add_argument(
+ "--test-file",
+ type=Path,
+ required=True,
+ help="Path to longmemeval JSON file (required)",
+ )
+
+ parser.add_argument(
+ "--honcho-url",
+ type=str,
+ default="http://localhost:8000",
+ help="URL of the running Honcho instance (default: http://localhost:8000)",
+ )
+
+ parser.add_argument(
+ "--anthropic-api-key",
+ type=str,
+ help="Anthropic API key for response judging (optional)",
+ )
+
+ parser.add_argument(
+ "--timeout",
+ type=int,
+ default=None,
+ help="Timeout for deriver queue to empty in seconds (default: 10 minutes)",
+ )
+
+ parser.add_argument(
+ "--batch-size",
+ type=int,
+ default=10,
+ help="Number of questions to run concurrently in each batch (default: 10)",
+ )
+
+ parser.add_argument(
+ "--json-output",
+ type=Path,
+ help="Path to write JSON summary results for analytics (optional)",
+ )
+
+ parser.add_argument(
+ "--merge-sessions",
+ action="store_true",
+ help="Merge all sessions within a question into a single session (default: False)",
+ )
+
+ args = parser.parse_args()
+
+ # Validate arguments
+ if not args.test_file.exists():
+ print(f"Error: Test file {args.test_file} does not exist")
+ return 1
+
+ if args.batch_size <= 0:
+ print(f"Error: Batch size must be positive, got {args.batch_size}")
+ return 1
+
+ # Create test runner
+ runner = LongMemEvalRunner(
+ honcho_url=args.honcho_url,
+ anthropic_api_key=args.anthropic_api_key,
+ timeout_seconds=args.timeout,
+ merge_sessions=args.merge_sessions,
+ )
+
+ try:
+ # Run all questions
+ results, total_elapsed = await runner.run_all_questions(
+ args.test_file, args.batch_size
+ )
+ runner.print_summary(results, total_elapsed_seconds=total_elapsed)
+
+ # Print metrics summary
+ runner.metrics_collector.print_summary()
+
+ # Generate JSON output if requested
+ if args.json_output:
+ runner.generate_json_summary(
+ results, args.test_file, total_elapsed, args.json_output
+ )
+ else:
+ # Always generate a default JSON output file with timestamp
+ default_output = Path(
+ f"tests/bench/eval_results/longmemeval_results_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
+ )
+ runner.generate_json_summary(
+ results, args.test_file, total_elapsed, default_output
+ )
+
+ # Export metrics to JSON file
+ metrics_output = Path(
+ f"tests/bench/perf_metrics/{datetime.now().strftime('%Y%m%d_%H%M%S')}.json"
+ )
+ runner.metrics_collector.export_to_json(metrics_output)
+ runner.metrics_collector.cleanup_collection()
+
+ # Return exit code based on results
+ all_passed = all(r.get("passed", False) for r in results)
+ return 0 if all_passed else 1
+
+ except KeyboardInterrupt:
+ print("\nTest execution interrupted by user")
+ return 1
+ except Exception as e:
+ print(f"Error running tests: {e}")
+ return 1
+
+
+if __name__ == "__main__":
+ exit_code = asyncio.run(main())
+ exit(exit_code)
diff --git a/tests/bench/peer_card_bench.py b/tests/bench/peer_card_bench.py
index a2233d7d..8119a795 100644
--- a/tests/bench/peer_card_bench.py
+++ b/tests/bench/peer_card_bench.py
@@ -22,6 +22,7 @@ import os
import time
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
+from datetime import datetime, timezone
from pathlib import Path
from typing import Any, cast
@@ -30,7 +31,8 @@ from anthropic import AsyncAnthropic
from src.config import settings
from src.deriver.prompts import peer_card_prompt
from src.utils.clients import honcho_llm_call
-from src.utils.shared_models import PeerCardQuery
+from src.utils.peer_card import PeerCardQuery
+from src.utils.representation import ExplicitObservation, Representation
COLOR_GREEN = "\033[32m"
COLOR_RED = "\033[31m"
@@ -134,23 +136,26 @@ def deduplicate_preserve_order(items: list[Candidate]) -> list[Candidate]:
def build_peer_card_caller(
candidate: Candidate,
-) -> Callable[[list[str] | None, list[str]], Coroutine[Any, Any, PeerCardQuery]]:
+) -> Callable[[list[str] | None, Representation], Coroutine[Any, Any, PeerCardQuery]]:
"""Create an async callable that invokes the peer card prompt with a specific provider/model."""
resolved_provider = (
"openai" if candidate.provider == "custom" else candidate.provider
)
- async def call(old_peer_card: list[str] | None, new_observations: list[str]) -> Any:
+ async def call(
+ old_peer_card: list[str] | None, new_observations: Representation
+ ) -> PeerCardQuery:
prompt = peer_card_prompt(
- old_peer_card=old_peer_card, new_observations=new_observations
+ old_peer_card=old_peer_card,
+ new_observations=new_observations.str_no_timestamps(),
)
response = await honcho_llm_call(
provider=cast(Any, resolved_provider),
model=candidate.model,
prompt=prompt,
- max_tokens=settings.DERIVER.PEER_CARD_MAX_OUTPUT_TOKENS,
+ max_tokens=settings.PEER_CARD.MAX_OUTPUT_TOKENS,
response_model=PeerCardQuery,
json_mode=True,
reasoning_effort="minimal",
@@ -317,11 +322,22 @@ async def run_benchmark(candidates: list[Candidate], cases: list[Case]) -> int:
async def run_case(
case: Case,
_caller: Callable[
- [list[str] | None, list[str]], Coroutine[Any, Any, PeerCardQuery]
+ [list[str] | None, Representation], Coroutine[Any, Any, PeerCardQuery]
] = caller,
) -> tuple[Case, dict[str, Any]]:
card: PeerCardQuery = await _caller(
- case.old_peer_card, case.new_observations
+ case.old_peer_card,
+ Representation(
+ explicit=[
+ ExplicitObservation(
+ content=o,
+ created_at=datetime.now(timezone.utc),
+ message_ids=[(0, 0)],
+ session_name=case.name,
+ )
+ for o in case.new_observations
+ ]
+ ),
)
new_card = card.card
if new_card is None or new_card == []:
diff --git a/tests/bench/run_tests.py b/tests/bench/run_tests.py
index 9168fe7e..6544a909 100644
--- a/tests/bench/run_tests.py
+++ b/tests/bench/run_tests.py
@@ -45,7 +45,7 @@ class QueryResult(TypedDict):
actual_response: str
session: str | None
observer: str | None
- target: str | None
+ observed: str | None
judgment: dict[str, Any]
@@ -416,8 +416,8 @@ Evaluate whether the actual response contains the core correct information from
query: str = query_data["query"]
expected_response: str = query_data["expected_response"]
session_name: str | None = query_data.get("session")
- observer_name: str | None = query_data.get("observer")
- target_name: str | None = query_data.get("target")
+ observer: str | None = query_data.get("observer")
+ observed: str | None = query_data.get("observed")
# Wait for deriver queue to be empty for this session
queue_empty = await self.wait_for_deriver_queue_empty(
@@ -431,17 +431,17 @@ Evaluate whether the actual response contains the core correct information from
context_parts: list[str] = []
if session_name:
context_parts.append(f"session: {session_name}")
- if observer_name:
- context_parts.append(f"observer: {observer_name}")
- if target_name:
- context_parts.append(f"target: {target_name}")
+ if observer:
+ context_parts.append(f"observer: {observer}")
+ if observed:
+ context_parts.append(f"target: {observed}")
if context_parts:
output_lines.append(" " + ", ".join(context_parts))
try:
# Determine which peer to use for the query (observer)
- if observer_name:
- query_peer = peers[observer_name]
+ if observer:
+ query_peer = peers[observer]
else:
# Use the first peer from the first session
first_session_data = list(sessions.values())[0]
@@ -449,19 +449,19 @@ Evaluate whether the actual response contains the core correct information from
query_peer = peers[first_peer_name]
# Execute chat query
- if session_name and target_name:
+ if session_name and observed:
response_text = await query_peer.chat(
query,
session_id=session_name,
- target=peers[target_name],
+ target=peers[observed],
)
elif session_name:
response_text = await query_peer.chat(
query, session_id=session_name
)
- elif target_name:
+ elif observed:
response_text = await query_peer.chat(
- query, target=peers[target_name]
+ query, target=peers[observed]
)
else:
response_text = await query_peer.chat(query)
@@ -480,8 +480,8 @@ Evaluate whether the actual response contains the core correct information from
"expected_response": expected_response,
"actual_response": actual_response,
"session": session_name,
- "observer": observer_name,
- "target": target_name,
+ "observer": observer,
+ "observed": observed,
"judgment": judgment,
}
@@ -515,8 +515,8 @@ Evaluate whether the actual response contains the core correct information from
expected_response=expected_response,
actual_response=f"ERROR: {e}",
session=session_name,
- observer=observer_name,
- target=target_name,
+ observer=observer,
+ observed=observed,
judgment={
"passed": False,
"reasoning": f"Query execution failed: {e}",
diff --git a/tests/conftest.py b/tests/conftest.py
index 480ccf7d..58da6027 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -322,26 +322,29 @@ def mock_llm_call_functions():
patch(
"src.dialectic.chat.dialectic_stream", new_callable=AsyncMock
) as mock_dialectic_stream,
- patch(
- "src.dialectic.utils.generate_semantic_queries", new_callable=AsyncMock
- ) as mock_semantic_queries,
):
# Import the required models for proper mocking
- from src.utils.shared_models import DeductiveObservation, SemanticQueries
+ from src.utils.representation import (
+ DeductiveObservationBase,
+ ExplicitObservationBase,
+ PromptRepresentation,
+ )
# Mock return values for different function types
mock_short_summary.return_value = "Test short summary content"
mock_long_summary.return_value = "Test long summary content"
# Mock critical_analysis_call to return a proper object with _response attribute
- mock_critical_analysis_result = MagicMock()
- mock_critical_analysis_result.explicit = ["Test explicit observation"]
- mock_critical_analysis_result.deductive = [
- DeductiveObservation(
- conclusion="Test deductive conclusion",
- premises=["Test premise 1", "Test premise 2"],
- )
- ]
+ _rep = PromptRepresentation(
+ explicit=[ExplicitObservationBase(content="Test explicit observation")],
+ deductive=[
+ DeductiveObservationBase(
+ conclusion="Test deductive conclusion",
+ premises=["Test premise 1", "Test premise 2"],
+ )
+ ],
+ )
+ mock_critical_analysis_result = MagicMock(wraps=_rep)
# Add the _response attribute that contains thinking (used in the actual code)
mock_response = MagicMock()
mock_response.thinking = "Test thinking content"
@@ -355,18 +358,12 @@ def mock_llm_call_functions():
mock_dialectic_stream.return_value = AsyncMock()
- # Mock semantic query generation
- mock_semantic_queries.return_value = SemanticQueries(
- queries=["test query 1", "test query 2"]
- )
-
yield {
"short_summary": mock_short_summary,
"long_summary": mock_long_summary,
"critical_analysis": mock_critical_analysis,
"dialectic_call": mock_dialectic_call,
"dialectic_stream": mock_dialectic_stream,
- "semantic_queries": mock_semantic_queries,
}
@@ -375,11 +372,10 @@ def mock_honcho_llm_call():
"""Generic mock for the honcho_llm_call decorator to avoid actual LLM calls during tests"""
from unittest.mock import AsyncMock, MagicMock
- from src.utils.shared_models import (
- DeductiveObservation,
- ReasoningResponse,
- ReasoningResponseWithThinking,
- SemanticQueries,
+ from src.utils.representation import (
+ DeductiveObservationBase,
+ ExplicitObservationBase,
+ PromptRepresentation,
)
def create_mock_response(
@@ -396,34 +392,22 @@ def mock_honcho_llm_call():
elif response_model:
# For structured responses, create appropriate mock objects
if getattr(response_model, "__name__", "") == "ReasoningResponse":
- mock_response = MagicMock(spec=ReasoningResponse)
- mock_response.explicit = ["Test explicit observation"]
- mock_response.deductive = [
- DeductiveObservation(
- conclusion="Test deductive conclusion",
- premises=["Test premise 1", "Test premise 2"],
- )
- ]
+ _rep = PromptRepresentation(
+ explicit=[
+ ExplicitObservationBase(content="Test explicit observation")
+ ],
+ deductive=[
+ DeductiveObservationBase(
+ conclusion="Test deductive conclusion",
+ premises=["Test premise 1", "Test premise 2"],
+ ),
+ ],
+ )
+ mock_response = MagicMock(wraps=_rep)
# Add the _response attribute that contains thinking (used in the actual code)
mock_response._response = MagicMock()
mock_response._response.thinking = "Test thinking content"
return mock_response
- elif (
- getattr(response_model, "__name__", "")
- == "ReasoningResponseWithThinking"
- ):
- mock_response = MagicMock(spec=ReasoningResponseWithThinking)
- mock_response.thinking = "Test thinking content"
- mock_response.explicit = ["Test explicit observation"]
- mock_response.deductive = [
- DeductiveObservation(
- conclusion="Test deductive conclusion",
- premises=["Test premise 1", "Test premise 2"],
- )
- ]
- return mock_response
- elif getattr(response_model, "__name__", "") == "SemanticQueries":
- return SemanticQueries(queries=["test query 1", "test query 2"])
else:
# Generic response model mock
mock_response = MagicMock(spec=response_model)
@@ -508,6 +492,8 @@ def mock_tracked_db(db_session: AsyncSession):
with (
patch("src.dependencies.tracked_db", mock_tracked_db_context),
patch("src.deriver.queue_manager.tracked_db", mock_tracked_db_context),
+ patch("src.routers.sessions.tracked_db", mock_tracked_db_context),
+ patch("src.crud.representation.tracked_db", mock_tracked_db_context),
):
yield
@@ -522,14 +508,14 @@ def mock_crud_collection_operations():
async def mock_get_or_create_collection(
_: AsyncSession,
workspace_name: str,
- collection_name: str,
- peer_name: str | None = None,
+ observer: str,
+ observed: str,
):
# Create a mock collection object that doesn't require database commit
mock_collection = models.Collection(
- name=collection_name,
+ observer=observer,
+ observed=observed,
workspace_name=workspace_name,
- peer_name=peer_name,
)
mock_collection.id = generate_nanoid()
return mock_collection
diff --git a/tests/crud/test_peer_card.py b/tests/crud/test_peer_card.py
new file mode 100644
index 00000000..c057323b
--- /dev/null
+++ b/tests/crud/test_peer_card.py
@@ -0,0 +1,451 @@
+import pytest
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from src import models
+from src.crud.peer_card import construct_peer_card_label, get_peer_card, set_peer_card
+from src.exceptions import ResourceNotFoundException
+
+
+@pytest.mark.asyncio
+async def test_peer_card_get_set_roundtrip(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Roundtrip set/get for peer card on a valid peer."""
+ workspace, peer = sample_data
+
+ # Initially absent
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer.name, observed=peer.name
+ )
+ is None
+ )
+
+ # Set and read back
+ value_1 = ["Initial peer card text"]
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ value_1,
+ observer=peer.name,
+ observed=peer.name,
+ )
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer.name, observed=peer.name
+ )
+ == value_1
+ )
+
+ # Update and read back
+ value_2 = ["Updated peer card text", "Another line"]
+ await set_peer_card(
+ db_session, workspace.name, value_2, observer=peer.name, observed=peer.name
+ )
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer.name, observed=peer.name
+ )
+ == value_2
+ )
+
+
+@pytest.mark.asyncio
+async def test_set_peer_card_missing_peer_raises(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Setting a peer card for a non-existent peer should raise ResourceNotFoundException."""
+ workspace, _existing_peer = sample_data
+ with pytest.raises(ResourceNotFoundException):
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ ["card"],
+ observer="missing-peer",
+ observed="missing-peer",
+ )
+
+
+@pytest.mark.asyncio
+async def test_get_peer_card_missing_peer_returns_none(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Getting a peer card for a non-existent peer should return None."""
+ workspace, _existing_peer = sample_data
+ result = await get_peer_card(
+ db_session,
+ workspace.name,
+ observer="missing-peer",
+ observed="missing-peer",
+ )
+ assert result is None
+
+
+@pytest.mark.asyncio
+async def test_get_peer_card_missing_workspace_returns_none(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Getting a peer card for a non-existent workspace should return None."""
+ _workspace, peer = sample_data
+ result = await get_peer_card(
+ db_session,
+ "missing-workspace",
+ observer=peer.name,
+ observed=peer.name,
+ )
+ assert result is None
+
+
+@pytest.mark.asyncio
+async def test_set_peer_card_missing_workspace_raises(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Setting a peer card for a non-existent workspace should raise ResourceNotFoundException."""
+ _workspace, peer = sample_data
+ with pytest.raises(ResourceNotFoundException):
+ await set_peer_card(
+ db_session,
+ "missing-workspace",
+ ["card"],
+ observer=peer.name,
+ observed=peer.name,
+ )
+
+
+@pytest.mark.asyncio
+async def test_peer_card_empty_list(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Setting and getting an empty peer card should work."""
+ workspace, peer = sample_data
+
+ empty_card: list[str] = []
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ empty_card,
+ observer=peer.name,
+ observed=peer.name,
+ )
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer.name, observed=peer.name
+ )
+ == empty_card
+ )
+
+
+@pytest.mark.asyncio
+async def test_peer_card_multiple_lines(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Peer card should support multiple lines."""
+ workspace, peer = sample_data
+
+ multi_line_card = [
+ "First observation about the peer",
+ "Second observation about the peer",
+ "Third observation about the peer",
+ "Fourth observation about the peer",
+ ]
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ multi_line_card,
+ observer=peer.name,
+ observed=peer.name,
+ )
+ result = await get_peer_card(
+ db_session, workspace.name, observer=peer.name, observed=peer.name
+ )
+ assert result is not None
+ assert result == multi_line_card
+ assert len(result) == 4
+
+
+@pytest.mark.asyncio
+async def test_peer_card_different_observer_observed(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Test peer card with different observer and observed peers."""
+ workspace, peer1 = sample_data
+
+ # Create second peer
+ peer2 = models.Peer(name="test-peer-2", workspace_name=workspace.name)
+ db_session.add(peer2)
+ await db_session.flush()
+
+ # Peer1 observes peer2
+ card_1_to_2 = ["Peer 1's observation of peer 2"]
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ card_1_to_2,
+ observer=peer1.name,
+ observed=peer2.name,
+ )
+
+ # Peer2 observes peer1
+ card_2_to_1 = ["Peer 2's observation of peer 1"]
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ card_2_to_1,
+ observer=peer2.name,
+ observed=peer1.name,
+ )
+
+ # Verify each card is independent
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer1.name, observed=peer2.name
+ )
+ == card_1_to_2
+ )
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer2.name, observed=peer1.name
+ )
+ == card_2_to_1
+ )
+
+
+@pytest.mark.asyncio
+async def test_peer_card_self_and_other_observations(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Test that a peer can have both a self-observation and observations of others."""
+ workspace, peer1 = sample_data
+
+ # Create second peer
+ peer2 = models.Peer(name="test-peer-2", workspace_name=workspace.name)
+ db_session.add(peer2)
+ await db_session.flush()
+
+ # Peer1's self-observation
+ self_card = ["Self observation"]
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ self_card,
+ observer=peer1.name,
+ observed=peer1.name,
+ )
+
+ # Peer1's observation of peer2
+ other_card = ["Observation of other peer"]
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ other_card,
+ observer=peer1.name,
+ observed=peer2.name,
+ )
+
+ # Both should be retrievable independently
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer1.name, observed=peer1.name
+ )
+ == self_card
+ )
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer1.name, observed=peer2.name
+ )
+ == other_card
+ )
+
+
+@pytest.mark.asyncio
+async def test_peer_card_multiple_observers_same_observed(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Test that multiple peers can observe the same peer with different cards."""
+ workspace, peer1 = sample_data
+
+ # Create two more peers
+ peer2 = models.Peer(name="test-peer-2", workspace_name=workspace.name)
+ peer3 = models.Peer(name="test-peer-3", workspace_name=workspace.name)
+ db_session.add_all([peer2, peer3])
+ await db_session.flush()
+
+ # Both peer2 and peer3 observe peer1
+ card_2_to_1 = ["Peer 2's view of peer 1"]
+ card_3_to_1 = ["Peer 3's view of peer 1"]
+
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ card_2_to_1,
+ observer=peer2.name,
+ observed=peer1.name,
+ )
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ card_3_to_1,
+ observer=peer3.name,
+ observed=peer1.name,
+ )
+
+ # Each observer should have their own independent card
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer2.name, observed=peer1.name
+ )
+ == card_2_to_1
+ )
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer3.name, observed=peer1.name
+ )
+ == card_3_to_1
+ )
+
+
+@pytest.mark.asyncio
+async def test_peer_card_update_does_not_affect_others(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Test that updating one peer card doesn't affect other peer cards."""
+ workspace, peer1 = sample_data
+
+ # Create second peer
+ peer2 = models.Peer(name="test-peer-2", workspace_name=workspace.name)
+ db_session.add(peer2)
+ await db_session.flush()
+
+ # Set initial cards
+ card_1 = ["Peer 1 self observation"]
+ card_2 = ["Peer 1's observation of peer 2"]
+
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ card_1,
+ observer=peer1.name,
+ observed=peer1.name,
+ )
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ card_2,
+ observer=peer1.name,
+ observed=peer2.name,
+ )
+
+ # Update one card
+ new_card_1 = ["Updated self observation"]
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ new_card_1,
+ observer=peer1.name,
+ observed=peer1.name,
+ )
+
+ # Verify only the updated card changed
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer1.name, observed=peer1.name
+ )
+ == new_card_1
+ )
+ assert (
+ await get_peer_card(
+ db_session, workspace.name, observer=peer1.name, observed=peer2.name
+ )
+ == card_2
+ )
+
+
+@pytest.mark.asyncio
+async def test_peer_card_special_characters(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Test that peer cards handle special characters correctly."""
+ workspace, peer = sample_data
+
+ special_card = [
+ "Line with special chars: @#$%^&*()",
+ "Line with unicode: 你好世界 🌍",
+ "Line with quotes: \"double\" and 'single'",
+ "Line with newlines embedded\\n",
+ ]
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ special_card,
+ observer=peer.name,
+ observed=peer.name,
+ )
+ result = await get_peer_card(
+ db_session, workspace.name, observer=peer.name, observed=peer.name
+ )
+ assert result == special_card
+
+
+@pytest.mark.asyncio
+async def test_peer_card_large_content(
+ db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
+):
+ """Test that peer cards can handle large amounts of content."""
+ workspace, peer = sample_data
+
+ # Create a large peer card with many observations
+ large_card = [f"Observation number {i}" for i in range(100)]
+ await set_peer_card(
+ db_session,
+ workspace.name,
+ large_card,
+ observer=peer.name,
+ observed=peer.name,
+ )
+ result = await get_peer_card(
+ db_session, workspace.name, observer=peer.name, observed=peer.name
+ )
+ assert result == large_card
+ assert result is not None
+ assert len(result) == 100
+
+
+def test_construct_helpers_labels():
+ """Helper label constructors return expected values."""
+ assert construct_peer_card_label(observer="a", observed="a") == "peer_card"
+ assert construct_peer_card_label(observer="a", observed="b") == "b_peer_card"
+
+
+def test_construct_peer_card_label_with_special_chars():
+ """Test label construction with special characters in peer names."""
+ # Test with dashes and underscores
+ assert (
+ construct_peer_card_label(observer="peer-1", observed="peer-2")
+ == "peer-2_peer_card"
+ )
+ assert (
+ construct_peer_card_label(observer="peer_1", observed="peer_2")
+ == "peer_2_peer_card"
+ )
+
+ # Test with same observer and observed with special chars
+ assert (
+ construct_peer_card_label(observer="peer-1", observed="peer-1") == "peer_card"
+ )
+
+
+def test_construct_peer_card_label_edge_cases():
+ """Test label construction edge cases."""
+ # Test with empty strings (though this shouldn't happen in practice)
+ assert construct_peer_card_label(observer="", observed="") == "peer_card"
+ assert construct_peer_card_label(observer="a", observed="") == "_peer_card"
+
+ # Test with very long peer names
+ long_name = "a" * 100
+ assert (
+ construct_peer_card_label(observer=long_name, observed=long_name) == "peer_card"
+ )
+ assert (
+ construct_peer_card_label(observer="a", observed=long_name)
+ == f"{long_name}_peer_card"
+ )
diff --git a/tests/deriver/README.md b/tests/deriver/README.md
index d7f4ebfa..b2f040e8 100644
--- a/tests/deriver/README.md
+++ b/tests/deriver/README.md
@@ -27,7 +27,7 @@ This directory contains tests for the deriver system, which handles background p
- `mock_critical_analysis_call` - Mocks the critical analysis LLM call
- `mock_queue_manager` - Mocks the queue manager for testing
-- `mock_embedding_store` - Mocks the embedding store operations
+- `mock_representation_manager` - Mocks the representation manager operations
## Testing Patterns
@@ -38,8 +38,8 @@ This directory contains tests for the deriver system, which handles background p
payload = create_queue_payload(
message=message,
task_type="representation",
- sender_name=message.peer_name,
- target_name=observer_peer.name,
+ observer=observer_peer.name,
+ observed=message.peer_name
)
# Add to queue
@@ -52,11 +52,11 @@ queue_items = await add_queue_items([payload], session.id)
# Create a work unit
work_unit = WorkUnit(
session_id=session.id,
- sender_name=sender.name,
- target_name=target.name,
task_type="representation",
+ observer=observer,
+ observed=observed
)
# Test string representation
-assert str(work_unit) == f"({session.id}, {sender.name}, {target.name}, representation)"
+assert str(work_unit) == f"({session.id}, {observed.name}, {observer.name}, representation)"
```
diff --git a/tests/deriver/conftest.py b/tests/deriver/conftest.py
index 5578152d..9ade5ef7 100644
--- a/tests/deriver/conftest.py
+++ b/tests/deriver/conftest.py
@@ -10,8 +10,8 @@ from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models, schemas
-from src.deriver.queue_payload import create_payload
-from src.deriver.utils import get_work_unit_key
+from src.utils.queue_payload import create_payload
+from src.utils.work_unit import get_work_unit_key
@pytest.fixture
@@ -124,8 +124,8 @@ def create_queue_payload() -> Callable[..., Any]:
def _create_payload(
message: models.Message,
task_type: Literal["representation", "summary"],
- sender_name: str | None = None,
- target_name: str | None = None,
+ observer: str | None = None,
+ observed: str | None = None,
message_seq_in_session: int | None = None,
) -> dict[str, Any]:
"""Create a queue payload for testing"""
@@ -141,9 +141,9 @@ def create_queue_payload() -> Callable[..., Any]:
return create_payload(
message=message_dict,
task_type=task_type,
- sender_name=sender_name,
- target_name=target_name,
message_seq_in_session=message_seq_in_session,
+ observer=observer,
+ observed=observed,
)
return _create_payload
@@ -163,7 +163,7 @@ async def add_queue_items(
for payload in payloads:
# Generate work_unit_key from the payload
task_type = payload.get("task_type", "unknown")
- work_unit_key = get_work_unit_key(task_type, payload)
+ work_unit_key = get_work_unit_key(payload)
queue_item = models.QueueItem(
session_id=session_id,
@@ -208,8 +208,8 @@ async def sample_queue_items(
payload1 = create_queue_payload(
message=message,
task_type="representation",
- sender_name=message.peer_name,
- target_name=message.peer_name,
+ observer=message.peer_name,
+ observed=message.peer_name,
)
payloads.append(payload1)
@@ -217,8 +217,8 @@ async def sample_queue_items(
payload2 = create_queue_payload(
message=message,
task_type="representation",
- sender_name=message.peer_name,
- target_name=peer2.name, # peer2 observes others
+ observer=peer2.name, # peer2 observes others
+ observed=message.peer_name,
)
payloads.append(payload2)
@@ -285,12 +285,12 @@ def mock_queue_manager(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: # pyright
@pytest.fixture
-def mock_embedding_store(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: # pyright: ignore[reportUnusedParameter]
- """Mock the embedding store to avoid actual embedding operations"""
- from src.utils.embedding_store import EmbeddingStore
+def mock_representation_manager(monkeypatch: pytest.MonkeyPatch) -> AsyncMock: # pyright: ignore[reportUnusedParameter]
+ """Mock the representation manager to avoid actual embedding operations"""
+ from src.crud.representation import RepresentationManager
- mock_store = AsyncMock(spec=EmbeddingStore)
- mock_store.save_unified_observations = AsyncMock()
- mock_store.get_relevant_observations = AsyncMock(return_value=MagicMock())
+ mock_manager = AsyncMock(spec=RepresentationManager)
+ mock_manager.save_representation.return_value = 0
+ mock_manager.get_relevant_observations = AsyncMock(return_value=MagicMock())
- return mock_store
+ return mock_manager
diff --git a/tests/deriver/test_deriver_processing.py b/tests/deriver/test_deriver_processing.py
index 14636185..dc986ae2 100644
--- a/tests/deriver/test_deriver_processing.py
+++ b/tests/deriver/test_deriver_processing.py
@@ -8,7 +8,7 @@ import pytest
from src import models
from src.deriver.deriver import process_representation_tasks_batch
-from src.utils.shared_models import ReasoningResponseWithThinking
+from src.utils.representation import Representation
@pytest.mark.asyncio
@@ -32,7 +32,7 @@ class TestDeriverProcessing:
sample_session_with_peers: tuple[models.Session, list[models.Peer]],
):
"""Test that work unit keys are generated correctly"""
- from src.deriver.utils import get_work_unit_key
+ from src.utils.work_unit import get_work_unit_key
session, peers = sample_session_with_peers
peer1, peer2, _ = peers
@@ -41,19 +41,19 @@ class TestDeriverProcessing:
representation_payload = {
"workspace_name": "workspace1",
"session_name": session.name,
- "sender_name": peer1.name,
- "target_name": peer2.name,
+ "observer": peer2.name,
+ "observed": peer1.name,
"task_type": "representation",
}
# Generate work unit key for representation
- work_unit_key = get_work_unit_key("representation", representation_payload)
+ work_unit_key = get_work_unit_key(representation_payload)
expected_key = (
- f"representation:workspace1:{session.name}:{peer1.name}:{peer2.name}"
+ f"representation:workspace1:{session.name}:{peer2.name}:{peer1.name}"
)
assert work_unit_key == expected_key
- # Create a payload for summary task (sender_name and target_name should be None)
+ # Create a payload for summary task
summary_payload = {
"workspace_name": "workspace1",
"session_name": session.name,
@@ -61,7 +61,7 @@ class TestDeriverProcessing:
}
# Generate work unit key for summary
- summary_work_unit_key = get_work_unit_key("summary", summary_payload)
+ summary_work_unit_key = get_work_unit_key(summary_payload)
expected_summary_key = f"summary:workspace1:{session.name}:None:None"
assert summary_work_unit_key == expected_summary_key
@@ -89,19 +89,21 @@ class TestDeriverProcessing:
mock_queue_manager.initialize.assert_called_once() # type: ignore[attr-defined]
mock_queue_manager.shutdown.assert_called_once() # type: ignore[attr-defined]
- async def test_mock_embedding_store(
+ async def test_mock_representation_manager(
self,
- mock_embedding_store: Any, # AsyncMock object
+ mock_representation_manager: Any, # AsyncMock object
):
- """Test that the embedding store is properly mocked"""
- assert mock_embedding_store is not None
+ """Test that the representation manager is properly mocked"""
+ assert mock_representation_manager is not None
# Verify we can call the mocked methods
- await mock_embedding_store.save_unified_observations([])
- mock_embedding_store.get_relevant_observations.return_value = [] # type: ignore[attr-defined]
+ await mock_representation_manager.save_representation(
+ Representation(explicit=[], deductive=[])
+ )
+ mock_representation_manager.get_relevant_observations.return_value = [] # type: ignore[attr-defined]
# Verify the methods were called
- assert mock_embedding_store.save_unified_observations.called # type: ignore[attr-defined]
+ assert mock_representation_manager.save_representation.called # type: ignore[attr-defined]
async def test_representation_batch_uses_earliest_cutoff(
self,
@@ -122,24 +124,22 @@ class TestDeriverProcessing:
# Provide a stub working representation so embedding lookups are skipped.
monkeypatch.setattr(
- "src.deriver.deriver.crud.get_working_representation_data",
+ "src.crud.get_working_representation",
AsyncMock(
- return_value={
- "final_observations": {
- "explicit": ["existing"],
- "deductive": [],
- }
- }
+ return_value=Representation(
+ explicit=[],
+ deductive=[],
+ )
),
)
# Avoid DB access for collection and peer card
monkeypatch.setattr(
- "src.deriver.deriver.crud.get_or_create_collection",
+ "src.crud.get_or_create_collection",
AsyncMock(return_value=type("Collection", (), {"name": "dummy"})()),
)
monkeypatch.setattr(
- "src.deriver.deriver.crud.get_peer_card",
+ "src.crud.get_peer_card",
AsyncMock(return_value=[]),
)
# Short-circuit tracked_db context manager
@@ -154,17 +154,7 @@ class TestDeriverProcessing:
# Avoid executing the full reasoning pipeline; we only care about cutoff behavior.
monkeypatch.setattr(
"src.deriver.deriver.CertaintyReasoner.reason",
- AsyncMock(
- return_value=ReasoningResponseWithThinking(
- thinking=None, explicit=[], deductive=[]
- )
- ),
- )
-
- # Skip persisting results back to the database.
- monkeypatch.setattr(
- "src.deriver.deriver.save_working_representation_to_peer",
- AsyncMock(),
+ AsyncMock(return_value=Representation(explicit=[], deductive=[])),
)
# Create test messages with different IDs (earlier message has lower ID)
@@ -185,7 +175,7 @@ class TestDeriverProcessing:
)
await process_representation_tasks_batch(
- sender_name="alice", target_name="alice", messages=messages
+ observer="alice", observed="alice", messages=messages
)
# Verify that the earliest message ID was used as the cutoff
diff --git a/tests/deriver/test_queue_processing.py b/tests/deriver/test_queue_processing.py
index 8309f04c..3b631a3f 100644
--- a/tests/deriver/test_queue_processing.py
+++ b/tests/deriver/test_queue_processing.py
@@ -10,7 +10,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
from src import models
from src.config import settings
from src.deriver.queue_manager import QueueManager, WorkerOwnership
-from src.deriver.utils import get_work_unit_key
+from src.utils.work_unit import get_work_unit_key
@pytest.mark.asyncio
@@ -143,8 +143,8 @@ class TestQueueProcessing:
payload = create_queue_payload( # type: ignore[reportUnknownArgumentType]
message=message,
task_type="representation",
- sender_name=peer.name,
- target_name=peer.name,
+ observed=peer.name,
+ observer=peer.name,
)
payloads.append(payload)
@@ -306,8 +306,8 @@ class TestQueueProcessing:
create_queue_payload( # type: ignore[reportUnknownArgumentType]
message=msg,
task_type="representation",
- sender_name=peer.name,
- target_name=peer.name,
+ observed=peer.name,
+ observer=peer.name,
)
for msg in messages
]
@@ -315,7 +315,7 @@ class TestQueueProcessing:
queue_items: list[models.QueueItem] = []
for payload in payloads:
task_type = payload.get("task_type", "unknown")
- work_unit_key = get_work_unit_key(task_type, payload)
+ work_unit_key = get_work_unit_key(payload)
queue_item = models.QueueItem(
session_id=session.id,
@@ -336,8 +336,8 @@ class TestQueueProcessing:
async def mock_process_representation_batch(
messages: list[models.Message],
- sender_name: str | None = None, # pyright: ignore[reportUnusedParameter]
- target_name: str | None = None, # pyright: ignore[reportUnusedParameter]
+ observed: str | None = None, # pyright: ignore[reportUnusedParameter]
+ observer: str | None = None, # pyright: ignore[reportUnusedParameter]
) -> None:
processed_batches.append(
{
@@ -424,10 +424,10 @@ class TestQueueProcessing:
payload = create_queue_payload( # type: ignore[reportUnknownArgumentType]
message=message,
task_type="representation",
- sender_name=peer.name,
- target_name=target.name,
+ observed=peer.name,
+ observer=target.name,
)
- work_unit_key = get_work_unit_key("representation", payload)
+ work_unit_key = get_work_unit_key(payload)
queue_item = models.QueueItem(
session_id=session.id,
@@ -480,9 +480,7 @@ class TestQueueProcessing:
assert alice_message_ids == expected_batch_ids
# Ensure items are only for alice
- assert all(
- qi.payload.get("sender_name") == alice.name for qi in alice_items
- )
+ assert all(qi.payload.get("observed") == alice.name for qi in alice_items)
# Test bob's work unit - starts at message 2 for per-work-unit anchoring
bob_work_unit_key = bob_queue_items[0].work_unit_key
@@ -508,7 +506,7 @@ class TestQueueProcessing:
}
assert bob_message_ids == expected_bob_ids
# Ensure items are only for bob
- assert all(qi.payload.get("sender_name") == bob.name for qi in bob_items)
+ assert all(qi.payload.get("observed") == bob.name for qi in bob_items)
# Test steve's work unit - starts at message 3 for per-work-unit anchoring
steve_work_unit_key = steve_queue_items[0].work_unit_key
@@ -535,9 +533,7 @@ class TestQueueProcessing:
}
assert steve_message_ids == expected_steve_ids
# Ensure items are only for steve
- assert all(
- qi.payload.get("sender_name") == steve.name for qi in steve_items
- )
+ assert all(qi.payload.get("observed") == steve.name for qi in steve_items)
@pytest.mark.asyncio
async def test_per_work_unit_anchoring_with_token_limits(
@@ -593,10 +589,10 @@ class TestQueueProcessing:
payload = create_queue_payload( # type: ignore[reportUnknownArgumentType]
message=message,
task_type="representation",
- sender_name=peer.name,
- target_name=target.name,
+ observed=peer.name,
+ observer=target.name,
)
- work_unit_key = get_work_unit_key("representation", payload)
+ work_unit_key = get_work_unit_key(payload)
queue_item = models.QueueItem(
session_id=session.id,
@@ -727,7 +723,7 @@ class TestQueueProcessing:
)
payload["token_count"] = token_counts[i]
- work_unit_key = get_work_unit_key("summary", payload)
+ work_unit_key = get_work_unit_key(payload)
queue_item = models.QueueItem(
session_id=session.id,
@@ -839,18 +835,17 @@ class TestQueueProcessing:
create_queue_payload( # type: ignore[reportUnknownArgumentType]
message=msg,
task_type="representation",
- sender_name=peer.name,
- target_name=peer.name,
+ observed=peer.name,
+ observer=peer.name,
)
for msg in messages
]
# Add items to queue
-
queue_items: list[models.QueueItem] = []
for payload in payloads:
task_type = payload.get("task_type", "unknown")
- work_unit_key = get_work_unit_key(task_type, payload)
+ work_unit_key = get_work_unit_key(payload)
queue_item = models.QueueItem(
session_id=session.id,
@@ -871,8 +866,8 @@ class TestQueueProcessing:
async def mock_process_representation_batch(
messages: list[models.Message],
- sender_name: str | None = None, # pyright: ignore[reportUnusedParameter]
- target_name: str | None = None, # pyright: ignore[reportUnusedParameter]
+ observed: str | None = None, # pyright: ignore[reportUnusedParameter]
+ observer: str | None = None, # pyright: ignore[reportUnusedParameter]
) -> None:
processed_batches.append(
{
@@ -950,18 +945,17 @@ class TestQueueProcessing:
create_queue_payload( # type: ignore[reportUnknownArgumentType]
message=msg,
task_type="representation",
- sender_name=peer.name,
- target_name=peer.name,
+ observed=peer.name,
+ observer=peer.name,
)
for msg in messages
]
# Add items to queue
-
queue_items: list[models.QueueItem] = []
for payload in payloads:
task_type = payload.get("task_type", "unknown")
- work_unit_key = get_work_unit_key(task_type, payload)
+ work_unit_key = get_work_unit_key(payload)
queue_item = models.QueueItem(
session_id=session.id,
@@ -982,8 +976,8 @@ class TestQueueProcessing:
async def mock_process_representation_batch(
messages: list[models.Message],
- sender_name: str | None = None, # pyright: ignore[reportUnusedParameter]
- target_name: str | None = None, # pyright: ignore[reportUnusedParameter]
+ observed: str | None = None, # pyright: ignore[reportUnusedParameter]
+ observer: str | None = None, # pyright: ignore[reportUnusedParameter]
) -> None:
processed_batches.append(
{
diff --git a/tests/deriver/test_representation_crud.py b/tests/deriver/test_representation_crud.py
index febd2f9b..fe4d8e95 100644
--- a/tests/deriver/test_representation_crud.py
+++ b/tests/deriver/test_representation_crud.py
@@ -1,593 +1,97 @@
import datetime
-from typing import Any
-import pytest
-from nanoid import generate as generate_nanoid
-from sqlalchemy import select
-from sqlalchemy.ext.asyncio import AsyncSession
-
-from src import crud, models, schemas
-from src.config import settings
-from src.crud.representation import (
- get_peer_card,
- get_working_representation,
- get_working_representation_data,
- set_peer_card,
- set_working_representation,
+from src.utils.representation import (
+ DeductiveObservation,
+ DeductiveObservationBase,
+ ExplicitObservation,
+ ExplicitObservationBase,
+ PromptRepresentation,
+ Representation,
)
-from src.exceptions import ResourceNotFoundException
-@pytest.mark.asyncio
-async def test_peer_card_get_set_roundtrip(
- db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
-):
- """Roundtrip set/get for peer card on a valid peer."""
- workspace, peer = sample_data
-
- # Initially absent
- assert await get_peer_card(db_session, workspace.name, peer.name, peer.name) is None
-
- # Set and read back
- value_1 = ["Initial peer card text"]
- await set_peer_card(db_session, workspace.name, peer.name, peer.name, value_1)
- assert (
- await get_peer_card(db_session, workspace.name, peer.name, peer.name) == value_1
+def test_representation_is_empty_and_diff():
+ """is_empty and diff_representation behave per the new definitions."""
+ now = datetime.datetime.now(datetime.timezone.utc)
+ shared_time = now - datetime.timedelta(seconds=10)
+ exp_shared_1 = ExplicitObservation(
+ content="A",
+ created_at=shared_time,
+ message_ids=[(1, 1)],
+ session_name="s",
+ )
+ exp_shared_2 = ExplicitObservation(
+ content="B",
+ created_at=shared_time,
+ message_ids=[(1, 1)],
+ session_name="s",
+ )
+ rep1 = Representation(explicit=[exp_shared_1], deductive=[])
+ rep2 = Representation(
+ explicit=[
+ ExplicitObservation(
+ content="A",
+ created_at=shared_time,
+ message_ids=[(1, 1)],
+ session_name="s",
+ ),
+ exp_shared_2,
+ ]
)
- # Update and read back
- value_2 = ["Updated peer card text", "Another line"]
- await set_peer_card(db_session, workspace.name, peer.name, peer.name, value_2)
- assert (
- await get_peer_card(db_session, workspace.name, peer.name, peer.name) == value_2
+ assert not rep1.is_empty()
+ assert Representation().is_empty()
+
+ diff = rep1.diff_representation(rep2)
+ assert [e.content for e in diff.explicit] == ["B"]
+ assert diff.deductive == []
+
+
+def test_representation_formatting_methods():
+ """__str__ and format_as_markdown produce expected section headers and content."""
+ now = datetime.datetime.now(datetime.timezone.utc)
+ e = ExplicitObservation(
+ content="has a dog",
+ created_at=now,
+ message_ids=[(1, 1)],
+ session_name="s",
)
-
-
-@pytest.mark.asyncio
-async def test_set_peer_card_missing_peer_raises(
- db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
-):
- """Setting a peer card for a non-existent peer should raise ResourceNotFoundException."""
- workspace, _existing_peer = sample_data
- with pytest.raises(ResourceNotFoundException):
- await set_peer_card(
- db_session, workspace.name, "missing-peer", "missing-peer", ["card"]
- )
-
-
-async def _create_session_with_peers(
- db_session: AsyncSession, workspace: models.Workspace
-) -> tuple[models.Session, models.Peer, models.Peer]:
- """Create a session with two peers and return (session, observer, observed)."""
- observer = models.Peer(name=str(generate_nanoid()), workspace_name=workspace.name)
- observed = models.Peer(name=str(generate_nanoid()), workspace_name=workspace.name)
- db_session.add_all([observer, observed])
- await db_session.flush()
-
- session_name = str(generate_nanoid())
- session = await crud.get_or_create_session(
- db_session,
- schemas.SessionCreate(
- name=session_name,
- peers={
- observer.name: schemas.SessionPeerConfig(),
- observed.name: schemas.SessionPeerConfig(),
- },
- ),
- workspace.name,
+ d = DeductiveObservation(
+ created_at=now,
+ message_ids=[(1, 1)],
+ session_name="s",
+ conclusion="owns a pet",
+ premises=[e.content],
)
- await db_session.commit()
- return session, observer, observed
+ rep = Representation(explicit=[e], deductive=[d])
+
+ s = str(rep)
+ assert "EXPLICIT:" in s
+ assert "DEDUCTIVE:" in s
+ assert "owns a pet" in s
+
+ md = rep.format_as_markdown()
+ assert "## Explicit Observations" in md
+ assert "## Deductive Observations" in md
+ assert "**Conclusion**: owns a pet" in md
-def _make_wr_payload(
- *,
- explicit: list[str],
- deductive: list[dict[str, Any]] | None = None,
- thinking: str | None = None,
- message_id: str = "m-new",
- created_at: str | None = None,
-) -> dict[str, Any]:
- """Build a structured working representation dict payload."""
- return {
- "final_observations": {
- "explicit": explicit,
- "deductive": deductive or [],
- "thinking": thinking,
- },
- "message_id": message_id,
- "created_at": created_at
- or datetime.datetime.now(datetime.timezone.utc).isoformat(),
- }
-
-
-@pytest.mark.asyncio
-async def test_working_representation_self_merge_and_trim(
- db_session: AsyncSession,
- sample_data: tuple[models.Workspace, models.Peer],
- monkeypatch: pytest.MonkeyPatch,
-):
- """Merging a self working representation appends and trims to 25, and updates metadata fields."""
- workspace, peer = sample_data
-
- # Align limit with current implementation and future config usage
- monkeypatch.setattr(
- settings.DERIVER, "WORKING_REPRESENTATION_MAX_OBSERVATIONS", 25, raising=False
+def test_prompt_representation_conversion():
+ """PromptRepresentation.to_representation maps strings to observation objects."""
+ pr = PromptRepresentation(
+ explicit=[ExplicitObservationBase(content="A")],
+ deductive=[DeductiveObservationBase(conclusion="C", premises=["P1"])],
)
- LIMIT = settings.DERIVER.WORKING_REPRESENTATION_MAX_OBSERVATIONS
-
- # Create a session with the single peer
- session = await crud.get_or_create_session(
- db_session,
- schemas.SessionCreate(
- name=str(generate_nanoid()), peers={peer.name: schemas.SessionPeerConfig()}
- ),
- workspace.name,
+ timestamp = datetime.datetime(2025, 1, 1, 12, 0, 0, tzinfo=datetime.timezone.utc)
+ rep = Representation.from_prompt_representation(
+ pr,
+ message_ids=(1, 1),
+ session_name="s",
+ created_at=timestamp,
)
- await db_session.commit()
-
- # Existing explicit: 24 items 0..23
- existing = _make_wr_payload(
- explicit=[f"E{i}" for i in range(24)],
- deductive=[{"conclusion": f"D{i}", "premises": [f"P{i}"]} for i in range(10)],
- thinking="old-think",
- message_id="m-old",
- created_at=datetime.datetime(
- 2024, 1, 1, tzinfo=datetime.timezone.utc
- ).isoformat(),
- )
- await set_working_representation(
- db_session, existing, workspace.name, peer.name, peer.name, session.name
- )
-
- # New explicit: 3 items 24..26, deductive two new items
- new_payload = _make_wr_payload(
- explicit=["E24", "E25", "E26"],
- deductive=[
- {"conclusion": "D_new1", "premises": []},
- {"conclusion": "D_new2", "premises": ["PX"]},
- ],
- thinking="new-think",
- message_id="m-new",
- created_at=datetime.datetime(
- 2025, 1, 1, tzinfo=datetime.timezone.utc
- ).isoformat(),
- )
- await set_working_representation(
- db_session, new_payload, workspace.name, peer.name, peer.name, session.name
- )
-
- # Verify merged raw data
- raw = await get_working_representation_data(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- assert isinstance(raw, dict)
- final = raw["final_observations"]
-
- # Explicit should be last LIMIT of 24+3 (drop oldest overflow)
- explicit = final["explicit"]
- total_explicit = 24 + 3
- expected_len = min(LIMIT, total_explicit)
- dropped = max(0, total_explicit - LIMIT)
- assert len(explicit) == expected_len
- assert explicit[0] == f"E{dropped}"
- assert explicit[-1] == "E26"
-
- # Deductive should be appended and capped to LIMIT
- deductive = final["deductive"]
- assert len(deductive) == min(LIMIT, 12)
- assert deductive[-2]["conclusion"] == "D_new1"
- assert deductive[-1]["conclusion"] == "D_new2"
-
- # Thinking and metadata should reflect latest
- assert final.get("thinking") == "new-think"
- assert raw.get("message_id") == "m-new"
- created_at_value = raw.get("created_at")
- assert created_at_value is not None
- assert created_at_value.startswith("2025-01-01")
-
- # Formatted string getter returns sections and bullets
- formatted = await get_working_representation(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- assert "EXPLICIT OBSERVATIONS:" in formatted
- assert "DEDUCTIVE OBSERVATIONS:" in formatted
- assert "- E26" in formatted
-
-
-@pytest.mark.asyncio
-async def test_working_representation_directional_merge_and_keys(
- db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
-):
- """Directional working representations are stored under observer_observed key and merge correctly."""
- workspace, _ = sample_data
- session, observer, observed = await _create_session_with_peers(
- db_session, workspace
- )
-
- # Store initial
- first = _make_wr_payload(explicit=["A", "B"], thinking="first")
- await set_working_representation(
- db_session, first, workspace.name, observer.name, observed.name, session.name
- )
-
- # Merge second
- second = _make_wr_payload(
- explicit=["C"],
- deductive=[{"conclusion": "Z", "premises": ["p1", "p2"]}],
- thinking="second",
- )
- await set_working_representation(
- db_session, second, workspace.name, observer.name, observed.name, session.name
- )
-
- # Fetch raw and assert structure
- raw = await get_working_representation_data(
- db_session, workspace.name, observer.name, observed.name, session.name
- )
- assert isinstance(raw, dict)
- final = raw["final_observations"]
- assert final["explicit"] == ["A", "B", "C"]
- assert final["deductive"][-1]["conclusion"] == "Z"
- assert final.get("thinking") == "second"
-
- # Validate it's stored under the derived key in SessionPeer.internal_metadata
- derived_key = f"{observer.name}_{observed.name}"
- stmt = select(models.SessionPeer).where(
- models.SessionPeer.peer_name == observer.name,
- models.SessionPeer.session_name == session.name,
- models.SessionPeer.workspace_name == workspace.name,
- )
- result = await db_session.execute(stmt)
- sp = result.scalar_one()
- assert derived_key in sp.internal_metadata
- assert sp.internal_metadata[derived_key]["final_observations"]["explicit"] == [
- "A",
- "B",
- "C",
- ]
-
-
-@pytest.mark.asyncio
-async def test_wr_string_roundtrip_then_structured_overrides(
- db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
-):
- """String WR stores as-is and later structured merge does not incorporate old string content."""
- workspace, peer = sample_data
- session = await crud.get_or_create_session(
- db_session,
- schemas.SessionCreate(
- name=str(generate_nanoid()), peers={peer.name: schemas.SessionPeerConfig()}
- ),
- workspace.name,
- )
- await db_session.commit()
-
- # Store legacy string representation first
- raw_string = "legacy working rep text"
- await set_working_representation(
- db_session, raw_string, workspace.name, peer.name, peer.name, session.name
- )
-
- # Ensure get returns the same string
- assert (
- await get_working_representation(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- == raw_string
- )
-
- # Now store structured representation; merge should ignore old string and just store new structured
- structured = _make_wr_payload(explicit=["X", "Y"], deductive=[])
- await set_working_representation(
- db_session, structured, workspace.name, peer.name, peer.name, session.name
- )
-
- raw = await get_working_representation_data(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- assert isinstance(raw, dict)
- assert raw["final_observations"]["explicit"] == ["X", "Y"]
-
-
-@pytest.mark.asyncio
-async def test_wr_missing_levels_and_empty_new_lists(
- db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
-):
- """Merge handles missing levels and empty new lists; result formatting is empty string."""
- workspace, peer = sample_data
- session = await crud.get_or_create_session(
- db_session,
- schemas.SessionCreate(
- name=str(generate_nanoid()), peers={peer.name: schemas.SessionPeerConfig()}
- ),
- workspace.name,
- )
- await db_session.commit()
-
- # Existing has only deductive; explicit missing
- existing = {
- "final_observations": {
- "deductive": [{"conclusion": "old", "premises": ["p"]}],
- },
- "message_id": "m-old",
- "created_at": datetime.datetime(
- 2024, 1, 2, tzinfo=datetime.timezone.utc
- ).isoformat(),
- }
- await set_working_representation(
- db_session, existing, workspace.name, peer.name, peer.name, session.name
- )
-
- # New has empty lists and no thinking
- new_payload = _make_wr_payload(
- explicit=[], deductive=[], thinking=None, message_id="m-new"
- )
- await set_working_representation(
- db_session, new_payload, workspace.name, peer.name, peer.name, session.name
- )
-
- raw = await get_working_representation_data(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- assert isinstance(raw, dict)
- final = raw["final_observations"]
- # Deductive remains as old (no additions), explicit remains missing/empty and thinking None
- assert final["deductive"] == [{"conclusion": "old", "premises": ["p"]}]
- assert final["explicit"] == []
- assert final.get("thinking") is None
- # Formatter includes the remaining deductive observation
- formatted = await get_working_representation(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- assert "DEDUCTIVE OBSERVATIONS:" in formatted
- assert "- old (based on: p)" in formatted
-
-
-@pytest.mark.asyncio
-async def test_wr_trim_boundary_exact_25_plus_one(
- db_session: AsyncSession,
- sample_data: tuple[models.Workspace, models.Peer],
- monkeypatch: pytest.MonkeyPatch,
-):
- """Exactly 25 existing + 1 new yields last 25, preserving order and including the new last element."""
- workspace, peer = sample_data
- monkeypatch.setattr(
- settings.DERIVER, "WORKING_REPRESENTATION_MAX_OBSERVATIONS", 25, raising=False
- )
- LIMIT = settings.DERIVER.WORKING_REPRESENTATION_MAX_OBSERVATIONS
- session = await crud.get_or_create_session(
- db_session,
- schemas.SessionCreate(
- name=str(generate_nanoid()), peers={peer.name: schemas.SessionPeerConfig()}
- ),
- workspace.name,
- )
- await db_session.commit()
-
- existing = _make_wr_payload(
- explicit=[f"E{i}" for i in range(LIMIT)],
- deductive=[{"conclusion": f"D{i}", "premises": []} for i in range(LIMIT)],
- )
- await set_working_representation(
- db_session, existing, workspace.name, peer.name, peer.name, session.name
- )
-
- new_payload = _make_wr_payload(
- explicit=[f"E{LIMIT}"],
- deductive=[{"conclusion": f"D{LIMIT}", "premises": []}],
- thinking="t2",
- message_id="m2",
- )
- await set_working_representation(
- db_session, new_payload, workspace.name, peer.name, peer.name, session.name
- )
-
- raw = await get_working_representation_data(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- assert isinstance(raw, dict)
- final = raw["final_observations"]
- assert final["explicit"] == [f"E{i}" for i in range(1, LIMIT + 1)]
- assert final["explicit"][-1] == f"E{LIMIT}"
- assert len(final["deductive"]) == LIMIT
- assert final["deductive"][-1]["conclusion"] == f"D{LIMIT}"
-
-
-@pytest.mark.asyncio
-async def test_wr_formatting_rules_mixed_observation_types(
- db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
-):
- """Formatting includes section headers, bullets, premise display, and content fallback."""
- workspace, peer = sample_data
- session = await crud.get_or_create_session(
- db_session,
- schemas.SessionCreate(
- name=str(generate_nanoid()), peers={peer.name: schemas.SessionPeerConfig()}
- ),
- workspace.name,
- )
- await db_session.commit()
-
- payload = {
- "final_observations": {
- "explicit": ["likes pizza", {"content": "runs daily"}],
- "deductive": [
- {
- "conclusion": "is healthy",
- "premises": ["runs daily", "eats veggies"],
- },
- {"content": "fallback without conclusion"},
- ],
- "thinking": "t",
- },
- "message_id": "m1",
- "created_at": datetime.datetime.now(datetime.timezone.utc).isoformat(),
- }
- await set_working_representation(
- db_session, payload, workspace.name, peer.name, peer.name, session.name
- )
-
- formatted = await get_working_representation(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- # Section headers present
- assert formatted.splitlines()[0] == "EXPLICIT OBSERVATIONS:"
- assert "DEDUCTIVE OBSERVATIONS:" in formatted
- # Bullets present for strings and dict content fallback
- assert "- likes pizza" in formatted
- assert "- runs daily" in formatted
- # Deductive with premises shows based on
- assert "is healthy (based on: runs daily; eats veggies)" in formatted
- # Deductive dict without conclusion falls back to content
- assert "- fallback without conclusion" in formatted
-
-
-@pytest.mark.asyncio
-async def test_wr_legacy_key_fallback_for_self_representation(
- db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
-):
- """If only the legacy key is present, data retrieval falls back appropriately."""
- workspace, peer = sample_data
- session = await crud.get_or_create_session(
- db_session,
- schemas.SessionCreate(
- name=str(generate_nanoid()), peers={peer.name: schemas.SessionPeerConfig()}
- ),
- workspace.name,
- )
- await db_session.commit()
-
- # Manually write legacy key on SessionPeer
- stmt = select(models.SessionPeer).where(
- models.SessionPeer.peer_name == peer.name,
- models.SessionPeer.session_name == session.name,
- models.SessionPeer.workspace_name == workspace.name,
- )
- result = await db_session.execute(stmt)
- sp: models.SessionPeer = result.scalar_one()
- # Assign a new dict so JSON mutation is tracked & persisted
- sp.internal_metadata = {
- **(sp.internal_metadata or {}),
- "global_representation": "legacy-global",
- }
- await db_session.commit()
-
- # Retrieval should see legacy value
- raw = await get_working_representation_data(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- assert raw == "legacy-global"
-
-
-@pytest.mark.asyncio
-async def test_wr_empty_both_levels_formats_empty_string(
- db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
-):
- """When both explicit and deductive are empty after merge, formatted string is empty."""
- workspace, peer = sample_data
- session = await crud.get_or_create_session(
- db_session,
- schemas.SessionCreate(
- name=str(generate_nanoid()), peers={peer.name: schemas.SessionPeerConfig()}
- ),
- workspace.name,
- )
- await db_session.commit()
-
- empty_payload = _make_wr_payload(explicit=[], deductive=[], thinking=None)
- await set_working_representation(
- db_session, empty_payload, workspace.name, peer.name, peer.name, session.name
- )
-
- formatted = await get_working_representation(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- assert formatted == ""
-
-
-@pytest.mark.asyncio
-async def test_wr_existing_dict_without_final_observations(
- db_session: AsyncSession, sample_data: tuple[models.Workspace, models.Peer]
-):
- """Existing dict lacking final_observations merges cleanly with new structured payload."""
- workspace, peer = sample_data
- session = await crud.get_or_create_session(
- db_session,
- schemas.SessionCreate(
- name=str(generate_nanoid()), peers={peer.name: schemas.SessionPeerConfig()}
- ),
- workspace.name,
- )
- await db_session.commit()
-
- await set_working_representation(
- db_session,
- {"some": "field"},
- workspace.name,
- peer.name,
- peer.name,
- session.name,
- )
- new_payload = _make_wr_payload(
- explicit=["n1"], deductive=[{"conclusion": "c1", "premises": []}]
- )
- await set_working_representation(
- db_session, new_payload, workspace.name, peer.name, peer.name, session.name
- )
-
- raw = await get_working_representation_data(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- assert isinstance(raw, dict)
- final = raw["final_observations"]
- assert final["explicit"] == ["n1"]
- assert final["deductive"][0]["conclusion"] == "c1"
-
-
-@pytest.mark.asyncio
-async def test_wr_deductive_trim_when_no_new_items(
- db_session: AsyncSession,
- sample_data: tuple[models.Workspace, models.Peer],
- monkeypatch: pytest.MonkeyPatch,
-):
- """If existing deductive > 25 and no new entries, merge still trims to last 25."""
- workspace, peer = sample_data
- monkeypatch.setattr(
- settings.DERIVER, "WORKING_REPRESENTATION_MAX_OBSERVATIONS", 25, raising=False
- )
- LIMIT = settings.DERIVER.WORKING_REPRESENTATION_MAX_OBSERVATIONS
- session = await crud.get_or_create_session(
- db_session,
- schemas.SessionCreate(
- name=str(generate_nanoid()), peers={peer.name: schemas.SessionPeerConfig()}
- ),
- workspace.name,
- )
- await db_session.commit()
-
- over = LIMIT + 5
- existing = _make_wr_payload(
- explicit=[],
- deductive=[{"conclusion": f"D{i}", "premises": []} for i in range(over)],
- )
- await set_working_representation(
- db_session, existing, workspace.name, peer.name, peer.name, session.name
- )
-
- # Merge empty new; should trim to last 25 existing
- new_payload = _make_wr_payload(explicit=[], deductive=[])
- await set_working_representation(
- db_session, new_payload, workspace.name, peer.name, peer.name, session.name
- )
-
- raw = await get_working_representation_data(
- db_session, workspace.name, peer.name, peer.name, session.name
- )
- assert isinstance(raw, dict)
- final = raw["final_observations"]
- assert len(final["deductive"]) == LIMIT
- # Oldest retained index is over - LIMIT
- oldest_kept = over - LIMIT
- assert final["deductive"][0]["conclusion"] == f"D{oldest_kept}"
- assert final["deductive"][-1]["conclusion"] == f"D{over - 1}"
+ assert isinstance(rep, Representation)
+ assert [e.content for e in rep.explicit] == ["A"]
+ assert rep.deductive[0].conclusion == "C"
+ assert rep.deductive[0].premises == ["P1"]
+ assert rep.explicit[0].created_at == timestamp
+ assert rep.deductive[0].created_at == timestamp
diff --git a/tests/dreamer/__init__.py b/tests/dreamer/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/tests/dreamer/test_dream_scheduler.py b/tests/dreamer/test_dream_scheduler.py
new file mode 100644
index 00000000..92331b95
--- /dev/null
+++ b/tests/dreamer/test_dream_scheduler.py
@@ -0,0 +1,64 @@
+"""Tests for DreamScheduler singleton pattern."""
+
+import asyncio
+import contextlib
+
+from src.dreamer.dream_scheduler import DreamScheduler
+
+
+def test_dream_scheduler_singleton():
+ """Test that DreamScheduler implements proper singleton pattern."""
+ # Reset singleton state
+ DreamScheduler.reset_singleton()
+
+ # Create first instance
+ scheduler1 = DreamScheduler()
+
+ # Create second instance
+ scheduler2 = DreamScheduler()
+
+ # Both should be the same instance
+ assert scheduler1 is scheduler2
+
+ # Should share the same pending_dreams dict
+ assert scheduler1.pending_dreams is scheduler2.pending_dreams
+
+
+async def test_dream_scheduler_initialized_once():
+ """Test that DreamScheduler is only initialized once."""
+ # Reset singleton state
+ DreamScheduler.reset_singleton()
+
+ # Create first instance
+ scheduler1 = DreamScheduler()
+
+ # Create a dummy task to add to pending_dreams
+ async def dummy_task():
+ pass
+
+ task = asyncio.create_task(dummy_task())
+ scheduler1.pending_dreams["test_key"] = task
+
+ # Create second instance
+ scheduler2 = DreamScheduler()
+
+ # Second instance should have the same data as first
+ assert "test_key" in scheduler2.pending_dreams
+ assert scheduler2.pending_dreams["test_key"] is task
+
+ # Cleanup
+ task.cancel()
+ with contextlib.suppress(asyncio.CancelledError):
+ await task
+
+
+def test_dream_scheduler_multiple_instances():
+ """Test that creating multiple instances doesn't reset state."""
+ # Reset singleton state
+ DreamScheduler.reset_singleton()
+
+ instances = [DreamScheduler() for _ in range(5)]
+
+ # All instances should be the same
+ for instance in instances[1:]:
+ assert instance is instances[0]
diff --git a/tests/dreamer/test_queue_manager_singleton.py b/tests/dreamer/test_queue_manager_singleton.py
new file mode 100644
index 00000000..cdcd14b8
--- /dev/null
+++ b/tests/dreamer/test_queue_manager_singleton.py
@@ -0,0 +1,56 @@
+"""Tests that QueueManager instances share the same DreamScheduler singleton."""
+
+import asyncio
+import contextlib
+
+from src.deriver.queue_manager import QueueManager
+from src.dreamer.dream_scheduler import DreamScheduler
+
+
+def test_queue_manager_shares_dream_scheduler():
+ """Test that multiple QueueManager instances share the same DreamScheduler."""
+ # Reset singleton state
+ DreamScheduler.reset_singleton()
+
+ # Create first QueueManager
+ manager1 = QueueManager()
+
+ # Create second QueueManager
+ manager2 = QueueManager()
+
+ # Both should have the same DreamScheduler instance
+ assert manager1.dream_scheduler is manager2.dream_scheduler
+
+ # Should share the same pending_dreams dict
+ assert (
+ manager1.dream_scheduler.pending_dreams
+ is manager2.dream_scheduler.pending_dreams
+ )
+
+
+async def test_queue_manager_preserves_dream_scheduler_state():
+ """Test that creating a new QueueManager doesn't reset DreamScheduler state."""
+ # Reset singleton state
+ DreamScheduler.reset_singleton()
+
+ # Create first QueueManager and modify scheduler state
+ manager1 = QueueManager()
+
+ # Create a dummy task to add to pending_dreams
+ async def dummy_task():
+ pass
+
+ task = asyncio.create_task(dummy_task())
+ manager1.dream_scheduler.pending_dreams["test_key"] = task
+
+ # Create second QueueManager
+ manager2 = QueueManager()
+
+ # Second manager should see the first manager's state
+ assert "test_key" in manager2.dream_scheduler.pending_dreams
+ assert manager2.dream_scheduler.pending_dreams["test_key"] is task
+
+ # Cleanup
+ task.cancel()
+ with contextlib.suppress(asyncio.CancelledError):
+ await task
diff --git a/tests/integration/test_enqueue.py b/tests/integration/test_enqueue.py
index 5e07d3d6..0f541704 100644
--- a/tests/integration/test_enqueue.py
+++ b/tests/integration/test_enqueue.py
@@ -223,15 +223,15 @@ class TestEnqueueFunction:
for _ in range(NUM_MESSAGES):
expected_payloads.append(
{
- "sender_name": test_peer1.name,
- "target_name": test_peer1.name,
+ "observed": test_peer1.name,
+ "observer": test_peer1.name,
"task_type": "representation",
}
)
actual_payloads = [
{
- "sender_name": item.payload.get("sender_name"),
- "target_name": item.payload.get("target_name"),
+ "observed": item.payload.get("observed"),
+ "observer": item.payload.get("observer"),
"task_type": item.payload.get("task_type"),
}
for item in queue_items
@@ -302,22 +302,22 @@ class TestEnqueueFunction:
for _ in range(NUM_MESSAGES):
expected_payloads.append(
{
- "sender_name": test_peer1.name,
- "target_name": test_peer1.name,
+ "observed": test_peer1.name,
+ "observer": test_peer1.name,
"task_type": "representation",
}
)
expected_payloads.append(
{
- "sender_name": test_peer1.name,
- "target_name": test_peer2.name,
+ "observed": test_peer1.name,
+ "observer": test_peer2.name,
"task_type": "representation",
}
)
actual_payloads = [
{
- "sender_name": item.payload.get("sender_name"),
- "target_name": item.payload.get("target_name"),
+ "observed": item.payload.get("observed"),
+ "observer": item.payload.get("observer"),
"task_type": item.payload.get("task_type"),
}
for item in queue_items
@@ -396,22 +396,22 @@ class TestEnqueueFunction:
for _ in range(NUM_MESSAGES):
expected_payloads.append(
{
- "sender_name": test_peer1.name,
- "target_name": test_peer1.name,
+ "observed": test_peer1.name,
+ "observer": test_peer1.name,
"task_type": "representation",
}
)
expected_payloads.append(
{
- "sender_name": test_peer1.name,
- "target_name": observing_peer.name,
+ "observed": test_peer1.name,
+ "observer": observing_peer.name,
"task_type": "representation",
}
)
actual_payloads = [
{
- "sender_name": item.payload.get("sender_name"),
- "target_name": item.payload.get("target_name"),
+ "observed": item.payload.get("observed"),
+ "observer": item.payload.get("observer"),
"task_type": item.payload.get("task_type"),
}
for item in queue_items
@@ -423,7 +423,7 @@ class TestEnqueueFunction:
assert expected in actual_payloads
assert unobserving_peer.name not in [
- item.payload.get("target_name") for item in queue_items
+ item.payload.get("observer") for item in queue_items
]
@pytest.mark.asyncio
@@ -545,25 +545,25 @@ class TestEnqueueFunction:
expected_payloads.append(
{
"task_type": "representation",
- "sender_name": sender,
- "target_name": sender,
+ "observed": sender,
+ "observer": sender,
}
)
# representation for observer_peer (observe_others=True)
expected_payloads.append(
{
"task_type": "representation",
- "sender_name": sender,
- "target_name": observer_peer.name,
+ "observed": sender,
+ "observer": observer_peer.name,
}
)
- # Extract actual payloads (task_type, sender_name, target_name) from queue_items
+ # Extract actual payloads (task_type, observed, observer) from queue_items
actual_payloads = [
{
"task_type": item[0].payload.get("task_type"),
- "sender_name": item[0].payload.get("sender_name"),
- "target_name": item[0].payload.get("target_name"),
+ "observed": item[0].payload.get("observed"),
+ "observer": item[0].payload.get("observer"),
}
for item in queue_items
]
@@ -645,20 +645,20 @@ class TestEnqueueFunction:
expected_payloads = [
{
- "sender_name": sender_peer.name,
- "target_name": sender_peer.name,
+ "observed": sender_peer.name,
+ "observer": sender_peer.name,
"task_type": "representation",
},
{
- "sender_name": sender_peer.name,
- "target_name": observer_peer.name,
+ "observed": sender_peer.name,
+ "observer": observer_peer.name,
"task_type": "representation",
},
]
actual_payloads = [
{
- "sender_name": item.payload.get("sender_name"),
- "target_name": item.payload.get("target_name"),
+ "observed": item.payload.get("observed"),
+ "observer": item.payload.get("observer"),
"task_type": item.payload.get("task_type"),
}
for item in queue_items
@@ -745,10 +745,10 @@ class TestEnqueueFunction:
queue_items = result.scalars().all()
# Verify observer_who_left is NOT in the target names
- target_names = [item.payload.get("target_name") for item in queue_items]
- assert observer_who_left.name not in target_names
- assert observer_who_stayed.name in target_names
- assert sender_peer.name in target_names
+ observers = [item.payload.get("observer") for item in queue_items]
+ assert observer_who_left.name not in observers
+ assert observer_who_stayed.name in observers
+ assert sender_peer.name in observers
@pytest.mark.asyncio
@patch("src.deriver.enqueue.tracked_db")
@@ -807,20 +807,20 @@ class TestEnqueueFunction:
expected_payloads = [
{
- "sender_name": existing_peer.name,
- "target_name": existing_peer.name,
+ "observed": existing_peer.name,
+ "observer": existing_peer.name,
"task_type": "representation",
},
{
- "sender_name": existing_peer.name,
- "target_name": observer_peer.name,
+ "observed": existing_peer.name,
+ "observer": observer_peer.name,
"task_type": "representation",
},
]
actual_payloads = [
{
- "sender_name": item.payload.get("sender_name"),
- "target_name": item.payload.get("target_name"),
+ "observed": item.payload.get("observed"),
+ "observer": item.payload.get("observer"),
"task_type": item.payload.get("task_type"),
}
for item in queue_items
@@ -932,20 +932,20 @@ class TestEnqueueFunction:
expected_payloads = [
{
- "sender_name": sender_peer.name,
- "target_name": sender_peer.name,
+ "observed": sender_peer.name,
+ "observer": sender_peer.name,
"task_type": "representation",
},
{
- "sender_name": sender_peer.name,
- "target_name": active_observer.name,
+ "observed": sender_peer.name,
+ "observer": active_observer.name,
"task_type": "representation",
},
]
actual_payloads = [
{
- "sender_name": item.payload.get("sender_name"),
- "target_name": item.payload.get("target_name"),
+ "observed": item.payload.get("observed"),
+ "observer": item.payload.get("observer"),
"task_type": item.payload.get("task_type"),
}
for item in queue_items
@@ -956,10 +956,10 @@ class TestEnqueueFunction:
assert expected in actual_payloads
# Verify inactive peers are not in target names
- target_names = [item.payload.get("target_name") for item in queue_items]
- assert inactive_observer.name not in target_names
- assert inactive_non_observer.name not in target_names
- assert active_non_observer.name not in target_names
+ observers = [item.payload.get("observer") for item in queue_items]
+ assert inactive_observer.name not in observers
+ assert inactive_non_observer.name not in observers
+ assert active_non_observer.name not in observers
class TestGetEffectiveObserveMeFunction:
@@ -1058,14 +1058,14 @@ class TestGetEffectiveObserveMeFunction:
if peer_config is None:
# Test missing sender
peers_with_configuration = {}
- sender_name = "missing_sender"
+ observed = "missing_sender"
else:
peers_with_configuration = {
f"sender_{i}": [peer_config or {}, session_config or {}]
}
- sender_name = f"sender_{i}"
+ observed = f"sender_{i}"
- result = get_effective_observe_me(sender_name, peers_with_configuration)
+ result = get_effective_observe_me(observed, peers_with_configuration)
assert (
result == expected
), f"Test case {i} failed: peer_config={peer_config}, session_config={session_config}, expected={expected}, got={result}"
@@ -1192,8 +1192,8 @@ class TestAdvancedEnqueueEdgeCases:
queue_items = result.scalars().all()
assert len(queue_items) == 1
- assert queue_items[0].payload["sender_name"] == sender_peer.name
- assert queue_items[0].payload["target_name"] == sender_peer.name
+ assert queue_items[0].payload["observed"] == sender_peer.name
+ assert queue_items[0].payload["observer"] == sender_peer.name
assert queue_items[0].payload["task_type"] == "representation"
@pytest.mark.asyncio
@@ -1278,8 +1278,8 @@ class TestAdvancedEnqueueEdgeCases:
queue_items = result.scalars().all()
assert len(queue_items) == 1
- assert queue_items[0].payload["sender_name"] == sender_peer.name
- assert queue_items[0].payload["target_name"] == sender_peer.name
+ assert queue_items[0].payload["observed"] == sender_peer.name
+ assert queue_items[0].payload["observer"] == sender_peer.name
assert queue_items[0].payload["task_type"] == "representation"
@pytest.mark.asyncio
@@ -1337,20 +1337,20 @@ class TestAdvancedEnqueueEdgeCases:
expected_payloads = [
{
- "sender_name": existing_peer.name,
- "target_name": existing_peer.name,
+ "observed": existing_peer.name,
+ "observer": existing_peer.name,
"task_type": "representation",
},
{
- "sender_name": existing_peer.name,
- "target_name": observer_peer.name,
+ "observed": existing_peer.name,
+ "observer": observer_peer.name,
"task_type": "representation",
},
]
actual_payloads = [
{
- "sender_name": item.payload.get("sender_name"),
- "target_name": item.payload.get("target_name"),
+ "observed": item.payload.get("observed"),
+ "observer": item.payload.get("observer"),
"task_type": item.payload.get("task_type"),
}
for item in queue_items
diff --git a/tests/integration/test_representation.py b/tests/integration/test_representation.py
new file mode 100644
index 00000000..15d8dd37
--- /dev/null
+++ b/tests/integration/test_representation.py
@@ -0,0 +1,533 @@
+"""
+Integration tests for the Representation class and related workflows.
+
+This test suite covers the full representation workflow including:
+- Document creation with embedding and duplicate detection
+- Representation building from documents
+- Representation merging and diffing operations
+- Working representation retrieval with different strategies
+"""
+
+from datetime import datetime, timezone
+from unittest.mock import patch
+
+import pytest
+from nanoid import generate as generate_nanoid
+from sqlalchemy import select
+from sqlalchemy.ext.asyncio import AsyncSession
+
+from src import crud, models
+from src.utils.representation import (
+ DeductiveObservation,
+ DeductiveObservationBase,
+ ExplicitObservation,
+ ExplicitObservationBase,
+ PromptRepresentation,
+ Representation,
+)
+
+
+@pytest.fixture
+def fixed_embedding_vector() -> list[float]:
+ """
+ Fixture providing a deterministic embedding vector for hermetic tests.
+
+ Returns a 1536-dimensional vector with predictable values to avoid
+ network calls to external embedding services during testing.
+ """
+ # Create a deterministic 1536-dimensional embedding vector
+ # Using a simple pattern that's easy to verify in tests
+ return [0.1 * (i % 10) for i in range(1536)]
+
+
+@pytest.mark.asyncio
+class TestRepresentationWorkflow:
+ """Test suite for the complete representation workflow"""
+
+ async def create_test_workspace_and_peer(
+ self,
+ db_session: AsyncSession,
+ workspace_name: str | None = None,
+ peer_name: str | None = None,
+ ) -> tuple[models.Workspace, models.Peer]:
+ """Helper to create test workspace and peer"""
+ workspace_name = workspace_name or generate_nanoid()
+ peer_name = peer_name or generate_nanoid()
+
+ workspace = models.Workspace(name=workspace_name)
+ db_session.add(workspace)
+ await db_session.flush()
+
+ peer = models.Peer(name=peer_name, workspace_name=workspace_name)
+ db_session.add(peer)
+ await db_session.flush()
+
+ return workspace, peer
+
+ async def create_test_session(
+ self,
+ db_session: AsyncSession,
+ workspace: models.Workspace,
+ session_name: str | None = None,
+ ) -> models.Session:
+ """Helper to create test session"""
+ session_name = session_name or generate_nanoid()
+
+ session = models.Session(
+ name=session_name,
+ workspace_name=workspace.name,
+ )
+ db_session.add(session)
+ await db_session.flush()
+
+ return session
+
+ async def test_representation_class_creation_and_operations(self):
+ """Test basic Representation class operations"""
+ # Create explicit observations
+ explicit_obs1 = ExplicitObservation(
+ content="User likes dogs",
+ created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="test_session",
+ )
+ explicit_obs2 = ExplicitObservation(
+ content="User has a pet named Rover",
+ created_at=datetime(2025, 1, 1, 12, 1, 0, tzinfo=timezone.utc),
+ message_ids=[(2, 2)],
+ session_name="test_session",
+ )
+
+ # Create deductive observations
+ deductive_obs1 = DeductiveObservation(
+ conclusion="User probably has a dog named Rover",
+ premises=["User likes dogs", "User has a pet named Rover"],
+ created_at=datetime(2025, 1, 1, 12, 2, 0, tzinfo=timezone.utc),
+ message_ids=[(3, 3)],
+ session_name="test_session",
+ )
+
+ # Create representation
+ representation = Representation(
+ explicit=[explicit_obs1, explicit_obs2], deductive=[deductive_obs1]
+ )
+
+ # Test basic properties
+ assert not representation.is_empty()
+ assert len(representation.explicit) == 2
+ assert len(representation.deductive) == 1
+
+ # Test string formatting
+ str_output = str(representation)
+ assert "EXPLICIT:" in str_output
+ assert "DEDUCTIVE:" in str_output
+ assert "User likes dogs" in str_output
+ assert "User probably has a dog named Rover" in str_output
+
+ # Test no-timestamp formatting
+ no_timestamp_output = representation.str_no_timestamps()
+ assert "User likes dogs" in no_timestamp_output
+ assert "[" not in no_timestamp_output
+
+ # Test markdown formatting
+ markdown_output = representation.format_as_markdown()
+ assert "## Explicit Observations" in markdown_output
+ assert "## Deductive Observations" in markdown_output
+ assert "**Conclusion**:" in markdown_output
+ assert "**Premises**:" in markdown_output
+
+ async def test_representation_merging_and_diffing(self):
+ """Test representation merge and diff operations"""
+ # Create first representation
+ rep1 = Representation(
+ explicit=[
+ ExplicitObservation(
+ content="User likes cats",
+ created_at=datetime(2025, 1, 1, 10, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="session1",
+ )
+ ]
+ )
+
+ # Create second representation with some overlap
+ rep2 = Representation(
+ explicit=[
+ ExplicitObservation(
+ content="User likes cats", # Duplicate
+ created_at=datetime(2025, 1, 1, 10, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="session1",
+ ),
+ ExplicitObservation(
+ content="User likes dogs", # New
+ created_at=datetime(2025, 1, 1, 11, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(2, 2)],
+ session_name="session1",
+ ),
+ ]
+ )
+
+ # Test diff - should only return new observations
+ diff = rep1.diff_representation(rep2)
+ assert len(diff.explicit) == 1
+ assert diff.explicit[0].content == "User likes dogs"
+
+ # Test merge - should deduplicate and preserve order
+ rep1.merge_representation(rep2)
+ assert len(rep1.explicit) == 2
+ # Should be sorted by created_at
+ assert rep1.explicit[0].content == "User likes cats"
+ assert rep1.explicit[1].content == "User likes dogs"
+
+ # Test merge with max_observations limit
+ rep3 = Representation(
+ explicit=[
+ ExplicitObservation(
+ content="User likes birds",
+ created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(3, 3)],
+ session_name="session1",
+ )
+ ]
+ )
+
+ rep1.merge_representation(rep3, max_observations=2)
+ assert len(rep1.explicit) == 2
+ # Should keep the most recent observations (FIFO)
+ assert rep1.explicit[0].content == "User likes dogs"
+ assert rep1.explicit[1].content == "User likes birds"
+
+
+@pytest.mark.asyncio
+class TestDocumentCreationWorkflow:
+ """Test document creation with embedding and duplicate detection"""
+
+ async def test_get_working_representation_with_semantic_query(
+ self, db_session: AsyncSession
+ ):
+ """Test working representation retrieval with semantic query"""
+ workspace, observer_peer = await self.create_test_workspace_and_peer(db_session)
+ _, observed_peer = await self.create_test_workspace_and_peer(
+ db_session, workspace.name
+ )
+
+ # Mock semantic search to return specific documents
+ from src.models import Document
+
+ mock_document = Document(
+ id="test_doc_id",
+ workspace_name=workspace.name,
+ observer=observer_peer.name,
+ observed=observed_peer.name,
+ content="User likes dogs",
+ session_name="test_session",
+ internal_metadata={
+ "level": "explicit",
+ "message_ids": [(1, 1)],
+ "session_name": "test_session",
+ },
+ created_at=datetime.now(timezone.utc),
+ )
+
+ with patch(
+ "src.crud.representation.RepresentationManager._query_documents_semantic"
+ ) as mock_semantic:
+ mock_semantic.return_value = [mock_document]
+
+ representation = await crud.get_working_representation(
+ workspace.name,
+ include_semantic_query="pets dogs animals",
+ observer=observer_peer.name,
+ observed=observed_peer.name,
+ )
+
+ # Should have called semantic search
+ mock_semantic.assert_called_once()
+ assert len(representation.explicit) == 1
+ assert representation.explicit[0].content == "User likes dogs"
+
+ async def test_get_working_representation_with_most_derived(
+ self, db_session: AsyncSession
+ ):
+ """Test working representation retrieval prioritizing most derived observations"""
+ workspace, observer_peer = await self.create_test_workspace_and_peer(db_session)
+ _, observed_peer = await self.create_test_workspace_and_peer(
+ db_session, workspace.name
+ )
+
+ session = await self.create_test_session(db_session, workspace)
+
+ # Create collection first - need to do it directly since mock doesn't persist to DB
+ collection = models.Collection(
+ observer=observer_peer.name,
+ observed=observed_peer.name,
+ workspace_name=workspace.name,
+ )
+ db_session.add(collection)
+ await db_session.flush()
+
+ # Create document with high times_derived count
+ highly_derived_doc = models.Document(
+ workspace_name=workspace.name,
+ observer=observer_peer.name,
+ observed=observed_peer.name,
+ content="Highly derived observation",
+ session_name=session.name,
+ internal_metadata={"level": "explicit", "times_derived": 5},
+ embedding=[0.1] * 1536,
+ )
+ db_session.add(highly_derived_doc)
+
+ # Create document with lower times_derived count
+ less_derived_doc = models.Document(
+ workspace_name=workspace.name,
+ observer=observer_peer.name,
+ observed=observed_peer.name,
+ content="Less derived observation",
+ session_name=session.name,
+ internal_metadata={"level": "explicit", "times_derived": 2},
+ embedding=[0.2] * 1536,
+ )
+ db_session.add(less_derived_doc)
+
+ await db_session.commit()
+
+ # Retrieve with most_derived=True
+ representation = await crud.get_working_representation(
+ workspace.name,
+ include_most_derived=True,
+ observer=observer_peer.name,
+ observed=observed_peer.name,
+ )
+
+ # Should prioritize highly derived observation
+ assert len(representation.explicit) >= 1
+ # The highly derived observation should be included
+ contents = [obs.content for obs in representation.explicit]
+ assert "Highly derived observation" in contents
+
+ async def test_representation_from_documents(self):
+ """Test converting documents to representation"""
+ # Create test documents
+ explicit_doc = models.Document(
+ workspace_name="test_workspace",
+ observer="test_peer",
+ observed="test_peer",
+ content="User said they like programming",
+ internal_metadata={
+ "level": "explicit",
+ "message_ids": [(1, 1)],
+ },
+ session_name="test_session",
+ embedding=[0.1] * 1536,
+ created_at=datetime(2025, 1, 1, 10, 0, 0, tzinfo=timezone.utc),
+ )
+
+ deductive_doc = models.Document(
+ workspace_name="test_workspace",
+ observer="test_peer",
+ observed="test_peer",
+ content="User is likely a software developer",
+ internal_metadata={
+ "level": "deductive",
+ "message_ids": [(1, 1)],
+ "premises": ["User said they like programming"],
+ },
+ session_name="test_session",
+ embedding=[0.2] * 1536,
+ created_at=datetime(2025, 1, 1, 10, 1, 0, tzinfo=timezone.utc),
+ )
+
+ # Convert to representation
+ representation = Representation.from_documents([explicit_doc, deductive_doc])
+
+ assert len(representation.explicit) == 1
+ assert len(representation.deductive) == 1
+
+ explicit_obs = representation.explicit[0]
+ assert explicit_obs.content == "User said they like programming"
+ assert explicit_obs.message_ids == [(1, 1)]
+ assert explicit_obs.session_name == "test_session"
+
+ deductive_obs = representation.deductive[0]
+ assert deductive_obs.conclusion == "User is likely a software developer"
+ assert deductive_obs.premises == ["User said they like programming"]
+ assert deductive_obs.message_ids == [(1, 1)]
+ assert deductive_obs.session_name == "test_session"
+
+ async def create_test_workspace_and_peer(
+ self, db_session: AsyncSession, workspace_name: str | None = None
+ ) -> tuple[models.Workspace, models.Peer]:
+ """Helper to create test workspace and peer"""
+ workspace_name = workspace_name or generate_nanoid()
+ peer_name = generate_nanoid()
+
+ # Check if workspace already exists to avoid uniqueness constraint
+ workspace = (
+ await db_session.execute(
+ select(models.Workspace).where(models.Workspace.name == workspace_name)
+ )
+ ).scalar_one_or_none()
+
+ if workspace is None:
+ workspace = models.Workspace(name=workspace_name)
+ db_session.add(workspace)
+ await db_session.flush()
+
+ peer = models.Peer(name=peer_name, workspace_name=workspace_name)
+ db_session.add(peer)
+ await db_session.flush()
+
+ return workspace, peer
+
+ async def create_test_session(
+ self, db_session: AsyncSession, workspace: models.Workspace
+ ) -> models.Session:
+ """Helper to create test session"""
+ session_name = generate_nanoid()
+
+ session = models.Session(
+ name=session_name,
+ workspace_name=workspace.name,
+ )
+ db_session.add(session)
+ await db_session.flush()
+
+ return session
+
+
+@pytest.mark.asyncio
+class TestPromptRepresentationConversion:
+ """Test conversion between PromptRepresentation and Representation"""
+
+ async def test_prompt_representation_to_representation(self):
+ """Test converting PromptRepresentation to Representation"""
+ prompt_rep = PromptRepresentation(
+ explicit=[
+ ExplicitObservationBase(content="User likes coffee"),
+ ExplicitObservationBase(content="User works remotely"),
+ ],
+ deductive=[
+ DeductiveObservationBase(
+ conclusion="User probably works from a coffee shop sometimes",
+ premises=["User likes coffee", "User works remotely"],
+ )
+ ],
+ )
+
+ timestamp = datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc)
+
+ representation = Representation.from_prompt_representation(
+ prompt_rep,
+ message_ids=(123, 123),
+ session_name="test_session",
+ created_at=timestamp,
+ )
+
+ assert len(representation.explicit) == 2
+ assert len(representation.deductive) == 1
+
+ # Check explicit observations
+ assert representation.explicit[0].content == "User likes coffee"
+ assert representation.explicit[0].message_ids == [(123, 123)]
+ assert representation.explicit[0].session_name == "test_session"
+ assert representation.explicit[1].content == "User works remotely"
+ assert representation.explicit[0].created_at == timestamp
+
+ # Check deductive observation
+ deductive_obs = representation.deductive[0]
+ assert (
+ deductive_obs.conclusion
+ == "User probably works from a coffee shop sometimes"
+ )
+ assert deductive_obs.premises == ["User likes coffee", "User works remotely"]
+ assert deductive_obs.message_ids == [(123, 123)]
+ assert deductive_obs.session_name == "test_session"
+ assert deductive_obs.created_at == timestamp
+
+ async def test_empty_prompt_representation_conversion(self):
+ """Test converting empty PromptRepresentation"""
+ empty_prompt_rep = PromptRepresentation()
+ representation = Representation.from_prompt_representation(
+ empty_prompt_rep,
+ message_ids=(1, 1),
+ session_name="test",
+ created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+ )
+
+ assert representation.is_empty()
+ assert len(representation.explicit) == 0
+ assert len(representation.deductive) == 0
+
+
+@pytest.mark.asyncio
+class TestRepresentationHashingAndEquality:
+ """Test hashing and equality behavior for observations"""
+
+ async def test_explicit_observation_equality(self):
+ """Test ExplicitObservation equality and hashing"""
+ obs1 = ExplicitObservation(
+ content="Test content",
+ created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="session1",
+ )
+
+ obs2 = ExplicitObservation(
+ content="Test content",
+ created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="session1",
+ )
+
+ obs3 = ExplicitObservation(
+ content="Different content",
+ created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="session1",
+ )
+
+ # Test equality
+ assert obs1 == obs2
+ assert obs1 != obs3
+ assert obs1 != "not an observation"
+
+ # Test hashing (should be able to use in sets)
+ obs_set = {obs1, obs2, obs3}
+ assert len(obs_set) == 2 # obs1 and obs2 are duplicates
+
+ async def test_deductive_observation_equality(self):
+ """Test DeductiveObservation equality and hashing"""
+ obs1 = DeductiveObservation(
+ conclusion="Test conclusion",
+ premises=["premise1", "premise2"],
+ created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="session1",
+ )
+
+ obs2 = DeductiveObservation(
+ conclusion="Test conclusion",
+ premises=["premise1", "premise2"],
+ created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="session1",
+ )
+
+ obs3 = DeductiveObservation(
+ conclusion="Different conclusion",
+ premises=["premise1", "premise2"],
+ created_at=datetime(2025, 1, 1, 12, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="session1",
+ )
+
+ # Test equality
+ assert obs1 == obs2
+ assert obs1 != obs3
+ assert obs1 != "not an observation"
+
+ # Test hashing
+ obs_set = {obs1, obs2, obs3}
+ assert len(obs_set) == 2
diff --git a/tests/routes/test_peers.py b/tests/routes/test_peers.py
index 2b4f3c1e..2341e15e 100644
--- a/tests/routes/test_peers.py
+++ b/tests/routes/test_peers.py
@@ -626,9 +626,9 @@ async def test_get_peer_card_with_data(
await crud.set_peer_card(
db_session,
test_workspace.name,
- observer_peer.name,
- observer_peer.name,
self_card_content,
+ observer=observer_peer.name,
+ observed=observer_peer.name,
)
# Set a card for the observer describing the target peer
@@ -636,9 +636,9 @@ async def test_get_peer_card_with_data(
await crud.set_peer_card(
db_session,
test_workspace.name,
- target_peer_name,
- observer_peer.name,
target_card_content,
+ observer=observer_peer.name,
+ observed=target_peer_name,
)
# Test getting observer's own card
diff --git a/tests/routes/test_queue_status.py b/tests/routes/test_queue_status.py
index 441a5a0a..85fd048f 100644
--- a/tests/routes/test_queue_status.py
+++ b/tests/routes/test_queue_status.py
@@ -3,7 +3,7 @@ from fastapi.testclient import TestClient
from sqlalchemy.ext.asyncio import AsyncSession
from src import models
-from src.deriver.utils import get_work_unit_key
+from src.utils.work_unit import get_work_unit_key
@pytest.mark.asyncio
@@ -148,8 +148,8 @@ class TestDeriverStatusEndpoint:
queue_items: list[models.QueueItem] = []
for _ in range(5):
payload = {
- "sender_name": peer.name,
- "target_name": peer.name,
+ "observed": peer.name,
+ "observer": peer.name,
"task_type": "representation",
"workspace_name": workspace.name,
"session_name": session.name,
@@ -157,7 +157,7 @@ class TestDeriverStatusEndpoint:
queue_item = models.QueueItem(
session_id=session.id,
task_type="representation",
- work_unit_key=get_work_unit_key("representation", payload),
+ work_unit_key=get_work_unit_key(payload),
payload=payload,
processed=False,
)
@@ -222,8 +222,8 @@ class TestDeriverStatusEndpoint:
queue_items: list[models.QueueItem] = []
for _ in range(i + 1): # 1,2,3 items respectively
payload = {
- "sender_name": peer.name,
- "target_name": peer.name,
+ "observed": peer.name,
+ "observer": peer.name,
"task_type": "representation",
"workspace_name": workspace.name,
"session_name": session.name,
@@ -231,7 +231,7 @@ class TestDeriverStatusEndpoint:
queue_item = models.QueueItem(
session_id=session.id,
task_type="representation",
- work_unit_key=get_work_unit_key("representation", payload),
+ work_unit_key=get_work_unit_key(payload),
payload=payload,
processed=False,
)
@@ -283,8 +283,8 @@ class TestDeriverStatusEndpoint:
await db_session.commit()
await db_session.refresh(session)
payload = {
- "sender_name": peer.name,
- "target_name": peer.name,
+ "observed": peer.name,
+ "observer": peer.name,
"task_type": "representation",
"workspace_name": workspace.name,
"session_name": session.name,
@@ -292,7 +292,7 @@ class TestDeriverStatusEndpoint:
queue_item = models.QueueItem(
session_id=session.id,
task_type="representation",
- work_unit_key=get_work_unit_key("representation", payload),
+ work_unit_key=get_work_unit_key(payload),
payload=payload,
processed=False,
)
diff --git a/tests/test_datetime_parsing.py b/tests/test_datetime_parsing.py
index 69ac2059..13f6078c 100644
--- a/tests/test_datetime_parsing.py
+++ b/tests/test_datetime_parsing.py
@@ -237,14 +237,14 @@ class TestFormatDatetimeUTC:
assert result == "2023-01-01T12:00:00Z" # Should be 12 PM UTC
def test_format_with_microseconds(self):
- """Test formatting datetimes with microseconds."""
+ """Test formatting datetimes with microseconds. They should be removed."""
dt = datetime(2023, 1, 1, 12, 0, 0, 123456, tzinfo=timezone.utc)
result = format_datetime_utc(dt)
- assert result == "2023-01-01T12:00:00.123456Z"
+ assert result == "2023-01-01T12:00:00Z"
def test_roundtrip_consistency(self):
"""Test that format -> parse -> format is consistent."""
- original_dt = datetime(2023, 1, 1, 12, 30, 45, 123456, tzinfo=timezone.utc)
+ original_dt = datetime(2023, 1, 1, 12, 30, 45, 0, tzinfo=timezone.utc)
# Format to string
formatted = format_datetime_utc(original_dt)
@@ -271,16 +271,6 @@ class TestUTCNowISO:
assert isinstance(parsed, datetime)
assert parsed.tzinfo == timezone.utc
- def test_returns_recent_time(self):
- """Test that utc_now_iso returns a recent time (within last few seconds)."""
- before = datetime.now(timezone.utc)
- result_str = utc_now_iso()
- after = datetime.now(timezone.utc)
-
- result = parse_datetime_iso(result_str)
-
- assert before <= result <= after
-
def test_format_consistency(self):
"""Test that utc_now_iso uses consistent Z format."""
result = utc_now_iso()
@@ -395,24 +385,6 @@ class TestDatetimeEdgeCasesIntegration:
assert isinstance(result, datetime)
assert result.tzinfo == timezone.utc
- def test_time_precision_limits(self):
- """Test handling of time precision edge cases."""
- precision_cases = [
- "2023-01-01T12:00:00.000000Z", # No microseconds
- "2023-01-01T12:00:00.000001Z", # Minimum microseconds
- "2023-01-01T12:00:00.999999Z", # Maximum microseconds
- "2023-01-01T12:00:00.123Z", # Partial microseconds (should pad)
- ]
-
- for precision_case in precision_cases:
- result = parse_datetime_iso(precision_case)
- assert isinstance(result, datetime)
-
- # Should be able to format and parse back consistently
- formatted = format_datetime_utc(result)
- reparsed = parse_datetime_iso(formatted)
- assert result == reparsed
-
def test_extreme_date_values(self):
"""Test handling of extreme date values within reasonable bounds."""
# Python datetime has limits: 1-01-01 to 9999-12-31
diff --git a/tests/test_llm_mock.py b/tests/test_llm_mock.py
index 306b8c31..3cb877ab 100644
--- a/tests/test_llm_mock.py
+++ b/tests/test_llm_mock.py
@@ -5,6 +5,11 @@ from unittest.mock import MagicMock
import pytest
from src.models import Message
+from src.utils.representation import (
+ DeductiveObservation,
+ ExplicitObservation,
+ Representation,
+)
@pytest.mark.asyncio
@@ -18,7 +23,25 @@ async def test_generic_honcho_llm_call_mock():
peer_id="test_peer_id",
peer_card=["test_peer_card"],
message_created_at=datetime(2023, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
- working_representation="test working representation",
+ working_representation=Representation(
+ explicit=[
+ ExplicitObservation(
+ content="test explicit observation",
+ created_at=datetime(2023, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="test_session",
+ )
+ ],
+ deductive=[
+ DeductiveObservation(
+ conclusion="test deductive conclusion",
+ premises=["test premise 1", "test premise 2"],
+ created_at=datetime(2023, 1, 1, 0, 0, 0, tzinfo=timezone.utc),
+ message_ids=[(1, 1)],
+ session_name="test_session",
+ )
+ ],
+ ),
history="test history",
new_turns=["test new turn"],
)
diff --git a/tests/test_schema_validations.py b/tests/test_schema_validations.py
index 8ac14ff2..86ad1b37 100644
--- a/tests/test_schema_validations.py
+++ b/tests/test_schema_validations.py
@@ -3,6 +3,7 @@ from pydantic import ValidationError
from src.schemas import (
DocumentCreate,
+ DocumentMetadata,
MessageCreate,
PeerCreate,
SessionCreate,
@@ -92,18 +93,49 @@ class TestMessageValidations:
class TestDocumentValidations:
def test_valid_document_create(self):
- doc = DocumentCreate(content="test content", metadata={})
+ metadata = DocumentMetadata(
+ message_ids=[(1, 1)],
+ level="explicit",
+ premises=[],
+ message_created_at="2021-01-01T00:00:00Z",
+ )
+ doc = DocumentCreate(
+ content="test content",
+ session_name="test",
+ metadata=metadata,
+ embedding=[0.1, 0.2, 0.3],
+ )
assert doc.content == "test content"
- assert doc.metadata == {}
+ assert doc.metadata == metadata
def test_document_content_too_short(self):
with pytest.raises(ValidationError) as exc_info:
- DocumentCreate(content="", metadata={})
+ DocumentCreate(
+ content="",
+ session_name="test",
+ metadata=DocumentMetadata(
+ message_ids=[(1, 1)],
+ level="explicit",
+ premises=[],
+ message_created_at="2021-01-01T00:00:00Z",
+ ),
+ embedding=[0.1, 0.2, 0.3],
+ )
error_dict = exc_info.value.errors()[0]
assert error_dict["type"] == "string_too_short"
def test_document_content_too_long(self):
with pytest.raises(ValidationError) as exc_info:
- DocumentCreate(content="a" * 100001, metadata={})
+ DocumentCreate(
+ content="a" * 100001,
+ session_name="test",
+ metadata=DocumentMetadata(
+ message_ids=[(1, 1)],
+ level="explicit",
+ premises=[],
+ message_created_at="2021-01-01T00:00:00Z",
+ ),
+ embedding=[0.1, 0.2, 0.3],
+ )
error_dict = exc_info.value.errors()[0]
assert error_dict["type"] == "string_too_long"
diff --git a/tests/utils/test_clients.py b/tests/utils/test_clients.py
index be582ebd..93816c85 100644
--- a/tests/utils/test_clients.py
+++ b/tests/utils/test_clients.py
@@ -780,6 +780,8 @@ class TestGroqClient:
from groq import AsyncGroq
mock_client = AsyncMock(spec=AsyncGroq)
+ # Mock JSON response that matches SampleTestModel structure
+ json_content = '{"name": "Bob", "age": 30, "active": true}'
mock_response = ChatCompletion(
id="test-id",
object="chat.completion",
@@ -788,7 +790,9 @@ class TestGroqClient:
choices=[
Choice(
index=0,
- message=ChatCompletionMessage(role="assistant", content="Bob"),
+ message=ChatCompletionMessage(
+ role="assistant", content=json_content
+ ),
finish_reason="stop",
)
],
@@ -799,7 +803,7 @@ class TestGroqClient:
mock_client.chat.completions.create = AsyncMock(return_value=mock_response)
with patch.dict(CLIENTS, {"groq": mock_client}):
- _response = await honcho_llm_call_inner(
+ response = await honcho_llm_call_inner(
provider="groq",
model="llama-3.1-70b",
prompt="Generate a person",
@@ -807,7 +811,15 @@ class TestGroqClient:
response_model=SampleTestModel,
)
- # Verify response_format was set to the model
+ # Verify the response contains the parsed model
+ assert isinstance(response.content, SampleTestModel)
+ assert response.content.name == "Bob"
+ assert response.content.age == 30
+ assert response.content.active is True
+ assert response.output_tokens == 12
+ assert response.finish_reasons == ["stop"]
+
+ # Verify the response format was set to the model
mock_client.chat.completions.create.assert_called_once()
call_args = mock_client.chat.completions.create.call_args
assert call_args.kwargs["response_format"] == SampleTestModel
diff --git a/uv.lock b/uv.lock
index f47926b5..f2396942 100644
--- a/uv.lock
+++ b/uv.lock
@@ -683,6 +683,7 @@ dependencies = [
{ name = "greenlet" },
{ name = "groq" },
{ name = "httpx" },
+ { name = "json-repair" },
{ name = "langfuse" },
{ name = "nanoid" },
{ name = "openai" },
@@ -725,6 +726,7 @@ requires-dist = [
{ name = "greenlet", specifier = ">=3.0.3" },
{ name = "groq", specifier = ">=0.31.0" },
{ name = "httpx", specifier = ">=0.27.0" },
+ { name = "json-repair", specifier = ">=0.49.0" },
{ name = "langfuse", specifier = ">=3.3.2" },
{ name = "nanoid", specifier = ">=2.0.0" },
{ name = "openai", specifier = ">=1.99.7" },
@@ -1006,6 +1008,15 @@ wheels = [
{ url = "https://files.pythonhosted.org/packages/b3/4a/4175a563579e884192ba6e81725fc0448b042024419be8d83aa8a80a3f44/jiter-0.10.0-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:3aa96f2abba33dc77f79b4cf791840230375f9534e5fac927ccceb58c5e604a5", size = 354213, upload-time = "2025-05-18T19:04:41.894Z" },
]
+[[package]]
+name = "json-repair"
+version = "0.51.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/4f/3a/f30f3c92da3a285dcbe469c50b058f2d349dc9a20fc1b60c3219befda53f/json_repair-0.51.0.tar.gz", hash = "sha256:487e00042d5bc5cc4897ea9c3cccd4f6641e926b732cc09f98691a832485098a", size = 35289, upload-time = "2025-09-19T04:23:16.745Z" }
+wheels = [
+ { url = "https://files.pythonhosted.org/packages/d0/fc/eb15e39547b29dbf2b786bbbd1e79e7f1d87ec4e7c9ea61786f093181481/json_repair-0.51.0-py3-none-any.whl", hash = "sha256:871f7651ee82abf72efc50a80d3a9af0ade8abf5b4541b418eeeabe4e677e314", size = 26263, upload-time = "2025-09-19T04:23:15.064Z" },
+]
+
[[package]]
name = "langfuse"
version = "3.3.2"