unique-web-search 1.9.0__py3-none-any.whl → 1.10.0__py3-none-any.whl

unique_web_search/config.py

@@ -35,6 +35,7 @@ from unique_web_search.services.executors.configs import (
     WebSearchV2Config,
     get_default_web_search_mode_config,
 )
+from unique_web_search.services.query_elicitation import QueryElicitationConfig
 from unique_web_search.services.search_engine import (
     get_default_search_engine_config,
     get_search_engine_config_types_from_names,
@@ -75,7 +76,11 @@ class AnswerGenerationConfig(BaseModel):
     )
 
 
-class ExperimentalFeatures(FeatureExtendedSourceSerialization): ...
+class ExperimentalFeatures(FeatureExtendedSourceSerialization):
+    query_elicitation_config: QueryElicitationConfig = Field(
+        default_factory=QueryElicitationConfig,
+        description="Query elicitation configuration",
+    )
 
 
 class WebSearchConfig(BaseToolConfig):
@@ -158,7 +163,11 @@ class WebSearchConfig(BaseToolConfig):
         description="Tool format information for system prompt. This is used to format the tool response for the system prompt.",
     )
 
-    experimental_features: SkipJsonSchema[ExperimentalFeatures] = ExperimentalFeatures()
+    experimental_features: ExperimentalFeatures = Field(
+        default_factory=ExperimentalFeatures,
+        description="Experimental features",
+        title="Experimental Features",
+    )
 
     debug: bool = Field(
         default=False,

unique_web_search/service.py

@@ -1,3 +1,4 @@
+import logging
 from time import time
 
 from typing_extensions import override
@@ -10,31 +11,32 @@ from unique_toolkit.agentic.tools.tool import (
     Tool,
 )
 from unique_toolkit.agentic.tools.tool_progress_reporter import ProgressState
-from unique_toolkit.chat.schemas import (
-    MessageLog,
-    MessageLogDetails,
-    MessageLogEvent,
-    MessageLogStatus,
-)
-from unique_toolkit.content import ContentReference
 from unique_toolkit.language_model.schemas import (
     LanguageModelFunction,
     LanguageModelToolDescription,
 )
 
 from unique_web_search.config import WebSearchConfig
-from unique_web_search.schema import StepType, WebSearchPlan, WebSearchToolParameters
+from unique_web_search.schema import WebSearchPlan, WebSearchToolParameters
 from unique_web_search.services.content_processing import ContentProcessor, WebPageChunk
 from unique_web_search.services.crawlers import get_crawler_service
 from unique_web_search.services.executors import (
     WebSearchV1Executor,
     WebSearchV2Executor,
 )
-from unique_web_search.services.executors.base_executor import WebSearchLogEntry
 from unique_web_search.services.executors.configs import WebSearchMode
+from unique_web_search.services.executors.context import (
+    ExecutorCallbacks,
+    ExecutorConfiguration,
+    ExecutorServiceContext,
+)
+from unique_web_search.services.message_log import WebSearchMessageLogger
+from unique_web_search.services.query_elicitation import QueryElicitationService
 from unique_web_search.services.search_engine import get_search_engine_service
 from unique_web_search.utils import WebSearchDebugInfo, reduce_sources_to_token_limit
 
+_LOGGER = logging.getLogger(__name__)
+
 
 class WebSearchTool(Tool[WebSearchConfig]):
     name = "WebSearch"
@@ -59,7 +61,6 @@ class WebSearchTool(Tool[WebSearchConfig]):
         )
         self.debug = self.config.debug
         self._display_name = kwargs.get("display_name", "Web Search")
-        self._active_message_log: MessageLog | None = None
 
         def content_reducer(web_page_chunks: list[WebPageChunk]) -> list[WebPageChunk]:
             return reduce_sources_to_token_limit(
@@ -106,35 +107,37 @@ class WebSearchTool(Tool[WebSearchConfig]):
 
     @override
     async def run(self, tool_call: LanguageModelFunction) -> ToolCallResponse:
-        self.logger.info("Running the WebSearch tool")
+        _LOGGER.info("Running the WebSearch tool")
         start_time = time()
         parameters = self.tool_parameter_calls.model_validate(
             tool_call.arguments,
         )
 
         debug_info = WebSearchDebugInfo(parameters=parameters.model_dump())
-        executor = self._get_executor(tool_call, parameters, debug_info)
+
+        web_search_message_logger = WebSearchMessageLogger(
+            message_step_logger=self._message_step_logger,
+            tool_display_name=self.display_name(),
+        )
+        executor = self._get_executor(
+            tool_call, parameters, debug_info, web_search_message_logger
+        )
+
+        notify_from_tool_call = self._ff_tool_progress_reporter_callback()
 
         try:
-            content_chunks, queries_for_log = await executor.run()
+            content_chunks = await executor.run()
             debug_info.num_chunks_in_final_prompts = len(content_chunks)
             debug_info.execution_time = time() - start_time
 
-            self._active_message_log = self.create_or_update_active_message_log(
-                queries_for_log=queries_for_log,
-                status=MessageLogStatus.COMPLETED,
-            )
+            await web_search_message_logger.finished()
 
-            if (
-                not feature_flags.is_new_answers_ui_enabled(self.company_id)
-                and self.tool_progress_reporter
-            ):
-                await self.tool_progress_reporter.notify_from_tool_call(
-                    tool_call=tool_call,
-                    name=executor.notify_name,
-                    message=executor.notify_message,
-                    state=ProgressState.FINISHED,
-                )
+            await notify_from_tool_call(
+                tool_call=tool_call,
+                name=executor.notify_name,
+                message=executor.notify_message,
+                state=ProgressState.FINISHED,
+            )
 
             return ToolCallResponse(
                 id=tool_call.id,  # type: ignore
@@ -143,22 +146,16 @@ class WebSearchTool(Tool[WebSearchConfig]):
                 content_chunks=content_chunks,
             )
         except Exception as e:
-            self.logger.exception(f"Error executing WebSearch tool: {e}")
+            _LOGGER.exception(f"Error executing WebSearch tool: {e}")
 
-            self._active_message_log = self.create_or_update_active_message_log(
-                status=MessageLogStatus.COMPLETED,
-            )
+            await web_search_message_logger.failed()
 
-            if (
-                not feature_flags.is_new_answers_ui_enabled(self.company_id)
-                and self.tool_progress_reporter
-            ):
-                await self.tool_progress_reporter.notify_from_tool_call(
-                    tool_call=tool_call,
-                    name=executor.notify_name,
-                    message=executor.notify_message,
-                    state=ProgressState.FAILED,
-                )
+            await notify_from_tool_call(
+                tool_call=tool_call,
+                name=executor.notify_name,
+                message=executor.notify_message,
+                state=ProgressState.FAILED,
+            )
 
             return ToolCallResponse(
                 id=tool_call.id,  # type: ignore
@@ -172,63 +169,59 @@ class WebSearchTool(Tool[WebSearchConfig]):
         tool_call: LanguageModelFunction,
         parameters: WebSearchPlan | WebSearchToolParameters,
         debug_info: WebSearchDebugInfo,
+        web_search_message_logger: WebSearchMessageLogger,
     ) -> WebSearchV1Executor | WebSearchV2Executor:
-        # Create callback that wraps the service's create_or_update_active_message_log method
-        def message_log_callback(
-            *,
-            progress_message: str | None = None,
-            queries_for_log: list[WebSearchLogEntry] | list[str] | None = None,
-            status: MessageLogStatus | None = None,
-        ) -> MessageLog | None:
-            return self.create_or_update_active_message_log(
-                progress_message=progress_message,
-                queries_for_log=queries_for_log,
-                status=status,
-            )
+        # Initialize query elicitation service and get callbacks
+        elicitation_service = QueryElicitationService(
+            chat_service=self._chat_service,
+            display_name=self.display_name(),
+            config=self.config.experimental_features.query_elicitation_config,
+        )
+
+        # Build context objects
+        services = ExecutorServiceContext(
+            search_engine_service=self.search_engine_service,
+            crawler_service=self.crawler_service,
+            content_processor=self.content_processor,
+            language_model_service=self.language_model_service,
+            chunk_relevancy_sorter=self.chunk_relevancy_sorter,
+        )
+
+        config = ExecutorConfiguration(
+            language_model=self.language_model,
+            chunk_relevancy_sort_config=self.config.chunk_relevancy_sort_config,
+            company_id=self.company_id,
+            debug_info=debug_info,
+        )
+
+        callbacks = ExecutorCallbacks(
+            message_log_callback=web_search_message_logger,
+            content_reducer=self.content_reducer,
+            query_elicitation=elicitation_service,
+            tool_progress_reporter=self._ff_tool_progress_reporter(),
+        )
 
         if isinstance(parameters, WebSearchPlan):
             assert self.config.web_search_mode_config.mode == WebSearchMode.V2
             return WebSearchV2Executor(
-                search_service=self.search_engine_service,
-                language_model_service=self.language_model_service,
-                language_model=self.language_model,
-                crawler_service=self.crawler_service,
+                services=services,
+                config=config,
+                callbacks=callbacks,
                 tool_call=tool_call,
                 tool_parameters=parameters,
-                company_id=self.company_id,
-                content_processor=self.content_processor,
-                chunk_relevancy_sorter=self.chunk_relevancy_sorter,
-                chunk_relevancy_sort_config=self.config.chunk_relevancy_sort_config,
-                tool_progress_reporter=self.tool_progress_reporter
-                if not feature_flags.is_new_answers_ui_enabled(self.company_id)
-                else None,
-                content_reducer=self.content_reducer,
                 max_steps=self.config.web_search_mode_config.max_steps,
-                debug_info=debug_info,
-                message_log_callback=message_log_callback,
             )
         elif isinstance(parameters, WebSearchToolParameters):
             assert self.config.web_search_mode_config.mode == WebSearchMode.V1
             return WebSearchV1Executor(
-                search_service=self.search_engine_service,
-                language_model_service=self.language_model_service,
-                language_model=self.language_model,
-                crawler_service=self.crawler_service,
+                services=services,
+                config=config,
+                callbacks=callbacks,
                 tool_call=tool_call,
                 tool_parameters=parameters,
-                company_id=self.company_id,
+                refine_query_system_prompt=self.config.web_search_mode_config.refine_query_mode.system_prompt,
                 mode=self.config.web_search_mode_config.refine_query_mode.mode,
                 max_queries=self.config.web_search_mode_config.max_queries,
-                content_processor=self.content_processor,
-                chunk_relevancy_sorter=self.chunk_relevancy_sorter,
-                chunk_relevancy_sort_config=self.config.chunk_relevancy_sort_config,
-                tool_progress_reporter=self.tool_progress_reporter
-                if not feature_flags.is_new_answers_ui_enabled(self.company_id)
-                else None,
-                content_reducer=self.content_reducer,
-                refine_query_system_prompt=self.config.web_search_mode_config.refine_query_mode.system_prompt,
-                debug_info=debug_info,
-                message_log_callback=message_log_callback,
             )
         else:
             raise ValueError(f"Invalid parameters: {parameters}")
@@ -244,65 +237,29 @@ class WebSearchTool(Tool[WebSearchConfig]):
             return []
         return evaluation_check_list
 
-    def _prepare_message_logs_entries(
-        self, queries_for_log: list[WebSearchLogEntry] | list[str]
-    ) -> tuple[MessageLogDetails, list[ContentReference]]:
-        # Handle list of strings (query strings only)
-        if queries_for_log and isinstance(queries_for_log[0], str):
-            queries_for_log = [
-                WebSearchLogEntry(
-                    type=StepType.SEARCH,
-                    message=query,
-                    web_search_results=[],
-                )
-                for query in queries_for_log
-            ]
-
-        details = MessageLogDetails(
-            data=[
-                MessageLogEvent(
-                    type="WebSearch",
-                    text=query_for_log.message,
+    def _ff_tool_progress_reporter(self):
+        if not feature_flags.enable_new_answers_ui_un_14411.is_enabled(self.company_id):
+            return self.tool_progress_reporter
+        return None
+
+    def _ff_tool_progress_reporter_callback(self):
+        async def notify_from_tool_call(
+            tool_call: LanguageModelFunction,
+            name: str,
+            message: str,
+            state: ProgressState,
+        ):
+            if (
+                not feature_flags.enable_new_answers_ui_un_14411.is_enabled(
+                    self.company_id
                 )
-                for query_for_log in queries_for_log
-            ]
-        )
-        references = []
-        sequence_number = 0
-        for query_for_log in queries_for_log:
-            for web_search_result in query_for_log.web_search_results:
-                references.append(
-                    web_search_result.to_content_reference(sequence_number)
+                and self.tool_progress_reporter is not None
+            ):
+                await self.tool_progress_reporter.notify_from_tool_call(
+                    tool_call, name, message, state
                 )
-                sequence_number += 1
-
-        return details, references
 
-    def create_or_update_active_message_log(
-        self,
-        *,
-        progress_message: str | None = None,
-        queries_for_log: list[WebSearchLogEntry] | list[str] | None = None,
-        status: MessageLogStatus | None = None,
-    ) -> MessageLog | None:
-        details = MessageLogDetails(data=[])
-        reference_list = []
-        if queries_for_log is not None:
-            details, reference_list = self._prepare_message_logs_entries(
-                queries_for_log
-            )
-
-        self._active_message_log = (
-            self._message_step_logger.create_or_update_message_log(
-                active_message_log=self._active_message_log,
-                header=self._display_name,
-                progress_message=progress_message,
-                details=details,
-                references=reference_list,
-                **({"status": status} if status is not None else {}),
-            )
-        )
-        return self._active_message_log
+        return notify_from_tool_call
 
 
 ToolFactory.register_tool(WebSearchTool, WebSearchConfig)

unique_web_search/services/executors/__init__.py

@@ -1,3 +1,8 @@
+from unique_web_search.services.executors.context import (
+    ExecutorCallbacks,
+    ExecutorConfiguration,
+    ExecutorServiceContext,
+)
 from unique_web_search.services.executors.web_search_v1_executor import (
     WebSearchV1Executor,
 )
@@ -5,4 +10,10 @@ from unique_web_search.services.executors.web_search_v2_executor import (
     WebSearchV2Executor,
 )
 
-__all__ = ["WebSearchV1Executor", "WebSearchV2Executor"]
+__all__ = [
+    "WebSearchV1Executor",
+    "WebSearchV2Executor",
+    "ExecutorServiceContext",
+    "ExecutorConfiguration",
+    "ExecutorCallbacks",
+]

unique_web_search/services/executors/base_executor.py

@@ -1,87 +1,65 @@
 import logging
 from abc import ABC, abstractmethod
 from time import time
-from typing import Callable, Optional, Protocol
 
-from pydantic import BaseModel
-from unique_toolkit import LanguageModelService
-from unique_toolkit._common.chunk_relevancy_sorter.config import (
-    ChunkRelevancySortConfig,
-)
-from unique_toolkit._common.chunk_relevancy_sorter.service import ChunkRelevancySorter
-from unique_toolkit._common.validators import LMI
+from unique_toolkit.agentic.feature_flags import feature_flags
 from unique_toolkit.agentic.tools.tool_progress_reporter import (
     ProgressState,
-    ToolProgressReporter,
 )
-from unique_toolkit.chat.schemas import MessageLog, MessageLogStatus
 from unique_toolkit.content import ContentChunk
 from unique_toolkit.language_model import LanguageModelFunction
 
-from unique_web_search.schema import StepType, WebSearchPlan, WebSearchToolParameters
-from unique_web_search.services.content_processing import ContentProcessor, WebPageChunk
-from unique_web_search.services.crawlers import CrawlerTypes
-from unique_web_search.services.search_engine import SearchEngineTypes
+from unique_web_search.schema import WebSearchPlan, WebSearchToolParameters
+from unique_web_search.services.content_processing import WebPageChunk
+from unique_web_search.services.executors.context import (
+    ExecutorCallbacks,
+    ExecutorConfiguration,
+    ExecutorServiceContext,
+)
 from unique_web_search.services.search_engine.schema import (
     WebSearchResult,
 )
-from unique_web_search.utils import StepDebugInfo, WebSearchDebugInfo
+from unique_web_search.utils import StepDebugInfo
 
 _LOGGER = logging.getLogger(__name__)
 
 
-class WebSearchLogEntry(BaseModel):
-    type: StepType
-    message: str
-    web_search_results: list[WebSearchResult]
-
-
-class MessageLogCallback(Protocol):
-    def __call__(
-        self,
-        *,
-        progress_message: str | None = None,
-        queries_for_log: list[WebSearchLogEntry] | list[str] | None = None,
-        status: MessageLogStatus | None = None,
-    ) -> MessageLog | None: ...
-
-
 class BaseWebSearchExecutor(ABC):
     def __init__(
         self,
-        search_service: SearchEngineTypes,
-        language_model_service: LanguageModelService,
-        language_model: LMI,
-        crawler_service: CrawlerTypes,
+        services: ExecutorServiceContext,
+        config: ExecutorConfiguration,
+        callbacks: ExecutorCallbacks,
         tool_call: LanguageModelFunction,
         tool_parameters: WebSearchPlan | WebSearchToolParameters,
-        company_id: str,
-        content_processor: ContentProcessor,
-        message_log_callback: MessageLogCallback,
-        chunk_relevancy_sorter: ChunkRelevancySorter | None,
-        chunk_relevancy_sort_config: ChunkRelevancySortConfig,
-        debug_info: WebSearchDebugInfo,
-        content_reducer: Callable[[list[WebPageChunk]], list[WebPageChunk]],
-        tool_progress_reporter: Optional[ToolProgressReporter] = None,
     ):
-        self.company_id = company_id
-        self.search_service = search_service
-        self.crawler_service = crawler_service
-        self.language_model_service = language_model_service
-        self.tool_progress_reporter = tool_progress_reporter
-        self.language_model = language_model
-        self.content_processor = content_processor
-        self.chunk_relevancy_sorter = chunk_relevancy_sorter
-        self.chunk_relevancy_sort_config = chunk_relevancy_sort_config
+        # Extract from service context
+        self.search_service = services.search_engine_service
+        self.crawler_service = services.crawler_service
+        self.content_processor = services.content_processor
+        self.language_model_service = services.language_model_service
+        self.chunk_relevancy_sorter = services.chunk_relevancy_sorter
+
+        # Extract from configuration
+        self.language_model = config.language_model
+        self.chunk_relevancy_sort_config = config.chunk_relevancy_sort_config
+        self.company_id = config.company_id
+        self.debug_info = config.debug_info
+
+        # Extract from callbacks
+        self._message_log_callback = callbacks.message_log_callback
+        self.content_reducer = callbacks.content_reducer
+        self.query_elicitation = callbacks.query_elicitation
+        self.tool_progress_reporter = callbacks.tool_progress_reporter
+
+        # Store tool parameters
         self.tool_call = tool_call
         self.tool_parameters = tool_parameters
-        self.content_reducer = content_reducer
+
+        # Initialize notification state
         self._notify_name = ""
         self._notify_message = ""
 
-        self.debug_info = debug_info
-        self._message_log_callback = message_log_callback
-
         async def notify_callback() -> None:
             _LOGGER.debug(f"{self.notify_name}: {self.notify_message}")
             if self.tool_progress_reporter:
@@ -111,9 +89,7 @@ class BaseWebSearchExecutor(ABC):
         self._notify_message = value
 
     @abstractmethod
-    async def run(
-        self,
-    ) -> tuple[list[ContentChunk], list[WebSearchLogEntry]]:
+    async def run(self) -> list[ContentChunk]:
         raise NotImplementedError("Subclasses must implement this method.")
 
     async def _content_processing(
@@ -167,3 +143,12 @@ class BaseWebSearchExecutor(ABC):
             return sorted_chunks.content_chunks
         else:
             return content
+
+    async def _ff_elicitate_queries(self, queries: list[str]) -> list[str]:
+        if not feature_flags.enable_elicitation_un_15809.is_enabled(self.company_id):
+            return queries
+
+        # Create a query elicitation
+        elicitation = await self.query_elicitation(queries)
+
+        return elicitation

unique_web_search/services/executors/context.py

@@ -0,0 +1,104 @@
+"""Context classes for executor dependency injection.
+
+This module provides container classes that group related dependencies
+for web search executors, reducing parameter redundancy in __init__ methods.
+"""
+
+from dataclasses import dataclass
+from typing import Optional, Protocol
+
+from unique_toolkit import LanguageModelService
+from unique_toolkit._common.chunk_relevancy_sorter.config import (
+    ChunkRelevancySortConfig,
+)
+from unique_toolkit._common.chunk_relevancy_sorter.service import ChunkRelevancySorter
+from unique_toolkit._common.validators import LMI
+from unique_toolkit.agentic.tools.tool_progress_reporter import (
+    ToolProgressReporter,
+)
+
+from unique_web_search.services.content_processing import ContentProcessor, WebPageChunk
+from unique_web_search.services.crawlers import CrawlerTypes
+from unique_web_search.services.search_engine import SearchEngineTypes
+from unique_web_search.services.search_engine.schema import WebSearchResult
+from unique_web_search.utils import WebSearchDebugInfo
+
+
+@dataclass
+class ExecutorServiceContext:
+    """Container for all service dependencies used by executors.
+
+    Groups together all the external services that executors depend on,
+    making it easier to pass dependencies and reducing parameter redundancy.
+
+    Attributes:
+        search_service: Service for performing web searches
+        crawler_service: Service for crawling and extracting web page content
+        content_processor: Service for processing and analyzing content
+        language_model_service: Service for language model operations
+        chunk_relevancy_sorter: Optional service for sorting content chunks by relevance
+    """
+
+    search_engine_service: SearchEngineTypes
+    crawler_service: CrawlerTypes
+    content_processor: ContentProcessor
+    language_model_service: LanguageModelService
+    chunk_relevancy_sorter: ChunkRelevancySorter | None
+
+
+@dataclass
+class ExecutorConfiguration:
+    """Container for all configuration parameters used by executors.
+
+    Groups together all configuration values and metadata that executors need,
+    providing a clean separation between services and configuration.
+
+    Attributes:
+        language_model: Language model identifier and configuration
+        chunk_relevancy_sort_config: Configuration for chunk relevancy sorting
+        company_id: Identifier for the company/organization
+        debug_info: Container for debug information and metrics
+    """
+
+    language_model: LMI
+    chunk_relevancy_sort_config: ChunkRelevancySortConfig
+    company_id: str
+    debug_info: WebSearchDebugInfo
+
+
+class ContentReducer(Protocol):
+    def __call__(self, web_page_chunks: list[WebPageChunk]) -> list[WebPageChunk]: ...
+
+
+class MessageLogCallback(Protocol):
+    async def log_progress(self, progress_message: str) -> None: ...
+
+    async def log_queries(self, queries: list[str]) -> None: ...
+
+    async def log_web_search_results(
+        self, web_search_results: list[WebSearchResult]
+    ) -> None: ...
+
+
+class QueryElicitationProtocol(Protocol):
+    async def __call__(self, queries: list[str]) -> list[str]: ...
+
+
+@dataclass
+class ExecutorCallbacks:
+    """Container for callback functions used by executors.
+
+    Groups together all callback functions and optional reporters,
+    making the callback interface clear and extensible.
+
+    Attributes:
+        message_log_callback: Callback for logging messages and progress
+        content_reducer: Function to reduce content chunks based on token limits
+        query_elicitation: Function to elicit queries from the user
+        tool_progress_reporter: Optional reporter for tool execution progress
+    """
+
+    message_log_callback: MessageLogCallback
+    content_reducer: ContentReducer
+    query_elicitation: QueryElicitationProtocol
+    tool_progress_reporter: Optional[ToolProgressReporter] = None