mirror of
https://github.com/hanzoai/python-sdk.git
synced 2026-10-23 20:23:00 +00:00
feat: Add CLI tools support to hanzo-mcp with batch execution
- Implement unified CLI tools module for all AI providers - Add support for claude/cc, codex, gemini, grok, openhands/oh, hanzo-dev, cline, aider - Enable parallel batch execution of CLI tools - Add comprehensive test suite with 16 passing tests - Fix typing issues with mypy strict mode - Add GitHub Actions CI/CD workflow with optional integration tests - Support authentication flow for all providers - Follow DRY principle with single source of truth
This commit is contained in:
+1
-1
@@ -6,7 +6,7 @@ import logging
|
||||
from typing import TYPE_CHECKING, Iterator, AsyncIterator
|
||||
|
||||
# Add pkg directory to path for hanzoai import
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(__file__)), 'pkg'))
|
||||
sys.path.insert(0, os.path.join(os.path.dirname(os.path.dirname(__file__)), "pkg"))
|
||||
|
||||
import pytest
|
||||
from pytest_asyncio import is_async_test
|
||||
|
||||
@@ -38,125 +38,133 @@ from hanzo_network.local_network import check_local_llm_status
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class HanzoNetE2ETest:
|
||||
"""End-to-end test for Hanzo Net orchestration."""
|
||||
|
||||
|
||||
def __init__(self):
|
||||
self.hanzo_net_process = None
|
||||
self.hanzo_dev_process = None
|
||||
self.test_results = []
|
||||
self.models_to_test = [
|
||||
"qwen3", # Qwen 3
|
||||
"llama-3.2-3b", # Llama 3.2 3B
|
||||
"deepseek-v3", # DeepSeek V3
|
||||
"mistral-7b", # Mistral 7B
|
||||
"qwen3", # Qwen 3
|
||||
"llama-3.2-3b", # Llama 3.2 3B
|
||||
"deepseek-v3", # DeepSeek V3
|
||||
"mistral-7b", # Mistral 7B
|
||||
]
|
||||
self.api_models = [
|
||||
"claude-3-5-sonnet-20241022",
|
||||
"gpt-4-turbo-preview",
|
||||
"gemini-pro"
|
||||
"gpt-4-turbo-preview",
|
||||
"gemini-pro",
|
||||
]
|
||||
|
||||
|
||||
async def setup(self):
|
||||
"""Set up test environment."""
|
||||
logger.info("Setting up E2E test environment...")
|
||||
|
||||
|
||||
# Check if hanzo net is installed
|
||||
result = subprocess.run(["which", "hanzo"], capture_output=True, text=True)
|
||||
if result.returncode != 0:
|
||||
raise RuntimeError("hanzo CLI not found. Please install: pip install -e pkg/hanzo/")
|
||||
|
||||
raise RuntimeError(
|
||||
"hanzo CLI not found. Please install: pip install -e pkg/hanzo/"
|
||||
)
|
||||
|
||||
# Check for required environment variables
|
||||
self.check_api_keys()
|
||||
|
||||
|
||||
# Create test workspace
|
||||
self.workspace = Path.home() / ".hanzo" / "e2e-test"
|
||||
self.workspace.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
logger.info(f"Test workspace: {self.workspace}")
|
||||
|
||||
|
||||
def check_api_keys(self):
|
||||
"""Check for required API keys."""
|
||||
required_keys = {
|
||||
"OPENAI_API_KEY": "OpenAI (GPT-4/Codex)",
|
||||
"ANTHROPIC_API_KEY": "Anthropic (Claude)",
|
||||
"GOOGLE_API_KEY": "Google (Gemini)"
|
||||
"GOOGLE_API_KEY": "Google (Gemini)",
|
||||
}
|
||||
|
||||
|
||||
missing_keys = []
|
||||
for key, service in required_keys.items():
|
||||
if not os.getenv(key):
|
||||
missing_keys.append(f"{key} ({service})")
|
||||
|
||||
|
||||
if missing_keys:
|
||||
logger.warning(f"Missing API keys: {', '.join(missing_keys)}")
|
||||
logger.warning("Some API models will not be available")
|
||||
|
||||
|
||||
async def start_hanzo_net(self, model: str = "qwen3", port: int = 52415) -> bool:
|
||||
"""Start hanzo net with specified model."""
|
||||
logger.info(f"Starting hanzo net with model: {model} on port {port}")
|
||||
|
||||
|
||||
# Check if port is already in use
|
||||
if self.is_port_open("localhost", port):
|
||||
logger.warning(f"Port {port} already in use, attempting to use existing instance")
|
||||
logger.warning(
|
||||
f"Port {port} already in use, attempting to use existing instance"
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
try:
|
||||
cmd = [
|
||||
"hanzo", "net",
|
||||
"--name", f"e2e-test-{model}",
|
||||
"--port", str(port),
|
||||
"--models", model,
|
||||
"--network", "local",
|
||||
"--max-jobs", "10"
|
||||
"hanzo",
|
||||
"net",
|
||||
"--name",
|
||||
f"e2e-test-{model}",
|
||||
"--port",
|
||||
str(port),
|
||||
"--models",
|
||||
model,
|
||||
"--network",
|
||||
"local",
|
||||
"--max-jobs",
|
||||
"10",
|
||||
]
|
||||
|
||||
|
||||
logger.info(f"Command: {' '.join(cmd)}")
|
||||
|
||||
|
||||
self.hanzo_net_process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True
|
||||
cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
|
||||
)
|
||||
|
||||
|
||||
# Wait for hanzo net to start
|
||||
for i in range(30): # 30 seconds timeout
|
||||
await asyncio.sleep(1)
|
||||
if self.is_port_open("localhost", port):
|
||||
logger.info(f"✓ hanzo net started successfully on port {port}")
|
||||
return True
|
||||
|
||||
|
||||
# Check if process has died
|
||||
if self.hanzo_net_process.poll() is not None:
|
||||
stdout, stderr = self.hanzo_net_process.communicate()
|
||||
logger.error(f"hanzo net failed to start:\nSTDOUT: {stdout}\nSTDERR: {stderr}")
|
||||
logger.error(
|
||||
f"hanzo net failed to start:\nSTDOUT: {stdout}\nSTDERR: {stderr}"
|
||||
)
|
||||
return False
|
||||
|
||||
|
||||
logger.error("Timeout waiting for hanzo net to start")
|
||||
return False
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to start hanzo net: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def is_port_open(self, host: str, port: int) -> bool:
|
||||
"""Check if a port is open."""
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
result = sock.connect_ex((host, port))
|
||||
sock.close()
|
||||
return result == 0
|
||||
|
||||
|
||||
async def create_local_orchestrator(self, model: str, port: int = 52415):
|
||||
"""Create a local AI orchestrator."""
|
||||
logger.info(f"Creating local orchestrator with {model}")
|
||||
|
||||
|
||||
orchestrator = create_local_agent(
|
||||
name=f"{model}_orchestrator",
|
||||
description=f"Local {model} orchestrator for E2E test",
|
||||
@@ -176,21 +184,21 @@ class HanzoNetE2ETest:
|
||||
Always prefer local models when possible to reduce costs.""",
|
||||
local_model=model,
|
||||
base_url=f"http://localhost:{port}",
|
||||
tools=[]
|
||||
tools=[],
|
||||
)
|
||||
|
||||
|
||||
return orchestrator
|
||||
|
||||
|
||||
async def create_agent_network(self, orchestrator_model: str = "qwen3"):
|
||||
"""Create a full agent network with local orchestrator and mixed workers."""
|
||||
logger.info("Creating agent network...")
|
||||
|
||||
|
||||
agents = []
|
||||
|
||||
|
||||
# 1. Create local orchestrator
|
||||
orchestrator = await self.create_local_orchestrator(orchestrator_model)
|
||||
agents.append(orchestrator)
|
||||
|
||||
|
||||
# 2. Create local worker agents
|
||||
for i in range(2):
|
||||
worker = create_local_agent(
|
||||
@@ -198,11 +206,11 @@ class HanzoNetE2ETest:
|
||||
description=f"Local worker {i} for simple tasks",
|
||||
system="You are a local worker handling simple code tasks efficiently.",
|
||||
local_model=orchestrator_model,
|
||||
base_url="http://localhost:52415"
|
||||
base_url="http://localhost:52415",
|
||||
)
|
||||
agents.append(worker)
|
||||
logger.info(f" Created local worker {i}")
|
||||
|
||||
|
||||
# 3. Create API-based worker agents (if keys available)
|
||||
if os.getenv("ANTHROPIC_API_KEY"):
|
||||
claude_worker = create_agent(
|
||||
@@ -211,13 +219,13 @@ class HanzoNetE2ETest:
|
||||
model=ModelConfig(
|
||||
provider=ModelProvider.ANTHROPIC,
|
||||
model="claude-3-5-sonnet-20241022",
|
||||
api_key=os.getenv("ANTHROPIC_API_KEY")
|
||||
api_key=os.getenv("ANTHROPIC_API_KEY"),
|
||||
),
|
||||
system="You are Claude, specialized in complex code implementation."
|
||||
system="You are Claude, specialized in complex code implementation.",
|
||||
)
|
||||
agents.append(claude_worker)
|
||||
logger.info(" Created Claude worker")
|
||||
|
||||
|
||||
if os.getenv("OPENAI_API_KEY"):
|
||||
gpt_worker = create_agent(
|
||||
name="gpt4_worker",
|
||||
@@ -225,13 +233,13 @@ class HanzoNetE2ETest:
|
||||
model=ModelConfig(
|
||||
provider=ModelProvider.OPENAI,
|
||||
model="gpt-4-turbo-preview",
|
||||
api_key=os.getenv("OPENAI_API_KEY")
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
),
|
||||
system="You are GPT-4, specialized in code analysis and system design."
|
||||
system="You are GPT-4, specialized in code analysis and system design.",
|
||||
)
|
||||
agents.append(gpt_worker)
|
||||
logger.info(" Created GPT-4 worker")
|
||||
|
||||
|
||||
# Also create a Codex-style worker
|
||||
codex_worker = create_agent(
|
||||
name="codex_worker",
|
||||
@@ -240,13 +248,13 @@ class HanzoNetE2ETest:
|
||||
provider=ModelProvider.OPENAI,
|
||||
model="gpt-4-turbo-preview", # Using GPT-4 as Codex replacement
|
||||
api_key=os.getenv("OPENAI_API_KEY"),
|
||||
temperature=0.2 # Lower temperature for code
|
||||
temperature=0.2, # Lower temperature for code
|
||||
),
|
||||
system="You are a code completion specialist. Generate precise, efficient code."
|
||||
system="You are a code completion specialist. Generate precise, efficient code.",
|
||||
)
|
||||
agents.append(codex_worker)
|
||||
logger.info(" Created Codex-style worker")
|
||||
|
||||
|
||||
if os.getenv("GOOGLE_API_KEY"):
|
||||
gemini_worker = create_agent(
|
||||
name="gemini_worker",
|
||||
@@ -254,13 +262,13 @@ class HanzoNetE2ETest:
|
||||
model=ModelConfig(
|
||||
provider=ModelProvider.GOOGLE,
|
||||
model="gemini-pro",
|
||||
api_key=os.getenv("GOOGLE_API_KEY")
|
||||
api_key=os.getenv("GOOGLE_API_KEY"),
|
||||
),
|
||||
system="You are Gemini, capable of handling diverse tasks including code and analysis."
|
||||
system="You are Gemini, capable of handling diverse tasks including code and analysis.",
|
||||
)
|
||||
agents.append(gemini_worker)
|
||||
logger.info(" Created Gemini worker")
|
||||
|
||||
|
||||
# 4. Create critic agents for quality assurance
|
||||
critic = create_local_agent(
|
||||
name="local_critic",
|
||||
@@ -275,11 +283,11 @@ class HanzoNetE2ETest:
|
||||
|
||||
Provide specific, actionable feedback.""",
|
||||
local_model=orchestrator_model,
|
||||
base_url="http://localhost:52415"
|
||||
base_url="http://localhost:52415",
|
||||
)
|
||||
agents.append(critic)
|
||||
logger.info(" Created local critic")
|
||||
|
||||
|
||||
# 5. Create intelligent router
|
||||
router = create_routing_agent(
|
||||
agent=orchestrator,
|
||||
@@ -291,315 +299,320 @@ class HanzoNetE2ETest:
|
||||
Multimodal or diverse tasks → gemini_worker
|
||||
Code review → local_critic
|
||||
|
||||
Optimize for cost by preferring local models when possible."""
|
||||
Optimize for cost by preferring local models when possible.""",
|
||||
)
|
||||
|
||||
|
||||
# 6. Create the network
|
||||
network = create_network(
|
||||
agents=agents,
|
||||
router=router,
|
||||
default_agent=orchestrator.name
|
||||
agents=agents, router=router, default_agent=orchestrator.name
|
||||
)
|
||||
|
||||
|
||||
logger.info(f"✓ Created agent network with {len(agents)} agents")
|
||||
return network
|
||||
|
||||
|
||||
async def test_task_routing(self, network):
|
||||
"""Test that tasks are routed correctly between local and API models."""
|
||||
logger.info("\n=== Testing Task Routing ===")
|
||||
|
||||
|
||||
test_cases = [
|
||||
{
|
||||
"task": "List all Python files in the current directory",
|
||||
"expected_agent": "local_worker",
|
||||
"reason": "Simple file listing task"
|
||||
"reason": "Simple file listing task",
|
||||
},
|
||||
{
|
||||
"task": "Format this JSON: {'key':'value','nested':{'a':1}}",
|
||||
"expected_agent": "local_worker",
|
||||
"reason": "Simple formatting task"
|
||||
"reason": "Simple formatting task",
|
||||
},
|
||||
{
|
||||
"task": "Implement a binary search tree with insert, delete, and search operations",
|
||||
"expected_agent": "claude_worker|gpt4_worker",
|
||||
"reason": "Complex implementation task"
|
||||
"reason": "Complex implementation task",
|
||||
},
|
||||
{
|
||||
"task": "Complete this function: def fibonacci(n): # Calculate nth Fibonacci number",
|
||||
"expected_agent": "codex_worker|gpt4_worker",
|
||||
"reason": "Code completion task"
|
||||
"reason": "Code completion task",
|
||||
},
|
||||
{
|
||||
"task": "Review this code for security issues: eval(user_input)",
|
||||
"expected_agent": "local_critic",
|
||||
"reason": "Code review task"
|
||||
}
|
||||
"reason": "Code review task",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
results = []
|
||||
for test in test_cases:
|
||||
logger.info(f"\nTest: {test['task'][:50]}...")
|
||||
|
||||
|
||||
try:
|
||||
# Create state and run network
|
||||
state = NetworkState()
|
||||
result = await network.run(
|
||||
prompt=test["task"],
|
||||
state=state
|
||||
)
|
||||
|
||||
result = await network.run(prompt=test["task"], state=state)
|
||||
|
||||
# Check which agent handled it
|
||||
last_message = state.messages[-1] if state.messages else None
|
||||
agent_used = last_message.get("agent", "unknown") if last_message else "unknown"
|
||||
|
||||
agent_used = (
|
||||
last_message.get("agent", "unknown") if last_message else "unknown"
|
||||
)
|
||||
|
||||
# Validate routing
|
||||
expected_agents = test["expected_agent"].split("|")
|
||||
success = any(expected in agent_used for expected in expected_agents)
|
||||
|
||||
results.append({
|
||||
"task": test["task"][:50],
|
||||
"expected": test["expected_agent"],
|
||||
"actual": agent_used,
|
||||
"success": success,
|
||||
"reason": test["reason"]
|
||||
})
|
||||
|
||||
|
||||
results.append(
|
||||
{
|
||||
"task": test["task"][:50],
|
||||
"expected": test["expected_agent"],
|
||||
"actual": agent_used,
|
||||
"success": success,
|
||||
"reason": test["reason"],
|
||||
}
|
||||
)
|
||||
|
||||
logger.info(f" Agent used: {agent_used}")
|
||||
logger.info(f" Success: {'✓' if success else '✗'}")
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f" Error: {e}")
|
||||
results.append({
|
||||
"task": test["task"][:50],
|
||||
"error": str(e),
|
||||
"success": False
|
||||
})
|
||||
|
||||
results.append(
|
||||
{"task": test["task"][:50], "error": str(e), "success": False}
|
||||
)
|
||||
|
||||
# Print summary
|
||||
logger.info("\n=== Routing Test Summary ===")
|
||||
successful = sum(1 for r in results if r.get("success"))
|
||||
logger.info(f"Passed: {successful}/{len(results)}")
|
||||
|
||||
|
||||
for result in results:
|
||||
status = "✓" if result.get("success") else "✗"
|
||||
logger.info(f" {status} {result.get('task', 'Unknown')}")
|
||||
if not result.get("success"):
|
||||
logger.info(f" Expected: {result.get('expected', 'N/A')}")
|
||||
logger.info(f" Actual: {result.get('actual', result.get('error', 'N/A'))}")
|
||||
|
||||
logger.info(
|
||||
f" Actual: {result.get('actual', result.get('error', 'N/A'))}"
|
||||
)
|
||||
|
||||
return results
|
||||
|
||||
|
||||
async def test_cost_optimization(self, network):
|
||||
"""Test that the system optimizes for cost by using local models when possible."""
|
||||
logger.info("\n=== Testing Cost Optimization ===")
|
||||
|
||||
|
||||
# Track which models handle which tasks
|
||||
task_distribution = {
|
||||
"local": 0,
|
||||
"api": 0
|
||||
}
|
||||
|
||||
task_distribution = {"local": 0, "api": 0}
|
||||
|
||||
# Run a series of mixed tasks
|
||||
tasks = [
|
||||
"Check if file exists: config.json",
|
||||
"Validate this email: [email protected]",
|
||||
"Validate this email: [email protected]",
|
||||
"Count lines in a file",
|
||||
"Convert string to uppercase",
|
||||
"Parse a simple CSV line",
|
||||
# These should go to API models
|
||||
"Design a microservices architecture for an e-commerce platform",
|
||||
"Debug this complex concurrency issue in a distributed system",
|
||||
"Optimize this machine learning pipeline for better performance"
|
||||
"Optimize this machine learning pipeline for better performance",
|
||||
]
|
||||
|
||||
|
||||
for task in tasks:
|
||||
try:
|
||||
state = NetworkState()
|
||||
result = await network.run(prompt=task, state=state)
|
||||
|
||||
|
||||
# Check which type of agent handled it
|
||||
last_message = state.messages[-1] if state.messages else None
|
||||
agent_used = last_message.get("agent", "") if last_message else ""
|
||||
|
||||
|
||||
if "local" in agent_used.lower():
|
||||
task_distribution["local"] += 1
|
||||
else:
|
||||
task_distribution["api"] += 1
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing task: {e}")
|
||||
|
||||
|
||||
# Calculate cost savings
|
||||
local_ratio = task_distribution["local"] / len(tasks) if tasks else 0
|
||||
cost_savings = local_ratio * 100
|
||||
|
||||
|
||||
logger.info(f"\nTask Distribution:")
|
||||
logger.info(f" Local models: {task_distribution['local']}/{len(tasks)}")
|
||||
logger.info(f" API models: {task_distribution['api']}/{len(tasks)}")
|
||||
logger.info(f" Cost savings: ~{cost_savings:.1f}%")
|
||||
|
||||
|
||||
return {
|
||||
"distribution": task_distribution,
|
||||
"cost_savings_percent": cost_savings,
|
||||
"success": cost_savings > 50 # Expect >50% to use local models
|
||||
"success": cost_savings > 50, # Expect >50% to use local models
|
||||
}
|
||||
|
||||
|
||||
async def test_hanzo_dev_integration(self, model: str = "qwen3"):
|
||||
"""Test full hanzo dev integration with local orchestrator."""
|
||||
logger.info("\n=== Testing Hanzo Dev Integration ===")
|
||||
|
||||
|
||||
try:
|
||||
# Start hanzo dev with local orchestrator
|
||||
cmd = [
|
||||
"hanzo", "dev",
|
||||
"--orchestrator", f"local:{model}",
|
||||
"--instances", "3",
|
||||
"hanzo",
|
||||
"dev",
|
||||
"--orchestrator",
|
||||
f"local:{model}",
|
||||
"--instances",
|
||||
"3",
|
||||
"--use-hanzo-net",
|
||||
"--workspace", str(self.workspace),
|
||||
"--workspace",
|
||||
str(self.workspace),
|
||||
"--no-repl", # Non-interactive mode
|
||||
"--no-monitor"
|
||||
"--no-monitor",
|
||||
]
|
||||
|
||||
|
||||
logger.info(f"Starting hanzo dev: {' '.join(cmd)}")
|
||||
|
||||
|
||||
# Run hanzo dev for a short time to test initialization
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True
|
||||
cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
|
||||
)
|
||||
|
||||
|
||||
# Give it time to initialize
|
||||
await asyncio.sleep(10)
|
||||
|
||||
|
||||
# Check if it's still running
|
||||
if process.poll() is None:
|
||||
logger.info("✓ hanzo dev started successfully")
|
||||
|
||||
|
||||
# Terminate gracefully
|
||||
process.terminate()
|
||||
await asyncio.sleep(2)
|
||||
|
||||
|
||||
if process.poll() is None:
|
||||
process.kill()
|
||||
|
||||
|
||||
return True
|
||||
else:
|
||||
stdout, stderr = process.communicate()
|
||||
logger.error(f"hanzo dev failed:\nSTDOUT: {stdout}\nSTDERR: {stderr}")
|
||||
return False
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to test hanzo dev: {e}")
|
||||
return False
|
||||
|
||||
|
||||
async def run_full_test(self):
|
||||
"""Run the complete E2E test suite."""
|
||||
logger.info("\n" + "="*60)
|
||||
logger.info("\n" + "=" * 60)
|
||||
logger.info("HANZO NET E2E TEST - LOCAL AI ORCHESTRATION")
|
||||
logger.info("="*60)
|
||||
|
||||
logger.info("=" * 60)
|
||||
|
||||
all_results = {}
|
||||
|
||||
|
||||
try:
|
||||
# Setup
|
||||
await self.setup()
|
||||
|
||||
|
||||
# Test with different local models
|
||||
for model in self.models_to_test:
|
||||
logger.info(f"\n{'='*60}")
|
||||
logger.info(f"Testing with model: {model}")
|
||||
logger.info(f"{'='*60}")
|
||||
|
||||
|
||||
model_results = {}
|
||||
|
||||
|
||||
# Start hanzo net with the model
|
||||
if await self.start_hanzo_net(model):
|
||||
# Create agent network
|
||||
network = await self.create_agent_network(model)
|
||||
|
||||
|
||||
# Run tests
|
||||
model_results["routing"] = await self.test_task_routing(network)
|
||||
model_results["cost_optimization"] = await self.test_cost_optimization(network)
|
||||
model_results["hanzo_dev"] = await self.test_hanzo_dev_integration(model)
|
||||
|
||||
model_results["cost_optimization"] = (
|
||||
await self.test_cost_optimization(network)
|
||||
)
|
||||
model_results["hanzo_dev"] = await self.test_hanzo_dev_integration(
|
||||
model
|
||||
)
|
||||
|
||||
# Stop hanzo net
|
||||
if self.hanzo_net_process:
|
||||
self.hanzo_net_process.terminate()
|
||||
await asyncio.sleep(2)
|
||||
self.hanzo_net_process = None
|
||||
else:
|
||||
logger.warning(f"Skipping tests for {model} - failed to start hanzo net")
|
||||
logger.warning(
|
||||
f"Skipping tests for {model} - failed to start hanzo net"
|
||||
)
|
||||
model_results["error"] = "Failed to start hanzo net"
|
||||
|
||||
|
||||
all_results[model] = model_results
|
||||
|
||||
|
||||
# Brief pause between models
|
||||
await asyncio.sleep(5)
|
||||
|
||||
|
||||
# Print final summary
|
||||
self.print_summary(all_results)
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Test failed: {e}")
|
||||
raise
|
||||
finally:
|
||||
# Cleanup
|
||||
await self.cleanup()
|
||||
|
||||
|
||||
def print_summary(self, results: Dict):
|
||||
"""Print test summary."""
|
||||
logger.info("\n" + "="*60)
|
||||
logger.info("\n" + "=" * 60)
|
||||
logger.info("TEST SUMMARY")
|
||||
logger.info("="*60)
|
||||
|
||||
logger.info("=" * 60)
|
||||
|
||||
for model, model_results in results.items():
|
||||
logger.info(f"\n{model}:")
|
||||
|
||||
|
||||
if "error" in model_results:
|
||||
logger.info(f" ✗ {model_results['error']}")
|
||||
continue
|
||||
|
||||
|
||||
# Routing tests
|
||||
if "routing" in model_results:
|
||||
routing = model_results["routing"]
|
||||
passed = sum(1 for r in routing if r.get("success", False))
|
||||
logger.info(f" Routing: {passed}/{len(routing)} passed")
|
||||
|
||||
|
||||
# Cost optimization
|
||||
if "cost_optimization" in model_results:
|
||||
cost = model_results["cost_optimization"]
|
||||
if cost.get("success"):
|
||||
logger.info(f" Cost Optimization: ✓ ({cost['cost_savings_percent']:.1f}% savings)")
|
||||
logger.info(
|
||||
f" Cost Optimization: ✓ ({cost['cost_savings_percent']:.1f}% savings)"
|
||||
)
|
||||
else:
|
||||
logger.info(f" Cost Optimization: ✗")
|
||||
|
||||
|
||||
# Hanzo dev integration
|
||||
if "hanzo_dev" in model_results:
|
||||
if model_results["hanzo_dev"]:
|
||||
logger.info(f" Hanzo Dev Integration: ✓")
|
||||
else:
|
||||
logger.info(f" Hanzo Dev Integration: ✗")
|
||||
|
||||
|
||||
async def cleanup(self):
|
||||
"""Clean up test resources."""
|
||||
logger.info("\nCleaning up...")
|
||||
|
||||
|
||||
# Stop hanzo net
|
||||
if self.hanzo_net_process:
|
||||
self.hanzo_net_process.terminate()
|
||||
await asyncio.sleep(2)
|
||||
if self.hanzo_net_process.poll() is None:
|
||||
self.hanzo_net_process.kill()
|
||||
|
||||
|
||||
# Stop hanzo dev
|
||||
if self.hanzo_dev_process:
|
||||
self.hanzo_dev_process.terminate()
|
||||
await asyncio.sleep(2)
|
||||
if self.hanzo_dev_process.poll() is None:
|
||||
self.hanzo_dev_process.kill()
|
||||
|
||||
|
||||
logger.info("Cleanup complete")
|
||||
|
||||
|
||||
@@ -610,4 +623,4 @@ async def main():
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
asyncio.run(main())
|
||||
|
||||
@@ -18,104 +18,109 @@ from pathlib import Path
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
||||
level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class SimpleHanzoNetTest:
|
||||
"""Simplified E2E test for hanzo net orchestration."""
|
||||
|
||||
|
||||
def __init__(self):
|
||||
self.hanzo_net_process = None
|
||||
self.hanzo_dev_process = None
|
||||
self.workspace = Path.home() / ".hanzo" / "e2e-test"
|
||||
self.workspace.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
|
||||
def is_port_open(self, host: str, port: int) -> bool:
|
||||
"""Check if a port is open."""
|
||||
sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
||||
result = sock.connect_ex((host, port))
|
||||
sock.close()
|
||||
return result == 0
|
||||
|
||||
|
||||
async def test_hanzo_net_startup(self, model: str = "llama-3.2-3b") -> bool:
|
||||
"""Test that hanzo net can start with a local model."""
|
||||
logger.info(f"\n=== Testing Hanzo Net Startup with {model} ===")
|
||||
|
||||
|
||||
port = 52415
|
||||
|
||||
|
||||
# Check if already running
|
||||
if self.is_port_open("localhost", port):
|
||||
logger.info(f"✓ Port {port} already in use (hanzo net may be running)")
|
||||
return True
|
||||
|
||||
|
||||
try:
|
||||
# Start hanzo net
|
||||
cmd = [
|
||||
"hanzo", "net",
|
||||
"--name", "e2e-test",
|
||||
"--port", str(port),
|
||||
"--models", model,
|
||||
"--network", "local"
|
||||
"hanzo",
|
||||
"net",
|
||||
"--name",
|
||||
"e2e-test",
|
||||
"--port",
|
||||
str(port),
|
||||
"--models",
|
||||
model,
|
||||
"--network",
|
||||
"local",
|
||||
]
|
||||
|
||||
|
||||
logger.info(f"Starting: {' '.join(cmd)}")
|
||||
|
||||
|
||||
self.hanzo_net_process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True
|
||||
cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True
|
||||
)
|
||||
|
||||
|
||||
# Wait for startup
|
||||
for i in range(30):
|
||||
await asyncio.sleep(1)
|
||||
if self.is_port_open("localhost", port):
|
||||
logger.info(f"✓ hanzo net started on port {port}")
|
||||
return True
|
||||
|
||||
|
||||
# Check if process died
|
||||
if self.hanzo_net_process.poll() is not None:
|
||||
stdout, stderr = self.hanzo_net_process.communicate()
|
||||
logger.error(f"hanzo net failed to start")
|
||||
logger.error(f"STDERR: {stderr}")
|
||||
return False
|
||||
|
||||
|
||||
logger.error("Timeout waiting for hanzo net")
|
||||
return False
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error starting hanzo net: {e}")
|
||||
return False
|
||||
|
||||
|
||||
async def test_hanzo_dev_with_local(self, model: str = "llama-3.2-3b") -> bool:
|
||||
"""Test hanzo dev with local orchestrator."""
|
||||
logger.info(f"\n=== Testing Hanzo Dev with Local Orchestrator ===")
|
||||
|
||||
|
||||
try:
|
||||
# Start hanzo dev with local orchestrator
|
||||
cmd = [
|
||||
"hanzo", "dev",
|
||||
"--orchestrator", f"local:{model}",
|
||||
"--instances", "2",
|
||||
"hanzo",
|
||||
"dev",
|
||||
"--orchestrator",
|
||||
f"local:{model}",
|
||||
"--instances",
|
||||
"2",
|
||||
"--use-hanzo-net",
|
||||
"--workspace", str(self.workspace)
|
||||
"--workspace",
|
||||
str(self.workspace),
|
||||
]
|
||||
|
||||
|
||||
logger.info(f"Starting: {' '.join(cmd)}")
|
||||
|
||||
|
||||
# Start process and capture output
|
||||
process = subprocess.Popen(
|
||||
cmd,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.STDOUT,
|
||||
text=True,
|
||||
bufsize=1
|
||||
bufsize=1,
|
||||
)
|
||||
|
||||
|
||||
# Read output for a few seconds
|
||||
start_time = time.time()
|
||||
output_lines = []
|
||||
@@ -123,63 +128,73 @@ class SimpleHanzoNetTest:
|
||||
"Agent network initialized",
|
||||
"Created local",
|
||||
"orchestrator",
|
||||
"hanzo/net started"
|
||||
"hanzo/net started",
|
||||
]
|
||||
|
||||
|
||||
while time.time() - start_time < 15: # 15 second timeout
|
||||
line = process.stdout.readline()
|
||||
if line:
|
||||
output_lines.append(line.strip())
|
||||
logger.info(f" {line.strip()}")
|
||||
|
||||
|
||||
# Check for success indicators
|
||||
for indicator in success_indicators:
|
||||
if indicator.lower() in line.lower():
|
||||
logger.info(f"✓ Found: {indicator}")
|
||||
|
||||
|
||||
# Check if process died
|
||||
if process.poll() is not None:
|
||||
break
|
||||
|
||||
|
||||
await asyncio.sleep(0.1)
|
||||
|
||||
|
||||
# Terminate the process
|
||||
process.terminate()
|
||||
await asyncio.sleep(2)
|
||||
if process.poll() is None:
|
||||
process.kill()
|
||||
|
||||
|
||||
# Check if we saw success indicators
|
||||
full_output = "\n".join(output_lines)
|
||||
if any(indicator in full_output.lower() for indicator in ["network initialized", "orchestrator", "started"]):
|
||||
if any(
|
||||
indicator in full_output.lower()
|
||||
for indicator in ["network initialized", "orchestrator", "started"]
|
||||
):
|
||||
logger.info("✓ hanzo dev initialized successfully")
|
||||
return True
|
||||
else:
|
||||
logger.warning("Could not confirm successful initialization")
|
||||
return False
|
||||
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"Error testing hanzo dev: {e}")
|
||||
return False
|
||||
|
||||
|
||||
async def test_cost_optimization_config(self) -> bool:
|
||||
"""Test that cost optimization is configured correctly."""
|
||||
logger.info("\n=== Testing Cost Optimization Configuration ===")
|
||||
|
||||
|
||||
# Check if the dev.py file has cost optimization logic
|
||||
dev_file = Path(__file__).parent.parent.parent.parent / "pkg" / "hanzo" / "src" / "hanzo" / "dev.py"
|
||||
|
||||
dev_file = (
|
||||
Path(__file__).parent.parent.parent.parent
|
||||
/ "pkg"
|
||||
/ "hanzo"
|
||||
/ "src"
|
||||
/ "hanzo"
|
||||
/ "dev.py"
|
||||
)
|
||||
|
||||
if dev_file.exists():
|
||||
content = dev_file.read_text()
|
||||
|
||||
|
||||
checks = {
|
||||
"CostOptimizedRouter": "Cost-optimized router class",
|
||||
"local_workers": "Local worker configuration",
|
||||
"simple_keywords": "Simple task detection",
|
||||
"complex_keywords": "Complex task detection",
|
||||
"local models preferred": "Local model preference"
|
||||
"local models preferred": "Local model preference",
|
||||
}
|
||||
|
||||
|
||||
results = []
|
||||
for key, description in checks.items():
|
||||
if key.lower() in content.lower():
|
||||
@@ -188,54 +203,54 @@ class SimpleHanzoNetTest:
|
||||
else:
|
||||
logger.warning(f" ✗ {description} not found")
|
||||
results.append(False)
|
||||
|
||||
|
||||
success = all(results)
|
||||
if success:
|
||||
logger.info("✓ Cost optimization properly configured")
|
||||
else:
|
||||
logger.warning("⚠ Some cost optimization features missing")
|
||||
|
||||
|
||||
return success
|
||||
else:
|
||||
logger.error(f"dev.py not found at {dev_file}")
|
||||
return False
|
||||
|
||||
|
||||
async def run_all_tests(self):
|
||||
"""Run all tests."""
|
||||
logger.info("\n" + "="*60)
|
||||
logger.info("\n" + "=" * 60)
|
||||
logger.info("HANZO NET SIMPLE E2E TEST")
|
||||
logger.info("="*60)
|
||||
|
||||
logger.info("=" * 60)
|
||||
|
||||
results = {}
|
||||
|
||||
|
||||
try:
|
||||
# Test 1: Hanzo net startup
|
||||
results["hanzo_net"] = await self.test_hanzo_net_startup()
|
||||
|
||||
|
||||
# Test 2: Hanzo dev with local orchestrator
|
||||
results["hanzo_dev"] = await self.test_hanzo_dev_with_local()
|
||||
|
||||
|
||||
# Test 3: Cost optimization configuration
|
||||
results["cost_optimization"] = await self.test_cost_optimization_config()
|
||||
|
||||
|
||||
# Summary
|
||||
logger.info("\n" + "="*60)
|
||||
logger.info("\n" + "=" * 60)
|
||||
logger.info("TEST SUMMARY")
|
||||
logger.info("="*60)
|
||||
|
||||
logger.info("=" * 60)
|
||||
|
||||
for test_name, success in results.items():
|
||||
status = "✓ PASSED" if success else "✗ FAILED"
|
||||
logger.info(f"{test_name}: {status}")
|
||||
|
||||
|
||||
all_passed = all(results.values())
|
||||
|
||||
|
||||
if all_passed:
|
||||
logger.info("\n✓ All tests passed!")
|
||||
else:
|
||||
logger.info("\n✗ Some tests failed")
|
||||
|
||||
|
||||
return all_passed
|
||||
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
if self.hanzo_net_process:
|
||||
@@ -243,7 +258,7 @@ class SimpleHanzoNetTest:
|
||||
await asyncio.sleep(1)
|
||||
if self.hanzo_net_process.poll() is None:
|
||||
self.hanzo_net_process.kill()
|
||||
|
||||
|
||||
if self.hanzo_dev_process:
|
||||
self.hanzo_dev_process.terminate()
|
||||
await asyncio.sleep(1)
|
||||
@@ -259,4 +274,4 @@ async def main():
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
asyncio.run(main())
|
||||
|
||||
+86
-55
@@ -13,7 +13,7 @@ def mock_api(request):
|
||||
if "respx_mock" in request.fixturenames:
|
||||
yield None
|
||||
return
|
||||
|
||||
|
||||
with respx.mock(
|
||||
base_url="http://127.0.0.1:4010", assert_all_called=False
|
||||
) as respx_mock:
|
||||
@@ -139,9 +139,11 @@ def mock_api(request):
|
||||
respx_mock.get("/global/spend/tags").mock(return_value=Response(200, json=[]))
|
||||
respx_mock.post("/global/spend/reset").mock(return_value=Response(200, json={}))
|
||||
respx_mock.get("/global/spend/report").mock(return_value=Response(200, json=[]))
|
||||
|
||||
|
||||
# Guardrails
|
||||
respx_mock.get("/guardrails/list").mock(return_value=Response(200, json={"guardrails": []}))
|
||||
respx_mock.get("/guardrails/list").mock(
|
||||
return_value=Response(200, json={"guardrails": []})
|
||||
)
|
||||
|
||||
# Images
|
||||
respx_mock.post("/images/generations").mock(
|
||||
@@ -199,14 +201,12 @@ def mock_api(request):
|
||||
respx_mock.get("/organization/info/deprecated").mock(
|
||||
return_value=Response(200, json={})
|
||||
)
|
||||
respx_mock.post("/organization/new").mock(return_value=Response(200, json=org_response))
|
||||
respx_mock.post("/organization/new").mock(
|
||||
return_value=Response(200, json=org_response)
|
||||
)
|
||||
respx_mock.post("/organization").mock(return_value=Response(200, json={}))
|
||||
respx_mock.get("/organization/list").mock(
|
||||
return_value=Response(200, json=[])
|
||||
)
|
||||
respx_mock.get("/organization").mock(
|
||||
return_value=Response(200, json=[])
|
||||
)
|
||||
respx_mock.get("/organization/list").mock(return_value=Response(200, json=[]))
|
||||
respx_mock.get("/organization").mock(return_value=Response(200, json=[]))
|
||||
org_update_response = {
|
||||
"budget_id": "budget-123",
|
||||
"created_at": "2024-01-01T00:00:00Z",
|
||||
@@ -216,17 +216,23 @@ def mock_api(request):
|
||||
"updated_by": "user-123",
|
||||
"organization_id": "org-123",
|
||||
"members": [],
|
||||
"teams": []
|
||||
"teams": [],
|
||||
}
|
||||
respx_mock.patch("/organization/update").mock(return_value=Response(200, json=org_update_response))
|
||||
respx_mock.put("/organization/update").mock(return_value=Response(200, json=org_update_response))
|
||||
respx_mock.patch("/organization/update").mock(
|
||||
return_value=Response(200, json=org_update_response)
|
||||
)
|
||||
respx_mock.put("/organization/update").mock(
|
||||
return_value=Response(200, json=org_update_response)
|
||||
)
|
||||
respx_mock.put("/organization").mock(return_value=Response(200, json={}))
|
||||
respx_mock.delete("/organization/delete").mock(return_value=Response(200, json=[org_response]))
|
||||
respx_mock.delete("/organization/delete").mock(
|
||||
return_value=Response(200, json=[org_response])
|
||||
)
|
||||
respx_mock.delete("/organization").mock(return_value=Response(200, json={}))
|
||||
org_member_response = {
|
||||
"organization_id": "org-123",
|
||||
"updated_organization_memberships": [],
|
||||
"updated_users": []
|
||||
"updated_users": [],
|
||||
}
|
||||
respx_mock.post("/organization/member_add").mock(
|
||||
return_value=Response(200, json=org_member_response)
|
||||
@@ -244,7 +250,7 @@ def mock_api(request):
|
||||
"created_at": "2024-01-01T00:00:00Z",
|
||||
"organization_id": "org-123",
|
||||
"updated_at": "2024-01-01T00:00:00Z",
|
||||
"user_id": "user-123"
|
||||
"user_id": "user-123",
|
||||
}
|
||||
respx_mock.patch("/organization/member_update").mock(
|
||||
return_value=Response(200, json=org_update_member_response)
|
||||
@@ -266,7 +272,7 @@ def mock_api(request):
|
||||
"team_id": "team-123",
|
||||
"models": [],
|
||||
"members": [],
|
||||
"admins": []
|
||||
"admins": [],
|
||||
}
|
||||
respx_mock.route(method="GET", path__regex=r"/team/.*/callback").mock(
|
||||
return_value=Response(200, json={})
|
||||
@@ -276,35 +282,56 @@ def mock_api(request):
|
||||
)
|
||||
respx_mock.post("/team/model").mock(return_value=Response(200, json={}))
|
||||
respx_mock.delete("/team/model").mock(return_value=Response(200, json={}))
|
||||
respx_mock.post("/team/new").mock(return_value=Response(200, json=team_response))
|
||||
respx_mock.post("/team/new").mock(
|
||||
return_value=Response(200, json=team_response)
|
||||
)
|
||||
respx_mock.post("/team").mock(return_value=Response(200, json={}))
|
||||
respx_mock.get("/team/list").mock(return_value=Response(200, json=[]))
|
||||
respx_mock.get("/team").mock(return_value=Response(200, json=[]))
|
||||
respx_mock.get("/team/info").mock(return_value=Response(200, json=team_response))
|
||||
respx_mock.patch("/team/update").mock(return_value=Response(200, json=team_response))
|
||||
respx_mock.put("/team/update").mock(return_value=Response(200, json=team_response))
|
||||
respx_mock.get("/team/info").mock(
|
||||
return_value=Response(200, json=team_response)
|
||||
)
|
||||
respx_mock.patch("/team/update").mock(
|
||||
return_value=Response(200, json=team_response)
|
||||
)
|
||||
respx_mock.put("/team/update").mock(
|
||||
return_value=Response(200, json=team_response)
|
||||
)
|
||||
respx_mock.put("/team").mock(return_value=Response(200, json={}))
|
||||
respx_mock.delete("/team/delete").mock(return_value=Response(200, json=[team_response]))
|
||||
respx_mock.delete("/team/delete").mock(
|
||||
return_value=Response(200, json=[team_response])
|
||||
)
|
||||
respx_mock.delete("/team").mock(return_value=Response(200, json={}))
|
||||
team_add_member_response = {
|
||||
"team_id": "team-123",
|
||||
"updated_team_memberships": [],
|
||||
"updated_users": []
|
||||
"updated_users": [],
|
||||
}
|
||||
team_update_member_response = {
|
||||
"team_id": "team-123",
|
||||
"user_id": "user-123"
|
||||
}
|
||||
respx_mock.post("/team/member_add").mock(return_value=Response(200, json=team_add_member_response))
|
||||
team_update_member_response = {"team_id": "team-123", "user_id": "user-123"}
|
||||
respx_mock.post("/team/member_add").mock(
|
||||
return_value=Response(200, json=team_add_member_response)
|
||||
)
|
||||
respx_mock.post("/team/members").mock(return_value=Response(200, json={}))
|
||||
respx_mock.delete("/team/member_delete").mock(return_value=Response(200, json=team_response))
|
||||
respx_mock.delete("/team/member_delete").mock(
|
||||
return_value=Response(200, json=team_response)
|
||||
)
|
||||
respx_mock.delete("/team/members").mock(return_value=Response(200, json={}))
|
||||
respx_mock.post("/team/member_update").mock(return_value=Response(200, json=team_update_member_response))
|
||||
respx_mock.patch("/team/member_update").mock(return_value=Response(200, json=team_update_member_response))
|
||||
respx_mock.put("/team/member_update").mock(return_value=Response(200, json=team_update_member_response))
|
||||
respx_mock.post("/team/member_update").mock(
|
||||
return_value=Response(200, json=team_update_member_response)
|
||||
)
|
||||
respx_mock.patch("/team/member_update").mock(
|
||||
return_value=Response(200, json=team_update_member_response)
|
||||
)
|
||||
respx_mock.put("/team/member_update").mock(
|
||||
return_value=Response(200, json=team_update_member_response)
|
||||
)
|
||||
respx_mock.put("/team/members").mock(return_value=Response(200, json={}))
|
||||
respx_mock.post("/team/block").mock(return_value=Response(200, json=team_response))
|
||||
respx_mock.post("/team/unblock").mock(return_value=Response(200, json=team_response))
|
||||
respx_mock.post("/team/block").mock(
|
||||
return_value=Response(200, json=team_response)
|
||||
)
|
||||
respx_mock.post("/team/unblock").mock(
|
||||
return_value=Response(200, json=team_response)
|
||||
)
|
||||
|
||||
# Other endpoints
|
||||
respx_mock.get("/active/callbacks").mock(
|
||||
@@ -449,19 +476,15 @@ def mock_api(request):
|
||||
|
||||
# Spend
|
||||
respx_mock.get("/spend/calculate").mock(return_value=Response(200, json={}))
|
||||
respx_mock.get("/spend/logs").mock(
|
||||
return_value=Response(200, json=[])
|
||||
)
|
||||
respx_mock.get("/spend/tags").mock(
|
||||
return_value=Response(200, json=[])
|
||||
)
|
||||
respx_mock.get("/spend/logs").mock(return_value=Response(200, json=[]))
|
||||
respx_mock.get("/spend/tags").mock(return_value=Response(200, json=[]))
|
||||
|
||||
# Test
|
||||
respx_mock.get("/test").mock(return_value=Response(200, json={}))
|
||||
|
||||
# Root endpoint for client tests
|
||||
respx_mock.get("/").mock(return_value=Response(200, json={}))
|
||||
|
||||
|
||||
# Threads
|
||||
respx_mock.post("/threads").mock(return_value=Response(200, json={}))
|
||||
respx_mock.get("/threads").mock(return_value=Response(200, json={"data": []}))
|
||||
@@ -488,18 +511,23 @@ def mock_api(request):
|
||||
)
|
||||
|
||||
# User
|
||||
user_response = {
|
||||
"key": "sk-user-123",
|
||||
"user_id": "user-123"
|
||||
}
|
||||
respx_mock.post("/user/new").mock(return_value=Response(200, json=user_response))
|
||||
user_response = {"key": "sk-user-123", "user_id": "user-123"}
|
||||
respx_mock.post("/user/new").mock(
|
||||
return_value=Response(200, json=user_response)
|
||||
)
|
||||
respx_mock.post("/user").mock(return_value=Response(200, json={}))
|
||||
respx_mock.get("/user/list").mock(return_value=Response(200, json=[]))
|
||||
respx_mock.get("/user").mock(return_value=Response(200, json=[]))
|
||||
respx_mock.get("/user/info").mock(return_value=Response(200, json=user_response))
|
||||
respx_mock.put("/user/update").mock(return_value=Response(200, json=user_response))
|
||||
respx_mock.get("/user/info").mock(
|
||||
return_value=Response(200, json=user_response)
|
||||
)
|
||||
respx_mock.put("/user/update").mock(
|
||||
return_value=Response(200, json=user_response)
|
||||
)
|
||||
respx_mock.put("/user").mock(return_value=Response(200, json={}))
|
||||
respx_mock.delete("/user/delete").mock(return_value=Response(200, json=user_response))
|
||||
respx_mock.delete("/user/delete").mock(
|
||||
return_value=Response(200, json=user_response)
|
||||
)
|
||||
respx_mock.delete("/user").mock(return_value=Response(200, json={}))
|
||||
|
||||
# Utils
|
||||
@@ -507,12 +535,15 @@ def mock_api(request):
|
||||
return_value=Response(200, json={})
|
||||
)
|
||||
respx_mock.post("/utils/token_counter").mock(
|
||||
return_value=Response(200, json={
|
||||
"model_used": "gpt-3.5-turbo",
|
||||
"request_model": "gpt-3.5-turbo",
|
||||
"tokenizer_type": "cl100k_base",
|
||||
"total_tokens": 10
|
||||
})
|
||||
return_value=Response(
|
||||
200,
|
||||
json={
|
||||
"model_used": "gpt-3.5-turbo",
|
||||
"request_model": "gpt-3.5-turbo",
|
||||
"tokenizer_type": "cl100k_base",
|
||||
"total_tokens": 10,
|
||||
},
|
||||
)
|
||||
)
|
||||
respx_mock.post("/utils/transform_request").mock(
|
||||
return_value=Response(200, json={})
|
||||
|
||||
+22
-20
@@ -12,10 +12,10 @@ class MockHandler(BaseHTTPRequestHandler):
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
|
||||
|
||||
# Default response
|
||||
response = {}
|
||||
|
||||
|
||||
# Handle specific endpoints
|
||||
if "/files/" in self.path and "/content" in self.path:
|
||||
response = "file content"
|
||||
@@ -32,7 +32,7 @@ class MockHandler(BaseHTTPRequestHandler):
|
||||
"created_by": "user-123",
|
||||
"models": ["gpt-3.5-turbo"],
|
||||
"updated_at": "2024-01-01T00:00:00Z",
|
||||
"updated_by": "user-123"
|
||||
"updated_by": "user-123",
|
||||
}
|
||||
elif self.path == "/config/pass_through_endpoint":
|
||||
response = {"endpoints": []}
|
||||
@@ -40,25 +40,25 @@ class MockHandler(BaseHTTPRequestHandler):
|
||||
response = {}
|
||||
else:
|
||||
response = {"data": []}
|
||||
|
||||
|
||||
self.wfile.write(json.dumps(response).encode())
|
||||
|
||||
|
||||
def do_POST(self):
|
||||
"""Handle POST requests."""
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
|
||||
|
||||
# Default response
|
||||
response = {}
|
||||
|
||||
|
||||
# Handle specific endpoints
|
||||
if self.path == "/utils/token_counter":
|
||||
response = {
|
||||
"model_used": "gpt-3.5-turbo",
|
||||
"request_model": "gpt-3.5-turbo",
|
||||
"tokenizer_type": "cl100k_base",
|
||||
"total_tokens": 10
|
||||
"total_tokens": 10,
|
||||
}
|
||||
elif self.path == "/chat/completions":
|
||||
response = {
|
||||
@@ -66,36 +66,38 @@ class MockHandler(BaseHTTPRequestHandler):
|
||||
"object": "chat.completion",
|
||||
"created": 1677652288,
|
||||
"model": "gpt-3.5-turbo",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "Hello!"},
|
||||
"finish_reason": "stop"
|
||||
}]
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "Hello!"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
self.wfile.write(json.dumps(response).encode())
|
||||
|
||||
|
||||
def do_PUT(self):
|
||||
"""Handle PUT requests."""
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
self.wfile.write(b"{}")
|
||||
|
||||
|
||||
def do_DELETE(self):
|
||||
"""Handle DELETE requests."""
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", "application/json")
|
||||
self.end_headers()
|
||||
|
||||
|
||||
# Handle specific endpoints
|
||||
if self.path == "/config/pass_through_endpoint":
|
||||
response = {"endpoints": []}
|
||||
else:
|
||||
response = {}
|
||||
|
||||
|
||||
self.wfile.write(json.dumps(response).encode())
|
||||
|
||||
|
||||
def log_message(self, format, *args):
|
||||
pass # Suppress logs
|
||||
|
||||
@@ -111,4 +113,4 @@ if __name__ == "__main__":
|
||||
sys.exit(0)
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
sys.exit(1)
|
||||
sys.exit(1)
|
||||
|
||||
Reference in New Issue
Block a user