From a5bf1736a9be1095f27bb0f400bf1613f3a4d6ee Mon Sep 17 00:00:00 2001 From: Travis Fischer Date: Sun, 2 Jun 2024 18:31:09 -0500 Subject: [PATCH] feat: remove AIToolSet in favor of just AIFunctionSet --- examples/openai/weather.ts | 8 +- package.json | 1 - src/ai-function-set.ts | 22 ++-- src/ai-tool-set.ts | 100 ------------------ ...ion.test.ts => create-ai-function.test.ts} | 2 +- src/fns.ts | 12 +-- src/index.ts | 2 +- src/sdks/ai-sdk.ts | 10 +- src/sdks/dexter.ts | 10 +- src/sdks/genkit.ts | 8 +- src/sdks/langchain.ts | 8 +- src/services/serpapi-client.ts | 4 +- src/services/serper-client.ts | 4 +- src/services/weather-client.test.ts | 11 -- src/services/weather-client.ts | 4 +- src/types.ts | 5 + 16 files changed, 46 insertions(+), 165 deletions(-) delete mode 100644 src/ai-tool-set.ts rename src/{ai-function.test.ts => create-ai-function.test.ts} (94%) delete mode 100644 src/services/weather-client.test.ts diff --git a/examples/openai/weather.ts b/examples/openai/weather.ts index 7c86588..853a985 100644 --- a/examples/openai/weather.ts +++ b/examples/openai/weather.ts @@ -18,22 +18,20 @@ async function main() { { role: 'user', content: 'What is the weather in San Francisco?' } ] - const tools = weather.tools - { // First call to OpenAI to invoke the weather tool const res = await openai.chat.completions.create({ messages, model: 'gpt-4o', temperature: 0, - tools: tools.specs, + tools: weather.functions.toolSpecs, tool_choice: 'required' }) const message = res.choices[0]?.message! console.log(JSON.stringify(message, null, 2)) assert(message.tool_calls?.[0]?.function?.name === 'get_current_weather') - const fn = tools.get('get_current_weather')!.function + const fn = weather.functions.get('get_current_weather')! assert(fn) const toolParams = message.tool_calls[0].function.arguments @@ -53,7 +51,7 @@ async function main() { messages, model: 'gpt-4o', temperature: 0, - tools: tools.specs + tools: weather.functions.toolSpecs }) const message = res.choices[0].message console.log(JSON.stringify(message, null, 2)) diff --git a/package.json b/package.json index ba061f3..4479219 100644 --- a/package.json +++ b/package.json @@ -45,7 +45,6 @@ "scripts": { "preinstall": "npx only-allow pnpm", "build": "tsc", - "dev": "tsup --watch", "clean": "del dist", "prebuild": "run-s clean", "predev": "run-s clean", diff --git a/src/ai-function-set.ts b/src/ai-function-set.ts index c011e68..50aaec5 100644 --- a/src/ai-function-set.ts +++ b/src/ai-function-set.ts @@ -1,13 +1,19 @@ -import type { AIToolSet } from './ai-tool-set.js' import type * as types from './types.ts' +import { AIFunctionsProvider } from './fns.js' export class AIFunctionSet implements Iterable { protected readonly _map: Map - constructor(functions?: readonly types.AIFunction[]) { - this._map = new Map( - functions ? functions.map((fn) => [fn.spec.name, fn]) : null + constructor(aiFunctionLikeObjects?: types.AIFunctionLike[]) { + const fns = aiFunctionLikeObjects?.flatMap((fn) => + fn instanceof AIFunctionsProvider + ? [...fn.functions] + : fn instanceof AIFunctionSet + ? [...fn] + : [fn] ) + + this._map = new Map(fns ? fns.map((fn) => [fn.spec.name, fn]) : null) } get size(): number { @@ -76,12 +82,4 @@ export class AIFunctionSet implements Iterable { [Symbol.iterator](): Iterator { return this.entries } - - static fromAIToolSet(tools: AIToolSet): AIFunctionSet { - return new AIFunctionSet( - Array.from(tools) - .filter((tool) => tool.spec.type === 'function') - .map((tool) => tool.function) - ) - } } diff --git a/src/ai-tool-set.ts b/src/ai-tool-set.ts deleted file mode 100644 index d7b21c6..0000000 --- a/src/ai-tool-set.ts +++ /dev/null @@ -1,100 +0,0 @@ -import type * as types from './types.js' -import { AIFunctionSet } from './ai-function-set.js' - -export class AIToolSet implements Iterable { - protected _map: Map - - constructor(tools?: readonly types.AITool[]) { - this._map = new Map( - tools ? tools.map((tool) => [tool.spec.function.name, tool]) : [] - ) - } - - get size(): number { - return this._map.size - } - - add(tool: types.AITool): this { - this._map.set(tool.spec.function.name, tool) - return this - } - - get(name: string): types.AITool | undefined { - return this._map.get(name) - } - - set(name: string, tool: types.AITool): this { - this._map.set(name, tool) - return this - } - - has(name: string): boolean { - return this._map.has(name) - } - - clear(): void { - this._map.clear() - } - - delete(name: string): boolean { - return this._map.delete(name) - } - - pick(...keys: string[]): AIToolSet { - const keysToIncludeSet = new Set(keys) - return new AIToolSet( - Array.from(this).filter((tool) => - keysToIncludeSet.has(tool.spec.function.name) - ) - ) - } - - omit(...keys: string[]): AIToolSet { - const keysToExcludeSet = new Set(keys) - return new AIToolSet( - Array.from(this).filter( - (tool) => !keysToExcludeSet.has(tool.spec.function.name) - ) - ) - } - - map(fn: (fn: types.AITool) => T): T[] { - return [...this.entries].map(fn) - } - - get functionSpecs(): types.AIFunctionSpec[] { - return this.map((fn) => fn.function.spec) - } - - get specs(): types.AIToolSpec[] { - return this.map((fn) => fn.spec) - } - - get entries(): IterableIterator { - return this._map.values() - } - - [Symbol.iterator](): Iterator { - return this.entries - } - - static fromAIFunctionSet(functions: AIFunctionSet): AIToolSet { - return new AIToolSet( - Array.from(functions).map((fn) => ({ - function: fn, - spec: { - type: 'function' as const, - function: fn.spec - } - })) - ) - } - - static fromFunctions(functions: types.AIFunction[]): AIToolSet { - return AIToolSet.fromAIFunctionSet(new AIFunctionSet(functions)) - } - - static fromTools(tools: types.AITool[]): AIToolSet { - return new AIToolSet(tools) - } -} diff --git a/src/ai-function.test.ts b/src/create-ai-function.test.ts similarity index 94% rename from src/ai-function.test.ts rename to src/create-ai-function.test.ts index 8b5262a..94c3a62 100644 --- a/src/ai-function.test.ts +++ b/src/create-ai-function.test.ts @@ -1,7 +1,7 @@ import { describe, expect, it } from 'vitest' import { z } from 'zod' -import { createAIFunction } from './ai-function.js' +import { createAIFunction } from './create-ai-function.js' const fullName = createAIFunction( { diff --git a/src/fns.ts b/src/fns.ts index cd8d2ef..1543e15 100644 --- a/src/fns.ts +++ b/src/fns.ts @@ -4,7 +4,6 @@ import type { z } from 'zod' import type * as types from './types.js' import { AIFunctionSet } from './ai-function-set.js' -import { AIToolSet } from './ai-tool-set.js' import { createAIFunction } from './create-ai-function.js' import { assert } from './utils.js' @@ -15,18 +14,9 @@ export interface Invocable { methodName: string } -export abstract class AIToolsProvider { - private _tools?: AIToolSet +export abstract class AIFunctionsProvider { private _functions?: AIFunctionSet - get tools(): AIToolSet { - if (!this._tools) { - this._tools = AIToolSet.fromAIFunctionSet(this.functions) - } - - return this._tools - } - get functions(): AIFunctionSet { if (!this._functions) { const metadata = this.constructor[Symbol.metadata] diff --git a/src/index.ts b/src/index.ts index 2032dad..269386e 100644 --- a/src/index.ts +++ b/src/index.ts @@ -1,5 +1,5 @@ export * from './ai-function-set.js' -export * from './ai-tool-set.js' +export * from './create-ai-function.js' export * from './create-ai-function.js' export * from './errors.js' export * from './fns.js' diff --git a/src/sdks/ai-sdk.ts b/src/sdks/ai-sdk.ts index 870a987..7edaa5f 100644 --- a/src/sdks/ai-sdk.ts +++ b/src/sdks/ai-sdk.ts @@ -1,17 +1,17 @@ import { tool } from 'ai' -import type { AIFunctionSet } from '../ai-function-set.js' -import { AIToolsProvider } from '../fns.js' +import type { AIFunctionLike } from '../types.js' +import { AIFunctionSet } from '../ai-function-set.js' /** * Converts a set of Agentic stdlib AI functions to an object compatible with * the Vercel AI SDK's `tools` parameter. */ -export function createAISDKTools(tools: AIToolsProvider | AIFunctionSet) { - const fns = tools instanceof AIToolsProvider ? tools.functions : tools +export function createAISDKTools(...aiFunctionLikeTools: AIFunctionLike[]) { + const fns = new AIFunctionSet(aiFunctionLikeTools) return Object.fromEntries( - [...fns].map((fn) => [ + fns.map((fn) => [ fn.spec.name, tool({ description: fn.spec.description, diff --git a/src/sdks/dexter.ts b/src/sdks/dexter.ts index 14110b5..8e8df06 100644 --- a/src/sdks/dexter.ts +++ b/src/sdks/dexter.ts @@ -1,14 +1,16 @@ import { createAIFunction } from '@dexaai/dexter' -import type { AIFunctionSet } from '../ai-function-set.js' -import { AIToolsProvider } from '../fns.js' +import type { AIFunctionLike } from '../types.js' +import { AIFunctionSet } from '../ai-function-set.js' /** * Converts a set of Agentic stdlib AI functions to an array of Dexter- * compatible AI functions. */ -export function createDexterFunctions(input: AIToolsProvider | AIFunctionSet) { - const fns = input instanceof AIToolsProvider ? input.functions : input +export function createDexterFunctions( + ...aiFunctionLikeTools: AIFunctionLike[] +) { + const fns = new AIFunctionSet(aiFunctionLikeTools) return fns.map((fn) => createAIFunction( diff --git a/src/sdks/genkit.ts b/src/sdks/genkit.ts index c38b53a..a9c454b 100644 --- a/src/sdks/genkit.ts +++ b/src/sdks/genkit.ts @@ -1,15 +1,15 @@ import { defineTool } from '@genkit-ai/ai' import { z } from 'zod' -import type { AIFunctionSet } from '../ai-function-set.js' -import { AIToolsProvider } from '../fns.js' +import type { AIFunctionLike } from '../types.js' +import { AIFunctionSet } from '../ai-function-set.js' /** * Converts a set of Agentic stdlib AI functions to an array of Genkit- * compatible tools. */ -export function createGenkitTools(input: AIToolsProvider | AIFunctionSet) { - const fns = input instanceof AIToolsProvider ? input.functions : input +export function createGenkitTools(...aiFunctionLikeTools: AIFunctionLike[]) { + const fns = new AIFunctionSet(aiFunctionLikeTools) return fns.map((fn) => defineTool( diff --git a/src/sdks/langchain.ts b/src/sdks/langchain.ts index b5650dc..b273c91 100644 --- a/src/sdks/langchain.ts +++ b/src/sdks/langchain.ts @@ -1,15 +1,15 @@ import { DynamicStructuredTool } from '@langchain/core/tools' -import type { AIFunctionSet } from '../ai-function-set.js' -import { AIToolsProvider } from '../fns.js' +import type { AIFunctionLike } from '../types.js' +import { AIFunctionSet } from '../ai-function-set.js' import { stringifyForModel } from '../stringify-for-model.js' /** * Converts a set of Agentic stdlib AI functions to an array of LangChain- * compatible tools. */ -export function createLangChainTools(input: AIToolsProvider | AIFunctionSet) { - const fns = input instanceof AIToolsProvider ? input.functions : input +export function createLangChainTools(...aiFunctionLikeTools: AIFunctionLike[]) { + const fns = new AIFunctionSet(aiFunctionLikeTools) return fns.map( (fn) => diff --git a/src/services/serpapi-client.ts b/src/services/serpapi-client.ts index af5f130..f351f07 100644 --- a/src/services/serpapi-client.ts +++ b/src/services/serpapi-client.ts @@ -1,7 +1,7 @@ import defaultKy, { type KyInstance } from 'ky' import { z } from 'zod' -import { aiFunction, AIToolsProvider } from '../fns.js' +import { aiFunction, AIFunctionsProvider } from '../fns.js' import { assert, getEnv } from '../utils.js' /** @@ -633,7 +633,7 @@ export namespace serpapi { * * @see https://serpapi.com/search-api */ -export class SerpAPIClient extends AIToolsProvider { +export class SerpAPIClient extends AIFunctionsProvider { protected ky: KyInstance protected apiKey: string protected apiBaseUrl: string diff --git a/src/services/serper-client.ts b/src/services/serper-client.ts index 064f56d..556dfb8 100644 --- a/src/services/serper-client.ts +++ b/src/services/serper-client.ts @@ -1,7 +1,7 @@ import defaultKy, { type KyInstance } from 'ky' import { z } from 'zod' -import { aiFunction, AIToolsProvider } from '../fns.js' +import { aiFunction, AIFunctionsProvider } from '../fns.js' import { assert, getEnv } from '../utils.js' export namespace serper { @@ -194,7 +194,7 @@ export namespace serper { * * @see https://serper.dev */ -export class SerperClient extends AIToolsProvider { +export class SerperClient extends AIFunctionsProvider { readonly ky: KyInstance readonly apiKey: string readonly apiBaseUrl: string diff --git a/src/services/weather-client.test.ts b/src/services/weather-client.test.ts deleted file mode 100644 index 53ad2d1..0000000 --- a/src/services/weather-client.test.ts +++ /dev/null @@ -1,11 +0,0 @@ -import { expect, test } from 'vitest' - -import { WeatherClient } from './weather-client.js' - -test('WeatherClient.functions', () => { - const weather = new WeatherClient({ - apiKey: 'sk-test' - }) - - expect(weather.functions.get('get_current_weather')).toBeTruthy() -}) diff --git a/src/services/weather-client.ts b/src/services/weather-client.ts index 4f03235..c5cc94a 100644 --- a/src/services/weather-client.ts +++ b/src/services/weather-client.ts @@ -1,7 +1,7 @@ import defaultKy, { type KyInstance } from 'ky' import { z } from 'zod' -import { aiFunction, AIToolsProvider } from '../fns.js' +import { aiFunction, AIFunctionsProvider } from '../fns.js' import { assert, getEnv } from '../utils.js' export namespace weatherapi { @@ -74,7 +74,7 @@ export namespace weatherapi { } } -export class WeatherClient extends AIToolsProvider { +export class WeatherClient extends AIFunctionsProvider { readonly ky: KyInstance readonly apiKey: string readonly apiBaseUrl: string diff --git a/src/types.ts b/src/types.ts index 17a5215..653b16b 100644 --- a/src/types.ts +++ b/src/types.ts @@ -1,6 +1,9 @@ import type { Jsonifiable } from 'type-fest' import type { z } from 'zod' +import type { AIFunctionSet } from './ai-function-set.js' +import type { AIFunctionsProvider } from './fns.js' + export type { KyInstance } from 'ky' export type { ThrottledFunction } from 'p-throttle' @@ -28,6 +31,8 @@ export type AIFunctionImpl = Omit< 'name' | 'toString' | 'arguments' | 'caller' | 'prototype' | 'length' > +export type AIFunctionLike = AIFunctionsProvider | AIFunction | AIFunctionSet + /** * A function meant to be used with LLM function calling. */