File size: 6,304 Bytes
2a8a6f9 abe21b9 2a8a6f9 abe21b9 2a8a6f9 | 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 | # app.py
import os
import json
import streamlit as st
import requests
from requests.exceptions import RequestException
BACKEND_URL = os.getenv("BACKEND_URL", "http://127.0.0.1:8000/chat")
DEFAULT_SYSTEM_PROMPT = "You are a helpful assistant."
# Some example OpenRouter model IDs; adjust as needed
PRESET_MODELS = [
"openai/gpt-4o-mini",
"openai/gpt-4o",
"meta-llama/llama-3.1-8b-instruct:free",
"anthropic/claude-3.5-sonnet",
"google/gemini-1.5-flash",
"mistralai/mistral-medium",
"cohere/command-r-plus",
"custom", # sentinel for custom entry
]
st.set_page_config(page_title="OpenRouter Chat", page_icon="💬", layout="centered")
# Session state initialization
if "messages" not in st.session_state:
st.session_state.messages = [{"role": "system", "content": DEFAULT_SYSTEM_PROMPT}]
if "api_key" not in st.session_state:
st.session_state.api_key = "" # kept only in browser memory
if "model_choice" not in st.session_state:
st.session_state.model_choice = PRESET_MODELS[0]
if "custom_model" not in st.session_state:
st.session_state.custom_model = ""
if "temperature" not in st.session_state:
st.session_state.temperature = 0.7
if "max_tokens" not in st.session_state:
st.session_state.max_tokens = 0
st.title("💬 OpenRouter Chat")
with st.expander("Settings", expanded=False):
# API key in browser memory only
st.caption("API key is stored only in your browser session and sent with each request.")
api_key = st.text_input("OpenRouter API Key", type="password", value=st.session_state.api_key)
st.session_state.api_key = api_key
# System prompt editor
current_system = next((m for m in st.session_state.messages if m["role"] == "system"), None)
sys_prompt = st.text_area("System Message", value=current_system["content"] if current_system else DEFAULT_SYSTEM_PROMPT, height=120)
# Model selector
model_choice = st.selectbox("Model (preset)", PRESET_MODELS, index=PRESET_MODELS.index(st.session_state.model_choice) if st.session_state.model_choice in PRESET_MODELS else 0)
st.session_state.model_choice = model_choice
custom_model = ""
if model_choice == "custom":
custom_model = st.text_input("Custom model id (OpenRouter id)", value=st.session_state.custom_model)
st.session_state.custom_model = custom_model
# Sampling params
temperature = st.slider("Temperature", 0.0, 1.5, st.session_state.temperature, 0.1)
max_tokens = st.number_input("Max tokens (0 = auto)", min_value=0, value=st.session_state.max_tokens)
col1, col2 = st.columns(2)
with col1:
apply_clicked = st.button("Apply settings")
with col2:
reset_chat = st.button("Reset conversation")
if apply_clicked:
# Update system message at index 0 or insert if missing
if st.session_state.messages and st.session_state.messages[0]["role"] == "system":
st.session_state.messages[0]["content"] = sys_prompt
else:
st.session_state.messages.insert(0, {"role": "system", "content": sys_prompt})
st.session_state.temperature = float(temperature)
st.session_state.max_tokens = int(max_tokens)
st.success("Settings applied.")
if reset_chat:
st.session_state.messages = [{"role": "system", "content": sys_prompt}]
st.experimental_rerun()
# Resolve model to use
def get_selected_model():
if st.session_state.model_choice == "custom":
return st.session_state.custom_model.strip() or PRESET_MODELS[0]
return st.session_state.model_choice
# Show chat history (exclude system)
for m in st.session_state.messages:
if m["role"] == "user":
with st.chat_message("user"):
st.markdown(m["content"])
elif m["role"] == "assistant":
with st.chat_message("assistant"):
st.markdown(m["content"])
prompt = st.chat_input("Type your message")
def stream_chat(messages, model, temperature, max_tokens, api_key):
payload = {
"messages": messages,
"model": model,
"temperature": temperature,
}
if max_tokens and max_tokens > 0:
payload["max_tokens"] = int(max_tokens)
headers = {}
# Pass API key via header; never persisted server-side
if api_key:
headers["X-OpenRouter-Api-Key"] = api_key
# Optional attribution header
headers["HTTP-Referer"] = "http://localhost:8501"
try:
with requests.post(BACKEND_URL, json=payload, headers=headers, stream=True, timeout=300) as r:
r.raise_for_status()
for line in r.iter_lines(decode_unicode=True):
if not line:
continue
if line.startswith("data: "):
data = line[len("data: "):]
if data.strip() == "done":
break
try:
obj = json.loads(data)
if "content" in obj:
yield obj["content"]
except json.JSONDecodeError:
continue
except RequestException as e:
yield f"\n[Connection error: {e}]"
if prompt:
if not get_selected_model():
st.error("Please select or enter a model.")
else:
# Append user message
st.session_state.messages.append({"role": "user", "content": prompt})
with st.chat_message("assistant"):
ai_area = st.empty()
streamed = ""
for chunk in stream_chat(
st.session_state.messages,
model=get_selected_model(),
temperature=st.session_state.temperature,
max_tokens=st.session_state.max_tokens,
api_key=st.session_state.api_key,
):
streamed += chunk
ai_area.markdown(streamed)
st.session_state.messages.append({"role": "assistant", "content": streamed})
# Sidebar quick actions
with st.sidebar:
if st.button("Clear chat"):
sys_msg = next((m for m in st.session_state.messages if m["role"] == "system"), None)
content = sys_msg["content"] if sys_msg else DEFAULT_SYSTEM_PROMPT
st.session_state.messages = [{"role": "system", "content": content}]
st.experimental_rerun() |