app.py

import gradio as gr
import pandas as pd
import numpy as np
import os
import base64
from together import Together

def extract_medicines(api_key, image):
    """
    Extract medicine names from a prescription image using Together AI's Llama-Vision-Free model
    """
    # Check if API key is provided
    if not api_key:
        return "Please enter your Together API key."
    
    if image is None:
        return "Please upload an image."
    
    try:
        # Initialize Together client with the provided API key
        client = Together(api_key=api_key)
        
        # Convert image to base64
        with open(image, "rb") as img_file:
            img_data = img_file.read()
            b64_img = base64.b64encode(img_data).decode('utf-8')
        
        # Make API call with base64 encoded image
        response = client.chat.completions.create(
            model="meta-llama/Llama-Vision-Free",
            messages=[
                {
                    "role": "system",
                    "content": "You are an expert in identifying medicine names from prescription images."
                },
                {
                    "role": "user",
                    "content": [
                        {
                            "type": "text",
                            "text": "Please extract the names of the medicines only."
                        },
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{b64_img}"
                            }
                        }
                    ]
                }
            ]
        )
        
        # Extract medicine names from response
        medicine_list = response.choices[0].message.content
        return medicine_list
    
    except Exception as e:
        return f"Error: {str(e)}"

def recommend_medicine(api_key, medicine_name, csv_file=None):
    """
    Use Together API to recommend alternative medicines based on input medicine name
    using data from the provided CSV file with specific column structure
    """
    try:
        # If CSV file is provided, use it; otherwise use default
        if csv_file is not None:
            # Read the uploaded CSV
            if isinstance(csv_file, str):  # Path to default CSV
                df = pd.read_csv(csv_file)
            else:  # Uploaded file
                df = pd.read_csv(csv_file.name)
        else:
            # Use the default medicine_dataset.csv in the current directory
            df = pd.read_csv("medicine_dataset.csv")
        
        # Verify the medicine name exists in the dataset
        if medicine_name not in df['name'].values:
            return f"Error: Medicine '{medicine_name}' not found in the dataset. Please check the spelling or try another medicine."
        
        # Create system prompt with CSV data and column structure information
        system_prompt = f"""Develop an expert system to recommend alternative medicines for {medicine_name} based on the medicine dataset. The dataset has the following columns:
- name: Medicine name
- substitute0 through substitute4: Potential substitute medicines
- sideEffect0 through sideEffect41: Possible side effects
- use0 through use4: Medical uses
- Chemical Class: The chemical classification
- Habit Forming: Whether the medicine is habit-forming
- Therapeutic Class: The therapeutic classification
- Action Class: How the medicine works

Your task is to:
1. Find the row in the dataset where name matches exactly "{medicine_name}"
2. Find alternatives by:
   - Using the substitute0-substitute4 values as primary alternatives
   - Finding other medicines with similar Chemical Class, Therapeutic Class, or Action Class

For each recommended alternative, provide:
- Name of the alternative medicine
- All side effects (from relevant sideEffect columns)
- All uses (from relevant use columns)
- Chemical Class, Habit Forming status, Therapeutic Class, and Action Class
- A similarity score (0-1) indicating how similar it is to the original medicine

Format the response clearly with headings for "Recommended Medicines", "Medicine Details", and "Similarity Score".
"""

        # Extract the specific row containing the medicine data to give more context
        medicine_data = df[df['name'] == medicine_name]
        if not medicine_data.empty:
            # Convert the specific medicine data to a string representation
            medicine_info = medicine_data.to_string(index=False)
            system_prompt += f"\n\nThe specific data for {medicine_name} is:\n{medicine_info}\n\n"
            
            # Extract substitute information for better recommendations
            substitutes = []
            for i in range(5):  # substitute0 through substitute4
                col_name = f"substitute{i}"
                if col_name in medicine_data.columns:
                    sub_value = medicine_data[col_name].iloc[0]
                    if pd.notna(sub_value) and sub_value != "":
                        substitutes.append(sub_value)
            
            if substitutes:
                system_prompt += f"The primary substitutes for {medicine_name} are: {', '.join(substitutes)}\n\n"
        
        # Include a sample of other medicines for comparison
        other_medicines = df[df['name'] != medicine_name].sample(min(10, len(df)-1)) if len(df) > 1 else pd.DataFrame()
        if not other_medicines.empty:
            system_prompt += "Here's a sample of other medicines in the dataset for comparison:\n"
            for idx, row in other_medicines.iterrows():
                system_prompt += f"- {row['name']}: Chemical Class: {row['Chemical Class']}, Therapeutic Class: {row['Therapeutic Class']}, Action Class: {row['Action Class']}\n"

        # Initialize Together client with the API key
        client = Together(api_key=api_key)
        
        # Make API call
        response = client.chat.completions.create(
            model="meta-llama/Llama-3.3-70B-Instruct-Turbo-Free",
            messages=[
                {
                    "role": "system",
                    "content": system_prompt
                },
                {
                    "role": "user",
                    "content": f"Please recommend alternatives for {medicine_name} based on the dataset. Include detailed information about each alternative."
                }
            ],
            max_tokens=2000
        )
        
        # Return the generated recommendations
        return response.choices[0].message.content
    
    except Exception as e:
        return f"Error: {str(e)}"

def send_medicine_to_recommender(api_key, medicine_names, csv_file):
    """
    Takes medicine names extracted from prescription and gets recommendations
    """
    if not medicine_names or medicine_names.startswith("Error") or medicine_names.startswith("Please"):
        return "Please extract valid medicine names first"
    
    # Extract the first medicine name from the list (assuming it's the first line or first item)
    medicine_lines = medicine_names.strip().split('\n')
    if not medicine_lines:
        return "No valid medicine name found in extraction results"
    
    # Get the first medicine name (remove any bullet points or numbers)
    first_medicine = medicine_lines[0]
    # Clean up the medicine name (remove bullets, numbers, etc.)
    first_medicine = first_medicine.lstrip('•-*0123456789. ').strip()
    
    # Check if we have a valid medicine name
    if not first_medicine:
        return "Could not identify a valid medicine name from extraction"
    
    # Call the recommend medicine function with the first extracted medicine
    return recommend_medicine(api_key, first_medicine, csv_file)

# Create Gradio interface with tabs for both functionalities
with gr.Blocks(title="Medicine Assistant") as app:
    gr.Markdown("# Medicine Assistant")
    gr.Markdown("This application helps you extract medicine names from prescriptions and find alternative medicines.")
    
    # API key input (shared between tabs)
    api_key_input = gr.Textbox(
        label="Together API Key", 
        placeholder="Enter your Together API key here...",
        type="password"
    )
    
    # Create a file input for CSV that can be shared between tabs
    # Fixed the 'type' parameter to use 'filepath' instead of 'file'
    csv_file_input = gr.File(
        label="Upload Medicine CSV (Optional)", 
        file_types=[".csv"],
        type="filepath"  # Changed from 'file' to 'filepath'
    )
    gr.Markdown("If no CSV is uploaded, the app will use the default 'medicine_dataset.csv' file.")
    
    with gr.Tabs():
        with gr.Tab("Prescription Medicine Extractor"):
            gr.Markdown("## Prescription Medicine Extractor")
            gr.Markdown("Upload a prescription image to extract medicine names using Together AI's Llama-Vision-Free model.")
            
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(type="filepath", label="Upload Prescription Image")
                    extract_btn = gr.Button("Extract Medicines")
                    recommend_from_extract_btn = gr.Button("Get Recommendations for First Medicine")
                    
                with gr.Column():
                    extracted_output = gr.Textbox(label="Extracted Medicines", lines=10)
                    recommendation_from_extract_output = gr.Markdown(label="Recommendations")
            
            # Connect the buttons to functions
            extract_btn.click(
                fn=extract_medicines, 
                inputs=[api_key_input, image_input], 
                outputs=extracted_output
            )
            
            recommend_from_extract_btn.click(
                fn=send_medicine_to_recommender,
                inputs=[api_key_input, extracted_output, csv_file_input],
                outputs=recommendation_from_extract_output
            )
            
            gr.Markdown("""
            ### How to use:
            1. Enter your Together API key
            2. Upload a clear image of a prescription
            3. Click 'Extract Medicines' to see the identified medicines
            4. Optionally upload a custom medicine dataset CSV
            5. Click 'Get Recommendations for First Medicine' to find alternatives
            
            ### Note:
            - Your API key is used only for the current session
            - For best results, ensure the prescription image is clear and readable
            """)
        
        with gr.Tab("Medicine Alternative Recommender"):
            gr.Markdown("## Medicine Alternative Recommender")
            gr.Markdown("This tool recommends alternative medicines based on an input medicine name using the Together API.")
            
            with gr.Row():
                with gr.Column():
                    medicine_name = gr.Textbox(
                        label="Medicine Name", 
                        placeholder="Enter a medicine name exactly as it appears in the dataset"
                    )
                    submit_btn = gr.Button("Get Recommendations", variant="primary")
                
                with gr.Column():
                    recommendation_output = gr.Markdown(label="Recommendations")
            
            submit_btn.click(
                recommend_medicine,
                inputs=[api_key_input, medicine_name, csv_file_input],
                outputs=recommendation_output
            )

            gr.Markdown("""
            ## How to use this tool:
            1. Enter your Together API key (same key used across the application)
            2. Enter a medicine name **exactly as it appears** in the CSV file
            3. Click "Get Recommendations" to see alternatives
            
            ### CSV Format Requirements:
            The app expects a CSV with these columns:
            - `name`: Medicine name
            - `substitute0` through `substitute4`: Potential substitute medicines
            - `sideEffect0` through `sideEffect41`: Possible side effects
            - `use0` through `use4`: Medical uses
            - `Chemical Class`: The chemical classification
            - `Habit Forming`: Whether the medicine is habit-forming
            - `Therapeutic Class`: The therapeutic classification
            - `Action Class`: How the medicine works
            """)
    
    gr.Markdown("""
    ## About This Application
    
    This Medicine Assistant application combines two powerful tools:
    
    1. **Prescription Medicine Extractor**: Uses computer vision AI to identify medicine names from prescription images
    2. **Medicine Alternative Recommender**: Provides detailed information about alternative medications
    
    Both tools utilize the Together AI platform for advanced AI capabilities. Your API key is not stored and is only used to make API calls during your active session.
    
    ### Important Note
    
    This application is for informational purposes only. Always consult with a healthcare professional before making any changes to your medication regimen.
    """)

# Launch the app
if __name__ == "__main__":
    app.launch()