Tracing LangChain Agents

pip install netra-sdk langchain langchain-openai openai

export NETRA_API_KEY="your-netra-api-key"
export NETRA_OTLP_ENDPOINT="your-netra-otlp-endpoint"
export OPENAI_API_KEY="your-openai-api-key"

from typing import Dict, List
from langchain.tools import tool
from langchain_openai import ChatOpenAI
from langchain.agents import create_react_agent

# Mock databases
TICKETS = {
    "TKT-001": {"id": "TKT-001", "subject": "Return policy question", "status": "open"},
    "TKT-002": {"id": "TKT-002", "subject": "Damaged item", "status": "open", "order_id": "ORD-12345"},
}

ORDERS = {
    "ORD-12345": {"id": "ORD-12345", "status": "delivered", "items": ["Headphones"], "total": 79.99},
}

KNOWLEDGE_BASE = [
    {"title": "Return Policy", "content": "Items can be returned within 30 days."},
    {"title": "Refund Processing", "content": "Refunds processed in 5-7 business days."},
]

@tool
def lookup_ticket(ticket_id: str) -> str:
    """Look up a ticket by its ID to get details about the issue."""
    ticket = TICKETS.get(ticket_id.upper())
    if not ticket:
        return f"No ticket found with ID: {ticket_id}"
    return f"Ticket {ticket['id']}: {ticket['subject']} (Status: {ticket['status']})"

@tool
def search_kb(query: str) -> str:
    """Search the knowledge base for information about policies or procedures."""
    query_lower = query.lower()
    results = [a for a in KNOWLEDGE_BASE if query_lower in a["title"].lower()]
    if not results:
        return "No relevant articles found."
    return "\n".join([f"**{a['title']}**: {a['content']}" for a in results])

@tool
def check_order_status(order_id: str) -> str:
    """Check the status of an order including shipping information."""
    order = ORDERS.get(order_id.upper())
    if not order:
        return f"No order found with ID: {order_id}"
    return f"Order {order['id']}: {order['status']}, Items: {order['items']}, Total: ${order['total']}"

@tool
def escalate_to_human(ticket_id: str, reason: str) -> str:
    """Escalate a ticket to a human operator for complex issues."""
    return f"Ticket {ticket_id} escalated. Reason: {reason}. A specialist will respond within 1 hour."

# Initialize the LLM
model = ChatOpenAI(model="gpt-4o-mini", temperature=0)

# Define tools
tools = [lookup_ticket, search_kb, check_order_status, escalate_to_human]

# Create the agent
agent = create_react_agent(
    model,
    tools,
    prompt="""You are TaskBot, an AI assistant for ShopFlow e-commerce platform.

You help users with:
- Order status and tracking
- Return and refund requests
- Policy questions
- Escalating complex issues

Use tools to look up information before responding.
Escalate to human operators when the user is frustrated or you cannot resolve the issue."""
)

from netra import Netra
from netra.instrumentation.instruments import InstrumentSet

# Initialize Netra with LangChain and OpenAI instrumentation
Netra.init(
    app_name="taskbot",
    environment="development",
    trace_content=True,
    instruments={InstrumentSet.OPENAI, InstrumentSet.LANGCHAIN},
)

from netra.decorators import agent

@agent(name="taskbot-agent")
def handle_request(query: str, user_id: str = None) -> dict:
    """Handle a user request with full tracing."""

    # Set user context if provided
    if user_id:
        Netra.set_user_id(user_id)

    # Execute the agent
    result = agent.invoke({
        "messages": [{"role": "user", "content": query}]
    })

    return {
        "query": query,
        "response": result["messages"][-1].content,
    }

from netra import Netra, SpanType

@tool
def lookup_ticket_traced(ticket_id: str) -> str:
    """Look up a ticket with custom span attributes."""
    with Netra.start_span("ticket-lookup", as_type=SpanType.TOOL) as span:
        span.set_attribute("ticket_id", ticket_id)

        ticket = TICKETS.get(ticket_id.upper())

        if ticket:
            span.set_attribute("ticket_status", ticket["status"])
            span.set_attribute("ticket_priority", ticket.get("priority", "normal"))
            span.set_attribute("found", True)
        else:
            span.set_attribute("found", False)

        if not ticket:
            return f"No ticket found with ID: {ticket_id}"

        return f"Ticket {ticket['id']}: {ticket['subject']} (Status: {ticket['status']})"

# Single-tool query - should use search_kb
response = handle_request(
    query="What is your return policy?",
    user_id="user-001",
)
print(response["response"])

# Order status query - should use check_order_status
response = handle_request(
    query="Where is my order ORD-12345?",
    user_id="user-002",
)
print(response["response"])

# Multi-step workflow - should use multiple tools
response = handle_request(
    query="I have ticket TKT-002 about a damaged item. Can you check the order status?",
    user_id="user-003",
)
print(response["response"])

# Escalation scenario - should detect urgency
response = handle_request(
    query="I've been waiting 3 weeks and need urgent help! I want to speak to someone immediately!",
    user_id="user-004",
)
print(response["response"])