Google ADK Research Assistant

pip install mantisdk google-adk openinference-instrumentation-google-adk python-dotenv

# .env file
INSIGHT_HOST=https://insight.withmetis.ai
INSIGHT_PUBLIC_KEY=pk-lf-your-key
INSIGHT_SECRET_KEY=sk-lf-your-key
GOOGLE_API_KEY=your-google-api-key

#!/usr/bin/env python3
"""Google ADK Research Assistant with Mantis SDK Tracing."""
 
import argparse
import asyncio
import datetime
from typing import Optional
from zoneinfo import ZoneInfo
from dotenv import load_dotenv
load_dotenv()
 
# Initialize tracing FIRST
import mantisdk.tracing as tracing
from mantisdk.tracing.exporters.insight import is_insight_configured
 
 
def setup_tracing(trace_to_console: bool = False):
    """Initialize mantisdk tracing with Google ADK instrumentation."""
    exporters = []
 
    if is_insight_configured():
        exporters.append(tracing.insight_exporter())
        print("[Tracing] Insight exporter configured")
    else:
        print("[Tracing] No Insight credentials found")
 
    if trace_to_console:
        from opentelemetry.sdk.trace.export import ConsoleSpanExporter
        exporters.append(ConsoleSpanExporter())
 
    tracing.init(
        service_name="google-adk-research-assistant",
        exporters=exporters if exporters else None,
        instrument_default=False,
    )
 
    tracing.instrument(names=["google_adk"])
    print("[Tracing] Google ADK instrumentation enabled")
 
 
# NOW import Google ADK
from google.adk.agents import Agent
from google.adk.runners import InMemoryRunner
from google.genai import types
 
 
# Tool implementations
def get_current_time(city: str) -> dict:
    """Get the current time in a specified city."""
    timezone_map = {
        "new york": "America/New_York",
        "los angeles": "America/Los_Angeles",
        "london": "Europe/London",
        "paris": "Europe/Paris",
        "tokyo": "Asia/Tokyo",
        "sydney": "Australia/Sydney",
    }
 
    city_lower = city.lower().strip()
    tz_id = timezone_map.get(city_lower)
 
    if not tz_id:
        return {
            "status": "error",
            "error_message": f"Unknown city: {city}. "
                           f"Available: {', '.join(timezone_map.keys())}"
        }
 
    tz = ZoneInfo(tz_id)
    now = datetime.datetime.now(tz)
    return {
        "status": "success",
        "report": f"Current time in {city.title()}: {now.strftime('%Y-%m-%d %H:%M:%S %Z')}"
    }
 
 
def get_weather(city: str) -> dict:
    """Get the current weather for a city (simulated)."""
    weather_data = {
        "new york": {"temp": 22, "condition": "Partly cloudy", "humidity": 65},
        "london": {"temp": 15, "condition": "Rainy", "humidity": 80},
        "tokyo": {"temp": 28, "condition": "Sunny", "humidity": 70},
        "paris": {"temp": 18, "condition": "Overcast", "humidity": 72},
    }
 
    city_lower = city.lower().strip()
    if city_lower not in weather_data:
        return {"status": "error", "error_message": f"No weather data for {city}"}
 
    data = weather_data[city_lower]
    return {
        "status": "success",
        "report": f"Weather in {city.title()}: {data['condition']}, "
                 f"{data['temp']}°C, Humidity: {data['humidity']}%"
    }
 
 
def calculate(expression: str) -> dict:
    """Evaluate a mathematical expression."""
    import math
    safe_dict = {
        "sqrt": math.sqrt, "sin": math.sin, "cos": math.cos,
        "tan": math.tan, "log": math.log, "exp": math.exp,
        "pi": math.pi, "e": math.e, "abs": abs, "round": round,
    }
    try:
        result = eval(expression.replace("^", "**"),
                     {"__builtins__": {}}, safe_dict)
        return {"status": "success", "result": result}
    except Exception as e:
        return {"status": "error", "error_message": str(e)}
 
 
def search_knowledge(query: str, topic: Optional[str] = None) -> dict:
    """Search a knowledge base (simulated)."""
    knowledge = {
        "python": {
            "topic": "technology",
            "content": "Python is a high-level programming language "
                      "known for simplicity and readability."
        },
        "machine learning": {
            "topic": "technology",
            "content": "ML enables systems to learn from data "
                      "without explicit programming."
        },
        "quantum computing": {
            "topic": "technology",
            "content": "Quantum computing uses quantum mechanics "
                      "for computation."
        },
    }
 
    query_lower = query.lower()
    for key, value in knowledge.items():
        if query_lower in key:
            if topic is None or value["topic"] == topic.lower():
                return {"status": "success", "result": value["content"]}
 
    return {"status": "not_found", "message": f"No results for '{query}'"}
 
 
def create_research_assistant(model: str = "gemini-2.0-flash") -> Agent:
    """Create the research assistant agent."""
    return Agent(
        name="research_assistant",
        model=model,
        description="A research assistant with time, weather, "
                   "calculator, and knowledge tools.",
        instruction="""You are a helpful research assistant with tools:
1. get_current_time - Get time in various cities
2. get_weather - Get weather information
3. calculate - Evaluate math expressions
4. search_knowledge - Search a knowledge base
 
Use these tools to answer user questions. Be helpful and concise.""",
        tools=[get_current_time, get_weather, calculate, search_knowledge],
    )
 
 
async def run_interactive_session(agent: Agent):
    """Run an interactive CLI session with the agent."""
    import uuid
 
    runner = InMemoryRunner(agent=agent, app_name="research_app")
    session_service = runner.session_service
 
    user_id = f"user_{uuid.uuid4().hex[:8]}"
    session_id = f"session_{uuid.uuid4().hex[:8]}"
 
    await session_service.create_session(
        app_name="research_app",
        user_id=user_id,
        session_id=session_id,
    )
 
    print("\n" + "="*60)
    print("  Research Assistant (Google ADK + Mantis Tracing)")
    print("="*60)
    print("\nI can help you with:")
    print("  - Time in various cities")
    print("  - Weather information")
    print("  - Mathematical calculations")
    print("  - Knowledge base searches")
    print("\nType 'quit' to exit.")
    print("="*60 + "\n")
 
    while True:
        try:
            user_input = input("You: ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\n\nGoodbye!")
            break
 
        if not user_input:
            continue
        if user_input.lower() in ("quit", "exit", "q"):
            print("\nGoodbye!")
            break
 
        message = types.Content(
            role="user",
            parts=[types.Part(text=user_input)]
        )
 
        print("\nAssistant: ", end="", flush=True)
 
        try:
            # This run is auto-instrumented
            async for event in runner.run_async(
                user_id=user_id,
                session_id=session_id,
                new_message=message,
            ):
                if event.is_final_response() and event.content:
                    for part in event.content.parts:
                        if hasattr(part, "text") and part.text:
                            print(part.text)
        except Exception as e:
            print(f"\n[Error] {e}")
        print()
 
 
async def run_single_query(agent: Agent, query: str):
    """Run a single query against the agent."""
    import uuid
 
    runner = InMemoryRunner(agent=agent, app_name="research_app")
    session_service = runner.session_service
 
    user_id = f"user_{uuid.uuid4().hex[:8]}"
    session_id = f"session_{uuid.uuid4().hex[:8]}"
 
    await session_service.create_session(
        app_name="research_app",
        user_id=user_id,
        session_id=session_id,
    )
 
    message = types.Content(
        role="user",
        parts=[types.Part(text=query)]
    )
 
    async for event in runner.run_async(
        user_id=user_id,
        session_id=session_id,
        new_message=message,
    ):
        if event.is_final_response() and event.content:
            for part in event.content.parts:
                if hasattr(part, "text") and part.text:
                    print(part.text)
 
 
def main():
    parser = argparse.ArgumentParser(description="Google ADK Research Assistant")
    parser.add_argument("--model", default="gemini-2.0-flash")
    parser.add_argument("--query", "-q", help="Single query")
    parser.add_argument("--trace-to-console", action="store_true")
    parser.add_argument("--no-tracing", action="store_true")
    args = parser.parse_args()
 
    if not args.no_tracing:
        setup_tracing(trace_to_console=args.trace_to_console)
 
    agent = create_research_assistant(model=args.model)
 
    try:
        if args.query:
            asyncio.run(run_single_query(agent, args.query))
        else:
            asyncio.run(run_interactive_session(agent))
    finally:
        if not args.no_tracing:
            tracing.shutdown()
 
 
if __name__ == "__main__":
    main()

python research_assistant_agent.py

python research_assistant_agent.py --query "What time is it in Tokyo?"

You: What time is it in New York?
Assistant: The current time in New York is 2024-01-15 14:30:45 EST.

You: Calculate sqrt(144) + 10
Assistant: The result is 22.0.

You: Tell me about machine learning
Assistant: ML enables systems to learn from data without explicit programming.

run_async (auto-instrumented)
├── agent: research_assistant
├── user_id: user_abc123
├── session_id: session_xyz789
├── model: gemini-2.0-flash
└── tool_calls
    └── get_current_time
        ├── input: {"city": "Tokyo"}
        └── output: {"status": "success", "report": "..."}

agent = Agent(
    name="research_assistant",
    model="gemini-2.0-flash",
    tools=[get_current_time, get_weather, calculate, search_knowledge],
)

await session_service.create_session(
    app_name="research_app",
    user_id=user_id,
    session_id=session_id,
)

async for event in runner.run_async(...):
    if event.is_final_response():
        print(event.content.parts[0].text)