Update app.py

app.py

@@ -263,8 +263,9 @@ async def chat_completion(request: ChatRequest, authorization: str = Header(None
                     return response.model_dump()
             else:
                 gemini_messages = convert_messages_to_gemini_format(request.messages)
+                model = request.model
                 payload = {}
-                if request.model == "gemini-2.0-flash-exp-search":
+                if model == "gemini-2.0-flash-exp-search":
                     payload = {
                         "contents": gemini_messages,
                         "generationConfig": {
@@ -276,6 +277,7 @@ async def chat_completion(request: ChatRequest, authorization: str = Header(None
                               }
                         ]
                     }
+                    model = "gemini-2.0-flash-exp"
                 else:
                     payload = {
                         "contents": gemini_messages,
@@ -293,7 +295,7 @@ async def chat_completion(request: ChatRequest, authorization: str = Header(None
                         while retries < MAX_RETRIES:
                             try:
                                 async with httpx.AsyncClient() as client:
-                                    stream_url = f"https://generativelanguage.googleapis.com/v1beta/models/{request.model}:streamGenerateContent?alt=sse&key={api_key}"
+                                    stream_url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:streamGenerateContent?alt=sse&key={api_key}"
                                     async with client.stream("POST", stream_url, json=payload, timeout=60.0) as response:
                                         if response.status_code == 429:
                                             logger.warning(f"Rate limit reached for key: {api_key}")
@@ -377,7 +379,7 @@ async def chat_completion(request: ChatRequest, authorization: str = Header(None
                     return StreamingResponse(content=generate(), media_type="text/event-stream")
                 else:
                     async with httpx.AsyncClient() as client:
-                        non_stream_url = f"https://generativelanguage.googleapis.com/v1beta/models/{request.model}:generateContent?key={api_key}"
+                        non_stream_url = f"https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent?key={api_key}"
                         response = await client.post(non_stream_url, json=payload)
                         
                         if response.status_code != 200:
@@ -397,7 +399,7 @@ async def chat_completion(request: ChatRequest, authorization: str = Header(None
 
                         gemini_response = response.json()
                         logger.info("Chat completion successful")
-                        return await convert_gemini_response_to_openai(gemini_response, request.model)
+                        return await convert_gemini_response_to_openai(gemini_response, model)
 
         except Exception as e:
             logger.error(f"Error in chat completion: {str(e)}")