llm_handler.py

from openai import OpenAI
from params import OPENAI_MODEL, OPENAI_API_KEY
import llamanet

# Create an instance of the OpenAI class
client = OpenAI(api_key=OPENAI_API_KEY)

# Initialize LlamaNet client
llamanet_client = llamanet.Client()

def send_to_chatgpt(msg_list):
    try:
        completion = client.chat.completions.create(
            model=OPENAI_MODEL,
            messages=msg_list,
            temperature=0.6,
            stream=True
        )
        
        chatgpt_response = ""
        for chunk in completion:
            if chunk.choices[0].delta.content is not None:
                chatgpt_response += chunk.choices[0].delta.content
        
        # Note: Usage information might not be available with LlamaNet
        chatgpt_usage = None
        return chatgpt_response, chatgpt_usage
    except Exception as e:
        print(f"Error in send_to_chatgpt: {str(e)}")
        return f"Error: {str(e)}", None

def send_to_llamanet(msg_list):
    try:
        # Convert msg_list to the format expected by LlamaNet
        llamanet_messages = [{"role": msg["role"], "content": msg["content"]} for msg in msg_list]
        
        # Send request to LlamaNet
        response = llamanet_client.chat.completions.create(
            model="llamanet",
            messages=llamanet_messages,
            stream=True
        )
        
        llamanet_response = ""
        for chunk in response:
            if chunk.choices[0].delta.content is not None:
                llamanet_response += chunk.choices[0].delta.content
        
        # LlamaNet doesn't provide usage information
        llamanet_usage = None
        return llamanet_response, llamanet_usage
    except Exception as e:
        print(f"Error in send_to_llamanet: {str(e)}")
        return f"Error: {str(e)}", None

def send_to_llm(provider, msg_list):
    if provider == "llamanet":
        return send_to_llamanet(msg_list)
    elif provider == "openai":
        return send_to_chatgpt(msg_list)
    else:
        raise ValueError(f"Unknown provider: {provider}")