solana-agent 27.2.0__py3-none-any.whl → 27.3.1__py3-none-any.whl

solana_agent/services/knowledge_base.py

@@ -0,0 +1,771 @@
+from datetime import datetime as dt
+from typing import Dict, List, Any, Optional, Union
+import uuid
+import asyncio
+import io
+
+# PDF and Chunking Imports
+import pypdf
+from llama_index.core import Document as LlamaDocument
+from llama_index.core.node_parser import SemanticSplitterNodeParser
+from llama_index.embeddings.openai import OpenAIEmbedding
+
+from solana_agent.adapters.pinecone_adapter import PineconeAdapter
+from solana_agent.adapters.mongodb_adapter import MongoDBAdapter
+from solana_agent.interfaces.services.knowledge_base import KnowledgeBaseService as KnowledgeBaseInterface
+
+
+class KnowledgeBaseService(KnowledgeBaseInterface):
+    """
+    Knowledge Base service using Pinecone for vector search and MongoDB for metadata/full document storage.
+    Supports text documents and PDF semantic chunking using OpenAI embeddings via LlamaIndex.
+    """
+
+    def __init__(
+        self,
+        pinecone_adapter: PineconeAdapter,
+        mongodb_adapter: MongoDBAdapter,
+        openai_api_key: str,
+        openai_model_name: str = "text-embedding-3-large",
+        collection_name: str = "knowledge_documents",
+        rerank_results: bool = False,
+        rerank_top_k: int = 3,
+        # Semantic Splitter Config
+        splitter_buffer_size: int = 1,
+        splitter_breakpoint_percentile: int = 95,
+    ):
+        """
+        Initialize the Knowledge Base service.
+
+        Args:
+            pinecone_adapter: Configured PineconeAdapter instance.
+            mongodb_adapter: Configured MongoDBAdapter instance.
+            openai_api_key: OpenAI API key for embedding.
+            openai_model_name: OpenAI embedding model name.
+            collection_name: MongoDB collection for storing document metadata and full PDFs.
+            rerank_results: Whether PineconeAdapter should rerank results.
+            rerank_top_k: Number of results to return after reranking (used by PineconeAdapter).
+            splitter_buffer_size: Buffer size for SemanticSplitterNodeParser.
+            splitter_breakpoint_percentile: Breakpoint percentile for SemanticSplitterNodeParser.
+        """
+        self.pinecone = pinecone_adapter
+        self.mongo = mongodb_adapter
+        self.collection = collection_name
+        # Store rerank config for reference, but PineconeAdapter applies it
+        self.rerank_results = rerank_results
+        self.rerank_top_k = rerank_top_k
+
+        # --- Initialize Semantic Splitter with OpenAIEmbedding ---
+        api_key = openai_api_key
+        if not api_key:
+            raise ValueError(
+                "OpenAI API key not provided via argument or OPENAI_API_KEY environment variable.")
+
+        # Determine expected embedding dimensions based on model name
+        if openai_model_name == "text-embedding-3-large":
+            openai_dimensions = 3072
+        elif openai_model_name == "text-embedding-3-small":
+            openai_dimensions = 1536
+        else:
+            # Attempt to get dimension from Pinecone config if available, otherwise raise error
+            openai_dimensions = getattr(
+                self.pinecone, 'embedding_dimensions', 0)
+            if openai_dimensions <= 0:
+                raise ValueError(
+                    f"Cannot determine dimension for unknown OpenAI model '{openai_model_name}' and Pinecone dimension not configured.")
+            print(
+                f"Warning: Unknown OpenAI model '{openai_model_name}'. Using dimension {openai_dimensions} from Pinecone config. Ensure this is correct.")
+
+        # Instantiate OpenAIEmbedding
+        # Note: LlamaIndex OpenAIEmbedding doesn't directly support reducing dimensions via 'dimensions' param during init
+        # like the raw OpenAI client does. It uses the model's default.
+        try:
+            llama_embed_model = OpenAIEmbedding(
+                model=openai_model_name,
+                api_key=api_key,
+                # embed_batch_size=10 # Optional: Adjust batch size if needed
+            )
+        except Exception as e:
+            print(f"Error initializing OpenAIEmbedding: {e}")
+            raise
+
+        self.semantic_splitter = SemanticSplitterNodeParser(
+            buffer_size=splitter_buffer_size,
+            breakpoint_percentile_threshold=splitter_breakpoint_percentile,
+            embed_model=llama_embed_model  # Use the OpenAIEmbedding instance
+        )
+        # Store model name for logging/reference
+        self.openai_model_name = openai_model_name
+        # --- End Semantic Splitter Init ---
+
+        self._ensure_collection()
+
+    def _ensure_collection(self) -> None:
+        """Set up MongoDB collection with appropriate indexes."""
+        if not self.mongo.collection_exists(self.collection):
+            self.mongo.create_collection(self.collection)
+            print(f"Created MongoDB collection: {self.collection}")
+
+        # Indexes for retrieval and filtering
+        self.mongo.create_index(
+            self.collection, [("document_id", 1)], unique=True)
+        self.mongo.create_index(self.collection, [("parent_document_id", 1)])
+        self.mongo.create_index(self.collection, [("source", 1)])
+        self.mongo.create_index(self.collection, [("created_at", -1)])
+        self.mongo.create_index(self.collection, [("tags", 1)])
+        self.mongo.create_index(self.collection, [("is_chunk", 1)])
+        print(
+            f"Ensured indexes exist for MongoDB collection: {self.collection}")
+
+    async def add_document(
+        self,
+        text: str,
+        metadata: Dict[str, Any],
+        document_id: Optional[str] = None,
+        namespace: Optional[str] = None
+    ) -> str:
+        """
+        Add a plain text document to the knowledge base. Embeds using OpenAI.
+
+        Args:
+            text: Document text content.
+            metadata: Document metadata.
+            document_id: Optional document ID.
+            namespace: Optional Pinecone namespace.
+
+        Returns:
+            The document ID.
+        """
+        doc_id = document_id or str(uuid.uuid4())
+
+        # Store metadata and content in MongoDB
+        mongo_doc = {
+            "document_id": doc_id,
+            "content": text,
+            "is_chunk": False,
+            "parent_document_id": None,
+            **metadata,
+            # Use timezone aware datetime
+            "created_at": metadata.get("created_at", dt.now(tz=dt.now().astimezone().tzinfo)),
+            "updated_at": dt.now(tz=dt.now().astimezone().tzinfo)
+        }
+        try:
+            self.mongo.insert_one(self.collection, mongo_doc)
+        except Exception as e:
+            print(f"Error inserting document {doc_id} into MongoDB: {e}")
+            raise
+
+        # Embed text using OpenAIEmbedding
+        embed_model: OpenAIEmbedding = self.semantic_splitter.embed_model
+        try:
+            embedding = await embed_model.aget_text_embedding(text)
+        except Exception as e:
+            print(
+                f"Error embedding document {doc_id} using {self.openai_model_name}: {e}")
+            # Decide how to handle - Mongo insert succeeded, embedding failed
+            raise  # Re-raise for now
+
+        # Prepare Pinecone metadata
+        pinecone_metadata = {
+            "document_id": doc_id,
+            "is_chunk": False,
+            "source": metadata.get("source", "unknown"),
+            "tags": metadata.get("tags", [])
+        }
+        # Add text itself if Pinecone adapter reranking is used
+        if self.pinecone.use_reranking:
+            pinecone_metadata[self.pinecone.rerank_text_field] = text
+
+        # Upsert vector to Pinecone using the generic upsert method
+        try:
+            await self.pinecone.upsert(
+                vectors=[{"id": doc_id, "values": embedding,
+                          "metadata": pinecone_metadata}],
+                namespace=namespace
+            )
+        except Exception as e:
+            print(f"Error upserting vector for {doc_id} to Pinecone: {e}")
+            # Decide how to handle - Mongo insert succeeded, Pinecone failed
+            raise  # Re-raise for now
+
+        return doc_id
+
+    async def add_pdf_document(
+        self,
+        pdf_data: Union[bytes, str],
+        metadata: Dict[str, Any],
+        document_id: Optional[str] = None,
+        namespace: Optional[str] = None,
+        chunk_batch_size: int = 50
+    ) -> str:
+        """
+        Add a PDF document, performs semantic chunking using OpenAI embeddings,
+        stores full PDF in Mongo, and chunk vectors in Pinecone.
+
+        Args:
+            pdf_data: PDF content as bytes or a path to the PDF file.
+            metadata: Document metadata (applies to the parent PDF).
+            document_id: Optional parent document ID.
+            namespace: Optional Pinecone namespace for chunks.
+            chunk_batch_size: Batch size for embedding and upserting chunks.
+
+        Returns:
+            The parent document ID.
+        """
+        parent_doc_id = document_id or str(uuid.uuid4())
+        pdf_bytes: bytes
+
+        # --- 1. Read PDF and Extract Text ---
+        try:
+            if isinstance(pdf_data, str):
+                with open(pdf_data, "rb") as f:
+                    pdf_bytes = f.read()
+            elif isinstance(pdf_data, bytes):
+                pdf_bytes = pdf_data
+            else:
+                raise ValueError(
+                    "pdf_data must be bytes or a file path string.")
+
+            reader = pypdf.PdfReader(io.BytesIO(pdf_bytes))
+            extracted_text = "".join(
+                page.extract_text() or "" for page in reader.pages)
+            if not extracted_text.strip():
+                print(f"Warning: No text extracted from PDF {parent_doc_id}.")
+        except Exception as e:
+            print(
+                f"Error reading or extracting text from PDF {parent_doc_id}: {e}")
+            raise
+
+        # --- 2. Store Full PDF and Metadata in MongoDB ---
+        mongo_parent_doc = {
+            "document_id": parent_doc_id,
+            "content": extracted_text,
+            "pdf_data": pdf_bytes,
+            "is_chunk": False,
+            "parent_document_id": None,
+            **metadata,
+            "created_at": metadata.get("created_at", dt.now(tz=dt.now().astimezone().tzinfo)),
+            "updated_at": dt.now(tz=dt.now().astimezone().tzinfo)
+        }
+        try:
+            self.mongo.insert_one(self.collection, mongo_parent_doc)
+            print(f"Stored full PDF {parent_doc_id} in MongoDB.")
+        except Exception as e:  # pragma: no cover
+            print(
+                f"Error inserting parent PDF {parent_doc_id} into MongoDB: {e}")  # pragma: no cover
+            raise  # pragma: no cover
+
+        # --- 3. Semantic Chunking ---
+        if not extracted_text.strip():
+            print(
+                f"Skipping chunking for PDF {parent_doc_id} due to no extracted text.")
+            return parent_doc_id
+
+        try:
+            llama_doc = LlamaDocument(text=extracted_text)
+            # Run synchronous splitter in thread pool
+            nodes = await asyncio.to_thread(
+                self.semantic_splitter.get_nodes_from_documents,
+                [llama_doc]
+            )
+            print(
+                f"Generated {len(nodes)} semantic chunks for PDF {parent_doc_id}.")
+        except Exception as e:
+            print(
+                f"Error during semantic chunking for PDF {parent_doc_id}: {e}")
+            raise
+
+        # --- 4. Embed Chunks and Batch Upsert to Pinecone ---
+        if not nodes:
+            return parent_doc_id  # No chunks generated
+
+        print(
+            f"Embedding {len(nodes)} chunks using {self.openai_model_name}...")
+        chunk_texts = [node.get_content() for node in nodes]
+        embed_model: OpenAIEmbedding = self.semantic_splitter.embed_model
+        all_chunk_embeddings = []
+
+        # Embed chunks in batches (using embed_model's internal batching)
+        try:
+            # Use aget_text_embedding_batch for async embedding
+            # Note: LlamaIndex OpenAIEmbedding might handle batch size internally.
+            # If large number of nodes, consider explicit batching here if needed.
+            all_chunk_embeddings = await embed_model.aget_text_embedding_batch(
+                chunk_texts, show_progress=True
+            )
+        except Exception as e:
+            print(f"Error embedding chunks for PDF {parent_doc_id}: {e}")
+            raise  # Stop if embedding fails
+
+        print("Embedding complete. Preparing vectors for Pinecone.")
+        pinecone_vectors = []
+        for i, node in enumerate(nodes):
+            chunk_id = f"{parent_doc_id}_chunk_{i}"
+            chunk_metadata = {
+                "document_id": chunk_id,
+                "parent_document_id": parent_doc_id,
+                "chunk_index": i,
+                "is_chunk": True,
+                "source": metadata.get("source", "unknown"),
+                "tags": metadata.get("tags", [])
+            }
+            # Add chunk text itself if Pinecone adapter reranking is used
+            if self.pinecone.use_reranking:
+                chunk_metadata[self.pinecone.rerank_text_field] = chunk_texts[i]
+
+            pinecone_vectors.append({
+                "id": chunk_id,
+                "values": all_chunk_embeddings[i],
+                "metadata": chunk_metadata
+            })
+
+        # Upsert vectors in batches using the generic upsert method
+        print(
+            f"Upserting {len(pinecone_vectors)} vectors to Pinecone in batches of {chunk_batch_size}...")
+        upsert_tasks = []
+        for i in range(0, len(pinecone_vectors), chunk_batch_size):
+            batch_vectors = pinecone_vectors[i:i+chunk_batch_size]
+            # Create task for each batch upsert
+            upsert_tasks.append(
+                self.pinecone.upsert(
+                    vectors=batch_vectors,
+                    namespace=namespace
+                )
+            )
+            # Optional: Add a small delay between initiating tasks if rate limiting is a concern
+            # await asyncio.sleep(0.05)
+
+        # Run upsert tasks concurrently
+        results = await asyncio.gather(*upsert_tasks, return_exceptions=True)
+
+        # Check for errors during upsert
+        for idx, result in enumerate(results):
+            if isinstance(result, Exception):
+                print(
+                    f"Error upserting vector batch {idx + 1} to Pinecone: {result}")
+                # Decide on error handling: log, raise, etc.
+
+        print(f"Finished processing PDF {parent_doc_id}.")
+        return parent_doc_id
+
+    async def query(
+        self,
+        query_text: str,
+        filter: Optional[Dict[str, Any]] = None,
+        top_k: int = 5,
+        namespace: Optional[str] = None,
+        include_content: bool = True,
+        include_metadata: bool = True
+    ) -> List[Dict[str, Any]]:
+        """
+        Query the knowledge base using semantic search with OpenAI embeddings.
+
+        Args:
+            query_text: The query text.
+            filter: Optional Pinecone metadata filter.
+            top_k: Number of results to retrieve initially.
+            namespace: Optional Pinecone namespace.
+            include_content: Whether to include document/chunk content in results.
+            include_metadata: Whether to include document/chunk metadata in results.
+
+        Returns:
+            List of result dictionaries, potentially reranked by PineconeAdapter.
+        """
+        # Determine initial K for Pinecone query, considering reranking multiplier
+        initial_k = top_k
+        if self.pinecone.use_reranking:
+            initial_k = top_k * self.pinecone.initial_query_top_k_multiplier
+
+        # --- Embed Query using OpenAIEmbedding ---
+        embed_model: OpenAIEmbedding = self.semantic_splitter.embed_model
+        try:
+            query_vector = await embed_model.aget_query_embedding(query_text)
+        except Exception as e:
+            print(f"Error embedding query text '{query_text}': {e}")
+            return []
+
+        # --- Query Pinecone using the vector ---
+        try:
+            # Use the generic query method with the vector
+            # PineconeAdapter handles reranking internally if configured
+            pinecone_results = await self.pinecone.query(
+                vector=query_vector,
+                filter=filter,
+                top_k=initial_k,  # Fetch more initially if reranking
+                namespace=namespace,
+                include_values=False,  # Don't need embeddings back
+                include_metadata=True  # Need metadata for linking
+            )
+        except Exception as e:
+            print(f"Error querying Pinecone: {e}")
+            return []
+
+        if not pinecone_results:
+            return []
+
+        # Extract IDs, scores, and metadata from Pinecone results
+        # PineconeAdapter might have already reranked and truncated to final top_k
+        result_ids = [res['id'] for res in pinecone_results]
+        scores = {res['id']: res['score'] for res in pinecone_results}
+        pinecone_metadatas = {res['id']: res.get(
+            'metadata', {}) for res in pinecone_results}
+
+        # --- Fetch corresponding data from MongoDB ---
+        mongo_docs_map = {}
+        parent_ids_to_fetch = set()
+        for res_id in result_ids:
+            meta = pinecone_metadatas.get(res_id, {})
+            if meta.get("is_chunk"):
+                parent_id = meta.get("parent_document_id")
+                if parent_id:
+                    parent_ids_to_fetch.add(parent_id)
+            else:
+                # If it's not a chunk, its own ID might be in Mongo
+                parent_ids_to_fetch.add(res_id)
+
+        if parent_ids_to_fetch:
+            try:
+                mongo_docs = self.mongo.find(
+                    self.collection,
+                    {"document_id": {"$in": list(parent_ids_to_fetch)}}
+                )
+                mongo_docs_map = {doc["document_id"]
+                    : doc for doc in mongo_docs}
+            except Exception as e:
+                print(f"Error fetching documents from MongoDB: {e}")
+                # Proceed with potentially missing Mongo data
+
+        # --- Combine Results ---
+        results = []
+        for res_id in result_ids:  # Iterate in Pinecone's score order
+            pinecone_meta = pinecone_metadatas.get(res_id, {})
+            is_chunk = pinecone_meta.get("is_chunk", False)
+            parent_doc_id = pinecone_meta.get("parent_document_id")
+
+            # Determine which Mongo doc holds the relevant info
+            mongo_doc_for_meta = None
+            mongo_doc_for_content = None
+            if is_chunk and parent_doc_id:
+                mongo_doc_for_meta = mongo_docs_map.get(parent_doc_id)
+                mongo_doc_for_content = mongo_doc_for_meta  # Parent holds full content
+            else:  # Not a chunk
+                mongo_doc_for_meta = mongo_docs_map.get(res_id)
+                mongo_doc_for_content = mongo_doc_for_meta
+
+            result = {
+                "document_id": res_id,
+                "score": scores.get(res_id, 0.0),
+                "is_chunk": is_chunk,
+                "parent_document_id": parent_doc_id,
+            }
+
+            if include_content:
+                content = None
+                # Priority 1: Reranking field in Pinecone metadata (holds chunk text)
+                if self.pinecone.use_reranking and self.pinecone.rerank_text_field in pinecone_meta:
+                    content = pinecone_meta[self.pinecone.rerank_text_field]
+                # Priority 2: Get content from the relevant Mongo doc
+                elif mongo_doc_for_content:
+                    content = mongo_doc_for_content.get("content")
+                result["content"] = content or ""
+
+            if include_metadata:
+                combined_meta = {}
+                # Merge metadata from the relevant Mongo doc (parent or self)
+                if mongo_doc_for_meta:
+                    combined_meta = {
+                        k: v for k, v in mongo_doc_for_meta.items()
+                        if k not in ["_id", "document_id", "content", "pdf_data", "is_chunk", "parent_document_id", "created_at", "updated_at"]
+                    }
+                # Add/overwrite with chunk-specific info from Pinecone meta
+                combined_meta.update({
+                    k: v for k, v in pinecone_meta.items()
+                    # Avoid redundancy
+                    if k not in ["document_id", self.pinecone.rerank_text_field]
+                })
+                result["metadata"] = combined_meta
+
+            results.append(result)
+
+        return results
+
+    async def delete_document(
+        self,
+        document_id: str,
+        namespace: Optional[str] = None
+    ) -> bool:
+        """
+        Delete a document (plain text or PDF) and all its associated chunks.
+
+        Args:
+            document_id: ID of the parent document (or plain text document).
+            namespace: Optional Pinecone namespace.
+
+        Returns:
+            True if deletion was successful (or partially successful).
+        """
+        print(
+            f"Attempting to delete document and associated data for ID: {document_id}")
+        mongo_deleted_count = 0
+        pinecone_deleted = False
+
+        # --- 1. Find all associated document IDs in MongoDB ---
+        # This includes the parent doc and potentially chunk metadata if we stored it
+        # We primarily need the IDs to delete from Pinecone.
+        # A more robust way might be to query Pinecone directly for vectors with parent_document_id == document_id
+        # For now, assume IDs in Mongo cover what needs deletion.
+        docs_to_delete_mongo = []
+        mongo_ids_to_delete = set([document_id])  # Start with the main ID
+        try:
+            docs_to_delete_mongo = list(self.mongo.find(
+                self.collection,
+                {"$or": [{"document_id": document_id}, {
+                    "parent_document_id": document_id}]}
+            ))
+            for doc in docs_to_delete_mongo:
+                mongo_ids_to_delete.add(doc["document_id"])
+        except Exception as e:
+            print(
+                f"Warning: Error finding documents in MongoDB for deletion ({document_id}): {e}. Proceeding with main ID only.")
+
+        pinecone_ids_to_delete = list(mongo_ids_to_delete)
+
+        # --- 2. Delete from Pinecone ---
+        if pinecone_ids_to_delete:
+            try:
+                await self.pinecone.delete(ids=pinecone_ids_to_delete, namespace=namespace)
+                print(
+                    f"Deleted {len(pinecone_ids_to_delete)} vectors from Pinecone for parent {document_id}.")
+                pinecone_deleted = True
+            except Exception as e:
+                print(
+                    f"Error deleting vectors from Pinecone for {document_id}: {e}")
+
+        # --- 3. Delete from MongoDB ---
+        # Use the IDs confirmed to be in Mongo
+        mongo_ids_found_in_db = [doc["document_id"]
+                                 for doc in docs_to_delete_mongo]
+        if mongo_ids_found_in_db:
+            try:
+                delete_result = self.mongo.delete_many(
+                    self.collection,
+                    {"document_id": {"$in": mongo_ids_found_in_db}}
+                )
+                mongo_deleted_count = delete_result.deleted_count
+                print(
+                    f"Deleted {mongo_deleted_count} documents from MongoDB for parent {document_id}.")
+            except Exception as e:
+                print(
+                    f"Error deleting documents from MongoDB for {document_id}: {e}")
+
+        return pinecone_deleted or mongo_deleted_count > 0
+
+    async def update_document(
+        self,
+        document_id: str,
+        text: Optional[str] = None,
+        metadata: Optional[Dict[str, Any]] = None,
+        namespace: Optional[str] = None
+    ) -> bool:
+        """
+        Update an existing plain text document or metadata. Embeds using OpenAI.
+        Updating PDF content requires deleting and re-adding.
+
+        Args:
+            document_id: ID of document to update.
+            text: Optional new text content (for plain text docs only).
+            metadata: Optional metadata to update.
+            namespace: Optional Pinecone namespace.
+
+        Returns:
+            True if successful.
+        """
+        current_doc = self.mongo.find_one(
+            self.collection, {"document_id": document_id})
+        if not current_doc:
+            print(f"Document {document_id} not found for update.")
+            return False
+
+        if current_doc.get("is_chunk"):
+            print(f"Cannot update chunk {document_id} directly.")
+            return False
+        if current_doc.get("pdf_data") and text is not None:
+            print(f"Cannot update PDF content via this method. Delete and re-add.")
+            return False
+
+        update_text = text is not None and not current_doc.get("pdf_data")
+        text_content = text if update_text else current_doc.get("content", "")
+
+        # --- Update MongoDB ---
+        mongo_update = {}
+        if metadata:
+            mongo_update.update(metadata)
+        if update_text:
+            mongo_update["content"] = text_content
+        mongo_update["updated_at"] = dt.now(tz=dt.now().astimezone().tzinfo)
+
+        mongo_updated = False
+        if mongo_update:  # Only update if there are changes
+            try:
+                update_result = self.mongo.update_one(
+                    self.collection, {"document_id": document_id}, {
+                        "$set": mongo_update}
+                )
+                mongo_updated = update_result.modified_count > 0
+            except Exception as e:
+                print(f"Error updating document {document_id} in MongoDB: {e}")
+                # Decide if we should proceed to Pinecone update if Mongo failed
+                return False  # Return False if Mongo update fails
+
+        # --- Update Pinecone (only if text changed) ---
+        pinecone_updated = False
+        if update_text:
+            # Embed updated text
+            embed_model: OpenAIEmbedding = self.semantic_splitter.embed_model
+            try:
+                embedding = await embed_model.aget_text_embedding(text_content)
+            except Exception as e:
+                print(f"Error embedding updated text for {document_id}: {e}")
+                # Mongo update might have succeeded, but embedding failed
+                return mongo_updated  # Return based on Mongo success
+
+            # Prepare Pinecone metadata
+            final_metadata = {**current_doc, **
+                              mongo_update}  # Use updated data
+            pinecone_metadata = {"document_id": document_id, "is_chunk": False}
+            for key, value in final_metadata.items():
+                if key not in ["_id", "content", "pdf_data", "created_at", "updated_at", "document_id", "is_chunk", "parent_document_id"]:
+                    pinecone_metadata[key] = value
+            if self.pinecone.use_reranking:
+                pinecone_metadata[self.pinecone.rerank_text_field] = text_content
+
+            # Upsert vector to Pinecone
+            try:
+                await self.pinecone.upsert(
+                    vectors=[{"id": document_id, "values": embedding,
+                              "metadata": pinecone_metadata}],
+                    namespace=namespace
+                )
+                pinecone_updated = True
+            except Exception as e:
+                print(
+                    f"Error upserting updated vector in Pinecone for {document_id}: {e}")
+                # Mongo update succeeded, Pinecone failed
+
+        return mongo_updated or pinecone_updated
+
+    async def add_documents_batch(
+        self,
+        # Expects {'text': ..., 'metadata': ...}
+        documents: List[Dict[str, Any]],
+        namespace: Optional[str] = None,
+        batch_size: int = 50
+    ) -> List[str]:
+        """
+        Add multiple plain text documents in batches using OpenAI embeddings.
+
+        Args:
+            documents: List of documents, each with 'text' and 'metadata'.
+            namespace: Optional Pinecone namespace.
+            batch_size: Number of documents per embedding/upsert batch.
+
+        Returns:
+            List of added document IDs.
+        """
+        all_doc_ids = []
+        embed_model: OpenAIEmbedding = self.semantic_splitter.embed_model
+
+        for i in range(0, len(documents), batch_size):
+            batch_docs_input = documents[i:i+batch_size]
+            batch_texts = [doc['text'] for doc in batch_docs_input]
+            batch_metadatas = [doc['metadata'] for doc in batch_docs_input]
+            # Generate IDs if not provided in metadata
+            batch_doc_ids = [doc['metadata'].get('document_id') or str(
+                uuid.uuid4()) for doc in batch_docs_input]
+            all_doc_ids.extend(batch_doc_ids)
+
+            # Prepare MongoDB docs
+            mongo_batch = []
+            for idx, text in enumerate(batch_texts):
+                doc_id = batch_doc_ids[idx]
+                metadata = batch_metadatas[idx]
+                mongo_doc = {
+                    "document_id": doc_id, "content": text, "is_chunk": False,
+                    "parent_document_id": None, **metadata,
+                    "created_at": metadata.get("created_at", dt.now(tz=dt.now().astimezone().tzinfo)),
+                    "updated_at": dt.now(tz=dt.now().astimezone().tzinfo)
+                }
+                # Ensure generated ID is in the doc for Mongo
+                if 'document_id' not in metadata:
+                    mongo_doc['document_id'] = doc_id
+                mongo_batch.append(mongo_doc)
+
+            # Insert into MongoDB
+            if mongo_batch:
+                try:
+                    self.mongo.insert_many(self.collection, mongo_batch)
+                except Exception as e:
+                    print(
+                        f"Error inserting batch {i//batch_size + 1} into MongoDB: {e}")
+                    # Decide if we should skip Pinecone for this batch
+                    continue  # Skip to next batch
+
+            # Embed batch using OpenAIEmbedding
+            try:
+                batch_embeddings = await embed_model.aget_text_embedding_batch(batch_texts, show_progress=True)
+            except Exception as e:
+                print(
+                    f"Error embedding batch {i//batch_size + 1} using {self.openai_model_name}: {e}")
+                continue  # Skip Pinecone upsert for this batch
+
+            # Prepare Pinecone vectors
+            pinecone_vectors = []
+            for idx, doc_id in enumerate(batch_doc_ids):
+                metadata = batch_metadatas[idx]
+                pinecone_meta = {
+                    "document_id": doc_id, "is_chunk": False,
+                    "source": metadata.get("source", "unknown"),
+                    "tags": metadata.get("tags", [])
+                }
+                if self.pinecone.use_reranking:
+                    pinecone_meta[self.pinecone.rerank_text_field] = batch_texts[idx]
+
+                pinecone_vectors.append({
+                    "id": doc_id,
+                    "values": batch_embeddings[idx],
+                    "metadata": pinecone_meta
+                })
+
+            # Upsert vectors to Pinecone
+            if pinecone_vectors:
+                try:
+                    await self.pinecone.upsert(
+                        vectors=pinecone_vectors,
+                        namespace=namespace
+                    )
+                except Exception as e:
+                    print(
+                        f"Error upserting vector batch {i//batch_size + 1} to Pinecone: {e}")
+
+            # Optional delay
+            if i + batch_size < len(documents):
+                await asyncio.sleep(0.1)
+
+        return all_doc_ids
+
+    async def get_full_document(self, document_id: str) -> Optional[Dict[str, Any]]:
+        """
+        Retrieve the full document entry (including PDF data if applicable) from MongoDB.
+
+        Args:
+            document_id: The ID of the document (parent ID if it was a PDF).
+
+        Returns:
+            The document dictionary from MongoDB, or None if not found.
+        """
+        try:
+            return self.mongo.find_one(self.collection, {"document_id": document_id})
+        except Exception as e:
+            print(
+                f"Error retrieving full document {document_id} from MongoDB: {e}")
+            return None