jaf-py 2.2.4__py3-none-any.whl → 2.3.1__py3-none-any.whl

jaf/core/engine.py

@@ -94,7 +94,9 @@ async def run(
             config.on_event(RunStartEvent(data=to_event_data(RunStartEventData(
                 run_id=initial_state.run_id,
                 trace_id=initial_state.trace_id,
-                session_id=config.conversation_id
+                session_id=config.conversation_id,
+                context=initial_state.context,
+                messages=initial_state.messages
             ))))
 
         state_with_memory = await _load_conversation_history(initial_state, config)
@@ -293,7 +295,9 @@ async def _run_internal(
             agent_name=current_agent.name,
             model=model,
             trace_id=state.trace_id,
-            run_id=state.run_id
+            run_id=state.run_id,
+            context=state.context,
+            messages=state.messages
         ))))
 
     # Get completion from model provider

jaf/core/tracing.py

@@ -358,15 +358,98 @@ class LangfuseTraceCollector:
             if event.type == "run_start":
                 # Start a new trace for the entire run
                 print(f"[LANGFUSE] Starting trace for run: {trace_id}")
+                
+                # Extract user query from the run_start data
+                user_query = None
+                user_id = None
+                
+                # Debug: Print the event data structure to understand what we're working with
+                if event.data.get("context"):
+                    context = event.data["context"]
+                    print(f"[LANGFUSE DEBUG] Context type: {type(context)}")
+                    print(f"[LANGFUSE DEBUG] Context attributes: {dir(context) if hasattr(context, '__dict__') else 'Not an object'}")
+                    if hasattr(context, '__dict__'):
+                        print(f"[LANGFUSE DEBUG] Context dict: {context.__dict__}")
+                
+                # Try to extract from context first
+                context = event.data.get("context")
+                if context:
+                    # Try direct attribute access
+                    if hasattr(context, 'query'):
+                        user_query = context.query
+                        print(f"[LANGFUSE DEBUG] Found user_query from context.query: {user_query}")
+                    
+                    # Try to extract from combined_history
+                    if hasattr(context, 'combined_history') and context.combined_history:
+                        history = context.combined_history
+                        print(f"[LANGFUSE DEBUG] Found combined_history with {len(history)} messages")
+                        for i, msg in enumerate(reversed(history)):
+                            print(f"[LANGFUSE DEBUG] History message {i}: {msg}")
+                            if isinstance(msg, dict) and msg.get("role") == "user":
+                                user_query = msg.get("content", "")
+                                print(f"[LANGFUSE DEBUG] Found user_query from history: {user_query}")
+                                break
+                    
+                    # Try to extract user_id from token_response
+                    if hasattr(context, 'token_response'):
+                        token_response = context.token_response
+                        print(f"[LANGFUSE DEBUG] Found token_response: {type(token_response)}")
+                        if isinstance(token_response, dict):
+                            user_id = token_response.get("email") or token_response.get("username")
+                            print(f"[LANGFUSE DEBUG] Extracted user_id: {user_id}")
+                        elif hasattr(token_response, 'email'):
+                            user_id = token_response.email
+                            print(f"[LANGFUSE DEBUG] Extracted user_id from attr: {user_id}")
+                
+                # Fallback: try to extract from messages if context didn't work
+                if not user_query and event.data.get("messages"):
+                    print(f"[LANGFUSE DEBUG] Trying fallback from messages")
+                    messages = event.data["messages"]
+                    print(f"[LANGFUSE DEBUG] Found {len(messages)} messages")
+                    # Find the last user message which should be the current query
+                    for i, msg in enumerate(reversed(messages)):
+                        print(f"[LANGFUSE DEBUG] Message {i}: {msg}")
+                        if isinstance(msg, dict) and msg.get("role") == "user":
+                            user_query = msg.get("content", "")
+                            print(f"[LANGFUSE DEBUG] Found user_query from messages: {user_query}")
+                            break
+                        elif hasattr(msg, 'role') and msg.role == 'user':
+                            user_query = msg.content
+                            print(f"[LANGFUSE DEBUG] Found user_query from message attr: {user_query}")
+                            break
+                
+                print(f"[LANGFUSE DEBUG] Final extracted - user_query: {user_query}, user_id: {user_id}")
+                
+                # Create comprehensive input data for the trace
+                trace_input = {
+                    "user_query": user_query,
+                    "run_id": str(trace_id),
+                    "agent_name": event.data.get("agent_name", "analytics_agent_jaf"),
+                    "session_info": {
+                        "session_id": event.data.get("session_id"),
+                        "user_id": user_id or event.data.get("user_id")
+                    }
+                }
+                
                 trace = self.langfuse.trace(
                     name=f"jaf-run-{trace_id}",
-                    user_id=event.data.get("user_id"),
+                    user_id=user_id or event.data.get("user_id"),
                     session_id=event.data.get("session_id"),
-                    input=event.data,
-                    metadata={"framework": "jaf", "event_type": "run_start", "trace_id": str(trace_id)}
+                    input=trace_input,
+                    metadata={
+                        "framework": "jaf", 
+                        "event_type": "run_start", 
+                        "trace_id": str(trace_id),
+                        "user_query": user_query,
+                        "user_id": user_id or event.data.get("user_id"),
+                        "agent_name": event.data.get("agent_name", "analytics_agent_jaf")
+                    }
                 )
                 self.trace_spans[trace_id] = trace
-                print(f"[LANGFUSE] Created trace: {trace}")
+                # Store user_id and user_query for later use in generations
+                trace._user_id = user_id or event.data.get("user_id")
+                trace._user_query = user_query
+                print(f"[LANGFUSE] Created trace with user query: {user_query[:100] if user_query else 'None'}...")
                 
             elif event.type == "run_end":
                 if trace_id in self.trace_spans:
@@ -386,10 +469,21 @@ class LangfuseTraceCollector:
                 # Start a generation for LLM calls
                 model = event.data.get("model", "unknown")
                 print(f"[LANGFUSE] Starting generation for LLM call with model: {model}")
-                generation = self.trace_spans[trace_id].generation(
+                
+                # Get stored user information from the trace
+                trace = self.trace_spans[trace_id]
+                user_id = getattr(trace, '_user_id', None)
+                user_query = getattr(trace, '_user_query', None)
+                
+                generation = trace.generation(
                     name=f"llm-call-{model}",
                     input=event.data.get("messages"),
-                    metadata={"agent_name": event.data.get("agent_name"), "model": model}
+                    metadata={
+                        "agent_name": event.data.get("agent_name"), 
+                        "model": model,
+                        "user_id": user_id,
+                        "user_query": user_query
+                    }
                 )
                 span_id = self._get_span_id(event)
                 self.active_spans[span_id] = generation
@@ -405,47 +499,111 @@ class LangfuseTraceCollector:
                     
                     # Extract usage from the event data
                     usage = event.data.get("usage", {})
-
-                    # Convert to Langfuse v2 format
+                    
+                    # Extract model information from choice data or event data
+                    model = choice.get("model", "unknown")
+                    if model == "unknown":
+                        # Try to get model from the choice response structure
+                        if isinstance(choice, dict):
+                            model = choice.get("model") or choice.get("id", "unknown")
+                    
+                    # Convert to Langfuse v2 format - let Langfuse handle cost calculation automatically
                     langfuse_usage = None
                     if usage:
+                        prompt_tokens = usage.get("prompt_tokens", 0)
+                        completion_tokens = usage.get("completion_tokens", 0)
+                        total_tokens = usage.get("total_tokens", 0)
+                        
                         langfuse_usage = {
-                            "input": usage.get("prompt_tokens", 0),
-                            "output": usage.get("completion_tokens", 0),
-                            "total": usage.get("total_tokens", 0),
+                            "input": prompt_tokens,
+                            "output": completion_tokens,
+                            "total": total_tokens,
                             "unit": "TOKENS"
                         }
-
-                    generation.end(output=choice, usage=langfuse_usage)
+                        
+                        print(f"[LANGFUSE] Usage data for automatic cost calculation: {langfuse_usage}")
+
+                    # Include model information in the generation end - Langfuse will calculate costs automatically
+                    generation.end(
+                        output=choice, 
+                        usage=langfuse_usage,
+                        model=model,  # Pass model directly for automatic cost calculation
+                        metadata={
+                            "model": model,
+                            "system_fingerprint": choice.get("system_fingerprint"),
+                            "created": choice.get("created"),
+                            "response_id": choice.get("id")
+                        }
+                    )
 
                     # Clean up the span reference
                     del self.active_spans[span_id]
-                    print(f"[LANGFUSE] Generation ended")
+                    print(f"[LANGFUSE] Generation ended with cost tracking")
                 else:
                     print(f"[LANGFUSE] No generation found for llm_call_end: {span_id}")
                     
             elif event.type == "tool_call_start":
-                # Start a span for tool calls
-                print(f"[LANGFUSE] Starting span for tool call: {event.data.get('tool_name')}")
+                # Start a span for tool calls with detailed input information
+                tool_name = event.data.get('tool_name', 'unknown')
+                tool_args = event.data.get("args", {})
+                
+                print(f"[LANGFUSE] Starting span for tool call: {tool_name}")
+                
+                # Create comprehensive input data for the tool call
+                tool_input = {
+                    "tool_name": tool_name,
+                    "arguments": tool_args,
+                    "call_id": event.data.get("call_id"),
+                    "timestamp": datetime.now().isoformat()
+                }
+                
                 span = self.trace_spans[trace_id].span(
-                    name=f"tool-{event.data.get('tool_name', 'unknown')}",
-                    input=event.data.get("args"),
-                    metadata={"tool_name": event.data.get("tool_name")}
+                    name=f"tool-{tool_name}",
+                    input=tool_input,
+                    metadata={
+                        "tool_name": tool_name,
+                        "call_id": event.data.get("call_id"),
+                        "framework": "jaf",
+                        "event_type": "tool_call"
+                    }
                 )
                 span_id = self._get_span_id(event)
                 self.active_spans[span_id] = span
-                print(f"[LANGFUSE] Created tool span: {span}")
+                print(f"[LANGFUSE] Created tool span for {tool_name} with args: {str(tool_args)[:100]}...")
                     
             elif event.type == "tool_call_end":
                 span_id = self._get_span_id(event)
                 if span_id in self.active_spans:
-                    print(f"[LANGFUSE] Ending span for tool call")
-                    # End the span
+                    tool_name = event.data.get('tool_name', 'unknown')
+                    tool_result = event.data.get("result")
+                    
+                    print(f"[LANGFUSE] Ending span for tool call: {tool_name}")
+                    
+                    # Create comprehensive output data for the tool call
+                    tool_output = {
+                        "tool_name": tool_name,
+                        "result": tool_result,
+                        "call_id": event.data.get("call_id"),
+                        "timestamp": datetime.now().isoformat(),
+                        "status": "completed"
+                    }
+                    
+                    # End the span with detailed output
                     span = self.active_spans[span_id]
-                    span.end(output=event.data.get("result"))
+                    span.end(
+                        output=tool_output,
+                        metadata={
+                            "tool_name": tool_name,
+                            "call_id": event.data.get("call_id"),
+                            "result_length": len(str(tool_result)) if tool_result else 0,
+                            "framework": "jaf",
+                            "event_type": "tool_call_end"
+                        }
+                    )
+                    
                     # Clean up the span reference
                     del self.active_spans[span_id]
-                    print(f"[LANGFUSE] Tool span ended")
+                    print(f"[LANGFUSE] Tool span ended for {tool_name} with result length: {len(str(tool_result)) if tool_result else 0}")
                 else:
                     print(f"[LANGFUSE] No tool span found for tool_call_end: {span_id}")
                     
@@ -490,7 +648,6 @@ class LangfuseTraceCollector:
                 return TraceId(event.data['runId'])
         
         # Debug: print what's actually in the event data
-        print(f"[LANGFUSE] Event data keys: {list(event.data.keys()) if hasattr(event, 'data') and event.data else 'No data'}")
         return None
 
     def _get_span_id(self, event: TraceEvent) -> str:
@@ -502,19 +659,10 @@ class LangfuseTraceCollector:
             tool_name = event.data.get('tool_name') or event.data.get('toolName', 'unknown')
             return f"tool-{tool_name}-{trace_id}"
         elif event.type.startswith('llm_call'):
-            # For LLM calls, get the model name from the event data
-            model = event.data.get('model')
-            
-            # For llm_call_end events, the model might be in the choice object
-            if not model and event.type == 'llm_call_end':
-                choice = event.data.get('choice', {})
-                if isinstance(choice, dict):
-                    model = choice.get('model')
-            
-            # Handle case where model might be empty or None
-            if not model or model == '':
-                model = 'unknown'
-            return f"llm-{model}-{trace_id}"
+            # For LLM calls, use a simpler consistent ID that matches between start and end
+            # Get run_id for more consistent matching
+            run_id = event.data.get('run_id') or event.data.get('runId', trace_id)
+            return f"llm-{run_id}"
         else:
             return f"{event.type}-{trace_id}"
 

jaf/core/types.py

@@ -379,6 +379,8 @@ class RunStartEventData:
     run_id: RunId
     trace_id: TraceId
     session_id: Optional[str] = None
+    context: Optional[Any] = None
+    messages: Optional[List[Message]] = None
 
 @dataclass(frozen=True)
 class RunStartEvent:
@@ -392,6 +394,8 @@ class LLMCallStartEventData:
     model: str
     trace_id: TraceId
     run_id: RunId
+    context: Optional[Any] = None
+    messages: Optional[List[Message]] = None
 
 @dataclass(frozen=True)
 class LLMCallStartEvent:

{jaf_py-2.2.4.dist-info → jaf_py-2.3.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaf-py
-Version: 2.2.4
+Version: 2.3.1
 Summary: A purely functional agent framework with immutable state and composable tools - Python implementation
 Author: JAF Contributors
 Maintainer: JAF Contributors

{jaf_py-2.2.4.dist-info → jaf_py-2.3.1.dist-info}/RECORD

@@ -42,7 +42,7 @@ jaf/core/__init__.py,sha256=rBvP_7TGbJICDJnA7a3qyX8yQErCDWaGAn5WzpyH4gU,1339
 jaf/core/agent_tool.py,sha256=8TcBuSxGmDTW5F_GhBU_m5S43nYqkjO4qTrNERraAig,11656
 jaf/core/analytics.py,sha256=NrUfOLLTDIhOzdfc65ZqS9AJ4ZAP9BtNtga69q0YdYw,23265
 jaf/core/composition.py,sha256=IVxRO1Q9nK7JRH32qQ4p8WMIUu66BhqPNrlTNMGFVwE,26317
-jaf/core/engine.py,sha256=ydlF2m9GWyHXWAP6dng8sOwxEWcaaH8etkImorfY0aY,26954
+jaf/core/engine.py,sha256=v6yXp-zJF9Cj_9qlaW1rmCz7_N3YuDU02xxlzl-howM,27122
 jaf/core/errors.py,sha256=5fwTNhkojKRQ4wZj3lZlgDnAsrYyjYOwXJkIr5EGNUc,5539
 jaf/core/performance.py,sha256=jedQmTEkrKMD6_Aw1h8PdG-5TsdYSFFT7Or6k5dmN2g,9974
 jaf/core/proxy.py,sha256=_WM3cpRlSQLYpgSBrnY30UPMe2iZtlqDQ65kppE-WY0,4609
@@ -50,8 +50,8 @@ jaf/core/proxy_helpers.py,sha256=i7a5fAX9rLmO4FMBX51-yRkTFwfWedzQNgnLmeLUd_A,437
 jaf/core/streaming.py,sha256=c5o9iqpjoYV2LrUpG6qLWCYrWcP-DCcZsvMbyqKunp8,16089
 jaf/core/tool_results.py,sha256=-bTOqOX02lMyslp5Z4Dmuhx0cLd5o7kgR88qK2HO_sw,11323
 jaf/core/tools.py,sha256=SbJRRr4y_xxNYNTulZg6OiyNaHBlo_qXWYY510jxQEs,16489
-jaf/core/tracing.py,sha256=r96LCi4dV7BSoqVuFMXI7j72G9jOwpbEjnhYqubuQzc,23455
-jaf/core/types.py,sha256=fF7_UMFTPUzGQQxZ9tGoWdjF5LSI02RHoJ9vvSe3W8I,16905
+jaf/core/tracing.py,sha256=3ByTbpYMWHJku-NEasK5ncyMOdv7vHwmG6ybJP19pxE,31659
+jaf/core/types.py,sha256=9UXrPkepw7opgv1VGbPAC1Zx80RP4-ouRxtb6kVTA-A,17063
 jaf/core/workflows.py,sha256=Ul-82gzjIXtkhnSMSPv-8igikjkMtW1EBo9yrfodtvI,26294
 jaf/memory/__init__.py,sha256=-L98xlvihurGAzF0DnXtkueDVvO_wV2XxxEwAWdAj50,1400
 jaf/memory/factory.py,sha256=Fh6JyvQtCKe38DZV5-NnC9vPRCvzBgSSPFIGaX7Nt5E,2958
@@ -79,9 +79,9 @@ jaf/visualization/functional_core.py,sha256=zedMDZbvjuOugWwnh6SJ2stvRNQX1Hlkb9Ab
 jaf/visualization/graphviz.py,sha256=WTOM6UP72-lVKwI4_SAr5-GCC3ouckxHv88ypCDQWJ0,12056
 jaf/visualization/imperative_shell.py,sha256=GpMrAlMnLo2IQgyB2nardCz09vMvAzaYI46MyrvJ0i4,2593
 jaf/visualization/types.py,sha256=QQcbVeQJLuAOXk8ynd08DXIS-PVCnv3R-XVE9iAcglw,1389
-jaf_py-2.2.4.dist-info/licenses/LICENSE,sha256=LXUQBJxdyr-7C4bk9cQBwvsF_xwA-UVstDTKabpcjlI,1063
-jaf_py-2.2.4.dist-info/METADATA,sha256=LRWLY4HHpn57BSrHesTA1i5gXSP95zA7Wt1vjIx5P7w,27613
-jaf_py-2.2.4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-jaf_py-2.2.4.dist-info/entry_points.txt,sha256=OtIJeNJpb24kgGrqRx9szGgDx1vL9ayq8uHErmu7U5w,41
-jaf_py-2.2.4.dist-info/top_level.txt,sha256=Xu1RZbGaM4_yQX7bpalo881hg7N_dybaOW282F15ruE,4
-jaf_py-2.2.4.dist-info/RECORD,,
+jaf_py-2.3.1.dist-info/licenses/LICENSE,sha256=LXUQBJxdyr-7C4bk9cQBwvsF_xwA-UVstDTKabpcjlI,1063
+jaf_py-2.3.1.dist-info/METADATA,sha256=0o3VxXPHt1FO6lHDttZyFGNDk-WPH5nxaX9zXoAvcLg,27613
+jaf_py-2.3.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+jaf_py-2.3.1.dist-info/entry_points.txt,sha256=OtIJeNJpb24kgGrqRx9szGgDx1vL9ayq8uHErmu7U5w,41
+jaf_py-2.3.1.dist-info/top_level.txt,sha256=Xu1RZbGaM4_yQX7bpalo881hg7N_dybaOW282F15ruE,4
+jaf_py-2.3.1.dist-info/RECORD,,