Update code

model_client.py

@@ -12,6 +12,7 @@ class ModelClient:
         self.timeout = settings.MODEL_TIMEOUT_SECONDS
         self.temperature = settings.DEFAULT_TEMPERATURE
         self.top_p = settings.DEFAULT_TOP_P
+        self.max_tokens = settings.MAX_OUTPUT_TOKENS
         self.hf_token = getattr(settings, "HUGGINGFACE_API_TOKEN", "")
 
     def _create_client(self) -> InferenceClient:
@@ -20,7 +21,31 @@ class ModelClient:
             timeout=self.timeout,
         )
 
+    def _extract_content(self, response) -> str:
+        if not response or not getattr(response, "choices", None):
+            raise RuntimeError("Empty response from model.")
+
+        first_choice = response.choices[0]
+        if not first_choice or not getattr(first_choice, "message", None):
+            raise RuntimeError("Model returned an invalid response structure.")
+
+        message = first_choice.message
+        content = getattr(message, "content", None)
+
+        if content is None:
+            raise RuntimeError("Model returned no content.")
+
+        cleaned = str(content).strip()
+        if not cleaned:
+            raise RuntimeError("Model returned empty content.")
+
+        return cleaned
+
     def _call_model(self, prompt: str, model_name: str) -> str:
+        cleaned_prompt = str(prompt or "").strip()
+        if not cleaned_prompt:
+            raise RuntimeError("Prompt is empty.")
+
         client = self._create_client()
 
         response = client.chat.completions.create(
@@ -28,35 +53,30 @@ class ModelClient:
             messages=[
                 {
                     "role": "user",
-                    "content": prompt,
+                    "content": cleaned_prompt,
                 }
             ],
             temperature=self.temperature,
             top_p=self.top_p,
-            max_tokens=1200,
+            max_tokens=self.max_tokens,
         )
 
-        if not response or not response.choices:
-            raise RuntimeError("Empty response from model.")
-
-        message = response.choices[0].message
-        if not message or not message.content:
-            raise RuntimeError("Model returned no content.")
-
-        return str(message.content).strip()
+        return self._extract_content(response)
 
     def generate(self, prompt: str) -> Tuple[str, str, bool]:
         try:
             output = self._call_model(prompt, self.primary_model)
             return output, self.primary_model, False
+
         except Exception as primary_error:
-            print(f"Primary model failed: {primary_error}")
+            print(f"Primary model failed: {primary_error}", flush=True)
 
             try:
                 output = self._call_model(prompt, self.fallback_model)
                 return output, self.fallback_model, True
+
             except Exception as fallback_error:
-                print(f"Fallback model failed: {fallback_error}")
+                print(f"Fallback model failed: {fallback_error}", flush=True)
                 raise RuntimeError("Both primary and fallback models failed.")