RouteLLM Chat

import os
from dotenv import load_dotenv
from routellm.controller import Controller

load_dotenv()

# Initialize RouteLLM client
def get_routellm_client():
    nebius_api_key = os.getenv("NEBIUS_API_KEY")
    openai_api_key = os.getenv("OPENAI_API_KEY")
    
    if not nebius_api_key or not openai_api_key:
        return None
    
    # Create controller with model routing
    client = Controller(
        routers=["mf"],  # Model Forwarding router
        strong_model="gpt-4o-mini",
        weak_model="meta-llama/Meta-Llama-3.1-70B-Instruct",
    )
    return client

import streamlit as st

st.set_page_config(
    page_title="RouteLLM Chat", 
    layout="wide", 
    page_icon="🤖"
)

# Chat input and response handling
if prompt := st.chat_input("Type your message here..."):
    # Add user message to history
    st.session_state.messages.append({
        "role": "user", 
        "content": prompt
    })
    
    # Get RouteLLM client
    client = get_routellm_client()
    
    # Generate response
    response = client.chat.completions.create(
        model="router-mf-0.11593",
        messages=st.session_state.messages
    )
    
    # Display response with model badge
    message_content = response.choices[0].message.content
    model_name = getattr(response, "model", "Unknown")
    
    st.session_state.messages.append({
        "role": "assistant",
        "content": message_content,
        "model": model_name
    })

client = Controller(
    routers=["mf"],
    strong_model="gpt-4o-mini",  # Change to any OpenAI model
    weak_model="meta-llama/Meta-Llama-3.1-70B-Instruct",  # Change model
)

client = Controller(
    routers=["sw-ranking"],  # Try different router
    strong_model="gpt-4o-mini",
    weak_model="meta-llama/Meta-Llama-3.1-70B-Instruct",
)

model_badge_color = (
    "#667eea" if "gpt" in model_name.lower() else "#764ba2"
)