unique_toolkit 1.38.3__py3-none-any.whl → 1.39.0__py3-none-any.whl

unique_toolkit/agentic/loop_runner/__init__.py

@@ -1,25 +1,35 @@
+from unique_toolkit.agentic.loop_runner._iteration_handler_utils import (
+    handle_forced_tools_iteration,
+    handle_last_iteration,
+    handle_normal_iteration,
+)
 from unique_toolkit.agentic.loop_runner.base import LoopIterationRunner
 from unique_toolkit.agentic.loop_runner.middleware import (
-    QWEN_FORCED_TOOL_CALL_PROMPT_INSTRUCTION,
     PlanningConfig,
     PlanningMiddleware,
     PlanningSchemaConfig,
-    QwenForcedToolCallMiddleware,
-    is_qwen_model,
 )
 from unique_toolkit.agentic.loop_runner.runners import (
+    QWEN_FORCED_TOOL_CALL_INSTRUCTION,
+    QWEN_LAST_ITERATION_INSTRUCTION,
     BasicLoopIterationRunner,
     BasicLoopIterationRunnerConfig,
+    QwenLoopIterationRunner,
+    is_qwen_model,
 )
 
 __all__ = [
     "LoopIterationRunner",
-    "QWEN_FORCED_TOOL_CALL_PROMPT_INSTRUCTION",
     "PlanningConfig",
     "PlanningMiddleware",
     "PlanningSchemaConfig",
-    "QwenForcedToolCallMiddleware",
+    "QwenLoopIterationRunner",
     "is_qwen_model",
     "BasicLoopIterationRunnerConfig",
     "BasicLoopIterationRunner",
+    "handle_forced_tools_iteration",
+    "handle_last_iteration",
+    "handle_normal_iteration",
+    "QWEN_FORCED_TOOL_CALL_INSTRUCTION",
+    "QWEN_LAST_ITERATION_INSTRUCTION",
 ]

unique_toolkit/agentic/loop_runner/_iteration_handler_utils.py

@@ -0,0 +1,95 @@
+import logging
+from typing import Protocol, Unpack, cast
+
+from unique_toolkit.agentic.loop_runner._stream_handler_utils import stream_response
+from unique_toolkit.agentic.loop_runner.base import (
+    _LoopIterationRunnerKwargs,
+)
+from unique_toolkit.chat.functions import LanguageModelStreamResponse
+
+_LOGGER = logging.getLogger(__name__)
+
+
+class PrepareForcedToolIterationKwargs(Protocol):
+    def __call__(
+        self,
+        func_name: str | None,
+        per_choice_kwargs: _LoopIterationRunnerKwargs,
+    ) -> _LoopIterationRunnerKwargs: ...
+
+
+async def handle_last_iteration(
+    **kwargs: Unpack[_LoopIterationRunnerKwargs],
+) -> LanguageModelStreamResponse:
+    _LOGGER.info("Reached last iteration, removing tools and producing final response")
+
+    return await stream_response(
+        loop_runner_kwargs=kwargs,
+        tools=None,
+    )
+
+
+async def handle_normal_iteration(
+    **kwargs: Unpack[_LoopIterationRunnerKwargs],
+) -> LanguageModelStreamResponse:
+    _LOGGER.info("Running loop iteration %d", kwargs["iteration_index"])
+
+    return await stream_response(loop_runner_kwargs=kwargs)
+
+
+async def handle_forced_tools_iteration(
+    **kwargs: Unpack[_LoopIterationRunnerKwargs],
+) -> LanguageModelStreamResponse:
+    return await run_forced_tools_iteration(loop_runner_kwargs=kwargs)
+
+
+async def run_forced_tools_iteration(
+    *,
+    loop_runner_kwargs: _LoopIterationRunnerKwargs,
+    prepare_loop_runner_kwargs: PrepareForcedToolIterationKwargs | 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).
+    """
+    assert "tool_choices" in loop_runner_kwargs
+
+    tool_choices = loop_runner_kwargs["tool_choices"]
+    assert len(tool_choices) > 0, (
+        "run_forced_tools_iteration requires at least one tool choice"
+    )
+    _LOGGER.info("Forcing tools calls: %s", tool_choices)
+
+    responses: list[LanguageModelStreamResponse] = []
+
+    available_tools = {t.name: t for t in loop_runner_kwargs.get("tools") or []}
+
+    for opt in tool_choices:
+        func_name = opt.get("function", {}).get("name")
+
+        per_choice_kwargs = cast(_LoopIterationRunnerKwargs, dict(loop_runner_kwargs))
+        if prepare_loop_runner_kwargs:
+            per_choice_kwargs = prepare_loop_runner_kwargs(func_name, per_choice_kwargs)
+
+        limited_tool = available_tools.get(func_name) if func_name else None
+        stream_kwargs = {"loop_runner_kwargs": per_choice_kwargs, "tool_choice": opt}
+        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:
+        if r.tool_calls:
+            tool_calls.extend(r.tool_calls)
+        references.extend(r.message.references)
+
+    response = responses[0]
+    response.tool_calls = tool_calls if len(tool_calls) > 0 else None
+    response.message.references = references
+
+    return response

unique_toolkit/agentic/loop_runner/middleware/__init__.py

@@ -3,17 +3,9 @@ from unique_toolkit.agentic.loop_runner.middleware.planning import (
     PlanningMiddleware,
     PlanningSchemaConfig,
 )
-from unique_toolkit.agentic.loop_runner.middleware.qwen_forced_tool_call import (
-    QWEN_FORCED_TOOL_CALL_PROMPT_INSTRUCTION,
-    QwenForcedToolCallMiddleware,
-    is_qwen_model,
-)
 
 __all__ = [
     "PlanningConfig",
     "PlanningMiddleware",
     "PlanningSchemaConfig",
-    "QWEN_FORCED_TOOL_CALL_PROMPT_INSTRUCTION",
-    "QwenForcedToolCallMiddleware",
-    "is_qwen_model",
 ]

unique_toolkit/agentic/loop_runner/runners/__init__.py

@@ -2,5 +2,18 @@ from unique_toolkit.agentic.loop_runner.runners.basic import (
     BasicLoopIterationRunner,
     BasicLoopIterationRunnerConfig,
 )
+from unique_toolkit.agentic.loop_runner.runners.qwen import (
+    QWEN_FORCED_TOOL_CALL_INSTRUCTION,
+    QWEN_LAST_ITERATION_INSTRUCTION,
+    QwenLoopIterationRunner,
+    is_qwen_model,
+)
 
-__all__ = ["BasicLoopIterationRunnerConfig", "BasicLoopIterationRunner"]
+__all__ = [
+    "BasicLoopIterationRunnerConfig",
+    "BasicLoopIterationRunner",
+    "QwenLoopIterationRunner",
+    "QWEN_FORCED_TOOL_CALL_INSTRUCTION",
+    "QWEN_LAST_ITERATION_INSTRUCTION",
+    "is_qwen_model",
+]

unique_toolkit/agentic/loop_runner/runners/basic.py

@@ -4,7 +4,11 @@ from typing import Unpack, override
 from pydantic import BaseModel
 
 from unique_toolkit._common.pydantic_helpers import get_configuration_dict
-from unique_toolkit.agentic.loop_runner._stream_handler_utils import stream_response
+from unique_toolkit.agentic.loop_runner._iteration_handler_utils import (
+    handle_forced_tools_iteration,
+    handle_last_iteration,
+    handle_normal_iteration,
+)
 from unique_toolkit.agentic.loop_runner.base import (
     LoopIterationRunner,
     _LoopIterationRunnerKwargs,
@@ -26,60 +30,6 @@ class BasicLoopIterationRunner(LoopIterationRunner):
     def __init__(self, config: BasicLoopIterationRunnerConfig) -> None:
         self._config = config
 
-    async def _handle_last_iteration(
-        self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
-    ) -> LanguageModelStreamResponse:
-        _LOGGER.info(
-            "Reached last iteration, removing tools and producing final response"
-        )
-
-        return await stream_response(
-            loop_runner_kwargs=kwargs,
-            tools=None,
-        )
-
-    async def _handle_normal_iteration(
-        self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
-    ) -> LanguageModelStreamResponse:
-        _LOGGER.info("Running loop iteration %d", kwargs["iteration_index"])
-
-        return await stream_response(loop_runner_kwargs=kwargs)
-
-    async def _handle_forced_tools_iteration(
-        self,
-        **kwargs: Unpack[_LoopIterationRunnerKwargs],
-    ) -> LanguageModelStreamResponse:
-        assert "tool_choices" in kwargs
-
-        tool_choices = kwargs["tool_choices"]
-        _LOGGER.info("Forcing tools calls: %s", tool_choices)
-
-        responses: list[LanguageModelStreamResponse] = []
-
-        available_tools = {t.name: t for t in kwargs.get("tools") or []}
-
-        for opt in tool_choices:
-            func_name = opt.get("function", {}).get("name")
-            limited_tool = available_tools.get(func_name) if func_name else None
-            stream_kwargs = {"loop_runner_kwargs": kwargs, "tool_choice": opt}
-            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:
-            if r.tool_calls:
-                tool_calls.extend(r.tool_calls)
-            references.extend(r.message.references)
-
-        response = responses[0]
-        response.tool_calls = tool_calls if len(tool_calls) > 0 else None
-        response.message.references = references
-
-        return response
-
     @override
     async def __call__(
         self,
@@ -89,8 +39,8 @@ class BasicLoopIterationRunner(LoopIterationRunner):
         iteration_index = kwargs["iteration_index"]
 
         if len(tool_choices) > 0 and iteration_index == 0:
-            return await self._handle_forced_tools_iteration(**kwargs)
+            return await handle_forced_tools_iteration(**kwargs)
         elif iteration_index == self._config.max_loop_iterations - 1:
-            return await self._handle_last_iteration(**kwargs)
+            return await handle_last_iteration(**kwargs)
         else:
-            return await self._handle_normal_iteration(**kwargs)
+            return await handle_normal_iteration(**kwargs)

unique_toolkit/agentic/loop_runner/runners/qwen/__init__.py

@@ -0,0 +1,15 @@
+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,
+    QwenLoopIterationRunner,
+)
+
+__all__ = [
+    "QwenLoopIterationRunner",
+    "is_qwen_model",
+    "QWEN_FORCED_TOOL_CALL_INSTRUCTION",
+    "QWEN_LAST_ITERATION_INSTRUCTION",
+]

unique_toolkit/agentic/loop_runner/{middleware/qwen_forced_tool_call → runners/qwen}/helpers.py

@@ -1,5 +1,6 @@
 from unique_toolkit.language_model.infos import LanguageModelInfo
 from unique_toolkit.language_model.schemas import (
+    LanguageModelAssistantMessage,
     LanguageModelMessageRole,
     LanguageModelMessages,
 )
@@ -31,3 +32,17 @@ def append_qwen_forced_tool_call_instruction(
             )
             break
     return LanguageModelMessages(root=messages_list)
+
+
+def append_qwen_last_iteration_assistant_message(
+    *,
+    messages: LanguageModelMessages,
+    last_iteration_instruction: str,
+) -> LanguageModelMessages:
+    """Append an assistant message at the end to indicate no further tool calls are allowed."""
+    messages_list = list(messages)
+    assistant_message = LanguageModelAssistantMessage(
+        content=last_iteration_instruction,
+    )
+    messages_list.append(assistant_message)
+    return LanguageModelMessages(root=messages_list)

unique_toolkit/agentic/loop_runner/runners/qwen/qwen_runner.py

@@ -0,0 +1,118 @@
+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.runners.qwen.helpers import (
+    append_qwen_forced_tool_call_instruction,
+    append_qwen_last_iteration_assistant_message,
+)
+from unique_toolkit.chat.service import ChatService, LanguageModelStreamResponse
+
+_LOGGER = logging.getLogger(__name__)
+
+QWEN_FORCED_TOOL_CALL_INSTRUCTION = (
+    "**Tool Call Instruction:** \nYou MUST call the tool {TOOL_NAME}. "
+    "You must start the response with <tool_call>. "
+    "Do NOT provide natural language explanations, summaries, or any text outside the <tool_call> block."
+)
+
+QWEN_LAST_ITERATION_INSTRUCTION = "The maximum number of loop iteration have been reached. Not further tool calls are allowed. Based on the found information, an answer should be generated"
+
+
+class QwenLoopIterationRunner(LoopIterationRunner):
+    def __init__(
+        self,
+        *,
+        qwen_forced_tool_call_instruction: str,
+        qwen_last_iteration_instruction: str,
+        max_loop_iterations: int,
+        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
+        self._chat_service = chat_service
+
+    async def __call__(
+        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(
+                TOOL_NAME=func_name or ""
+            )
+            per_choice_kwargs["messages"] = append_qwen_forced_tool_call_instruction(
+                messages=original_messages,
+                forced_tool_call_instruction=prompt_instruction,
+            )
+            return per_choice_kwargs
+
+        response = await run_forced_tools_iteration(
+            loop_runner_kwargs=kwargs,
+            prepare_loop_runner_kwargs=_prepare,
+        )
+        return self._process_response(response)
+
+    async def _qwen_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,
+        )
+
+        response = await handle_last_iteration(
+            **kwargs,
+        )
+
+        return self._process_response(response)
+
+    async def _qwen_handle_normal_iteration(
+        self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
+    ) -> LanguageModelStreamResponse:
+        response = await 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."
+            )
+            self._chat_service.modify_assistant_message(content="")
+            response.message.text = ""
+
+        return response

unique_toolkit/app/schemas.py

@@ -120,11 +120,15 @@ class McpServer(BaseModel):
 class ChatEventUserMessage(BaseModel):
     model_config = model_config
 
-    id: str
-    text: str
-    original_text: str
-    created_at: str
-    language: str
+    id: str = Field(
+        description="The id of the user message. On an event this corresponds to the user message that created the event."
+    )
+    text: str = Field(description="The text of the user message.")
+    original_text: str = Field(description="The original text of the user message.")
+    created_at: str = Field(
+        description="The creation date and time of the user message."
+    )
+    language: str = Field(description="The language of the user message.")
 
 
 @deprecated(
@@ -140,8 +144,12 @@ class EventUserMessage(ChatEventUserMessage):
 class ChatEventAssistantMessage(BaseModel):
     model_config = model_config
 
-    id: str
-    created_at: str
+    id: str = Field(
+        description="The id of the assistant message. On an event this corresponds to the assistant message that will be returned by the process handling the event."
+    )
+    created_at: str = Field(
+        description="The creation date and time of the assistant message."
+    )
 
 
 @deprecated(