Update orchestrator/gemini.py

orchestrator/gemini.py

@@ -1,30 +1,58 @@
 import os
 import streamlit as st
+
+# Try to import the Gemini client
 try:
     import google.generativeai as genai
 except ImportError:
     genai = None
 
 def get_gemini_api_key():
+    """
+    Return Gemini API key from Streamlit secrets or environment variable.
+    """
+    # Try Streamlit Secrets (recommended on Spaces)
     if hasattr(st.secrets, "GEMINI_API_KEY"):
         return st.secrets["GEMINI_API_KEY"]
+    # Fallback to environment variable
     return os.getenv("GEMINI_API_KEY")
 
-def gemini_generate_code(prompt, model="gemini-1.5-pro"):
+def ensure_genai():
+    if genai is None:
+        raise ImportError(
+            "google-generativeai is not installed. Please add it to requirements.txt and re-deploy."
+        )
+
+def gemini_generate(prompt, model="gemini-1.5-pro"):
+    """
+    Simple Gemini text completion for Q&A or summarization.
+    """
+    ensure_genai()
     api_key = get_gemini_api_key()
     if not api_key:
-        raise RuntimeError("Gemini API key not set!")
-    if genai is None:
-        raise ImportError("Please install google-generativeai!")
+        raise RuntimeError("Gemini API key not set in Streamlit secrets or environment variable.")
     genai.configure(api_key=api_key)
     model_obj = genai.GenerativeModel(model)
     response = model_obj.generate_content(prompt)
+    # Try to return the 'text' property (for Gemini's Python client), fallback to string
     return response.text if hasattr(response, "text") else str(response)
 
-def gemini_explain_code(code, model="gemini-1.5-pro"):
-    # Ask Gemini to explain code in plain English
-    explanation_prompt = (
-        "Explain step-by-step, in clear language, what this Python MCP server code does, and draw a high-level block diagram in Mermaid format:\n\n" +
-        code
-    )
-    return gemini_generate_code(explanation_prompt, model)
+def gemini_generate_code(system_instruction, user_prompt, model="gemini-1.5-pro"):
+    """
+    Use Gemini for code generation with a system prompt and a user prompt.
+    Returns: generated code as string.
+    """
+    ensure_genai()
+    api_key = get_gemini_api_key()
+    if not api_key:
+        raise RuntimeError("Gemini API key not set in Streamlit secrets or environment variable.")
+    genai.configure(api_key=api_key)
+    model_obj = genai.GenerativeModel(model)
+    # Compose prompt as role-annotated list for best model behavior
+    prompt = [
+        {"role": "system", "parts": [system_instruction]},
+        {"role": "user", "parts": [user_prompt]},
+    ]
+    response = model_obj.generate_content(prompt)
+    # Try to return the 'text' property (for Gemini's Python client), fallback to string
+    return response.text if hasattr(response, "text") else str(response)