Nyraxis AI

LangGraph

Add governance guardrails to LangGraph stateful agents.

LangGraph Integration

Installation

pip install "nyraxis-sdk[instrument]"

Setup

import nyraxis_sdk
from langgraph.graph import StateGraph, START, END
from langchain_openai import ChatOpenAI

nyraxis_sdk.init(
    api_key="nyx_...",
    instrument=True,
    enforce=True,
    agent_name="my-agent",
)

llm = ChatOpenAI(model="gpt-4o")

Grouping spans into a single trace

LangGraph agents make multiple calls per user interaction (LLM, retrieval, tools). Without a parent span, each call becomes a separate trace in the dashboard.

Wrap your graph invocation with nyraxis_sdk.trace() to group everything into one trace:

def handle_message(user_id: str, message: str):
    with nyraxis_sdk.trace("chat", user_id=user_id):
        result = graph.invoke({"messages": [{"role": "user", "content": message}]})
    return result["messages"][-1].content

All LLM calls, retrieval queries, and tool executions inside the with block become child spans of a single trace.

Decorator alternative

@nyraxis_sdk.traced("handle-message")
def handle(user_id: str, message: str):
    return graph.invoke({"messages": [{"role": "user", "content": message}]})

Streaming

def handle_stream(user_id: str, message: str):
    with nyraxis_sdk.trace("chat-stream", user_id=user_id):
        messages = prepare_messages(message)
        for token in llm.stream(messages):
            yield token

StateGraph example

from typing import TypedDict, Annotated
from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages

class State(TypedDict):
    messages: Annotated[list, add_messages]

def chatbot(state: State):
    response = llm.invoke(state["messages"])
    return {"messages": [response]}

graph = StateGraph(State)
graph.add_node("chatbot", chatbot)
graph.add_edge(START, "chatbot")
graph.add_edge("chatbot", END)

app = graph.compile()

with nyraxis_sdk.trace("user-query"):
    result = app.invoke({"messages": [{"role": "user", "content": "Hello!"}]})

Multi-node graphs

def researcher(state):
    return {"messages": [llm.invoke(state["messages"] + [{"role": "system", "content": "You are a researcher."}])]}

def writer(state):
    return {"messages": [llm.invoke(state["messages"] + [{"role": "system", "content": "You are a writer."}])]}

graph = StateGraph(State)
graph.add_node("researcher", researcher)
graph.add_node("writer", writer)
graph.add_edge(START, "researcher")
graph.add_edge("researcher", "writer")
graph.add_edge("writer", END)

app = graph.compile()

with nyraxis_sdk.trace("research-and-write", user_id="u123"):
    result = app.invoke({"messages": [{"role": "user", "content": "Write about AI safety"}]})

What gets traced

ComponentSpan TypeNotes
nyraxis_sdk.trace() wrapperparentGroups all spans into one trace
LLM calls (OpenAI, Anthropic, etc.)llmModel name, tokens, latency
Retrieval / vector DB queriesretrievalQuery, results count
Tool callstoolTool name, input/output

Governance enforcement

When enforce=True, every LLM call within your graph is checked against your policies before the response is returned. If a violation is detected:

  • Blocking mode: NyraxisBlockedError is raised — the LLM call never happens
  • Monitor mode: The call proceeds but the violation is logged

System prompts are automatically excluded from governance checks — only user and tool messages are evaluated.

If you don't wrap with trace(), each LLM call and retrieval becomes a separate trace in the dashboard. Always use trace() in production.

On this page