wassname/stampy-chat
1
0
Fork 0
mirror of https://github.com/wassname/stampy-chat.git synced 2026-09-11 12:50:34 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1 from FraserLee/streaming

Browse Source 
Streaming
This commit is contained in:
FraserLee
2023-04-16 21:14:34 -04:00
committed by GitHub
parent d924e18a5c dee78258d6
commit 887af41c77
5 changed files with 260 additions and 111 deletions
Download Patch File Download Diff File Expand all files Collapse all files
api/Procfile
+1 -1
View File
@@ -1 +1 @@
web: gunicorn main:app
web: gunicorn main:app --worker-class eventlet --threads 4
api/chat.py
+21 -9
View File
@@ -24,7 +24,7 @@ CONTEXT_FRACTION = 0.5 # the (approximate) fraction of num_tokens to use for co
ENCODER = tiktoken.get_encoding("cl100k_base")
DEBUG_PRINT = False
DEBUG_PRINT = True
# --------------------------------- prompt code --------------------------------
@@ -116,32 +116,44 @@ def construct_prompt(query: str, history: List[Dict[str, str]], context: List[Bl
return prompt
# ------------------------------- completion code -------------------------------
import time
import json
# returns either (True, reply string, top_k_blocks)) or (False, error message string, None)
def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = STANDARD_K):
try:
# 1. Find the most relevant blocks from the Alignment Research Dataset
yield json.dumps({"state": "loading", "phase": "semantic"})
top_k_blocks = get_top_k_blocks(index, query, k)
yield json.dumps({"state": "loading", "phase": "semantic", 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in top_k_blocks]})
# 2. Generate a prompt
yield json.dumps({"state": "loading", "phase": "prompt"})
prompt = construct_prompt(query, history, top_k_blocks)
# 3. Count number of tokens left for completion (-50 for a buffer)
max_tokens_completion = NUM_TOKENS - sum([len(ENCODER.encode(message["content"]) + ENCODER.encode(message["role"])) for message in prompt]) - 50
# 4. Answer the user query
yield json.dumps({"state": "loading", "phase": "llm"})
t1 = time.time()
response = openai.ChatCompletion.create(
response = ''
for chunk in openai.ChatCompletion.create(
model=COMPLETIONS_MODEL,
messages=prompt,
max_tokens=max_tokens_completion
)["choices"][0]["message"]["content"]
max_tokens=max_tokens_completion,
stream=True
):
res = chunk["choices"][0]["delta"]
if res is not None and res.get("content") is not None:
response += res["content"]
yield json.dumps({"state": "streaming", "content": res["content"]})
t2 = time.time()
print("Time to get response: ", t2 - t1)
if DEBUG_PRINT:
print('\n' * 10)
@@ -155,9 +167,9 @@ def talk_to_robot(index, query: str, history: List[Dict[str, str]], k: int = STA
print(" ------------------------------ response: -----------------------------")
print(response)
return (True, response, top_k_blocks)
yield json.dumps({"state": "done"})
except Exception as e:
print(e)
return (False, "Error: " + str(e), None)
yield json.dumps({"state": "error", "error": str(e)})
api/main.py
+10 -9
View File
@@ -1,4 +1,4 @@
from flask import Flask, jsonify, request
from flask import Flask, jsonify, request, Response
from flask_cors import CORS, cross_origin
from get_blocks import get_top_k_blocks
from chat import talk_to_robot
@@ -34,6 +34,12 @@ app = Flask(__name__)
cors = CORS(app)
app.config['CORS_HEADERS'] = 'Content-Type'
# ---------------------------------- sse stuff ---------------------------------
def stream(src):
yield from ('data: ' + '\ndata: '.join(message.splitlines()) + '\n\n' for message in src)
yield 'event: close\n\n'
# ------------------------------- semantic search ------------------------------
@@ -50,20 +56,15 @@ def semantic():
@app.route('/chat', methods=['POST'])
@cross_origin()
def chat():
query = request.json['query']
history = request.json['history']
is_valid, response, context = talk_to_robot(index, query, history)
if is_valid:
return jsonify({'response': response, 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in context]})
else:
return jsonify({'error': response})
return Response(stream(talk_to_robot(index, query, history)), mimetype='text/event-stream')
# ------------------------------------------------------------------------------
if __name__ == '__main__':
app.run(debug=True, port=3000)
app.run(debug=True, port=3000)
web/src/pages/index.tsx
+227 -91
View File
@@ -23,7 +23,8 @@ type UserEntry = {
type AssistantEntry = {
role: "assistant";
content: string;
citations: Map<number, Citation>;
citations: Citation[];
base_count: number; // the number to start counting citations at
}
type ErrorMessage = {
@@ -82,35 +83,109 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
if (entry.role === "user") {
return ( <p className="border border-gray-300 px-1 text-right"> {entry.content} </p>);
}
// error message
if (entry.role === "error") {
return ( <p className="border bg-red-100 border-red-500 text-red-800 px-1"> {entry.content} </p>);
}
// todo: memoize this if too slow.
const ProcessText: (text: string, base_count: number) => [string, Map<string, number>] = (text, base_count) => {
// ---------------------- normalize citation form ----------------------
// transform all things that look like [a, b, c] into [a][b][c]
let response = text.replace(
/\[((?:[a-z]+,\s*)*[a-z]+)\]/g, // identify groups of this form
(block: string) => block.split(',')
.map((x) => x.trim())
.join("][")
)
// transform all things that look like [(a), (b), (c)] into [(a)][(b)][(c)]
response = response.replace(
/\[((?:\([a-z]+\),\s*)*\([a-z]+\))\]/g, // identify groups of this form
(block: string) => block.split(',')
.map((x) => x.trim())
.join("][")
)
// transform all things that look like [(a)] into [a]
response = response.replace(
/\[\(([a-z]+)\)\]/g,
(_match: string, x: string) => `[${x}]`
)
// transform all things that look like [ a ] into [a]
response = response.replace(
/\[\s*([a-z]+)\s*\]/g,
(_match: string, x: string) => `[${x}]`
)
// -------------- map citations from strings into numbers --------------
// figure out what citations are in the response, and map them appropriately
const cite_map = new Map<string, number>();
let cite_count = 0;
// scan a regex for [x] over the response. If x isn't in the map, add it.
const regex = /\[([a-z]+)\]/g;
let match;
let response_copy = ""
while ((match = regex.exec(response)) !== null) {
if (!cite_map.has(match[1]!)) {
cite_map.set(match[1]!, base_count + cite_count++);
}
// replace [x] with [i]
response_copy += response.slice(response_copy.length, match.index) + `[${cite_map.get(match[1]!)! + 1}]`;
}
response = response_copy + response.slice(response_copy.length);
return [response, cite_map]
}
const ShowAssistantEntry: React.FC<{entry: AssistantEntry}> = ({entry}) => {
const in_text_citation_regex = /\[([0-9]+)\]/g;
// system reply
let [response, cite_map] = ProcessText(entry.content, entry.base_count);
// ----------------- create the ordered citation array -----------------
const citations = new Map<number, Citation>();
cite_map.forEach((value, key) => {
const index = key.charCodeAt(0) - 'a'.charCodeAt(0);
if (index >= entry.citations.length) {
console.log("invalid citation index: " + index);
} else {
citations.set(value, entry.citations[index]!);
}
});
return (
<div className="mt-3 mb-8">
{ // split into paragraphs
entry.content.split("\n").map(paragraph => ( <p> {
response.split("\n").map(paragraph => ( <p> {
paragraph.split(in_text_citation_regex).map((text, i) => {
if (i % 2 === 0) {
return text.trim();
}
i = parseInt(text) - 1;
if (!entry.citations.has(i)) return `[${text}]`;
const citation = entry.citations.get(i)!;
if (!citations.has(i)) return `[${text}]`;
const citation = citations.get(i)!;
return (
<ShowInTextCitation citation={citation} i={i} />
);
})
} </p>))
}
<ul>
<ul className="mt-5">
{ // show citations
Array.from(entry.citations.entries()).map(([i, citation]) => (
Array.from(citations.entries()).map(([i, citation]) => (
<li key={i}>
<ShowCitation citation={citation} i={i} />
</li>
@@ -121,29 +196,55 @@ const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => {
);
};
type State = {
state: "idle";
} | {
state: "loading";
phase: "semantic" | "prompt" | "llm";
citations: Citation[];
} | {
state: "streaming";
response: AssistantEntry;
};
const Home: NextPage = () => {
const [ entries, setEntries ] = useState<Entry[]>([]);
const [ runningIndex, setRunningIndex ] = useState(0);
const [ loadState, setLoadState ] = useState<State>({state: "idle"});
const search = async (
query: string,
setQuery: (query: string) => void,
setLoading: (loading: boolean) => void
) => {
// clear the query box, append to entries
const old_entries = entries;
const new_entries: Entry[] = [...old_entries, {role: "user", content: query}];
setEntries(new_entries);
setQuery("");
setLoading(true);
// do SSE on a POST request.
const res = await fetch(API_URL + "/chat", {
method: "POST",
headers: { "Content-Type": "application/json", "Allow-Control-Allow-Origin": "*" },
body: JSON.stringify({query: query, history:
cache: "no-cache",
keepalive: true,
headers: {
"Content-Type": "application/json",
"Accept": "text/event-stream",
"Allow-Control-Allow-Origin": "*"
},
body: JSON.stringify({query: query, history:
old_entries.filter((entry) => entry.role !== "error")
.map((entry) => {
return {
@@ -151,98 +252,105 @@ const Home: NextPage = () => {
"content" : entry.content.trim(),
}
})
})
})
}),
});
if (!res.ok) {
setLoading(false);
console.log("load failure: " + res.status);
setLoadState({state: "idle"});
setEntries([...new_entries, {role: "error", content: "POST Error: " + res.status}]);
return;
}
const data = await res.json();
// read back the SSE stream
// -------------------------- error checking ---------------------------
const reader = res.body!.getReader();
var message = "";
read: while (true) {
if (data.error) {
setEntries([...new_entries, {role: "error", content: data.error}]);
setLoading(false);
return;
}
const {done, value} = await reader.read();
// ---------------------- normalize citation form ----------------------
if (done) break;
const chunk = new TextDecoder("utf-8").decode(value);
if (chunk.startsWith("event: close\n")) break;
// transform all things that look like [a, b, c] into [a][b][c]
let response = data.response.replace(
// note: this form isn't even remotely close to optimal in terms of network usage.
/\[((?:[a-z]+,\s*)*[a-z]+)\]/g, // identify groups of this form
for (const line of chunk.split('\n')) {
(block: string) => block.split(',')
.map((x) => x.trim())
.join("][")
)
// Most times, it seems that a single read() call will be one SSE "message",
// but I'll do the proper aggregation spec thing in case that's not always true.
// transform all things that look like [(a), (b), (c)] into [(a)][(b)][(c)]
response = response.replace(
/\[((?:\([a-z]+\),\s*)*\([a-z]+\))\]/g, // identify groups of this form
if (line.startsWith("data: ")) message += line.slice(6);
if (line === "") {
if (message !== "") {
const data = JSON.parse(message);
(block: string) => block.split(',')
.map((x) => x.trim())
.join("][")
)
switch (data.state) {
// transform all things that look like [(a)] into [a]
response = response.replace(
/\[\(([a-z]+)\)\]/g,
(_match: string, x: string) => `[${x}]`
)
case "loading":
// transform all things that look like [ a ] into [a]
response = response.replace(
/\[\s*([a-z]+)\s*\]/g,
(_match: string, x: string) => `[${x}]`
)
// display loading phases, once citations are available toss them
// into the loading state.
// -------------- map citations from strings into numbers --------------
setLoadState((s) => {
var citations = s.state === "loading" ? s.citations : [];
if (data.citations !== undefined) {
citations = data.citations;
}
return {state: "loading", phase: data.phase, citations: citations};
});
// figure out what citations are in the response, and map them appropriately
const cite_map = new Map<string, number>();
let cite_count = runningIndex;
break;
// scan a regex for [x] over the response. If x isn't in the map, add it.
const regex = /\[([a-z]+)\]/g;
let match;
let response_copy = ""
while ((match = regex.exec(response)) !== null) {
if (!cite_map.has(match[1]!)) {
cite_map.set(match[1]!, cite_count++);
case "streaming":
// incrementally build up the response
setLoadState((s) => {
const response = s.state === "streaming" ? s.response :
{role: "assistant",
content: "",
citations: s.state === "loading" ? s.citations : [],
base_count: runningIndex
};
return {state: "streaming", response: {
role: "assistant",
content: response.content + data.content,
citations: response.citations,
base_count: response.base_count
}};
});
break;
case "done":
// append the response to the entries, reset to normal
setLoadState((s) => {
if (s.state === "streaming") {
setEntries([...new_entries, s.response]);
setRunningIndex((i) => (i + ProcessText(s.response.content, 0)[1].size));
}
return {state: "idle"};
});
break read;
case "error":
setEntries([...new_entries, {role: "error", content: data.error}]);
break read;
}
}
message = "";
}
}
// replace [x] with [i]
response_copy += response.slice(response_copy.length, match.index) + `[${cite_map.get(match[1]!)! + 1}]`;
}
setRunningIndex(cite_count);
response = response_copy + response.slice(response_copy.length);
// ----------------- create the ordered citation array -----------------
const citations = new Map<number, Citation>();
cite_map.forEach((value, key) => {
const index = key.charCodeAt(0) - 'a'.charCodeAt(0);
if (index >= data.citations.length) {
console.log("invalid citation index: " + index);
} else {
citations.set(value, data.citations[index]);
}
});
setEntries([...new_entries, {role: "assistant",
content: response,
citations: citations}]);
setLoading(false);
setLoadState({state: "idle"});
};
@@ -254,13 +362,41 @@ const Home: NextPage = () => {
<main>
<Header page="index" />
<ul>
{entries.map((entry, i) => (
<li key={i}>
<ShowEntry entry={entry} />
</li>
))}
{entries.map((entry, i) => {
if (entry.role === "user") {
return <li key={i}>
<p className="border border-gray-300 px-1 text-right"> {entry.content} </p>
</li>
}
if (entry.role === "error") {
return <li key={i}>
<p className="border bg-red-100 border-red-500 text-red-800 px-1"> {entry.content} </p>
</li>
}
if (entry.role === "assistant") {
return <li key={i}>
<ShowAssistantEntry entry={entry}/>
</li>
}
return <></>
})}
<SearchBox search={search} />
{(() => {
if (loadState.state === "loading") {
switch (loadState.phase) {
case "semantic": return <p>Loading: Performing semantic search...</p>;
case "prompt": return <p>Loading: Creating prompt...</p>;
case "llm": return <p>Loading: Waiting for LLM...</p>;
}
} else if (loadState.state === "streaming") {
return <ShowAssistantEntry entry={loadState.response}/>;
}
return <></>;
})()}
</ul>
<SearchBox search={search} />
</main>
</>
);
web/src/searchbox.tsx
+1 -1
View File
@@ -20,7 +20,7 @@ const SearchBox: React.FC<{search: (
if (!loading) inputRef.current?.focus();
}, [loading]);
if (loading) return <p>loading...</p>;
if (loading) return <></>;
return (<>
<form className="flex mb-2" onSubmit={async (e) => {
e.preventDefault();