diff --git a/pyproject.toml b/pyproject.toml index 74252985..5844d51f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -52,7 +52,7 @@ select = [ # isort "I", ] -ignore = ["E501"] +ignore = ["E501", "B008"] [tool.ruff.lint.flake8-bugbear] extend-immutable-calls = ["fastapi.Depends"] diff --git a/src/agent.py b/src/agent.py index 8223d5ce..591604b4 100644 --- a/src/agent.py +++ b/src/agent.py @@ -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""" {self.agent_input} {self.user_representation} {self.chat_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""" {self.agent_input} {self.user_representation} {self.chat_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 \ No newline at end of file + return representation diff --git a/src/crud.py b/src/crud.py index 55992a0f..918d1ee3 100644 --- a/src/crud.py +++ b/src/crud.py @@ -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 diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index 255cdf79..5e6d787e 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -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") diff --git a/src/deriver/queue.py b/src/deriver/queue.py index d77bbdcb..33d87e5a 100644 --- a/src/deriver/queue.py +++ b/src/deriver/queue.py @@ -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") diff --git a/src/deriver/tom/__init__.py b/src/deriver/tom/__init__.py index 310604d0..5d6a6452 100644 --- a/src/deriver/tom/__init__.py +++ b/src/deriver/tom/__init__.py @@ -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}") diff --git a/src/deriver/tom/embeddings.py b/src/deriver/tom/embeddings.py index 4ed29601..5c95df20 100644 --- a/src/deriver/tom/embeddings.py +++ b/src/deriver/tom/embeddings.py @@ -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 \ No newline at end of file + + return unique_facts diff --git a/src/deriver/tom/long_term.py b/src/deriver/tom/long_term.py index d08707d3..7cf5f004 100644 --- a/src/deriver/tom/long_term.py +++ b/src/deriver/tom/long_term.py @@ -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 \ No newline at end of file + return facts diff --git a/src/deriver/tom/single_prompt.py b/src/deriver/tom/single_prompt.py index 624d3a54..b72f9fbf 100644 --- a/src/deriver/tom/single_prompt.py +++ b/src/deriver/tom/single_prompt.py @@ -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 diff --git a/src/routers/collections.py b/src/routers/collections.py index bfe36455..7cd5d0ee 100644 --- a/src/routers/collections.py +++ b/src/routers/collections.py @@ -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" diff --git a/src/routers/documents.py b/src/routers/documents.py index e8111a4e..0c5926ec 100644 --- a/src/routers/documents.py +++ b/src/routers/documents.py @@ -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""" diff --git a/src/routers/keys.py b/src/routers/keys.py index 45888e54..87705559 100644 --- a/src/routers/keys.py +++ b/src/routers/keys.py @@ -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: diff --git a/src/routers/messages.py b/src/routers/messages.py index ec376628..602508ed 100644 --- a/src/routers/messages.py +++ b/src/routers/messages.py @@ -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""" diff --git a/src/routers/metamessages.py b/src/routers/metamessages.py index 4e280d2a..b87cf130 100644 --- a/src/routers/metamessages.py +++ b/src/routers/metamessages.py @@ -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""" diff --git a/src/routers/sessions.py b/src/routers/sessions.py index 6a0d5baf..5d49eec9 100644 --- a/src/routers/sessions.py +++ b/src/routers/sessions.py @@ -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""" diff --git a/src/routers/users.py b/src/routers/users.py index 721d1f99..da1661d4 100644 --- a/src/routers/users.py +++ b/src/routers/users.py @@ -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, ): diff --git a/src/utils/__init__.py b/src/utils/__init__.py index 343f4bce..aead6593 100644 --- a/src/utils/__init__.py +++ b/src/utils/__init__.py @@ -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}>(.*?)" match = re.search(pattern, text, re.DOTALL) - return match.group(1).strip() if match else "" \ No newline at end of file + return match.group(1).strip() if match else "" diff --git a/src/utils/model_client.py b/src/utils/model_client.py index cf9b85e9..8979c334 100644 --- a/src/utils/model_client.py +++ b/src/utils/model_client.py @@ -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 \ No newline at end of file + + return stream diff --git a/tests/conftest.py b/tests/conftest.py index 1bc41881..0d111596 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -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 \ No newline at end of file + yield mock_create diff --git a/tests/routes/test_users.py b/tests/routes/test_users.py index f0e4d4dd..6cce77e9 100644 --- a/tests/routes/test_users.py +++ b/tests/routes/test_users.py @@ -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 diff --git a/tests/routes/test_validation_api.py b/tests/routes/test_validation_api.py index d07f66d7..02adb5a5 100644 --- a/tests/routes/test_validation_api.py +++ b/tests/routes/test_validation_api.py @@ -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) diff --git a/tests/utils/test_model_client.py b/tests/utils/test_model_client.py index 6e533bfd..2b51fc50 100644 --- a/tests/utils/test_model_client.py +++ b/tests/utils/test_model_client.py @@ -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"}} ] -