Streaming

"""A scripted, streaming fake chat model for the streaming examples."""
from collections.abc import Iterator
from typing import Any

from langchain_core.callbacks import CallbackManagerForLLMRun
from langchain_core.language_models.chat_models import BaseChatModel
from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage, ToolMessage
from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult


class ScriptedChat(BaseChatModel):
    """Calls search_docs first, then answers with the tool result, streaming word by word."""

    @property
    def _llm_type(self) -> str:
        return "scripted-fake"

    def bind_tools(self, tools, **kwargs):  # tools are ignored: the script decides
        return self

    def _reply(self, messages: list[BaseMessage]) -> AIMessage:
        last = messages[-1]
        if isinstance(last, ToolMessage):
            return AIMessage(content=f"According to the docs: {last.content}")
        return AIMessage(
            content="",
            tool_calls=[{"name": "search_docs", "args": {"query": "HNSW ef_search"}, "id": "call_1"}],
        )

    def _generate(self, messages, stop=None, run_manager=None, **kwargs) -> ChatResult:
        return ChatResult(generations=[ChatGeneration(message=self._reply(messages))])

    def _stream(self, messages: list[BaseMessage], stop=None,
                run_manager: CallbackManagerForLLMRun | None = None, **kwargs: Any) -> Iterator[ChatGenerationChunk]:
        reply = self._reply(messages)
        if reply.tool_calls:
            yield ChatGenerationChunk(message=AIMessageChunk(content="", tool_call_chunks=[
                {"name": "search_docs", "args": '{"query": "HNSW ef_search"}', "id": "call_1", "index": 0}]))
            return
        for word in reply.text.split(" "):
            yield ChatGenerationChunk(message=AIMessageChunk(content=word + " "))

from langchain.agents import create_agent
from langchain.tools import tool
from langgraph.config import get_stream_writer

from fake import ScriptedChat


@tool
def search_docs(query: str) -> str:
    """Search the documentation and return the best matching passage."""
    writer = get_stream_writer()
    writer({"status": "searching", "query": query})
    writer({"status": "found", "hits": 3})
    return "ef_search is the query-time candidate list size (default 40)."


agent = create_agent(ScriptedChat(), tools=[search_docs], system_prompt="You are DocsAssistant.")

import warnings
warnings.filterwarnings("ignore")
from agent import agent

inputs = {"messages": [{"role": "user", "content": "What does ef_search do?"}]}
stream = agent.stream_events(inputs, version="v3")
for kind, item in stream.interleave("messages", "tool_calls"):
    if kind == "messages":
        print("model:", str(item.text) or "(tool call)")
    else:
        print("tool: ", item.tool_name, item.input)
print("final:", stream.output["messages"][-1].text)

model: (tool call)
tool:  search_docs {'query': 'HNSW ef_search'}
model: According to the docs: ef_search is the query-time candidate list size (default 40).
final: According to the docs: ef_search is the query-time candidate list size (default 40).

import asyncio

from agent import agent


async def main() -> None:
    inputs = {"messages": [{"role": "user", "content": "What does ef_search do?"}]}
    async for event in agent.astream_events(inputs, version="v2"):
        kind = event["event"]
        if kind == "on_tool_start":
            print("tool start:", event["name"], event["data"]["input"])
        elif kind == "on_tool_end":
            print("tool end:  ", event["name"], event["data"]["output"].content)
        elif kind == "on_chat_model_stream" and event["data"]["chunk"].text:
            print(event["data"]["chunk"].text, end="|", flush=True)
    print()


asyncio.run(main())

tool start: search_docs {'query': 'HNSW ef_search'}
tool end:   search_docs ef_search is the query-time candidate list size (default 40).
According |to |the |docs: |ef_search |is |the |query-time |candidate |list |size |(default |40). |

import asyncio
from langchain_core.callbacks.manager import adispatch_custom_event
from langchain_core.runnables import RunnableLambda

async def step(x: str) -> str:
    await adispatch_custom_event("progress", {"stage": "retrieving"})
    return x.upper()

async def main():
    async for ev in RunnableLambda(step).astream_events("hi", version="v2"):
        print(ev["event"], ev["name"], ev["data"])
asyncio.run(main())

on_chain_start step {'input': 'hi'}
on_custom_event progress {'stage': 'retrieving'}
on_chain_stream step {'chunk': 'HI'}
on_chain_end step {'output': 'HI'}

from collections.abc import AsyncIterable

from fastapi import FastAPI
from fastapi.sse import EventSourceResponse, ServerSentEvent
from langchain.agents import create_agent
from langchain_core.messages import AIMessage
from langgraph.checkpoint.memory import InMemorySaver
from pydantic import BaseModel

from agent import search_docs
from fake import ScriptedChat  # in production: ChatOllama(model="llama3.1:8b")

agent = create_agent(ScriptedChat(), tools=[search_docs], checkpointer=InMemorySaver())
app = FastAPI()


class ChatRequest(BaseModel):
    message: str
    thread_id: str


@app.post("/chat", response_class=EventSourceResponse)
async def chat(req: ChatRequest) -> AsyncIterable[ServerSentEvent]:
    config = {"configurable": {"thread_id": req.thread_id}}
    inputs = {"messages": [{"role": "user", "content": req.message}]}
    try:
        async for part in agent.astream(
            inputs, config, stream_mode=["messages", "updates", "custom"], version="v2"
        ):
            if part["type"] == "messages":
                chunk, _meta = part["data"]
                if isinstance(chunk, AIMessage) and chunk.text:
                    yield ServerSentEvent(event="token", data={"text": chunk.text})
            elif part["type"] == "custom":
                yield ServerSentEvent(event="status", data=part["data"])
            elif part["type"] == "updates":
                for node, update in part["data"].items():
                    yield ServerSentEvent(event="step", data={"node": node, "messages": len(update["messages"])})
        yield ServerSentEvent(event="done", data={"thread_id": req.thread_id})
    except Exception as exc:  # tell the client, then close the stream cleanly
        yield ServerSentEvent(event="error", data={"type": type(exc).__name__})

event: step
data: {"node": "model", "messages": 1}

event: status
data: {"status": "searching", "query": "HNSW ef_search"}

event: status
data: {"status": "found", "hits": 3}

event: step
data: {"node": "tools", "messages": 1}

event: token
data: {"text": "According "}

event: token
data: {"text": "to "}

event: token
data: {"text": "the "}

event: token
data: {"text": "docs: "}

event: token
data: {"text": "ef_search "}

event: token
data: {"text": "is "}

event: token
data: {"text": "the "}

event: token
data: {"text": "query-time "}

event: token
data: {"text": "candidate "}

event: token
data: {"text": "list "}

event: token
data: {"text": "size "}

event: token
data: {"text": "(default "}

event: token
data: {"text": "40). "}

event: step
data: {"node": "model", "messages": 1}

event: done
data: {"thread_id": "t-1"}

import json

import httpx

answer = []
with httpx.stream(
    "POST",
    "http://localhost:8765/chat",
    json={"message": "What does ef_search do?", "thread_id": "t-2"},
    timeout=None,
) as response:
    event = None
    for line in response.iter_lines():
        if line.startswith("event:"):
            event = line.removeprefix("event:").strip()
        elif line.startswith("data:"):
            data = json.loads(line.removeprefix("data:"))
            if event == "token":
                answer.append(data["text"])
            else:
                print(f"{event:7} {data}")
print("answer:", "".join(answer))

step    {'node': 'model', 'messages': 1}
status  {'status': 'searching', 'query': 'HNSW ef_search'}
status  {'status': 'found', 'hits': 3}
step    {'node': 'tools', 'messages': 1}
step    {'node': 'model', 'messages': 1}
done    {'thread_id': 't-2'}
answer: According to the docs: ef_search is the query-time candidate list size (default 40).

event: error
data: {"type": "RuntimeError"}

export type ChatEvent =
  | { event: "token"; data: { text: string } }
  | { event: "status"; data: Record<string, unknown> }
  | { event: "step"; data: { node: string; messages: number } }
  | { event: "done"; data: { thread_id: string } }
  | { event: "error"; data: { type: string } };

/** POST to the FastAPI endpoint and yield its SSE events (EventSource cannot POST). */
export async function* chat(
  message: string,
  threadId: string,
  signal?: AbortSignal,
): AsyncGenerator<ChatEvent> {
  const res = await fetch("/chat", {
    method: "POST",
    headers: { "content-type": "application/json" },
    body: JSON.stringify({ message, thread_id: threadId }),
    signal,
  });
  if (!res.ok || !res.body) throw new Error(`HTTP ${res.status}`);

  const reader = res.body.pipeThrough(new TextDecoderStream()).getReader();
  let buffer = "";
  for (;;) {
    const { done, value } = await reader.read();
    if (done) return;
    buffer += value;
    let end: number;
    while ((end = buffer.indexOf("\n\n")) >= 0) {
      const block = buffer.slice(0, end);
      buffer = buffer.slice(end + 2);
      let event = "message";
      let data = "";
      for (const line of block.split("\n")) {
        if (line.startsWith("event:")) event = line.slice(6).trim();
        else if (line.startsWith("data:")) data += line.slice(5).trim();
      }
      if (data) yield { event, data: JSON.parse(data) } as ChatEvent;
    }
  }
}

step {"node":"model","messages":1}
status {"status":"searching","query":"HNSW ef_search"}
status {"status":"found","hits":3}
step {"node":"tools","messages":1}
step {"node":"model","messages":1}
done {"thread_id":"t-js"}
TEXT: According to the docs: ef_search is the query-time candidate list size (default 40).