ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter3/retrieval-pipeline/document_store.py

@@ -0,0 +1,98 @@
+"""Document store for the retrieval pipeline."""
+
+from typing import Dict, Any, List, Optional
+from datetime import datetime
+import logging
+
+logger = logging.getLogger(__name__)
+
+class DocumentStore:
+    """In-memory document store for educational purposes."""
+    
+    def __init__(self):
+        self.documents: Dict[str, Dict[str, Any]] = {}
+        self.metadata_index: Dict[str, List[str]] = {}  # Index by metadata fields
+        
+    def add_document(self, doc_id: str, text: str, metadata: Optional[Dict[str, Any]] = None) -> None:
+        """Add a document to the store."""
+        new_metadata = metadata or {}
+        if doc_id in self.documents:
+            old_metadata = self.documents[doc_id].get("metadata") or {}
+            for key in list(old_metadata.keys()):
+                if key not in new_metadata:
+                    if key in self.metadata_index and doc_id in self.metadata_index[key]:
+                        self.metadata_index[key].remove(doc_id)
+                        if not self.metadata_index[key]:
+                            del self.metadata_index[key]
+
+        self.documents[doc_id] = {
+            "doc_id": doc_id,
+            "text": text,
+            "metadata": new_metadata,
+            "indexed_at": datetime.now().isoformat()
+        }
+        
+        # Update metadata index
+        for key in new_metadata:
+            if key not in self.metadata_index:
+                self.metadata_index[key] = []
+            if doc_id not in self.metadata_index[key]:
+                self.metadata_index[key].append(doc_id)
+                    
+        logger.debug(f"Added document {doc_id} to store")
+        
+    def get_document(self, doc_id: str) -> Optional[Dict[str, Any]]:
+        """Get a document by ID."""
+        return self.documents.get(doc_id)
+    
+    def get_documents(self, doc_ids: List[str]) -> List[Dict[str, Any]]:
+        """Get multiple documents by IDs."""
+        docs = []
+        for doc_id in doc_ids:
+            doc = self.get_document(doc_id)
+            if doc:
+                docs.append(doc)
+        return docs
+    
+    def delete_document(self, doc_id: str) -> bool:
+        """Delete a document from the store."""
+        if doc_id in self.documents:
+            doc = self.documents[doc_id]
+            
+            # Remove from metadata index
+            if doc.get("metadata"):
+                for key in doc["metadata"]:
+                    if key in self.metadata_index and doc_id in self.metadata_index[key]:
+                        self.metadata_index[key].remove(doc_id)
+                        if not self.metadata_index[key]:
+                            del self.metadata_index[key]
+            
+            del self.documents[doc_id]
+            logger.debug(f"Deleted document {doc_id} from store")
+            return True
+        return False
+    
+    def list_documents(self, limit: int = 100, offset: int = 0) -> List[Dict[str, Any]]:
+        """List documents with pagination."""
+        doc_ids = list(self.documents.keys())[offset:offset + limit]
+        return [self.documents[doc_id] for doc_id in doc_ids]
+    
+    def clear(self) -> None:
+        """Clear all documents."""
+        self.documents.clear()
+        self.metadata_index.clear()
+        logger.info("Cleared all documents from store")
+        
+    def size(self) -> int:
+        """Get the number of documents."""
+        return len(self.documents)
+    
+    def get_stats(self) -> Dict[str, Any]:
+        """Get store statistics."""
+        return {
+            "total_documents": self.size(),
+            "metadata_fields": list(self.metadata_index.keys()),
+            "metadata_distribution": {
+                key: len(values) for key, values in self.metadata_index.items()
+            }
+        }