unique_orchestrator 0.0.1__py3-none-any.whl

unique_orchestrator/prompts/user_message_prompt.jinja2

@@ -0,0 +1,18 @@
+{# Comments for the user message template
+    - This template is used to format the user message for the LLM
+    - Variables available:
+    - query: The original user query
+    - model_info: Information about the language model being used
+    - date_string: The current date in formatted string
+    - mcp_server_user_prompts: List of unique server-wide user prompts from MCP servers
+    - tool_descriptions_with_user_prompts: List of UniqueToolDescription objects with user prompts
+#}{{ query }}
+
+{%- if mcp_server_user_prompts and mcp_server_user_prompts|length > 0 %}
+
+## MCP Server Context
+{%- for server_prompt in mcp_server_user_prompts %}
+
+{{ server_prompt }}
+{%- endfor %}
+{%- endif %}

unique_orchestrator/tests/test_config.py

@@ -0,0 +1,115 @@
+from unique_orchestrator.config import (
+    SearchAgentConfig,
+    search_agent_config_to_unique_ai_space_config,
+)
+
+
+def test_ai_tools_conversion():
+    """
+    AI-authored test: Validates the transfer of tool configurations.
+
+    This test ensures that:
+    1. All tools from the old config are present in the new config
+    2. Tool names are preserved exactly
+    3. Tool configurations are maintained without modification
+
+    Written by AI Assistant to verify tool configuration preservation.
+    """
+    old_config = SearchAgentConfig()
+    new_config = search_agent_config_to_unique_ai_space_config(old_config)
+
+    assert len(new_config.space.tools) == len(old_config.tools)
+    for old_tool, new_tool in zip(old_config.tools, new_config.space.tools):
+        assert old_tool.name == new_tool.name
+        assert old_tool.configuration == new_tool.configuration
+
+
+def test_ai_services_conversion():
+    """
+    AI-authored test: Verifies the conversion of service configurations.
+
+    This test checks that all service configurations are properly transferred:
+    1. Follow-up questions configuration
+    2. Evaluation configuration
+    3. Stock ticker configuration
+    4. Reference manager configuration
+
+    Written by AI Assistant to ensure service configuration integrity.
+    """
+    old_config = SearchAgentConfig()
+    new_config = search_agent_config_to_unique_ai_space_config(old_config)
+
+    services = new_config.agent.services
+    assert services.follow_up_questions_config == old_config.follow_up_questions_config
+    assert services.evaluation_config == old_config.evaluation_config
+    assert services.stock_ticker_config == old_config.stock_ticker_config
+
+
+def test_ai_experimental_config_conversion():
+    """
+    AI-authored test: Checks the conversion of experimental features.
+
+    This test verifies that:
+    1. Experimental features like thinking_steps_display are properly transferred
+    2. Boolean values are preserved accurately
+
+    Written by AI Assistant to ensure experimental feature preservation.
+    """
+    old_config = SearchAgentConfig()
+    old_config.thinking_steps_display = True
+    new_config = search_agent_config_to_unique_ai_space_config(old_config)
+
+    assert (
+        new_config.agent.experimental.thinking_steps_display
+        == old_config.thinking_steps_display
+    )
+
+
+def test_ai_force_checks_conversion():
+    """
+    AI-authored test: Validates the conversion of force checks configuration.
+
+    This test ensures that:
+    1. Force checks for stream response references are properly transferred
+    2. The configuration maintains its integrity during conversion
+
+    Written by AI Assistant to verify force checks preservation.
+    """
+    old_config = SearchAgentConfig()
+    new_config = search_agent_config_to_unique_ai_space_config(old_config)
+
+    assert (
+        new_config.agent.experimental.force_checks_on_stream_response_references
+        == old_config.force_checks_on_stream_response_references
+    )
+
+
+def test_ai_custom_values_conversion():
+    """
+    AI-authored test: Verifies the conversion of custom configuration values.
+
+    This test validates that custom values are properly transferred:
+    1. Project name
+    2. Custom instructions
+    3. Temperature settings
+    4. Loop iteration limits
+    5. Additional LLM options
+
+    Written by AI Assistant to ensure custom configuration preservation.
+    """
+    old_config = SearchAgentConfig(
+        project_name="Custom Project",
+        custom_instructions="Custom Instructions",
+        temperature=0.8,
+        max_loop_iterations=5,
+        additional_llm_options={"some_option": "value"},
+    )
+    new_config = search_agent_config_to_unique_ai_space_config(old_config)
+
+    assert new_config.space.project_name == "Custom Project"
+    assert new_config.space.custom_instructions == "Custom Instructions"
+    assert new_config.agent.experimental.temperature == 0.8
+    assert new_config.agent.max_loop_iterations == 5
+    assert new_config.agent.experimental.additional_llm_options == {
+        "some_option": "value"
+    }

unique_orchestrator/tests/test_unique_ai_reference_order.py

@@ -0,0 +1,127 @@
+from unittest.mock import AsyncMock, MagicMock
+
+import pytest
+
+
+@pytest.mark.asyncio
+async def test_history_updated_before_reference_extraction(monkeypatch):
+    # Lazy import to avoid heavy dependencies at module import time
+    from unique_orchestrator.unique_ai import UniqueAI
+
+    # Create a minimal UniqueAI instance with mocked dependencies
+    mock_logger = MagicMock()
+
+    class DummyEvent:
+        class Payload:
+            class AssistantMessage:
+                id = "assist_1"
+
+            assistant_message = AssistantMessage()
+            user_message = MagicMock()
+            user_message.text = "query"
+
+        payload = Payload()
+
+    dummy_event = MagicMock()
+    dummy_event.payload = DummyEvent.Payload()
+
+    mock_config = MagicMock()
+    mock_config.agent.max_loop_iterations = 1
+    mock_config.space.language_model.name = "dummy-model"
+    mock_config.agent.experimental.temperature = 0.0
+    mock_config.agent.experimental.additional_llm_options = {}
+
+    # Managers
+    mock_history_manager = MagicMock()
+    mock_history_manager.has_no_loop_messages.return_value = True
+    mock_history_manager._append_tool_calls_to_history = MagicMock()
+    mock_history_manager.add_tool_call_results = MagicMock()
+
+    mock_reference_manager = MagicMock()
+    mock_reference_manager.extract_referenceable_chunks = MagicMock()
+    mock_reference_manager.get_chunks.return_value = []
+
+    mock_thinking_manager = MagicMock()
+    mock_debug_info_manager = MagicMock()
+    mock_debug_info_manager.get.return_value = {}
+    mock_debug_info_manager.extract_tool_debug_info = MagicMock()
+
+    mock_tool_manager = MagicMock()
+    mock_tool_manager.get_forced_tools.return_value = []
+    mock_tool_manager.get_tool_definitions.return_value = []
+    mock_tool_manager.execute_selected_tools = AsyncMock(return_value=[MagicMock()])
+    mock_tool_manager.does_a_tool_take_control.return_value = False
+
+    class DummyStreamResponse:
+        def __init__(self):
+            self.tool_calls = [MagicMock()]
+            self.message = MagicMock()
+            self.message.references = []
+            self.message.text = ""
+
+        def is_empty(self):
+            return False
+
+    mock_chat_service = MagicMock()
+    mock_chat_service.complete_with_references_async = AsyncMock(
+        return_value=DummyStreamResponse()
+    )
+    mock_chat_service.modify_assistant_message_async = AsyncMock(return_value=None)
+    mock_chat_service.create_assistant_message_async = AsyncMock(
+        return_value=MagicMock(id="assist_new")
+    )
+    mock_content_service = MagicMock()
+    mock_history_manager.get_history_for_model_call = AsyncMock(
+        return_value=MagicMock()
+    )
+
+    # Instantiate
+    ua = UniqueAI(
+        logger=mock_logger,
+        event=dummy_event,
+        config=mock_config,
+        chat_service=mock_chat_service,
+        content_service=mock_content_service,
+        debug_info_manager=mock_debug_info_manager,
+        reference_manager=mock_reference_manager,
+        thinking_manager=mock_thinking_manager,
+        tool_manager=mock_tool_manager,
+        history_manager=mock_history_manager,
+        evaluation_manager=MagicMock(),
+        postprocessor_manager=MagicMock(),
+        mcp_servers=[],
+    )
+
+    # Bypass Jinja template compilation by stubbing prompt renderers
+    ua._render_user_prompt = AsyncMock(return_value="user")  # type: ignore
+    ua._render_system_prompt = AsyncMock(return_value="system")  # type: ignore
+    # Avoid creating new assistant messages path
+    ua._thinking_manager.thinking_is_displayed = MagicMock(return_value=True)  # type: ignore
+
+    # Spy on call order by recording sequence
+    call_sequence = []
+
+    def record_history_add(results):
+        call_sequence.append("history_add")
+
+    def record_reference_extract(results):
+        call_sequence.append("reference_extract")
+
+    def record_debug_extract(results):
+        call_sequence.append("debug_extract")
+
+    mock_history_manager.add_tool_call_results.side_effect = record_history_add
+    mock_reference_manager.extract_referenceable_chunks.side_effect = (
+        record_reference_extract
+    )
+    mock_debug_info_manager.extract_tool_debug_info.side_effect = record_debug_extract
+
+    # Run a single iteration
+    await ua.run()
+
+    # Verify order: history first, then references, then debug
+    assert call_sequence[:3] == [
+        "history_add",
+        "reference_extract",
+        "debug_extract",
+    ]

unique_orchestrator/unique_ai.py

@@ -0,0 +1,375 @@
+from datetime import datetime
+from logging import Logger
+
+import jinja2
+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.debug_info_manager.debug_info_manager import (
+    DebugInfoManager,
+)
+from unique_toolkit.evals.evaluation_manager import EvaluationManager
+from unique_toolkit.history_manager.history_manager import HistoryManager
+from unique_toolkit.language_model.schemas import (
+    LanguageModelAssistantMessage,
+    LanguageModelMessages,
+    LanguageModelStreamResponse,
+)
+from unique_toolkit.postprocessor.postprocessor_manager import (
+    PostprocessorManager,
+)
+from unique_toolkit.reference_manager.reference_manager import ReferenceManager
+from unique_toolkit.thinking_manager.thinking_manager import ThinkingManager
+from unique_toolkit.tools.tool_manager import ToolManager
+
+from unique_orchestrator.config import UniqueAIConfig
+
+EMPTY_MESSAGE_WARNING = (
+    "⚠️ **The language model was unable to produce an output.**\n"
+    "It did not generate any content or perform a tool call in response to your request. "
+    "This is a limitation of the language model itself.\n\n"
+    "**Please try adapting or simplifying your prompt.** "
+    "Rewording your input can often help the model respond successfully."
+)
+
+
+class UniqueAI:
+    start_text = ""
+    current_iteration_index = 0
+
+    def __init__(
+        self,
+        logger: Logger,
+        event: ChatEvent,
+        config: UniqueAIConfig,
+        chat_service: ChatService,
+        content_service: ContentService,
+        debug_info_manager: DebugInfoManager,
+        reference_manager: ReferenceManager,
+        thinking_manager: ThinkingManager,
+        tool_manager: ToolManager,
+        history_manager: HistoryManager,
+        evaluation_manager: EvaluationManager,
+        postprocessor_manager: PostprocessorManager,
+        mcp_servers: list[McpServer],
+    ):
+        self._logger = logger
+        self._event = event
+        self._config = config
+        self._chat_service = chat_service
+        self._content_service = content_service
+
+        self._debug_info_manager = debug_info_manager
+        self._reference_manager = reference_manager
+        self._thinking_manager = thinking_manager
+        self._tool_manager = tool_manager
+
+        self._history_manager = history_manager
+
+        self._evaluation_manager = evaluation_manager
+        self._postprocessor_manager = postprocessor_manager
+        self._latest_assistant_id: str = event.payload.assistant_message.id
+        self._mcp_servers = mcp_servers
+
+    ############################################################
+    # Override of base methods
+    ############################################################
+    # @track(name="loop_agent_run")  # Group traces together
+    async def run(self):
+        """
+        Main loop of the agent. The agent will iterate through the loop, runs the plan and
+        processes tool calls if any are returned.
+        """
+        self._logger.info("Start LoopAgent...")
+
+        if self._history_manager.has_no_loop_messages():  # TODO: why do we even need to check its always no loop messages on this when its called.
+            self._chat_service.modify_assistant_message(
+                content="Starting agentic loop..."  # TODO: this must be more informative
+            )
+
+        ## Loop iteration
+        for i in range(self._config.agent.max_loop_iterations):
+            self.current_iteration_index = i
+            self._logger.info(f"Starting iteration {i + 1}...")
+
+            # Plan execution
+            loop_response = await self._plan_or_execute()
+            self._logger.info("Done with _plan_or_execute")
+
+            self._reference_manager.add_references(loop_response.message.references)
+            self._logger.info("Done with adding references")
+
+            # Update tool progress reporter
+            self._thinking_manager.update_tool_progress_reporter(loop_response)
+
+            # Execute the plan
+            exit_loop = await self._process_plan(loop_response)
+            self._logger.info("Done with _process_plan")
+
+            if exit_loop:
+                self._thinking_manager.close_thinking_steps(loop_response)
+                self._logger.info("Exiting loop.")
+                break
+
+            if i == self._config.agent.max_loop_iterations - 1:
+                self._logger.error("Max iterations reached.")
+                await self._chat_service.modify_assistant_message_async(
+                    content="I have reached the maximum number of self-reflection iterations. Please clarify your request and try again...",
+                )
+                break
+
+            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
+            )
+
+        await self._chat_service.modify_assistant_message_async(
+            set_completed_at=True,
+        )
+
+    # @track()
+    async def _plan_or_execute(self) -> LanguageModelStreamResponse:
+        self._logger.info("Planning or executing the loop.")
+        messages = await self._compose_message_plan_execution()
+
+        self._logger.info("Done composing message plan execution.")
+
+        # Forces tool calls only in first iteration
+        if (
+            len(self._tool_manager.get_forced_tools()) > 0
+            and self.current_iteration_index == 0
+        ):
+            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(
+                    messages=messages,
+                    model_name=self._config.space.language_model.name,
+                    tools=self._tool_manager.get_tool_definitions(),
+                    content_chunks=self._reference_manager.get_chunks(),
+                    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},
+                )
+                for opt in self._tool_manager.get_forced_tools()
+            ]
+
+            # Merge responses and refs:
+            tool_calls = []
+            references = []
+            for r in responses:
+                if r.tool_calls:
+                    tool_calls.extend(r.tool_calls)
+                references.extend(r.message.references)
+
+            stream_response = responses[0]
+            stream_response.tool_calls = tool_calls if len(tool_calls) > 0 else None
+            stream_response.message.references = references
+        elif self.current_iteration_index == self._config.agent.max_loop_iterations - 1:
+            self._logger.info(
+                "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(
+                messages=messages,
+                model_name=self._config.space.language_model.name,
+                content_chunks=self._reference_manager.get_chunks(),
+                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,
+            )
+
+        else:
+            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(
+                messages=messages,
+                model_name=self._config.space.language_model.name,
+                tools=self._tool_manager.get_tool_definitions(),
+                content_chunks=self._reference_manager.get_chunks(),
+                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,
+            )
+
+        return stream_response
+
+    async def _process_plan(self, loop_response: LanguageModelStreamResponse) -> bool:
+        self._logger.info(
+            "Processing the plan, executing the tools and checking for loop exit conditions once all is done."
+        )
+
+        if loop_response.is_empty():
+            self._logger.debug("Empty model response, exiting loop.")
+            self._chat_service.modify_assistant_message(content=EMPTY_MESSAGE_WARNING)
+            return True
+
+        call_tools = len(loop_response.tool_calls or []) > 0
+        if call_tools:
+            self._logger.debug(
+                "Tools were called we process them and do not exit the loop"
+            )
+
+            return await self._handle_tool_calls(loop_response)
+
+        self._logger.debug("No tool calls. we might exit the loop")
+
+        return await self._handle_no_tool_calls(loop_response)
+
+    async def _compose_message_plan_execution(self) -> LanguageModelMessages:
+        original_user_message = self._event.payload.user_message.text
+        rendered_user_message_string = await self._render_user_prompt()
+        rendered_system_message_string = await self._render_system_prompt()
+
+        messages = await self._history_manager.get_history_for_model_call(
+            original_user_message,
+            rendered_user_message_string,
+            rendered_system_message_string,
+            self._postprocessor_manager.remove_from_text,
+        )
+        return messages
+
+    async def _render_user_prompt(self) -> str:
+        user_message_template = jinja2.Template(
+            self._config.agent.prompt_config.user_message_prompt_template
+        )
+
+        tool_descriptions_with_user_prompts = [
+            prompts.tool_user_prompt
+            for prompts in self._tool_manager.get_tool_prompts()
+        ]
+
+        used_tools = [m.name for m in self._tool_manager.get_tools()]
+
+        mcp_server_user_prompts = [
+            mcp_server.user_prompt for mcp_server in self._mcp_servers
+        ]
+
+        tool_descriptions = self._tool_manager.get_tool_prompts()
+
+        query = self._event.payload.user_message.text
+
+        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,
+        )
+        return user_msg
+
+    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()]
+
+        system_prompt_template = jinja2.Template(
+            self._config.agent.prompt_config.system_prompt_template
+        )
+
+        date_string = datetime.now().strftime("%A %B %d, %Y")
+
+        mcp_server_system_prompts = [
+            mcp_server.system_prompt for mcp_server in self._mcp_servers
+        ]
+
+        system_message = system_prompt_template.render(
+            model_info=self._config.space.language_model.model_dump(mode="json"),
+            date_string=date_string,
+            tool_descriptions=tool_descriptions,
+            used_tools=used_tools,
+            project_name=self._config.space.project_name,
+            custom_instructions=self._config.space.custom_instructions,
+            max_tools_per_iteration=self._config.agent.experimental.loop_configuration.max_tool_calls_per_iteration,
+            max_loop_iterations=self._config.agent.max_loop_iterations,
+            current_iteration=self.current_iteration_index + 1,
+            mcp_server_system_prompts=mcp_server_system_prompts,
+        )
+        return system_message
+
+    async def _handle_no_tool_calls(
+        self, loop_response: LanguageModelStreamResponse
+    ) -> bool:
+        """Handle the case where no tool calls are returned."""
+        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
+        )
+
+        await self._postprocessor_manager.run_postprocessors(loop_response)
+
+        if not all(result.is_positive for result in evaluation_results):
+            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
+
+    async def _handle_tool_calls(
+        self, loop_response: LanguageModelStreamResponse
+    ) -> bool:
+        """Handle the case where tool calls are returned."""
+        self._logger.info("Processing tool calls")
+
+        tool_calls = loop_response.tool_calls or []
+
+        # Append function call to history
+        self._history_manager._append_tool_calls_to_history(tool_calls)
+
+        # Execute tool calls
+        tool_call_responses = await self._tool_manager.execute_selected_tools(
+            tool_calls
+        )
+
+        # Process results with error handling
+        # Add tool call results to history first to stabilize source numbering,
+        # 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)
+
+        return self._tool_manager.does_a_tool_take_control(tool_calls)
+
+    async def _create_new_assistant_message_if_loop_response_contains_content(
+        self, loop_response: LanguageModelStreamResponse
+    ) -> None:
+        if self._thinking_manager.thinking_is_displayed():
+            return
+        if not loop_response.message.text:
+            return
+
+        # if anything sets the start text the model did not produce content.
+        # So we need to remove that text from the message.
+        message_text_without_start_text = loop_response.message.text.replace(
+            self.start_text.strip(), ""
+        ).strip()
+        if message_text_without_start_text == "":
+            return
+
+        ###
+        # ToDo: Once references on existing assistant messages can be deleted, we will switch from creating a new assistant message to modifying the existing one (with previous references deleted)
+        ###
+        new_assistant_message = await self._chat_service.create_assistant_message_async(
+            content=""
+        )
+
+        # the new message must have an id that is valid else we use the old one
+        self._latest_assistant_id = (
+            new_assistant_message.id or self._latest_assistant_id
+        )
+
+        self._history_manager.add_assistant_message(
+            LanguageModelAssistantMessage(
+                content=loop_response.message.original_text or "",
+            )
+        )