From 1ff16a487f8211c54e3a53be7afb94ccd4490465 Mon Sep 17 00:00:00 2001 From: Wlad Paiva Date: Wed, 1 Nov 2023 16:25:12 -0300 Subject: [PATCH] add generic types prep --- .changeset/violet-paws-invite.md | 5 +++ src/index.test.ts | 2 +- src/index.ts | 67 +++++++++++++++++++++----------- src/plugins/cli.ts | 2 +- src/plugins/file-history.ts | 2 +- src/plugins/web-browsing.ts | 2 +- 6 files changed, 53 insertions(+), 27 deletions(-) create mode 100644 .changeset/violet-paws-invite.md diff --git a/.changeset/violet-paws-invite.md b/.changeset/violet-paws-invite.md new file mode 100644 index 0000000..4895f96 --- /dev/null +++ b/.changeset/violet-paws-invite.md @@ -0,0 +1,5 @@ +--- +'aibitat': patch +--- + +AIbitat prep to accept generic types diff --git a/src/index.test.ts b/src/index.test.ts index e922524..1b68d3b 100644 --- a/src/index.test.ts +++ b/src/index.test.ts @@ -218,7 +218,7 @@ describe('direct message', () => { describe('as a group', () => { const members = ['🐶', '😸', '🐭'] - let aibitat: AIbitat + let aibitat: AIbitat beforeEach(() => { ai.complete.mockImplementation(x => { diff --git a/src/index.ts b/src/index.ts index d0bf6c0..1872c95 100644 --- a/src/index.ts +++ b/src/index.ts @@ -6,18 +6,23 @@ import * as Providers from './providers' export * from './providers' +/** + * Available providers + */ +type Provider = 'openai' | 'anthropic' | Providers.Provider + /** * The provider config to use for the AI. */ -export type ProviderConfig = - // FIX: should only show Openai models when there's no provider - | { +export type ProviderConfig = T extends 'openai' + ? { /** The OpenAI API provider */ provider?: 'openai' /** The model to use with the OpenAI */ model?: Providers.OpenAIProviderConfig['model'] } - | { + : T extends 'anthropic' + ? { /** The custom AI provider */ provider: 'anthropic' /** @@ -26,7 +31,7 @@ export type ProviderConfig = */ model?: Providers.AnthropicProviderConfig['model'] } - | { + : { /** The custom AI provider */ provider: Providers.Provider } @@ -34,7 +39,7 @@ export type ProviderConfig = /** * Base config for AIbitat agents. */ -export type AgentConfig = ProviderConfig & { +export type AgentConfig = ProviderConfig & { /** The role this agent will play in the conversation */ role?: string @@ -54,7 +59,7 @@ export type AgentConfig = ProviderConfig & { /** * The channel configuration for the AIbitat. */ -export type ChannelConfig = ProviderConfig & { +export type ChannelConfig = ProviderConfig & { /** The role this agent will play in the conversation */ role?: string @@ -76,9 +81,16 @@ type Route = { /** * A chat message. */ -type Message = Route & { - content: string -} +type Message = Prettify< + Route & { + content: string + } +> + +// Prettify a type +type Prettify = { + [K in keyof T]: T[K] +} & {} /** * A chat message that is saved in the history. @@ -96,7 +108,7 @@ type History = Array /** * AIbitat props. */ -export type AIbitatProps = ProviderConfig & { +export type AIbitatProps = ProviderConfig & { /** * Chat history between all agents. * @default [] @@ -122,19 +134,19 @@ export type AIbitatProps = ProviderConfig & { * * Guiding the chat through a graph of agents. */ -export class AIbitat { +export class AIbitat { private emitter = new EventEmitter() - private defaultProvider: ProviderConfig + private defaultProvider: ProviderConfig private defaultInterrupt private maxRounds private _chats - private agents = new Map() - private channels = new Map() + private agents = new Map>() + private channels = new Map>() private functions = new Map() - constructor(props: AIbitatProps = {}) { + constructor(props: AIbitatProps = {} as AIbitatProps) { const { chats = [], interrupt = 'NEVER', @@ -162,7 +174,7 @@ export class AIbitat { /** * Install a plugin. */ - use(plugin: AIbitat.Plugin) { + use(plugin: AIbitat.Plugin) { plugin.setup(this) return this } @@ -174,7 +186,10 @@ export class AIbitat { * @param config * @returns */ - public agent(name: string, config: AgentConfig = {}) { + public agent( + name: string, + config: AgentConfig = {} as AgentConfig, + ) { this.agents.set(name, config) return this } @@ -187,7 +202,11 @@ export class AIbitat { * @param config * @returns */ - public channel(name: string, members: string[], config: ChannelConfig = {}) { + public channel( + name: string, + members: string[], + config: ChannelConfig = {} as ChannelConfig, + ) { this.channels.set(name, { members, ...config, @@ -378,7 +397,9 @@ export class AIbitat { * @param listener * @returns */ - public onStart(listener: (chat: Chat, aibitat: AIbitat) => void) { + public onStart( + listener: (chat: Chat, aibitat: AIbitat) => void, + ) { this.emitter.on('start', listener) return this } @@ -805,7 +826,7 @@ ${this.getHistory({to: route.to}) * * @param config The provider configuration. */ - private getProviderForConfig(config: ProviderConfig) { + private getProviderForConfig(config: ProviderConfig) { if (typeof config.provider === 'object') { return config.provider } @@ -838,7 +859,7 @@ export namespace AIbitat { /** * Plugin to use with the aibitat */ - export type Plugin = { + export type Plugin = { /** * The name of the plugin. This will be used to identify the plugin. * If the plugin is already installed, it will replace the old plugin. @@ -848,7 +869,7 @@ export namespace AIbitat { /** * The setup function to be called when the plugin is installed. */ - setup: (aibitat: AIbitat) => void + setup: (aibitat: AIbitat) => void } export type FunctionConfig = Function.FunctionConfig diff --git a/src/plugins/cli.ts b/src/plugins/cli.ts index 7f5e50a..1991acd 100644 --- a/src/plugins/cli.ts +++ b/src/plugins/cli.ts @@ -70,7 +70,7 @@ function cli({ await aibitat.continue(feedback) }) }, - } as AIbitat.Plugin + } as AIbitat.Plugin } /** diff --git a/src/plugins/file-history.ts b/src/plugins/file-history.ts index 1fd6bfa..d69903a 100644 --- a/src/plugins/file-history.ts +++ b/src/plugins/file-history.ts @@ -37,5 +37,5 @@ export function fileHistory({ }) }) }, - } as AIbitat.Plugin + } as AIbitat.Plugin } diff --git a/src/plugins/web-browsing.ts b/src/plugins/web-browsing.ts index 4937150..ff7bd5d 100644 --- a/src/plugins/web-browsing.ts +++ b/src/plugins/web-browsing.ts @@ -174,5 +174,5 @@ export function experimental_webBrowsing({}: {} = {}) { }, }) }, - } as AIbitat.Plugin + } as AIbitat.Plugin }