Misc improvements (#179)

* openwebui stream better escaping

* fix dialog agent  finish msg rendering

* ResearchDialogAgent was not included

Author: virrius

examples/sgr_deep_research/agents.py

@@ -10,7 +10,9 @@
 from openai import AsyncOpenAI, pydantic_function_tool
 from openai.types.chat import ChatCompletionFunctionToolParam, ChatCompletionMessageParam
 
+from sgr_agent_core import AnswerTool
 from sgr_agent_core.agent_definition import AgentConfig
+from sgr_agent_core.agents.dialog_agent import DialogAgent
 from sgr_agent_core.agents.sgr_agent import SGRAgent
 from sgr_agent_core.agents.sgr_tool_calling_agent import SGRToolCallingAgent
 from sgr_agent_core.agents.tool_calling_agent import ToolCallingAgent
@@ -152,3 +154,50 @@ async def _prepare_tools(self) -> list[ChatCompletionFunctionToolParam]:
                 WebSearchTool,
             }
         return [pydantic_function_tool(tool, name=tool.tool_name, description="") for tool in tools]
+
+
+class ResearchDialogAgent(DialogAgent):
+    """Dialog research agent: supports intermediate results (AnswerTool) and
+    research toolkit with iteration/clarification/search limits.
+    """
+
+    def __init__(
+        self,
+        task_messages: list[ChatCompletionMessageParam],
+        openai_client: AsyncOpenAI,
+        agent_config: AgentConfig,
+        toolkit: list[Type[BaseTool]],
+        def_name: str | None = None,
+        **kwargs: dict,
+    ):
+        research_toolkit = [WebSearchTool, ExtractPageContentTool, CreateReportTool, AnswerTool]
+        merged = research_toolkit + [t for t in toolkit if t not in research_toolkit]
+        super().__init__(
+            task_messages=task_messages,
+            openai_client=openai_client,
+            agent_config=agent_config,
+            toolkit=merged,
+            def_name=def_name,
+            **kwargs,
+        )
+
+    async def _prepare_tools(self) -> list[ChatCompletionFunctionToolParam]:
+        """Prepare available tools with research limits (iterations,
+        clarifications, searches)."""
+        tools = set(self.toolkit)
+        if self._context.iteration >= self.config.execution.max_iterations:
+            tools = {
+                ReasoningTool,
+                CreateReportTool,
+                AnswerTool,
+            }
+        if self._context.clarifications_used >= self.config.execution.max_clarifications:
+            tools -= {
+                ClarificationTool,
+            }
+        search_config = self.get_tool_config(WebSearchTool)
+        if self._context.searches_used >= search_config.max_searches:
+            tools -= {
+                WebSearchTool,
+            }
+        return [pydantic_function_tool(tool, name=tool.tool_name, description="") for tool in tools]

sgr_agent_core/agents/dialog_agent.py

@@ -61,7 +61,9 @@ async def _after_action_phase(self, action_tool: BaseTool, result: str) -> None:
             self._context.execution_result = result
             self.logger.info("\n⏸️  Research paused - please answer questions")
             self._context.state = AgentStatesEnum.WAITING_FOR_CLARIFICATION
-            self.streaming_generator.finish()
+            self.streaming_generator.finish(
+                phase_id="{self._context.iteration}-final", content=self._context.execution_result
+            )
             self._context.clarification_received.clear()
             await self._context.clarification_received.wait()
             return
@@ -70,6 +72,8 @@ async def _after_action_phase(self, action_tool: BaseTool, result: str) -> None:
             self._context.execution_result = result
             self.logger.info("\n💬 Dialog shared - agent waiting for response")
             self._context.state = AgentStatesEnum.WAITING_FOR_CLARIFICATION
-            self.streaming_generator.finish(result)
+            self.streaming_generator.finish(
+                phase_id="{self._context.iteration}-final", content=self._context.execution_result
+            )
             self._context.clarification_received.clear()
             await self._context.clarification_received.wait()

sgr_agent_core/stream.py

@@ -1,4 +1,5 @@
 import asyncio
+import html
 import json
 import time
 from typing import Any, AsyncGenerator
@@ -159,8 +160,9 @@ def _wrap_in_code_block(self, content: str, language: str = "") -> str:
         """Wraps content in a Markdown code block."""
         if not content:  # otherwise confusing placeholder will show
             return "{}"
+        safe = html.escape(content, quote=False).replace("`", "\\`")
         lang_suffix = f" {language}" if language else ""
-        return f"```{lang_suffix}\n{content}\n```"
+        return f"```{lang_suffix}\n{safe}\n```"
 
     def add_tool_call(self, phase_id: str, tool: BaseTool) -> None:
         """Formats tool/reasoning and sends in <details>."""

tests/test_streaming.py

@@ -601,6 +601,46 @@ async def test_add_tool_result_plain_wraps_in_code_block(self):
         assert "Tool Result: custom" in content
         assert "plain text result" in content
 
+    @pytest.mark.asyncio
+    async def test_wrap_in_code_block_escapes_inner_backticks(self):
+        """Content with triple backticks is escaped so markdown block does not
+        break."""
+        generator = OpenWebUIStreamingGenerator(agent_id="test-id")
+        report_with_code = "# Report\n\n```python\nprint(1)\n```\n"
+        generator.add_tool_result(TEST_PHASE_ID, report_with_code, tool_name="CreateReport")
+        generator.finish(TEST_PHASE_ID)
+
+        items = []
+        async for item in generator.stream():
+            items.append(item)
+
+        content_chunk = items[0]
+        data = json.loads(content_chunk[6:].strip())
+        content = data["choices"][0]["delta"]["content"]
+        # Backticks in content are escaped so inner ``` do not close the block
+        assert "\\`" in content
+        assert "print(1)" in content
+
+    @pytest.mark.asyncio
+    async def test_wrap_in_code_block_escapes_html(self):
+        """Content with </details> or </summary> is HTML-escaped so it does not
+        break the page."""
+        generator = OpenWebUIStreamingGenerator(agent_id="test-id")
+        dangerous = "Summary: </summary> and </details> in text"
+        generator.add_tool_result(TEST_PHASE_ID, dangerous, tool_name="Report")
+        generator.finish(TEST_PHASE_ID)
+
+        items = []
+        async for item in generator.stream():
+            items.append(item)
+
+        content_chunk = items[0]
+        data = json.loads(content_chunk[6:].strip())
+        content = data["choices"][0]["delta"]["content"]
+        assert "&lt;/summary&gt;" in content
+        assert "&lt;/details&gt;" in content
+        assert "<details>" in content  # our real tag stays
+
     @pytest.mark.asyncio
     async def test_finish_inherited_works(self):
         """Test that finish() still produces [DONE] and final chunk."""