local_llm_serving¶
第2章 · 上下文工程 · 配套项目
chapter2/local_llm_serving
项目说明¶
Universal Tool Calling Demo¶
A cross-platform demonstration of LLM tool calling using standard OpenAI-compatible APIs. Works seamlessly on Windows, macOS, and Linux by automatically selecting the best backend for your system.
🌟 Features¶
- Universal Compatibility: Single entry point (
main.py) that works on all platforms - Automatic Backend Selection:
- Uses vLLM on Linux/Windows with NVIDIA GPU
- Uses Ollama on macOS, Windows, or Linux without GPU
- Standard Tool Calling: Only uses OpenAI-compatible tool calling format
- Built-in Tools: Weather, calculator, time, and easy to add custom tools
- Interactive & Example Modes: Test with examples or chat interactively
- 🆕 Streaming Support: Real-time display of thinking process, tool calls, and responses
🚀 Quick Start¶
# 1. Clone the repository
git clone <repository>
cd chapter2/local_llm_serving
# 2. Install dependencies
pip install -r requirements.txt
# 3. Check your system compatibility
python check_compatibility.py
# 4. Run the main script (auto-detects best backend)
python main.py
That's it! The script automatically detects your platform and uses the appropriate backend.
📋 Prerequisites¶
All Platforms¶
- Python 3.10+
pip install -r requirements.txt
Platform-Specific Setup¶
🍎 macOS¶
# Install Ollama
brew install ollama
# Start Ollama service (in separate terminal)
ollama serve
# Download a model
ollama pull qwen3:0.6b
🪟 Windows¶
With NVIDIA GPU: - CUDA toolkit installed - NVIDIA drivers 452.39+ - vLLM will be used automatically
Without GPU:
# Download and install Ollama
# From: https://ollama.com/download/windows
# Pull a model
ollama pull qwen3:0.6b
🐧 Linux¶
With NVIDIA GPU: - CUDA toolkit installed - vLLM will be used automatically
Without GPU:
# Install Ollama
curl -fsSL https://ollama.com/install.sh | sh
# Start service
systemctl start ollama
# Pull a model
ollama pull qwen3:0.6b
🎮 Usage¶
Basic Usage¶
# Run with auto-detection (recommended)
python main.py
# Run examples only
python main.py --mode examples
# Run interactive mode only
python main.py --mode interactive
# Force specific backend
python main.py --backend ollama # Force Ollama
python main.py --backend vllm # Force vLLM (requires GPU)
# Show system info
python main.py --info
Using in Your Code¶
from main import ToolCallingAgent
# Initialize (auto-detects best backend)
agent = ToolCallingAgent()
# Send a message
response = agent.chat("What's the weather in Tokyo?")
print(response)
# Disable tools for a query
response = agent.chat("Tell me a joke", use_tools=False)
# Reset conversation
agent.reset_conversation()
Adding Custom Tools¶
from tools import ToolRegistry
# Get the tool registry
registry = ToolRegistry()
# Define your tool function
def my_custom_tool(param1: str, param2: int) -> str:
return f"Processed {param1} with {param2}"
# Register it
registry.register_tool(
name="my_custom_tool",
function=my_custom_tool,
description="My custom tool description",
parameters={
"type": "object",
"properties": {
"param1": {"type": "string", "description": "First parameter"},
"param2": {"type": "integer", "description": "Second parameter"}
},
"required": ["param1", "param2"]
}
)
📁 Project Structure¶
local_llm_serving/
├── main.py # Main entry point (auto-detects backend)
├── benchmark.py # Serving benchmark: throughput / TTFT / KV cache / batching
├── agent.py # vLLM agent implementation
├── ollama_native.py # Ollama native tool calling
├── tools.py # Tool implementations
├── config.py # Configuration settings
├── server.py # vLLM server manager
├── check_compatibility.py # System compatibility checker
├── requirements.txt # Python dependencies
├── env.example # Environment variables template
└── README.md # This file
🛠️ Available Tools¶
- get_current_temperature: Get real-time weather information using Open-Meteo API (no API key required)
- get_current_time: Get current time in different timezones
- convert_currency: Convert between different currencies (simulated rates)
- parse_pdf: Parse PDF documents from URL or local file
- code_interpreter: Execute Python code for complex calculations and data processing
🎬 Streaming Mode¶
The agents now support streaming responses, which displays: - 🧠 Internal thinking process (shown in gray) - 🔧 Tool calls as they happen - ✓ Tool results in real-time - 📝 Final response streamed character by character
Using Streaming¶
Interactive Mode (Default)¶
# Streaming is enabled by default
python main.py
# Disable streaming
python main.py --no-stream
# Toggle streaming during chat with /stream command
Programmatic Usage¶
from main import ToolCallingAgent
# Initialize agent
agent = ToolCallingAgent()
# Stream response
for chunk in agent.chat("What's the weather in Tokyo?", stream=True):
chunk_type = chunk.get("type")
content = chunk.get("content", "")
if chunk_type == "thinking":
print(f"Thinking: {content}")
elif chunk_type == "tool_call":
print(f"Tool: {content['name']}")
elif chunk_type == "tool_result":
print(f"Result: {content}")
elif chunk_type == "content":
print(content, end="", flush=True)
Test Streaming¶
# Run streaming demo
python demo_streaming.py
# Compare streaming vs regular mode
python test_streaming.py --mode compare
📈 Serving Benchmark (benchmark.py)¶
benchmark.py 是实验 2-1 的配套基准，用于测量本地部署的小模型在 服务（serving） 层面的核心指标，帮助建立对吞吐 / 延迟 / 批处理 / KV Cache 的直觉。它通过 OpenAI 兼容接口工作，vLLM 与 Ollama 均可。
所有数字都来自真实服务端的实测，脚本本身不产生任何合成数据。 如果服务端尚未启动，可用 --dry-run 离线查看每个场景将要发出的请求配置。
场景（--scenario）¶
| 场景 | 说明 | 对应书中要点 |
|---|---|---|
throughput |
单流解码吞吐（tok/s）与首 token 延迟（TTFT） | 实验 2-1 第 2 点：M2 上 >100 tok/s |
kv-cache |
前缀缓存 命中 vs 未命中 的 TTFT 对比 | 实验 2-1 第 5 点：改动系统提示词开头 → 缓存失效、需重算整个前缀 |
batching |
不同并发度下的聚合吞吐（批处理权衡） | 连续批处理如何提升系统吞吐 |
all |
依次运行以上全部场景（默认） | — |
用法¶
# 1. 先启动服务端（二选一）
python server.py # vLLM（需要 NVIDIA GPU）
ollama serve && ollama pull qwen3:0.6b # Ollama（Mac / 无 GPU）
# 2. 运行基准
python benchmark.py --scenario all --output results.json # 跑全部并保存
python benchmark.py --scenario kv-cache --backend ollama # 只看 KV Cache TTFT 对比
python benchmark.py --scenario batching --concurrency 1,2,4,8 # 批处理吞吐扫描
# 离线查看计划（不访问服务端），可用于验证参数
python benchmark.py --dry-run
python benchmark.py --help
主要参数¶
--backend {vllm,ollama}：推断默认地址与模型名（vLLMQwen3-0.6B@:8000/v1，Ollamaqwen3:0.6b@:11434/v1）--base-url/--model/--api-key：覆盖默认连接配置--repeats：throughput/kv-cache的重复次数（默认 5）--max-tokens/--temperature：生成参数--prefix-tokens：kv-cache场景共享前缀的近似长度，越长缓存效果越明显（默认 1024）--concurrency：batching场景的并发度列表，逗号分隔（默认1,2,4,8）--output：将结果写入 JSON 文件
说明：
kv-cache依赖服务端的前缀缓存（vLLM 的 automatic prefix caching 默认开启）。命中组保持系统提示词逐字节不变；未命中组每次只在系统提示词开头插入一个唯一计数串，前缀被改写导致缓存全部失效——这正是书中「系统提示词一旦定下来就不要改」的实测演示。
🔧 Configuration¶
Copy env.example to .env and customize:
# For vLLM (if you have GPU)
MODEL_NAME=Qwen/Qwen3-0.6B
VLLM_HOST=localhost
VLLM_PORT=8000
# Logging
LOG_LEVEL=INFO
📊 Tool Calling Format¶
This project uses standard OpenAI-compatible tool calling:
{
"tool_calls": [{
"id": "call_123",
"type": "function",
"function": {
"name": "get_weather",
"arguments": {"location": "Tokyo"}
}
}]
}
No ad-hoc parsing or custom formats - just the standard that works across platforms.
🐛 Troubleshooting¶
"Ollama not found"¶
- Mac:
brew install ollama && ollama serve - Windows: Download from ollama.com
- Linux:
curl -fsSL https://ollama.com/install.sh | sh
"No models installed"¶
"CUDA not available" (Linux/Windows)¶
- Install NVIDIA drivers and CUDA toolkit
- Or the script will automatically use Ollama instead
Check System Compatibility¶
🤝 Supported Models¶
Default Model:¶
- Qwen3 (0.6B) - Default model used by this project. Small size with decent tool calling support.
Other Compatible Models for Tool Calling:¶
- Qwen3 (8B+) - Good tool support
- Llama 3.1/3.2 (8B+) - Good tool support
- Mistral Nemo - Great tool calling
For vLLM:¶
- Uses Qwen3-0.6B by default
- Any model supported by vLLM can be configured
📚 How It Works¶
- Platform Detection:
main.pydetects your OS and GPU availability - Backend Selection:
- Has NVIDIA GPU? → Uses vLLM for best performance
- No GPU or on Mac? → Uses Ollama for local inference
- Tool Execution: Both backends use standard OpenAI tool calling format
- Response Generation: Tools are executed and results fed back to the model
🔗 References¶
📄 License¶
This demo project is provided as-is for educational purposes.
源代码¶
agent.py¶
"""
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="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 ""
# Check for tool calls in the response
tool_calls = []
if use_tools and content:
tool_calls = self._parse_tool_calls(content)
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="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 = []
tool_calls_buffer = ""
tool_calls_detected = False
pending_tool_calls = []
# 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)
buffered_this_chunk = False
# Check if this is internal thinking (between <think> tags)
if '<think>' in content_chunk or tool_calls_buffer:
tool_calls_buffer += content_chunk
buffered_this_chunk = True
if '</think>' in tool_calls_buffer:
# Extract and yield thinking
import re
thinking_match = re.search(r'<think>(.*?)</think>', tool_calls_buffer, re.DOTALL)
if thinking_match:
# Stream thinking character by character
for char in thinking_match.group(1).strip():
yield {"type": "thinking", "content": char}
tool_calls_buffer = re.sub(r'<think>.*?</think>', '', tool_calls_buffer, flags=re.DOTALL)
# Check for tool calls
if '<tool_call>' in content_chunk or tool_calls_buffer:
if not buffered_this_chunk:
tool_calls_buffer += content_chunk
# Collect all complete tool calls; they are executed
# in parallel once the stream finishes
while '</tool_call>' in tool_calls_buffer:
tool_calls_detected = True
import re
tool_match = re.search(r'<tool_call>(.*?)</tool_call>', tool_calls_buffer, re.DOTALL)
if not tool_match:
break
try:
tool_data = json.loads(tool_match.group(1).strip())
pending_tool_calls.append(tool_data)
yield {"type": "tool_call", "content": tool_data}
except Exception as e:
error_msg = f"❌ Tool call parse exception: {str(e)}"
logger.error(f"Tool call parse error: {e}")
yield {"type": "tool_error", "content": error_msg}
# Add error to history so agent can see it
self.conversation_history.append({
"role": "user",
"content": f'<tool_response>\n{error_msg}\n</tool_response>',
"name": "unknown"
})
# Keep any trailing partial tool call for the next chunk
tool_calls_buffer = tool_calls_buffer[tool_match.end():]
else:
# Regular content
yield {"type": "content", "content": content_chunk}
# Execute all tool calls from this turn in parallel (they are
# independent by construction, since the model generated all of
# them without seeing any result)
if pending_tool_calls:
if 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"]
})
# Save complete response to history
complete_response = ''.join(collected_content)
if tool_calls_detected:
# Add assistant's message to history
self.conversation_history.append({
"role": "assistant",
"content": complete_response
})
# 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()
benchmark.py¶
#!/usr/bin/env python3
"""
本地 LLM 服务性能基准（实验 2-1 配套）
本脚本通过 OpenAI 兼容接口（vLLM 或 Ollama 均可）测量本地部署的小模型在
「服务（serving）」层面的三个核心指标，帮助读者建立对吞吐 / 延迟 / 批处理 /
KV Cache 的直觉：
1. throughput —— 单流解码吞吐（tokens/s）与首 token 延迟（TTFT）
2. kv-cache —— 前缀缓存命中 vs 未命中的 TTFT 对比
（对应实验 2-1 第 5 点：系统提示词不变时缓存命中更快，
修改系统提示词开头几个字符导致缓存失效、需重算整个前缀）
3. batching —— 不同并发度下的聚合吞吐，直观展示批处理带来的吞吐提升
所有数字均来自真实服务端的实测，脚本本身不产生任何合成数据。
如果尚未启动服务端，可用 --dry-run 离线查看每个场景将要发出的请求配置。
示例：
# 先启动服务端（二选一）
python server.py # vLLM（需要 NVIDIA GPU）
ollama serve && ollama pull qwen3:0.6b # Ollama（Mac / 无 GPU）
# 跑全部场景并保存结果
python benchmark.py --scenario all --output results.json
# 只看 KV Cache 命中 / 未命中的 TTFT 对比
python benchmark.py --scenario kv-cache --backend ollama
# 批处理吞吐扫描
python benchmark.py --scenario batching --concurrency 1,2,4,8
"""
import argparse
import json
import logging
import statistics
import sys
import time
from concurrent.futures import ThreadPoolExecutor
from typing import Any, Dict, List, Optional
logging.basicConfig(
level=logging.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
logger = logging.getLogger("benchmark")
# 各后端的默认 OpenAI 兼容地址
BACKEND_DEFAULTS = {
"vllm": {"base_url": "http://localhost:8000/v1", "model": "Qwen3-0.6B"},
"ollama": {"base_url": "http://localhost:11434/v1", "model": "qwen3:0.6b"},
}
# 一段确定性的填充文本，用于把共享前缀撑长，让 KV Cache 的效果更明显
_FILLER_SENTENCE = (
"You are a meticulous assistant that follows the operating manual precisely. "
)
def build_padded_system_prompt(target_tokens: int) -> str:
"""构造一个约含 target_tokens 个 token 的系统提示词（用重复句子填充）。
这里用「4 字符 ≈ 1 token」的粗略估计来控制长度，只需保证前缀足够长、
可复现即可，不追求精确的 token 数。
"""
header = (
"# Operating Manual\n"
"You are a helpful local assistant deployed for the AI Agent book experiment.\n\n"
)
approx_chars = max(0, target_tokens * 4 - len(header))
repeats = approx_chars // len(_FILLER_SENTENCE) + 1
body = _FILLER_SENTENCE * repeats
return header + body
def make_client(base_url: str, api_key: str):
"""创建 OpenAI 兼容客户端。"""
try:
from openai import OpenAI
except ImportError:
logger.error("缺少依赖 openai，请先执行：pip install openai")
sys.exit(1)
return OpenAI(base_url=base_url, api_key=api_key)
def stream_once(
client,
model: str,
messages: List[Dict[str, str]],
max_tokens: int,
temperature: float,
) -> Dict[str, float]:
"""发起一次流式请求，返回 TTFT、总时长、输出 token 数与解码吞吐。
- ttft：从发起请求到收到第一个内容分片的时间（秒）
- total：整个响应的墙钟时间（秒）
- output_tokens：优先取服务端返回的 usage.completion_tokens，
否则用收到的内容分片数量作为近似
- decode_tps：解码阶段吞吐 = 输出 token / (总时长 - TTFT)
"""
start = time.perf_counter()
ttft: Optional[float] = None
chunk_count = 0
usage_tokens: Optional[int] = None
stream = client.chat.completions.create(
model=model,
messages=messages,
max_tokens=max_tokens,
temperature=temperature,
stream=True,
stream_options={"include_usage": True},
)
for chunk in stream:
# 最后一个分片可能只携带 usage 而没有 choices
if getattr(chunk, "usage", None) is not None:
try:
usage_tokens = chunk.usage.completion_tokens
except AttributeError:
pass
if not chunk.choices:
continue
delta = chunk.choices[0].delta
if getattr(delta, "content", None):
if ttft is None:
ttft = time.perf_counter() - start
chunk_count += 1
total = time.perf_counter() - start
if ttft is None:
ttft = total
output_tokens = usage_tokens if usage_tokens is not None else chunk_count
decode_time = max(total - ttft, 1e-6)
decode_tps = output_tokens / decode_time if output_tokens else 0.0
return {
"ttft": ttft,
"total": total,
"output_tokens": float(output_tokens),
"decode_tps": decode_tps,
}
# --------------------------------------------------------------------------- #
# 场景实现
# --------------------------------------------------------------------------- #
def scenario_throughput(client, model, args) -> Dict[str, Any]:
"""单流吞吐 + TTFT：连续发起若干次解码密集的请求并汇总统计。"""
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{
"role": "user",
"content": "Write a detailed explanation of how KV Cache works in transformer inference.",
},
]
runs = []
for i in range(args.repeats):
r = stream_once(client, model, messages, args.max_tokens, args.temperature)
runs.append(r)
logger.info(
"throughput 第 %d/%d 次: TTFT=%.3fs, 解码=%.1f tok/s, 输出=%d tok",
i + 1, args.repeats, r["ttft"], r["decode_tps"], int(r["output_tokens"]),
)
return {
"scenario": "throughput",
"repeats": args.repeats,
"ttft_mean_s": statistics.fmean(x["ttft"] for x in runs),
"decode_tps_mean": statistics.fmean(x["decode_tps"] for x in runs),
"output_tokens_mean": statistics.fmean(x["output_tokens"] for x in runs),
"runs": runs,
}
def scenario_kv_cache(client, model, args) -> Dict[str, Any]:
"""KV Cache 命中 vs 未命中的 TTFT 对比（实验 2-1 第 5 点）。
- 命中组：系统提示词逐字节不变，重复发送同一请求，服务端前缀缓存命中，
prefill 几乎可以跳过 → TTFT 明显更低。
- 未命中组：每次只在系统提示词「开头」插入一个不同的计数串，前缀被改写，
缓存全部失效，服务端必须重算整个前缀 → TTFT 明显更高。
两组的提示词长度基本一致，因此差异主要来自前缀缓存是否命中。
"""
base_prompt = build_padded_system_prompt(args.prefix_tokens)
user_msg = {"role": "user", "content": "In one short sentence, say hello."}
# 预热：先发一次把缓存写入（这一次一定是冷启动，不计入统计）
warm_msgs = [{"role": "system", "content": base_prompt}, user_msg]
stream_once(client, model, warm_msgs, args.max_tokens, args.temperature)
hit_ttfts, miss_ttfts = [], []
for i in range(args.repeats):
# 命中：完全相同的前缀
hit = stream_once(client, model, warm_msgs, args.max_tokens, args.temperature)
hit_ttfts.append(hit["ttft"])
# 未命中：在开头插入唯一前缀，使缓存失效
mutated = f"[req-{i}-{time.time_ns()}] " + base_prompt
miss_msgs = [{"role": "system", "content": mutated}, user_msg]
miss = stream_once(client, model, miss_msgs, args.max_tokens, args.temperature)
miss_ttfts.append(miss["ttft"])
logger.info(
"kv-cache 第 %d/%d 次: 命中 TTFT=%.3fs, 未命中 TTFT=%.3fs",
i + 1, args.repeats, hit["ttft"], miss["ttft"],
)
hit_mean = statistics.fmean(hit_ttfts)
miss_mean = statistics.fmean(miss_ttfts)
return {
"scenario": "kv-cache",
"prefix_tokens_approx": args.prefix_tokens,
"repeats": args.repeats,
"ttft_hit_mean_s": hit_mean,
"ttft_miss_mean_s": miss_mean,
"speedup": (miss_mean / hit_mean) if hit_mean > 0 else None,
"ttft_hit_s": hit_ttfts,
"ttft_miss_s": miss_ttfts,
}
def scenario_batching(client, model, args) -> Dict[str, Any]:
"""批处理：在不同并发度下并发发起请求，测量聚合吞吐。
连续批处理（continuous batching）是本地 serving 的核心优化：并发越高，
GPU 利用率越充分，系统聚合吞吐（所有请求合计 tok/s）通常显著上升，
但单个请求的延迟可能上升。此场景把这个权衡直接量化出来。
"""
messages = [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": "Explain what a large language model is."},
]
levels = args.concurrency
rows = []
for level in levels:
start = time.perf_counter()
with ThreadPoolExecutor(max_workers=level) as pool:
futures = [
pool.submit(
stream_once, client, model, messages, args.max_tokens, args.temperature
)
for _ in range(level)
]
results = [f.result() for f in futures]
wall = time.perf_counter() - start
total_tokens = sum(r["output_tokens"] for r in results)
agg_tps = total_tokens / wall if wall > 0 else 0.0
per_req_tps = agg_tps / level if level else 0.0
rows.append(
{
"concurrency": level,
"wall_s": wall,
"total_output_tokens": total_tokens,
"aggregate_tps": agg_tps,
"per_request_tps": per_req_tps,
"ttft_mean_s": statistics.fmean(r["ttft"] for r in results),
}
)
logger.info(
"batching 并发=%d: 聚合吞吐=%.1f tok/s, 单请求=%.1f tok/s, 墙钟=%.2fs",
level, agg_tps, per_req_tps, wall,
)
return {"scenario": "batching", "levels": rows}
# --------------------------------------------------------------------------- #
# 结果表格
# --------------------------------------------------------------------------- #
def print_report(results: List[Dict[str, Any]]) -> None:
print("\n" + "=" * 68)
print("本地 LLM 服务基准结果")
print("=" * 68)
for res in results:
s = res["scenario"]
if s == "throughput":
print("\n[throughput] 单流吞吐 / 首 token 延迟")
print(f" 次数 : {res['repeats']}")
print(f" 平均 TTFT : {res['ttft_mean_s']:.3f} s")
print(f" 平均解码吞吐 : {res['decode_tps_mean']:.1f} tok/s")
print(f" 平均输出长度 : {res['output_tokens_mean']:.0f} tok")
elif s == "kv-cache":
print("\n[kv-cache] 前缀缓存命中 vs 未命中（TTFT）")
print(f" 前缀长度(约) : {res['prefix_tokens_approx']} tok")
print(f" 命中平均 TTFT : {res['ttft_hit_mean_s']:.3f} s")
print(f" 未命中平均TTFT : {res['ttft_miss_mean_s']:.3f} s")
if res.get("speedup"):
print(f" 缓存加速比 : {res['speedup']:.2f}x")
elif s == "batching":
print("\n[batching] 并发度对聚合吞吐的影响")
print(f" {'并发':>4} | {'聚合tok/s':>10} | {'单请求tok/s':>12} | {'平均TTFT(s)':>11} | {'墙钟(s)':>8}")
print(f" {'-'*4}-+-{'-'*10}-+-{'-'*12}-+-{'-'*11}-+-{'-'*8}")
for row in res["levels"]:
print(
f" {row['concurrency']:>4} | {row['aggregate_tps']:>10.1f} | "
f"{row['per_request_tps']:>12.1f} | {row['ttft_mean_s']:>11.3f} | {row['wall_s']:>8.2f}"
)
print("\n" + "=" * 68)
def describe_dry_run(args) -> None:
"""离线打印将要执行的场景配置，不访问服务端。"""
print("=" * 68)
print("DRY RUN —— 仅打印计划，不访问服务端")
print("=" * 68)
print(f"后端 : {args.backend}")
print(f"base_url : {args.base_url}")
print(f"模型 : {args.model}")
print(f"重复次数 : {args.repeats}")
print(f"max_tokens : {args.max_tokens}")
print(f"temperature : {args.temperature}")
scenarios = ["throughput", "kv-cache", "batching"] if args.scenario == "all" else [args.scenario]
print(f"待运行场景 : {', '.join(scenarios)}")
if "kv-cache" in scenarios:
prompt = build_padded_system_prompt(args.prefix_tokens)
print(f" kv-cache : 填充前缀约 {args.prefix_tokens} tok（实际 {len(prompt)} 字符）")
if "batching" in scenarios:
print(f" batching : 并发扫描 {args.concurrency}")
print("=" * 68)
def parse_concurrency(value: str) -> List[int]:
try:
levels = [int(x) for x in value.split(",") if x.strip()]
except ValueError:
raise argparse.ArgumentTypeError("--concurrency 需为逗号分隔的正整数，例如 1,2,4,8")
if not levels or any(x <= 0 for x in levels):
raise argparse.ArgumentTypeError("--concurrency 中的并发度必须为正整数")
return levels
def build_parser() -> argparse.ArgumentParser:
parser = argparse.ArgumentParser(
description="本地 LLM 服务性能基准：吞吐 / 延迟 / KV Cache / 批处理（实验 2-1 配套）",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog=(
"场景说明:\n"
" throughput 单流解码吞吐(tok/s)与首 token 延迟(TTFT)\n"
" kv-cache 前缀缓存命中 vs 未命中的 TTFT 对比\n"
" batching 不同并发度下的聚合吞吐（批处理权衡）\n"
" all 依次运行以上全部场景\n"
),
)
parser.add_argument(
"--scenario",
choices=["throughput", "kv-cache", "batching", "all"],
default="all",
help="要运行的基准场景（默认: all）",
)
parser.add_argument(
"--backend",
choices=["vllm", "ollama"],
default="vllm",
help="服务端类型，用于推断默认地址与模型名（默认: vllm）",
)
parser.add_argument(
"--base-url",
type=str,
default=None,
help="OpenAI 兼容接口地址，覆盖后端默认值（如 http://localhost:8000/v1）",
)
parser.add_argument(
"--model",
type=str,
default=None,
help="模型名，覆盖后端默认值（vLLM 默认 Qwen3-0.6B，Ollama 默认 qwen3:0.6b）",
)
parser.add_argument(
"--api-key",
type=str,
default="EMPTY",
help="API Key，本地服务端一般无需真实值（默认: EMPTY）",
)
parser.add_argument(
"--repeats",
type=int,
default=5,
help="throughput / kv-cache 场景的重复次数（默认: 5）",
)
parser.add_argument(
"--max-tokens",
type=int,
default=256,
help="每次请求的最大生成 token 数（默认: 256）",
)
parser.add_argument(
"--temperature",
type=float,
default=0.7,
help="采样温度（默认: 0.7）",
)
parser.add_argument(
"--prefix-tokens",
type=int,
default=1024,
help="kv-cache 场景中共享前缀的近似 token 长度，越长缓存效果越明显（默认: 1024）",
)
parser.add_argument(
"--concurrency",
type=parse_concurrency,
default=[1, 2, 4, 8],
help="batching 场景的并发度列表，逗号分隔（默认: 1,2,4,8）",
)
parser.add_argument(
"--output",
type=str,
default=None,
help="将结果以 JSON 写入指定文件",
)
parser.add_argument(
"--dry-run",
action="store_true",
help="离线打印计划而不访问服务端，用于验证配置",
)
return parser
def main() -> int:
parser = build_parser()
args = parser.parse_args()
# 用后端默认值补全 base_url / model
defaults = BACKEND_DEFAULTS[args.backend]
if args.base_url is None:
args.base_url = defaults["base_url"]
if args.model is None:
args.model = defaults["model"]
print("=" * 68)
print("🚀 本地 LLM 服务性能基准（实验 2-1）")
print("=" * 68)
if args.dry_run:
describe_dry_run(args)
return 0
client = make_client(args.base_url, args.api_key)
logger.info("连接服务端: %s（模型: %s）", args.base_url, args.model)
scenarios = (
["throughput", "kv-cache", "batching"]
if args.scenario == "all"
else [args.scenario]
)
dispatch = {
"throughput": scenario_throughput,
"kv-cache": scenario_kv_cache,
"batching": scenario_batching,
}
results: List[Dict[str, Any]] = []
try:
for name in scenarios:
logger.info("开始场景: %s", name)
results.append(dispatch[name](client, args.model, args))
except Exception as e: # noqa: BLE001
logger.error("基准执行失败: %s", e)
logger.info(
"请确认服务端已启动：vLLM 用 `python server.py`，"
"Ollama 用 `ollama serve` 并已 `ollama pull %s`",
args.model,
)
return 1
print_report(results)
if args.output:
payload = {
"backend": args.backend,
"base_url": args.base_url,
"model": args.model,
"results": results,
}
with open(args.output, "w", encoding="utf-8") as f:
json.dump(payload, f, ensure_ascii=False, indent=2)
logger.info("结果已写入: %s", args.output)
return 0
if __name__ == "__main__":
sys.exit(main())
check_compatibility.py¶
#!/usr/bin/env python3
"""
Check system compatibility for running vLLM tool calling demo
"""
import sys
import platform
import subprocess
import shutil
def check_system():
"""Check system compatibility"""
print("="*60)
print("🔍 System Compatibility Check")
print("="*60)
# Get system info
system = platform.system()
machine = platform.machine()
python_version = sys.version_info
print(f"\n📊 System Information:")
print(f" OS: {system} ({platform.platform()})")
print(f" Architecture: {machine}")
print(f" Python: {python_version.major}.{python_version.minor}.{python_version.micro}")
# Check for CUDA
cuda_available = False
gpu_info = None
print(f"\n🎮 GPU Check:")
if system == "Darwin": # macOS
print(" ❌ macOS detected - No CUDA support available")
print(" ℹ️ Macs use Metal (Apple Silicon) or AMD/Intel GPUs")
return False, "darwin"
# Check for NVIDIA GPU
if shutil.which("nvidia-smi"):
try:
result = subprocess.run(
["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"],
capture_output=True,
text=True
)
if result.returncode == 0:
gpu_info = result.stdout.strip()
print(f" ✅ NVIDIA GPU found: {gpu_info}")
cuda_available = True
else:
print(" ⚠️ nvidia-smi found but couldn't query GPU")
except Exception as e:
print(f" ⚠️ Error checking GPU: {e}")
else:
print(" ❌ No NVIDIA GPU detected (nvidia-smi not found)")
# Check PyTorch CUDA
print(f"\n🔥 PyTorch CUDA Check:")
try:
import torch
if torch.cuda.is_available():
print(f" ✅ PyTorch CUDA is available")
print(f" CUDA version: {torch.version.cuda}")
print(f" Number of GPUs: {torch.cuda.device_count()}")
if torch.cuda.device_count() > 0:
print(f" GPU 0: {torch.cuda.get_device_name(0)}")
else:
print(" ❌ PyTorch CUDA is not available")
cuda_available = False
except ImportError:
print(" ⚠️ PyTorch not installed")
return cuda_available, system.lower()
def provide_recommendations(cuda_available, system):
"""Provide recommendations based on system"""
print("\n" + "="*60)
print("💡 Recommendations")
print("="*60)
if cuda_available:
print("\n✅ Your system supports vLLM!")
print("\nNext steps:")
print("1. Install requirements: pip install -r requirements.txt")
print("2. Run the main script: python main.py")
print("3. The script will automatically use vLLM")
elif system == "darwin" or system.lower() == "darwin": # macOS
print("\n🍎 You're on macOS - will use Ollama")
print("\n📋 Setup steps:\n")
print("1️⃣ Install Ollama:")
print(" brew install ollama")
print(" ollama serve # Run in separate terminal\n")
print("2️⃣ Install a model with tool support:")
print(" ollama pull qwen3:0.6b # Default model for this project\n")
print("3️⃣ Run the main script:")
print(" python main.py")
print(" # Will automatically use Ollama")
elif system.lower() == "windows": # Windows
if not cuda_available:
print("\n🪟 You're on Windows without CUDA - will use Ollama")
print("\n📋 Setup steps:\n")
print("1️⃣ Install Ollama:")
print(" Download from: https://ollama.com/download/windows")
print(" Run OllamaSetup.exe\n")
print("2️⃣ Install a model:")
print(" ollama pull qwen3:0.6b # Default model for this project\n")
print("3️⃣ Run the main script:")
print(" python main.py")
print(" # Will automatically use Ollama")
else: # Linux without CUDA
print("\n🐧 You're on Linux without CUDA - will use Ollama")
print("\n📋 Setup steps:\n")
print("1️⃣ Install Ollama:")
print(" curl -fsSL https://ollama.com/install.sh | sh")
print(" systemctl start ollama # Or: ollama serve\n")
print("2️⃣ Install a model:")
print(" ollama pull qwen3:0.6b # Default model for this project\n")
print("3️⃣ Run the main script:")
print(" python main.py")
print(" # Will automatically use Ollama")
def main():
"""Main compatibility check"""
cuda_available, system = check_system()
provide_recommendations(cuda_available, system)
print("\n" + "="*60)
print("For more details, see README.md")
print("="*60)
if __name__ == "__main__":
main()
config.py¶
"""
Configuration for vLLM Tool Calling Demo
"""
import os
from pathlib import Path
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
# Model Configuration
MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3-0.6B") # Can use ModelScope path or HuggingFace
MODEL_PATH = os.getenv("MODEL_PATH", None) # Optional: local model path
VLLM_PORT = int(os.getenv("VLLM_PORT", 8000))
VLLM_HOST = os.getenv("VLLM_HOST", "localhost")
# vLLM Server Configuration
VLLM_SERVER_CONFIG = {
"model": MODEL_NAME,
"port": VLLM_PORT,
"host": VLLM_HOST,
"enable_auto_tool_choice": True,
"tool_call_parser": "hermes",
"chat_template": "tool_use", # Use tool-specific chat template
"max_model_len": 8192,
"gpu_memory_utilization": 0.9,
"dtype": "auto",
"enforce_eager": False, # Set to True if you encounter issues
}
# OpenAI Client Configuration (for connecting to vLLM)
OPENAI_API_BASE = f"http://{VLLM_HOST}:{VLLM_PORT}/v1"
OPENAI_API_KEY = "EMPTY" # vLLM doesn't require a real key
# Tool Configuration
ENABLE_WEATHER_TOOL = True
ENABLE_CALCULATOR_TOOL = True
ENABLE_SEARCH_TOOL = True
# Logging
LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
LOG_FILE = Path("logs") / "vllm_tool_demo.log"
demo_streaming.py¶
#!/usr/bin/env python3
"""
Simple demo showing how to use streaming with the chat template agents
"""
import sys
import time
def print_with_typing_effect(text, delay=0.03):
"""Print text with a typing effect"""
for char in text:
print(char, end="", flush=True)
time.sleep(delay)
print()
def demo_vllm_streaming():
"""Demo streaming with vLLM backend"""
from agent import VLLMToolAgent
from config import OPENAI_API_BASE, OPENAI_API_KEY
print("="*60)
print("🚀 vLLM Streaming Demo")
print("="*60)
try:
agent = VLLMToolAgent(
api_base=OPENAI_API_BASE,
api_key=OPENAI_API_KEY
)
query = "What's the weather in New York and calculate 32°F in Celsius?"
print(f"\n📝 Query: {query}\n")
print("Streaming response:\n")
print("-"*40)
for chunk in agent.chat_stream(query):
chunk_type = chunk.get("type")
content = chunk.get("content", "")
if chunk_type == "thinking":
print(f"\n💭 [Thinking]: \033[90m{content}\033[0m")
elif chunk_type == "tool_call":
print(f"\n🔧 [Tool Call]: {content['name']}({content['arguments']})")
elif chunk_type == "tool_result":
print(f" ✓ Result: {content}")
elif chunk_type == "content":
# Stream content character by character
print(content, end="", flush=True)
print("\n" + "-"*40)
except Exception as e:
print(f"Error: {e}")
print("Make sure vLLM server is running!")
def demo_ollama_streaming():
"""Demo streaming with Ollama backend"""
from ollama_native import OllamaNativeAgent
import ollama
print("="*60)
print("🦙 Ollama Streaming Demo")
print("="*60)
try:
# Check available models
client = ollama.Client()
models = [m['name'] for m in client.list()['models']]
# Use qwen3:0.6b as the default model
model = "qwen3:0.6b"
if model not in models:
print(f"⚠️ Recommended model {model} not found")
print("Install with: ollama pull qwen3:0.6b")
if models:
model = models[0]
print(f"Using fallback model: {model}")
else:
print("❌ No Ollama models found. Install with: ollama pull qwen3:0.6b")
return
print(f"Using model: {model}")
agent = OllamaNativeAgent(model=model)
query = "What's 15 * 23? Also get the current time in Tokyo."
print(f"\n📝 Query: {query}\n")
print("Streaming response:\n")
print("-"*40)
for chunk in agent.chat_stream(query):
chunk_type = chunk.get("type")
content = chunk.get("content", "")
if chunk_type == "thinking":
print(f"\n💭 [Thinking]: \033[90m{content}\033[0m")
elif chunk_type == "tool_call":
print(f"\n🔧 [Tool Call]: {content['name']}({content['arguments']})")
elif chunk_type == "tool_result":
print(f" ✓ Result: {content}")
elif chunk_type == "content":
# Stream content
print(content, end="", flush=True)
print("\n" + "-"*40)
except Exception as e:
print(f"Error: {e}")
print("Make sure Ollama is running: ollama serve")
def demo_unified_streaming():
"""Demo with unified ToolCallingAgent that auto-selects backend"""
from main import ToolCallingAgent
print("="*60)
print("🎯 Unified Streaming Demo (Auto-detect Backend)")
print("="*60)
# Initialize agent (auto-detects best backend)
print("\n⚙️ Initializing agent...")
agent = ToolCallingAgent()
print(f"✅ Using {agent.backend_type} backend\n")
# Example queries
queries = [
"Calculate the compound interest on $1000 at 5% for 3 years",
"What's the weather in London and what time is it there?",
"Convert 50 EUR to USD and JPY"
]
for i, query in enumerate(queries, 1):
print(f"\n{'='*60}")
print(f"Query {i}: {query}")
print("-"*60)
# Track what sections we've shown
sections_shown = set()
last_chunk_type = None
for chunk in agent.chat(query, stream=True):
chunk_type = chunk.get("type")
content = chunk.get("content", "")
if chunk_type == "thinking":
if "thinking" not in sections_shown:
print("\n💭 Thinking: ", end="", flush=True)
sections_shown.add("thinking")
# Stream thinking character by character in gray
print(f"\033[90m{content}\033[0m", end="", flush=True)
elif chunk_type == "tool_call":
if "tools" not in sections_shown:
print("\n🔧 Tool Calls:")
sections_shown.add("tools")
print(f" → {content['name']}: {content['arguments']}")
# Remove response section so it shows again after tools
sections_shown.discard("response")
elif chunk_type == "tool_result":
result_str = str(content)
print(f" ✓ {result_str}")
# Remove response section so it shows again after tools
sections_shown.discard("response")
elif chunk_type == "content":
if "response" not in sections_shown:
if last_chunk_type in ["tool_result", "tool_call"]:
print("\n📝 Response (after tools):")
else:
print("\n📝 Response:")
sections_shown.add("response")
print(" ", end="")
print(content, end="", flush=True)
last_chunk_type = chunk_type
print() # New line after response
# Reset for next query
agent.reset_conversation()
print("\n" + "="*60)
print("✅ Demo completed!")
print("="*60)
def main():
"""Main demo function"""
import argparse
parser = argparse.ArgumentParser(description="Streaming Demo for Chat Template Agents")
parser.add_argument(
"--backend",
choices=["vllm", "ollama", "auto"],
default="auto",
help="Backend to use for demo"
)
args = parser.parse_args()
if args.backend == "vllm":
demo_vllm_streaming()
elif args.backend == "ollama":
demo_ollama_streaming()
else:
demo_unified_streaming()
if __name__ == "__main__":
main()
main.py¶
#!/usr/bin/env python3
"""
Main Entry Point for Tool Calling Demo
Automatically selects the best backend based on your platform:
- Linux/Windows with NVIDIA GPU: Uses vLLM
- Mac/Windows/Linux without GPU: Uses Ollama
"""
import os
import sys
import platform
import logging
from typing import Optional, Dict, Any, List
import json
logging.basicConfig(
level=logging.INFO,
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)
class ToolCallingAgent:
"""
Universal tool calling agent that works on all platforms
Automatically selects vLLM (if GPU available) or Ollama
"""
def __init__(self, backend: Optional[str] = None):
"""
Initialize with automatic backend detection
Args:
backend: Force a specific backend ('vllm', 'ollama', or None for auto)
"""
self.agent = None
self.backend_type = backend or self._detect_best_backend()
logger.info(f"Initializing on {platform.system()} with {self.backend_type}")
self._initialize_backend()
def _detect_best_backend(self) -> str:
"""Detect the best backend for current platform"""
system = platform.system()
# Check for CUDA support (Linux/Windows with NVIDIA GPU)
if system in ["Linux", "Windows"]:
try:
import torch
if torch.cuda.is_available():
logger.info("CUDA detected - will use vLLM")
return "vllm"
except ImportError:
pass
# Default to Ollama for Mac or systems without CUDA
logger.info(f"Using Ollama on {system}")
return "ollama"
def _initialize_backend(self):
"""Initialize the selected backend"""
if self.backend_type == "vllm":
self._init_vllm()
else:
self._init_ollama()
def _init_vllm(self):
"""Initialize vLLM backend"""
try:
# Check if vLLM server is running
import requests
from config import VLLM_HOST, VLLM_PORT
server_url = f"http://{VLLM_HOST}:{VLLM_PORT}/health"
try:
response = requests.get(server_url, timeout=1)
if response.status_code != 200:
raise ConnectionError("vLLM server not responding")
except:
# Try to start the server
logger.info("Starting vLLM server...")
from server import VLLMServer
server = VLLMServer()
server.start(wait_for_ready=True)
# Initialize vLLM agent
from agent import VLLMToolAgent
from config import OPENAI_API_BASE, OPENAI_API_KEY
self.agent = VLLMToolAgent(
api_base=OPENAI_API_BASE,
api_key=OPENAI_API_KEY
)
logger.info("✅ vLLM agent initialized")
except Exception as e:
logger.warning(f"Failed to initialize vLLM: {e}")
logger.info("Falling back to Ollama")
self.backend_type = "ollama"
self._init_ollama()
def _init_ollama(self):
"""Initialize Ollama backend"""
try:
import ollama
from ollama_native import OllamaNativeAgent
# Check if Ollama is running
client = ollama.Client()
try:
models_response = client.list()
available_models = []
if hasattr(models_response, 'models'):
available_models = [m.model for m in models_response.models]
if not available_models:
logger.error("No Ollama models installed")
logger.info("Install a model with: ollama pull qwen3:0.6b")
sys.exit(1)
# Use qwen3:0.6b as the default model
model = "qwen3:0.6b"
# Check if qwen3:0.6b is available
if model not in available_models:
logger.warning(f"Recommended model {model} not found in available models")
logger.info("Install with: ollama pull qwen3:0.6b")
# Fall back to first available model if qwen3:0.6b is not installed
model = available_models[0]
logger.info(f"Using fallback model: {model}")
logger.info(f"Using Ollama model: {model}")
self.agent = OllamaNativeAgent(model=model)
except Exception as e:
logger.error(f"Ollama is not running: {e}")
logger.info("\nPlease start Ollama:")
system = platform.system()
if system == "Darwin": # Mac
logger.info(" brew services start ollama")
logger.info(" or: ollama serve")
elif system == "Windows":
logger.info(" Start Ollama from the system tray")
logger.info(" or run: ollama serve")
else: # Linux
logger.info(" systemctl start ollama")
logger.info(" or: ollama serve")
sys.exit(1)
except ImportError:
logger.error("Ollama not installed")
logger.info("Install with: pip install ollama")
sys.exit(1)
def chat(self, message: str, use_tools: bool = True, stream: bool = False, **kwargs) -> str:
"""
Send a message to the agent
Args:
message: User message
use_tools: Whether to enable tool calling
stream: Whether to stream the response
**kwargs: Additional backend-specific parameters
Returns:
Agent response (or generator if streaming)
"""
if not self.agent:
raise RuntimeError("Agent not initialized")
return self.agent.chat(message, use_tools=use_tools, stream=stream, **kwargs)
def reset_conversation(self):
"""Reset conversation history"""
if hasattr(self.agent, 'reset_conversation'):
self.agent.reset_conversation()
def get_sample_tasks() -> List[Dict[str, str]]:
"""Get sample tasks for demonstration"""
return [
{
"name": "🕐 Current Time Check",
"description": "Get the current time in a specific city",
"task": "What is the current time in Vancouver?"
},
{
"name": "☀️ Simple Weather Check",
"description": "Get current weather for a single city",
"task": "What's the weather like in Vancouver right now?"
},
{
"name": "☀️ Time and Weather Check",
"description": "Get current time and weather for a single city",
"task": "What's the current time and weather like in Vancouver right now?"
},
{
"name": "💵 Compound Interest Calculation",
"description": "Calculate compound interest using code interpreter",
"task": "Calculate the compound interest on $5,000 invested at 6% annual interest rate for 30 years, compounded monthly."
},
{
"name": "🌡️ Multi-City Weather Analysis",
"description": "Compare weather across multiple cities using real-time data",
"task": """Get the current weather for Tokyo, New York, London, Sydney, and Dubai.
Then:
1. Which city has the highest temperature?
2. Which city has the lowest humidity?
3. Convert all temperatures to Fahrenheit for comparison
4. Calculate the average temperature across all cities"""
},
{
"name": "💰 Complex Financial Analysis",
"description": "Multi-step financial calculation with currency conversion",
"task": """A company has the following quarterly revenues:
- Q1: $2,500,000 USD
- Q2: €2,100,000 EUR
- Q3: £1,800,000 GBP
- Q4: ¥380,000,000 JPY
Please:
1. Convert all revenues to USD
2. Calculate the total annual revenue in USD
3. Determine the average quarterly revenue
4. Find which quarter had the highest revenue
5. If the company has a 20% profit margin, calculate the annual profit in USD"""
},
{
"name": "⏰ Global Time Zone Coordination",
"description": "Coordinate meeting times across time zones",
"task": """We need to schedule a global meeting with offices in:
- San Francisco (PST)
- New York (EST)
- London (GMT/BST)
- Tokyo (JST)
- Sydney (AEST)
If the meeting is at 2 PM London time:
1. What time would it be in each city?
2. Is this during normal business hours (9 AM - 5 PM) for each location?
3. Suggest a better time that works for most offices"""
},
]
def run_single_task(agent: ToolCallingAgent, task: str, stream: bool = True):
"""Run a single task with optional streaming"""
print("\n" + "="*60)
print("TASK EXECUTION")
print("="*60)
print(f"\n📋 Task: {task}")
print("-"*60)
try:
if stream:
print("\n⏳ Processing (streaming)...\n")
response_chunks = []
thinking_shown = False
tools_shown = False
response_started = False
last_chunk_type = None
for chunk in agent.chat(task, stream=True):
chunk_type = chunk.get("type")
content = chunk.get("content", "")
if chunk_type == "thinking":
if not thinking_shown:
print("🧠 Thinking: ", end="", flush=True)
thinking_shown = True
# Stream thinking character by character in gray
print(f"\033[90m{content}\033[0m", end="", flush=True)
elif chunk_type == "tool_call":
if not tools_shown:
print("\n\n🔧 Tool Calls:")
tools_shown = True
# Display tool call info
tool_info = content
print(f" → {tool_info.get('name', 'unknown')}: {tool_info.get('arguments', {})}")
# Reset response_started flag after tool calls
response_started = False
elif chunk_type == "tool_result":
# Display tool result
result_str = str(content)
print(f" ✓ {result_str}")
# Reset response_started flag after tool results
response_started = False
elif chunk_type == "content":
if not response_started:
# Check if this is content after tool execution
if last_chunk_type in ["tool_result", "tool_call"]:
print("\n🤖 Assistant: ", end="", flush=True)
elif thinking_shown or tools_shown:
print("\n\n🤖 Assistant: ", end="", flush=True)
else:
print("🤖 Assistant: ", end="", flush=True)
response_started = True
# Stream the actual response content
print(content, end="", flush=True)
response_chunks.append(content)
elif chunk_type == "error":
print(f"\n❌ Error: {content}")
last_chunk_type = chunk_type
print("\n" + "-"*40)
else:
print("\n⏳ Processing...")
response = agent.chat(task, stream=False)
print("\n✅ Response:")
print("-"*40)
print(response)
print("-"*40)
except Exception as e:
print(f"\n❌ Error: {e}")
logger.exception("Task execution failed")
def interactive_mode(agent: ToolCallingAgent, stream: bool = True):
"""Run interactive chat mode with optional streaming"""
print("\n" + "="*60)
print("💬 INTERACTIVE MODE" + (" (STREAMING)" if stream else ""))
print("="*60)
print("\nYou can now chat with the AI agent. It has access to various tools:")
# Show available tools
from tools import ToolRegistry
registry = ToolRegistry()
tools = registry.get_tool_schemas()
print("\n📦 Available Tools:")
for i, tool in enumerate(tools, 1):
func = tool["function"]
print(f" {i}. {func['name']}: {func['description']}")
print("\n💡 Commands:")
print(" /reset - Reset conversation")
print(" /tools - Show available tools")
print(" /samples - Show sample tasks")
print(" /sample <n> - Run sample task number n")
print(" /stream - Toggle streaming mode")
print(" /help - Show this help")
print(" /exit - Exit the program")
print("-"*60)
streaming_enabled = stream
while True:
try:
user_input = input("\n👤 You: ").strip()
if not user_input:
continue
# Handle commands
if user_input.lower() == "/exit" or user_input.lower() == "quit":
print("👋 Goodbye!")
break
elif user_input.lower() == "/reset":
agent.reset_conversation()
print("✅ Conversation reset")
continue
elif user_input.lower() == "/tools":
print("\n📦 Available Tools:")
for i, tool in enumerate(tools, 1):
func = tool["function"]
print(f" {i}. {func['name']}: {func['description']}")
continue
elif user_input.lower() == "/samples":
print("\n📋 Sample Tasks:")
sample_tasks = get_sample_tasks()
for i, sample in enumerate(sample_tasks, 1):
print(f" {i}. {sample['name']}")
# Show first 100 chars of task for readability
task_preview = sample['task'].replace('\n', ' ')[:100]
if len(sample['task']) > 100:
task_preview += "..."
print(f" {task_preview}")
print("\n💡 Tip: Use /sample <n> to run a specific sample (e.g., /sample 1)")
continue
elif user_input.lower().startswith("/sample "):
# Extract the sample number
try:
sample_num = int(user_input.split()[1])
sample_tasks = get_sample_tasks()
if 1 <= sample_num <= len(sample_tasks):
selected_sample = sample_tasks[sample_num - 1]
print(f"\n🎯 Running Sample: {selected_sample['name']}")
print("-"*60)
print(f"Task: {selected_sample['task']}")
print("-"*60)
# Process the sample task as regular input
user_input = selected_sample['task']
# Don't continue - let it fall through to normal processing
else:
print(f"❌ Invalid sample number. Please choose between 1 and {len(sample_tasks)}")
print("Use /samples to see available samples")
continue
except (ValueError, IndexError):
print("❌ Invalid format. Use: /sample <number> (e.g., /sample 1)")
continue
elif user_input.lower() == "/help":
print("\n💡 Commands:")
print(" /reset - Reset conversation")
print(" /tools - Show available tools")
print(" /samples - Show sample tasks")
print(" /sample <n> - Run sample task number n")
print(" /stream - Toggle streaming mode")
print(" /help - Show this help")
print(" /exit - Exit the program")
continue
elif user_input.lower() == "/stream":
streaming_enabled = not streaming_enabled
print(f"✅ Streaming {'enabled' if streaming_enabled else 'disabled'}")
continue
# Process user input
if streaming_enabled:
print("\n⏳ Processing (streaming)...\n")
response_chunks = []
thinking_shown = False
tools_shown = False
response_started = False
last_chunk_type = None
for chunk in agent.chat(user_input, stream=True):
chunk_type = chunk.get("type")
content = chunk.get("content", "")
if chunk_type == "thinking":
if not thinking_shown:
print("🧠 Thinking: ", end="", flush=True)
thinking_shown = True
# Stream thinking character by character in gray
print(f"\033[90m{content}\033[0m", end="", flush=True)
elif chunk_type == "tool_call":
if not tools_shown:
print("\n🔧 Tool Calls:")
tools_shown = True
tool_info = content
print(f" → {tool_info.get('name', 'unknown')}: {tool_info.get('arguments', {})}")
# Reset response_started flag after tool calls so the next content gets a label
response_started = False
elif chunk_type == "tool_result":
result_str = str(content)
print(f" ✓ {result_str}")
# Reset response_started flag after tool results
response_started = False
elif chunk_type == "content":
# If we're starting a new content section after tool results
if not response_started:
if last_chunk_type in ["tool_result", "tool_call"]:
# This is a response after tool execution
print("\n🤖 Assistant: ", end="", flush=True)
elif thinking_shown or tools_shown:
print("\n🤖 Assistant: ", end="", flush=True)
else:
print("🤖 Assistant: ", end="", flush=True)
response_started = True
print(content, end="", flush=True)
response_chunks.append(content)
elif chunk_type == "error":
print(f"\n❌ Error: {content}")
last_chunk_type = chunk_type
print() # New line after streaming
else:
print("\n⏳ Processing...")
response = agent.chat(user_input, stream=False)
print(f"🤖 Assistant: {response}")
except KeyboardInterrupt:
print("\n\n👋 Goodbye!")
break
except Exception as e:
print(f"❌ Error: {e}")
logger.exception("Error in interactive mode")
def main():
"""Main function"""
import argparse
parser = argparse.ArgumentParser(
description="Universal Tool Calling Agent - Works on all platforms"
)
parser.add_argument(
"--mode",
choices=["single", "interactive"],
default="interactive",
help="Execution mode (default: interactive)"
)
parser.add_argument(
"--task",
type=str,
help="Task to execute (for single mode)"
)
parser.add_argument(
"--backend",
choices=["vllm", "ollama", "auto"],
default="auto",
help="Backend to use (default: auto-detect)"
)
parser.add_argument(
"--info",
action="store_true",
help="Show system information and exit"
)
parser.add_argument(
"--stream",
action="store_true",
default=True,
help="Enable streaming mode (default: True)"
)
parser.add_argument(
"--no-stream",
action="store_true",
help="Disable streaming mode"
)
args = parser.parse_args()
# Header
print("="*60)
print("🚀 Universal Tool Calling Agent")
print("="*60)
# Show system info if requested
if args.info:
print("\n📊 System Information:")
print(f" Platform: {platform.system()} {platform.release()}")
print(f" Architecture: {platform.machine()}")
print(f" Python: {sys.version.split()[0]}")
# Check CUDA
try:
import torch
cuda_available = torch.cuda.is_available()
if cuda_available:
print(f" CUDA: ✅ Available (GPU: {torch.cuda.get_device_name(0)})")
else:
print(" CUDA: ❌ Not available")
except ImportError:
print(" CUDA: ❌ PyTorch not installed")
# Check Ollama
try:
import ollama
print(" Ollama: ✅ Package installed")
except ImportError:
print(" Ollama: ❌ Package not installed")
return 0
# Initialize agent
print("\n⚙️ Initializing agent...")
backend = None if args.backend == "auto" else args.backend
try:
agent = ToolCallingAgent(backend=backend)
except SystemExit:
return 1
except Exception as e:
print(f"❌ Failed to initialize: {e}")
return 1
print(f"✅ Agent ready! Using {agent.backend_type} backend")
# Execute based on mode
if args.mode == "single":
if not args.task:
# Show sample tasks for selection
print("\n" + "="*60)
print("SINGLE TASK MODE - No task provided")
print("="*60)
sample_tasks = get_sample_tasks()
print("\n📋 Available sample tasks:")
for i, sample in enumerate(sample_tasks, 1):
print(f"\n{i}. {sample['name']}")
print(f" {sample['description']}")
print("\n" + "="*60)
try:
choice = input(f"\nSelect a task number (1-{len(sample_tasks)}) or 'q' to quit: ").strip()
if choice.lower() == 'q':
return 0
task_num = int(choice)
if 1 <= task_num <= len(sample_tasks):
selected_task = sample_tasks[task_num - 1]
print(f"\n✅ Selected: {selected_task['name']}")
print("\nTask details:")
print("-"*40)
print(selected_task['task'])
print("-"*40)
confirm = input("\nRun this task? (y/n): ").strip().lower()
if confirm == 'y':
stream_enabled = not args.no_stream if hasattr(args, 'no_stream') else True
run_single_task(agent, selected_task['task'], stream=stream_enabled)
else:
print("Task cancelled.")
else:
print(f"Invalid selection. Please choose 1-{len(sample_tasks)}")
return 1
except (ValueError, KeyboardInterrupt):
print("\nExiting...")
return 0
else:
stream_enabled = not args.no_stream if hasattr(args, 'no_stream') else True
run_single_task(agent, args.task, stream=stream_enabled)
else: # interactive mode
stream_enabled = not args.no_stream if hasattr(args, 'no_stream') else True
interactive_mode(agent, stream=stream_enabled)
return 0
if __name__ == "__main__":
sys.exit(main())
ollama_native.py¶
"""
Ollama Native Tool Calling Implementation
Uses Ollama's standard tool calling API (requires compatible models)
"""
import json
import logging
from concurrent.futures import ThreadPoolExecutor
from typing import List, Dict, Any, Optional
import ollama
from tools import ToolRegistry
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class OllamaNativeAgent:
"""Agent using Ollama's native tool calling support"""
def __init__(self, model: str = "qwen3:0.6b"):
"""
Initialize with a model that supports tool calling
"""
self.model = model
self.client = ollama.Client()
self.tool_registry = ToolRegistry()
self.conversation_history = []
# Check if Ollama is running
try:
self.client.list()
logger.info(f"✅ Connected to Ollama with model: {model}")
except Exception as e:
logger.error(f"❌ Failed to connect to Ollama: {e}")
logger.info("Please start Ollama with: ollama serve")
def _convert_tools_to_ollama_format(self) -> List[Dict]:
"""Convert tool registry to Ollama's expected format"""
tools = []
for tool_def in self.tool_registry.get_tool_schemas():
# Ollama expects the same format as OpenAI
tools.append(tool_def)
return tools
def _execute_tool_calls(self, tool_calls: List[Dict[str, Any]]) -> List[str]:
"""
Execute tool calls and return results in order.
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).
"""
def run_one(tool_call: Dict[str, Any]) -> str:
function = tool_call.get('function', {})
tool_name = function.get('name')
tool_args = function.get('arguments')
# Parse arguments if they're a string
if isinstance(tool_args, str):
try:
tool_args = json.loads(tool_args)
except json.JSONDecodeError:
logger.error(f"Failed to parse tool arguments: {tool_args}")
tool_args = {}
# Execute the tool
logger.info(f"Executing tool: {tool_name} with args: {tool_args}")
return self.tool_registry.execute_tool(tool_name, tool_args)
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 chat(self, message: str, use_tools: bool = True,
temperature: float = 0.3, stream: bool = False) -> str:
"""
Send a message using Ollama's native tool calling
Args:
message: User message
use_tools: Whether to enable tool calling
temperature: Sampling temperature
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)
# Original non-streaming implementation continues below...
# Add user message to history
self.conversation_history.append({
"role": "user",
"content": message
})
# Prepare tools if enabled
tools = self._convert_tools_to_ollama_format() if use_tools else None
try:
# Call Ollama with tools
response = self.client.chat(
model=self.model,
messages=self.conversation_history,
tools=tools,
options={"temperature": temperature}
)
# Check if model made tool calls
message_content = response.get('message', {})
# Handle tool calls if present
if 'tool_calls' in message_content:
tool_calls = message_content['tool_calls']
logger.info(f"Model requested {len(tool_calls)} tool call(s)")
# Add assistant's message with tool calls to history
self.conversation_history.append({
"role": "assistant",
"content": message_content.get('content', ''),
"tool_calls": tool_calls
})
# Execute the tool calls (independent calls run in parallel)
results = self._execute_tool_calls(tool_calls)
# Add tool results to conversation
for result in results:
self.conversation_history.append({
"role": "tool",
"content": result
})
# Get final response with tool results (still include tools!)
final_response = self.client.chat(
model=self.model,
messages=self.conversation_history,
tools=tools, # IMPORTANT: Keep tools available
options={"temperature": temperature}
)
final_content = final_response.get('message', {}).get('content', '')
# Clean response (remove <think> tags if present)
import re
final_content = re.sub(r'<think>.*?</think>', '', final_content, flags=re.DOTALL).strip()
# Add final response to history
self.conversation_history.append({
"role": "assistant",
"content": final_content
})
return final_content
else:
# No tool calls, just return the response
content = message_content.get('content', '')
self.conversation_history.append({
"role": "assistant",
"content": content
})
return content
except Exception as e:
logger.error(f"Error in chat: {e}")
return f"Error: {e}"
def chat_stream(self, message: str, use_tools: bool = True,
temperature: float = 0.3):
"""
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 if enabled
tools = self._convert_tools_to_ollama_format() 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
try:
# Get response from model
stream_response = self.client.chat(
model=self.model,
messages=self.conversation_history,
tools=tools,
options={"temperature": temperature},
stream=True
)
collected_content = []
tool_calls_detected = False
pending_tool_calls = []
thinking_buffer = ""
in_thinking = False
# Process the stream
for chunk in stream_response:
# Extract message content from chunk
message_chunk = chunk.get('message', {})
content_chunk = message_chunk.get('content', '')
if content_chunk:
collected_content.append(content_chunk)
# Handle thinking content
if '<think>' in content_chunk:
in_thinking = True
thinking_buffer = content_chunk
# Extract any content before <think>
import re
before_think = content_chunk.split('<think>')[0]
if before_think:
yield {"type": "content", "content": before_think}
# Extract thinking content from this chunk
if '</think>' in content_chunk:
# Complete thinking in one chunk
thinking_match = re.search(r'<think>(.*?)</think>', content_chunk, re.DOTALL)
if thinking_match:
thinking_content = thinking_match.group(1).strip()
# Stream thinking content character by character
for char in thinking_content:
yield {"type": "thinking", "content": char}
# Check for content after </think>
after_think = content_chunk.split('</think>')[-1]
if after_think:
yield {"type": "content", "content": after_think}
in_thinking = False
thinking_buffer = ""
else:
# Partial thinking, extract what we have so far
partial_thinking = content_chunk.split('<think>')[-1]
for char in partial_thinking:
yield {"type": "thinking", "content": char}
elif in_thinking:
thinking_buffer += content_chunk
if '</think>' in content_chunk:
# End of thinking
before_end = content_chunk.split('</think>')[0]
for char in before_end:
yield {"type": "thinking", "content": char}
# Check for content after </think>
after_think = content_chunk.split('</think>')[-1]
if after_think:
yield {"type": "content", "content": after_think}
in_thinking = False
thinking_buffer = ""
else:
# Continue streaming thinking
for char in content_chunk:
yield {"type": "thinking", "content": char}
else:
# Regular content - yield as-is
yield {"type": "content", "content": content_chunk}
# Check for tool calls in the chunk
if 'tool_calls' in message_chunk:
tool_calls_detected = True
for tool_call in message_chunk['tool_calls']:
function = tool_call.get('function', {})
tool_name = function.get('name')
tool_args = function.get('arguments')
# Parse arguments if they're a string
if isinstance(tool_args, str):
try:
tool_args = json.loads(tool_args)
except json.JSONDecodeError:
tool_args = {}
# Collect the tool call; execution happens after the
# stream finishes so calls can run in parallel.
# Skip duplicates: some servers stream the accumulated
# tool_calls list in every chunk.
if not any(
tc.get('function', {}).get('name') == tool_name
and tc.get('function', {}).get('arguments') == function.get('arguments')
for tc in pending_tool_calls
):
pending_tool_calls.append(tool_call)
yield {"type": "tool_call", "content": {"name": tool_name, "arguments": tool_args}}
# Execute all tool calls from this turn in parallel
if pending_tool_calls:
results = self._execute_tool_calls(pending_tool_calls)
for result in results:
# Yield tool result
yield {"type": "tool_result", "content": result}
# Add tool result to conversation
self.conversation_history.append({
"role": "tool",
"content": result
})
# Save complete response to history
complete_response = ''.join(collected_content)
if tool_calls_detected:
# Add assistant's message to history
self.conversation_history.append({
"role": "assistant",
"content": complete_response if complete_response else ""
})
# Continue the ReAct loop - let the model decide what to do next
# The loop will continue and get the next response
else:
# No tool calls - we have a final response
self.conversation_history.append({
"role": "assistant",
"content": complete_response
})
# Exit the ReAct loop
break
except Exception as e:
logger.error(f"Error in chat stream: {e}")
yield {"type": "error", "content": str(e)}
break
# Check if we hit max iterations
if iteration >= max_iterations:
yield {"type": "error", "content": "Maximum iterations reached in ReAct loop"}
def reset_conversation(self):
"""Reset the conversation history"""
self.conversation_history = []
logger.info("Conversation history reset")
class OllamaOpenAICompatible:
"""Use Ollama through its OpenAI-compatible endpoint"""
def __init__(self, model: str = "qwen3:0.6b",
base_url: str = "http://localhost:11434/v1"):
"""
Initialize using Ollama's OpenAI-compatible API
This provides better compatibility with tool calling
"""
from openai import OpenAI
self.model = model
self.client = OpenAI(
base_url=base_url,
api_key="ollama" # Ollama doesn't need a real key
)
self.tool_registry = ToolRegistry()
self.conversation_history = []
logger.info(f"✅ Initialized Ollama OpenAI-compatible client with {model}")
def chat(self, message: str, use_tools: bool = True,
temperature: float = 0.3) -> str:
"""
Chat using OpenAI-compatible endpoint
"""
# Add user message
self.conversation_history.append({
"role": "user",
"content": message
})
# Prepare tools
tools = self.tool_registry.get_tool_schemas() if use_tools else None
try:
# Call with tools
response = self.client.chat.completions.create(
model=self.model,
messages=self.conversation_history,
tools=tools,
tool_choice="auto" if tools else None,
temperature=temperature
)
assistant_message = response.choices[0].message
# Check for tool calls
if assistant_message.tool_calls:
logger.info(f"Model requested {len(assistant_message.tool_calls)} tool(s)")
# Add assistant message to history
self.conversation_history.append({
"role": "assistant",
"content": assistant_message.content or "",
"tool_calls": [
{
"id": tc.id,
"type": tc.type,
"function": {
"name": tc.function.name,
"arguments": tc.function.arguments
}
} for tc in assistant_message.tool_calls
]
})
# Execute tool calls (independent calls run in parallel;
# executor.map preserves order)
def run_one(tool_call):
# Parse arguments
try:
args = json.loads(tool_call.function.arguments)
except json.JSONDecodeError:
args = {}
# Execute tool
return self.tool_registry.execute_tool(
tool_call.function.name,
args
)
tool_calls_list = list(assistant_message.tool_calls)
if len(tool_calls_list) <= 1:
results = [run_one(tc) for tc in tool_calls_list]
else:
with ThreadPoolExecutor(max_workers=len(tool_calls_list)) as executor:
results = list(executor.map(run_one, tool_calls_list))
# Add tool results
for tool_call, result in zip(tool_calls_list, results):
self.conversation_history.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": result
})
# Get final response
final_response = self.client.chat.completions.create(
model=self.model,
messages=self.conversation_history,
tools=tools, # IMPORTANT: Keep tools available
temperature=temperature
)
final_content = final_response.choices[0].message.content
# Add to history
self.conversation_history.append({
"role": "assistant",
"content": final_content
})
return final_content
else:
# No tool calls
content = assistant_message.content
self.conversation_history.append({
"role": "assistant",
"content": content
})
return content
except Exception as e:
logger.error(f"Error: {e}")
return f"Error: {e}"
def reset_conversation(self):
"""Reset conversation history"""
self.conversation_history = []
logger.info("Conversation reset")
def test_native_tools():
"""Test Ollama's native tool calling"""
print("="*60)
print("🔧 Testing Ollama Native Tool Calling")
print("="*60)
# Test with default model
models_to_test = [
"qwen3:0.6b", # Default model for this project
]
for model_name in models_to_test:
print(f"\n📦 Testing with {model_name}")
print("-"*40)
try:
# Check if model is available
client = ollama.Client()
available_models = [m['name'] for m in client.list()['models']]
if not any(model_name in m for m in available_models):
print(f"⚠️ Model {model_name} not installed")
print(f" Install with: ollama pull {model_name}")
continue
# Test the model
agent = OllamaNativeAgent(model=model_name)
test_queries = [
"What's 15 * 23?",
"What's the weather in London?",
]
for query in test_queries:
print(f"\n👤 User: {query}")
response = agent.chat(query)
print(f"🤖 Assistant: {response[:200]}...") # Truncate long responses
agent.reset_conversation()
except Exception as e:
print(f"❌ Error testing {model_name}: {e}")
print("\n" + "="*60)
print("💡 Note:")
print("This project uses qwen3:0.6b as the default model.")
print("Install with: ollama pull qwen3:0.6b")
print("="*60)
def demo():
"""Interactive demo with proper tool calling"""
print("="*60)
print("🎯 Ollama Standard Tool Calling Demo")
print("="*60)
# Let user choose implementation
print("\nChoose implementation:")
print("1. Native Ollama API (recommended)")
print("2. OpenAI-compatible API")
choice = input("\nEnter choice (1 or 2): ").strip()
if choice == "2":
print("\nUsing OpenAI-compatible endpoint...")
agent = OllamaOpenAICompatible()
else:
print("\nUsing native Ollama API...")
# Check for best available model
try:
client = ollama.Client()
models = [m['name'] for m in client.list()['models']]
# Use qwen3:0.6b as the default model
model = "qwen3:0.6b"
if model in models:
print(f"Using recommended model: {model}")
else:
print(f"Recommended model {model} not found")
print("Install with: ollama pull qwen3:0.6b")
# Fall back to first available model
model = models[0] if models else "qwen3:0.6b"
print(f"Using fallback model: {model}")
agent = OllamaNativeAgent(model=model)
except Exception as e:
print(f"Error: {e}")
return
# Interactive loop
print("\n💬 Chat with the assistant (type 'exit' to quit)")
print("-"*40)
while True:
user_input = input("\n👤 You: ").strip()
if user_input.lower() in ['exit', 'quit']:
break
response = agent.chat(user_input)
print(f"🤖 Assistant: {response}")
print("\n👋 Goodbye!")
if __name__ == "__main__":
import sys
if len(sys.argv) > 1 and sys.argv[1] == "test":
test_native_tools()
else:
demo()
quick_test_streaming.py¶
server.py¶
"""
vLLM Server Launcher for Qwen3 with Tool Calling Support
"""
import os
import sys
import subprocess
import time
import requests
from pathlib import Path
from config import VLLM_SERVER_CONFIG, VLLM_HOST, VLLM_PORT
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class VLLMServer:
"""Manager for vLLM server process"""
def __init__(self, config: dict = None):
"""Initialize server manager with configuration"""
self.config = config or VLLM_SERVER_CONFIG
self.process = None
self.server_url = f"http://{self.config['host']}:{self.config['port']}"
def _build_command(self) -> list:
"""Build vLLM server command with arguments"""
cmd = [
sys.executable, "-m", "vllm.entrypoints.openai.api_server",
"--model", self.config["model"],
"--port", str(self.config["port"]),
"--host", self.config["host"],
]
# Add tool-specific arguments
if self.config.get("enable_auto_tool_choice"):
cmd.append("--enable-auto-tool-choice")
if self.config.get("tool_call_parser"):
cmd.extend(["--tool-call-parser", self.config["tool_call_parser"]])
if self.config.get("chat_template"):
cmd.extend(["--chat-template", self.config["chat_template"]])
# Add performance arguments
if self.config.get("max_model_len"):
cmd.extend(["--max-model-len", str(self.config["max_model_len"])])
if self.config.get("gpu_memory_utilization"):
cmd.extend(["--gpu-memory-utilization", str(self.config["gpu_memory_utilization"])])
if self.config.get("dtype"):
cmd.extend(["--dtype", self.config["dtype"]])
if self.config.get("enforce_eager"):
cmd.append("--enforce-eager")
# Add tensor parallel size if multiple GPUs
if self.config.get("tensor_parallel_size"):
cmd.extend(["--tensor-parallel-size", str(self.config["tensor_parallel_size"])])
return cmd
def start(self, wait_for_ready: bool = True, timeout: int = 120):
"""
Start the vLLM server
Args:
wait_for_ready: Wait for server to be ready
timeout: Maximum time to wait for server startup
"""
if self.is_running():
logger.info("vLLM server is already running")
return
# Create logs directory
log_dir = Path("logs")
log_dir.mkdir(exist_ok=True)
# Build command
cmd = self._build_command()
logger.info(f"Starting vLLM server with command: {' '.join(cmd)}")
# Start server process
log_file = log_dir / "vllm_server.log"
with open(log_file, "w") as f:
self.process = subprocess.Popen(
cmd,
stdout=f,
stderr=subprocess.STDOUT,
env=os.environ.copy()
)
logger.info(f"vLLM server process started with PID: {self.process.pid}")
logger.info(f"Server logs are being written to: {log_file}")
if wait_for_ready:
self._wait_for_ready(timeout)
def _wait_for_ready(self, timeout: int = 120):
"""Wait for server to be ready"""
start_time = time.time()
health_url = f"{self.server_url}/health"
logger.info(f"Waiting for vLLM server to be ready at {health_url}...")
while time.time() - start_time < timeout:
try:
response = requests.get(health_url, timeout=1)
if response.status_code == 200:
logger.info("vLLM server is ready!")
# Test model availability
models_url = f"{self.server_url}/v1/models"
models_response = requests.get(models_url)
if models_response.status_code == 200:
models = models_response.json()
logger.info(f"Available models: {models}")
return
except requests.exceptions.RequestException:
pass
# Check if process is still running
if self.process and self.process.poll() is not None:
raise RuntimeError(f"vLLM server process died with code: {self.process.returncode}")
time.sleep(2)
raise TimeoutError(f"vLLM server did not start within {timeout} seconds")
def stop(self):
"""Stop the vLLM server"""
if self.process:
logger.info("Stopping vLLM server...")
self.process.terminate()
try:
self.process.wait(timeout=10)
except subprocess.TimeoutExpired:
logger.warning("Server did not stop gracefully, forcing kill...")
self.process.kill()
self.process.wait()
self.process = None
logger.info("vLLM server stopped")
def is_running(self) -> bool:
"""Check if server is running"""
if not self.process:
return False
# Check if process is still alive
if self.process.poll() is not None:
return False
# Try to connect to health endpoint
try:
response = requests.get(f"{self.server_url}/health", timeout=1)
return response.status_code == 200
except requests.exceptions.RequestException:
return False
def restart(self):
"""Restart the server"""
logger.info("Restarting vLLM server...")
self.stop()
time.sleep(2)
self.start()
def download_model_from_modelscope():
"""
Download Qwen3-0.6B model from ModelScope
This is optional - vLLM can download from HuggingFace automatically
"""
try:
from modelscope import snapshot_download
model_dir = snapshot_download(
'Qwen/Qwen3-0.6B',
cache_dir='./models'
)
logger.info(f"Model downloaded to: {model_dir}")
return model_dir
except ImportError:
logger.warning("ModelScope not installed. Install with: pip install modelscope")
logger.info("vLLM will download from HuggingFace instead")
return None
def main():
"""Main function to start vLLM server"""
import argparse
parser = argparse.ArgumentParser(description="Start vLLM server with Qwen3 model")
parser.add_argument("--download", action="store_true",
help="Download model from ModelScope first")
parser.add_argument("--model", type=str, default=None,
help="Model name or path (overrides config)")
parser.add_argument("--port", type=int, default=None,
help="Server port (overrides config)")
parser.add_argument("--host", type=str, default=None,
help="Server host (overrides config)")
args = parser.parse_args()
# Download model if requested
if args.download:
model_path = download_model_from_modelscope()
if model_path:
VLLM_SERVER_CONFIG["model"] = model_path
# Override config with command line arguments
if args.model:
VLLM_SERVER_CONFIG["model"] = args.model
if args.port:
VLLM_SERVER_CONFIG["port"] = args.port
if args.host:
VLLM_SERVER_CONFIG["host"] = args.host
# Create and start server
server = VLLMServer(VLLM_SERVER_CONFIG)
try:
server.start(wait_for_ready=True)
logger.info(f"vLLM server is running at {server.server_url}")
logger.info("Press Ctrl+C to stop the server")
# Keep the server running
while True:
time.sleep(1)
if not server.is_running():
logger.error("Server stopped unexpectedly!")
break
except KeyboardInterrupt:
logger.info("\nShutting down...")
except Exception as e:
logger.error(f"Error: {e}")
finally:
server.stop()
if __name__ == "__main__":
main()
test_code_interpreter_full.py¶
"""
Test the full Python environment code interpreter with error handling
"""
import json
from tools import ToolRegistry
def test_successful_execution():
"""Test that code executes successfully with full Python environment"""
print("=" * 60)
print("Test 1: Successful execution with full Python environment")
print("=" * 60)
registry = ToolRegistry()
# Test with various Python features that would fail in a sandbox
test_cases = [
{
"name": "Complex calculation",
"code": "import numpy as np\nresult = np.array([1, 2, 3, 4, 5]).mean()"
},
{
"name": "File operations (simulated)",
"code": "import os\nresult = os.getcwd()"
},
{
"name": "Dict comprehension",
"code": "result = {i: i**2 for i in range(5)}"
},
{
"name": "Lambda and map",
"code": "result = list(map(lambda x: x**2, [1, 2, 3, 4, 5]))"
}
]
for test in test_cases:
print(f"\n{test['name']}:")
try:
result = registry.execute_tool("code_interpreter", {"code": test["code"]})
result_dict = json.loads(result)
if result_dict.get("success"):
print(f" ✓ Success: {result_dict.get('result')}")
else:
print(f" ✗ Failed: {result_dict.get('error')}")
except Exception as e:
print(f" ✗ Exception: {e}")
def test_error_handling():
"""Test that errors are properly captured and formatted"""
print("\n" + "=" * 60)
print("Test 2: Error handling and reporting")
print("=" * 60)
registry = ToolRegistry()
error_cases = [
{
"name": "Syntax Error",
"code": "if True\n print('missing colon')"
},
{
"name": "Name Error",
"code": "result = undefined_variable + 10"
},
{
"name": "Type Error",
"code": "result = '5' + 5"
},
{
"name": "Division by Zero",
"code": "result = 10 / 0"
},
{
"name": "Import Error",
"code": "import nonexistent_module\nresult = 42"
}
]
for test in error_cases:
print(f"\n{test['name']}:")
result = registry.execute_tool("code_interpreter", {"code": test["code"]})
result_dict = json.loads(result)
if not result_dict.get("success"):
print(f" ✓ Error properly caught:")
print(f" Error Type: {result_dict.get('error_type')}")
print(f" Error Message: {result_dict.get('error')}")
if result_dict.get('traceback'):
print(f" Traceback: {result_dict.get('traceback')[:100]}...")
else:
print(f" ✗ Error not caught - this shouldn't happen")
def test_full_environment_access():
"""Test that the code interpreter has access to full Python environment"""
print("\n" + "=" * 60)
print("Test 3: Full Python environment access")
print("=" * 60)
registry = ToolRegistry()
# Test access to various Python features that would be blocked in a sandbox
full_env_tests = [
{
"name": "Access to all builtins",
"code": "result = [callable(eval), callable(exec), callable(compile), callable(__import__)]"
},
{
"name": "Dynamic import",
"code": "import sys\nresult = f'Python {sys.version_info.major}.{sys.version_info.minor}'"
},
{
"name": "List comprehension with filter",
"code": "result = [x for x in range(20) if x % 2 == 0 and x % 3 == 0]"
},
{
"name": "Multiple variable assignment",
"code": "a, b, c = 1, 2, 3\nresult = a + b + c"
}
]
for test in full_env_tests:
print(f"\n{test['name']}:")
result = registry.execute_tool("code_interpreter", {"code": test["code"]})
result_dict = json.loads(result)
if result_dict.get("success"):
print(f" ✓ Success: {result_dict.get('result')}")
if result_dict.get('output'):
print(f" Output: {result_dict.get('output')}")
else:
print(f" ✗ Failed: {result_dict.get('error')}")
def test_agent_error_propagation():
"""Test that errors are properly formatted for the agent"""
print("\n" + "=" * 60)
print("Test 4: Agent error message formatting")
print("=" * 60)
from agent import VLLMToolAgent
# This test would require a running vLLM server, so we'll just show
# how errors would be formatted
registry = ToolRegistry()
# Simulate an error
result = registry.execute_tool("code_interpreter", {
"code": "result = 1 / 0"
})
result_dict = json.loads(result)
# Format as the agent would
if not result_dict.get("success"):
error_msg = f"❌ Tool 'code_interpreter' 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"
print("\nFormatted error message that would be sent to agent:")
print("-" * 60)
print(error_msg)
print("-" * 60)
print("\nThe agent will receive this error message and can:")
print(" 1. Try to fix the code")
print(" 2. Ask the user for clarification")
print(" 3. Provide an alternative solution")
if __name__ == "__main__":
test_successful_execution()
test_error_handling()
test_full_environment_access()
test_agent_error_propagation()
print("\n" + "=" * 60)
print("All tests completed!")
print("=" * 60)
print("\nSummary:")
print(" ✓ Full Python environment is available (no sandbox restrictions)")
print(" ✓ Errors are properly caught and detailed traceback provided")
print(" ✓ Error messages are formatted clearly for the agent")
print(" ✓ Agent can receive and process error information")
test_ollama_streaming_final.py¶
test_parallel_tools.py¶
"""
End-to-end test for parallel tool calling with a real LLM (Ollama).
Covers:
1. Deterministic proof of parallel execution: two 2s-sleep tools must finish
in ~2s (not ~4s) through each agent's tool-execution path.
2. Real-model runs of the book's "Vancouver time + weather" example through:
- OllamaNativeAgent.chat / chat_stream (native tool calling)
- OllamaOpenAICompatible.chat (OpenAI-compatible endpoint, native tools)
- VLLMToolAgent.chat / chat_stream (<tool_call> XML parsed from content;
the OpenAI-compatible endpoint of Ollama stands in for the vLLM server,
so the `tools` kwarg is dropped to make the model emit XML in content)
Run from this directory: python3 test_parallel_tools.py
Requires: ollama serve + ollama pull qwen2.5:7b-instruct-q8_0
"""
import json
import logging
import time
from types import SimpleNamespace
from ollama_native import OllamaNativeAgent, OllamaOpenAICompatible
from agent import VLLMToolAgent
logging.basicConfig(level=logging.WARNING)
SLEEP = 2
QUERY = "What time is it in Vancouver and what's the current weather in Vancouver?"
# qwen3:0.6b is the project default but flaky at native tool calling;
# qwen2.5:7b-instruct-q8_0 (also pulled locally) calls tools reliably.
REAL_MODEL = "qwen2.5:7b-instruct-q8_0"
MAX_ATTEMPTS = 4 # small models occasionally skip tool calling; retry
# ---------------------------------------------------------------- helpers
def _sleep_tool(name):
def fn(tag: str = "") -> dict:
time.sleep(SLEEP)
return {"tool": name, "slept": SLEEP, "success": True}
return fn
SLEEP_SCHEMA = {
"type": "object",
"properties": {"tag": {"type": "string", "description": "optional tag"}},
"required": [],
}
def _register_sleep_tools(registry):
registry.register_tool("sleep_a", _sleep_tool("sleep_a"), "Sleep tool A", SLEEP_SCHEMA)
registry.register_tool("sleep_b", _sleep_tool("sleep_b"), "Sleep tool B", SLEEP_SCHEMA)
def _check_parallel(label, elapsed, n=2):
status = "PARALLEL OK" if elapsed < SLEEP * 1.8 else "NOT PARALLEL"
print(f" [{label}] {n} x {SLEEP}s tools finished in {elapsed:.1f}s -> {status}")
assert elapsed < SLEEP * 1.8, f"{label}: tool calls were not executed in parallel"
# ------------------------------------- 1. deterministic parallel checks
def test_native_execute_parallel():
print("\n== OllamaNativeAgent._execute_tool_calls (sleep tools) ==")
agent = OllamaNativeAgent(model=REAL_MODEL)
_register_sleep_tools(agent.tool_registry)
calls = [
{"function": {"name": "sleep_a", "arguments": {}}},
{"function": {"name": "sleep_b", "arguments": {}}},
]
start = time.time()
results = agent._execute_tool_calls(calls)
_check_parallel("native _execute_tool_calls", time.time() - start)
assert all(json.loads(r)["success"] for r in results)
def _make_chunk(text):
return SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content=text))])
def test_vllm_stream_parallel():
print("\n== VLLMToolAgent.chat_stream (fake stream with 2 tool calls) ==")
agent = VLLMToolAgent(api_base="http://localhost:11434/v1", api_key="ollama")
_register_sleep_tools(agent.tool_registry)
tool_xml = (
'<tool_call>{"name": "sleep_a", "arguments": {}}</tool_call>'
'<tool_call>{"name": "sleep_b", "arguments": {}}</tool_call>'
)
state = {"n": 0}
def fake_create(**kwargs):
state["n"] += 1
if state["n"] == 1:
# split the two tool calls across chunks to exercise buffering
parts = [tool_xml[:25], tool_xml[25:90], tool_xml[90:]]
return iter([_make_chunk(p) for p in parts])
return iter([_make_chunk("done")])
agent.client = SimpleNamespace(
chat=SimpleNamespace(completions=SimpleNamespace(create=fake_create))
)
start = time.time()
events = list(agent.chat_stream("run both sleep tools"))
elapsed = time.time() - start
types = [e["type"] for e in events]
print(f" event types: {types}")
assert types.count("tool_call") == 2, f"expected 2 tool_call events: {types}"
assert types.count("tool_result") == 2, f"expected 2 tool_result events: {types}"
_check_parallel("vllm chat_stream", elapsed)
# ------------------------------------- 2. real-model end-to-end runs
def test_native_real_model():
print(f"\n== OllamaNativeAgent.chat (real {REAL_MODEL}) ==")
agent = OllamaNativeAgent(model=REAL_MODEL)
response, tool_msgs = "", []
for attempt in range(1, MAX_ATTEMPTS + 1):
agent.reset_conversation()
response = agent.chat(QUERY)
tool_msgs = [m for m in agent.conversation_history if m.get("role") == "tool"]
if tool_msgs:
break
print(f" attempt {attempt}: model made no tool call, retrying")
assistant_tc = [m for m in agent.conversation_history if m.get("tool_calls")]
batch = len(assistant_tc[0]["tool_calls"]) if assistant_tc else 0
print(f" tool calls in one turn: {batch} ({'PARALLEL BATCH' if batch > 1 else 'single'})")
for m in tool_msgs:
print(f" - {m['content'][:120]}")
print(f" final response: {response[:300]}")
assert tool_msgs, "model did not call any tool"
print(f"\n== OllamaNativeAgent.chat_stream (real {REAL_MODEL}) ==")
events = []
for attempt in range(1, MAX_ATTEMPTS + 1):
agent.reset_conversation()
events = list(agent.chat_stream(QUERY))
if any(e["type"] == "tool_result" for e in events):
break
print(f" attempt {attempt}: no tool call, retrying")
tool_calls = [e for e in events if e["type"] == "tool_call"]
tool_results = [e for e in events if e["type"] == "tool_result"]
print(f" tool_call events: {len(tool_calls)}, tool_result events: {len(tool_results)}")
for e in tool_calls:
print(f" - {e['content']['name']}({e['content']['arguments']})")
final = "".join(e["content"] for e in events if e["type"] == "content")
print(f" final content: {final[:300]}")
assert tool_results, "streaming path produced no tool results"
def test_openai_compat_real_model():
print(f"\n== OllamaOpenAICompatible.chat (real {REAL_MODEL}) ==")
agent = OllamaOpenAICompatible(model=REAL_MODEL)
response, tool_msgs = "", []
for attempt in range(1, MAX_ATTEMPTS + 1):
agent.reset_conversation()
response = agent.chat(QUERY)
tool_msgs = [m for m in agent.conversation_history if m.get("role") == "tool"]
if tool_msgs:
break
print(f" attempt {attempt}: model made no tool call, retrying")
print(f" tool results in history: {len(tool_msgs)}")
print(f" final response: {response[:300]}")
assert tool_msgs, "model did not call any tool"
def test_vllm_agent_real_model():
print(f"\n== VLLMToolAgent.chat (real {REAL_MODEL} via OpenAI endpoint) ==")
agent = VLLMToolAgent(api_base="http://localhost:11434/v1", api_key="ollama")
# Ollama's OpenAI endpoint stands in for the vLLM server. The agent parses
# <tool_call> XML from content, so drop the native `tools` kwarg (which
# would make Ollama return structured tool_calls instead) and fix the
# hardcoded vLLM model name to the Ollama one.
real_create = agent.client.chat.completions.create
def create_xml_mode(**kwargs):
kwargs.pop("tools", None)
kwargs.pop("tool_choice", None)
kwargs["model"] = REAL_MODEL
return real_create(**kwargs)
agent.client = SimpleNamespace(
chat=SimpleNamespace(completions=SimpleNamespace(create=create_xml_mode))
)
response, tool_msgs, batch = "", [], 0
for attempt in range(1, MAX_ATTEMPTS + 1):
agent.reset_conversation()
# temperature 0.7 sampling can degenerate into a tool-call loop with
# this handcrafted XML format; 0.3 is stable
response = agent.chat(QUERY, temperature=0.3)
tool_msgs = [m for m in agent.conversation_history if m.get("name")]
assistant_tc = [m for m in agent.conversation_history if m.get("tool_calls")]
batch = len(assistant_tc[0]["tool_calls"]) if assistant_tc else 0
if tool_msgs and batch <= 10:
break
print(f" attempt {attempt}: no tool call or degenerate run ({batch} calls), retrying")
print(f" tool calls in one turn: {batch} ({'PARALLEL BATCH' if batch > 1 else 'single'})")
for m in tool_msgs:
print(f" - {m['name']}: {m['content'][:100]}")
print(f" final response: {response[:300]}")
assert tool_msgs, "model did not emit <tool_call> XML"
print(f"\n== VLLMToolAgent.chat_stream (real {REAL_MODEL}) ==")
events = []
for attempt in range(1, MAX_ATTEMPTS + 1):
agent.reset_conversation()
# temperature 0.7 sampling can degenerate into a tool-call loop with
# this handcrafted XML format; 0.3 is stable
events = list(agent.chat_stream(QUERY, temperature=0.3))
n_calls = sum(1 for e in events if e["type"] == "tool_call")
if any(e["type"] == "tool_result" for e in events) and n_calls <= 10:
break
print(f" attempt {attempt}: no tool call or degenerate run ({n_calls} calls), retrying")
tool_calls = [e for e in events if e["type"] == "tool_call"]
tool_results = [e for e in events if e["type"] == "tool_result"]
print(f" tool_call events: {len(tool_calls)}, tool_result events: {len(tool_results)}")
for e in tool_calls:
print(f" - {e['content']}")
final = "".join(e["content"] for e in events if e["type"] == "content")
print(f" final content: {final[:300]}")
assert tool_results, "streaming path produced no tool results"
if __name__ == "__main__":
test_native_execute_parallel()
test_vllm_stream_parallel()
test_native_real_model()
test_openai_compat_real_model()
test_vllm_agent_real_model()
print("\nALL TESTS PASSED")
test_sample_command.py¶
test_streaming.py¶
#!/usr/bin/env python3
"""
Test script to demonstrate streaming functionality
for both vLLM and Ollama backends
"""
import sys
import platform
from main import ToolCallingAgent
def test_streaming():
"""Test streaming functionality with various queries"""
print("="*60)
print("🚀 STREAMING TEST DEMO")
print("="*60)
print(f"Platform: {platform.system()}")
print("="*60)
# Initialize agent
print("\n⚙️ Initializing agent...")
agent = ToolCallingAgent()
print(f"✅ Using {agent.backend_type} backend")
# Test queries that will demonstrate streaming features
test_queries = [
{
"name": "Simple Calculation with Thinking",
"query": "Calculate 15 * 23 + sqrt(144). Think through the steps."
},
{
"name": "Tool Usage with Weather",
"query": "What's the weather in Tokyo? If it's hot (above 25°C), suggest some cooling tips."
},
{
"name": "Multiple Tools",
"query": "Convert 100 USD to EUR and tell me the current time in London."
}
]
for test in test_queries:
print("\n" + "="*60)
print(f"📋 Test: {test['name']}")
print("="*60)
print(f"Query: {test['query']}")
print("-"*60)
try:
print("\n🔄 Streaming response:\n")
thinking_shown = False
tools_shown = False
response_shown = False
# Stream the response
for chunk in agent.chat(test['query'], stream=True):
chunk_type = chunk.get("type")
content = chunk.get("content", "")
if chunk_type == "thinking":
if not thinking_shown:
print("🧠 Internal Thinking:")
print(" ", end="")
thinking_shown = True
# Show thinking in gray/dim text
print(f"\033[90m{content}\033[0m")
elif chunk_type == "tool_call":
if not tools_shown:
print("\n🔧 Tool Calls:")
tools_shown = True
tool_info = content
print(f" 📦 Calling: {tool_info.get('name', 'unknown')}")
print(f" Arguments: {tool_info.get('arguments', {})}")
elif chunk_type == "tool_result":
result_str = str(content)
print(f" ✓ Result: {result_str}")
elif chunk_type == "content":
if not response_shown:
print("\n🤖 Assistant Response:")
print(" ", end="")
response_shown = True
# Stream the content character by character
print(content, end="", flush=True)
elif chunk_type == "error":
print(f"\n❌ Error: {content}")
print("\n") # New line after response
except Exception as e:
print(f"\n❌ Error during test: {e}")
# Reset conversation for next test
agent.reset_conversation()
# Ask user if they want to continue
if test != test_queries[-1]:
cont = input("\nPress Enter to continue to next test (or 'q' to quit): ")
if cont.lower() == 'q':
break
print("\n" + "="*60)
print("✅ Streaming test completed!")
print("="*60)
def compare_streaming_vs_regular():
"""Compare streaming vs regular responses"""
print("="*60)
print("📊 STREAMING VS REGULAR COMPARISON")
print("="*60)
# Initialize agent
agent = ToolCallingAgent()
test_query = "What's the weather in Paris and convert 20°C to Fahrenheit?"
print(f"\n📋 Test Query: {test_query}")
print("="*60)
# Regular mode
print("\n1️⃣ REGULAR MODE (No Streaming):")
print("-"*40)
print("⏳ Processing...")
response = agent.chat(test_query, stream=False)
print(f"🤖 Response: {response}")
agent.reset_conversation()
# Streaming mode
print("\n2️⃣ STREAMING MODE:")
print("-"*40)
print("⏳ Processing (you'll see content as it arrives)...\n")
for chunk in agent.chat(test_query, stream=True):
chunk_type = chunk.get("type")
content = chunk.get("content", "")
if chunk_type == "thinking":
print(f"[THINKING] \033[90m{content}\033[0m")
elif chunk_type == "tool_call":
print(f"[TOOL CALL] {content}")
elif chunk_type == "tool_result":
print(f"[TOOL RESULT] {content}")
elif chunk_type == "content":
print(content, end="", flush=True)
print("\n\n" + "="*60)
print("✅ Comparison complete!")
print("💡 Streaming mode shows intermediate steps in real-time")
print("="*60)
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description="Test streaming functionality")
parser.add_argument(
"--mode",
choices=["demo", "compare"],
default="demo",
help="Test mode: demo (full demo) or compare (streaming vs regular)"
)
args = parser.parse_args()
if args.mode == "demo":
test_streaming()
else:
compare_streaming_vs_regular()
test_streaming_fix.py¶
test_weather.py¶
tools.py¶
"""
Sample tools for demonstrating vLLM tool calling functionality
"""
import json
import math
import random
import io
import contextlib
from typing import Dict, Any, List
from datetime import datetime
import requests
from io import BytesIO
import PyPDF2
class ToolRegistry:
"""Registry for managing available tools"""
def __init__(self):
self.tools = {}
self._register_default_tools()
def _register_default_tools(self):
"""Register default tools"""
self.register_tool(
name="get_current_temperature",
function=self.get_current_temperature,
description="Get the current temperature for a specific location",
parameters={
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "The city and country, e.g., 'Paris, France'"
},
"unit": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
"description": "The temperature unit to use (by default, celsius)"
}
},
"required": ["location", "unit"]
}
)
self.register_tool(
name="get_current_time",
function=self.get_current_time,
description="Get the current date and time in a specific timezone",
parameters={
"type": "object",
"properties": {
"timezone": {
"type": "string",
"description": "Timezone name (e.g., 'America/New_York', 'Europe/London', 'Asia/Tokyo'). Use standard IANA timezone names.",
"default": "UTC"
}
},
"required": []
}
)
self.register_tool(
name="convert_currency",
function=self.convert_currency,
description="Convert an amount from one currency to another. You MUST use this tool to convert currencies in order to get the latest exchange rate.",
parameters={
"type": "object",
"properties": {
"amount": {
"type": "number",
"description": "Amount to convert"
},
"from_currency": {
"type": "string",
"description": "Source currency code (e.g., 'USD', 'EUR')"
},
"to_currency": {
"type": "string",
"description": "Target currency code (e.g., 'USD', 'EUR')"
}
},
"required": ["amount", "from_currency", "to_currency"]
}
)
self.register_tool(
name="code_interpreter",
function=self.code_interpreter,
description="Execute Python code for calculations and data processing. You MUST use this tool to perform any complex calculations or data processing.",
parameters={
"type": "object",
"properties": {
"code": {
"type": "string",
"description": "Python code to execute"
}
},
"required": ["code"]
}
)
def register_tool(self, name: str, function: callable, description: str, parameters: Dict):
"""Register a new tool"""
self.tools[name] = {
"function": function,
"description": description,
"parameters": parameters
}
def get_tool_schemas(self) -> List[Dict]:
"""Get OpenAI-compatible tool schemas"""
schemas = []
for name, tool in self.tools.items():
schemas.append({
"type": "function",
"function": {
"name": name,
"description": tool["description"],
"parameters": tool["parameters"]
}
})
return schemas
def execute_tool(self, name: str, arguments: Dict[str, Any]) -> str:
"""Execute a tool by name with given arguments"""
if name not in self.tools:
return json.dumps({"error": f"Tool '{name}' not found"})
try:
result = self.tools[name]["function"](**arguments)
return json.dumps(result) if isinstance(result, (dict, list)) else str(result)
except Exception as e:
return json.dumps({"error": str(e)})
# Tool implementations
@staticmethod
def get_current_temperature(location: str, unit: str = "celsius") -> Dict:
"""
Get current temperature using Open-Meteo free weather API
No API key required - https://open-meteo.com/
"""
try:
# First, geocode the location to get coordinates
geocoding_url = "https://geocoding-api.open-meteo.com/v1/search"
geo_params = {
"name": location,
"count": 1,
"language": "en",
"format": "json"
}
geo_response = requests.get(geocoding_url, params=geo_params, timeout=5)
geo_data = geo_response.json()
if not geo_data.get("results"):
return {
"location": location,
"error": f"Location '{location}' not found",
"timestamp": datetime.now().isoformat()
}
# Get coordinates from first result
result = geo_data["results"][0]
latitude = result["latitude"]
longitude = result["longitude"]
location_name = f"{result.get('name', location)}, {result.get('country', '')}"
# Get current weather from Open-Meteo
weather_url = "https://api.open-meteo.com/v1/forecast"
# Determine temperature unit
temp_unit = "fahrenheit" if unit.lower() == "fahrenheit" else "celsius"
weather_params = {
"latitude": latitude,
"longitude": longitude,
"current": "temperature_2m,relative_humidity_2m,weather_code,wind_speed_10m",
"temperature_unit": temp_unit,
"timezone": "auto"
}
weather_response = requests.get(weather_url, params=weather_params, timeout=5)
weather_data = weather_response.json()
if "current" not in weather_data:
return {
"location": location_name,
"error": "Weather data not available",
"timestamp": datetime.now().isoformat()
}
current = weather_data["current"]
# Map weather codes to conditions
weather_codes = {
0: "clear sky",
1: "mainly clear", 2: "partly cloudy", 3: "overcast",
45: "foggy", 48: "foggy",
51: "light drizzle", 53: "moderate drizzle", 55: "dense drizzle",
61: "light rain", 63: "moderate rain", 65: "heavy rain",
71: "light snow", 73: "moderate snow", 75: "heavy snow",
77: "snow grains",
80: "light rain showers", 81: "moderate rain showers", 82: "heavy rain showers",
85: "light snow showers", 86: "heavy snow showers",
95: "thunderstorm", 96: "thunderstorm with light hail", 99: "thunderstorm with heavy hail"
}
weather_code = current.get("weather_code", 0)
conditions = weather_codes.get(weather_code, "unknown")
unit_symbol = "°F" if unit.lower() == "fahrenheit" else "°C"
return {
"location": location_name,
"temperature": round(current["temperature_2m"], 1),
"unit": unit_symbol,
"conditions": conditions,
"humidity": current.get("relative_humidity_2m"),
"wind_speed": round(current.get("wind_speed_10m", 0), 1),
"wind_unit": "km/h",
"coordinates": {"latitude": latitude, "longitude": longitude},
"timestamp": current.get("time", datetime.now().isoformat()),
"source": "Open-Meteo"
}
except requests.RequestException as e:
# Fallback to simulated data if API fails
import logging
logging.warning(f"Open-Meteo API error: {e}. Using simulated data.")
# Simulated fallback
base_temp = 20 + random.uniform(-10, 10)
if unit == "fahrenheit":
temp = base_temp * 9/5 + 32
unit_symbol = "°F"
else:
temp = base_temp
unit_symbol = "°C"
return {
"location": location,
"temperature": round(temp, 1),
"unit": unit_symbol,
"conditions": random.choice(["sunny", "cloudy", "partly cloudy", "rainy"]),
"timestamp": datetime.now().isoformat(),
"note": "Simulated data (API unavailable)"
}
except Exception as e:
return {
"location": location,
"error": str(e),
"timestamp": datetime.now().isoformat()
}
@staticmethod
def get_current_time(timezone: str = "UTC") -> Dict:
"""
Get current date and time in specified timezone using zoneinfo (Python 3.9+)
"""
from datetime import datetime
from zoneinfo import ZoneInfo
# Common abbreviation mappings to IANA timezone names
timezone_aliases = {
"EST": "America/New_York",
"EDT": "America/New_York",
"PST": "America/Los_Angeles",
"PDT": "America/Los_Angeles",
"CST": "America/Chicago",
"CDT": "America/Chicago",
"MST": "America/Denver",
"MDT": "America/Denver",
"GMT": "Europe/London",
"BST": "Europe/London",
"CET": "Europe/Paris",
"CEST": "Europe/Paris",
"JST": "Asia/Tokyo",
"IST": "Asia/Kolkata",
"AEST": "Australia/Sydney",
"AEDT": "Australia/Sydney",
"SGT": "Asia/Singapore",
"HKT": "Asia/Hong_Kong",
"UTC+1": "Etc/GMT-1", # Note: signs are inverted in Etc/GMT
"UTC-1": "Etc/GMT+1",
"UTC+8": "Etc/GMT-8",
"UTC-8": "Etc/GMT+8"
}
# Convert abbreviation to IANA name if needed
tz_name = timezone_aliases.get(timezone.upper(), timezone)
try:
tz = ZoneInfo(tz_name)
current_time = datetime.now(tz)
return {
"timezone": tz_name,
"datetime": current_time.strftime("%Y-%m-%d %H:%M:%S"),
"date": current_time.strftime("%Y-%m-%d"),
"time": current_time.strftime("%H:%M:%S"),
"day_of_week": current_time.strftime("%A"),
"utc_offset": current_time.strftime("%z"),
"timestamp": current_time.isoformat()
}
except Exception as e:
# Fallback to UTC if timezone not found
try:
tz_utc = ZoneInfo("UTC")
current_time = datetime.now(tz_utc)
return {
"timezone": "UTC",
"datetime": current_time.strftime("%Y-%m-%d %H:%M:%S"),
"date": current_time.strftime("%Y-%m-%d"),
"time": current_time.strftime("%H:%M:%S"),
"day_of_week": current_time.strftime("%A"),
"utc_offset": "+0000",
"timestamp": current_time.isoformat(),
"note": f"Invalid timezone '{timezone}', using UTC as fallback"
}
except Exception as fallback_error:
return {
"error": str(e),
"fallback_error": str(fallback_error),
"timezone": timezone,
"timestamp": datetime.utcnow().isoformat()
}
@staticmethod
def convert_currency(amount: float, from_currency: str, to_currency: str) -> Dict:
"""
Convert currency using live exchange rates (simulated)
"""
# Normalize currency codes
from_currency = from_currency.upper().replace("S$", "SGD").replace("$", "USD")
to_currency = to_currency.upper().replace("S$", "SGD").replace("$", "USD")
# Simulated exchange rates
exchange_rates = {
"USD": 1.0,
"EUR": 0.92,
"GBP": 0.79,
"JPY": 149.50,
"CNY": 7.24,
"CAD": 1.36,
"AUD": 1.53,
"CHF": 0.88,
"INR": 83.12,
"SGD": 1.34,
"KRW": 1330.50,
"MXN": 17.10
}
if from_currency not in exchange_rates or to_currency not in exchange_rates:
return {"error": f"Unsupported currency: {from_currency} or {to_currency}"}
# Convert to USD first, then to target currency
usd_amount = amount / exchange_rates[from_currency]
converted_amount = usd_amount * exchange_rates[to_currency]
return {
"original_amount": amount,
"from_currency": from_currency,
"to_currency": to_currency,
"converted_amount": round(converted_amount, 2),
"exchange_rate": round(exchange_rates[to_currency] / exchange_rates[from_currency], 4),
"timestamp": datetime.now().isoformat()
}
@staticmethod
def parse_pdf(url: str) -> Dict:
"""
Parse a PDF document from URL or local file
"""
try:
# Check if it's a local file
if url.startswith('file://') or url.startswith('/') or url.startswith('./'):
# Local file
file_path = url.replace('file://', '')
with open(file_path, 'rb') as f:
pdf_content = f.read()
else:
# Remote URL
response = requests.get(url, timeout=30)
response.raise_for_status()
pdf_content = response.content
# Parse PDF
pdf_file = BytesIO(pdf_content)
pdf_reader = PyPDF2.PdfReader(pdf_file)
text_content = []
for page_num, page in enumerate(pdf_reader.pages, 1):
text = page.extract_text()
text_content.append({
"page": page_num,
"text": text[:1000] # Limit text per page
})
return {
"url": url,
"num_pages": len(pdf_reader.pages),
"content": text_content[:5], # Limit to first 5 pages
"success": True
}
except Exception as e:
return {"error": str(e), "success": False}
@staticmethod
def code_interpreter(code: str) -> Dict:
"""
Execute Python code in a full Python environment.
This provides unrestricted access to Python's built-in functions and standard library.
"""
try:
# Strip markdown code blocks and other formatting
import re
# Remove ```python or ```py or ``` blocks
code = re.sub(r'^```(?:python|py)?\s*\n', '', code.strip())
code = re.sub(r'\n```\s*$', '', code)
code = re.sub(r'^```\s*', '', code)
code = re.sub(r'\s*```$', '', code)
# Also strip any leading/trailing whitespace
code = code.strip()
# Convert common mathematical notation to Python syntax
# Replace ^ with ** for exponentiation
code = re.sub(r'\^', '**', code)
# Create a full Python namespace with all builtins available
# This gives the agent access to the complete Python environment
import sys
namespace = {
'__builtins__': __builtins__,
'math': math,
'random': random,
'datetime': datetime,
'sys': sys,
're': re,
'json': json
}
# Capture both stdout and stderr
output_buffer = io.StringIO()
error_buffer = io.StringIO()
with contextlib.redirect_stdout(output_buffer), contextlib.redirect_stderr(error_buffer):
exec(code, namespace)
# Get output and any error messages
printed_output = output_buffer.getvalue()
error_output = error_buffer.getvalue()
# Try to get result from common variable names
result = namespace.get('result', None)
if result is None:
for var_name in ['A', 'total', 'sum', 'output', 'answer', 'final', 'value']:
if var_name in namespace:
result = namespace[var_name]
break
response = {
"result": result,
"output": printed_output if printed_output else None,
"stderr": error_output if error_output else None,
"success": True
}
return response
except SyntaxError as e:
error_msg = f"Syntax Error on line {e.lineno}: {e.msg}\n{e.text}"
return {
"error": error_msg,
"error_type": "SyntaxError",
"success": False
}
except Exception as e:
import traceback
error_trace = traceback.format_exc()
return {
"error": str(e),
"error_type": type(e).__name__,
"traceback": error_trace,
"success": False
}
def format_tool_response(tool_name: str, tool_result: str) -> Dict:
"""Format tool response for the chat model"""
return {
"role": "tool",
"name": tool_name,
"content": tool_result
}
setup.sh¶
#!/bin/bash
# vLLM Tool Calling Demo - Setup Script
# This script helps set up the environment for the demo
echo "======================================"
echo "vLLM Tool Calling Demo Setup"
echo "======================================"
# Check Python version
echo -e "\n1. Checking Python version..."
python_version=$(python3 --version 2>&1 | grep -Po '(?<=Python )\d+\.\d+')
required_version="3.8"
if [ "$(printf '%s\n' "$required_version" "$python_version" | sort -V | head -n1)" = "$required_version" ]; then
echo "✅ Python $python_version is installed (>= $required_version required)"
else
echo "❌ Python $python_version is too old. Please install Python >= $required_version"
exit 1
fi
# Check CUDA availability
echo -e "\n2. Checking CUDA availability..."
if command -v nvidia-smi &> /dev/null; then
echo "✅ NVIDIA GPU detected:"
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader
else
echo "⚠️ No NVIDIA GPU detected. vLLM requires a CUDA-capable GPU."
echo " You can still install dependencies but won't be able to run the model locally."
fi
# Create virtual environment
echo -e "\n3. Setting up Python virtual environment..."
if [ ! -d "venv" ]; then
python3 -m venv venv
echo "✅ Virtual environment created"
else
echo "✅ Virtual environment already exists"
fi
# Activate virtual environment
source venv/bin/activate
# Upgrade pip
echo -e "\n4. Upgrading pip..."
pip install --upgrade pip
# Install requirements
echo -e "\n5. Installing requirements..."
pip install -r requirements.txt
# Check PyTorch CUDA
echo -e "\n6. Checking PyTorch CUDA support..."
python3 -c "import torch; print('✅ PyTorch CUDA available' if torch.cuda.is_available() else '❌ PyTorch CUDA not available')"
# Create .env file if it doesn't exist
echo -e "\n7. Setting up environment configuration..."
if [ ! -f ".env" ]; then
cp env.example .env
echo "✅ Created .env file from template"
else
echo "✅ .env file already exists"
fi
# Create logs directory
echo -e "\n8. Creating logs directory..."
mkdir -p logs
echo "✅ Logs directory created"
# Optional: Install ModelScope for downloading from Chinese mirror
echo -e "\n9. Optional packages..."
read -p "Install ModelScope for downloading models from Chinese mirror? (y/n): " -n 1 -r
echo
if [[ $REPLY =~ ^[Yy]$ ]]; then
pip install modelscope
echo "✅ ModelScope installed"
fi
echo -e "\n======================================"
echo "Setup Complete!"
echo "======================================"
echo ""
echo "Next steps:"
echo "1. Activate the virtual environment: source venv/bin/activate"
echo "2. Check system compatibility: python check_compatibility.py"
echo "3. Run the main script: python main.py"
echo ""
echo "The script will automatically detect your platform and use:"
echo " - vLLM if you have an NVIDIA GPU"
echo " - Ollama on Mac or systems without GPU"
echo ""
echo "For more information, see README.md"