xgae 0.1.5__py3-none-any.whl → 0.1.6__py3-none-any.whl

xgae/engine/responser/{xga_stream_responser.py → stream_responser.py}

@@ -5,9 +5,10 @@ import uuid
 
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from typing import List, Dict, Any, Optional, AsyncGenerator, override
+from typing import List, Dict, Any, Optional, AsyncGenerator, override, Literal
 
-from xgae.engine.responser.xga_responser_base import TaskResponseProcessor, ToolExecutionStrategy, XmlAddingStrategy, TaskResponseContext,TaskRunContinuousState
+from xgae.utils import langfuse
+from xgae.engine.responser.responser_base import TaskResponseProcessor, TaskResponserContext,TaskRunContinuousState,XmlAddingStrategy,ToolExecutionStrategy
 from xgae.utils.json_helpers import (
     ensure_dict, safe_json_parse,
     to_json_string, format_for_yield
@@ -55,7 +56,7 @@ class ProcessorConfig:
 
 
 class StreamTaskResponser(TaskResponseProcessor):
-    def __init__(self, response_context: TaskResponseContext):
+    def __init__(self, response_context: TaskResponserContext):
         super().__init__(response_context)
 
     @override
@@ -129,14 +130,14 @@ class StreamTaskResponser(TaskResponseProcessor):
             # --- Save and Yield Start Events (only if not auto-continuing) ---
             if auto_continue_count == 0:
                 start_content = {"status_type": "thread_run_start", "thread_run_id": thread_run_id}
-                start_msg_obj = await self.add_message(
+                start_msg_obj = await self.add_response_message(
                    type="status", content=start_content,
                     is_llm_message=False, metadata={"thread_run_id": thread_run_id}
                 )
                 if start_msg_obj: yield format_for_yield(start_msg_obj)
 
                 assist_start_content = {"status_type": "assistant_response_start"}
-                assist_start_msg_obj = await self.add_message(
+                assist_start_msg_obj = await self.add_response_message(
                     type="status", content=assist_start_content,
                     is_llm_message=False, metadata={"thread_run_id": thread_run_id}
                 )
@@ -204,7 +205,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                             __sequence += 1
                         else:
                             logging.info("XML tool call limit reached - not yielding more content chunks")
-                            self.trace.event(name="xml_tool_call_limit_reached", level="DEFAULT", status_message=(
+                            langfuse.create_event(trace_context=self.trace_context, name="xml_tool_call_limit_reached", level="DEFAULT", status_message=(
                                 f"XML tool call limit reached - not yielding more content chunks"))
 
                         # --- Process XML Tool Calls (if enabled and limit not reached) ---
@@ -226,8 +227,7 @@ class StreamTaskResponser(TaskResponseProcessor):
 
                                     if config.execute_tools and config.execute_on_stream:
                                         # Save and Yield tool_started status
-                                        started_msg_obj = await self._yield_and_save_tool_started(context, thread_id,
-                                                                                                  thread_run_id)
+                                        started_msg_obj = await self._add_tool_start_message(context)
                                         if started_msg_obj: yield format_for_yield(started_msg_obj)
                                         yielded_tool_indices.add(tool_index)  # Mark status as yielded
 
@@ -304,8 +304,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                                 )
 
                                 # Save and Yield tool_started status
-                                started_msg_obj = await self._yield_and_save_tool_started(context, thread_id,
-                                                                                          thread_run_id)
+                                started_msg_obj = await self._add_tool_start_message(context)
                                 if started_msg_obj: yield format_for_yield(started_msg_obj)
                                 yielded_tool_indices.add(tool_index)  # Mark status as yielded
 
@@ -318,7 +317,7 @@ class StreamTaskResponser(TaskResponseProcessor):
 
                 if finish_reason == "xml_tool_limit_reached":
                     logging.info("Stopping stream processing after loop due to XML tool call limit")
-                    self.trace.event(name="stopping_stream_processing_after_loop_due_to_xml_tool_call_limit",
+                    langfuse.create_event(trace_context=self.trace_context, name="stopping_stream_processing_after_loop_due_to_xml_tool_call_limit",
                                      level="DEFAULT", status_message=(
                             f"Stopping stream processing after loop due to XML tool call limit"))
                     break
@@ -353,18 +352,18 @@ class StreamTaskResponser(TaskResponseProcessor):
                     #     f"🔥 Estimated tokens – prompt: {prompt_tokens}, "
                     #     f"completion: {completion_tokens}, total: {prompt_tokens + completion_tokens}"
                     # )
-                    self.trace.event(name="usage_calculated_with_litellm_token_counter", level="DEFAULT",
+                    langfuse.create_event(trace_context=self.trace_context, name="usage_calculated_with_litellm_token_counter", level="DEFAULT",
                                      status_message=(f"Usage calculated with litellm.token_counter"))
                 except Exception as e:
                     logging.warning(f"Failed to calculate usage: {str(e)}")
-                    self.trace.event(name="failed_to_calculate_usage", level="WARNING",
+                    langfuse.create_event(trace_context=self.trace_context, name="failed_to_calculate_usage", level="WARNING",
                                      status_message=(f"Failed to calculate usage: {str(e)}"))
 
             # Wait for pending tool executions from streaming phase
             tool_results_buffer = []  # Stores (tool_call, result, tool_index, context)
             if pending_tool_executions:
                 logging.info(f"Waiting for {len(pending_tool_executions)} pending streamed tool executions")
-                self.trace.event(name="waiting_for_pending_streamed_tool_executions", level="DEFAULT", status_message=(
+                langfuse.create_event(trace_context=self.trace_context, name="waiting_for_pending_streamed_tool_executions", level="DEFAULT", status_message=(
                     f"Waiting for {len(pending_tool_executions)} pending streamed tool executions"))
                 # ... (asyncio.wait logic) ...
                 pending_tasks = [execution["task"] for execution in pending_tool_executions]
@@ -388,24 +387,24 @@ class StreamTaskResponser(TaskResponseProcessor):
                                 if tool_name in ['ask', 'complete']:
                                     logging.info(
                                         f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag.")
-                                    self.trace.event(name="terminating_tool_completed_during_streaming",
+                                    langfuse.create_event(trace_context=self.trace_context, name="terminating_tool_completed_during_streaming",
                                                      level="DEFAULT", status_message=(
                                             f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag."))
                                     agent_should_terminate = True
 
                             else:  # Should not happen with asyncio.wait
                                 logging.warning(f"Task for tool index {tool_idx} not done after wait.")
-                                self.trace.event(name="task_for_tool_index_not_done_after_wait", level="WARNING",
+                                langfuse.create_event(trace_context=self.trace_context, name="task_for_tool_index_not_done_after_wait", level="WARNING",
                                                  status_message=(
                                                      f"Task for tool index {tool_idx} not done after wait."))
                         except Exception as e:
                             logging.error(f"Error getting result for pending tool execution {tool_idx}: {str(e)}")
-                            self.trace.event(name="error_getting_result_for_pending_tool_execution", level="ERROR",
+                            langfuse.create_event(trace_context=self.trace_context, name="error_getting_result_for_pending_tool_execution", level="ERROR",
                                              status_message=(
                                                  f"Error getting result for pending tool execution {tool_idx}: {str(e)}"))
                             context.error = e
                             # Save and Yield tool error status message (even if started was yielded)
-                            error_msg_obj = await self._yield_and_save_tool_error(context, thread_id, thread_run_id)
+                            error_msg_obj = await self._add_tool_error_message(context)
                             if error_msg_obj: yield format_for_yield(error_msg_obj)
                         continue  # Skip further status yielding for this tool index
 
@@ -420,40 +419,39 @@ class StreamTaskResponser(TaskResponseProcessor):
                             if tool_name in ['ask', 'complete']:
                                 logging.info(
                                     f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag.")
-                                self.trace.event(name="terminating_tool_completed_during_streaming", level="DEFAULT",
+                                langfuse.create_event(trace_context=self.trace_context, name="terminating_tool_completed_during_streaming", level="DEFAULT",
                                                  status_message=(
                                                      f"Terminating tool '{tool_name}' completed during streaming. Setting termination flag."))
                                 agent_should_terminate = True
 
                             # Save and Yield tool completed/failed status
-                            completed_msg_obj = await self._yield_and_save_tool_completed(
-                                context, None, thread_id, thread_run_id
-                            )
+                            completed_msg_obj = await self._add_tool_completed_message(
+                                context, None)
                             if completed_msg_obj: yield format_for_yield(completed_msg_obj)
                             yielded_tool_indices.add(tool_idx)
                     except Exception as e:
                         logging.error(
                             f"Error getting result/yielding status for pending tool execution {tool_idx}: {str(e)}")
-                        self.trace.event(name="error_getting_result_yielding_status_for_pending_tool_execution",
+                        langfuse.create_event(trace_context=self.trace_context, name="error_getting_result_yielding_status_for_pending_tool_execution",
                                          level="ERROR", status_message=(
                                 f"Error getting result/yielding status for pending tool execution {tool_idx}: {str(e)}"))
                         context.error = e
                         # Save and Yield tool error status
-                        error_msg_obj = await self._yield_and_save_tool_error(context, thread_id, thread_run_id)
+                        error_msg_obj = await self._add_tool_error_message(context)
                         if error_msg_obj: yield format_for_yield(error_msg_obj)
                         yielded_tool_indices.add(tool_idx)
 
             # Save and yield finish status if limit was reached
             if finish_reason == "xml_tool_limit_reached":
                 finish_content = {"status_type": "finish", "finish_reason": "xml_tool_limit_reached"}
-                finish_msg_obj = await self.add_message(
+                finish_msg_obj = await self.add_response_message(
                      type="status", content=finish_content,
                     is_llm_message=False, metadata={"thread_run_id": thread_run_id}
                 )
                 if finish_msg_obj: yield format_for_yield(finish_msg_obj)
                 logging.info(
                     f"Stream finished with reason: xml_tool_limit_reached after {xml_tool_call_count} XML tool calls")
-                self.trace.event(name="stream_finished_with_reason_xml_tool_limit_reached_after_xml_tool_calls",
+                langfuse.create_event(trace_context=self.trace_context, name="stream_finished_with_reason_xml_tool_limit_reached_after_xml_tool_calls",
                                  level="DEFAULT", status_message=(
                         f"Stream finished with reason: xml_tool_limit_reached after {xml_tool_call_count} XML tool calls"))
 
@@ -489,7 +487,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                     "tool_calls": complete_native_tool_calls or None
                 }
 
-                last_assistant_message_object = await self._add_message_with_agent_info(type="assistant", content=message_data,
+                last_assistant_message_object = await self.add_response_message(type="assistant", content=message_data,
                     is_llm_message=True, metadata={"thread_run_id": thread_run_id}
                 )
 
@@ -503,12 +501,12 @@ class StreamTaskResponser(TaskResponseProcessor):
                     yield format_for_yield(yield_message)
                 else:
                     logging.error(f"Failed to save final assistant message for thread {thread_id}")
-                    self.trace.event(name="failed_to_save_final_assistant_message_for_thread", level="ERROR",
+                    langfuse.create_event(trace_context=self.trace_context, name="failed_to_save_final_assistant_message_for_thread", level="ERROR",
                                      status_message=(f"Failed to save final assistant message for thread {thread_id}"))
                     # Save and yield an error status
                     err_content = {"role": "system", "status_type": "error",
                                    "message": "Failed to save final assistant message"}
-                    err_msg_obj = await self.add_message(
+                    err_msg_obj = await self.add_response_message(
                         type="status", content=err_content,
                         is_llm_message=False, metadata={"thread_run_id": thread_run_id}
                     )
@@ -571,7 +569,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                 # Populate from buffer if executed on stream
                 if config.execute_on_stream and tool_results_buffer:
                     logging.info(f"Processing {len(tool_results_buffer)} buffered tool results")
-                    self.trace.event(name="processing_buffered_tool_results", level="DEFAULT",
+                    langfuse.create_event(trace_context=self.trace_context, name="processing_buffered_tool_results", level="DEFAULT",
                                      status_message=(f"Processing {len(tool_results_buffer)} buffered tool results"))
                     for tool_call, result, tool_idx, context in tool_results_buffer:
                         if last_assistant_message_object: context.assistant_message_id = last_assistant_message_object[
@@ -582,7 +580,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                 elif final_tool_calls_to_process and not config.execute_on_stream:
                     logging.info(
                         f"Executing {len(final_tool_calls_to_process)} tools ({config.tool_execution_strategy}) after stream")
-                    self.trace.event(name="executing_tools_after_stream", level="DEFAULT", status_message=(
+                    langfuse.create_event(trace_context=self.trace_context, name="executing_tools_after_stream", level="DEFAULT", status_message=(
                         f"Executing {len(final_tool_calls_to_process)} tools ({config.tool_execution_strategy}) after stream"))
                     results_list = await self._execute_tools(final_tool_calls_to_process,
                                                              config.tool_execution_strategy)
@@ -600,14 +598,14 @@ class StreamTaskResponser(TaskResponseProcessor):
                             tool_results_map[current_tool_idx] = (tc, res, context)
                         else:
                             logging.warning(f"Could not map result for tool index {current_tool_idx}")
-                            self.trace.event(name="could_not_map_result_for_tool_index", level="WARNING",
+                            langfuse.create_event(trace_context=self.trace_context, name="could_not_map_result_for_tool_index", level="WARNING",
                                              status_message=(f"Could not map result for tool index {current_tool_idx}"))
                         current_tool_idx += 1
 
                 # Save and Yield each result message
                 if tool_results_map:
                     logging.info(f"Saving and yielding {len(tool_results_map)} final tool result messages")
-                    self.trace.event(name="saving_and_yielding_final_tool_result_messages", level="DEFAULT",
+                    langfuse.create_event(trace_context=self.trace_context, name="saving_and_yielding_final_tool_result_messages", level="DEFAULT",
                                      status_message=(
                                          f"Saving and yielding {len(tool_results_map)} final tool result messages"))
                     for tool_idx in sorted(tool_results_map.keys()):
@@ -618,21 +616,19 @@ class StreamTaskResponser(TaskResponseProcessor):
 
                         # Yield start status ONLY IF executing non-streamed (already yielded if streamed)
                         if not config.execute_on_stream and tool_idx not in yielded_tool_indices:
-                            started_msg_obj = await self._yield_and_save_tool_started(context, thread_id, thread_run_id)
+                            started_msg_obj = await self._add_tool_start_message(context)
                             if started_msg_obj: yield format_for_yield(started_msg_obj)
                             yielded_tool_indices.add(tool_idx)  # Mark status yielded
 
                         # Save the tool result message to DB
-                        saved_tool_result_object = await self._add_tool_result(  # Returns full object or None
-                            thread_id, tool_call, result, config.xml_adding_strategy,
-                            context.assistant_message_id, context.parsing_details
-                        )
+                        saved_tool_result_object = await self._add_tool_messsage(tool_call, result, config.xml_adding_strategy,
+                                                                                 context.assistant_message_id, context.parsing_details
+                                                                                 )
 
                         # Yield completed/failed status (linked to saved result ID if available)
-                        completed_msg_obj = await self._yield_and_save_tool_completed(
+                        completed_msg_obj = await self._add_tool_completed_message(
                             context,
-                            saved_tool_result_object['message_id'] if saved_tool_result_object else None,
-                            thread_id, thread_run_id
+                            saved_tool_result_object['message_id'] if saved_tool_result_object else None
                         )
                         if completed_msg_obj: yield format_for_yield(completed_msg_obj)
                         # Don't add to yielded_tool_indices here, completion status is separate yield
@@ -644,7 +640,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                         else:
                             logging.error(
                                 f"Failed to save tool result for index {tool_idx}, not yielding result message.")
-                            self.trace.event(name="failed_to_save_tool_result_for_index", level="ERROR",
+                            langfuse.create_event(trace_context=self.trace_context, name="failed_to_save_tool_result_for_index", level="ERROR",
                                              status_message=(
                                                  f"Failed to save tool result for index {tool_idx}, not yielding result message."))
                             # Optionally yield error status for saving failure?
@@ -652,7 +648,7 @@ class StreamTaskResponser(TaskResponseProcessor):
             # --- Final Finish Status ---
             if finish_reason and finish_reason != "xml_tool_limit_reached":
                 finish_content = {"status_type": "finish", "finish_reason": finish_reason}
-                finish_msg_obj = await self.add_message(
+                finish_msg_obj = await self.add_response_message(
                     type="status", content=finish_content,
                     is_llm_message=False, metadata={"thread_run_id": thread_run_id}
                 )
@@ -662,7 +658,7 @@ class StreamTaskResponser(TaskResponseProcessor):
             if agent_should_terminate:
                 logging.info(
                     "Agent termination requested after executing ask/complete tool. Stopping further processing.")
-                self.trace.event(name="agent_termination_requested", level="DEFAULT",
+                langfuse.create_event(trace_context=self.trace_context, name="agent_termination_requested", level="DEFAULT",
                                  status_message="Agent termination requested after executing ask/complete tool. Stopping further processing.")
 
                 # Set finish reason to indicate termination
@@ -670,7 +666,7 @@ class StreamTaskResponser(TaskResponseProcessor):
 
                 # Save and yield termination status
                 finish_content = {"status_type": "finish", "finish_reason": "agent_terminated"}
-                finish_msg_obj = await self.add_message(
+                finish_msg_obj = await self.add_response_message(
                     type="status", content=finish_content,
                     is_llm_message=False, metadata={"thread_run_id": thread_run_id}
                 )
@@ -714,7 +710,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                         if streaming_metadata.get("response_ms"):
                             assistant_end_content["response_ms"] = streaming_metadata["response_ms"]
 
-                        await self.add_message(
+                        await self.add_response_message(
                             type="assistant_response_end",
                             content=assistant_end_content,
                             is_llm_message=False,
@@ -723,7 +719,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                         logging.info("Assistant response end saved for stream (before termination)")
                     except Exception as e:
                         logging.error(f"Error saving assistant response end for stream (before termination): {str(e)}")
-                        self.trace.event(name="error_saving_assistant_response_end_for_stream_before_termination",
+                        langfuse.create_event(trace_context=self.trace_context, name="error_saving_assistant_response_end_for_stream_before_termination",
                                          level="ERROR", status_message=(
                                 f"Error saving assistant response end for stream (before termination): {str(e)}"))
 
@@ -770,7 +766,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                         if streaming_metadata.get("response_ms"):
                             assistant_end_content["response_ms"] = streaming_metadata["response_ms"]
 
-                        await self.add_message(
+                        await self.add_response_message(
                             type="assistant_response_end",
                             content=assistant_end_content,
                             is_llm_message=False,
@@ -779,18 +775,18 @@ class StreamTaskResponser(TaskResponseProcessor):
                         logging.info("Assistant response end saved for stream")
                     except Exception as e:
                         logging.error(f"Error saving assistant response end for stream: {str(e)}")
-                        self.trace.event(name="error_saving_assistant_response_end_for_stream", level="ERROR",
+                        langfuse.create_event(trace_context=self.trace_context, name="error_saving_assistant_response_end_for_stream", level="ERROR",
                                          status_message=(f"Error saving assistant response end for stream: {str(e)}"))
 
         except Exception as e:
             logging.error(f"Error processing stream: {str(e)}", exc_info=True)
-            self.trace.event(name="error_processing_stream", level="ERROR",
+            langfuse.create_event(trace_context=self.trace_context, name="error_processing_stream", level="ERROR",
                              status_message=(f"Error processing stream: {str(e)}"))
             # Save and yield error status message
 
             err_content = {"role": "system", "status_type": "error", "message": str(e)}
             if (not "AnthropicException - Overloaded" in str(e)):
-                err_msg_obj = await self.add_message(
+                err_msg_obj = await self.add_response_message(
                     type="status", content=err_content,
                     is_llm_message=False,
                     metadata={"thread_run_id": thread_run_id if 'thread_run_id' in locals() else None}
@@ -798,12 +794,12 @@ class StreamTaskResponser(TaskResponseProcessor):
                 if err_msg_obj: yield format_for_yield(err_msg_obj)  # Yield the saved error message
                 # Re-raise the same exception (not a new one) to ensure proper error propagation
                 logging.critical(f"Re-raising error to stop further processing: {str(e)}")
-                self.trace.event(name="re_raising_error_to_stop_further_processing", level="ERROR",
+                langfuse.create_event(trace_context=self.trace_context, name="re_raising_error_to_stop_further_processing", level="ERROR",
                                  status_message=(f"Re-raising error to stop further processing: {str(e)}"))
             else:
                 logging.error(f"AnthropicException - Overloaded detected - Falling back to OpenRouter: {str(e)}",
                              exc_info=True)
-                self.trace.event(name="anthropic_exception_overloaded_detected", level="ERROR", status_message=(
+                langfuse.create_event(trace_context=self.trace_context, name="anthropic_exception_overloaded_detected", level="ERROR", status_message=(
                     f"AnthropicException - Overloaded detected - Falling back to OpenRouter: {str(e)}"))
             raise  # Use bare 'raise' to preserve the original exception with its traceback
 
@@ -818,7 +814,7 @@ class StreamTaskResponser(TaskResponseProcessor):
                 # Save and Yield the final thread_run_end status (only if not auto-continuing and finish_reason is not 'length')
                 try:
                     end_content = {"status_type": "thread_run_end"}
-                    end_msg_obj = await self.add_message(
+                    end_msg_obj = await self.add_response_message(
                         type="status", content=end_content,
                         is_llm_message=False,
                         metadata={"thread_run_id": thread_run_id if 'thread_run_id' in locals() else None}
@@ -826,5 +822,5 @@ class StreamTaskResponser(TaskResponseProcessor):
                     if end_msg_obj: yield format_for_yield(end_msg_obj)
                 except Exception as final_e:
                     logging.error(f"Error in finally block: {str(final_e)}", exc_info=True)
-                    self.trace.event(name="error_in_finally_block", level="ERROR",
+                    langfuse.create_event(trace_context=self.trace_context, name="error_in_finally_block", level="ERROR",
                                      status_message=(f"Error in finally block: {str(final_e)}"))