From 1005f6a60c638600fe5748c54ed8003a6fa08117 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Tue, 8 Apr 2025 23:39:38 -0400 Subject: [PATCH] fix: Langfuse tracing --- src/agent.py | 6 ++++-- src/deriver/consumer.py | 6 ++++-- src/deriver/tom/conversational.py | 4 ++-- src/deriver/tom/long_term.py | 4 ++-- src/deriver/tom/single_prompt.py | 4 ++-- src/utils/model_client.py | 7 +++++-- 6 files changed, 19 insertions(+), 12 deletions(-) diff --git a/src/agent.py b/src/agent.py index 591604b4..697f4f61 100644 --- a/src/agent.py +++ b/src/agent.py @@ -69,7 +69,7 @@ class Dialectic: self.system_prompt = """You are operating as a context service that helps maintain psychological understanding of users across applications. Alongside a query, you'll receive: 1) previously collected psychological context about the user that I've maintained, 2) a series of long-term facts about the user, and 3) their current conversation/interaction from the requesting application. Your goal is to analyze this information and provide theory-of-mind insights that help applications personalize their responses. Please respond in a brief, matter-of-fact, and appropriate manner to convey as much relevant information to the application based on its query and the user's most recent message. You are encouraged to provide any context from the provided resources that helps provide a more complete or nuanced understanding of the user, as long as it is somewhat relevant to the query. If the context provided doesn't help address the query, write absolutely NOTHING but "None".""" @ai_track("Dialectic Call") - @observe(as_type="generation") + @observe() async def call(self): with sentry_sdk.start_transaction( op="dialectic-inference", name="Dialectic API Response" @@ -107,7 +107,7 @@ class Dialectic: return [{"text": response}] @ai_track("Dialectic Call") - @observe(as_type="generation") + @observe() async def stream(self): with sentry_sdk.start_transaction( op="dialectic-inference", name="Dialectic API Response" @@ -553,6 +553,8 @@ RELEVANT LONG-TERM FACTS ABOUT THE USER: logger.error(f"Message with ID {message_id} does not exist") else: metamessage = models.Metamessage( + user_id=user_id, + session_id=session_id, message_id=message_id, metamessage_type="user_representation", content=representation, diff --git a/src/deriver/consumer.py b/src/deriver/consumer.py index 5e6d787e..a2ad7686 100644 --- a/src/deriver/consumer.py +++ b/src/deriver/consumer.py @@ -81,7 +81,7 @@ async def process_item(db: AsyncSession, payload: dict): @sentry_sdk.trace -@observe() +# @observe() async def process_ai_message( content: str, app_id: str, @@ -129,7 +129,9 @@ async def process_user_message( # 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") + collection = await crud.get_or_create_user_protected_collection( + db=db, app_id=app_id, user_id=user_id + ) embedding_store = CollectionEmbeddingStore( db=db, app_id=app_id, diff --git a/src/deriver/tom/conversational.py b/src/deriver/tom/conversational.py index bc604adf..98b91bfe 100644 --- a/src/deriver/tom/conversational.py +++ b/src/deriver/tom/conversational.py @@ -15,7 +15,7 @@ anthropic = Anthropic( @ai_track("Tom Inference") -@observe(as_type="generation") +@observe() async def get_tom_inference_conversational( chat_history: str, session_id: str, user_representation: str = "None" ) -> str: @@ -81,7 +81,7 @@ async def get_tom_inference_conversational( @ai_track("User Representation") -@observe(as_type="generation") +@observe() async def get_user_representation_conversational( chat_history: str, session_id: str, diff --git a/src/deriver/tom/long_term.py b/src/deriver/tom/long_term.py index 7cf5f004..54328205 100644 --- a/src/deriver/tom/long_term.py +++ b/src/deriver/tom/long_term.py @@ -25,7 +25,7 @@ MAX_FACT_DISTANCE = 0.85 @ai_track("User Representation") -@observe(as_type="generation") +@observe() async def get_user_representation_long_term( chat_history: str, session_id: str, @@ -118,7 +118,7 @@ UPDATES: @ai_track("Fact Extraction") -@observe(as_type="generation") +@observe() async def extract_facts_long_term(chat_history: str) -> list[str]: logger.debug("Starting fact extraction from chat history") extract_start = time.time() diff --git a/src/deriver/tom/single_prompt.py b/src/deriver/tom/single_prompt.py index b72f9fbf..3696eb5e 100644 --- a/src/deriver/tom/single_prompt.py +++ b/src/deriver/tom/single_prompt.py @@ -102,7 +102,7 @@ UPDATES: @ai_track("Tom Inference") -@observe(as_type="generation") +@observe() async def get_tom_inference_single_prompt( chat_history: str, session_id: str, @@ -150,7 +150,7 @@ async def get_tom_inference_single_prompt( @ai_track("User Representation") -@observe(as_type="generation") +@observe() async def get_user_representation_single_prompt( chat_history: str, session_id: str, diff --git a/src/utils/model_client.py b/src/utils/model_client.py index 8979c334..b599cddf 100644 --- a/src/utils/model_client.py +++ b/src/utils/model_client.py @@ -10,7 +10,9 @@ import sentry_sdk from anthropic import AsyncAnthropic from dotenv import load_dotenv from langfuse.decorators import langfuse_context, observe -from openai import AsyncOpenAI + +# from openai import AsyncOpenAI +from langfuse.openai import AsyncOpenAI # Load environment variables load_dotenv() @@ -108,7 +110,6 @@ class ModelClient: # 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]], @@ -156,6 +157,7 @@ class ModelClient: else: raise ValueError(f"Unsupported provider: {self.provider}") + @observe(as_type="generation") async def _generate_anthropic( self, messages: list[dict[str, Any]], @@ -200,6 +202,7 @@ class ModelClient: return str(content_block) return "" + @observe(as_type="generation") async def _generate_openai( self, messages: list[dict[str, str]],