ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter2/local_llm_serving/agent.py

@@ -0,0 +1,591 @@
+"""
+vLLM Tool Calling Agent Implementation
+Demonstrates how to use vLLM with Qwen3 for tool calling
+"""
+
+import json
+import uuid
+import logging
+from concurrent.futures import ThreadPoolExecutor
+from typing import List, Dict, Any, Optional, Tuple
+from openai import OpenAI
+from tools import ToolRegistry
+from config import OPENAI_API_BASE, OPENAI_API_KEY, LOG_LEVEL
+
+# Set up logging
+logging.basicConfig(level=LOG_LEVEL)
+logger = logging.getLogger(__name__)
+
+
+class VLLMToolAgent:
+    """Agent that uses vLLM for tool calling with Qwen3 model"""
+    
+    def __init__(self, api_base: str = OPENAI_API_BASE, api_key: str = OPENAI_API_KEY):
+        """
+        Initialize the agent with vLLM server connection
+        
+        Args:
+            api_base: Base URL for vLLM server
+            api_key: API key (not required for vLLM, use "EMPTY")
+        """
+        self.client = OpenAI(
+            api_key=api_key,
+            base_url=api_base
+        )
+        self.tool_registry = ToolRegistry()
+        self.conversation_history = []
+        logger.info(f"Initialized VLLMToolAgent with server at {api_base}")
+    
+    def _format_system_prompt_with_tools(self) -> str:
+        """
+        Format the system prompt with available tools in Qwen3 format
+        """
+        tools_json = json.dumps(self.tool_registry.get_tool_schemas(), indent=2)
+        
+        system_prompt = f"""# Tools
+
+You may call one or more functions to assist with the user query.
+
+You are provided with function signatures within <tools></tools> XML tags:
+<tools>
+{tools_json}
+</tools>
+
+For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:
+<tool_call>
+{{"name": <function-name>, "arguments": <args-json-object>}}
+</tool_call>
+
+You are a helpful assistant that can use tools to answer questions and perform tasks.
+When you need to use a tool, generate the appropriate tool call.
+After receiving tool results, use them to provide a comprehensive answer to the user."""
+        
+        return system_prompt
+    
+    def _parse_tool_calls(self, content: str) -> List[Dict[str, Any]]:
+        """
+        Parse tool calls from model output
+        Extracts content between <tool_call> tags
+        """
+        tool_calls = []
+        
+        # Find all tool call blocks
+        import re
+        pattern = r'<tool_call>(.*?)</tool_call>'
+        matches = re.findall(pattern, content, re.DOTALL)
+        
+        for match in matches:
+            try:
+                tool_call = json.loads(match.strip())
+                if "name" in tool_call and "arguments" in tool_call:
+                    tool_calls.append({
+                        "id": str(uuid.uuid4())[:8],  # Generate short ID
+                        "type": "function",
+                        "function": {
+                            "name": tool_call["name"],
+                            "arguments": tool_call["arguments"]
+                        }
+                    })
+                    logger.debug(f"Parsed tool call: {tool_call['name']}")
+            except json.JSONDecodeError as e:
+                logger.warning(f"Failed to parse tool call JSON: {e}")
+                logger.debug(f"Content was: {match}")
+        
+        return tool_calls
+    
+    def _execute_tool_calls(self, tool_calls: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """
+        Execute tool calls and return results.
+        Multiple tool calls in the same turn are executed in parallel (they are
+        independent by construction, since the model generated all of them
+        without seeing any result). Ensures that error messages from failed
+        tool executions are properly formatted.
+        """
+        def run_one(tool_call: Dict[str, Any]) -> Dict[str, Any]:
+            tool_name = tool_call["function"]["name"]
+            tool_args = tool_call["function"]["arguments"]
+            tool_id = tool_call["id"]
+            
+            logger.info(f"Executing tool: {tool_name} with args: {tool_args}")
+            
+            # Execute the tool
+            result = self.tool_registry.execute_tool(tool_name, tool_args)
+            
+            # Check if the result indicates an error
+            try:
+                result_dict = json.loads(result) if isinstance(result, str) else result
+                if isinstance(result_dict, dict) and not result_dict.get("success", True):
+                    # Tool execution failed - format error message clearly
+                    error_msg = f"❌ Tool '{tool_name}' execution failed:\n"
+                    if "error" in result_dict:
+                        error_msg += f"Error: {result_dict['error']}\n"
+                    if "error_type" in result_dict:
+                        error_msg += f"Type: {result_dict['error_type']}\n"
+                    if "traceback" in result_dict:
+                        error_msg += f"Traceback:\n{result_dict['traceback']}\n"
+                    
+                    logger.error(f"Tool {tool_name} failed: {result_dict.get('error', 'Unknown error')}")
+                    result = error_msg
+                else:
+                    logger.debug(f"Tool {tool_name} returned: {result}")
+            except (json.JSONDecodeError, TypeError):
+                # Result is not JSON, just pass it through
+                logger.debug(f"Tool {tool_name} returned: {result}")
+            
+            # Format the result
+            return {
+                "role": "tool",
+                "tool_call_id": tool_id,
+                "name": tool_name,
+                "content": result if isinstance(result, str) else str(result)
+            }
+        
+        if len(tool_calls) <= 1:
+            return [run_one(tc) for tc in tool_calls]
+        
+        # Independent tool calls run concurrently; executor.map preserves order
+        with ThreadPoolExecutor(max_workers=len(tool_calls)) as executor:
+            return list(executor.map(run_one, tool_calls))
+    
+    def _execute_single_tool(self, tool_data: Dict[str, Any]) -> Tuple[str, bool]:
+        """
+        Execute one parsed tool call ({"name": ..., "arguments": ...}).
+        Returns (result_text, is_error) with error messages formatted clearly.
+        """
+        tool_name = tool_data["name"]
+        try:
+            result = self.tool_registry.execute_tool(tool_name, tool_data["arguments"])
+        except Exception as e:
+            logger.error(f"Tool execution error: {e}")
+            return f"❌ Tool execution exception: {str(e)}", True
+        
+        # Check if the result indicates an error
+        try:
+            result_dict = json.loads(result) if isinstance(result, str) else result
+            if isinstance(result_dict, dict) and not result_dict.get("success", True):
+                # Tool execution failed - format error message clearly
+                error_msg = f"❌ Tool '{tool_name}' execution failed:\n"
+                if "error" in result_dict:
+                    error_msg += f"Error: {result_dict['error']}\n"
+                if "error_type" in result_dict:
+                    error_msg += f"Type: {result_dict['error_type']}\n"
+                if "traceback" in result_dict:
+                    error_msg += f"Traceback:\n{result_dict['traceback']}\n"
+                
+                logger.error(f"Tool {tool_name} failed: {result_dict.get('error', 'Unknown error')}")
+                return error_msg, True
+        except (json.JSONDecodeError, TypeError):
+            # Result is not JSON, just pass it through
+            pass
+        return result, False
+    
+    def chat(self, message: str, use_tools: bool = True,  
+             temperature: float = 0.3, max_tokens: int = 2048, 
+             stream: bool = False) -> str:
+        """
+        Send a message to the model and handle tool calls in a ReAct loop
+        
+        Args:
+            message: User message
+            use_tools: Whether to enable tool calling
+            temperature: Sampling temperature
+            max_tokens: Maximum tokens to generate
+            stream: Whether to stream the response
+            
+        Returns:
+            Final response from the model (or generator if streaming)
+        """
+        if stream:
+            return self.chat_stream(message, use_tools, temperature, max_tokens)
+        
+        # Add user message to history
+        self.conversation_history.append({"role": "user", "content": message})
+        
+        # Prepare messages with system prompt if using tools
+        messages = []
+        if use_tools:
+            messages.append({
+                "role": "system",
+                "content": self._format_system_prompt_with_tools()
+            })
+        else:
+            messages.append({
+                "role": "system",
+                "content": "You are a helpful assistant."
+            })
+        
+        # Add conversation history
+        messages.extend(self.conversation_history)
+        
+        # Prepare tools for the API call
+        tools = self.tool_registry.get_tool_schemas() if use_tools else None
+        
+        # ReAct loop - keep going until no more tool calls are needed
+        max_iterations = 10  # Prevent infinite loops
+        iteration = 0
+        final_response = ""
+        
+        while iteration < max_iterations:
+            iteration += 1
+            logger.info(f"ReAct iteration {iteration}")
+            
+            # Prepare messages for this iteration
+            messages = []
+            if use_tools:
+                messages.append({
+                    "role": "system",
+                    "content": self._format_system_prompt_with_tools()
+                })
+            else:
+                messages.append({
+                    "role": "system",
+                    "content": "You are a helpful assistant."
+                })
+            messages.extend(self.conversation_history)
+            
+            # Call the model
+            response = self.client.chat.completions.create(
+                model="Qwen/Qwen3-0.6B",
+                messages=messages,
+                tools=tools,
+                tool_choice="auto" if use_tools else None,
+                temperature=temperature,
+                max_tokens=max_tokens
+            )
+            
+            assistant_message = response.choices[0].message
+            content = assistant_message.content or ""
+            
+            # Read tool calls from the structured field. With
+            # enable_auto_tool_choice + the hermes parser, vLLM extracts the
+            # <tool_call> tags out of the text and returns them here instead of
+            # leaving them in `content` (which only holds <think> and final text).
+            tool_calls = []
+            if use_tools and assistant_message.tool_calls:
+                for tc in assistant_message.tool_calls:
+                    raw_args = tc.function.arguments
+                    try:
+                        parsed_args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
+                    except json.JSONDecodeError as e:
+                        logger.warning(f"Failed to parse tool arguments for {tc.function.name}: {e}")
+                        parsed_args = {}
+                    logger.info(f"Model requested tool call: {tc.function.name}({raw_args})")
+                    tool_calls.append({
+                        "id": tc.id,
+                        "type": "function",
+                        "function": {
+                            "name": tc.function.name,
+                            "arguments": parsed_args,
+                        },
+                    })
+            
+            if tool_calls:
+                logger.info(f"Model requested {len(tool_calls)} tool call(s)")
+                
+                # Add assistant message with tool calls to history
+                # (arguments must be a JSON string per the OpenAI API spec)
+                self.conversation_history.append({
+                    "role": "assistant",
+                    "content": content,
+                    "tool_calls": [
+                        {
+                            **tc,
+                            "function": {
+                                **tc["function"],
+                                "arguments": tc["function"]["arguments"]
+                                if isinstance(tc["function"]["arguments"], str)
+                                else json.dumps(tc["function"]["arguments"])
+                            }
+                        }
+                        for tc in tool_calls
+                    ]
+                })
+                
+                # Execute tool calls
+                tool_results = self._execute_tool_calls(tool_calls)
+                
+                # Add tool results to conversation
+                for result in tool_results:
+                    # Format tool response for Qwen3
+                    tool_response = f'<tool_response>\n{result["content"]}\n</tool_response>'
+                    self.conversation_history.append({
+                        "role": "user",  # Tool responses are treated as user messages in Qwen3
+                        "content": tool_response,
+                        "name": result.get("name", "tool")
+                    })
+                
+                # Continue the ReAct loop
+                continue
+            else:
+                # No tool calls - we have a final response
+                self.conversation_history.append({
+                    "role": "assistant",
+                    "content": content
+                })
+                final_response = content
+                break
+        
+        # Check if we hit max iterations
+        if iteration >= max_iterations:
+            logger.warning("Maximum iterations reached in ReAct loop")
+            final_response = "I've reached the maximum number of reasoning steps. " + final_response
+        
+        return final_response
+    
+    def reset_conversation(self):
+        """Reset the conversation history"""
+        self.conversation_history = []
+        logger.info("Conversation history reset")
+    
+    def chat_stream(self, message: str, use_tools: bool = True,
+                    temperature: float = 0.3, max_tokens: int = 2048):
+        """
+        Stream a message to the model and handle tool calls in a ReAct loop
+        
+        Yields chunks that include:
+        - type: 'thinking', 'tool_call', 'tool_result', 'content'
+        - content: The actual content
+        """
+        # Add user message to history
+        self.conversation_history.append({"role": "user", "content": message})
+        
+        # Prepare tools for the API call
+        tools = self.tool_registry.get_tool_schemas() if use_tools else None
+        
+        # ReAct loop - keep going until no more tool calls are needed
+        max_iterations = 10  # Prevent infinite loops
+        iteration = 0
+        
+        while iteration < max_iterations:
+            iteration += 1
+            logger.info(f"ReAct stream iteration {iteration}")
+            
+            # Prepare messages for this iteration
+            messages = []
+            if use_tools:
+                messages.append({
+                    "role": "system",
+                    "content": self._format_system_prompt_with_tools()
+                })
+            else:
+                messages.append({
+                    "role": "system",
+                    "content": "You are a helpful assistant."
+                })
+            messages.extend(self.conversation_history)
+            
+            # Stream response from model
+            stream_response = self.client.chat.completions.create(
+                model="Qwen/Qwen3-0.6B",
+                messages=messages,
+                tools=tools,
+                tool_choice="auto" if use_tools else None,
+                temperature=temperature,
+                max_tokens=max_tokens,
+                stream=True
+            )
+            
+            collected_content = []
+            thinking_buffer = ""
+            tool_call_parts = {}
+            
+            # Process the stream
+            for chunk in stream_response:
+                if chunk.choices and chunk.choices[0].delta:
+                    delta = chunk.choices[0].delta
+                    if delta.content:
+                        content_chunk = delta.content
+                        collected_content.append(content_chunk)
+                        
+                        # Check if this is internal thinking (between <think> tags)
+                        if '<think>' in content_chunk or thinking_buffer:
+                            thinking_buffer += content_chunk
+                            if '</think>' in thinking_buffer:
+                                # Extract and yield thinking
+                                import re
+                                thinking_match = re.search(r'<think>(.*?)</think>', thinking_buffer, re.DOTALL)
+                                if thinking_match:
+                                    # Stream thinking character by character
+                                    for char in thinking_match.group(1).strip():
+                                        yield {"type": "thinking", "content": char}
+                                remaining = re.sub(
+                                    r'<think>.*?</think>', '', thinking_buffer, flags=re.DOTALL
+                                )
+                                thinking_buffer = ""
+                                if remaining:
+                                    yield {"type": "content", "content": remaining}
+                        else:
+                            # Regular content
+                            yield {"type": "content", "content": content_chunk}
+
+                    # vLLM streams structured tool calls in fragments. Calls
+                    # are keyed by index because ids, names, and arguments may
+                    # arrive in separate chunks.
+                    for fragment in getattr(delta, "tool_calls", None) or []:
+                        index = getattr(fragment, "index", None)
+                        if index is None:
+                            index = 0
+                        buffered = tool_call_parts.setdefault(index, {
+                            "id": None,
+                            "type": "function",
+                            "function": {"name": "", "arguments": ""},
+                        })
+                        if getattr(fragment, "id", None):
+                            buffered["id"] = fragment.id
+                        if getattr(fragment, "type", None):
+                            buffered["type"] = fragment.type
+                        function = getattr(fragment, "function", None)
+                        if function:
+                            if getattr(function, "name", None):
+                                buffered["function"]["name"] += function.name
+                            if getattr(function, "arguments", None):
+                                buffered["function"]["arguments"] += function.arguments
+
+            # Save complete response and structured calls to history before
+            # adding tool results, matching the non-streaming message order.
+            complete_response = ''.join(collected_content)
+
+            if tool_call_parts:
+                pending_tool_calls = []
+                parse_errors = []
+                assistant_tool_calls = []
+
+                for index in sorted(tool_call_parts):
+                    buffered = tool_call_parts[index]
+                    call_id = buffered["id"] or str(uuid.uuid4())[:8]
+                    tool_name = buffered["function"]["name"] or "unknown"
+                    raw_args = buffered["function"]["arguments"] or "{}"
+                    assistant_tool_calls.append({
+                        "id": call_id,
+                        "type": buffered["type"],
+                        "function": {
+                            "name": tool_name,
+                            "arguments": raw_args,
+                        },
+                    })
+
+                    try:
+                        parsed_args = json.loads(raw_args)
+                    except json.JSONDecodeError as e:
+                        error_msg = f"❌ Tool call parse exception: {str(e)}"
+                        logger.error(f"Tool call parse error: {e}")
+                        parse_errors.append((tool_name, error_msg))
+                        yield {"type": "tool_error", "content": error_msg}
+                        continue
+
+                    tool_data = {
+                        "id": call_id,
+                        "name": tool_name,
+                        "arguments": parsed_args,
+                    }
+                    pending_tool_calls.append(tool_data)
+                    yield {
+                        "type": "tool_call",
+                        "content": {
+                            "name": tool_name,
+                            "arguments": parsed_args,
+                        },
+                    }
+
+                self.conversation_history.append({
+                    "role": "assistant",
+                    "content": complete_response,
+                    "tool_calls": assistant_tool_calls,
+                })
+
+                for tool_name, error_msg in parse_errors:
+                    self.conversation_history.append({
+                        "role": "user",
+                        "content": f'<tool_response>\n{error_msg}\n</tool_response>',
+                        "name": tool_name,
+                    })
+
+                # Execute all valid tool calls from this turn in parallel.
+                if not pending_tool_calls:
+                    outcomes = []
+                elif len(pending_tool_calls) == 1:
+                    outcomes = [self._execute_single_tool(pending_tool_calls[0])]
+                else:
+                    with ThreadPoolExecutor(max_workers=len(pending_tool_calls)) as executor:
+                        outcomes = list(executor.map(self._execute_single_tool, pending_tool_calls))
+                
+                for tool_data, (result, is_error) in zip(pending_tool_calls, outcomes):
+                    if is_error:
+                        yield {"type": "tool_error", "content": result}
+                    else:
+                        yield {"type": "tool_result", "content": result}
+                    
+                    # Add to history
+                    self.conversation_history.append({
+                        "role": "user",
+                        "content": f'<tool_response>\n{result}\n</tool_response>',
+                        "name": tool_data["name"]
+                    })
+
+                # Continue the ReAct loop - let the model decide what to do next
+                continue
+            else:
+                # No tool calls - we have a final response
+                self.conversation_history.append({
+                    "role": "assistant",
+                    "content": complete_response
+                })
+                # Exit the ReAct loop
+                break
+        
+        # Check if we hit max iterations
+        if iteration >= max_iterations:
+            yield {"type": "error", "content": "Maximum iterations reached in ReAct loop"}
+    
+    def get_conversation_history(self) -> List[Dict[str, Any]]:
+        """Get the current conversation history"""
+        return self.conversation_history
+    
+    def add_custom_tool(self, name: str, function: callable, 
+                       description: str, parameters: Dict):
+        """
+        Add a custom tool to the registry
+        
+        Args:
+            name: Tool name
+            function: Callable function
+            description: Tool description
+            parameters: OpenAI-style parameter schema
+        """
+        self.tool_registry.register_tool(name, function, description, parameters)
+        logger.info(f"Added custom tool: {name}")
+
+
+def demonstrate_tool_calling():
+    """Demonstrate the tool calling functionality"""
+    print("=" * 60)
+    print("vLLM Tool Calling Demo with Qwen3")
+    print("=" * 60)
+    
+    # Initialize agent
+    agent = VLLMToolAgent()
+    
+    # Test cases
+    test_queries = [
+        "What's the current temperature in Paris, France?",
+        "Calculate 15 * 23 + sqrt(144)",
+        "What time is it in Tokyo (JST)?",
+        "Search for information about vLLM tool calling",
+        "What's the weather in Dubai and what's 100 fahrenheit in celsius?",
+    ]
+    
+    for i, query in enumerate(test_queries, 1):
+        print(f"\n--- Test {i} ---")
+        print(f"User: {query}")
+        
+        response = agent.chat(query)
+        print(f"Assistant: {response}")
+        
+        # Reset conversation for next test
+        agent.reset_conversation()
+        print("-" * 40)
+
+
+if __name__ == "__main__":
+    # Run demonstration
+    demonstrate_tool_calling()