From 2e06e9019cd89d44ba52e2ebfe798f53df58b96d Mon Sep 17 00:00:00 2001 From: Konstantin Mandrika Date: Thu, 12 Sep 2024 12:45:03 -0400 Subject: [PATCH 1/4] Record output properly --- js/plugins/google-cloud/src/telemetry/flow.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/js/plugins/google-cloud/src/telemetry/flow.ts b/js/plugins/google-cloud/src/telemetry/flow.ts index c242689c5..a2b18a6a7 100644 --- a/js/plugins/google-cloud/src/telemetry/flow.ts +++ b/js/plugins/google-cloud/src/telemetry/flow.ts @@ -84,7 +84,7 @@ class FlowsTelemetry implements Telemetry { } if (output && logIO) { - this.recordIO(span, 'Output', name, path, input, projectId); + this.recordIO(span, 'Output', name, path, output, projectId); } if (state === 'success') { From 4ecf720ca896a10574a969172a4958df5340515c Mon Sep 17 00:00:00 2001 From: Konstantin Mandrika Date: Thu, 12 Sep 2024 14:36:43 -0400 Subject: [PATCH 2/4] [WIP] Allow plugins to implicitly provide functionality --- js/core/src/config.ts | 39 ++++++++++++++++++++++++- js/core/src/plugin.ts | 17 ++++++++++- js/core/src/registry.ts | 26 +++++++++++++++-- js/plugins/google-cloud/src/index.ts | 9 +++++- js/testapps/flow-simple-ai/src/index.ts | 39 ++++++++++++------------- 5 files changed, 105 insertions(+), 25 deletions(-) diff --git a/js/core/src/config.ts b/js/core/src/config.ts index 65391d374..1c21b6733 100644 --- a/js/core/src/config.ts +++ b/js/core/src/config.ts @@ -20,7 +20,7 @@ import path from 'path'; import { FlowStateStore } from './flowTypes.js'; import { LocalFileFlowStateStore } from './localFileFlowStateStore.js'; import { logger } from './logging.js'; -import { PluginProvider } from './plugin.js'; +import { PluginProvider, PluginProvidesType } from './plugin.js'; import * as registry from './registry.js'; import { AsyncProvider } from './registry.js'; import { @@ -111,6 +111,9 @@ class Config { logger.info(`Initializing plugin ${plugin.name}:`); return await plugin.initializer(); }, + provides() { + return plugin.provides(); + }, }); }); @@ -120,6 +123,23 @@ class Config { logger.debug(` - all environments: ${loggerPluginName}`); this.loggerConfig = async () => this.resolveLoggerConfig(loggerPluginName); + } else { + const capablePlugins = registry.lookupPluginsByAbility( + PluginProvidesType.TELEMETRY + ); + if (capablePlugins) { + if (capablePlugins.length != 1) { + logger.error(`More than one plugin implicitly provides telemetry.`); + logger.error( + `Disambiguate by adding a telemetry {} block to your configuration and specifying one of ${capablePlugins.map((plugin) => plugin.name)}.` + ); + } else { + logger.debug('Registering logging exporters...'); + logger.debug(` - all environments: ${capablePlugins[0].name}`); + this.loggerConfig = async () => + this.resolveLoggerConfig(capablePlugins[0].name); + } + } } if (this.options.telemetry?.instrumentation) { @@ -128,6 +148,23 @@ class Config { logger.debug(` - all environments: ${telemetryPluginName}`); this.telemetryConfig = async () => this.resolveTelemetryConfig(telemetryPluginName); + } else { + const capablePlugins = registry.lookupPluginsByAbility( + PluginProvidesType.TELEMETRY + ); + if (capablePlugins) { + if (capablePlugins.length != 1) { + logger.error(`More than one plugin implicitly provides telemetry.`); + logger.error( + `Disambiguate by adding a telemetry {} block to your configuration and specifying one of ${capablePlugins.map((plugin) => plugin.name)}.` + ); + } else { + logger.debug('Registering telemetry exporters...'); + logger.debug(` - all environments: ${capablePlugins[0].name}`); + this.telemetryConfig = async () => + this.resolveTelemetryConfig(capablePlugins[0].name); + } + } } logger.debug('Registering flow state stores...'); diff --git a/js/core/src/plugin.ts b/js/core/src/plugin.ts index a21a3a5f3..8664b5db3 100644 --- a/js/core/src/plugin.ts +++ b/js/core/src/plugin.ts @@ -25,12 +25,25 @@ export interface Provider { value: T; } +export enum PluginProvidesType { + UNSPECIFIED = 0x0, + MODEL = 0x1, + RETRIEVER = 0x2, + EMBEDDER = 0x3, + INDEXER = 0x4, + EVALUATOR = 0x5, + FLOW_STATE_STORE = 0x6, + TRACE_STORE = 0x7, + TELEMETRY = 0x8, +} + export interface PluginProvider { name: string; initializer: () => | InitializedPlugin | void | Promise; + provides: () => PluginProvidesType; } export interface InitializedPlugin { @@ -58,7 +71,8 @@ export type Plugin = (...args: T) => PluginProvider; */ export function genkitPlugin( pluginName: string, - initFn: T + initFn: T, + providesFn: () => PluginProvidesType = () => PluginProvidesType.UNSPECIFIED ): Plugin> { return (...args: Parameters) => ({ name: pluginName, @@ -67,6 +81,7 @@ export function genkitPlugin( validatePluginActions(pluginName, initializedPlugin); return initializedPlugin; }, + provides: () => providesFn(), }); } diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts index 94647c92b..7ad037cef 100644 --- a/js/core/src/registry.ts +++ b/js/core/src/registry.ts @@ -18,7 +18,7 @@ import * as z from 'zod'; import { Action } from './action.js'; import { FlowStateStore } from './flowTypes.js'; import { logger } from './logging.js'; -import { PluginProvider } from './plugin.js'; +import { PluginProvider, PluginProvidesType } from './plugin.js'; import { startReflectionApi } from './reflectionApi.js'; import { JSONSchema } from './schema.js'; import { TraceStore } from './tracing/types.js'; @@ -29,6 +29,7 @@ const ACTIONS_BY_ID = 'genkit__ACTIONS_BY_ID'; const TRACE_STORES_BY_ENV = 'genkit__TRACE_STORES_BY_ENV'; const FLOW_STATE_STORES_BY_ENV = 'genkit__FLOW_STATE_STORES_BY_ENV'; const PLUGINS_BY_NAME = 'genkit__PLUGINS_BY_NAME'; +const PLUGINS_BY_ABILITY = 'genkit__PLUGINS_BY_ABILITY'; const SCHEMAS_BY_NAME = 'genkit__SCHEMAS_BY_NAME'; function actionsById(): Record> { @@ -55,6 +56,12 @@ function pluginsByName(): Record { } return global[PLUGINS_BY_NAME]; } +function pluginsByAbility(): Record { + if (global[PLUGINS_BY_ABILITY] === undefined) { + global[PLUGINS_BY_ABILITY] = {}; + } + return global[PLUGINS_BY_ABILITY]; +} function schemasByName(): Record< string, { schema?: z.ZodTypeAny; jsonSchema?: JSONSchema } @@ -196,7 +203,8 @@ export async function lookupFlowStateStore( */ export function registerPluginProvider(name: string, provider: PluginProvider) { let cached; - pluginsByName()[name] = { + + const plugin = { name: provider.name, initializer: () => { if (cached) { @@ -205,13 +213,27 @@ export function registerPluginProvider(name: string, provider: PluginProvider) { cached = provider.initializer(); return cached; }, + provides: () => { + return provider.provides(); + }, }; + + pluginsByName()[plugin.name] = plugin; + + if (plugin.provides() !== PluginProvidesType.UNSPECIFIED) { + const plugins = pluginsByAbility()[plugin.provides()] || []; + pluginsByAbility()[plugin.provides()] = [...plugins, plugin]; + } } export function lookupPlugin(name: string) { return pluginsByName()[name]; } +export function lookupPluginsByAbility(ability: PluginProvidesType) { + return pluginsByAbility()[ability]; +} + /** * */ diff --git a/js/plugins/google-cloud/src/index.ts b/js/plugins/google-cloud/src/index.ts index b0a2d2803..6972be332 100644 --- a/js/plugins/google-cloud/src/index.ts +++ b/js/plugins/google-cloud/src/index.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { genkitPlugin, Plugin } from '@genkit-ai/core'; +import { genkitPlugin, Plugin, PluginProvidesType } from '@genkit-ai/core'; import { credentialsFromEnvironment } from './auth.js'; import { GcpLogger } from './gcpLogger.js'; import { GcpOpenTelemetry } from './gcpOpenTelemetry.js'; @@ -29,6 +29,13 @@ export const googleCloud: Plugin<[GcpPluginOptions] | []> = genkitPlugin( async (options?: GcpPluginOptions) => build(options) ); +export const googleCloudWithTelemetry: Plugin<[GcpPluginOptions] | []> = + genkitPlugin( + 'googleCloudWithTelemetry', + async (options?: GcpPluginOptions) => build(options), + () => PluginProvidesType.TELEMETRY + ); + /** * Configures and builds the plugin. * Not normally needed, but exposed for use by the firebase plugin. diff --git a/js/testapps/flow-simple-ai/src/index.ts b/js/testapps/flow-simple-ai/src/index.ts index 934fda755..becc36ee0 100644 --- a/js/testapps/flow-simple-ai/src/index.ts +++ b/js/testapps/flow-simple-ai/src/index.ts @@ -20,7 +20,7 @@ import { configureGenkit } from '@genkit-ai/core'; import { dotprompt, prompt } from '@genkit-ai/dotprompt'; import { defineFirestoreRetriever, firebase } from '@genkit-ai/firebase'; import { defineFlow, run } from '@genkit-ai/flow'; -import { googleCloud } from '@genkit-ai/google-cloud'; +import { googleCloudWithTelemetry } from '@genkit-ai/google-cloud'; import { gemini15Flash, googleAI, @@ -32,7 +32,6 @@ import { textEmbeddingGecko, vertexAI, } from '@genkit-ai/vertexai'; -import { AlwaysOnSampler } from '@opentelemetry/sdk-trace-base'; import { initializeApp } from 'firebase-admin/app'; import { getFirestore } from 'firebase-admin/firestore'; import { Allow, parse } from 'partial-json'; @@ -43,22 +42,22 @@ configureGenkit({ firebase(), googleAI(), vertexAI(), - googleCloud({ + googleCloudWithTelemetry({ // These are configured for demonstration purposes. Sensible defaults are // in place in the event that telemetryConfig is absent. - telemetryConfig: { - // Forces telemetry export in 'dev' - forceDevExport: true, - sampler: new AlwaysOnSampler(), - autoInstrumentation: true, - autoInstrumentationConfig: { - '@opentelemetry/instrumentation-fs': { enabled: false }, - '@opentelemetry/instrumentation-dns': { enabled: false }, - '@opentelemetry/instrumentation-net': { enabled: false }, - }, - metricExportIntervalMillis: 5_000, - metricExportTimeoutMillis: 5_000, - }, + // telemetryConfig: { + // // Forces telemetry export in 'dev' + // forceDevExport: true, + // sampler: new AlwaysOnSampler(), + // autoInstrumentation: true, + // autoInstrumentationConfig: { + // '@opentelemetry/instrumentation-fs': { enabled: false }, + // '@opentelemetry/instrumentation-dns': { enabled: false }, + // '@opentelemetry/instrumentation-net': { enabled: false }, + // }, + // metricExportIntervalMillis: 5_000, + // metricExportTimeoutMillis: 5_000, + // }, }), dotprompt(), ], @@ -66,10 +65,10 @@ configureGenkit({ traceStore: 'firebase', enableTracingAndMetrics: true, logLevel: 'debug', - telemetry: { - instrumentation: 'googleCloud', - logger: 'googleCloud', - }, + // telemetry: { + // instrumentation: 'googleCloud', + // logger: 'googleCloud', + // }, }); const app = initializeApp(); From 76e9591e2423d2d774eb2d9b7948f81c98d83d2f Mon Sep 17 00:00:00 2001 From: Konstantin Mandrika Date: Fri, 13 Sep 2024 13:41:30 -0400 Subject: [PATCH 3/4] Clean up --- js/core/src/config.ts | 138 +++++++++++++--------- js/core/src/index.ts | 1 + js/core/src/plugin.ts | 13 +- js/core/src/registry.ts | 36 +++++- js/core/tests/registry_test.ts | 16 +++ js/plugins/dotprompt/tests/prompt_test.ts | 4 + js/plugins/firebase/src/index.ts | 13 +- 7 files changed, 155 insertions(+), 66 deletions(-) diff --git a/js/core/src/config.ts b/js/core/src/config.ts index 1c21b6733..fddf5b1ff 100644 --- a/js/core/src/config.ts +++ b/js/core/src/config.ts @@ -95,6 +95,28 @@ class Config { return this.telemetryConfig(); } + private registerImplicitPlugin( + ability: PluginProvidesType, + block: string, + registerFn: (name: string) => void + ): boolean { + const capablePlugins = registry.lookupPluginsByAbility(ability); + if (!capablePlugins) { + return false; + } + + if (capablePlugins.length != 1) { + logger.error(`More than one plugin implicitly provides ${ability}`); + logger.error( + `Disambiguate by adding a ${block} block to your configuration and specifying one of ${capablePlugins.map((plugin) => plugin.name)}.` + ); + return false; + } + + registerFn(capablePlugins[0].name); + return true; + } + /** * Configures the system. */ @@ -117,54 +139,40 @@ class Config { }); }); - if (this.options.telemetry?.logger) { - const loggerPluginName = this.options.telemetry.logger; + const registerLogger = (name: string) => { logger.debug('Registering logging exporters...'); - logger.debug(` - all environments: ${loggerPluginName}`); - this.loggerConfig = async () => - this.resolveLoggerConfig(loggerPluginName); + logger.debug(` - all environments: ${name}`); + this.loggerConfig = async () => this.resolveLoggerConfig(name); + }; + + if (this.options.telemetry?.logger) { + registerLogger(this.options.telemetry.logger); } else { - const capablePlugins = registry.lookupPluginsByAbility( - PluginProvidesType.TELEMETRY - ); - if (capablePlugins) { - if (capablePlugins.length != 1) { - logger.error(`More than one plugin implicitly provides telemetry.`); - logger.error( - `Disambiguate by adding a telemetry {} block to your configuration and specifying one of ${capablePlugins.map((plugin) => plugin.name)}.` - ); - } else { - logger.debug('Registering logging exporters...'); - logger.debug(` - all environments: ${capablePlugins[0].name}`); - this.loggerConfig = async () => - this.resolveLoggerConfig(capablePlugins[0].name); + this.registerImplicitPlugin( + PluginProvidesType.TELEMETRY, + `telementry {}`, + (name: string) => { + registerLogger(name); } - } + ); } - if (this.options.telemetry?.instrumentation) { - const telemetryPluginName = this.options.telemetry.instrumentation; + const registerTelemetry = (name: string) => { logger.debug('Registering telemetry exporters...'); - logger.debug(` - all environments: ${telemetryPluginName}`); - this.telemetryConfig = async () => - this.resolveTelemetryConfig(telemetryPluginName); + logger.debug(` - all environments: ${name}`); + this.telemetryConfig = async () => this.resolveTelemetryConfig(name); + }; + + if (this.options.telemetry?.instrumentation) { + registerTelemetry(this.options.telemetry.instrumentation); } else { - const capablePlugins = registry.lookupPluginsByAbility( - PluginProvidesType.TELEMETRY - ); - if (capablePlugins) { - if (capablePlugins.length != 1) { - logger.error(`More than one plugin implicitly provides telemetry.`); - logger.error( - `Disambiguate by adding a telemetry {} block to your configuration and specifying one of ${capablePlugins.map((plugin) => plugin.name)}.` - ); - } else { - logger.debug('Registering telemetry exporters...'); - logger.debug(` - all environments: ${capablePlugins[0].name}`); - this.telemetryConfig = async () => - this.resolveTelemetryConfig(capablePlugins[0].name); + this.registerImplicitPlugin( + PluginProvidesType.TELEMETRY, + 'telemetry {}', + (name: string) => { + registerTelemetry(name); } - } + ); } logger.debug('Registering flow state stores...'); @@ -175,12 +183,24 @@ class Config { ); logger.debug('Registered dev flow state store.'); } - if (this.options.flowStateStore) { - const flowStorePluginName = this.options.flowStateStore; - logger.debug(` - prod: ${flowStorePluginName}`); + + const registerFlowStateStore = (name: string) => { + logger.debug(` - prod: ${name}`); this.configuredEnvs.add('prod'); registry.registerFlowStateStore('prod', () => - this.resolveFlowStateStore(flowStorePluginName) + this.resolveFlowStateStore(name) + ); + }; + + if (this.options.flowStateStore) { + registerFlowStateStore(this.options.flowStateStore); + } else { + this.registerImplicitPlugin( + PluginProvidesType.FLOW_STATE_STORE, + 'flowStateStore', + (name: string) => { + registerFlowStateStore(name); + } ); } @@ -189,22 +209,34 @@ class Config { registry.registerTraceStore('dev', async () => new LocalFileTraceStore()); logger.debug('Registered dev trace store.'); } - if (this.options.traceStore) { - const traceStorePluginName = this.options.traceStore; - logger.debug(` - prod: ${traceStorePluginName}`); + + const registerTraceStore = (name: string) => { + logger.debug(` - prod: ${name}`); this.configuredEnvs.add('prod'); - registry.registerTraceStore('prod', () => - this.resolveTraceStore(traceStorePluginName) - ); + registry.registerTraceStore('prod', () => this.resolveTraceStore(name)); + }; + + if (this.options.traceStore) { + registerTraceStore(this.options.traceStore); if (isDevEnv()) { logger.info( 'In dev mode `traceStore` is defaulted to local file store.' ); } } else { - logger.info( - '`traceStore` is not specified in the config; Traces are not going to be persisted in prod.' - ); + if ( + !this.registerImplicitPlugin( + PluginProvidesType.TRACE_STORE, + 'traceStore', + (name: string) => { + registerTraceStore(name); + } + ) + ) { + logger.info( + '`traceStore` is not specified in the config; Traces are not going to be persisted in prod.' + ); + } } } diff --git a/js/core/src/index.ts b/js/core/src/index.ts index a8b004fa9..8df7ac775 100644 --- a/js/core/src/index.ts +++ b/js/core/src/index.ts @@ -23,5 +23,6 @@ export * from './action.js'; export * from './config.js'; export { GenkitError } from './error.js'; export * from './flowTypes.js'; +export { PluginProvidesType } from './plugin.js'; export { defineJsonSchema, defineSchema } from './schema.js'; export * from './telemetryTypes.js'; diff --git a/js/core/src/plugin.ts b/js/core/src/plugin.ts index 8664b5db3..176dfcb33 100644 --- a/js/core/src/plugin.ts +++ b/js/core/src/plugin.ts @@ -26,15 +26,10 @@ export interface Provider { } export enum PluginProvidesType { - UNSPECIFIED = 0x0, - MODEL = 0x1, - RETRIEVER = 0x2, - EMBEDDER = 0x3, - INDEXER = 0x4, - EVALUATOR = 0x5, - FLOW_STATE_STORE = 0x6, - TRACE_STORE = 0x7, - TELEMETRY = 0x8, + UNSPECIFIED = 0, + FLOW_STATE_STORE = 1 << 0, + TRACE_STORE = 1 << 1, + TELEMETRY = 1 << 2, } export interface PluginProvider { diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts index 7ad037cef..a3ca8197d 100644 --- a/js/core/src/registry.ts +++ b/js/core/src/registry.ts @@ -221,8 +221,11 @@ export function registerPluginProvider(name: string, provider: PluginProvider) { pluginsByName()[plugin.name] = plugin; if (plugin.provides() !== PluginProvidesType.UNSPECIFIED) { - const plugins = pluginsByAbility()[plugin.provides()] || []; - pluginsByAbility()[plugin.provides()] = [...plugins, plugin]; + getDiscreteAbilities(plugin.provides()).forEach((ability) => { + console.log(`>>>>>>> plugin: ${plugin.name} provides: ${ability}`); + const plugins = pluginsByAbility()[ability] || []; + pluginsByAbility()[ability] = [...plugins, plugin]; + }); } } @@ -274,3 +277,32 @@ export function __hardResetRegistryForTesting() { function deleteAll(map: Record) { Object.keys(map).forEach((key) => delete map[key]); } + +function getDiscreteAbilities( + provides: PluginProvidesType +): PluginProvidesType[] { + let abilities: PluginProvidesType[] = []; + + if ( + (provides & PluginProvidesType.FLOW_STATE_STORE) === + PluginProvidesType.FLOW_STATE_STORE + ) { + abilities.push(PluginProvidesType.FLOW_STATE_STORE); + } + + if ( + (provides & PluginProvidesType.TRACE_STORE) === + PluginProvidesType.TRACE_STORE + ) { + abilities.push(PluginProvidesType.TRACE_STORE); + } + + if ( + (provides & PluginProvidesType.TELEMETRY) === + PluginProvidesType.TELEMETRY + ) { + abilities.push(PluginProvidesType.TELEMETRY); + } + + return abilities; +} diff --git a/js/core/tests/registry_test.ts b/js/core/tests/registry_test.ts index c969eba29..1bbd7f156 100644 --- a/js/core/tests/registry_test.ts +++ b/js/core/tests/registry_test.ts @@ -16,6 +16,7 @@ import assert from 'node:assert'; import { beforeEach, describe, it } from 'node:test'; +import { PluginProvidesType } from '../lib/config.js'; import { action } from '../src/action.js'; import { __hardResetRegistryForTesting, @@ -54,6 +55,9 @@ describe('registry', () => { registerAction('model', fooSomethingAction); return {}; }, + provides() { + return PluginProvidesType.UNSPECIFIED; + }, }); const fooSomethingAction = action( { @@ -70,6 +74,9 @@ describe('registry', () => { registerAction('model', barSomethingAction); return {}; }, + provides() { + return PluginProvidesType.UNSPECIFIED; + }, }); const barSomethingAction = action( { @@ -97,6 +104,9 @@ describe('registry', () => { fooInitialized = true; return {}; }, + provides() { + return PluginProvidesType.UNSPECIFIED; + }, }); let barInitialized = false; registerPluginProvider('bar', { @@ -105,6 +115,9 @@ describe('registry', () => { barInitialized = true; return {}; }, + provides() { + return PluginProvidesType.UNSPECIFIED; + }, }); await lookupAction('/model/foo/something'); @@ -148,6 +161,9 @@ describe('registry', () => { registerAction('model', somethingAction); return {}; }, + provides() { + return PluginProvidesType.UNSPECIFIED; + }, }); const somethingAction = action( { diff --git a/js/plugins/dotprompt/tests/prompt_test.ts b/js/plugins/dotprompt/tests/prompt_test.ts index c9a386141..a0be3f607 100644 --- a/js/plugins/dotprompt/tests/prompt_test.ts +++ b/js/plugins/dotprompt/tests/prompt_test.ts @@ -18,6 +18,7 @@ import assert from 'node:assert'; import { describe, it } from 'node:test'; import { defineModel } from '@genkit-ai/ai/model'; +import { PluginProvidesType } from '@genkit-ai/core'; import { toJsonSchema, ValidationError } from '@genkit-ai/core/schema'; import z from 'zod'; import { registerPluginProvider } from '../../../core/src/registry.js'; @@ -31,6 +32,9 @@ function registerDotprompt() { async initializer() { return {}; }, + provides() { + return PluginProvidesType.UNSPECIFIED; + }, }); } diff --git a/js/plugins/firebase/src/index.ts b/js/plugins/firebase/src/index.ts index 282fffcf2..f557b8734 100644 --- a/js/plugins/firebase/src/index.ts +++ b/js/plugins/firebase/src/index.ts @@ -14,7 +14,12 @@ * limitations under the License. */ -import { genkitPlugin, isDevEnv, Plugin } from '@genkit-ai/core'; +import { + genkitPlugin, + isDevEnv, + Plugin, + PluginProvidesType, +} from '@genkit-ai/core'; import { logger } from '@genkit-ai/core/logging'; import { FirestoreStateStore } from '@genkit-ai/flow'; import { @@ -85,5 +90,9 @@ export const firebase: Plugin<[FirestorePluginParams] | []> = genkitPlugin( }, }, }; - } + }, + () => + PluginProvidesType.FLOW_STATE_STORE | + PluginProvidesType.TRACE_STORE | + PluginProvidesType.TELEMETRY ); From 0e22cde20fe8ca54d8af04ee958d1ac241e7ab2e Mon Sep 17 00:00:00 2001 From: Konstantin Mandrika Date: Mon, 16 Sep 2024 14:32:00 -0400 Subject: [PATCH 4/4] Rename to a more sensible type --- js/core/src/config.ts | 16 ++++++----- js/core/src/index.ts | 2 +- js/core/src/plugin.ts | 6 ++--- js/core/src/registry.ts | 33 +++++++++++------------ js/core/tests/registry_test.ts | 12 ++++----- js/plugins/dotprompt/tests/prompt_test.ts | 4 +-- js/plugins/firebase/src/index.ts | 8 +++--- js/plugins/google-cloud/src/index.ts | 4 +-- 8 files changed, 44 insertions(+), 41 deletions(-) diff --git a/js/core/src/config.ts b/js/core/src/config.ts index fddf5b1ff..29b9d3c27 100644 --- a/js/core/src/config.ts +++ b/js/core/src/config.ts @@ -20,7 +20,7 @@ import path from 'path'; import { FlowStateStore } from './flowTypes.js'; import { LocalFileFlowStateStore } from './localFileFlowStateStore.js'; import { logger } from './logging.js'; -import { PluginProvider, PluginProvidesType } from './plugin.js'; +import { PluginAbilityType, PluginProvider } from './plugin.js'; import * as registry from './registry.js'; import { AsyncProvider } from './registry.js'; import { @@ -95,8 +95,12 @@ class Config { return this.telemetryConfig(); } + /** + * Registers plugins that provide capabilities not explicitly specified in + * the config. + */ private registerImplicitPlugin( - ability: PluginProvidesType, + ability: PluginAbilityType, block: string, registerFn: (name: string) => void ): boolean { @@ -149,7 +153,7 @@ class Config { registerLogger(this.options.telemetry.logger); } else { this.registerImplicitPlugin( - PluginProvidesType.TELEMETRY, + PluginAbilityType.TELEMETRY, `telementry {}`, (name: string) => { registerLogger(name); @@ -167,7 +171,7 @@ class Config { registerTelemetry(this.options.telemetry.instrumentation); } else { this.registerImplicitPlugin( - PluginProvidesType.TELEMETRY, + PluginAbilityType.TELEMETRY, 'telemetry {}', (name: string) => { registerTelemetry(name); @@ -196,7 +200,7 @@ class Config { registerFlowStateStore(this.options.flowStateStore); } else { this.registerImplicitPlugin( - PluginProvidesType.FLOW_STATE_STORE, + PluginAbilityType.FLOW_STATE_STORE, 'flowStateStore', (name: string) => { registerFlowStateStore(name); @@ -226,7 +230,7 @@ class Config { } else { if ( !this.registerImplicitPlugin( - PluginProvidesType.TRACE_STORE, + PluginAbilityType.TRACE_STORE, 'traceStore', (name: string) => { registerTraceStore(name); diff --git a/js/core/src/index.ts b/js/core/src/index.ts index 8df7ac775..92f68b542 100644 --- a/js/core/src/index.ts +++ b/js/core/src/index.ts @@ -23,6 +23,6 @@ export * from './action.js'; export * from './config.js'; export { GenkitError } from './error.js'; export * from './flowTypes.js'; -export { PluginProvidesType } from './plugin.js'; +export { PluginAbilityType } from './plugin.js'; export { defineJsonSchema, defineSchema } from './schema.js'; export * from './telemetryTypes.js'; diff --git a/js/core/src/plugin.ts b/js/core/src/plugin.ts index 176dfcb33..10abded60 100644 --- a/js/core/src/plugin.ts +++ b/js/core/src/plugin.ts @@ -25,7 +25,7 @@ export interface Provider { value: T; } -export enum PluginProvidesType { +export enum PluginAbilityType { UNSPECIFIED = 0, FLOW_STATE_STORE = 1 << 0, TRACE_STORE = 1 << 1, @@ -38,7 +38,7 @@ export interface PluginProvider { | InitializedPlugin | void | Promise; - provides: () => PluginProvidesType; + provides: () => PluginAbilityType; } export interface InitializedPlugin { @@ -67,7 +67,7 @@ export type Plugin = (...args: T) => PluginProvider; export function genkitPlugin( pluginName: string, initFn: T, - providesFn: () => PluginProvidesType = () => PluginProvidesType.UNSPECIFIED + providesFn: () => PluginAbilityType = () => PluginAbilityType.UNSPECIFIED ): Plugin> { return (...args: Parameters) => ({ name: pluginName, diff --git a/js/core/src/registry.ts b/js/core/src/registry.ts index a3ca8197d..8706bd3ac 100644 --- a/js/core/src/registry.ts +++ b/js/core/src/registry.ts @@ -18,7 +18,7 @@ import * as z from 'zod'; import { Action } from './action.js'; import { FlowStateStore } from './flowTypes.js'; import { logger } from './logging.js'; -import { PluginProvider, PluginProvidesType } from './plugin.js'; +import { PluginAbilityType, PluginProvider } from './plugin.js'; import { startReflectionApi } from './reflectionApi.js'; import { JSONSchema } from './schema.js'; import { TraceStore } from './tracing/types.js'; @@ -56,7 +56,7 @@ function pluginsByName(): Record { } return global[PLUGINS_BY_NAME]; } -function pluginsByAbility(): Record { +function pluginsByAbility(): Record { if (global[PLUGINS_BY_ABILITY] === undefined) { global[PLUGINS_BY_ABILITY] = {}; } @@ -220,9 +220,8 @@ export function registerPluginProvider(name: string, provider: PluginProvider) { pluginsByName()[plugin.name] = plugin; - if (plugin.provides() !== PluginProvidesType.UNSPECIFIED) { + if (plugin.provides() !== PluginAbilityType.UNSPECIFIED) { getDiscreteAbilities(plugin.provides()).forEach((ability) => { - console.log(`>>>>>>> plugin: ${plugin.name} provides: ${ability}`); const plugins = pluginsByAbility()[ability] || []; pluginsByAbility()[ability] = [...plugins, plugin]; }); @@ -233,7 +232,7 @@ export function lookupPlugin(name: string) { return pluginsByName()[name]; } -export function lookupPluginsByAbility(ability: PluginProvidesType) { +export function lookupPluginsByAbility(ability: PluginAbilityType) { return pluginsByAbility()[ability]; } @@ -279,29 +278,29 @@ function deleteAll(map: Record) { } function getDiscreteAbilities( - provides: PluginProvidesType -): PluginProvidesType[] { - let abilities: PluginProvidesType[] = []; + provides: PluginAbilityType +): PluginAbilityType[] { + let abilities: PluginAbilityType[] = []; if ( - (provides & PluginProvidesType.FLOW_STATE_STORE) === - PluginProvidesType.FLOW_STATE_STORE + (provides & PluginAbilityType.FLOW_STATE_STORE) === + PluginAbilityType.FLOW_STATE_STORE ) { - abilities.push(PluginProvidesType.FLOW_STATE_STORE); + abilities.push(PluginAbilityType.FLOW_STATE_STORE); } if ( - (provides & PluginProvidesType.TRACE_STORE) === - PluginProvidesType.TRACE_STORE + (provides & PluginAbilityType.TRACE_STORE) === + PluginAbilityType.TRACE_STORE ) { - abilities.push(PluginProvidesType.TRACE_STORE); + abilities.push(PluginAbilityType.TRACE_STORE); } if ( - (provides & PluginProvidesType.TELEMETRY) === - PluginProvidesType.TELEMETRY + (provides & PluginAbilityType.TELEMETRY) === + PluginAbilityType.TELEMETRY ) { - abilities.push(PluginProvidesType.TELEMETRY); + abilities.push(PluginAbilityType.TELEMETRY); } return abilities; diff --git a/js/core/tests/registry_test.ts b/js/core/tests/registry_test.ts index 1bbd7f156..57df72c63 100644 --- a/js/core/tests/registry_test.ts +++ b/js/core/tests/registry_test.ts @@ -16,8 +16,8 @@ import assert from 'node:assert'; import { beforeEach, describe, it } from 'node:test'; -import { PluginProvidesType } from '../lib/config.js'; import { action } from '../src/action.js'; +import { PluginAbilityType } from '../src/config.js'; import { __hardResetRegistryForTesting, listActions, @@ -56,7 +56,7 @@ describe('registry', () => { return {}; }, provides() { - return PluginProvidesType.UNSPECIFIED; + return PluginAbilityType.UNSPECIFIED; }, }); const fooSomethingAction = action( @@ -75,7 +75,7 @@ describe('registry', () => { return {}; }, provides() { - return PluginProvidesType.UNSPECIFIED; + return PluginAbilityType.UNSPECIFIED; }, }); const barSomethingAction = action( @@ -105,7 +105,7 @@ describe('registry', () => { return {}; }, provides() { - return PluginProvidesType.UNSPECIFIED; + return PluginAbilityType.UNSPECIFIED; }, }); let barInitialized = false; @@ -116,7 +116,7 @@ describe('registry', () => { return {}; }, provides() { - return PluginProvidesType.UNSPECIFIED; + return PluginAbilityType.UNSPECIFIED; }, }); @@ -162,7 +162,7 @@ describe('registry', () => { return {}; }, provides() { - return PluginProvidesType.UNSPECIFIED; + return PluginAbilityType.UNSPECIFIED; }, }); const somethingAction = action( diff --git a/js/plugins/dotprompt/tests/prompt_test.ts b/js/plugins/dotprompt/tests/prompt_test.ts index a0be3f607..32e4b512a 100644 --- a/js/plugins/dotprompt/tests/prompt_test.ts +++ b/js/plugins/dotprompt/tests/prompt_test.ts @@ -18,7 +18,7 @@ import assert from 'node:assert'; import { describe, it } from 'node:test'; import { defineModel } from '@genkit-ai/ai/model'; -import { PluginProvidesType } from '@genkit-ai/core'; +import { PluginAbilityType } from '@genkit-ai/core'; import { toJsonSchema, ValidationError } from '@genkit-ai/core/schema'; import z from 'zod'; import { registerPluginProvider } from '../../../core/src/registry.js'; @@ -33,7 +33,7 @@ function registerDotprompt() { return {}; }, provides() { - return PluginProvidesType.UNSPECIFIED; + return PluginAbilityType.UNSPECIFIED; }, }); } diff --git a/js/plugins/firebase/src/index.ts b/js/plugins/firebase/src/index.ts index f557b8734..486e10659 100644 --- a/js/plugins/firebase/src/index.ts +++ b/js/plugins/firebase/src/index.ts @@ -18,7 +18,7 @@ import { genkitPlugin, isDevEnv, Plugin, - PluginProvidesType, + PluginAbilityType, } from '@genkit-ai/core'; import { logger } from '@genkit-ai/core/logging'; import { FirestoreStateStore } from '@genkit-ai/flow'; @@ -92,7 +92,7 @@ export const firebase: Plugin<[FirestorePluginParams] | []> = genkitPlugin( }; }, () => - PluginProvidesType.FLOW_STATE_STORE | - PluginProvidesType.TRACE_STORE | - PluginProvidesType.TELEMETRY + PluginAbilityType.FLOW_STATE_STORE | + PluginAbilityType.TRACE_STORE | + PluginAbilityType.TELEMETRY ); diff --git a/js/plugins/google-cloud/src/index.ts b/js/plugins/google-cloud/src/index.ts index 6972be332..e1ba860ce 100644 --- a/js/plugins/google-cloud/src/index.ts +++ b/js/plugins/google-cloud/src/index.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { genkitPlugin, Plugin, PluginProvidesType } from '@genkit-ai/core'; +import { genkitPlugin, Plugin, PluginAbilityType } from '@genkit-ai/core'; import { credentialsFromEnvironment } from './auth.js'; import { GcpLogger } from './gcpLogger.js'; import { GcpOpenTelemetry } from './gcpOpenTelemetry.js'; @@ -33,7 +33,7 @@ export const googleCloudWithTelemetry: Plugin<[GcpPluginOptions] | []> = genkitPlugin( 'googleCloudWithTelemetry', async (options?: GcpPluginOptions) => build(options), - () => PluginProvidesType.TELEMETRY + () => PluginAbilityType.TELEMETRY ); /**