unique_toolkit 1.50.4__tar.gz → 1.52.0__tar.gz

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: unique-toolkit
-Version: 1.50.4
+Version: 1.52.0
 Summary: 
 Author: Cedric Klinkert, Martin Fadler, Sadique Sheik, Fabian Schläpfer, Pascal Hauri, Ejiroghene Oki-Peter, Rami Azouz, Ricardo Visini, Sanzio Monti
 Author-email: Cedric Klinkert <cedric.klinkert@unique.ch>, Martin Fadler <martin.fadler@unique.ch>, Sadique Sheik <sadique@unique.ch>, Fabian Schläpfer <fabian@unique.ch>, Pascal Hauri <pascal@unique.ch>, Ejiroghene Oki-Peter <oki-peter.ejiroghene@unique.ai>, Rami Azouz <rami.ext@unique.ch>, Ricardo Visini <ricardo.visini@unique.ch>, Sanzio Monti <sanzio.monti@unique.ai>

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "unique_toolkit"
-version = "1.50.4"
+version = "1.52.0"
 description = ""
 readme = "README.md"
 requires-python = ">=3.12"

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/unique_toolkit/agentic/loop_runner/__init__.py

@@ -18,9 +18,9 @@ from unique_toolkit.agentic.loop_runner.runners import (
     QWEN_MAX_LOOP_ITERATIONS,
     BasicLoopIterationRunner,
     BasicLoopIterationRunnerConfig,
+    MistralLoopIterationRunner,
     QwenLoopIterationRunner,
     ResponsesBasicLoopIterationRunner,
-    is_qwen_model,
 )
 
 __all__ = [
@@ -28,10 +28,10 @@ __all__ = [
     "PlanningConfig",
     "PlanningMiddleware",
     "PlanningSchemaConfig",
-    "QwenLoopIterationRunner",
-    "is_qwen_model",
     "BasicLoopIterationRunnerConfig",
     "BasicLoopIterationRunner",
+    "MistralLoopIterationRunner",
+    "QwenLoopIterationRunner",
     "handle_forced_tools_iteration",
     "handle_last_iteration",
     "handle_normal_iteration",

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/unique_toolkit/agentic/loop_runner/_iteration_handler_utils.py

@@ -47,13 +47,17 @@ async def run_forced_tools_iteration(
     *,
     loop_runner_kwargs: _LoopIterationRunnerKwargs,
     prepare_loop_runner_kwargs: PrepareForcedToolIterationKwargs | None = None,
+    tool_choice_override: str | None = None,
 ) -> LanguageModelStreamResponse:
     """
     Execute a "forced tools" iteration by running one stream_response per tool choice,
     then merging tool calls and references into a single response.
 
-    Some models (e.g. Qwen) need per-tool-choice message rewriting; this can be done
-    via prepare_loop_runner_kwargs (called once per tool choice, on a copy of kwargs).
+    Args:
+        prepare_loop_runner_kwargs: Optional callback to transform kwargs per tool choice
+            (e.g. Qwen rewrites messages with tool-specific instructions).
+        tool_choice_override: If set, replaces the per-tool named dict with this value
+            (e.g. Mistral needs "any" instead of the named format).
     """
     assert "tool_choices" in loop_runner_kwargs
 
@@ -74,13 +78,16 @@ async def run_forced_tools_iteration(
         if prepare_loop_runner_kwargs:
             per_choice_kwargs = prepare_loop_runner_kwargs(func_name, per_choice_kwargs)
 
+        effective_tool_choice = tool_choice_override or opt
         limited_tool = available_tools.get(func_name) if func_name else None
-        stream_kwargs = {"loop_runner_kwargs": per_choice_kwargs, "tool_choice": opt}
+        stream_kwargs = {
+            "loop_runner_kwargs": per_choice_kwargs,
+            "tool_choice": effective_tool_choice,
+        }
         if limited_tool:
             stream_kwargs["tools"] = [limited_tool]
         responses.append(await stream_response(**stream_kwargs))
 
-    # Merge responses and refs:
     tool_calls = []
     references = []
     for r in responses:

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/unique_toolkit/agentic/loop_runner/_stream_handler_utils.py

@@ -1,6 +1,6 @@
 from typing import Any, Required
 
-from openai.types.chat import ChatCompletionNamedToolChoiceParam
+from openai.types.chat import ChatCompletionToolChoiceOptionParam
 from typing_extensions import TypedDict
 
 from unique_toolkit import LanguageModelToolDescription
@@ -18,7 +18,7 @@ class _StreamingHandlerKwargs(TypedDict, total=False):
     start_text: str
     debug_info: dict[str, Any]
     temperature: float
-    tool_choice: ChatCompletionNamedToolChoiceParam
+    tool_choice: ChatCompletionToolChoiceOptionParam
     other_options: dict[str, Any]
 
 

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/unique_toolkit/agentic/loop_runner/runners/__init__.py

@@ -3,21 +3,23 @@ from unique_toolkit.agentic.loop_runner.runners.basic import (
     BasicLoopIterationRunnerConfig,
     ResponsesBasicLoopIterationRunner,
 )
+from unique_toolkit.agentic.loop_runner.runners.mistral import (
+    MistralLoopIterationRunner,
+)
 from unique_toolkit.agentic.loop_runner.runners.qwen import (
     QWEN_FORCED_TOOL_CALL_INSTRUCTION,
     QWEN_LAST_ITERATION_INSTRUCTION,
     QWEN_MAX_LOOP_ITERATIONS,
     QwenLoopIterationRunner,
-    is_qwen_model,
 )
 
 __all__ = [
     "BasicLoopIterationRunnerConfig",
     "BasicLoopIterationRunner",
     "ResponsesBasicLoopIterationRunner",
+    "MistralLoopIterationRunner",
     "QwenLoopIterationRunner",
     "QWEN_FORCED_TOOL_CALL_INSTRUCTION",
     "QWEN_LAST_ITERATION_INSTRUCTION",
     "QWEN_MAX_LOOP_ITERATIONS",
-    "is_qwen_model",
 ]

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/unique_toolkit/agentic/loop_runner/runners/basic.py

@@ -5,9 +5,9 @@ from pydantic import BaseModel
 
 from unique_toolkit._common.pydantic_helpers import get_configuration_dict
 from unique_toolkit.agentic.loop_runner._iteration_handler_utils import (
-    handle_forced_tools_iteration,
     handle_last_iteration,
     handle_normal_iteration,
+    run_forced_tools_iteration,
 )
 from unique_toolkit.agentic.loop_runner._responses_iteration_handler_utils import (
     handle_responses_forced_tools_iteration,
@@ -44,11 +44,29 @@ class BasicLoopIterationRunner(LoopIterationRunner):
         iteration_index = kwargs["iteration_index"]
 
         if len(tool_choices) > 0 and iteration_index == 0:
-            return await handle_forced_tools_iteration(**kwargs)
+            return await self._handle_forced_tools(**kwargs)
         elif iteration_index == self._config.max_loop_iterations - 1:
-            return await handle_last_iteration(**kwargs)
+            return await self._handle_last_iteration(**kwargs)
         else:
-            return await handle_normal_iteration(**kwargs)
+            return await self._handle_normal_iteration(**kwargs)
+
+    async def _handle_forced_tools(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse:
+        return await run_forced_tools_iteration(loop_runner_kwargs=kwargs)
+
+    async def _handle_last_iteration(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse:
+        return await handle_last_iteration(**kwargs)
+
+    async def _handle_normal_iteration(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse:
+        return await handle_normal_iteration(**kwargs)
 
 
 class ResponsesBasicLoopIterationRunner(ResponsesLoopIterationRunner):

unique_toolkit-1.52.0/unique_toolkit/agentic/loop_runner/runners/mistral/__init__.py

@@ -0,0 +1,7 @@
+from unique_toolkit.agentic.loop_runner.runners.mistral.mistral_runner import (
+    MistralLoopIterationRunner,
+)
+
+__all__ = [
+    "MistralLoopIterationRunner",
+]

unique_toolkit-1.52.0/unique_toolkit/agentic/loop_runner/runners/mistral/mistral_runner.py

@@ -0,0 +1,30 @@
+from typing import Unpack
+
+from unique_toolkit.agentic.loop_runner._iteration_handler_utils import (
+    run_forced_tools_iteration,
+)
+from unique_toolkit.agentic.loop_runner.base import _LoopIterationRunnerKwargs
+from unique_toolkit.agentic.loop_runner.runners.basic import (
+    BasicLoopIterationRunner,
+)
+from unique_toolkit.chat.functions import LanguageModelStreamResponse
+
+
+class MistralLoopIterationRunner(BasicLoopIterationRunner):
+    """Runner for Mistral models.
+
+    Mistral doesn't support named tool_choice (``{"type": "function", ...}``).
+    This runner overrides forced-tool handling to send ``"any"`` instead,
+    which forces the model to call a tool. Since each forced iteration already
+    restricts the tools list to a single tool, ``"any"`` effectively forces
+    the correct one.
+    """
+
+    async def _handle_forced_tools(
+        self,
+        **kwargs: Unpack[_LoopIterationRunnerKwargs],
+    ) -> LanguageModelStreamResponse:
+        return await run_forced_tools_iteration(
+            loop_runner_kwargs=kwargs,
+            tool_choice_override="any",
+        )

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/unique_toolkit/agentic/loop_runner/runners/qwen/__init__.py

@@ -1,6 +1,3 @@
-from unique_toolkit.agentic.loop_runner.runners.qwen.helpers import (
-    is_qwen_model,
-)
 from unique_toolkit.agentic.loop_runner.runners.qwen.qwen_runner import (
     QWEN_FORCED_TOOL_CALL_INSTRUCTION,
     QWEN_LAST_ITERATION_INSTRUCTION,
@@ -10,7 +7,6 @@ from unique_toolkit.agentic.loop_runner.runners.qwen.qwen_runner import (
 
 __all__ = [
     "QwenLoopIterationRunner",
-    "is_qwen_model",
     "QWEN_FORCED_TOOL_CALL_INSTRUCTION",
     "QWEN_LAST_ITERATION_INSTRUCTION",
     "QWEN_MAX_LOOP_ITERATIONS",

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/unique_toolkit/agentic/loop_runner/runners/qwen/helpers.py

@@ -1,4 +1,3 @@
-from unique_toolkit.language_model.infos import LanguageModelInfo
 from unique_toolkit.language_model.schemas import (
     LanguageModelAssistantMessage,
     LanguageModelMessageRole,
@@ -6,17 +5,6 @@ from unique_toolkit.language_model.schemas import (
 )
 
 
-def is_qwen_model(*, model: str | LanguageModelInfo | None) -> bool:
-    """Check if the model is a Qwen model."""
-    if isinstance(model, LanguageModelInfo):
-        name = model.name
-        # name is an Enum with a .value attribute
-        return "qwen" in str(getattr(name, "value", name)).lower()
-    elif isinstance(model, str):
-        return "qwen" in model.lower()
-    return False
-
-
 def append_qwen_forced_tool_call_instruction(
     *,
     messages: LanguageModelMessages,

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/unique_toolkit/agentic/loop_runner/runners/qwen/qwen_runner.py

@@ -2,13 +2,12 @@ import logging
 from typing import Unpack
 
 from unique_toolkit.agentic.loop_runner._iteration_handler_utils import (
-    handle_last_iteration,
-    handle_normal_iteration,
     run_forced_tools_iteration,
 )
-from unique_toolkit.agentic.loop_runner.base import (
-    LoopIterationRunner,
-    _LoopIterationRunnerKwargs,
+from unique_toolkit.agentic.loop_runner.base import _LoopIterationRunnerKwargs
+from unique_toolkit.agentic.loop_runner.runners.basic import (
+    BasicLoopIterationRunner,
+    BasicLoopIterationRunnerConfig,
 )
 from unique_toolkit.agentic.loop_runner.runners.qwen.helpers import (
     append_qwen_forced_tool_call_instruction,
@@ -30,45 +29,30 @@ QWEN_LAST_ITERATION_INSTRUCTION = """The maximum number of loop iteration have b
 Based on the found information, an answer should be generated"""
 
 
-class QwenLoopIterationRunner(LoopIterationRunner):
+class QwenLoopIterationRunner(BasicLoopIterationRunner):
     def __init__(
         self,
         *,
-        qwen_forced_tool_call_instruction: str,
-        qwen_last_iteration_instruction: str,
-        max_loop_iterations: int,
+        config: BasicLoopIterationRunnerConfig,
+        forced_tool_call_instruction: str,
+        last_iteration_instruction: str,
         chat_service: ChatService,
     ) -> None:
-        self._qwen_forced_tool_call_instruction = qwen_forced_tool_call_instruction
-        self._qwen_last_iteration_instruction = qwen_last_iteration_instruction
-        self._max_loop_iterations = max_loop_iterations
+        super().__init__(config)
+        self._forced_tool_call_instruction = forced_tool_call_instruction
+        self._last_iteration_instruction = last_iteration_instruction
         self._chat_service = chat_service
 
-    async def __call__(
+    async def _handle_forced_tools(
         self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
     ) -> LanguageModelStreamResponse:
-        tool_choices = kwargs.get("tool_choices") or []
-        iteration_index = kwargs["iteration_index"]
-
-        if len(tool_choices) > 0 and iteration_index == 0:
-            return await self._qwen_handle_forced_tools_iteration(**kwargs)
-        elif iteration_index == self._max_loop_iterations - 1:
-            return await self._qwen_handle_last_iteration(**kwargs)
-        else:
-            return await self._qwen_handle_normal_iteration(**kwargs)
-
-    async def _qwen_handle_forced_tools_iteration(
-        self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
-    ) -> LanguageModelStreamResponse:
-        # For Qwen models, append tool call instruction to the last user message. These models ignore the parameter tool_choice.
-        # As the message has to be modified for each tool call instruction, the function from the basic runner cant be used.
         original_messages = kwargs["messages"].model_copy(deep=True)
 
         def _prepare(
             func_name: str | None,
             per_choice_kwargs: _LoopIterationRunnerKwargs,
         ) -> _LoopIterationRunnerKwargs:
-            prompt_instruction = self._qwen_forced_tool_call_instruction.format(
+            prompt_instruction = self._forced_tool_call_instruction.format(
                 TOOL_NAME=func_name or ""
             )
             per_choice_kwargs["messages"] = append_qwen_forced_tool_call_instruction(
@@ -83,34 +67,28 @@ class QwenLoopIterationRunner(LoopIterationRunner):
         )
         return self._process_response(response)
 
-    async def _qwen_handle_last_iteration(
+    async def _handle_last_iteration(
         self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
     ) -> LanguageModelStreamResponse:
-        # For Qwen models, append an assistant message with instructions to not call any tool in this iteration.
         _LOGGER.info(
             "Reached last iteration, removing tools. Appending assistant message with instructions to not call any tool in this iteration."
         )
         kwargs["messages"] = append_qwen_last_iteration_assistant_message(
             messages=kwargs["messages"],
-            last_iteration_instruction=self._qwen_last_iteration_instruction,
+            last_iteration_instruction=self._last_iteration_instruction,
         )
-
-        response = await handle_last_iteration(
-            **kwargs,
-        )
-
+        response = await super()._handle_last_iteration(**kwargs)
         return self._process_response(response)
 
-    async def _qwen_handle_normal_iteration(
+    async def _handle_normal_iteration(
         self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
     ) -> LanguageModelStreamResponse:
-        response = await handle_normal_iteration(**kwargs)
+        response = await super()._handle_normal_iteration(**kwargs)
         return self._process_response(response)
 
     def _process_response(
         self, response: LanguageModelStreamResponse
     ) -> LanguageModelStreamResponse:
-        # Check if content of response is </tool_call>
         if "</tool_call>" == response.message.text.strip():
             _LOGGER.warning(
                 "Response contains only <tool_call>. This is not allowed. Returning empty response."

{unique_toolkit-1.50.4 → unique_toolkit-1.52.0}/unique_toolkit/agentic/tools/config.py

@@ -1,5 +1,5 @@
 from enum import StrEnum
-from typing import Annotated, Any
+from typing import Annotated, Any, Generic
 
 from pydantic import (
     BaseModel,
@@ -9,10 +9,13 @@ from pydantic import (
     field_serializer,
     model_validator,
 )
+from typing_extensions import TypeVar
 
 from unique_toolkit._common.pydantic_helpers import get_configuration_dict
 from unique_toolkit.agentic.tools.schemas import BaseToolConfig
 
+T = TypeVar("T", bound=BaseToolConfig, default=BaseToolConfig)
+
 
 class ToolIcon(StrEnum):
     ANALYTICS = "IconAnalytics"
@@ -45,12 +48,12 @@ def handle_undefined_icon(value: Any) -> ToolIcon:
         return ToolIcon.BOOK
 
 
-class ToolBuildConfig(BaseModel):
+class ToolBuildConfig(BaseModel, Generic[T]):
     model_config = get_configuration_dict()
     """Main tool configuration"""
 
     name: str
-    configuration: BaseToolConfig
+    configuration: T
     display_name: str = ""
     icon: Annotated[ToolIcon, BeforeValidator(handle_undefined_icon)] = Field(
         default=ToolIcon.BOOK,
@@ -124,5 +127,5 @@ class ToolBuildConfig(BaseModel):
         return value
 
     @field_serializer("configuration")
-    def serialize_config(self, value: BaseToolConfig) -> dict[str, Any]:
+    def serialize_config(self, value: T) -> dict[str, Any]:
         return value.__class__.model_dump(value)