85 lines
2.5 KiB
Python
85 lines
2.5 KiB
Python
import chainlit as cl
|
|
from dotenv import load_dotenv
|
|
load_dotenv()
|
|
|
|
from core.config import DEFAULT_MODEL, MAX_HISTORY
|
|
from core.init_agent import assistant_agent
|
|
from core.stream_response import stream_response
|
|
|
|
from api.fetch_api_data import set_log_callback
|
|
|
|
STARTERS = [
|
|
("What legal data can the agent find?","magnifying_glass"),
|
|
("What is the agent not allowed to do or use?","ban"),
|
|
("What are the details of your AI model?","hexagon"),
|
|
("What data sources does the agent rely on?","database"),
|
|
]
|
|
|
|
PROFILES = [
|
|
("qwen3.5:cloud","Qwen 3.5 CLOUD"),
|
|
("gpt-oss:20b-cloud","GPT-OSS 20B CLOUD"),
|
|
]
|
|
|
|
@cl.set_starters
|
|
async def set_starters():
|
|
return [
|
|
cl.Starter(label=label, message=label, icon=f"/public/icon/{icon}.svg")
|
|
for label, icon in STARTERS
|
|
]
|
|
|
|
@cl.set_chat_profiles
|
|
async def chat_profile():
|
|
return [
|
|
cl.ChatProfile(name=name, markdown_description=f"Uses **{desc}**")
|
|
for name, desc in PROFILES
|
|
]
|
|
|
|
@cl.on_chat_start
|
|
async def start():
|
|
model_name = cl.user_session.get("chat_profile") or DEFAULT_MODEL
|
|
cl.user_session.set("agent", assistant_agent(model_name))
|
|
cl.user_session.set("message_history", [])
|
|
cl.user_session.set("current_model", model_name)
|
|
|
|
|
|
@cl.on_message
|
|
async def main(message: cl.Message):
|
|
current_profile = cl.user_session.get("chat_profile") or DEFAULT_MODEL
|
|
current_agent = cl.user_session.get("agent")
|
|
current_model = cl.user_session.get("current_model")
|
|
|
|
if current_model != current_profile:
|
|
current_agent = assistant_agent(current_profile)
|
|
cl.user_session.set("agent", current_agent)
|
|
cl.user_session.set("current_model", current_profile)
|
|
cl.user_session.set("message_history", [])
|
|
|
|
agent = current_agent
|
|
history = cl.user_session.get("message_history")
|
|
|
|
history.append({"role": "user", "content": message.content})
|
|
|
|
if len(history) > MAX_HISTORY:
|
|
history = history[-MAX_HISTORY:]
|
|
|
|
async with cl.Step(name="🔍 Fetching data...") as step:
|
|
log_lines = []
|
|
|
|
def on_log(line: str):
|
|
log_lines.append(line)
|
|
step.output = "\n".join(log_lines)
|
|
|
|
set_log_callback(on_log)
|
|
|
|
msg = cl.Message(content="")
|
|
async for token in stream_response(agent, history):
|
|
await msg.stream_token(token)
|
|
await msg.update()
|
|
|
|
set_log_callback(None)
|
|
await step.update()
|
|
|
|
if msg.content:
|
|
history.append({"role": "assistant", "content": msg.content})
|
|
|
|
cl.user_session.set("message_history", history) |