Update src/main.py

src/main.py

@@ -1,46 +1,44 @@
 import os
+import io
+import base64
 from fastapi import FastAPI, HTTPException, Header
 from pydantic import BaseModel
-from groq import Groq
+from google import genai
+from PIL import Image
 
 app = FastAPI()
 
 # Get the key from environment variables
-GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
-client = Groq(api_key=GROQ_API_KEY)
+GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
+
+# Initialize Gemini Client
+client = genai.Client(api_key=GEMINI_API_KEY)
 
 class ImageRequest(BaseModel):
     base64_image: str
 
 @app.post("/explain-chart")
 async def explain_chart(request: ImageRequest, authorization: str = Header(None)):
-    
-    if not authorization or authorization != f"Bearer {GROQ_API_KEY}":
+    # Security check (matching the key sent from Streamlit)
+    if not authorization or authorization != f"Bearer {GEMINI_API_KEY}":
         raise HTTPException(status_code=401, detail="Unauthorized")
-
+    
     try:
-        completion = client.chat.completions.create(
-            # Using the vision-capable model
-            model="llama-3.2-11b-vision-preview", 
-            messages=[
-                {
-                    "role": "user",
-                    "content": [
-                        {"type": "text", "text": "Explain this IPL data chart. What are the key insights?"},
-                        {
-                            "type": "image_url",
-                            "image_url": {"url": f"data:image/png;base64,{request.base64_image}"}
-                        }
-                    ]
-                }
+        # 1. Decode the base64 string back into an image
+        image_data = base64.b64decode(request.base64_image)
+        img = Image.open(io.BytesIO(image_data))
+
+        # 2. Call Gemini 1.5 Flash (optimized for vision/charts)
+        response = client.models.generate_content(
+            model="gemini-1.5-flash",
+            contents=[
+                "You are an expert IPL data analyst. Explain this chart and provide the key insights in concise bullet points.",
+                img
             ]
         )
-        return {"explanation": completion.choices[0].message.content}
+
+        return {"explanation": response.text}
+
     except Exception as e:
         print(f"Error detail: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-    
-# after clicking the button in the Streamlit app, it will send a POST request to this endpoint 
-# with the base64 image, and the FastAPI server will process it and return the AI-generated explanation.
-# this is happening in the background, so the user experience is seamless. 
-
+        raise HTTPException(status_code=500, detail=str(e))