Remove env and flow state APIs

firebase · Nov 7, 2024 · 1259b1f · 1259b1f
1 parent c358582
commit 1259b1f
Show file tree

Hide file tree

Showing 10 changed files with 14 additions and 160 deletions.
diff --git 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('<flowName>', 'name of the flow to run')
-  .option('--env <env>', 'environment (dev/prod)', 'dev')
   .option(
     '--output <filename>',
     '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
@@ -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
@@ -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<FlowState> {
-  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<FlowState> {
-  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
@@ -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
@@ -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
@@ -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<apis.ListTracesResponse> {
-    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<TraceData> {
-    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
@@ -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<typeof EnvTypesSchema>;
-
 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<typeof ListTracesResponseSchema>;
 
 export const GetTraceRequestSchema = z.object({
-  env: EnvTypesSchema,
   traceId: z.string().describe('ID of the trace.'),
 });
 
 export type GetTraceRequest = z.infer<typeof GetTraceRequestSchema>;
 
-export const ListFlowStatesRequestSchema = z.object({
-  env: EnvTypesSchema.optional(),
-  limit: z.number().optional(),
-  continuationToken: z.string().optional(),
-});
-
-export type ListFlowStatesRequest = z.infer<typeof ListFlowStatesRequestSchema>;
-
-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<typeof GetFlowStateRequestSchema>;
-
 export const RunActionRequestSchema = z.object({
   key: z
     .string()

diff --git 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
@@ -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
@@ -33,9 +33,7 @@ configureGenkit({
     ]),
   ],
   enableTracingAndMetrics: true,
-  flowStateStore: 'firebase',
   logLevel: 'debug',
-  traceStore: 'firebase',
 });
 
 // Export all of the example prompts and flows