python-sdk/examples/using_grok.py

#!/usr/bin/env python
"""Example of using xAI Grok with Hanzo Python SDK.

This example demonstrates:
1. Direct API usage with Grok
2. Streaming responses
3. Batch operations with Grok
4. Consensus with multiple models including Grok
"""

import os
import sys
import asyncio

# Add paths for development (remove in production)
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "pkg", "hanzoai"))
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "pkg", "hanzo", "src"))

try:
    from hanzoai import Hanzo, AsyncHanzo
except ImportError:
    print("Error: hanzoai not installed")
    print("Install with: pip install hanzoai")
    sys.exit(1)

# Note: Requires XAI_API_KEY environment variable to be set
# Get your API key from https://x.ai/api


def basic_grok_usage():
    """Basic usage of Grok through Hanzo SDK."""
    client = Hanzo()  # Uses HANZO_API_KEY env var

    # Simple completion
    response = client.chat.completions.create(
        model="grok-4",  # or "grok" or "xai-grok"
        messages=[
            {"role": "system", "content": "You have real-time knowledge."},
            {"role": "user", "content": "What are the latest AI developments this week?"},
        ],
        temperature=0.7,
        max_tokens=500,
    )

    print("Grok Response:")
    print(response.choices[0].message.content)
    print("\n" + "=" * 50 + "\n")


def streaming_grok():
    """Stream responses from Grok for real-time output."""
    client = Hanzo()

    print("Streaming from Grok:")
    stream = client.chat.completions.create(
        model="grok-4", messages=[{"role": "user", "content": "Explain quantum computing in simple terms"}], stream=True
    )

    for chunk in stream:
        if chunk.choices[0].delta.content:
            print(chunk.choices[0].delta.content, end="", flush=True)

    print("\n\n" + "=" * 50 + "\n")


async def async_grok_usage():
    """Async usage of Grok for concurrent operations."""
    client = AsyncHanzo()

    # Concurrent requests to different models
    tasks = [
        client.chat.completions.create(
            model="grok-4", messages=[{"role": "user", "content": "What's happening in tech today?"}]
        ),
        client.chat.completions.create(
            model="claude-3-5-sonnet", messages=[{"role": "user", "content": "Explain machine learning"}]
        ),
        client.chat.completions.create(model="gpt-4", messages=[{"role": "user", "content": "Write a haiku about AI"}]),
    ]

    print("Concurrent requests to Grok, Claude, and GPT-4:")
    responses = await asyncio.gather(*tasks)

    for i, (model, response) in enumerate(zip(["Grok", "Claude", "GPT-4"], responses)):
        print(f"\n{model} Response:")
        print(response.choices[0].message.content[:200] + "...")

    print("\n" + "=" * 50 + "\n")


def batch_operations_with_grok():
    """Demonstrate batch operations using Grok."""
    from hanzo.batch_orchestrator import BatchConfig, BatchOrchestrator

    print("Batch Operations with Grok:")

    # Parse different batch configurations
    configs = [
        "batch:5 agent:grok analyze code quality",
        "consensus:3 agent:grok,claude,gemini review architecture",
        "critic:2 agent:grok,gpt-4 chain:true security audit",
    ]

    for cmd in configs:
        config = BatchConfig.from_command(cmd)
        print(f"\nCommand: {cmd}")
        print(f"  - Batch size: {config.batch_size}")
        print(f"  - Agent model: {config.agent_model}")
        if config.consensus_mode:
            print(f"  - Consensus models: {config.consensus_models}")
        if config.critic_mode:
            print(f"  - Critic chain: {config.critic_chain}")

    print("\n" + "=" * 50 + "\n")


def grok_with_tools():
    """Use Grok with MCP tools."""
    try:
        from hanzo_mcp.tools.agent.cli_tools import GrokCLITool

        print("Grok CLI Tool Configuration:")
        grok_tool = GrokCLITool()

        print(f"Tool name: {grok_tool.name}")
        print(f"Description: {grok_tool.description}")
        print(f"Default model: {grok_tool.default_model}")
        print(f"API key env: {grok_tool.api_key_env}")

        # Get auth environment
        env = grok_tool.get_auth_env()
        if "XAI_API_KEY" in env:
            print("✅ XAI_API_KEY is configured")
        else:
            print("⚠️ XAI_API_KEY not found in environment")

    except ImportError:
        print("hanzo-mcp not installed. Install with: pip install hanzo-mcp")

    print("\n" + "=" * 50 + "\n")


def main():
    """Run all Grok examples."""
    print("=" * 50)
    print("xAI Grok Examples with Hanzo SDK")
    print("=" * 50 + "\n")

    # Check for API key
    if not os.environ.get("XAI_API_KEY"):
        print("⚠️ Warning: XAI_API_KEY not set")
        print("Get your API key from https://x.ai/api")
        print("Set it with: export XAI_API_KEY='your-key-here'")
        print("\nRunning examples that don't require actual API calls...\n")

        # Run non-API examples
        batch_operations_with_grok()
        grok_with_tools()
        return

    # Run all examples
    try:
        # Basic usage
        basic_grok_usage()

        # Streaming
        streaming_grok()

        # Async operations
        asyncio.run(async_grok_usage())

        # Batch operations
        batch_operations_with_grok()

        # MCP tools
        grok_with_tools()

        print("✅ All Grok examples completed successfully!")

    except Exception as e:
        print(f"Error: {e}")
        print("\nMake sure you have:")
        print("1. Set XAI_API_KEY environment variable")
        print("2. Set HANZO_API_KEY environment variable (or pass api_key to client)")
        print("3. Installed hanzoai: pip install hanzoai")
        print("4. (Optional) Installed hanzo-mcp: pip install hanzo-mcp")


if __name__ == "__main__":
    main()