Spaces:
Sleeping
Sleeping
| # medicine.py - FIXED VERSION | |
| import os | |
| import io | |
| from flask import Flask, request, jsonify | |
| from PIL import Image | |
| from dotenv import load_dotenv | |
| import google.generativeai as genai | |
| import json | |
| # --- INITIAL SETUP --- | |
| load_dotenv() | |
| api_key = os.getenv("GOOGLE_API_KEY") | |
| if not api_key: | |
| raise ValueError("GOOGLE_API_KEY not found. Please set it in your .env file.") | |
| genai.configure(api_key=api_key) | |
| app = Flask(__name__) | |
| # --- CONFIGURATION --- | |
| TEXT_FILES_DIR = "MEDICINE_TXT" | |
| ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp'} | |
| # Check for available knowledge base files | |
| try: | |
| os.makedirs(TEXT_FILES_DIR, exist_ok=True) | |
| AVAILABLE_FILES = [f for f in os.listdir(TEXT_FILES_DIR) if f.endswith('.txt')] | |
| if not AVAILABLE_FILES: | |
| print(f"Warning: No .txt files found in '{TEXT_FILES_DIR}'. Running without knowledge base.") | |
| except Exception as e: | |
| print(f"Warning: Error accessing '{TEXT_FILES_DIR}': {e}") | |
| AVAILABLE_FILES = [] | |
| # --- HELPER FUNCTIONS --- | |
| def allowed_file(filename): | |
| return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS | |
| def find_relevant_file(topic: str) -> str: | |
| """Find the most relevant file for a given topic using Gemini""" | |
| if not AVAILABLE_FILES: | |
| return None | |
| try: | |
| model = genai.GenerativeModel(os.getenv("MODEL","gemini-2.5-flash-lite")) | |
| prompt = f""" | |
| From the following list of files, which one is most relevant for: "{topic}"? | |
| Respond with ONLY the filename, nothing else. | |
| Files: {', '.join(AVAILABLE_FILES)} | |
| """ | |
| response = model.generate_content(prompt) | |
| filename = response.text.strip().replace("`", "").replace('"', '') | |
| if filename in AVAILABLE_FILES: | |
| print(f"Found relevant file: {filename} for topic: {topic}") | |
| return filename | |
| else: | |
| print(f"No matching file found for: {topic}") | |
| return None | |
| except Exception as e: | |
| print(f"Error finding relevant file: {e}") | |
| return None | |
| def get_context_from_file(filename: str) -> str: | |
| """Read content from a text file""" | |
| if not filename: | |
| return None | |
| filepath = os.path.join(TEXT_FILES_DIR, filename) | |
| try: | |
| with open(filepath, 'r', encoding='utf-8') as f: | |
| content = f.read() | |
| print(f"Successfully loaded context from {filename}") | |
| return content | |
| except Exception as e: | |
| print(f"Error reading file {filename}: {e}") | |
| return None | |
| # --- MAIN API ENDPOINT --- | |
| def handle_query(): | |
| """ | |
| Handle medicine queries with optional image upload | |
| Accepts both JSON and FormData | |
| """ | |
| # FIXED: Handle both JSON and FormData | |
| user_query = None | |
| medicine_topic = None | |
| # Check if request is JSON | |
| if request.is_json: | |
| data = request.get_json() | |
| user_query = data.get('main_query') | |
| else: | |
| # Handle FormData (when image is uploaded) | |
| user_query = request.form.get('main_query') | |
| if not user_query: | |
| return jsonify({"error": "Missing 'main_query' in request"}), 400 | |
| print(f"Received query: {user_query}") | |
| # Handle image upload if present | |
| if 'file' in request.files: | |
| file = request.files['file'] | |
| if file and file.filename != '' and allowed_file(file.filename): | |
| try: | |
| print(f"Processing uploaded image: {file.filename}") | |
| # Process image with Gemini Vision | |
| img = Image.open(file.stream) | |
| vision_model = genai.GenerativeModel(os.getenv("MODEL","gemini-2.5-flash-lite")) | |
| vision_prompt = [ | |
| """Identify the medicine from this image. Look for: | |
| - Medicine name or brand | |
| - Active ingredients | |
| - Rx number or formula | |
| - Any text on packaging or pills | |
| Respond with just the medicine name or main component.""", | |
| img | |
| ] | |
| response = vision_model.generate_content(vision_prompt) | |
| medicine_topic = response.text.strip() | |
| print(f"Identified from image: {medicine_topic}") | |
| except Exception as e: | |
| print(f"Error processing image: {e}") | |
| # Continue without image data | |
| # If no medicine identified from image, extract from query text | |
| if not medicine_topic: | |
| try: | |
| model = genai.GenerativeModel(os.getenv("MODEL","gemini-2.5-flash-lite")) | |
| extract_prompt = f""" | |
| From this query: "{user_query}" | |
| Extract the main medicine or medical topic being asked about. | |
| Respond with ONLY the medicine/topic name (e.g., 'Paracetamol', 'Antibiotics') | |
| """ | |
| response = model.generate_content(extract_prompt) | |
| medicine_topic = response.text.strip() | |
| print(f"Extracted topic from query: {medicine_topic}") | |
| except Exception as e: | |
| print(f"Error extracting topic: {e}") | |
| medicine_topic = "general medicine" | |
| # Find relevant knowledge base file | |
| context = None | |
| source_file = None | |
| if AVAILABLE_FILES: | |
| relevant_file = find_relevant_file(medicine_topic) | |
| if relevant_file: | |
| context = get_context_from_file(relevant_file) | |
| source_file = relevant_file | |
| # Generate response | |
| try: | |
| model = genai.GenerativeModel('gemini-2.0-flash-exp') | |
| # Build prompt based on available context | |
| if context: | |
| final_prompt = f""" | |
| You are a medical information assistant. | |
| CONTEXT FROM KNOWLEDGE BASE: | |
| {context} | |
| IDENTIFIED MEDICINE/TOPIC: {medicine_topic} | |
| USER QUESTION: {user_query} | |
| Instructions: | |
| - Answer based on the context if available | |
| - Use simple language | |
| - Keep response under 200 words | |
| - Include dosage, usage, and warnings if relevant | |
| - If context doesn't have the info, use general knowledge | |
| - Respond in the same language as the user query | |
| """ | |
| else: | |
| final_prompt = f""" | |
| You are a medical information assistant. | |
| MEDICINE/TOPIC: {medicine_topic} | |
| USER QUESTION: {user_query} | |
| Provide accurate medical information about this topic. | |
| - Use simple language | |
| - Keep response under 200 words | |
| - Include dosage, usage, side effects if relevant | |
| - Add standard medical disclaimers | |
| - Respond in the same language as the user query | |
| """ | |
| response = model.generate_content(final_prompt) | |
| return jsonify({ | |
| "status": "success", | |
| "response": response.text.strip(), | |
| "identified_topic": medicine_topic, | |
| "source_file": source_file if source_file else "general_knowledge", | |
| "knowledge_base_available": len(AVAILABLE_FILES) > 0 | |
| }) | |
| except Exception as e: | |
| print(f"Error generating response: {e}") | |
| return jsonify({ | |
| "error": "Failed to generate response", | |
| "details": str(e) | |
| }), 500 | |
| def health_check(): | |
| """Health check endpoint""" | |
| return jsonify({ | |
| "status": "running", | |
| "service": "medicine_info", | |
| "knowledge_base_files": len(AVAILABLE_FILES), | |
| "port": 5002 | |
| }) | |
| def index(): | |
| """Basic info endpoint""" | |
| return jsonify({ | |
| "service": "Medicine Information API", | |
| "endpoint": "/api/query", | |
| "methods": ["POST"], | |
| "accepts": "JSON or FormData with 'main_query' and optional 'file'", | |
| "knowledge_base": f"{len(AVAILABLE_FILES)} files available" | |
| }) | |
| if __name__ == '__main__': | |
| print("="*50) | |
| print("Starting Medicine Information Service") | |
| print(f"Knowledge base: {len(AVAILABLE_FILES)} files in '{TEXT_FILES_DIR}'") | |
| if AVAILABLE_FILES: | |
| print(f"Available files: {', '.join(AVAILABLE_FILES[:5])}") | |
| print("="*50) | |
| app.run(host='0.0.0.0', port=5002, debug=True) |