python-sdk/pkg/hanzo-mcp/examples/run_local_agents.py

#!/usr/bin/env python
"""Run agents with local hanzo/net inference (using dummy engine for demo)."""

import asyncio
import sys
from pathlib import Path

# Add hanzo-network to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent / "hanzo-network" / "src"))

from hanzo_network import create_agent, create_distributed_network, create_tool
from hanzo_network.core.agent import ModelConfig, ModelProvider


# Create real tools
async def search_code(query: str) -> str:
    """Search for code patterns."""
    return f"Found 3 matches for '{query}' in codebase:\n1. main.py:42\n2. utils.py:15\n3. test_main.py:8"


async def analyze_function(function_name: str) -> str:
    """Analyze a function's implementation."""
    return f"""Analysis of '{function_name}':
- Parameters: 2 (x: int, y: int)
- Returns: int
- Complexity: O(1)
- Purpose: Adds two numbers"""


async def generate_test(code: str) -> str:
    """Generate unit tests for code."""
    return f"""import pytest

def test_{code.split()[0].lower()}():
    # Test generated by hanzo/net
    assert {code.split()[0]}(2, 3) == 5
    assert {code.split()[0]}(0, 0) == 0
    assert {code.split()[0]}(-1, 1) == 0"""


async def explain_concept(concept: str) -> str:
    """Explain a programming concept."""
    return f"""Explanation of '{concept}':
{concept} is a fundamental programming concept that involves...
Key points:
1. Definition and purpose
2. Common use cases
3. Best practices"""


async def main():
    """Run agents with hanzo/net local inference."""
    print("🚀 Running Local Agents with Hanzo/Net")
    print("=" * 60)

    # Create specialized agents
    print("\n🤖 Creating specialized agents...")

    # Code Search Agent
    search_agent = create_agent(
        name="search_agent",
        description="Searches through codebase",
        model=ModelConfig(
            provider=ModelProvider.LOCAL, model="llama3.2", temperature=0.3
        ),
        system="""You are a code search specialist. Use the search_code tool to find patterns in code.
Be precise and helpful in locating code elements.""",
        tools=[
            create_tool(
                name="search_code",
                description="Search for code patterns",
                handler=search_code,
            )
        ],
    )

    # Code Analyzer Agent
    analyzer_agent = create_agent(
        name="analyzer",
        description="Analyzes code structure and quality",
        model=ModelConfig(
            provider=ModelProvider.LOCAL, model="llama3.2", temperature=0.5
        ),
        system="""You are a code analysis expert. Use your tools to analyze functions and provide insights.
Focus on code quality, performance, and best practices.""",
        tools=[
            create_tool(
                name="analyze_function",
                description="Analyze a function",
                handler=analyze_function,
            )
        ],
    )

    # Test Generator Agent
    test_agent = create_agent(
        name="test_generator",
        description="Generates unit tests",
        model=ModelConfig(
            provider=ModelProvider.LOCAL, model="llama3.2", temperature=0.7
        ),
        system="""You are a test generation specialist. Create comprehensive unit tests for given code.
Follow pytest conventions and ensure good coverage.""",
        tools=[
            create_tool(
                name="generate_test",
                description="Generate unit tests",
                handler=generate_test,
            )
        ],
    )

    # Teacher Agent
    teacher_agent = create_agent(
        name="teacher",
        description="Explains programming concepts",
        model=ModelConfig(
            provider=ModelProvider.LOCAL, model="llama3.2", temperature=0.8
        ),
        system="""You are a programming teacher. Explain concepts clearly and provide examples.
Make complex topics accessible to learners.""",
        tools=[
            create_tool(
                name="explain_concept",
                description="Explain a concept",
                handler=explain_concept,
            )
        ],
    )

    # Create distributed network
    network = create_distributed_network(
        agents=[search_agent, analyzer_agent, test_agent, teacher_agent],
        name="dev-network",
        listen_port=15750,
        broadcast_port=15750,
    )

    print("\n🌐 Starting distributed network...")
    await network.start(wait_for_peers=0)

    # Show network status
    status = network.get_network_status()
    print("\n📊 Network Status:")
    print(f"  Node: {status['node_id']}")
    print(f"  Agents: {len(status['local_agents'])} agents")
    for agent in status["local_agents"]:
        print(f"    - {agent}")
    print("  Inference: hanzo/net (local)")

    # Test 1: Code Search
    print("\n🔍 Test 1: Code Search")
    result = await network.run(
        prompt="Search for all functions that handle authentication",
        initial_agent=search_agent,
    )
    print(f"Result: {result['final_output']}")

    # Test 2: Code Analysis
    print("\n📊 Test 2: Code Analysis")
    result = await network.run(
        prompt="Analyze the add function implementation", initial_agent=analyzer_agent
    )
    print(f"Result: {result['final_output']}")

    # Test 3: Test Generation
    print("\n🧪 Test 3: Test Generation")
    result = await network.run(
        prompt="Generate tests for an add function", initial_agent=test_agent
    )
    print(f"Result: {result['final_output']}")

    # Test 4: Concept Explanation
    print("\n📚 Test 4: Concept Explanation")
    result = await network.run(
        prompt="Explain what recursion is", initial_agent=teacher_agent
    )
    print(f"Result: {result['final_output']}")

    # Test 5: Multi-Agent Collaboration
    print("\n🤝 Test 5: Multi-Agent Collaboration")
    result = await network.run(
        prompt="First search for sorting algorithms, then analyze bubble sort, and finally generate tests for it"
    )
    print(f"Result: {result['final_output']}")
    print(f"Agents involved: {result['iterations']}")

    # Test 6: Complex Query
    print("\n🎯 Test 6: Complex Development Task")
    result = await network.run(
        prompt="I need help understanding and testing a binary search implementation"
    )
    print(f"Result: {result['final_output']}")

    print("\n✅ Local agent demo complete!")
    print("   - All agents running with hanzo/net local inference")
    print("   - No external API calls needed")
    print("   - Distributed across network nodes")

    await network.stop()


if __name__ == "__main__":
    print("\nThis demo shows local agents running with hanzo/net inference.")
    print("Using the dummy engine for demonstration purposes.\n")
    asyncio.run(main())