-
Notifications
You must be signed in to change notification settings - Fork 106
/
polling.py
71 lines (61 loc) · 3.07 KB
/
polling.py
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
from fasthtml.common import *
from claudette import *
# Set up the app, including daisyui and tailwind for the chat component
tlink = Script(src="https://cdn.tailwindcss.com"),
dlink = Link(rel="stylesheet", href="https://cdn.jsdelivr.net/npm/[email protected]/dist/full.min.css")
app = FastHTML(hdrs=(tlink, dlink, picolink))
# Set up a chat model client and list of messages (https://claudette.answer.ai/)
cli = Client(models[-1])
sp = """You are a helpful and concise assistant."""
messages = []
# Chat message component, polling if message is still being generated
def ChatMessage(msg_idx):
msg = messages[msg_idx]
text = "..." if msg['content'] == "" else msg['content']
bubble_class = "chat-bubble-primary" if msg['role']=='user' else 'chat-bubble-secondary'
chat_class = "chat-end" if msg['role']=='user' else 'chat-start'
generating = 'generating' in messages[msg_idx] and messages[msg_idx]['generating']
stream_args = {"hx_trigger":"every 0.1s", "hx_swap":"outerHTML", "hx_get":f"/chat_message/{msg_idx}"}
return Div(Div(msg['role'], cls="chat-header"),
Div(text, cls=f"chat-bubble {bubble_class}"),
cls=f"chat {chat_class}", id=f"chat-message-{msg_idx}",
**stream_args if generating else {})
# Route that gets polled while streaming
@app.get("/chat_message/{msg_idx}")
def get_chat_message(msg_idx:int):
if msg_idx >= len(messages): return ""
return ChatMessage(msg_idx)
# The input field for the user message. Also used to clear the
# input field after sending a message via an OOB swap
def ChatInput():
return Input(type="text", name='msg', id='msg-input',
placeholder="Type a message",
cls="input input-bordered w-full", hx_swap_oob='true')
# The main screen
@app.route("/")
def get():
page = Body(H1('Chatbot Demo'),
Div(*[ChatMessage(i) for i in range(len(messages))],
id="chatlist", cls="chat-box h-[73vh] overflow-y-auto"),
Form(Group(ChatInput(), Button("Send", cls="btn btn-primary")),
hx_post="/", hx_target="#chatlist", hx_swap="beforeend",
cls="flex space-x-2 mt-2",
), cls="p-4 max-w-lg mx-auto")
return Title('Chatbot Demo'), page
# Run the chat model in a separate thread
@threaded
def get_response(r, idx):
for chunk in r: messages[idx]["content"] += chunk
messages[idx]["generating"] = False
# Handle the form submission
@app.post("/")
def post(msg:str):
idx = len(messages)
messages.append({"role":"user", "content":msg.rstrip()})
r = cli(messages, sp=sp, stream=True) # Send message to chat model (with streaming)
messages.append({"role":"assistant", "generating":True, "content":""}) # Response initially blank
get_response(r, idx+1) # Start a new thread to fill in content
return (ChatMessage(idx), # The user's message
ChatMessage(idx+1), # The chatbot's response
ChatInput()) # And clear the input field via an OOB swap
if __name__ == '__main__': uvicorn.run("polling:app", host='0.0.0.0', port=8000, reload=True)