r2rpy/py/core/agent/base.py

import asyncio
import json
import logging
import re
from abc import ABCMeta
from typing import AsyncGenerator, Optional, Tuple

from core.base import AsyncSyncMeta, LLMChatCompletion, Message, syncable
from core.base.agent import Agent, Conversation
from core.utils import (
    CitationTracker,
    SearchResultsCollector,
    SSEFormatter,
    convert_nonserializable_objects,
    dump_obj,
    find_new_citation_spans,
)

logger = logging.getLogger()


class CombinedMeta(AsyncSyncMeta, ABCMeta):
    pass


def sync_wrapper(async_gen):
    loop = asyncio.get_event_loop()

    def wrapper():
        try:
            while True:
                try:
                    yield loop.run_until_complete(async_gen.__anext__())
                except StopAsyncIteration:
                    break
        finally:
            loop.run_until_complete(async_gen.aclose())

    return wrapper()


class R2RAgent(Agent, metaclass=CombinedMeta):
    def __init__(self, *args, **kwargs):
        self.search_results_collector = SearchResultsCollector()
        super().__init__(*args, **kwargs)
        self._reset()

    async def _generate_llm_summary(self, iterations_count: int) -> str:
        """
        Generate a summary of the conversation using the LLM when max iterations are exceeded.

        Args:
            iterations_count: The number of iterations that were completed

        Returns:
            A string containing the LLM-generated summary
        """
        try:
            # Get all messages in the conversation
            all_messages = await self.conversation.get_messages()

            # Create a prompt for the LLM to summarize
            summary_prompt = {
                "role": "user",
                "content": (
                    f"The conversation has reached the maximum limit of {iterations_count} iterations "
                    f"without completing the task. Please provide a concise summary of: "
                    f"1) The key information you've gathered that's relevant to the original query, "
                    f"2) What you've attempted so far and why it's incomplete, and "
                    f"3) A specific recommendation for how to proceed. "
                    f"Keep your summary brief (3-4 sentences total) and focused on the most valuable insights. If it is possible to answer the original user query, then do so now instead."
                    f"Start with '⚠️ **Maximum iterations exceeded**'"
                ),
            }

            # Create a new message list with just the conversation history and summary request
            summary_messages = all_messages + [summary_prompt]

            # Get a completion for the summary
            generation_config = self.get_generation_config(summary_prompt)
            response = await self.llm_provider.aget_completion(
                summary_messages,
                generation_config,
            )

            return response.choices[0].message.content
        except Exception as e:
            logger.error(f"Error generating LLM summary: {str(e)}")
            # Fall back to basic summary if LLM generation fails
            return (
                "⚠️ **Maximum iterations exceeded**\n\n"
                "The agent reached the maximum iteration limit without completing the task. "
                "Consider breaking your request into smaller steps or refining your query."
            )

    def _reset(self):
        self._completed = False
        self.conversation = Conversation()

    @syncable
    async def arun(
        self,
        messages: list[Message],
        system_instruction: Optional[str] = None,
        *args,
        **kwargs,
    ) -> list[dict]:
        self._reset()
        await self._setup(system_instruction)

        if messages:
            for message in messages:
                await self.conversation.add_message(message)
        iterations_count = 0
        while (
            not self._completed
            and iterations_count < self.config.max_iterations
        ):
            iterations_count += 1
            messages_list = await self.conversation.get_messages()
            generation_config = self.get_generation_config(messages_list[-1])
            response = await self.llm_provider.aget_completion(
                messages_list,
                generation_config,
            )
            logger.debug(f"R2RAgent response: {response}")
            await self.process_llm_response(response, *args, **kwargs)

        if not self._completed:
            # Generate a summary of the conversation using the LLM
            summary = await self._generate_llm_summary(iterations_count)
            await self.conversation.add_message(
                Message(role="assistant", content=summary)
            )

        # Return final content
        all_messages: list[dict] = await self.conversation.get_messages()
        all_messages.reverse()

        output_messages = []
        for message_2 in all_messages:
            if (
                # message_2.get("content")
                message_2.get("content") != messages[-1].content
            ):
                output_messages.append(message_2)
            else:
                break
        output_messages.reverse()

        return output_messages

    async def process_llm_response(
        self, response: LLMChatCompletion, *args, **kwargs
    ) -> None:
        if not self._completed:
            message = response.choices[0].message
            finish_reason = response.choices[0].finish_reason

            if finish_reason == "stop":
                self._completed = True

            # Determine which provider we're using
            using_anthropic = (
                "anthropic" in self.rag_generation_config.model.lower()
            )

            # OPENAI HANDLING
            if not using_anthropic:
                if message.tool_calls:
                    assistant_msg = Message(
                        role="assistant",
                        content="",
                        tool_calls=[msg.dict() for msg in message.tool_calls],
                    )
                    await self.conversation.add_message(assistant_msg)

                    # If there are multiple tool_calls, call them sequentially here
                    for tool_call in message.tool_calls:
                        await self.handle_function_or_tool_call(
                            tool_call.function.name,
                            tool_call.function.arguments,
                            tool_id=tool_call.id,
                            *args,
                            **kwargs,
                        )
                else:
                    await self.conversation.add_message(
                        Message(role="assistant", content=message.content)
                    )
                    self._completed = True

            else:
                # First handle thinking blocks if present
                if (
                    hasattr(message, "structured_content")
                    and message.structured_content
                ):
                    # Check if structured_content contains any tool_use blocks
                    has_tool_use = any(
                        block.get("type") == "tool_use"
                        for block in message.structured_content
                    )

                    if not has_tool_use and message.tool_calls:
                        # If it has thinking but no tool_use, add a separate message with structured_content
                        assistant_msg = Message(
                            role="assistant",
                            structured_content=message.structured_content,  # Use structured_content field
                        )
                        await self.conversation.add_message(assistant_msg)

                        # Add explicit tool_use blocks in a separate message
                        tool_uses = []
                        for tool_call in message.tool_calls:
                            # Safely parse arguments if they're a string
                            try:
                                if isinstance(
                                    tool_call.function.arguments, str
                                ):
                                    input_args = json.loads(
                                        tool_call.function.arguments
                                    )
                                else:
                                    input_args = tool_call.function.arguments
                            except json.JSONDecodeError:
                                logger.error(
                                    f"Failed to parse tool arguments: {tool_call.function.arguments}"
                                )
                                input_args = {
                                    "_raw": tool_call.function.arguments
                                }

                            tool_uses.append(
                                {
                                    "type": "tool_use",
                                    "id": tool_call.id,
                                    "name": tool_call.function.name,
                                    "input": input_args,
                                }
                            )

                        # Add tool_use blocks as a separate assistant message with structured content
                        if tool_uses:
                            await self.conversation.add_message(
                                Message(
                                    role="assistant",
                                    structured_content=tool_uses,
                                    content="",
                                )
                            )
                    else:
                        # If it already has tool_use or no tool_calls, preserve original structure
                        assistant_msg = Message(
                            role="assistant",
                            structured_content=message.structured_content,
                        )
                        await self.conversation.add_message(assistant_msg)

                elif message.content:
                    # For regular text content
                    await self.conversation.add_message(
                        Message(role="assistant", content=message.content)
                    )

                    # If there are tool calls, add them as structured content
                    if message.tool_calls:
                        tool_uses = []
                        for tool_call in message.tool_calls:
                            # Same safe parsing as above
                            try:
                                if isinstance(
                                    tool_call.function.arguments, str
                                ):
                                    input_args = json.loads(
                                        tool_call.function.arguments
                                    )
                                else:
                                    input_args = tool_call.function.arguments
                            except json.JSONDecodeError:
                                logger.error(
                                    f"Failed to parse tool arguments: {tool_call.function.arguments}"
                                )
                                input_args = {
                                    "_raw": tool_call.function.arguments
                                }

                            tool_uses.append(
                                {
                                    "type": "tool_use",
                                    "id": tool_call.id,
                                    "name": tool_call.function.name,
                                    "input": input_args,
                                }
                            )

                        await self.conversation.add_message(
                            Message(
                                role="assistant", structured_content=tool_uses
                            )
                        )

                # NEW CASE: Handle tool_calls with no content or structured_content
                elif message.tool_calls:
                    # Create tool_uses for the message with only tool_calls
                    tool_uses = []
                    for tool_call in message.tool_calls:
                        try:
                            if isinstance(tool_call.function.arguments, str):
                                input_args = json.loads(
                                    tool_call.function.arguments
                                )
                            else:
                                input_args = tool_call.function.arguments
                        except json.JSONDecodeError:
                            logger.error(
                                f"Failed to parse tool arguments: {tool_call.function.arguments}"
                            )
                            input_args = {"_raw": tool_call.function.arguments}

                        tool_uses.append(
                            {
                                "type": "tool_use",
                                "id": tool_call.id,
                                "name": tool_call.function.name,
                                "input": input_args,
                            }
                        )

                    # Add tool_use blocks as a message before processing tools
                    if tool_uses:
                        await self.conversation.add_message(
                            Message(
                                role="assistant",
                                structured_content=tool_uses,
                            )
                        )

                # Process the tool calls
                if message.tool_calls:
                    for tool_call in message.tool_calls:
                        await self.handle_function_or_tool_call(
                            tool_call.function.name,
                            tool_call.function.arguments,
                            tool_id=tool_call.id,
                            *args,
                            **kwargs,
                        )


class R2RStreamingAgent(R2RAgent):
    """
    Base class for all streaming agents with core streaming functionality.
    Supports emitting messages, tool calls, and results as SSE events.
    """

    # These two regexes will detect bracket references and then find short IDs.
    BRACKET_PATTERN = re.compile(r"\[([^\]]+)\]")
    SHORT_ID_PATTERN = re.compile(
        r"[A-Za-z0-9]{7,8}"
    )  # 7-8 chars, for example

    def __init__(self, *args, **kwargs):
        # Force streaming on
        if hasattr(kwargs.get("config", {}), "stream"):
            kwargs["config"].stream = True
        super().__init__(*args, **kwargs)

    async def arun(
        self,
        system_instruction: str | None = None,
        messages: list[Message] | None = None,
        *args,
        **kwargs,
    ) -> AsyncGenerator[str, None]:
        """
        Main streaming entrypoint: returns an async generator of SSE lines.
        """
        self._reset()
        await self._setup(system_instruction)

        if messages:
            for m in messages:
                await self.conversation.add_message(m)

        # Initialize citation tracker for this run
        citation_tracker = CitationTracker()

        # Dictionary to store citation payloads by ID
        citation_payloads = {}

        # Track all citations emitted during streaming for final persistence
        self.streaming_citations: list[dict] = []

        async def sse_generator() -> AsyncGenerator[str, None]:
            pending_tool_calls = {}
            partial_text_buffer = ""
            iterations_count = 0

            try:
                # Keep streaming until we complete
                while (
                    not self._completed
                    and iterations_count < self.config.max_iterations
                ):
                    iterations_count += 1
                    # 1) Get current messages
                    msg_list = await self.conversation.get_messages()
                    gen_cfg = self.get_generation_config(
                        msg_list[-1], stream=True
                    )

                    accumulated_thinking = ""
                    thinking_signatures = {}  # Map thinking content to signatures

                    # 2) Start streaming from LLM
                    llm_stream = self.llm_provider.aget_completion_stream(
                        msg_list, gen_cfg
                    )
                    async for chunk in llm_stream:
                        delta = chunk.choices[0].delta
                        finish_reason = chunk.choices[0].finish_reason

                        if hasattr(delta, "thinking") and delta.thinking:
                            # Accumulate thinking for later use in messages
                            accumulated_thinking += delta.thinking

                            # Emit SSE "thinking" event
                            async for (
                                line
                            ) in SSEFormatter.yield_thinking_event(
                                delta.thinking
                            ):
                                yield line

                        # Add this new handler for thinking signatures
                        if hasattr(delta, "thinking_signature"):
                            thinking_signatures[accumulated_thinking] = (
                                delta.thinking_signature
                            )
                            accumulated_thinking = ""

                        # 3) If new text, accumulate it
                        if delta.content:
                            partial_text_buffer += delta.content

                            # (a) Now emit the newly streamed text as a "message" event
                            async for line in SSEFormatter.yield_message_event(
                                delta.content
                            ):
                                yield line

                            # (b) Find new citation spans in the accumulated text
                            new_citation_spans = find_new_citation_spans(
                                partial_text_buffer, citation_tracker
                            )

                            # Process each new citation span
                            for cid, spans in new_citation_spans.items():
                                for span in spans:
                                    # Check if this is the first time we've seen this citation ID
                                    is_new_citation = (
                                        citation_tracker.is_new_citation(cid)
                                    )

                                    # Get payload if it's a new citation
                                    payload = None
                                    if is_new_citation:
                                        source_obj = self.search_results_collector.find_by_short_id(
                                            cid
                                        )
                                        if source_obj:
                                            # Store payload for reuse
                                            payload = dump_obj(source_obj)
                                            citation_payloads[cid] = payload

                                    # Create citation event payload
                                    citation_data = {
                                        "id": cid,
                                        "object": "citation",
                                        "is_new": is_new_citation,
                                        "span": {
                                            "start": span[0],
                                            "end": span[1],
                                        },
                                    }

                                    # Only include full payload for new citations
                                    if is_new_citation and payload:
                                        citation_data["payload"] = payload

                                    # Add to streaming citations for final answer
                                    self.streaming_citations.append(
                                        citation_data
                                    )

                                    # Emit the citation event
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_citation_event(
                                        citation_data
                                    ):
                                        yield line

                        if delta.tool_calls:
                            for tc in delta.tool_calls:
                                idx = tc.index
                                if idx not in pending_tool_calls:
                                    pending_tool_calls[idx] = {
                                        "id": tc.id,
                                        "name": tc.function.name or "",
                                        "arguments": tc.function.arguments
                                        or "",
                                    }
                                else:
                                    # Accumulate partial name/arguments
                                    if tc.function.name:
                                        pending_tool_calls[idx]["name"] = (
                                            tc.function.name
                                        )
                                    if tc.function.arguments:
                                        pending_tool_calls[idx][
                                            "arguments"
                                        ] += tc.function.arguments

                        # 5) If the stream signals we should handle "tool_calls"
                        if finish_reason == "tool_calls":
                            # Handle thinking if present
                            await self._handle_thinking(
                                thinking_signatures, accumulated_thinking
                            )

                            calls_list = []
                            for idx in sorted(pending_tool_calls.keys()):
                                cinfo = pending_tool_calls[idx]
                                calls_list.append(
                                    {
                                        "tool_call_id": cinfo["id"]
                                        or f"call_{idx}",
                                        "name": cinfo["name"],
                                        "arguments": cinfo["arguments"],
                                    }
                                )

                            # (a) Emit SSE "tool_call" events
                            for c in calls_list:
                                tc_data = self._create_tool_call_data(c)
                                async for (
                                    line
                                ) in SSEFormatter.yield_tool_call_event(
                                    tc_data
                                ):
                                    yield line

                            # (b) Add an assistant message capturing these calls
                            await self._add_tool_calls_message(
                                calls_list, partial_text_buffer
                            )

                            # (c) Execute each tool call in parallel
                            await asyncio.gather(
                                *[
                                    self.handle_function_or_tool_call(
                                        c["name"],
                                        c["arguments"],
                                        tool_id=c["tool_call_id"],
                                    )
                                    for c in calls_list
                                ]
                            )

                            # Reset buffer & calls
                            pending_tool_calls.clear()
                            partial_text_buffer = ""

                        elif finish_reason == "stop":
                            # Handle thinking if present
                            await self._handle_thinking(
                                thinking_signatures, accumulated_thinking
                            )

                            # 6) The LLM is done. If we have any leftover partial text,
                            #    finalize it in the conversation
                            if partial_text_buffer:
                                # Create the final message with metadata including citations
                                final_message = Message(
                                    role="assistant",
                                    content=partial_text_buffer,
                                    metadata={
                                        "citations": self.streaming_citations
                                    },
                                )

                                # Add it to the conversation
                                await self.conversation.add_message(
                                    final_message
                                )

                            # (a) Prepare final answer with optimized citations
                            consolidated_citations = []
                            # Group citations by ID with all their spans
                            for (
                                cid,
                                spans,
                            ) in citation_tracker.get_all_spans().items():
                                if cid in citation_payloads:
                                    consolidated_citations.append(
                                        {
                                            "id": cid,
                                            "object": "citation",
                                            "spans": [
                                                {"start": s[0], "end": s[1]}
                                                for s in spans
                                            ],
                                            "payload": citation_payloads[cid],
                                        }
                                    )

                            # Create final answer payload
                            final_evt_payload = {
                                "id": "msg_final",
                                "object": "agent.final_answer",
                                "generated_answer": partial_text_buffer,
                                "citations": consolidated_citations,
                            }

                            # Emit final answer event
                            async for (
                                line
                            ) in SSEFormatter.yield_final_answer_event(
                                final_evt_payload
                            ):
                                yield line

                            # (b) Signal the end of the SSE stream
                            yield SSEFormatter.yield_done_event()
                            self._completed = True
                            break

                # If we exit the while loop due to hitting max iterations
                if not self._completed:
                    # Generate a summary using the LLM
                    summary = await self._generate_llm_summary(
                        iterations_count
                    )

                    # Send the summary as a message event
                    async for line in SSEFormatter.yield_message_event(
                        summary
                    ):
                        yield line

                    # Add summary to conversation with citations metadata
                    await self.conversation.add_message(
                        Message(
                            role="assistant",
                            content=summary,
                            metadata={"citations": self.streaming_citations},
                        )
                    )

                    # Create and emit a final answer payload with the summary
                    final_evt_payload = {
                        "id": "msg_final",
                        "object": "agent.final_answer",
                        "generated_answer": summary,
                        "citations": consolidated_citations,
                    }

                    async for line in SSEFormatter.yield_final_answer_event(
                        final_evt_payload
                    ):
                        yield line

                    # Signal the end of the SSE stream
                    yield SSEFormatter.yield_done_event()
                    self._completed = True

            except Exception as e:
                logger.error(f"Error in streaming agent: {str(e)}")
                # Emit error event for client
                async for line in SSEFormatter.yield_error_event(
                    f"Agent error: {str(e)}"
                ):
                    yield line
                # Send done event to close the stream
                yield SSEFormatter.yield_done_event()

        # Finally, we return the async generator
        async for line in sse_generator():
            yield line

    async def _handle_thinking(
        self, thinking_signatures, accumulated_thinking
    ):
        """Process any accumulated thinking content"""
        if accumulated_thinking:
            structured_content = [
                {
                    "type": "thinking",
                    "thinking": accumulated_thinking,
                    # Anthropic will validate this in their API
                    "signature": "placeholder_signature",
                }
            ]

            assistant_msg = Message(
                role="assistant",
                structured_content=structured_content,
            )
            await self.conversation.add_message(assistant_msg)

        elif thinking_signatures:
            for (
                accumulated_thinking,
                thinking_signature,
            ) in thinking_signatures.items():
                structured_content = [
                    {
                        "type": "thinking",
                        "thinking": accumulated_thinking,
                        # Anthropic will validate this in their API
                        "signature": thinking_signature,
                    }
                ]

                assistant_msg = Message(
                    role="assistant",
                    structured_content=structured_content,
                )
                await self.conversation.add_message(assistant_msg)

    async def _add_tool_calls_message(self, calls_list, partial_text_buffer):
        """Add a message with tool calls to the conversation"""
        assistant_msg = Message(
            role="assistant",
            content=partial_text_buffer or "",
            tool_calls=[
                {
                    "id": c["tool_call_id"],
                    "type": "function",
                    "function": {
                        "name": c["name"],
                        "arguments": c["arguments"],
                    },
                }
                for c in calls_list
            ],
        )
        await self.conversation.add_message(assistant_msg)

    def _create_tool_call_data(self, call_info):
        """Create tool call data structure from call info"""
        return {
            "tool_call_id": call_info["tool_call_id"],
            "name": call_info["name"],
            "arguments": call_info["arguments"],
        }

    def _create_citation_payload(self, short_id, payload):
        """Create citation payload for a short ID"""
        # This will be overridden in RAG subclasses
        # check if as_dict is on payload
        if hasattr(payload, "as_dict"):
            payload = payload.as_dict()
        if hasattr(payload, "dict"):
            payload = payload.dict
        if hasattr(payload, "to_dict"):
            payload = payload.to_dict()

        return {
            "id": f"{short_id}",
            "object": "citation",
            "payload": dump_obj(payload),  # Will be populated in RAG agents
        }

    def _create_final_answer_payload(self, answer_text, citations):
        """Create the final answer payload"""
        # This will be extended in RAG subclasses
        return {
            "id": "msg_final",
            "object": "agent.final_answer",
            "generated_answer": answer_text,
            "citations": citations,
        }


class R2RXMLStreamingAgent(R2RStreamingAgent):
    """
    A streaming agent that parses XML-formatted responses with special handling for:
     - <think> or <Thought> blocks for chain-of-thought reasoning
     - <Action>, <ToolCalls>, <ToolCall> blocks for tool execution
    """

    # We treat <think> or <Thought> as the same token boundaries
    THOUGHT_OPEN = re.compile(r"<(Thought|think)>", re.IGNORECASE)
    THOUGHT_CLOSE = re.compile(r"</(Thought|think)>", re.IGNORECASE)

    # Regexes to parse out <Action>, <ToolCalls>, <ToolCall>, <Name>, <Parameters>, <Response>
    ACTION_PATTERN = re.compile(
        r"<Action>(.*?)</Action>", re.IGNORECASE | re.DOTALL
    )
    TOOLCALLS_PATTERN = re.compile(
        r"<ToolCalls>(.*?)</ToolCalls>", re.IGNORECASE | re.DOTALL
    )
    TOOLCALL_PATTERN = re.compile(
        r"<ToolCall>(.*?)</ToolCall>", re.IGNORECASE | re.DOTALL
    )
    NAME_PATTERN = re.compile(r"<Name>(.*?)</Name>", re.IGNORECASE | re.DOTALL)
    PARAMS_PATTERN = re.compile(
        r"<Parameters>(.*?)</Parameters>", re.IGNORECASE | re.DOTALL
    )
    RESPONSE_PATTERN = re.compile(
        r"<Response>(.*?)</Response>", re.IGNORECASE | re.DOTALL
    )

    async def arun(
        self,
        system_instruction: str | None = None,
        messages: list[Message] | None = None,
        *args,
        **kwargs,
    ) -> AsyncGenerator[str, None]:
        """
        Main streaming entrypoint: returns an async generator of SSE lines.
        """
        self._reset()
        await self._setup(system_instruction)

        if messages:
            for m in messages:
                await self.conversation.add_message(m)

        # Initialize citation tracker for this run
        citation_tracker = CitationTracker()

        # Dictionary to store citation payloads by ID
        citation_payloads = {}

        # Track all citations emitted during streaming for final persistence
        self.streaming_citations: list[dict] = []

        async def sse_generator() -> AsyncGenerator[str, None]:
            iterations_count = 0

            try:
                # Keep streaming until we complete
                while (
                    not self._completed
                    and iterations_count < self.config.max_iterations
                ):
                    iterations_count += 1
                    # 1) Get current messages
                    msg_list = await self.conversation.get_messages()
                    gen_cfg = self.get_generation_config(
                        msg_list[-1], stream=True
                    )

                    # 2) Start streaming from LLM
                    llm_stream = self.llm_provider.aget_completion_stream(
                        msg_list, gen_cfg
                    )

                    # Create state variables for each iteration
                    iteration_buffer = ""
                    yielded_first_event = False
                    in_action_block = False
                    is_thinking = False
                    accumulated_thinking = ""
                    thinking_signatures = {}

                    async for chunk in llm_stream:
                        delta = chunk.choices[0].delta
                        finish_reason = chunk.choices[0].finish_reason

                        # Handle thinking if present
                        if hasattr(delta, "thinking") and delta.thinking:
                            # Accumulate thinking for later use in messages
                            accumulated_thinking += delta.thinking

                            # Emit SSE "thinking" event
                            async for (
                                line
                            ) in SSEFormatter.yield_thinking_event(
                                delta.thinking
                            ):
                                yield line

                        # Add this new handler for thinking signatures
                        if hasattr(delta, "thinking_signature"):
                            thinking_signatures[accumulated_thinking] = (
                                delta.thinking_signature
                            )
                            accumulated_thinking = ""

                        # 3) If new text, accumulate it
                        if delta.content:
                            iteration_buffer += delta.content

                            # Check if we have accumulated enough text for a `<Thought>` block
                            if len(iteration_buffer) < len("<Thought>"):
                                continue

                            # Check if we have yielded the first event
                            if not yielded_first_event:
                                # Emit the first chunk
                                if self.THOUGHT_OPEN.findall(iteration_buffer):
                                    is_thinking = True
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_thinking_event(
                                        iteration_buffer
                                    ):
                                        yield line
                                else:
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_message_event(
                                        iteration_buffer
                                    ):
                                        yield line

                                # Mark as yielded
                                yielded_first_event = True
                                continue

                            # Check if we are in a thinking block
                            if is_thinking:
                                # Still thinking, so keep yielding thinking events
                                if not self.THOUGHT_CLOSE.findall(
                                    iteration_buffer
                                ):
                                    # Emit SSE "thinking" event
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_thinking_event(
                                        delta.content
                                    ):
                                        yield line

                                    continue
                                # Done thinking, so emit the last thinking event
                                else:
                                    is_thinking = False
                                    thought_text = delta.content.split(
                                        "</Thought>"
                                    )[0].split("</think>")[0]
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_thinking_event(
                                        thought_text
                                    ):
                                        yield line
                                    post_thought_text = delta.content.split(
                                        "</Thought>"
                                    )[-1].split("</think>")[-1]
                                    delta.content = post_thought_text

                            # (b) Find new citation spans in the accumulated text
                            new_citation_spans = find_new_citation_spans(
                                iteration_buffer, citation_tracker
                            )

                            # Process each new citation span
                            for cid, spans in new_citation_spans.items():
                                for span in spans:
                                    # Check if this is the first time we've seen this citation ID
                                    is_new_citation = (
                                        citation_tracker.is_new_citation(cid)
                                    )

                                    # Get payload if it's a new citation
                                    payload = None
                                    if is_new_citation:
                                        source_obj = self.search_results_collector.find_by_short_id(
                                            cid
                                        )
                                        if source_obj:
                                            # Store payload for reuse
                                            payload = dump_obj(source_obj)
                                            citation_payloads[cid] = payload

                                    # Create citation event payload
                                    citation_data = {
                                        "id": cid,
                                        "object": "citation",
                                        "is_new": is_new_citation,
                                        "span": {
                                            "start": span[0],
                                            "end": span[1],
                                        },
                                    }

                                    # Only include full payload for new citations
                                    if is_new_citation and payload:
                                        citation_data["payload"] = payload

                                    # Add to streaming citations for final answer
                                    self.streaming_citations.append(
                                        citation_data
                                    )

                                    # Emit the citation event
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_citation_event(
                                        citation_data
                                    ):
                                        yield line

                            # Now prepare to emit the newly streamed text as a "message" event
                            if (
                                iteration_buffer.count("<")
                                and not in_action_block
                            ):
                                in_action_block = True

                            if (
                                in_action_block
                                and len(
                                    self.ACTION_PATTERN.findall(
                                        iteration_buffer
                                    )
                                )
                                < 2
                            ):
                                continue

                            elif in_action_block:
                                in_action_block = False
                                # Emit the post action block text, if it is there
                                post_action_text = iteration_buffer.split(
                                    "</Action>"
                                )[-1]
                                if post_action_text:
                                    async for (
                                        line
                                    ) in SSEFormatter.yield_message_event(
                                        post_action_text
                                    ):
                                        yield line

                            else:
                                async for (
                                    line
                                ) in SSEFormatter.yield_message_event(
                                    delta.content
                                ):
                                    yield line

                        elif finish_reason == "stop":
                            break

                    # Process any accumulated thinking
                    await self._handle_thinking(
                        thinking_signatures, accumulated_thinking
                    )

                    # 6) The LLM is done. If we have any leftover partial text,
                    #    finalize it in the conversation
                    if iteration_buffer:
                        # Create the final message with metadata including citations
                        final_message = Message(
                            role="assistant",
                            content=iteration_buffer,
                            metadata={"citations": self.streaming_citations},
                        )

                        # Add it to the conversation
                        await self.conversation.add_message(final_message)

                    # --- 4) Process any <Action>/<ToolCalls> blocks, or mark completed
                    action_matches = self.ACTION_PATTERN.findall(
                        iteration_buffer
                    )

                    if len(action_matches) > 0:
                        # Process each ToolCall
                        xml_toolcalls = "<ToolCalls>"

                        for action_block in action_matches:
                            tool_calls_text = []
                            # Look for ToolCalls wrapper, or use the raw action block
                            calls_wrapper = self.TOOLCALLS_PATTERN.findall(
                                action_block
                            )
                            if calls_wrapper:
                                for tw in calls_wrapper:
                                    tool_calls_text.append(tw)
                            else:
                                tool_calls_text.append(action_block)

                            for calls_region in tool_calls_text:
                                calls_found = self.TOOLCALL_PATTERN.findall(
                                    calls_region
                                )
                                for tc_block in calls_found:
                                    tool_name, tool_params = (
                                        self._parse_single_tool_call(tc_block)
                                    )
                                    if tool_name:
                                        # Emit SSE event for tool call
                                        tool_call_id = (
                                            f"call_{abs(hash(tc_block))}"
                                        )
                                        call_evt_data = {
                                            "tool_call_id": tool_call_id,
                                            "name": tool_name,
                                            "arguments": json.dumps(
                                                tool_params
                                            ),
                                        }
                                        async for line in (
                                            SSEFormatter.yield_tool_call_event(
                                                call_evt_data
                                            )
                                        ):
                                            yield line

                                        try:
                                            tool_result = await self.handle_function_or_tool_call(
                                                tool_name,
                                                json.dumps(tool_params),
                                                tool_id=tool_call_id,
                                                save_messages=False,
                                            )
                                            result_content = tool_result.llm_formatted_result
                                        except Exception as e:
                                            result_content = f"Error in tool '{tool_name}': {str(e)}"

                                        xml_toolcalls += (
                                            f"<ToolCall>"
                                            f"<Name>{tool_name}</Name>"
                                            f"<Parameters>{json.dumps(tool_params)}</Parameters>"
                                            f"<Result>{result_content}</Result>"
                                            f"</ToolCall>"
                                        )

                                        # Emit SSE tool result for non-result tools
                                        result_data = {
                                            "tool_call_id": tool_call_id,
                                            "role": "tool",
                                            "content": json.dumps(
                                                convert_nonserializable_objects(
                                                    result_content
                                                )
                                            ),
                                        }
                                        async for line in SSEFormatter.yield_tool_result_event(
                                            result_data
                                        ):
                                            yield line

                        xml_toolcalls += "</ToolCalls>"
                        pre_action_text = iteration_buffer[
                            : iteration_buffer.find(action_block)
                        ]
                        post_action_text = iteration_buffer[
                            iteration_buffer.find(action_block)
                            + len(action_block) :
                        ]
                        iteration_text = (
                            pre_action_text + xml_toolcalls + post_action_text
                        )

                        # Update the conversation with tool results
                        await self.conversation.add_message(
                            Message(
                                role="assistant",
                                content=iteration_text,
                                metadata={
                                    "citations": self.streaming_citations
                                },
                            )
                        )
                    else:
                        # (a) Prepare final answer with optimized citations
                        consolidated_citations = []
                        # Group citations by ID with all their spans
                        for (
                            cid,
                            spans,
                        ) in citation_tracker.get_all_spans().items():
                            if cid in citation_payloads:
                                consolidated_citations.append(
                                    {
                                        "id": cid,
                                        "object": "citation",
                                        "spans": [
                                            {"start": s[0], "end": s[1]}
                                            for s in spans
                                        ],
                                        "payload": citation_payloads[cid],
                                    }
                                )

                        # Create final answer payload
                        final_evt_payload = {
                            "id": "msg_final",
                            "object": "agent.final_answer",
                            "generated_answer": iteration_buffer,
                            "citations": consolidated_citations,
                        }

                        # Emit final answer event
                        async for (
                            line
                        ) in SSEFormatter.yield_final_answer_event(
                            final_evt_payload
                        ):
                            yield line

                        # (b) Signal the end of the SSE stream
                        yield SSEFormatter.yield_done_event()
                        self._completed = True

                # If we exit the while loop due to hitting max iterations
                if not self._completed:
                    # Generate a summary using the LLM
                    summary = await self._generate_llm_summary(
                        iterations_count
                    )

                    # Send the summary as a message event
                    async for line in SSEFormatter.yield_message_event(
                        summary
                    ):
                        yield line

                    # Add summary to conversation with citations metadata
                    await self.conversation.add_message(
                        Message(
                            role="assistant",
                            content=summary,
                            metadata={"citations": self.streaming_citations},
                        )
                    )

                    # Create and emit a final answer payload with the summary
                    final_evt_payload = {
                        "id": "msg_final",
                        "object": "agent.final_answer",
                        "generated_answer": summary,
                        "citations": consolidated_citations,
                    }

                    async for line in SSEFormatter.yield_final_answer_event(
                        final_evt_payload
                    ):
                        yield line

                    # Signal the end of the SSE stream
                    yield SSEFormatter.yield_done_event()
                    self._completed = True

            except Exception as e:
                logger.error(f"Error in streaming agent: {str(e)}")
                # Emit error event for client
                async for line in SSEFormatter.yield_error_event(
                    f"Agent error: {str(e)}"
                ):
                    yield line
                # Send done event to close the stream
                yield SSEFormatter.yield_done_event()

        # Finally, we return the async generator
        async for line in sse_generator():
            yield line

    def _parse_single_tool_call(
        self, toolcall_text: str
    ) -> Tuple[Optional[str], dict]:
        """
        Parse a ToolCall block to extract the name and parameters.

        Args:
            toolcall_text: The text content of a ToolCall block

        Returns:
            Tuple of (tool_name, tool_parameters)
        """
        name_match = self.NAME_PATTERN.search(toolcall_text)
        if not name_match:
            return None, {}
        tool_name = name_match.group(1).strip()

        params_match = self.PARAMS_PATTERN.search(toolcall_text)
        if not params_match:
            return tool_name, {}

        raw_params = params_match.group(1).strip()
        try:
            # Handle potential JSON parsing issues
            # First try direct parsing
            tool_params = json.loads(raw_params)
        except json.JSONDecodeError:
            # If that fails, try to clean up the JSON string
            try:
                # Replace escaped quotes that might cause issues
                cleaned_params = raw_params.replace('\\"', '"')
                # Try again with the cleaned string
                tool_params = json.loads(cleaned_params)
            except json.JSONDecodeError:
                # If all else fails, treat as a plain string value
                tool_params = {"value": raw_params}

        return tool_name, tool_params


class R2RXMLToolsAgent(R2RAgent):
    """
    A non-streaming agent that:
     - parses <think> or <Thought> blocks as chain-of-thought
     - filters out XML tags related to tool calls and actions
     - processes <Action><ToolCalls><ToolCall> blocks
     - properly extracts citations when they appear in the text
    """

    # We treat <think> or <Thought> as the same token boundaries
    THOUGHT_OPEN = re.compile(r"<(Thought|think)>", re.IGNORECASE)
    THOUGHT_CLOSE = re.compile(r"</(Thought|think)>", re.IGNORECASE)

    # Regexes to parse out <Action>, <ToolCalls>, <ToolCall>, <Name>, <Parameters>, <Response>
    ACTION_PATTERN = re.compile(
        r"<Action>(.*?)</Action>", re.IGNORECASE | re.DOTALL
    )
    TOOLCALLS_PATTERN = re.compile(
        r"<ToolCalls>(.*?)</ToolCalls>", re.IGNORECASE | re.DOTALL
    )
    TOOLCALL_PATTERN = re.compile(
        r"<ToolCall>(.*?)</ToolCall>", re.IGNORECASE | re.DOTALL
    )
    NAME_PATTERN = re.compile(r"<Name>(.*?)</Name>", re.IGNORECASE | re.DOTALL)
    PARAMS_PATTERN = re.compile(
        r"<Parameters>(.*?)</Parameters>", re.IGNORECASE | re.DOTALL
    )
    RESPONSE_PATTERN = re.compile(
        r"<Response>(.*?)</Response>", re.IGNORECASE | re.DOTALL
    )

    async def process_llm_response(self, response, *args, **kwargs):
        """
        Override the base process_llm_response to handle XML structured responses
        including thoughts and tool calls.
        """
        if self._completed:
            return

        message = response.choices[0].message
        finish_reason = response.choices[0].finish_reason

        if not message.content:
            # If there's no content, let the parent class handle the normal tool_calls flow
            return await super().process_llm_response(
                response, *args, **kwargs
            )

        # Get the response content
        content = message.content

        # HACK for gemini
        content = content.replace("```action", "")
        content = content.replace("```tool_code", "")
        content = content.replace("```", "")

        if (
            not content.startswith("<")
            and "deepseek" in self.rag_generation_config.model
        ):  # HACK - fix issues with adding `<think>` to the beginning
            content = "<think>" + content

        # Process any tool calls in the content
        action_matches = self.ACTION_PATTERN.findall(content)
        if action_matches:
            xml_toolcalls = "<ToolCalls>"
            for action_block in action_matches:
                tool_calls_text = []
                # Look for ToolCalls wrapper, or use the raw action block
                calls_wrapper = self.TOOLCALLS_PATTERN.findall(action_block)
                if calls_wrapper:
                    for tw in calls_wrapper:
                        tool_calls_text.append(tw)
                else:
                    tool_calls_text.append(action_block)

                # Process each ToolCall
                for calls_region in tool_calls_text:
                    calls_found = self.TOOLCALL_PATTERN.findall(calls_region)
                    for tc_block in calls_found:
                        tool_name, tool_params = self._parse_single_tool_call(
                            tc_block
                        )
                        if tool_name:
                            tool_call_id = f"call_{abs(hash(tc_block))}"
                            try:
                                tool_result = (
                                    await self.handle_function_or_tool_call(
                                        tool_name,
                                        json.dumps(tool_params),
                                        tool_id=tool_call_id,
                                        save_messages=False,
                                    )
                                )

                                # Add tool result to XML
                                xml_toolcalls += (
                                    f"<ToolCall>"
                                    f"<Name>{tool_name}</Name>"
                                    f"<Parameters>{json.dumps(tool_params)}</Parameters>"
                                    f"<Result>{tool_result.llm_formatted_result}</Result>"
                                    f"</ToolCall>"
                                )

                            except Exception as e:
                                logger.error(f"Error in tool call: {str(e)}")
                                # Add error to XML
                                xml_toolcalls += (
                                    f"<ToolCall>"
                                    f"<Name>{tool_name}</Name>"
                                    f"<Parameters>{json.dumps(tool_params)}</Parameters>"
                                    f"<Result>Error: {str(e)}</Result>"
                                    f"</ToolCall>"
                                )

            xml_toolcalls += "</ToolCalls>"
            pre_action_text = content[: content.find(action_block)]
            post_action_text = content[
                content.find(action_block) + len(action_block) :
            ]
            iteration_text = pre_action_text + xml_toolcalls + post_action_text

            # Create the assistant message
            await self.conversation.add_message(
                Message(role="assistant", content=iteration_text)
            )
        else:
            # Create an assistant message with the content as-is
            await self.conversation.add_message(
                Message(role="assistant", content=content)
            )

        # Only mark as completed if the finish_reason is "stop" or there are no action calls
        # This allows the agent to continue the conversation when tool calls are processed
        if finish_reason == "stop":
            self._completed = True

    def _parse_single_tool_call(
        self, toolcall_text: str
    ) -> Tuple[Optional[str], dict]:
        """
        Parse a ToolCall block to extract the name and parameters.

        Args:
            toolcall_text: The text content of a ToolCall block

        Returns:
            Tuple of (tool_name, tool_parameters)
        """
        name_match = self.NAME_PATTERN.search(toolcall_text)
        if not name_match:
            return None, {}
        tool_name = name_match.group(1).strip()

        params_match = self.PARAMS_PATTERN.search(toolcall_text)
        if not params_match:
            return tool_name, {}

        raw_params = params_match.group(1).strip()
        try:
            # Handle potential JSON parsing issues
            # First try direct parsing
            tool_params = json.loads(raw_params)
        except json.JSONDecodeError:
            # If that fails, try to clean up the JSON string
            try:
                # Replace escaped quotes that might cause issues
                cleaned_params = raw_params.replace('\\"', '"')
                # Try again with the cleaned string
                tool_params = json.loads(cleaned_params)
            except json.JSONDecodeError:
                # If all else fails, treat as a plain string value
                tool_params = {"value": raw_params}

        return tool_name, tool_params