341 lines
12 KiB
TypeScript

import { useState, useRef } from 'react';
import Divider from '@mui/material/Divider';
import Accordion from '@mui/material/Accordion';
import AccordionSummary from '@mui/material/AccordionSummary';
import AccordionDetails from '@mui/material/AccordionDetails';
import Card from '@mui/material/Card';
import Table from '@mui/material/Table';
import TableBody from '@mui/material/TableBody';
import TableCell from '@mui/material/TableCell';
import TableContainer from '@mui/material/TableContainer';
import TableHead from '@mui/material/TableHead';
import TableRow from '@mui/material/TableRow';
import Box from '@mui/material/Box';
import Button from '@mui/material/Button';
import CardContent from '@mui/material/CardContent';
import CardActions from '@mui/material/CardActions';
import Collapse from '@mui/material/Collapse';
import ExpandMoreIcon from '@mui/icons-material/ExpandMore';
import { ExpandMore } from './ExpandMore';
import { SxProps, Theme } from '@mui/material';
import JsonView from '@uiw/react-json-view';
import { ChatBubble } from './ChatBubble';
import { StyledMarkdown } from '../NewApp/Components/StyledMarkdown';
import { VectorVisualizer } from './VectorVisualizer';
import { SetSnackType } from './Snack';
import { CopyBubble } from './CopyBubble';
import { Scrollable } from './Scrollable';
import { BackstoryElementProps } from './BackstoryTab';
type MessageRoles =
'assistant' |
'content' |
'error' |
'fact-check' |
'info' |
'job-description' |
'job-requirements' |
'processing' |
'qualifications' |
'resume' |
'status' |
'streaming' |
'system' |
'thinking' |
'user';
type BackstoryMessage = {
// Only two required fields
role: MessageRoles,
content: string,
// Rest are optional
prompt?: string;
preamble?: {};
status?: string;
remaining_time?: number;
full_content?: string;
response?: string; // Set when status === 'done', 'partial', or 'error'
chunk?: string; // Used when status === 'streaming'
timestamp?: number;
disableCopy?: boolean,
user?: string,
title?: string,
origin?: string,
display?: string, /* Messages generated on the server for filler should not be shown */
id?: string,
isProcessing?: boolean,
actions?: string[],
metadata?: MessageMetaData,
expanded?: boolean,
expandable?: boolean,
};
interface MessageMetaData {
query?: {
query_embedding: number[];
vector_embedding: number[];
},
origin: string,
rag: any[],
tools?: {
tool_calls: any[],
},
eval_count: number,
eval_duration: number,
prompt_eval_count: number,
prompt_eval_duration: number,
connectionBase: string,
setSnack: SetSnackType,
}
type MessageList = BackstoryMessage[];
interface MessageProps extends BackstoryElementProps {
sx?: SxProps<Theme>,
message: BackstoryMessage,
expanded?: boolean,
onExpand?: (open: boolean) => void,
className?: string,
};
interface MessageMetaProps {
metadata: MessageMetaData,
messageProps: MessageProps
};
const MessageMeta = (props: MessageMetaProps) => {
const {
/* MessageData */
rag,
tools,
eval_count,
eval_duration,
prompt_eval_count,
prompt_eval_duration,
} = props.metadata || {};
const message: any = props.messageProps.message;
let llm_submission: string = "<|system|>\n"
llm_submission += message.system_prompt + "\n\n"
llm_submission += message.context_prompt
return (<>
{
prompt_eval_duration !== 0 && eval_duration !== 0 && <>
<TableContainer component={Card} className="PromptStats" sx={{ mb: 1 }}>
<Table aria-label="prompt stats" size="small">
<TableHead>
<TableRow>
<TableCell></TableCell>
<TableCell align="right" >Tokens</TableCell>
<TableCell align="right">Time (s)</TableCell>
<TableCell align="right">TPS</TableCell>
</TableRow>
</TableHead>
<TableBody>
<TableRow key="prompt" sx={{ '&:last-child td, &:last-child th': { border: 0 } }}>
<TableCell component="th" scope="row">Prompt</TableCell>
<TableCell align="right">{prompt_eval_count}</TableCell>
<TableCell align="right">{Math.round(prompt_eval_duration / 10 ** 7) / 100}</TableCell>
<TableCell align="right">{Math.round(prompt_eval_count * 10 ** 9 / prompt_eval_duration)}</TableCell>
</TableRow>
<TableRow key="response" sx={{ '&:last-child td, &:last-child th': { border: 0 } }}>
<TableCell component="th" scope="row">Response</TableCell>
<TableCell align="right">{eval_count}</TableCell>
<TableCell align="right">{Math.round(eval_duration / 10 ** 7) / 100}</TableCell>
<TableCell align="right">{Math.round(eval_count * 10 ** 9 / eval_duration)}</TableCell>
</TableRow>
<TableRow key="total" sx={{ '&:last-child td, &:last-child th': { border: 0 } }}>
<TableCell component="th" scope="row">Total</TableCell>
<TableCell align="right">{prompt_eval_count + eval_count}</TableCell>
<TableCell align="right">{Math.round((prompt_eval_duration + eval_duration) / 10 ** 7) / 100}</TableCell>
<TableCell align="right">{Math.round((prompt_eval_count + eval_count) * 10 ** 9 / (prompt_eval_duration + eval_duration))}</TableCell>
</TableRow>
</TableBody>
</Table>
</TableContainer>
</>
}
{
tools && tools.tool_calls && tools.tool_calls.length !== 0 &&
<Accordion sx={{ boxSizing: "border-box" }}>
<AccordionSummary expandIcon={<ExpandMoreIcon />}>
<Box sx={{ fontSize: "0.8rem" }}>
Tools queried
</Box>
</AccordionSummary>
<AccordionDetails>
{
tools.tool_calls.map((tool: any, index: number) =>
<Box key={index} sx={{ m: 0, p: 1, pt: 0, display: "flex", flexDirection: "column", border: "1px solid #e0e0e0" }}>
{index !== 0 && <Divider />}
<Box sx={{ fontSize: "0.75rem", display: "flex", flexDirection: "column", mt: 1, mb: 1, fontWeight: "bold" }}>
{tool.name}
</Box>
{tool.content !== "null" &&
<JsonView
displayDataTypes={false}
objectSortKeys={true}
collapsed={1} value={JSON.parse(tool.content)} style={{ fontSize: "0.8rem", maxHeight: "20rem", overflow: "auto" }}>
<JsonView.String
render={({ children, ...reset }) => {
if (typeof (children) === "string" && children.match("\n")) {
return <pre {...reset} style={{ display: "flex", border: "none", ...reset.style }}>{children}</pre>
}
}}
/>
</JsonView>
}
{tool.content === "null" && "No response from tool call"}
</Box>)
}
</AccordionDetails>
</Accordion>
}
{
rag.map((collection: any) => (
<Accordion key={collection.name}>
<AccordionSummary expandIcon={<ExpandMoreIcon />}>
<Box sx={{ fontSize: "0.8rem" }}>
Top {collection.ids.length} RAG matches from {collection.size} entries using an embedding vector of {collection.query_embedding.length} dimensions
</Box>
</AccordionSummary>
<AccordionDetails>
<VectorVisualizer inline
{...props.messageProps} {...props.metadata}
rag={collection} />
{/* { ...rag, query: message.prompt }} /> */}
</AccordionDetails>
</Accordion>
))
}
<Accordion>
<AccordionSummary expandIcon={<ExpandMoreIcon />}>
<Box sx={{ fontSize: "0.8rem" }}>
Full Response Details
</Box>
</AccordionSummary>
<AccordionDetails>
<Box sx={{ pb: 1 }}>Copy LLM submission: <CopyBubble content={llm_submission} /></Box>
<JsonView displayDataTypes={false} objectSortKeys={true} collapsed={1} value={message} style={{ fontSize: "0.8rem", maxHeight: "20rem", overflow: "auto" }}>
<JsonView.String
render={({ children, ...reset }) => {
if (typeof (children) === "string" && children.match("\n")) {
return <pre {...reset} style={{ display: "inline", border: "none", ...reset.style }}>{children.trim()}</pre>
}
}}
/>
</JsonView>
</AccordionDetails>
</Accordion>
</>);
};
const Message = (props: MessageProps) => {
const { message, submitQuery, sx, className, onExpand, setSnack, sessionId, expanded } = props;
const [metaExpanded, setMetaExpanded] = useState<boolean>(false);
const textFieldRef = useRef(null);
const backstoryProps = {
submitQuery,
sessionId,
setSnack
};
const handleMetaExpandClick = () => {
setMetaExpanded(!metaExpanded);
};
if (message === undefined) {
return (<></>);
}
if (message.content === undefined) {
console.info("Message content is undefined");
return (<></>);
}
const formattedContent = message.content.trim();
if (formattedContent === "") {
return (<></>);
}
return (
<ChatBubble
className={`${className || ""} Message Message-${message.role}`}
{...message}
expanded={expanded}
onExpand={onExpand}
sx={{
display: "flex",
flexDirection: "column",
pb: message.metadata ? 0 : "8px",
m: 0,
mt: 1,
marginBottom: "0px !important", // Remove whitespace from expanded Accordion
// overflowX: "auto"
...sx,
}}>
<CardContent ref={textFieldRef} sx={{ position: "relative", display: "flex", flexDirection: "column", overflowX: "auto", m: 0, p: 0, paddingBottom: '0px !important' }}>
<Scrollable
className="MessageContent"
autoscroll
fallbackThreshold={0.5}
sx={{
p: 0,
m: 0,
// maxHeight: (message.role === "streaming") ? "20rem" : "unset",
display: "flex",
flexGrow: 1,
overflow: "auto", /* Handles scrolling for the div */
}}
>
<StyledMarkdown streaming={message.role === "streaming"} content={formattedContent} {...backstoryProps} />
</Scrollable>
</CardContent>
<CardActions disableSpacing sx={{ display: "flex", flexDirection: "row", justifyContent: "space-between", alignItems: "center", width: "100%", p: 0, m: 0 }}>
{(message.disableCopy === undefined || message.disableCopy === false) && <CopyBubble content={message.content} />}
{message.metadata && (
<Box sx={{ display: "flex", alignItems: "center", gap: 1 }}>
<Button variant="text" onClick={handleMetaExpandClick} sx={{ color: "darkgrey", p: 0 }}>
LLM information for this query
</Button>
<ExpandMore
expand={metaExpanded}
onClick={handleMetaExpandClick}
aria-expanded={message.expanded}
aria-label="show more"
>
<ExpandMoreIcon />
</ExpandMore>
</Box>
)}
</CardActions>
{message.metadata && <>
<Collapse in={metaExpanded} timeout="auto" unmountOnExit>
<CardContent>
<MessageMeta messageProps={props} metadata={message.metadata} />
</CardContent>
</Collapse>
</>}
</ChatBubble>
);
};
export type {
MessageProps,
MessageList,
BackstoryMessage,
MessageMetaData,
MessageRoles,
};
export {
Message,
MessageMeta,
};