"""Small local streaming graph used for exact-match fixed answers. The graph deliberately uses a local ``BaseChatModel`` implementation so LangGraph emits the same token/message stream events as a normal model-backed answer, while making no network or model-provider request. """ from __future__ import annotations import asyncio import math import re import uuid from collections.abc import AsyncIterator, Sequence from typing import Any from langchain_core.callbacks import AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import AIMessage, AIMessageChunk, BaseMessage from langchain_core.outputs import ChatGeneration, ChatGenerationChunk, ChatResult from langchain_core.runnables import RunnableConfig from langgraph.graph import END, START, StateGraph from deerflow.agents.thread_state import ThreadState _LOCAL_TOKEN_PATTERN = re.compile( r"\s*(?:[\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff]|\w{1,4}|[^\w\s])|\s+", re.UNICODE, ) _TARGET_CHUNKS_PER_SECOND = 10 _MAX_TOKENS_PER_CHUNK = 20 _MAX_TOKENS_PER_SECOND = 1_000 def _split_local_tokens(text: str) -> list[str]: """Split text into lossless, token-like units without loading model assets. Fixed answers only need a stable unit for stream pacing. Using a model tokenizer here is both model-specific and unsafe for offline deployments: ``tiktoken.get_encoding`` downloads its vocabulary when the host cache is empty. This local approximation treats each CJK character, short word fragment, or punctuation mark as one unit and attaches surrounding whitespace so the emitted chunks reconstruct the configured answer exactly. Limiting word fragments also prevents a long URL or identifier from bypassing the configured pace as one oversized unit. """ if not text: return [] tokens = _LOCAL_TOKEN_PATTERN.findall(text) if "".join(tokens) != text: # The expression is intentionally exhaustive, but keep a lossless # standard-library-only fallback for unusual Unicode input. return list(text) return tokens class FixedAnswerChatModel(BaseChatModel): """A local streaming model that only replays an administrator answer.""" answer: str message_id: str fixed_question_id: str tokens_per_second: int = 100 @property def _llm_type(self) -> str: return "fixed-answer" @property def _identifying_params(self) -> dict[str, Any]: return { "fixed_question_id": self.fixed_question_id, "tokens_per_second": self.tokens_per_second, } def _message(self) -> AIMessage: return AIMessage( content=self.answer, id=self.message_id, additional_kwargs={ "fixed_answer": True, "fixed_question_id": self.fixed_question_id, "tokens_per_second": self.tokens_per_second, }, ) def _generate( self, messages: list[BaseMessage], stop: list[str] | None = None, run_manager: CallbackManagerForLLMRun | None = None, **kwargs: Any, ) -> ChatResult: del messages, stop, run_manager, kwargs return ChatResult(generations=[ChatGeneration(message=self._message())]) async def _astream( self, messages: list[BaseMessage], stop: list[str] | None = None, run_manager: AsyncCallbackManagerForLLMRun | None = None, **kwargs: Any, ) -> AsyncIterator[ChatGenerationChunk]: del messages, stop, run_manager, kwargs answer = self.answer if not answer: yield ChatGenerationChunk(message=AIMessageChunk(content="", id=self.message_id)) return tokens = _split_local_tokens(answer) tokens_per_second = max( 1, min(_MAX_TOKENS_PER_SECOND, self.tokens_per_second), ) chunk_size = max( 1, min( _MAX_TOKENS_PER_CHUNK, math.ceil(tokens_per_second / _TARGET_CHUNKS_PER_SECOND), ), ) token_groups = [ tokens[:1], *(tokens[offset : offset + chunk_size] for offset in range(1, len(tokens), chunk_size)), ] for index, chunk_tokens in enumerate(token_groups): if index > 0: await asyncio.sleep(len(chunk_tokens) / tokens_per_second) content = "".join(chunk_tokens) chunk_kwargs = ( { "fixed_answer": True, "fixed_question_id": self.fixed_question_id, "tokens_per_second": tokens_per_second, } if index == 0 else {} ) yield ChatGenerationChunk( message=AIMessageChunk( content=content, id=self.message_id, additional_kwargs=chunk_kwargs, ) ) def _fallback_title(question: str) -> str: if len(question) > 50: return question[:50].rstrip() + "..." return question or "New Conversation" def make_fixed_answer_agent( *, question: str, answer: str, fixed_question_id: str, tokens_per_second: int = 100, ): """Build a one-node graph that streams and persists a fixed answer.""" tokens_per_second = max( 1, min(_MAX_TOKENS_PER_SECOND, int(tokens_per_second)), ) message_id = f"fixed-answer-{uuid.uuid4()}" model = FixedAnswerChatModel( answer=answer, message_id=message_id, fixed_question_id=fixed_question_id, tokens_per_second=tokens_per_second, ).with_config(tags=["fixed-answer"]) async def respond( state: ThreadState, config: RunnableConfig, ) -> dict[str, Any]: messages: Sequence[BaseMessage] = state.get("messages", []) # Calling the local model through ``astream`` makes the runtime publish # normal messages-mode chunks. The full AIMessage returned below is the # durable checkpoint value used on refresh/history reads. async for _chunk in model.astream(list(messages), config=config): pass update: dict[str, Any] = { "messages": [ AIMessage( content=answer, id=message_id, additional_kwargs={ "fixed_answer": True, "fixed_question_id": fixed_question_id, "tokens_per_second": tokens_per_second, }, ) ] } if not state.get("title"): update["title"] = _fallback_title(question) update["title_provisional"] = False elif state.get("title_provisional") is True: # A fixed-answer run must not trigger the worker's background LLM # title finalizer either; keep the already-visible local title. update["title_provisional"] = False return update graph = StateGraph(ThreadState) graph.add_node("fixed_answer", respond) graph.add_edge(START, "fixed_answer") graph.add_edge("fixed_answer", END) return graph.compile()