Update src/models.py

src/models.py

@@ -23,6 +23,10 @@ CLAUDE_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
 
 # Define a helper function to perform the direct API call
 def hf_inference_request(token, model, prompt, max_new_tokens=200):
+    # Ensure the prompt is a string
+    if not isinstance(prompt, str):
+        prompt = str(prompt)
+        
     url = f"https://api-inference.huggingface.co/models/{model}"
     headers = {"Authorization": f"Bearer {token}"}
     payload = {"inputs": prompt, "parameters": {"max_new_tokens": max_new_tokens}}
@@ -34,6 +38,11 @@ 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)
         return response
@@ -43,6 +52,11 @@ 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)
         return response