💄 style: add build-in web search support for Wenxin & Hunyuan (#6617)

*  feat: add build-in web search support for Wenxin

* 🐛 fix: fix web_search calling issue

*  feat: add support wenxin `search_results` stream

*  feat: add `search_info` stream support for Hunyuan

* 🔨 chore: minor logic

* 🔨 chore: add unit test

* 🐛 fix: try to fix ci error

* 🐛 fix: fix ci error
This commit is contained in:
Zhijie He
2025-03-03 22:51:13 +08:00
committed by GitHub
parent 6f7cf45493
commit dfd1f093fe
8 changed files with 335 additions and 20 deletions
+32
View File
@@ -4,6 +4,7 @@ const wenxinChatModels: AIChatModelCard[] = [
{
abilities: {
functionCall: true,
search: true,
},
contextWindowTokens: 8192,
description:
@@ -16,11 +17,15 @@ const wenxinChatModels: AIChatModelCard[] = [
input: 0.8,
output: 2,
},
settings: {
searchImpl: 'params',
},
type: 'chat',
},
{
abilities: {
functionCall: true,
search: true,
},
contextWindowTokens: 8192,
description:
@@ -32,11 +37,15 @@ const wenxinChatModels: AIChatModelCard[] = [
input: 0.8,
output: 2,
},
settings: {
searchImpl: 'params',
},
type: 'chat',
},
{
abilities: {
functionCall: true,
search: true,
},
contextWindowTokens: 128_000,
description:
@@ -49,11 +58,15 @@ const wenxinChatModels: AIChatModelCard[] = [
input: 0.8,
output: 2,
},
settings: {
searchImpl: 'params',
},
type: 'chat',
},
{
abilities: {
functionCall: true,
search: true,
},
contextWindowTokens: 8192,
description:
@@ -66,11 +79,15 @@ const wenxinChatModels: AIChatModelCard[] = [
input: 30,
output: 90,
},
settings: {
searchImpl: 'params',
},
type: 'chat',
},
{
abilities: {
functionCall: true,
search: true,
},
contextWindowTokens: 8192,
description:
@@ -82,11 +99,15 @@ const wenxinChatModels: AIChatModelCard[] = [
input: 30,
output: 90,
},
settings: {
searchImpl: 'params',
},
type: 'chat',
},
{
abilities: {
functionCall: true,
search: true,
},
contextWindowTokens: 8192,
description:
@@ -99,11 +120,15 @@ const wenxinChatModels: AIChatModelCard[] = [
input: 20,
output: 60,
},
settings: {
searchImpl: 'params',
},
type: 'chat',
},
{
abilities: {
functionCall: true,
search: true,
},
contextWindowTokens: 128_000,
description:
@@ -116,11 +141,15 @@ const wenxinChatModels: AIChatModelCard[] = [
input: 20,
output: 60,
},
settings: {
searchImpl: 'params',
},
type: 'chat',
},
{
abilities: {
functionCall: true,
search: true,
},
contextWindowTokens: 8192,
description:
@@ -132,6 +161,9 @@ const wenxinChatModels: AIChatModelCard[] = [
input: 20,
output: 60,
},
settings: {
searchImpl: 'params',
},
type: 'chat',
},
{
+137 -2
View File
@@ -1,7 +1,9 @@
// @vitest-environment node
import { ModelProvider } from '@/libs/agent-runtime';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { LobeOpenAICompatibleRuntime, ModelProvider } from '@/libs/agent-runtime';
import { testProvider } from '@/libs/agent-runtime/providerTestUtils';
import { testProvider } from '../providerTestUtils';
import { LobeHunyuanAI } from './index';
testProvider({
@@ -11,3 +13,136 @@ testProvider({
chatDebugEnv: 'DEBUG_HUNYUAN_CHAT_COMPLETION',
chatModel: 'hunyuan-lite',
});
// Mock the console.error to avoid polluting test output
vi.spyOn(console, 'error').mockImplementation(() => {});
let instance: LobeOpenAICompatibleRuntime;
beforeEach(() => {
instance = new LobeHunyuanAI({ apiKey: 'test' });
// 使用 vi.spyOn 来模拟 chat.completions.create 方法
vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
new ReadableStream() as any,
);
});
describe('LobeHunyuanAI', () => {
describe('chat', () => {
it('should with search citations', async () => {
const data = [
{
id: "939fbdb8dbb9b4c5944cbbe687c977c2",
object: "chat.completion.chunk",
created: 1741000456,
model: "hunyuan-turbo",
system_fingerprint: "",
choices: [
{
index: 0,
delta: { role: "assistant", content: "为您" },
finish_reason: null
}
],
note: "以上内容为AI生成,不代表开发者立场,请勿删除或修改本标记",
search_info: {
search_results: [
{
index: 1,
title: "公务员考试时政热点【2025年3月3日】_公务员考试网_华图教育",
url: "http://www.huatu.com/2025/0303/2803685.html",
icon: "https://hunyuan-img-1251316161.cos.ap-guangzhou.myqcloud.com/%2Fpublic/img/63ce96deffe0119827f12deaa5ffe7ef.jpg",
text: "华图教育官网"
},
{
index: 2,
title: "外交部新闻(2025年3月3日)",
url: "https://view.inews.qq.com/a/20250303A02NLC00?scene=qqsearch",
icon: "https://hunyuan-img-1251316161.cos.ap-guangzhou.myqcloud.com/%2Fpublic/img/00ce40298870d1accb7920d641152722.jpg",
text: "腾讯网"
}
]
}
},
{
id: "939fbdb8dbb9b4c5944cbbe687c977c2",
object: "chat.completion.chunk",
created: 1741000456,
model: "hunyuan-turbo",
system_fingerprint: "",
choices: [
{
index: 0,
delta: { role: "assistant", content: "找到" },
finish_reason: null
}
],
note: "以上内容为AI生成,不代表开发者立场,请勿删除或修改本标记",
search_info: {
search_results: [
{
index: 1,
title: "公务员考试时政热点【2025年3月3日】_公务员考试网_华图教育",
url: "http://www.huatu.com/2025/0303/2803685.html",
icon: "https://hunyuan-img-1251316161.cos.ap-guangzhou.myqcloud.com/%2Fpublic/img/63ce96deffe0119827f12deaa5ffe7ef.jpg",
text: "华图教育官网"
},
{
index: 2,
title: "外交部新闻(2025年3月3日)",
url: "https://view.inews.qq.com/a/20250303A02NLC00?scene=qqsearch",
icon: "https://hunyuan-img-1251316161.cos.ap-guangzhou.myqcloud.com/%2Fpublic/img/00ce40298870d1accb7920d641152722.jpg",
text: "腾讯网"
}
]
}
}
];
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: 939fbdb8dbb9b4c5944cbbe687c977c2',
'event: grounding',
'data: {"citations":[{"title":"公务员考试时政热点【2025年3月3日】_公务员考试网_华图教育","url":"http://www.huatu.com/2025/0303/2803685.html"},{"title":"外交部新闻(2025年3月3日)","url":"https://view.inews.qq.com/a/20250303A02NLC00?scene=qqsearch"}]}\n',
'id: 939fbdb8dbb9b4c5944cbbe687c977c2',
'event: text',
'data: "为您"\n',
'id: 939fbdb8dbb9b4c5944cbbe687c977c2',
'event: text',
'data: "找到"\n',
].map((line) => `${line}\n`),
);
expect((await reader.read()).done).toBe(true);
});
});
});
+5 -4
View File
@@ -15,14 +15,15 @@ export const LobeHunyuanAI = LobeOpenAICompatibleFactory({
return {
...rest,
stream: true,
...(enabledSearch && {
/*
citation: true,
enable_multimedia: true,
search_info: true
*/
enable_enhancement: true,
/*
enable_multimedia: true,
*/
enable_speed_search: process.env.HUNYUAN_ENABLE_SPEED_SEARCH === '1',
search_info: true,
}),
} as any;
},
+1 -1
View File
@@ -35,7 +35,7 @@ export const LobeQwenAI = LobeOpenAICompatibleFactory({
: presence_penalty !== undefined && presence_penalty >= -2 && presence_penalty <= 2
? presence_penalty
: undefined,
stream: !payload.tools,
stream: true,
temperature:
temperature !== undefined && temperature >= 0 && temperature < 2
? temperature
+27 -11
View File
@@ -127,19 +127,35 @@ 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 (!streamContext?.returnedCitation) {
const citations =
// in Perplexity api, the citation is in every chunk, but we only need to return it once
('citations' in chunk && chunk.citations) ||
// in Hunyuan api, the citation is in every chunk
('search_info' in chunk && (chunk.search_info as any)?.search_results) ||
// in Wenxin api, the citation is in the first and last chunk
('search_results' in chunk && chunk.search_results);
if ('citations' in chunk && !!chunk.citations && !streamContext?.returnedPplxCitation) {
streamContext.returnedPplxCitation = true;
if (citations) {
streamContext.returnedCitation = 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: 'grounding' },
{ data: content, id: chunk.id, type: 'text' },
];
return [
{
data: {
citations: (citations as any[]).map(
(item) =>
({
title: typeof item === 'string' ? item : item.title,
url: typeof item === 'string' ? item : item.url,
}) as CitationItem
),
},
id: chunk.id,
type: 'grounding',
},
{ data: content, id: chunk.id, type: 'text' },
];
}
}
return { data: content, id: chunk.id, type: 'text' };
@@ -9,9 +9,10 @@ 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
* this flag is used to check if the pplx citation is returned,and then not return it again.
* Same as Hunyuan and Wenxin
*/
returnedPplxCitation?: boolean;
returnedCitation?: boolean;
thinking?: {
id: string;
name: string;
+113
View File
@@ -0,0 +1,113 @@
// @vitest-environment node
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
import { LobeOpenAICompatibleRuntime, ModelProvider } from '@/libs/agent-runtime';
import { testProvider } from '@/libs/agent-runtime/providerTestUtils';
import { LobeWenxinAI } from './index';
testProvider({
Runtime: LobeWenxinAI,
provider: ModelProvider.Wenxin,
defaultBaseURL: 'https://qianfan.baidubce.com/v2',
chatDebugEnv: 'DEBUG_WENXIN_CHAT_COMPLETION',
chatModel: 'ernie-speed-128k',
});
// Mock the console.error to avoid polluting test output
vi.spyOn(console, 'error').mockImplementation(() => {});
let instance: LobeOpenAICompatibleRuntime;
beforeEach(() => {
instance = new LobeWenxinAI({ apiKey: 'test' });
// 使用 vi.spyOn 来模拟 chat.completions.create 方法
vi.spyOn(instance['client'].chat.completions, 'create').mockResolvedValue(
new ReadableStream() as any,
);
});
describe('LobeWenxinAI', () => {
describe('chat', () => {
it('should with search citations', async () => {
const data = [
{
id: "as-bhrxwy5fq1",
object: "chat.completion.chunk",
created: 1741000028,
model: "ernie-4.0-8k-latest",
choices: [
{
index: 0,
delta: { content: "今天是**", role: "assistant" },
flag: 0
}
],
search_results: [
{ index: 1, url: "http://www.mnw.cn/news/shehui/", title: "社会新闻" },
{ index: 2, url: "https://www.chinanews.com.cn/sh/2025/03-01/10376297.shtml", title: "中越边民共庆“春龙节”" },
{ index: 3, url: "https://www.chinanews.com/china/index.shtml", title: "中国新闻网_时政" }
]
},
{
id: "as-bhrxwy5fq1",
object: "chat.completion.chunk",
created: 1741000028,
model: "ernie-4.0-8k-latest",
choices: [
{
index: 0,
delta: { content: "20" },
flag: 0
}
]
}
];
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: as-bhrxwy5fq1',
'event: grounding',
'data: {"citations":[{"title":"社会新闻","url":"http://www.mnw.cn/news/shehui/"},{"title":"中越边民共庆“春龙节”","url":"https://www.chinanews.com.cn/sh/2025/03-01/10376297.shtml"},{"title":"中国新闻网_时政","url":"https://www.chinanews.com/china/index.shtml"}]}\n',
'id: as-bhrxwy5fq1',
'event: text',
'data: "今天是**"\n',
'id: as-bhrxwy5fq1',
'event: text',
'data: "20"\n',
].map((line) => `${line}\n`),
);
expect((await reader.read()).done).toBe(true);
});
});
});
+17
View File
@@ -3,6 +3,23 @@ import { LobeOpenAICompatibleFactory } from '../utils/openaiCompatibleFactory';
export const LobeWenxinAI = LobeOpenAICompatibleFactory({
baseURL: 'https://qianfan.baidubce.com/v2',
chatCompletion: {
handlePayload: (payload) => {
const { enabledSearch, ...rest } = payload;
return {
...rest,
stream: true,
...(enabledSearch && {
web_search: {
enable: true,
enable_citation: true,
enable_trace: true,
}
}),
} as any;
},
},
debug: {
chatCompletion: () => process.env.DEBUG_WENXIN_CHAT_COMPLETION === '1',
},