check model before send request

app/services/chat/api_client.py

@@ -24,12 +24,18 @@ class GeminiApiClient(ApiClient):
         self.base_url = base_url
         self.timeout = timeout
 
-    async def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
-        timeout = httpx.Timeout(self.timeout, read=self.timeout)
+    def _get_real_model(self, model: str) -> str:
         if model.endswith("-search"):
             model = model[:-7]
         if model.endswith("-image"):
             model = model[:-6]
+
+        return model
+
+    async def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
+        timeout = httpx.Timeout(self.timeout, read=self.timeout)
+        model = self._get_real_model(model)
+
         async with httpx.AsyncClient(timeout=timeout) as client:
             url = f"{self.base_url}/models/{model}:generateContent?key={api_key}"
             response = await client.post(url, json=payload)
@@ -40,10 +46,8 @@ class GeminiApiClient(ApiClient):
 
     async def stream_generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> AsyncGenerator[str, None]:
         timeout = httpx.Timeout(self.timeout, read=self.timeout)
-        if model.endswith("-search"):
-            model = model[:-7]
-        if model.endswith("-image"):
-            model = model[:-6]
+        model = self._get_real_model(model)
+        
         async with httpx.AsyncClient(timeout=timeout) as client:
             url = f"{self.base_url}/models/{model}:streamGenerateContent?alt=sse&key={api_key}"
             async with client.stream(method="POST", url=url, json=payload) as response:

app/services/model_service.py

@@ -68,3 +68,17 @@ class ModelService:
             image_model["id"] = f"{settings.CREATE_IMAGE_MODEL}-chat"
             openai_format["data"].append(image_model)
         return openai_format
+
+    def check_model_support(self, model: str) -> bool:
+        if not model or not isinstance(model, str):
+            return False
+
+        model = model.strip()
+        if model.endswith("-search"):
+            model = model[:-7]
+            return model in settings.MODEL_SEARCH
+        if model.endswith("-image"):
+            model = model[:-6]
+            return model in settings.MODEL_IMAGE
+
+        return True