react_agent.py

"""Example: ReAct agent with real web search and Temporal durability.

This example demonstrates a DSPy ReAct agent that combines:
- Real web search using Tavily API
- Mathematical computation
- Durable execution with automatic retries and checkpointing

The agent answers complex questions requiring multiple tool calls, with each
tool execution becoming a Temporal activity for fault tolerance.

Usage:
    # Start Temporal server first: docker-compose up -d
    # Set environment variables:
    #   export OPENAI_API_KEY=your-key-here
    #   export TAVILY_API_KEY=your-key-here
    # Then run: python examples/react_agent.py
"""

from __future__ import annotations

import asyncio
import os
import uuid

import dspy
from dotenv import load_dotenv
from tavily import TavilyClient
from temporalio import workflow
from temporalio.client import Client
from temporalio.worker import Worker

from dspy_temporal import DSPyPlugin, TemporalModule
from dspy_temporal.sandbox import get_default_sandbox_runner

load_dotenv()


# ============================================================================
# Tool Definitions (standard DSPy functions)
# ============================================================================


def evaluate_math(expression: str) -> str:
    """Safely evaluate a basic math expression."""

    import math

    allowed_names = {
        "abs": abs,
        "round": round,
        "min": min,
        "max": max,
        "sum": sum,
        "pow": pow,
        "sqrt": math.sqrt,
        "sin": math.sin,
        "cos": math.cos,
        "tan": math.tan,
        "log": math.log,
        "log10": math.log10,
        "pi": math.pi,
        "e": math.e,
    }

    for char in expression:
        if char not in "0123456789+-*/.() ," and not any(name in expression for name in allowed_names):
            return f"Error: Invalid character '{char}' in expression"

    try:
        result = eval(expression, {"__builtins__": {}}, allowed_names)  # noqa: S307
    except Exception as exc:
        print(f"Error: {exc}")
        return f"Error: {exc}"
    return str(result)


def search_web(query: str) -> str:
    """Search the web using Tavily API.

    Args:
        query: The search query.

    Returns:
        Formatted search results with title, content, and source URLs.
    """
    api_key = os.getenv("TAVILY_API_KEY")
    if not api_key:
        return f"[Search failed for '{query}' - TAVILY_API_KEY not set]"

    try:
        client = TavilyClient(api_key)
        response = client.search(query=query, max_results=3)
        # Response format: {"results": [{"title": ..., "url": ..., "content": ..., "score": ...}]}
        results = response.get("results", [])

        if not results:
            return f"No results found for '{query}'"

        # Format results for the agent
        passages = []
        for r in results:
            title = r.get("title", "")
            content = r.get("content", "")
            url = r.get("url", "")
            passages.append(f"{title}\n{content}\nSource: {url}")

        return "\n\n".join(passages)
    except Exception as e:
        print(f"Tavily error: {e}")
        return f"[Search failed for '{query}' - {e}]"


# ============================================================================
# DSPy ReAct Agent
# ============================================================================


react_agent = dspy.ReAct(
    "question -> answer",
    tools=[evaluate_math, search_web],
)

# TemporalModule automatically wraps the tools as activities
# and exposes them alongside the ReAct LLM calls.
temporal_react = TemporalModule(react_agent, name="react_agent")


# ============================================================================
# Temporal Workflow
# ============================================================================


@workflow.defn
class ReActWorkflow:
    """Temporal workflow that executes the ReAct agent."""

    @workflow.run
    async def run(self, question: str) -> str:
        result = await temporal_react.run(question=question)
        return result.answer


# ============================================================================
# Main Execution
# ============================================================================


async def main() -> None:
    # Optional: Enable debug logging to see detailed execution traces
    # import logging
    # logging.basicConfig(level=logging.DEBUG)
    # logging.getLogger('dspy_temporal').setLevel(logging.DEBUG)

    # Check for required API keys
    if not os.environ.get("OPENAI_API_KEY"):
        print("ERROR: OPENAI_API_KEY environment variable not set")
        print("Please set it: export OPENAI_API_KEY=your-key-here")
        return

    if not os.environ.get("TAVILY_API_KEY"):
        print("ERROR: TAVILY_API_KEY environment variable not set")
        print("Please set it: export TAVILY_API_KEY=your-key-here")
        return

    dspy.configure(lm=dspy.LM("openai/gpt-5-mini"))

    temporal_url = os.getenv("TEMPORAL_SERVICE_URL", "localhost:7233")
    client = await Client.connect(temporal_url)

    async with Worker(
        client,
        task_queue="react-agent-queue",
        workflows=[ReActWorkflow],
        plugins=[DSPyPlugin(temporal_react)],
        workflow_runner=get_default_sandbox_runner(),
    ):
        print("Worker started. Executing ReAct agent workflow...")

        workflow_id = f"react-agent-{uuid.uuid4().hex[:8]}"
        result = await client.execute_workflow(
            ReActWorkflow.run,
            args=[
                "What is the current population of Ankara, and how much would it cost to buy a cup of tea for everyone in Ankara?"
            ],
            id=workflow_id,
            task_queue="react-agent-queue",
        )
        print(f"\nAnswer: {result}")


if __name__ == "__main__":
    asyncio.run(main())