unique_orchestrator 0.0.4__py3-none-any.whl → 1.7.16__py3-none-any.whl

unique_orchestrator/unique_ai.py

@@ -1,4 +1,5 @@
-from datetime import datetime
+import asyncio
+from datetime import datetime, timezone
 from logging import Logger
 
 import jinja2
@@ -7,20 +8,29 @@ from unique_toolkit.agentic.debug_info_manager.debug_info_manager import (
 )
 from unique_toolkit.agentic.evaluation.evaluation_manager import EvaluationManager
 from unique_toolkit.agentic.history_manager.history_manager import HistoryManager
+from unique_toolkit.agentic.message_log_manager.service import MessageStepLogger
 from unique_toolkit.agentic.postprocessor.postprocessor_manager import (
     PostprocessorManager,
 )
 from unique_toolkit.agentic.reference_manager.reference_manager import ReferenceManager
 from unique_toolkit.agentic.thinking_manager.thinking_manager import ThinkingManager
-from unique_toolkit.agentic.tools.tool_manager import ToolManager
+from unique_toolkit.agentic.tools.tool_manager import (
+    ResponsesApiToolManager,
+    SafeTaskExecutor,
+    ToolManager,
+)
 from unique_toolkit.app.schemas import ChatEvent, McpServer
 from unique_toolkit.chat.service import ChatService
 from unique_toolkit.content.service import ContentService
+from unique_toolkit.language_model import LanguageModelAssistantMessage
 from unique_toolkit.language_model.schemas import (
-    LanguageModelAssistantMessage,
     LanguageModelMessages,
     LanguageModelStreamResponse,
 )
+from unique_toolkit.protocols.support import (
+    ResponsesSupportCompleteWithReferences,
+    SupportCompleteWithReferences,
+)
 
 from unique_orchestrator.config import UniqueAIConfig
 
@@ -45,12 +55,14 @@ class UniqueAI:
         chat_service: ChatService,
         content_service: ContentService,
         debug_info_manager: DebugInfoManager,
+        streaming_handler: SupportCompleteWithReferences,
         reference_manager: ReferenceManager,
         thinking_manager: ThinkingManager,
         tool_manager: ToolManager,
         history_manager: HistoryManager,
         evaluation_manager: EvaluationManager,
         postprocessor_manager: PostprocessorManager,
+        message_step_logger: MessageStepLogger,
         mcp_servers: list[McpServer],
     ):
         self._logger = logger
@@ -70,6 +82,11 @@ class UniqueAI:
         self._postprocessor_manager = postprocessor_manager
         self._latest_assistant_id: str = event.payload.assistant_message.id
         self._mcp_servers = mcp_servers
+        self._streaming_handler = streaming_handler
+
+        self._message_step_logger = message_step_logger
+        # Helper variable to support control loop
+        self._tool_took_control = False
 
     ############################################################
     # Override of base methods
@@ -121,12 +138,10 @@ class UniqueAI:
             self.start_text = self._thinking_manager.update_start_text(
                 self.start_text, loop_response
             )
-            await self._create_new_assistant_message_if_loop_response_contains_content(
-                loop_response
-            )
 
+        # Only set completed_at if no tool took control. Tools that take control will set the message state to completed themselves.
         await self._chat_service.modify_assistant_message_async(
-            set_completed_at=True,
+            set_completed_at=not self._tool_took_control,
         )
 
     # @track()
@@ -144,7 +159,7 @@ class UniqueAI:
             self._logger.info("Its needs forced tool calls.")
             self._logger.info(f"Forced tools: {self._tool_manager.get_forced_tools()}")
             responses = [
-                await self._chat_service.complete_with_references_async(
+                await self._streaming_handler.complete_with_references_async(
                     messages=messages,
                     model_name=self._config.space.language_model.name,
                     tools=self._tool_manager.get_tool_definitions(),
@@ -152,8 +167,8 @@ class UniqueAI:
                     start_text=self.start_text,
                     debug_info=self._debug_info_manager.get(),
                     temperature=self._config.agent.experimental.temperature,
-                    other_options=self._config.agent.experimental.additional_llm_options
-                    | {"toolChoice": opt},
+                    tool_choice=opt,
+                    other_options=self._config.agent.experimental.additional_llm_options,
                 )
                 for opt in self._tool_manager.get_forced_tools()
             ]
@@ -174,7 +189,7 @@ class UniqueAI:
                 "we are in the last iteration we need to produce an answer now"
             )
             # No tool calls in last iteration
-            stream_response = await self._chat_service.complete_with_references_async(
+            stream_response = await self._streaming_handler.complete_with_references_async(
                 messages=messages,
                 model_name=self._config.space.language_model.name,
                 content_chunks=self._reference_manager.get_chunks(),
@@ -188,7 +203,7 @@ class UniqueAI:
             self._logger.info(
                 f"we are in the iteration {self.current_iteration_index} asking the model to tell if we should use tools or if it will just stream"
             )
-            stream_response = await self._chat_service.complete_with_references_async(
+            stream_response = await self._streaming_handler.complete_with_references_async(
                 messages=messages,
                 model_name=self._config.space.language_model.name,
                 tools=self._tool_manager.get_tool_definitions(),
@@ -216,6 +231,9 @@ class UniqueAI:
             self._logger.debug(
                 "Tools were called we process them and do not exit the loop"
             )
+            await self._create_new_assistant_message_if_loop_response_contains_content(
+                loop_response
+            )
 
             return await self._handle_tool_calls(loop_response)
 
@@ -246,32 +264,52 @@ class UniqueAI:
             for prompts in self._tool_manager.get_tool_prompts()
         ]
 
-        used_tools = [m.name for m in self._tool_manager.get_tools()]
+        used_tools = [t.name for t in self._history_manager.get_tool_calls()]
+        sub_agent_calls = self._tool_manager.filter_tool_calls(
+            self._history_manager.get_tool_calls(), ["subagent"]
+        )
 
         mcp_server_user_prompts = [
             mcp_server.user_prompt for mcp_server in self._mcp_servers
         ]
 
+        user_metadata = self._get_filtered_user_metadata()
+
         tool_descriptions = self._tool_manager.get_tool_prompts()
 
         query = self._event.payload.user_message.text
 
+        if (
+            self._config.agent.experimental.sub_agents_config.referencing_config
+            is not None
+            and len(sub_agent_calls) > 0
+        ):
+            use_sub_agent_references = True
+            sub_agent_referencing_instructions = self._config.agent.experimental.sub_agents_config.referencing_config.referencing_instructions_for_user_prompt
+        else:
+            use_sub_agent_references = False
+            sub_agent_referencing_instructions = None
+
         user_msg = user_message_template.render(
             query=query,
             tool_descriptions=tool_descriptions,
             used_tools=used_tools,
             mcp_server_user_prompts=list(mcp_server_user_prompts),
             tool_descriptions_with_user_prompts=tool_descriptions_with_user_prompts,
+            use_sub_agent_references=use_sub_agent_references,
+            sub_agent_referencing_instructions=sub_agent_referencing_instructions,
+            user_metadata=user_metadata,
         )
         return user_msg
 
-    async def _render_system_prompt(
-        self,
-    ) -> str:
+    async def _render_system_prompt(self) -> str:
         # TODO: Collect tool information here and adapt to system prompt
         tool_descriptions = self._tool_manager.get_tool_prompts()
 
-        used_tools = [m.name for m in self._tool_manager.get_tools()]
+        used_tools = [t.name for t in self._history_manager.get_tool_calls()]
+        sub_agent_calls = self._tool_manager.filter_tool_calls(
+            self._history_manager.get_tool_calls(), ["subagent"]
+        )
 
         system_prompt_template = jinja2.Template(
             self._config.agent.prompt_config.system_prompt_template
@@ -279,10 +317,31 @@ class UniqueAI:
 
         date_string = datetime.now().strftime("%A %B %d, %Y")
 
+        user_metadata = self._get_filtered_user_metadata()
+
         mcp_server_system_prompts = [
             mcp_server.system_prompt for mcp_server in self._mcp_servers
         ]
 
+        if (
+            self._config.agent.experimental.sub_agents_config.referencing_config
+            is not None
+            and len(sub_agent_calls) > 0
+        ):
+            use_sub_agent_references = True
+            sub_agent_referencing_instructions = self._config.agent.experimental.sub_agents_config.referencing_config.referencing_instructions_for_system_prompt
+        else:
+            use_sub_agent_references = False
+            sub_agent_referencing_instructions = None
+
+        uploaded_documents = self._content_service.get_documents_uploaded_to_chat()
+        uploaded_documents_expired = [
+            doc
+            for doc in uploaded_documents
+            if doc.expired_at is not None
+            and doc.expired_at <= datetime.now(timezone.utc)
+        ]
+
         system_message = system_prompt_template.render(
             model_info=self._config.space.language_model.model_dump(mode="json"),
             date_string=date_string,
@@ -294,6 +353,10 @@ class UniqueAI:
             max_loop_iterations=self._config.agent.max_loop_iterations,
             current_iteration=self.current_iteration_index + 1,
             mcp_server_system_prompts=mcp_server_system_prompts,
+            use_sub_agent_references=use_sub_agent_references,
+            sub_agent_referencing_instructions=sub_agent_referencing_instructions,
+            user_metadata=user_metadata,
+            uploaded_documents_expired=uploaded_documents_expired,
         )
         return system_message
 
@@ -301,20 +364,62 @@ class UniqueAI:
         self, loop_response: LanguageModelStreamResponse
     ) -> bool:
         """Handle the case where no tool calls are returned."""
+        task_executor = SafeTaskExecutor(
+            logger=self._logger,
+        )
+
         selected_evaluation_names = self._tool_manager.get_evaluation_check_list()
-        evaluation_results = await self._evaluation_manager.run_evaluations(
-            selected_evaluation_names, loop_response, self._latest_assistant_id
+        evaluation_results = task_executor.execute_async(
+            self._evaluation_manager.run_evaluations,
+            selected_evaluation_names,
+            loop_response,
+            self._latest_assistant_id,
         )
 
-        await self._postprocessor_manager.run_postprocessors(loop_response)
+        postprocessor_result = task_executor.execute_async(
+            self._postprocessor_manager.run_postprocessors,
+            loop_response.model_copy(deep=True),
+        )
 
-        if not all(result.is_positive for result in evaluation_results):
+        _, evaluation_results = await asyncio.gather(
+            postprocessor_result,
+            evaluation_results,
+        )
+
+        if evaluation_results.success and not all(
+            result.is_positive for result in evaluation_results.unpack()
+        ):
             self._logger.warning(
                 "we should add here the retry counter add an instruction and retry the loop for now we just exit the loop"
             )  # TODO: add retry counter and instruction
 
         return True
 
+    def _log_tool_calls(self, tool_calls: list) -> None:
+        # Create dictionary mapping tool names to display names for efficient lookup
+        all_tools_dict: dict[str, str] = {
+            tool.name: tool.display_name()
+            for tool in self._tool_manager.available_tools
+        }
+
+        # Tool names that should not be logged in the message steps
+        tool_names_not_to_log = ["DeepResearch"]
+
+        tool_string: str = ""
+        used_tools = []
+        for tool_call in tool_calls:
+            if tool_call.name in all_tools_dict.keys() and tool_call.name not in tool_names_not_to_log:
+                if tool_call.name not in used_tools:
+                    used_tools.append(tool_call.name)
+                    tool_name = (all_tools_dict[tool_call.name]) or tool_call.name
+                    tool_string += f"\n• {tool_name}"
+            self._history_manager.add_tool_call(tool_call)
+
+        if tool_string:
+            self._message_step_logger.create_message_log_entry(
+                text=f"**Triggered Tool Calls:**\n {tool_string}", references=[]
+            )
+
     async def _handle_tool_calls(
         self, loop_response: LanguageModelStreamResponse
     ) -> bool:
@@ -323,9 +428,11 @@ class UniqueAI:
 
         tool_calls = loop_response.tool_calls or []
 
-        # Append function call to history
+        # Append function calls to history
         self._history_manager._append_tool_calls_to_history(tool_calls)
 
+        # Log tool calls
+        self._log_tool_calls(tool_calls)
         # Execute tool calls
         tool_call_responses = await self._tool_manager.execute_selected_tools(
             tool_calls
@@ -336,9 +443,14 @@ class UniqueAI:
         # then extract referenceable chunks and debug info
         self._history_manager.add_tool_call_results(tool_call_responses)
         self._reference_manager.extract_referenceable_chunks(tool_call_responses)
-        self._debug_info_manager.extract_tool_debug_info(tool_call_responses)
+        self._debug_info_manager.extract_tool_debug_info(
+            tool_call_responses, self.current_iteration_index
+        )
 
-        return self._tool_manager.does_a_tool_take_control(tool_calls)
+        self._tool_took_control = self._tool_manager.does_a_tool_take_control(
+            tool_calls
+        )
+        return self._tool_took_control
 
     async def _create_new_assistant_message_if_loop_response_contains_content(
         self, loop_response: LanguageModelStreamResponse
@@ -373,3 +485,61 @@ class UniqueAI:
                 content=loop_response.message.original_text or "",
             )
         )
+
+    def _get_filtered_user_metadata(self) -> dict[str, str]:
+        """
+        Filter user metadata to only include keys specified in the agent's prompt config.
+
+        Returns:
+            Dictionary containing only the metadata keys that are configured to be included.
+        """
+        user_metadata = {}
+        if (
+            self._config.agent.prompt_config.user_metadata
+            and self._event.payload.user_metadata is not None
+        ):
+            # Filter metadata to only include selected keys
+            user_metadata = {
+                k: str(v)
+                for k, v in self._event.payload.user_metadata.items()
+                if k in self._config.agent.prompt_config.user_metadata
+            }
+        return user_metadata
+
+
+class UniqueAIResponsesApi(UniqueAI):
+    def __init__(
+        self,
+        logger: Logger,
+        event: ChatEvent,
+        config: UniqueAIConfig,
+        chat_service: ChatService,
+        content_service: ContentService,
+        debug_info_manager: DebugInfoManager,
+        streaming_handler: ResponsesSupportCompleteWithReferences,
+        reference_manager: ReferenceManager,
+        thinking_manager: ThinkingManager,
+        tool_manager: ResponsesApiToolManager,
+        history_manager: HistoryManager,
+        evaluation_manager: EvaluationManager,
+        postprocessor_manager: PostprocessorManager,
+        message_step_logger: MessageStepLogger,
+        mcp_servers: list[McpServer],
+    ) -> None:
+        super().__init__(
+            logger,
+            event=event,
+            config=config,
+            chat_service=chat_service,
+            content_service=content_service,
+            debug_info_manager=debug_info_manager,
+            streaming_handler=streaming_handler,  # type: ignore
+            reference_manager=reference_manager,
+            thinking_manager=thinking_manager,
+            tool_manager=tool_manager,  # type: ignore
+            history_manager=history_manager,
+            evaluation_manager=evaluation_manager,
+            postprocessor_manager=postprocessor_manager,
+            message_step_logger=message_step_logger,
+            mcp_servers=mcp_servers,
+        )