Update src/models.py

src/models.py

@@ -30,29 +30,16 @@ CLAUDE_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
 # Call OpenAI's GPT model with prompt and system message
 def get_gpt_completion(prompt, system_message):
     try:
-        response = openai.chat.completions.create(
-            model=OPENAI_MODEL,
-            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
 
-# Call Anthropic's Claude model with prompt and system message
 def get_claude_completion(prompt, system_message):
     try:
-        result = claude.messages.create(
-            model=CLAUDE_MODEL,
-            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