feat(storage): add selectable SQLite and PostgreSQL persistence
- Select one storage backend per instance via AGENT_STORAGE, retaining memory as the backward-compatible default. Manage the selected async repository, official checkpointer and connection pools in lifespan; fail startup on unavailable or uninitialized storage without fallback. - Add SQLite and PostgreSQL adapters, schema initialization CLI, sample PostgreSQL compose configuration and a Windows-compatible serve entry. SQLite enforces single-instance file ownership; PostgreSQL uses a dedicated pool for session advisory locks plus version/epoch fencing. - Make repository IO and commit callbacks awaitable. Persist business fields, a complete committed graph snapshot and turn result/version in one business transaction before success. Rebuild execution checkpoints from that snapshot, discarding uncommitted progress after interruption. - Support optional Idempotency-Key on durable chat requests. Replay committed results and return 409 on changed request content; delete removes business/history data and replay records while retaining only an epoch/version tombstone to reject stale commits. - Document migration, database switching boundaries, backup/deployment expectations and crash recovery. Ignore generated database files. Validation: 315 tests passed with real SQLite and isolated PostgreSQL, including abrupt child-process exits before/after business commit; both backends passed CLI migration and actual HTTP server restart smoke. Existing API tests were rechecked after isolating them from user storage. An initial Windows CreateProcess permission failure did not recur in the full diagnostic run; failure assertions were not skipped or relaxed.
This commit is contained in:
+14
-2
@@ -1,4 +1,5 @@
|
||||
from fastapi import APIRouter, HTTPException, Depends
|
||||
from typing import Annotated
|
||||
from fastapi import APIRouter, HTTPException, Depends, Header
|
||||
from fastapi.responses import StreamingResponse
|
||||
from ..schemas.models import ProcessRequest_chat, ProcessResponse_chat, ProcessRequest_get, ProcessResponse_get, ProcessRequest_set, ProcessResponse_set, ProcessResponse_delete_session, ProcessRequest_delete_session
|
||||
from fastgpt_client import AsyncChatClient
|
||||
@@ -6,6 +7,7 @@ from ..backends.chat import (
|
||||
ChatBackend,
|
||||
ChatBackendAPIError,
|
||||
ChatBackendAuthenticationError,
|
||||
ChatBackendConflictError,
|
||||
ChatBackendRateLimitError,
|
||||
ChatInput,
|
||||
FormUpdate,
|
||||
@@ -156,7 +158,8 @@ def create_sse_event(event: str, data: dict) -> str:
|
||||
async def chat(
|
||||
request: ProcessRequest_chat,
|
||||
stream: bool = False,
|
||||
backend: ChatBackend = Depends(get_chat_backend)
|
||||
backend: ChatBackend = Depends(get_chat_backend),
|
||||
idempotency_key: Annotated[str | None, Header(alias="Idempotency-Key", min_length=1, max_length=128)] = None,
|
||||
):
|
||||
"""Handle chat completion request.
|
||||
|
||||
@@ -172,6 +175,7 @@ async def chat(
|
||||
text=json_data['text'],
|
||||
need_form_update=need_form_update,
|
||||
client_mode=client_mode,
|
||||
request_key=idempotency_key,
|
||||
)
|
||||
# 应用服务:langgraph 模式下包装字段注入/写回;fastgpt 模式透传。
|
||||
# 单测直调函数(未走 lifespan)时回退 backend 直调。
|
||||
@@ -419,6 +423,8 @@ async def chat(
|
||||
code, msg = "401", "认证失败"
|
||||
elif isinstance(e, ChatBackendRateLimitError):
|
||||
code, msg = "429", "请求过于频繁,请稍后重试"
|
||||
elif isinstance(e, ChatBackendConflictError):
|
||||
code, msg = "409", "请求幂等键或会话版本冲突"
|
||||
else:
|
||||
code, msg = "500", "大模型服务器无响应"
|
||||
yield create_sse_event("error", {"msg": msg, "code": code})
|
||||
@@ -458,6 +464,12 @@ async def chat(
|
||||
code="429",
|
||||
msg="请求过于频繁,请稍后重试"
|
||||
)
|
||||
except ChatBackendConflictError:
|
||||
return ProcessResponse_chat(
|
||||
sessionId=json_data['sessionId'], timeStamp=json_data['timeStamp'],
|
||||
outputText="", nextStage="", nextStageCode="", code="409",
|
||||
msg="请求幂等键或会话版本冲突",
|
||||
)
|
||||
except ChatBackendAPIError as e:
|
||||
logger.error(f"API error: {e}")
|
||||
return ProcessResponse_chat(
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""Backend-neutral chat contract."""
|
||||
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Protocol
|
||||
|
||||
@@ -18,7 +18,8 @@ class ChatInput:
|
||||
# FastGPT 后端忽略该字段。
|
||||
injected_fields: dict | None = None
|
||||
# 图提交后、公开结果之前同步业务字段;不进入图状态或 HTTP schema。
|
||||
on_commit: Callable[[dict], None] | None = field(default=None, repr=False, compare=False)
|
||||
on_commit: Callable[..., Awaitable[None] | None] | None = field(default=None, repr=False, compare=False)
|
||||
request_key: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -69,6 +70,10 @@ class ChatBackendError(Exception):
|
||||
"""Base error raised by a chat backend adapter."""
|
||||
|
||||
|
||||
class ChatBackendConflictError(ChatBackendError):
|
||||
"""请求幂等键或会话版本冲突。"""
|
||||
|
||||
|
||||
class ChatBackendAuthenticationError(ChatBackendError):
|
||||
"""The backend rejected its configured credentials."""
|
||||
|
||||
|
||||
@@ -15,6 +15,7 @@ test/backends/test_langgraph_streaming.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from contextlib import suppress
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
|
||||
@@ -103,6 +104,10 @@ class LangGraphBackend(ChatBackend):
|
||||
self._graph = graph if graph is not None else build_graph()
|
||||
self._llm = llm if llm is not None else workflow_fake_llm()
|
||||
|
||||
async def reset_checkpoint(self, session_id):
|
||||
from ..storage.checkpoints import delete_checkpoint
|
||||
await delete_checkpoint(self._graph, session_id)
|
||||
|
||||
async def complete(self, chat_input: ChatInput) -> ChatResult:
|
||||
"""非流式:执行一轮图并翻译结果。
|
||||
|
||||
@@ -135,15 +140,13 @@ class LangGraphBackend(ChatBackend):
|
||||
LLMAuthError -> ChatBackendAuthenticationError(API 层 401);
|
||||
LLMRateLimitError -> ChatBackendRateLimitError(API 层 429)。
|
||||
"""
|
||||
graph_input: dict = {
|
||||
graph_input: dict = dict(chat_input.injected_fields or {})
|
||||
# 恢复快照中的上一轮 input_text 不得覆盖本次请求。
|
||||
graph_input.update({
|
||||
"session_id": chat_input.session_id,
|
||||
"input_text": chat_input.text,
|
||||
"client_mode": chat_input.client_mode,
|
||||
}
|
||||
# 作业 6:应用服务注入的最新业务字段覆盖本轮输入,
|
||||
# 使 set_info 写入后本轮 chat 立即可见(输入覆盖 checkpoint 同键)。
|
||||
if chat_input.injected_fields:
|
||||
graph_input.update(chat_input.injected_fields)
|
||||
})
|
||||
try:
|
||||
result = await self._graph.ainvoke(
|
||||
graph_input,
|
||||
@@ -152,7 +155,9 @@ class LangGraphBackend(ChatBackend):
|
||||
)
|
||||
# 即使消费者在图完成后断开,业务 repository 仍与 checkpoint 一致。
|
||||
if chat_input.on_commit is not None:
|
||||
chat_input.on_commit(result.get("turn_patch") or {})
|
||||
pending = chat_input.on_commit(result.get("turn_patch") or {}, state=result)
|
||||
if inspect.isawaitable(pending):
|
||||
await pending
|
||||
return result
|
||||
except LLMAuthError as exc:
|
||||
raise ChatBackendAuthenticationError(str(exc)) from exc
|
||||
|
||||
@@ -19,6 +19,14 @@ class Config:
|
||||
# "fastgpt" 保持现有行为;"langgraph" 使用内存图 fake 回复。
|
||||
AGENT_BACKEND = os.getenv("AGENT_BACKEND", "fastgpt")
|
||||
|
||||
# 每个实例只选一个会话存储;与通用 DATABASE_URL 独立。
|
||||
AGENT_STORAGE = os.getenv("AGENT_STORAGE", "memory").strip().lower()
|
||||
AGENT_SQLITE_PATH = os.getenv("AGENT_SQLITE_PATH", "./data/agent.sqlite3")
|
||||
AGENT_POSTGRES_DSN = os.getenv("AGENT_POSTGRES_DSN", "")
|
||||
AGENT_DB_AUTO_MIGRATE = os.getenv("AGENT_DB_AUTO_MIGRATE", "false").lower() in ("true", "1")
|
||||
AGENT_DB_POOL_SIZE = int(os.getenv("AGENT_DB_POOL_SIZE", "10"))
|
||||
AGENT_DB_LOCK_TIMEOUT_SECONDS = float(os.getenv("AGENT_DB_LOCK_TIMEOUT_SECONDS", "30"))
|
||||
|
||||
# 命名 LLM 提供商列表(均为 OpenAI 兼容接口),两种写法二选一:
|
||||
# ① JSON 字典(推荐,单变量管理):
|
||||
# LLM_PROVIDERS={"deepseek": {"base_url": "https://...", "api_key": "...",
|
||||
@@ -138,6 +146,17 @@ class Config:
|
||||
raise ValueError(
|
||||
f"AGENT_BACKEND must be 'fastgpt' or 'langgraph', got {cls.AGENT_BACKEND!r}"
|
||||
)
|
||||
if cls.AGENT_BACKEND == "langgraph":
|
||||
if cls.AGENT_STORAGE not in ("memory", "sqlite", "postgresql"):
|
||||
raise ValueError("AGENT_STORAGE must be memory, sqlite or postgresql")
|
||||
if cls.AGENT_STORAGE == "sqlite" and (
|
||||
not cls.AGENT_SQLITE_PATH.strip() or cls.AGENT_SQLITE_PATH == ":memory:"
|
||||
):
|
||||
raise ValueError("AGENT_SQLITE_PATH must be a persistent file path")
|
||||
if cls.AGENT_STORAGE == "postgresql" and not cls.AGENT_POSTGRES_DSN.strip():
|
||||
raise ValueError("AGENT_POSTGRES_DSN is required for postgresql storage")
|
||||
if cls.AGENT_DB_POOL_SIZE < 1 or cls.AGENT_DB_LOCK_TIMEOUT_SECONDS <= 0:
|
||||
raise ValueError("Database pool size and lock timeout must be positive")
|
||||
profiles = cls.llm_profiles() # JSON 非法时在此抛出明确错误
|
||||
incomplete = [n for n in cls._declared_provider_names() if n not in profiles]
|
||||
if incomplete:
|
||||
|
||||
+38
-66
@@ -1,107 +1,79 @@
|
||||
"""FastGPT client and chat backend dependency injection."""
|
||||
from contextlib import asynccontextmanager
|
||||
"""后端与单一会话存储的生命周期、依赖注入。"""
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastgpt_client import AsyncChatClient
|
||||
|
||||
from ..agent.graph import build_graph
|
||||
from ..backends.chat import ChatBackend
|
||||
from ..backends.fastgpt import FastGPTBackend
|
||||
from ..backends.langgraph import LangGraphBackend
|
||||
from ..services.app_service import AccidentAppService
|
||||
from ..services.repository import InMemorySessionRepository
|
||||
from ..storage.checkpoints import delete_checkpoint
|
||||
from .config import Config
|
||||
|
||||
|
||||
_fastgpt_client: AsyncChatClient | None = None
|
||||
_chat_backend: ChatBackend | None = None
|
||||
_langgraph_graph = None
|
||||
_app_service: AccidentAppService | None = None
|
||||
|
||||
# 两后端缺省 60 秒、最多额外重试 3 次,1/2/4 秒退避。
|
||||
# FastGPT SDK 的 max_retries 实为总尝试数,注入时转换为 1 + 重试次数。
|
||||
# LLM 关闭 SDK 自动重试;流式发出文本后不再重试。
|
||||
# LLM 预算按单次模型调用计算,FastGPT timeout 为 HTTP SDK 超时,
|
||||
# 均不承诺多次模型调用的整个聊天回合具有相同总期限。
|
||||
FASTGPT_TIMEOUT_SECONDS = 60.0
|
||||
FASTGPT_MAX_RETRIES = 3
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
"""Manage chat backend resource lifecycle."""
|
||||
global _chat_backend, _fastgpt_client
|
||||
global _chat_backend, _fastgpt_client, _langgraph_graph, _app_service
|
||||
Config.validate()
|
||||
llm = None
|
||||
async with AsyncExitStack() as resources:
|
||||
try:
|
||||
if Config.AGENT_BACKEND == 'langgraph':
|
||||
from ..agent.llm import build_llm_service
|
||||
from ..storage.runtime import StorageSettings, open_storage
|
||||
|
||||
if Config.AGENT_BACKEND == "langgraph":
|
||||
# 作业 2:LangGraph fake 模式不要求 FastGPT 密钥。
|
||||
# TODO(作业 2): 若语音仍启用,需在此单独初始化并校验 FastGPT 配置。
|
||||
from ..agent.llm import build_llm_service
|
||||
|
||||
global _langgraph_graph, _app_service
|
||||
_langgraph_graph = build_graph()
|
||||
llm = build_llm_service()
|
||||
_chat_backend = LangGraphBackend(_langgraph_graph, llm=llm)
|
||||
# 作业 6:应用服务 + 内存 repository(作业 7 换 PostgreSQL)。
|
||||
_app_service = AccidentAppService(
|
||||
_chat_backend, InMemorySessionRepository()
|
||||
)
|
||||
else:
|
||||
_fastgpt_client = AsyncChatClient(
|
||||
api_key=Config.FASTGPT_API_KEY,
|
||||
base_url=Config.FASTGPT_BASE_URL,
|
||||
timeout=FASTGPT_TIMEOUT_SECONDS,
|
||||
max_retries=FASTGPT_MAX_RETRIES + 1,
|
||||
retry_delay=1.0,
|
||||
enable_logging=Config.DEBUG,
|
||||
)
|
||||
await _fastgpt_client.__aenter__()
|
||||
_chat_backend = FastGPTBackend(_fastgpt_client)
|
||||
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
_chat_backend = None
|
||||
_langgraph_graph = None
|
||||
_app_service = None
|
||||
if llm is not None and hasattr(llm, "aclose"):
|
||||
await llm.aclose()
|
||||
if _fastgpt_client:
|
||||
await _fastgpt_client.__aexit__(None, None, None)
|
||||
_fastgpt_client = None
|
||||
storage = await resources.enter_async_context(open_storage(StorageSettings.from_config(Config)))
|
||||
_langgraph_graph = build_graph(storage.checkpointer)
|
||||
llm = build_llm_service()
|
||||
_chat_backend = LangGraphBackend(_langgraph_graph, llm=llm)
|
||||
_app_service = AccidentAppService(_chat_backend, storage.repository, locks=storage.locks)
|
||||
else:
|
||||
_fastgpt_client = AsyncChatClient(
|
||||
api_key=Config.FASTGPT_API_KEY, base_url=Config.FASTGPT_BASE_URL,
|
||||
timeout=FASTGPT_TIMEOUT_SECONDS, max_retries=FASTGPT_MAX_RETRIES + 1,
|
||||
retry_delay=1.0, enable_logging=Config.DEBUG,
|
||||
)
|
||||
await _fastgpt_client.__aenter__()
|
||||
_chat_backend = FastGPTBackend(_fastgpt_client)
|
||||
yield
|
||||
finally:
|
||||
_chat_backend = None
|
||||
_langgraph_graph = None
|
||||
_app_service = None
|
||||
if llm is not None and hasattr(llm, 'aclose'):
|
||||
await llm.aclose()
|
||||
if _fastgpt_client:
|
||||
await _fastgpt_client.__aexit__(None, None, None)
|
||||
_fastgpt_client = None
|
||||
|
||||
|
||||
async def delete_chat_session(thread_id: str) -> None:
|
||||
"""删除图 checkpoint;缺少清理能力或存储出错必须报告失败。"""
|
||||
if _langgraph_graph is None:
|
||||
raise RuntimeError("LangGraph backend not initialized")
|
||||
checkpointer = getattr(_langgraph_graph, "checkpointer", None)
|
||||
async_delete = getattr(checkpointer, "adelete_thread", None)
|
||||
if callable(async_delete):
|
||||
await async_delete(thread_id)
|
||||
return
|
||||
sync_delete = getattr(checkpointer, "delete_thread", None)
|
||||
if callable(sync_delete):
|
||||
sync_delete(thread_id)
|
||||
return
|
||||
raise RuntimeError("LangGraph checkpointer does not support thread deletion")
|
||||
await delete_checkpoint(_langgraph_graph, thread_id)
|
||||
|
||||
|
||||
def get_fastgpt_client() -> AsyncChatClient:
|
||||
"""Get the initialized FastGPT client."""
|
||||
if _fastgpt_client is None:
|
||||
raise RuntimeError("FastGPT client not initialized")
|
||||
raise RuntimeError('FastGPT client not initialized')
|
||||
return _fastgpt_client
|
||||
|
||||
|
||||
def get_app_service() -> AccidentAppService:
|
||||
"""Get the initialized application service(langgraph 模式)。"""
|
||||
if _app_service is None:
|
||||
raise RuntimeError("Application service not initialized")
|
||||
raise RuntimeError('Application service not initialized')
|
||||
return _app_service
|
||||
|
||||
|
||||
def get_chat_backend() -> ChatBackend:
|
||||
"""Get the backend-neutral chat service."""
|
||||
if _chat_backend is None:
|
||||
raise RuntimeError("Chat backend not initialized")
|
||||
raise RuntimeError('Chat backend not initialized')
|
||||
return _chat_backend
|
||||
|
||||
@@ -0,0 +1,23 @@
|
||||
"""Portable single-worker entry point: python -m src.serve --port 8000."""
|
||||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
|
||||
import uvicorn
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument('--host', default='127.0.0.1')
|
||||
parser.add_argument('--port', default=8000, type=int)
|
||||
args = parser.parse_args()
|
||||
# psycopg 不支持 Windows ProactorEventLoop。直接运行 Server.serve
|
||||
# 使用已创建的 Selector loop,不被 uvicorn 默认 loop factory 覆盖。
|
||||
factory = asyncio.SelectorEventLoop if os.name == 'nt' else None
|
||||
config = uvicorn.Config('src.main:app', host=args.host, port=args.port, loop='asyncio')
|
||||
with asyncio.Runner(loop_factory=factory) as runner:
|
||||
runner.run(uvicorn.Server(config).serve())
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
+131
-167
@@ -1,65 +1,41 @@
|
||||
"""事故处理应用服务(作业 6-1):四业务接口共享同一份业务状态。
|
||||
"""四接口应用服务:同会话串行,持久化模式按已提交快照恢复。"""
|
||||
|
||||
职责边界(对应计划作业 6):
|
||||
- set_info / get_info / delete_session **不调用 LLM、不创建辅助对话**
|
||||
(对照旧 FastGPT 实现的“辅助对话 + 删记录”hack);
|
||||
- chat 每轮从 repository 读最新业务字段注入图(set_info 写入后
|
||||
下一轮 chat 必须可见),图产出字段差量后写回 repository;
|
||||
- chat / stream / set_info / get_info / delete_session 都在同会话锁内。
|
||||
锁的引用计数包含等待者,取消和异常都会归还引用;没有使用者时回收。
|
||||
不同会话独立执行,同会话 set_info 不会被在途聊天的旧结果覆盖。
|
||||
- 未知 key 兼容策略:get 对未知 key 返回 JSON 空串 '""'(现状兼容);
|
||||
set 对未知 key 拒绝写入(白名单保护,差异见兼容清单);
|
||||
- includeInputInfo 继续被 endpoint 接受,但本层不再使用
|
||||
——新后端无需触发辅助对话。
|
||||
import hashlib
|
||||
import json
|
||||
from contextlib import aclosing, asynccontextmanager, suppress
|
||||
from dataclasses import dataclass, replace
|
||||
|
||||
endpoint 只负责校验、协议(SSE / 响应模型)与依赖注入。
|
||||
"""
|
||||
|
||||
from contextlib import aclosing
|
||||
from dataclasses import replace
|
||||
|
||||
from ..backends.chat import ChatBackend, ChatInput, ChatResult, FormUpdate
|
||||
from ..backends.chat import ChatBackendConflictError, ChatInput, ChatResult, FormUpdate, TextDelta
|
||||
from .fields_registry import (
|
||||
FIELD_REGISTRY,
|
||||
GROUP_KEYS,
|
||||
BUSINESS_TO_FORM_KEY,
|
||||
FieldGroup,
|
||||
business_patch_from_form,
|
||||
business_value,
|
||||
canonical_key,
|
||||
parse_field_value,
|
||||
to_external_string,
|
||||
FIELD_REGISTRY, GROUP_KEYS, BUSINESS_TO_FORM_KEY, FieldGroup,
|
||||
business_patch_from_form, business_value, canonical_key,
|
||||
parse_field_value, to_external_string,
|
||||
)
|
||||
from .repository import CheckpointNotFoundError, SessionRepository
|
||||
from .repository import CheckpointNotFoundError
|
||||
from .session_locks import SessionLocks
|
||||
|
||||
# set_info 对未知/受保护 key 的拒绝码(与现有“入参不规范”文案对齐)
|
||||
SET_INFO_INVALID_KEY_MSG = "入参不规范:未知或受保护的字段 key"
|
||||
SET_INFO_INVALID_KEY_MSG = '入参不规范:未知或受保护的字段 key'
|
||||
|
||||
|
||||
@dataclass
|
||||
class _Turn:
|
||||
injected: dict
|
||||
ticket: object = None
|
||||
committed: bool = False
|
||||
|
||||
|
||||
class AccidentAppService:
|
||||
"""四个业务接口的应用服务(langgraph 后端)。"""
|
||||
|
||||
def __init__(self, backend: ChatBackend, repository: SessionRepository):
|
||||
def __init__(self, backend, repository, *, locks=None):
|
||||
self._backend = backend
|
||||
self._repository = repository
|
||||
self._locks = SessionLocks()
|
||||
self._locks = locks if locks is not None else SessionLocks()
|
||||
self._durable = hasattr(repository, 'begin_turn')
|
||||
|
||||
# ---- chat:图每轮读取最新业务状态,产出差量写回 ----
|
||||
|
||||
def _injected_fields(self, session_id: str) -> dict:
|
||||
"""repository 业务字段 -> 图输入的分组注入。
|
||||
|
||||
acd 业务键经 BUSINESS_TO_FORM_KEY 映射回图表单键
|
||||
(无表单来源的 bjrjs/ywmtc 不注入);人员业务键与图的
|
||||
person_fields 同构,原样注入。
|
||||
"""
|
||||
business = self._repository.load(session_id)
|
||||
async def _injected_fields(self, session_id):
|
||||
business = await self._repository.load(session_id)
|
||||
if not business:
|
||||
return {}
|
||||
injected: dict = {"accident_fields": {}, "person_fields": {}}
|
||||
# 注入同义键的一致视图,覆盖 checkpoint 里的旧值。
|
||||
injected = {'accident_fields': {}, 'person_fields': {}}
|
||||
for key, spec in FIELD_REGISTRY.items():
|
||||
value = business_value(business, key)
|
||||
if value is None:
|
||||
@@ -67,141 +43,129 @@ class AccidentAppService:
|
||||
if spec.group is FieldGroup.ACD:
|
||||
form_key = BUSINESS_TO_FORM_KEY.get(key)
|
||||
if form_key is not None:
|
||||
injected["accident_fields"][form_key] = value
|
||||
injected['accident_fields'][form_key] = value
|
||||
else:
|
||||
injected["person_fields"][key] = value
|
||||
injected['person_fields'][key] = value
|
||||
return injected
|
||||
|
||||
def _writeback_patch(self, session_id: str, patch: dict) -> None:
|
||||
"""图产出的差量(表单键/人员键)经映射落 repository 业务键。"""
|
||||
async def _writeback_patch(self, session_id, patch):
|
||||
business_patch = business_patch_from_form(patch)
|
||||
if business_patch:
|
||||
self._repository.apply_update(session_id, business_patch)
|
||||
await self._repository.apply_update(session_id, business_patch)
|
||||
|
||||
async def chat(self, chat_input: ChatInput) -> ChatResult:
|
||||
"""非流式一轮对话:注入最新业务字段 -> 执行图 -> 写回字段差量。"""
|
||||
async with self._locks.hold(chat_input.session_id):
|
||||
return await self._chat_locked(chat_input)
|
||||
|
||||
async def _chat_locked(self, chat_input: ChatInput) -> ChatResult:
|
||||
injected = self._injected_fields(chat_input.session_id)
|
||||
committed = False
|
||||
|
||||
def commit(patch):
|
||||
nonlocal committed
|
||||
self._writeback_patch(chat_input.session_id, patch)
|
||||
committed = True
|
||||
|
||||
result = await self._backend.complete(
|
||||
replace(chat_input, injected_fields=injected or None, on_commit=commit)
|
||||
)
|
||||
if not committed:
|
||||
patch = result.committed_patch
|
||||
if patch is None:
|
||||
patch = result.form_update
|
||||
if isinstance(patch, dict):
|
||||
commit(patch)
|
||||
return replace(
|
||||
result, form_update=result.form_update if chat_input.need_form_update else {},
|
||||
)
|
||||
|
||||
async def stream(self, chat_input: ChatInput):
|
||||
"""流式代理:字段注入同 chat();FormUpdate 事件捕获写回后透传。"""
|
||||
async with self._locks.hold(chat_input.session_id):
|
||||
async with aclosing(self._stream_locked(chat_input)) as events:
|
||||
async for event in events:
|
||||
yield event
|
||||
|
||||
async def _stream_locked(self, chat_input: ChatInput):
|
||||
injected = self._injected_fields(chat_input.session_id)
|
||||
committed = False
|
||||
|
||||
def commit(patch):
|
||||
nonlocal committed
|
||||
self._writeback_patch(chat_input.session_id, patch)
|
||||
committed = True
|
||||
|
||||
result_input = replace(chat_input, injected_fields=injected or None, on_commit=commit)
|
||||
async with aclosing(self._backend.stream(result_input)) as events:
|
||||
async for event in events:
|
||||
if isinstance(event, FormUpdate):
|
||||
if not committed and isinstance(event.data, dict):
|
||||
commit(event.data)
|
||||
if not chat_input.need_form_update:
|
||||
continue
|
||||
yield event
|
||||
|
||||
# ---- set_info / get_info:零 LLM、零辅助对话 ----
|
||||
|
||||
async def set_info(self, session_id: str, key: str, value: str) -> tuple[str, str]:
|
||||
async with self._locks.hold(session_id):
|
||||
return self._set_info(session_id, key, value)
|
||||
|
||||
def _set_info(self, session_id: str, key: str, value: str) -> tuple[str, str]:
|
||||
"""白名单校验后写入 repository。返回 (code, msg)。"""
|
||||
spec = FIELD_REGISTRY.get(key)
|
||||
if spec is None or not spec.writable:
|
||||
return "500", SET_INFO_INVALID_KEY_MSG
|
||||
@asynccontextmanager
|
||||
async def _turn(self, request):
|
||||
ticket = None
|
||||
turn = None
|
||||
try:
|
||||
parsed = parse_field_value(spec, value)
|
||||
except ValueError:
|
||||
return "500", "入参不规范:字段值无法解析"
|
||||
self._repository.apply_update(session_id, {canonical_key(key): parsed})
|
||||
return "200", ""
|
||||
if self._durable:
|
||||
payload = {'text': request.text, 'client_mode': request.client_mode,
|
||||
'need_form_update': request.need_form_update}
|
||||
digest = hashlib.sha256(json.dumps(payload, sort_keys=True, ensure_ascii=False).encode()).hexdigest()
|
||||
ticket = await self._repository.begin_turn(request.session_id, request.request_key, digest)
|
||||
if ticket.replay is None:
|
||||
# 官方 checkpoint 可能先写成功,业务事务却未提交。
|
||||
# 丢弃该执行进度;始终从业务表的提交快照恢复。
|
||||
await self._backend.reset_checkpoint(request.session_id)
|
||||
elif request.request_key:
|
||||
raise ChatBackendConflictError('Idempotency-Key requires persistent storage')
|
||||
injected = dict(ticket.graph_state) if ticket else {}
|
||||
injected.update(await self._injected_fields(request.session_id))
|
||||
if ticket:
|
||||
injected['state_version'] = ticket.version
|
||||
turn = _Turn(injected, ticket)
|
||||
yield turn
|
||||
except BaseException:
|
||||
if ticket is not None and ticket.replay is None and not (turn and turn.committed):
|
||||
# 清理失败时保留 pending;下次持锁访问会标记 aborted 并恢复快照。
|
||||
with suppress(Exception):
|
||||
await self._repository.abort_turn(ticket)
|
||||
raise
|
||||
|
||||
async def get_info(self, session_id: str, key: str) -> str:
|
||||
async def _commit(self, request, turn, patch, state):
|
||||
if turn.ticket:
|
||||
if state is None:
|
||||
raise RuntimeError('Persistent turns require a complete graph state')
|
||||
await self._repository.commit_turn(turn.ticket, business_patch_from_form(patch), state)
|
||||
else:
|
||||
await self._writeback_patch(request.session_id, patch)
|
||||
turn.committed = True
|
||||
|
||||
@staticmethod
|
||||
def _replay(request, state):
|
||||
patch = state.get('turn_patch') or {}
|
||||
return ChatResult(content=state['response_text'], status_code=state.get('stage_code'),
|
||||
form_update=patch if request.need_form_update else {}, committed_patch=patch)
|
||||
|
||||
async def chat(self, request: ChatInput):
|
||||
async with self._locks.hold(request.session_id):
|
||||
async with self._turn(request) as turn:
|
||||
if turn.ticket and turn.ticket.replay is not None:
|
||||
return self._replay(request, turn.ticket.replay)
|
||||
async def commit(patch, *, state=None):
|
||||
await self._commit(request, turn, patch, state)
|
||||
result = await self._backend.complete(
|
||||
replace(request, injected_fields=turn.injected or None, on_commit=commit))
|
||||
if not turn.committed:
|
||||
patch = result.committed_patch if result.committed_patch is not None else result.form_update
|
||||
await commit(patch if isinstance(patch, dict) else {})
|
||||
return replace(result, form_update=result.form_update if request.need_form_update else {})
|
||||
|
||||
async def stream(self, request: ChatInput):
|
||||
async with self._locks.hold(request.session_id):
|
||||
async with self._turn(request) as turn:
|
||||
if turn.ticket and turn.ticket.replay is not None:
|
||||
result = self._replay(request, turn.ticket.replay)
|
||||
yield TextDelta(result.content)
|
||||
if result.form_update:
|
||||
yield FormUpdate(result.form_update)
|
||||
return
|
||||
async def commit(patch, *, state=None):
|
||||
await self._commit(request, turn, patch, state)
|
||||
forwarded = replace(request, injected_fields=turn.injected or None, on_commit=commit)
|
||||
async with aclosing(self._backend.stream(forwarded)) as events:
|
||||
async for event in events:
|
||||
if isinstance(event, FormUpdate):
|
||||
if not turn.committed and isinstance(event.data, dict):
|
||||
await commit(event.data)
|
||||
if not request.need_form_update:
|
||||
continue
|
||||
yield event
|
||||
|
||||
async def set_info(self, session_id, key, value):
|
||||
async with self._locks.hold(session_id):
|
||||
return self._get_info(session_id, key)
|
||||
spec = FIELD_REGISTRY.get(key)
|
||||
if spec is None or not spec.writable:
|
||||
return '500', SET_INFO_INVALID_KEY_MSG
|
||||
try:
|
||||
parsed = parse_field_value(spec, value)
|
||||
except ValueError:
|
||||
return '500', '入参不规范:字段值无法解析'
|
||||
await self._repository.apply_update(session_id, {canonical_key(key): parsed})
|
||||
return '200', ''
|
||||
|
||||
def _get_info(self, session_id: str, key: str) -> str:
|
||||
"""按分组或单键读出,返回 JSON 字符串(兼容现有 /get_info 契约)。"""
|
||||
import json
|
||||
async def get_info(self, session_id, key):
|
||||
async with self._locks.hold(session_id):
|
||||
business = await self._repository.load(session_id)
|
||||
def group_value(group):
|
||||
return {k: to_external_string(business_value(business, k)) for k in GROUP_KEYS[group]}
|
||||
if key == 'all':
|
||||
value = {group.value: group_value(group) for group in FieldGroup}
|
||||
else:
|
||||
try:
|
||||
value = group_value(FieldGroup(key))
|
||||
except ValueError:
|
||||
value = to_external_string(business_value(business, key)) if key in FIELD_REGISTRY else ''
|
||||
return json.dumps(value, ensure_ascii=False)
|
||||
|
||||
business = self._repository.load(session_id)
|
||||
if key == "all":
|
||||
return json.dumps(
|
||||
{
|
||||
group.value: {
|
||||
k: to_external_string(business_value(business, k))
|
||||
for k in GROUP_KEYS[group]
|
||||
}
|
||||
for group in FieldGroup
|
||||
},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
try:
|
||||
group = FieldGroup(key)
|
||||
except ValueError:
|
||||
group = None
|
||||
if group is not None:
|
||||
return json.dumps(
|
||||
{k: to_external_string(business_value(business, k)) for k in GROUP_KEYS[group]},
|
||||
ensure_ascii=False,
|
||||
)
|
||||
# 单键:未知 key 返回 JSON 空串(兼容现状);已知 key 返回其外部表示
|
||||
if key not in FIELD_REGISTRY:
|
||||
return json.dumps("", ensure_ascii=False)
|
||||
return json.dumps(to_external_string(business_value(business, key)), ensure_ascii=False)
|
||||
|
||||
# ---- delete_session:业务状态 + 图 checkpoint 一起清理 ----
|
||||
|
||||
async def delete_session(self, session_id: str, drop_checkpoint) -> None:
|
||||
"""删除会话全部数据。
|
||||
|
||||
drop_checkpoint:由 core 层注入的 checkpoint 清理函数
|
||||
(图生命周期归 lifespan 管理,本层不持有图引用)。
|
||||
与会话锁互斥:删除不与在途 chat/set/get 交错。
|
||||
释放锁引用时保留等待者使用的锁,最后一个使用者退出后回收。
|
||||
checkpoint 清理成功(或明确不存在)后才删除业务字段。
|
||||
其他异常向上传播,由 endpoint 返回失败,保留业务字段供重试。
|
||||
"""
|
||||
async def delete_session(self, session_id, drop_checkpoint):
|
||||
async with self._locks.hold(session_id):
|
||||
if drop_checkpoint is not None:
|
||||
try:
|
||||
await drop_checkpoint(session_id)
|
||||
except CheckpointNotFoundError:
|
||||
pass
|
||||
self._repository.delete(session_id)
|
||||
await self._repository.delete(session_id)
|
||||
|
||||
|
||||
__all__ = ["AccidentAppService", "SET_INFO_INVALID_KEY_MSG"]
|
||||
__all__ = ['AccidentAppService', 'SET_INFO_INVALID_KEY_MSG']
|
||||
|
||||
+12
-28
@@ -1,55 +1,39 @@
|
||||
"""会话业务状态 repository(作业 6-1/6-2)。
|
||||
|
||||
抽象出 repository 接口,使应用服务不依赖具体存储;
|
||||
作业 6 用内存实现跑通四接口,作业 7 再替换为 PostgreSQL 实现
|
||||
(对话 checkpoint 已由 LangGraph checkpointer 承担,本层只管
|
||||
业务字段——单一事实来源,图每轮读取最新值)。
|
||||
|
||||
语义约定:
|
||||
- 业务状态按外部业务键(hphm1、ywrysw...)扁平存储,值为内部强类型;
|
||||
- 不存在的会话视为空状态(读取返回 {}),与首轮 checkpoint 无字段一致;
|
||||
- delete 幂等:重复删除固定成功,不抛异常(行为进测试)。
|
||||
"""
|
||||
"""异步业务字段存储接口;字段值保留 bool/int/str 的区别。"""
|
||||
|
||||
from typing import Protocol
|
||||
|
||||
|
||||
class CheckpointNotFoundError(Exception):
|
||||
"""存储适配器明确确认线程不存在;不能用一般存储异常代替。"""
|
||||
"""存储适配器明确确认线程不存在。"""
|
||||
|
||||
|
||||
class SessionRepository(Protocol):
|
||||
"""业务字段存储接口(作业 7 由 PostgreSQL 实现替换)。"""
|
||||
|
||||
def load(self, session_id: str) -> dict:
|
||||
"""读取会话业务字段;无记录时返回空 dict。"""
|
||||
async def load(self, session_id: str) -> dict:
|
||||
"""读取业务字段;不存在返回空 dict。"""
|
||||
...
|
||||
|
||||
def apply_update(self, session_id: str, patch: dict) -> dict:
|
||||
"""合并差量并返回合并后的完整业务字段。"""
|
||||
async def apply_update(self, session_id: str, patch: dict) -> dict:
|
||||
"""提交差量后返回完整业务字段。"""
|
||||
...
|
||||
|
||||
def delete(self, session_id: str) -> None:
|
||||
"""删除会话业务状态;幂等,重复删除不抛异常。"""
|
||||
async def delete(self, session_id: str) -> None:
|
||||
"""幂等删除;真实存储错误必须传播。"""
|
||||
...
|
||||
|
||||
|
||||
class InMemorySessionRepository:
|
||||
"""内存实现:字典存储,仅用于单进程开发与测试。"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
def __init__(self):
|
||||
self._store: dict[str, dict] = {}
|
||||
|
||||
def load(self, session_id: str) -> dict:
|
||||
async def load(self, session_id: str) -> dict:
|
||||
return dict(self._store.get(session_id) or {})
|
||||
|
||||
def apply_update(self, session_id: str, patch: dict) -> dict:
|
||||
async def apply_update(self, session_id: str, patch: dict) -> dict:
|
||||
state = self._store.setdefault(session_id, {})
|
||||
state.update(patch)
|
||||
return dict(state)
|
||||
|
||||
def delete(self, session_id: str) -> None:
|
||||
# 幂等:pop 的缺省值吞掉不存在的情况,重复删除行为固定为成功
|
||||
async def delete(self, session_id: str) -> None:
|
||||
self._store.pop(session_id, None)
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1 @@
|
||||
"""数据库适配器:短事务、连接池生命周期与 SQLite 单实例文件锁。"""
|
||||
@@ -0,0 +1,16 @@
|
||||
"""图 checkpoint 清理的统一边界;不把缺少清理能力当成功。"""
|
||||
|
||||
|
||||
async def delete_checkpoint(graph, session_id):
|
||||
if graph is None:
|
||||
raise RuntimeError('LangGraph backend not initialized')
|
||||
saver = getattr(graph, 'checkpointer', None)
|
||||
async_delete = getattr(saver, 'adelete_thread', None)
|
||||
if callable(async_delete):
|
||||
await async_delete(session_id)
|
||||
return
|
||||
sync_delete = getattr(saver, 'delete_thread', None)
|
||||
if callable(sync_delete):
|
||||
sync_delete(session_id)
|
||||
return
|
||||
raise RuntimeError('LangGraph checkpointer does not support thread deletion')
|
||||
@@ -0,0 +1,93 @@
|
||||
import asyncio
|
||||
import os
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
|
||||
|
||||
class SqliteInstanceLock:
|
||||
"""拒绝第二个进程/实例打开同一 SQLite 文件;锁由操作系统随退出释放。"""
|
||||
|
||||
def __init__(self, path):
|
||||
self.path = Path(str(path) + '.instance.lock')
|
||||
self.file = None
|
||||
|
||||
def __enter__(self):
|
||||
self.path.parent.mkdir(parents=True, exist_ok=True)
|
||||
self.file = self.path.open('a+b')
|
||||
try:
|
||||
self.file.seek(0, os.SEEK_END)
|
||||
if self.file.tell() == 0:
|
||||
self.file.write(b'0')
|
||||
self.file.flush()
|
||||
self.file.seek(0)
|
||||
if os.name == 'nt':
|
||||
import msvcrt
|
||||
msvcrt.locking(self.file.fileno(), msvcrt.LK_NBLCK, 1)
|
||||
else:
|
||||
import fcntl
|
||||
fcntl.flock(self.file.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||||
except OSError as exc:
|
||||
self.file.close()
|
||||
self.file = None
|
||||
raise RuntimeError('SQLite storage already has an active instance; use PostgreSQL for multiple workers') from exc
|
||||
return self
|
||||
|
||||
def __exit__(self, *_):
|
||||
if self.file is not None:
|
||||
self.file.seek(0)
|
||||
if os.name == 'nt':
|
||||
import msvcrt
|
||||
msvcrt.locking(self.file.fileno(), msvcrt.LK_UNLCK, 1)
|
||||
else:
|
||||
import fcntl
|
||||
fcntl.flock(self.file.fileno(), fcntl.LOCK_UN)
|
||||
self.file.close()
|
||||
self.file = None
|
||||
|
||||
|
||||
class Connection:
|
||||
def __init__(self, raw, postgres=False):
|
||||
self.raw = raw
|
||||
self.postgres = postgres
|
||||
|
||||
async def execute(self, sql, params=()):
|
||||
if self.postgres:
|
||||
sql = sql.replace('?', '%s')
|
||||
return await self.raw.execute(sql, params)
|
||||
|
||||
async def one(self, sql, params=()):
|
||||
cursor = await self.execute(sql, params)
|
||||
try:
|
||||
row = await cursor.fetchone()
|
||||
return dict(row) if row is not None else None
|
||||
finally:
|
||||
await cursor.close()
|
||||
|
||||
|
||||
class SqliteDatabase:
|
||||
def __init__(self, connection):
|
||||
self.connection = connection
|
||||
self.lock = asyncio.Lock()
|
||||
|
||||
@asynccontextmanager
|
||||
async def transaction(self):
|
||||
# 不跨 LLM 调用持有事务;仅保护一次业务提交。
|
||||
async with self.lock:
|
||||
await self.connection.execute('BEGIN IMMEDIATE')
|
||||
try:
|
||||
yield Connection(self.connection)
|
||||
await self.connection.commit()
|
||||
except BaseException:
|
||||
await self.connection.rollback()
|
||||
raise
|
||||
|
||||
|
||||
class PostgresDatabase:
|
||||
def __init__(self, pool):
|
||||
self.pool = pool
|
||||
|
||||
@asynccontextmanager
|
||||
async def transaction(self):
|
||||
async with self.pool.connection() as connection:
|
||||
async with connection.transaction():
|
||||
yield Connection(connection, postgres=True)
|
||||
@@ -0,0 +1,61 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
from contextlib import asynccontextmanager
|
||||
|
||||
from ..backends.chat import ChatBackendAPIError
|
||||
from ..services.session_locks import SessionLocks
|
||||
|
||||
|
||||
class PostgresSessionLocks:
|
||||
"""本进程锁 + 独立连接上的 session advisory lock;不跨 LLM 持有行事务。"""
|
||||
|
||||
def __init__(self, pool, timeout=30):
|
||||
self.pool = pool
|
||||
self.timeout = timeout
|
||||
self.local = SessionLocks()
|
||||
|
||||
@staticmethod
|
||||
def key(session_id):
|
||||
digest = hashlib.blake2b(('znjj-agent:'+session_id).encode(), digest_size=8).digest()
|
||||
return int.from_bytes(digest, 'big', signed=True)
|
||||
|
||||
async def _release(self, connection, key):
|
||||
try:
|
||||
async with asyncio.timeout(5):
|
||||
await connection.execute('SELECT pg_advisory_unlock(%s)', (key,))
|
||||
except Exception:
|
||||
# 不把仍持锁或失联的连接放回池。
|
||||
await connection.close()
|
||||
|
||||
@asynccontextmanager
|
||||
async def hold(self, session_id):
|
||||
async with self.local.hold(session_id):
|
||||
async with self.pool.connection(timeout=self.timeout) as connection:
|
||||
acquired = False
|
||||
key = self.key(session_id)
|
||||
try:
|
||||
try:
|
||||
async with asyncio.timeout(self.timeout):
|
||||
while True:
|
||||
cursor = await connection.execute('SELECT pg_try_advisory_lock(%s) AS locked', (key,))
|
||||
row = await cursor.fetchone()
|
||||
await cursor.close()
|
||||
if row['locked']:
|
||||
acquired = True
|
||||
break
|
||||
await asyncio.sleep(0.05)
|
||||
except TimeoutError as exc:
|
||||
raise ChatBackendAPIError('Session is busy; retry later') from exc
|
||||
yield
|
||||
finally:
|
||||
if not acquired:
|
||||
# 查询取消时服务端可能已加锁但客户端尚未收到结果。
|
||||
# 关闭此连接,避免把潜在持锁连接放回池。
|
||||
await connection.close()
|
||||
if acquired:
|
||||
release = asyncio.create_task(self._release(connection, key))
|
||||
try:
|
||||
await asyncio.shield(release)
|
||||
except asyncio.CancelledError:
|
||||
await release
|
||||
raise
|
||||
@@ -0,0 +1,34 @@
|
||||
"""python -m src.storage.migrate --backend sqlite|postgresql"""
|
||||
import argparse
|
||||
import asyncio
|
||||
import os
|
||||
from dataclasses import replace
|
||||
|
||||
from ..core.config import Config
|
||||
from .runtime import StorageSettings, open_storage
|
||||
|
||||
|
||||
async def migrate(settings):
|
||||
if settings.kind == 'memory':
|
||||
raise ValueError('Select --backend sqlite or postgresql')
|
||||
async with open_storage(replace(settings, auto_migrate=True)):
|
||||
print(f'Initialized {settings.kind} business schema and LangGraph checkpoint tables.')
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Initialize the selected agent storage (no data transfer between backends).')
|
||||
parser.add_argument('--backend', choices=['sqlite', 'postgresql'])
|
||||
parser.add_argument('--sqlite-path')
|
||||
args = parser.parse_args()
|
||||
settings = StorageSettings.from_config(Config)
|
||||
if args.backend:
|
||||
settings = replace(settings, kind=args.backend)
|
||||
if args.sqlite_path:
|
||||
settings = replace(settings, sqlite_path=args.sqlite_path)
|
||||
if os.name == 'nt' and settings.kind == 'postgresql':
|
||||
asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
|
||||
asyncio.run(migrate(settings))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
@@ -0,0 +1,142 @@
|
||||
"""SQLite/PostgreSQL 共用的业务表、版本检查与轮次提交日志。"""
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from uuid import uuid4
|
||||
|
||||
from ..backends.chat import ChatBackendConflictError
|
||||
|
||||
SCHEMA_VERSION = 1
|
||||
DDL = [
|
||||
'CREATE TABLE IF NOT EXISTS agent_schema (version INTEGER PRIMARY KEY)',
|
||||
'''CREATE TABLE IF NOT EXISTS agent_sessions (
|
||||
session_id TEXT PRIMARY KEY, version BIGINT NOT NULL, epoch BIGINT NOT NULL,
|
||||
fields TEXT NOT NULL, graph_state TEXT NOT NULL, deleted INTEGER NOT NULL,
|
||||
updated_at TEXT NOT NULL)''',
|
||||
'''CREATE TABLE IF NOT EXISTS agent_turns (
|
||||
turn_id TEXT PRIMARY KEY, session_id TEXT NOT NULL, epoch BIGINT NOT NULL,
|
||||
request_key TEXT, request_hash TEXT NOT NULL, base_version BIGINT NOT NULL,
|
||||
status TEXT NOT NULL, result TEXT, created_at TEXT NOT NULL,
|
||||
UNIQUE(session_id, epoch, request_key))''',
|
||||
'CREATE INDEX IF NOT EXISTS agent_turns_session ON agent_turns(session_id)',
|
||||
]
|
||||
|
||||
|
||||
def encode(value):
|
||||
return json.dumps(value, ensure_ascii=False, separators=(',', ':'))
|
||||
|
||||
|
||||
def now():
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
|
||||
@dataclass
|
||||
class TurnTicket:
|
||||
session_id: str
|
||||
turn_id: str
|
||||
version: int
|
||||
epoch: int
|
||||
graph_state: dict
|
||||
replay: dict | None = None
|
||||
|
||||
|
||||
class SqlSessionRepository:
|
||||
def __init__(self, database):
|
||||
self.database = database
|
||||
|
||||
async def migrate(self):
|
||||
async with self.database.transaction() as conn:
|
||||
for sql in DDL:
|
||||
await conn.execute(sql)
|
||||
await conn.execute('INSERT INTO agent_schema(version) VALUES (?) ON CONFLICT DO NOTHING', (SCHEMA_VERSION,))
|
||||
|
||||
async def verify(self):
|
||||
async with self.database.transaction() as conn:
|
||||
row = await conn.one('SELECT MAX(version) AS version FROM agent_schema')
|
||||
if not row or row['version'] != SCHEMA_VERSION:
|
||||
raise RuntimeError('Unsupported agent storage schema; run the storage migration command')
|
||||
|
||||
async def _ensure(self, conn, session_id):
|
||||
await conn.execute('''INSERT INTO agent_sessions
|
||||
(session_id,version,epoch,fields,graph_state,deleted,updated_at)
|
||||
VALUES (?,0,0,'{}','{}',0,?) ON CONFLICT DO NOTHING''', (session_id, now()))
|
||||
return await conn.one('SELECT * FROM agent_sessions WHERE session_id=?', (session_id,))
|
||||
|
||||
async def load(self, session_id):
|
||||
async with self.database.transaction() as conn:
|
||||
row = await conn.one('SELECT fields FROM agent_sessions WHERE session_id=?', (session_id,))
|
||||
return json.loads(row['fields']) if row else {}
|
||||
|
||||
async def apply_update(self, session_id, patch):
|
||||
async with self.database.transaction() as conn:
|
||||
row = await self._ensure(conn, session_id)
|
||||
fields = json.loads(row['fields'])
|
||||
fields.update(patch)
|
||||
cursor = await conn.execute('''UPDATE agent_sessions SET fields=?,version=version+1,
|
||||
deleted=0,updated_at=? WHERE session_id=? AND version=?''',
|
||||
(encode(fields), now(), session_id, row['version']))
|
||||
if cursor.rowcount != 1:
|
||||
raise ChatBackendConflictError('Session version changed; retry the request')
|
||||
return fields
|
||||
|
||||
async def delete(self, session_id):
|
||||
async with self.database.transaction() as conn:
|
||||
row = await self._ensure(conn, session_id)
|
||||
cursor = await conn.execute('''UPDATE agent_sessions SET fields='{}',graph_state='{}',
|
||||
deleted=1,version=version+1,epoch=epoch+1,updated_at=?
|
||||
WHERE session_id=? AND version=?''', (now(), session_id, row['version']))
|
||||
if cursor.rowcount != 1:
|
||||
raise ChatBackendConflictError('Session version changed during deletion')
|
||||
await conn.execute('DELETE FROM agent_turns WHERE session_id=?', (session_id,))
|
||||
|
||||
async def begin_turn(self, session_id, request_key, request_hash):
|
||||
async with self.database.transaction() as conn:
|
||||
row = await self._ensure(conn, session_id)
|
||||
previous = None
|
||||
if request_key:
|
||||
previous = await conn.one('''SELECT * FROM agent_turns
|
||||
WHERE session_id=? AND epoch=? AND request_key=?''',
|
||||
(session_id, row['epoch'], request_key))
|
||||
if previous and previous['request_hash'] != request_hash:
|
||||
raise ChatBackendConflictError('Idempotency-Key was used with a different request')
|
||||
if previous and previous['status'] == 'committed':
|
||||
return TurnTicket(session_id, previous['turn_id'], row['version'], row['epoch'],
|
||||
json.loads(row['graph_state']), json.loads(previous['result']))
|
||||
# 同会话锁已持有:遗留 pending 只能来自中断,下一次从提交快照重建。
|
||||
await conn.execute("UPDATE agent_turns SET status='aborted' WHERE session_id=? AND status='pending'", (session_id,))
|
||||
turn_id = previous['turn_id'] if previous else uuid4().hex
|
||||
if previous:
|
||||
await conn.execute("UPDATE agent_turns SET status='pending',base_version=?,result=NULL WHERE turn_id=?",
|
||||
(row['version'], turn_id))
|
||||
else:
|
||||
await conn.execute('''INSERT INTO agent_turns
|
||||
(turn_id,session_id,epoch,request_key,request_hash,base_version,status,result,created_at)
|
||||
VALUES (?,?,?,?,?,?,'pending',NULL,?)''',
|
||||
(turn_id, session_id, row['epoch'], request_key, request_hash, row['version'], now()))
|
||||
return TurnTicket(session_id, turn_id, row['version'], row['epoch'], json.loads(row['graph_state']))
|
||||
|
||||
async def commit_turn(self, ticket, patch, graph_state):
|
||||
async with self.database.transaction() as conn:
|
||||
row = await conn.one('SELECT * FROM agent_sessions WHERE session_id=?', (ticket.session_id,))
|
||||
if not row or row['epoch'] != ticket.epoch or row['version'] != ticket.version:
|
||||
raise ChatBackendConflictError('Session changed before turn commit')
|
||||
fields = json.loads(row['fields'])
|
||||
fields.update(patch)
|
||||
state = dict(graph_state)
|
||||
state['state_version'] = ticket.version + 1
|
||||
state['turn_patch'] = graph_state.get('turn_patch') or {}
|
||||
cursor = await conn.execute('''UPDATE agent_sessions SET fields=?,graph_state=?,deleted=0,
|
||||
version=version+1,updated_at=? WHERE session_id=? AND version=? AND epoch=?''',
|
||||
(encode(fields), encode(state), now(), ticket.session_id, ticket.version, ticket.epoch))
|
||||
if cursor.rowcount != 1:
|
||||
raise ChatBackendConflictError('Session version changed during commit')
|
||||
result = {key: state.get(key) for key in ('response_text', 'turn_patch', 'stage_code')}
|
||||
cursor = await conn.execute("UPDATE agent_turns SET status='committed',result=? WHERE turn_id=? AND status='pending'",
|
||||
(encode(result), ticket.turn_id))
|
||||
if cursor.rowcount != 1:
|
||||
raise ChatBackendConflictError('Turn is no longer pending')
|
||||
|
||||
async def abort_turn(self, ticket):
|
||||
async with self.database.transaction() as conn:
|
||||
await conn.execute("UPDATE agent_turns SET status='aborted' WHERE turn_id=? AND status='pending'", (ticket.turn_id,))
|
||||
@@ -0,0 +1,92 @@
|
||||
from contextlib import AsyncExitStack, asynccontextmanager
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
from langgraph.checkpoint.memory import InMemorySaver
|
||||
|
||||
from ..services.repository import InMemorySessionRepository
|
||||
from ..services.session_locks import SessionLocks
|
||||
from .database import PostgresDatabase, SqliteDatabase, SqliteInstanceLock
|
||||
from .repository import SqlSessionRepository
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StorageSettings:
|
||||
kind: str = 'memory'
|
||||
sqlite_path: str = './data/agent.sqlite3'
|
||||
postgres_dsn: str = ''
|
||||
auto_migrate: bool = False
|
||||
pool_size: int = 10
|
||||
lock_timeout: float = 30
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, config):
|
||||
return cls(config.AGENT_STORAGE, config.AGENT_SQLITE_PATH, config.AGENT_POSTGRES_DSN,
|
||||
config.AGENT_DB_AUTO_MIGRATE, config.AGENT_DB_POOL_SIZE,
|
||||
config.AGENT_DB_LOCK_TIMEOUT_SECONDS)
|
||||
|
||||
|
||||
@dataclass
|
||||
class StorageRuntime:
|
||||
repository: object
|
||||
checkpointer: object
|
||||
locks: object
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def open_storage(settings: StorageSettings):
|
||||
if settings.kind not in ('memory', 'sqlite', 'postgresql'):
|
||||
raise ValueError('Unknown AGENT_STORAGE')
|
||||
async with AsyncExitStack() as stack:
|
||||
if settings.kind == 'memory':
|
||||
yield StorageRuntime(InMemorySessionRepository(), InMemorySaver(), SessionLocks())
|
||||
return
|
||||
if settings.kind == 'sqlite':
|
||||
import aiosqlite
|
||||
from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
|
||||
|
||||
if not settings.sqlite_path.strip() or settings.sqlite_path == ':memory:':
|
||||
raise ValueError('SQLite persistence requires a file path')
|
||||
path = Path(settings.sqlite_path).resolve()
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
stack.enter_context(SqliteInstanceLock(path))
|
||||
business = await stack.enter_async_context(aiosqlite.connect(path, isolation_level=None))
|
||||
business.row_factory = aiosqlite.Row
|
||||
await business.execute('PRAGMA journal_mode=WAL')
|
||||
await business.execute('PRAGMA busy_timeout=5000')
|
||||
checkpoint_connection = await stack.enter_async_context(aiosqlite.connect(path))
|
||||
await checkpoint_connection.execute('PRAGMA busy_timeout=5000')
|
||||
checkpointer = AsyncSqliteSaver(checkpoint_connection)
|
||||
database = SqliteDatabase(business)
|
||||
locks = SessionLocks()
|
||||
else:
|
||||
from psycopg.rows import dict_row
|
||||
from psycopg_pool import AsyncConnectionPool
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from .locks import PostgresSessionLocks
|
||||
|
||||
if not settings.postgres_dsn.strip():
|
||||
raise ValueError('AGENT_POSTGRES_DSN is required')
|
||||
async def pool():
|
||||
value = AsyncConnectionPool(
|
||||
settings.postgres_dsn, min_size=1, max_size=settings.pool_size, open=False,
|
||||
timeout=settings.lock_timeout,
|
||||
kwargs={'autocommit': True, 'prepare_threshold': 0, 'row_factory': dict_row, 'connect_timeout': 5},
|
||||
)
|
||||
await stack.enter_async_context(value)
|
||||
await value.wait(timeout=15)
|
||||
return value
|
||||
business_pool = await pool()
|
||||
checkpoint_pool = await pool()
|
||||
lock_pool = await pool()
|
||||
database = PostgresDatabase(business_pool)
|
||||
checkpointer = AsyncPostgresSaver(checkpoint_pool)
|
||||
locks = PostgresSessionLocks(lock_pool, settings.lock_timeout)
|
||||
repository = SqlSessionRepository(database)
|
||||
if settings.auto_migrate:
|
||||
await repository.migrate()
|
||||
await checkpointer.setup()
|
||||
await repository.verify()
|
||||
# 启动时同时检查 saver 表,不能仅检查业务表后偷偷用内存 saver。
|
||||
await checkpointer.aget_tuple({'configurable': {'thread_id': '__storage_health__'}})
|
||||
yield StorageRuntime(repository, checkpointer, locks)
|
||||
Reference in New Issue
Block a user