Update main.py

main.py

@@ -1,3 +1,5 @@
+# main.py - FIXED VERSION
+
 import os
 import uuid
 import json
@@ -8,88 +10,75 @@ import google.generativeai as genai
 from dotenv import load_dotenv
 from flask_cors import CORS
 
-
-# Step 1: API Key aur Environment Setup
+# Step 1: API Key and Environment Setup
 load_dotenv()
 try:
     genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
 except TypeError:
-    print("ERROR: Google API Key nahi mila.")
-    print("Ek .env file banayein aur usmein 'GOOGLE_API_KEY=your_key_here' likhein.")
+    print("ERROR: Google API Key not found.")
+    print("Create a .env file and add 'GOOGLE_API_KEY=your_key_here'")
     exit()
 
 app = Flask(__name__, static_folder='.', static_url_path='')
 CORS(app)
 
-# Serve static files (CSS, JS, images)
+# Serve static files
 @app.route('/<path:path>')
 def serve_static(path):
     return send_from_directory('.', path)
 
-# Serve the main page
 @app.route('/')
 def serve_index():
     return send_from_directory('.', 'index.html')
 
-# Step 2: API Endpoints aur Session Storage
+# Step 2: FIXED API Endpoints
 API_ENDPOINTS = {
-    "skin_disease": "https://your-api-domain.com/skin-disease-detection",
+    "skin_disease": "http://localhost:5003/api/skin",  # Update when ready
     "medicine_info": "http://localhost:5002/api/query",
-    "report_reading": "https://your-api-domain.com/report-reading",
+    "report_reading": "http://localhost:5004/api/report",  # Update when ready
     "disease_query": "http://localhost:5001/ask"
 }
+
+# Session storage
 SESSIONS = {}
 
-# Step 3: Gemini se Query Classify karne ka Function
+# Step 3: Query Classification Function
 def classify_query_with_gemini(query: str):
-    """User ki query ko Gemini API ka istemal karke classify karta hai."""
+    """Classify user query using Gemini API"""
     model = genai.GenerativeModel('gemini-2.5-flash-lite')
     
-    # *** PROMPT HAS BEEN IMPROVED ***
     prompt = f"""
     Analyze the user's medical query and classify it into one of the following categories:
     - skin_disease: For queries about skin conditions, rashes, moles, spots, or any visible symptoms on the skin.
-    - medicine_info: For query about medicine(like how to use it, side effects, etc.) and also questions about a specific Medicine shown in attached image (optional) .
-    - report_reading: For queries asking to interpret or explain a medical report, lab test, or blood work from an image.
+    - medicine_info: For queries about medicines (usage, side effects, identification from image).
+    - report_reading: For queries asking to interpret medical reports, lab tests, or blood work from images.
     - disease_query: For general questions about diseases, symptoms, causes, or treatments.
 
-    Based on the classification, determine if an image is essential to answer the query accurately.
-    Generate response in English only.
-
-    The user query is: 
-    ---START OF QUERY---
-    {query}
-    ---END OF QUERY---
-
-    Provide the output ONLY in a valid JSON format with two keys: "category" (string) and "image_required" (boolean).
-
-    Example 1:
-    Query: "what are the symptoms of typhoid"
-    Output: {{"category": "disease_query", "image_required": false}}
-
-    Example 2:
-    Query: "I have a red circular rash on my arm, what is it?"
-    Output: {{"category": "skin_disease", "image_required": true}}
-
-    Example 3:
-    Query: "Can you tell me what this lab report says?"
-    Output: {{"category": "report_reading", "image_required": true}}
+    Determine if an image is essential to answer the query accurately.
+    
+    User query: "{query}"
 
-    Example 4:
-    Query: "What is this white pill with 'IP 204' written on it?"
-    Output: {{"category": "medicine_info", "image_required": true}}
+    Output ONLY valid JSON with "category" (string) and "image_required" (boolean).
+    
+    Examples:
+    "what are the symptoms of typhoid" -> {{"category": "disease_query", "image_required": false}}
+    "I have a red rash on my arm" -> {{"category": "skin_disease", "image_required": true}}
+    "What does this lab report say?" -> {{"category": "report_reading", "image_required": true}}
+    "What is this white pill?" -> {{"category": "medicine_info", "image_required": true}}
     """
     
     try:
         response = model.generate_content(prompt)
         cleaned_text = response.text.strip().replace('```json', '').replace('```', '')
         result = json.loads(cleaned_text)
+        print(f"Classification result: {result}")
         return result
     except Exception as e:
-        print(f"Gemini API call ya JSON parsing mein error: {e}")
-        return None
+        print(f"Classification error: {e}")
+        # Default fallback
+        return {"category": "disease_query", "image_required": False}
 
-# Step 4: API Routes (Endpoints)
+# Step 4: API Routes
 @app.route('/start_session', methods=['POST'])
 def start_session():
     session_id = str(uuid.uuid4())
@@ -111,37 +100,58 @@ def process_query():
     print(f"Session {session_id}: Query received: '{query}'")
     classification = classify_query_with_gemini(query)
     
-    if not classification:
-        return jsonify({"error": "Could not classify the query."}), 500
-
     SESSIONS[session_id]['classification'] = classification
     SESSIONS[session_id]['query'] = query
     
     if classification.get('image_required'):
-        print(f"Session {session_id}: Image required for category '{classification.get('category')}'")
+        print(f"Session {session_id}: Image required for '{classification.get('category')}'")
         return jsonify({
             "status": "image_required",
-            "message": "Please send the request to /process_with_image with the required photo."
+            "message": "Please upload an image via /process_with_image endpoint",
+            "category": classification.get('category')
         }), 200
     else:
-        print(f"Session {session_id}: No image required. Forwarding to '{classification.get('category')}' API.")
-        # Asli API ko call karein (Abhi ke liye mock response)
-        # Session se query aur classification nikalein
-        query = SESSIONS[session_id].get('query')
-        classification = SESSIONS[session_id].get('classification')
+        # Process without image
         category = classification['category']
         endpoint_url = API_ENDPOINTS.get(category)
-        if category == "medicine_info":
-            response = requests.post(endpoint_url, json={"main_query": query})
-        else:
-            response = requests.post(endpoint_url, data={"main_query": query})
-        del SESSIONS[session_id]
-        print(f"Session {session_id} closed.")
-        return jsonify({
-            "status": "success",
-            "response": response.json(),
-            "data": f"Information about '{query}': This is a tuned response from the {classification.get('category')} service."
-        })
+        
+        if not endpoint_url:
+            return jsonify({"error": f"No endpoint configured for category: {category}"}), 500
+        
+        try:
+            print(f"Forwarding to {category} API at {endpoint_url}")
+            
+            # FIXED: Proper request format for each service
+            if category == "disease_query":
+                # disease.py expects 'query' field
+                response = requests.post(endpoint_url, json={"query": query})
+            elif category == "medicine_info":
+                # medicine.py expects 'main_query' field
+                response = requests.post(endpoint_url, json={"main_query": query})
+            else:
+                # Default format for future services
+                response = requests.post(endpoint_url, json={"query": query})
+            
+            response.raise_for_status()
+            result = response.json()
+            
+            # Clean up session
+            del SESSIONS[session_id]
+            print(f"Session {session_id} closed successfully")
+            
+            return jsonify({
+                "status": "success",
+                "category": category,
+                "response": result
+            })
+            
+        except requests.exceptions.RequestException as e:
+            print(f"Error calling {category} API: {e}")
+            del SESSIONS[session_id]
+            return jsonify({
+                "error": f"Failed to reach {category} service",
+                "details": str(e)
+            }), 503
 
 @app.route('/process_with_image', methods=['POST'])
 def process_with_image():
@@ -151,52 +161,82 @@ def process_with_image():
         return jsonify({"error": "Invalid or missing session_id"}), 400
     
     if 'photo' not in request.files:
-        return jsonify({"error": "No photo file found in the request"}), 400
+        return jsonify({"error": "No photo file found in request"}), 400
         
     file = request.files['photo']
     if file.filename == '':
-        return jsonify({"error": "No selected file"}), 400
+        return jsonify({"error": "No file selected"}), 400
 
-    # Session se query aur classification nikalein
+    # Get session data
     query = SESSIONS[session_id].get('query')
     classification = SESSIONS[session_id].get('classification')
-    category = classification['category']
+    category = classification.get('category')
     endpoint_url = API_ENDPOINTS.get(category)
     
-    print(f"Session {session_id}: Image received. Preparing to forward to '{category}' API.")
-
-    # *** NEW: FORWARDING LOGIC THAT MATCHES YOUR CURL COMMAND ***
-    # The file object from Flask needs its stream to be readable by `requests`
-    # We pass the file stream, filename, and mimetype to requests
-    # The dictionary key 'file' matches the '-F file=@...' part of your curl command
-    files_payload = {'file': (file.filename, file.stream, file.mimetype)}
+    if not endpoint_url:
+        del SESSIONS[session_id]
+        return jsonify({"error": f"No endpoint configured for category: {category}"}), 500
     
-    # The dictionary key 'query' matches the '-F query=...' part
-    data_payload = {'query': query}
+    print(f"Session {session_id}: Processing image for '{category}'")
 
     try:
-        # NOTE: Neeche di gayi line asli API call hai.
-        # Jab aapka backend service (e.g., http://localhost:5002/api/query) taiyaar ho,
-        # to is line ko uncomment kar dein.
+        # FIXED: Proper request format based on category
+        if category == "medicine_info":
+            # medicine.py expects 'file' and 'main_query'
+            files = {'file': (file.filename, file.stream, file.mimetype)}
+            data = {'main_query': query}
+            response = requests.post(endpoint_url, files=files, data=data)
+            
+        elif category == "disease_query":
+            # disease.py doesn't typically need images, but if it did:
+            files = {'file': (file.filename, file.stream, file.mimetype)}
+            data = {'query': query}
+            response = requests.post(endpoint_url, files=files, data=data)
+            
+        else:
+            # Default format for skin_disease, report_reading
+            files = {'file': (file.filename, file.stream, file.mimetype)}
+            data = {'query': query}
+            response = requests.post(endpoint_url, files=files, data=data)
         
-        response_from_service = requests.post(endpoint_url, files=files_payload, data=data_payload)
-        response_from_service.raise_for_status() # Agar 4xx/5xx error ho to exception raise karega
-        tuned_response = response_from_service.json() # Assume service returns JSON
-
-        # Abhi ke liye, hum ek mock response bhej rahe hain
-        mock_response = {
-            "status": "success",
-            "response": tuned_response,
-            "data": f"Analysis for '{query}' based on your image: This is a tuned MOCK response from the {category} service."
-        }
+        response.raise_for_status()
+        result = response.json()
         
-        # Session close karein
+        # Clean up session
         del SESSIONS[session_id]
-        print(f"Session {session_id} closed.")
+        print(f"Session {session_id} closed successfully")
         
-        return jsonify(mock_response)
+        return jsonify({
+            "status": "success",
+            "category": category,
+            "response": result
+        })
 
     except requests.exceptions.RequestException as e:
+        print(f"Error calling {category} API with image: {e}")
         del SESSIONS[session_id]
-        print(f"Session {session_id} closed after failed API call.")
-        return jsonify({"status": "error", "message": f"Backend service call failed: {e}"}), 503
+        return jsonify({
+            "error": f"Failed to reach {category} service",
+            "details": str(e)
+        }), 503
+
+@app.route('/health', methods=['GET'])
+def health_check():
+    """Health check endpoint to verify services"""
+    service_status = {}
+    
+    for service, url in API_ENDPOINTS.items():
+        try:
+            # Try to connect to each service
+            if "localhost" in url:
+                test_url = url.replace('/ask', '').replace('/api/query', '')
+                requests.get(test_url, timeout=1)
+                service_status[service] = "✅ Running"
+        except:
+            service_status[service] = "❌ Not reachable"
+    
+    return jsonify({
+        "main_service": "✅ Running",
+        "connected_services": service_status,
+        "sessions_active": len(SESSIONS)
+    })