feat(lightspeed): improved chat experience with added context from pr…

…evious responses (#2064)
janus-idp · Aug 20, 2024 · 177ece0 · 177ece0
1 parent 8f8ad35
commit 177ece0
Show file tree

Hide file tree

Showing 4 changed files with 659 additions and 66 deletions.
diff --git a/plugins/lightspeed/package.json b/plugins/lightspeed/package.json
@@ -40,7 +40,9 @@
     "@material-ui/icons": "^4.11.3",
     "@material-ui/lab": "^4.0.0-alpha.61",
     "@mui/icons-material": "^5.15.18",
+    "@patternfly/react-core": "6.0.0-alpha.101",
     "openai": "^4.52.6",
+    "react-markdown": "^9.0.1",
     "react-use": "^17.2.4"
   },
   "peerDependencies": {

diff --git a/plugins/lightspeed/src/components/LightspeedPage.tsx b/plugins/lightspeed/src/components/LightspeedPage.tsx
@@ -41,28 +41,46 @@ export const LightspeedPage = () => {
   const lightspeedApi = useApi(lightspeedApiRef);
 
   const [, setChunkIndex] = React.useState(0);
-  const [prompts, setPrompts] = React.useState<string[]>([]);
-  const [completions, setCompletions] = React.useState<{
-    [key: string]: string;
+  const [convoIndex, setConvoIndex] = React.useState<number[]>([0]);
+  const [messages, setMessages] = React.useState<{
+    [key: string]: { user: string; assistant: string };
   }>({});
 
+  const convoCounter = convoIndex.length - 1;
+
+  const context_template = Object.values(messages)
+    .map(convo => `USER: ${convo.user}\nASSISTANT: ${convo.assistant}`)
+    .join('\n\n');
+
   const handleInputPrompt = React.useCallback(
     async (prompt: string) => {
-      setPrompts(p => [...p, prompt]);
       setChunkIndex(0);
+      setMessages(m => {
+        m[convoCounter] = {
+          user: prompt,
+          assistant: '',
+        };
+        return m;
+      });
 
-      const result = await lightspeedApi.createChatCompletions(prompt);
+      const latestQuestion = `Respond to the users current message\n\nCONTEXT: ${context_template}\n\nCURRENT QUESTION: ${prompt}`;
 
+      const result = await lightspeedApi.createChatCompletions(
+        `${latestQuestion}`,
+      );
       for await (const chunk of result) {
         setChunkIndex(index => index + 1);
-        setCompletions(c => {
-          c[prompt] =
-            `${c[prompt] || ''}${chunk.choices[0]?.delta?.content || ''}`;
-          return c;
+        setMessages(m => {
+          m[convoCounter] = {
+            ...m[convoCounter],
+            assistant: `${m[`${convoCounter}`]?.assistant || ''}${chunk.choices[0]?.delta?.content || ''}`,
+          };
+          return m;
         });
       }
+      setConvoIndex(conIndex => [...conIndex, convoCounter + 1]);
     },
-    [lightspeedApi],
+    [context_template, convoCounter, lightspeedApi],
   );
 
   return (
@@ -77,11 +95,13 @@ export const LightspeedPage = () => {
       <Content className={classes.container}>
         <Paper className={classes.paper} elevation={2}>
           <Paper id="style-1" className={classes.messagesBody}>
-            {prompts.map(prompt => (
+            {convoIndex.map(id => (
               <>
-                {prompt && <UserMessage message={prompt} />}
-                {completions[prompt] && (
-                  <SystemMessage message={completions[prompt]} />
+                {messages[id]?.user && (
+                  <UserMessage message={messages[id].user} />
+                )}
+                {messages[id]?.assistant && (
+                  <SystemMessage message={messages[id].assistant} />
                 )}
               </>
             ))}

diff --git a/plugins/lightspeed/src/components/Message.tsx b/plugins/lightspeed/src/components/Message.tsx
@@ -1,8 +1,10 @@
 import React from 'react';
+import Markdown from 'react-markdown';
 
 import { createStyles, makeStyles } from '@material-ui/core/styles';
+import { CodeBlock, CodeBlockCode } from '@patternfly/react-core';
 
-const useStyles = makeStyles(() =>
+const useStyles = makeStyles(theme =>
   createStyles({
     messageRow: {
       display: 'flex',
@@ -16,66 +18,19 @@ const useStyles = makeStyles(() =>
       marginLeft: '20px',
       marginBottom: '10px',
       padding: '10px',
-      backgroundColor: '#476373',
       width: '60%',
       textAlign: 'left',
-      border: '1px solid #97C6E3',
-      borderRadius: '10px',
-      '&:after': {
-        content: "''",
-        position: 'absolute',
-        width: '0',
-        height: '0',
-        borderTop: '15px solid #476373',
-        borderLeft: '15px solid transparent',
-        borderRight: '15px solid transparent',
-        top: '0',
-        left: '-15px',
-      },
-      '&:before': {
-        content: "''",
-        position: 'absolute',
-        width: '0',
-        height: '0',
-        borderTop: '17px solid #97C6E3',
-        borderLeft: '16px solid transparent',
-        borderRight: '16px solid transparent',
-        top: '-1px',
-        left: '-17px',
-      },
     },
     messageOrange: {
       position: 'relative',
       marginRight: '20px',
       marginBottom: '15px',
       padding: '15px',
-      backgroundColor: '#665f3b',
+      backgroundColor: theme.palette.background.default,
       width: '60%',
       textAlign: 'left',
-      border: '1px solid #dfd087',
+      border: `1px solid ${theme.palette.grey[400]}`,
       borderRadius: '10px',
-      '&:after': {
-        content: "''",
-        position: 'absolute',
-        width: '0',
-        height: '0',
-        borderTop: '15px solid #665f3b',
-        borderLeft: '15px solid transparent',
-        borderRight: '15px solid transparent',
-        top: '0',
-        right: '-15px',
-      },
-      '&:before': {
-        content: "''",
-        position: 'absolute',
-        width: '0',
-        height: '0',
-        borderTop: '17px solid #dfd087',
-        borderLeft: '16px solid transparent',
-        borderRight: '16px solid transparent',
-        top: '-1px',
-        right: '-17px',
-      },
     },
 
     messageContent: {
@@ -94,6 +49,33 @@ const useStyles = makeStyles(() =>
   }),
 );
 
+const useCodeStyles = makeStyles(theme => ({
+  codeBlock: { backgroundColor: theme.palette.grey[900] },
+  codeBlockCode: {
+    maxHeight: '20rem',
+    overflow: 'auto',
+    whiteSpace: 'pre',
+    color: '#ffffff',
+    padding: theme.spacing(2),
+  },
+}));
+
+const Code = ({ children }: { children?: React.ReactNode }) => {
+  const classes = useCodeStyles();
+
+  if (!String(children).includes('\n')) {
+    return <code>{children}</code>;
+  }
+
+  return (
+    <CodeBlock className={classes.codeBlock}>
+      <CodeBlockCode className={classes.codeBlockCode}>
+        {children}
+      </CodeBlockCode>
+    </CodeBlock>
+  );
+};
+
 type SystemMessageProps = {
   message: string;
   timestamp?: string;
@@ -108,7 +90,7 @@ export const SystemMessage: React.FC<SystemMessageProps> = ({
     <>
       <div className={classes.messageRow}>
         <div className={classes.messageBlue}>
-          <p className={classes.messageContent}>{message}</p>
+          <Markdown components={{ code: Code }}>{message}</Markdown>
           {timestamp && (
             <div className={classes.messageTimeStampRight}>{timestamp}</div>
           )}