Update src/models.py

src/models.py

@@ -21,33 +21,31 @@ if not anthropic_api_key:
 OPENAI_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
 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}}
-    
-    response = requests.post(url, headers=headers, json=payload)
-    # Raise an HTTPError if the API response was unsuccessful
-    response.raise_for_status()
-    return response.json()
-
 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.chat.completions.create(
+            model=OPENAI_MODEL,
+            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
 
+# Call Anthropic's Claude model with prompt and system message
 def get_claude_completion(prompt, system_message):
     try:
-        response = claude.text_generation(prompt=f"{system_message}\n{prompt}", max_new_tokens=200)
-        return response
+        result = claude.messages.create(
+            model=CLAUDE_MODEL,
+            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