Update src/models.py

src/models.py

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