unique_toolkit 1.35.4__py3-none-any.whl → 1.37.0__py3-none-any.whl

unique_toolkit/_common/referencing.py

@@ -15,7 +15,7 @@ def _iter_ref_numbers(text: str) -> Generator[int, None, None]:
 
 
 @functools.cache
-def _get_detection_pattern_for_ref(ref_number: int) -> re.Pattern[str]:
+def get_detection_pattern_for_ref(ref_number: int) -> re.Pattern[str]:
     return re.compile(rf"<sup>\s*{ref_number}\s*</sup>")
 
 
@@ -35,11 +35,11 @@ def replace_ref_number(text: str, ref_number: int, replacement: int | str) -> st
     if isinstance(replacement, int):
         replacement = get_reference_pattern(replacement)
 
-    return _get_detection_pattern_for_ref(ref_number).sub(replacement, text)
+    return get_detection_pattern_for_ref(ref_number).sub(replacement, text)
 
 
 def remove_ref_number(text: str, ref_number: int) -> str:
-    return _get_detection_pattern_for_ref(ref_number).sub("", text)
+    return get_detection_pattern_for_ref(ref_number).sub("", text)
 
 
 def remove_all_refs(text: str) -> str:

unique_toolkit/agentic/loop_runner/__init__.py

@@ -1,8 +1,11 @@
 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 (
     BasicLoopIterationRunner,
@@ -11,9 +14,12 @@ from unique_toolkit.agentic.loop_runner.runners import (
 
 __all__ = [
     "LoopIterationRunner",
+    "QWEN_FORCED_TOOL_CALL_PROMPT_INSTRUCTION",
     "PlanningConfig",
     "PlanningMiddleware",
     "PlanningSchemaConfig",
+    "QwenForcedToolCallMiddleware",
+    "is_qwen_model",
     "BasicLoopIterationRunnerConfig",
     "BasicLoopIterationRunner",
 ]

unique_toolkit/agentic/loop_runner/middleware/__init__.py

@@ -3,5 +3,17 @@ 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"]
+__all__ = [
+    "PlanningConfig",
+    "PlanningMiddleware",
+    "PlanningSchemaConfig",
+    "QWEN_FORCED_TOOL_CALL_PROMPT_INSTRUCTION",
+    "QwenForcedToolCallMiddleware",
+    "is_qwen_model",
+]

unique_toolkit/agentic/loop_runner/middleware/qwen_forced_tool_call/__init__.py

@@ -0,0 +1,13 @@
+from unique_toolkit.agentic.loop_runner.middleware.qwen_forced_tool_call.helpers import (
+    is_qwen_model,
+)
+from unique_toolkit.agentic.loop_runner.middleware.qwen_forced_tool_call.qwen_forced_tool_call import (
+    QWEN_FORCED_TOOL_CALL_PROMPT_INSTRUCTION,
+    QwenForcedToolCallMiddleware,
+)
+
+__all__ = [
+    "QwenForcedToolCallMiddleware",
+    "QWEN_FORCED_TOOL_CALL_PROMPT_INSTRUCTION",
+    "is_qwen_model",
+]

unique_toolkit/agentic/loop_runner/middleware/qwen_forced_tool_call/helpers.py

@@ -0,0 +1,33 @@
+from unique_toolkit.language_model.infos import LanguageModelInfo
+from unique_toolkit.language_model.schemas import (
+    LanguageModelMessageRole,
+    LanguageModelMessages,
+)
+
+
+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,
+    forced_tool_call_instruction: str,
+) -> LanguageModelMessages:
+    """Append tool call instruction to the last user message for Qwen models."""
+    messages_list = list(messages)
+    for i in range(len(messages_list) - 1, -1, -1):
+        msg = messages_list[i]
+        if msg.role == LanguageModelMessageRole.USER and isinstance(msg.content, str):
+            messages_list[i] = msg.model_copy(
+                update={"content": msg.content + "\n" + forced_tool_call_instruction}
+            )
+            break
+    return LanguageModelMessages(root=messages_list)

unique_toolkit/agentic/loop_runner/middleware/qwen_forced_tool_call/qwen_forced_tool_call.py

@@ -0,0 +1,50 @@
+import logging
+from typing import Unpack
+
+from unique_toolkit.agentic.loop_runner.base import (
+    LoopIterationRunner,
+    _LoopIterationRunnerKwargs,
+)
+from unique_toolkit.agentic.loop_runner.middleware.qwen_forced_tool_call.helpers import (
+    append_qwen_forced_tool_call_instruction,
+)
+from unique_toolkit.chat.service import LanguageModelStreamResponse
+
+_LOGGER = logging.getLogger(__name__)
+
+QWEN_FORCED_TOOL_CALL_PROMPT_INSTRUCTION = (
+    "Tool Call Instruction: \nYou always have to return a tool call. "
+    "You must start the response with <tool_call> and end with </tool_call>. "
+    "Do NOT provide natural language explanations, summaries, or any text outside the <tool_call> block."
+)
+
+
+class QwenForcedToolCallMiddleware(LoopIterationRunner):
+    def __init__(
+        self,
+        *,
+        loop_runner: LoopIterationRunner,
+        qwen_forced_tool_call_prompt_instruction: str,
+    ) -> None:
+        self._qwen_forced_tool_call_prompt_instruction = (
+            qwen_forced_tool_call_prompt_instruction
+        )
+        self._loop_runner = loop_runner
+
+    async def __call__(
+        self, **kwargs: Unpack[_LoopIterationRunnerKwargs]
+    ) -> LanguageModelStreamResponse:
+        tool_choices = kwargs.get("tool_choices") or []
+        iteration_index = kwargs["iteration_index"]
+
+        # For Qwen models, append tool call instruction to the last user message. These models ignore the parameter tool_choice.
+        if len(tool_choices) > 0 and iteration_index == 0 and kwargs.get("messages"):
+            _LOGGER.info(
+                "Appending tool call instruction to the last user message for Qwen models to force tool calls."
+            )
+            kwargs["messages"] = append_qwen_forced_tool_call_instruction(
+                messages=kwargs["messages"],
+                forced_tool_call_instruction=self._qwen_forced_tool_call_prompt_instruction,
+            )
+
+        return await self._loop_runner(**kwargs)

unique_toolkit/agentic/loop_runner/runners/basic.py

@@ -56,13 +56,15 @@ class BasicLoopIterationRunner(LoopIterationRunner):
 
         responses: list[LanguageModelStreamResponse] = []
 
+        available_tools = {t.name: t for t in kwargs.get("tools") or []}
+
         for opt in tool_choices:
-            responses.append(
-                await stream_response(
-                    loop_runner_kwargs=kwargs,
-                    tool_choice=opt,
-                )
-            )
+            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 = []

unique_toolkit/agentic/tools/a2a/postprocessing/config.py

@@ -65,7 +65,7 @@ class SubAgentDisplayConfig(BaseModel):
     )
     force_include_references: bool = Field(
         default=False,
-        description="If set, the sub agent references will be added to the main agent response references even in not mentioned in the main agent response text.",
+        description="If set, the sub agent references will be added to the main agent response references even if not mentioned in the main agent response text.",
     )
 
     answer_substrings_config: list[SubAgentAnswerSubstringConfig] = Field(

unique_toolkit/agentic/tools/a2a/prompts.py

@@ -37,6 +37,8 @@ References: <sup><name>SubAgentName 4</name>2</sup><sup><name>SubAgentName 4</na
 
 6. Fact repetition: If you reuse a fact from SubAgentName, you MUST reference it again inline with the correct format.
 
+7. You can ONLY use references if they are present in the subagent response! You must NOT create any references!
+
 Reminder:
 Inline = directly next to the fact, inside the same sentence or bullet.
 """.strip()

unique_toolkit/agentic/tools/a2a/tool/config.py

@@ -13,6 +13,7 @@ class SubAgentSystemReminderType(StrEnum):
     FIXED = "fixed"
     REGEXP = "regexp"
     REFERENCE = "reference"
+    NO_REFERENCE = "no_reference"
 
 
 T = TypeVar("T", bound=SubAgentSystemReminderType)
@@ -31,12 +32,24 @@ The reminder to add to the tool response. The reminder can be a Jinja template a
 """.strip()
 
 
+class NoReferenceSystemReminderConfig(SystemReminderConfig):
+    """A system reminder that is only added if the sub agent response does not contain any references."""
+
+    type: Literal[SubAgentSystemReminderType.NO_REFERENCE] = (
+        SubAgentSystemReminderType.NO_REFERENCE
+    )
+    reminder: str = Field(
+        default="Do NOT create any references from this sub agent in your response! The sub agent response does not contain any references.",
+        description=_SYSTEM_REMINDER_FIELD_DESCRIPTION,
+    )
+
+
 class ReferenceSystemReminderConfig(SystemReminderConfig):
     type: Literal[SubAgentSystemReminderType.REFERENCE] = (
         SubAgentSystemReminderType.REFERENCE
     )
     reminder: str = Field(
-        default="Rememeber to properly reference EACH fact from sub agent {{ display_name }}'s response with the correct format INLINE.",
+        default="Rememeber to properly reference EACH fact from sub agent {{ display_name }}'s response with the correct format INLINE. You MUST COPY THE REFERENCE AS PRESENT IN THE SUBAGENT RESPONSE.",
         description=_SYSTEM_REMINDER_FIELD_DESCRIPTION,
     )
 
@@ -69,6 +82,13 @@ class RegExpDetectedSystemReminderConfig(SystemReminderConfig):
     )
 
 
+SystemReminderConfigType = (
+    FixedSystemReminderConfig
+    | RegExpDetectedSystemReminderConfig
+    | ReferenceSystemReminderConfig
+    | NoReferenceSystemReminderConfig
+)
+
 DEFAULT_PARAM_DESCRIPTION_SUB_AGENT_USER_MESSAGE = """
 This is the message that will be sent to the sub-agent.
 """.strip()
@@ -147,9 +167,7 @@ class SubAgentToolConfig(BaseToolConfig):
 
     system_reminders_config: list[
         Annotated[
-            FixedSystemReminderConfig
-            | RegExpDetectedSystemReminderConfig
-            | ReferenceSystemReminderConfig,
+            SystemReminderConfigType,
             Field(discriminator="type"),
         ]
     ] = Field(

unique_toolkit/agentic/tools/a2a/tool/service.py

@@ -12,7 +12,7 @@ from unique_sdk.utils.chat_in_space import send_message_and_wait_for_completion
 
 from unique_toolkit._common.referencing import (
     get_all_ref_numbers,
-    remove_all_refs,
+    get_detection_pattern_for_ref,
     replace_ref_number,
 )
 from unique_toolkit._common.utils.jinja.render import render_template
@@ -29,6 +29,7 @@ from unique_toolkit.agentic.tools.a2a.tool.config import (
     RegExpDetectedSystemReminderConfig,
     SubAgentSystemReminderType,
     SubAgentToolConfig,
+    SystemReminderConfigType,
 )
 from unique_toolkit.agentic.tools.factory import ToolFactory
 from unique_toolkit.agentic.tools.schemas import ToolCallResponse
@@ -201,15 +202,39 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             if response["text"] is None:
                 raise ValueError("No response returned from sub agent")
 
+            has_refs = False
+            content = ""
+            content_chunks = None
             if self.config.returns_content_chunks:
-                content = ""
                 content_chunks = _ContentChunkList.validate_json(response["text"])
             else:
-                content = self._prepare_response_references(
+                has_refs = self.config.use_sub_agent_references and _response_has_refs(
+                    response
+                )
+                content = response["text"]
+                if has_refs:
+                    refs = response["references"]
+                    assert refs is not None  # Checked in _response_has_refs
+                    content = _prepare_sub_agent_response_refs(
+                        response=content,
+                        name=self.name,
+                        sequence_number=sequence_number,
+                        refs=refs,
+                    )
+                    content = _remove_extra_refs(content, refs=refs)
+                else:
+                    content = _remove_extra_refs(content, refs=[])
+
+            system_reminders = []
+            if not self.config.returns_content_chunks:
+                system_reminders = _get_sub_agent_system_reminders(
                     response=response["text"],
+                    configs=self.config.system_reminders_config,
+                    name=self.name,
+                    display_name=self.display_name(),
                     sequence_number=sequence_number,
+                    has_refs=has_refs,
                 )
-                content_chunks = None
 
             await self._notify_progress(
                 tool_call=tool_call,
@@ -223,58 +248,11 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
                 content=_format_response(
                     tool_name=self.name,
                     text=content,
-                    system_reminders=self._get_system_reminders(response),
+                    system_reminders=system_reminders,
                 ),
                 content_chunks=content_chunks,
             )
 
-    def _get_system_reminders(self, message: unique_sdk.Space.Message) -> list[str]:
-        reminders = []
-        for reminder_config in self.config.system_reminders_config:
-            if reminder_config.type == SubAgentSystemReminderType.FIXED:
-                reminders.append(
-                    render_template(
-                        reminder_config.reminder,
-                        display_name=self.display_name(),
-                        tool_name=self.name,
-                    )
-                )
-            elif (
-                reminder_config.type == SubAgentSystemReminderType.REFERENCE
-                and self.config.use_sub_agent_references
-                and message["references"] is not None
-                and len(message["references"]) > 0
-            ):
-                reminders.append(
-                    render_template(
-                        reminder_config.reminder,
-                        display_name=self.display_name(),
-                        tool_name=self.name,
-                    )
-                )
-            elif (
-                reminder_config.type == SubAgentSystemReminderType.REGEXP
-                and message["text"] is not None
-            ):
-                reminder_config = cast(
-                    RegExpDetectedSystemReminderConfig, reminder_config
-                )
-                text_matches = [
-                    match.group(0)
-                    for match in reminder_config.regexp.finditer(message["text"])
-                ]
-                if len(text_matches) > 0:
-                    reminders.append(
-                        render_template(
-                            reminder_config.reminder,
-                            display_name=self.display_name(),
-                            tool_name=self.name,
-                            text_matches=text_matches,
-                        )
-                    )
-
-        return reminders
-
     async def _get_chat_id(self) -> str | None:
         if not self.config.reuse_chat:
             return None
@@ -290,23 +268,6 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
 
         return None
 
-    def _prepare_response_references(self, response: str, sequence_number: int) -> str:
-        if not self.config.use_sub_agent_references:
-            # Remove all references from the response
-            response = remove_all_refs(response)
-            return response
-
-        for ref_number in get_all_ref_numbers(response):
-            reference = self.get_sub_agent_reference_format(
-                name=self.name,
-                sequence_number=sequence_number,
-                reference_number=ref_number,
-            )
-            response = replace_ref_number(
-                text=response, ref_number=ref_number, replacement=reference
-            )
-        return response
-
     async def _save_chat_id(self, chat_id: str) -> None:
         if not self.config.reuse_chat:
             return
@@ -390,4 +351,93 @@ def _format_response(tool_name: str, text: str, system_reminders: list[str]) ->
     return json.dumps(response, indent=2)
 
 
+def _response_has_refs(response: unique_sdk.Space.Message) -> bool:
+    if (
+        response["text"] is None
+        or response["references"] is None
+        or len(response["references"]) == 0
+    ):
+        return False
+
+    for ref in response["references"]:
+        if (
+            re.search(
+                get_detection_pattern_for_ref(ref["sequenceNumber"]), response["text"]
+            )
+            is not None
+        ):
+            return True
+
+    return False
+
+
+def _remove_extra_refs(response: str, refs: list[unique_sdk.Space.Reference]) -> str:
+    text_ref_numbers = set(get_all_ref_numbers(response))
+    extra_ref_numbers = text_ref_numbers - set(ref["sequenceNumber"] for ref in refs)
+
+    for ref_num in extra_ref_numbers:
+        response = get_detection_pattern_for_ref(ref_num).sub("", response)
+
+    return response
+
+
+def _prepare_sub_agent_response_refs(
+    response: str,
+    name: str,
+    sequence_number: int,
+    refs: list[unique_sdk.Space.Reference],
+) -> str:
+    for ref in refs:
+        ref_number = ref["sequenceNumber"]
+        reference = SubAgentTool.get_sub_agent_reference_format(
+            name=name, sequence_number=sequence_number, reference_number=ref_number
+        )
+        response = replace_ref_number(
+            text=response, ref_number=ref_number, replacement=reference
+        )
+
+    return response
+
+
+def _get_sub_agent_system_reminders(
+    response: str,
+    configs: list[SystemReminderConfigType],
+    name: str,
+    display_name: str,
+    sequence_number: int,
+    has_refs: bool,
+) -> list[str]:
+    reminders = []
+
+    for reminder_config in configs:
+        render_kwargs = {}
+        render_kwargs["display_name"] = display_name
+        render_kwargs["tool_name"] = name
+        template = None
+
+        if reminder_config.type == SubAgentSystemReminderType.FIXED:
+            template = reminder_config.reminder
+        elif (
+            reminder_config.type == SubAgentSystemReminderType.REFERENCE and has_refs
+        ) or (
+            reminder_config.type == SubAgentSystemReminderType.NO_REFERENCE
+            and not has_refs
+        ):
+            render_kwargs["tool_name"] = f"{name} {sequence_number}"
+            template = reminder_config.reminder
+        elif reminder_config.type == SubAgentSystemReminderType.REGEXP:
+            reminder_config = cast(RegExpDetectedSystemReminderConfig, reminder_config)
+            text_matches = [
+                match.group(0) for match in reminder_config.regexp.finditer(response)
+            ]
+            if len(text_matches) > 0:
+                template = reminder_config.reminder
+                render_kwargs["text_matches"] = text_matches
+
+        if template is not None:
+            reminders.append(render_template(template, **render_kwargs))
+
+    return reminders
+
+
 ToolFactory.register_tool(SubAgentTool, SubAgentToolConfig)