465 lines
21 KiB
Python
465 lines
21 KiB
Python
|
"""
|
||
|
|
User Manual Agent node for the Agentic RAG system.
|
||
|
|
This module contains the autonomous user manual agent that can use tools and generate responses.
|
||
|
|
"""
|
||
|
|
|
||
|
|
import logging
|
||
|
|
from typing import Dict, Any, List, Optional, Callable, Literal
|
||
|
|
from contextvars import ContextVar
|
||
|
|
from langchain_core.messages import AIMessage, SystemMessage, BaseMessage, ToolMessage, HumanMessage
|
||
|
|
from langchain_core.runnables import RunnableConfig
|
||
|
|
|
||
|
|
from .state import AgentState
|
||
|
|
from .user_manual_tools import get_user_manual_tool_schemas, get_user_manual_tools_by_name
|
||
|
|
from .message_trimmer import create_conversation_trimmer
|
||
|
|
from ..llm_client import LLMClient
|
||
|
|
from ..config import get_config
|
||
|
|
from ..sse import (
|
||
|
|
create_tool_start_event,
|
||
|
|
create_tool_result_event,
|
||
|
|
create_tool_error_event,
|
||
|
|
create_token_event,
|
||
|
|
create_error_event
|
||
|
|
)
|
||
|
|
from ..utils.error_handler import (
|
||
|
|
StructuredLogger, ErrorCategory, ErrorCode,
|
||
|
|
handle_async_errors, get_user_message
|
||
|
|
)
|
||
|
|
|
||
|
|
logger = StructuredLogger(__name__)
|
||
|
|
|
||
|
|
# Cache configuration at module level to avoid repeated get_config() calls
|
||
|
|
_cached_config = None
|
||
|
|
|
||
|
|
def get_cached_config():
|
||
|
|
"""Get cached configuration, loading it if not already cached"""
|
||
|
|
global _cached_config
|
||
|
|
if _cached_config is None:
|
||
|
|
_cached_config = get_config()
|
||
|
|
return _cached_config
|
||
|
|
|
||
|
|
|
||
|
|
# User Manual Agent node (autonomous function calling agent)
|
||
|
|
async def user_manual_agent_node(state: AgentState, config: Optional[RunnableConfig] = None) -> Dict[str, Any]:
|
||
|
|
"""
|
||
|
|
User Manual Agent node that autonomously uses user manual tools and generates final answer.
|
||
|
|
Implements "detect-first-then-stream" strategy for optimal multi-round behavior:
|
||
|
|
1. Always start with non-streaming detection to check for tool needs
|
||
|
|
2. If tool_calls exist → return immediately for routing to tools
|
||
|
|
3. If no tool_calls → temporarily disable tools and perform streaming final synthesis
|
||
|
|
"""
|
||
|
|
app_config = get_cached_config()
|
||
|
|
llm_client = LLMClient()
|
||
|
|
|
||
|
|
# Get stream callback from context variable
|
||
|
|
from .graph import stream_callback_context
|
||
|
|
stream_callback = stream_callback_context.get()
|
||
|
|
|
||
|
|
# Get user manual tool schemas and bind tools for planning phase
|
||
|
|
tool_schemas = get_user_manual_tool_schemas()
|
||
|
|
llm_client.bind_tools(tool_schemas, force_tool_choice=True)
|
||
|
|
|
||
|
|
# Create conversation trimmer for managing context length
|
||
|
|
trimmer = create_conversation_trimmer()
|
||
|
|
|
||
|
|
# Prepare messages with user manual system prompt
|
||
|
|
messages = state["messages"].copy()
|
||
|
|
if not messages or not isinstance(messages[0], SystemMessage):
|
||
|
|
rag_prompts = app_config.get_rag_prompts()
|
||
|
|
user_manual_prompt = rag_prompts.get("user_manual_prompt", "")
|
||
|
|
if not user_manual_prompt:
|
||
|
|
raise ValueError("user_manual_prompt is null")
|
||
|
|
|
||
|
|
# For user manual agent, we need to format the prompt with placeholders
|
||
|
|
# Extract current query and conversation history
|
||
|
|
current_query = ""
|
||
|
|
for message in reversed(messages):
|
||
|
|
if isinstance(message, HumanMessage):
|
||
|
|
current_query = message.content
|
||
|
|
break
|
||
|
|
|
||
|
|
conversation_history = ""
|
||
|
|
if len(messages) > 1:
|
||
|
|
conversation_history = render_conversation_history(messages[:-1]) # Exclude current query
|
||
|
|
|
||
|
|
# Format system prompt (initially with empty context, tools will provide it)
|
||
|
|
formatted_system_prompt = user_manual_prompt.format(
|
||
|
|
conversation_history=conversation_history,
|
||
|
|
context_content="", # Will be filled by tools
|
||
|
|
current_query=current_query
|
||
|
|
)
|
||
|
|
|
||
|
|
messages = [SystemMessage(content=formatted_system_prompt)] + messages
|
||
|
|
|
||
|
|
# Track tool rounds
|
||
|
|
current_round = state.get("tool_rounds", 0)
|
||
|
|
# Get max_tool_rounds_user_manual from state, fallback to config if not set
|
||
|
|
max_rounds = state.get("max_tool_rounds_user_manual", None)
|
||
|
|
if max_rounds is None:
|
||
|
|
max_rounds = app_config.app.max_tool_rounds_user_manual
|
||
|
|
|
||
|
|
# Only apply trimming at the start of a new conversation turn (when tool_rounds = 0)
|
||
|
|
# This prevents trimming current turn's tool results during multi-round tool calling
|
||
|
|
if current_round == 0:
|
||
|
|
# Trim conversation history to manage context length (only for previous conversation turns)
|
||
|
|
if trimmer.should_trim(messages):
|
||
|
|
messages = trimmer.trim_conversation_history(messages)
|
||
|
|
logger.info("Applied conversation history trimming for context management (new conversation turn)")
|
||
|
|
else:
|
||
|
|
logger.info(f"Skipping trimming during tool round {current_round} to preserve current turn's context")
|
||
|
|
|
||
|
|
logger.info(f"User Manual Agent node: tool_rounds={current_round}, max_tool_rounds={max_rounds}")
|
||
|
|
|
||
|
|
# Check if this should be final synthesis (max rounds reached)
|
||
|
|
has_tool_messages = any(isinstance(msg, ToolMessage) for msg in messages)
|
||
|
|
is_final_synthesis = has_tool_messages and current_round >= max_rounds
|
||
|
|
|
||
|
|
if is_final_synthesis:
|
||
|
|
logger.info("Starting final synthesis phase - no more tool calls allowed")
|
||
|
|
# ✅ STEP 1: Final synthesis with tools disabled from the start
|
||
|
|
# Disable tools to prevent any tool calling during synthesis
|
||
|
|
try:
|
||
|
|
original_tools = llm_client.bind_tools([], force_tool_choice=False) # Disable tools
|
||
|
|
|
||
|
|
if not stream_callback:
|
||
|
|
# No streaming callback, generate final response without tools
|
||
|
|
draft = await llm_client.ainvoke(list(messages))
|
||
|
|
return {"messages": [draft]}
|
||
|
|
|
||
|
|
# ✅ STEP 2: Streaming final synthesis with improved HTML comment filtering
|
||
|
|
response_content = ""
|
||
|
|
accumulated_content = ""
|
||
|
|
|
||
|
|
async for token in llm_client.astream(list(messages)):
|
||
|
|
accumulated_content += token
|
||
|
|
response_content += token
|
||
|
|
|
||
|
|
# Check for complete HTML comments in accumulated content
|
||
|
|
while "<!--" in accumulated_content and "-->" in accumulated_content:
|
||
|
|
comment_start = accumulated_content.find("<!--")
|
||
|
|
comment_end = accumulated_content.find("-->", comment_start)
|
||
|
|
|
||
|
|
if comment_start >= 0 and comment_end >= 0:
|
||
|
|
# Send content before comment
|
||
|
|
before_comment = accumulated_content[:comment_start]
|
||
|
|
if stream_callback and before_comment:
|
||
|
|
await stream_callback(create_token_event(before_comment))
|
||
|
|
|
||
|
|
# Skip the comment and continue with content after
|
||
|
|
accumulated_content = accumulated_content[comment_end + 3:]
|
||
|
|
else:
|
||
|
|
break
|
||
|
|
|
||
|
|
# Send accumulated content if no pending comment
|
||
|
|
if "<!--" not in accumulated_content:
|
||
|
|
if stream_callback and accumulated_content:
|
||
|
|
await stream_callback(create_token_event(accumulated_content))
|
||
|
|
accumulated_content = ""
|
||
|
|
|
||
|
|
# Send any remaining content (if not in middle of comment)
|
||
|
|
if accumulated_content and "<!--" not in accumulated_content:
|
||
|
|
if stream_callback:
|
||
|
|
await stream_callback(create_token_event(accumulated_content))
|
||
|
|
|
||
|
|
return {"messages": [AIMessage(content=response_content)]}
|
||
|
|
|
||
|
|
finally:
|
||
|
|
# ✅ STEP 3: Restore tool binding for next interaction
|
||
|
|
llm_client.bind_tools(tool_schemas, force_tool_choice=True)
|
||
|
|
|
||
|
|
else:
|
||
|
|
logger.info(f"User Manual tool calling round {current_round + 1}/{max_rounds}")
|
||
|
|
|
||
|
|
# ✅ STEP 1: Non-streaming detection to check for tool needs
|
||
|
|
draft = await llm_client.ainvoke_with_tools(list(messages))
|
||
|
|
|
||
|
|
# ✅ STEP 2: If draft has tool_calls, execute them within this node
|
||
|
|
if isinstance(draft, AIMessage) and hasattr(draft, 'tool_calls') and draft.tool_calls:
|
||
|
|
logger.info(f"Detected {len(draft.tool_calls)} tool calls, executing within user manual agent")
|
||
|
|
|
||
|
|
# Create a new state with the tool call message added
|
||
|
|
tool_call_state = state.copy()
|
||
|
|
updated_messages = state["messages"].copy()
|
||
|
|
updated_messages.append(draft)
|
||
|
|
tool_call_state["messages"] = updated_messages
|
||
|
|
|
||
|
|
# Execute the tools using the existing streaming tool execution function
|
||
|
|
tool_results = await run_user_manual_tools_with_streaming(tool_call_state)
|
||
|
|
tool_messages = tool_results.get("messages", [])
|
||
|
|
|
||
|
|
# Increment tool round counter for next iteration
|
||
|
|
new_tool_rounds = current_round + 1
|
||
|
|
logger.info(f"Incremented user manual tool_rounds to {new_tool_rounds}")
|
||
|
|
|
||
|
|
# Continue with another round if under max rounds
|
||
|
|
if new_tool_rounds < max_rounds:
|
||
|
|
# Recursive call for next round with all messages
|
||
|
|
final_messages = updated_messages + tool_messages
|
||
|
|
recursive_state = state.copy()
|
||
|
|
recursive_state["messages"] = final_messages
|
||
|
|
recursive_state["tool_rounds"] = new_tool_rounds
|
||
|
|
return await user_manual_agent_node(recursive_state)
|
||
|
|
else:
|
||
|
|
# Max rounds reached, force final synthesis
|
||
|
|
logger.info("Max tool rounds reached, forcing final synthesis")
|
||
|
|
# Update messages for final synthesis
|
||
|
|
messages = updated_messages + tool_messages
|
||
|
|
# Continue to final synthesis below
|
||
|
|
|
||
|
|
# ✅ STEP 3: No tool_calls needed or max rounds reached → Enter final synthesis with streaming
|
||
|
|
# Temporarily disable tools to prevent accidental tool calling during synthesis
|
||
|
|
try:
|
||
|
|
llm_client.bind_tools([], force_tool_choice=False) # Disable tools
|
||
|
|
|
||
|
|
if not stream_callback:
|
||
|
|
# No streaming callback, use the draft we already have
|
||
|
|
return {"messages": [draft]}
|
||
|
|
|
||
|
|
# ✅ STEP 4: Streaming final synthesis with improved HTML comment filtering
|
||
|
|
response_content = ""
|
||
|
|
accumulated_content = ""
|
||
|
|
|
||
|
|
async for token in llm_client.astream(list(messages)):
|
||
|
|
accumulated_content += token
|
||
|
|
response_content += token
|
||
|
|
|
||
|
|
# Check for complete HTML comments in accumulated content
|
||
|
|
while "<!--" in accumulated_content and "-->" in accumulated_content:
|
||
|
|
comment_start = accumulated_content.find("<!--")
|
||
|
|
comment_end = accumulated_content.find("-->", comment_start)
|
||
|
|
|
||
|
|
if comment_start >= 0 and comment_end >= 0:
|
||
|
|
# Send content before comment
|
||
|
|
before_comment = accumulated_content[:comment_start]
|
||
|
|
if stream_callback and before_comment:
|
||
|
|
await stream_callback(create_token_event(before_comment))
|
||
|
|
|
||
|
|
# Skip the comment and continue with content after
|
||
|
|
accumulated_content = accumulated_content[comment_end + 3:]
|
||
|
|
else:
|
||
|
|
break
|
||
|
|
|
||
|
|
# Send accumulated content if no pending comment
|
||
|
|
if "<!--" not in accumulated_content:
|
||
|
|
if stream_callback and accumulated_content:
|
||
|
|
await stream_callback(create_token_event(accumulated_content))
|
||
|
|
accumulated_content = ""
|
||
|
|
|
||
|
|
# Send any remaining content (if not in middle of comment)
|
||
|
|
if accumulated_content and "<!--" not in accumulated_content:
|
||
|
|
if stream_callback:
|
||
|
|
await stream_callback(create_token_event(accumulated_content))
|
||
|
|
|
||
|
|
return {"messages": [AIMessage(content=response_content)]}
|
||
|
|
|
||
|
|
finally:
|
||
|
|
# ✅ STEP 5: Restore tool binding for next interaction
|
||
|
|
llm_client.bind_tools(tool_schemas, force_tool_choice=True)
|
||
|
|
|
||
|
|
|
||
|
|
def render_conversation_history(messages, max_messages: int = 10) -> str:
|
||
|
|
"""Render conversation history for context"""
|
||
|
|
recent_messages = messages[-max_messages:] if len(messages) > max_messages else messages
|
||
|
|
lines = []
|
||
|
|
for msg in recent_messages:
|
||
|
|
if hasattr(msg, 'content'):
|
||
|
|
content = msg.content
|
||
|
|
if isinstance(content, str):
|
||
|
|
# Determine message type by class name or other attributes
|
||
|
|
if 'Human' in str(type(msg)):
|
||
|
|
lines.append(f"<user>{content}</user>")
|
||
|
|
elif 'AI' in str(type(msg)):
|
||
|
|
lines.append(f"<ai>{content}</ai>")
|
||
|
|
elif isinstance(content, list):
|
||
|
|
content_str = " ".join([str(item) for item in content if isinstance(item, str)])
|
||
|
|
if 'Human' in str(type(msg)):
|
||
|
|
lines.append(f"<user>{content_str}</user>")
|
||
|
|
elif 'AI' in str(type(msg)):
|
||
|
|
lines.append(f"<ai>{content_str}</ai>")
|
||
|
|
return "\n".join(lines)
|
||
|
|
|
||
|
|
|
||
|
|
# User Manual Tools routing condition
|
||
|
|
def user_manual_should_continue(state: AgentState) -> Literal["user_manual_tools", "user_manual_agent", "post_process"]:
|
||
|
|
"""
|
||
|
|
Routing logic for user manual agent:
|
||
|
|
- has tool_calls → route to user_manual_tools
|
||
|
|
- no tool_calls → route to post_process (final synthesis already completed)
|
||
|
|
"""
|
||
|
|
messages = state["messages"]
|
||
|
|
if not messages:
|
||
|
|
logger.info("user_manual_should_continue: No messages, routing to post_process")
|
||
|
|
return "post_process"
|
||
|
|
|
||
|
|
last_message = messages[-1]
|
||
|
|
current_round = state.get("tool_rounds", 0)
|
||
|
|
# Get max_tool_rounds_user_manual from state, fallback to config if not set
|
||
|
|
max_rounds = state.get("max_tool_rounds_user_manual", None)
|
||
|
|
if max_rounds is None:
|
||
|
|
app_config = get_cached_config()
|
||
|
|
max_rounds = app_config.app.max_tool_rounds_user_manual
|
||
|
|
|
||
|
|
logger.info(f"user_manual_should_continue: Last message type: {type(last_message)}, tool_rounds: {current_round}/{max_rounds}")
|
||
|
|
|
||
|
|
# If last message is AI message with tool calls, route to tools
|
||
|
|
if isinstance(last_message, AIMessage):
|
||
|
|
has_tool_calls = hasattr(last_message, 'tool_calls') and last_message.tool_calls
|
||
|
|
logger.info(f"user_manual_should_continue: AI message has tool_calls: {has_tool_calls}")
|
||
|
|
|
||
|
|
if has_tool_calls:
|
||
|
|
logger.info("user_manual_should_continue: Routing to user_manual_tools")
|
||
|
|
return "user_manual_tools"
|
||
|
|
else:
|
||
|
|
# No tool calls = final synthesis already completed in user_manual_agent_node
|
||
|
|
logger.info("user_manual_should_continue: No tool calls, routing to post_process")
|
||
|
|
return "post_process"
|
||
|
|
|
||
|
|
# If last message is tool message(s), continue with agent for next round or final synthesis
|
||
|
|
if isinstance(last_message, ToolMessage):
|
||
|
|
logger.info("user_manual_should_continue: Tool message completed, continuing to user_manual_agent")
|
||
|
|
return "user_manual_agent"
|
||
|
|
|
||
|
|
logger.info("user_manual_should_continue: Routing to post_process")
|
||
|
|
return "post_process"
|
||
|
|
|
||
|
|
|
||
|
|
# User Manual Tools node with streaming support
|
||
|
|
async def run_user_manual_tools_with_streaming(state: AgentState, config: Optional[RunnableConfig] = None) -> Dict[str, Any]:
|
||
|
|
"""Execute user manual tools with streaming events - supports parallel execution"""
|
||
|
|
messages = state["messages"]
|
||
|
|
last_message = messages[-1]
|
||
|
|
|
||
|
|
# Get stream callback from context variable
|
||
|
|
from .graph import stream_callback_context
|
||
|
|
stream_callback = stream_callback_context.get()
|
||
|
|
|
||
|
|
if not isinstance(last_message, AIMessage) or not hasattr(last_message, 'tool_calls'):
|
||
|
|
return {"messages": []}
|
||
|
|
|
||
|
|
tool_calls = last_message.tool_calls or []
|
||
|
|
tool_results = []
|
||
|
|
new_messages = []
|
||
|
|
|
||
|
|
# User manual tools mapping
|
||
|
|
tools_map = get_user_manual_tools_by_name()
|
||
|
|
|
||
|
|
async def execute_single_tool(tool_call):
|
||
|
|
"""Execute a single user manual tool call with enhanced error handling"""
|
||
|
|
# Get stream callback from context
|
||
|
|
from .graph import stream_callback_context
|
||
|
|
stream_callback = stream_callback_context.get()
|
||
|
|
|
||
|
|
# Apply error handling decorator
|
||
|
|
@handle_async_errors(
|
||
|
|
ErrorCategory.TOOL,
|
||
|
|
ErrorCode.TOOL_ERROR,
|
||
|
|
stream_callback,
|
||
|
|
tool_call.get("id", "unknown") if isinstance(tool_call, dict) else "unknown"
|
||
|
|
)
|
||
|
|
async def _execute():
|
||
|
|
# Validate tool_call format
|
||
|
|
if not isinstance(tool_call, dict):
|
||
|
|
raise ValueError(f"Tool call must be dict, got {type(tool_call)}")
|
||
|
|
|
||
|
|
tool_name = tool_call.get("name")
|
||
|
|
tool_args = tool_call.get("args", {})
|
||
|
|
tool_id = tool_call.get("id", "unknown")
|
||
|
|
|
||
|
|
if not tool_name:
|
||
|
|
raise ValueError("Tool call missing 'name' field")
|
||
|
|
|
||
|
|
if tool_name not in tools_map:
|
||
|
|
available_tools = list(tools_map.keys())
|
||
|
|
raise ValueError(f"Tool '{tool_name}' not found. Available user manual tools: {available_tools}")
|
||
|
|
|
||
|
|
tool_func = tools_map[tool_name]
|
||
|
|
|
||
|
|
# Stream tool start event
|
||
|
|
if stream_callback:
|
||
|
|
await stream_callback(create_tool_start_event(tool_id, tool_name, tool_args))
|
||
|
|
|
||
|
|
import time
|
||
|
|
start_time = time.time()
|
||
|
|
|
||
|
|
try:
|
||
|
|
# Execute the user manual tool
|
||
|
|
result = await tool_func.ainvoke(tool_args)
|
||
|
|
|
||
|
|
# Calculate execution time
|
||
|
|
took_ms = int((time.time() - start_time) * 1000)
|
||
|
|
|
||
|
|
# Stream tool result event
|
||
|
|
if stream_callback:
|
||
|
|
await stream_callback(create_tool_result_event(tool_id, tool_name, result, took_ms))
|
||
|
|
|
||
|
|
# Create tool message
|
||
|
|
tool_message = ToolMessage(
|
||
|
|
content=str(result),
|
||
|
|
tool_call_id=tool_id,
|
||
|
|
name=tool_name
|
||
|
|
)
|
||
|
|
|
||
|
|
return tool_message, {"name": tool_name, "result": result, "took_ms": took_ms}
|
||
|
|
|
||
|
|
except Exception as e:
|
||
|
|
took_ms = int((time.time() - start_time) * 1000)
|
||
|
|
error_msg = get_user_message(ErrorCategory.TOOL)
|
||
|
|
|
||
|
|
# Stream tool error event
|
||
|
|
if stream_callback:
|
||
|
|
await stream_callback(create_tool_error_event(tool_id, tool_name, error_msg))
|
||
|
|
|
||
|
|
# Create error tool message
|
||
|
|
tool_message = ToolMessage(
|
||
|
|
content=f"Error executing {tool_name}: {error_msg}",
|
||
|
|
tool_call_id=tool_id,
|
||
|
|
name=tool_name
|
||
|
|
)
|
||
|
|
|
||
|
|
return tool_message, {"name": tool_name, "error": error_msg, "took_ms": took_ms}
|
||
|
|
|
||
|
|
return await _execute()
|
||
|
|
|
||
|
|
# Execute user manual tools (typically just one for user manual retrieval)
|
||
|
|
import asyncio
|
||
|
|
tasks = [execute_single_tool(tool_call) for tool_call in tool_calls]
|
||
|
|
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||
|
|
|
||
|
|
for i, result in enumerate(results):
|
||
|
|
if isinstance(result, Exception):
|
||
|
|
# Handle execution exception
|
||
|
|
tool_call = tool_calls[i]
|
||
|
|
tool_id = tool_call.get("id", f"error_{i}") or f"error_{i}"
|
||
|
|
tool_name = tool_call.get("name", "unknown")
|
||
|
|
error_msg = get_user_message(ErrorCategory.TOOL)
|
||
|
|
|
||
|
|
if stream_callback:
|
||
|
|
await stream_callback(create_tool_error_event(tool_id, tool_name, error_msg))
|
||
|
|
|
||
|
|
error_message = ToolMessage(
|
||
|
|
content=f"Error executing {tool_name}: {error_msg}",
|
||
|
|
tool_call_id=tool_id,
|
||
|
|
name=tool_name
|
||
|
|
)
|
||
|
|
new_messages.append(error_message)
|
||
|
|
elif isinstance(result, tuple) and len(result) == 2:
|
||
|
|
# result is a tuple: (tool_message, tool_result)
|
||
|
|
tool_message, tool_result = result
|
||
|
|
new_messages.append(tool_message)
|
||
|
|
tool_results.append(tool_result)
|
||
|
|
else:
|
||
|
|
# Unexpected result format
|
||
|
|
logger.error(f"Unexpected tool execution result format: {type(result)}")
|
||
|
|
continue
|
||
|
|
|
||
|
|
return {"messages": new_messages, "tool_results": tool_results}
|
||
|
|
|
||
|
|
|
||
|
|
# Legacy function for backward compatibility
|
||
|
|
async def user_manual_rag_node(state: AgentState, config: Optional[RunnableConfig] = None) -> Dict[str, Any]:
|
||
|
|
"""
|
||
|
|
Legacy user manual RAG node - redirects to new agent-based implementation
|
||
|
|
"""
|
||
|
|
logger.info("📚 USER_MANUAL_RAG_NODE: Redirecting to user_manual_agent_node")
|
||
|
|
return await user_manual_agent_node(state, config)
|