Spaces:
Sleeping
Sleeping
File size: 2,468 Bytes
367cc0b 25f7a75 483491a 367cc0b 25f7a75 367cc0b 0e97127 367cc0b 70650cf 25f7a75 a6361a2 367cc0b 25f7a75 367cc0b a6361a2 367cc0b a6361a2 367cc0b 25f7a75 367cc0b 25f7a75 367cc0b 25f7a75 483491a 25f7a75 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 |
# 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):
# Load environment variables
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,
prompt: str,
system_message: str = "",
model_id: str = "openai/gpt-4o-mini",
conversation_id: str = "",
user_id: str = "string"
) -> str:
"""
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.
model_id (str): The model ID to use (default: "openai/gpt-4o-mini").
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.
"""
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(
extra_headers={
"HTTP-Referer": self.site_url, # Optional
"X-Title": self.site_name, # Optional
},
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)}") |