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

initial attempt

This commit is contained in:
hyusap 2025-06-06 14:44:28 -04:00
parent 67d3a3e9f1
commit 36454834d8
10 changed files with 1249 additions and 2135 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
pyproject.toml
View File

@ -6,7 +6,7 @@ authors = [
{name = "Plastic Labs", email = "hello@plasticlabs.ai"},
]
readme = "README.md"
requires-python = ">=3.9"
requires-python = ">=3.10"
dependencies = [
"fastapi[standard]>=0.111.0",
"python-dotenv>=1.0.0",
@ -22,9 +22,10 @@ dependencies = [
"anthropic>=0.36.0",
"nanoid>=2.0.0",
"alembic>=1.14.0",
"langfuse>=2.57.1",
"langfuse<3",
"pyjwt>=2.10.0",
"google-genai>=1.10.0",
"mirascope[anthropic,google,groq]>=1.24.2",
]
[tool.uv]
dev-dependencies = [

251
src/agent.py
View File

@ -4,141 +4,75 @@ import logging
import os
from collections.abc import Iterable
from typing import Any, Optional
from inspect import cleandoc as c
import sentry_sdk
from anthropic import MessageStreamManager
from dotenv import load_dotenv
from langfuse.decorators import langfuse_context, observe
from sentry_sdk.ai.monitoring import ai_track
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from src import crud, models, schemas
from src import crud, models
from src.dependencies import tracked_db
from src.deriver.tom import get_tom_inference
from src.deriver.tom.embeddings import CollectionEmbeddingStore
from src.deriver.tom.long_term import get_user_representation_long_term
from src.utils import history, parse_xml_content
from src.utils.model_client import ModelClient, ModelProvider
from mirascope import llm, prompt_template
from mirascope.integrations.langfuse import with_langfuse
# Configure logging
logger = logging.getLogger(__name__)
USER_REPRESENTATION_METAMESSAGE_TYPE = "honcho_user_representation"
DEF_DIALECTIC_PROVIDER = ModelProvider.ANTHROPIC
DEF_DIALECTIC_MODEL = "claude-3-7-sonnet-20250219"
DEF_QUERY_GENERATION_PROVIDER = ModelProvider.GROQ
DEF_QUERY_GENERATION_MODEL = "llama-3.1-8b-instant"
QUERY_GENERATION_SYSTEM = """Given this query about a user, generate 3 focused search queries that would help retrieve relevant facts about the user.
Each query should focus on a specific aspect related to the original query, rephrased to maximize semantic search effectiveness.
For example, if the original query asks "what does the user like to eat?", generated queries might include "user's food preferences", "user's favorite cuisine", etc.
Format your response as a JSON array of strings, with each string being a search query.
Respond only in valid JSON, without markdown formatting or quotes, and nothing else.
Example:
["query about interests", "query about personality", "query about experiences"]"""
load_dotenv()
class Dialectic:
def __init__(self, agent_input: str, user_representation: str, chat_history: str):
self.agent_input = agent_input
self.user_representation = user_representation
self.chat_history = chat_history
self.client = ModelClient(
provider=DEF_DIALECTIC_PROVIDER, model=DEF_DIALECTIC_MODEL
)
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"."""
@prompt_template()
def dialectic_prompt(query: str, user_representation: str, chat_history: str) -> str:
return c(
f"""
You are operating as a context service that helps maintain psychological understanding of users across applications. Alongside a query, you'll receive:
@ai_track("Dialectic Call")
@observe()
async def call(self):
with sentry_sdk.start_transaction(
op="dialectic-inference", name="Dialectic API Response"
):
logger.debug(
f"Starting call() method with query length: {len(self.agent_input)}"
)
call_start = asyncio.get_event_loop().time()
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.
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"
)
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".
# Create a properly formatted message
message: dict[str, Any] = {"role": "user", "content": prompt}
# Generate the response
logger.debug("Calling model for generation")
model_start = asyncio.get_event_loop().time()
response = await self.client.generate(
messages=[message], system=self.system_prompt, 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"
)
total_time = asyncio.get_event_loop().time() - call_start
logger.debug(f"call() completed in {total_time:.2f}s")
return [{"text": response}]
@ai_track("Dialectic Call")
@observe()
async def stream(self):
with sentry_sdk.start_transaction(
op="dialectic-inference", name="Dialectic API Response"
):
logger.debug(
f"Starting stream() method with query length: {len(self.agent_input)}"
)
stream_start = asyncio.get_event_loop().time()
prompt = f"""
<query>{self.agent_input}</query>
<context>{self.user_representation}</context>
<conversation_history>{self.chat_history}</conversation_history>
"""
logger.debug(
f"Prompt constructed with context length: {len(self.user_representation)} chars"
)
# Create a properly formatted message
message: dict[str, Any] = {"role": "user", "content": prompt}
# Stream the response
logger.debug("Calling model for streaming")
model_start = asyncio.get_event_loop().time()
stream = await self.client.stream(
messages=[message], system=self.system_prompt, 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
<query>{query}</query>
<context>{user_representation}</context>
<conversation_history>{chat_history}</conversation_history>
"""
)
@ai_track("Dialectic Call")
@with_langfuse()
@llm.call(provider="anthropic", model="claude-3-7-sonnet-20250219")
async def dialectic_call(query: str, user_representation: str, chat_history: str):
return dialectic_prompt(query, user_representation, chat_history)
@ai_track("Dialectic Stream")
@with_langfuse()
@llm.call(provider="anthropic", model="claude-3-7-sonnet-20250219", stream=True)
async def dialectic_stream(query: str, user_representation: str, chat_history: str):
return dialectic_prompt(query, user_representation, chat_history)
@observe()
async def chat(
app_id: str,
user_id: str,
session_id: str,
queries: str | list[str],
stream: bool = False,
) -> schemas.DialecticResponse | MessageStreamManager:
) -> llm.Stream | llm.CallResponse:
"""
Chat with the Dialectic API using on-demand user representation generation.
Chat with the Dialectic API usingx on-demand user representation generation.
This function:
1. Sets up resources needed (embedding store, latest message ID)
@ -223,40 +157,21 @@ async def chat(
f"User representation generated: {len(user_representation)} characters"
)
# Create a Dialectic chain with the fresh user representation
chain = Dialectic(
agent_input=final_query,
user_representation=user_representation,
chat_history=chat_history,
)
generation_time = asyncio.get_event_loop().time() - start_time
logger.debug(f"User representation generation completed in {generation_time:.2f}s")
langfuse_context.update_current_trace(
session_id=session_id,
user_id=user_id,
release=os.getenv("SENTRY_RELEASE"),
metadata={"environment": os.getenv("SENTRY_ENVIRONMENT")},
)
# Use streaming or non-streaming response based on the request
logger.debug(f"Calling Dialectic with streaming={stream}")
query_start_time = asyncio.get_event_loop().time()
if stream:
response_stream = await chain.stream()
logger.debug(
f"Dialectic stream started after {asyncio.get_event_loop().time() - query_start_time:.2f}s"
logger.debug("Calling Dialectic with streaming")
response = await dialectic_stream(
final_query, user_representation, chat_history
)
return response_stream
response = await chain.call()
query_time = asyncio.get_event_loop().time() - query_start_time
total_time = asyncio.get_event_loop().time() - start_time
logger.debug(
f"Dialectic response received in {query_time:.2f}s (total: {total_time:.2f}s)"
)
return schemas.DialecticResponse(content=response[0]["text"])
else:
logger.debug("Calling Dialectic with non-streaming")
query_start_time = asyncio.get_event_loop().time()
response = await dialectic_call(final_query, user_representation, chat_history)
query_time = asyncio.get_event_loop().time() - query_start_time
logger.debug(f"Dialectic response received in {query_time:.2f}s")
return response
async def get_long_term_facts(
@ -344,67 +259,23 @@ async def run_tom_inference(chat_history: str, session_id: str) -> str:
return prediction
async def generate_semantic_queries(query: str) -> list[str]:
"""
Generate multiple semantically relevant queries based on the original query using LLM.
This helps retrieve more diverse and relevant facts from the vector store.
@with_langfuse()
@llm.call(provider="groq", model="llama-3.1-8b-instant", response_model=list[str])
async def generate_semantic_queries(query: str):
return c(
f"""
Given this query about a user, generate 3 focused search queries that would help retrieve relevant facts about the user.
Each query should focus on a specific aspect related to the original query, rephrased to maximize semantic search effectiveness.
For example, if the original query asks "what does the user like to eat?", generated queries might include "user's food preferences", "user's favorite cuisine", etc.
Format your response as a list of strings, with each string being a search query.
Example:
["some query about interests", "some query about personality", "some query about experiences"]
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
<query>{query}</query>
"""
)
# Prepare the messages for Anthropic
messages: list[dict[str, Any]] = [{"role": "user", "content": query}]
# Generate the response
try:
result = await client.generate(
messages=messages,
system=QUERY_GENERATION_SYSTEM,
max_tokens=1000,
use_caching=True, # Likely not caching because the system prompt is under 1000 tokens
)
llm_time = asyncio.get_event_loop().time() - llm_start
logger.debug(f"LLM response received in {llm_time:.2f}s: {result[:100]}...")
# Parse the JSON response to get a list of queries
try:
queries = json.loads(result)
if not isinstance(queries, list):
# Fallback if response is not a valid list
logger.debug("LLM response not a list, using as single query")
queries = [result]
except json.JSONDecodeError:
# Fallback if response is not valid JSON
logger.debug("Failed to parse JSON response, using raw response as query")
queries = [query] # Fall back to the original query
# Ensure we always include the original query
if query not in queries:
logger.debug("Adding original query to results")
queries.append(query)
total_time = asyncio.get_event_loop().time() - query_start
logger.debug(f"Generated {len(queries)} queries in {total_time:.2f}s")
return queries
except Exception as e:
logger.error(f"Error during API call: {str(e)}")
raise
async def generate_user_representation(
app_id: str,
@ -466,10 +337,6 @@ async def generate_user_representation(
logger.debug(f"get_user_representation completed in {gen_time:.2f}s")
# Extract the representation from the response
representation = parse_xml_content(
user_representation_response, "representation"
)
logger.debug(f"Extracted representation: {len(representation)} characters")
else:
representation = f"""
PREDICTION ABOUT THE USER'S CURRENT MENTAL STATE:

5
src/deriver/consumer.py
View File

@ -108,9 +108,10 @@ async def process_user_message(
# Extract facts from chat history
logger.debug("Extracting facts from chat history")
extract_start = os.times()[4]
facts = await extract_facts_long_term(chat_history_str)
fact_extraction = await extract_facts_long_term(chat_history_str)
facts = fact_extraction.facts
extract_time = os.times()[4] - extract_start
console.print(f"Extracted Facts: {facts}", style="bright_blue")
console.print(f"Extracted Facts: {fact_extraction.facts}", style="bright_blue")
logger.debug(f"Extracted {len(facts)} facts in {extract_time:.2f}s")
# Save the facts to the collection

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

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

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

@ -1,52 +1,27 @@
import logging
from typing import Any, Optional
from typing import Optional
from inspect import cleandoc as c
from enum import Enum
from pydantic import BaseModel
import sentry_sdk
from langfuse.decorators import langfuse_context, observe
from sentry_sdk.ai.monitoring import ai_track
from mirascope import llm
from mirascope.integrations.langfuse import with_langfuse
from src.utils.model_client import ModelClient, ModelProvider
logger = logging.getLogger(__name__)
DEF_PROVIDER = ModelProvider.GROQ
DEF_MODEL = "llama-3.3-70b-versatile"
TOM_SYSTEM_PROMPT = """You are a system for analyzing conversations to make evidence-based inferences about user mental states.
# Enums for strongly typed fields
class InfoType(str, Enum):
STYLE = "STYLE"
STATEMENT = "STATEMENT"
REQUIREMENTS:
1. Only make inferences that are directly supported by conversation evidence
2. For each inference, cite the specific message that supports it
3. Use uncertainty qualifiers (may, might, possibly) for speculative inferences
4. Do not make assumptions about demographics unless explicitly stated
5. Focus on current mental state and immediate context
6. Consider your own knowledge gaps and violations of expectations (what would surprise you)
7. Always wrap your prediction in <prediction> tags.
OUTPUT FORMAT:
<prediction>
CURRENT STATE:
- Immediate Context: User's current situation
- Active Goals: What user is trying to achieve
- Present Mood: Observable emotional state
class CertaintyLevel(str, Enum):
LIKELY = "LIKELY"
POTENTIAL = "POTENTIAL"
SPECULATIVE = "SPECULATIVE"
SUPPORTED OBSERVATIONS:
- List only behaviors/preferences with direct evidence
- Format: "OBSERVATION: [detail] (SOURCE: [exact message])"
TENTATIVE INFERENCES:
- List possible but uncertain interpretations
- Format: "POSSIBLE: [interpretation] (BASIS: [supporting message])"
KNOWLEDGE GAPS:
- List important unknown information
- Format: "UNKNOWN: [topic/question]"
EXPECTATION VIOLATIONS:
- Based on the above information, if the next message were to surprise you, what could it contain?
- Format: "POTENTIAL SURPRISE: [possible content] [reason] [confidence level]"
- Include 3-5 possible surprises
</prediction>"""
USER_REPRESENTATION_SYSTEM_PROMPT = """You are a system for maintaining factual user representations based on conversation history and theory of mind analysis.
@ -101,98 +76,224 @@ UPDATES:
</representation>"""
@ai_track("Tom Inference")
@observe()
class CurrentState(BaseModel):
immediate_context: str
active_goals: str
present_mood: str
class SupportedObservation(BaseModel):
detail: str
source: str
class TentativeInference(BaseModel):
interpretation: str
basis: str
class KnowledgeGap(BaseModel):
topic: str
class ExpectationViolation(BaseModel):
possible_surprise: str
reason: str
confidence_level: float
class TomInferenceOutput(BaseModel):
current_state: CurrentState
tentative_inferences: list[TentativeInference]
knowledge_gaps: list[KnowledgeGap]
expectation_violations: list[ExpectationViolation]
# User Representation Output Models
class SourcedInfo(BaseModel):
detail: str
source: str
class UserCurrentState(BaseModel):
active_context: SourcedInfo
temporary_conditions: SourcedInfo
present_mood_activity: SourcedInfo
class PersistentInfo(BaseModel):
detail: str
source: str
info_type: InfoType
class TentativePattern(BaseModel):
pattern: str
source: str
certainty_level: CertaintyLevel
class UserKnowledgeGap(BaseModel):
missing_info: str
class UserExpectationViolation(BaseModel):
potential_surprise: str
reason: str
confidence_level: float
class UpdateSection(BaseModel):
new_information: list[SourcedInfo]
changes: list[SourcedInfo]
removals: list[SourcedInfo]
class UserRepresentationOutput(BaseModel):
current_state: UserCurrentState
persistent_information: list[PersistentInfo]
tentative_patterns: list[TentativePattern]
knowledge_gaps: list[UserKnowledgeGap]
expectation_violations: list[UserExpectationViolation]
updates: UpdateSection
@with_langfuse()
@llm.call(
provider="groq", model="llama-3.3-70b-versatile", response_model=TomInferenceOutput
)
async def tom_inference(
chat_history: str,
user_representation: Optional[str] = None,
):
return c(
f"""
You are a system for analyzing conversations to make evidence-based inferences about user mental states.
REQUIREMENTS:
1. Only make inferences that are directly supported by conversation evidence
2. For each inference, cite the specific message that supports it
3. Use uncertainty qualifiers (may, might, possibly) for speculative inferences
4. Do not make assumptions about demographics unless explicitly stated
5. Focus on current mental state and immediate context
6. Consider your own knowledge gaps and violations of expectations (what would surprise you)
7. Always wrap your prediction in <prediction> tags.
OUTPUT FORMAT:
current_state:
- immediate_context: User's current situation
- active_goals: What user is trying to achieve
- present_mood: Observable emotional state
tentative_inferences: list of objects with:
- interpretation: Possible but uncertain interpretation
- basis: Supporting message or evidence
knowledge_gaps: list of objects with:
- topic: Important unknown information or question
expectation_violations: list of objects with:
- possible_surprise: What content could surprise you in the next message
- reason: Why this would be surprising based on current information
- confidence_level: Float between 0.0 and 1.0 indicating confidence
- Include 3-5 possible surprises
<conversation>
{chat_history or "Not provided"}
</conversation>
<user_representation>
{user_representation or "Not provided"}
</user_representation>
"""
)
async def get_tom_inference_single_prompt(
chat_history: str,
session_id: str,
user_representation: Optional[str] = None,
**kwargs,
) -> str:
with sentry_sdk.start_transaction(op="tom-inference", name="ToM Inference"):
# Create a new model client
client = ModelClient(provider=DEF_PROVIDER, model=DEF_MODEL)
# Prepare the messages
messages: list[dict[str, Any]] = [
{
"role": "user",
"content": f"Please analyze this conversation and provide a prediction following the format above:\n{chat_history}",
}
]
# Add existing user representation if available
if user_representation:
messages.append(
{
"role": "user",
"content": f"Consider this existing user representation for context, but focus on current state:\n{user_representation}",
}
)
langfuse_context.update_current_observation(input=messages, model=DEF_MODEL)
# Generate the response with caching enabled
try:
response = await client.generate(
messages=messages,
system=TOM_SYSTEM_PROMPT,
max_tokens=1000,
temperature=0,
use_caching=True, # Enable caching for the system prompt
)
except Exception as e:
sentry_sdk.capture_exception(e)
logger.error(f"Error generating Tom inference: {e}")
raise e
return response
):
inference = await tom_inference(chat_history, user_representation)
return inference.model_dump_json()
@ai_track("User Representation")
@observe()
async def get_user_representation_single_prompt(
@with_langfuse()
@llm.call(
provider="groq",
model="llama-3.3-70b-versatile",
response_model=UserRepresentationOutput,
)
async def user_representation_inference(
chat_history: str,
session_id: str,
user_representation: Optional[str] = None,
tom_inference: Optional[str] = None,
**kwargs,
) -> str:
with sentry_sdk.start_transaction(
op="user-representation-inference", name="User Representation"
):
# Create a new model client
client = ModelClient(provider=DEF_PROVIDER, model=DEF_MODEL)
):
return c(
f"""
You are a system for maintaining factual user representations based on conversation history and theory of mind analysis.
# Build the context message
context_str = f"CONVERSATION:\n{chat_history}\n\n"
if tom_inference:
context_str += f"PREDICTION OF USER MENTAL STATE - MIGHT BE INCORRECT:\n{tom_inference}\n\n"
if user_representation:
context_str += f"EXISTING USER REPRESENTATION - INCOMPLETE, TO BE UPDATED:\n{user_representation}"
Your job is to update the existing user representation (if provided) with the new information from the conversation history and theory of mind analysis.
# Prepare the messages
messages: list[dict[str, Any]] = [
{
"role": "user",
"content": f"Please analyze this information and provide an updated user representation:\n{context_str}",
}
]
Copy over information as-is from the existing user representation. Add new information as needed. Only remove content from this section if new information contradicts it. This is especially important for Persistent Information and Tentative Patterns.
langfuse_context.update_current_observation(input=messages, model=DEF_MODEL)
REQUIREMENTS:
1. Distinguish between temporary states and persistent patterns
2. Only incorporate verified information into core profile
3. Track certainty levels for all information
4. Maintain areas of uncertainty explicitly
5. Update representation incrementally
# Generate the response with caching enabled
try:
response = await client.generate(
messages=messages,
system=USER_REPRESENTATION_SYSTEM_PROMPT,
max_tokens=1000,
temperature=0,
use_caching=True, # Enable caching for the system prompt
)
except Exception as e:
sentry_sdk.capture_exception(e)
logger.error(f"Error generating user representation: {e}")
raise e
OUTPUT FORMAT:
current_state:
- active_context: object with "detail" (current situation/activity/location) and "source" (exact message)
- temporary_conditions: object with "detail" (immediate circumstances) and "source" (exact message)
- present_mood_activity: object with "detail" (what user is doing right now) and "source" (exact message)
return response
persistent_information: list of objects with:
- detail: The specific information or pattern
- source: Exact message that supports this
- info_type: Must be exactly "STYLE" for communication patterns or "STATEMENT" for explicit facts
tentative_patterns: list of objects with:
- pattern: The observed pattern
- source: Supporting evidence from specific message
- certainty_level: Must be exactly "LIKELY" (almost certain), "POTENTIAL" (possible), or "SPECULATIVE" (uncertain)
knowledge_gaps: list of objects with:
- missing_info: Key information that is missing or needs clarification
expectation_violations: list of objects with:
- potential_surprise: What could surprise you in the next message
- reason: Why this would be surprising based on current information
- confidence_level: Float between 0.0 and 1.0
- Include 3-5 possible surprises
updates:
- new_information: List of objects with "detail" (recent observation) and "source" (supporting message)
- changes: List of objects with "detail" (modified interpretation) and "source" (supporting message)
- removals: List of objects with "detail" (information no longer supported) and "source" (contradicting message)
<conversation>
{chat_history or "Not provided"}
</conversation>
<existing_user_representation>
{user_representation or "Not provided"}
</existing_user_representation>
<tom_analysis>
{tom_inference or "Not provided"}
</tom_analysis>
"""
)
async def get_user_representation_single_prompt(
chat_history: str,
user_representation: Optional[str] = None,
tom_inference: Optional[str] = None,
):
representation = await user_representation_inference(
chat_history, user_representation, tom_inference
)
return representation.model_dump_json()

21
src/routers/sessions.py
View File

@ -1,7 +1,7 @@
import logging
from typing import Optional
from anthropic import AsyncMessageStreamManager
from mirascope.llm import Stream
from fastapi import APIRouter, Body, Depends, Path, Query
from fastapi.exceptions import HTTPException
from fastapi.responses import StreamingResponse
@ -94,11 +94,11 @@ async def get_sessions(
is_active_param = False # Default to None, meaning no filter on is_active
if options:
if hasattr(options, 'filter') and options.filter:
if hasattr(options, "filter") and options.filter:
filter_param = options.filter
if filter_param == {}: # Explicitly check for empty dict
if filter_param == {}: # Explicitly check for empty dict
filter_param = None
if hasattr(options, 'is_active'): # Check if is_active is present
if hasattr(options, "is_active"): # Check if is_active is present
is_active_param = options.is_active
return await paginate(
@ -217,15 +217,17 @@ async def chat(
..., description="Dialectic Endpoint Parameters"
),
):
"""Chat with the Dialectic API"""
if not options.stream:
return await agent.chat(
response = await agent.chat(
app_id=app_id,
user_id=user_id,
session_id=session_id,
queries=options.queries,
)
return schemas.DialecticResponse(
content=response.content,
)
else:
async def parse_stream():
@ -237,10 +239,9 @@ async def chat(
queries=options.queries,
stream=True,
)
if type(stream) is AsyncMessageStreamManager:
async with stream as stream_manager:
async for text in stream_manager.text_stream:
yield text
if type(stream) is Stream:
async for chunk, _ in stream:
yield chunk.content
except Exception as e:
logger.error(f"Error in stream: {str(e)}")
raise HTTPException(status_code=500, detail=str(e)) from e

182
src/utils/history.py
View File

@ -4,8 +4,11 @@ from typing import Optional, Union, cast
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from inspect import cleandoc as c
from mirascope import llm
from mirascope.integrations.langfuse import with_langfuse
from src.utils.model_client import ModelClient, ModelProvider
from .. import models
@ -37,11 +40,6 @@ class SummaryType(Enum):
LONG = "honcho_chat_summary_long"
# Default model settings for summary generation
DEFAULT_PROVIDER = ModelProvider.GEMINI
DEFAULT_MODEL = "gemini-2.0-flash-lite"
async def get_session_summaries(
db: AsyncSession,
session_id: str,
@ -123,99 +121,93 @@ async def get_messages_since_message(
return list(result.scalars().all())
@with_langfuse()
@llm.call(
provider="google",
model="gemini-2.0-flash-lite",
call_params={"max_tokens": 1000},
)
async def create_short_summary(
messages: list[models.Message],
previous_summary: Optional[str] = None,
):
return c(
f"""
You are a system that summarizes parts of a conversation to create a concise and accurate summary.
Focus on capturing:
1. Key facts and information shared
2. User preferences, opinions, and questions
3. Important context and requests
4. Core topics discussed
5. User's apparent emotional state
It is very important that you clearly distinguish between the user's messages and the assistant's messages, and that only the user's literal words are attributed to them.
Provide a concise, factual summary that captures the essence of the conversation.
Your summary should be detailed enough to serve as context for future messages,
but brief enough to be helpful.
Return only the summary without any explanation or meta-commentary.
<conversation>
{messages}
</conversation>
<previous_summary>
{previous_summary}
</previous_summary>
"""
)
@with_langfuse()
@llm.call(
provider="google",
model="gemini-2.0-flash-lite",
call_params={"max_tokens": 2000},
)
async def create_long_summary(
messages: list[models.Message],
previous_summary: Optional[str] = None,
):
return c(
f"""
You are a system that creates comprehensive summaries of conversations.
Focus on capturing:
1. Key facts and information shared
2. User preferences, opinions, and questions
3. Important context and requests
4. Core topics discussed in detail
5. User's apparent emotional state and personality traits
6. Important themes and patterns across the conversation
It is very important that you clearly distinguish between the user's messages and the assistant's messages, and that only the user's literal words are attributed to them.
Provide a thorough and detailed summary that captures the essence of the conversation.
Your summary should serve as a comprehensive record of the important information in this conversation.
Return only the summary without any explanation or meta-commentary.
<conversation>
{messages}
</conversation>
<previous_summary>
{previous_summary}
</previous_summary>
"""
)
async def create_summary(
messages: list[models.Message],
previous_summary: Optional[str] = None,
summary_type: SummaryType = SummaryType.SHORT,
) -> str:
"""
Generate a summary of the provided messages using an LLM.
Args:
messages: List of messages to summarize
previous_summary: Optional previous summary to provide context
summary_type: Type of summary to create ("short" or "long")
Returns:
A summary of the conversation
"""
# Combine messages into a conversation format
conversation = "\n".join(
[
f"{'human' if msg.is_user else 'assistant'}: {msg.content}"
for msg in messages
]
)
# Adjust system prompt based on summary type
if summary_type == SummaryType.LONG:
system_prompt = """You are a system that creates comprehensive summaries of conversations.
Focus on capturing:
1. Key facts and information shared
2. User preferences, opinions, and questions
3. Important context and requests
4. Core topics discussed in detail
5. User's apparent emotional state and personality traits
6. Important themes and patterns across the conversation
It is very important that you clearly distinguish between the user's messages and the assistant's messages, and that only the user's literal words are attributed to them.
Provide a thorough and detailed summary that captures the essence of the conversation.
Your summary should serve as a comprehensive record of the important information in this conversation.
Return only the summary without any explanation or meta-commentary."""
else: # short summary
system_prompt = """You are a system that summarizes parts of a conversation to create a concise and accurate summary.
Focus on capturing:
1. Key facts and information shared
2. User preferences, opinions, and questions
3. Important context and requests
4. Core topics discussed
5. User's apparent emotional state
It is very important that you clearly distinguish between the user's messages and the assistant's messages, and that only the user's literal words are attributed to them.
Provide a concise, factual summary that captures the essence of the conversation.
Your summary should be detailed enough to serve as context for future messages,
but brief enough to be helpful.
Return only the summary without any explanation or meta-commentary."""
# Include previous summary if available
if previous_summary:
user_prompt = f"""Here is a previous summary of the conversation:
{previous_summary}
Now please summarize these additional messages, incorporating the context from the previous summary.
Your summary should summarize the entire conversation in a self-contained way, such that someone could read it and understand the entire conversation.
{conversation}
Provide a {"comprehensive" if summary_type == SummaryType.LONG else "concise"} summary that captures both the previous context and the new information."""
else:
user_prompt = f"""Please summarize the following conversation segment:
{conversation}
Provide a {"comprehensive" if summary_type == SummaryType.LONG else "concise"} summary that captures the key points and context."""
# Create a model client
client = ModelClient(provider=DEFAULT_PROVIDER, model=DEFAULT_MODEL)
# Generate the summary
llm_messages = [{"role": "user", "content": user_prompt}]
try:
summary = await client.generate(
messages=llm_messages,
system=system_prompt,
max_tokens=1000
if summary_type == SummaryType.SHORT
else 2000, # Allow longer responses for long summaries
temperature=0.0,
use_caching=True,
)
return summary
except Exception as e:
logger.error(f"Error generating summary: {str(e)}")
# Fallback to a basic summary in case of error
return f"Conversation with {len(messages)} messages about {messages[-1].content[:30]}..."
):
if summary_type == SummaryType.SHORT:
return await create_short_summary(messages, previous_summary)
elif summary_type == SummaryType.LONG:
return await create_long_summary(messages, previous_summary)
async def save_summary_metamessage(

541
src/utils/model_client.py
View File

@ -1,541 +0,0 @@
"""
Utility functions for interacting with various language model APIs.
"""
import os
from enum import Enum
from typing import Any, Optional, Protocol
import sentry_sdk
from anthropic import AsyncAnthropic
from dotenv import load_dotenv
from google import genai
from google.genai import types as genai_types
from langfuse.decorators import langfuse_context, observe
# from openai import AsyncOpenAI
from langfuse.openai import AsyncOpenAI
# Load environment variables
load_dotenv()
# Supported model providers
class ModelProvider(str, Enum):
ANTHROPIC = "anthropic"
OPENAI = "openai"
OPENROUTER = "openrouter"
CEREBRAS = "cerebras"
GROQ = "groq"
GEMINI = "gemini"
# Add other providers as needed
# Default models for each provider
DEFAULT_MODELS = {
ModelProvider.ANTHROPIC: "claude-3-7-sonnet-20250219",
ModelProvider.OPENAI: "gpt-4o",
ModelProvider.OPENROUTER: "meta-llama/Llama-3.3-70B-Instruct",
ModelProvider.CEREBRAS: "llama-3.3-70b",
ModelProvider.GROQ: "llama-3.3-70b-versatile",
ModelProvider.GEMINI: "gemini-2.0-flash-lite",
}
OPENAI_COMPATIBLE_PROVIDERS = [
ModelProvider.OPENAI,
ModelProvider.OPENROUTER,
ModelProvider.CEREBRAS,
ModelProvider.GROQ,
]
DEFAULT_TEMPERATURE = 0.0
DEFAULT_MAX_TOKENS = 1000
class Message(Protocol):
"""Protocol for a message that works with any provider."""
role: str
content: str
class ModelClient:
"""A client for interacting with various language model APIs."""
def __init__(
self,
provider: ModelProvider = ModelProvider.ANTHROPIC,
model: Optional[str] = None,
api_key: Optional[str] = None,
base_url: Optional[str] = None,
):
"""
Initialize the model client.
Args:
provider: The model provider to use
model: The specific model to use, or None to use the default model for the provider
api_key: The API key to use, or None to read from environment variables
base_url: Custom base URL for the API endpoints (used for OpenRouter)
"""
self.provider = provider
self.model = model or DEFAULT_MODELS[provider]
self.base_url = base_url
self.openai_client = None
self.gemini_client = None
# Setup provider-specific clients
if provider == ModelProvider.ANTHROPIC:
self.api_key = api_key or os.getenv("ANTHROPIC_API_KEY")
if not self.api_key:
raise ValueError("Anthropic API key is required")
self.client = AsyncAnthropic(api_key=self.api_key)
elif provider in OPENAI_COMPATIBLE_PROVIDERS:
self.api_key = api_key or os.getenv("OPENAI_COMPATIBLE_API_KEY")
self.base_url = base_url or os.getenv("OPENAI_COMPATIBLE_BASE_URL")
if not self.api_key:
raise ValueError("OpenAI-compatible API key is required")
self.openai_client = AsyncOpenAI(
api_key=self.api_key, base_url=self.base_url
)
elif provider == ModelProvider.GEMINI:
self.api_key = api_key or os.getenv("GEMINI_API_KEY")
if not self.api_key:
raise ValueError("Gemini API key is required")
self.gemini_client = genai.Client(api_key=self.api_key)
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()
async def generate(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False,
) -> str:
"""
Generate a response using the configured model.
Args:
messages: The conversation history
system: Optional system prompt
max_tokens: Maximum number of tokens to generate
temperature: Temperature for generation
extra_headers: Optional headers to add to the request
use_caching: Whether to use provider-side caching for the response
Returns:
The generated text
"""
with sentry_sdk.start_transaction(
op="llm-api", name=f"{self.provider} API Call"
):
# Log to langfuse
langfuse_context.update_current_observation(
input=messages, model=self.model
)
if self.provider == ModelProvider.ANTHROPIC:
return await self._generate_anthropic(
messages,
system,
max_tokens,
temperature,
extra_headers,
use_caching,
)
elif self.provider in OPENAI_COMPATIBLE_PROVIDERS:
return await self._generate_openai(
messages, system, max_tokens, temperature
)
elif self.provider == ModelProvider.GEMINI:
return await self._generate_gemini(
messages, system, max_tokens, temperature
)
else:
raise ValueError(f"Unsupported provider: {self.provider}")
@observe(as_type="generation")
async def _generate_anthropic(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False,
) -> str:
"""Generate a response using the Anthropic API."""
if not self.client:
raise ValueError("Anthropic client not initialized.")
params = {
"model": self.model,
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
}
# Handle system prompt with caching if enabled
if system:
if use_caching:
params["system"] = [
{
"type": "text",
"text": system,
"cache_control": {"type": "ephemeral"},
}
]
else:
params["system"] = system
langfuse_context.update_current_observation(input=messages, model=self.model)
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 (
content_block
and hasattr(content_block, "type")
and content_block.type == "text"
):
return content_block.text
return str(content_block)
return ""
@observe(as_type="generation")
async def _generate_openai(
self,
messages: list[dict[str, str]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
) -> str:
"""Generate text using OpenAI or OpenRouter API."""
if not self.openai_client:
raise ValueError("OpenAI client not initialized")
# Prepare messages
formatted_messages = []
# Add system message if provided
if system:
formatted_messages.append({"role": "system", "content": system})
# Add the rest of the messages
formatted_messages.extend(messages)
# Make the API call with the OpenAI client
response = await self.openai_client.chat.completions.create(
model=self.model,
messages=formatted_messages,
max_tokens=max_tokens,
temperature=temperature,
)
# Extract the generated text
choice = response.choices[0]
if choice and choice.message and choice.message.content:
return choice.message.content
return ""
@observe()
async def stream(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False,
) -> Any:
"""
Stream a response using the configured model.
Args:
messages: The conversation history
system: Optional system prompt
max_tokens: Maximum number of tokens to generate
temperature: Temperature for generation
extra_headers: Optional headers to add to the request
use_caching: Whether to use provider-side caching for the response
Returns:
A streaming response from the provider
"""
with sentry_sdk.start_transaction(
op="llm-api-stream", name=f"{self.provider} API Stream"
):
# Log to langfuse
langfuse_context.update_current_observation(
input=messages, model=self.model
)
if self.provider == ModelProvider.ANTHROPIC:
return await self._stream_anthropic(
messages,
system,
max_tokens,
temperature,
extra_headers,
use_caching,
)
elif self.provider in OPENAI_COMPATIBLE_PROVIDERS:
return await self._stream_openai(
messages, system, max_tokens, temperature
)
elif self.provider == ModelProvider.GEMINI:
return await self._stream_gemini(
messages, system, max_tokens, temperature
)
else:
raise ValueError(f"Unsupported provider: {self.provider}")
@observe(as_type="generation")
async def _stream_anthropic(
self,
messages: list[dict[str, Any]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
extra_headers: Optional[dict[str, str]] = None,
use_caching: bool = False,
) -> Any:
"""Stream text using Anthropic API."""
if not self.client:
raise ValueError("Anthropic client not initialized.")
params = {
"model": self.model,
"messages": messages,
"max_tokens": max_tokens,
"temperature": temperature,
}
# Handle system prompt with caching if enabled
if system:
if use_caching:
params["system"] = [
{
"type": "text",
"text": system,
"cache_control": {"type": "ephemeral"},
}
]
else:
params["system"] = system
langfuse_context.update_current_observation(input=messages, model=self.model)
# Return the stream directly without awaiting it
return self.client.messages.stream(**params)
@observe(as_type="generation")
async def _stream_openai(
self,
messages: list[dict[str, str]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
) -> Any:
"""Stream text using OpenAI or OpenRouter API."""
if not self.openai_client:
raise ValueError("OpenAI client not initialized")
# Prepare messages
formatted_messages = []
# Add system message if provided
if system:
formatted_messages.append({"role": "system", "content": system})
# Add the rest of the messages
formatted_messages.extend(messages)
# Make the API call with the OpenAI client
stream = await self.openai_client.chat.completions.create(
model=self.model,
messages=formatted_messages,
max_tokens=max_tokens,
temperature=temperature,
stream=True,
)
return stream
@observe(as_type="generation")
async def _generate_gemini(
self,
messages: list[dict[str, str]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
) -> str:
"""Generate text using Gemini API."""
if not self.gemini_client:
raise ValueError("Gemini client not initialized")
# Format messages for Gemini
gemini_messages = []
# Convert messages to Gemini format
for message in messages:
role = message["role"]
# Map roles to what Gemini expects
if role == "user":
gemini_role = "user"
elif role == "assistant":
gemini_role = "model"
else:
# Skip system messages as they're handled through config
continue
gemini_messages.append(
genai_types.Content(
role=gemini_role,
parts=[genai_types.Part.from_text(text=message["content"])],
)
)
# Set generation config
generate_content_config = genai_types.GenerateContentConfig(
temperature=temperature,
max_output_tokens=max_tokens,
response_mime_type="text/plain",
)
# Add system instruction if provided
if system:
generate_content_config.system_instruction = system
# Make the API call
if not gemini_messages:
# If we have no messages but have a system prompt, create a default user message
if system:
default_content = genai_types.Content(
role="user",
parts=[
genai_types.Part.from_text(
text="Please respond based on the system instructions."
)
],
)
# model = self.gemini_client.get_model(self.model)
response = await self.gemini_client.aio.models.generate_content(
model=self.model,
contents=default_content,
config=generate_content_config,
)
else:
raise ValueError("No messages provided for Gemini generation")
else:
# Normal case with messages
# model = get_model(self.model)
response = await self.gemini_client.aio.models.generate_content(
model=self.model,
contents=gemini_messages
if len(gemini_messages) > 1
else gemini_messages[0],
config=generate_content_config,
)
# Extract text from response
if response and response.text:
return response.text
return ""
@observe(as_type="generation")
async def _stream_gemini(
self,
messages: list[dict[str, str]],
system: Optional[str] = None,
max_tokens: int = DEFAULT_MAX_TOKENS,
temperature: float = DEFAULT_TEMPERATURE,
) -> Any:
"""Stream text using Gemini API."""
if not self.gemini_client:
raise ValueError("Gemini client not initialized")
# Format messages for Gemini
gemini_messages = []
# Convert messages to Gemini format
for message in messages:
role = message["role"]
# Map roles to what Gemini expects
if role == "user":
gemini_role = "user"
elif role == "assistant":
gemini_role = "model"
else:
# Skip system messages as they're handled through config
continue
gemini_messages.append(
genai_types.Content(
role=gemini_role,
parts=[genai_types.Part.from_text(text=message["content"])],
)
)
# Set generation config
generate_content_config = genai_types.GenerateContentConfig(
temperature=temperature,
max_output_tokens=max_tokens,
response_mime_type="text/plain",
)
# Add system instruction if provided
if system:
generate_content_config.system_instruction = system
# Make the streaming API call
if not gemini_messages:
# If we have no messages but have a system prompt, create a default user message
if system:
default_content = genai_types.Content(
role="user",
parts=[
genai_types.Part.from_text(
text="Please respond based on the system instructions."
)
],
)
stream = await self.gemini_client.aio.models.generate_content_stream(
model=self.model,
contents=default_content,
config=generate_content_config,
)
else:
raise ValueError("No messages provided for Gemini streaming")
else:
# Normal case with messages
stream = await self.gemini_client.aio.models.generate_content_stream(
model=self.model,
contents=gemini_messages
if len(gemini_messages) > 1
else gemini_messages[0],
config=generate_content_config,
)
return stream

191
tests/utils/test_model_client.py
View File

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

1562
uv.lock
View File

File diff suppressed because it is too large Load Diff