workflow.py

# Required imports
import asyncio
from typing import Dict, List, Any, Optional
from llama_index.core import Settings
from llama_index.core.tools import BaseTool
from llama_index.core.llms import ChatMessage
from llama_index.core.llms.llm import ToolSelection, LLM
from llama_index.core.workflow import (
    Workflow,
    Event,
    StartEvent,
    StopEvent,
    step,
    Context,
)


#####################################
# Define Router Agent Workflow
#####################################
class InputEvent(Event):
    """Input event."""


class GatherToolsEvent(Event):
    """Gather Tools Event"""

    tool_calls: Any


class ToolCallEvent(Event):
    """Tool Call event"""

    tool_call: ToolSelection


class ToolCallEventResult(Event):
    """Tool call event result."""

    msg: ChatMessage


class RouterOutputAgentWorkflow(Workflow):
    """Custom router output agent workflow."""

    def __init__(
        self,
        tools: List[BaseTool],
        timeout: Optional[float] = 10.0,
        disable_validation: bool = False,
        verbose: bool = False,
        llm: Optional[LLM] = None,
        chat_history: Optional[List[ChatMessage]] = None,
    ):
        """Constructor."""
        super().__init__(
            timeout=timeout, disable_validation=disable_validation, verbose=verbose
        )
        self.tools: List[BaseTool] = tools
        self.tools_dict: Optional[Dict[str, BaseTool]] = {
            tool.metadata.name: tool for tool in self.tools
        }
        # Use provided LLM or fall back to Settings.llm
        self.llm: LLM = llm or Settings.llm
        if self.llm is None:
            raise ValueError("No LLM provided and Settings.llm is not initialized")
        self.chat_history: List[ChatMessage] = chat_history or []

    def reset(self) -> None:
        """Resets Chat History"""
        self.chat_history = []

    @step()
    async def prepare_chat(self, ev: StartEvent) -> InputEvent:
        message = ev.get("message")
        if message is None:
            raise ValueError("'message' field is required.")

        # Add message to chat history
        chat_history = self.chat_history
        chat_history.append(ChatMessage(role="user", content=message))
        return InputEvent()

    @step()
    async def chat(self, ev: InputEvent) -> GatherToolsEvent | StopEvent:
        """Appends msg to chat history, then gets tool calls."""
        try:
            # Put message into LLM with tools included
            chat_res = await self.llm.achat_with_tools(
                self.tools,
                chat_history=self.chat_history,
                verbose=self._verbose,
                allow_parallel_tool_calls=True,
            )
            tool_calls = self.llm.get_tool_calls_from_response(
                chat_res, error_on_no_tool_call=False
            )

            ai_message = chat_res.message
            self.chat_history.append(ai_message)
            if self._verbose:
                print(f"Chat message: {ai_message.content}")

            # No tool calls, return chat message.
            if not tool_calls:
                return StopEvent(result=ai_message.content)

            return GatherToolsEvent(tool_calls=tool_calls)
        except asyncio.CancelledError:
            print("Chat operation was cancelled")
            return StopEvent(result="The operation was cancelled. Please try again.")
        except Exception as e:
            error_msg = f"Error during chat: {str(e)}"
            print(error_msg)
            return StopEvent(
                result="I'm sorry, I encountered an issue processing your request. Could you try asking in a different way?"
            )

    @step(pass_context=True)
    async def dispatch_calls(self, ctx: Context, ev: GatherToolsEvent) -> ToolCallEvent:
        """Dispatches calls."""
        tool_calls = ev.tool_calls
        await ctx.set("num_tool_calls", len(tool_calls))

        # Trigger tool call events
        for tool_call in tool_calls:
            ctx.send_event(ToolCallEvent(tool_call=tool_call))

        return None

    @step()
    async def call_tool(self, ev: ToolCallEvent) -> ToolCallEventResult:
        """Calls tool."""
        try:
            tool_call = ev.tool_call
            # Get tool ID and function call
            id_ = tool_call.tool_id

            if self._verbose:
                print(
                    f"Calling function {tool_call.tool_name} with msg {tool_call.tool_kwargs}"
                )

            # Call function and put result into a chat message
            tool = self.tools_dict[tool_call.tool_name]
            output = await tool.acall(**tool_call.tool_kwargs)
            
            # Check if output is a dictionary (response, trust_score) for document tool
            if isinstance(output, dict) and "response" in output:
                response = output.get("response", "")
                trust_score = output.get("trust_score")
                # Ensure response is a string
                content = str(response) if response is not None else ""
                # Store additional metadata
                additional_kwargs = {
                    "tool_call_id": id_, 
                    "name": tool_call.tool_name,
                    "trust_score": trust_score,
                    "tool_used": tool_call.tool_name
                }
                if self._verbose:
                    print(f"Tool {tool_call.tool_name} returned dict: response='{content}', trust_score={trust_score}")
            else:
                content = str(output) if output is not None else ""
                additional_kwargs = {
                    "tool_call_id": id_, 
                    "name": tool_call.tool_name,
                    "tool_used": tool_call.tool_name
                }
                if self._verbose:
                    print(f"Tool {tool_call.tool_name} returned: '{content}'")
            
            msg = ChatMessage(
                name=tool_call.tool_name,
                content=content,
                role="tool",
                additional_kwargs=additional_kwargs,
            )

            return ToolCallEventResult(msg=msg)
            
        except asyncio.CancelledError:
            print(f"Tool call {tool_call.tool_name} was cancelled")
            # Return a dummy result to avoid workflow breakdown
            msg = ChatMessage(
                name=tool_call.tool_name,
                content="Tool execution was cancelled",
                role="tool",
                additional_kwargs={"tool_call_id": id_, "name": tool_call.tool_name, "tool_used": tool_call.tool_name},
            )
            return ToolCallEventResult(msg=msg)
        except Exception as e:
            print(f"Error in tool call {tool_call.tool_name}: {str(e)}")
            # Return an error result instead of failing
            msg = ChatMessage(
                name=tool_call.tool_name,
                content=f"Error executing tool: {str(e)}",
                role="tool",
                additional_kwargs={"tool_call_id": id_, "name": tool_call.tool_name, "tool_used": tool_call.tool_name},
            )
            return ToolCallEventResult(msg=msg)

    @step(pass_context=True)
    async def gather(self, ctx: Context, ev: ToolCallEventResult) -> StopEvent | None:
        """Gathers tool calls."""
        try:
            # Wait for all tool call events to finish.
            tool_events = ctx.collect_events(
                ev, [ToolCallEventResult] * await ctx.get("num_tool_calls")
            )
            if not tool_events:
                return None

            for tool_event in tool_events:
                # Append tool call chat messages to history
                self.chat_history.append(tool_event.msg)

            # After all tool calls finish, pass input event back, restart agent loop
            return InputEvent()
        except Exception as e:
            print(f"Error in gather step: {str(e)}")
            # Return a stop event instead of continuing the loop if there's an error
            return StopEvent(result="I encountered an issue processing the tool responses. Please try again.")