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

chore: run ruff formatter

This commit is contained in:
Vineeth Voruganti 2025-04-07 23:00:01 -04:00
parent 62f04a14d5
commit 225e2b8ce9
22 changed files with 496 additions and 350 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
pyproject.toml
View File

@ -52,7 +52,7 @@ select = [
# isort
"I",
]
ignore = ["E501"]
ignore = ["E501", "B008"]
[tool.ruff.lint.flake8-bugbear]
extend-immutable-calls = ["fastapi.Depends"]

247
src/agent.py
View File

@ -10,7 +10,7 @@ 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 func, select
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models, schemas
@ -63,7 +63,9 @@ class Dialectic:
self.agent_input = agent_input
self.user_representation = user_representation
self.chat_history = chat_history
self.client = ModelClient(provider=ModelProvider.ANTHROPIC, 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")
@ -72,33 +74,34 @@ class Dialectic:
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)}")
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")
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
}
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,
max_tokens=1000
messages=[message], system=self.system_prompt, max_tokens=1000
)
model_time = asyncio.get_event_loop().time() - model_start
logger.debug(f"Model response received in {model_time:.2f}s: {len(response)} chars")
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}]
@ -109,33 +112,32 @@ class Dialectic:
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)}")
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>
"""
logger.debug(f"Prompt constructed with context length: {len(self.user_representation)} chars")
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
}
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,
max_tokens=1000
messages=[message], system=self.system_prompt, 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
@ -147,16 +149,16 @@ async def get_chat_history(app_id: str, user_id: str, session_id: str) -> str:
stmt = await crud.get_messages(db, app_id, user_id, session_id)
results = await db.execute(stmt)
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
@ -164,10 +166,13 @@ async def get_chat_history(app_id: str, user_id: str, session_id: str) -> str:
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")
logger.debug(
f"Constructed history with {user_count} user messages and {assistant_count} assistant messages"
)
return history
@observe()
async def chat(
app_id: str,
@ -178,7 +183,7 @@ async def chat(
) -> 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:
@ -191,26 +196,30 @@ async def chat(
# 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:
# 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)
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
collection_id=collection.public_id, # type: ignore
)
logger.debug(f"Created embedding store with collection_id: {collection.public_id if collection else None}")
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)
@ -228,15 +237,14 @@ async def chat(
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'))
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
@ -245,13 +253,10 @@ async def chat(
short_term_task = run_tom_inference(history, session_id)
# Wait for both tasks to complete
facts, tom_inference = await asyncio.gather(
long_term_task,
short_term_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(
@ -264,9 +269,11 @@ async def chat(
embedding_store=embedding_store,
db=db,
message_id=latest_message_id,
with_inference=False
with_inference=False,
)
logger.debug(
f"User representation generated: {len(user_representation)} characters"
)
logger.debug(f"User representation generated: {len(user_representation)} characters")
# Create a Dialectic chain with the fresh user representation
chain = Dialectic(
@ -274,7 +281,7 @@ async def chat(
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")
@ -290,98 +297,97 @@ async def chat(
query_start_time = asyncio.get_event_loop().time()
if stream:
response_stream = await chain.stream()
logger.debug(f"Dialectic stream started after {asyncio.get_event_loop().time() - query_start_time:.2f}s")
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)")
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
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}")
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
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")
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)]
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")
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:
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=""
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
@ -389,41 +395,38 @@ 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)
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
}
]
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
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)
@ -435,15 +438,15 @@ async def generate_semantic_queries(query: str) -> list[str]:
# 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)}")
@ -460,25 +463,28 @@ async def generate_user_representation(
embedding_store: CollectionEmbeddingStore,
db: AsyncSession,
message_id: Optional[str] = None,
with_inference: bool = False
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.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")
@ -488,13 +494,15 @@ async def generate_user_representation(
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
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"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()
@ -508,9 +516,11 @@ async def generate_user_representation(
)
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")
representation = parse_xml_content(
user_representation_response, "representation"
)
logger.debug(f"Extracted representation: {len(representation)} characters")
else:
representation = f"""
@ -533,13 +543,12 @@ RELEVANT LONG-TERM FACTS ABOUT THE USER:
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_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:
@ -558,7 +567,7 @@ RELEVANT LONG-TERM FACTS ABOUT THE USER:
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
return representation

4
src/crud.py
View File

@ -375,9 +375,7 @@ async def get_session(
result = await db.execute(stmt)
session = result.scalar_one_or_none()
if session is None:
logger.warning(
f"Session with ID '{session_id}' not found for user {user_id}"
)
logger.warning(f"Session with ID '{session_id}' not found for user {user_id}")
raise ResourceNotFoundException("Session not found or does not belong to user")
return session

20
src/deriver/consumer.py
View File

@ -59,7 +59,9 @@ async def get_chat_history(db, session_id, message_id, limit: int = 10) -> str:
async def process_item(db: AsyncSession, payload: dict):
logger.debug(f"process_item received payload: {payload['message_id']} is_user={payload['is_user']}")
logger.debug(
f"process_item received payload: {payload['message_id']} is_user={payload['is_user']}"
)
processing_args = [
payload["content"],
payload["app_id"],
@ -124,7 +126,7 @@ async def process_user_message(
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")
@ -132,16 +134,18 @@ async def process_user_message(
db=db,
app_id=app_id,
user_id=user_id,
collection_id=collection.public_id # type: ignore
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")
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")
@ -151,8 +155,8 @@ async def process_user_message(
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")

29
src/deriver/queue.py
View File

@ -1,9 +1,8 @@
import asyncio
import logging
import os
import signal
from logging import getLogger
from datetime import datetime, timedelta
from logging import getLogger
import sentry_sdk
from dotenv import load_dotenv
@ -59,7 +58,7 @@ 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)
@ -171,7 +170,9 @@ class QueueManager:
# Track this session
self.track_session(session_id)
logger.debug(f"Claimed session {session_id} for processing")
logger.debug(
f"Claimed session {session_id} for processing"
)
# Create a new task for processing this session
if not self.shutdown_event.is_set():
@ -181,7 +182,9 @@ class QueueManager:
self.add_task(task)
except IntegrityError:
await db.rollback()
logger.debug(f"Failed to claim session {session_id}, already owned")
logger.debug(
f"Failed to claim session {session_id}, already owned"
)
else:
self.queue_empty_flag.set()
await asyncio.sleep(1)
@ -211,9 +214,11 @@ class QueueManager:
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})")
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}"
@ -234,7 +239,9 @@ class QueueManager:
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}")
logger.debug(
f"Shutdown requested, stopping processing for session {session_id}"
)
break
# Update last_updated timestamp to show this session is still being processed
@ -244,8 +251,10 @@ class QueueManager:
.values(last_updated=func.now())
)
await db.commit()
logger.debug(f"Completed processing session {session_id}, processed {message_count} messages")
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")

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

@ -1,33 +1,52 @@
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 .long_term import get_user_representation_long_term
from .single_prompt import (
get_tom_inference_single_prompt,
get_user_representation_single_prompt,
)
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)
return await get_user_representation_long_term(
chat_history, session_id, user_representation, tom_inference, **kwargs
)
else:
raise ValueError(f"Invalid method: {method}")

47
src/deriver/tom/embeddings.py
View File

@ -6,6 +6,7 @@ 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
@ -13,9 +14,14 @@ class CollectionEmbeddingStore:
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:
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
@ -30,20 +36,23 @@ class CollectionEmbeddingStore:
app_id=self.app_id,
user_id=self.user_id,
collection_id=self.collection_id,
duplicate_threshold=1-similarity_threshold # Convert similarity to distance
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]:
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
"""
@ -54,23 +63,25 @@ class CollectionEmbeddingStore:
collection_id=self.collection_id,
query=query,
max_distance=max_distance,
top_k=top_k
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]:
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
@ -80,18 +91,20 @@ class CollectionEmbeddingStore:
user_id=self.user_id,
collection_id=self.collection_id,
content=fact,
similarity_threshold=similarity_threshold
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}")
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
return unique_facts

50
src/deriver/tom/long_term.py
View File

@ -23,14 +23,15 @@ 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,
chat_history: str,
session_id: str,
embedding_store: CollectionEmbeddingStore,
user_representation: str = "None",
tom_inference: str = "None",
user_representation: str = "None",
tom_inference: str = "None",
facts: Optional[list[str]] = None,
) -> str:
if facts is None:
@ -88,21 +89,25 @@ UPDATES:
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}"
}]
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)
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
use_caching=True, # Enable caching for the system prompt
)
# Inject the facts into the response
@ -117,7 +122,7 @@ UPDATES:
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.
@ -175,30 +180,25 @@ Example of the expected output format:
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
}
]
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
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")
@ -210,7 +210,7 @@ Remember to focus on clear, concise statements that capture key information abou
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
return facts

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

@ -104,12 +104,15 @@ UPDATES:
@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
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]] = [
{
@ -127,10 +130,8 @@ async def get_tom_inference_single_prompt(
}
)
langfuse_context.update_current_observation(
input=messages, model=DEF_MODEL
)
langfuse_context.update_current_observation(input=messages, model=DEF_MODEL)
# Generate the response with caching enabled
try:
response = await client.generate(
@ -138,13 +139,13 @@ async def get_tom_inference_single_prompt(
system=TOM_SYSTEM_PROMPT,
max_tokens=1000,
temperature=0,
use_caching=True # Enable caching for the system prompt
)
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
@ -162,7 +163,7 @@ async def get_user_representation_single_prompt(
):
# 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:
@ -178,10 +179,8 @@ async def get_user_representation_single_prompt(
}
]
langfuse_context.update_current_observation(
input=messages, model=DEF_MODEL
)
langfuse_context.update_current_observation(input=messages, model=DEF_MODEL)
# Generate the response with caching enabled
try:
response = await client.generate(
@ -189,11 +188,11 @@ async def get_user_representation_single_prompt(
system=USER_REPRESENTATION_SYSTEM_PROMPT,
max_tokens=1000,
temperature=0,
use_caching=True # Enable caching for the system prompt
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

20
src/routers/collections.py
View File

@ -1,12 +1,12 @@
from typing import Optional
from fastapi import APIRouter, Depends, Query, Path, Body
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
from src import crud, schemas
from src.dependencies import db
from src.exceptions import AuthenticationException, ResourceNotFoundException
from src.exceptions import AuthenticationException
from src.security import JWTParams, require_auth
router = APIRouter(
@ -71,8 +71,12 @@ async def get_collection(
async def get_collections(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
options: schemas.CollectionGet = Body(..., description="Filtering options for the collections list"),
reverse: Optional[bool] = Query(False, description="Whether to reverse the order of results"),
options: schemas.CollectionGet = Body(
..., description="Filtering options for the collections list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""Get All Collections for a User"""
@ -110,7 +114,9 @@ async def get_collection_by_name(
async def create_collection(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection: schemas.CollectionCreate = Body(..., description="Collection creation parameters"),
collection: schemas.CollectionCreate = Body(
..., description="Collection creation parameters"
),
db=db,
):
"""Create a new Collection"""
@ -136,7 +142,9 @@ async def update_collection(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection to update"),
collection: schemas.CollectionUpdate = Body(..., description="Updated collection parameters"),
collection: schemas.CollectionUpdate = Body(
..., description="Updated collection parameters"
),
db=db,
):
"Update a Collection's name or metadata"

22
src/routers/documents.py
View File

@ -2,7 +2,7 @@ import logging
from collections.abc import Sequence
from typing import Optional
from fastapi import APIRouter, Depends, Query, Path, Body
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
@ -31,8 +31,12 @@ async def get_documents(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection"),
options: schemas.DocumentGet = Body(..., description="Filtering options for the documents list"),
reverse: Optional[bool] = Query(False, description="Whether to reverse the order of results"),
options: schemas.DocumentGet = Body(
..., description="Filtering options for the documents list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""Get all of the Documents in a Collection"""
@ -80,7 +84,9 @@ async def query_documents(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection"),
options: schemas.DocumentQuery = Body(..., description="Query parameters for document search"),
options: schemas.DocumentQuery = Body(
..., description="Query parameters for document search"
),
db=db,
):
"""Cosine Similarity Search for Documents"""
@ -115,7 +121,9 @@ async def create_document(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection"),
document: schemas.DocumentCreate = Body(..., description="Document creation parameters"),
document: schemas.DocumentCreate = Body(
..., description="Document creation parameters"
),
db=db,
):
"""Embed text as a vector and create a Document"""
@ -147,7 +155,9 @@ async def update_document(
user_id: str = Path(..., description="ID of the user"),
collection_id: str = Path(..., description="ID of the collection"),
document_id: str = Path(..., description="ID of the document to update"),
document: schemas.DocumentUpdate = Body(..., description="Updated document parameters"),
document: schemas.DocumentUpdate = Body(
..., description="Updated document parameters"
),
db=db,
):
"""Update the content and/or the metadata of a Document"""

8
src/routers/keys.py
View File

@ -25,8 +25,12 @@ router = APIRouter(
async def create_key(
app_id: str | None = Query(None, description="ID of the app to scope the key to"),
user_id: str | None = Query(None, description="ID of the user to scope the key to"),
session_id: str | None = Query(None, description="ID of the session to scope the key to"),
collection_id: str | None = Query(None, description="ID of the collection to scope the key to"),
session_id: str | None = Query(
None, description="ID of the session to scope the key to"
),
collection_id: str | None = Query(
None, description="ID of the collection to scope the key to"
),
):
"""Create a new Key"""
if not USE_AUTH:

32
src/routers/messages.py
View File

@ -2,7 +2,7 @@ import logging
import os
from typing import List, Optional
from fastapi import APIRouter, BackgroundTasks, Depends, Query, Path, Body
from fastapi import APIRouter, BackgroundTasks, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
from sqlalchemy.sql import insert
@ -19,11 +19,11 @@ logger = logging.getLogger(__name__)
router = APIRouter(
prefix="/apps/{app_id}/users/{user_id}/sessions/{session_id}/messages",
tags=["messages"],
dependencies=[Depends(require_auth(
app_id="app_id",
user_id="user_id",
session_id="session_id"
))],
dependencies=[
Depends(
require_auth(app_id="app_id", user_id="user_id", session_id="session_id")
)
],
)
@ -155,7 +155,9 @@ async def create_message_for_session(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session"),
message: schemas.MessageCreate = Body(..., description="Message creation parameters"),
message: schemas.MessageCreate = Body(
..., description="Message creation parameters"
),
db=db,
):
"""Adds a message to a session"""
@ -193,7 +195,9 @@ async def create_batch_messages_for_session(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session"),
batch: schemas.MessageBatchCreate = Body(..., description="Batch of messages to create"),
batch: schemas.MessageBatchCreate = Body(
..., description="Batch of messages to create"
),
db=db,
):
"""Bulk create messages for a session while maintaining order. Maximum 100 messages per batch."""
@ -239,8 +243,12 @@ async def get_messages(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session"),
options: schemas.MessageGet = Body(..., description="Filtering options for the messages list"),
reverse: Optional[bool] = Query(False, description="Whether to reverse the order of results"),
options: schemas.MessageGet = Body(
..., description="Filtering options for the messages list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""Get all messages for a session"""
@ -288,7 +296,9 @@ async def update_message(
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session"),
message_id: str = Path(..., description="ID of the message to update"),
message: schemas.MessageUpdate = Body(..., description="Updated message parameters"),
message: schemas.MessageUpdate = Body(
..., description="Updated message parameters"
),
db=db,
):
"""Update the metadata of a Message"""

23
src/routers/metamessages.py
View File

@ -1,7 +1,7 @@
import logging
from typing import Optional
from fastapi import APIRouter, Depends, Query, Path, Body
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
@ -15,10 +15,7 @@ logger = logging.getLogger(__name__)
router = APIRouter(
prefix="/apps/{app_id}/users/{user_id}/metamessages",
tags=["metamessages"],
dependencies=[Depends(require_auth(
app_id="app_id",
user_id="user_id"
))],
dependencies=[Depends(require_auth(app_id="app_id", user_id="user_id"))],
)
@ -26,7 +23,9 @@ router = APIRouter(
async def create_metamessage(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
metamessage: schemas.MetamessageCreate = Body(..., description="Metamessage creation parameters"),
metamessage: schemas.MetamessageCreate = Body(
..., description="Metamessage creation parameters"
),
db=db,
):
"""
@ -53,8 +52,12 @@ async def create_metamessage(
async def get_metamessages(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
options: schemas.MetamessageGet = Body(..., description="Filtering options for the metamessages list"),
reverse: Optional[bool] = Query(False, description="Whether to reverse the order of results"),
options: schemas.MetamessageGet = Body(
..., description="Filtering options for the metamessages list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""
@ -117,7 +120,9 @@ async def update_metamessage(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
metamessage_id: str = Path(..., description="ID of the metamessage to update"),
metamessage: schemas.MetamessageUpdate = Body(..., description="Updated metamessage parameters"),
metamessage: schemas.MetamessageUpdate = Body(
..., description="Updated metamessage parameters"
),
db=db,
):
"""Update a metamessage's metadata, type, or relationships"""

22
src/routers/sessions.py
View File

@ -2,7 +2,7 @@ import logging
from typing import Optional
from anthropic import MessageStreamManager
from fastapi import APIRouter, Depends, Query, Path, Body
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi.responses import StreamingResponse
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
@ -80,8 +80,12 @@ async def get_session(
async def get_sessions(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
options: schemas.SessionGet = Body(..., description="Filtering and pagination options for the sessions list"),
reverse: Optional[bool] = Query(False, description="Whether to reverse the order of results"),
options: schemas.SessionGet = Body(
..., description="Filtering and pagination options for the sessions list"
),
reverse: Optional[bool] = Query(
False, description="Whether to reverse the order of results"
),
db=db,
):
"""Get All Sessions for a User"""
@ -106,7 +110,9 @@ async def get_sessions(
async def create_session(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session: schemas.SessionCreate = Body(..., description="Session creation parameters"),
session: schemas.SessionCreate = Body(
..., description="Session creation parameters"
),
db=db,
):
"""Create a Session for a User"""
@ -134,7 +140,9 @@ async def update_session(
app_id: str = Path(..., description="ID of the app"),
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session to update"),
session: schemas.SessionUpdate = Body(..., description="Updated session parameters"),
session: schemas.SessionUpdate = Body(
..., description="Updated session parameters"
),
db=db,
):
"""Update the metadata of a Session"""
@ -252,7 +260,9 @@ async def clone_session(
user_id: str = Path(..., description="ID of the user"),
session_id: str = Path(..., description="ID of the session to clone"),
db=db,
message_id: Optional[str] = Query(None, description="Message ID to cut off the clone at"),
message_id: Optional[str] = Query(
None, description="Message ID to cut off the clone at"
),
deep_copy: bool = Query(False, description="Whether to deep copy metamessages"),
):
"""Clone a session, optionally up to a specific message"""

6
src/routers/users.py
View File

@ -1,7 +1,7 @@
import logging
from typing import Optional
from fastapi import APIRouter, Depends, Query, Path, Body
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi_pagination import Page
from fastapi_pagination.ext.sqlalchemy import paginate
@ -43,7 +43,9 @@ async def create_user(
)
async def get_users(
app_id: str = Path(..., description="ID of the app"),
options: schemas.UserGet = Body(..., description="Filtering options for the users list"),
options: schemas.UserGet = Body(
..., description="Filtering options for the users list"
),
reverse: bool = Query(False, description="Whether to reverse the order of results"),
db=db,
):

10
src/utils/__init__.py
View File

@ -1,19 +1,21 @@
"""
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 ""
return match.group(1).strip() if match else ""

158
src/utils/model_client.py
View File

@ -1,6 +1,7 @@
"""
Utility functions for interacting with various language model APIs.
"""
import os
from enum import Enum
from typing import Any, Optional, Protocol
@ -14,6 +15,7 @@ from openai import AsyncOpenAI
# Load environment variables
load_dotenv()
# Supported model providers
class ModelProvider(str, Enum):
ANTHROPIC = "anthropic"
@ -23,6 +25,7 @@ class ModelProvider(str, Enum):
GROQ = "groq"
# Add other providers as needed
# Default models for each provider
DEFAULT_MODELS = {
ModelProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
@ -36,30 +39,33 @@ OPENAI_COMPATIBLE_PROVIDERS = [
ModelProvider.OPENAI,
ModelProvider.OPENROUTER,
ModelProvider.CEREBRAS,
ModelProvider.GROQ
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,
self,
provider: ModelProvider = ModelProvider.ANTHROPIC,
model: Optional[str] = None,
api_key: Optional[str] = None,
base_url: 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
@ -70,7 +76,7 @@ class ModelClient:
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")
@ -82,37 +88,39 @@ class ModelClient:
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)
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,
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
use_caching: bool = False,
) -> str:
"""
Generate a response using the configured model.
Args:
messages: The conversation history
system: Optional system prompt
@ -120,23 +128,34 @@ class ModelClient:
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"):
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)
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)
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]],
@ -144,19 +163,19 @@ class ModelClient:
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False
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:
@ -164,71 +183,71 @@ class ModelClient:
{
"type": "text",
"text": system,
"cache_control": {"type": "ephemeral"}
"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":
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
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
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,
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
use_caching: bool = False,
) -> Any:
"""
Stream a response using the configured model.
Args:
messages: The conversation history
system: Optional system prompt
@ -236,23 +255,34 @@ class ModelClient:
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"):
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)
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)
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]],
@ -260,19 +290,19 @@ class ModelClient:
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False
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
"temperature": temperature,
}
# Handle system prompt with caching if enabled
if system:
if use_caching:
@ -280,43 +310,43 @@ class ModelClient:
{
"type": "text",
"text": system,
"cache_control": {"type": "ephemeral"}
"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
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
stream=True,
)
return stream
return stream

30
tests/conftest.py
View File

@ -1,6 +1,5 @@
import logging # noqa: I001
import os
import sys
import jwt
from nanoid import generate as generate_nanoid
from unittest.mock import patch, MagicMock, AsyncMock
@ -23,28 +22,35 @@ from src.exceptions import HonchoException
from src.security import create_admin_jwt, create_jwt, JWTParams
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",
handlers=[test_handler]
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", "postgresql+psycopg://postgres:postgres@localhost:5432/postgres"))
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"))
@ -215,22 +221,24 @@ async def sample_data(db_session):
@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:
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):
@ -245,4 +253,4 @@ def mock_openai_embeddings():
mock_response = AsyncMock()
mock_response.data = [MagicMock(embedding=[0.1] * 1536)]
mock_create.return_value = mock_response
yield mock_create
yield mock_create

4
tests/routes/test_users.py
View File

@ -17,7 +17,9 @@ def test_create_user(client, sample_data):
def test_get_user_by_id(client, sample_data):
test_app, test_user = sample_data
response = client.get(f"/v1/apps/{test_app.public_id}/users?user_id={test_user.public_id}")
response = client.get(
f"/v1/apps/{test_app.public_id}/users?user_id={test_user.public_id}"
)
assert response.status_code == 200
data = response.json()
assert data["name"] == test_user.name

21
tests/routes/test_validation_api.py
View File

@ -374,25 +374,25 @@ 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
# 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"
@ -418,11 +418,16 @@ def test_agent_query_validations_api(client, sample_data, monkeypatch):
return MockStream()
# Apply the monkeypatches
monkeypatch.setattr("src.crud.get_or_create_user_protected_collection", mock_get_or_create_collection)
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.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)

1
tests/utils/test_model_client.py
View File

@ -189,4 +189,3 @@ async def test_generate_with_caching(mock_anthropic_client, mock_anthropic_respo
assert call_args["system"] == [
{"type": "text", "text": system, "cache_control": {"type": "ephemeral"}}
]