Something went wrong. Try again.
ATProto-based AI Assistant Chat interface.
Something went wrong. Try again.
4.5 kB · 120 lines
TypeScript
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121import { type Citation, type Ingester, type ModelProvider, matchIngester, type Plugin, type PluginContext, type Tool, type ToolContext,} from './index.ts'
/** Run a plugin's setup against a recording context, for plugin tests. */export async function setupForTest( plugin: Plugin, options: { /** Ingesters from other plugins, which the plugin can reach through `ctx.ingesters`. */ ingesters?: Ingester[] userSettings?: (user: string) => unknown app?: { name: string; publicUrl: string } /** DIDs that have an account, which `roles.syncMembers` adds when listed. */ accounts?: string[] } = {},) { const providers: ModelProvider[] = [] const tools: Tool<unknown>[] = [] const ingesters: Ingester[] = [] const roleMembers = new Map<string, string[]>() const accountDids = new Set(options.accounts) const accountChanges: { action: 'suspend' | 'restore'; did: string; reason?: string }[] = [] const hooks: { name: string; handler: unknown }[] = [] const closers: (() => void | Promise<void>)[] = [] const reachable = () => [...ingesters, ...(options.ingesters ?? [])] const noop = () => {} const ctx = { providers: { register: (provider: ModelProvider) => void providers.push(provider) }, tools: { register: (tool: Tool<unknown>) => void tools.push(tool) }, toolSources: { register: noop }, roles: { syncMembers: async (role: string, dids: string[]) => { const listed = new Set(dids) const current = roleMembers.get(role) ?? [] const kept = current.filter((did) => listed.has(did)) const added = dids.filter((did) => accountDids.has(did) && !current.includes(did)) roleMembers.set(role, [...kept, ...added]) return { added: added.length, removed: current.length - kept.length } }, addMember: async (role: string, did: string) => { const current = roleMembers.get(role) ?? [] if (!current.includes(did)) roleMembers.set(role, [...current, did]) }, }, accounts: { suspend: async (did: string, options?: { reason?: string }) => { accountChanges.push({ action: 'suspend', did, ...options }) return true }, restore: async (did: string) => { accountChanges.push({ action: 'restore', did }) return true }, }, ingesters: { register: (ingester: Ingester) => void ingesters.push(ingester), accepts: (mimeType: string) => matchIngester(reachable(), mimeType) !== undefined, ingest: async (file: Parameters<Ingester['ingest']>[0]) => matchIngester(reachable(), file.mimeType)?.ingest(file), }, hooks: { on: (name: string, handler: unknown) => void hooks.push({ name, handler }) }, models: { generateText: async () => ({ text: '', finishReason: 'stop' }) }, conversations: { updateInfo: async () => {} }, userSettings: async (user: string) => options.userSettings?.(user) ?? {}, logger: { debug: noop, info: noop, warn: noop, error: noop }, app: options.app ?? { name: 'Test', publicUrl: 'http://127.0.0.1:3000' }, onClose: (fn: () => void | Promise<void>) => void closers.push(fn), } as unknown as PluginContext await plugin.setup(ctx) /** Run the plugin's onClose handlers, as when its runtime is swapped out. */ const close = async () => { for (const fn of closers) await fn() } /** Run the plugin's handlers for an action hook, as core would. */ const fire = async (name: string, payload: unknown) => { for (const hook of hooks.filter((h) => h.name === name)) await (hook.handler as (p: unknown) => unknown)(payload) } return { providers, tools, ingesters, roleMembers, accountChanges, hooks, fire, ctx, close, }}
/** A tool context that records citations, for tool tests. */export function toolContextForTest( options: { fetch?: typeof globalThis.fetch user?: string conversation?: string roles?: string[] signal?: AbortSignal } = {},) { const citations: Citation[] = [] const context: ToolContext = { user: options.user ?? 'did:plc:alice', conversation: options.conversation ?? 'at://did:plc:alice/space/c/1', roles: options.roles ?? ['user'], signal: options.signal ?? new AbortController().signal, fetch: options.fetch ?? (() => Promise.reject(new Error('No fetch in this test'))), cite: (source) => void citations.push(source), turnCache: new Map(), } return { context, citations }}