diff --git a/genkit-tools/cli/src/commands/eval-extract-data.ts b/genkit-tools/cli/src/commands/eval-extract-data.ts index c5a4449a4..3679091a6 100644 --- a/genkit-tools/cli/src/commands/eval-extract-data.ts +++ b/genkit-tools/cli/src/commands/eval-extract-data.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { EnvTypes, EvalInput, TraceData } from '@genkit-ai/tools-common'; +import { EvalInput, TraceData } from '@genkit-ai/tools-common'; import { generateTestCaseId, getEvalExtractors, @@ -25,7 +25,6 @@ import { writeFile } from 'fs/promises'; import { runWithManager } from '../utils/manager-utils'; interface EvalDatasetOptions { - env: EnvTypes; output?: string; maxRows: string; label?: string; @@ -35,7 +34,6 @@ interface EvalDatasetOptions { export const evalExtractData = new Command('eval:extractData') .description('extract evaludation data for a given flow from the trace store') .argument('', 'name of the flow to run') - .option('--env ', 'environment (dev/prod)', 'dev') .option( '--output ', 'name of the output file to store the extracted data' @@ -51,7 +49,6 @@ export const evalExtractData = new Command('eval:extractData') let continuationToken = undefined; while (dataset.length < parseInt(options.maxRows)) { const response = await manager.listTraces({ - env: options.env, limit: parseInt(options.maxRows), continuationToken, }); diff --git a/genkit-tools/cli/src/commands/flow-batch-run.ts b/genkit-tools/cli/src/commands/flow-batch-run.ts index d803ff13d..9178262b0 100644 --- a/genkit-tools/cli/src/commands/flow-batch-run.ts +++ b/genkit-tools/cli/src/commands/flow-batch-run.ts @@ -22,7 +22,7 @@ import { import { logger } from '@genkit-ai/tools-common/utils'; import { Command } from 'commander'; import { readFile, writeFile } from 'fs/promises'; -import { runWithManager, waitForFlowToComplete } from '../utils/manager-utils'; +import { runWithManager } from '../utils/manager-utils'; interface FlowBatchRunOptions { wait?: boolean; @@ -76,14 +76,6 @@ export const flowBatchRun = new Command('flow:batchRun') }) ).result as FlowState; - if (!state.operation.done && options.wait) { - logger.info('Started flow run, waiting for it to complete...'); - state = await waitForFlowToComplete( - manager, - flowName, - state.flowId - ); - } logger.info( 'Flow operation:\n' + JSON.stringify(state.operation, undefined, ' ') diff --git a/genkit-tools/cli/src/utils/manager-utils.ts b/genkit-tools/cli/src/utils/manager-utils.ts index ab4625475..cecbb98d1 100644 --- a/genkit-tools/cli/src/utils/manager-utils.ts +++ b/genkit-tools/cli/src/utils/manager-utils.ts @@ -18,11 +18,7 @@ import { LocalFileTraceStore, startTelemetryServer, } from '@genkit-ai/telemetry-server'; -import { - FlowInvokeEnvelopeMessage, - FlowState, - Status, -} from '@genkit-ai/tools-common'; +import { Status } from '@genkit-ai/tools-common'; import { GenkitToolsError, RuntimeManager, @@ -89,43 +85,3 @@ export async function runWithManager( logger.error(`${error.stack}`); } } - -/** - * Poll and wait for the flow to fully complete. - */ -export async function waitForFlowToComplete( - manager: RuntimeManager, - flowName: string, - flowId: string -): Promise { - let state; - // eslint-disable-next-line no-constant-condition - while (true) { - state = await getFlowState(manager, flowName, flowId); - if (state.operation.done) { - break; - } - await new Promise((r) => setTimeout(r, 1000)); - } - return state; -} - -/** - * Retrieve the flow state. - */ -export async function getFlowState( - manager: RuntimeManager, - flowName: string, - flowId: string -): Promise { - return ( - await manager.runAction({ - key: `/flow/${flowName}`, - input: { - state: { - flowId, - }, - } as FlowInvokeEnvelopeMessage, - }) - ).result as FlowState; -} diff --git a/genkit-tools/common/src/api/reflection.ts b/genkit-tools/common/src/api/reflection.ts index 4393634cf..02eebd672 100644 --- a/genkit-tools/common/src/api/reflection.ts +++ b/genkit-tools/common/src/api/reflection.ts @@ -78,8 +78,8 @@ registry.registerPath({ }); registry.registerPath({ method: 'get', - path: '/api/envs/{env}/traces', - summary: 'Retrieves all traces for a given environment (e.g. dev or prod).', + path: '/api/traces', + summary: 'Retrieves all traces', request: { params: apis.ListTracesRequestSchema, }, @@ -96,8 +96,8 @@ registry.registerPath({ }); registry.registerPath({ method: 'get', - path: '/api/envs/{env}/traces/{traceId}', - summary: 'Retrieves traces for the given environment.', + path: '/api/traces/{traceId}', + summary: 'Retrieves a trace by traceId', request: { params: apis.GetTraceRequestSchema, }, @@ -112,43 +112,6 @@ registry.registerPath({ }, }, }); -registry.registerPath({ - method: 'get', - path: '/api/envs/{env}/flowStates', - summary: - 'Retrieves all flow states for a given environment (e.g. dev or prod).', - request: { - params: apis.ListFlowStatesRequestSchema, - }, - responses: { - '200': { - description: 'Success', - content: { - 'application/json': { - schema: z.array(FlowStateSchema), - }, - }, - }, - }, -}); -registry.registerPath({ - method: 'get', - path: '/api/envs/{env}/flowStates/{flowId}', - summary: 'Retrieves a flow state for the given ID.', - request: { - params: apis.GetFlowStateRequestSchema, - }, - responses: { - '200': { - description: 'Success', - content: { - 'application/json': { - schema: FlowStateSchema, - }, - }, - }, - }, -}); const generator = new OpenApiGeneratorV3(registry.definitions); const document = generator.generateDocument({ diff --git a/genkit-tools/common/src/eval/evaluate.ts b/genkit-tools/common/src/eval/evaluate.ts index a0727a95b..28654afaf 100644 --- a/genkit-tools/common/src/eval/evaluate.ts +++ b/genkit-tools/common/src/eval/evaluate.ts @@ -333,10 +333,6 @@ async function gatherEvalInput(params: { } const trace = await manager.getTrace({ - // TODO: We should consider making this a argument and using it to - // to control which tracestore environment is being used when - // running a flow. - env: 'dev', traceId, }); diff --git a/genkit-tools/common/src/manager/manager.ts b/genkit-tools/common/src/manager/manager.ts index b203d3604..9bcfb3c9f 100644 --- a/genkit-tools/common/src/manager/manager.ts +++ b/genkit-tools/common/src/manager/manager.ts @@ -223,12 +223,12 @@ export class RuntimeManager { } /** - * Retrieves all traces for a given environment (e.g. dev or prod). + * Retrieves all traces */ async listTraces( input: apis.ListTracesRequest ): Promise { - const { env, limit, continuationToken } = input; + const { limit, continuationToken } = input; let query = ''; if (limit) { query += `limit=${limit}`; @@ -243,10 +243,7 @@ export class RuntimeManager { const response = await axios .get(`${this.telemetryServerUrl}/api/traces?${query}`) .catch((err) => - this.httpErrorHandler( - err, - `Error listing traces for env='${env}', query='${query}'.` - ) + this.httpErrorHandler(err, `Error listing traces for query='${query}'.`) ); return apis.ListTracesResponseSchema.parse(response.data); @@ -256,13 +253,13 @@ export class RuntimeManager { * Retrieves a trace for a given ID. */ async getTrace(input: apis.GetTraceRequest): Promise { - const { env, traceId } = input; + const { traceId } = input; const response = await axios .get(`${this.telemetryServerUrl}/api/traces/${traceId}`) .catch((err) => this.httpErrorHandler( err, - `Error getting trace for traceId='${traceId}', env='${env}'.` + `Error getting trace for traceId='${traceId}'` ) ); diff --git a/genkit-tools/common/src/types/apis.ts b/genkit-tools/common/src/types/apis.ts index f1889b050..aebce4c36 100644 --- a/genkit-tools/common/src/types/apis.ts +++ b/genkit-tools/common/src/types/apis.ts @@ -20,7 +20,6 @@ import { EvalInferenceInputSchema, EvalRunKeySchema, } from './eval'; -import { FlowStateSchema } from './flow'; import { GenerationCommonConfigSchema, MessageSchema, @@ -33,14 +32,7 @@ import { TraceDataSchema } from './trace'; * It's used directly in the generation of the Reflection API OpenAPI spec. */ -export const EnvTypesSchema = z - .enum(['dev', 'prod']) - .describe('Supported environments in the runtime.'); - -export type EnvTypes = z.infer; - export const ListTracesRequestSchema = z.object({ - env: EnvTypesSchema.optional(), limit: z.number().optional(), continuationToken: z.string().optional(), }); @@ -55,36 +47,11 @@ export const ListTracesResponseSchema = z.object({ export type ListTracesResponse = z.infer; export const GetTraceRequestSchema = z.object({ - env: EnvTypesSchema, traceId: z.string().describe('ID of the trace.'), }); export type GetTraceRequest = z.infer; -export const ListFlowStatesRequestSchema = z.object({ - env: EnvTypesSchema.optional(), - limit: z.number().optional(), - continuationToken: z.string().optional(), -}); - -export type ListFlowStatesRequest = z.infer; - -export const ListFlowStatesResponseSchema = z.object({ - flowStates: z.array(FlowStateSchema), - continuationToken: z.string().optional(), -}); - -export type ListFlowStatesResponse = z.infer< - typeof ListFlowStatesResponseSchema ->; - -export const GetFlowStateRequestSchema = z.object({ - env: EnvTypesSchema, - flowId: z.string().describe('ID of the flow state.'), -}); - -export type GetFlowStateRequest = z.infer; - export const RunActionRequestSchema = z.object({ key: z .string() diff --git a/go/genkit/servers.go b/go/genkit/servers.go index 64f71d2b7..c82520054 100644 --- a/go/genkit/servers.go +++ b/go/genkit/servers.go @@ -233,9 +233,8 @@ func newDevServeMux(s *devServer) *http.ServeMux { }) handle(mux, "POST /api/runAction", s.handleRunAction) handle(mux, "GET /api/actions", s.handleListActions) - handle(mux, "GET /api/envs/{env}/traces/{traceID}", s.handleGetTrace) - handle(mux, "GET /api/envs/{env}/traces", s.handleListTraces) - handle(mux, "GET /api/envs/{env}/flowStates", s.handleListFlowStates) + handle(mux, "GET /api/traces/{traceID}", s.handleGetTrace) + handle(mux, "GET /api/traces", s.handleListTraces) return mux } @@ -373,15 +372,6 @@ type listTracesResult struct { ContinuationToken string `json:"continuationToken"` } -func (s *devServer) handleListFlowStates(w http.ResponseWriter, r *http.Request) error { - return writeJSON(r.Context(), w, listFlowStatesResult{[]base.FlowStater{}, ""}) -} - -type listFlowStatesResult struct { - FlowStates []base.FlowStater `json:"flowStates"` - ContinuationToken string `json:"continuationToken"` -} - // NewFlowServeMux constructs a [net/http.ServeMux]. // If flows is non-empty, the each of the named flows is registered as a route. // Otherwise, all defined flows are registered. diff --git a/samples/js-coffee-shop/src/index.ts b/samples/js-coffee-shop/src/index.ts index e083b14a0..63c855714 100644 --- a/samples/js-coffee-shop/src/index.ts +++ b/samples/js-coffee-shop/src/index.ts @@ -28,9 +28,7 @@ configureGenkit({ dotprompt(), ], enableTracingAndMetrics: true, - flowStateStore: 'firebase', logLevel: 'debug', - traceStore: 'firebase', }); // This example generates greetings for a customer at our new AI-powered coffee shop, diff --git a/samples/js-menu/src/index.ts b/samples/js-menu/src/index.ts index 364737e2e..faf0c3efa 100644 --- a/samples/js-menu/src/index.ts +++ b/samples/js-menu/src/index.ts @@ -33,9 +33,7 @@ configureGenkit({ ]), ], enableTracingAndMetrics: true, - flowStateStore: 'firebase', logLevel: 'debug', - traceStore: 'firebase', }); // Export all of the example prompts and flows