unique_toolkit 1.18.0__py3-none-any.whl → 1.18.1__py3-none-any.whl

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

@@ -95,17 +95,29 @@ class SubAgentResponsesPostprocessor(Postprocessor):
                 "Main agent message is not set, the `run` method must be called first"
             )
 
+        _add_sub_agent_references_in_place(
+            loop_response=loop_response,
+            sub_agent_infos=list(self._assistant_id_to_tool_info.values()),
+        )
+
         existing_refs = {
             ref.source_id: ref.sequence_number
             for ref in loop_response.message.references
         }
 
+        displayed_sub_agent_infos = dict(
+            (assistant_id, sub_agent_info)
+            for assistant_id, sub_agent_info in self._assistant_id_to_tool_info.items()
+            if sub_agent_info["display_config"].mode
+            != SubAgentResponseDisplayMode.HIDDEN
+        )
+
         _consolidate_references_in_place(
-            list(self._assistant_id_to_tool_info.values()), existing_refs, loop_response
+            list(displayed_sub_agent_infos.values()), existing_refs
         )
 
         answers = []
-        for assistant_id in self._assistant_id_to_tool_info.keys():
+        for assistant_id in displayed_sub_agent_infos.keys():
             messages = self._assistant_id_to_tool_info[assistant_id]["responses"]
 
             for sequence_number in sorted(messages):
@@ -140,9 +152,12 @@ class SubAgentResponsesPostprocessor(Postprocessor):
                     for ref in message["references"]
                 )
 
-        loop_response.message.text = (
-            "<br>\n\n".join(answers) + "<br>\n\n" + loop_response.message.text.strip()
-        )
+        if len(answers) > 0:
+            loop_response.message.text = (
+                "<br>\n\n".join(answers)
+                + "<br>\n\n"
+                + loop_response.message.text.strip()
+            )
 
         return True
 
@@ -163,13 +178,6 @@ class SubAgentResponsesPostprocessor(Postprocessor):
     def register_sub_agent_tool(
         self, tool: SubAgentTool, display_config: SubAgentDisplayConfig
     ) -> None:
-        if display_config.mode == SubAgentResponseDisplayMode.HIDDEN:
-            logger.info(
-                "Sub agent tool %s has display mode `hidden`, responses will be ignored.",
-                tool.config.assistant_id,
-            )
-            return
-
         if tool.config.assistant_id not in self._assistant_id_to_tool_info:
             tool.subscribe(self)
             self._assistant_id_to_tool_info[tool.config.assistant_id] = (
@@ -216,10 +224,73 @@ class SubAgentResponsesPostprocessor(Postprocessor):
         }
 
 
+def _add_sub_agent_references_in_place(
+    loop_response: LanguageModelStreamResponse,
+    sub_agent_infos: list[_SubAgentToolInfo],
+) -> None:
+    message_text = loop_response.message.text
+
+    existing_refs = {
+        ref.source_id: ref.sequence_number
+        for ref in loop_response.message.references or []
+    }
+    start_index = max(existing_refs.values(), default=0) + 1
+
+    for sub_agent_info in sub_agent_infos:
+        sub_agent_name = sub_agent_info["name"]
+        for sequence_number in sorted(sub_agent_info["responses"]):
+            sub_agent_response = sub_agent_info["responses"][sequence_number]
+
+            for reference in sorted(
+                sub_agent_response["references"], key=lambda r: r["sequenceNumber"]
+            ):
+                ref_num = reference["sequenceNumber"]
+                source_id = reference["sourceId"]
+
+                reference_str = SubAgentTool.get_sub_agent_reference_format(
+                    name=sub_agent_name,
+                    sequence_number=sequence_number,
+                    reference_number=ref_num,
+                )
+
+                if reference_str not in message_text:
+                    # Reference not used
+                    continue
+
+                if source_id in existing_refs:
+                    new_ref_num = existing_refs[source_id]
+
+                else:
+                    new_ref_num = start_index
+                    existing_refs[source_id] = new_ref_num
+                    start_index += 1
+
+                    loop_response.message.references.append(
+                        ContentReference(
+                            name=reference["name"],
+                            url=reference["url"] or "",
+                            sequence_number=new_ref_num,
+                            source_id=source_id,
+                            source=reference["source"],
+                        )
+                    )
+
+                message_text = re.sub(
+                    rf"\s*{re.escape(reference_str)}",
+                    f" <sup>{new_ref_num}</sup>",
+                    message_text,
+                )
+
+    loop_response.message.text = message_text
+    # Remove spaces between consecutive references
+    loop_response.message.text = re.sub(
+        r"</sup>\s*<sup>", "</sup><sup>", loop_response.message.text
+    )
+
+
 def _consolidate_references_in_place(
     messages: list[_SubAgentToolInfo],
     existing_refs: dict[str, int],
-    loop_response: LanguageModelStreamResponse,
 ) -> None:
     start_index = max(existing_refs.values(), default=0) + 1
 
@@ -253,29 +324,5 @@ 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/agentic/tools/a2a/postprocessing/test/test_consolidate_references.py

@@ -22,9 +22,8 @@ class TestConsolidateReferencesInPlace:
         """Test with empty messages list."""
         messages = []
         existing_refs = {}
-        loop_response = self._create_mock_loop_response()
 
-        _consolidate_references_in_place(messages, existing_refs, loop_response)
+        _consolidate_references_in_place(messages, existing_refs)
 
         assert existing_refs == {}
 
@@ -53,9 +52,8 @@ class TestConsolidateReferencesInPlace:
             }
         ]
         existing_refs = {}
-        loop_response = self._create_mock_loop_response()
 
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         # Should start from index 1 since existing_refs is empty
         assert existing_refs == {"source1": 1}
@@ -90,9 +88,8 @@ class TestConsolidateReferencesInPlace:
             }
         ]
         existing_refs = {"existing_source": 5}
-        loop_response = self._create_mock_loop_response()
 
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         # Should start from 6 (max existing + 1)
         assert existing_refs == {"existing_source": 5, "new_source": 6}
@@ -141,8 +138,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         # Both should map to the same consolidated reference number
         assert existing_refs == {"shared_source": 1}
@@ -187,8 +183,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         assert existing_refs == {"source1": 1, "source2": 2}
         # Both references should be in the message since they're unique
@@ -240,8 +235,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         # Should be processed in order 1, 2, 3 and assigned consecutive numbers
         assert existing_refs == {"source1": 1, "source2": 2, "source3": 3}
@@ -262,8 +256,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         assert existing_refs == {}
         assert messages[0]["responses"][1]["text"] == "Text with no references"
@@ -314,8 +307,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         assert existing_refs == {"source1": 1, "source2": 2}
         assert messages[0]["responses"][1]["text"] == "Assistant 1 text <sup>1</sup>"
@@ -368,8 +360,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {"existing": 10}  # Start from 11
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         # sourceB gets 11 (first in sorted order), sourceA gets 12
         assert existing_refs == {"existing": 10, "sourceB": 11, "sourceA": 12}
@@ -416,8 +407,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # 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})
@@ -446,8 +436,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         # The original reference object should be modified in place
         assert original_ref["sequenceNumber"] == 1
@@ -488,8 +477,7 @@ class TestConsolidateReferencesInPlace:
         ]
 
         existing_refs_copy = existing_refs.copy()
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs_copy, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs_copy)  # type: ignore
 
         assert existing_refs_copy["new_source"] == expected_start
 
@@ -505,8 +493,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         assert existing_refs == {}
 
@@ -550,8 +537,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         # Only valid messages should be processed
         assert existing_refs == {"source1": 1, "source3": 2}
@@ -611,8 +597,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         # Should be processed in order 1, 2, 3 based on response sequence numbers
         assert existing_refs == {"source1": 1, "source2": 2, "source3": 3}
@@ -646,8 +631,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {"existing_source": 99}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         # Should use existing reference number and not add to new refs
         assert existing_refs == {"existing_source": 99}
@@ -685,8 +669,7 @@ class TestConsolidateReferencesInPlace:
         ]
         existing_refs = {"zero": 0, "negative": -5, "positive": 3}
 
-        loop_response = self._create_mock_loop_response()
-        _consolidate_references_in_place(messages, existing_refs, loop_response)  # type: ignore
+        _consolidate_references_in_place(messages, existing_refs)  # type: ignore
 
         # Should start from max(0, -5, 3) + 1 = 4
         assert existing_refs["new_source"] == 4

{unique_toolkit-1.18.0.dist-info → unique_toolkit-1.18.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: unique_toolkit
-Version: 1.18.0
+Version: 1.18.1
 Summary: 
 License: Proprietary
 Author: Cedric Klinkert
@@ -118,6 +118,9 @@ 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.18.1] - 2025-10-28
+- Fix bug where sub agent references were not properly displayed in the main agent response when the sub agent response was hidden.
+
 ## [1.18.0] - 2025-10-27
 - Temporary fix to rendering of sub agent responses.
 - Add config option `stop_condition` to `SubAgentToolConfig`

{unique_toolkit-1.18.0.dist-info → unique_toolkit-1.18.1.dist-info}/RECORD

@@ -68,8 +68,8 @@ unique_toolkit/agentic/tools/a2a/postprocessing/__init__.py,sha256=R90CSecxJrKH7
 unique_toolkit/agentic/tools/a2a/postprocessing/_display.py,sha256=rdWk-6M6V27v7G3dqaaDj1vXcgsFvDrHswFuSKQfd2I,5186
 unique_toolkit/agentic/tools/a2a/postprocessing/_utils.py,sha256=JsWwylR2Ao_L0wk1UlhqeN2fTxPnrbhoi1klYHVBnLk,750
 unique_toolkit/agentic/tools/a2a/postprocessing/config.py,sha256=hqo3vQZX63yXoLT7wN13vf0rC7qKiozKKfdkicYCQno,1061
-unique_toolkit/agentic/tools/a2a/postprocessing/postprocessor.py,sha256=lEjJ1u-CTpoyDT7OFaBBOklONUCUU-CGsosHJRYUfkI,10335
-unique_toolkit/agentic/tools/a2a/postprocessing/test/test_consolidate_references.py,sha256=0hlTbtqClyMqHNnZRRxETgjVhYrtRu3xUnZx0JSE-uo,28030
+unique_toolkit/agentic/tools/a2a/postprocessing/postprocessor.py,sha256=MMy3NeN4koHWmamaXvsEnrRbsh7XBaWhD0PRcr1i7C4,11931
+unique_toolkit/agentic/tools/a2a/postprocessing/test/test_consolidate_references.py,sha256=HvisnbhPUwuR9uELGunxLjNHvCXheoD5eW52L0tvUXk,26789
 unique_toolkit/agentic/tools/a2a/postprocessing/test/test_display.py,sha256=Hn2GOsVhpCkpFV9Asvy_IyiYvCd88rxDzi0E_zt2kWc,37553
 unique_toolkit/agentic/tools/a2a/postprocessing/test/test_postprocessor_reference_functions.py,sha256=GxSkkY-Xgd61Bk8sIRfEtkT9hqL1VgPLWrq-6XoB0rA,11360
 unique_toolkit/agentic/tools/a2a/prompts.py,sha256=0ILHL_RAcT04gFm2d470j4Gho7PoJXdCJy-bkZgf_wk,2401
@@ -165,7 +165,7 @@ unique_toolkit/short_term_memory/service.py,sha256=5PeVBu1ZCAfyDb2HLVvlmqSbyzBBu
 unique_toolkit/smart_rules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 unique_toolkit/smart_rules/compile.py,sha256=Ozhh70qCn2yOzRWr9d8WmJeTo7AQurwd3tStgBMPFLA,1246
 unique_toolkit/test_utilities/events.py,sha256=_mwV2bs5iLjxS1ynDCjaIq-gjjKhXYCK-iy3dRfvO3g,6410
-unique_toolkit-1.18.0.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
-unique_toolkit-1.18.0.dist-info/METADATA,sha256=6XAUmwya1a8ILqkVKETFxYEaCF4p4Q3GfhRmgWO6tcI,38528
-unique_toolkit-1.18.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-unique_toolkit-1.18.0.dist-info/RECORD,,
+unique_toolkit-1.18.1.dist-info/LICENSE,sha256=GlN8wHNdh53xwOPg44URnwag6TEolCjoq3YD_KrWgss,193
+unique_toolkit-1.18.1.dist-info/METADATA,sha256=JM7lTRErEGHbrIrH_VsVvp4NC2snqXYAzq89D6qOrF8,38686
+unique_toolkit-1.18.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+unique_toolkit-1.18.1.dist-info/RECORD,,