Update src/models.py

src/models.py

@@ -38,13 +38,7 @@ def hf_inference_request(token, model, prompt, max_new_tokens=200):
 
 def get_gpt_completion(prompt, system_message):
     try:
-        # Ensure both system_message and prompt are strings
-        if not isinstance(system_message, str):
-            system_message = str(system_message)
-        if not isinstance(prompt, str):
-            prompt = str(prompt)
-        full_prompt = f"{system_message}\n{prompt}"
-        response = hf_inference_request(openai_api_key, OPENAI_MODEL, full_prompt, max_new_tokens=200)
+        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}")
@@ -52,13 +46,7 @@ def get_gpt_completion(prompt, system_message):
 
 def get_claude_completion(prompt, system_message):
     try:
-        # Ensure both system_message and prompt are strings
-        if not isinstance(system_message, str):
-            system_message = str(system_message)
-        if not isinstance(prompt, str):
-            prompt = str(prompt)
-        full_prompt = f"{system_message}\n{prompt}"
-        response = hf_inference_request(anthropic_api_key, CLAUDE_MODEL, full_prompt, max_new_tokens=200)
+        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}")