初始化提交
This commit is contained in:
chenxin committed
commit
887c1f449e
963 files changed
+166869
No files matched your search
@@ -0,0 +1,304 @@
|
||||
"""记忆服务。
|
||||
|
||||
集成 ChromaDB 向量存储：
|
||||
- create_memory / 大/小周天自动创建的 Memory 写入后异步生成 embedding 存向量
|
||||
- retrieve_memories 支持向量相似度检索，失败时回退 SQL 排序检索
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.config import settings
|
||||
from app.core.errors import AppError
|
||||
from app.core.logging import logger
|
||||
from app.core.time import utc_now
|
||||
from app.db.session import AsyncSessionLocal
|
||||
from app.models.memory import Memory
|
||||
from app.modules.memory.context import MemoryRetrievalResult
|
||||
from app.repositories import memory as memory_repo
|
||||
from app.repositories.json_utils import dumps
|
||||
from app.schemas.common import Page
|
||||
from app.schemas.memory import MemoryCreateRequest, MemoryResponse, MemoryUpdateRequest
|
||||
from app.services.audit import record_audit_event
|
||||
from app.services.characters import get_active_character_model
|
||||
from app.services.serializers import memory_to_response
|
||||
|
||||
|
||||
def _get_embedding_service(manager):
|
||||
"""尝试创建 EmbeddingService，失败时返回 None。"""
|
||||
try:
|
||||
models = manager.model_registry.list_by_capability("embedding")
|
||||
if not models:
|
||||
return None
|
||||
from app.ai_provider.embedding import EmbeddingService
|
||||
return EmbeddingService(manager)
|
||||
except Exception: # noqa: BLE001
|
||||
return None
|
||||
|
||||
|
||||
async def _embed_and_store(memory: Memory) -> None:
|
||||
"""异步后台任务：为记忆生成 embedding 并存入向量库。
|
||||
|
||||
embedding 服务和 ChromaDB 均不可用时静默跳过——不影响主流程。
|
||||
"""
|
||||
text = memory.content or memory.title or ""
|
||||
if not text:
|
||||
return
|
||||
try:
|
||||
from app.ai_provider import ProviderManagerFactory
|
||||
|
||||
manager = ProviderManagerFactory.create()
|
||||
service = _get_embedding_service(manager)
|
||||
if service is None:
|
||||
return
|
||||
|
||||
import asyncio as _asyncio
|
||||
|
||||
response = await _asyncio.wait_for(
|
||||
service.embed(text),
|
||||
timeout=10.0,
|
||||
)
|
||||
if not response.data or not response.data[0].embedding:
|
||||
return
|
||||
embedding = response.data[0].embedding
|
||||
|
||||
from app.integrations.vector_store import get_vector_store
|
||||
|
||||
store = get_vector_store()
|
||||
if store is None:
|
||||
return
|
||||
store.add_memory(
|
||||
character_id=memory.character_id,
|
||||
memory_id=memory.id,
|
||||
text=text,
|
||||
embedding=embedding,
|
||||
metadata={
|
||||
"type": memory.type,
|
||||
"source": memory.source,
|
||||
"title": memory.title or "",
|
||||
},
|
||||
)
|
||||
except Exception: # noqa: BLE001
|
||||
logger.warning(
|
||||
"memory_embed_failed",
|
||||
memory_id=memory.id,
|
||||
character_id=memory.character_id,
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
async def embed_memory_async(memory: Memory) -> None:
|
||||
"""公开接口：为单条记忆生成 embedding 并存入向量库（供大/小周天后台任务调用）。"""
|
||||
await _embed_and_store(memory)
|
||||
# ── CRUD ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
async def list_memory_page(
|
||||
*,
|
||||
page: int = 1,
|
||||
page_size: int = 20,
|
||||
memory_type: str | None = None,
|
||||
source: str | None = None,
|
||||
owner_scope: str = settings.default_owner_scope,
|
||||
) -> Page[MemoryResponse]:
|
||||
async with AsyncSessionLocal() as session:
|
||||
character = await get_active_character_model(session, owner_scope)
|
||||
safe_page_size = min(max(page_size, 1), 100)
|
||||
safe_page = max(page, 1)
|
||||
items, total = await memory_repo.list_memories(
|
||||
session,
|
||||
character.id,
|
||||
memory_type=memory_type,
|
||||
source=source,
|
||||
page=safe_page,
|
||||
page_size=safe_page_size,
|
||||
)
|
||||
return Page(
|
||||
items=[memory_to_response(item) for item in items],
|
||||
page=safe_page,
|
||||
page_size=safe_page_size,
|
||||
total=total,
|
||||
)
|
||||
|
||||
|
||||
async def create_memory(
|
||||
session: AsyncSession,
|
||||
payload: MemoryCreateRequest,
|
||||
*,
|
||||
owner_scope: str = settings.default_owner_scope,
|
||||
) -> MemoryResponse:
|
||||
character = await get_active_character_model(session, owner_scope)
|
||||
character_id = payload.character_id or character.id
|
||||
if character_id != character.id:
|
||||
raise AppError("not_found", "人物不存在。")
|
||||
memory = Memory(
|
||||
character_id=character_id,
|
||||
type=payload.type,
|
||||
title=payload.title,
|
||||
content=payload.content,
|
||||
source=payload.source,
|
||||
confidence=payload.confidence,
|
||||
importance=payload.importance,
|
||||
sensitivity=payload.sensitivity,
|
||||
evidence_ids_json=dumps(payload.evidence_ids),
|
||||
status="active",
|
||||
)
|
||||
session.add(memory)
|
||||
await session.flush()
|
||||
# 异步 embedding
|
||||
if settings.vector_db_enabled:
|
||||
import asyncio
|
||||
|
||||
asyncio.ensure_future(_embed_and_store(memory))
|
||||
await record_audit_event(
|
||||
session,
|
||||
action="memory.created",
|
||||
character_id=character_id,
|
||||
actor="user",
|
||||
decision="auto",
|
||||
input_summary=payload.title,
|
||||
resource_type="memory",
|
||||
resource_id=memory.id,
|
||||
)
|
||||
await session.commit()
|
||||
await session.refresh(memory)
|
||||
return memory_to_response(memory)
|
||||
|
||||
|
||||
async def update_memory(
|
||||
memory_id: str, payload: MemoryUpdateRequest, *, owner_scope: str = settings.default_owner_scope
|
||||
) -> MemoryResponse:
|
||||
async with AsyncSessionLocal() as session:
|
||||
character = await get_active_character_model(session, owner_scope)
|
||||
memory = await memory_repo.get_memory(session, memory_id)
|
||||
if memory is None or memory.character_id != character.id or memory.deleted_at is not None:
|
||||
raise AppError("not_found", "记忆不存在。")
|
||||
updates = payload.model_dump(exclude_unset=True)
|
||||
evidence_ids = updates.pop("evidence_ids", None)
|
||||
for key, value in updates.items():
|
||||
setattr(memory, key, value)
|
||||
if evidence_ids is not None:
|
||||
memory.evidence_ids_json = dumps(evidence_ids)
|
||||
memory.updated_at = utc_now()
|
||||
# 如果内容变了，异步更新向量
|
||||
if ("content" in updates or "title" in updates) and settings.vector_db_enabled:
|
||||
import asyncio
|
||||
|
||||
asyncio.ensure_future(_embed_and_store(memory))
|
||||
await record_audit_event(
|
||||
session,
|
||||
action="memory.updated",
|
||||
character_id=character.id,
|
||||
actor="user",
|
||||
decision="auto",
|
||||
input_summary=memory.title,
|
||||
resource_type="memory",
|
||||
resource_id=memory.id,
|
||||
)
|
||||
await session.commit()
|
||||
await session.refresh(memory)
|
||||
return memory_to_response(memory)
|
||||
|
||||
|
||||
async def delete_memory(
|
||||
memory_id: str, *, owner_scope: str = settings.default_owner_scope
|
||||
) -> MemoryResponse:
|
||||
async with AsyncSessionLocal() as session:
|
||||
character = await get_active_character_model(session, owner_scope)
|
||||
memory = await memory_repo.get_memory(session, memory_id)
|
||||
if memory is None or memory.character_id != character.id or memory.deleted_at is not None:
|
||||
raise AppError("not_found", "记忆不存在。")
|
||||
memory.status = "deleted"
|
||||
memory.deleted_at = utc_now()
|
||||
memory.updated_at = utc_now()
|
||||
# 同步删除向量
|
||||
if settings.vector_db_enabled:
|
||||
try:
|
||||
from app.integrations.vector_store import get_vector_store
|
||||
|
||||
store = get_vector_store()
|
||||
if store is not None:
|
||||
store.delete_memory(character.id, memory_id)
|
||||
except Exception: # noqa: BLE001
|
||||
pass
|
||||
await record_audit_event(
|
||||
session,
|
||||
action="memory.deleted",
|
||||
character_id=character.id,
|
||||
actor="user",
|
||||
decision="auto",
|
||||
input_summary=memory.title,
|
||||
resource_type="memory",
|
||||
resource_id=memory.id,
|
||||
)
|
||||
await session.commit()
|
||||
await session.refresh(memory)
|
||||
return memory_to_response(memory)
|
||||
|
||||
|
||||
async def memory_context(
|
||||
session: AsyncSession,
|
||||
character_id: str,
|
||||
*,
|
||||
limit: int,
|
||||
_owner_scope: str = settings.default_owner_scope,
|
||||
) -> str:
|
||||
"""获取记忆上下文文本（供叩心枢开场、辰正等场景使用）。"""
|
||||
memories = await memory_repo.list_injectable_memories(session, character_id, limit=limit)
|
||||
if not memories:
|
||||
return ""
|
||||
lines = [f"- {item.title}: {item.content}" for item in memories]
|
||||
return "相关记忆：\n" + "\n".join(lines)
|
||||
|
||||
|
||||
async def retrieve_memories(
|
||||
session: AsyncSession,
|
||||
character_id: str,
|
||||
*,
|
||||
limit: int,
|
||||
_owner_scope: str = settings.default_owner_scope,
|
||||
) -> MemoryRetrievalResult:
|
||||
"""检索可注入上下文的相关记忆。
|
||||
|
||||
优先尝试向量相似度检索，失败时回退 SQL 排序检索。
|
||||
"""
|
||||
# 尝试向量检索（5s 超时，避免 embedding endpoint 不可达时阻塞消息流）
|
||||
if settings.vector_db_enabled:
|
||||
try:
|
||||
from app.ai_provider import ProviderManagerFactory
|
||||
|
||||
manager = ProviderManagerFactory.create()
|
||||
models = manager.model_registry.list_by_capability("embedding")
|
||||
if models:
|
||||
import asyncio as _asyncio
|
||||
|
||||
from app.ai_provider.embedding import EmbeddingService
|
||||
from app.integrations.vector_store import get_vector_store
|
||||
|
||||
query_text = f"检索 {character_id} 的最新相关记忆"
|
||||
service = EmbeddingService(manager)
|
||||
response = await _asyncio.wait_for(
|
||||
service.embed(query_text),
|
||||
timeout=5.0,
|
||||
)
|
||||
if response.data and response.data[0].embedding:
|
||||
store = get_vector_store()
|
||||
if store is not None:
|
||||
vector_results = store.search(
|
||||
character_id,
|
||||
response.data[0].embedding,
|
||||
limit=limit,
|
||||
)
|
||||
if vector_results:
|
||||
ids = [r["id"] for r in vector_results]
|
||||
memories = await memory_repo.get_memories_by_ids(session, ids)
|
||||
if memories:
|
||||
return MemoryRetrievalResult(memories=memories)
|
||||
except Exception: # noqa: BLE001
|
||||
logger.warning("vector_retrieval_failed，回退 SQL 检索")
|
||||
|
||||
# 回退 SQL
|
||||
memories = await memory_repo.list_injectable_memories(session, character_id, limit=limit)
|
||||
return MemoryRetrievalResult(memories=memories)
|
||||
Reference in new issue
Block a user