Update vector_store_retriever.py

vector_store_retriever.py

@@ -1,195 +1,258 @@
-import json
 import os
-import gradio as gr
-import time
-import langchain
-
-from pydantic import BaseModel, Field
-from typing import Any, Optional, Dict, List, Union
-from huggingface_hub import InferenceClient
-from langchain.llms.base import LLM
-#from langchain.Images import Images
-
-from langchain.llms.base import LLM
-#from langchain_core.embeddings import EmbeddingFunction, Embeddings
-
-from langchain.embeddings import HuggingFaceInstructEmbeddings
-#from langchain import [all]
-#from langchain.Documents import Documents
-from langchain.vectorstores import Chroma
-from dotenv import load_dotenv
-from transformers import AutoTokenizer, AutoModel, Tool
-
-load_dotenv()
-
-path_work = "."
-hf_token = os.getenv("HF")
-
-class HuggingFaceInstructEmbeddings(HuggingFaceInstructEmbeddings):
-    def __init__(self, model_name: str, model_kwargs: Optional[Dict[str, Any]] = None):
-        self.model = AutoModel.from_pretrained(model_name, **(model_kwargs or {}))
-        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+from typing import Dict, List, Optional, Union, Any
+from smolagents import Tool
+from langchain.vectorstores import FAISS, Chroma
+from langchain.embeddings import HuggingFaceBgeEmbeddings, HuggingFaceEmbeddings
+from langchain.document_loaders import PyPDFLoader, TextLoader, DirectoryLoader
+from langchain.text_splitter import CharacterTextSplitter, RecursiveCharacterTextSplitter
+from PyPDF2 import PdfReader
+import json
 
-    def __call__(self, input: Union[Documents]) -> HuggingFaceInstructEmbeddings:
-        if isinstance(input, Documents):
-            texts = [doc.text for doc in input]
-            embeddings = self._embed_text(texts)
+class RAGTool(Tool):
+    name = "rag_retriever"
+    description = """
+    Advanced RAG (Retrieval-Augmented Generation) tool that searches in vector stores based on given prompts.
+    This tool allows you to query documents stored in vector databases using semantic similarity.
+    It supports various configurations including different embedding models, vector stores, and document types.
+    """
+    inputs = {
+        "query": {
+            "type": "string",
+            "description": "The search query to retrieve relevant information from the document store",
+        },
+        "top_k": {
+            "type": "integer",
+            "description": "Number of most relevant documents to retrieve (default: 3)",
+        }
+    }
+    output_type = "string"
+    
+    def __init__(self, 
+                 documents_path: str = "./documents",
+                 embedding_model: str = "BAAI/bge-small-en-v1.5",
+                 vector_store_type: str = "faiss",
+                 chunk_size: int = 1000,
+                 chunk_overlap: int = 200,
+                 persist_directory: str = "./vector_store",
+                 device: str = "cpu"):
+        """
+        Initialize the RAG Tool with configurable parameters.
+        
+        Args:
+            documents_path: Path to documents or folder containing documents
+            embedding_model: HuggingFace model ID for embeddings
+            vector_store_type: Type of vector store ('faiss' or 'chroma')
+            chunk_size: Size of text chunks for splitting documents
+            chunk_overlap: Overlap between text chunks
+            persist_directory: Directory to persist vector store
+            device: Device to run embedding model on ('cpu' or 'cuda')
+        """
+        super().__init__()
+        self.documents_path = documents_path
+        self.embedding_model = embedding_model
+        self.vector_store_type = vector_store_type
+        self.chunk_size = chunk_size
+        self.chunk_overlap = chunk_overlap
+        self.persist_directory = persist_directory
+        self.device = device
+        
+        # Create the vector store if it doesn't exist
+        os.makedirs(persist_directory, exist_ok=True)
+        self._setup_vector_store()
+        
+    def _setup_vector_store(self):
+        """Set up the vector store with documents if it doesn't exist"""
+        # Check if we need to create a new vector store
+        if not os.path.exists(os.path.join(self.persist_directory, "index.faiss")) and \
+           not os.path.exists(os.path.join(self.persist_directory, "chroma")):
+            # Check if documents path exists
+            if not os.path.exists(self.documents_path):
+                print(f"Warning: Documents path {self.documents_path} does not exist.")
+                return
+            
+            # Load and process documents
+            documents = self._load_documents()
+            if not documents:
+                print("No documents loaded. Vector store not created.")
+                return
+            
+            # Create the vector store
+            self._create_vector_store(documents)
         else:
-            # Handle image embeddings if needed
-            pass
-
-        return embeddings
-
-    def _embed_text(self, texts: List[str]) -> Embeddings:
-        # Your existing logic for text embeddings using Hugging Face models...
-        inputs = self.tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
-        with torch.no_grad():
-            outputs = self.model(**inputs)
-        embeddings = outputs.last_hidden_state.mean(dim=1)  # Adjust this based on your specific model
-
-        return embeddings
-
-
-vectordb = Chroma(
-    persist_directory=path_work + '/new_papers',
-    embedding_function=HuggingFaceInstructEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": "cpu"})
-)
-
-retriever = vectordb.as_retriever(search_kwargs={"k": 2})#5
-
-
-class KwArgsModel(BaseModel):
-    kwargs: Dict[str, Any] = Field(default_factory=dict)
-
-class CustomInferenceClient(LLM, KwArgsModel):
-    model_name: str
-    inference_client: InferenceClient
-
-    def __init__(self, model_name: str, hf_token: str, kwargs: Optional[Dict[str, Any]] = None):
-        inference_client = InferenceClient(model=model_name, token=hf_token)
-        super().__init__(
-            model_name=model_name,
-            hf_token=hf_token,
-            kwargs=kwargs,
-            inference_client=inference_client
-        )
-
-    def _call(
-        self,
-        prompt: str,
-        stop: Optional[List[str]] = None
-    ) -> str:
-        if stop is not None:
-            raise ValueError("stop kwargs are not permitted.")
-        response_gen = self.inference_client.text_generation(prompt, **self.kwargs, stream=True)
-        response = ''.join(response_gen)
-        return response
-
-    @property
-    def _llm_type(self) -> str:
-        return "custom"
-
-    @property
-    def _identifying_params(self) -> dict:
-        return {"model_name": self.model_name}
-
-kwargs = {"max_new_tokens": 256, "temperature": 0.9, "top_p": 0.6, "repetition_penalty": 1.3, "do_sample": True}
-
-model_list = [
-    "meta-llama/Llama-2-13b-chat-hf",
-    "HuggingFaceH4/zephyr-7b-alpha",
-    "meta-llama/Llama-2-70b-chat-hf",
-    "tiiuae/falcon-180B-chat"
-]
-
-qa_chain = None
-
-def load_model(model_selected):
-    global qa_chain
-    model_name = model_selected
-    llm = CustomInferenceClient(model_name=model_name, hf_token=hf_token, kwargs=kwargs)
-
-    from langchain.chains import RetrievalQA
-    qa_chain = RetrievalQA.from_chain_type(
-        llm=llm,
-        chain_type="stuff",
-        retriever=retriever,
-        return_source_documents=True,
-        verbose=True,
-    )
-    return qa_chain
-
-load_model("meta-llama/Llama-2-70b-chat-hf")
-
-##########
-#####
-#########
-
-
-###
-###
-###
-
-def predict(message, temperature=0.9, max_new_tokens=512, top_p=0.6, repetition_penalty=1.3):
-    temperature = float(temperature)
-    if temperature < 1e-2: temperature = 1e-2
-    top_p = float(top_p)
-
-    llm_response = qa_chain(message)
-    res_result = llm_response['result']
-
-    res_relevant_doc = [source.metadata['source'] for source in llm_response["source_documents"]]
-    response = f"{res_result}" + "\n\n" + "[Answer Source Documents (Ctrl + Click!)] :" + "\n" + f" \n {res_relevant_doc}"
-    print("response: =====> \n", response, "\n\n")
-    tokens = response.split('\n')
-    token_list = []
-    for idx, token in enumerate(tokens):
-        token_dict = {"id": idx + 1, "text": token}
-        token_list.append(token_dict)
-    response = {"data": {"token": token_list}}
-    response = json.dumps(response, indent=4)
-
-    response = json.loads(response)
-    data_dict = response.get('data', {})
-    token_list = data_dict.get('token', [])
-
-    partial_message = ""
-    for token_entry in token_list:
-        if token_entry:
+            print(f"Vector store already exists at {self.persist_directory}")
+            self._load_vector_store()
+    
+    def _get_embeddings(self):
+        """Get embedding model based on configuration"""
+        try:
+            if "bge" in self.embedding_model.lower():
+                encode_kwargs = {"normalize_embeddings": True}
+                return HuggingFaceBgeEmbeddings(
+                    model_name=self.embedding_model,
+                    encode_kwargs=encode_kwargs,
+                    model_kwargs={"device": self.device}
+                )
+            else:
+                return HuggingFaceEmbeddings(
+                    model_name=self.embedding_model,
+                    model_kwargs={"device": self.device}
+                )
+        except Exception as e:
+            print(f"Error loading embedding model: {e}")
+            # Fallback to a reliable model
+            print("Falling back to sentence-transformers/all-MiniLM-L6-v2")
+            return HuggingFaceEmbeddings(
+                model_name="sentence-transformers/all-MiniLM-L6-v2",
+                model_kwargs={"device": self.device}
+            )
+            
+    def _load_documents(self):
+        """Load documents from the documents path"""
+        documents = []
+        
+        # Check if documents_path is a file or directory
+        if os.path.isfile(self.documents_path):
+            # Load single file
+            if self.documents_path.lower().endswith('.pdf'):
+                try:
+                    loader = PyPDFLoader(self.documents_path)
+                    documents = loader.load()
+                except Exception as e:
+                    print(f"Error loading PDF: {e}")
+                    # Fallback to using PdfReader
+                    try:
+                        text = self._extract_text_from_pdf(self.documents_path)
+                        splitter = CharacterTextSplitter(
+                            separator="\n", 
+                            chunk_size=self.chunk_size, 
+                            chunk_overlap=self.chunk_overlap
+                        )
+                        documents = splitter.create_documents([text])
+                    except Exception as e2:
+                        print(f"Error with fallback PDF extraction: {e2}")
+            elif self.documents_path.lower().endswith(('.txt', '.md', '.html')):
+                loader = TextLoader(self.documents_path)
+                documents = loader.load()
+        elif os.path.isdir(self.documents_path):
+            # Load all supported files in directory
             try:
-                # Handle missing 'id' key gracefully
-                token_id = token_entry.get('id', None)
-                token_text = token_entry.get('text', None)
-
-                if token_text:
-                    for char in token_text:
-                        partial_message += char
-                        yield partial_message
-                        time.sleep(0.01)
-                else:
-                    print(f"Warning ==> The key 'text' does not exist or is None in this token entry: {token_entry}")
-                    pass
-
-            except KeyError as e:
-                print(f"KeyError: {e} occurred for token entry: {token_entry}")
-                continue
-
-
-class TextGeneratorTool(Tool):
-    name = "vector_retriever"
-    description = "This tool searches in a vector store based on a given prompt."
-    inputs = ["prompt"]
-    outputs = ["text"]
-
-
-    def __init__(self):
-        #self.retriever = db.as_retriever(search_kwargs={"k": 1})
-        pass  # You might want to add some initialization logic here
-		
-    def __call__(self, prompt: str):
-        result = predict(prompt,  0.9, 512, 0.6, 1.4)
-        return result
-
-
-
-
+                loader = DirectoryLoader(
+                    self.documents_path,
+                    glob="**/*.*",
+                    loader_cls=TextLoader,
+                    loader_kwargs={"autodetect_encoding": True}
+                )
+                documents = loader.load()
+            except Exception as e:
+                print(f"Error loading directory: {e}")
+        
+        # Split documents into chunks if they exist
+        if documents:
+            splitter = RecursiveCharacterTextSplitter(
+                chunk_size=self.chunk_size,
+                chunk_overlap=self.chunk_overlap
+            )
+            return splitter.split_documents(documents)
+        return []
+    
+    def _extract_text_from_pdf(self, pdf_path):
+        """Extract text from PDF using PyPDF2"""
+        text = ""
+        pdf_reader = PdfReader(pdf_path)
+        for page in pdf_reader.pages:
+            text += page.extract_text()
+        return text
+            
+    def _create_vector_store(self, documents):
+        """Create a new vector store from documents"""
+        embeddings = self._get_embeddings()
+        
+        if self.vector_store_type.lower() == "faiss":
+            vector_store = FAISS.from_documents(documents, embeddings)
+            vector_store.save_local(self.persist_directory)
+            print(f"Created FAISS vector store at {self.persist_directory}")
+        else:  # Default to Chroma
+            vector_store = Chroma.from_documents(
+                documents, 
+                embeddings,
+                persist_directory=self.persist_directory
+            )
+            vector_store.persist()
+            print(f"Created Chroma vector store at {self.persist_directory}")
+        
+        self.vector_store = vector_store
+    
+    def _load_vector_store(self):
+        """Load an existing vector store"""
+        embeddings = self._get_embeddings()
+        
+        try:
+            if self.vector_store_type.lower() == "faiss":
+                self.vector_store = FAISS.load_local(self.persist_directory, embeddings)
+                print(f"Loaded FAISS vector store from {self.persist_directory}")
+            else:  # Default to Chroma
+                self.vector_store = Chroma(
+                    persist_directory=self.persist_directory,
+                    embedding_function=embeddings
+                )
+                print(f"Loaded Chroma vector store from {self.persist_directory}")
+        except Exception as e:
+            print(f"Error loading vector store: {e}")
+            print("Creating a new vector store...")
+            documents = self._load_documents()
+            if documents:
+                self._create_vector_store(documents)
+            else:
+                print("No documents available. Cannot create vector store.")
+                self.vector_store = None
+
+    def forward(self, query: str, top_k: int = 3) -> str:
+        """
+        Retrieve relevant documents based on the query.
+        
+        Args:
+            query: The search query
+            top_k: Number of results to return
+            
+        Returns:
+            String with formatted search results
+        """
+        if not hasattr(self, 'vector_store') or self.vector_store is None:
+            return "Vector store is not initialized. Please check your configuration."
+        
+        try:
+            # Perform similarity search
+            results = self.vector_store.similarity_search(query, k=top_k)
+            
+            # Format results
+            formatted_results = []
+            for i, doc in enumerate(results):
+                content = doc.page_content
+                metadata = doc.metadata
+                
+                # Format metadata nicely
+                meta_str = ""
+                if metadata:
+                    meta_str = "\nSource: "
+                    if "source" in metadata:
+                        meta_str += metadata["source"]
+                    if "page" in metadata:
+                        meta_str += f", Page: {metadata['page']}"
+                
+                formatted_results.append(f"Document {i+1}:\n{content}{meta_str}\n")
+            
+            if formatted_results:
+                return "Retrieved relevant information:\n\n" + "\n".join(formatted_results)
+            else:
+                return "No relevant information found for the query."
+        except Exception as e:
+            return f"Error retrieving information: {str(e)}"
+
+# Example usage:
+# rag_tool = RAGTool(
+#     documents_path="./my_docs",
+#     embedding_model="sentence-transformers/all-MiniLM-L6-v2",
+#     vector_store_type="faiss",
+#     chunk_size=1000,
+#     chunk_overlap=200
+# )