Update app.py with Open Router endpoint

app.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, Depends
+from fastapi import FastAPI, Depends, Query
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
 from typing import Annotated
@@ -6,9 +6,11 @@ from mistralai import Mistral
 from google import genai
 from google.genai import types
 from auth import verify_token
+from enum import Enum
 import os
 import httpx
 import base64
+import json
 
 app = FastAPI()
 
@@ -18,6 +20,8 @@ mistral_client = Mistral(api_key=mistral)
 gemini = os.environ.get('GEMINI_KEY', '')
 gemini_client = genai.Client(api_key=gemini)
 
+open_router = os.environ.get('OPEN_ROUTER_KEY', '')
+
 @app.get("/")
 def hello():
     return {"Hello": "World!"}
@@ -76,4 +80,36 @@ async def gemini_multimodal(request: GeminiMultimodalRequest, token: Annotated[s
         contents=[request.prompt, image]
     )
 
-    return {"response": response.text}
+    return {"response": response.text}
+
+class ModelName(str, Enum):
+    deepseek_r1 = "deepseek/deepseek-r1:free"
+    gemini_2_flash_lite = "google/gemini-2.0-flash-lite-preview-02-05:free"
+    gemini_2_pro = "google/gemini-2.0-pro-exp-02-05:free"
+    llama_3_3 = "meta-llama/llama-3.3-70b-instruct:free"
+    mistral_small_3 ="mistralai/mistral-small-24b-instruct-2501:free"
+    
+@app.post("/open-router")
+async def open_router(model: ModelName = Query(..., description="Select a model"), prompt: str, token: Annotated[str, Depends(verify_token)]):
+    async with httpx.AsyncClient() as client:
+        response = await client.post(
+            url="https://openrouter.ai/api/v1/chat/completions",
+            headers={
+                "Authorization": f"Bearer {open_router}",
+                "Content-Type": "application/json",
+                "HTTP-Referer": "<YOUR_SITE_URL>",  # Optional
+                "X-Title": "<YOUR_SITE_NAME>",  # Optional
+            },
+            json={
+                "model": model,
+                "messages": [
+                    {
+                        "role": "user",
+                        "content": prompt,
+                    }
+                ],
+            }
+        )
+
+        response.raise_for_status()  # Raise HTTPError for bad responses (4xx or 5xx)
+        return response.json()