SUTRA with OpenAI Agents SDK

!pip install "openai-agents[litellm]"

import os
from google.colab import userdata

# Set the API key from Colab secrets
os.environ["SUTRA_API_KEY"] = userdata.get("SUTRA_API_KEY")

import logging
logging.getLogger("openai.agents").setLevel(logging.ERROR)

colab="{"base_uri":"https://localhost:8080/"}"
id="biPIsngSOd5r" outputId="6fe28960-9efa-46ef-bcd9-bcb197670a92">

```python
import nest_asyncio
import asyncio
import os
from agents import Agent, Runner
from agents.extensions.models.litellm_model import LitellmModel

# Patch the running event loop
nest_asyncio.apply()

# Main async logic
async def main():
    agent = Agent(
        name="Sutra Agent",
        instructions="You are a helpful assistant that responds in Hindi.",
        model=LitellmModel(
            model="openai/sutra-v2",
            api_key=os.environ.get("SUTRA_API_KEY"),
            base_url="https://api.two.ai/v2"
        )
    )

    result = await Runner.run(agent, "इक्विटी क्या होता है?")
    print(result.final_output)

# Run the async function safely
await main()

import asyncio
import os
from openai.types.responses import ResponseTextDeltaEvent
from agents import Agent, Runner
from agents.extensions.models.litellm_model import LitellmModel

async def main():
    # Create an Agent using Sutra (via LiteLLM)
    agent = Agent(
        name="Sutra Stream Agent",
        instructions="तुम एक सहायक हो जो हिंदी में जवाब देता है।",  # Instructions in Hindi
        model=LitellmModel(
            model="openai/sutra-v2",
            api_key=os.environ.get("SUTRA_API_KEY"),
            base_url="https://api.two.ai/v2"
        ),
    )

    # Start streaming the response
    result = Runner.run_streamed(agent, input="भारत का इतिहास बताइए।")

    print("🟢 Streaming started...\n")
    async for event in result.stream_events():
        # Print LLM tokens as they're streamed
        if event.type == "raw_response_event" and isinstance(event.data, ResponseTextDeltaEvent):
            print(event.data.delta, end="", flush=True)

    print("\n✅ Streaming complete.")

# 🔁 Run in a proper async environment
await main()

from agents import function_tool

@function_tool
def get_weather(city: str) -> str:
    """Get weather details for a given city."""
    return f"{city} में मौसम बहुत ही सुहावना है।"

from typing import Any
from pydantic import BaseModel
from agents import FunctionTool, RunContextWrapper

class UserInfo(BaseModel):
    name: str
    age: int

async def process_user(ctx: RunContextWrapper[Any], args: str) -> str:
    data = UserInfo.model_validate_json(args)
    return f"{data.name} की उम्र {data.age} वर्ष है।"

process_user_tool = FunctionTool(
    name="process_user_info",
    description="उपयोगकर्ता की जानकारी को संसाधित करता है।",
    params_json_schema=UserInfo.model_json_schema(),
    on_invoke_tool=process_user
)

from agents import Agent

hindi_translator = Agent(
    name="Hindi Translator",
    instructions="आप उपयोगकर्ता के संदेश को हिंदी में अनुवाद करते हैं।",
)

# Turn agent into a tool
translator_tool = hindi_translator.as_tool(
    tool_name="translate_to_hindi",
    tool_description="संदेश को हिंदी में अनुवाद करें।"
)

import nest_asyncio
import asyncio
import os
from agents import Agent, Runner
from agents.extensions.models.litellm_model import LitellmModel

nest_asyncio.apply()

async def main():
    agent = Agent(
        name="Sutra Agent",
        instructions="आप एक सहायक हैं जो हिंदी में जवाब देता है।",
        model=LitellmModel(
            model="openai/sutra-v2",
            api_key=os.environ["SUTRA_API_KEY"],
            base_url="https://api.two.ai/v2"
        ),
        tools=[get_weather, process_user_tool, translator_tool],
    )

    result = await Runner.run(agent, "मुझे जयपुर के मौसम के बारे में बताइए।")
    print(result.final_output)

await main()

import nest_asyncio
import asyncio
import os
from pydantic import BaseModel
from agents import (
    Agent,
    Runner,
    function_tool,
    GuardrailFunctionOutput,
    RunContextWrapper,
    InputGuardrailTripwireTriggered,
    input_guardrail,
    TResponseInputItem,
)
from agents.extensions.models.litellm_model import LitellmModel

# Patch event loop for Jupyter or other async environments
nest_asyncio.apply()


# Define the guardrail output model
class MathHomeworkOutput(BaseModel):
    is_math_homework: bool
    reasoning: str

# Setup guardrail agent - MUST output strict JSON only!
guardrail_agent = Agent(
    name="Math Homework Guardrail",
    instructions=(
        "You are a JSON API. Determine if the input asks for math homework help.\n"
        "Output ONLY a JSON object with fields:\n"
        "- is_math_homework (boolean)\n"
        "- reasoning (string)\n"
        "Do NOT output any extra text."
    ),
    output_type=MathHomeworkOutput,
    model=LitellmModel(
        model="openai/sutra-v2",
        api_key=os.environ.get("SUTRA_API_KEY"),
        base_url="https://api.two.ai/v2",

    ),
)

# Guardrail function using input_guardrail decorator
@input_guardrail
async def math_guardrail(
    ctx: RunContextWrapper[None],
    agent: Agent,
    input: str | list[TResponseInputItem]
) -> GuardrailFunctionOutput:
    result = await Runner.run(guardrail_agent, input, context=ctx.context)
    return GuardrailFunctionOutput(
        output_info=result.final_output,
        tripwire_triggered=result.final_output.is_math_homework,
    )

# Main agent setup with Sutra model and input guardrail
main_agent = Agent(
    name="Sutra Agent",
    instructions="You are a helpful assistant that responds in Hindi.",
    model=LitellmModel(
        model="openai/sutra-v2",
        api_key=os.environ.get("SUTRA_API_KEY"),
        base_url="https://api.two.ai/v2",
    ),
    tools=[get_weather],
    input_guardrails=[math_guardrail],
)

async def main():
    # Example 1: Math homework input (should trigger guardrail)
    try:
        response = await Runner.run(main_agent, "मुझे 2x + 3 = 11 का हल निकालने में मदद करें।")
        print("Example 1 - Agent response:", response.final_output)
    except InputGuardrailTripwireTriggered:
        print("Example 1 - ⚠️ Guardrail triggered: Math homework detected! Agent run stopped.")

    # Example 2: Non-math input (should NOT trigger guardrail)
    try:
        response = await Runner.run(main_agent, "मुझे दिल्ली के मौसम के बारे में बताओ।")
        print("Example 2 - Agent response:", response.final_output)
    except InputGuardrailTripwireTriggered:
        print("Example 2 - ⚠️ Guardrail triggered unexpectedly!")

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

import nest_asyncio
import asyncio
import os
from agents import Agent, Runner, function_tool
from agents.extensions.models.litellm_model import LitellmModel

nest_asyncio.apply()

# Function tool example (optional)
@function_tool
def search_web(query: str) -> str:
    # Placeholder: pretend to search web and return snippet
    return f"Search results for '{query}': Important info about {query}."

# Setup Research Agent
research_agent = Agent(
    name="Research Agent",
    instructions="You research topics and return concise summaries.",
    model=LitellmModel(
        model="openai/sutra-v2",
        api_key=os.environ.get("SUTRA_API_KEY"),
        base_url="https://api.two.ai/v2",
    ),
    tools=[search_web],
)

# Setup Writer Agent
writer_agent = Agent(
    name="Writer Agent",
    instructions="You write a blog post based on the provided research notes.",
    model=LitellmModel(
        model="openai/sutra-v2",
        api_key=os.environ.get("SUTRA_API_KEY"),
        base_url="https://api.two.ai/v2",
    ),
)

# Orchestrator Agent: will delegate via handoffs
orchestrator_agent = Agent(
    name="Orchestrator Agent",
    instructions=(
        "You plan a blog post workflow. "
        "First, delegate research tasks to the Research Agent, "
        "then delegate writing to the Writer Agent using research output."
    ),
    model=LitellmModel(
        model="openai/sutra-v2",
        api_key=os.environ.get("SUTRA_API_KEY"),
        base_url="https://api.two.ai/v2",
    ),
    handoffs=[research_agent, writer_agent],  # allows delegating tasks
)

# -------- ORCHESTRATION VIA CODE --------
async def orchestrate_via_code(topic: str):
    # Step 1: Run Research Agent
    research_result = await Runner.run(research_agent, topic)
    research_notes = research_result.final_output
    print("\n[Research Agent Output]:", research_notes)

    # Step 2: Run Writer Agent with research notes as input
    writer_result = await Runner.run(writer_agent, research_notes)
    blog_post = writer_result.final_output
    print("\n[Writer Agent Output]:", blog_post)

# -------- ORCHESTRATION VIA LLM HANDOFFS --------
async def orchestrate_via_llm(topic: str):
    # Just call orchestrator agent with the topic and it handles handoffs internally
    result = await Runner.run(orchestrator_agent, topic)
    print("\n[Orchestrator Agent Output]:", result.final_output)

async def main():
    topic = "Climate change impact on agriculture"

    print("===== Orchestration via code =====")
    await orchestrate_via_code(topic)

    print("\n===== Orchestration via LLM handoffs =====")
    await orchestrate_via_llm(topic)

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