Update main.py

main.py

@@ -1,128 +1,81 @@
 import os
 import shutil
-from fastapi import FastAPI, File, UploadFile, HTTPException, Request
+import tempfile
+from typing import List
+
+from fastapi import FastAPI, File, UploadFile, Form, HTTPException, Request
 from fastapi.responses import JSONResponse, HTMLResponse
 from fastapi.templating import Jinja2Templates
 
-# Import the required modules from google.genai
-from google import genai
-from google.genai import types
-
-app = FastAPI()
-templates = Jinja2Templates(directory="templates")
-
-def generate_question(image_path: str) -> str:
-    # Initialize the client using your GEMINI_API_KEY environment variable.
-    client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
-
-    # Upload the image file
-    files = [client.files.upload(file=image_path)]
-    model = "gemini-2.0-flash"
-
-    # Prepare the content using the image URI and provided sample text
-    contents = [
-        types.Content(
-            role="user",
-            parts=[
-                types.Part.from_uri(
-                    file_uri=files[0].uri,
-                    mime_type=files[0].mime_type,
-                ),
-            ],
-        ),
-        types.Content(
-            role="model",
-            parts=[
-                types.Part.from_text(text="""You are absolutely correct! My apologies for the repeated errors. The correct answer is 24 degrees.
-
-```json
-{
-  \"question\": \"This horizontal bar chart displays the average daily temperatures recorded for each day of the week. If the temperature on Thursday was increased by 1/4th times of the second highest temperature of the week, what temperature would be experienced on Thursday?\",
-  \"options\": [
-    \"a. 35°\",
-    \"b. 24°\",
-    \"c. 18°\",
-    \"d. 30°\"
-  ],
-  \"answer\": \"b\"
-}
-```"""),
-            ],
-        ),
-        types.Content(
-            role="user",
-            parts=[
-                types.Part.from_text(text="INSERT_INPUT_HERE"),
-            ],
-        ),
-    ]
-    tools = [types.Tool(function_declarations=[])]
-    generate_content_config = types.GenerateContentConfig(
-        temperature=1,
-        top_p=0.95,
-        top_k=40,
-        max_output_tokens=8192,
-        tools=tools,
-        response_mime_type="text/plain",
-        system_instruction=[
-            types.Part.from_text(text="""analyse question , analyse option , and give me question , it's options and answer in json format 
+import google.generativeai as genai
 
-fixed json format
+# Configure the API key from environment variables.
+if not os.environ.get("GOOGLE_API_KEY"):
+    raise Exception("GOOGLE_API_KEY environment variable not set.")
+genai.configure(api_key=os.environ.get("GOOGLE_API_KEY"))
 
-below is example of json format ,the json structure is same like below
+# Choose the Gemini model.
+model = genai.GenerativeModel(model_name="gemini-1.5-pro-latest")
 
-{
-  \"question\": \"What is the capital of France?\",
-  \"options\": [
-    \"a. Berlin\",
-    \"b. Madrid\",
-    \"c. Paris\",
-    \"d. Rome\"
-  ],
-  \"answer\": \"c\"
-}
+app = FastAPI()
+templates = Jinja2Templates(directory="templates")
 
-only json format nothing else
-"""),
-        ],
-    )
 
-    # Generate content stream and collect the response
-    response_text = ""
-    for chunk in client.models.generate_content_stream(
-        model=model,
-        contents=contents,
-        config=generate_content_config,
-    ):
-        if chunk.text:
-            response_text += chunk.text
-        elif chunk.function_calls:
-            response_text += chunk.function_calls[0]
+def generate_content(file_paths: List[str], prompt: str) -> str:
+    """
+    Given a list of image file paths and a prompt, call the Gemini model
+    to generate content. The list provided to generate_content will contain
+    the image file paths followed by the prompt.
+    """
+    # Build the contents list: first the file paths, then the prompt.
+    contents = file_paths.copy()
+    contents.append(prompt)
+    
+    try:
+        response = model.generate_content(contents)
+        if not response or not response.text:
+            raise Exception("No response text received.")
+        return response.text
+    except Exception as e:
+        print("Error generating content:", e)
+        raise e
 
-    return response_text
 
 @app.get("/", response_class=HTMLResponse)
 async def index(request: Request):
     return templates.TemplateResponse("index.html", {"request": request})
 
+
 @app.post("/upload", response_class=JSONResponse)
-async def upload_image(file: UploadFile = File(...)):
-    if not file.content_type.startswith("image/"):
-        raise HTTPException(status_code=400, detail="Only image files are accepted.")
+async def upload_images(
+    prompt: str = Form(...),
+    files: List[UploadFile] = File(...),
+):
+    # Validate that each uploaded file is an image.
+    for file in files:
+        if not file.content_type.startswith("image/"):
+            raise HTTPException(status_code=400, detail="Only image files are accepted.")
     
-    # Save uploaded file temporarily
-    temp_dir = "temp"
-    os.makedirs(temp_dir, exist_ok=True)
-    file_path = os.path.join(temp_dir, file.filename)
-    with open(file_path, "wb") as buffer:
-        shutil.copyfileobj(file.file, buffer)
+    temp_dir = tempfile.mkdtemp()
+    saved_paths = []
     
     try:
-        result = generate_question(file_path)
+        # Save each uploaded image temporarily.
+        for file in files:
+            file_path = os.path.join(temp_dir, file.filename)
+            with open(file_path, "wb") as buffer:
+                shutil.copyfileobj(file.file, buffer)
+            saved_paths.append(file_path)
+        
+        # Generate the content using the images and prompt.
+        result = generate_content(saved_paths, prompt)
+        return {"result": result}
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"detail": f"Error processing the image(s): {str(e)}"})
     finally:
-        os.remove(file_path)
-    
-    return {"result": result}
+        # Clean up temporary files.
+        shutil.rmtree(temp_dir, ignore_errors=True)
+
 
 if __name__ == "__main__":
     import uvicorn