Update routers/gemini.py

routers/gemini.py

@@ -1,101 +1,101 @@
-from fastapi import APIRouter, Depends, HTTPException
-from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
-from typing import Annotated
-from google import genai
-from google.genai import types
-from auth import verify_token
-import os
-import httpx
-import base64
-
-router = APIRouter(prefix="/gemini", tags=["gemini"])
-
-gemini_key = os.environ.get('GEMINI_KEY', '') # Changed variable name for clarity
-# Configure the genai client safely
-if gemini_key:
-    genai.configure(api_key=gemini_key)
-else:
-    print("Warning: GEMINI_KEY environment variable not set.")
-    # Optionally raise an error or handle the missing key appropriately
-    # raise ValueError("GEMINI_KEY environment variable is required.")
-
-
-class LLMRequest(BaseModel):
-    model: str
-    prompt: str
-
-class GeminiMultimodalRequest(BaseModel):
-    model: str
-    prompt: str
-    image: str # url or base64
-
-@router.post("/")
-async def gemini_chat(request: LLMRequest, token: Annotated[str, Depends(verify_token)]):
-    if not gemini_key:
-         raise HTTPException(status_code=500, detail="Gemini API key not configured")
-    model = genai.GenerativeModel(request.model)
-    async def generate():
-        # Use generate_content_async for async streaming
-        response = await model.generate_content_async(
-            contents=[request.prompt], stream=True)
-
-        async for chunk in response:
-            if chunk.text:
-                yield chunk.text
-
-    return StreamingResponse(generate(), media_type="text/plain")
-
-
-@router.post("/multimodal")
-async def gemini_multimodal(request: GeminiMultimodalRequest, token: Annotated[str, Depends(verify_token)]):
-    if not gemini_key:
-         raise HTTPException(status_code=500, detail="Gemini API key not configured")
-    model = genai.GenerativeModel(request.model)
-    image_part = None
-    try:
-        if request.image.startswith('http'):
-            async with httpx.AsyncClient() as client:
-                img_response = await client.get(request.image)
-                img_response.raise_for_status() # Check if the request was successful
-                # Determine mime type dynamically if possible, default to jpeg
-                content_type = img_response.headers.get('Content-Type', 'image/jpeg')
-                image_part = types.Part.from_bytes(data=img_response.content, mime_type=content_type)
-        elif request.image.startswith('data:image'):
-             # Handle base64 data URI
-             header, encoded = request.image.split(',', 1)
-             mime_type = header.split(':')[1].split(';')[0]
-             image_data = base64.b64decode(encoded)
-             image_part = types.Part.from_bytes(data=image_data, mime_type=mime_type)
-        else:
-             # Assume raw base64
-             image_part = types.Part.from_bytes(data=base64.b64decode(request.image), mime_type="image/jpeg") # Default mime or raise error
-
-    except httpx.HTTPStatusError as e:
-         raise HTTPException(status_code=400, detail=f"Failed to fetch image from URL: {e.response.status_code}")
-    except (ValueError, TypeError, base64.binascii.Error) as e:
-         raise HTTPException(status_code=400, detail=f"Invalid image data: {e}")
-    except Exception as e: # Catch other potential errors during image processing
-         raise HTTPException(status_code=500, detail=f"Error processing image: {e}")
-
-
-    if image_part is None:
-        raise HTTPException(status_code=400, detail="Could not process image input.")
-
-    try:
-        response = await model.generate_content_async(
-            contents=[request.prompt, image_part]
-        )
-        # Access the text content safely
-        response_text = response.text if hasattr(response, 'text') else "No text content generated."
-        return {"response": response_text}
-    except Exception as e:
-        # Log the error for debugging
-        print(f"Error during Gemini content generation: {e}")
-        # Provide a more informative error message
-        error_detail = f"Error generating content with Gemini model: {e}"
-        # Check for specific API errors if the SDK provides them
-        if hasattr(e, 'message'):
-            error_detail = f"Gemini API Error: {e.message}"
-        raise HTTPException(status_code=500, detail=error_detail)
-
+from fastapi import APIRouter, Depends, HTTPException
+from fastapi.responses import StreamingResponse
+from pydantic import BaseModel
+from typing import Annotated
+from google import genai
+from google.genai import types
+from auth import verify_token
+import os
+import httpx
+import base64
+
+router = APIRouter(prefix="/gemini", tags=["gemini"])
+
+gemini_key = os.environ.get('GEMINI_KEY', '') # Changed variable name for clarity
+# Configure the genai client safely
+if gemini_key:
+    genai.Client(api_key=gemini_key)
+else:
+    print("Warning: GEMINI_KEY environment variable not set.")
+    # Optionally raise an error or handle the missing key appropriately
+    # raise ValueError("GEMINI_KEY environment variable is required.")
+
+
+class LLMRequest(BaseModel):
+    model: str
+    prompt: str
+
+class GeminiMultimodalRequest(BaseModel):
+    model: str
+    prompt: str
+    image: str # url or base64
+
+@router.post("/")
+async def gemini_chat(request: LLMRequest, token: Annotated[str, Depends(verify_token)]):
+    if not gemini_key:
+         raise HTTPException(status_code=500, detail="Gemini API key not configured")
+    model = genai.GenerativeModel(request.model)
+    async def generate():
+        # Use generate_content_async for async streaming
+        response = await model.generate_content_async(
+            contents=[request.prompt], stream=True)
+
+        async for chunk in response:
+            if chunk.text:
+                yield chunk.text
+
+    return StreamingResponse(generate(), media_type="text/plain")
+
+
+@router.post("/multimodal")
+async def gemini_multimodal(request: GeminiMultimodalRequest, token: Annotated[str, Depends(verify_token)]):
+    if not gemini_key:
+         raise HTTPException(status_code=500, detail="Gemini API key not configured")
+    model = genai.GenerativeModel(request.model)
+    image_part = None
+    try:
+        if request.image.startswith('http'):
+            async with httpx.AsyncClient() as client:
+                img_response = await client.get(request.image)
+                img_response.raise_for_status() # Check if the request was successful
+                # Determine mime type dynamically if possible, default to jpeg
+                content_type = img_response.headers.get('Content-Type', 'image/jpeg')
+                image_part = types.Part.from_bytes(data=img_response.content, mime_type=content_type)
+        elif request.image.startswith('data:image'):
+             # Handle base64 data URI
+             header, encoded = request.image.split(',', 1)
+             mime_type = header.split(':')[1].split(';')[0]
+             image_data = base64.b64decode(encoded)
+             image_part = types.Part.from_bytes(data=image_data, mime_type=mime_type)
+        else:
+             # Assume raw base64
+             image_part = types.Part.from_bytes(data=base64.b64decode(request.image), mime_type="image/jpeg") # Default mime or raise error
+
+    except httpx.HTTPStatusError as e:
+         raise HTTPException(status_code=400, detail=f"Failed to fetch image from URL: {e.response.status_code}")
+    except (ValueError, TypeError, base64.binascii.Error) as e:
+         raise HTTPException(status_code=400, detail=f"Invalid image data: {e}")
+    except Exception as e: # Catch other potential errors during image processing
+         raise HTTPException(status_code=500, detail=f"Error processing image: {e}")
+
+
+    if image_part is None:
+        raise HTTPException(status_code=400, detail="Could not process image input.")
+
+    try:
+        response = await model.generate_content_async(
+            contents=[request.prompt, image_part]
+        )
+        # Access the text content safely
+        response_text = response.text if hasattr(response, 'text') else "No text content generated."
+        return {"response": response_text}
+    except Exception as e:
+        # Log the error for debugging
+        print(f"Error during Gemini content generation: {e}")
+        # Provide a more informative error message
+        error_detail = f"Error generating content with Gemini model: {e}"
+        # Check for specific API errors if the SDK provides them
+        if hasattr(e, 'message'):
+            error_detail = f"Gemini API Error: {e.message}"
+        raise HTTPException(status_code=500, detail=error_detail)
+