Changes in app.py

app.py

@@ -1,102 +1,114 @@
 import gradio as gr
-import requests
-import json
 import time
+import json
+
+# Import RAG system components
+from rag_system.vector_store import VectorStore
+from rag_system.retriever import SQLRetriever
+from rag_system.prompt_engine import PromptEngine
+from rag_system.sql_generator import SQLGenerator
+
+# Initialize RAG system components
+print("Initializing RAG system...")
+try:
+    vector_store = VectorStore()
+    retriever = SQLRetriever(vector_store)
+    prompt_engine = PromptEngine()
+    sql_generator = SQLGenerator(retriever, prompt_engine)
+    print("RAG system initialized successfully!")
+except Exception as e:
+    print(f"Error initializing RAG system: {e}")
+    sql_generator = None
 
 def generate_sql(question, table_headers):
-    """Generate SQL using the RAG API."""
+    """Generate SQL using the RAG system directly."""
+    if sql_generator is None:
+        return "❌ Error: RAG system not initialized"
+    
     try:
-        # Prepare the request
-        data = {
-            "question": question,
-            "table_headers": [h.strip() for h in table_headers.split(",") if h.strip()]
-        }
+        start_time = time.time()
+        
+        # Generate SQL using RAG system
+        result = sql_generator.generate_sql(question, table_headers)
         
-        # Make API call to the RAG system
-        response = requests.post("http://localhost:8000/predict", json=data)
+        processing_time = time.time() - start_time
         
-        if response.status_code == 200:
-            result = response.json()
-            return f"""
+        return f"""
 **Generated SQL:**
 ```sql
 {result['sql_query']}
 ```
 
 **Model Used:** {result['model_used']}
-**Processing Time:** {result['processing_time']:.2f}s
+**Processing Time:** {processing_time:.2f}s
 **Status:** {result['status']}
 **Retrieved Examples:** {len(result['retrieved_examples'])} examples used for RAG
 """
-        else:
-            return f"❌ Error: {response.status_code} - {response.text}"
-            
     except Exception as e:
         return f"❌ Error: {str(e)}"
 
 def batch_generate_sql(questions_text, table_headers):
     """Generate SQL for multiple questions."""
+    if sql_generator is None:
+        return "❌ Error: RAG system not initialized"
+    
     try:
         # Parse questions
         questions = [q.strip() for q in questions_text.split("\n") if q.strip()]
         
-        # Prepare batch request
-        data = {
-            "queries": [
-                {
-                    "question": q,
-                    "table_headers": [h.strip() for h in table_headers.split(",") if h.strip()]
-                }
-                for q in questions
-            ]
-        }
+        output = f"**Batch Results:**\n"
+        output += f"Total Queries: {len(questions)}\n"
+        successful_count = 0
         
-        # Make API call
-        response = requests.post("http://localhost:8000/batch", json=data)
+        for i, question in enumerate(questions):
+            try:
+                start_time = time.time()
+                result = sql_generator.generate_sql(question, table_headers)
+                processing_time = time.time() - start_time
+                
+                output += f"\n**Query {i+1}:** {question}\n"
+                output += f"```sql\n{result['sql_query']}\n```\n"
+                output += f"Model: {result['model_used']} | Time: {processing_time:.2f}s\n"
+                
+                if result['status'] == 'success':
+                    successful_count += 1
+                    
+            except Exception as e:
+                output += f"\n**Query {i+1}:** {question}\n"
+                output += f"❌ Error: {str(e)}\n"
+        
+        output += f"\n**Successful:** {successful_count}/{len(questions)}"
+        return output
         
-        if response.status_code == 200:
-            result = response.json()
-            output = f"**Batch Results:**\n"
-            output += f"Total Queries: {result['total_queries']}\n"
-            output += f"Successful: {result['successful_queries']}\n\n"
-            
-            for i, res in enumerate(result['results']):
-                output += f"**Query {i+1}:** {res['question']}\n"
-                output += f"```sql\n{res['sql_query']}\n```\n"
-                output += f"Model: {res['model_used']} | Time: {res['processing_time']:.2f}s\n\n"
-            
-            return output
-        else:
-            return f"❌ Error: {response.status_code} - {response.text}"
-            
     except Exception as e:
         return f"❌ Error: {str(e)}"
 
 def check_system_health():
     """Check the health of the RAG system."""
     try:
-        response = requests.get("http://localhost:8000/health")
-        if response.status_code == 200:
-            health_data = response.json()
-            return f"""
+        if sql_generator is None:
+            return "❌ System Status: RAG system not initialized"
+        
+        # Get model info
+        model_info = sql_generator.get_model_info()
+        
+        return f"""
 **System Health:**
-- **Status:** {health_data['status']}
-- **System Loaded:** {health_data['system_loaded']}
-- **System Loading:** {health_data['system_loading']}
-- **Error:** {health_data['system_error'] or 'None'}
-- **Timestamp:** {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(health_data['timestamp']))}
+- **Status:** ✅ Healthy
+- **System Loaded:** ✅ Yes
+- **System Loading:** ❌ No
+- **Error:** None
+- **Timestamp:** {time.strftime('%Y-%m-%d %H:%M:%S')}
 
 **Model Info:**
-{json.dumps(health_data.get('model_info', {}), indent=2) if health_data.get('model_info') else 'Not available'}
+{json.dumps(model_info, indent=2) if model_info else 'Not available'}
 """
-        else:
-            return f"❌ Health check failed: {response.status_code}"
     except Exception as e:
         return f"❌ Health check error: {str(e)}"
 
 # Create Gradio interface
 with gr.Blocks(title="Text-to-SQL RAG with CodeLlama", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🚀 Text-to-SQL RAG with CodeLlama")
+    gr.Markdown("#Text-to-SQL RAG with CodeLlama")
     gr.Markdown("Generate SQL queries from natural language using **RAG (Retrieval-Augmented Generation)** and **CodeLlama** models.")
     gr.Markdown("**Features:** RAG-enhanced generation, CodeLlama integration, Vector-based retrieval, Advanced prompt engineering")
     
@@ -113,7 +125,7 @@ with gr.Blocks(title="Text-to-SQL RAG with CodeLlama", theme=gr.themes.Soft()) a
                     placeholder="e.g., id, name, salary, department",
                     value="id, name, salary, department"
                 )
-                generate_btn = gr.Button("🚀 Generate SQL", variant="primary", size="lg")
+                generate_btn = gr.Button("Generate SQL", variant="primary", size="lg")
             
             with gr.Column(scale=1):
                 output = gr.Markdown(label="Result")
@@ -131,14 +143,14 @@ with gr.Blocks(title="Text-to-SQL RAG with CodeLlama", theme=gr.themes.Soft()) a
                     placeholder="e.g., id, name, salary, department",
                     value="id, name, salary, department"
                 )
-                batch_btn = gr.Button("🚀 Generate Batch SQL", variant="primary", size="lg")
+                batch_btn = gr.Button("Generate Batch SQL", variant="primary", size="lg")
             
             with gr.Column(scale=1):
                 batch_output = gr.Markdown(label="Batch Results")
     
     with gr.Tab("System Health"):
         with gr.Row():
-            health_btn = gr.Button("🔍 Check System Health", variant="secondary", size="lg")
+            health_btn = gr.Button("Check System Health", variant="secondary", size="lg")
             health_output = gr.Markdown(label="Health Status")
     
     # Event handlers
@@ -161,22 +173,22 @@ with gr.Blocks(title="Text-to-SQL RAG with CodeLlama", theme=gr.themes.Soft()) a
     
     gr.Markdown("---")
     gr.Markdown("""
-    ## 🎯 How It Works
+    ## How It Works
     
     1. **RAG System**: Retrieves relevant SQL examples from vector database
     2. **CodeLlama**: Generates SQL using retrieved examples as context
     3. **Vector Search**: Finds similar questions and their SQL solutions
     4. **Enhanced Generation**: Combines retrieval + generation for better accuracy
     
-    ## 🛠️ Technology Stack
+    ## Technology Stack
     
-    - **Backend**: FastAPI + Python
+    - **Backend**: Direct RAG system integration
     - **LLM**: CodeLlama-7B-Python-GGUF (primary)
     - **Vector DB**: ChromaDB with sentence transformers
     - **Frontend**: Gradio interface
     - **Hosting**: Hugging Face Spaces
     
-    ## 📊 Performance
+    ## Performance
     
     - **Model**: CodeLlama-7B-Python-GGUF
     - **Response Time**: < 5 seconds