references/lobe-chat
2
0
Fork 0
mirror of https://github.com/lobehub/lobe-chat.git synced 2026-06-18 13:25:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

✨ feat: support pplx search grounding (#6331)

Browse Source
This commit is contained in:
Arvin Xu
2025-02-20 12:31:08 +08:00
committed by GitHub
parent 057f45586b
commit ccb0003bff
33 changed files with 4125 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
locales/en-US/components.json
+3 -3
View File
@@ -111,8 +111,8 @@
}
},
"Thinking": {
"thinking": "Deep in thought...",
"thought": "Deeply thought (took {{duration}} seconds)",
"thoughtWithDuration": "Deeply thought"
"thinking": "Deep Thinking...",
"thought": "Deeply Thought (in {{duration}} seconds)",
"thoughtWithDuration": "Deeply Thought"
}
}
package.json
+1 -1
View File
@@ -129,7 +129,7 @@
"@lobehub/chat-plugins-gateway": "^1.9.0",
"@lobehub/icons": "^1.73.1",
"@lobehub/tts": "^1.28.0",
"@lobehub/ui": "^1.164.15",
"@lobehub/ui": "^1.165.0",
"@neondatabase/serverless": "^0.10.4",
"@next/third-parties": "^15.1.4",
"@react-spring/web": "^9.7.5",
src/components/ModelSelect/index.tsx
+24 -2
View File
@@ -2,7 +2,14 @@ import { IconAvatarProps, ModelIcon, ProviderIcon } from '@lobehub/icons';
import { Avatar, Icon, Tooltip } from '@lobehub/ui';
import { Typography } from 'antd';
import { createStyles } from 'antd-style';
import { Infinity, AtomIcon, LucideEye, LucidePaperclip, ToyBrick } from 'lucide-react';
import {
Infinity,
AtomIcon,
LucideEye,
LucideGlobe,
LucidePaperclip,
ToyBrick,
} from 'lucide-react';
import numeral from 'numeral';
import { rgba } from 'polished';
import { FC, memo } from 'react';
@@ -14,7 +21,7 @@ import { AiProviderSourceType } from '@/types/aiProvider';
import { ChatModelCard } from '@/types/llm';
import { formatTokenNumber } from '@/utils/format';
const useStyles = createStyles(({ css, token }) => ({
const useStyles = createStyles(({ css, token, isDarkMode }) => ({
custom: css`
width: 36px;
height: 20px;
@@ -41,6 +48,10 @@ const useStyles = createStyles(({ css, token }) => ({
color: ${token.geekblue};
background: ${token.geekblue1};
`,
tagCyan: css`
color: ${isDarkMode ? token.cyan7 : token.cyan10};
background: ${isDarkMode ? token.cyan1 : token.cyan2};
`,
tagGreen: css`
color: ${token.green};
background: ${token.green1};
@@ -122,6 +133,17 @@ export const ModelInfoTags = memo<ModelInfoTagsProps>(
</div>
</Tooltip>
)}
{model.search && (
<Tooltip
placement={placement}
styles={{ root: { pointerEvents: 'none' } }}
title={t('ModelSelect.featureTag.search')}
>
<div className={cx(styles.tag, styles.tagCyan)} style={{ cursor: 'pointer' }} title="">
<Icon icon={LucideGlobe} />
</div>
</Tooltip>
)}
{typeof model.contextWindowTokens === 'number' && (
<Tooltip
placement={placement}
src/components/Thinking/index.tsx
+7 -2
View File
@@ -7,6 +7,8 @@ import { CSSProperties, memo, useEffect, useState } from 'react';
import { useTranslation } from 'react-i18next';
import { Flexbox } from 'react-layout-kit';
import { CitationItem } from '@/types/message';
const useStyles = createStyles(({ css, token, isDarkMode }) => ({
container: css`
width: fit-content;
@@ -59,13 +61,14 @@ const useStyles = createStyles(({ css, token, isDarkMode }) => ({
}));
interface ThinkingProps {
citations?: CitationItem[];
content?: string;
duration?: number;
style?: CSSProperties;
thinking?: boolean;
}
const Thinking = memo<ThinkingProps>(({ content, duration, thinking, style }) => {
const Thinking = memo<ThinkingProps>(({ content, duration, thinking, style, citations }) => {
const { t } = useTranslation(['components', 'common']);
const { styles, cx } = useStyles();
@@ -135,7 +138,9 @@ const Thinking = memo<ThinkingProps>(({ content, duration, thinking, style }) =>
}}
>
{typeof content === 'string' ? (
<Markdown variant={'chat'}>{content}</Markdown>
<Markdown citations={citations} variant={'chat'}>
{content}
</Markdown>
) : (
content
)}
src/config/aiModels/jina.ts
+7 -5
View File
@@ -4,19 +4,21 @@ const jinaChatModels: AIChatModelCard[] = [
{
abilities: {
reasoning: true,
search: true,
},
contextWindowTokens: 1_000_000,
description: '深度搜索结合了网络搜索、阅读和推理，可进行全面调查。您可以将其视为一个代理，接受您的研究任务 - 它会进行广泛搜索并经过多次迭代，然后才能给出答案。这个过程涉及持续的研究、推理和从各个角度解决问题。这与直接从预训练数据生成答案的标准大模型以及依赖一次性表面搜索的传统 RAG 系统有着根本的不同。',
description:
'深度搜索结合了网络搜索、阅读和推理，可进行全面调查。您可以将其视为一个代理，接受您的研究任务 - 它会进行广泛搜索并经过多次迭代，然后才能给出答案。这个过程涉及持续的研究、推理和从各个角度解决问题。这与直接从预训练数据生成答案的标准大模型以及依赖一次性表面搜索的传统 RAG 系统有着根本的不同。',
displayName: 'Jina DeepSearch v1',
enabled: true,
id: 'jina-deepsearch-v1',
pricing: {
input: 0.02,
output: 0.02
output: 0.02,
},
type: 'chat'
}
]
type: 'chat',
},
];
export const allModels = [...jinaChatModels];
src/config/aiModels/perplexity.ts
+8
View File
@@ -4,6 +4,7 @@ const perplexityChatModels: AIChatModelCard[] = [
{
abilities: {
reasoning: true,
search: true,
},
contextWindowTokens: 127_072,
description: '由 DeepSeek 推理模型提供支持的新 API 产品。',
@@ -16,6 +17,7 @@ const perplexityChatModels: AIChatModelCard[] = [
{
abilities: {
reasoning: true,
search: true,
},
contextWindowTokens: 127_072,
description: '由 DeepSeek 推理模型提供支持的新 API 产品。',
@@ -26,6 +28,9 @@ const perplexityChatModels: AIChatModelCard[] = [
type: 'chat',
},
{
abilities: {
search: true,
},
contextWindowTokens: 200_000,
description: '支持搜索上下文的高级搜索产品，支持高级查询和跟进。',
displayName: 'Sonar Pro',
@@ -34,6 +39,9 @@ const perplexityChatModels: AIChatModelCard[] = [
type: 'chat',
},
{
abilities: {
search: true,
},
contextWindowTokens: 127_072,
description: '基于搜索上下文的轻量级搜索产品，比 Sonar Pro 更快、更便宜。',
displayName: 'Sonar',
src/database/client/migrations.json
+12 -8
View File
@@ -223,10 +223,7 @@
"hash": "9646161fa041354714f823d726af27247bcd6e60fa3be5698c0d69f337a5700b"
},
{
"sql": [
"DROP TABLE \"user_budgets\";",
"\nDROP TABLE \"user_subscriptions\";"
],
"sql": ["DROP TABLE \"user_budgets\";", "\nDROP TABLE \"user_subscriptions\";"],
"bps": true,
"folderMillis": 1729699958471,
"hash": "7dad43a2a25d1aec82124a4e53f8d82f8505c3073f23606c1dc5d2a4598eacf9"
@@ -298,11 +295,18 @@
"hash": "845a692ceabbfc3caf252a97d3e19a213bc0c433df2689900135f9cfded2cf49"
},
{
"sql": [
"ALTER TABLE \"messages\" ADD COLUMN \"reasoning\" jsonb;"
],
"sql": ["ALTER TABLE \"messages\" ADD COLUMN \"reasoning\" jsonb;"],
"bps": true,
"folderMillis": 1737609172353,
"hash": "2cb36ae4fcdd7b7064767e04bfbb36ae34518ff4bb1b39006f2dd394d1893868"
},
{
"sql": [
"ALTER TABLE \"messages\" ADD COLUMN \"search\" jsonb;",
"\nALTER TABLE \"messages\" ADD COLUMN \"metadata\" jsonb;"
],
"bps": true,
"folderMillis": 1739901891891,
"hash": "78d8fefd8c58938d7bc3da2295a73b35ce2e8d7cb2820f8e817acdb8dd5bebb2"
}
]
]
src/database/migrations/0015_add_message_search_metadata.sql
+2
View File
@@ -0,0 +1,2 @@
ALTER TABLE "messages" ADD COLUMN "search" jsonb;--> statement-breakpoint
ALTER TABLE "messages" ADD COLUMN "metadata" jsonb;
src/database/migrations/meta/0015_snapshot.json
+3616
View File
File diff suppressed because it is too large Load Diff
src/database/migrations/meta/_journal.json
+7
View File
@@ -105,6 +105,13 @@
"when": 1737609172353,
"tag": "0014_add_message_reasoning",
"breakpoints": true
},
{
"idx": 15,
"version": "7",
"when": 1739901891891,
"tag": "0015_add_message_search_metadata",
"breakpoints": true
}
],
"version": "6"
src/database/schemas/message.ts
+3 -1
View File
@@ -13,7 +13,7 @@ import {
import { createSelectSchema } from 'drizzle-zod';
import { idGenerator } from '@/database/utils/idGenerator';
import { ModelReasoning } from '@/types/message';
import { GroundingSearch, ModelReasoning } from '@/types/message';
import { timestamps } from './_helpers';
import { agents } from './agent';
@@ -34,6 +34,8 @@ export const messages = pgTable(
role: text('role', { enum: ['user', 'system', 'assistant', 'tool'] }).notNull(),
content: text('content'),
reasoning: jsonb('reasoning').$type<ModelReasoning>(),
search: jsonb('search').$type<GroundingSearch>(),
metadata: jsonb('metadata'),
model: text('model'),
provider: text('provider'),
src/database/server/models/message.ts
+2
View File
@@ -74,6 +74,8 @@ export class MessageModel {
role: messages.role,
content: messages.content,
reasoning: messages.reasoning,
search: messages.search,
metadata: messages.metadata,
error: messages.error,
model: messages.model,
src/features/Conversation/components/ChatItem/index.tsx
+10 -1
View File
@@ -172,12 +172,21 @@ const Item = memo<ChatListItemProps>(
const markdownProps = useMemo(
() => ({
citations: item?.role === 'user' ? undefined : item?.search?.citations,
components,
customRender: markdownCustomRender,
rehypePlugins: item?.role === 'user' ? undefined : rehypePlugins,
remarkPlugins: item?.role === 'user' ? undefined : remarkPlugins,
showCitations:
item?.role === 'user'
? undefined
: item?.search?.citations &&
// if the citations are all empty, we should not show the citations
item?.search?.citations.length > 0 &&
// if the citations's url and title are all the same, we should not show the citations
item?.search?.citations.every((item) => item.title !== item.url),
}),
[components, markdownCustomRender, item?.role],
[components, markdownCustomRender, item?.role, item?.search],
);
const onChange = useCallback((value: string) => updateMessageContent(id, value), [id]);
src/features/Conversation/components/MarkdownElements/Thinking/Render.tsx
+5 -1
View File
@@ -18,10 +18,14 @@ const Render = memo<MarkdownElementProps>(({ children, id }) => {
const message = chatSelectors.getMessageById(id)(s);
return [!isThinkingClosed(message?.content)];
});
const citations = useChatStore((s) => {
const message = chatSelectors.getMessageById(id)(s);
return message?.search?.citations;
});
if (!isGenerating && !children) return;
return <Thinking content={children as string} thinking={isGenerating} />;
return <Thinking citations={citations} content={children as string} thinking={isGenerating} />;
});
export default Render;
src/features/Conversation/components/MarkdownElements/remarkPlugins/createRemarkCustomTagPlugin.ts
+1
View File
@@ -32,6 +32,7 @@ export const createRemarkCustomTagPlugin = (tag: string) => () => {
);
// 转换为 Markdown 字符串
const content = treeNodeToString(contentNodes);
// 创建自定义节点
src/features/Conversation/components/MarkdownElements/remarkPlugins/getNodeContent.test.ts
+107
View File
@@ -393,4 +393,111 @@ describe('treeNodeToString', () => {
1. 求 $a_2$ 和 $a_3$，根据前三项的规律猜想该数列的通项公式
2. 用数学归纳法证明你的猜想。`);
});
describe('link node', () => {
it('with url', () => {
const nodes = [
{
type: 'paragraph',
children: [
{
type: 'link',
title: null,
url: 'citation-1',
children: [
{
type: 'text',
value: '#citation-1',
position: {
start: {
line: 5,
column: 26,
offset: 78,
},
end: {
line: 5,
column: 37,
offset: 89,
},
},
},
],
position: {
start: {
line: 5,
column: 25,
offset: 77,
},
end: {
line: 5,
column: 50,
offset: 102,
},
},
},
],
position: {
start: {
line: 5,
column: 1,
offset: 53,
},
end: {
line: 5,
column: 220,
offset: 272,
},
},
},
];
const result = treeNodeToString(nodes as Parent[]);
expect(result).toEqual(`[#citation-1](citation-1)`);
});
it('handle error case', () => {
const nodes = [
{
type: 'paragraph',
children: [
{
type: 'link',
title: null,
url: 'citation-1',
children: [],
position: {
start: {
line: 5,
column: 25,
offset: 77,
},
end: {
line: 5,
column: 50,
offset: 102,
},
},
},
],
position: {
start: {
line: 5,
column: 1,
offset: 53,
},
end: {
line: 5,
column: 220,
offset: 272,
},
},
},
];
const result = treeNodeToString(nodes as Parent[]);
expect(result).toEqual(`[](citation-1)`);
});
});
});
src/features/Conversation/components/MarkdownElements/remarkPlugins/getNodeContent.ts
+6
View File
@@ -7,6 +7,12 @@ const processNode = (node: any): string => {
return `$${node.value}$`;
}
if (node.type === 'link') {
const text = node.children?.[0] ? processNode(node.children?.[0]) : '';
return `[${text}](${node.url})`;
}
// 处理带有子节点的容器
if (node.children) {
const content = node.children.map((element: Parent) => processNode(element)).join('');
src/libs/agent-runtime/perplexity/index.test.ts
+156 -12
View File
@@ -28,17 +28,9 @@ beforeEach(() => {
);
});
afterEach(() => {
vi.clearAllMocks();
});
describe('LobePerplexityAI', () => {
describe('chat', () => {
it('should call chat method with temperature', async () => {
vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
new ReadableStream() as any,
);
await instance.chat({
messages: [{ content: 'Hello', role: 'user' }],
model: 'text-davinci-003',
@@ -56,10 +48,6 @@ describe('LobePerplexityAI', () => {
});
it('should be undefined when temperature >= 2', async () => {
vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
new ReadableStream() as any,
);
await instance.chat({
messages: [{ content: 'Hello', role: 'user' }],
model: 'text-davinci-003',
@@ -75,5 +63,161 @@ describe('LobePerplexityAI', () => {
expect.any(Object),
);
});
it('should with search citations', async () => {
const data = [
{
id: '506d64fb-e7f2-4d94-b80f-158369e9446d',
model: 'sonar-pro',
created: 1739896615,
usage: {
prompt_tokens: 4,
completion_tokens: 3,
total_tokens: 7,
citation_tokens: 2217,
num_search_queries: 1,
},
citations: [
'https://www.weather.com.cn/weather/101210101.shtml',
'https://tianqi.moji.com/weather/china/zhejiang/hangzhou',
'https://weather.cma.cn/web/weather/58457.html',
'https://tianqi.so.com/weather/101210101',
'https://www.accuweather.com/zh/cn/hangzhou/106832/weather-forecast/106832',
'https://www.hzqx.com',
'https://www.hzqx.com/pc/hztq/',
],
object: 'chat.completion',
choices: [
{
index: 0,
finish_reason: null,
message: {
role: 'assistant',
content: '杭州今',
},
delta: {
role: 'assist',
content: '杭州今',
},
},
],
},
{
id: '506d64fb-e7f2-4d94-b80f-158369e9446d',
model: 'sonar-pro',
created: 1739896615,
usage: {
prompt_tokens: 4,
completion_tokens: 9,
total_tokens: 13,
citation_tokens: 2217,
num_search_queries: 1,
},
citations: [
'https://www.weather.com.cn/weather/101210101.shtml',
'https://tianqi.moji.com/weather/china/zhejiang/hangzhou',
'https://weather.cma.cn/web/weather/58457.html',
'https://tianqi.so.com/weather/101210101',
'https://www.accuweather.com/zh/cn/hangzhou/106832/weather-forecast/106832',
'https://www.hzqx.com',
'https://www.hzqx.com/pc/hztq/',
],
object: 'chat.completion',
choices: [
{
index: 0,
finish_reason: null,
message: {
role: 'assistant',
content: '杭州今天和未来几天的',
},
delta: {
role: 'assistant',
content: '天和未来几天的',
},
},
],
},
{
id: '506d64fb-e7f2-4d94-b80f-158369e9446d',
model: 'sonar-pro',
created: 1739896615,
usage: {
prompt_tokens: 4,
completion_tokens: 14,
total_tokens: 18,
citation_tokens: 2217,
num_search_queries: 1,
},
citations: [
'https://www.weather.com.cn/weather/101210101.shtml',
'https://tianqi.moji.com/weather/china/zhejiang/hangzhou',
'https://weather.cma.cn/web/weather/58457.html',
'https://tianqi.so.com/weather/101210101',
'https://www.accuweather.com/zh/cn/hangzhou/106832/weather-forecast/106832',
'https://www.hzqx.com',
'https://www.hzqx.com/pc/hztq/',
],
object: 'chat.completion',
choices: [
{
index: 0,
finish_reason: null,
message: {
role: 'assistant',
content: '杭州今天和未来几天的天气预报如',
},
},
],
},
];
const mockStream = new ReadableStream({
start(controller) {
data.forEach((chunk) => {
controller.enqueue(chunk);
});
controller.close();
},
});
vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(mockStream as any);
const result = await instance.chat({
messages: [{ content: 'Hello', role: 'user' }],
model: 'mistralai/mistral-7b-instruct:free',
temperature: 0,
});
const decoder = new TextDecoder();
const reader = result.body!.getReader();
const stream: string[] = [];
while (true) {
const { value, done } = await reader.read();
if (done) break;
stream.push(decoder.decode(value));
}
expect(stream).toEqual(
[
'id: 506d64fb-e7f2-4d94-b80f-158369e9446d',
'event: citations',
'data: [{"title":"https://www.weather.com.cn/weather/101210101.shtml","url":"https://www.weather.com.cn/weather/101210101.shtml"},{"title":"https://tianqi.moji.com/weather/china/zhejiang/hangzhou","url":"https://tianqi.moji.com/weather/china/zhejiang/hangzhou"},{"title":"https://weather.cma.cn/web/weather/58457.html","url":"https://weather.cma.cn/web/weather/58457.html"},{"title":"https://tianqi.so.com/weather/101210101","url":"https://tianqi.so.com/weather/101210101"},{"title":"https://www.accuweather.com/zh/cn/hangzhou/106832/weather-forecast/106832","url":"https://www.accuweather.com/zh/cn/hangzhou/106832/weather-forecast/106832"},{"title":"https://www.hzqx.com","url":"https://www.hzqx.com"},{"title":"https://www.hzqx.com/pc/hztq/","url":"https://www.hzqx.com/pc/hztq/"}]\n',
'id: 506d64fb-e7f2-4d94-b80f-158369e9446d',
'event: text',
'data: "杭州今"\n',
'id: 506d64fb-e7f2-4d94-b80f-158369e9446d',
'event: text',
'data: "天和未来几天的"\n',
'id: 506d64fb-e7f2-4d94-b80f-158369e9446d',
'event: data',
'data: {"id":"506d64fb-e7f2-4d94-b80f-158369e9446d","index":0}\n',
].map((line) => `${line}\n`),
);
expect((await reader.read()).done).toBe(true);
});
});
});
src/libs/agent-runtime/utils/streams/anthropic.ts
+3 -3
View File
@@ -3,9 +3,9 @@ import type { Stream } from '@anthropic-ai/sdk/streaming';
import { ChatStreamCallbacks } from '../../types';
import {
StreamContext,
StreamProtocolChunk,
StreamProtocolToolCallChunk,
StreamStack,
StreamToolCallChunkData,
convertIterableToStream,
createCallbacksTransformer,
@@ -14,7 +14,7 @@ import {
export const transformAnthropicStream = (
chunk: Anthropic.MessageStreamEvent,
stack: StreamStack,
stack: StreamContext,
): StreamProtocolChunk => {
// maybe need another structure to add support for multiple choices
switch (chunk.type) {
@@ -100,7 +100,7 @@ export const AnthropicStream = (
stream: Stream<Anthropic.MessageStreamEvent> | ReadableStream,
callbacks?: ChatStreamCallbacks,
) => {
const streamStack: StreamStack = { id: '' };
const streamStack: StreamContext = { id: '' };
const readableStream =
stream instanceof ReadableStream ? stream : convertIterableToStream(stream);
src/libs/agent-runtime/utils/streams/bedrock/claude.ts
+6 -2
View File
@@ -4,14 +4,18 @@ import { nanoid } from '@/utils/uuid';
import { ChatStreamCallbacks } from '../../../types';
import { transformAnthropicStream } from '../anthropic';
import { StreamStack, createCallbacksTransformer, createSSEProtocolTransformer } from '../protocol';
import {
StreamContext,
createCallbacksTransformer,
createSSEProtocolTransformer,
} from '../protocol';
import { createBedrockStream } from './common';
export const AWSBedrockClaudeStream = (
res: InvokeModelWithResponseStreamResponse | ReadableStream,
cb?: ChatStreamCallbacks,
): ReadableStream<string> => {
const streamStack: StreamStack = { id: 'chat_' + nanoid() };
const streamStack: StreamContext = { id: 'chat_' + nanoid() };
const stream = res instanceof ReadableStream ? res : createBedrockStream(res);
src/libs/agent-runtime/utils/streams/bedrock/llama.ts
+3 -3
View File
@@ -4,8 +4,8 @@ import { nanoid } from '@/utils/uuid';
import { ChatStreamCallbacks } from '../../../types';
import {
StreamContext,
StreamProtocolChunk,
StreamStack,
createCallbacksTransformer,
createSSEProtocolTransformer,
} from '../protocol';
@@ -27,7 +27,7 @@ interface BedrockLlamaStreamChunk {
export const transformLlamaStream = (
chunk: BedrockLlamaStreamChunk,
stack: StreamStack,
stack: StreamContext,
): StreamProtocolChunk => {
// maybe need another structure to add support for multiple choices
if (chunk.stop_reason) {
@@ -41,7 +41,7 @@ export const AWSBedrockLlamaStream = (
res: InvokeModelWithResponseStreamResponse | ReadableStream,
cb?: ChatStreamCallbacks,
): ReadableStream<string> => {
const streamStack: StreamStack = { id: 'chat_' + nanoid() };
const streamStack: StreamContext = { id: 'chat_' + nanoid() };
const stream = res instanceof ReadableStream ? res : createBedrockStream(res);
src/libs/agent-runtime/utils/streams/google-ai.ts
+3 -3
View File
@@ -4,8 +4,8 @@ import { nanoid } from '@/utils/uuid';
import { ChatStreamCallbacks } from '../../types';
import {
StreamContext,
StreamProtocolChunk,
StreamStack,
StreamToolCallChunkData,
createCallbacksTransformer,
createSSEProtocolTransformer,
@@ -14,7 +14,7 @@ import {
const transformGoogleGenerativeAIStream = (
chunk: EnhancedGenerateContentResponse,
stack: StreamStack,
stack: StreamContext,
): StreamProtocolChunk => {
// maybe need another structure to add support for multiple choices
const functionCalls = chunk.functionCalls();
@@ -49,7 +49,7 @@ export const GoogleGenerativeAIStream = (
rawStream: ReadableStream<EnhancedGenerateContentResponse>,
callbacks?: ChatStreamCallbacks,
) => {
const streamStack: StreamStack = { id: 'chat_' + nanoid() };
const streamStack: StreamContext = { id: 'chat_' + nanoid() };
return rawStream
.pipeThrough(createSSEProtocolTransformer(transformGoogleGenerativeAIStream, streamStack))
src/libs/agent-runtime/utils/streams/ollama.ts
+3 -3
View File
@@ -4,14 +4,14 @@ import { ChatStreamCallbacks } from '@/libs/agent-runtime';
import { nanoid } from '@/utils/uuid';
import {
StreamContext,
StreamProtocolChunk,
StreamStack,
createCallbacksTransformer,
createSSEProtocolTransformer,
generateToolCallId,
} from './protocol';
const transformOllamaStream = (chunk: ChatResponse, stack: StreamStack): StreamProtocolChunk => {
const transformOllamaStream = (chunk: ChatResponse, stack: StreamContext): StreamProtocolChunk => {
// maybe need another structure to add support for multiple choices
if (chunk.done && !chunk.message.content) {
return { data: 'finished', id: stack.id, type: 'stop' };
@@ -39,7 +39,7 @@ export const OllamaStream = (
res: ReadableStream<ChatResponse>,
cb?: ChatStreamCallbacks,
): ReadableStream<string> => {
const streamStack: StreamStack = { id: 'chat_' + nanoid() };
const streamStack: StreamContext = { id: 'chat_' + nanoid() };
return res
.pipeThrough(createSSEProtocolTransformer(transformOllamaStream, streamStack))
src/libs/agent-runtime/utils/streams/openai.ts
+26 -8
View File
@@ -1,15 +1,15 @@
import OpenAI from 'openai';
import type { Stream } from 'openai/streaming';
import { ChatMessageError } from '@/types/message';
import { ChatMessageError, CitationItem } from '@/types/message';
import { AgentRuntimeErrorType, ILobeAgentRuntimeErrorType } from '../../error';
import { ChatStreamCallbacks } from '../../types';
import {
FIRST_CHUNK_ERROR_KEY,
StreamContext,
StreamProtocolChunk,
StreamProtocolToolCallChunk,
StreamStack,
StreamToolCallChunkData,
convertIterableToStream,
createCallbacksTransformer,
@@ -20,8 +20,8 @@ import {
export const transformOpenAIStream = (
chunk: OpenAI.ChatCompletionChunk,
stack?: StreamStack,
): StreamProtocolChunk => {
streamContext: StreamContext,
): StreamProtocolChunk | StreamProtocolChunk[] => {
// handle the first chunk error
if (FIRST_CHUNK_ERROR_KEY in chunk) {
delete chunk[FIRST_CHUNK_ERROR_KEY];
@@ -48,8 +48,8 @@ export const transformOpenAIStream = (
if (typeof item.delta?.tool_calls === 'object' && item.delta.tool_calls?.length > 0) {
return {
data: item.delta.tool_calls.map((value, index): StreamToolCallChunkData => {
if (stack && !stack.tool) {
stack.tool = { id: value.id!, index: value.index, name: value.function!.name! };
if (streamContext && !streamContext.tool) {
streamContext.tool = { id: value.id!, index: value.index, name: value.function!.name! };
}
return {
@@ -57,7 +57,10 @@ export const transformOpenAIStream = (
arguments: value.function?.arguments ?? '{}',
name: value.function?.name ?? null,
},
id: value.id || stack?.tool?.id || generateToolCallId(index, value.function?.name),
id:
value.id ||
streamContext?.tool?.id ||
generateToolCallId(index, value.function?.name),
// mistral's tool calling don't have index and function field, it's data like:
// [{"id":"xbhnmTtY7","function":{"name":"lobe-image-designer____text2image____builtin","arguments":"{\"prompts\": [\"A photo of a small, fluffy dog with a playful expression and wagging tail.\", \"A watercolor painting of a small, energetic dog with a glossy coat and bright eyes.\", \"A vector illustration of a small, adorable dog with a short snout and perky ears.\", \"A drawing of a small, scruffy dog with a mischievous grin and a wagging tail.\"], \"quality\": \"standard\", \"seeds\": [123456, 654321, 111222, 333444], \"size\": \"1024x1024\", \"style\": \"vivid\"}"}}]
@@ -114,6 +117,21 @@ export const transformOpenAIStream = (
}
if (typeof content === 'string') {
// in Perplexity api, the citation is in every chunk, but we only need to return it once
if ('citations' in chunk && !streamContext?.returnedPplxCitation) {
streamContext.returnedPplxCitation = true;
const citations = (chunk.citations as any[]).map((item) =>
typeof item === 'string' ? ({ title: item, url: item } as CitationItem) : item,
);
return [
{ data: citations, id: chunk.id, type: 'citations' },
{ data: content, id: chunk.id, type: 'text' },
];
}
return { data: content, id: chunk.id, type: 'text' };
}
}
@@ -164,7 +182,7 @@ export const OpenAIStream = (
stream: Stream<OpenAI.ChatCompletionChunk> | ReadableStream,
{ callbacks, provider, bizErrorTypeTransformer }: OpenAIStreamOptions = {},
) => {
const streamStack: StreamStack = { id: '' };
const streamStack: StreamContext = { id: '' };
const readableStream =
stream instanceof ReadableStream ? stream : convertIterableToStream(stream);
src/libs/agent-runtime/utils/streams/protocol.ts
+33 -8
View File
@@ -2,8 +2,16 @@ import { ChatStreamCallbacks } from '@/libs/agent-runtime';
import { AgentRuntimeErrorType } from '../../error';
export interface StreamStack {
/**
* context in the stream to save temporarily data
*/
export interface StreamContext {
id: string;
/**
* As pplx citations is in every chunk, but we only need to return it once
* this flag is used to check if the pplx citation is returned,and then not return it again
*/
returnedPplxCitation?: boolean;
tool?: {
id: string;
index: number;
@@ -15,7 +23,20 @@ export interface StreamStack {
export interface StreamProtocolChunk {
data: any;
id?: string;
type: 'text' | 'tool_calls' | 'data' | 'stop' | 'error' | 'reasoning';
type: // pure text
| 'text'
// Tools use
| 'tool_calls'
// Model Thinking
| 'reasoning'
// Search or Grounding
| 'citations'
// stop signal
| 'stop'
// Error
| 'error'
// unknown data result
| 'data';
}
export interface StreamToolCallChunkData {
@@ -85,16 +106,20 @@ export const convertIterableToStream = <T>(stream: AsyncIterable<T>) => {
* Create a transformer to convert the response into an SSE format
*/
export const createSSEProtocolTransformer = (
transformer: (chunk: any, stack: StreamStack) => StreamProtocolChunk,
streamStack?: StreamStack,
transformer: (chunk: any, stack: StreamContext) => StreamProtocolChunk | StreamProtocolChunk[],
streamStack?: StreamContext,
) =>
new TransformStream({
transform: (chunk, controller) => {
const { type, id, data } = transformer(chunk, streamStack || { id: '' });
const result = transformer(chunk, streamStack || { id: '' });
controller.enqueue(`id: ${id}\n`);
controller.enqueue(`event: ${type}\n`);
controller.enqueue(`data: ${JSON.stringify(data)}\n\n`);
const buffers = Array.isArray(result) ? result : [result];
buffers.forEach(({ type, id, data }) => {
controller.enqueue(`id: ${id}\n`);
controller.enqueue(`event: ${type}\n`);
controller.enqueue(`data: ${JSON.stringify(data)}\n\n`);
});
},
});
src/libs/agent-runtime/utils/streams/vertex-ai.ts
+3 -3
View File
@@ -4,8 +4,8 @@ import { nanoid } from '@/utils/uuid';
import { ChatStreamCallbacks } from '../../types';
import {
StreamContext,
StreamProtocolChunk,
StreamStack,
createCallbacksTransformer,
createSSEProtocolTransformer,
generateToolCallId,
@@ -13,7 +13,7 @@ import {
const transformVertexAIStream = (
chunk: GenerateContentResponse,
stack: StreamStack,
stack: StreamContext,
): StreamProtocolChunk => {
// maybe need another structure to add support for multiple choices
const candidates = chunk.candidates;
@@ -67,7 +67,7 @@ export const VertexAIStream = (
rawStream: ReadableStream<EnhancedGenerateContentResponse>,
callbacks?: ChatStreamCallbacks,
) => {
const streamStack: StreamStack = { id: 'chat_' + nanoid() };
const streamStack: StreamContext = { id: 'chat_' + nanoid() };
return rawStream
.pipeThrough(createSSEProtocolTransformer(transformVertexAIStream, streamStack))
src/locales/default/components.ts
+1
View File
@@ -79,6 +79,7 @@ export default {
file: '该模型支持上传文件读取与识别',
functionCall: '该模型支持函数调用（Function Call）',
reasoning: '该模型支持深度思考',
search: '该模型支持联网搜索',
tokens: '该模型单个会话最多支持 {{tokens}} Tokens',
vision: '该模型支持视觉识别',
},
src/store/chat/slices/aiChat/actions/generateAIChat.ts
+17 -6
View File
@@ -455,7 +455,7 @@ export const generateAIChat: StateCreator<
await messageService.updateMessageError(messageId, error);
await refreshMessages();
},
onFinish: async (content, { traceId, observationId, toolCalls, reasoning }) => {
onFinish: async (content, { traceId, observationId, toolCalls, reasoning, citations }) => {
// if there is traceId, update it
if (traceId) {
msgTraceId = traceId;
@@ -470,15 +470,26 @@ export const generateAIChat: StateCreator<
}
// update the content after fetch result
await internal_updateMessageContent(
messageId,
content,
await internal_updateMessageContent(messageId, content, {
toolCalls,
!!reasoning ? { content: reasoning, duration } : undefined,
);
reasoning: !!reasoning ? { content: reasoning, duration } : undefined,
search: !!citations ? { citations } : undefined,
});
},
onMessageHandle: async (chunk) => {
switch (chunk.type) {
case 'citations': {
// if there is no citations, then stop
if (!chunk.citations || chunk.citations.length <= 0) return;
internal_dispatchMessage({
id: messageId,
type: 'updateMessage',
value: { search: { citations: chunk.citations } },
});
break;
}
case 'text': {
output += chunk.text;
src/store/chat/slices/message/action.ts
+12 -7
View File
@@ -16,6 +16,7 @@ import {
ChatMessage,
ChatMessageError,
CreateMessageParams,
GroundingSearch,
MessageToolCall,
ModelReasoning,
} from '@/types/message';
@@ -73,8 +74,11 @@ export interface ChatMessageAction {
internal_updateMessageContent: (
id: string,
content: string,
toolCalls?: MessageToolCall[],
reasoning?: ModelReasoning,
extra?: {
toolCalls?: MessageToolCall[];
reasoning?: ModelReasoning;
search?: GroundingSearch;
},
) => Promise<void>;
/**
* update the message error with optimistic update
@@ -272,17 +276,17 @@ export const chatMessage: StateCreator<
await messageService.updateMessage(id, { error });
await get().refreshMessages();
},
internal_updateMessageContent: async (id, content, toolCalls, reasoning) => {
internal_updateMessageContent: async (id, content, extra) => {
const { internal_dispatchMessage, refreshMessages, internal_transformToolCalls } = get();
// Due to the async update method and refresh need about 100ms
// we need to update the message content at the frontend to avoid the update flick
// refs: https://medium.com/@kyledeguzmanx/what-are-optimistic-updates-483662c3e171
if (toolCalls) {
if (extra?.toolCalls) {
internal_dispatchMessage({
id,
type: 'updateMessage',
value: { tools: internal_transformToolCalls(toolCalls) },
value: { tools: internal_transformToolCalls(extra?.toolCalls) },
});
} else {
internal_dispatchMessage({ id, type: 'updateMessage', value: { content } });
@@ -290,8 +294,9 @@ export const chatMessage: StateCreator<
await messageService.updateMessage(id, {
content,
tools: toolCalls ? internal_transformToolCalls(toolCalls) : undefined,
reasoning,
tools: extra?.toolCalls ? internal_transformToolCalls(extra?.toolCalls) : undefined,
reasoning: extra?.reasoning,
search: extra?.search,
});
await refreshMessages();
},
src/types/aiModel.ts
+5
View File
@@ -34,6 +34,11 @@ export interface ModelAbilities {
* whether model supports reasoning
*/
reasoning?: boolean;
/**
* whether model supports search web
*/
search?: boolean;
/**
* whether model supports vision
*/
src/types/message/base.ts
+13
View File
@@ -1,3 +1,15 @@
export interface CitationItem {
id?: string;
onlyUrl?: boolean;
title?: string;
url: string;
}
export interface GroundingSearch {
citations?: CitationItem[];
searchQueries?: string[];
}
export interface ModelReasoning {
content?: string;
duration?: number;
@@ -20,6 +32,7 @@ export interface MessageItem {
quotaId: string | null;
reasoning: ModelReasoning | null;
role: string;
search: GroundingSearch | null;
sessionId: string | null;
threadId: string | null;
// jsonb type
src/types/message/chat.ts
+3 -2
View File
@@ -2,7 +2,7 @@ import { IPluginErrorType } from '@lobehub/chat-plugin-sdk';
import { ILobeAgentRuntimeErrorType } from '@/libs/agent-runtime';
import { ErrorType } from '@/types/fetch';
import { MessageRoleType, ModelReasoning } from '@/types/message/base';
import { GroundingSearch, MessageRoleType, ModelReasoning } from '@/types/message/base';
import { ChatPluginPayload, ChatToolPayload } from '@/types/message/tools';
import { Translate } from '@/types/message/translate';
import { MetaData } from '@/types/meta';
@@ -100,11 +100,12 @@ export interface ChatMessage {
ragRawQuery?: string | null;
reasoning?: ModelReasoning | null;
/**
* message role type
*/
role: MessageRoleType;
search?: GroundingSearch | null;
sessionId?: string;
threadId?: string | null;
tool_call_id?: string;
src/utils/fetch/fetchSSE.ts
+17 -1
View File
@@ -6,6 +6,7 @@ import { ChatErrorType } from '@/types/fetch';
import { SmoothingParams } from '@/types/llm';
import {
ChatMessageError,
CitationItem,
MessageToolCall,
MessageToolCallChunk,
MessageToolCallSchema,
@@ -20,6 +21,7 @@ type SSEFinishType = 'done' | 'error' | 'abort';
export type OnFinishHandler = (
text: string,
context: {
citations?: CitationItem[];
observationId?: string | null;
reasoning?: string;
toolCalls?: MessageToolCall[];
@@ -38,6 +40,11 @@ export interface MessageReasoningChunk {
type: 'reasoning';
}
export interface MessageCitationsChunk {
citations: CitationItem[];
type: 'citations';
}
interface MessageToolCallsChunk {
isAnimationActives?: boolean[];
tool_calls: MessageToolCall[];
@@ -50,7 +57,7 @@ export interface FetchSSEOptions {
onErrorHandle?: (error: ChatMessageError) => void;
onFinish?: OnFinishHandler;
onMessageHandle?: (
chunk: MessageTextChunk | MessageToolCallsChunk | MessageReasoningChunk,
chunk: MessageTextChunk | MessageToolCallsChunk | MessageReasoningChunk | MessageCitationsChunk,
) => void;
smoothing?: SmoothingParams | boolean;
}
@@ -279,6 +286,7 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio
startSpeed: smoothingSpeed,
});
let citations: CitationItem[] | undefined = undefined;
await fetchEventSource(url, {
body: options.body,
fetch: options?.fetcher,
@@ -350,6 +358,13 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio
break;
}
case 'citations': {
citations = data;
options.onMessageHandle?.({ citations: data, type: 'citations' });
break;
}
case 'reasoning': {
if (textSmoothing) {
thinkingController.pushToQueue(data);
@@ -419,6 +434,7 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio
}
await options?.onFinish?.(output, {
citations,
observationId,
reasoning: !!thinking ? thinking : undefined,
toolCalls,