unique_toolkit 1.10.0__tar.gz → 1.11.0__tar.gz

{unique_toolkit-1.10.0 → unique_toolkit-1.11.0}/CHANGELOG.md

@@ -5,8 +5,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.11.0] - 2025-10-07
+- Add sub-agent response referencing.
 
-## 1.10.0 - 2025-10-07
+## [1.10.0] - 2025-10-07
 - Introduce future proof knowledgebase service decoupled from chat
 - Extend chat service to download contents in the chat
 - Update documentation
@@ -733,4 +735,4 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   - `utils.py` with utility functions for parsing language model output.
 
 ## [0.0.2] - 2024-07-10
-- Initial release of `unique_toolkit`.
+- Initial release of `unique_toolkit`.

{unique_toolkit-1.10.0 → unique_toolkit-1.11.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 1.10.0
+Version: 1.11.0
 Summary: 
 License: Proprietary
 Author: Cedric Klinkert
@@ -118,8 +118,10 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), 
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.11.0] - 2025-10-07
+- Add sub-agent response referencing.
 
-## 1.10.0 - 2025-10-07
+## [1.10.0] - 2025-10-07
 - Introduce future proof knowledgebase service decoupled from chat
 - Extend chat service to download contents in the chat
 - Update documentation
@@ -847,3 +849,4 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [0.0.2] - 2024-07-10
 - Initial release of `unique_toolkit`.
+

{unique_toolkit-1.10.0 → unique_toolkit-1.11.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "unique_toolkit"
-version = "1.10.0"
+version = "1.11.0"
 description = ""
 authors = [
     "Cedric Klinkert <cedric.klinkert@unique.ch>",

{unique_toolkit-1.10.0 → unique_toolkit-1.11.0}/unique_toolkit/agentic/tools/a2a/__init__.py

@@ -7,6 +7,10 @@ from unique_toolkit.agentic.tools.a2a.manager import A2AManager
 from unique_toolkit.agentic.tools.a2a.postprocessing import (
     SubAgentResponsesPostprocessor,
 )
+from unique_toolkit.agentic.tools.a2a.prompts import (
+    REFERENCING_INSTRUCTIONS_FOR_SYSTEM_PROMPT,
+    REFERENCING_INSTRUCTIONS_FOR_USER_PROMPT,
+)
 from unique_toolkit.agentic.tools.a2a.tool import SubAgentTool, SubAgentToolConfig
 
 __all__ = [
@@ -17,4 +21,6 @@ __all__ = [
     "ExtendedSubAgentToolConfig",
     "SubAgentEvaluationServiceConfig",
     "SubAgentEvaluationService",
+    "REFERENCING_INSTRUCTIONS_FOR_SYSTEM_PROMPT",
+    "REFERENCING_INSTRUCTIONS_FOR_USER_PROMPT",
 ]

{unique_toolkit-1.10.0 → unique_toolkit-1.11.0}/unique_toolkit/agentic/tools/a2a/postprocessing/postprocessor.py

@@ -1,4 +1,5 @@
 import logging
+import re
 from typing import TypedDict, override
 
 import unique_sdk
@@ -30,6 +31,7 @@ class _SubAgentMessageInfo(TypedDict):
 
 
 class _SubAgentToolInfo(TypedDict):
+    name: str
     display_name: str
     display_config: SubAgentDisplayConfig
     responses: dict[int, _SubAgentMessageInfo]
@@ -83,7 +85,7 @@ class SubAgentResponsesPostprocessor(Postprocessor):
         }
 
         _consolidate_references_in_place(
-            list(self._assistant_id_to_tool_info.values()), existing_refs
+            list(self._assistant_id_to_tool_info.values()), existing_refs, loop_response
         )
 
         answers = []
@@ -151,6 +153,7 @@ class SubAgentResponsesPostprocessor(Postprocessor):
             self._assistant_id_to_tool_info[tool.config.assistant_id] = (
                 _SubAgentToolInfo(
                     display_config=display_config,
+                    name=tool.name,
                     display_name=tool.display_name(),
                     responses={},
                 )
@@ -192,7 +195,9 @@ class SubAgentResponsesPostprocessor(Postprocessor):
 
 
 def _consolidate_references_in_place(
-    messages: list[_SubAgentToolInfo], existing_refs: dict[str, int]
+    messages: list[_SubAgentToolInfo],
+    existing_refs: dict[str, int],
+    loop_response: LanguageModelStreamResponse,
 ) -> None:
     start_index = max(existing_refs.values(), default=0) + 1
 
@@ -226,5 +231,29 @@ def _consolidate_references_in_place(
                 ref_map[reference["sequenceNumber"]] = reference_num
                 reference["sequenceNumber"] = reference_num
 
+            loop_response.message.text = (
+                _replace_sub_agent_references_in_main_agent_message(
+                    loop_response.message.text,
+                    assistant_tool_info["name"],
+                    sequence_number,
+                    ref_map,
+                )
+            )
             message["text"] = _replace_references_in_text(message["text"], ref_map)
             message["references"] = message_new_refs
+
+
+def _replace_sub_agent_references_in_main_agent_message(
+    message: str, sub_agent_name: str, sequence_number: int, ref_map: dict[int, int]
+) -> str:
+    for old_seq_num, new_seq_num in ref_map.items():
+        reference = SubAgentTool.get_sub_agent_reference_format(
+            name=sub_agent_name,
+            sequence_number=sequence_number,
+            reference_number=old_seq_num,
+        )
+        message = re.sub(rf"\s*{reference}", f" <sup>{new_seq_num}</sup>", message)
+
+    # Remove spaces between consecutive references
+    message = re.sub(r"</sup>\s*<sup>", "</sup><sup>", message)
+    return message

{unique_toolkit-1.10.0 → unique_toolkit-1.11.0}/unique_toolkit/agentic/tools/a2a/postprocessing/test/test_consolidate_references.py

@@ -1,4 +1,4 @@
-from unittest.mock import patch
+from unittest.mock import MagicMock, patch
 
 import pytest
 
@@ -10,12 +10,21 @@ from unique_toolkit.agentic.tools.a2a.postprocessing.postprocessor import (
 class TestConsolidateReferencesInPlace:
     """Test cases for _consolidate_references_in_place function."""
 
+    def _create_mock_loop_response(self, text: str = "Test response") -> MagicMock:
+        """Create a mock LanguageModelStreamResponse object."""
+        mock_response = MagicMock()
+        mock_message = MagicMock()
+        mock_message.text = text
+        mock_response.message = mock_message
+        return mock_response
+
     def test_empty_messages_list(self):
         """Test with empty messages list."""
         messages = []
         existing_refs = {}
+        loop_response = self._create_mock_loop_response()
 
-        _consolidate_references_in_place(messages, existing_refs)
+        _consolidate_references_in_place(messages, existing_refs, loop_response)
 
         assert existing_refs == {}
 
@@ -23,6 +32,7 @@ class TestConsolidateReferencesInPlace:
         """Test with empty existing references."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -43,8 +53,9 @@ class TestConsolidateReferencesInPlace:
             }
         ]
         existing_refs = {}
+        loop_response = self._create_mock_loop_response()
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # Should start from index 1 since existing_refs is empty
         assert existing_refs == {"source1": 1}
@@ -58,6 +69,7 @@ class TestConsolidateReferencesInPlace:
         """Test with existing references."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -78,8 +90,9 @@ class TestConsolidateReferencesInPlace:
             }
         ]
         existing_refs = {"existing_source": 5}
+        loop_response = self._create_mock_loop_response()
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # Should start from 6 (max existing + 1)
         assert existing_refs == {"existing_source": 5, "new_source": 6}
@@ -93,6 +106,7 @@ class TestConsolidateReferencesInPlace:
         """Test with duplicate source IDs across different messages."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -127,7 +141,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # Both should map to the same consolidated reference number
         assert existing_refs == {"shared_source": 1}
@@ -142,6 +157,7 @@ class TestConsolidateReferencesInPlace:
         """Test with multiple references in a single message."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -171,7 +187,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         assert existing_refs == {"source1": 1, "source2": 2}
         # Both references should be in the message since they're unique
@@ -185,6 +202,7 @@ class TestConsolidateReferencesInPlace:
         """Test that references are sorted by sequence number before processing."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -222,7 +240,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # Should be processed in order 1, 2, 3 and assigned consecutive numbers
         assert existing_refs == {"source1": 1, "source2": 2, "source3": 3}
@@ -235,6 +254,7 @@ class TestConsolidateReferencesInPlace:
         """Test with empty references list."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {1: {"text": "Text with no references", "references": []}},
@@ -242,16 +262,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        with patch(
-            "unique_toolkit.agentic.tools.a2a.postprocessing.postprocessor.logger"
-        ) as mock_logger:
-            _consolidate_references_in_place(messages, existing_refs)  # type: ignore
-
-            # Should log debug message about no references
-            mock_logger.debug.assert_called_once()
-            assert (
-                "does not contain any references" in mock_logger.debug.call_args[0][0]
-            )
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         assert existing_refs == {}
         assert messages[0]["responses"][1]["text"] == "Text with no references"
@@ -260,6 +272,7 @@ class TestConsolidateReferencesInPlace:
         """Test with multiple assistants."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -279,6 +292,7 @@ class TestConsolidateReferencesInPlace:
                 },
             },
             {
+                "name": "Assistant2",
                 "display_name": "Assistant2",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -300,7 +314,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         assert existing_refs == {"source1": 1, "source2": 2}
         assert messages[0]["responses"][1]["text"] == "Assistant 1 text <sup>1</sup>"
@@ -310,6 +325,7 @@ class TestConsolidateReferencesInPlace:
         """Test complex scenario with overlapping sequence numbers and mixed source IDs."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -352,7 +368,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {"existing": 10}  # Start from 11
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # sourceB gets 11 (first in sorted order), sourceA gets 12
         assert existing_refs == {"existing": 10, "sourceB": 11, "sourceA": 12}
@@ -377,6 +394,7 @@ class TestConsolidateReferencesInPlace:
 
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -398,7 +416,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # Should call _replace_references_in_text with the original text and ref mapping
         mock_replace.assert_called_once_with("Original text <sup>1</sup>", {1: 1})
@@ -417,6 +436,7 @@ class TestConsolidateReferencesInPlace:
 
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -426,7 +446,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # The original reference object should be modified in place
         assert original_ref["sequenceNumber"] == 1
@@ -445,6 +466,7 @@ class TestConsolidateReferencesInPlace:
         """Test that start_index is calculated correctly from existing_refs."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -466,7 +488,8 @@ class TestConsolidateReferencesInPlace:
         ]
 
         existing_refs_copy = existing_refs.copy()
-        _consolidate_references_in_place(messages, existing_refs_copy)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs_copy, loop_response)  # type: ignore
 
         assert existing_refs_copy["new_source"] == expected_start
 
@@ -474,6 +497,7 @@ class TestConsolidateReferencesInPlace:
         """Test with assistant that has no responses."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {},
@@ -481,7 +505,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         assert existing_refs == {}
 
@@ -489,6 +514,7 @@ class TestConsolidateReferencesInPlace:
         """Test with mix of valid messages and messages with no references."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -524,13 +550,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        with patch(
-            "unique_toolkit.agentic.tools.a2a.postprocessing.postprocessor.logger"
-        ) as mock_logger:
-            _consolidate_references_in_place(messages, existing_refs)  # type: ignore
-
-            # Should log once for message 2
-            mock_logger.debug.assert_called_once()
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # Only valid messages should be processed
         assert existing_refs == {"source1": 1, "source3": 2}
@@ -542,6 +563,7 @@ class TestConsolidateReferencesInPlace:
         """Test that responses are processed in sorted sequence number order."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -589,7 +611,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # Should be processed in order 1, 2, 3 based on response sequence numbers
         assert existing_refs == {"source1": 1, "source2": 2, "source3": 3}
@@ -601,6 +624,7 @@ class TestConsolidateReferencesInPlace:
         """Test when source ID already exists in existing_refs."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -622,7 +646,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {"existing_source": 99}
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # Should use existing reference number and not add to new refs
         assert existing_refs == {"existing_source": 99}
@@ -638,6 +663,7 @@ class TestConsolidateReferencesInPlace:
         """Test edge case with zero and negative values in existing_refs."""
         messages = [
             {
+                "name": "Assistant1",
                 "display_name": "Assistant1",
                 "display_config": {"mode": "expanded"},
                 "responses": {
@@ -659,7 +685,8 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {"zero": 0, "negative": -5, "positive": 3}
 
-        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
+        loop_response = self._create_mock_loop_response()
+        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
 
         # Should start from max(0, -5, 3) + 1 = 4
         assert existing_refs["new_source"] == 4

unique_toolkit-1.11.0/unique_toolkit/agentic/tools/a2a/prompts.py

@@ -0,0 +1,46 @@
+REFERENCING_INSTRUCTIONS_FOR_SYSTEM_PROMPT = """
+Whenever you use a fact from a sub-agent response in yours, you MUST always reference it. 
+
+CRITICAL INSTRUCTION: References must always appear immediately after the fact they support.
+Do NOT collect, group, or move references into a list at the end.
+
+Rules:
+
+1. Inline placement: After every fact from SubAgentName, immediately attach its reference(s) inline.
+Example:
+“The stock price of Apple Inc. is $150” <sup><name>SubAgentName 2</name>1</sup>.
+
+2. No separate reference list: Do not place references in footnotes, bibliographies, or at the bottom.
+Wrong:
+“The stock price of Apple Inc. is $150.”
+References: <sup><name>SubAgentName 2</name>1</sup>
+Correct:
+“The stock price of Apple Inc. is $150” <sup><name>SubAgentName 2</name>1</sup>.
+
+3. Exact copy: Copy references character-for-character from SubAgentName’s message.
+Do not alter numbering, labels, or order.
+
+4. Multiple references: If more than one reference supports a single fact, include all of them inline, in the same sentence, in the original order.
+Example:
+“MSFT would be a good investment” <sup><name>SubAgentName 3</name>4</sup><sup><name>SubAgentName 3</name>8</sup>.
+Wrong:
+“MSFT would be a good investment” <sup><name>SubAgentName 3</name>8</sup><sup><name>SubAgentName 3</name>4</sup>. (order changed)
+
+5. Never at the bottom: References must always stay attached inline with the fact.
+Multi-fact Example (Correct):
+“Tesla delivered 400,000 cars in Q2” <sup><name>SubAgentName 4</name>2</sup>.
+“Its revenue for the quarter was $24B” <sup><name>SubAgentName 4</name>5</sup>.
+“The company also expanded its Berlin factory capacity” <sup><name>SubAgentName 4</name>7</sup>.
+Wrong Multi-fact Example:
+“Tesla delivered 400,000 cars in Q2. Its revenue for the quarter was $24B. The company also expanded its Berlin factory capacity.”
+References: <sup><name>SubAgentName 4</name>2</sup><sup><name>SubAgentName 4</name>5</sup><sup><name>SubAgentName 4</name>7</sup>
+
+6. Fact repetition: If you reuse a fact from SubAgentName, you MUST reference it again inline with the correct format.
+
+Reminder:
+Inline = directly next to the fact, inside the same sentence or bullet.
+""".strip()
+
+REFERENCING_INSTRUCTIONS_FOR_USER_PROMPT = """
+Rememeber to properly reference EACH fact from a sub agent's response with the correct format INLINE.
+""".strip()

{unique_toolkit-1.10.0 → unique_toolkit-1.11.0}/unique_toolkit/agentic/tools/a2a/tool/config.py

@@ -7,10 +7,6 @@ DEFAULT_PARAM_DESCRIPTION_SUB_AGENT_USER_MESSAGE = """
 This is the message that will be sent to the sub-agent.
 """.strip()
 
-DEFAULT_FORMAT_INFORMATION_SUB_AGENT_SYSTEM_MESSAGE = """
-NEVER mention any references from sub-agent answers in your response.
-""".strip()
-
 
 class SubAgentToolConfig(BaseToolConfig):
     model_config = get_configuration_dict()
@@ -27,6 +23,10 @@ class SubAgentToolConfig(BaseToolConfig):
         default=True,
         description="Whether to reuse the existing chat or create a new one for each sub-agent call.",
     )
+    use_sub_agent_references: bool = Field(
+        default=True,
+        description="Whether this sub agent's references should be used in the main agent's response.",
+    )
 
     tool_description_for_system_prompt: str = Field(
         default="",
@@ -41,7 +41,7 @@ class SubAgentToolConfig(BaseToolConfig):
         description="Description of the user message parameter that will be sent to the model.",
     )
     tool_format_information_for_system_prompt: str = Field(
-        default=DEFAULT_FORMAT_INFORMATION_SUB_AGENT_SYSTEM_MESSAGE,
+        default="",
         description="Format information that will be included in the system prompt to guide response formatting.",
     )
     tool_description_for_user_prompt: str = Field(

{unique_toolkit-1.10.0 → unique_toolkit-1.11.0}/unique_toolkit/agentic/tools/a2a/tool/service.py

@@ -1,5 +1,6 @@
 import asyncio
 import contextlib
+import re
 from typing import Protocol, override
 
 import unique_sdk
@@ -79,6 +80,12 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
         self._sequence_number = 1
         self._lock = asyncio.Lock()
 
+    @staticmethod
+    def get_sub_agent_reference_format(
+        name: str, sequence_number: int, reference_number: int
+    ) -> str:
+        return f"<sup><name>{name} {sequence_number}</name>{reference_number}</sup>"
+
     def subscribe(self, subscriber: SubAgentResponseSubscriber) -> None:
         self._subscribers.append(subscriber)
 
@@ -175,6 +182,11 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             if response["text"] is None:
                 raise ValueError("No response returned from sub agent")
 
+            response_text_with_references = self._prepare_response_references(
+                response=response["text"],
+                sequence_number=sequence_number,
+            )
+
             await self._notify_progress(
                 tool_call=tool_call,
                 message=tool_input.user_message,
@@ -184,7 +196,7 @@ class SubAgentTool(Tool[SubAgentToolConfig]):
             return ToolCallResponse(
                 id=tool_call.id,  # type: ignore
                 name=tool_call.name,
-                content=response["text"],
+                content=response_text_with_references,
             )
 
     @override
@@ -209,6 +221,21 @@ 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 = re.sub(r"<sup>\s*\d+\s*</sup>", "", response)
+            return response
+
+        for ref_number in re.findall(r"<sup>(\d+)</sup>", response):
+            reference = self.get_sub_agent_reference_format(
+                name=self.name,
+                sequence_number=sequence_number,
+                reference_number=ref_number,
+            )
+            response = response.replace(f"<sup>{ref_number}</sup>", reference)
+        return response
+
     async def _save_chat_id(self, chat_id: str) -> None:
         if not self.config.reuse_chat:
             return