463 lines
16 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 Button from '@mui/material/Button';
import CardContent from '@mui/material/CardContent';
import CardActions from '@mui/material/CardActions';
import Collapse from '@mui/material/Collapse';
import { ExpandMore } from './ExpandMore';
import JsonView from '@uiw/react-json-view';
import React from 'react';
import { Box } from '@mui/material';
import { useTheme } from '@mui/material/styles';
import { SxProps, Theme } from '@mui/material';
import ExpandMoreIcon from '@mui/icons-material/ExpandMore';
import LocationSearchingIcon from '@mui/icons-material/LocationSearching';
import { ErrorOutline, InfoOutline, Memory, Psychology, /* Stream, */ } from '@mui/icons-material';
import { StyledMarkdown } from './StyledMarkdown';
import { VectorVisualizer } from './VectorVisualizer';
import { SetSnackType } from './Snack';
import { CopyBubble } from './CopyBubble';
import { Scrollable } from './Scrollable';
import { BackstoryElementProps } from './BackstoryTab';
import { ChatMessage, ChatSession, ChatMessageMetaData, ChromaDBGetResponse, ApiActivityType, ChatMessageUser, ChatMessageError, ChatMessageStatus, ChatSenderType } from 'types/types';
const getStyle = (theme: Theme, type: ApiActivityType | ChatSenderType | "error"): any => {
const defaultRadius = '16px';
const defaultStyle = {
padding: theme.spacing(1, 2),
fontSize: '0.875rem',
alignSelf: 'flex-start',
maxWidth: '100%',
minWidth: '100%',
height: 'fit-content',
'& > *': {
color: 'inherit',
overflow: 'hidden',
m: 0,
},
'& > :last-child': {
mb: 0,
m: 0,
p: 0,
},
};
const styles: any = {
assistant: {
...defaultStyle,
backgroundColor: theme.palette.primary.main,
border: `1px solid ${theme.palette.secondary.main}`,
borderRadius: `${defaultRadius} ${defaultRadius} ${defaultRadius} 0`,
color: theme.palette.primary.contrastText,
},
content: {
...defaultStyle,
backgroundColor: '#F5F2EA',
border: `1px solid ${theme.palette.custom.highlight}`,
borderRadius: 0,
alignSelf: 'center',
color: theme.palette.text.primary,
padding: '8px 8px',
marginBottom: '0px',
boxShadow: '0 2px 4px rgba(0, 0, 0, 0.05)',
fontSize: '0.9rem',
lineHeight: '1.3',
fontFamily: theme.typography.fontFamily,
},
error: {
...defaultStyle,
backgroundColor: '#F8E7E7',
border: `1px solid #D83A3A`,
borderRadius: defaultRadius,
maxWidth: '90%',
minWidth: '90%',
alignSelf: 'center',
color: '#8B2525',
padding: '10px 16px',
boxShadow: '0 1px 3px rgba(216, 58, 58, 0.15)',
},
'fact-check': 'qualifications',
generating: 'status',
'job-description': 'content',
'job-requirements': 'qualifications',
information: {
...defaultStyle,
backgroundColor: '#BFD8D8',
border: `1px solid ${theme.palette.secondary.main}`,
borderRadius: defaultRadius,
color: theme.palette.text.primary,
opacity: 0.95,
},
preparing: 'status',
processing: 'status',
qualifications: {
...defaultStyle,
backgroundColor: theme.palette.primary.light,
border: `1px solid ${theme.palette.secondary.main}`,
borderRadius: `${defaultRadius} ${defaultRadius} ${defaultRadius} 0`,
color: theme.palette.primary.contrastText,
},
resume: 'content',
searching: 'status',
status: {
...defaultStyle,
backgroundColor: 'rgba(74, 122, 125, 0.15)',
border: `1px solid ${theme.palette.secondary.light}`,
borderRadius: '4px',
maxWidth: '75%',
minWidth: '75%',
alignSelf: 'center',
color: theme.palette.secondary.dark,
fontWeight: 500,
fontSize: '0.95rem',
padding: '8px 12px',
opacity: 0.9,
transition: 'opacity 0.3s ease-in-out',
},
streaming: 'response',
system: {
...defaultStyle,
backgroundColor: '#EDEAE0',
border: `1px dashed ${theme.palette.custom.highlight}`,
borderRadius: defaultRadius,
maxWidth: '90%',
minWidth: '90%',
alignSelf: 'center',
color: theme.palette.text.primary,
fontStyle: 'italic',
},
thinking: 'status',
user: {
...defaultStyle,
backgroundColor: theme.palette.background.default,
border: `1px solid ${theme.palette.custom.highlight}`,
borderRadius: `${defaultRadius} ${defaultRadius} 0 ${defaultRadius}`,
alignSelf: 'flex-end',
color: theme.palette.primary.main,
},
};
// Resolve string references in styles
for (const [key, value] of Object.entries(styles)) {
if (typeof value === 'string') {
styles[key] = styles[value];
}
}
if (!(type in styles)) {
console.log(`Style does not exist for: ${type}`);
}
return styles[type];
}
const getIcon = (activityType: ApiActivityType | ChatSenderType | "error"): React.ReactNode | null => {
const icons: any = {
error: <ErrorOutline color="error" />,
generating: <LocationSearchingIcon />,
information: <InfoOutline color="info" />,
preparing: <LocationSearchingIcon />,
processing: <LocationSearchingIcon />,
system: <Memory />,
thinking: <Psychology />,
tooling: <LocationSearchingIcon />,
};
return icons[activityType] || null;
}
interface MessageProps extends BackstoryElementProps {
message: ChatMessageUser | ChatMessage | ChatMessageError | ChatMessageStatus,
title?: string,
chatSession?: ChatSession,
className?: string,
sx?: SxProps<Theme>,
expandable?: boolean,
expanded?: boolean,
onExpand?: (open: boolean) => void,
};
interface MessageMetaProps {
metadata: ChatMessageMetaData,
messageProps: MessageProps
};
const MessageMeta = (props: MessageMetaProps) => {
const {
/* MessageData */
ragResults = [],
tools = null,
evalCount = 0,
evalDuration = 0,
promptEvalCount = 0,
promptEvalDuration = 0,
} = 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 (<>
{
promptEvalDuration !== 0 && evalDuration !== 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">{promptEvalCount}</TableCell>
<TableCell align="right">{Math.round(promptEvalDuration / 10 ** 7) / 100}</TableCell>
<TableCell align="right">{Math.round(promptEvalCount * 10 ** 9 / promptEvalDuration)}</TableCell>
</TableRow>
<TableRow key="response" sx={{ '&:last-child td, &:last-child th': { border: 0 } }}>
<TableCell component="th" scope="row">Response</TableCell>
<TableCell align="right">{evalCount}</TableCell>
<TableCell align="right">{Math.round(evalDuration / 10 ** 7) / 100}</TableCell>
<TableCell align="right">{Math.round(evalCount * 10 ** 9 / evalDuration)}</TableCell>
</TableRow>
<TableRow key="total" sx={{ '&:last-child td, &:last-child th': { border: 0 } }}>
<TableCell component="th" scope="row">Total</TableCell>
<TableCell align="right">{promptEvalCount + evalCount}</TableCell>
<TableCell align="right">{Math.round((promptEvalDuration + evalDuration) / 10 ** 7) / 100}</TableCell>
<TableCell align="right">{Math.round((promptEvalCount + evalCount) * 10 ** 9 / (promptEvalDuration + evalDuration))}</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>
}
{
ragResults.map((collection: ChromaDBGetResponse) => (
<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.queryEmbedding?.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>
</>);
};
interface MessageContainerProps {
type: ApiActivityType | ChatSenderType | "error",
metadataView?: React.ReactNode | null,
messageView?: React.ReactNode | null,
sx?: SxProps<Theme>,
copyContent?: string,
};
const MessageContainer = (props: MessageContainerProps) => {
const { type, sx, messageView, metadataView, copyContent } = props;
const icon = getIcon(type);
return <Box
className={`Message Message-${type}`}
sx={{
display: "flex",
flexDirection: "column",
m: 0,
mt: 1,
marginBottom: "0px !important", // Remove whitespace from expanded Accordion
gap: 1,
...sx,
}}>
<Box sx={{ display: "flex", flexDirection: 'row', alignItems: 'center', gap: 1 }}>
{icon !== null && icon}
{messageView}
</Box>
<Box flex={{ display: "flex", position: "relative", flexDirection: "row", justifyContent: "flex-end", alignItems: "center" }}>
{copyContent && <CopyBubble content={copyContent} sx={{ position: "absolute", top: "11px", left: 0 }} />}
{metadataView}
</Box>
</Box>;
};
const Message = (props: MessageProps) => {
const { message, title, sx, className, chatSession, onExpand, expanded, expandable } = props;
const [metaExpanded, setMetaExpanded] = useState<boolean>(false);
const theme = useTheme();
const type: ApiActivityType | ChatSenderType | "error" = ('activity' in message) ? message.activity : ('error' in message) ? 'error' : (message as ChatMessage).role;
const style: any = getStyle(theme, type);
const handleMetaExpandClick = () => {
setMetaExpanded(!metaExpanded);
};
let content;
if (typeof (message.content) === "string") {
content = message.content.trim();
} else {
console.error(`message content is not a string`);
return (<></>)
}
if (!content) {
return (<></>)
};
const messageView = (
<StyledMarkdown chatSession={chatSession} streaming={message.status === "streaming"} content={content} />
);
let metadataView = (<></>);
const metadata: ChatMessageMetaData | null = ('metadata' in message) ? (message.metadata as ChatMessageMetaData || null) : null;
if (metadata) {
metadataView = (
<Box sx={{ display: "flex", flexDirection: "column", width: "100%" }}>
<Box sx={{ display: "flex", alignItems: "center", gap: 1, flexDirection: "row" }}>
<Box sx={{ display: "flex", flexGrow: 1 }} />
<Button variant="text" onClick={handleMetaExpandClick} sx={{ flexShrink: 1, color: "darkgrey", p: 0 }}>
LLM information for this query
</Button>
<ExpandMore
sx={{ flexShrink: 1 }}
expand={metaExpanded}
onClick={handleMetaExpandClick}
aria-expanded={true /*message.expanded*/}
aria-label="show more">
<ExpandMoreIcon />
</ExpandMore>
</Box>
<Collapse in={metaExpanded} timeout="auto" unmountOnExit>
<CardContent>
<MessageMeta messageProps={props} metadata={metadata} />
</CardContent>
</Collapse>
</Box>);
}
const copyContent = (type === 'assistant') ? message.content : undefined;
if (!expandable) {
/* When not expandable, the styles are applied directly to MessageContainer */
return (<>
{messageView && <MessageContainer copyContent={copyContent} type={type} {...{ messageView, metadataView }} sx={{ ...style, ...sx }} />}
</>);
}
// Determine if Accordion is controlled
const isControlled = typeof expanded === 'boolean' && typeof onExpand === 'function';
return (
<Accordion
expanded={isControlled ? expanded : undefined} // Omit expanded prop for uncontrolled
defaultExpanded={expanded} // Default to collapsed for uncontrolled Accordion
className={className}
onChange={(_event, newExpanded) => { isControlled && onExpand && onExpand(newExpanded) }}
sx={{ ...sx, ...style }}>
<AccordionSummary
expandIcon={<ExpandMoreIcon />}
slotProps={{
content: {
sx: {
display: 'flex',
justifyItems: 'center',
m: 0, p: 0,
fontWeight: 'bold',
fontSize: '1.1rem',
},
},
}}>
{title || ''}
</AccordionSummary>
<AccordionDetails sx={{ mt: 0, mb: 0, p: 0, pl: 2, pr: 2 }}>
<MessageContainer copyContent={copyContent} type={type} {...{ messageView, metadataView }} />
</AccordionDetails>
</Accordion>
);
}
export type {
MessageProps,
};
export {
Message,
MessageMeta,
};