fix: Langfuse tracing

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,

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,

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,

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()

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,

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]],