AIエージェントのメモリ管理：長期記憶を持つシステム設計

from typing import List, Dict, Optional
from datetime import datetime

class WorkingMemory:
    """作業メモリ：現在のセッション内の情報を管理"""

    def __init__(self, max_size: int = 10):
        self.messages: List[Dict[str, str]] = []
        self.max_size = max_size
        self.current_task: Optional[str] = None

    def add_message(self, role: str, content: str) -> None:
        """新しいメッセージを追加し、容量超過時は古いメッセージを削除"""
        self.messages.append({"role": role, "content": content})
        if len(self.messages) > self.max_size:
            self.messages.pop(0)

    def get_context(self) -> List[Dict[str, str]]:
        """現在のコンテキストを返す"""
        return self.messages.copy()

    def clear(self) -> None:
        """セッション終了時にクリア"""
        self.messages.clear()
        self.current_task = None

import chromadb
from chromadb.config import Settings
from sentence_transformers import SentenceTransformer
import logging
from datetime import datetime
from typing import List, Dict, Optional

logger = logging.getLogger(__name__)

class SemanticMemory:
    """セマンティックメモリ：ベクトルデータベースを活用した長期記憶"""

    def __init__(
        self,
        collection_name: str = "agent_memories",
        embedding_model: str = "sentence-transformers/all-MiniLM-L6-v2",
        persist_directory: Optional[str] = None
    ):
        self.embedding_model = SentenceTransformer(embedding_model)
        self.dimension = 384  # all-MiniLM-L6-v2 の出力次元数

        if persist_directory:
            self.client = chromadb.PersistentClient(path=persist_directory)
        else:
            self.client = chromadb.Client()

        try:
            self.collection = self.client.get_collection(name=collection_name)
            logger.info(f"既存のコレクション '{collection_name}' をロードしました")
        except Exception:
            self.collection = self.client.create_collection(
                name=collection_name,
                metadata={"dimension": self.dimension}
            )
            logger.info(f"新しいコレクション '{collection_name}' を作成しました")

    def store(
        self,
        content: str,
        metadata: Optional[Dict[str, str]] = None,
        importance: float = 1.0
    ) -> str:
        """新しい記憶を保存する"""
        try:
            embedding = self.embedding_model.encode(content).tolist()
            doc_id = f"mem_{datetime.now().strftime('%Y%m%d%H%M%S%f')}"

            self.collection.add(
                ids=[doc_id],
                embeddings=[embedding],
                documents=[content],
                metadatas=[{
                    **(metadata or {}),
                    "importance": importance,
                    "created_at": datetime.now().isoformat()
                }]
            )

            logger.info(f"記憶を保存しました: {doc_id}")
            return doc_id

        except Exception as e:
            logger.error(f"記憶の保存に失敗しました: {e}")
            raise MemoryStoreError(f"Failed to store memory: {e}")

    def retrieve(
        self,
        query: str,
        top_k: int = 5,
        filter_metadata: Optional[Dict[str, str]] = None
    ) -> List[Dict]:
        """クエリに関連する記憶を検索する"""
        try:
            query_embedding = self.embedding_model.encode(query).tolist()
            where_clause = filter_metadata if filter_metadata else None

            results = self.collection.query(
                query_embeddings=[query_embedding],
                n_results=top_k,
                where=where_clause,
                include=["documents", "metadatas", "distances"]
            )

            memories = []
            for i, doc in enumerate(results["documents"][0]):
                memories.append({
                    "content": doc,
                    "metadata": results["metadatas"][0][i],
                    "similarity": 1 - results["distances"][0][i],
                    "id": results["ids"][0][i]
                })

            logger.info(f"{len(memories)} 件の関連記憶を検索しました")
            return memories

        except Exception as e:
            logger.error(f"記憶の検索に失敗しました: {e}")
            raise MemoryRetrieveError(f"Failed to retrieve memories: {e}")

from enum import Enum
from dataclasses import dataclass, field
from typing import List, Dict, Optional, Any
from datetime import datetime

class MemoryType(Enum):
    EPISODIC = "episodic"
    SEMANTIC = "semantic"
    PROCEDURAL = "procedural"

@dataclass
class ProceduralMemory:
    """手続き記憶：タスク実行のパターンを管理"""

    patterns: Dict[str, Dict[str, Any]] = field(default_factory=dict)

    def register_pattern(
        self,
        task_type: str,
        approach: str,
        success_rate: float,
        execution_history: List[Dict] = None
    ) -> None:
        """成功したタスク実行パターンを登録"""
        if task_type not in self.patterns:
            self.patterns[task_type] = {
                "approaches": [],
                "total_executions": 0,
                "successful_executions": 0
            }

        pattern_entry = {
            "approach": approach,
            "last_used": datetime.now().isoformat(),
            "success_count": 0
        }

        self.patterns[task_type]["approaches"].append(pattern_entry)
        self.patterns[task_type]["total_executions"] += 1
        self.patterns[task_type]["successful_executions"] += (
            1 if success_rate > 0.7 else 0
        )

    def get_best_approach(self, task_type: str) -> Optional[str]:
        """成功率の高いアプローチを提案"""
        if task_type not in self.patterns:
            return None

        approaches = self.patterns[task_type]["approaches"]
        if not approaches:
            return None

        best = max(approaches, key=lambda x: x.get("success_count", 0))
        return best["approach"]

import os
from typing import List, Dict, Optional
from dataclasses import dataclass
import logging
from datetime import datetime

from .working_memory import WorkingMemory
from .semantic_memory import SemanticMemory
from .procedural_memory import ProceduralMemory, MemoryType

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

@dataclass
class AgentConfig:
    """エージェント設定"""
    max_context_messages: int = 10
    memory_top_k: int = 5
    memory_importance_threshold: float = 0.5
    semantic_memory_path: str = "./data/semantic_memory"

class LongTermMemoryAgent:
    """長期記憶を持つAIエージェント"""

    def __init__(self, llm_client, config: Optional[AgentConfig] = None):
        self.config = config or AgentConfig()
        self.llm = llm_client

        self.working_memory = WorkingMemory(max_size=self.config.max_context_messages)
        self.semantic_memory = SemanticMemory(persist_directory=self.config.semantic_memory_path)
        self.procedural_memory = ProceduralMemory()

        self.session_id = datetime.now().strftime("%Y%m%d%H%M%S")
        logger.info(f"エージェントを初期化しました: セッション {self.session_id}")

    async def process(
        self,
        user_input: str,
        user_id: Optional[str] = None,
        task_type: Optional[str] = None
    ) -> str:
        """ユーザー入力処理のメインエントリーポイント"""
        try:
            self.working_memory.add_message("user", user_input)

            relevant_memories = self.semantic_memory.retrieve(
                query=user_input,
                top_k=self.config.memory_top_k,
                filter_metadata={"user_id": user_id} if user_id else None
            )

            best_approach = None
            if task_type:
                best_approach = self.procedural_memory.get_best_approach(task_type)

            prompt = self._build_prompt(
                user_input=user_input,
                memories=relevant_memories,
                best_approach=best_approach
            )

            response = await self.llm.generate(prompt)

            self.working_memory.add_message("assistant", response)

            if self._should_memorize(user_input, response):
                self._store_important_memory(
                    user_input=user_input,
                    response=response,
                    user_id=user_id,
                    task_type=task_type
                )

            return response

        except Exception as e:
            logger.error(f"処理中にエラーが発生しました: {e}")
            return f"申し訳ありません。処理中にエラーが発生しました: {str(e)}"