RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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
This commit is contained in:
doria 2025-10-07 15:28:44 -04:00 committed by GitHub
parent accdc79fdb
commit f988aae996
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
79 changed files with 7254 additions and 3643 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
.env.template
View File

@ -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

2
.gitignore vendored
View File

@ -185,4 +185,6 @@ config.toml
CRUSH.md
.crush/
metrics.jsonl
AGENTS.md

30
CLAUDE.md
View File

@ -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

12
config.toml.example
View File

@ -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"

1
docker-compose.yml.example
View File

@ -30,6 +30,7 @@ services:
restart: always
ports:
- 5432:5432
command: ["postgres", "-c", "max_connections=800"]
environment:
- POSTGRES_DB=honcho
- POSTGRES_USER=testuser

564
migrations/versions/08894082221a_replace_collection_name_with_observer_.py Normal file
View File

@ -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)

104
migrations/versions/564ba40505c5_add_session_name_column_to_documents.py Normal file
View File

@ -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)

1
pyproject.toml
View File

@ -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 = [

12
scripts/generate_message_embeddings.py
View File

@ -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(

2
sdks/python/CHANGELOG.md
View File

@ -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

2
sdks/typescript/CHANGELOG.md
View File

@ -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

24
sdks/typescript/src/session_context.ts
View File

@ -138,9 +138,9 @@ export class SessionContext {
const assistantId = typeof assistant === 'string' ? assistant : assistant.id
const summaryMessage = this.summary
? {
role: 'system',
content: `<summary>${this.summary.content}</summary>`,
}
role: 'system',
content: `<summary>${this.summary.content}</summary>`,
}
: 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: `<summary>${this.summary.content}</summary>`,
}
role: 'user',
content: `<summary>${this.summary.content}</summary>`,
}
: 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
}

52
src/config.py
View File

@ -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:

24
src/crud/__init__.py
View File

@ -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",

36
src/crud/collection.py
View File

@ -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
)

34
src/crud/deriver.py
View File

@ -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

183
src/crud/document.py
View File

@ -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)

96
src/crud/peer_card.py Normal file
View File

@ -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"

886
src/crud/representation.py
View File

@ -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"

5
src/dependencies.py
View File

@ -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)

59
src/deriver/consumer.py
View File

@ -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
)

687
src/deriver/deriver.py
View File

@ -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

74
src/deriver/enqueue.py
View File

@ -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),
)

72
src/deriver/prompts.py
View File

@ -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}:
<current_context>
{working_representation}
{str(working_representation)}
</current_context>
"""
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=[],
)

77
src/deriver/queue_manager.py
View File

@ -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()

82
src/deriver/utils.py
View File

@ -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

253
src/dialectic/chat.py
View File

@ -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)

88
src/dialectic/prompts.py
View File

@ -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>{query}</query>
<working_representation>{working_representation}</working_representation>
{f"<global_context>{additional_context}</global_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

246
src/dialectic/utils.py
View File

@ -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

13
src/dreamer/__init__.py Normal file
View File

@ -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",
]

354
src/dreamer/dream_scheduler.py Normal file
View File

@ -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

195
src/dreamer/dreamer.py Normal file
View File

@ -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

26
src/dreamer/prompts.py Normal file
View File

@ -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}
"""
)

207
src/embedding_client.py
View File

@ -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()

50
src/models.py
View File

@ -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

28
src/routers/peers.py
View File

@ -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)

167
src/routers/sessions.py
View File

@ -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,
)

4
src/routers/workspaces.py
View File

@ -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)}")

48
src/schemas.py
View File

@ -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):

170
src/utils/clients.py
View File

@ -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

458
src/utils/embedding_store.py
View File

@ -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

189
src/utils/formatting.py
View File

@ -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

378
src/utils/json_parser.py Normal file
View File

@ -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'(?<!\\)"(?![,\]\}:\s]|$)', r'\\"', fixed)
# Remove incomplete key-value pairs at the end
fixed = re.sub(r',\s*"[^"]*"?\s*:?\s*$', "", fixed)
# Try to parse the fixed version
try:
json.loads(fixed)
return fixed
except json.JSONDecodeError:
pass
# If that didn't work, try closing it
return try_contextual_closure_repair(fixed)
def simple_bracket_repair(json_str: str) -> 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

40
src/utils/langfuse_client.py Normal file
View File

@ -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}'")

195
src/utils/logging.py
View File

@ -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)

364
src/utils/metrics_collector.py Normal file
View File

@ -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

25
src/utils/peer_card.py Normal file
View File

@ -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.",
)

50
src/deriver/queue_payload.py → src/utils/queue_payload.py
View File

@ -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:

337
src/utils/representation.py Normal file
View File

@ -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)

205
src/utils/shared_models.py
View File

@ -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

4
src/utils/summarizer.py
View File

@ -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,

15
src/utils/tokens.py Normal file
View File

@ -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

2
src/utils/types.py
View File

@ -1,3 +1,3 @@
from typing import Literal
SupportedProviders = Literal["anthropic", "openai", "google", "groq", "custom"]
SupportedProviders = Literal["anthropic", "openai", "google", "groq", "custom", "vllm"]

110
src/utils/work_unit.py Normal file
View File

@ -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}")

13
src/webhooks/events.py
View File

@ -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,

2
src/webhooks/webhook_delivery.py
View File

@ -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__)

3
tests/bench/.gitignore vendored Normal file
View File

@ -0,0 +1,3 @@
longmemeval_data
eval_results
perf_metrics

1151
tests/bench/longmem.py Normal file
View File

File diff suppressed because it is too large Load Diff

30
tests/bench/peer_card_bench.py
View File

@ -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 == []:

34
tests/bench/run_tests.py
View File

@ -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}",

88
tests/conftest.py
View File

@ -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

451
tests/crud/test_peer_card.py Normal file
View File

@ -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"
)

12
tests/deriver/README.md
View File

@ -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)"
```

36
tests/deriver/conftest.py
View File

@ -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

62
tests/deriver/test_deriver_processing.py
View File

@ -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

62
tests/deriver/test_queue_processing.py
View File

@ -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(
{

664
tests/deriver/test_representation_crud.py
View File

@ -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

0
tests/dreamer/__init__.py Normal file
View File

64
tests/dreamer/test_dream_scheduler.py Normal file
View File

@ -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]

56
tests/dreamer/test_queue_manager_singleton.py Normal file
View File

@ -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

126
tests/integration/test_enqueue.py
View File

@ -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

533
tests/integration/test_representation.py Normal file
View File

@ -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

8
tests/routes/test_peers.py
View File

@ -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

20
tests/routes/test_queue_status.py
View File

@ -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,
)

34
tests/test_datetime_parsing.py
View File

@ -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

25
tests/test_llm_mock.py
View File

@ -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"],
)

40
tests/test_schema_validations.py
View File

@ -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"

18
tests/utils/test_clients.py
View File

@ -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

11
uv.lock
View File

@ -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"