loom-agent 0.0.1__py3-none-any.whl

loom/builtin/memory/persistent_memory.py

@@ -0,0 +1,278 @@
+"""US6: Three-Tier Memory System
+
+Provides a practical memory system with automatic persistence for agent conversations.
+
+Tiers:
+1. Short-term: In-memory message array (current session)
+2. Mid-term: Compression summaries with metadata (managed by CompressionManager)
+3. Long-term: JSON file persistence for cross-session recall
+
+Design goals:
+- Simple API for developers
+- Automatic backup and recovery
+- Zero-config defaults with customization options
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from pathlib import Path
+from typing import List, Optional
+from datetime import datetime
+import asyncio
+
+from loom.core.types import Message
+from loom.interfaces.memory import BaseMemory
+
+
+class PersistentMemory(BaseMemory):
+    """Three-tier memory with automatic persistence.
+
+    Example:
+        # Simple usage - auto-creates .loom directory
+        memory = PersistentMemory()
+
+        # Custom persistence path
+        memory = PersistentMemory(persist_dir=".my_agent_memory")
+
+        # Disable persistence
+        memory = PersistentMemory(enable_persistence=False)
+    """
+
+    def __init__(
+        self,
+        persist_dir: str = ".loom",
+        session_id: Optional[str] = None,
+        enable_persistence: bool = True,
+        auto_backup: bool = True,
+        max_backup_files: int = 5,
+    ):
+        """Initialize persistent memory.
+
+        Args:
+            persist_dir: Directory for persisting memory (default: .loom)
+            session_id: Session identifier (default: auto-generated timestamp)
+            enable_persistence: Enable file persistence (default: True)
+            auto_backup: Create backup before overwriting (default: True)
+            max_backup_files: Maximum backup files to keep (default: 5)
+        """
+        self.persist_dir = Path(persist_dir)
+        self.session_id = session_id or datetime.now().strftime("%Y%m%d_%H%M%S")
+        self.enable_persistence = enable_persistence
+        self.auto_backup = auto_backup
+        self.max_backup_files = max_backup_files
+
+        # Tier 1: Short-term (in-memory)
+        self._messages: List[Message] = []
+
+        # Tier 2: Mid-term (compression metadata - managed externally)
+        self._compression_metadata: List[dict] = []
+
+        # Setup persistence
+        if self.enable_persistence:
+            self._ensure_persist_dir()
+            self._load_from_disk()
+
+        self._lock = asyncio.Lock()
+
+    def _ensure_persist_dir(self) -> None:
+        """Create persistence directory if it doesn't exist."""
+        self.persist_dir.mkdir(parents=True, exist_ok=True)
+
+    def _get_memory_file(self) -> Path:
+        """Get path to memory file."""
+        return self.persist_dir / f"session_{self.session_id}.json"
+
+    def _get_backup_file(self, index: int) -> Path:
+        """Get path to backup file."""
+        return self.persist_dir / f"session_{self.session_id}.backup{index}.json"
+
+    def _load_from_disk(self) -> None:
+        """Load memory from disk if exists."""
+        memory_file = self._get_memory_file()
+        if not memory_file.exists():
+            return
+
+        try:
+            with open(memory_file, 'r', encoding='utf-8') as f:
+                data = json.load(f)
+
+            # Load messages
+            self._messages = [
+                Message(**msg_data) for msg_data in data.get('messages', [])
+            ]
+
+            # Load compression metadata
+            self._compression_metadata = data.get('compression_metadata', [])
+
+        except Exception as e:
+            # Try to recover from backup
+            if self._recover_from_backup():
+                return
+            # If recovery fails, start fresh
+            print(f"Warning: Failed to load memory from disk: {e}")
+            self._messages = []
+            self._compression_metadata = []
+
+    def _save_to_disk(self) -> None:
+        """Save memory to disk with optional backup."""
+        if not self.enable_persistence:
+            return
+
+        memory_file = self._get_memory_file()
+
+        try:
+            # Create backup if file exists
+            if self.auto_backup and memory_file.exists():
+                self._create_backup()
+
+            # Save current state
+            data = {
+                'session_id': self.session_id,
+                'timestamp': datetime.now().isoformat(),
+                'messages': [self._message_to_dict(m) for m in self._messages],
+                'compression_metadata': self._compression_metadata,
+            }
+
+            with open(memory_file, 'w', encoding='utf-8') as f:
+                json.dump(data, f, indent=2, ensure_ascii=False)
+
+        except Exception as e:
+            print(f"Warning: Failed to save memory to disk: {e}")
+
+    def _message_to_dict(self, message: Message) -> dict:
+        """Convert Message to JSON-serializable dict."""
+        return {
+            'role': message.role,
+            'content': message.content,
+            'tool_call_id': message.tool_call_id,
+            'metadata': message.metadata,
+        }
+
+    def _create_backup(self) -> None:
+        """Create backup of current memory file."""
+        memory_file = self._get_memory_file()
+        if not memory_file.exists():
+            return
+
+        # Rotate existing backups
+        for i in range(self.max_backup_files - 1, 0, -1):
+            old_backup = self._get_backup_file(i)
+            new_backup = self._get_backup_file(i + 1)
+            if old_backup.exists():
+                old_backup.rename(new_backup)
+
+        # Create new backup
+        backup_file = self._get_backup_file(1)
+        memory_file.rename(backup_file)
+
+        # Clean up old backups
+        self._cleanup_old_backups()
+
+    def _cleanup_old_backups(self) -> None:
+        """Remove backups exceeding max_backup_files."""
+        for i in range(self.max_backup_files + 1, self.max_backup_files + 10):
+            backup_file = self._get_backup_file(i)
+            if backup_file.exists():
+                backup_file.unlink()
+
+    def _recover_from_backup(self) -> bool:
+        """Attempt to recover from most recent backup.
+
+        Returns:
+            True if recovery successful, False otherwise
+        """
+        for i in range(1, self.max_backup_files + 1):
+            backup_file = self._get_backup_file(i)
+            if not backup_file.exists():
+                continue
+
+            try:
+                with open(backup_file, 'r', encoding='utf-8') as f:
+                    data = json.load(f)
+
+                self._messages = [
+                    Message(**msg_data) for msg_data in data.get('messages', [])
+                ]
+                self._compression_metadata = data.get('compression_metadata', [])
+
+                print(f"Successfully recovered from backup {i}")
+                return True
+
+            except Exception as e:
+                print(f"Failed to recover from backup {i}: {e}")
+                continue
+
+        return False
+
+    async def add_message(self, message: Message) -> None:
+        """Add message to memory and persist."""
+        async with self._lock:
+            self._messages.append(message)
+            self._save_to_disk()
+
+    async def get_messages(self) -> List[Message]:
+        """Get all messages from memory."""
+        async with self._lock:
+            return self._messages.copy()
+
+    async def clear(self) -> None:
+        """Clear all messages from memory."""
+        async with self._lock:
+            self._messages.clear()
+            self._compression_metadata.clear()
+            self._save_to_disk()
+
+    async def set_messages(self, messages: List[Message]) -> None:
+        """Replace all messages in memory.
+
+        Used by CompressionManager when compressing history.
+        """
+        async with self._lock:
+            self._messages = messages.copy()
+            self._save_to_disk()
+
+    def add_compression_metadata(self, metadata: dict) -> None:
+        """Add compression metadata (Tier 2).
+
+        Called by CompressionManager to track compression events.
+        """
+        self._compression_metadata.append({
+            'timestamp': datetime.now().isoformat(),
+            **metadata
+        })
+        self._save_to_disk()
+
+    def get_compression_history(self) -> List[dict]:
+        """Get compression history metadata."""
+        return self._compression_metadata.copy()
+
+    def get_persistence_info(self) -> dict:
+        """Get information about persistence state.
+
+        Useful for debugging and monitoring.
+        """
+        memory_file = self._get_memory_file()
+
+        backup_files = []
+        for i in range(1, self.max_backup_files + 1):
+            backup = self._get_backup_file(i)
+            if backup.exists():
+                backup_files.append({
+                    'index': i,
+                    'path': str(backup),
+                    'size_bytes': backup.stat().st_size,
+                    'modified': datetime.fromtimestamp(backup.stat().st_mtime).isoformat(),
+                })
+
+        return {
+            'enabled': self.enable_persistence,
+            'session_id': self.session_id,
+            'persist_dir': str(self.persist_dir),
+            'memory_file': str(memory_file),
+            'memory_file_exists': memory_file.exists(),
+            'message_count': len(self._messages),
+            'compression_event_count': len(self._compression_metadata),
+            'backups': backup_files,
+        }

loom/builtin/retriever/__init__.py

@@ -0,0 +1,9 @@
+"""内置检索器实现"""
+
+from loom.builtin.retriever.in_memory import InMemoryRetriever
+
+try:
+    from loom.builtin.retriever.vector_store import VectorStoreRetriever
+    __all__ = ["InMemoryRetriever", "VectorStoreRetriever"]
+except ImportError:
+    __all__ = ["InMemoryRetriever"]

loom/builtin/retriever/chroma_store.py

@@ -0,0 +1,265 @@
+"""ChromaDB 向量存储适配器"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional, Tuple
+import uuid
+
+from loom.interfaces.retriever import Document
+from loom.interfaces.vector_store import BaseVectorStore
+
+try:
+    import chromadb
+    from chromadb.config import Settings
+    CHROMA_AVAILABLE = True
+except ImportError:
+    CHROMA_AVAILABLE = False
+
+
+class ChromaVectorStore(BaseVectorStore):
+    """
+    ChromaDB 向量存储适配器
+
+    特点:
+    - ✅ 开源嵌入式向量数据库
+    - ✅ 极简 API
+    - ✅ 支持本地持久化
+    - ✅ 自带 Embedding 功能（可选）
+    - ✅ 适合快速原型开发
+
+    示例:
+        from loom.builtin.retriever.chroma_store import ChromaVectorStore
+        from loom.builtin.retriever.vector_store_config import ChromaConfig
+
+        # 本地持久化模式
+        config = ChromaConfig.create_local(
+            persist_directory="./chroma_db",
+            collection_name="loom_docs"
+        )
+
+        # 远程服务模式
+        config = ChromaConfig.create_remote(
+            host="localhost",
+            port=8000,
+            collection_name="loom_docs"
+        )
+
+        vector_store = ChromaVectorStore(config)
+        await vector_store.initialize()
+    """
+
+    def __init__(self, config: Dict[str, Any] | Any):
+        """
+        Parameters:
+            config: ChromaConfig 对象或配置字典
+        """
+        if not CHROMA_AVAILABLE:
+            raise ImportError(
+                "ChromaDB is not installed. "
+                "Install with: pip install chromadb"
+            )
+
+        # 支持字典或 Pydantic 模型
+        if hasattr(config, "model_dump"):
+            self.config = config.model_dump()
+        else:
+            self.config = config
+
+        self.collection_name = self.config.get("collection_name", "loom_documents")
+        self.dimension = self.config.get("dimension", 1536)
+        self.client_type = self.config.get("client_type", "local")
+        self.persist_directory = self.config.get("persist_directory")
+        self.host = self.config.get("host")
+        self.port = self.config.get("port", 8000)
+
+        self.client: Optional[Any] = None
+        self.collection: Optional[Any] = None
+        self._initialized = False
+
+    async def initialize(self) -> None:
+        """初始化 ChromaDB 连接和集合"""
+        if self._initialized:
+            return
+
+        # 初始化客户端
+        if self.client_type == "local":
+            # 本地持久化模式
+            if self.persist_directory:
+                self.client = chromadb.PersistentClient(
+                    path=self.persist_directory
+                )
+            else:
+                # 内存模式（不持久化）
+                self.client = chromadb.Client()
+        else:
+            # 远程 HTTP 模式
+            self.client = chromadb.HttpClient(
+                host=self.host,
+                port=self.port
+            )
+
+        # 获取或创建集合
+        try:
+            self.collection = self.client.get_collection(
+                name=self.collection_name
+            )
+        except Exception:
+            # 集合不存在，创建新集合
+            self.collection = self.client.create_collection(
+                name=self.collection_name,
+                metadata={"dimension": self.dimension}
+            )
+
+        self._initialized = True
+
+    async def add_vectors(
+        self,
+        vectors: List[List[float]],
+        documents: List[Document]
+    ) -> None:
+        """
+        添加向量到 ChromaDB
+
+        Parameters:
+            vectors: 向量列表
+            documents: 对应的文档列表
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        # 构建 ChromaDB 数据格式
+        ids = []
+        embeddings = []
+        documents_text = []
+        metadatas = []
+
+        for i, (vector, doc) in enumerate(zip(vectors, documents)):
+            # 生成或使用文档 ID
+            doc_id = doc.doc_id or str(uuid.uuid4())
+
+            ids.append(doc_id)
+            embeddings.append(vector)
+            documents_text.append(doc.content)
+
+            # 构建元数据
+            metadata = doc.metadata or {}
+            if doc.score is not None:
+                metadata["score"] = doc.score
+            metadatas.append(metadata)
+
+        # 批量添加（ChromaDB 自动处理批量）
+        self.collection.add(
+            ids=ids,
+            embeddings=embeddings,
+            documents=documents_text,
+            metadatas=metadatas
+        )
+
+    async def search(
+        self,
+        query_vector: List[float],
+        top_k: int = 5,
+        filters: Optional[Dict[str, Any]] = None
+    ) -> List[Tuple[Document, float]]:
+        """
+        搜索相似向量
+
+        Parameters:
+            query_vector: 查询向量
+            top_k: 返回结果数量
+            filters: 元数据过滤条件
+
+        Returns:
+            [(Document, score), ...] 列表
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        # 构建 where 过滤条件（ChromaDB 格式）
+        where = None
+        if filters:
+            where = self._build_chroma_filter(filters)
+
+        # 执行查询
+        results = self.collection.query(
+            query_embeddings=[query_vector],
+            n_results=top_k,
+            where=where,
+            include=["documents", "metadatas", "distances"]
+        )
+
+        # 转换结果
+        documents_with_scores = []
+
+        if results["ids"]:
+            for i, doc_id in enumerate(results["ids"][0]):
+                content = results["documents"][0][i]
+                metadata = results["metadatas"][0][i] or {}
+                distance = results["distances"][0][i]
+
+                # 转换距离到相似度分数（ChromaDB 返回的是距离）
+                # 对于余弦距离，相似度 = 1 - distance
+                score = 1.0 - distance
+
+                doc = Document(
+                    content=content,
+                    metadata=metadata,
+                    score=score,
+                    doc_id=doc_id
+                )
+                documents_with_scores.append((doc, score))
+
+        return documents_with_scores
+
+    async def delete(self, doc_ids: List[str]) -> None:
+        """
+        删除文档
+
+        Parameters:
+            doc_ids: 文档 ID 列表
+        """
+        if not self._initialized:
+            await self.initialize()
+
+        self.collection.delete(ids=doc_ids)
+
+    async def clear(self) -> None:
+        """清空集合"""
+        if not self._initialized:
+            await self.initialize()
+
+        # ChromaDB: 删除并重建集合
+        self.client.delete_collection(name=self.collection_name)
+        self.collection = self.client.create_collection(
+            name=self.collection_name,
+            metadata={"dimension": self.dimension}
+        )
+
+    def _build_chroma_filter(self, filters: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        构建 ChromaDB where 过滤器
+
+        ChromaDB 过滤语法:
+        {
+            "field": "value",
+            "numeric_field": {"$gte": 10}
+        }
+
+        支持的操作符: $eq, $ne, $gt, $gte, $lt, $lte, $in, $nin
+        """
+        where = {}
+
+        for key, value in filters.items():
+            if isinstance(value, dict):
+                # 复杂查询（例如: {"price": {"$gte": 100}}）
+                where[key] = value
+            else:
+                # 简单相等查询
+                where[key] = value
+
+        return where
+
+    async def close(self) -> None:
+        """关闭连接"""
+        # ChromaDB 客户端自动管理连接
+        self._initialized = False

loom/builtin/retriever/in_memory.py

@@ -0,0 +1,106 @@
+"""简单的内存检索器 - 无需外部依赖"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from loom.interfaces.retriever import BaseRetriever, Document
+
+
+class InMemoryRetriever(BaseRetriever):
+    """
+    简单的内存检索器 - 基于关键词匹配
+
+    特点:
+    - 无需外部依赖
+    - 适用于开发/测试
+    - 基于简单的关键词匹配 (非向量检索)
+
+    示例:
+        retriever = InMemoryRetriever()
+        await retriever.add_texts([
+            "Python is a programming language",
+            "JavaScript is used for web development"
+        ])
+
+        docs = await retriever.retrieve("programming")
+        # 返回包含 "programming" 的文档
+    """
+
+    def __init__(self):
+        self.documents: List[Document] = []
+
+    async def retrieve(
+        self,
+        query: str,
+        top_k: int = 5,
+        filters: Optional[Dict[str, Any]] = None,
+    ) -> List[Document]:
+        """
+        基于关键词匹配检索文档
+
+        简单算法:
+        1. 将查询分词
+        2. 计算每个文档包含的关键词数量
+        3. 按匹配度排序
+        """
+        if not self.documents:
+            return []
+
+        # 简单分词 (按空格)
+        query_terms = set(query.lower().split())
+
+        # 计算每个文档的匹配分数
+        scored_docs = []
+        for doc in self.documents:
+            # 应用过滤器
+            if filters and not self._match_filters(doc, filters):
+                continue
+
+            # 计算匹配分数
+            doc_terms = set(doc.content.lower().split())
+            matches = query_terms.intersection(doc_terms)
+            score = len(matches) / len(query_terms) if query_terms else 0.0
+
+            if score > 0:
+                # 创建副本并设置分数
+                doc_with_score = Document(
+                    content=doc.content,
+                    metadata=doc.metadata,
+                    score=score,
+                    doc_id=doc.doc_id
+                )
+                scored_docs.append(doc_with_score)
+
+        # 按分数排序
+        scored_docs.sort(key=lambda d: d.score or 0, reverse=True)
+
+        return scored_docs[:top_k]
+
+    async def add_documents(self, documents: List[Document]) -> None:
+        """添加文档到内存"""
+        for doc in documents:
+            # 分配 ID (如果没有)
+            if doc.doc_id is None:
+                doc.doc_id = str(len(self.documents))
+
+            self.documents.append(doc)
+
+    def _match_filters(self, doc: Document, filters: Dict[str, Any]) -> bool:
+        """检查文档是否匹配过滤条件"""
+        if not doc.metadata:
+            return False
+
+        for key, value in filters.items():
+            if doc.metadata.get(key) != value:
+                return False
+
+        return True
+
+    def clear(self) -> None:
+        """清空所有文档"""
+        self.documents.clear()
+
+    def __len__(self) -> int:
+        """返回文档数量"""
+        return len(self.documents)