Update main.py

main.py

@@ -501,6 +501,8 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
 
     env_vars = get_env_vars()
 
+    target_url_path = "/v1/chat/completions" # Default path
+
     if model_to_use in mistral_models:
         endpoint = env_vars['mistral_api']
         custom_headers = {
@@ -524,6 +526,7 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
         custom_headers = {
             "Authorization": f"Bearer {env_vars['gemini_key']}"
         }
+        target_url_path = "/chat/completions" # Use /chat/completions for Gemini
     else:
         endpoint = env_vars['secret_api_endpoint']
         custom_headers = {
@@ -532,12 +535,12 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
             "Referer": header_url
         }
 
-    print(f"Using endpoint: {endpoint} for model: {model_to_use}")
+    print(f"Using endpoint: {endpoint} with path: {target_url_path} for model: {model_to_use}")
 
     async def real_time_stream_generator():
         try:
             async with httpx.AsyncClient(timeout=60.0) as client:
-                async with client.stream("POST", f"{endpoint}/v1/chat/completions", json=payload_dict, headers=custom_headers) as response:
+                async with client.stream("POST", f"{endpoint}{target_url_path}", json=payload_dict, headers=custom_headers) as response:
                     if response.status_code >= 400:
                         error_messages = {
                             422: "Unprocessable entity. Check your payload.",
@@ -576,7 +579,6 @@ async def get_completion(payload: Payload, request: Request, authenticated: bool
         async for chunk in real_time_stream_generator():
             response_content.append(chunk)
         return JSONResponse(content=json.loads(''.join(response_content)))
-
 @app.post("/images/generations")
 async def create_image(payload: ImageGenerationPayload, authenticated: bool = Depends(verify_api_key)):
     if not server_status: