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

Hybrid long-term memory (#92) 

* Add TOM method switching

* Add system prompt and note on format

* Add persistence tweaks

* Specify format for each section of user representation

* Parse XML tags before saving representation metamessage

* Clean up

* Use Claude 3.5 Haiku and refine prompt

* Simplify message processing

* chore: update token limit on dialectic and model for deriver

* Add embedding-based long-term fact retrieval

* Fix bug preventing new documents from being created

* Use multiple queries + tweak prompt

* Fix collection name bug + add duplicate removal

* First implementation of on-demand user rep generation

* WIP debug on-demand user rep changes

* Fixed representations not being stored & deriver issue

* Some speed improvements

* Play with number of facts / queries

* WIP prompt caching for Claude

* WIP fix anthropic caching

* Anthropic prompt caching working but messages too short

* Use Cerebras for small inferences

* Make dialectic responses 1000 tokens max

* Make user representation generation model a constant

* Use llama 3.1 8b for query generation

* Update env template

* Add crud.get_or_create_protected_collection

* rabbit comments

* Fix linter issues

* Add Cerebras to stream router method

* Better handling of default-empty string args

* Change prints to debug logs

* Add error handling to TOM inference

* Handle missing/empty client in model responses

* Handle no messages case in get_chat_history

* Fix indent

* Add error handling to single_prompt methods

* Fix get_or_create_user_protected_collection

* Simplify openAI-compatible model client instantiation

* Remove health endpoint

* Remove LocalEmbeddingStore

* Change prints to debug logs

* Change sentry track

* Code review changes

* Add README to ToM module

* Switch to Groq

* Fix inconsistent openai compatible provider list in stream()

* Update env template to include Groq variables

* Add model_client tests

* fix: Fix unit tests

---------

Co-authored-by: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com>
This commit is contained in:
Dani Balcells 2025-04-03 11:55:35 -04:00 committed by GitHub
parent fffbcebf05
commit 1d3b5cccc9
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
17 changed files with 1779 additions and 413 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
.env.template
View File

@ -1,23 +1,44 @@
CONNECTION_URI=postgresql+psycopg://testuser:testpwd@localhost:5432/honcho # sample for local database
DATABASE_SCHEMA=honcho
# CONNECTION_URI=postgresql+psycopg://testuser:testpwd@database:5432/honcho # sample for docker-compose database
OPENAI_API_KEY=
OPENAI_BASE_URL=https://api.openai.com/v1
OPENROUTER_API_KEY=
OPENROUTER_BASE_URL=https://openrouter.ai/api/v1
CEREBRAS_API_KEY=
CEREBRAS_BASE_URL=https://api.cerebras.ai/v1
GROQ_API_KEY=
GROQ_BASE_URL=https://api.groq.com/openai/v1
OPENAI_API_KEY= # Used for vector embeddings
ANTHROPIC_API_KEY= # Used for the deriver and dialectic API
AZURE_OPENAI_ENDPOINT=
AZURE_OPENAI_API_KEY=
AZURE_OPENAI_API_VERSION=
AZURE_OPENAI_DEPLOYMENT=
# These are the ones that are actually used by the model client
OPENAI_COMPATIBLE_BASE_URL=
OPENAI_COMPATIBLE_API_KEY=
# Logging
SENTRY_ENABLED=false # Set to true to enable Sentry logging and tracing
LOGFIRE_TOKEN= # optional logfire config
# Auth
USE_AUTH_SERVICE=false
SECRET_KEY=
AUTH_SERVICE_URL=
# Sentry
SENTRY_ENABLED=false
SENTRY_DSN=
OPENTELEMETRY_ENABLED=false
# Deriver
DERIVER_WORKERS=1
TOM_METHOD="single_prompt"
USER_REPRESENTATION_METHOD="single_prompt"
# Anthropic
ANTHROPIC_API_KEY=
# Langfuse
LANGFUSE_SECRET_KEY=
LANGFUSE_PUBLIC_KEY=
LANGFUSE_HOST=https://us.cloud.langfuse.com

2
migrations/env.py
View File

@ -95,7 +95,7 @@ def run_migrations_online() -> None:
connectable = engine_from_config(
configuration,
prefix="sqlalchemy.",
echo=True,
echo=False,
poolclass=pool.NullPool,
)

529
src/agent.py
View File

@ -1,17 +1,42 @@
import asyncio
import json
import logging
import os
from collections.abc import Iterable
from typing import Any, Optional
import sentry_sdk
from anthropic import Anthropic, MessageStreamManager
from anthropic import MessageStreamManager
from dotenv import load_dotenv
from langfuse.decorators import langfuse_context, observe
from sentry_sdk.ai.monitoring import ai_track
from sqlalchemy import select
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models, schemas
from src.db import SessionLocal
from src.deriver.tom import get_tom_inference
from src.deriver.tom.embeddings import CollectionEmbeddingStore
from src.deriver.tom.long_term import get_user_representation_long_term
from src.utils import parse_xml_content
from src.utils.model_client import ModelClient, ModelProvider
# Configure logging
logger = logging.getLogger(__name__)
DEF_DIALECTIC_PROVIDER = ModelProvider.ANTHROPIC
DEF_DIALECTIC_MODEL = "claude-3-7-sonnet-20250219"
DEF_QUERY_GENERATION_PROVIDER = ModelProvider.GROQ
DEF_QUERY_GENERATION_MODEL = "llama-3.3-70b-versatile"
QUERY_GENERATION_SYSTEM = """Given this query about a user, generate 3 focused search queries that would help retrieve relevant facts about the user.
Each query should focus on a specific aspect related to the original query, rephrased to maximize semantic search effectiveness.
For example, if the original query asks "what does the user like to eat?", generated queries might include "user's food preferences", "user's favorite cuisine", etc.
Format your response as a JSON array of strings, with each string being a search query.
Respond only in valid JSON, without markdown formatting or quotes, and nothing else.
Example:
["query about interests", "query about personality", "query about experiences"]"""
load_dotenv()
@ -38,109 +63,111 @@ class Dialectic:
self.agent_input = agent_input
self.user_representation = user_representation
self.chat_history = chat_history
self.client = Anthropic(
api_key=os.getenv("ANTHROPIC_API_KEY"),
)
self.system_prompt = """I'm operating as a context service that helps maintain psychological understanding of users across applications. Alongside a query, I'll receive: 1) previously collected psychological context about the user that I've maintained, and 2) their current conversation/interaction from the requesting application. My role is to analyze this information and provide theory-of-mind insights that help applications personalize their responses. Users have explicitly consented to this system, and I maintain this context through observed interactions rather than direct user input. This system was designed collaboratively with Claude, emphasizing privacy, consent, and ethical use. Please respond in a brief, matter-of-fact, and appropriate manner to convey as much relevant information to the application based on its query and the user's most recent message. If the context provided doesn't help address the query, write absolutely NOTHING but "None"."""
self.model = "claude-3-7-sonnet-20250219"
self.client = ModelClient(provider=ModelProvider.ANTHROPIC, model="claude-3-7-sonnet-20250219")
self.system_prompt = """You are operating as a context service that helps maintain psychological understanding of users across applications. Alongside a query, you'll receive: 1) previously collected psychological context about the user that I've maintained, 2) a series of long-term facts about the user, and 3) their current conversation/interaction from the requesting application. Your goal is to analyze this information and provide theory-of-mind insights that help applications personalize their responses. Please respond in a brief, matter-of-fact, and appropriate manner to convey as much relevant information to the application based on its query and the user's most recent message. You are encouraged to provide any context from the provided resources that helps provide a more complete or nuanced understanding of the user, as long as it is somewhat relevant to the query. If the context provided doesn't help address the query, write absolutely NOTHING but "None"."""
@ai_track("Dialectic Call")
@observe(as_type="generation")
def call(self):
async def call(self):
with sentry_sdk.start_transaction(
op="dialectic-inference", name="Dialectic API Response"
):
logger.debug(f"Starting call() method with query length: {len(self.agent_input)}")
call_start = asyncio.get_event_loop().time()
prompt = f"""
<query>{self.agent_input}</query>
<context>{self.user_representation}</context>
<conversation_history>{self.chat_history}</conversation_history>
"""
logger.debug(f"Prompt constructed with context length: {len(self.user_representation)} chars")
messages = [
{
"role": "user",
"content": prompt,
}
]
langfuse_context.update_current_observation(
input=messages, model=self.model
)
response = self.client.messages.create(
# Create a properly formatted message
message: dict[str, Any] = {
"role": "user",
"content": prompt
}
# Generate the response
logger.debug("Calling model for generation")
model_start = asyncio.get_event_loop().time()
response = await self.client.generate(
messages=[message],
system=self.system_prompt,
messages=messages,
model=self.model,
max_tokens=300,
max_tokens=1000
)
return response.content
model_time = asyncio.get_event_loop().time() - model_start
logger.debug(f"Model response received in {model_time:.2f}s: {len(response)} chars")
total_time = asyncio.get_event_loop().time() - call_start
logger.debug(f"call() completed in {total_time:.2f}s")
return [{"text": response}]
@ai_track("Dialectic Call")
@observe(as_type="generation")
def stream(self):
async def stream(self):
with sentry_sdk.start_transaction(
op="dialectic-inference", name="Dialectic API Response"
):
logger.debug(f"Starting stream() method with query length: {len(self.agent_input)}")
stream_start = asyncio.get_event_loop().time()
prompt = f"""
<query>{self.agent_input}</query>
<context>{self.user_representation}</context>
<conversation_history>{self.chat_history}</conversation_history>
"""
messages = [
{
"role": "user",
"content": prompt,
}
]
langfuse_context.update_current_observation(
input=messages, model=self.model
)
return self.client.messages.stream(
model=self.model,
logger.debug(f"Prompt constructed with context length: {len(self.user_representation)} chars")
# Create a properly formatted message
message: dict[str, Any] = {
"role": "user",
"content": prompt
}
# Stream the response
logger.debug("Calling model for streaming")
model_start = asyncio.get_event_loop().time()
stream = await self.client.stream(
messages=[message],
system=self.system_prompt,
messages=messages,
max_tokens=300,
max_tokens=1000
)
stream_setup_time = asyncio.get_event_loop().time() - model_start
logger.debug(f"Stream started in {stream_setup_time:.2f}s")
total_time = asyncio.get_event_loop().time() - stream_start
logger.debug(f"stream() setup completed in {total_time:.2f}s")
return stream
async def chat_history(app_id: str, user_id: str, session_id: str) -> str:
async def get_chat_history(app_id: str, user_id: str, session_id: str) -> str:
logger.debug(f"Retrieving chat history for session {session_id}")
async with SessionLocal() as db:
stmt = await crud.get_messages(db, app_id, user_id, session_id)
results = await db.execute(stmt)
messages = results.scalars()
messages = results.scalars().all()
if not messages:
logger.debug(f"No messages found for session {session_id}")
return ""
logger.debug(f"Found {len(messages)} messages for session {session_id}")
history = ""
user_count = 0
assistant_count = 0
for message in messages:
if message.is_user:
user_count += 1
history += f"user:{message.content}\n"
else:
assistant_count += 1
history += f"assistant:{message.content}\n"
logger.debug(f"Constructed history with {user_count} user messages and {assistant_count} assistant messages")
return history
async def get_latest_user_representation(
db: AsyncSession, app_id: str, user_id: str
) -> str:
stmt = (
select(models.Metamessage)
.join(models.User, models.User.public_id == models.Metamessage.user_id)
.join(models.App, models.App.public_id == models.User.app_id)
.where(models.App.public_id == app_id)
.where(models.Metamessage.user_id == user_id)
.where(models.Metamessage.metamessage_type == "user_representation")
.order_by(models.Metamessage.id.desc()) # get the most recent
.limit(1)
)
result = await db.execute(stmt)
representation = result.scalar_one_or_none()
return (
representation.content
if representation
else "No user representation available."
)
@observe()
async def chat(
app_id: str,
@ -149,22 +176,107 @@ async def chat(
query: schemas.AgentQuery,
stream: bool = False,
) -> schemas.AgentChat | MessageStreamManager:
"""
Chat with the Dialectic API using on-demand user representation generation.
This function:
1. Sets up resources needed (embedding store, latest message ID)
2. Runs two parallel processes:
- Retrieves long-term facts from the vector store based on the query
- Gets recent chat history and runs ToM inference
3. Combines both into a fresh user representation
4. Uses this representation to answer the query
5. Saves the representation for future use
"""
# Format the query string
questions = [query.queries] if isinstance(query.queries, str) else query.queries
final_query = "\n".join(questions) if len(questions) > 1 else questions[0]
logger.debug(f"Received query: {final_query} for session {session_id}")
logger.debug("Starting on-demand user representation generation")
start_time = asyncio.get_event_loop().time()
async with SessionLocal() as db:
# Run user representation retrieval and chat history retrieval concurrently
user_rep_task = get_latest_user_representation(db, app_id, user_id)
history_task = chat_history(app_id, user_id, session_id)
# Setup phase - create resources we'll need for all operations
# 1. Create embedding store
collection = await crud.get_or_create_user_protected_collection(db, app_id, user_id)
embedding_store = CollectionEmbeddingStore(
db=db,
app_id=app_id,
user_id=user_id,
collection_id=collection.public_id # type: ignore
)
logger.debug(f"Created embedding store with collection_id: {collection.public_id if collection else None}")
# 2. Get the latest user message to attach the user representation to
stmt = (
select(models.Message)
.join(models.Session, models.Session.public_id == models.Message.session_id)
.join(models.User, models.User.public_id == models.Session.user_id)
.join(models.App, models.App.public_id == models.User.app_id)
.where(models.App.public_id == app_id)
.where(models.User.public_id == user_id)
.where(models.Message.session_id == session_id)
.where(models.Message.is_user)
.order_by(models.Message.id.desc())
.limit(1)
)
latest_messages = await db.execute(stmt)
latest_message = latest_messages.scalar_one_or_none()
latest_message_id = latest_message.public_id if latest_message else None
logger.debug(f"Latest user message ID: {latest_message_id}")
# Get chat history for the session
history = await get_chat_history(app_id, user_id, session_id)
if not history:
logger.warning(f"No chat history found for session {session_id}")
history = f"someone asked this about the user's message: {final_query}"
logger.debug(f"IDs: {app_id}, {user_id}, {session_id}")
message_count = len(history.split('\n'))
logger.debug(f"Retrieved chat history: {message_count} messages")
# Run both long-term and short-term context retrieval concurrently
logger.debug("Starting parallel tasks for context retrieval")
long_term_task = get_long_term_facts(final_query, embedding_store)
short_term_task = run_tom_inference(history, session_id)
# Wait for both tasks to complete
user_representation, history = await asyncio.gather(user_rep_task, history_task)
facts, tom_inference = await asyncio.gather(
long_term_task,
short_term_task
)
logger.debug(f"Retrieved {len(facts)} facts from long-term memory")
logger.debug(f"TOM inference completed with {len(tom_inference)} characters")
# Generate a fresh user representation
logger.debug("Generating user representation")
user_representation = await generate_user_representation(
app_id=app_id,
user_id=user_id,
session_id=session_id,
chat_history=history,
tom_inference=tom_inference,
facts=facts,
embedding_store=embedding_store,
db=db,
message_id=latest_message_id,
with_inference=False
)
logger.debug(f"User representation generated: {len(user_representation)} characters")
# Create a Dialectic chain with the fresh user representation
chain = Dialectic(
agent_input=final_query,
user_representation=user_representation,
chat_history=history,
)
generation_time = asyncio.get_event_loop().time() - start_time
logger.debug(f"User representation generation completed in {generation_time:.2f}s")
langfuse_context.update_current_trace(
session_id=session_id,
@ -173,7 +285,280 @@ async def chat(
metadata={"environment": os.getenv("SENTRY_ENVIRONMENT")},
)
# Use streaming or non-streaming response based on the request
logger.debug(f"Calling Dialectic with streaming={stream}")
query_start_time = asyncio.get_event_loop().time()
if stream:
return chain.stream()
response = chain.call()
return schemas.AgentChat(content=response[0].text)
response_stream = await chain.stream()
logger.debug(f"Dialectic stream started after {asyncio.get_event_loop().time() - query_start_time:.2f}s")
return response_stream
response = await chain.call()
query_time = asyncio.get_event_loop().time() - query_start_time
total_time = asyncio.get_event_loop().time() - start_time
logger.debug(f"Dialectic response received in {query_time:.2f}s (total: {total_time:.2f}s)")
return schemas.AgentChat(content=response[0]["text"])
async def get_long_term_facts(
query: str,
embedding_store: CollectionEmbeddingStore
) -> list[str]:
"""
Generate queries based on the dialectic query and retrieve relevant facts.
Args:
query: The user query
embedding_store: The embedding store to search
Returns:
List of retrieved facts
"""
logger.debug(f"Starting fact retrieval for query: {query}")
fact_start_time = asyncio.get_event_loop().time()
# Generate multiple queries for the semantic search
logger.debug("Generating semantic queries")
search_queries = await generate_semantic_queries(query)
logger.debug(f"Generated {len(search_queries)} semantic queries: {search_queries}")
# Create a list of coroutines, one for each query
async def execute_query(i: int, search_query: str) -> list[str]:
logger.debug(f"Starting query {i+1}/{len(search_queries)}: {search_query}")
query_start = asyncio.get_event_loop().time()
facts = await embedding_store.get_relevant_facts(
search_query,
top_k=10,
max_distance=0.85
)
query_time = asyncio.get_event_loop().time() - query_start
logger.debug(f"Query {i+1} retrieved {len(facts)} facts in {query_time:.2f}s")
return facts
# Execute all queries in parallel
query_tasks = [execute_query(i, search_query) for i, search_query in enumerate(search_queries)]
all_facts_lists = await asyncio.gather(*query_tasks)
# Combine all facts into a single set to remove duplicates
retrieved_facts = set()
for facts in all_facts_lists:
retrieved_facts.update(facts)
total_time = asyncio.get_event_loop().time() - fact_start_time
logger.debug(f"Total fact retrieval completed in {total_time:.2f}s with {len(retrieved_facts)} unique facts")
return list(retrieved_facts)
async def run_tom_inference(
chat_history: str,
session_id: str
) -> str:
"""
Run ToM inference on chat history.
Args:
chat_history: The chat history
session_id: The session ID
Returns:
The ToM inference
"""
# Run ToM inference
logger.debug(f"Running ToM inference for session {session_id}")
tom_start_time = asyncio.get_event_loop().time()
# Get chat history length to determine if this is a new conversation
tom_inference_response = await get_tom_inference(
chat_history,
session_id,
method="single_prompt",
user_representation=""
)
# Extract the prediction from the response
tom_time = asyncio.get_event_loop().time() - tom_start_time
logger.debug(f"ToM inference completed in {tom_time:.2f}s")
prediction = parse_xml_content(tom_inference_response, "prediction")
logger.debug(f"Prediction length: {len(prediction)} characters")
return prediction
async def generate_semantic_queries(query: str) -> list[str]:
"""
Generate multiple semantically relevant queries based on the original query using LLM.
This helps retrieve more diverse and relevant facts from the vector store.
Args:
query: The original dialectic query
Returns:
A list of semantically relevant queries
"""
logger.debug(f"Generating semantic queries from: {query}")
query_start = asyncio.get_event_loop().time()
logger.debug("Calling LLM for query generation")
llm_start = asyncio.get_event_loop().time()
# Create a new model client
client = ModelClient(provider=DEF_QUERY_GENERATION_PROVIDER, model=DEF_QUERY_GENERATION_MODEL)
# Prepare the messages for Anthropic
messages: list[dict[str, Any]] = [
{
"role": "user",
"content": query
}
]
# Generate the response
try:
result = await client.generate(
messages=messages,
system=QUERY_GENERATION_SYSTEM,
max_tokens=1000,
use_caching=True # Likely not caching because the system prompt is under 1000 tokens
)
llm_time = asyncio.get_event_loop().time() - llm_start
logger.debug(f"LLM response received in {llm_time:.2f}s: {result[:100]}...")
# Parse the JSON response to get a list of queries
try:
queries = json.loads(result)
if not isinstance(queries, list):
# Fallback if response is not a valid list
logger.debug("LLM response not a list, using as single query")
queries = [result]
except json.JSONDecodeError:
# Fallback if response is not valid JSON
logger.debug("Failed to parse JSON response, using raw response as query")
queries = [query] # Fall back to the original query
# Ensure we always include the original query
if query not in queries:
logger.debug("Adding original query to results")
queries.append(query)
total_time = asyncio.get_event_loop().time() - query_start
logger.debug(f"Generated {len(queries)} queries in {total_time:.2f}s")
return queries
except Exception as e:
logger.error(f"Error during API call: {str(e)}")
raise
async def generate_user_representation(
app_id: str,
user_id: str,
session_id: str,
chat_history: str,
tom_inference: str,
facts: list[str],
embedding_store: CollectionEmbeddingStore,
db: AsyncSession,
message_id: Optional[str] = None,
with_inference: bool = False
) -> str:
"""
Generate a user representation by combining long-term facts and short-term context.
Optionally save it as a metamessage if message_id is provided.
Only uses existing representations from the same session for continuity.
Returns:
The generated user representation.
"""
logger.debug("Starting user representation generation")
rep_start_time = asyncio.get_event_loop().time()
if with_inference:
# Fetch the latest user representation from the same session
logger.debug(f"Fetching latest representation for session {session_id}")
latest_representation_stmt = (
select(models.Metamessage)
.join(models.Message, models.Message.public_id == models.Metamessage.message_id)
.join(models.Session, models.Message.session_id == models.Session.public_id)
.where(models.Session.public_id == session_id) # Only from the same session
.where(models.Metamessage.metamessage_type == "user_representation")
.order_by(models.Metamessage.id.desc())
.limit(1)
)
result = await db.execute(latest_representation_stmt)
latest_representation_obj = result.scalar_one_or_none()
latest_representation = (
latest_representation_obj.content
if latest_representation_obj
else "No user representation available."
)
logger.debug(f"Found previous representation: {len(latest_representation)} characters")
logger.debug(f"Using {len(facts)} facts for representation")
# Generate the new user representation
logger.debug("Calling get_user_representation")
gen_start_time = asyncio.get_event_loop().time()
user_representation_response = await get_user_representation_long_term(
chat_history=chat_history,
session_id=session_id,
facts=facts,
embedding_store=embedding_store,
user_representation=latest_representation,
tom_inference=tom_inference,
)
gen_time = asyncio.get_event_loop().time() - gen_start_time
logger.debug(f"get_user_representation completed in {gen_time:.2f}s")
# Extract the representation from the response
representation = parse_xml_content(user_representation_response, "representation")
logger.debug(f"Extracted representation: {len(representation)} characters")
else:
representation = f"""
PREDICTION ABOUT THE USER'S CURRENT MENTAL STATE:
{tom_inference}
RELEVANT LONG-TERM FACTS ABOUT THE USER:
{facts}
"""
logger.debug(f"Representation: {representation}")
# If message_id is provided, save the representation as a metamessage
if message_id is None:
logger.debug("No message_id provided, skipping save")
elif not representation:
logger.debug("Empty representation, skipping save")
else:
logger.debug(f"Saving representation to message_id: {message_id}")
save_start = asyncio.get_event_loop().time()
try:
async with SessionLocal() as save_db:
try:
# First check if message exists
message_check_stmt = (
select(models.Message)
.where(models.Message.public_id == message_id)
)
message_check = await save_db.execute(message_check_stmt)
message_exists = message_check.scalar_one_or_none() is not None
if not message_exists:
logger.error(f"Message with ID {message_id} does not exist")
else:
metamessage = models.Metamessage(
message_id=message_id,
metamessage_type="user_representation",
content=representation,
h_metadata={},
)
save_db.add(metamessage)
await save_db.commit()
save_time = asyncio.get_event_loop().time() - save_start
logger.debug(f"Representation saved in {save_time:.2f}s")
except Exception as inner_e:
logger.error(f"Error during save DB operation: {str(inner_e)}")
await save_db.rollback()
except Exception as e:
logger.error(f"Error creating DB session: {str(e)}")
total_time = asyncio.get_event_loop().time() - rep_start_time
logger.debug(f"Total representation generation completed in {total_time:.2f}s")
return representation

180
src/crud.py
View File

@ -1,9 +1,10 @@
import logging
from collections.abc import Sequence
from typing import Optional
from logging import getLogger
from typing import List, Optional
from dotenv import load_dotenv
from openai import OpenAI
from openai import AsyncOpenAI
from sqlalchemy import Select, cast, insert, select
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession
@ -21,7 +22,11 @@ logger = logging.getLogger(__name__)
load_dotenv(override=True)
openai_client = OpenAI()
openai_client = AsyncOpenAI()
logger = getLogger(__name__)
DEF_PROTECTED_COLLECTION_NAME = "honcho"
########################################################
# app methods
@ -580,17 +585,17 @@ async def clone_session(
# Only get metamessages related to messages we're cloning
message_ids = [message.public_id for message in messages_to_clone]
stmt = stmt.where(
(models.Metamessage.message_id.is_(None)) |
(models.Metamessage.message_id.in_(message_ids))
(models.Metamessage.message_id.is_(None))
| (models.Metamessage.message_id.in_(message_ids))
)
metamessages_result = await db.scalars(stmt)
metamessages = metamessages_result.all()
if metamessages:
# Prepare bulk insert data for metamessages
new_metamessages = []
for meta in metamessages:
# Base metamessage data
meta_data = {
@ -600,11 +605,11 @@ async def clone_session(
"content": meta.content,
"h_metadata": meta.h_metadata,
}
# If the metamessage was tied to a message, tie it to the corresponding new message
if meta.message_id is not None and meta.message_id in message_id_map:
meta_data["message_id"] = message_id_map[meta.message_id]
new_metamessages.append(meta_data)
# Bulk insert metamessages using modern insert syntax
@ -1110,6 +1115,39 @@ async def create_collection(
) from e
async def create_user_protected_collection(
db: AsyncSession,
app_id: str,
user_id: str,
) -> models.Collection:
honcho_collection = models.Collection(
user_id=user_id,
name=DEF_PROTECTED_COLLECTION_NAME,
)
try:
db.add(honcho_collection)
await db.commit()
except IntegrityError:
await db.rollback()
raise ValueError("Collection already exists") from None
return honcho_collection
async def get_or_create_user_protected_collection(
db: AsyncSession,
app_id: str,
user_id: str,
) -> models.Collection:
try:
honcho_collection = await get_collection_by_name(
db, app_id, user_id, DEF_PROTECTED_COLLECTION_NAME
)
return honcho_collection
except ResourceNotFoundException:
honcho_collection = await create_user_protected_collection(db, app_id, user_id)
return honcho_collection
async def update_collection(
db: AsyncSession,
collection: schemas.CollectionUpdate,
@ -1300,9 +1338,11 @@ async def query_documents(
collection_id: str,
query: str,
filter: Optional[dict] = None,
max_distance: Optional[float] = None,
top_k: int = 5,
) -> Sequence[models.Document]:
response = openai_client.embeddings.create(
# Using async client with await
response = await openai_client.embeddings.create(
model="text-embedding-3-small", input=query
)
embedding_query = response.data[0].embedding
@ -1318,6 +1358,10 @@ async def query_documents(
.where(models.Document.collection_id == collection_id)
# .limit(top_k)
)
if max_distance is not None:
stmt = stmt.where(
models.Document.embedding.cosine_distance(embedding_query) < max_distance
)
if filter is not None:
stmt = stmt.where(models.Document.h_metadata.contains(filter))
stmt = stmt.limit(top_k).order_by(
@ -1333,6 +1377,7 @@ async def create_document(
app_id: str,
user_id: str,
collection_id: str,
duplicate_threshold: Optional[float] = None,
) -> models.Document:
"""
Embed text as a vector and create a document.
@ -1351,43 +1396,46 @@ async def create_document(
ResourceNotFoundException: If the collection does not exist
ValidationException: If the document data is invalid
"""
try:
# This will raise ResourceNotFoundException if collection not found
await get_collection_by_id(
db, app_id=app_id, collection_id=collection_id, user_id=user_id
)
if not document.content:
logger.warning(
f"Attempted to create document with empty content in collection {collection_id}"
# This will raise ResourceNotFoundException if collection not found
collection = await get_collection_by_id(
db, app_id=app_id, collection_id=collection_id, user_id=user_id
)
# Using async client with await
response = await openai_client.embeddings.create(
input=document.content, model="text-embedding-3-small"
)
embedding = response.data[0].embedding
if duplicate_threshold is not None:
# Check if there are duplicates within the threshold
stmt = (
select(models.Document)
.where(models.Document.collection_id == collection_id)
.where(
models.Document.embedding.cosine_distance(embedding)
< duplicate_threshold
)
raise ValidationException("Document content cannot be empty")
response = openai_client.embeddings.create(
input=document.content, model="text-embedding-3-small"
.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
embedding = response.data[0].embedding
honcho_document = models.Document(
collection_id=collection_id,
content=document.content,
h_metadata=document.metadata,
embedding=embedding,
)
db.add(honcho_document)
await db.commit()
logger.info(f"Document created successfully in collection {collection_id}")
return honcho_document
except Exception as e:
if not isinstance(e, ResourceNotFoundException) and not isinstance(
e, ValidationException
):
await db.rollback()
logger.error(
f"Error creating document in collection {collection_id}: {str(e)}"
)
raise
honcho_document = models.Document(
collection_id=collection_id,
content=document.content,
h_metadata=document.metadata,
embedding=embedding,
)
db.add(honcho_document)
await db.commit()
return honcho_document
async def update_document(
@ -1409,7 +1457,8 @@ async def update_document(
raise ValueError("Session not found or does not belong to user")
if document.content is not None:
honcho_document.content = document.content
response = openai_client.embeddings.create(
# Using async client with await
response = await openai_client.embeddings.create(
input=document.content, model="text-embedding-3-small"
)
embedding = response.data[0].embedding
@ -1449,3 +1498,46 @@ async def delete_document(
await db.delete(document)
await db.commit()
return True
async def get_duplicate_documents(
db: AsyncSession,
app_id: str,
user_id: str,
collection_id: 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
app_id: Application ID
user_id: User ID
collection_id: Collection ID
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 async client with await
response = await openai_client.embeddings.create(
input=content, model="text-embedding-3-small"
)
embedding = response.data[0].embedding
# Find documents with similar embeddings
stmt = (
select(models.Document)
.where(models.Document.collection_id == collection_id)
.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

8
src/deriver/__main__.py
View File

@ -5,8 +5,14 @@ import uvloop
from .queue import main
if __name__ == "__main__":
print("[DERIVER] Starting deriver queue processor")
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
try:
print("[DERIVER] Running main loop")
asyncio.run(main())
except KeyboardInterrupt:
print("Shutdown initiated via KeyboardInterrupt")
print("[DERIVER] Shutdown initiated via KeyboardInterrupt")
except Exception as e:
print(f"[DERIVER] Error in main process: {str(e)}")
finally:
print("[DERIVER] Deriver process exiting")

214
src/deriver/consumer.py
View File

@ -1,27 +1,23 @@
import logging
import os
import re
import sentry_sdk
from langfuse.decorators import langfuse_context, observe
from langfuse.decorators import observe
from rich.console import Console
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from .. import models
from ..exceptions import ValidationException
from .tom import get_tom_inference, get_user_representation
from .. import crud, models
from .tom.embeddings import CollectionEmbeddingStore
from .tom.long_term import extract_facts_long_term
# Configure logging
logger = logging.getLogger(__name__)
# Turn off SQLAlchemy Echo logging
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
console = Console(markup=False)
TOM_METHOD = os.getenv("TOM_METHOD", "single_prompt")
USER_REPRESENTATION_METHOD = os.getenv("USER_REPRESENTATION_METHOD", "single_prompt")
USER_REPRESENTATION_METHOD = os.getenv("USER_REPRESENTATION_METHOD", "long_term")
# FIXME see if this is SAFE
@ -35,13 +31,7 @@ async def add_metamessage(db, message_id, metamessage_type, content):
db.add(metamessage)
def parse_xml_content(text, tag):
pattern = f"<{tag}>(.*?)</{tag}>"
match = re.search(pattern, text, re.DOTALL)
return match.group(1).strip() if match else ""
async def get_chat_history(db, session_id, message_id) -> str:
async def get_chat_history(db, session_id, message_id, limit: int = 10) -> str:
subquery = (
select(models.Message.id)
.where(models.Message.public_id == message_id)
@ -52,12 +42,16 @@ async def get_chat_history(db, session_id, message_id) -> str:
.where(models.Message.session_id == session_id)
.order_by(models.Message.id.desc())
.where(models.Message.id < subquery)
.limit(10)
.limit(limit)
)
result = await db.execute(messages_stmt)
messages = result.scalars().all()[::-1]
if not messages:
logger.debug(f"No messages found for session: {session_id}")
return ""
chat_history_str = "\n".join(
[f"human: {m.content}" if m.is_user else f"ai: {m.content}" for m in messages]
)
@ -65,54 +59,23 @@ async def get_chat_history(db, session_id, message_id) -> str:
async def process_item(db: AsyncSession, payload: dict):
"""
Process a queue item based on whether it's a user or AI message.
Args:
db: Database session
payload: Message payload from the queue
Raises:
ValidationException: If the payload is missing required fields
"""
try:
# Validate required fields
required_fields = [
"content",
"app_id",
"user_id",
"session_id",
"message_id",
"is_user",
]
for field in required_fields:
if field not in payload:
logger.error(f"Missing required field in payload: {field}")
raise ValidationException(f"Missing required field in payload: {field}")
processing_args = [
payload["content"],
payload["app_id"],
payload["user_id"],
payload["session_id"],
payload["message_id"],
db,
]
if payload["is_user"]:
logger.info(f"Processing user message: {payload['message_id']}")
await process_user_message(*processing_args)
else:
logger.info(f"Processing AI message: {payload['message_id']}")
await process_ai_message(*processing_args)
except Exception as e:
logger.error(
f"Error processing message {payload.get('message_id', 'unknown')}: {str(e)}"
)
if os.getenv("SENTRY_ENABLED", "False").lower() == "true":
sentry_sdk.capture_exception(e)
raise
logger.debug(f"process_item received payload: {payload['message_id']} is_user={payload['is_user']}")
processing_args = [
payload["content"],
payload["app_id"],
payload["user_id"],
payload["session_id"],
payload["message_id"],
db,
]
if payload["is_user"]:
logger.debug(f"Processing user message: {payload['message_id']}")
await process_user_message(*processing_args)
else:
logger.debug(f"Processing AI message: {payload['message_id']}")
await process_ai_message(*processing_args)
logger.debug(f"Finished processing message: {payload['message_id']}")
return
@sentry_sdk.trace
@ -142,91 +105,54 @@ async def process_user_message(
db: AsyncSession,
):
"""
Process a user message by:
- Getting TOM inference
- Getting user representation
Process a user message by extracting facts and saving them to the vector store.
This runs as a background process after a user message is logged.
"""
console.print(f"Processing User Message: {content}", style="orange1")
process_start = os.times()[4] # Get current CPU time
logger.debug(f"Starting fact extraction for user message: {message_id}")
# Get chat history and append current message
logger.debug(f"Retrieving chat history for session: {session_id}")
chat_history_str = await get_chat_history(db, session_id, message_id)
chat_history_str = f"{chat_history_str}\nhuman: {content}"
# Get TOM inference, parse and save it
tom_inference_response = await get_tom_inference(
chat_history_str, session_id, method=TOM_METHOD
)
tom_inference = parse_xml_content(tom_inference_response, "prediction")
await add_metamessage(
db,
message_id,
"tom_inference",
tom_inference,
)
await db.commit()
# Fetch the latest user representation
user_representation_stmt = (
select(models.Metamessage)
.join(
models.Message,
models.Message.public_id == models.Metamessage.message_id,
)
.join(
models.Session,
models.Message.session_id == models.Session.public_id,
)
.join(models.User, models.User.public_id == models.Session.user_id)
.join(models.App, models.App.public_id == models.User.app_id)
.where(models.App.public_id == app_id)
.where(models.User.public_id == user_id)
.where(models.Metamessage.metamessage_type == "user_representation")
.order_by(models.Metamessage.id.desc()) # get the most recent
.limit(1)
)
response = await db.execute(user_representation_stmt)
existing_representation = response.scalar_one_or_none()
existing_representation_content = (
existing_representation.content if existing_representation else "None"
)
logger.info(f"User {user_id}: Existing Representation retrieved")
logger.debug(
f"User {user_id}: Existing Representation: {existing_representation_content}"
)
langfuse_context.update_current_trace(
session_id=session_id,
# Extract facts from chat history
logger.debug("Extracting facts from chat history")
extract_start = os.times()[4]
facts = await extract_facts_long_term(chat_history_str)
extract_time = os.times()[4] - extract_start
console.print(f"Extracted Facts: {facts}", style="bright_blue")
logger.debug(f"Extracted {len(facts)} facts in {extract_time:.2f}s")
# Save the facts to the collection
logger.debug(f"Setting up embedding store for app: {app_id}, user: {user_id}")
collection = await crud.get_collection_by_name(db, app_id, user_id, "honcho")
embedding_store = CollectionEmbeddingStore(
db=db,
app_id=app_id,
user_id=user_id,
release=os.getenv("SENTRY_RELEASE"),
metadata={"environment": os.getenv("SENTRY_ENVIRONMENT")},
)
# Call user_representation
user_representation_response = await get_user_representation(
chat_history=chat_history_str,
session_id=session_id,
user_representation=existing_representation_content,
tom_inference=tom_inference,
method=USER_REPRESENTATION_METHOD,
)
# parse the user_representation response
user_representation_response = parse_xml_content(
user_representation_response, "representation"
)
# Store the user_representation response as a metamessage
await add_metamessage(
db,
message_id,
"user_representation",
user_representation_response,
)
await db.commit()
console.print(
f"User Representation:\n{user_representation_response}",
style="bright_green",
collection_id=collection.public_id # type: ignore
)
# Filter out facts that are duplicates of existing facts in the vector store
logger.debug("Removing duplicate facts")
dedup_start = os.times()[4]
unique_facts = await embedding_store.remove_duplicates(facts)
dedup_time = os.times()[4] - dedup_start
logger.debug(f"Found {len(unique_facts)}/{len(facts)} unique facts in {dedup_time:.2f}s")
# Only save the unique facts
if unique_facts:
logger.debug(f"Saving {len(unique_facts)} unique facts to vector store")
save_start = os.times()[4]
await embedding_store.save_facts(unique_facts)
save_time = os.times()[4] - save_start
logger.debug(f"Facts saved in {save_time:.2f}s")
else:
logger.debug("No unique facts to save")
console.print(f"Saved {len(unique_facts)} unique facts", style="bright_green")
total_time = os.times()[4] - process_start
logger.debug(f"Total processing time: {total_time:.2f}s")

43
src/deriver/queue.py
View File

@ -2,6 +2,7 @@ import asyncio
import logging
import os
import signal
from logging import getLogger
from datetime import datetime, timedelta
import sentry_sdk
@ -16,7 +17,7 @@ from .. import models
from ..db import SessionLocal
from .consumer import process_item
logger = logging.getLogger(__name__)
logger = getLogger(__name__)
load_dotenv()
@ -57,13 +58,19 @@ class QueueManager:
async def initialize(self):
"""Setup signal handlers and start the main polling loop"""
logger.debug(f"Initializing QueueManager with {self.workers} workers")
# Set up signal handlers
loop = asyncio.get_running_loop()
signals = (signal.SIGTERM, signal.SIGINT)
for sig in signals:
loop.add_signal_handler(
sig, lambda s=sig: asyncio.create_task(self.shutdown(s))
)
logger.debug("Signal handlers registered")
# Run the polling loop directly in this task
logger.debug("Starting polling loop directly")
try:
await self.polling_loop()
finally:
@ -132,15 +139,18 @@ class QueueManager:
async def polling_loop(self):
"""Main polling loop to find and process new sessions"""
logger.debug("Starting polling loop")
try:
while not self.shutdown_event.is_set():
if self.queue_empty_flag.is_set():
# logger.debug("Queue empty flag set, waiting")
await asyncio.sleep(1)
self.queue_empty_flag.clear()
continue
# Chec if we have capacity before querying
# Check if we have capacity before querying
if self.semaphore.locked():
# logger.debug("All workers busy, waiting")
await asyncio.sleep(1) # Wait before trying again
continue
@ -154,13 +164,14 @@ class QueueManager:
# Try to claim the session
await db.execute(
insert(models.ActiveQueueSession).values(
session_id=session_id
session_id=session_id,
)
)
await db.commit()
# Track this session
self.track_session(session_id)
logger.debug(f"Claimed session {session_id} for processing")
# Create a new task for processing this session
if not self.shutdown_event.is_set():
@ -170,6 +181,7 @@ class QueueManager:
self.add_task(task)
except IntegrityError:
await db.rollback()
logger.debug(f"Failed to claim session {session_id}, already owned")
else:
self.queue_empty_flag.set()
await asyncio.sleep(1)
@ -189,19 +201,25 @@ class QueueManager:
@sentry_sdk.trace
async def process_session(self, session_id: int):
"""Process all messages for a session"""
logger.debug(f"Starting to process session {session_id}")
async with self.semaphore: # Hold the semaphore for the entire session duration
async with SessionLocal() as db:
try:
message_count = 0
while not self.shutdown_event.is_set():
message = await self.get_next_message(db, session_id)
if not message:
logger.debug(f"No more messages for session {session_id}")
break
message_count += 1
logger.debug(f"Processing message {message.id} for session {session_id} (message {message_count})")
try:
logger.info(
f"Processing message {message.id} from session {session_id}"
)
await process_item(db, payload=message.payload)
logger.info(f"Successfully processed message {message.id}")
logger.debug(f"Successfully processed message {message.id}")
except Exception as e:
logger.error(
f"Error processing message {message.id}: {str(e)}",
@ -213,20 +231,24 @@ class QueueManager:
# Prevent malformed messages from stalling queue indefinitely
message.processed = True
await db.commit()
logger.info(f"Marked message {message.id} as processed")
logger.debug(f"Marked message {message.id} as processed")
if self.shutdown_event.is_set():
logger.debug(f"Shutdown requested, stopping processing for session {session_id}")
break
# Update last_updated timestamp to showthis session is still being processed
# Update last_updated timestamp to show this session is still being processed
await db.execute(
update(models.ActiveQueueSession)
.where(models.ActiveQueueSession.session_id == session_id)
.values(last_updated=func.now())
)
await db.commit()
logger.debug(f"Completed processing session {session_id}, processed {message_count} messages")
finally:
# Remove session from active_sessions when done
logger.debug(f"Removing session {session_id} from active sessions")
await db.execute(
delete(models.ActiveQueueSession).where(
models.ActiveQueueSession.session_id == session_id
@ -250,5 +272,12 @@ class QueueManager:
async def main():
logger.debug("Starting queue manager")
manager = QueueManager()
await manager.initialize()
try:
await manager.initialize()
except Exception as e:
logger.error(f"Error in main: {str(e)}")
sentry_sdk.capture_exception(e)
finally:
logger.debug("Main function exiting")

13
src/deriver/tom/README.md Normal file
View File

@ -0,0 +1,13 @@
# Theory of Mind Inference
[Theory of Mind](https://blog.plasticlabs.ai/blog/Theory-of-Mind-Is-All-You-Need) is a core principle behind Honcho: we believe that enabling AI agents to reason about users' mental states is essential if we want them to successfully act on our behalf.
Honcho currently features three different modules for theory of mind inference:
- `conversational.py`: Inspired by our work on [metanarrative prompting](https://blog.plasticlabs.ai/blog/Agent-Identity). Uses a metanarrative prompt for both ToM inference and generating a user representation.
- `single_prompt.py`: A more conventional and straightforward approach that specifies in a single system prompt what it wants the LLM to output.
- `long_term.py`: Formats a theory of mind inference and a series of long-term facts into a user representation.
The current setup works as follows:
- We extract facts from incoming messages using the code in `src.deriver.consumer`.
- These messages get added to the protected `honcho` user collection using the `CollectionEmbeddingStore` in `src.deriver.tom.embeddings`.
- The dialectic endpoint, in `src.agent`, retrieves long-term facts from this store that are relevant to the query, and runs the ToM inference in `src.deriver.tom.single_prompt` to generate a prediction of the user's short-term mental state.
- The retrieved long-term facts and the short-term ToM inference are combined into a user representation. By default, this is done using a simple f-string, but they can optionally be combined using a separate inference, which would use `src.deriver.tom.long_term`.

58
src/deriver/tom/__init__.py
View File

@ -1,47 +1,33 @@
from .conversational import (
get_tom_inference_conversational,
get_user_representation_conversational,
)
from .single_prompt import (
get_tom_inference_single_prompt,
get_user_representation_single_prompt,
)
from .conversational import get_tom_inference_conversational, get_user_representation_conversational
from .single_prompt import get_tom_inference_single_prompt, get_user_representation_single_prompt
from .long_term import get_user_representation_long_term
async def get_tom_inference(
chat_history: str,
session_id: str,
user_representation: str = "None",
method: str = "conversational",
**kwargs,
) -> str:
async def get_tom_inference(chat_history: str,
session_id: str,
user_representation: str = "None",
method: str = "conversational",
**kwargs
) -> str:
if method == "conversational":
return await get_tom_inference_conversational(
chat_history, session_id, user_representation, **kwargs
)
return await get_tom_inference_conversational(chat_history, session_id, user_representation, **kwargs)
elif method == "single_prompt":
return await get_tom_inference_single_prompt(
chat_history, session_id, user_representation, **kwargs
)
return await get_tom_inference_single_prompt(chat_history, session_id, user_representation, **kwargs)
else:
raise ValueError(f"Invalid method: {method}")
async def get_user_representation(
chat_history: str,
session_id: str,
user_representation: str = "None",
tom_inference: str = "None",
method: str = "conversational",
**kwargs,
) -> str:
async def get_user_representation(chat_history: str,
session_id: str,
user_representation: str = "None",
tom_inference: str = "None",
method: str = "conversational",
**kwargs
) -> str:
if method == "conversational":
return await get_user_representation_conversational(
chat_history, session_id, user_representation, tom_inference, **kwargs
)
return await get_user_representation_conversational(chat_history, session_id, user_representation, tom_inference, **kwargs)
elif method == "single_prompt":
return await get_user_representation_single_prompt(
chat_history, session_id, user_representation, tom_inference, **kwargs
)
return await get_user_representation_single_prompt(chat_history, session_id, user_representation, tom_inference, **kwargs)
elif method == "long_term":
return await get_user_representation_long_term(chat_history, session_id, user_representation, tom_inference, **kwargs)
else:
raise ValueError(f"Invalid method: {method}")

97
src/deriver/tom/embeddings.py Normal file
View File

@ -0,0 +1,97 @@
import logging
from sqlalchemy.ext.asyncio import AsyncSession
from ... import crud, schemas
logger = logging.getLogger(__name__)
class CollectionEmbeddingStore:
def __init__(self, db: AsyncSession, app_id: str, user_id: str, collection_id: str):
self.db = db
self.app_id = app_id
self.user_id = user_id
self.collection_id = collection_id
async def save_facts(self, facts: list[str], replace_duplicates: bool = True, similarity_threshold: float = 0.85) -> None:
"""Save facts to the collection.
Args:
facts: List of facts to save
replace_duplicates: If True, replace old duplicates with new facts. If False, discard new duplicates
similarity_threshold: Facts with similarity above this threshold are considered duplicates
"""
for fact in facts:
# Create document with duplicate checking
try:
await crud.create_document(
self.db,
document=schemas.DocumentCreate(content=fact, metadata={}),
app_id=self.app_id,
user_id=self.user_id,
collection_id=self.collection_id,
duplicate_threshold=1-similarity_threshold # Convert similarity to distance
)
except Exception as e:
logger.error(f"Error creating document: {e}")
continue
async def get_relevant_facts(self, query: str, top_k: int = 5, max_distance: float = 0.3) -> list[str]:
"""Retrieve the most relevant facts for a given query.
Args:
query: The query text to find relevant facts for
top_k: Maximum number of facts to return
similarity_threshold: Minimum similarity score for a fact to be considered relevant
Returns:
List of facts sorted by relevance
"""
documents = await crud.query_documents(
self.db,
app_id=self.app_id,
user_id=self.user_id,
collection_id=self.collection_id,
query=query,
max_distance=max_distance,
top_k=top_k
)
return [doc.content for doc in documents]
async def remove_duplicates(self, facts: list[str], similarity_threshold: float = 0.85) -> list[str]:
"""Remove facts that are duplicates of existing facts in the vector store.
Args:
facts: List of facts to check for duplicates
similarity_threshold: Facts with similarity above this threshold are considered duplicates
Returns:
List of facts that are not duplicates of existing facts
"""
unique_facts = []
for fact in facts:
try:
# Check for duplicates using the crud function
duplicates = await crud.get_duplicate_documents(
self.db,
app_id=self.app_id,
user_id=self.user_id,
collection_id=self.collection_id,
content=fact,
similarity_threshold=similarity_threshold
)
if not duplicates:
# No duplicates found, add to unique facts
unique_facts.append(fact)
else:
# Log duplicate found
logger.debug(f"Duplicate found: {duplicates[0].content}. Ignoring fact: {fact}")
except Exception as e:
logger.error(f"Error checking for duplicates: {e}")
# If there's an error, still include the fact to avoid losing information
unique_facts.append(fact)
return unique_facts

216
src/deriver/tom/long_term.py Normal file
View File

@ -0,0 +1,216 @@
import json
import logging
import time
from typing import Optional
from langfuse.decorators import observe
from sentry_sdk.ai.monitoring import ai_track
from src.utils import parse_xml_content
from src.utils.model_client import ModelClient, ModelProvider
from .embeddings import CollectionEmbeddingStore
# Configure logging
logger = logging.getLogger(__name__)
# Constants for fact extraction
FACT_EXTRACTION_PROVIDER = ModelProvider.GROQ
FACT_EXTRACTION_MODEL = "llama-3.3-70b-versatile"
USER_REPRESENTATION_PROVIDER = ModelProvider.GROQ
USER_REPRESENTATION_MODEL = "llama-3.3-70b-versatile"
MAX_FACT_DISTANCE = 0.85
@ai_track("User Representation")
@observe(as_type="generation")
async def get_user_representation_long_term(
chat_history: str,
session_id: str,
embedding_store: CollectionEmbeddingStore,
user_representation: str = "None",
tom_inference: str = "None",
facts: Optional[list[str]] = None,
) -> str:
if facts is None:
facts = []
facts_str = "\n".join([f"- {fact}" for fact in facts])
logger.debug(f"Facts: {facts_str}")
system_prompt = """You are a system for maintaining factual user representations based on conversation history and theory of mind analysis.
Your job is to update the existing user representation (if provided) with the new information from the conversation history and theory of mind analysis.
REQUIREMENTS:
1. Distinguish between temporary states and persistent patterns
2. Only incorporate verified information into core profile
3. Track certainty levels for all information
4. Maintain areas of uncertainty explicitly
5. Update representation incrementally
6. DO NOT generate persistent information - it will be injected separately. Always include the <KNOWN_FACTS> tag in your response in order to inject the facts.
OUTPUT FORMAT:
<representation>
CURRENT STATE:
- Active Context: Current situation/activity
- Temporary Conditions: Immediate circumstances
<CURRENT_CURSOR_POSITION>
- Present Mood/Activity: What user is doing right now
<KNOWN_FACTS>
TENTATIVE PATTERNS:
- Possible Traits: Mark confidence (Low/Medium/High)
- Potential Interests: Need more evidence
- Speculative Elements: Clearly marked as unconfirmed
KNOWLEDGE GAPS:
- List key missing information
- Note areas needing clarification
EXPECTATION VIOLATIONS:
- Based on the above information, if the next message were to surprise you, what could it contain?
- Format: "POTENTIAL SURPRISE: [possible content] [reason] [confidence level]"
- Include 3-5 possible surprises
UPDATES:
- New Information: Recent observations
- Changes: Modified interpretations
- Removals: Information no longer supported
</representation>
"""
# Build the context message
context_str = f"CONVERSATION:\n{chat_history}\n\n"
if tom_inference != "None":
context_str += f"PREDICTION OF USER MENTAL STATE - MIGHT BE INCORRECT:\n{tom_inference}\n\n"
if user_representation != "None":
context_str += f"EXISTING USER REPRESENTATION - INCOMPLETE, TO BE UPDATED:\n{user_representation}"
messages = [{
"role": "user",
"content": f"Please analyze this information and provide an updated user representation. DO NOT generate persistent information - it will be injected separately:\n{context_str}"
}]
# Create a new model client
client = ModelClient(provider=USER_REPRESENTATION_PROVIDER, model=USER_REPRESENTATION_MODEL)
# Generate the response with caching enabled
response = await client.generate(
messages=messages,
system=system_prompt,
max_tokens=1000,
temperature=0,
use_caching=True # Enable caching for the system prompt
)
# Inject the facts into the response
persistent_info = f"""PERSISTENT INFORMATION:
{facts_str}"""
return response.replace("<KNOWN_FACTS>", persistent_info)
@ai_track("Fact Extraction")
@observe(as_type="generation")
async def extract_facts_long_term(chat_history: str) -> list[str]:
logger.debug("Starting fact extraction from chat history")
extract_start = time.time()
system_prompt = """
You are an AI assistant specialized in extracting and formatting relevant information about users from conversations. Your task is to analyze a given conversation and create a list of concise, factual statements about the user. These statements will be stored in a vector embedding database to enhance future interactions.
Here is the conversation you need to analyze:
<conversation>
{chat_history}
</conversation>
Instructions:
1. Carefully read through the conversation. Extract only new facts, from only the last message sent by the user - treat the rest of the conversation only as context. Ignore facts in the last message that are already stated in the conversation.
2. Identify key new pieces of information from the last message sent by the user that would be valuable for future interactions. Look for:
- Personal details (name, age, occupation, location, etc.)
- Preferences (likes, dislikes, interests, hobbies)
- Experiences (travel, education, work history)
- Expressive style (writing style, tone, etc.)
- Relationships (family, friends, pets)
- Goals or aspirations
- Challenges or problems they're facing
- Opinions or beliefs
3. For each piece of information you identify:
a. Verify that it is factual and explicitly stated in the conversation, not inferred.
b. Formulate it as a concise statement that would aid in semantic retrieval.
c. Ensure it is not similar to information previously stated in the conversation.
4. Before providing your final output, wrap your analysis in <information_extraction> tags. In this analysis:
- List each piece of information you've identified.
- For each piece of information:
* Quote the relevant part of the conversation.
* Categorize the information (e.g., personal detail, preference, experience).
* Explain why you've included this information.
* Show how you've formulated the fact for optimal semantic retrieval.
- Discuss any challenges you encountered in extracting or formatting the information.
5. After your analysis, provide your final output as a JSON array of strings. Each string should be a single fact about the user. Wrap the facts in <facts> tags.
Example of the expected output format:
<information_extraction>
[Analysis goes here]
</information_extraction>
<facts>
{{
"facts":
[
"User is 28 years old",
"User's friend Mary works as a software engineer",
"Favorite food is sushi"
]
}}
</facts>
Remember to focus on clear, concise statements that capture key information about the user. Each fact should be worded in a way that will aid its semantic retrieval from a vector embedding database. It's OK for this section to be quite long.
"""
message = system_prompt.format(chat_history=chat_history)
messages = [
{
"role": "user",
"content": message
}
]
logger.debug("Calling LLM for fact extraction")
llm_start = time.time()
# Create a new model client
client = ModelClient(provider=FACT_EXTRACTION_PROVIDER, model=FACT_EXTRACTION_MODEL)
# Generate the response with caching enabled
response = await client.generate(
messages=messages,
max_tokens=1000,
temperature=0.0,
use_caching=True # Enable caching for the system prompt
)
llm_time = time.time() - llm_start
logger.debug(f"LLM response received in {llm_time:.2f}s")
try:
logger.debug("Parsing JSON response")
facts_str = parse_xml_content(response, "facts")
response_data = json.loads(facts_str)
facts = response_data["facts"]
logger.debug(f"Extracted {len(facts)} facts")
if facts:
logger.debug(f"Sample facts: {facts[:3] if len(facts) > 3 else facts}")
except (json.JSONDecodeError, KeyError) as e:
logger.error(f"Error parsing response: {str(e)}")
facts = []
total_time = time.time() - extract_start
logger.debug(f"Total extraction completed in {total_time:.2f}s")
return facts

176
src/deriver/tom/single_prompt.py
View File

@ -1,28 +1,18 @@
import os
import logging
from typing import Any, Optional
import sentry_sdk
from anthropic import Anthropic
from langfuse.decorators import langfuse_context, observe
from sentry_sdk.ai.monitoring import ai_track
# Place the code below at the beginning of your application to initialize the tracer
from src.utils.model_client import ModelClient, ModelProvider
# Initialize the Anthropic client
anthropic = Anthropic(
api_key=os.getenv("ANTHROPIC_API_KEY"),
max_retries=5,
)
logger = logging.getLogger(__name__)
ANTHROPIC_MODEL = "claude-3-5-haiku-20241022"
DEF_PROVIDER = ModelProvider.GROQ
DEF_MODEL = "llama-3.3-70b-versatile"
@ai_track("Tom Inference")
@observe(as_type="generation")
async def get_tom_inference_single_prompt(
chat_history: str, session_id: str, user_representation: str = "None", **kwargs
) -> str:
with sentry_sdk.start_transaction(op="tom-inference", name="ToM Inference"):
system_prompt = """You are a system for analyzing conversations to make evidence-based inferences about user mental states.
TOM_SYSTEM_PROMPT = """You are a system for analyzing conversations to make evidence-based inferences about user mental states.
REQUIREMENTS:
1. Only make inferences that are directly supported by conversation evidence
@ -56,53 +46,9 @@ EXPECTATION VIOLATIONS:
- Based on the above information, if the next message were to surprise you, what could it contain?
- Format: "POTENTIAL SURPRISE: [possible content] [reason] [confidence level]"
- Include 3-5 possible surprises
</prediction>
"""
</prediction>"""
messages = [
{
"role": "user",
"content": f"Please analyze this conversation and provide a prediction following the format above:\n{chat_history}",
}
]
# Add existing user representation if available
if user_representation != "None":
messages.append(
{
"role": "user",
"content": f"Consider this existing user representation for context, but focus on current state:\n{user_representation}",
}
)
langfuse_context.update_current_observation(
input=messages, model=ANTHROPIC_MODEL
)
message = anthropic.messages.create(
model=ANTHROPIC_MODEL,
max_tokens=1000,
temperature=0,
messages=messages,
system=system_prompt,
)
print(f"tom_inference in single_prompt.py: {message.content[0].text=}")
message = message.content[0].text
return message
@ai_track("User Representation")
@observe(as_type="generation")
async def get_user_representation_single_prompt(
chat_history: str,
session_id: str,
user_representation: str = "None",
tom_inference: str = "None",
**kwargs,
) -> str:
with sentry_sdk.start_transaction(
op="user-representation-inference", name="User Representation"
):
system_prompt = """You are a system for maintaining factual user representations based on conversation history and theory of mind analysis.
USER_REPRESENTATION_SYSTEM_PROMPT = """You are a system for maintaining factual user representations based on conversation history and theory of mind analysis.
Your job is to update the existing user representation (if provided) with the new information from the conversation history and theory of mind analysis.
@ -152,36 +98,102 @@ UPDATES:
- New Information: Recent observations
- Changes: Modified interpretations
- Removals: Information no longer supported
</representation>
"""
</representation>"""
messages = []
print(f"in single_prompt.py: chat_history: {chat_history}")
print(f"in single_prompt.py: user_representation: {user_representation}")
@ai_track("Tom Inference")
@observe(as_type="generation")
async def get_tom_inference_single_prompt(
chat_history: str, session_id: str, user_representation: Optional[str] = None, **kwargs
) -> str:
with sentry_sdk.start_transaction(op="tom-inference", name="ToM Inference"):
# Create a new model client
client = ModelClient(provider=DEF_PROVIDER, model=DEF_MODEL)
# Prepare the messages
messages: list[dict[str, Any]] = [
{
"role": "user",
"content": f"Please analyze this conversation and provide a prediction following the format above:\n{chat_history}",
}
]
# Add existing user representation if available
if user_representation:
messages.append(
{
"role": "user",
"content": f"Consider this existing user representation for context, but focus on current state:\n{user_representation}",
}
)
langfuse_context.update_current_observation(
input=messages, model=DEF_MODEL
)
# Generate the response with caching enabled
try:
response = await client.generate(
messages=messages,
system=TOM_SYSTEM_PROMPT,
max_tokens=1000,
temperature=0,
use_caching=True # Enable caching for the system prompt
)
except Exception as e:
sentry_sdk.capture_exception(e)
logger.error(f"Error generating Tom inference: {e}")
raise e
return response
@ai_track("User Representation")
@observe(as_type="generation")
async def get_user_representation_single_prompt(
chat_history: str,
session_id: str,
user_representation: Optional[str] = None,
tom_inference: Optional[str] = None,
**kwargs,
) -> str:
with sentry_sdk.start_transaction(
op="user-representation-inference", name="User Representation"
):
# Create a new model client
client = ModelClient(provider=DEF_PROVIDER, model=DEF_MODEL)
# Build the context message
context_str = f"CONVERSATION:\n{chat_history}\n\n"
if tom_inference != "None":
if tom_inference:
context_str += f"PREDICTION OF USER MENTAL STATE - MIGHT BE INCORRECT:\n{tom_inference}\n\n"
if user_representation != "None":
if user_representation:
context_str += f"EXISTING USER REPRESENTATION - INCOMPLETE, TO BE UPDATED:\n{user_representation}"
messages.append(
# Prepare the messages
messages: list[dict[str, Any]] = [
{
"role": "user",
"content": f"Please analyze this information and provide an updated user representation:\n{context_str}",
}
)
]
langfuse_context.update_current_observation(
input=messages, model=ANTHROPIC_MODEL
input=messages, model=DEF_MODEL
)
message = anthropic.messages.create(
model=ANTHROPIC_MODEL,
max_tokens=1000,
temperature=0,
messages=messages,
system=system_prompt,
)
message = message.content[0].text
return message
# Generate the response with caching enabled
try:
response = await client.generate(
messages=messages,
system=USER_REPRESENTATION_SYSTEM_PROMPT,
max_tokens=1000,
temperature=0,
use_caching=True # Enable caching for the system prompt
)
except Exception as e:
sentry_sdk.capture_exception(e)
logger.error(f"Error generating user representation: {e}")
raise e
return response

19
src/utils/__init__.py Normal file
View File

@ -0,0 +1,19 @@
"""
Utility modules for the Honcho app.
"""
import re
def parse_xml_content(text: str, tag: str) -> str:
"""
Extract content from XML-like tags in a string.
Args:
text: The text containing XML-like tags
tag: The tag name to extract content from
Returns:
The content between the opening and closing tags, or an empty string if not found
"""
pattern = f"<{tag}>(.*?)</{tag}>"
match = re.search(pattern, text, re.DOTALL)
return match.group(1).strip() if match else ""

322
src/utils/model_client.py Normal file
View File

@ -0,0 +1,322 @@
"""
Utility functions for interacting with various language model APIs.
"""
import os
from enum import Enum
from typing import Any, Optional, Protocol
import sentry_sdk
from anthropic import AsyncAnthropic
from dotenv import load_dotenv
from langfuse.decorators import langfuse_context, observe
from openai import AsyncOpenAI
# Load environment variables
load_dotenv()
# Supported model providers
class ModelProvider(str, Enum):
ANTHROPIC = "anthropic"
OPENAI = "openai"
OPENROUTER = "openrouter"
CEREBRAS = "cerebras"
GROQ = "groq"
# Add other providers as needed
# Default models for each provider
DEFAULT_MODELS = {
ModelProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
ModelProvider.OPENAI: "gpt-4o",
ModelProvider.OPENROUTER: "meta-llama/Llama-3.3-70B-Instruct",
ModelProvider.CEREBRAS: "llama-3.3-70b",
ModelProvider.GROQ: "llama-3.3-70b-versatile",
}
OPENAI_COMPATIBLE_PROVIDERS = [
ModelProvider.OPENAI,
ModelProvider.OPENROUTER,
ModelProvider.CEREBRAS,
ModelProvider.GROQ
]
DEFAULT_TEMPERATURE = 0.0
DEFAULT_MAX_TOKENS = 1000
class Message(Protocol):
"""Protocol for a message that works with any provider."""
role: str
content: str
class ModelClient:
"""A client for interacting with various language model APIs."""
def __init__(
self,
provider: ModelProvider = ModelProvider.ANTHROPIC,
model: Optional[str] = None,
api_key: Optional[str] = None,
base_url: Optional[str] = None
):
"""
Initialize the model client.
Args:
provider: The model provider to use
model: The specific model to use, or None to use the default model for the provider
api_key: The API key to use, or None to read from environment variables
base_url: Custom base URL for the API endpoints (used for OpenRouter)
"""
self.provider = provider
self.model = model or DEFAULT_MODELS[provider]
self.base_url = base_url
self.openai_client = None
# Setup provider-specific clients
if provider == ModelProvider.ANTHROPIC:
self.api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
if not self.api_key:
raise ValueError("Anthropic API key is required")
self.client = AsyncAnthropic(api_key=self.api_key)
elif provider in OPENAI_COMPATIBLE_PROVIDERS:
self.api_key = api_key or os.getenv("OPENAI_COMPATIBLE_API_KEY")
self.base_url = base_url or os.getenv("OPENAI_COMPATIBLE_BASE_URL")
if not self.api_key:
raise ValueError("OpenAI-compatible API key is required")
self.openai_client = AsyncOpenAI(api_key=self.api_key, base_url=self.base_url)
else:
raise ValueError(f"Unsupported provider: {provider}")
def create_message(self, role: str, content: str) -> dict[str, str]:
"""
Create a message that works with the current provider.
Args:
role: The role of the message (e.g., "user", "assistant")
content: The message content
Returns:
A message compatible with the current provider
"""
# For now, just return a dictionary that works with both Anthropic and OpenAI
return {"role": role, "content": content}
@observe(as_type="generation")
async def generate(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False
) -> str:
"""
Generate a response using the configured model.
Args:
messages: The conversation history
system: Optional system prompt
max_tokens: Maximum number of tokens to generate
temperature: Temperature for generation
extra_headers: Optional headers to add to the request
use_caching: Whether to use provider-side caching for the response
Returns:
The generated text
"""
with sentry_sdk.start_transaction(op="llm-api", name=f"{self.provider} API Call"):
# Log to langfuse
langfuse_context.update_current_observation(
input=messages, model=self.model
)
if self.provider == ModelProvider.ANTHROPIC:
return await self._generate_anthropic(messages, system, max_tokens, temperature, extra_headers, use_caching)
elif self.provider in OPENAI_COMPATIBLE_PROVIDERS:
return await self._generate_openai(messages, system, max_tokens, temperature)
else:
raise ValueError(f"Unsupported provider: {self.provider}")
async def _generate_anthropic(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False
) -> str:
"""Generate a response using the Anthropic API."""
if not self.client:
raise ValueError("Anthropic client not initialized.")
params = {
"model": self.model,
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
}
# Handle system prompt with caching if enabled
if system:
if use_caching:
params["system"] = [
{
"type": "text",
"text": system,
"cache_control": {"type": "ephemeral"}
}
]
else:
params["system"] = system
response = await self.client.messages.create(**params)
# Extract the text from the response
if response.content and len(response.content) > 0:
content_block = response.content[0]
# Check content_block by checking for attribute 'type' instead of using isinstance
if hasattr(content_block, 'type') and content_block.type == "text":
return content_block.text
return str(content_block)
return ""
async def _generate_openai(
self,
messages: list[dict[str, str]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE
) -> str:
"""Generate text using OpenAI or OpenRouter API."""
if not self.openai_client:
raise ValueError("OpenAI client not initialized")
# Prepare messages
formatted_messages = []
# Add system message if provided
if system:
formatted_messages.append({"role": "system", "content": system})
# Add the rest of the messages
formatted_messages.extend(messages)
# Make the API call with the OpenAI client
response = await self.openai_client.chat.completions.create(
model=self.model,
messages=formatted_messages,
max_tokens=max_tokens,
temperature=temperature
)
# Extract the generated text
choice = response.choices[0]
if choice and choice.message and choice.message.content:
return choice.message.content
return ""
@observe(as_type="generation")
async def stream(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False
) -> Any:
"""
Stream a response using the configured model.
Args:
messages: The conversation history
system: Optional system prompt
max_tokens: Maximum number of tokens to generate
temperature: Temperature for generation
extra_headers: Optional headers to add to the request
use_caching: Whether to use provider-side caching for the response
Returns:
A streaming response from the provider
"""
with sentry_sdk.start_transaction(op="llm-api-stream", name=f"{self.provider} API Stream"):
# Log to langfuse
langfuse_context.update_current_observation(
input=messages, model=self.model
)
if self.provider == ModelProvider.ANTHROPIC:
return await self._stream_anthropic(messages, system, max_tokens, temperature, extra_headers, use_caching)
elif self.provider in OPENAI_COMPATIBLE_PROVIDERS:
return await self._stream_openai(messages, system, max_tokens, temperature)
else:
raise ValueError(f"Unsupported provider: {self.provider}")
async def _stream_anthropic(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False
) -> Any:
"""Stream text using Anthropic API."""
if not self.client:
raise ValueError("Anthropic client not initialized.")
params = {
"model": self.model,
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature
}
# Handle system prompt with caching if enabled
if system:
if use_caching:
params["system"] = [
{
"type": "text",
"text": system,
"cache_control": {"type": "ephemeral"}
}
]
else:
params["system"] = system
# Return the stream directly without awaiting it
return self.client.messages.stream(**params)
async def _stream_openai(
self,
messages: list[dict[str, str]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE
) -> Any:
"""Stream text using OpenAI or OpenRouter API."""
if not self.openai_client:
raise ValueError("OpenAI client not initialized")
# Prepare messages
formatted_messages = []
# Add system message if provided
if system:
formatted_messages.append({"role": "system", "content": system})
# Add the rest of the messages
formatted_messages.extend(messages)
# Make the API call with the OpenAI client
stream = await self.openai_client.chat.completions.create(
model=self.model,
messages=formatted_messages,
max_tokens=max_tokens,
temperature=temperature,
stream=True
)
return stream

42
tests/conftest.py
View File

@ -2,6 +2,7 @@ import logging # noqa: I001
import os
import sys
from nanoid import generate as generate_nanoid
from unittest.mock import patch, MagicMock
import pytest
import pytest_asyncio
@ -20,17 +21,28 @@ from src.dependencies import get_db
from src.exceptions import HonchoException
from src.main import app
# Create a custom handler that doesn't get closed prematurely
class TestHandler(logging.Handler):
def __init__(self):
super().__init__()
self.records = []
def emit(self, record):
self.records.append(record)
# Setup logging with our custom handler
test_handler = TestHandler()
logging.basicConfig(
level=logging.DEBUG,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
stream=sys.stdout, # This ensures the output goes to stdout
handlers=[test_handler]
)
logger = logging.getLogger(__name__)
logging.getLogger("sqlalchemy.engine.Engine").disabled = True
# Test database URL
# TODO use environment variable
CONNECTION_URI = make_url(os.getenv("CONNECTION_URI"))
CONNECTION_URI = make_url(os.getenv("CONNECTION_URI", "postgresql+psycopg://postgres:postgres@localhost:5432/postgres"))
TEST_DB_URL = CONNECTION_URI.set(database="test_db")
DEFAULT_DB_URL = str(CONNECTION_URI.set(database="postgres"))
@ -152,3 +164,29 @@ async def sample_data(db_session):
yield test_app, test_user
await db_session.rollback()
@pytest.fixture(autouse=True)
def mock_langfuse():
"""Mock Langfuse decorator and context during tests"""
with patch("langfuse.decorators.observe") as mock_observe, \
patch("langfuse.decorators.langfuse_context") as mock_context:
# Mock the decorator to just return the function
mock_observe.return_value = lambda func: func
# Mock the context object
mock_context_obj = MagicMock()
mock_context_obj.update_current_observation = MagicMock()
mock_context_obj.update_current_trace = MagicMock()
mock_context.return_value = mock_context_obj
# Disable httpx logging during tests
logging.getLogger("httpx").setLevel(logging.WARNING)
yield
# Clean up logging handlers
for handler in logging.getLogger().handlers[:]:
if isinstance(handler, TestHandler):
handler.close()
logging.getLogger().removeHandler(handler)

44
tests/routes/test_validation_api.py
View File

@ -374,22 +374,33 @@ def test_session_validations_api(client, sample_data):
def test_agent_query_validations_api(client, sample_data, monkeypatch):
# Mock the functions in agent.py that are causing the database issues
async def mock_get_user_representation(*args, **kwargs):
return "Mock user representation"
# Create a mock collection with a public_id
class MockCollection:
def __init__(self):
self.public_id = "mock_collection_id"
# Mock collection retrieval/creation function
async def mock_get_or_create_collection(*args, **kwargs):
return MockCollection()
async def mock_chat_history(*args, **kwargs):
return "Mock chat history"
async def mock_get_long_term_facts(*args, **kwargs):
return ["Mock fact 1", "Mock fact 2"]
async def mock_run_tom_inference(*args, **kwargs):
return "Mock TOM inference"
async def mock_generate_user_representation(*args, **kwargs):
return "Mock user representation"
# Mock the Dialectic.call method
def mock_dialectic_call(self):
# Create a mock response that will work with line 179 in agent.py:
# return schemas.AgentChat(content=response[0].text)
class MockText:
def __init__(self):
self.text = "Mock response"
# Return a list with MockText object at index 0
return [MockText()]
async def mock_dialectic_call(self):
# Create a mock response that will work with line 300 in agent.py:
# return schemas.AgentChat(content=response[0]["text"])
return [{"text": "Mock response"}]
# Mock the Dialectic.stream method
def mock_dialectic_stream(self):
@ -407,10 +418,11 @@ def test_agent_query_validations_api(client, sample_data, monkeypatch):
return MockStream()
# Apply the monkeypatches
monkeypatch.setattr(
"src.agent.get_latest_user_representation", mock_get_user_representation
)
monkeypatch.setattr("src.agent.chat_history", mock_chat_history)
monkeypatch.setattr("src.crud.get_or_create_user_protected_collection", mock_get_or_create_collection)
monkeypatch.setattr("src.agent.get_chat_history", mock_chat_history)
monkeypatch.setattr("src.agent.get_long_term_facts", mock_get_long_term_facts)
monkeypatch.setattr("src.agent.run_tom_inference", mock_run_tom_inference)
monkeypatch.setattr("src.agent.generate_user_representation", mock_generate_user_representation)
monkeypatch.setattr("src.agent.Dialectic.call", mock_dialectic_call)
monkeypatch.setattr("src.agent.Dialectic.stream", mock_dialectic_stream)

192
tests/utils/test_model_client.py Normal file
View File

@ -0,0 +1,192 @@
import os
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from src.utils.model_client import (
DEFAULT_MAX_TOKENS,
DEFAULT_MODELS,
DEFAULT_TEMPERATURE,
ModelClient,
ModelProvider,
)
@pytest.fixture
def mock_env(monkeypatch):
monkeypatch.setenv("ANTHROPIC_API_KEY", "mock-anthropic-api-key")
monkeypatch.setenv("OPENAI_COMPATIBLE_API_KEY", "mock-openai-api-key")
monkeypatch.setenv("OPENAI_API_KEY", "mock-openai-api-key")
monkeypatch.setenv("GROQ_API_KEY", "mock-groq-api-key")
monkeypatch.setenv("CEREBRAS_API_KEY", "mock-cerebras-api-key")
monkeypatch.setenv("OPENROUTER_API_KEY", "mock-openrouter-api-key")
# Test fixtures
@pytest.fixture
def mock_anthropic_response():
mock_response = MagicMock()
mock_content = MagicMock()
mock_content.type = "text"
mock_content.text = "Test response"
mock_response.content = [mock_content]
return mock_response
@pytest.fixture
def mock_openai_response():
mock_response = MagicMock()
mock_message = MagicMock()
mock_message.content = "Test response"
mock_choice = MagicMock()
mock_choice.message = mock_message
mock_response.choices = [mock_choice]
return mock_response
@pytest.fixture
def mock_anthropic_client():
with patch("src.utils.model_client.AsyncAnthropic") as mock:
client = MagicMock()
client.messages.create = AsyncMock()
client.messages.stream = AsyncMock()
mock.return_value = client
yield client
@pytest.fixture
def mock_openai_client():
with patch("src.utils.model_client.AsyncOpenAI") as mock:
client = MagicMock()
client.chat.completions.create = AsyncMock()
mock.return_value = client
yield client
# Initialization Tests
def test_default_initialization(mock_env):
"""Test default initialization with Anthropic provider."""
client = ModelClient()
assert client.provider == ModelProvider.ANTHROPIC
assert client.model == DEFAULT_MODELS[ModelProvider.ANTHROPIC]
assert client.base_url is None
def test_custom_model_initialization():
"""Test initialization with custom model name."""
custom_model = "custom-model"
client = ModelClient(model=custom_model)
assert client.model == custom_model
def test_custom_api_key_initialization():
"""Test initialization with custom API key."""
custom_key = "test-api-key"
client = ModelClient(api_key=custom_key)
assert client.api_key == custom_key
def test_custom_base_url_initialization():
"""Test initialization with custom base URL."""
custom_url = "https://custom-api.example.com"
client = ModelClient(base_url=custom_url)
assert client.base_url == custom_url
def test_unsupported_provider_initialization():
"""Test initialization with unsupported provider."""
with pytest.raises(ValueError, match="is not a valid ModelProvider"):
ModelClient(provider=ModelProvider("unsupported"))
def test_missing_api_key_initialization():
"""Test initialization without required API key."""
with patch.dict(os.environ, {}, clear=True):
with pytest.raises(ValueError, match="API key is required"):
ModelClient()
# Message Creation Tests
def test_create_message():
"""Test message creation with different roles."""
client = ModelClient()
message = client.create_message("user", "Hello")
assert message == {"role": "user", "content": "Hello"}
# Generation Tests
@pytest.mark.asyncio
async def test_generate_anthropic(mock_anthropic_client, mock_anthropic_response):
"""Test generation with Anthropic provider."""
mock_anthropic_client.messages.create.return_value = mock_anthropic_response
client = ModelClient(provider=ModelProvider.ANTHROPIC)
messages = [{"role": "user", "content": "Hello"}]
response = await client.generate(messages)
assert response == "Test response"
# Verify the API call
mock_anthropic_client.messages.create.assert_called_once()
call_args = mock_anthropic_client.messages.create.call_args[1]
assert call_args["model"] == DEFAULT_MODELS[ModelProvider.ANTHROPIC]
assert call_args["messages"] == messages
assert call_args["max_tokens"] == DEFAULT_MAX_TOKENS
assert call_args["temperature"] == DEFAULT_TEMPERATURE
@pytest.mark.asyncio
async def test_generate_openai(mock_openai_client, mock_openai_response, mock_env):
"""Test generation with OpenAI provider."""
mock_openai_client.chat.completions.create.return_value = mock_openai_response
client = ModelClient(provider=ModelProvider.OPENAI)
messages = [{"role": "user", "content": "Hello"}]
response = await client.generate(messages)
assert response == "Test response"
# Verify the API call
mock_openai_client.chat.completions.create.assert_called_once()
call_args = mock_openai_client.chat.completions.create.call_args[1]
assert call_args["model"] == DEFAULT_MODELS[ModelProvider.OPENAI]
assert call_args["messages"] == messages
assert call_args["max_tokens"] == DEFAULT_MAX_TOKENS
assert call_args["temperature"] == DEFAULT_TEMPERATURE
@pytest.mark.asyncio
async def test_generate_with_system_prompt(
mock_anthropic_client, mock_anthropic_response
):
"""Test generation with system prompt."""
mock_anthropic_client.messages.create.return_value = mock_anthropic_response
client = ModelClient(provider=ModelProvider.ANTHROPIC)
messages = [{"role": "user", "content": "Hello"}]
system = "You are a helpful assistant"
response = await client.generate(messages, system=system)
assert response == "Test response"
# Verify the API call
mock_anthropic_client.messages.create.assert_called_once()
call_args = mock_anthropic_client.messages.create.call_args[1]
assert call_args["system"] == system
@pytest.mark.asyncio
async def test_generate_with_caching(mock_anthropic_client, mock_anthropic_response):
"""Test generation with caching enabled."""
mock_anthropic_client.messages.create.return_value = mock_anthropic_response
client = ModelClient(provider=ModelProvider.ANTHROPIC)
messages = [{"role": "user", "content": "Hello"}]
system = "You are a helpful assistant"
response = await client.generate(messages, system=system, use_caching=True)
assert response == "Test response"
# Verify the API call
mock_anthropic_client.messages.create.assert_called_once()
call_args = mock_anthropic_client.messages.create.call_args[1]
assert call_args["system"] == [
{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}
]