langgraph streaming vs llm streaming

[Anand-puthiyapurayil]

@dosu

i build a langgraph
where i `class ProductState(TypedDict):
question: str
k: int # Dynamic number of documents to retrieve.
context: List[Document] # Holds up to k retrieved product documents.
final_ids_output: str # The final output after generating the recommendation.
loop_step: int

-------------------------------------------------------------------

4) Retrieval: Retrieve k Products in One Call

-------------------------------------------------------------------

def retrieve_products(state: ProductState, faiss_store) -> ProductState:
"""
Uses the FAISS store retriever to fetch up to k product documents.
"""
try:
k = state.get("k", 200) # Use the dynamic k value (default to 200)
retriever = faiss_store.as_retriever(search_kwargs={'k': k})
docs = retriever.invoke(state["question"])
logger.info(f"Retrieved {len(docs)} product documents with k={k}.")
return {"context": docs}
except Exception as e:
logger.error(f"Error retrieving {k} product docs: {e}", exc_info=True)
return {"context": []}

-------------------------------------------------------------------

5) Generate Recommendation and List Retrieved Product IDs

-------------------------------------------------------------------

def generate_recommendation(state: ProductState, llm) -> ProductState:
docs = state.get("context", [])
if not docs:
logger.warning("No product documents retrieved for recommendation generation.")
state["final_ids_output"] = "No relevant product IDs found."
return state

top25 = docs[:25]
formatted_context = "\n".join(
    f"Product ID: {doc.metadata.get('product_id', 'Unknown')}, "
    f"Name: {doc.metadata.get('product_name', 'Unknown')}, "
    f"Supplier: {doc.metadata.get('Supplier', 'Unknown')}, "
    f"City: {doc.metadata.get('city', 'Unknown')}, "
    f"State: {doc.metadata.get('state', 'Unknown')}, "
    f"Country: {doc.metadata.get('country', 'Unknown')}, "
    f"Variant: {doc.metadata.get('variant', 'Unknown')}, "
    f"Category: {doc.metadata.get('category', 'Unknown')}"
    for doc in top25
)
prompt = custom_recommendation_prompt.format(
    context=formatted_context,
    question=state["question"]
)
# Call the LLM's stream method. Ensure no tool bindings interfere with streaming.
state["llm_stream"] = llm.stream([HumanMessage(content=prompt)])
return state

-------------------------------------------------------------------

6) Build LangGraph Workflow

-------------------------------------------------------------------

def build_product_graph(faiss_store, llm):
"""
Build a LangGraph pipeline with two nodes:
1. Retrieve k product documents.
2. Generate a recommendation and list the retrieved product IDs.
"""
graph_builder = StateGraph(ProductState)
graph_builder.add_node("retrieve", lambda s: retrieve_products(s, faiss_store))
graph_builder.add_node("generate_recommendation", lambda s: generate_recommendation(s, llm))
graph_builder.add_edge(START, "retrieve")
graph_builder.add_edge("retrieve", "generate_recommendation")
compiled_graph = graph_builder.compile()
logger.info("LangGraph product pipeline compiled successfully.")
return compiled_graph`

im trying to stream to the llm ouput 

its not happening fastapi swagger , but its happening in my terminal 

[dosubot[bot]]

I found a similar open issue related to streaming responses using FastAPI with LangChain. The suggested solution involves setting up a streaming response using Server-Sent Events (SSE) with FastAPI. Here's a brief overview of the approach:

1. Import Required Libraries: Use FastAPI and asyncio for handling asynchronous operations.
2. Create a FastAPI Instance: Set up your FastAPI application.
3. Define an Async Generator: This function will yield data for the streaming response.
4. Create a Streaming Endpoint: Use FastAPI's StreamingResponse to stream data from the generator.

Here's a simplified example of how you might implement this:

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import asyncio

app = FastAPI()

async def stream_messages():
    while True:
        yield "data: Some message\n\n"
        await asyncio.sleep(1)

@app.get("/stream")
async def stream(request: Request):
    async def event_generator():
        async for message in stream_messages():
            if await request.is_disconnected():
                break
            yield message

    return StreamingResponse(event_generator(), media_type="text/event-stream")

To integrate this with your LangGraph setup, you can modify the on_llm_end method in the AsyncIteratorCallbackHandler class to handle the streaming of LLM results. This involves processing the response and sending it as a message in the SSE format ^[1].

If you're experiencing issues with streaming in FastAPI's Swagger UI but not in the terminal, ensure that your FastAPI setup is correctly configured to handle streaming responses and that the client (Swagger UI) supports SSE.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[dosubot[bot]]

I wasn't able to find specific information within the LangChain repository on loading SBERT models and cross-encoders from MLflow to use with the Faiss wrapper. However, I can provide a general approach on how you might achieve this:

1. Load the Model from MLflow:
   Use the mlflow.pyfunc.load_model function to load your registered SBERT or cross-encoder model from MLflow.

   import mlflow
   
   model_uri = "models:/your_model_name/production"
   model = mlflow.pyfunc.load_model(model_uri)

2. Extract Embeddings:
   Once the model is loaded, use it to generate embeddings for your data. Ensure that the embeddings are normalized if you're using cosine similarity.

   def get_embeddings(texts):
       return model.predict(texts)

3. Integrate with Faiss:
   Use the embeddings to build a Faiss index. If necessary, normalize the embeddings before indexing.

   import faiss
   import numpy as np
   
   # Example: Normalize embeddings
   def normalize_embeddings(embeddings):
       norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
       return embeddings / norms
   
   # Create Faiss index
   embeddings = get_embeddings(your_text_data)
   normalized_embeddings = normalize_embeddings(embeddings)
   index = faiss.IndexFlatL2(normalized_embeddings.shape[1])
   index.add(normalized_embeddings)

4. Query the Faiss Index:
   Normalize the query embeddings and use the Faiss index to retrieve similar items.

   query_embeddings = get_embeddings(["your query text"])
   normalized_query_embeddings = normalize_embeddings(query_embeddings)
   distances, indices = index.search(normalized_query_embeddings, k=5)

This approach provides a general workflow for loading models from MLflow and using them with Faiss for retrieval tasks. Adjust the specifics based on your model and data requirements.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[Anand-puthiyapurayil]

@dosu

"""
product4.py — LangGraph-powered product assistant (streaming version)

✅ State-machine agent with SQLite checkpoints
✅ Hybrid retrieval tool
✅ Streams ONLY the assistant’s tokens to the console
"""

from future import annotations
import os, json, sqlite3, uuid, torch
from typing import Annotated, List, TypedDict
from typing_extensions import TypedDict # for Python < 3.11

from elasticsearch import Elasticsearch
from sentence_transformers import CrossEncoder

from langchain_core.documents import Document
from langchain_core.messages import AIMessageChunk, HumanMessage
from langchain_core.prompts import (
ChatPromptTemplate,
SystemMessagePromptTemplate,
MessagesPlaceholder,
)
from langchain_core.tools import tool
from langchain_ollama import ChatOllama
from langchain.retrievers import EnsembleRetriever

from langgraph.graph import StateGraph, START
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.checkpoint.sqlite import SqliteSaver

local helpers

from utils.utils import load_config, initialize_embeddings, load_faiss_store
from hybrid.custom_sparse_retriever import CustomSparseRetriever
from logger.logger import get_logger

─────────────────────────── config / setup ───────────────────────────

logger = get_logger(name)
HERE = os.path.dirname(os.path.abspath(file))
CFG = load_config(os.path.join(HERE, "..", "..", "config.yaml"))
device = "cuda" if torch.cuda.is_available() else "cpu"

SYSTEM_PROMPT = (
"You are a helpful product-sourcing assistant.\n"
"When the user asks about a product you MAY call the run_retrieval tool. "
"Otherwise answer normally.\n\n"
"If you call a tool respond with the tool JSON only; "
"otherwise respond with helpful text."
)

────────────────────────────── LLM ────────────────────────────────────

_BASE_LLM = ChatOllama(
model="llama3.1:8b",
stream=True,
temperature=0.2,
top_k=40,
repeat_penalty=1.2,
num_ctx=4096,
seed=42,
)

───────────────────── retrieval infrastructure ───────────────────────

embeddings = initialize_embeddings(CFG["output_dir"])
faiss_store = load_faiss_store(CFG["product_store_path"], embeddings)
es_client = Elasticsearch(CFG["elasticsearch_url"])
cross_enc = CrossEncoder(CFG["base_cross_encoder"], device=device)

@tool
def run_retrieval(query: str, top_k: int = 25) -> str:
"""Hybrid FAISS + BM25 search with cross-encoder re-rank."""
dense = faiss_store.as_retriever(search_kwargs={"k": top_k})
sparse = CustomSparseRetriever(es_client, CFG["product_index_name"], size=top_k)
dense_docs = dense.invoke(query)
sparse_docs = sparse.invoke(query)
logger.info("Dense %d • Sparse %d", len(dense_docs), len(sparse_docs))

docs: List[Document] = EnsembleRetriever(
    retrievers=[dense, sparse], weights=[0.5, 0.5]
).invoke(query)

if not docs:
    return json.dumps({"hits": [], "notice": "No close matches"})

scores = cross_enc.predict([(query, d.page_content) for d in docs])
docs   = [d for d, _ in sorted(zip(docs, scores), key=lambda t: t[1], reverse=True)]

hits = [
    {
        "product_id":   d.metadata.get("product_id"),
        "product_name": d.metadata.get("product_name"),
        "supplier":     d.metadata.get("Supplier"),
        "city":         d.metadata.get("city"),
        "state":        d.metadata.get("state"),
        "country":      d.metadata.get("country"),
        "category":     d.metadata.get("category"),
    }
    for d in docs[: top_k]
]
return json.dumps({"hits": hits})

TOOLS = [run_retrieval]
prompt = ChatPromptTemplate.from_messages(
[
SystemMessagePromptTemplate.from_template(SYSTEM_PROMPT),
MessagesPlaceholder(variable_name="messages"),
]
)
assistant = prompt | _BASE_LLM.bind_tools(TOOLS)

─────────────────────────── LangGraph ────────────────────────────────

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

graph_builder = StateGraph(ChatState)
graph_builder.add_node("agent", assistant)
graph_builder.add_node("tools", ToolNode(TOOLS))
graph_builder.add_conditional_edges("agent", tools_condition)
graph_builder.add_edge("tools", "agent")
graph_builder.add_edge(START, "agent")

DB_PATH = os.path.join(HERE, "assistant.db")
conn = sqlite3.connect(DB_PATH, check_same_thread=False)
graph = graph_builder.compile(checkpointer=SqliteSaver(conn))

───────────────────────────── CLI ────────────────────────────────────

def run_cli() -> None:
print("Product assistant ready — stream-enabled. (exit/quit to leave)\n")
cfg_obj = {"configurable": {"thread_id": str(uuid.uuid4())}}

while True:
    try:
        user = input("🧑‍💻  ").strip()
    except (EOFError, KeyboardInterrupt):
        print("\nbye!")
        break
    if user.lower() in {"exit", "quit"}:
        break

    stream = graph.stream(
        {"messages": [HumanMessage(content=user)]},
        cfg_obj,
        stream_mode="messages",
    )

    print("🤖 ", end="", flush=True)
    for chunk, meta in stream:
        # keep ONLY tokens from the agent node
        if meta.get("langgraph_node") != "agent":
            continue
        if isinstance(chunk, AIMessageChunk) and chunk.content:
            print(chunk.content, end="", flush=True)
    print()

if name == "main":
run_cli()
the stream is not working ? any idea?

[dosubot[bot]]

To troubleshoot the issue with streaming not working in your LangGraph-powered product assistant, ensure the following:

1. Stream Mode: Confirm that you are using stream_mode="messages" when calling the graph.stream method. This is necessary for streaming tokens from chat model invocations.

2. Version Requirement: Make sure you have langchain-core>=0.3.37 installed, as certain functions like message.text() require this version or higher.

3. Node Filtering: Ensure that you are filtering the output to the correct node. In your case, you are already filtering to the "agent" node, which is correct for streaming the assistant's tokens.

4. Console Output: Since the streaming works in the terminal but not elsewhere, verify that your environment (e.g., terminal settings or console) supports real-time streaming output. If using a different interface, ensure it can handle streaming data.

If these conditions are met and the issue persists, consider checking for any additional logs or errors that might provide more insight into the problem.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[Anand-puthiyapurayil]

@dosu

from future import annotations
import os
from langchain_core.messages import HumanMessage
from langgraph.prebuilt import create_react_agent
from langchain_ollama import ChatOllama

import your retrieval tool (must be decorated with @tool)

from hybrid.products.product4 import run_retrieval

def main() -> None:
# 1. Define the tools your agent can call
tools = [run_retrieval]

# 2. Instantiate your LLM with streaming enabled
llm = ChatOllama(
    model="llama3.1:8b",
    stream=True,
    temperature=0.2,
    top_k=40,
    repeat_penalty=1.2,
    num_ctx=4096,
    seed=42,
)

# 3. Create the agent with tool support
agent_executor = create_react_agent(llm, tools)

print("Chat Agent ready! Type 'exit' or 'quit' to stop.")
while True:
    user_input = input("\n> ").strip()
    if user_input.lower() in {"exit", "quit"}:
        break

    # Streaming Messages
    for step, metadata in agent_executor.stream(
        {"messages": [HumanMessage(content=user_input)]},
        stream_mode="messages",
    ):
        if metadata["langgraph_node"] == "agent" and (text := step.text()):
            print(text, end="|")
    print()  # newline after the full response

if name == "main":
main()
why the stream is not happening?

[dosubot[bot]]

To address the issue with streaming not working in your setup, consider the following points:

1. Ensure Correct Stream Mode: Verify that you are using stream_mode="messages" correctly in the agent_executor.stream method. This mode is necessary for streaming messages from the agent.

2. Check Node Filtering: Ensure that the filtering condition metadata["langgraph_node"] == "agent" is correct and matches the node you expect to stream from. If the node name is different, adjust the condition accordingly.

3. Verify LLM Configuration: Confirm that the ChatOllama model is correctly configured for streaming. The parameters like stream=True, temperature, top_k, and others should be set as intended.

4. Console Support: Ensure that your console or terminal supports real-time streaming output. If using a different interface, make sure it can handle streaming data.

5. Debugging: Add logging or print statements to check if the agent_executor.stream method is being called and if any steps are being yielded. This can help identify where the process might be failing.

If these checks do not resolve the issue, consider reviewing any additional logs or errors that might provide more insight into the problem.

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}