54 lines
1.5 KiB
Python
54 lines
1.5 KiB
Python
import chainlit as cl
|
|
|
|
from core.init_agent import assistant_agent
|
|
from core.stream_response import stream_response
|
|
|
|
DEFAULT_MODEL = "qwen3.5:cloud"
|
|
|
|
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.on_message
|
|
async def main(message: cl.Message):
|
|
agent = cl.user_session.get("agent")
|
|
history = cl.user_session.get("message_history")
|
|
|
|
history.append({"role": "user", "content": message.content})
|
|
|
|
msg = cl.Message(content="")
|
|
async for token in stream_response(agent, history):
|
|
await msg.stream_token(token)
|
|
await msg.update()
|
|
|
|
if msg.content:
|
|
history.append({"role": "assistant", "content": msg.content}) |