212 lines
7.2 KiB
Python
212 lines
7.2 KiB
Python
"""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()
|