"""RAG (Retrieval-Augmented Generation) tools for ChemGraph. Provides tools to load documents (.txt and .pdf) into a FAISS vector store and query them for relevant context. Supports OpenAI and HuggingFace embeddings with automatic fallback. """ import os import logging from typing import Optional from langchain_core.tools import tool from pydantic import BaseModel, Field logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Module-level vector store registry # --------------------------------------------------------------------------- # Maps a document identifier (file path or user-provided name) to a # FAISS vector store instance so that documents loaded during a session # remain queryable across multiple tool calls. _vector_stores: dict = {} # --------------------------------------------------------------------------- # Pydantic schemas for tool inputs # --------------------------------------------------------------------------- class LoadDocumentInput(BaseModel): """Input schema for the load_document tool.""" file_path: str = Field( description="Absolute or relative path to a .txt or .pdf file to ingest." ) chunk_size: int = Field( default=1000, description="Maximum number of characters per text chunk.", ) chunk_overlap: int = Field( default=200, description="Number of overlapping characters between consecutive chunks.", ) embedding_provider: str = Field( default="openai", description=( "Embedding provider to use: 'openai' (requires OPENAI_API_KEY) " "or 'huggingface' (local, no API key needed). " "Falls back to huggingface if openai is unavailable." ), ) class QueryKnowledgeBaseInput(BaseModel): """Input schema for the query_knowledge_base tool.""" query: str = Field(description="The question or search query.") file_path: Optional[str] = Field( default=None, description=( "Path of a previously loaded document to search. " "If None, searches the most recently loaded document." ), ) top_k: int = Field( default=5, description="Number of most relevant chunks to retrieve.", ) # --------------------------------------------------------------------------- # Supported file types # --------------------------------------------------------------------------- _SUPPORTED_EXTENSIONS = {".txt", ".pdf"} # --------------------------------------------------------------------------- # PDF text extraction # --------------------------------------------------------------------------- def _extract_text_from_pdf(file_path: str) -> str: """Extract text content from a PDF file using PyMuPDF. Parameters ---------- file_path : str Absolute path to the PDF file. Returns ------- str Concatenated text from all pages, separated by newlines. Raises ------ ImportError If PyMuPDF (``fitz``) is not installed. """ try: import fitz # PyMuPDF except ImportError as exc: raise ImportError( "PyMuPDF is required for PDF support. " "Install the 'rag' extra: pip install chemgraphagent[rag]" ) from exc pages: list[str] = [] with fitz.open(file_path) as doc: for page_num, page in enumerate(doc): page_text = page.get_text() if page_text.strip(): pages.append(page_text) return "\n\n".join(pages) # --------------------------------------------------------------------------- # Embedding helpers # --------------------------------------------------------------------------- def _get_embeddings(provider: str = "openai"): """Return an embeddings instance for the requested provider. Supports OpenAI-compatible custom endpoints via OPENAI_BASE_URL. Falls back to HuggingFace if OpenAI embeddings are unavailable. Parameters ---------- provider : str, optional Preferred embedding provider. Returns ------- Embeddings LangChain-compatible embeddings object. """ if provider == "openai": try: from langchain_openai import OpenAIEmbeddings api_key = os.environ.get("OPENAI_API_KEY") base_url = os.environ.get("OPENAI_BASE_URL") if not api_key: raise EnvironmentError("OPENAI_API_KEY not set") kwargs = { "model": os.environ.get("OPENAI_EMBEDDING_MODEL", "text-embedding-3-large"), "api_key": api_key, "check_embedding_ctx_length":False, } if base_url: kwargs["base_url"] = base_url return OpenAIEmbeddings(**kwargs) except Exception as exc: logger.warning( "OpenAI embeddings unavailable (%s); falling back to HuggingFace.", exc, ) provider = "huggingface" try: from langchain_huggingface import HuggingFaceEmbeddings return HuggingFaceEmbeddings( model_name="all-MiniLM-L6-v2", model_kwargs={"device": "cpu"}, ) except ImportError as exc: raise ImportError( "Neither langchain-openai nor langchain-huggingface is installed. " "Install the 'rag' extra: pip install chemgraphagent[rag]" ) from exc # --------------------------------------------------------------------------- # Tools # --------------------------------------------------------------------------- @tool(args_schema=LoadDocumentInput) def load_document( file_path: str, chunk_size: int = 1000, chunk_overlap: int = 200, embedding_provider: str = "openai", ) -> dict: """Load a document (.txt or .pdf), split it into chunks, and index it in a FAISS vector store. The document remains available for querying via ``query_knowledge_base`` for the duration of the session. Parameters ---------- file_path : str Path to the ``.txt`` or ``.pdf`` file to ingest. chunk_size : int, optional Max characters per chunk, by default 1000. chunk_overlap : int, optional Overlap between consecutive chunks, by default 200. embedding_provider : str, optional ``"openai"`` or ``"huggingface"``, by default ``"openai"``. Returns ------- dict Status information including the number of chunks created. """ from langchain_text_splitters import RecursiveCharacterTextSplitter from langchain_community.vectorstores import FAISS resolved_path = os.path.abspath(file_path) if not os.path.isfile(resolved_path): return {"ok": False, "error": f"File not found: {resolved_path}"} _, ext = os.path.splitext(resolved_path) ext = ext.lower() if ext not in _SUPPORTED_EXTENSIONS: supported = ", ".join(sorted(_SUPPORTED_EXTENSIONS)) return { "ok": False, "error": (f"Unsupported file type '{ext}'. Supported formats: {supported}"), } # ----- Extract text based on file type ----- if ext == ".pdf": try: text = _extract_text_from_pdf(resolved_path) except ImportError as exc: return {"ok": False, "error": str(exc)} except Exception as exc: return { "ok": False, "error": f"Failed to extract text from PDF: {exc}", } else: # .txt with open(resolved_path, "r", encoding="utf-8") as fh: text = fh.read() if not text.strip(): return {"ok": False, "error": "File is empty or contains no extractable text."} # Split into chunks splitter = RecursiveCharacterTextSplitter( chunk_size=chunk_size, chunk_overlap=chunk_overlap, length_function=len, separators=["\n\n", "\n", ". ", " ", ""], ) chunks = splitter.create_documents( [text], metadatas=[{"source": resolved_path, "file_type": ext}], ) # Build FAISS index embeddings = _get_embeddings(provider=embedding_provider) vector_store = FAISS.from_documents(chunks, embeddings) # Register in module-level store _vector_stores[resolved_path] = vector_store # Also track the most-recently loaded path for convenience _vector_stores["__latest__"] = resolved_path logger.info( "Loaded '%s' (%s) into FAISS vector store (%d chunks, chunk_size=%d, overlap=%d).", resolved_path, ext, len(chunks), chunk_size, chunk_overlap, ) return { "ok": True, "file_path": resolved_path, "file_type": ext, "num_chunks": len(chunks), "chunk_size": chunk_size, "chunk_overlap": chunk_overlap, "embedding_provider": embedding_provider, } @tool(args_schema=QueryKnowledgeBaseInput) def query_knowledge_base( query: str, file_path: Optional[str] = None, top_k: int = 5, ) -> dict: """Search a previously loaded document for passages relevant to a query. Parameters ---------- query : str The natural-language question or search query. file_path : str, optional Path of a previously loaded document. If ``None``, the most recently loaded document is searched. top_k : int, optional Number of top-matching chunks to return, by default 5. Returns ------- dict A dict with ``"ok"``, ``"query"``, ``"num_results"``, and ``"results"`` (list of dicts with ``"content"`` and ``"metadata"``). """ # Resolve which vector store to query if file_path is not None: resolved_path = os.path.abspath(file_path) else: resolved_path = _vector_stores.get("__latest__") if resolved_path is None or resolved_path not in _vector_stores: available = [k for k in _vector_stores if k != "__latest__"] return { "ok": False, "error": ( "No document loaded yet. Use the load_document tool first." if not available else f"Document '{file_path}' not found. Available: {available}" ), } vector_store = _vector_stores[resolved_path] docs = vector_store.similarity_search(query, k=top_k) results = [ { "content": doc.page_content, "metadata": doc.metadata, } for doc in docs ] return { "ok": True, "query": query, "num_results": len(results), "results": results, } def get_loaded_documents() -> list[str]: """Return a list of file paths currently loaded in the vector store. This is a plain helper (not a tool) for programmatic access. """ return [k for k in _vector_stores if k != "__latest__"] def clear_vector_stores() -> None: """Remove all loaded vector stores. Useful for testing and cleanup.""" _vector_stores.clear()