Maas-Agent/apps/application/services/application_memory_service.py

# coding=utf-8
"""
智能体记忆服务
提供记忆的 CRUD 操作和检索功能
"""
from typing import List, Optional
from django.db.models import QuerySet
from application.models.application_memory import ApplicationMemory, ApplicationMemoryType
from common.utils.logger import maxkb_logger


class ApplicationMemoryService:
    """智能体记忆服务"""

    @staticmethod
    def get_memories(application_id: str, memory_type: str = None,
                     is_enabled: bool = True, limit: int = 100) -> QuerySet:
        """获取应用的记忆列表"""
        queryset = ApplicationMemory.objects.filter(
            application_id=application_id,
            is_enabled=is_enabled
        )
        if memory_type:
            queryset = queryset.filter(memory_type=memory_type)
        return queryset.order_by('-relevance_score', '-create_time')[:limit]

    @staticmethod
    def get_memory(memory_id: str) -> Optional[ApplicationMemory]:
        """获取单条记忆"""
        try:
            return ApplicationMemory.objects.get(id=memory_id)
        except ApplicationMemory.DoesNotExist:
            return None

    @staticmethod
    def create_memory(application_id: str, content: str,
                      memory_type: str = ApplicationMemoryType.DIALOGUE,
                      metadata: dict = None) -> ApplicationMemory:
        """创建记忆"""
        return ApplicationMemory.objects.create(
            application_id=application_id,
            content=content,
            memory_type=memory_type,
            metadata=metadata or {}
        )

    @staticmethod
    def update_memory(memory_id: str, content: str = None,
                      memory_type: str = None, is_enabled: bool = None,
                      relevance_score: float = None,
                      metadata: dict = None) -> Optional[ApplicationMemory]:
        """更新记忆"""
        memory = ApplicationMemoryService.get_memory(memory_id)
        if not memory:
            return None

        if content is not None:
            memory.content = content
        if memory_type is not None:
            memory.memory_type = memory_type
        if is_enabled is not None:
            memory.is_enabled = is_enabled
        if relevance_score is not None:
            memory.relevance_score = relevance_score
        if metadata is not None:
            memory.metadata = metadata

        memory.save()
        return memory

    @staticmethod
    def delete_memory(memory_id: str) -> bool:
        """删除记忆"""
        memory = ApplicationMemoryService.get_memory(memory_id)
        if not memory:
            return False
        memory.delete()
        return True

    @staticmethod
    def batch_delete_memories(memory_ids: List[str]) -> int:
        """批量删除记忆"""
        return ApplicationMemory.objects.filter(id__in=memory_ids).delete()[0]

    @staticmethod
    def search_memories(application_id: str, query: str,
                        limit: int = 10) -> List[ApplicationMemory]:
        """搜索记忆（简单文本匹配）"""
        return list(
            ApplicationMemory.objects.filter(
                application_id=application_id,
                is_enabled=True,
                content__icontains=query
            ).order_by('-relevance_score', '-create_time')[:limit]
        )

    @staticmethod
    def get_memory_context(application_id: str, query: str = None,
                           max_tokens: int = 2000) -> str:
        """获取记忆上下文（用于注入到对话 prompt）"""
        memories = ApplicationMemoryService.get_memories(
            application_id=application_id,
            is_enabled=True
        )

        if query:
            # 如果有查询，优先返回相关记忆
            relevant = ApplicationMemoryService.search_memories(
                application_id=application_id,
                query=query
            )
            if relevant:
                memories = relevant

        context_parts = []
        current_tokens = 0

        for memory in memories:
            # 简单估算 token 数（中文约 1.5 字符/token）
            estimated_tokens = len(memory.content) / 1.5
            if current_tokens + estimated_tokens > max_tokens:
                break
            context_parts.append(memory.content)
            current_tokens += estimated_tokens

        return "\n".join(context_parts)