# LangGraph Integration for Arzule SDK

## Overview

The Arzule SDK now includes comprehensive support for LangGraph, a framework for building stateful, multi-step agent applications. LangGraph is built on top of LangChain and provides a graph-based execution model with nodes, edges, and state management.

## Architecture

### Key Components

1. **`langgraph/__init__.py`** - Public API for LangGraph integration
2. **`langgraph/install.py`** - Instrumentation entrypoint
3. **`langgraph/callback_handler.py`** - Captures LangGraph execution events via LangChain's callback system
4. **`langgraph/normalize.py`** - Converts LangGraph events to Arzule TraceEvent format
5. **`langgraph/spanctx.py`** - Manages span context for proper parent-child relationships

### How It Works

LangGraph is built on LangChain and uses the same `BaseCallbackHandler` system for observability. Our integration:

1. **Extends LangChain's callback system** - The `ArzuleLangGraphHandler` inherits from `BaseCallbackHandler` and intercepts execution events
2. **Distinguishes graph vs node execution** - Uses serialized metadata to differentiate between:
   - Graph-level execution (Pregel/CompiledStateGraph)
   - Node-level execution (individual node functions)
3. **Tracks nested calls** - Captures LLM calls, tool invocations, and other LangChain components used within nodes
4. **Maps to Arzule event types**:
   - Graph execution → `graph.start` / `graph.end`
   - Node execution → `agent.start` / `agent.end` (nodes map to agents in Arzule's model)
   - Checkpoints → `checkpoint.save` / `checkpoint.load`
   - LLM calls → `llm.call.start` / `llm.call.end`
   - Tool calls → `tool.call.start` / `tool.call.end`

## Installation & Usage

### Automatic Instrumentation

When you call `arzule_ingest.init()` with `auto_instrument=True` (default), LangGraph will be automatically instrumented if installed:

```python
import arzule_ingest

# Auto-instruments LangGraph, LangChain, CrewAI, AutoGen (if available)
arzule_ingest.init()
```

### Manual Instrumentation

For more control, manually instrument LangGraph:

```python
from arzule_ingest.langgraph import instrument_langgraph

# Full instrumentation (recommended)
handler = instrument_langgraph()

# Minimal instrumentation (essential events only)
handler = instrument_langgraph(mode="minimal")

# Custom instrumentation
handler = instrument_langgraph(
    enable_graph_callbacks=True,      # Graph start/end
    enable_node_callbacks=True,       # Node start/end
    enable_llm_callbacks=True,        # LLM calls
    enable_tool_callbacks=True,       # Tool invocations
    enable_checkpoint_callbacks=True, # State persistence
)
```

### Using with LangGraph

```python
from langgraph.graph import START, StateGraph
from typing import TypedDict
import arzule_ingest
from arzule_ingest.langgraph import instrument_langgraph

# Initialize Arzule
arzule_ingest.init()

# Instrument LangGraph
handler = instrument_langgraph()

# Define your state
class State(TypedDict):
    text: str

# Define nodes
def node_a(state: State) -> dict:
    return {"text": state["text"] + "a"}

def node_b(state: State) -> dict:
    return {"text": state["text"] + "b"}

# Build graph
graph = StateGraph(State)
graph.add_node("node_a", node_a)
graph.add_node("node_b", node_b)
graph.add_edge(START, "node_a")
graph.add_edge("node_a", "node_b")

compiled = graph.compile()

# Run with Arzule tracking
with arzule_ingest.ArzuleRun(
    tenant_id="your-tenant-id",
    project_id="your-project-id",
    metadata={"workflow": "example"}
) as run:
    result = compiled.invoke(
        {"text": ""},
        config={"callbacks": [handler]}
    )
```

## Event Mapping

### LangGraph Events → Arzule TraceEvents

| LangGraph Event | Arzule Event Type | Notes |
|----------------|-------------------|-------|
| Graph starts | `graph.start` | Triggered when StateGraph.invoke() is called |
| Graph ends | `graph.end` | Includes final output or error |
| Node starts | `agent.start` | Each node is treated as an agent |
| Node ends | `agent.end` | Includes node output or error |
| LLM call (nested) | `llm.call.start` / `llm.call.end` | From LangChain LLMs used in nodes |
| Tool call (nested) | `tool.call.start` / `tool.call.end` | From LangChain tools used in nodes |
| Checkpoint save | `checkpoint.save` | State persistence |
| Checkpoint load | `checkpoint.load` | State recovery |

### Span Hierarchy

```
run_id: uuid
  └─ graph.start (span_id: xxx)
      ├─ agent.start (node_a, span_id: yyy, parent: xxx)
      │   ├─ llm.call.start (span_id: zzz, parent: yyy)
      │   └─ llm.call.end (span_id: zzz, parent: yyy)
      ├─ agent.end (node_a, span_id: yyy, parent: xxx)
      ├─ agent.start (node_b, span_id: aaa, parent: xxx)
      └─ agent.end (node_b, span_id: aaa, parent: xxx)
      └─ graph.end (span_id: xxx)
```

## Supported LangGraph Features

### ✅ Fully Supported

- ✅ **StateGraph execution** - Basic graph execution with state management
- ✅ **Node execution tracking** - Individual node start/end events
- ✅ **Conditional edges** - Branching logic based on state
- ✅ **Nested LLM calls** - LangChain LLMs called within nodes
- ✅ **Nested tool calls** - LangChain tools called within nodes
- ✅ **Error tracking** - Exceptions in graphs or nodes
- ✅ **Parallel execution** - Multiple nodes running concurrently
- ✅ **Stream mode support** - Works with all LangGraph stream modes

### 🚧 Partial Support

- 🚧 **Checkpoint tracking** - Events are defined but need checkpoint saver integration
- 🚧 **Subgraphs** - Basic support; nested graph tracking can be improved
- 🚧 **Human-in-the-loop** - Interrupt events are tracked via graph execution

### ❌ Not Yet Supported

- ❌ **Custom stream data** - LangGraph's custom stream mode
- ❌ **Message-level streaming** - Token-by-token LLM streaming (would be very noisy)

## Advanced Features

### Conditional Routing

LangGraph's conditional edges are automatically tracked:

```python
def router(state: State) -> str:
    """Route to different nodes based on state."""
    if state["value"] > 50:
        return "high_value_node"
    return "low_value_node"

graph.add_conditional_edges(
    "start",
    router,
    {
        "high_value_node": "high_value_node",
        "low_value_node": "low_value_node"
    }
)
```

Both the routing decision and the selected node execution are captured in traces.

### Checkpointing & State Persistence

When using LangGraph's checkpointing (e.g., `MemorySaver`, `SqliteSaver`):

```python
from langgraph.checkpoint.memory import MemorySaver

checkpointer = MemorySaver()
compiled = graph.compile(checkpointer=checkpointer)

# Checkpoint save/load events are captured
result = compiled.invoke(
    input_data,
    config={"configurable": {"thread_id": "thread-1"}, "callbacks": [handler]}
)
```

### Parallel Node Execution

LangGraph can execute multiple nodes in parallel. Each node gets its own span:

```python
# Nodes that can run in parallel
graph.add_node("parallel_1", func1)
graph.add_node("parallel_2", func2)
graph.add_edge(START, "parallel_1")
graph.add_edge(START, "parallel_2")
```

All parallel executions are properly tracked with correct parent-child relationships.

## Testing

Run the integration test:

```bash
cd arzuleIngestionWrapper
export ARZULE_API_KEY=your-key
export ARZULE_TENANT_ID=your-tenant-id
export ARZULE_PROJECT_ID=your-project-id

python test_langgraph_integration.py
```

## Coexistence with LangChain Integration

### Why Both Integrations Exist

**LangChain integration** captures:
- Pure LangChain usage (chains, agents, tools, retrievers)
- Direct LLM/chain composition
- RAG pipelines

**LangGraph integration** captures:
- Graph-level execution (StateGraph/Pregel)
- Node-level execution
- State transitions and checkpointing

### Preventing Duplicate Events

The integrations are designed to coexist without duplication:

1. **LangGraph handler** only processes LangGraph-specific callbacks:
   - Checks serialized metadata for "Pregel", "StateGraph", "CompiledStateGraph"
   - Returns "other" for regular LangChain chains (letting LangChain handler process them)

2. **LangChain handler** processes all LangChain callbacks:
   - Captures direct LangChain usage
   - Also captures nested LangChain calls within LangGraph nodes

3. **Result**: Clean event hierarchy
   ```
   graph.start (LangGraph handler)
     └─ agent.start node_a (LangGraph handler)
         ├─ llm.call.start (LangChain handler - nested)
         └─ llm.call.end (LangChain handler - nested)
     └─ agent.end node_a (LangGraph handler)
   graph.end (LangGraph handler)
   ```

### When to Use Which

| Scenario | LangChain | LangGraph | Both |
|----------|-----------|-----------|------|
| Simple LLM chains | ✅ | ❌ | ❌ |
| RAG pipelines | ✅ | ❌ | ❌ |
| Single agents | ✅ | ❌ | ❌ |
| StateGraph workflows | ❌ | ✅ | ✅ |
| LangGraph + LangChain LLMs | ❌ | ❌ | ✅ |
| Mixed applications | ❌ | ❌ | ✅ |

### Automatic Detection

When you call `arzule_ingest.init()`:
- Detects if LangChain is installed → instruments it
- Detects if LangGraph is installed → instruments it
- Both can be active simultaneously without conflicts

## Comparison with Other Integrations

| Feature | CrewAI | LangChain | LangGraph | AutoGen |
|---------|--------|-----------|-----------|---------|
| Auto-instrumentation | ✅ | ✅ | ✅ | ✅ |
| Graph/workflow tracking | ✅ | ❌ | ✅ | ✅ |
| Node/agent tracking | ✅ | ✅ | ✅ | ✅ |
| LLM call tracking | ✅ | ✅ | ✅ | ✅ |
| Tool tracking | ✅ | ✅ | ✅ | ✅ |
| State checkpointing | ❌ | ❌ | 🚧 | ❌ |
| Handoff detection | ✅ | ❌ | 🚧 | ✅ |
| Duplicate prevention | N/A | ✅ | ✅ | N/A |

## Implementation Notes

### Why Extend LangChain Callbacks?

LangGraph doesn't have its own callback system - it uses LangChain's `BaseCallbackHandler`. Our implementation:

1. Listens for `on_chain_start` / `on_chain_end` callbacks
2. Inspects the `serialized` metadata to determine if it's a graph or node
3. Maps appropriately to Arzule event types

### Node → Agent Mapping

In Arzule's data model, we map LangGraph nodes to "agents" because:
- Nodes are the primary execution units (like agents in CrewAI/AutoGen)
- Nodes perform specific tasks and can have their own identity
- The graph itself represents the workflow/orchestration layer

### Thread Safety

The handler uses:
- Thread-safe span tracking with `_run_spans` dictionary
- Cached run_id for fallback when ContextVars fail in spawned threads
- Lock-protected access to shared state

## Future Enhancements

1. **Enhanced checkpoint tracking** - Deep integration with checkpoint savers
2. **Subgraph visualization** - Better parent-child relationships for nested graphs
3. **Streaming optimization** - More efficient event batching for streaming mode
4. **Custom event support** - Allow users to emit custom events from nodes
5. **Memory tracking** - Capture long-term memory operations
6. **Time-travel debugging** - Leverage checkpoint data for replay

## Related Documentation

- [LangGraph Official Docs](https://docs.langchain.com/oss/python/langgraph/overview)
- [LangGraph GitHub](https://github.com/langchain-ai/langgraph)
- [Arzule SDK Documentation](./README.md)
- [LangChain Integration](./src/arzule_ingest/langchain/)

