Skip to content

Commit 9b75114

Browse files
bogdan-tsbogdan-ts
committed
Merge branch 'toolcall_and_thinking_block' into 'master'
Tool_calls and Thinking blocks folding See merge request postgres-ai/database-lab!981
2 parents 71cc0f7 + 2fd56d7 commit 9b75114

15 files changed

Lines changed: 726 additions & 79 deletions

File tree

‎ui/cspell.json‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,8 @@
202202
"SPARQL",
203203
"subtransactions",
204204
"mbox",
205-
"SIEM"
205+
"SIEM",
206+
"toolcall",
207+
"thinkblock"
206208
]
207209
}

ui/packages/platform/src/pages/Bot/Messages/Message/CodeBlock.tsx renamed to ui/packages/platform/src/pages/Bot/Messages/Message/CodeBlock/CodeBlock.tsx

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import ExpandMoreIcon from "@material-ui/icons/ExpandMore";
77
import FileCopyOutlinedIcon from '@material-ui/icons/FileCopyOutlined';
88
import CheckCircleOutlineIcon from '@material-ui/icons/CheckCircleOutline';
99
import CodeIcon from '@material-ui/icons/Code';
10-
import { formatLanguageName } from "../../utils";
10+
import { formatLanguageName } from "../../../utils";
1111

1212
const useStyles = makeStyles((theme) => ({
1313
container: {

ui/packages/platform/src/pages/Bot/Messages/Message/MermaidDiagram.tsx renamed to ui/packages/platform/src/pages/Bot/Messages/Message/MermaidDiagram/MermaidDiagram.tsx

File renamed without changes.

ui/packages/platform/src/pages/Bot/Messages/Message/MermaidDiagramControls.tsx renamed to ui/packages/platform/src/pages/Bot/Messages/Message/MermaidDiagram/MermaidDiagramControls.tsx

File renamed without changes.

‎ui/packages/platform/src/pages/Bot/Messages/Message/Message.tsx‎

Lines changed: 41 additions & 62 deletions
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,23 @@
11
import React, { useEffect, useMemo, useRef, useState } from 'react'
2-
import cn from "classnames";
32
import ReactMarkdown, { Components } from "react-markdown";
43
import rehypeRaw from "rehype-raw";
54
import remarkGfm from "remark-gfm";
65
import { makeStyles } from "@material-ui/core";
76
import { colors } from "@postgres.ai/shared/styles/colors";
87
import { icons } from "@postgres.ai/shared/styles/icons";
98
import { DebugDialog } from "../../DebugDialog/DebugDialog";
10-
import { CodeBlock } from "./CodeBlock";
11-
import { disallowedHtmlTagsForMarkdown, permalinkLinkBuilder } from "../../utils";
9+
import { CodeBlock } from "./CodeBlock/CodeBlock";
10+
import { disallowedHtmlTagsForMarkdown } from "../../utils";
1211
import { MessageStatus, StateMessage } from "../../../../types/api/entities/bot";
13-
import { MermaidDiagram } from "./MermaidDiagram";
12+
import { MermaidDiagram } from "./MermaidDiagram/MermaidDiagram";
1413
import { useAiBot } from "../../hooks";
14+
import { ToolCallRenderer } from "./ToolCallRenderer/ToolCallRenderer";
15+
import { transformAllCustomTags } from "../utils";
16+
import { ThinkBlockRenderer } from './ThinkingCard/ThinkingCard';
17+
import { MessageHeader } from "./MessageHeader/MessageHeader";
1518

1619

17-
type BaseMessageProps = {
20+
export type BaseMessageProps = {
1821
id: string | null;
1922
created_at?: string;
2023
content?: string;
@@ -249,7 +252,6 @@ const useStyles = makeStyles(
249252
'50%': { borderRightColor: 'black' },
250253
},
251254
}),
252-
253255
)
254256

255257
export const Message = React.memo((props: MessageProps) => {
@@ -302,12 +304,16 @@ export const Message = React.memo((props: MessageProps) => {
302304
};
303305
}, [id, updateMessageStatus, isCurrentStreamMessage, isAi, threadId, status]);
304306

305-
const contentToRender: string = content?.replace(/\n/g, ' \n') || ''
307+
const contentToRender = useMemo(() => {
308+
if (!content) return '';
309+
return transformAllCustomTags(content?.replace(/\n/g, ' \n'));
310+
}, [content]);
306311

307312
const toggleDebugDialog = () => {
308313
setDebugVisible(prevState => !prevState)
309314
}
310315

316+
311317
const renderers = useMemo<Components>(() => ({
312318
p: ({ node, ...props }) => <div {...props} />,
313319
img: ({ node, ...props }) => <img style={{ maxWidth: '60%' }} {...props} />,
@@ -325,6 +331,8 @@ export const Message = React.memo((props: MessageProps) => {
325331
return <code {...props}>{children}</code>
326332
}
327333
},
334+
toolcall: ToolCallRenderer,
335+
thinkblock: ThinkBlockRenderer,
328336
}), []);
329337

330338
return (
@@ -344,51 +352,17 @@ export const Message = React.memo((props: MessageProps) => {
344352
/>
345353
: icons.userChatIcon}
346354
</div>
347-
<div className={classes.messageHeader}>
348-
<span className={classes.messageAuthor}>
349-
{isAi ? 'Postgres.AI' : name}
350-
</span>
351-
{created_at && formattedTime &&
352-
<span
353-
className={cn(classes.messageInfo)}
354-
title={created_at}
355-
>
356-
{formattedTime}
357-
</span>}
358-
<div className={classes.additionalInfo}>
359-
{id && isPublic && <>
360-
<span className={classes.messageInfo}>|</span>
361-
<a
362-
className={cn(classes.messageInfo, classes.messageInfoActive)}
363-
href={permalinkLinkBuilder(id)}
364-
target="_blank"
365-
rel="noreferrer"
366-
>
367-
permalink
368-
</a>
369-
</>}
370-
{!isLoading && isAi && id && <>
371-
<span className={classes.messageInfo}>|</span>
372-
<button
373-
className={cn(classes.messageInfo, classes.messageInfoActive)}
374-
onClick={toggleDebugDialog}
375-
>
376-
debug info
377-
</button>
378-
</>}
379-
{
380-
aiModel && isAi && <>
381-
<span className={classes.messageInfo}>|</span>
382-
<span
383-
className={cn(classes.messageInfo)}
384-
title={aiModel}
385-
>
386-
{aiModel}
387-
</span>
388-
</>
389-
}
390-
</div>
391-
</div>
355+
<MessageHeader
356+
name={name}
357+
createdAt={created_at}
358+
formattedTime={formattedTime}
359+
id={id}
360+
isPublic={isPublic}
361+
isAi={isAi}
362+
isLoading={isLoading}
363+
toggleDebugDialog={toggleDebugDialog}
364+
aiModel={aiModel}
365+
/>
392366
<div>
393367
{isLoading
394368
?
@@ -397,16 +371,21 @@ export const Message = React.memo((props: MessageProps) => {
397371
{stateMessage && stateMessage.state ? stateMessage.state : 'Thinking'}
398372
</div>
399373
</div>
400-
: <ReactMarkdown
401-
className={classes.markdown}
402-
children={contentToRender || ''}
403-
rehypePlugins={isAi ? [rehypeRaw] : []}
404-
remarkPlugins={[remarkGfm]}
405-
linkTarget='_blank'
406-
components={renderers}
407-
disallowedElements={disallowedHtmlTagsForMarkdown}
408-
unwrapDisallowed
409-
/>
374+
: <>
375+
<ReactMarkdown
376+
className={classes.markdown}
377+
children={contentToRender || ''}
378+
rehypePlugins={isAi ? [rehypeRaw] : []}
379+
remarkPlugins={[remarkGfm]}
380+
linkTarget='_blank'
381+
components={renderers}
382+
disallowedElements={disallowedHtmlTagsForMarkdown}
383+
unwrapDisallowed
384+
/>
385+
{stateMessage && stateMessage.state && <div className={classes.loading}>
386+
{stateMessage.state}
387+
</div>}
388+
</>
410389
}
411390
</div>
412391
</div>

‎ui/packages/platform/src/pages/Bot/Messages/Message/MessageHeader/MessageHeader.tsx‎

Lines changed: 118 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,118 @@
1+
import React from "react";
2+
import cn from "classnames";
3+
import { permalinkLinkBuilder } from "../../../utils";
4+
import { makeStyles } from "@material-ui/core";
5+
import { colors } from "@postgres.ai/shared/styles/colors";
6+
import { BaseMessageProps } from "../Message";
7+
8+
9+
const useStyles = makeStyles(
10+
() => ({
11+
messageAuthor: {
12+
fontSize: 14,
13+
fontWeight: 'bold',
14+
},
15+
messageInfo: {
16+
display: 'inline-block',
17+
marginLeft: 10,
18+
padding: 0,
19+
fontSize: '0.75rem',
20+
color: colors.pgaiDarkGray,
21+
transition: '.2s ease',
22+
background: "none",
23+
border: "none",
24+
textDecoration: "none",
25+
'@media (max-width: 450px)': {
26+
'&:nth-child(1)': {
27+
display: 'none'
28+
}
29+
}
30+
},
31+
messageInfoActive: {
32+
borderBottom: '1px solid currentcolor',
33+
cursor: 'pointer',
34+
'&:hover': {
35+
color: '#404040'
36+
}
37+
},
38+
messageHeader: {
39+
height: '1.125rem',
40+
display: 'flex',
41+
flexWrap: 'wrap',
42+
alignItems: 'baseline',
43+
'@media (max-width: 450px)': {
44+
height: 'auto',
45+
}
46+
},
47+
additionalInfo: {
48+
'@media (max-width: 450px)': {
49+
width: '100%',
50+
marginTop: 4,
51+
marginLeft: -10,
52+
53+
}
54+
},
55+
}),
56+
)
57+
58+
type MessageHeaderProps = Pick<
59+
BaseMessageProps,
60+
'name' | 'id' | 'formattedTime' | 'isPublic' | 'isLoading' | 'aiModel'
61+
> & {
62+
isAi: boolean;
63+
toggleDebugDialog: () => void;
64+
createdAt: BaseMessageProps["created_at"];
65+
};
66+
67+
export const MessageHeader = (props: MessageHeaderProps) => {
68+
const {isAi, formattedTime, id, name, createdAt, isLoading, aiModel, toggleDebugDialog, isPublic} = props;
69+
const classes = useStyles();
70+
return (
71+
<div className={classes.messageHeader}>
72+
<span className={classes.messageAuthor}>
73+
{isAi ? 'Postgres.AI' : name}
74+
</span>
75+
{createdAt && formattedTime &&
76+
<span
77+
className={cn(classes.messageInfo)}
78+
title={createdAt}
79+
>
80+
{formattedTime}
81+
</span>
82+
}
83+
<div className={classes.additionalInfo}>
84+
{id && isPublic && <>
85+
<span className={classes.messageInfo}>|</span>
86+
<a
87+
className={cn(classes.messageInfo, classes.messageInfoActive)}
88+
href={permalinkLinkBuilder(id)}
89+
target="_blank"
90+
rel="noreferrer"
91+
>
92+
permalink
93+
</a>
94+
</>}
95+
{!isLoading && isAi && id && <>
96+
<span className={classes.messageInfo}>|</span>
97+
<button
98+
className={cn(classes.messageInfo, classes.messageInfoActive)}
99+
onClick={toggleDebugDialog}
100+
>
101+
debug info
102+
</button>
103+
</>}
104+
{
105+
aiModel && isAi && <>
106+
<span className={classes.messageInfo}>|</span>
107+
<span
108+
className={cn(classes.messageInfo)}
109+
title={aiModel}
110+
>
111+
{aiModel}
112+
</span>
113+
</>
114+
}
115+
</div>
116+
</div>
117+
)
118+
}

‎ui/packages/platform/src/pages/Bot/Messages/Message/ThinkingCard/ThinkingCard.tsx‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,66 @@
1+
import React, { useState } from "react";
2+
import { Button } from "@postgres.ai/shared/components/Button2";
3+
import ExpandMoreIcon from "@material-ui/icons/ExpandMore";
4+
import { CardContent, Collapse } from "@mui/material";
5+
import ReactMarkdown from "react-markdown";
6+
import remarkGfm from "remark-gfm";
7+
8+
type ThinkBlockProps = {
9+
'data-think'?: string;
10+
node?: {
11+
properties?: {
12+
'data-think'?: string;
13+
dataThink?: string;
14+
};
15+
};
16+
}
17+
18+
type ThinkingCardProps = {
19+
content: string;
20+
}
21+
22+
const ThinkingCard = ({ content }: ThinkingCardProps) => {
23+
const [expanded, setExpanded] = useState(true);
24+
// TODO: Add "again"
25+
// TODO: Replace with "reasoned for X seconds"
26+
return (
27+
<>
28+
<Button
29+
onClick={() => setExpanded(!expanded)}
30+
>
31+
Took a moment to think
32+
<ExpandMoreIcon />
33+
</Button>
34+
35+
<Collapse in={expanded}>
36+
<CardContent>
37+
<ReactMarkdown remarkPlugins={[remarkGfm]}>
38+
{content}
39+
</ReactMarkdown>
40+
</CardContent>
41+
</Collapse>
42+
</>
43+
)
44+
}
45+
46+
export const ThinkBlockRenderer = React.memo((props: ThinkBlockProps) => {
47+
const dataThink =
48+
props?.['data-think'] ||
49+
props?.node?.properties?.['data-think'] ||
50+
props?.node?.properties?.dataThink;
51+
52+
if (!dataThink) return null;
53+
54+
let rawText = '';
55+
try {
56+
rawText = JSON.parse(dataThink);
57+
} catch (err) {
58+
console.error('Failed to parse data-think JSON:', err);
59+
}
60+
61+
return (
62+
<ThinkingCard content={rawText}/>
63+
)
64+
}, (prevProps, nextProps) => {
65+
return prevProps['data-think'] === nextProps['data-think'];
66+
})

0 commit comments

Comments
 (0)