ag-ui-protocol / ag-ui-protocol/ag-ui

Vertex AI Gemini streaming does not work with LangGraphAgent - tokens buffered until completion

Abierto
#724 3 comentarios 0 reacciones 0 asignados Ver en GitHub
bug Integration
Lenguaje dominante
Python
Estrellas
15.9k
Forks
1.4k
Merge medio
1 d 17 h
PR fusionados (30 d)
163

Descripción

### Environment
- **ag-ui-langgraph**: 0.0.22
- **langchain-google-vertexai**: 3.0.0
- **langgraph**: 1.0.3+
- **Model**: Gemini 2.5 Pro via Vertex AI
- **Python**: 3.12

### Issue
Token-level streaming does not work when using `ChatVertexAI` (Gemini) with `LangGraphAgent`. All tokens are buffered and emitted only after the LLM completes generation, despite:
- Setting `streaming=True` on the LLM
- Using `await llm.ainvoke()` in the agent node
- Using `add_langgraph_fastapi_endpoint`

### Expected Behavior
Tokens should stream progressively as they are generated by the LLM, similar to how OpenAI models work with the same setup.

### Actual Behavior
- All tokens arrive in a single batch after ~3-6 seconds
- No `TEXT_MESSAGE_CHUNK` or `TEXT_MESSAGE_CONTENT` events are emitted during generation
- Only `MESSAGES_SNAPSHOT` is emitted after completion with the full response

### Reproduction

**Graph setup:**
```python
from langchain_google_vertexai import ChatVertexAI
from langgraph.graph import StateGraph, MessagesState

llm = ChatVertexAI(
model="gemini-2.5-pro",
streaming=True,
temperature=0.1
)

async def agent_node(state: MessagesState):
response = await llm.ainvoke(state["messages"])
return {"messages": [response]}

workflow = StateGraph(MessagesState)
workflow.add_node("agent", agent_node)
workflow.set_entry_point("agent")
workflow.add_edge("agent", END)

graph = workflow.compile(checkpointer=MemorySaver())
```

**Endpoint setup:**
```python
from ag_ui_langgraph import LangGraphAgent, add_langgraph_fastapi_endpoint

agent = LangGraphAgent(name="test_agent", graph=graph)
add_langgraph_fastapi_endpoint(app, agent, "/agent")
```

### Comparison with OpenAI
The **exact same code** works perfectly with OpenAI:
```python
from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini", streaming=True)
```
With OpenAI, tokens stream progressively and `TEXT_MESSAGE_CHUNK` events are emitted in real-time.

### Investigation
1. Direct testing with `llm.astream()` confirms Gemini **does** stream tokens properly
2. The issue appears to be with how `LangGraphAgent` uses `astream_events()` internally
3. Related to LangGraph issue: https://github.com/langchain-ai/langgraph/issues/4718 (streaming from subgraphs)
4. Custom events emitted via `adispatch_custom_event()` inside the node ARE captured, but not converted to text message events

### Questions
1. Is Vertex AI Gemini streaming officially supported with AG-UI + LangGraph?
2. Should custom events be automatically converted to text message events?
3. Is there a recommended way to enable token-level streaming with Vertex AI?
4. Should we manually intercept custom events and convert them, or is there a better approach?

Thank you!

Guía de contribución

Abrir la guía de contribución

Evaluación

Este issue todavía no se ha evaluado.

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.