diff --git a/api/chat.py b/api/chat.py index 5300818..8b1c8fe 100644 --- a/api/chat.py +++ b/api/chat.py @@ -84,7 +84,7 @@ def normal_completion(prompt: List[Dict[str, str]]) -> str: print(e) return "I'm sorry, I failed to process your query. Please try again. If the problem persists, please contact the administrator." -def talk_to_robot(query: str, history: List[Dict[str, str]] = [], k: int = 10) -> str: +def talk_to_robot(query: str, history: List[Dict[str, str]] = [], k: int = 10): # 1. Find the most relevant blocks from the Alignment Research Dataset top_k_blocks: List[Block] = get_top_k_blocks(query, k) @@ -93,4 +93,4 @@ def talk_to_robot(query: str, history: List[Dict[str, str]] = [], k: int = 10) - prompt: List[Dict[str, str]] = construct_prompt(query, history, top_k_blocks) # 3. Answer the user query - return normal_completion(prompt) + return (normal_completion(prompt), top_k_blocks) diff --git a/api/main.py b/api/main.py index 3cd0d12..7e4884b 100644 --- a/api/main.py +++ b/api/main.py @@ -35,10 +35,8 @@ def semantic(): @cross_origin() def chat(): query = request.json['query'] - return talk_to_robot(query) - - - + response, context = talk_to_robot(query) + return jsonify({'response': response, 'citations': [{'title': block.title, 'author': block.author, 'date': block.date, 'url': block.url} for block in context]}) # ------------------------------------------------------------------------------ diff --git a/web/src/pages/index.tsx b/web/src/pages/index.tsx index d2a9280..f564a6d 100644 --- a/web/src/pages/index.tsx +++ b/web/src/pages/index.tsx @@ -8,21 +8,76 @@ import { useState } from "react"; import Header from "../header"; import SearchBox from "../searchbox"; -type Entry = { - role: "user" | "assistant"; + +type UserEntry = { + role: "user"; content: string; +} + +type AssistantEntry = { + role: "assistant"; + content: string; + citations: Citation[]; +} + +type Entry = UserEntry | AssistantEntry; + +type Citation = { + title: string; + author: string; + date: string; + url: string; +} + +// const Colours = ["blue", "cyan", "teal", "green", "amber"].map(colour => `bg-${colour}-100 border-${colour}-300 text-${colour}-800`); +// this would be nice, but Tailwind needs te actual string of the class to be in +// the source file for it to be included in the build + +const Colours = [ + // "bg-teal-100 border-teal-300 text-teal-800", + "bg-red-100 border-red-300 text-red-800", + "bg-amber-100 border-amber-300 text-amber-800", + "bg-orange-100 border-orange-300 text-orange-800", + "bg-lime-100 border-lime-300 text-lime-800", + "bg-green-100 border-green-300 text-green-800", + "bg-cyan-100 border-cyan-300 text-cyan-800", + "bg-blue-100 border-blue-300 text-blue-800", +]; + +const ShowCitation: React.FC<{citation: Citation, i: number}> = ({citation, i}) => { + return ( + + [{i + 1}] +

+ {citation.title + " - " + citation.author + " - " + citation.date} +

+
+ ); }; const ShowEntry: React.FC<{entry: Entry}> = ({entry}) => { + + // user message if (entry.role === "user") { return (

{entry.content}

); } + // system reply return (
{ // split into paragraphs - entry.content.split("\n").map((paragraph, i) => (

{paragraph}

)) + entry.content.split("\n").map(paragraph => (

{paragraph}

)) } +
); }; @@ -58,11 +113,9 @@ const Home: NextPage = () => { return; } - const response = res.body!.getReader().read().then(({value}) => { - return new TextDecoder("utf-8").decode(value); - }); + const data = await res.json(); - setEntries([...new_entries, {role: "assistant", content: await response}]); + setEntries([...new_entries, {role: "assistant", content: await data.response, citations: await data.citations}]); setLoading(false);