Update src/models.py

src/models.py

@@ -18,28 +18,16 @@ claude = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.3", token=anthr
 
 def get_gpt_completion(prompt, system_message):
     try:
-        response = openai.text_generation(
-            model=openai,
-            messages=[
-                {"role": "system", "content": system_message},
-                {"role": "user", "content": prompt}
-            ],
-            stream=False,
-        )
-        return response.choices[0].message.content
+        response = openai.text_generation(prompt=f"{system_message}\n{prompt}", max_new_tokens=200)
+        return response
     except Exception as e:
         print(f"GPT error: {e}")
         raise
 
 def get_claude_completion(prompt, system_message):
     try:
-        response = claude.text_generation(
-            model=claude,
-            max_tokens=2000,
-            system=system_message,
-            messages=[{"role": "user", "content": prompt}]
-        )
-        return result.content[0].text
+        response = claude.text_generation(prompt=f"{system_message}\n{prompt}", max_new_tokens=200)
+        return response
     except Exception as e:
         print(f"Claude error: {e}")
         raise