wehub-resource-sync
224 lines
8.0 KiB
Python
224 lines
8.0 KiB
Python
# 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.")
|