plugins/memory/document_memory.py

#!/usr/bin/env python3
"""Document Memory Plugin"""
from typing import Dict, Any, Optional

class DocumentMemory:
    """Cache and manage uploaded documents."""
    def __init__(self):
        self.documents = {}
        self.metadata = {}
    
    def store_document(self, doc_id: str, content: Any, metadata: Dict[str, Any]):
        """Store document with metadata."""
        self.documents[doc_id] = content
        self.metadata[doc_id] = metadata
    
    def get_document(self, doc_id: str) -> Optional[Any]:
        """Retrieve document by ID."""
        return self.documents.get(doc_id)
    
    def list_documents(self) -> Dict[str, Dict[str, Any]]:
        """List all stored documents with metadata."""
        return {doc_id: self.metadata[doc_id] for doc_id in self.documents.keys()}
    
    def clear(self):
        """Clear all documents."""
        self.documents.clear()
        self.metadata.clear()