diff --git a/packages/tool-bailian-kb/src/tools.ts b/packages/tool-bailian-kb/src/tools.ts new file mode 100644 index 0000000..01f2d64 --- /dev/null +++ b/packages/tool-bailian-kb/src/tools.ts @@ -0,0 +1,237 @@ +/** + * The three model-facing knowledge tools. Schemas are static per deployment: a configured + * defaultAgentId downgrades agent_id to optional at build time (never a runtime fallback chain). + */ + +import { defineTool } from '@deepseek-ai/dsh-tools' +import type { SearchRequest, SearchResponse } from './api-types.js' +import { KbApiError, type KbClient } from './client.js' +import { consumeChatStream } from './chat.js' +import { KB_PATHS } from './endpoints.js' +import { listServices } from './services.js' + +/** Client-side chunk cap applied when the model omits top_k. */ +const DEFAULT_TOP_K = 5 + +export interface KbToolDeps { + client: KbClient + defaultAgentId?: string + chatTimeoutMs: number +} + +/** Format one service list into the error-hint / result text form. */ +function formatServices(services: { agent_id: string; name: string; scene: string; status: string }[]): string { + return services.map(s => `${s.agent_id} (${s.name}, scene=${s.scene}, ${s.status})`).join('; ') +} + +/** + * Append the current service list to a client error so the model can correct + * an invalid agent_id in one step. Auth failures (401/403) keep their own message. + * @param client - the shared knowledge API client used for best-effort discovery. + * @param err - the failure being enriched; always rethrown. + * @returns never; the original or enriched error is thrown. + */ +async function withServiceHint(client: KbClient, err: unknown): Promise { + if (err instanceof KbApiError && err.status !== undefined && err.status >= 400 && err.status !== 401 && err.status !== 403) { + let hint: string | undefined + try { + const { services } = await listServices(client, {}) + if (services.length > 0) hint = formatServices(services) + } catch { /* discovery is best-effort; the original error already carries the failure */ } + if (hint !== undefined) throw new KbApiError(`${err.message}. Available services: ${hint}`, err.status) + } + throw err +} + +/** + * Build the three tool definitions over one shared client. + * @param deps - client plus the deployment's explicit pinning and timeout choices. + * @returns definitions ready for `ctx.tools.register()`. + */ +export function createKbTools(deps: KbToolDeps) { + const { client, defaultAgentId, chatTimeoutMs } = deps + const agentIdParam = { + type: 'string' as const, + ...(defaultAgentId === undefined ? { required: true as const } : {}), + description: defaultAgentId === undefined + ? 'Retrieval/Q&A service id (find one via kb_service_list).' + : 'Retrieval/Q&A service id; omit to use this deployment\'s default service.', + } + const resolveAgentId = (supplied: string | undefined): string => { + const agentId = supplied ?? defaultAgentId + if (agentId === undefined) throw new Error('agent_id is required: discover services with kb_service_list') + return agentId + } + + const serviceList = defineTool({ + name: 'kb_service_list', + description: + 'List the Bailian knowledge retrieval/Q&A services available in this workspace. ' + + 'Each entry names the service id (agent_id) to pass to kb_search (scene=search) or kb_chat (scene=chat), ' + + 'its bound knowledge bases, and its status (prefer deployed). ' + + 'Omit scene to see both kinds; narrow large workspaces with name_filter.', + parameters: { + scene: { type: 'string', enum: ['chat', 'search'], description: 'Only list services for this scene; omitted lists both.' }, + name_filter: { type: 'string', description: 'Fuzzy match on the service name.' }, + }, + output: { + schema: { + type: 'object', + additionalProperties: false, + properties: { + services: { + type: 'array', + required: true, + items: { + type: 'object', + additionalProperties: false, + properties: { + agent_id: { type: 'string', required: true }, + name: { type: 'string', required: true }, + scene: { type: 'string', required: true }, + status: { type: 'string', required: true }, + knowledge_bases: { type: 'array', required: true, items: { type: 'string' } }, + }, + }, + }, + total: { type: 'integer', required: true }, + truncated: { type: 'boolean', required: true }, + }, + }, + render: (_args, value) => [{ + type: 'text', + text: value.services.length === 0 + ? 'No knowledge services found.' + : `${value.services.length} service(s): ${formatServices(value.services)}` + + (value.truncated ? ` — listed first ${value.services.length} of ${value.total}; narrow with name_filter.` : ''), + }], + }, + async execute(args) { + return await listServices(client, { + ...(args.scene === 'chat' || args.scene === 'search' ? { scene: args.scene } : {}), + ...(args.name_filter ? { nameFilter: args.name_filter } : {}), + }) + }, + presentCall: args => ({ card: 'generic', title: 'List knowledge services', kind: 'other', rawInput: args }), + }) + + const search = defineTool({ + name: 'kb_search', + description: + 'Semantic search over a Bailian knowledge base. Returns raw knowledge chunks with scores and source ' + + 'references for you to verify, cite, or combine with other context. Retrieval scope and strategy ' + + '(multi-KB weighting, routing, reranking) come from the service configuration. ' + + 'top_k caps how many chunks return (client-side cut of the score-ranked results). ' + + 'Use kb_chat instead when the user question can be answered by the knowledge base alone.', + parameters: { + query: { type: 'string', required: true, description: 'Search query text.' }, + agent_id: agentIdParam, + top_k: { type: 'integer', description: `Maximum chunks to return; defaults to ${DEFAULT_TOP_K}.` }, + images: { type: 'array', items: { type: 'string' }, description: 'Image URLs for multimodal retrieval.' }, + }, + output: { + schema: { + type: 'object', + additionalProperties: false, + properties: { + chunks: { + type: 'array', + required: true, + items: { + type: 'object', + additionalProperties: false, + properties: { + text: { type: 'string', required: true }, + score: { type: 'number', required: true }, + doc_name: { type: 'string' }, + doc_id: { type: 'string' }, + title: { type: 'string' }, + }, + }, + }, + total: { type: 'integer', required: true }, + }, + }, + render: (_args, value) => [{ + type: 'text', + text: value.chunks.length === 0 + ? 'No matching knowledge chunks.' + : value.chunks.map((c, i) => `[${i + 1}] (score ${c.score.toFixed(2)}${c.doc_name ? `, ${c.doc_name}` : ''}) ${c.text}`).join('\n'), + }], + }, + async execute(args) { + const topK = args.top_k ?? DEFAULT_TOP_K + const body: SearchRequest = { + query: args.query, + agent_id: resolveAgentId(args.agent_id), + ...(client.agentVersion ? { agent_version: client.agentVersion } : {}), + ...(args.images && args.images.length > 0 ? { images: args.images } : {}), + } + const res = await client.postJson(KB_PATHS.search, body).catch(err => withServiceHint(client, err)) + const nodes = (res.data?.nodes ?? []).slice(0, topK) + return { + chunks: nodes.map(n => ({ + text: n.text, + score: n.score, + ...(typeof n.metadata?.doc_name === 'string' ? { doc_name: n.metadata.doc_name } : {}), + ...(typeof n.metadata?.doc_id === 'string' ? { doc_id: n.metadata.doc_id } : {}), + ...(typeof n.metadata?.title === 'string' ? { title: n.metadata.title } : {}), + })), + total: res.data?.total ?? nodes.length, + } + }, + presentCall: args => ({ card: 'generic', title: 'Search knowledge base', kind: 'search', rawInput: args }), + }) + + const chat = defineTool({ + name: 'kb_chat', + description: + 'Ask the knowledge base directly and get a complete, domain-tuned answer from a specialized RAG pipeline ' + + '(multi-round retrieval + reranking + grounded generation). For knowledge Q&A this typically outperforms ' + + 'searching and synthesizing yourself when the question can be answered by the knowledge base alone; ' + + 'use kb_search instead when you need raw chunks to verify, cite, or combine with other work. ' + + 'The pipeline runs an internal analysis/retrieval loop and may take a few minutes.', + parameters: { + message: { type: 'string', required: true, description: 'The question to ask.' }, + agent_id: agentIdParam, + }, + output: { + schema: { + type: 'object', + additionalProperties: false, + properties: { + answer: { type: 'string', required: true }, + request_id: { type: 'string' }, + }, + }, + render: (_args, value) => [{ type: 'text', text: value.answer.length === 0 ? '(empty answer)' : value.answer }], + }, + async execute(args) { + const body = { + input: { messages: [{ role: 'user' as const, content: args.message }] }, + parameters: { agent_options: { + agent_id: resolveAgentId(args.agent_id), + ...(client.agentVersion ? { agent_version: client.agentVersion } : {}), + } }, + stream: true as const, + } + let res: Response + try { + res = await client.postSse(KB_PATHS.chat, body, AbortSignal.timeout(chatTimeoutMs)) + } catch (err) { + if (err instanceof Error && err.name === 'TimeoutError') { + throw new Error( + `knowledge chat timed out after ${chatTimeoutMs}ms; the pipeline runs a multi-round retrieval loop ` + + 'and long questions can exceed the deployment timeout. Retry, or use kb_search for raw chunks instead.', + ) + } + return await withServiceHint(client, err) + } + const { answer, requestId } = await consumeChatStream(res) + return { answer, ...(requestId ? { request_id: requestId } : {}) } + }, + presentCall: args => ({ card: 'generic', title: 'Ask knowledge base (may take a few minutes)', kind: 'fetch', rawInput: args }), + }) + + return [serviceList, search, chat] +} diff --git a/packages/tool-bailian-kb/tests/tools.test.ts b/packages/tool-bailian-kb/tests/tools.test.ts new file mode 100644 index 0000000..1ac7d1f --- /dev/null +++ b/packages/tool-bailian-kb/tests/tools.test.ts @@ -0,0 +1,81 @@ +import { describe, expect, it, vi } from 'vitest' +import { KbApiError, KbClient } from '../src/client.js' +import { createKbTools } from '../src/tools.js' + +const EXEC = {} as never + +function toolsWith(postJson: unknown, postSse?: unknown, defaultAgentId?: string) { + const client = { postJson, postSse, agentVersion: undefined } as unknown as KbClient + const list = createKbTools({ client, ...(defaultAgentId ? { defaultAgentId } : {}), chatTimeoutMs: 1000 }) + const byName = Object.fromEntries(list.map(t => [t.name, t])) + return { byName, list } +} + +const searchResponse = { + request_id: 'r1', + data: { total: 3, nodes: [ + { score: 0.9, text: 'A', metadata: { doc_name: 'd1' } }, + { score: 0.8, text: 'B', metadata: {} }, + { score: 0.7, text: 'C', metadata: {} }, + ] }, +} + +describe('createKbTools', () => { + it('registers exactly kb_service_list, kb_search, kb_chat', () => { + const { list } = toolsWith(vi.fn()) + expect(list.map(t => t.name).sort()).toEqual(['kb_chat', 'kb_search', 'kb_service_list']) + }) + + it('kb_search truncates nodes client-side to top_k and never sends top_k to the server', async () => { + const postJson = vi.fn(async (_path: string, _body: unknown) => searchResponse) + const { byName } = toolsWith(postJson) + const out = await byName.kb_search!.execute({ query: 'q', agent_id: 'aid-1', top_k: 2 }, EXEC) as { chunks: unknown[] } + expect(out.chunks).toHaveLength(2) + const body = postJson.mock.calls[0]![1] as Record + expect(body).not.toHaveProperty('top_k') + expect(body.agent_id).toBe('aid-1') + }) + + it('agent_id is required without defaultAgentId and optional with one', () => { + const withoutDefault = toolsWith(vi.fn()).byName.kb_search! + const withDefault = toolsWith(vi.fn(), undefined, 'aid-fixed').byName.kb_search! + // defineTool compiles the spec into JSON Schema: requiredness lives in the top-level `required` array. + const requiredList = (tool: { parameters: Record }) => + (tool.parameters.required ?? []) as string[] + expect(requiredList(withoutDefault)).toContain('agent_id') + expect(requiredList(withDefault)).not.toContain('agent_id') + }) + + it('kb_search falls back to defaultAgentId as an explicit resolve step', async () => { + const postJson = vi.fn(async (_path: string, _body: unknown) => searchResponse) + const { byName } = toolsWith(postJson, undefined, 'aid-fixed') + await byName.kb_search!.execute({ query: 'q' }, EXEC) + expect((postJson.mock.calls[0]![1] as Record).agent_id).toBe('aid-fixed') + }) + + it('a 4xx failure appends the current service list to the error', async () => { + const postJson = vi.fn(async (path: string) => { + if (path === '/api/v1/indices/knowledge/search') throw new KbApiError('agent not found', 400) + return { data: { total_count: 1, rows: [{ agent_id: 'aid-9', agent_name: 'faq', agent_scene: 'search', agent_status: 'deployed' }] } } + }) + const { byName } = toolsWith(postJson) + const err = await byName.kb_search!.execute({ query: 'q', agent_id: 'bad' }, EXEC).catch((e: unknown) => e) + expect((err as Error).message).toContain('aid-9') + }) + + it('kb_chat buffers the SSE stream into one answer', async () => { + const sse = 'data: {"output":{"choices":[{"message":{"content":"hi"},"finish_reason":"stop"}]},"request_id":"r2"}\n\ndata: [DONE]\n\n' + const postSse = vi.fn(async () => new Response(sse, { status: 200 })) + const { byName } = toolsWith(vi.fn(), postSse) + const out = await byName.kb_chat!.execute({ message: 'q', agent_id: 'aid-1' }, EXEC) as { answer: string } + expect(out.answer).toBe('hi') + }) + + it('kb_chat translates a timeout into retry-or-search guidance', async () => { + const timeout = Object.assign(new Error('operation timed out'), { name: 'TimeoutError' }) + const postSse = vi.fn(async () => { throw timeout }) + const { byName } = toolsWith(vi.fn(), postSse) + const err = await byName.kb_chat!.execute({ message: 'q', agent_id: 'aid-1' }, EXEC).catch((e: unknown) => e) + expect((err as Error).message).toMatch(/timed out.*kb_search/s) + }) +})