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

unique_orchestrator/config.py

@@ -2,14 +2,14 @@ from enum import StrEnum
 from pathlib import Path
 from typing import Annotated, Any, Generic, Literal, TypeVar
 
-from pydantic import BaseModel, Field, ValidationInfo, field_validator
+from pydantic import BaseModel, Field, ValidationInfo, field_validator, model_validator
 from unique_deep_research.config import DeepResearchToolConfig
 from unique_deep_research.service import DeepResearchTool
 from unique_follow_up_questions.config import FollowUpQuestionsConfig
 from unique_internal_search.config import InternalSearchConfig
 from unique_internal_search.service import InternalSearchTool
 from unique_stock_ticker.config import StockTickerConfig
-from unique_toolkit._common.default_language_model import DEFAULT_GPT_4o
+from unique_swot import SwotAnalysisTool, SwotAnalysisToolConfig
 from unique_toolkit._common.validators import (
     LMI,
     ClipInt,
@@ -22,11 +22,32 @@ from unique_toolkit.agentic.evaluation.schemas import EvaluationMetricName
 from unique_toolkit.agentic.history_manager.history_manager import (
     UploadedContentConfig,
 )
+from unique_toolkit.agentic.responses_api import (
+    DisplayCodeInterpreterFilesPostProcessorConfig,
+    ShowExecutedCodePostprocessorConfig,
+)
+from unique_toolkit.agentic.tools.a2a import (
+    REFERENCING_INSTRUCTIONS_FOR_SYSTEM_PROMPT,
+    REFERENCING_INSTRUCTIONS_FOR_USER_PROMPT,
+)
+from unique_toolkit.agentic.tools.a2a.evaluation import SubAgentEvaluationServiceConfig
 from unique_toolkit.agentic.tools.config import get_configuration_dict
+from unique_toolkit.agentic.tools.openai_builtin.manager import (
+    OpenAICodeInterpreterConfig,
+)
 from unique_toolkit.agentic.tools.tool import ToolBuildConfig
+from unique_toolkit.agentic.tools.tool_progress_reporter import (
+    ToolProgressReporterConfig,
+)
+from unique_toolkit.language_model.default_language_model import DEFAULT_GPT_4o
 from unique_web_search.config import WebSearchConfig
 from unique_web_search.service import WebSearchTool
 
+DeactivatedNone = Annotated[
+    None,
+    Field(title="Deactivated", description="None"),
+]
+
 
 class SpaceType(StrEnum):
     UNIQUE_CUSTOM = "unique_custom"
@@ -72,6 +93,10 @@ class SpaceConfigBase(BaseModel, Generic[T]):
                 name=DeepResearchTool.name,
                 configuration=DeepResearchToolConfig(),
             ),
+            ToolBuildConfig(
+                name=SwotAnalysisTool.name,
+                configuration=SwotAnalysisToolConfig(),
+            ),
         ],
     )
 
@@ -140,11 +165,11 @@ class UniqueAIPromptConfig(BaseModel):
         description="The user message prompt template as a Jinja2 template string.",
     )
 
-
-DeactivatedNone = Annotated[
-    None,
-    Field(title="Deactivated", description="None"),
-]
+    user_metadata: list[str] = Field(
+        default=[],
+        title="User Metadata",
+        description="User metadata fields to be ingested in the system prompt.",
+    )
 
 
 class UniqueAIServices(BaseModel):
@@ -182,6 +207,10 @@ class UniqueAIServices(BaseModel):
 
     uploaded_content_config: UploadedContentConfig = UploadedContentConfig()
 
+    tool_progress_reporter_config: ToolProgressReporterConfig = (
+        ToolProgressReporterConfig()
+    )
+
 
 class InputTokenDistributionConfig(BaseModel):
     model_config = get_configuration_dict(frozen=True)
@@ -197,6 +226,79 @@ class InputTokenDistributionConfig(BaseModel):
         return int(self.percent_for_history * max_input_token)
 
 
+class SubAgentsReferencingConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    referencing_instructions_for_system_prompt: str = Field(
+        default=REFERENCING_INSTRUCTIONS_FOR_SYSTEM_PROMPT,
+        description="Referencing instructions for the main agent's system prompt.",
+    )
+    referencing_instructions_for_user_prompt: str = Field(
+        default=REFERENCING_INSTRUCTIONS_FOR_USER_PROMPT,
+        description="Referencing instructions for the main agent's user prompt. Should correspond to a short reminder.",
+    )
+
+
+class SubAgentsConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    referencing_config: (
+        Annotated[SubAgentsReferencingConfig, Field(title="Active")] | DeactivatedNone
+    ) = SubAgentsReferencingConfig()
+    evaluation_config: (
+        Annotated[SubAgentEvaluationServiceConfig, Field(title="Active")]
+        | DeactivatedNone
+    ) = SubAgentEvaluationServiceConfig()
+
+    sleep_time_before_update: float = Field(
+        default=0.5,
+        description="Time to sleep before updating the main agent message to display the sub agent responses. Temporary fix to avoid rendering issues.",
+    )
+
+
+class CodeInterpreterExtendedConfig(BaseModel):
+    model_config = get_configuration_dict()
+
+    generated_files_config: DisplayCodeInterpreterFilesPostProcessorConfig = Field(
+        default=DisplayCodeInterpreterFilesPostProcessorConfig(),
+        title="Generated files config",
+        description="Display config for generated files",
+    )
+
+    executed_code_display_config: (
+        Annotated[
+            ShowExecutedCodePostprocessorConfig,
+            Field(title="Active"),
+        ]
+        | DeactivatedNone
+    ) = Field(
+        ShowExecutedCodePostprocessorConfig(),
+        description="If active, generated code will be prepended to the LLM answer",
+    )
+
+    tool_config: OpenAICodeInterpreterConfig = Field(
+        default=OpenAICodeInterpreterConfig(),
+        title="Tool config",
+    )
+
+
+class ResponsesApiConfig(BaseModel):
+    model_config = get_configuration_dict(frozen=True)
+
+    code_interpreter: (
+        Annotated[CodeInterpreterExtendedConfig, Field(title="Active")]
+        | DeactivatedNone
+    ) = Field(
+        default=None,
+        description="If active, the main agent will have acces to the OpenAI Code Interpreter tool",
+    )
+
+    use_responses_api: bool = Field(
+        default=False,
+        description="If set, the main agent will use the Responses API from OpenAI",
+    )
+
+
 class ExperimentalConfig(BaseModel):
     """Experimental features this part of the configuration might evolve in the future continuously"""
 
@@ -228,6 +330,10 @@ class ExperimentalConfig(BaseModel):
         max_tool_calls_per_iteration=5
     )
 
+    sub_agents_config: SubAgentsConfig = SubAgentsConfig()
+
+    responses_api_config: ResponsesApiConfig = ResponsesApiConfig()
+
 
 class UniqueAIAgentConfig(BaseModel):
     model_config = get_configuration_dict(frozen=True)
@@ -252,3 +358,9 @@ class UniqueAIConfig(BaseModel):
     space: UniqueAISpaceConfig = UniqueAISpaceConfig()
 
     agent: UniqueAIAgentConfig = UniqueAIAgentConfig()
+
+    @model_validator(mode="after")
+    def disable_sub_agent_referencing_if_not_used(self) -> "UniqueAIConfig":
+        if not any(tool.is_sub_agent for tool in self.space.tools):
+            self.agent.experimental.sub_agents_config.referencing_config = None
+        return self

unique_orchestrator/prompts/system_prompt.jinja2

@@ -7,6 +7,24 @@ You are Unique AI Chat a system based on large language models
 **Knowledge cutoff**: {{ model_info.info_cutoff_at | default('unknown') }}
 **Current date**: {{ date_string }}
 
+{#- Expired Uploaded Documents Section #}
+{% if uploaded_documents_expired and uploaded_documents_expired|length > 0 %}
+# Expired Uploaded Documents
+Here are the uploaded documents that are expired due to company retention policy. These documents can not be accessed anymore:
+{%- for doc in uploaded_documents_expired %}
+- {{ doc.title or doc.key }}
+{%- endfor %}
+{%- endif %}
+
+{#- User Metadata Section #}
+{% if user_metadata and user_metadata|length > 0 %}
+# User Information
+Here is some metadata about the user, which may help you write better queries, and help contextualize the information you retrieve:
+{%- for key, value in user_metadata.items() %}
+- {{ key.replace('_', ' ').title() }}: {{ value }}
+{%- endfor %}
+{%- endif %}
+
 Over the course of the conversation, you adapt to the user's tone and preference.
 Try to match the user's vibe, tone, and generally how they are speaking. You want the conversation to feel natural.
 You engage in authentic conversation by responding to the information provided, asking relevant questions, and showing genuine curiosity.
@@ -52,13 +70,13 @@ This tool is called {{ tool_description.display_name }} by the user.
 {%- set tool_messages = {
   'automations': {'display': 'automations', 'message': 'Cannot create reminders, recurring tasks, or scheduled prompts.'},
   'canmore': {'display': 'canmore', 'message': 'Cannot create or edit documents/canvas for writing or coding.'},
-  'InternalSearch': {'display': 'file_search', 'message': 'Cannot search across internal company sources (Google Drive, Slack, Notion, etc.).'},
+  'InternalSearch': {'display': 'Internal Search', 'message': 'Cannot search across internal company sources. If you cannot force it in this space, ask an admin to activate the tool.'},
   'gcal': {'display': 'gcal (Google Calendar)', 'message': 'Cannot show or search calendar events.'},
   'gcontacts': {'display': 'gcontacts (Google Contacts)', 'message': 'Cannot look up or retrieve contact information.'},
   'gmail': {'display': 'gmail', 'message': 'Cannot search, read, or summarize emails.'},
   'image_gen': {'display': 'image_gen', 'message': 'Cannot generate or edit images.'},
   'python': {'display': 'python', 'message': 'Cannot analyze data, process files, generate charts, or create/export different file formats.'},
-  'WebSearch': {'display': 'web', 'message': 'Cannot perform live web searches, fetch fresh news, or look up real-time information.'},
+  'WebSearch': {'display': 'Web Search', 'message': 'Cannot perform live web searches, fetch fresh news, or look up real-time information. If you cannot force it in this space, ask an admin to activate the tool.'},
   'recording_knowledge': {'display': 'recording_knowledge', 'message': 'Cannot access or summarize meeting transcripts from ChatGPT Record Mode.'}
 } -%}
 {%- set ns = namespace(any=false) -%}
@@ -73,7 +91,12 @@ This tool is called {{ tool_description.display_name }} by the user.
   - {{ meta.message }}
   {%- endif -%}
 {%- endfor -%}
-{%- endif -%}
+{%- endif %}
+
+{% if use_sub_agent_references and sub_agent_referencing_instructions | length > 0 %}
+# Special Instructions: Referencing Sub Agents Answers
+{{ sub_agent_referencing_instructions }}
+{%- endif %}
 
 {# Answer Style Section #}
 # Answer Style

unique_orchestrator/prompts/user_message_prompt.jinja2

@@ -6,13 +6,18 @@
     - 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 }}
+#}
+# User Query
+{{ query }}
 
-{%- if mcp_server_user_prompts and mcp_server_user_prompts|length > 0 %}
+{% if use_sub_agent_references and sub_agent_referencing_instructions | length > 0 %}
+# Sub Agents Referencing Reminder
+{{ sub_agent_referencing_instructions }}
+{%- endif %}
 
-## MCP Server Context
+{%- 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_unique_ai_get_filtered_user_metadata.py

@@ -0,0 +1,259 @@
+from unittest.mock import MagicMock
+
+import pytest
+
+
+class TestGetFilteredUserMetadata:
+    """Test suite for UniqueAI._get_filtered_user_metadata method"""
+
+    @pytest.fixture
+    def mock_unique_ai(self):
+        """Create a minimal UniqueAI instance with mocked dependencies"""
+        # Lazy import to avoid heavy dependencies at module import time
+        from unique_orchestrator.unique_ai import UniqueAI
+
+        mock_logger = MagicMock()
+
+        # Create minimal event structure
+        dummy_event = MagicMock()
+        dummy_event.payload.assistant_message.id = "assist_1"
+        dummy_event.payload.user_message.text = "query"
+
+        # Create minimal config structure
+        mock_config = MagicMock()
+        mock_config.agent.prompt_config.user_metadata = []
+
+        # Create minimal required dependencies
+        mock_chat_service = MagicMock()
+        mock_content_service = MagicMock()
+        mock_debug_info_manager = MagicMock()
+        mock_reference_manager = MagicMock()
+        mock_thinking_manager = MagicMock()
+        mock_tool_manager = MagicMock()
+        mock_history_manager = MagicMock()
+        mock_evaluation_manager = MagicMock()
+        mock_postprocessor_manager = MagicMock()
+        mock_streaming_handler = MagicMock()
+        mock_message_step_logger = MagicMock()
+
+        # Instantiate UniqueAI
+        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,
+            streaming_handler=mock_streaming_handler,
+            reference_manager=mock_reference_manager,
+            thinking_manager=mock_thinking_manager,
+            tool_manager=mock_tool_manager,
+            history_manager=mock_history_manager,
+            evaluation_manager=mock_evaluation_manager,
+            postprocessor_manager=mock_postprocessor_manager,
+            message_step_logger=mock_message_step_logger,
+            mcp_servers=[],
+        )
+
+        return ua
+
+    def test_returns_empty_dict_when_config_is_empty_list(self, mock_unique_ai):
+        """Test that empty dict is returned when config.user_metadata is an empty list"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = []
+        mock_unique_ai._event.payload.user_metadata = {
+            "department": "Engineering",
+            "role": "Developer",
+        }
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        assert result == {}
+        assert isinstance(result, dict)
+
+    def test_returns_empty_dict_when_user_metadata_is_none(self, mock_unique_ai):
+        """Test that empty dict is returned when user_metadata is None"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = [
+            "department",
+            "role",
+        ]
+        mock_unique_ai._event.payload.user_metadata = None
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        assert result == {}
+        assert isinstance(result, dict)
+
+    def test_returns_empty_dict_when_both_config_and_metadata_are_empty(
+        self, mock_unique_ai
+    ):
+        """Test that empty dict is returned when both config and user_metadata are empty/None"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = []
+        mock_unique_ai._event.payload.user_metadata = None
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        assert result == {}
+        assert isinstance(result, dict)
+
+    def test_filters_metadata_to_include_only_configured_keys(self, mock_unique_ai):
+        """Test that only keys specified in config are included in the result"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = [
+            "department",
+            "role",
+        ]
+        mock_unique_ai._event.payload.user_metadata = {
+            "department": "Engineering",
+            "role": "Developer",
+            "location": "San Francisco",
+            "salary": "100000",
+        }
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        assert result == {"department": "Engineering", "role": "Developer"}
+        assert "location" not in result
+        assert "salary" not in result
+        # Verify all values are strings
+        assert all(isinstance(v, str) for v in result.values())
+
+    def test_returns_only_existing_keys_from_user_metadata(self, mock_unique_ai):
+        """Test that keys in config but not in user_metadata are not included"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = [
+            "department",
+            "role",
+            "team",
+            "manager",
+        ]
+        mock_unique_ai._event.payload.user_metadata = {
+            "department": "Engineering",
+            "role": "Developer",
+        }
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        assert result == {"department": "Engineering", "role": "Developer"}
+        assert "team" not in result
+        assert "manager" not in result
+
+    def test_handles_single_key_in_config(self, mock_unique_ai):
+        """Test filtering with a single key in config"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = ["department"]
+        mock_unique_ai._event.payload.user_metadata = {
+            "department": "Engineering",
+            "role": "Developer",
+        }
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        assert result == {"department": "Engineering"}
+        assert isinstance(result["department"], str)
+
+    def test_handles_string_values(self, mock_unique_ai):
+        """Test that string values in user_metadata are preserved"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = [
+            "name",
+            "email",
+            "department",
+            "title",
+        ]
+        mock_unique_ai._event.payload.user_metadata = {
+            "name": "John Doe",
+            "email": "john.doe@example.com",
+            "department": "Engineering",
+            "title": "Senior Developer",
+            "ignored": "This should not appear",
+        }
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        assert result == {
+            "name": "John Doe",
+            "email": "john.doe@example.com",
+            "department": "Engineering",
+            "title": "Senior Developer",
+        }
+        assert "ignored" not in result
+        # Verify all values are strings
+        assert all(isinstance(v, str) for v in result.values())
+
+    def test_handles_empty_dict_user_metadata(self, mock_unique_ai):
+        """Test behavior when user_metadata is an empty dict"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = ["department"]
+        mock_unique_ai._event.payload.user_metadata = {}
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        assert result == {}
+
+    def test_handles_empty_string_values(self, mock_unique_ai):
+        """Test that empty string values in user_metadata are preserved if key is in config"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = [
+            "department",
+            "role",
+        ]
+        mock_unique_ai._event.payload.user_metadata = {
+            "department": "Engineering",
+            "role": "",
+        }
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        assert result == {"department": "Engineering", "role": ""}
+        assert isinstance(result["role"], str)
+
+    def test_preserves_original_metadata_unchanged(self, mock_unique_ai):
+        """Test that the original user_metadata dict is not modified"""
+        original_metadata = {
+            "department": "Engineering",
+            "role": "Developer",
+            "location": "San Francisco",
+        }
+        mock_unique_ai._config.agent.prompt_config.user_metadata = ["department"]
+        mock_unique_ai._event.payload.user_metadata = original_metadata.copy()
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        # Original should still have all keys
+        assert mock_unique_ai._event.payload.user_metadata == original_metadata
+        # Result should only have filtered key
+        assert result == {"department": "Engineering"}
+
+    def test_handles_special_characters_in_values(self, mock_unique_ai):
+        """Test that special characters in string values are preserved"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = [
+            "description",
+            "notes",
+        ]
+        mock_unique_ai._event.payload.user_metadata = {
+            "description": "User with special chars: @#$%^&*()",
+            "notes": "Multi-line\ntext\twith\ttabs",
+            "other": "excluded",
+        }
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        assert result == {
+            "description": "User with special chars: @#$%^&*()",
+            "notes": "Multi-line\ntext\twith\ttabs",
+        }
+        assert all(isinstance(v, str) for v in result.values())
+
+    def test_return_type_is_dict_str_str(self, mock_unique_ai):
+        """Test that return type is dict[str, str]"""
+        mock_unique_ai._config.agent.prompt_config.user_metadata = [
+            "department",
+            "role",
+        ]
+        mock_unique_ai._event.payload.user_metadata = {
+            "department": "Engineering",
+            "role": "Developer",
+        }
+
+        result = mock_unique_ai._get_filtered_user_metadata()
+
+        # Check it's a dict
+        assert isinstance(result, dict)
+        # Check all keys are strings
+        assert all(isinstance(k, str) for k in result.keys())
+        # Check all values are strings
+        assert all(isinstance(v, str) for v in result.values())