Dockery/Trilium
1
0
Fork 0
mirror of https://github.com/zadam/trilium.git synced 2026-05-09 12:45:48 +02:00
Code Issues Packages Projects Releases Activity

refactor(llm): get rid of context-aware tools

Browse Source
This commit is contained in:
Elian Doran
2026-04-03 18:48:59 +03:00
parent 57a299de8f
commit df0efc39d5
2 changed files with 9 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
apps/server/src/services/llm/providers/base_provider.ts
View File

@@ -10,7 +10,6 @@ import type { LlmMessage } from "@triliumnext/commons";
import becca from "../../../becca/becca.js";
import { getSkillsSummary } from "../skills/index.js";
import { allToolRegistries } from "../tools/index.js";
import type { ToolContext } from "../tools/tool_registry.js";
import type { LlmProvider, LlmProviderConfig, ModelInfo, ModelPricing, StreamResult } from "../types.js";
const DEFAULT_MAX_TOKENS = 8096;
@@ -130,11 +129,8 @@ export abstract class BaseProvider implements LlmProvider {
}
if (config.enableNoteTools) {
const context: ToolContext | undefined = config.contextNoteId
? { contextNoteId: config.contextNoteId }
: undefined;
for (const registry of allToolRegistries) {
Object.assign(tools, registry.toToolSet(context));
Object.assign(tools, registry.toToolSet());
}
}

49
apps/server/src/services/llm/tools/tool_registry.ts
View File

@@ -1,7 +1,6 @@
/**
* Lightweight wrapper around AI tool definitions that carries extra metadata
* (e.g. `mutates`, `needsContext`) while remaining compatible with the Vercel
* AI SDK ToolSet.
* (e.g. `mutates`) while remaining compatible with the Vercel AI SDK ToolSet.
*
* Each tool module calls `defineTools({ ... })` to declare its tools.
* Consumers can then:
@@ -13,32 +12,14 @@ import { tool } from "ai";
import type { z } from "zod";
import type { ToolSet } from "ai";
/** Context passed to tools that declare `needsContext: true`. */
export interface ToolContext {
contextNoteId: string;
}
interface ToolDefinitionBase {
export interface ToolDefinition {
description: string;
inputSchema: z.ZodType;
/** Whether this tool modifies data (needs CLS + transaction wrapping). */
mutates?: boolean;
}
/** A tool that does not require a note context. */
export interface StaticToolDefinition extends ToolDefinitionBase {
needsContext?: false;
execute: (args: any) => Promise<unknown>;
}
/** A tool that requires a note context (e.g. "current note"). */
export interface ContextToolDefinition extends ToolDefinitionBase {
needsContext: true;
execute: (args: any, context: ToolContext) => Promise<unknown>;
}
export type ToolDefinition = StaticToolDefinition | ContextToolDefinition;
/**
* A named collection of tool definitions that can be iterated or converted
* to an AI SDK ToolSet.
@@ -53,28 +34,15 @@ export class ToolRegistry implements Iterable<[string, ToolDefinition]> {
/**
* Convert to an AI SDK ToolSet for use with the LLM chat providers.
*
* If `context` is provided, context-aware tools are included with the
* context bound into their execute function. Otherwise they are skipped.
*/
toToolSet(context?: ToolContext): ToolSet {
toToolSet(): ToolSet {
const set: ToolSet = {};
for (const [name, def] of this) {
if (def.needsContext) {
if (!context) continue;
const boundExecute = (args: any) => def.execute(args, context);
set[name] = tool({
description: def.description,
inputSchema: def.inputSchema,
execute: boundExecute
});
} else {
set[name] = tool({
description: def.description,
inputSchema: def.inputSchema,
execute: def.execute
});
}
set[name] = tool({
description: def.description,
inputSchema: def.inputSchema,
execute: def.execute
});
}
return set;
}
@@ -86,7 +54,6 @@ export class ToolRegistry implements Iterable<[string, ToolDefinition]> {
* ```ts
* export const noteTools = defineTools({
* search_notes: { description: "...", inputSchema: z.object({...}), execute: async (args) => {...} },
* get_current_note: { description: "...", inputSchema: z.object({}), execute: async (args, ctx) => {...}, needsContext: true },
* });
* ```
*/