helpers/ai_client.py

# helpers/ai_client.py
import os
from openai import OpenAI
from typing import Optional
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class AIClient:
    def __init__(self):
        # Initialize OpenAI client
        self.client = OpenAI(api_key=os.getenv("OPENROUTER_API_KEY"), base_url="https://openrouter.ai/api/v1")
    def chat(
        self,
        prompt: str,
        system_message: str = "",
        model_id: str = "openai/gpt-4o-mini",
        conversation_id: str = "",
        user_id: str = "string"
    ) -> str:
        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})

            # Create completion
            completion = self.client.chat.completions.create(
                model=model_id,
                messages=messages
            )

            # Return the response text
            return completion.choices[0].message.content

        except Exception as e:
            logger.error(f"Error in chat completion: {str(e)}")
            raise Exception(f"Error from OpenRouter API: {str(e)}")