ai-agent/MemoryAgents/backend/app/main.py

from __future__ import annotations

from contextlib import asynccontextmanager

from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware

from .config import settings
from .db import repository
from .embeddings import embedding_client
from .registry import agents
from .schemas import ChatRequest, ChatResponse, ResetRequest
from .adapters.base import AdapterUnavailable


@asynccontextmanager
async def lifespan(_: FastAPI):
    """启动时初始化存储连接，退出时释放连接池。"""
    await repository.initialize()
    yield
    if repository.pool:
        await repository.pool.close()


app = FastAPI(title="Memory Agents Web API", version="0.1.0", lifespan=lifespan)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3000", "http://127.0.0.1:3000"],
    allow_origin_regex=(
        r"^https?://(?:192\.168\.\d+\.\d+|10\.\d+\.\d+|172\.(?:1[6-9]|2\d|3[0-1])\.\d+\.\d+):3000$"
    ),
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.get("/api/health")
async def health() -> dict:
    result = {
        "status": "ok",
        "storage": repository.storage_mode,
        "workspace_id": settings.workspace_id,
        "llm_configured": bool(settings.llm_api_key),
        "embedding_configured": embedding_client.configured,
    }
    if repository.initialization_error:
        result["storage_error"] = repository.initialization_error
    return result


@app.get("/api/systems")
async def systems():
    return {
        "storage": repository.storage_mode,
        "systems": [agent.status for agent in agents.values()],
    }


@app.post("/api/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """把对话请求交给指定记忆适配器处理。"""
    agent = agents[request.system]
    try:
        return await agent.chat(request.message)
    except AdapterUnavailable as exc:
        raise HTTPException(status_code=409, detail=str(exc)) from exc
    except Exception as exc:
        raise HTTPException(status_code=502, detail=f"{request.system} 执行失败：{exc}") from exc


@app.get("/api/memories")
async def memories(system: str = Query(...), query: str | None = None):
    if system not in agents:
        raise HTTPException(status_code=404, detail="未知记忆系统")
    return {"system": system, "memories": await agents[system].memories(query)}


@app.delete("/api/memories/{memory_id}")
async def delete_memory(memory_id: str, system: str = Query(...)):
    if system not in agents:
        raise HTTPException(status_code=404, detail="未知记忆系统")
    deleted = await agents[system].delete_memory(memory_id)
    if not deleted:
        raise HTTPException(status_code=404, detail="记忆不存在或已删除")
    return {"status": "ok", "system": system, "memory_id": memory_id}


@app.get("/api/audit")
async def audit(system: str | None = None):
    if system is not None and system not in agents:
        raise HTTPException(status_code=404, detail="未知记忆系统")
    if system:
        events = await agents[system].audit()
    else:
        events = await repository.list_audit()
    return {"system": system, "events": events}


@app.post("/api/reset")
async def reset(request: ResetRequest):
    if request.system is None:
        # 全量重置必须逐个调用适配器，ReMe 等实现还维护自己的文件和索引。
        for agent in agents.values():
            await agent.reset()
    elif request.system in agents:
        await agents[request.system].reset()
    else:
        raise HTTPException(status_code=404, detail="未知记忆系统")
    return {"status": "ok", "system": request.system}