Update src/models.py

src/models.py

@@ -12,22 +12,47 @@ load_dotenv(override=True)
 openai_api_key = os.getenv("OPENAI_API_KEY")
 anthropic_api_key = os.getenv("ANTHROPIC_API_KEY")
 
+# Warn if any API key is missing
+if not openai_api_key:
+    print("❌ OpenAI API Key is missing!")
+
+if not anthropic_api_key:
+    print("❌ Anthropic API Key is missing!")
+
 # Initialize disguised API clients
-openai = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.3", token=openai_api_key)
-claude = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.3", token=anthropic_api_key)
+openai = InferenceClient(token=openai_api_key)
+claude = InferenceClient(token=anthropic_api_key)
 
+# Model names
+OPENAI_MODEL = "mistralai/Mistral-7B-Instruct-v0.3"
+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.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
+        raise