Update model_client.py

model_client.py

@@ -1,9 +1,7 @@
 import requests
-from typing import Tuple
-
+from typing import Optional, Tuple
 from config import settings
 
-
 class ModelClient:
     def __init__(self):
         self.primary_model = settings.PRIMARY_CODE_MODEL
@@ -12,7 +10,7 @@ class ModelClient:
         self.temperature = settings.DEFAULT_TEMPERATURE
         self.top_p = settings.DEFAULT_TOP_P
 
-    def _build_payload(self, prompt: str) -> dict:
+    def _build_payload(self, prompt: str, model_name: str) -> dict:
         return {
             "inputs": prompt,
             "parameters": {
@@ -23,32 +21,33 @@ class ModelClient:
             "options": {
                 "wait_for_model": True,
                 "use_cache": False,
-            },
+            }
         }
 
     def _extract_text(self, response_json) -> str:
         if isinstance(response_json, list) and len(response_json) > 0:
-            first_item = response_json[0]
+            first_item = response_json
             if isinstance(first_item, dict) and "generated_text" in first_item:
                 return str(first_item["generated_text"]).strip()
-
+        
         if isinstance(response_json, dict):
             if "generated_text" in response_json:
                 return str(response_json["generated_text"]).strip()
             if "error" in response_json:
                 raise RuntimeError(str(response_json["error"]).strip())
-
+        
         raise RuntimeError("Invalid model response format.")
 
     def _call_huggingface_model(self, prompt: str, model_name: str) -> str:
+    
         api_url = f"https://api-inference.huggingface.co/models/{model_name}"
+        
         headers = {}
-
         hf_token = getattr(settings, "HUGGINGFACE_API_TOKEN", "")
         if hf_token:
             headers["Authorization"] = f"Bearer {hf_token}"
 
-        payload = self._build_payload(prompt)
+        payload = self._build_payload(prompt, model_name)
 
         response = requests.post(
             api_url,
@@ -56,18 +55,12 @@ class ModelClient:
             json=payload,
             timeout=self.timeout,
         )
+        
+       
+        if response.status_code == 404 or "no longer supported" in response.text:
+             api_url = f"https://api-inference.huggingface.co/models/{model_name}" 
 
-        try:
-            response.raise_for_status()
-        except requests.HTTPError:
-            try:
-                error_json = response.json()
-                if isinstance(error_json, dict) and "error" in error_json:
-                    raise RuntimeError(str(error_json["error"]).strip())
-            except ValueError:
-                pass
-            raise
-
+        response.raise_for_status()
         return self._extract_text(response.json())
 
     def generate(self, prompt: str) -> Tuple[str, str, bool]:
@@ -76,13 +69,11 @@ class ModelClient:
             return output, self.primary_model, False
         except Exception as primary_error:
             print(f"Primary model failed: {primary_error}")
-
             try:
                 output = self._call_huggingface_model(prompt, self.fallback_model)
                 return output, self.fallback_model, True
             except Exception as fallback_error:
                 print(f"Fallback model failed: {fallback_error}")
-                raise RuntimeError("Both primary and fallback models failed.")
-
+                raise RuntimeError(f"Both primary and fallback models failed.")
 
 model_client = ModelClient()