Add Gemini AI integration for fixing reaction data with API key input and fix button

app.py

@@ -6,7 +6,10 @@ from rdkit import Chem
 from rdkit.Chem import AllChem, Draw
 import io
 import tempfile
+import base64
 import os
+from google import genai
+from google.genai import types
 from reportlab.pdfgen import canvas
 from reportlab.lib.pagesizes import letter
 from reportlab.lib.styles import getSampleStyleSheet
@@ -205,6 +208,85 @@ def get_autocomplete_reactions(query):
     matches = process.extract(query, reaction_names, limit=10)
     return [m[0] for m in matches if m[1] > 60]
 
+def fix_reaction_with_gemini(reaction_name, api_key):
+    if not api_key:
+        return "Please provide a Gemini API key."
+    
+    try:
+        client = genai.Client(api_key=api_key)
+        
+        prompt = f"""Please provide detailed information about the organic reaction named "{reaction_name}". 
+        Include the correct reaction name, reactants, reagents, products, byproducts, reaction conditions, mechanism, and description.
+        Make sure to provide accurate chemical information."""
+        
+        contents = [
+            types.Content(
+                role="user",
+                parts=[types.Part.from_text(text=prompt)],
+            ),
+        ]
+        
+        generate_content_config = types.GenerateContentConfig(
+            thinking_config=types.ThinkingConfig(thinking_budget=-1),
+            response_mime_type="application/json",
+            response_schema=genai.types.Schema(
+                type=genai.types.Type.OBJECT,
+                required=["reaction name", "reactants", "reagents", "products", "byproducts", "conditions", "mechanism", "description"],
+                properties={
+                    "reaction name": genai.types.Schema(type=genai.types.Type.STRING),
+                    "reactants": genai.types.Schema(
+                        type=genai.types.Type.ARRAY,
+                        items=genai.types.Schema(type=genai.types.Type.STRING),
+                    ),
+                    "reagents": genai.types.Schema(
+                        type=genai.types.Type.ARRAY,
+                        items=genai.types.Schema(type=genai.types.Type.STRING),
+                    ),
+                    "products": genai.types.Schema(
+                        type=genai.types.Type.ARRAY,
+                        items=genai.types.Schema(type=genai.types.Type.STRING),
+                    ),
+                    "byproducts": genai.types.Schema(
+                        type=genai.types.Type.ARRAY,
+                        items=genai.types.Schema(type=genai.types.Type.STRING),
+                    ),
+                    "conditions": genai.types.Schema(type=genai.types.Type.STRING),
+                    "mechanism": genai.types.Schema(type=genai.types.Type.STRING),
+                    "description": genai.types.Schema(type=genai.types.Type.STRING),
+                },
+            ),
+        )
+        
+        response_text = ""
+        for chunk in client.models.generate_content_stream(
+            model="gemini-2.0-flash-exp",
+            contents=contents,
+            config=generate_content_config,
+        ):
+            response_text += chunk.text
+        
+        # Parse the JSON response
+        import json
+        gemini_data = json.loads(response_text)
+        
+        # Update the global df (this won't persist, but shows the updated data)
+        global df
+        mask = df['corrected_name'].str.lower() == reaction_name.lower()
+        if mask.any():
+            idx = df[mask].index[0]
+            df.at[idx, 'corrected_name'] = gemini_data.get('reaction name', reaction_name)
+            df.at[idx, 'general_reactants'] = ', '.join(gemini_data.get('reactants', []))
+            df.at[idx, 'general_reagents'] = ', '.join(gemini_data.get('reagents', []))
+            df.at[idx, 'general_products'] = ', '.join(gemini_data.get('products', []))
+            df.at[idx, 'description'] = gemini_data.get('description', df.at[idx, 'description'])
+            
+            return f"Successfully updated reaction '{reaction_name}' with AI-generated data:\n\n**New Name:** {gemini_data.get('reaction name')}\n**Reactants:** {', '.join(gemini_data.get('reactants', []))}\n**Reagents:** {', '.join(gemini_data.get('reagents', []))}\n**Products:** {', '.join(gemini_data.get('products', []))}\n**Description:** {gemini_data.get('description', '')[:200]}..."
+        else:
+            return f"Reaction '{reaction_name}' not found in database."
+            
+    except Exception as e:
+        return f"Error calling Gemini API: {str(e)}"
+
 def get_autocomplete_reactants(query):
     if not query:
         return unique_reactants[:10]
@@ -252,15 +334,27 @@ with gr.Blocks(title="Organic Reactions Search") as demo:
         pdf_btn.click(generate_all_reactions_pdf, outputs=pdf_output)
 
     with gr.Tab("View All Reactions (Table)"):
-        gr.Markdown("Browse all 828 reactions in a tabular format.")
+        gr.Markdown("Browse all 828 reactions in a tabular format. Use the AI Fix section below to improve reaction data.")
+        
+        # AI Fix section
+        with gr.Row():
+            api_key_input = gr.Textbox(label="Gemini API Key", type="password", placeholder="Enter your Gemini API key")
+            reaction_to_fix = gr.Dropdown(label="Select Reaction to Fix", choices=reaction_names)
+        
+        fix_button = gr.Button("Fix with AI")
+        ai_status = gr.Markdown(label="AI Fix Status")
+        
+        fix_button.click(fix_reaction_with_gemini, inputs=[reaction_to_fix, api_key_input], outputs=ai_status)
+        
+        gr.Markdown("---")
+        gr.Markdown("**Database Table:**")
         
-        # Create HTML table with buttons
+        # Create HTML table (read-only for browsing)
         def create_reactions_table():
             html = """
             <table style="width:100%; border-collapse: collapse;">
                 <thead>
                     <tr style="background-color: #f2f2f2;">
-                        <th style="border: 1px solid #ddd; padding: 8px;">Action</th>
                         <th style="border: 1px solid #ddd; padding: 8px;">Reaction Name</th>
                         <th style="border: 1px solid #ddd; padding: 8px;">Reactants</th>
                         <th style="border: 1px solid #ddd; padding: 8px;">Reactants SMILES</th>
@@ -289,7 +383,6 @@ with gr.Blocks(title="Organic Reactions Search") as demo:
                 
                 html += f"""
                     <tr>
-                        <td style="border: 1px solid #ddd; padding: 8px;"><button onclick="alert('Fix with AI functionality for {reaction_name} would be implemented here')">Fix with AI</button></td>
                         <td style="border: 1px solid #ddd; padding: 8px;">{reaction_name}</td>
                         <td style="border: 1px solid #ddd; padding: 8px;">{reactants}</td>
                         <td style="border: 1px solid #ddd; padding: 8px; font-family: monospace; font-size: 12px;">{reactants_smiles}</td>

requirements.txt

@@ -5,4 +5,5 @@ fuzzywuzzy
 python-levenshtein
 rdkit
 reportlab
-svglib
+svglib
+google-genai