Update helpers/ai_client.py

helpers/ai_client.py

@@ -1,7 +1,7 @@
 # helpers/ai_client.py
-import requests
 import os
-from typing import Optional, Dict, Any
+from openai import OpenAI
+from typing import Optional
 import logging
 
 # Set up logging
@@ -11,8 +11,16 @@ logger = logging.getLogger(__name__)
 class AIClient:
     def __init__(self):
         # Load environment variables
-        self.llm_api_url = os.getenv("LLM_API_URL")
-        self.api_key = os.getenv("X_API_KEY")
+        self.base_url = os.getenv("LLM_API_URL", "https://openrouter.ai/api/v1")
+        self.api_key = os.getenv("OPENROUTER_API_KEY")
+        self.site_url = os.getenv("SITE_URL", "")  # Your site URL
+        self.site_name = os.getenv("SITE_NAME", "")  # Your site name
+        
+        # Initialize OpenAI client
+        self.client = OpenAI(
+            base_url=self.base_url,
+            api_key=self.api_key
+        )
 
     def chat(
         self,
@@ -20,12 +28,10 @@ class AIClient:
         system_message: str = "",
         model_id: str = "openai/gpt-4o-mini",
         conversation_id: str = "",
-        user_id: str = "string",
-        api_key: Optional[str] = None
+        user_id: str = "string"
     ) -> str:
         """
-        Sends a prompt to the LLM API and returns the response as text.
-
+        Sends a prompt to the OpenRouter API and returns the response as text.
         Args:
             prompt (str): The user's input prompt.
             system_message (str): Optional system message for the LLM.
@@ -33,34 +39,34 @@ class AIClient:
             conversation_id (str): Unique ID for the conversation.
             user_id (str): Unique ID for the user.
             api_key (str): API key for authentication.
-
         Returns:
             str: The text response from the LLM API.
-
         Raises:
             Exception: If the API request fails.
         """
-        if api_key is None:
-            api_key = self.api_key
-
-        payload = {
-            "prompt": prompt,
-            "system_message": system_message,
-            "model_id": model_id,
-            "conversation_id": conversation_id,
-            "user_id": user_id
-        }
+        try:
+            messages = []
+            
+            # Add system message if provided
+            if system_message:
+                messages.append({"role": "system", "content": system_message})
+            
+            # Add user message
+            messages.append({"role": "user", "content": prompt})
 
-        headers = {
-            "accept": "application/json",
-            "X-API-Key": api_key,
-            "Content-Type": "application/json"
-        }
+            # Create completion
+            completion = self.client.chat.completions.create(
+                extra_headers={
+                    "HTTP-Referer": self.site_url,  # Optional
+                    "X-Title": self.site_name,      # Optional
+                },
+                model=model_id,
+                messages=messages
+            )
 
-        # Use requests to call the external API
-        response = requests.post(self.llm_api_url, json=payload, headers=headers)
-        if response.status_code != 200:
-            raise Exception(f"Error from LLM API: {response.status_code} - {response.text}")
+            # Return the response text
+            return completion.choices[0].message.content
 
-        # Return the response as text
-        return response.text
+        except Exception as e:
+            logger.error(f"Error in chat completion: {str(e)}")
+            raise Exception(f"Error from OpenRouter API: {str(e)}")