Update rag_system.py

rag_system.py

@@ -6,41 +6,42 @@ from langchain.text_splitter import CharacterTextSplitter
 from langchain.docstore.document import Document
 from transformers import pipeline
 from langchain.prompts import PromptTemplate
-from typing import List, Dict, Any, Optional
 
 class RAGSystem:
-    def __init__(self, sql_generator: SQLGenerator, csv_path: str = "apparel.csv"):
-        self.sql_generator = sql_generator
+    def __init__(self, csv_path="apparel.csv"):
         self.setup_system(csv_path)
         self.qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
         
-    def setup_system(self, csv_path: str):
+    def setup_system(self, csv_path):
         if not os.path.exists(csv_path):
             raise FileNotFoundError(f"CSV file not found at {csv_path}")
             
+        # Read the CSV file
         documents = pd.read_csv(csv_path)
         
+        # Create proper Document objects
         docs = [
             Document(
-                page_content=str(row['Title']),
+                page_content=str(row['Title']),  # Convert to string to ensure compatibility
                 metadata={'index': idx}
             ) 
             for idx, row in documents.iterrows()
         ]
         
+        # Split documents
         text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
         split_docs = text_splitter.split_documents(docs)
         
+        # Create embeddings and vector store
         embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
         self.vector_store = FAISS.from_documents(split_docs, embeddings)
         self.retriever = self.vector_store.as_retriever()
         
-    def process_query(self, query: str, execute_sql: bool = True) -> Dict[str, Any]:
-        """
-        Process a query through both RAG and SQL if needed
-        """
-        # Get relevant documents
-        retrieved_docs = self.retriever.get_relevant_documents(query)
+    def process_query(self, query):
+        # Retrieve documents based on the query
+        retrieved_docs = self.retriever.get_relevant_documents(query)  # Changed from invoke to get_relevant_documents
+        
+        # Properly access page_content from Document objects
         retrieved_text = "\n".join([doc.page_content for doc in retrieved_docs])[:1000]
         
         # Process with QA pipeline
@@ -48,42 +49,13 @@ class RAGSystem:
             "question": query,
             "context": retrieved_text
         }
-        qa_response = self.qa_pipeline(qa_input)
+        response = self.qa_pipeline(qa_input)
         
-        result = {
-            "qa_answer": qa_response['answer'],
-            "relevant_docs": [doc.page_content for doc in retrieved_docs[:3]],
-            "sql_results": None
-        }
-        
-        # If SQL execution is requested and SQL is detected in the query
-        if execute_sql and "SELECT" in query.upper():
-            if self.sql_generator.validate_query(query):
-                sql_results = self.sql_generator.execute_query(query)
-                result["sql_results"] = sql_results
-            
-        return result
+        return response['answer']
 
-    def get_similar_documents(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
+    def get_similar_documents(self, query, k=5):
         """
         Retrieve similar documents without processing through QA pipeline
         """
         docs = self.retriever.get_relevant_documents(query)
-        return [{'content': doc.page_content, 'metadata': doc.metadata} for doc in docs[:k]]
-
-# Example usage
-if __name__ == "__main__":
-    # Initialize the SQL generator
-    sql_gen = SQLGenerator("shopify.db")
-    
-    # Initialize the RAG system with the SQL generator
-    rag = RAGSystem(sql_gen, "apparel.csv")
-    
-    # Example query that might include SQL
-    query = "SELECT * FROM products LIMIT 5"
-    results = rag.process_query(query)
-    
-    # Access different parts of the results
-    print("QA Answer:", results["qa_answer"])
-    print("Relevant Documents:", results["relevant_docs"])
-    print("SQL Results:", results["sql_results"])
+        return [{'content': doc.page_content, 'metadata': doc.metadata} for doc in docs[:k]]