Update rag_system.py

rag_system.py

@@ -6,42 +6,41 @@ 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, csv_path="apparel.csv"):
+    def __init__(self, sql_generator: SQLGenerator, csv_path: str = "apparel.csv"):
+        self.sql_generator = sql_generator
         self.setup_system(csv_path)
         self.qa_pipeline = pipeline("question-answering", model="distilbert-base-cased-distilled-squad")
         
-    def setup_system(self, csv_path):
+    def setup_system(self, csv_path: str):
         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']),  # Convert to string to ensure compatibility
+                page_content=str(row['Title']),
                 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):
-        # 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
+    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)
         retrieved_text = "\n".join([doc.page_content for doc in retrieved_docs])[:1000]
         
         # Process with QA pipeline
@@ -49,13 +48,42 @@ class RAGSystem:
             "question": query,
             "context": retrieved_text
         }
-        response = self.qa_pipeline(qa_input)
+        qa_response = self.qa_pipeline(qa_input)
         
-        return response['answer']
+        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
 
-    def get_similar_documents(self, query, k=5):
+    def get_similar_documents(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
         """
         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]]
+        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"])