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

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

Đang mở
#724 3 bình luận 0 reaction 0 người được giao Xem trên GitHub
bug Integration
Ngôn ngữ chính
Python
Star
15.9k
Fork
1.4k
Merge trung bình
1 ngày 17 giờ
Pull request đã merge (30 ngày)
163

Mô tả

### 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!

Hướng dẫn đóng góp

Mở hướng dẫn đóng góp

Đánh giá

Issue này chưa được đánh giá.

Nhận issue mới trong hộp thư của bạn

Bản tóm tắt ngắn những issue GitHub phù hợp với người mới.