aiagent/backend/app/api/agent_memory.py

"""
Agent 记忆管理 API — 全局知识 / 知识实体 / 学习模式的 CRUD + 导入导出
"""
from __future__ import annotations

import logging
import uuid
from datetime import datetime
from typing import List, Optional

from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, Field
from sqlalchemy.orm import Session

from app.core.database import get_db
from app.api.auth import get_current_user
from app.models.user import User
from app.models.agent import Agent, GlobalKnowledge, KnowledgeEntity, KnowledgeRelation
from app.models.agent_learning_pattern import AgentLearningPattern
from app.models.agent_vector_memory import AgentVectorMemory

logger = logging.getLogger(__name__)

router = APIRouter(
    prefix="/api/v1/agents",
    tags=["agent-memory"],
)

SCOPE_AGENT = "agent"


def _mem_scope(user: User, agent_id: str) -> str:
    """构建记忆 scope_id，与 AgentRuntime 写入的格式保持一致（{user_id}:{agent_id}）。"""
    return f"{user.id}:{agent_id}"


def _check_agent(db: Session, agent_id: str, user: User):
    agent = db.query(Agent).filter(Agent.id == agent_id).first()
    if not agent:
        raise HTTPException(status_code=404, detail="Agent 不存在")
    if agent.user_id and agent.user_id != user.id and user.role != "admin":
        raise HTTPException(status_code=403, detail="无权访问该 Agent")
    return agent


# ─────────── Pydantic models ───────────


class GlobalKnowledgeItem(BaseModel):
    id: str
    content: str
    source_agent_id: Optional[str] = None
    source_user_id: Optional[str] = None
    tags: Optional[list] = None
    confidence: str = "medium"
    scope_kind: str = "agent"
    scope_id: str = ""
    expires_at: Optional[str] = None
    created_at: Optional[str] = None


class GlobalKnowledgeCreate(BaseModel):
    content: str
    tags: Optional[list] = None
    confidence: str = "medium"


class GlobalKnowledgeUpdate(BaseModel):
    content: Optional[str] = None
    tags: Optional[list] = None
    confidence: Optional[str] = None


class GlobalKnowledgeList(BaseModel):
    items: List[GlobalKnowledgeItem]
    total: int


class KnowledgeEntityItem(BaseModel):
    id: str
    name: str
    entity_type: str = "concept"
    description: Optional[str] = None
    source: str = "extracted"
    confidence: str = "medium"
    tags: Optional[list] = None
    scope_kind: str = "agent"
    scope_id: str = ""
    created_at: Optional[str] = None


class KnowledgeEntityCreate(BaseModel):
    name: str
    entity_type: str = "concept"
    description: Optional[str] = None
    tags: Optional[list] = None
    confidence: str = "medium"


class KnowledgeEntityUpdate(BaseModel):
    name: Optional[str] = None
    entity_type: Optional[str] = None
    description: Optional[str] = None
    tags: Optional[list] = None
    confidence: Optional[str] = None


class KnowledgeEntityList(BaseModel):
    items: List[KnowledgeEntityItem]
    total: int


class LearningPatternItem(BaseModel):
    id: str
    scope_kind: str
    scope_id: str
    task_category: str
    task_keywords: Optional[str] = None
    suggested_tools: Optional[str] = None
    effectiveness_score: float = 0.0
    total_runs: int = 0
    successful_runs: int = 0
    avg_iterations: float = 0.0
    avg_tool_calls: float = 0.0
    last_used_at: Optional[str] = None
    created_at: Optional[str] = None


class LearningPatternList(BaseModel):
    items: List[LearningPatternItem]
    total: int


class VectorMemoryItem(BaseModel):
    id: str
    scope_kind: str = "agent"
    scope_id: str = ""
    session_key: str = ""
    content_text: str
    metadata: Optional[dict] = None
    created_at: Optional[str] = None


class VectorMemoryList(BaseModel):
    items: List[VectorMemoryItem]
    total: int


class MemoryExportResponse(BaseModel):
    agent_id: str
    exported_at: str
    global_knowledge: List[dict] = []
    knowledge_entities: List[dict] = []
    knowledge_relations: List[dict] = []
    learning_patterns: List[dict] = []
    vector_memories: List[dict] = []


class MemoryImportRequest(BaseModel):
    global_knowledge: List[dict] = []
    knowledge_entities: List[dict] = []
    knowledge_relations: List[dict] = []
    learning_patterns: List[dict] = []
    vector_memories: List[dict] = []


# ─────────── 辅助 ───────────


def _gk_to_item(gk: GlobalKnowledge) -> GlobalKnowledgeItem:
    return GlobalKnowledgeItem(
        id=gk.id,
        content=gk.content,
        source_agent_id=gk.source_agent_id,
        source_user_id=gk.source_user_id,
        tags=gk.tags,
        confidence=gk.confidence or "medium",
        scope_kind=gk.scope_kind or SCOPE_AGENT,
        scope_id=gk.scope_id or "",
        expires_at=gk.expires_at.isoformat() if gk.expires_at else None,
        created_at=gk.created_at.isoformat() if gk.created_at else None,
    )


def _ke_to_item(ke: KnowledgeEntity) -> KnowledgeEntityItem:
    return KnowledgeEntityItem(
        id=ke.id,
        name=ke.name,
        entity_type=ke.entity_type or "concept",
        description=getattr(ke, "description", None),
        source=ke.source or "extracted",
        confidence=ke.confidence or "medium",
        tags=getattr(ke, "metadata_", None),
        scope_kind=ke.scope_kind or SCOPE_AGENT,
        scope_id=ke.scope_id or "",
        created_at=ke.created_at.isoformat() if ke.created_at else None,
    )


def _vm_to_item(vm: AgentVectorMemory) -> VectorMemoryItem:
    return VectorMemoryItem(
        id=vm.id,
        scope_kind=vm.scope_kind or SCOPE_AGENT,
        scope_id=vm.scope_id or "",
        session_key=vm.session_key or "",
        content_text=vm.content_text or "",
        metadata=vm.metadata_ or {},
        created_at=vm.created_at.isoformat() if vm.created_at else None,
    )


# ─────────── Global Knowledge CRUD ───────────


@router.get("/{agent_id}/memory/global-knowledge", response_model=GlobalKnowledgeList)
def list_global_knowledge(
    agent_id: str,
    skip: int = 0,
    limit: int = 50,
    search: Optional[str] = None,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    q = db.query(GlobalKnowledge).filter(
        GlobalKnowledge.scope_kind == SCOPE_AGENT,
        GlobalKnowledge.scope_id == scope_id,
    )
    if search:
        q = q.filter(GlobalKnowledge.content.contains(search))
    total = q.count()
    items = q.order_by(GlobalKnowledge.created_at.desc()).offset(skip).limit(limit).all()
    return GlobalKnowledgeList(items=[_gk_to_item(it) for it in items], total=total)


@router.post("/{agent_id}/memory/global-knowledge", response_model=GlobalKnowledgeItem)
def create_global_knowledge(
    agent_id: str,
    body: GlobalKnowledgeCreate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    gk = GlobalKnowledge(
        id=str(uuid.uuid4()),
        content=body.content,
        tags=body.tags,
        confidence=body.confidence,
        scope_kind=SCOPE_AGENT,
        scope_id=scope_id,
        source_user_id=current_user.id,
    )
    db.add(gk)
    db.commit()
    db.refresh(gk)
    return _gk_to_item(gk)


@router.put("/{agent_id}/memory/global-knowledge/{knowledge_id}", response_model=GlobalKnowledgeItem)
def update_global_knowledge(
    agent_id: str,
    knowledge_id: str,
    body: GlobalKnowledgeUpdate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    gk = db.query(GlobalKnowledge).filter(
        GlobalKnowledge.id == knowledge_id,
        GlobalKnowledge.scope_kind == SCOPE_AGENT,
        GlobalKnowledge.scope_id == scope_id,
    ).first()
    if not gk:
        raise HTTPException(status_code=404, detail="知识条目不存在")
    if body.content is not None:
        gk.content = body.content
    if body.tags is not None:
        gk.tags = body.tags
    if body.confidence is not None:
        gk.confidence = body.confidence
    db.commit()
    db.refresh(gk)
    return _gk_to_item(gk)


@router.delete("/{agent_id}/memory/global-knowledge/{knowledge_id}")
def delete_global_knowledge(
    agent_id: str,
    knowledge_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    gk = db.query(GlobalKnowledge).filter(
        GlobalKnowledge.id == knowledge_id,
        GlobalKnowledge.scope_kind == SCOPE_AGENT,
        GlobalKnowledge.scope_id == scope_id,
    ).first()
    if not gk:
        raise HTTPException(status_code=404, detail="知识条目不存在")
    db.delete(gk)
    db.commit()
    return {"ok": True}


# ─────────── Knowledge Entities CRUD ───────────


@router.get("/{agent_id}/memory/knowledge-entities", response_model=KnowledgeEntityList)
def list_knowledge_entities(
    agent_id: str,
    skip: int = 0,
    limit: int = 50,
    search: Optional[str] = None,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    q = db.query(KnowledgeEntity).filter(
        KnowledgeEntity.scope_kind == SCOPE_AGENT,
        KnowledgeEntity.scope_id == scope_id,
    )
    if search:
        q = q.filter(
            (KnowledgeEntity.name.contains(search))
            | (KnowledgeEntity.description.contains(search))
        )
    total = q.count()
    items = q.order_by(KnowledgeEntity.created_at.desc()).offset(skip).limit(limit).all()
    return KnowledgeEntityList(items=[_ke_to_item(it) for it in items], total=total)


@router.post("/{agent_id}/memory/knowledge-entities", response_model=KnowledgeEntityItem)
def create_knowledge_entity(
    agent_id: str,
    body: KnowledgeEntityCreate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    ke = KnowledgeEntity(
        id=str(uuid.uuid4()),
        name=body.name,
        entity_type=body.entity_type,
        description=body.description,
        source="manual",
        confidence=body.confidence,
        metadata_=body.tags,
        scope_kind=SCOPE_AGENT,
        scope_id=scope_id,
        user_id=current_user.id,
    )
    db.add(ke)
    db.commit()
    db.refresh(ke)
    return _ke_to_item(ke)


@router.put("/{agent_id}/memory/knowledge-entities/{entity_id}", response_model=KnowledgeEntityItem)
def update_knowledge_entity(
    agent_id: str,
    entity_id: str,
    body: KnowledgeEntityUpdate,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    ke = db.query(KnowledgeEntity).filter(
        KnowledgeEntity.id == entity_id,
        KnowledgeEntity.scope_kind == SCOPE_AGENT,
        KnowledgeEntity.scope_id == scope_id,
    ).first()
    if not ke:
        raise HTTPException(status_code=404, detail="实体不存在")
    if body.name is not None:
        ke.name = body.name
    if body.entity_type is not None:
        ke.entity_type = body.entity_type
    if body.description is not None:
        ke.description = body.description
    if body.tags is not None:
        ke.metadata_ = body.tags
    if body.confidence is not None:
        ke.confidence = body.confidence
    db.commit()
    db.refresh(ke)
    return _ke_to_item(ke)


@router.delete("/{agent_id}/memory/knowledge-entities/{entity_id}")
def delete_knowledge_entity(
    agent_id: str,
    entity_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    ke = db.query(KnowledgeEntity).filter(
        KnowledgeEntity.id == entity_id,
        KnowledgeEntity.scope_kind == SCOPE_AGENT,
        KnowledgeEntity.scope_id == scope_id,
    ).first()
    if not ke:
        raise HTTPException(status_code=404, detail="实体不存在")
    # Also remove related relations
    db.query(KnowledgeRelation).filter(
        (KnowledgeRelation.source_entity_id == entity_id)
        | (KnowledgeRelation.target_entity_id == entity_id)
    ).delete()
    db.delete(ke)
    db.commit()
    return {"ok": True}


# ─────────── Learning Patterns ───────────


@router.get("/{agent_id}/memory/learning-patterns", response_model=LearningPatternList)
def list_learning_patterns(
    agent_id: str,
    skip: int = 0,
    limit: int = 50,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    q = db.query(AgentLearningPattern).filter(
        AgentLearningPattern.scope_kind == SCOPE_AGENT,
        AgentLearningPattern.scope_id == scope_id,
    )
    total = q.count()
    items = q.order_by(AgentLearningPattern.updated_at.desc()).offset(skip).limit(limit).all()
    return LearningPatternList(
        items=[
            LearningPatternItem(
                id=lp.id,
                scope_kind=lp.scope_kind or SCOPE_AGENT,
                scope_id=lp.scope_id or "",
                task_category=lp.task_category or "",
                task_keywords=lp.task_keywords,
                suggested_tools=lp.suggested_tools,
                effectiveness_score=lp.effectiveness_score or 0.0,
                total_runs=lp.total_runs or 0,
                successful_runs=lp.successful_runs or 0,
                avg_iterations=lp.avg_iterations or 0.0,
                avg_tool_calls=lp.avg_tool_calls or 0.0,
                last_used_at=lp.last_used_at.isoformat() if lp.last_used_at else None,
                created_at=lp.created_at.isoformat() if lp.created_at else None,
            )
            for lp in items
        ],
        total=total,
    )


@router.delete("/{agent_id}/memory/learning-patterns/{pattern_id}")
def delete_learning_pattern(
    agent_id: str,
    pattern_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    lp = db.query(AgentLearningPattern).filter(
        AgentLearningPattern.id == pattern_id,
        AgentLearningPattern.scope_kind == SCOPE_AGENT,
        AgentLearningPattern.scope_id == scope_id,
    ).first()
    if not lp:
        raise HTTPException(status_code=404, detail="学习模式不存在")
    db.delete(lp)
    db.commit()
    return {"ok": True}


# ─────────── Vector Memories ───────────


@router.get("/{agent_id}/memory/vector-memories", response_model=VectorMemoryList)
def list_vector_memories(
    agent_id: str,
    skip: int = 0,
    limit: int = 50,
    search: Optional[str] = None,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    q = db.query(AgentVectorMemory).filter(
        AgentVectorMemory.scope_kind == SCOPE_AGENT,
        AgentVectorMemory.scope_id == scope_id,
    )
    if search:
        q = q.filter(AgentVectorMemory.content_text.contains(search))
    total = q.count()
    items = q.order_by(AgentVectorMemory.created_at.desc()).offset(skip).limit(limit).all()
    return VectorMemoryList(items=[_vm_to_item(it) for it in items], total=total)


@router.delete("/{agent_id}/memory/vector-memories/{memory_id}")
def delete_vector_memory(
    agent_id: str,
    memory_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    vm = (
        db.query(AgentVectorMemory)
        .filter(
            AgentVectorMemory.id == memory_id,
            AgentVectorMemory.scope_kind == SCOPE_AGENT,
            AgentVectorMemory.scope_id == scope_id,
        )
        .first()
    )
    if not vm:
        raise HTTPException(status_code=404, detail="向量记忆不存在")
    db.delete(vm)
    db.commit()
    return {"ok": True}


# ─────────── Import / Export ───────────


def _serialize(row):
    """Serialize a SQLAlchemy row to dict, handling datetime."""
    d = {}
    for c in row.__table__.columns:
        val = getattr(row, c.name)
        if isinstance(val, datetime):
            val = val.isoformat()
        d[c.name] = val
    return d


@router.post("/{agent_id}/memory/export", response_model=MemoryExportResponse)
def export_memory(
    agent_id: str,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    gk_list = (
        db.query(GlobalKnowledge)
        .filter(GlobalKnowledge.scope_kind == SCOPE_AGENT, GlobalKnowledge.scope_id == scope_id)
        .all()
    )
    ke_list = (
        db.query(KnowledgeEntity)
        .filter(KnowledgeEntity.scope_kind == SCOPE_AGENT, KnowledgeEntity.scope_id == scope_id)
        .all()
    )
    kr_list = (
        db.query(KnowledgeRelation)
        .filter(KnowledgeRelation.scope_kind == SCOPE_AGENT, KnowledgeRelation.scope_id == scope_id)
        .all()
    )
    lp_list = (
        db.query(AgentLearningPattern)
        .filter(AgentLearningPattern.scope_kind == SCOPE_AGENT, AgentLearningPattern.scope_id == scope_id)
        .all()
    )
    vm_list = (
        db.query(AgentVectorMemory)
        .filter(
            AgentVectorMemory.scope_kind == SCOPE_AGENT,
            AgentVectorMemory.scope_id == scope_id,
        )
        .all()
    )
    return MemoryExportResponse(
        agent_id=agent_id,
        exported_at=datetime.utcnow().isoformat(),
        global_knowledge=[_serialize(gk) for gk in gk_list],
        knowledge_entities=[_serialize(ke) for ke in ke_list],
        knowledge_relations=[_serialize(kr) for kr in kr_list],
        learning_patterns=[_serialize(lp) for lp in lp_list],
        vector_memories=[_serialize(vm) for vm in vm_list],
    )


@router.post("/{agent_id}/memory/import")
def import_memory(
    agent_id: str,
    body: MemoryImportRequest,
    db: Session = Depends(get_db),
    current_user: User = Depends(get_current_user),
):
    _check_agent(db, agent_id, current_user)
    scope_id = _mem_scope(current_user, agent_id)
    imported = {"global_knowledge": 0, "knowledge_entities": 0, "knowledge_relations": 0, "learning_patterns": 0}

    for item in body.global_knowledge:
        gk = GlobalKnowledge(
            id=str(uuid.uuid4()),
            content=item.get("content", ""),
            tags=item.get("tags"),
            confidence=item.get("confidence", "medium"),
            scope_kind=SCOPE_AGENT,
            scope_id=scope_id,
            source_agent_id=item.get("source_agent_id"),
            source_user_id=current_user.id,
        )
        db.add(gk)
        imported["global_knowledge"] += 1

    for item in body.knowledge_entities:
        ke = KnowledgeEntity(
            id=str(uuid.uuid4()),
            name=item.get("name", ""),
            entity_type=item.get("entity_type", "concept"),
            description=item.get("description"),
            source="imported",
            confidence=item.get("confidence", "medium"),
            metadata_=item.get("metadata") or item.get("tags"),
            scope_kind=SCOPE_AGENT,
            scope_id=scope_id,
            user_id=current_user.id,
        )
        db.add(ke)
        imported["knowledge_entities"] += 1

    for item in body.knowledge_relations:
        kr = KnowledgeRelation(
            id=str(uuid.uuid4()),
            source_entity_id=item.get("source_entity_id", ""),
            target_entity_id=item.get("target_entity_id", ""),
            relation_type=item.get("relation_type", "related_to"),
            description=item.get("description"),
            weight=item.get("weight", "1.0"),
            scope_kind=SCOPE_AGENT,
            scope_id=scope_id,
        )
        db.add(kr)
        imported["knowledge_relations"] += 1

    for item in body.learning_patterns:
        lp = AgentLearningPattern(
            id=str(uuid.uuid4()),
            scope_kind=SCOPE_AGENT,
            scope_id=scope_id,
            task_category=item.get("task_category", "general"),
            task_keywords=item.get("task_keywords", ""),
            suggested_tools=item.get("suggested_tools", "[]"),
            effectiveness_score=item.get("effectiveness_score", 0.0),
            total_runs=item.get("total_runs", 1),
            successful_runs=item.get("successful_runs", 1),
            avg_iterations=item.get("avg_iterations", 1.0),
            avg_tool_calls=item.get("avg_tool_calls", 1.0),
        )
        db.add(lp)
        imported["learning_patterns"] += 1

    imported["vector_memories"] = 0
    for item in body.vector_memories:
        vm = AgentVectorMemory(
            id=str(uuid.uuid4()),
            scope_kind=SCOPE_AGENT,
            scope_id=scope_id,
            session_key=item.get("session_key", ""),
            content_text=item.get("content_text", ""),
            metadata_=item.get("metadata"),
        )
        db.add(vm)
        imported["vector_memories"] += 1

    db.commit()
    return {"ok": True, "imported": imported}