From 52c32d72b3d1d8ca93614d4ded2861058f2d4f6a Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Thu, 12 Mar 2026 14:49:12 -0400 Subject: [PATCH 01/10] feat(js,cli): implemented reflection api v2 --- docs/reflection-v2-protocol.md | 234 ++++++++ .../cli/src/commands/dev-test-model.ts | 12 +- genkit-tools/cli/src/commands/start.ts | 11 +- genkit-tools/cli/src/mcp/utils.ts | 6 +- genkit-tools/cli/src/utils/manager-utils.ts | 25 +- genkit-tools/common/package.json | 6 +- genkit-tools/common/src/eval/evaluate.ts | 22 +- genkit-tools/common/src/eval/validate.ts | 6 +- genkit-tools/common/src/manager/index.ts | 2 +- genkit-tools/common/src/manager/manager-v2.ts | 525 ++++++++++++++++++ genkit-tools/common/src/manager/manager.ts | 395 +++++++------ genkit-tools/common/src/server/router.ts | 4 +- genkit-tools/common/src/server/server.ts | 107 +++- genkit-tools/common/src/types/index.ts | 1 + genkit-tools/common/src/types/reflection.ts | 116 ++++ genkit-tools/common/src/utils/eval.ts | 6 +- genkit-tools/common/tests/manager-v2_test.ts | 443 +++++++++++++++ genkit-tools/common/tests/server_test.ts | 165 ++++++ genkit-tools/genkit-schema.json | 249 +++++++++ genkit-tools/pnpm-lock.yaml | 44 ++ genkit-tools/scripts/schema-exporter.ts | 1 + genkit-tools/telemetry-server/package.json | 2 +- go/ai/gen.go | 80 +++ go/internal/cmd/jsonschemagen/jsonschema.go | 1 + .../cmd/jsonschemagen/jsonschemagen.go | 65 ++- js/ai/src/model-types.ts | 2 +- js/ai/tests/prompt/prompt_test.ts | 2 +- js/core/package.json | 4 +- js/core/src/reflection-types.ts | 120 ++++ js/core/src/reflection-v2.ts | 396 +++++++++++++ js/core/src/reflection.ts | 21 + js/core/tests/reflection-v2_test.ts | 377 +++++++++++++ js/pnpm-lock.yaml | 96 +++- 33 files changed, 3305 insertions(+), 241 deletions(-) create mode 100644 docs/reflection-v2-protocol.md create mode 100644 genkit-tools/common/src/manager/manager-v2.ts create mode 100644 genkit-tools/common/src/types/reflection.ts create mode 100644 genkit-tools/common/tests/manager-v2_test.ts create mode 100644 genkit-tools/common/tests/server_test.ts create mode 100644 js/core/src/reflection-types.ts create mode 100644 js/core/src/reflection-v2.ts create mode 100644 js/core/tests/reflection-v2_test.ts diff --git a/docs/reflection-v2-protocol.md b/docs/reflection-v2-protocol.md new file mode 100644 index 0000000000..6c9d7cd168 --- /dev/null +++ b/docs/reflection-v2-protocol.md @@ -0,0 +1,234 @@ +# Genkit Reflection Protocol V2 (WebSocket) + +This document outlines the design for the V2 Reflection API, which uses WebSockets for bidirectional communication between the Genkit CLI (Runtime Manager) and Genkit Runtimes (User Applications). + +## Overview + +In V2, the connection direction is reversed compared to V1: +- **Server**: The Genkit CLI (`RuntimeManagerV2`) starts a WebSocket server. +- **Client**: The Genkit Runtime connects to the CLI's WebSocket server. + +This architecture allows the CLI to easily manage multiple runtimes (e.g., for multi-service projects) and eliminates the need for runtimes to manage their own HTTP servers and ports for reflection. + +## Transport + +| Feature | Specification | +| :--- | :--- | +| **Protocol** | WebSocket | +| **Data Format** | JSON | +| **Message Structure** | JSON-RPC 2.0 (modified for streaming) | + +## Message Format + +All messages follow the JSON-RPC 2.0 specification. + +### Request +```json +{ + "jsonrpc": "2.0", + "method": "methodName", + "params": { ... }, + "id": 1 +} +``` +*Note: The `id` is generated by the sender (Manager). It can be a number (auto-incrementing) or a string (UUID). It must be unique for the pending request within the WebSocket session.* + +### Response (Success) +```json +{ + "jsonrpc": "2.0", + "result": { ... }, + "id": 1 +} +``` + +### Response (Error) +```json +{ + "jsonrpc": "2.0", + "error": { + "code": -32000, + "message": "Error message", + "data": { + "code": 13, + "message": "Error message", + "details": { + "traceId": "...", + "stack": "..." + } + } + }, + "id": 1 +} +``` + +The `data` field contains a `Status` object (matching V1 API) with: +- **`code`**: Genkit canonical status code (e.g., 13 for INTERNAL, 3 for INVALID_ARGUMENT). +- **`message`**: The error message. +- **`details`**: Additional context, including `traceId` and `stack` trace. + +### Notification +A request without an `id`. +```json +{ + "jsonrpc": "2.0", + "method": "methodName", + "params": { ... } +} +``` + +## Streaming Extension + +JSON-RPC 2.0 does not natively support streaming. We extend it by using Notifications from the Runtime to the Manager associated with a specific Request ID. + +| Message Type | Method | Direction | Description | +| :--- | :--- | :--- | :--- | +| **Stream Chunk** | `streamChunk` | Runtime -> Manager | Sent by the Runtime during a streaming `runAction` request. | +| **State Update** | `runActionState` | Runtime -> Manager | Sent by the Runtime to provide status updates (e.g., trace ID) before the result. | + +### Stream Chunk Notification +```json +{ + "jsonrpc": "2.0", + "method": "streamChunk", + "params": { + "requestId": 1, + "chunk": { ... } + } +} +``` + +### Run Action State Notification +```json +{ + "jsonrpc": "2.0", + "method": "runActionState", + "params": { + "requestId": 1, + "state": { "traceId": "..." } + } +} +``` + +## Protocol Methods Summary + +| Method | Direction | Type | Description | +| :--- | :--- | :--- | :--- | +| **`register`** | Runtime -> Manager | Request | Registers the runtime with the Manager. | +| **`configure`** | Manager -> Runtime | Notification | Pushes configuration updates to the Runtime. | +| **`listActions`** | Manager -> Runtime | Request | Retrieves the list of available actions. | +| **`listValues`** | Manager -> Runtime | Request | Retrieves the list of values (prompts, schemas, etc.). | +| **`runAction`** | Manager -> Runtime | Request | Executes an action. | +| **`cancelAction`** | Manager -> Runtime | Request | Cancels a running action. | + +## Detailed API + +### 1. Registration +**Direction:** Runtime -> Manager +**Type:** Request + +**Parameters:** +| Field | Type | Description | +| :--- | :--- | :--- | +| `id` | `string` | Unique Runtime ID. | +| `pid` | `number` | Process ID. | +| `name` | `string` | App name (optional). | +| `genkitVersion` | `string` | e.g., "0.9.0". | +| `reflectionApiSpecVersion` | `number` | Protocol version. | +| `envs` | `string[]` | Configured environments (optional). | + +**Result:** `void` + +### 2. Configuration +**Direction:** Manager -> Runtime +**Type:** Notification + +**Parameters:** +| Field | Type | Description | +| :--- | :--- | :--- | +| `telemetryServerUrl` | `string` | URL of the telemetry server (optional). | + +### 3. List Actions +**Direction:** Manager -> Runtime +**Type:** Request + +**Parameters:** `void` + +**Result:** +| Type | Description | +| :--- | :--- | +| `Record` | Map of action keys to Action definitions. (Same schema as V1 `/api/actions`) | + +### 4. List Values +**Direction:** Manager -> Runtime +**Type:** Request + +**Parameters:** +| Field | Type | Description | +| :--- | :--- | :--- | +| `type` | `string` | The type of value to list (e.g., "model", "prompt", "schema"). | + +**Result:** +| Type | Description | +| :--- | :--- | +| `Record` | Map of value keys to value definitions. | + +### 5. Run Action +**Direction:** Manager -> Runtime +**Type:** Request + +**Parameters:** +| Field | Type | Description | +| :--- | :--- | :--- | +| `key` | `string` | Action key (e.g., "/flow/myFlow"). | +| `input` | `any` | Input payload. | +| `context` | `any` | Context data (optional). | +| `telemetryLabels` | `Record` | Telemetry labels (optional). | +| `stream` | `boolean` | Whether to stream results. | +| `streamInput` | `boolean` | Whether to stream input (for bidi actions). | + +**Result (Non-Streaming):** +| Field | Type | Description | +| :--- | :--- | :--- | +| `result` | `any` | The return value. | +| `telemetry` | `object` | Telemetry metadata (e.g., `{ traceId: string }`). | + +**Streaming Flow:** +1. Runtime sends optional `runActionState` notifications. +2. Runtime sends `streamChunk` notifications. +3. Runtime sends final response with `result` (same structure as non-streaming). + +**Bidirectional Streaming Flow (if `streamInput: true`):** +1. Manager sends `sendInputStreamChunk` notifications. +2. Manager sends `endInputStream` notification. +3. Runtime behaves as per Streaming Flow. + +### 6. Cancel Action +**Direction:** Manager -> Runtime +**Type:** Request + +**Parameters:** +| Field | Type | Description | +| :--- | :--- | :--- | +| `traceId` | `string` | The trace ID of the action to cancel. | + +**Result:** +| Field | Type | Description | +| :--- | :--- | :--- | +| `message` | `string` | Confirmation message. | + +## Health Checks + +| Check Type | Description | +| :--- | :--- | +| **Connection State** | The WebSocket connection state itself serves as a basic health check. | +| **Heartbeats** | Standard WebSocket Ping/Pong frames should be used to maintain the connection and detect timeouts. | + +## Compatibility + +| Version | Architecture | +| :--- | :--- | +| **V1** | HTTP Server on Runtime, Polling/Request from CLI. | +| **V2** | WebSocket Server on CLI, Persistent Connection from Runtime. | + +The CLI will determine which mode to use based on configuration (e.g., `--experimental-reflection-v2`). diff --git a/genkit-tools/cli/src/commands/dev-test-model.ts b/genkit-tools/cli/src/commands/dev-test-model.ts index 3e071708cd..2fd34a7ec9 100644 --- a/genkit-tools/cli/src/commands/dev-test-model.ts +++ b/genkit-tools/cli/src/commands/dev-test-model.ts @@ -22,8 +22,8 @@ import { Part, } from '@genkit-ai/tools-common'; import { + BaseRuntimeManager, GenkitToolsError, - RuntimeManager, } from '@genkit-ai/tools-common/manager'; import { findProjectRoot, logger } from '@genkit-ai/tools-common/utils'; import { Command } from 'commander'; @@ -475,7 +475,7 @@ const TEST_CASES: Record = { }, }; -async function waitForRuntime(manager: RuntimeManager) { +async function waitForRuntime(manager: BaseRuntimeManager) { // Poll for runtimes for (let i = 0; i < 20; i++) { if (manager.listRuntimes().length > 0) return; @@ -490,7 +490,7 @@ async function waitForRuntime(manager: RuntimeManager) { * dispatching tests before the runtime has finished registering actions. */ async function waitForActions( - manager: RuntimeManager, + manager: BaseRuntimeManager, suites: TestSuite[] ): Promise { const requiredKeys = new Set(); @@ -535,7 +535,7 @@ async function waitForActions( } async function runTest( - manager: RuntimeManager, + manager: BaseRuntimeManager, model: string, testCase: TestCase ): Promise { @@ -591,7 +591,7 @@ async function runTest( } async function runTestSuite( - manager: RuntimeManager, + manager: BaseRuntimeManager, suite: TestSuite, defaultSupports: string[] ): Promise<{ passed: number; failed: number }> { @@ -665,7 +665,7 @@ export const devTestModel = new Command('dev:test-model') if (args) cmd = args; } - let manager: RuntimeManager; + let manager: BaseRuntimeManager; if (cmd.length > 0) { const result = await startDevProcessManager( diff --git a/genkit-tools/cli/src/commands/start.ts b/genkit-tools/cli/src/commands/start.ts index 074d36c46d..88ca318b14 100644 --- a/genkit-tools/cli/src/commands/start.ts +++ b/genkit-tools/cli/src/commands/start.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import type { RuntimeManager } from '@genkit-ai/tools-common/manager'; +import type { BaseRuntimeManager } from '@genkit-ai/tools-common/manager'; import { startServer } from '@genkit-ai/tools-common/server'; import { findProjectRoot, logger } from '@genkit-ai/tools-common/utils'; import { Command } from 'commander'; @@ -28,6 +28,7 @@ interface RunOptions { open?: boolean; disableRealtimeTelemetry?: boolean; corsOrigin?: string; + experimentalReflectionV2?: boolean; } /** Command to run code in dev mode and/or the Dev UI. */ @@ -44,6 +45,10 @@ export const start = new Command('start') '--cors-origin ', 'specify the allowed origin for CORS requests' ) + .option( + '--experimental-reflection-v2', + 'start the experimental reflection server (WebSocket)' + ) .action(async (options: RunOptions) => { const projectRoot = await findProjectRoot(); if (projectRoot.includes('/.Trash/')) { @@ -53,7 +58,7 @@ export const start = new Command('start') ); } // Always start the manager. - let manager: RuntimeManager; + let manager: BaseRuntimeManager; let processPromise: Promise | undefined; if (start.args.length > 0) { const result = await startDevProcessManager( @@ -63,6 +68,7 @@ export const start = new Command('start') { disableRealtimeTelemetry: options.disableRealtimeTelemetry, corsOrigin: options.corsOrigin, + experimentalReflectionV2: options.experimentalReflectionV2, } ); manager = result.manager; @@ -72,6 +78,7 @@ export const start = new Command('start') projectRoot, manageHealth: true, corsOrigin: options.corsOrigin, + experimentalReflectionV2: options.experimentalReflectionV2, }); processPromise = new Promise(() => {}); } diff --git a/genkit-tools/cli/src/mcp/utils.ts b/genkit-tools/cli/src/mcp/utils.ts index 2a97c24f5d..2c9e7d9756 100644 --- a/genkit-tools/cli/src/mcp/utils.ts +++ b/genkit-tools/cli/src/mcp/utils.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { RuntimeManager } from '@genkit-ai/tools-common/manager'; +import { BaseRuntimeManager } from '@genkit-ai/tools-common/manager'; import { z } from 'zod'; import { startDevProcessManager, startManager } from '../utils/manager-utils'; @@ -62,7 +62,7 @@ export function resolveProjectRoot( /** Genkit Runtime manager specifically for the MCP server. Allows lazy * initialization and dev process manangement. */ export class McpRuntimeManager { - private manager: RuntimeManager | undefined; + private manager: BaseRuntimeManager | undefined; private currentProjectRoot: string | undefined; async getManager(projectRoot: string) { @@ -86,7 +86,7 @@ export class McpRuntimeManager { args: string[]; explicitProjectRoot: boolean; timeout?: number; - }): Promise { + }): Promise { const { projectRoot, command, args, timeout, explicitProjectRoot } = params; if (this.manager) { await this.manager.stop(); diff --git a/genkit-tools/cli/src/utils/manager-utils.ts b/genkit-tools/cli/src/utils/manager-utils.ts index a4b2efa98e..36d2afa5b7 100644 --- a/genkit-tools/cli/src/utils/manager-utils.ts +++ b/genkit-tools/cli/src/utils/manager-utils.ts @@ -20,6 +20,7 @@ import { } from '@genkit-ai/telemetry-server'; import type { Status } from '@genkit-ai/tools-common'; import { + BaseRuntimeManager, ProcessManager, RuntimeEvent, RuntimeManager, @@ -60,12 +61,14 @@ export async function startManager(options: { projectRoot: string; manageHealth?: boolean; corsOrigin?: string; -}): Promise { + experimentalReflectionV2?: boolean; +}): Promise { const telemetryServerUrl = await resolveTelemetryServer(options); const manager = RuntimeManager.create({ telemetryServerUrl, manageHealth: options.manageHealth, projectRoot: options.projectRoot, + experimentalReflectionV2: options.experimentalReflectionV2, }); return manager; } @@ -77,6 +80,7 @@ export interface DevProcessManagerOptions { timeout?: number; cwd?: string; corsOrigin?: string; + experimentalReflectionV2?: boolean; } export async function startDevProcessManager( @@ -84,16 +88,25 @@ export async function startDevProcessManager( command: string, args: string[], options?: DevProcessManagerOptions -): Promise<{ manager: RuntimeManager; processPromise: Promise }> { +): Promise<{ manager: BaseRuntimeManager; processPromise: Promise }> { const telemetryServerUrl = await resolveTelemetryServer({ projectRoot, corsOrigin: options?.corsOrigin, }); const disableRealtimeTelemetry = options?.disableRealtimeTelemetry ?? false; + const experimentalReflectionV2 = options?.experimentalReflectionV2 ?? false; + + let reflectionV2Port: number | undefined; const envVars: Record = { GENKIT_TELEMETRY_SERVER: telemetryServerUrl, GENKIT_ENV: 'dev', }; + + if (experimentalReflectionV2) { + reflectionV2Port = await getPort({ port: makeRange(3200, 3400) }); + envVars.GENKIT_REFLECTION_V2_SERVER = `ws://localhost:${reflectionV2Port}`; + } + if (!disableRealtimeTelemetry) { envVars.GENKIT_ENABLE_REALTIME_TELEMETRY = 'true'; } @@ -105,6 +118,8 @@ export async function startDevProcessManager( projectRoot, processManager, disableRealtimeTelemetry, + experimentalReflectionV2, + reflectionV2Port, }); const processPromise = processManager.start({ ...options }); @@ -120,7 +135,7 @@ export async function startDevProcessManager( * Rejects if the process exits or if the timeout is reached. */ export async function waitForRuntime( - manager: RuntimeManager, + manager: BaseRuntimeManager, processPromise: Promise, timeoutMs: number = 30000 ): Promise { @@ -173,9 +188,9 @@ export async function waitForRuntime( */ export async function runWithManager( projectRoot: string, - fn: (manager: RuntimeManager) => Promise + fn: (manager: BaseRuntimeManager) => Promise ) { - let manager: RuntimeManager; + let manager: BaseRuntimeManager; try { manager = await startManager({ projectRoot, manageHealth: false }); // Don't manage health in this case. } catch (e) { diff --git a/genkit-tools/common/package.json b/genkit-tools/common/package.json index 2d473dd7a1..3c0fa357bd 100644 --- a/genkit-tools/common/package.json +++ b/genkit-tools/common/package.json @@ -22,6 +22,7 @@ "commander": "^11.1.0", "configstore": "^5.0.1", "cors": "^2.8.5", + "events": "^3.3.0", "express": "^4.21.0", "get-port": "5.1.1", "glob": "^10.3.12", @@ -34,7 +35,8 @@ "winston": "^3.11.0", "yaml": "^2.4.1", "zod": "^3.22.4", - "zod-to-json-schema": "^3.22.4" + "zod-to-json-schema": "^3.22.4", + "ws": "^8.18.3" }, "devDependencies": { "@jest/globals": "^29.7.0", @@ -43,6 +45,7 @@ "@types/cli-color": "^2.0.6", "@types/configstore": "^6.0.2", "@types/cors": "^2.8.19", + "@types/events": "^3.0.3", "@types/express": "^4.17.21", "@types/inquirer": "^8.1.3", "@types/jest": "^29.5.12", @@ -50,6 +53,7 @@ "@types/json-schema": "^7.0.15", "@types/node": "^20.11.19", "@types/uuid": "^9.0.8", + "@types/ws": "^8.18.1", "bun-types": "^1.2.16", "genversion": "^3.2.0", "jest": "^29.7.0", diff --git a/genkit-tools/common/src/eval/evaluate.ts b/genkit-tools/common/src/eval/evaluate.ts index e9677449ec..1b636d7064 100644 --- a/genkit-tools/common/src/eval/evaluate.ts +++ b/genkit-tools/common/src/eval/evaluate.ts @@ -17,7 +17,7 @@ import { randomUUID } from 'crypto'; import { z } from 'zod'; import { getDatasetStore, getEvalStore } from '.'; -import type { RuntimeManager } from '../manager/manager'; +import type { BaseRuntimeManager } from '../manager/manager'; import { GenerateActionOptions, GenerateResponseData, @@ -70,7 +70,7 @@ const GENERATE_ACTION_UTIL = '/util/generate'; * Starts a new evaluation run. Intended to be used via the reflection API. */ export async function runNewEvaluation( - manager: RuntimeManager, + manager: BaseRuntimeManager, request: RunNewEvaluationRequest ): Promise { const { dataSource, actionRef, evaluators } = request; @@ -139,7 +139,7 @@ export async function runNewEvaluation( /** Handles the Inference part of Inference-Evaluation cycle */ export async function runInference(params: { - manager: RuntimeManager; + manager: BaseRuntimeManager; actionRef: string; inferenceDataset: Dataset; context?: string; @@ -163,7 +163,7 @@ export async function runInference(params: { /** Handles the Evaluation part of Inference-Evaluation cycle */ export async function runEvaluation(params: { - manager: RuntimeManager; + manager: BaseRuntimeManager; evaluatorActions: Action[]; evalDataset: EvalInput[]; augments?: EvalKeyAugments; @@ -219,7 +219,7 @@ export async function runEvaluation(params: { } export async function getAllEvaluatorActions( - manager: RuntimeManager + manager: BaseRuntimeManager ): Promise { const allActions = await manager.listActions(); const allEvaluatorActions = []; @@ -232,7 +232,7 @@ export async function getAllEvaluatorActions( } export async function getMatchingEvaluatorActions( - manager: RuntimeManager, + manager: BaseRuntimeManager, evaluators?: string[] ): Promise { if (!evaluators) { @@ -251,7 +251,7 @@ export async function getMatchingEvaluatorActions( } async function bulkRunAction(params: { - manager: RuntimeManager; + manager: BaseRuntimeManager; actionRef: string; inferenceDataset: Dataset; context?: string; @@ -313,7 +313,7 @@ async function bulkRunAction(params: { } async function runFlowAction(params: { - manager: RuntimeManager; + manager: BaseRuntimeManager; actionRef: string; sample: FullInferenceSample; context?: any; @@ -345,7 +345,7 @@ async function runFlowAction(params: { } async function runModelAction(params: { - manager: RuntimeManager; + manager: BaseRuntimeManager; actionRef: string; sample: FullInferenceSample; modelConfig?: any; @@ -377,7 +377,7 @@ async function runModelAction(params: { } async function runPromptAction(params: { - manager: RuntimeManager; + manager: BaseRuntimeManager; actionRef: string; sample: FullInferenceSample; context?: any; @@ -464,7 +464,7 @@ async function runPromptAction(params: { } async function gatherEvalInput(params: { - manager: RuntimeManager; + manager: BaseRuntimeManager; actionRef: string; state: InferenceRunState; }): Promise { diff --git a/genkit-tools/common/src/eval/validate.ts b/genkit-tools/common/src/eval/validate.ts index 7453f8c507..3f204da3a8 100644 --- a/genkit-tools/common/src/eval/validate.ts +++ b/genkit-tools/common/src/eval/validate.ts @@ -17,7 +17,7 @@ import Ajv, { type ErrorObject, type JSONSchemaType } from 'ajv'; import addFormats from 'ajv-formats'; import { getDatasetStore } from '.'; -import type { RuntimeManager } from '../manager'; +import type { BaseRuntimeManager } from '../manager'; import { InferenceDatasetSchema, type Action, @@ -35,7 +35,7 @@ type JSONSchema = JSONSchemaType | any; * reflection API. */ export async function validateSchema( - manager: RuntimeManager, + manager: BaseRuntimeManager, request: ValidateDataRequest ): Promise { const { dataSource, actionRef } = request; @@ -125,7 +125,7 @@ function toErrorDetail(error: ErrorObject): ErrorDetail { } async function getAction( - manager: RuntimeManager, + manager: BaseRuntimeManager, actionRef: string ): Promise { const actions = await manager.listActions(); diff --git a/genkit-tools/common/src/manager/index.ts b/genkit-tools/common/src/manager/index.ts index d28014cc0b..b72c27125b 100644 --- a/genkit-tools/common/src/manager/index.ts +++ b/genkit-tools/common/src/manager/index.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -export { RuntimeManager } from './manager'; +export { BaseRuntimeManager, RuntimeManager } from './manager'; export { AppProcessStatus, ProcessManager, diff --git a/genkit-tools/common/src/manager/manager-v2.ts b/genkit-tools/common/src/manager/manager-v2.ts new file mode 100644 index 0000000000..286cbf0ead --- /dev/null +++ b/genkit-tools/common/src/manager/manager-v2.ts @@ -0,0 +1,525 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import EventEmitter from 'events'; +import getPort, { makeRange } from 'get-port'; +import { WebSocket, WebSocketServer } from 'ws'; +import { + Action, + RunActionResponse, + RunActionResponseSchema, +} from '../types/action'; +import * as apis from '../types/apis'; +import { + ReflectionCancelActionParamsSchema, + ReflectionCancelActionResponseSchema, + ReflectionConfigureParamsSchema, + ReflectionEndInputStreamParamsSchema, + ReflectionListActionsResponseSchema, + ReflectionListValuesParamsSchema, + ReflectionListValuesResponseSchema, + ReflectionRegisterParamsSchema, + ReflectionRunActionParamsSchema, + ReflectionRunActionStateParamsSchema, + ReflectionSendInputStreamChunkParamsSchema, + ReflectionStreamChunkParamsSchema, +} from '../types/reflection'; +import { logger } from '../utils/logger'; +import { DevToolsInfo } from '../utils/utils'; +import { BaseRuntimeManager, RuntimeManagerOptions } from './manager'; +import { ProcessManager } from './process-manager'; +import { + GenkitToolsError, + RuntimeEvent, + RuntimeInfo, + StreamingCallback, +} from './types'; + +interface JsonRpcRequest { + jsonrpc: '2.0'; + method: string; + params?: any; + id?: number | string; +} + +interface JsonRpcResponse { + jsonrpc: '2.0'; + result?: any; + error?: { + code: number; + message: string; + data?: any; + }; + id: number | string; +} + +type JsonRpcMessage = JsonRpcRequest | JsonRpcResponse; + +interface ConnectedRuntime { + ws: WebSocket; + info: RuntimeInfo; +} + +export class RuntimeManagerV2 extends BaseRuntimeManager { + private _port?: number; + private wss?: WebSocketServer; + private runtimes: Map = new Map(); + + get port(): number | undefined { + return this._port; + } + private pendingRequests: Map< + number | string, + { + method: string; + resolve: (value: any) => void; + reject: (reason?: any) => void; + } + > = new Map(); + private streamCallbacks: Map> = + new Map(); + private traceIdCallbacks: Map void> = + new Map(); + private eventEmitter = new EventEmitter(); + private requestIdCounter = 0; + + constructor( + telemetryServerUrl: string | undefined, + readonly manageHealth: boolean, + projectRoot: string, + processManager?: ProcessManager, + disableRealtimeTelemetry: boolean = false + ) { + super( + telemetryServerUrl, + projectRoot, + processManager, + disableRealtimeTelemetry + ); + } + + static async create( + options: RuntimeManagerOptions + ): Promise { + const manager = new RuntimeManagerV2( + options.telemetryServerUrl, + options.manageHealth ?? true, + options.projectRoot, + options.processManager, + options.disableRealtimeTelemetry + ); + await manager.startWebSocketServer(options.reflectionV2Port); + return manager; + } + + /** + * Starts a WebSocket server. + */ + private async startWebSocketServer(port?: number): Promise<{ port: number }> { + if (!port) { + port = await getPort({ port: makeRange(3200, 3400) }); + } + this.wss = new WebSocketServer({ port }); + + this._port = port; + logger.info(`Starting reflection server: ws://localhost:${port}`); + + this.wss.on('connection', (ws) => { + ws.on('error', (err) => logger.error(`WebSocket error: ${err}`)); + + ws.on('message', (data) => { + try { + const message = JSON.parse(data.toString()) as JsonRpcMessage; + this.handleMessage(ws, message); + } catch (error) { + logger.error('Failed to parse WebSocket message:', error); + } + }); + + ws.on('close', () => { + this.handleDisconnect(ws); + }); + }); + return { port }; + } + + private handleMessage(ws: WebSocket, message: JsonRpcMessage) { + if ('method' in message) { + this.handleRequest(ws, message as JsonRpcRequest); + } else { + this.handleResponse(message as JsonRpcResponse); + } + } + + private handleRequest(ws: WebSocket, request: JsonRpcRequest) { + switch (request.method) { + case 'register': + this.handleRegister(ws, request); + break; + case 'streamChunk': + this.handleStreamChunk(request); + break; + case 'runActionState': + this.handleRunActionState(request); + break; + default: + logger.warn(`Unknown method: ${request.method}`); + } + } + + private handleRegister(ws: WebSocket, request: JsonRpcRequest) { + const params = ReflectionRegisterParamsSchema.parse(request.params); + const runtimeInfo: RuntimeInfo = { + id: params.id, + pid: params.pid, + name: params.name, + genkitVersion: params.genkitVersion, + reflectionApiSpecVersion: params.reflectionApiSpecVersion, + reflectionServerUrl: `ws://localhost:${this.port}`, // Virtual URL for compatibility + timestamp: new Date().toISOString(), + projectName: params.name || 'Unknown', // Or derive from other means if needed + }; + + this.runtimes.set(runtimeInfo.id, { ws, info: runtimeInfo }); + this.eventEmitter.emit(RuntimeEvent.ADD, runtimeInfo); + + // Send success response + if (request.id) { + ws.send( + JSON.stringify({ + jsonrpc: '2.0', + result: null, + id: request.id, + }) + ); + } + + // Configure the runtime immediately + this.notifyRuntime(runtimeInfo.id); + } + + private handleStreamChunk(notification: JsonRpcRequest) { + const { requestId, chunk } = ReflectionStreamChunkParamsSchema.parse( + notification.params + ); + const callback = this.streamCallbacks.get(requestId); + if (callback) { + callback(chunk); + } + } + + private handleRunActionState(notification: JsonRpcRequest) { + const { requestId, state } = ReflectionRunActionStateParamsSchema.parse( + notification.params + ); + const callback = this.traceIdCallbacks.get(requestId); + if (callback && state?.traceId) { + callback(state.traceId); + } + } + + private handleResponse(response: JsonRpcResponse) { + const pending = this.pendingRequests.get(response.id); + if (pending) { + if (response.error) { + const errorData = response.error.data || {}; + const massagedData = { + ...errorData, + stack: errorData.details?.stack, + data: { + genkitErrorMessage: errorData.message, + genkitErrorDetails: errorData.details, + }, + }; + const error = new GenkitToolsError(response.error.message); + error.data = massagedData; + pending.reject(error); + } else { + let result = response.result; + if (pending.method === 'listActions') { + result = ReflectionListActionsResponseSchema.parse(result); + } else if (pending.method === 'listValues') { + result = ReflectionListValuesResponseSchema.parse(result); + } else if (pending.method === 'cancelAction') { + result = ReflectionCancelActionResponseSchema.parse(result); + } + pending.resolve(result); + } + this.pendingRequests.delete(response.id); + } else { + logger.warn(`Received response for unknown request ID ${response.id}`); + } + } + + private handleDisconnect(ws: WebSocket) { + for (const [id, runtime] of this.runtimes.entries()) { + if (runtime.ws === ws) { + this.runtimes.delete(id); + this.eventEmitter.emit(RuntimeEvent.REMOVE, runtime.info); + break; + } + } + } + + private async sendRequest( + runtimeId: string, + method: string, + params?: any + ): Promise { + const runtime = this.runtimes.get(runtimeId); + if (!runtime) { + throw new Error(`Runtime ${runtimeId} not found`); + } + + switch (method) { + case 'listValues': + ReflectionListValuesParamsSchema.parse(params); + break; + case 'runAction': + ReflectionRunActionParamsSchema.parse(params); + break; + case 'cancelAction': + ReflectionCancelActionParamsSchema.parse(params); + break; + } + + const id = (++this.requestIdCounter).toString(); + const message: JsonRpcRequest = { + jsonrpc: '2.0', + method, + params, + id, + }; + + return new Promise((resolve, reject) => { + const timeoutId = setTimeout(() => { + if (this.pendingRequests.has(id)) { + this.pendingRequests.delete(id); + reject(new Error(`Request ${id} timed out`)); + } + }, 30000); + + this.pendingRequests.set(id, { + method, + resolve: (value) => { + clearTimeout(timeoutId); + resolve(value); + }, + reject: (reason) => { + clearTimeout(timeoutId); + reject(reason); + }, + }); + + runtime.ws.send(JSON.stringify(message)); + }); + } + + private sendNotification(runtimeId: string, method: string, params?: any) { + const runtime = this.runtimes.get(runtimeId); + if (!runtime) { + logger.warn(`Runtime ${runtimeId} not found, cannot send notification`); + return; + } + + switch (method) { + case 'configure': + ReflectionConfigureParamsSchema.parse(params); + break; + case 'sendInputStreamChunk': + ReflectionSendInputStreamChunkParamsSchema.parse(params); + break; + case 'endInputStream': + ReflectionEndInputStreamParamsSchema.parse(params); + break; + } + + const message: JsonRpcRequest = { + jsonrpc: '2.0', + method, + params, + }; + runtime.ws.send(JSON.stringify(message)); + } + + private notifyRuntime(runtimeId: string) { + this.sendNotification(runtimeId, 'configure', { + telemetryServerUrl: this.telemetryServerUrl, + }); + } + + listRuntimes(): RuntimeInfo[] { + return Array.from(this.runtimes.values()).map((r) => r.info); + } + + getRuntimeById(id: string): RuntimeInfo | undefined { + return this.runtimes.get(id)?.info; + } + + getMostRecentRuntime(): RuntimeInfo | undefined { + const runtimes = this.listRuntimes(); + if (runtimes.length === 0) return undefined; + return runtimes[runtimes.length - 1]; + } + + getMostRecentDevUI(): DevToolsInfo | undefined { + // Not applicable for V2 yet + return undefined; + } + + onRuntimeEvent( + listener: (eventType: RuntimeEvent, runtime: RuntimeInfo) => void + ) { + const listeners: Array<{ event: string; fn: (rt: RuntimeInfo) => void }> = + []; + Object.values(RuntimeEvent).forEach((event) => { + const fn = (rt: RuntimeInfo) => listener(event, rt); + this.eventEmitter.on(event, fn); + listeners.push({ event, fn }); + }); + return () => { + listeners.forEach(({ event, fn }) => { + this.eventEmitter.off(event, fn); + }); + }; + } + + async listActions( + input?: apis.ListActionsRequest + ): Promise> { + const runtimeId = input?.runtimeId || this.getMostRecentRuntime()?.id; + if (!runtimeId) { + throw new Error( + input?.runtimeId + ? `No runtime found with ID ${input.runtimeId}.` + : 'No runtimes found. Make sure your app is running using the `start_runtime` MCP tool or the CLI: `genkit start -- ...`. See getting started documentation.' + ); + } + return this.sendRequest(runtimeId, 'listActions'); + } + + async listValues( + input: apis.ListValuesRequest + ): Promise> { + const runtimeId = input?.runtimeId || this.getMostRecentRuntime()?.id; + if (!runtimeId) { + throw new Error( + input?.runtimeId + ? `No runtime found with ID ${input.runtimeId}.` + : 'No runtimes found. Make sure your app is running using `genkit start -- ...`. See getting started documentation.' + ); + } + return this.sendRequest(runtimeId, 'listValues', { type: input.type }); + } + + async stop() { + if (this.wss) { + this.wss.close(); + } + if (this.processManager) { + await this.processManager.kill(); + } + } + + async runAction( + input: apis.RunActionRequest, + streamingCallback?: StreamingCallback, + onTraceId?: (traceId: string) => void, + inputStream?: AsyncIterable + ): Promise { + const runtimeId = input.runtimeId || this.getMostRecentRuntime()?.id; + if (!runtimeId) { + throw new Error( + 'No runtimes found. Make sure your app is running using the `start_runtime` MCP tool or the CLI: `genkit start -- ...`. See getting started documentation.' + ); + } + + const runtime = this.runtimes.get(runtimeId); + if (!runtime) { + throw new Error(`Runtime ${runtimeId} not found`); + } + + const id = (++this.requestIdCounter).toString(); + + if (streamingCallback) { + this.streamCallbacks.set(id, streamingCallback); + } + if (onTraceId) { + this.traceIdCallbacks.set(id, onTraceId!); + } + + const message: JsonRpcRequest = { + jsonrpc: '2.0', + method: 'runAction', + params: { + ...input, + stream: !!streamingCallback, + streamInput: !!inputStream, + }, + id, + }; + + const promise = new Promise((resolve, reject) => { + this.pendingRequests.set(id, { method: 'runAction', resolve, reject }); + runtime!.ws.send(JSON.stringify(message)); + }) + .then((result) => { + return RunActionResponseSchema.parse(result); + }) + .finally(() => { + if (streamingCallback) { + this.streamCallbacks.delete(id); + } + if (onTraceId) { + this.traceIdCallbacks.delete(id); + } + }); + + if (inputStream) { + (async () => { + try { + for await (const chunk of inputStream!) { + this.sendNotification(runtimeId!, 'sendInputStreamChunk', { + requestId: id, + chunk, + }); + } + this.sendNotification(runtimeId!, 'endInputStream', { + requestId: id, + }); + } catch (e) { + logger.error(`Error streaming input: ${e}`); + } + })(); + } + + return promise as Promise; + } + + async cancelAction(input: { + traceId: string; + runtimeId?: string; + }): Promise<{ message: string }> { + const runtimeId = input.runtimeId || this.getMostRecentRuntime()?.id; + if (!runtimeId) { + throw new Error('No runtime found'); + } + // Assuming cancelAction is a request that returns a message + return this.sendRequest(runtimeId, 'cancelAction', { + traceId: input.traceId, + }); + } +} diff --git a/genkit-tools/common/src/manager/manager.ts b/genkit-tools/common/src/manager/manager.ts index 7a245e02e8..80ee90eb4e 100644 --- a/genkit-tools/common/src/manager/manager.ts +++ b/genkit-tools/common/src/manager/manager.ts @@ -51,7 +51,7 @@ const STREAM_DELIMITER = '\n'; const HEALTH_CHECK_INTERVAL = 5000; export const GENKIT_REFLECTION_API_SPEC_VERSION = 1; -interface RuntimeManagerOptions { +export interface RuntimeManagerOptions { /** URL of the telemetry server. */ telemetryServerUrl?: string; /** Whether to clean up unhealthy runtimes. */ @@ -62,11 +62,210 @@ interface RuntimeManagerOptions { processManager?: ProcessManager; /** Whether to disable realtime telemetry streaming. Defaults to false. */ disableRealtimeTelemetry?: boolean; + /** Experimental Reflection V2 flag */ + experimentalReflectionV2?: boolean; + /** Reflection V2 Port */ + reflectionV2Port?: number; } -export class RuntimeManager { - readonly processManager?: ProcessManager; - readonly disableRealtimeTelemetry: boolean; +export abstract class BaseRuntimeManager { + constructor( + readonly telemetryServerUrl: string | undefined, + readonly projectRoot: string, + readonly processManager?: ProcessManager, + readonly disableRealtimeTelemetry: boolean = false + ) {} + + abstract listRuntimes(): RuntimeInfo[]; + abstract getRuntimeById(id: string): RuntimeInfo | undefined; + abstract getMostRecentRuntime(): RuntimeInfo | undefined; + abstract getMostRecentDevUI(): DevToolsInfo | undefined; + abstract onRuntimeEvent( + listener: (eventType: RuntimeEvent, runtime: RuntimeInfo) => void + ): () => void; + abstract listActions( + input?: apis.ListActionsRequest + ): Promise>; + abstract runAction( + input: apis.RunActionRequest, + streamingCallback?: StreamingCallback, + onTraceId?: (traceId: string) => void, + inputStream?: AsyncIterable + ): Promise; + abstract cancelAction(input: { + traceId: string; + runtimeId?: string; + }): Promise<{ message: string }>; + abstract listValues( + input: apis.ListValuesRequest + ): Promise>; + + abstract stop(): Promise; + + /** + * Retrieves all traces + */ + async listTraces( + input: apis.ListTracesRequest + ): Promise { + const { limit, continuationToken, filter } = input; + let query = ''; + if (limit) { + query += `limit=${limit}`; + } + if (continuationToken) { + if (query !== '') { + query += '&'; + } + query += `continuationToken=${continuationToken}`; + } + if (filter) { + if (query !== '') { + query += '&'; + } + query += `filter=${encodeURI(JSON.stringify(filter))}`; + } + + const response = await axios + .get(`${this.telemetryServerUrl}/api/traces?${query}`) + .catch((err) => + this.httpErrorHandler(err, `Error listing traces for query='${query}'.`) + ); + + return apis.ListTracesResponseSchema.parse(response.data); + } + + /** + * Retrieves a trace for a given ID. + */ + async getTrace(input: apis.GetTraceRequest): Promise { + const { traceId } = input; + const response = await axios + .get(`${this.telemetryServerUrl}/api/traces/${traceId}`) + .catch((err) => + this.httpErrorHandler( + err, + `Error getting trace for traceId='${traceId}'` + ) + ); + + return response.data as TraceData; + } + + /** + * Streams trace updates in real-time from the telemetry server. + * Connects to the telemetry server's SSE endpoint and forwards updates via callback. + */ + async streamTrace( + input: apis.StreamTraceRequest, + streamingCallback: StreamingCallback + ): Promise { + const { traceId } = input; + + if (!this.telemetryServerUrl) { + throw new Error( + 'Telemetry server URL not configured. Cannot stream trace updates.' + ); + } + + const response = await axios + .get(`${this.telemetryServerUrl}/api/traces/${traceId}/stream`, { + headers: { + Accept: 'text/event-stream', + }, + responseType: 'stream', + }) + .catch((err) => + this.httpErrorHandler( + err, + `Error streaming trace for traceId='${traceId}'` + ) + ); + + const stream = response.data; + let buffer = ''; + + // Return a promise that resolves when the stream ends + return new Promise((resolve, reject) => { + stream.on('data', (chunk: Buffer) => { + buffer += chunk.toString(); + + // Process complete messages (ending with \n\n) + while (buffer.includes('\n\n')) { + const messageEnd = buffer.indexOf('\n\n'); + const message = buffer.substring(0, messageEnd).trim(); + buffer = buffer.substring(messageEnd + 2); + + // Skip empty messages + if (!message) { + continue; + } + // Parse SSE data line - strip "data: " prefix + try { + const jsonData = message.startsWith('data: ') + ? message.slice(6) + : message; + const parsed = JSON.parse(jsonData); + streamingCallback(parsed); + } catch (err) { + logger.error(`Error parsing stream data: ${err}`); + } + } + }); + + stream.on('end', () => { + resolve(); + }); + + stream.on('error', (err: Error) => { + logger.error(`Stream error for traceId='${traceId}': ${err}`); + reject(err); + }); + }); + } + + /** + * Adds a trace to the trace store + */ + async addTrace(input: TraceData): Promise { + await axios + .post(`${this.telemetryServerUrl}/api/traces/`, input) + .catch((err) => + this.httpErrorHandler(err, 'Error writing trace to store.') + ); + } + + /** + * Handles an HTTP error. + */ + protected httpErrorHandler(error: AxiosError, message?: string): never { + const newError = new GenkitToolsError(message || 'Internal Error'); + + if (error.response) { + // we got a non-200 response; copy the payload and rethrow + newError.data = error.response.data as GenkitError; + newError.stack = (error.response?.data as any).message; + if ((error.response?.data as any).message) { + newError.data.data = { + ...newError.data.data, + genkitErrorMessage: message, + genkitErrorDetails: { + stack: (error.response?.data as any).message, + traceId: (error.response?.data as any).traceId, + }, + }; + } + throw newError; + } + + // We actually have an exception; wrap it and re-throw. + throw new GenkitToolsError(message || 'Internal Error', { + cause: error.cause, + }); + } +} + +export class RuntimeManager extends BaseRuntimeManager { private filenameToRuntimeMap: Record = {}; private filenameToDevUiMap: Record = {}; private idToFileMap: Record = {}; @@ -75,20 +274,31 @@ export class RuntimeManager { private healthCheckInterval?: NodeJS.Timeout; private constructor( - readonly telemetryServerUrl: string | undefined, + telemetryServerUrl: string | undefined, private manageHealth: boolean, - readonly projectRoot: string, + projectRoot: string, processManager?: ProcessManager, disableRealtimeTelemetry?: boolean ) { - this.processManager = processManager; - this.disableRealtimeTelemetry = disableRealtimeTelemetry ?? false; + super( + telemetryServerUrl, + projectRoot, + processManager, + disableRealtimeTelemetry + ); } /** * Creates a new runtime manager. */ - static async create(options: RuntimeManagerOptions) { + static async create( + options: RuntimeManagerOptions + ): Promise { + if (options.experimentalReflectionV2) { + const { RuntimeManagerV2 } = await import('./manager-v2'); + return RuntimeManagerV2.create(options); + } + const manager = new RuntimeManager( options.telemetryServerUrl, options.manageHealth ?? true, @@ -209,7 +419,7 @@ export class RuntimeManager { } /** - * Retrieves all valuess. + * Retrieves all values. */ async listValues( input: apis.ListValuesRequest @@ -252,7 +462,8 @@ export class RuntimeManager { async runAction( input: apis.RunActionRequest, streamingCallback?: StreamingCallback, - onTraceId?: (traceId: string) => void + onTraceId?: (traceId: string) => void, + inputStream?: AsyncIterable ): Promise { const runtime = input.runtimeId ? this.getRuntimeById(input.runtimeId) @@ -461,139 +672,6 @@ export class RuntimeManager { } } - /** - * Retrieves all traces - */ - async listTraces( - input: apis.ListTracesRequest - ): Promise { - const { limit, continuationToken, filter } = input; - let query = ''; - if (limit) { - query += `limit=${limit}`; - } - if (continuationToken) { - if (query !== '') { - query += '&'; - } - query += `continuationToken=${continuationToken}`; - } - if (filter) { - if (query !== '') { - query += '&'; - } - query += `filter=${encodeURI(JSON.stringify(filter))}`; - } - - const response = await axios - .get(`${this.telemetryServerUrl}/api/traces?${query}`) - .catch((err) => - this.httpErrorHandler(err, `Error listing traces for query='${query}'.`) - ); - - return apis.ListTracesResponseSchema.parse(response.data); - } - - /** - * Retrieves a trace for a given ID. - */ - async getTrace(input: apis.GetTraceRequest): Promise { - const { traceId } = input; - const response = await axios - .get(`${this.telemetryServerUrl}/api/traces/${traceId}`) - .catch((err) => - this.httpErrorHandler( - err, - `Error getting trace for traceId='${traceId}'` - ) - ); - - return response.data as TraceData; - } - - /** - * Streams trace updates in real-time from the telemetry server. - * Connects to the telemetry server's SSE endpoint and forwards updates via callback. - */ - async streamTrace( - input: apis.StreamTraceRequest, - streamingCallback: StreamingCallback - ): Promise { - const { traceId } = input; - - if (!this.telemetryServerUrl) { - throw new Error( - 'Telemetry server URL not configured. Cannot stream trace updates.' - ); - } - - const response = await axios - .get(`${this.telemetryServerUrl}/api/traces/${traceId}/stream`, { - headers: { - Accept: 'text/event-stream', - }, - responseType: 'stream', - }) - .catch((err) => - this.httpErrorHandler( - err, - `Error streaming trace for traceId='${traceId}'` - ) - ); - - const stream = response.data; - let buffer = ''; - - // Return a promise that resolves when the stream ends - return new Promise((resolve, reject) => { - stream.on('data', (chunk: Buffer) => { - buffer += chunk.toString(); - - // Process complete messages (ending with \n\n) - while (buffer.includes('\n\n')) { - const messageEnd = buffer.indexOf('\n\n'); - const message = buffer.substring(0, messageEnd).trim(); - buffer = buffer.substring(messageEnd + 2); - - // Skip empty messages - if (!message) { - continue; - } - // Parse SSE data line - strip "data: " prefix - try { - const jsonData = message.startsWith('data: ') - ? message.slice(6) - : message; - const parsed = JSON.parse(jsonData); - streamingCallback(parsed); - } catch (err) { - logger.error(`Error parsing stream data: ${err}`); - } - } - }); - - stream.on('end', () => { - resolve(); - }); - - stream.on('error', (err: Error) => { - logger.error(`Stream error for traceId='${traceId}': ${err}`); - reject(err); - }); - }); - } - - /** - * Adds a trace to the trace store - */ - async addTrace(input: TraceData): Promise { - await axios - .post(`${this.telemetryServerUrl}/api/traces/`, input) - .catch((err) => - this.httpErrorHandler(err, 'Error writing trace to store.') - ); - } - /** * Notifies the runtime of dependencies it may need (e.g. telemetry server URL). */ @@ -789,35 +867,6 @@ export class RuntimeManager { } } - /** - * Handles an HTTP error. - */ - private httpErrorHandler(error: AxiosError, message?: string): never { - const newError = new GenkitToolsError(message || 'Internal Error'); - - if (error.response) { - // we got a non-200 response; copy the payload and rethrow - newError.data = error.response.data as GenkitError; - newError.stack = (error.response?.data as any).message; - if ((error.response?.data as any).message) { - newError.data.data = { - ...newError.data.data, - genkitErrorMessage: message, - genkitErrorDetails: { - stack: (error.response?.data as any).message, - traceId: (error.response?.data as any).traceId, - }, - }; - } - throw newError; - } - - // We actually have an exception; wrap it and re-throw. - throw new GenkitToolsError(message || 'Internal Error', { - cause: error.cause, - }); - } - /** * Handles a stream error by reading the stream and then calling httpErrorHandler. */ diff --git a/genkit-tools/common/src/server/router.ts b/genkit-tools/common/src/server/router.ts index e85a374c07..e137ac53b1 100644 --- a/genkit-tools/common/src/server/router.ts +++ b/genkit-tools/common/src/server/router.ts @@ -21,7 +21,7 @@ import { runNewEvaluation, validateSchema, } from '../eval'; -import type { RuntimeManager } from '../manager/manager'; +import type { BaseRuntimeManager } from '../manager/manager'; import { AppProcessStatus } from '../manager/process-manager'; import { GenkitToolsError, type RuntimeInfo } from '../manager/types'; import { TraceDataSchema } from '../types'; @@ -125,7 +125,7 @@ const loggedProcedure = t.procedure.use(async (opts) => { }); // eslint-disable-next-line @typescript-eslint/explicit-function-return-type -export const TOOLS_SERVER_ROUTER = (manager: RuntimeManager) => +export const TOOLS_SERVER_ROUTER = (manager: BaseRuntimeManager) => t.router({ /** Retrieves all runnable actions. */ listActions: loggedProcedure diff --git a/genkit-tools/common/src/server/server.ts b/genkit-tools/common/src/server/server.ts index 5ef7caf3ff..de20bd41e8 100644 --- a/genkit-tools/common/src/server/server.ts +++ b/genkit-tools/common/src/server/server.ts @@ -23,7 +23,7 @@ import type { Server } from 'http'; import os from 'os'; import path from 'path'; import type { GenkitToolsError } from '../manager'; -import type { RuntimeManager } from '../manager/manager'; +import type { BaseRuntimeManager } from '../manager/manager'; import { writeToolsInfoFile } from '../utils'; import { logger } from '../utils/logger'; import { toolsPackage } from '../utils/package'; @@ -43,10 +43,50 @@ const UI_ASSETS_ROOT = path.resolve( const UI_ASSETS_SERVE_PATH = path.resolve(UI_ASSETS_ROOT, 'ui', 'browser'); const API_BASE_PATH = '/api'; +class PushableAsyncIterable implements AsyncIterable { + private queue: T[] = []; + private resolvers: ((value: IteratorResult) => void)[] = []; + private closed = false; + + [Symbol.asyncIterator]() { + return { + next: () => this.next(), + }; + } + + next(): Promise> { + if (this.queue.length > 0) { + return Promise.resolve({ value: this.queue.shift()!, done: false }); + } + if (this.closed) { + return Promise.resolve({ value: undefined as any, done: true }); + } + return new Promise((resolve) => this.resolvers.push(resolve)); + } + + push(value: T) { + if (this.closed) return; + if (this.resolvers.length > 0) { + this.resolvers.shift()!({ value, done: false }); + } else { + this.queue.push(value); + } + } + + close() { + this.closed = true; + while (this.resolvers.length > 0) { + this.resolvers.shift()!({ value: undefined as any, done: true }); + } + } +} + +const activeInputStreams = new Map>(); + /** * Starts up the Genkit Tools server which includes static files for the UI and the Tools API. */ -export function startServer(manager: RuntimeManager, port: number) { +export function startServer(manager: BaseRuntimeManager, port: number) { let server: Server; const app = express(); @@ -131,30 +171,77 @@ export function startServer(manager: RuntimeManager, port: number) { res.flushHeaders(); } + let inputStream: PushableAsyncIterable | undefined; + const bidi = req.query.bidi === 'true'; + if (bidi) { + inputStream = new PushableAsyncIterable(); + } + + let capturedTraceId: string | undefined; + try { - const onTraceIdCallback = !manager.disableRealtimeTelemetry - ? (traceId: string) => { - // Set trace ID header and flush - this fires before first chunk - res.setHeader('X-Genkit-Trace-Id', traceId); - res.flushHeaders(); - } - : undefined; + const onTraceIdCallback = (traceId: string) => { + capturedTraceId = traceId; + // Set trace ID header and flush - this fires before first chunk + if (!manager.disableRealtimeTelemetry) { + res.setHeader('X-Genkit-Trace-Id', traceId); + res.flushHeaders(); + } + if (bidi && inputStream) { + activeInputStreams.set(traceId, inputStream); + } + }; const result = await manager.runAction( req.body, (chunk) => { res.write(JSON.stringify(chunk) + '\n'); }, - onTraceIdCallback + onTraceIdCallback, + inputStream ); res.write(JSON.stringify(result)); } catch (err) { res.write(JSON.stringify({ error: (err as GenkitToolsError).data })); + } finally { + if (capturedTraceId) { + activeInputStreams.delete(capturedTraceId); + } } res.end(); } ); + app.post( + '/api/sendBidiInput', + bodyParser.json({ limit: MAX_PAYLOAD_SIZE }), + (req, res) => { + const { traceId, chunk } = req.body; + const stream = activeInputStreams.get(traceId); + if (stream) { + stream.push(chunk); + res.status(200).send('OK'); + } else { + res.status(404).send('Stream not found'); + } + } + ); + + app.post('/api/endBidiInput', bodyParser.json(), (req, res) => { + const { traceId } = req.body; + const stream = activeInputStreams.get(traceId); + if (stream) { + stream.close(); + // Don't delete here, wait for action to complete (finally block) + // or delete if we want to ensure no more writes. + // If we delete here, subsequent writes will fail, which is correct. + // But finally block handles cleanup anyway. + res.status(200).send('OK'); + } else { + res.status(404).send('Stream not found'); + } + }); + app.post( '/api/streamTrace', bodyParser.json({ limit: MAX_PAYLOAD_SIZE }), diff --git a/genkit-tools/common/src/types/index.ts b/genkit-tools/common/src/types/index.ts index ea12971f0e..acc8b6a11b 100644 --- a/genkit-tools/common/src/types/index.ts +++ b/genkit-tools/common/src/types/index.ts @@ -25,6 +25,7 @@ export * from './eval'; export * from './evaluator'; export * from './model'; export * from './prompt'; +export * from './reflection'; export * from './retriever'; export * from './status'; export * from './trace'; diff --git a/genkit-tools/common/src/types/reflection.ts b/genkit-tools/common/src/types/reflection.ts new file mode 100644 index 0000000000..9d46f4a4a1 --- /dev/null +++ b/genkit-tools/common/src/types/reflection.ts @@ -0,0 +1,116 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { z } from 'zod'; +import { ActionMetadataSchema } from './action'; +import { RunActionRequestSchema } from './apis'; + +export { ActionMetadataSchema }; + +/** + * ReflectionRegisterSchema is the payload for the 'register' method. + */ +export const ReflectionRegisterParamsSchema = z.object({ + id: z.string(), + pid: z.number(), + name: z.string().optional(), + genkitVersion: z.string().optional(), + reflectionApiSpecVersion: z.number().optional(), +}); + +/** + * ReflectionStreamChunkSchema is the payload for the 'streamChunk' method. + */ +export const ReflectionStreamChunkParamsSchema = z.object({ + requestId: z.string(), + chunk: z.any(), +}); + +/** + * ReflectionRunActionStateSchema is the payload for the 'runActionState' method. + */ +export const ReflectionRunActionStateParamsSchema = z.object({ + requestId: z.string(), + state: z + .object({ + traceId: z.string().optional(), + }) + .optional(), +}); + +/** + * ReflectionConfigureSchema is the payload for the 'configure' method. + */ +export const ReflectionConfigureParamsSchema = z.object({ + telemetryServerUrl: z.string().optional(), +}); + +/** + * ReflectionListValuesSchema is the payload for the 'listValues' method. + */ +export const ReflectionListValuesParamsSchema = z.object({ + type: z.string(), +}); + +/** + * ReflectionRunActionSchema is the payload for the 'runAction' method. + */ +export const ReflectionRunActionParamsSchema = RunActionRequestSchema.extend({ + stream: z.boolean().optional(), + streamInput: z.boolean().optional(), +}); + +/** + * ReflectionCancelActionSchema is the payload for the 'cancelAction' method. + */ +export const ReflectionCancelActionParamsSchema = z.object({ + traceId: z.string(), +}); + +/** + * ReflectionSendInputStreamChunkSchema is the payload for the 'sendInputStreamChunk' method. + */ +export const ReflectionSendInputStreamChunkParamsSchema = z.object({ + requestId: z.string(), + chunk: z.any(), +}); + +/** + * ReflectionEndInputStreamSchema is the payload for the 'endInputStream' method. + */ +export const ReflectionEndInputStreamParamsSchema = z.object({ + requestId: z.string(), +}); + +/** + * ReflectionListActionsResponseSchema is the result for the 'listActions' method. + */ +export const ReflectionListActionsResponseSchema = z.record( + z.string(), + ActionMetadataSchema +); + +/** + * ReflectionListValuesResponseSchema is the result for the 'listValues' method. + */ +export const ReflectionListValuesResponseSchema = z.record(z.any()); + +/** + * ReflectionCancelActionResponseSchema is the result for the 'cancelAction' method. + */ +export const ReflectionCancelActionResponseSchema = z.object({ + message: z.string(), +}); diff --git a/genkit-tools/common/src/utils/eval.ts b/genkit-tools/common/src/utils/eval.ts index 8f0eefe5e6..3992debf94 100644 --- a/genkit-tools/common/src/utils/eval.ts +++ b/genkit-tools/common/src/utils/eval.ts @@ -19,7 +19,7 @@ import { randomUUID } from 'crypto'; import { createReadStream } from 'fs'; import { readFile } from 'fs/promises'; import { createInterface } from 'readline'; -import type { RuntimeManager } from '../manager'; +import type { BaseRuntimeManager } from '../manager'; import { findToolsConfig, isEvalField, @@ -323,7 +323,7 @@ async function readLines(fileName: string): Promise { } export async function hasAction(params: { - manager: RuntimeManager; + manager: BaseRuntimeManager; actionRef: string; }): Promise { const { manager, actionRef } = { ...params }; @@ -333,7 +333,7 @@ export async function hasAction(params: { } export async function getAction(params: { - manager: RuntimeManager; + manager: BaseRuntimeManager; actionRef: string; }): Promise { const { manager, actionRef } = { ...params }; diff --git a/genkit-tools/common/tests/manager-v2_test.ts b/genkit-tools/common/tests/manager-v2_test.ts new file mode 100644 index 0000000000..e6ff03dd6f --- /dev/null +++ b/genkit-tools/common/tests/manager-v2_test.ts @@ -0,0 +1,443 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { + afterEach, + beforeEach, + describe, + expect, + it, + jest, +} from '@jest/globals'; +import WebSocket from 'ws'; +import { RuntimeManagerV2 } from '../src/manager/manager-v2'; +import { RuntimeEvent } from '../src/manager/types'; + +describe('RuntimeManagerV2', () => { + let manager: RuntimeManagerV2; + let wsClient: WebSocket; + let port: number; + + beforeEach(async () => { + manager = await RuntimeManagerV2.create({ + projectRoot: './', + }); + port = manager.port!; + }); + + afterEach(async () => { + if (wsClient) { + wsClient.close(); + } + // Clean up server + await manager.stop(); + }); + + it('should accept connections and handle registration', (done) => { + wsClient = new WebSocket(`ws://localhost:${port}`); + + const unsubscribe = manager.onRuntimeEvent((event, runtime) => { + if (event === RuntimeEvent.ADD) { + expect(runtime.id).toBe('test-runtime-1'); + expect(runtime.pid).toBe(1234); + expect(manager.listRuntimes().length).toBe(1); + unsubscribe(); + done(); + } + }); + + wsClient.on('open', () => { + const registerMessage = { + jsonrpc: '2.0', + method: 'register', + params: { + id: 'test-runtime-1', + pid: 1234, + name: 'Test Runtime', + genkitVersion: '0.0.1', + reflectionApiSpecVersion: 1, + }, + id: '1', + }; + wsClient.send(JSON.stringify(registerMessage)); + }); + }); + + it('should allow unsubscribing from runtime events', async () => { + wsClient = new WebSocket(`ws://localhost:${port}`); + const listener = jest.fn(); + + const unsubscribe = manager.onRuntimeEvent(listener); + + await new Promise((resolve) => { + wsClient.on('open', () => { + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'register', + params: { id: 'test-runtime-unsubscribe', pid: 1234 }, + id: '1', + }) + ); + setTimeout(resolve, 100); + }); + }); + + // Wait for event + await new Promise((resolve) => setTimeout(resolve, 100)); + expect(listener).toHaveBeenCalled(); + + unsubscribe(); + listener.mockClear(); + + // Trigger another event (e.g. disconnect) + wsClient.close(); + await new Promise((resolve) => setTimeout(resolve, 100)); + + expect(listener).not.toHaveBeenCalled(); + }); + + it('should send requests and handle responses', async () => { + wsClient = new WebSocket(`ws://localhost:${port}`); + + await new Promise((resolve) => { + wsClient.on('open', () => { + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'register', + params: { id: 'test-runtime-2', pid: 1234 }, + id: '1', + }) + ); + // Wait for server to acknowledge or just wait a bit + setTimeout(resolve, 100); + }); + }); + + // Mock runtime response to runAction + wsClient.on('message', (data) => { + const message = JSON.parse(data.toString()); + if (message.method === 'runAction') { + const response = { + jsonrpc: '2.0', + result: { + result: 'Hello World', + telemetry: { + traceId: '1234', + }, + }, + id: message.id, + }; + wsClient.send(JSON.stringify(response)); + } + }); + + const response = await manager.runAction({ + key: 'testAction', + input: {}, + }); + + expect(response.result).toBe('Hello World'); + expect(response.telemetry).toStrictEqual({ + traceId: '1234', + }); + }); + + it('should handle listValues', async () => { + wsClient = new WebSocket(`ws://localhost:${port}`); + + await new Promise((resolve) => { + wsClient.on('open', () => { + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'register', + params: { id: 'test-runtime-values', pid: 1234 }, + id: '1', + }) + ); + setTimeout(resolve, 100); + }); + }); + + wsClient.on('message', (data) => { + const message = JSON.parse(data.toString()); + if (message.method === 'listValues') { + const response = { + jsonrpc: '2.0', + result: { + 'my-prompt': { template: 'foo' }, + }, + id: message.id, + }; + wsClient.send(JSON.stringify(response)); + } + }); + + const values = await manager.listValues({ + type: 'prompt', + }); + + expect(values['my-prompt']).toBeDefined(); + expect(values['my-prompt']).toEqual({ template: 'foo' }); + }); + + it('should handle streaming', async () => { + wsClient = new WebSocket(`ws://localhost:${port}`); + + await new Promise((resolve) => { + wsClient.on('open', () => { + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'register', + params: { id: 'test-runtime-3', pid: 1234 }, + id: '1', + }) + ); + setTimeout(resolve, 100); + }); + }); + + wsClient.on('message', (data) => { + const message = JSON.parse(data.toString()); + if (message.method === 'runAction' && message.params.stream) { + // Send chunk 1 + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'streamChunk', + params: { requestId: message.id, chunk: { content: 'Hello' } }, + }) + ); + // Send chunk 2 + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'streamChunk', + params: { requestId: message.id, chunk: { content: ' World' } }, + }) + ); + // Send final result + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + result: { result: 'Hello World', telemetry: {} }, + id: message.id, + }) + ); + } + }); + + const chunks: any[] = []; + const response = await manager.runAction( + { + key: 'testAction', + input: {}, + }, + (chunk) => { + chunks.push(chunk); + } + ); + + expect(chunks).toHaveLength(2); + expect(chunks[0]).toEqual({ content: 'Hello' }); + expect(chunks[1]).toEqual({ content: ' World' }); + expect(response.result).toBe('Hello World'); + expect(response.telemetry).toBeDefined(); + }); + + it('should handle streaming errors and massage the error object', async () => { + wsClient = new WebSocket(`ws://localhost:${port}`); + + await new Promise((resolve) => { + wsClient.on('open', () => { + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'register', + params: { id: 'test-runtime-error', pid: 1234 }, + id: '1', + }) + ); + setTimeout(resolve, 100); + }); + }); + + wsClient.on('message', (data) => { + const message = JSON.parse(data.toString()); + if (message.method === 'runAction' && message.params.stream) { + // Send chunk 1 + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'streamChunk', + params: { requestId: message.id, chunk: { content: 'Hello' } }, + }) + ); + // Send error + const errorResponse = { + code: -32000, + message: 'Test Error', + data: { + code: 13, + message: 'Test Error', + details: { + stack: 'Error stack...', + traceId: 'trace-123', + }, + }, + }; + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + error: errorResponse, + id: message.id, + }) + ); + } + }); + + const chunks: any[] = []; + try { + await manager.runAction( + { + key: 'testAction', + input: {}, + }, + (chunk) => { + chunks.push(chunk); + } + ); + throw new Error('Should have thrown'); + } catch (err: any) { + expect(chunks).toHaveLength(1); + expect(chunks[0]).toEqual({ content: 'Hello' }); + expect(err.message).toBe('Test Error'); + expect(err.data).toBeDefined(); + expect(err.data.data.genkitErrorMessage).toBe('Test Error'); + expect(err.data.stack).toBe('Error stack...'); + expect(err.data.data.genkitErrorDetails).toEqual({ + stack: 'Error stack...', + traceId: 'trace-123', + }); + } + }); + + it('should send cancelAction request', async () => { + wsClient = new WebSocket(`ws://localhost:${port}`); + + await new Promise((resolve) => { + wsClient.on('open', () => { + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'register', + params: { id: 'test-runtime-cancel', pid: 1234 }, + id: '1', + }) + ); + setTimeout(resolve, 100); + }); + }); + + wsClient.on('message', (data) => { + const message = JSON.parse(data.toString()); + if (message.method === 'cancelAction') { + const response = { + jsonrpc: '2.0', + result: { + message: 'Action cancelled', + }, + id: message.id, + }; + wsClient.send(JSON.stringify(response)); + } + }); + + const response = await manager.cancelAction({ + traceId: '1234', + }); + + expect(response.message).toBe('Action cancelled'); + }); + + it('should handle runActionState for early trace info', async () => { + wsClient = new WebSocket(`ws://localhost:${port}`); + + await new Promise((resolve) => { + wsClient.on('open', () => { + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'register', + params: { id: 'test-runtime-trace', pid: 1234 }, + id: '1', + }) + ); + setTimeout(resolve, 100); + }); + }); + + wsClient.on('message', (data) => { + const message = JSON.parse(data.toString()); + if (message.method === 'runAction') { + // Send runActionState with traceId + wsClient.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'runActionState', + params: { + requestId: message.id, + state: { + traceId: 'early-trace-id', + }, + }, + }) + ); + + // Send final result + const response = { + jsonrpc: '2.0', + result: { + result: 'Hello World', + telemetry: { + traceId: 'early-trace-id', + }, + }, + id: message.id, + }; + wsClient.send(JSON.stringify(response)); + } + }); + + let capturedTraceId: string | undefined; + const response = await manager.runAction( + { + key: 'testAction', + input: {}, + }, + undefined, + (traceId) => { + capturedTraceId = traceId; + } + ); + + expect(capturedTraceId).toBe('early-trace-id'); + expect(response.result).toBe('Hello World'); + }); +}); diff --git a/genkit-tools/common/tests/server_test.ts b/genkit-tools/common/tests/server_test.ts new file mode 100644 index 0000000000..876657041b --- /dev/null +++ b/genkit-tools/common/tests/server_test.ts @@ -0,0 +1,165 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { + afterEach, + beforeEach, + describe, + expect, + it, + jest, +} from '@jest/globals'; +import axios from 'axios'; +import getPort from 'get-port'; +import { BaseRuntimeManager } from '../src/manager/manager'; +import { startServer } from '../src/server/server'; + +describe('Tools Server', () => { + let port: number; + let serverPromise: Promise; + let mockManager: any; + + beforeEach(async () => { + port = await getPort(); + mockManager = { + projectRoot: './', + disableRealtimeTelemetry: false, + runAction: jest.fn(), + streamTrace: jest.fn(), + listActions: jest.fn(), + listTraces: jest.fn(), + getTrace: jest.fn(), + getMostRecentRuntime: jest.fn(), + listRuntimes: jest.fn(), + onRuntimeEvent: jest.fn(), + cancelAction: jest.fn(), + }; + serverPromise = startServer(mockManager as BaseRuntimeManager, port); + }); + + afterEach(async () => { + const exitSpy = jest + .spyOn(process, 'exit') + .mockImplementation((code?: any) => { + return undefined as never; + }); + try { + await axios.post(`http://localhost:${port}/api/__quitquitquit`); + } catch (e) { + // Ignore + } + await serverPromise; + exitSpy.mockRestore(); + }); + + it('should handle runAction', async () => { + mockManager.runAction.mockResolvedValue({ result: 'bar' }); + + let response; + try { + response = await axios.post(`http://localhost:${port}/api/runAction`, { + key: 'foo', + input: 'bar', + }); + } catch (e: any) { + throw new Error(`runAction failed: ${e.message}`); + } + + expect(response.data.result).toBe('bar'); + expect(mockManager.runAction).toHaveBeenCalledWith( + expect.objectContaining({ key: 'foo' }), + undefined, + expect.any(Function) + ); + }); + + it('should handle bidi streaming', async () => { + let inputStream: AsyncIterable | undefined; + let finishAction: (() => void) | undefined; + + mockManager.runAction.mockImplementation( + async (input: any, cb: any, trace: any, stream: any) => { + inputStream = stream; + await new Promise((resolve) => { + finishAction = resolve; + }); + return { result: 'done' }; + } + ); + + const responsePromise = axios + .post( + `http://localhost:${port}/api/streamAction?bidi=true`, + { key: 'bidi' }, + { responseType: 'stream' } + ) + .catch((e) => { + throw new Error(`Stream action failed: ${e.message}`); + }); + + // Wait for runAction to be called + while (!inputStream) { + await new Promise((r) => setTimeout(r, 10)); + } + + const traceId = 'test-trace-id'; + // Get the onTraceId callback from the mock call args + const [inputArg, cb, onTraceIdCallback] = + mockManager.runAction.mock.calls[0]; + onTraceIdCallback(traceId); + + // Collect input chunks in background + const chunks: any[] = []; + const collectPromise = (async () => { + for await (const chunk of inputStream!) { + chunks.push(chunk); + } + })(); + + // Now send input + try { + await axios.post(`http://localhost:${port}/api/sendBidiInput`, { + traceId, + chunk: 'input1', + }); + + await axios.post(`http://localhost:${port}/api/endBidiInput`, { + traceId, + }); + } catch (e: any) { + throw new Error(`send/end input failed: ${e.message}`); + } + + await collectPromise; + expect(chunks).toEqual(['input1']); + + // Emit output chunk + if (cb) cb({ result: 'chunk1' }); + + // Finish action + finishAction!(); + + const response = await responsePromise; + const stream = response.data; + const outputChunks: string[] = []; + for await (const chunk of stream) { + outputChunks.push(chunk.toString()); + } + const output = outputChunks.join(''); + expect(output).toContain('chunk1'); + expect(output).toContain('done'); + }); +}); diff --git a/genkit-tools/genkit-schema.json b/genkit-tools/genkit-schema.json index 7ea72e3be3..9b8cccfa29 100644 --- a/genkit-tools/genkit-schema.json +++ b/genkit-tools/genkit-schema.json @@ -1460,6 +1460,255 @@ ], "additionalProperties": false }, + "ActionMetadata": { + "type": "object", + "properties": { + "actionType": { + "type": "string" + }, + "name": { + "type": "string" + }, + "description": { + "type": "string" + }, + "inputSchema": {}, + "inputJsonSchema": { + "anyOf": [ + { + "type": "object", + "properties": {}, + "additionalProperties": false, + "description": "A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object." + }, + { + "type": "null" + } + ], + "description": "A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object." + }, + "outputSchema": {}, + "outputJsonSchema": { + "$ref": "#/$defs/ActionMetadata/properties/inputJsonSchema", + "description": "A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object." + }, + "streamSchema": {}, + "metadata": { + "type": "object", + "additionalProperties": { + "anyOf": [ + { + "anyOf": [ + { + "type": "string" + }, + { + "type": "number" + }, + { + "type": "integer", + "format": "int64" + }, + { + "type": "boolean" + }, + { + "type": "object", + "properties": {}, + "additionalProperties": false + } + ] + }, + { + "type": "array", + "items": { + "$ref": "#/$defs/ActionMetadata/properties/metadata/additionalProperties/anyOf/0" + } + } + ] + } + } + }, + "required": [ + "name" + ], + "additionalProperties": false + }, + "ReflectionCancelActionParams": { + "type": "object", + "properties": { + "traceId": { + "type": "string" + } + }, + "required": [ + "traceId" + ], + "additionalProperties": false + }, + "ReflectionCancelActionResponse": { + "type": "object", + "properties": { + "message": { + "type": "string" + } + }, + "required": [ + "message" + ], + "additionalProperties": false + }, + "ReflectionConfigureParams": { + "type": "object", + "properties": { + "telemetryServerUrl": { + "type": "string" + } + }, + "additionalProperties": false + }, + "ReflectionEndInputStreamParams": { + "type": "object", + "properties": { + "requestId": { + "type": "string" + } + }, + "required": [ + "requestId" + ], + "additionalProperties": false + }, + "ReflectionListActionsResponse": { + "type": "object", + "additionalProperties": { + "$ref": "#/$defs/ActionMetadata" + } + }, + "ReflectionListValuesParams": { + "type": "object", + "properties": { + "type": { + "type": "string" + } + }, + "required": [ + "type" + ], + "additionalProperties": false + }, + "ReflectionListValuesResponse": { + "type": "object", + "additionalProperties": {} + }, + "ReflectionRegisterParams": { + "type": "object", + "properties": { + "id": { + "type": "string" + }, + "pid": { + "type": "number" + }, + "name": { + "type": "string" + }, + "genkitVersion": { + "type": "string" + }, + "reflectionApiSpecVersion": { + "type": "number" + } + }, + "required": [ + "id", + "pid" + ], + "additionalProperties": false + }, + "ReflectionRunActionParams": { + "type": "object", + "properties": { + "runtimeId": { + "type": "string", + "description": "ID of the Genkit runtime to run the action on. Typically $pid-$port." + }, + "key": { + "type": "string", + "description": "Action key that consists of the action type and ID." + }, + "input": { + "description": "An input with the type that this action expects." + }, + "context": { + "description": "Additional runtime context data (ex. auth context data)." + }, + "telemetryLabels": { + "type": "object", + "additionalProperties": { + "type": "string" + }, + "description": "Labels to be applied to telemetry data." + }, + "stream": { + "type": "boolean" + }, + "streamInput": { + "type": "boolean" + } + }, + "required": [ + "key" + ], + "additionalProperties": false + }, + "ReflectionRunActionStateParams": { + "type": "object", + "properties": { + "requestId": { + "type": "string" + }, + "state": { + "type": "object", + "properties": { + "traceId": { + "type": "string" + } + }, + "additionalProperties": false + } + }, + "required": [ + "requestId" + ], + "additionalProperties": false + }, + "ReflectionSendInputStreamChunkParams": { + "type": "object", + "properties": { + "requestId": { + "type": "string" + }, + "chunk": {} + }, + "required": [ + "requestId" + ], + "additionalProperties": false + }, + "ReflectionStreamChunkParams": { + "type": "object", + "properties": { + "requestId": { + "type": "string" + }, + "chunk": {} + }, + "required": [ + "requestId" + ], + "additionalProperties": false + }, "CommonRetrieverOptions": { "type": "object", "properties": { diff --git a/genkit-tools/pnpm-lock.yaml b/genkit-tools/pnpm-lock.yaml index 8928abdb16..7c475dcef2 100644 --- a/genkit-tools/pnpm-lock.yaml +++ b/genkit-tools/pnpm-lock.yaml @@ -144,6 +144,9 @@ importers: cors: specifier: ^2.8.5 version: 2.8.5 + events: + specifier: ^3.3.0 + version: 3.3.0 express: specifier: ^4.21.0 version: 4.21.2 @@ -174,6 +177,9 @@ importers: winston: specifier: ^3.11.0 version: 3.17.0 + ws: + specifier: ^8.18.3 + version: 8.18.3 yaml: specifier: ^2.4.1 version: 2.8.0 @@ -202,6 +208,9 @@ importers: '@types/cors': specifier: ^2.8.19 version: 2.8.19 + '@types/events': + specifier: ^3.0.3 + version: 3.0.3 '@types/express': specifier: ^4.17.21 version: 4.17.23 @@ -223,6 +232,9 @@ importers: '@types/uuid': specifier: ^9.0.8 version: 9.0.8 + '@types/ws': + specifier: ^8.18.1 + version: 8.18.1 bun-types: specifier: ^1.2.16 version: 1.2.16 @@ -1126,6 +1138,9 @@ packages: '@types/cors@2.8.19': resolution: {integrity: sha512-mFNylyeyqN93lfe/9CSxOGREz8cpzAhH+E93xJ4xWQf62V8sQ/24reV2nyzUWM6H6Xji+GGHpkbLe7pVoUEskg==} + '@types/events@3.0.3': + resolution: {integrity: sha512-trOc4AAUThEz9hapPtSd7wf5tiQKvTtu5b371UxXdTuqzIh0ArcRspRP0i0Viu+LXstIQ1z96t1nsPxT9ol01g==} + '@types/express-serve-static-core@4.19.0': resolution: {integrity: sha512-bGyep3JqPCRry1wq+O5n7oiBgGWmeIJXPjXXCo8EK0u8duZGSYar7cGqd3ML2JUsLGeB7fmc06KYo9fLGWqPvQ==} @@ -1207,6 +1222,9 @@ packages: '@types/uuid@9.0.8': resolution: {integrity: sha512-jg+97EGIcY9AGHJJRaaPVgetKDsrTgbRjQ5Msgjh/DQKEFl0DtyRr/VCOyD1T2R1MNeWPK/u7JoGhlDZnKBAfA==} + '@types/ws@8.18.1': + resolution: {integrity: sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==} + '@types/yargs-parser@21.0.3': resolution: {integrity: sha512-I4q9QU9MQv4oEOz4tAHJtNz1cwuLxn2F3xcc2iV5WdqLPpUnj30aUuxt1mAxYTG+oe8CZMV/+6rU4S4gRDzqtQ==} @@ -1819,6 +1837,10 @@ packages: resolution: {integrity: sha512-i/2XbnSz/uxRCU6+NdVJgKWDTM427+MqYbkQzD321DuCQJUqOuJKIA0IM2+W2xtYHdKOmZ4dR6fExsd4SXL+WQ==} engines: {node: '>=6'} + events@3.3.0: + resolution: {integrity: sha512-mQw+2fkQbALzQ7V0MY0IqdnXNOeTtP4r0lN9z7AAawCXgqea7bDii20AYrIBrFd/Hx0M2Ocz6S111CaFkUcb0Q==} + engines: {node: '>=0.8.x'} + eventsource-parser@3.0.3: resolution: {integrity: sha512-nVpZkTMM9rF6AQ9gPJpFsNAMt48wIzB5TQgiTLdHiuO8XEDhUgZEhqKlZWXbIzo9VmJ/HvysHqEaVeD5v9TPvA==} engines: {node: '>=20.0.0'} @@ -3483,6 +3505,18 @@ packages: resolution: {integrity: sha512-7KxauUdBmSdWnmpaGFg+ppNjKF8uNLry8LyzjauQDOVONfFLNKrKvQOxZ/VuTIcS/gge/YNahf5RIIQWTSarlg==} engines: {node: ^12.13.0 || ^14.15.0 || >=16.0.0} + ws@8.18.3: + resolution: {integrity: sha512-PEIGCY5tSlUt50cqyMXfCzX+oOPqN0vuGqWzbcJ2xvnkzkq46oOpz7dQaTDBdfICb4N14+GARUDw2XV2N4tvzg==} + engines: {node: '>=10.0.0'} + peerDependencies: + bufferutil: ^4.0.1 + utf-8-validate: '>=5.0.2' + peerDependenciesMeta: + bufferutil: + optional: true + utf-8-validate: + optional: true + xdg-basedir@4.0.0: resolution: {integrity: sha512-PSNhEJDejZYV7h50BohL09Er9VaIefr2LMAf3OEmpCkjOi34eYyQYAXUTjEQtZJTKcF0E2UKTh+osDLsgNim9Q==} engines: {node: '>=8'} @@ -4443,6 +4477,8 @@ snapshots: dependencies: '@types/node': 20.19.1 + '@types/events@3.0.3': {} + '@types/express-serve-static-core@4.19.0': dependencies: '@types/node': 20.19.1 @@ -4535,6 +4571,10 @@ snapshots: '@types/uuid@9.0.8': {} + '@types/ws@8.18.1': + dependencies: + '@types/node': 20.19.1 + '@types/yargs-parser@21.0.3': {} '@types/yargs@17.0.32': @@ -5239,6 +5279,8 @@ snapshots: event-target-shim@5.0.1: {} + events@3.3.0: {} + eventsource-parser@3.0.3: {} eventsource@3.0.7: @@ -7266,6 +7308,8 @@ snapshots: imurmurhash: 0.1.4 signal-exit: 3.0.7 + ws@8.18.3: {} + xdg-basedir@4.0.0: {} y18n@5.0.8: {} diff --git a/genkit-tools/scripts/schema-exporter.ts b/genkit-tools/scripts/schema-exporter.ts index 48df79b56a..1d7bedf119 100644 --- a/genkit-tools/scripts/schema-exporter.ts +++ b/genkit-tools/scripts/schema-exporter.ts @@ -29,6 +29,7 @@ const EXPORTED_TYPE_MODULES = [ '../common/src/types/model.ts', '../common/src/types/parts.ts', '../common/src/types/reranker.ts', + '../common/src/types/reflection.ts', '../common/src/types/retriever.ts', '../common/src/types/trace.ts', ]; diff --git a/genkit-tools/telemetry-server/package.json b/genkit-tools/telemetry-server/package.json index 39d206700b..bf45bd3eac 100644 --- a/genkit-tools/telemetry-server/package.json +++ b/genkit-tools/telemetry-server/package.json @@ -41,10 +41,10 @@ "cors": "^2.8.6" }, "devDependencies": { + "@types/cors": "^2.8.19", "@types/express": "~4.17.21", "@types/lockfile": "^1.0.4", "@types/node": "^20.11.30", - "@types/cors": "^2.8.19", "get-port": "^7.1.0", "npm-run-all": "^4.1.5", "rimraf": "^6.0.1", diff --git a/go/ai/gen.go b/go/ai/gen.go index a92a059475..d50edc3d36 100644 --- a/go/ai/gen.go +++ b/go/ai/gen.go @@ -18,6 +18,20 @@ package ai +type ActionMetadata struct { + ActionType string `json:"actionType,omitempty"` + Description string `json:"description,omitempty"` + // A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object. + InputJsonSchema any `json:"inputJsonSchema,omitempty"` + InputSchema any `json:"inputSchema,omitempty"` + Metadata map[string]any `json:"metadata,omitempty"` + Name string `json:"name,omitempty"` + // A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object. + OutputJsonSchema *ActionMetadata `json:"outputJsonSchema,omitempty"` + OutputSchema any `json:"outputSchema,omitempty"` + StreamSchema any `json:"streamSchema,omitempty"` +} + type customPart struct { // Custom contains custom key-value data specific to this part. Custom map[string]any `json:"custom,omitempty"` @@ -425,6 +439,72 @@ type reasoningPart struct { Reasoning string `json:"reasoning,omitempty"` } +type ReflectionCancelActionParams struct { + TraceID string `json:"traceId,omitempty"` +} + +type ReflectionCancelActionResponse struct { + Message string `json:"message,omitempty"` +} + +type ReflectionConfigureParams struct { + TelemetryServerUrl string `json:"telemetryServerUrl,omitempty"` +} + +type ReflectionEndInputStreamParams struct { + RequestID string `json:"requestId,omitempty"` +} + +type ReflectionListActionsResponse map[string]*ActionMetadata + +type ReflectionListValuesParams struct { + Type string `json:"type,omitempty"` +} + +type ReflectionListValuesResponse map[string]any + +type ReflectionRegisterParams struct { + GenkitVersion string `json:"genkitVersion,omitempty"` + Id string `json:"id,omitempty"` + Name string `json:"name,omitempty"` + Pid float64 `json:"pid,omitempty"` + ReflectionApiSpecVersion float64 `json:"reflectionApiSpecVersion,omitempty"` +} + +type ReflectionRunActionParams struct { + // Additional runtime context data (ex. auth context data). + Context any `json:"context,omitempty"` + // An input with the type that this action expects. + Input any `json:"input,omitempty"` + // Action key that consists of the action type and ID. + Key string `json:"key,omitempty"` + // ID of the Genkit runtime to run the action on. Typically $pid-$port. + RuntimeID string `json:"runtimeId,omitempty"` + Stream bool `json:"stream,omitempty"` + StreamInput bool `json:"streamInput,omitempty"` + // Labels to be applied to telemetry data. + TelemetryLabels map[string]string `json:"telemetryLabels,omitempty"` +} + +type ReflectionRunActionStateParams struct { + RequestID string `json:"requestId,omitempty"` + State *ReflectionRunActionStateParamsState `json:"state,omitempty"` +} + +type ReflectionRunActionStateParamsState struct { + TraceID string `json:"traceId,omitempty"` +} + +type ReflectionSendInputStreamChunkParams struct { + Chunk any `json:"chunk,omitempty"` + RequestID string `json:"requestId,omitempty"` +} + +type ReflectionStreamChunkParams struct { + Chunk any `json:"chunk,omitempty"` + RequestID string `json:"requestId,omitempty"` +} + // RerankerRequest represents a request to rerank documents based on relevance. type RerankerRequest struct { // Documents is the array of documents to rerank. diff --git a/go/internal/cmd/jsonschemagen/jsonschema.go b/go/internal/cmd/jsonschemagen/jsonschema.go index 0afbde5e74..b41d2e3880 100644 --- a/go/internal/cmd/jsonschemagen/jsonschema.go +++ b/go/internal/cmd/jsonschemagen/jsonschema.go @@ -29,6 +29,7 @@ import ( type Schema struct { SchemaVersion string `json:"$schema,omitempty"` ID string `json:"$id,omitempty"` + Format string `json:"format,omitempty"` Type *OneOf[string, []string] `json:"type,omitempty"` Description string `json:"description,omitempty"` Properties map[string]*Schema `json:"properties,omitempty"` diff --git a/go/internal/cmd/jsonschemagen/jsonschemagen.go b/go/internal/cmd/jsonschemagen/jsonschemagen.go index 6ed2d5a8e7..f8d55e5ea2 100644 --- a/go/internal/cmd/jsonschemagen/jsonschemagen.go +++ b/go/internal/cmd/jsonschemagen/jsonschemagen.go @@ -218,6 +218,11 @@ func adjustAdditionalProperties(x any) { } } } + if k == "properties" { + if pm, ok := v.(map[string]any); ok && len(pm) == 0 { + delete(m, k) + } + } // TODO: Fix this - causing schemagen issues if k == "uniqueItems" { delete(m, k) @@ -356,6 +361,19 @@ func (g *generator) generateType(name string) (err error) { switch typ { case "object": // a JSONSchema object corresponds to a Go struct + if s.Properties == nil && s.AdditionalProperties != nil { + typ, err := g.typeExpr(s) + if err != nil { + return err + } + g.generateDoc(s, tcfg) + goName := tcfg.name + if goName == "" { + goName = adjustIdentifier(name) + } + g.pr("type %s %s\n\n", goName, typ) + return nil + } if err := g.generateStruct(name, s, tcfg); err != nil { return err } @@ -483,13 +501,12 @@ func (g *generator) typeExpr(s *Schema) (string, error) { return "any", nil } if s.Ref != "" { - name, ok := strings.CutPrefix(s.Ref, refPrefix) - if !ok { - return "", fmt.Errorf("ref %q does not begin with prefix %q", s.Ref, refPrefix) + s2, name, err := g.resolveRef(s.Ref) + if err != nil { + return "", err } ic := g.cfg.configFor(name) - s2, ok := g.schemas[name] - if !ok { + if s2 == nil { // If there is no schema, perhaps there is a config value. if ic != nil && ic.name != "" { return ic.name, nil @@ -497,7 +514,7 @@ func (g *generator) typeExpr(s *Schema) (string, error) { return "", fmt.Errorf("unknown type in reference: %q", name) } // Apply a config that changes the name. - if ic := g.cfg.configFor(name); ic != nil && ic.name != "" { + if ic != nil && ic.name != "" { name = ic.name } if s2.Enum != nil { @@ -556,6 +573,42 @@ func (g *generator) typeExpr(s *Schema) (string, error) { } } +// resolveRef resolves a JSON schema reference. +// It handles simple references like "#/$defs/Action" and nested ones like +// "#/$defs/ActionMetadata/properties/inputJsonSchema". +func (g *generator) resolveRef(ref string) (*Schema, string, error) { + name, ok := strings.CutPrefix(ref, refPrefix) + if !ok { + return nil, "", fmt.Errorf("ref %q does not begin with prefix %q", ref, refPrefix) + } + parts := strings.Split(name, "/") + s, ok := g.schemas[parts[0]] + if !ok { + return nil, "", fmt.Errorf("unknown type in reference: %q", parts[0]) + } + for i := 1; i < len(parts); i++ { + switch parts[i] { + case "properties": + if i+1 >= len(parts) { + return nil, "", fmt.Errorf("invalid ref (ends in properties): %q", ref) + } + s = s.Properties[parts[i+1]] + i++ + case "additionalProperties": + s = s.AdditionalProperties + case "items": + s = s.Items + default: + return nil, "", fmt.Errorf("cannot handle ref segment %q in %q", parts[i], ref) + } + if s == nil { + return nil, "", fmt.Errorf("ref path not found: %q", ref) + } + } + // The caller mostly cares about the direct name in $defs (parts[0]). + return s, parts[0], nil +} + // adjustIdentifier returns name with the first letter capitalized // so it is exported, and makes other idiomatic Go adjustments. func adjustIdentifier(name string) string { diff --git a/js/ai/src/model-types.ts b/js/ai/src/model-types.ts index 345c78d9e1..03e90cee73 100644 --- a/js/ai/src/model-types.ts +++ b/js/ai/src/model-types.ts @@ -17,10 +17,10 @@ import { OperationSchema, z } from '@genkit-ai/core'; import { DocumentDataSchema } from './document.js'; import { - Part, PartSchema, ToolRequestPartSchema, ToolResponsePartSchema, + type Part, } from './parts.js'; export { Part, PartSchema }; diff --git a/js/ai/tests/prompt/prompt_test.ts b/js/ai/tests/prompt/prompt_test.ts index 553cc6ee95..7f6251f3f7 100644 --- a/js/ai/tests/prompt/prompt_test.ts +++ b/js/ai/tests/prompt/prompt_test.ts @@ -59,7 +59,7 @@ describe('prompt', () => { }, async (request, opts) => { // Store the abortSignal for verification - (defineModel as any).__test__lastAbortSignal = request.abortSignal; + (defineModel as any).__test__lastAbortSignal = opts.abortSignal; return { message: { role: 'model', diff --git a/js/core/package.json b/js/core/package.json index afaa1b04c5..f294ea0b31 100644 --- a/js/core/package.json +++ b/js/core/package.json @@ -43,9 +43,11 @@ "get-port": "^5.1.0", "json-schema": "^0.4.0", "zod": "^3.23.8", - "zod-to-json-schema": "^3.22.4" + "zod-to-json-schema": "^3.22.4", + "ws": "^8.18.0" }, "devDependencies": { + "@types/ws": "^8.5.10", "@types/express": "^4.17.21", "@types/node": "^20.11.30", "genversion": "^3.2.0", diff --git a/js/core/src/reflection-types.ts b/js/core/src/reflection-types.ts new file mode 100644 index 0000000000..fea0fd5075 --- /dev/null +++ b/js/core/src/reflection-types.ts @@ -0,0 +1,120 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { z } from 'zod'; +import { ActionMetadataSchema } from './action.js'; + +// NOTE: Keep this file in sync with genkit-tools/common/src/types/reflection.ts +// and genkit-tools/common/src/types/apis.ts. + +/** + * ReflectionRegisterSchema is the payload for the 'register' method. + */ +export const ReflectionRegisterParamsSchema = z.object({ + id: z.string(), + pid: z.number(), + name: z.string().optional(), + genkitVersion: z.string().optional(), + reflectionApiSpecVersion: z.number().optional(), +}); + +/** + * ReflectionStreamChunkSchema is the payload for the 'streamChunk' method. + */ +export const ReflectionStreamChunkParamsSchema = z.object({ + requestId: z.string(), + chunk: z.any(), +}); + +/** + * ReflectionRunActionStateSchema is the payload for the 'runActionState' method. + */ +export const ReflectionRunActionStateParamsSchema = z.object({ + requestId: z.string(), + state: z + .object({ + traceId: z.string().optional(), + }) + .optional(), +}); + +/** + * ReflectionConfigureSchema is the payload for the 'configure' method. + */ +export const ReflectionConfigureParamsSchema = z.object({ + telemetryServerUrl: z.string().optional(), +}); + +/** + * ReflectionListValuesSchema is the payload for the 'listValues' method. + */ +export const ReflectionListValuesParamsSchema = z.object({ + type: z.string(), +}); + +/** + * ReflectionRunActionSchema is the payload for the 'runAction' method. + */ +export const ReflectionRunActionParamsSchema = z.object({ + key: z.string(), + input: z.any().optional(), + context: z.any().optional(), + telemetryLabels: z.record(z.string(), z.string()).optional(), + stream: z.boolean().optional(), + streamInput: z.boolean().optional(), +}); + +/** + * ReflectionCancelActionSchema is the payload for the 'cancelAction' method. + */ +export const ReflectionCancelActionParamsSchema = z.object({ + traceId: z.string(), +}); + +/** + * ReflectionSendInputStreamChunkSchema is the payload for the 'sendInputStreamChunk' method. + */ +export const ReflectionSendInputStreamChunkParamsSchema = z.object({ + requestId: z.string(), + chunk: z.any(), +}); + +/** + * ReflectionEndInputStreamSchema is the payload for the 'endInputStream' method. + */ +export const ReflectionEndInputStreamParamsSchema = z.object({ + requestId: z.string(), +}); + +/** + * ReflectionListActionsResponseSchema is the result for the 'listActions' method. + */ +export const ReflectionListActionsResponseSchema = z.record( + z.string(), + ActionMetadataSchema +); + +/** + * ReflectionListValuesResponseSchema is the result for the 'listValues' method. + */ +export const ReflectionListValuesResponseSchema = z.record(z.any()); + +/** + * ReflectionCancelActionResponseSchema is the result for the 'cancelAction' method. + */ +export const ReflectionCancelActionResponseSchema = z.object({ + message: z.string(), +}); diff --git a/js/core/src/reflection-v2.ts b/js/core/src/reflection-v2.ts new file mode 100644 index 0000000000..ad05a841f8 --- /dev/null +++ b/js/core/src/reflection-v2.ts @@ -0,0 +1,396 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import WebSocket from 'ws'; +import { StatusCodes, type Status } from './action.js'; +import { GENKIT_REFLECTION_API_SPEC_VERSION, GENKIT_VERSION } from './index.js'; +import { logger } from './logging.js'; +import { + ReflectionCancelActionParamsSchema, + ReflectionConfigureParamsSchema, + ReflectionListActionsResponseSchema, + ReflectionListValuesParamsSchema, + ReflectionListValuesResponseSchema, + ReflectionRegisterParamsSchema, + ReflectionRunActionParamsSchema, + ReflectionRunActionStateParamsSchema, + ReflectionStreamChunkParamsSchema, +} from './reflection-types.js'; +import type { Registry } from './registry.js'; +import { toJsonSchema } from './schema.js'; +import { flushTracing, setTelemetryServerUrl } from './tracing.js'; + +let apiIndex = 0; + +interface JsonRpcRequest { + jsonrpc: '2.0'; + method: string; + params?: any; + id?: number | string; +} + +interface JsonRpcResponse { + jsonrpc: '2.0'; + result?: any; + error?: { + code: number; + message: string; + data?: any; + }; + id: number | string; +} + +type JsonRpcMessage = JsonRpcRequest | JsonRpcResponse; + +export interface ReflectionServerV2Options { + configuredEnvs?: string[]; + name?: string; + url: string; +} + +export class ReflectionServerV2 { + private registry: Registry; + private options: ReflectionServerV2Options; + private ws: WebSocket | null = null; + private url: string; + private index = apiIndex++; + private activeActions = new Map< + string, + { + abortController: AbortController; + startTime: Date; + } + >(); + + constructor(registry: Registry, options: ReflectionServerV2Options) { + this.registry = registry; + this.options = { + configuredEnvs: ['dev'], + ...options, + }; + // The URL should be provided via environment variable by the CLI manager + this.url = this.options.url; + } + + async start() { + logger.debug(`Connecting to Reflection V2 server at ${this.url}`); + this.ws = new WebSocket(this.url); + + this.ws.on('open', () => { + logger.debug('Connected to Reflection V2 server.'); + this.register(); + }); + + this.ws.on('message', async (data) => { + try { + const message = JSON.parse(data.toString()) as JsonRpcMessage; + if ('method' in message) { + await this.handleRequest(message as JsonRpcRequest); + } + } catch (error) { + logger.error(`Failed to parse message: ${error}`); + } + }); + + this.ws.on('error', (error) => { + logger.error(`Reflection V2 WebSocket error: ${error}`); + }); + + this.ws.on('close', () => { + logger.debug('Reflection V2 WebSocket closed.'); + }); + } + + async stop() { + if (this.ws) { + this.ws.close(); + this.ws = null; + } + } + + private send(message: JsonRpcMessage) { + if (this.ws && this.ws.readyState === WebSocket.OPEN) { + this.ws.send(JSON.stringify(message)); + } + } + + private sendResponse(id: number | string, result: any) { + this.send({ + jsonrpc: '2.0', + result, + id, + }); + } + + private sendError( + id: number | string, + code: number, + message: string, + data?: any + ) { + this.send({ + jsonrpc: '2.0', + error: { code, message, data }, + id, + }); + } + + private sendNotification(method: string, params: any) { + this.send({ + jsonrpc: '2.0', + method, + params, + }); + } + + private register() { + const params = ReflectionRegisterParamsSchema.parse({ + id: process.env.GENKIT_RUNTIME_ID || this.runtimeId, + pid: process.pid, + name: this.options.name || this.runtimeId, + genkitVersion: GENKIT_VERSION, + reflectionApiSpecVersion: GENKIT_REFLECTION_API_SPEC_VERSION, + envs: this.options.configuredEnvs, + }); + this.sendNotification('register', params); + } + + get runtimeId() { + return `${process.pid}${this.index ? `-${this.index}` : ''}`; + } + + private async handleRequest(request: JsonRpcRequest) { + try { + switch (request.method) { + case 'listActions': + await this.handleListActions(request); + break; + case 'listValues': + await this.handleListValues(request); + break; + case 'runAction': + await this.handleRunAction(request); + break; + case 'configure': + this.handleConfigure(request); + break; + case 'cancelAction': + await this.handleCancelAction(request); + break; + case 'sendInputStreamChunk': + this.handleSendInputStreamChunk(request); + break; + case 'endInputStream': + this.handleEndInputStream(request); + break; + default: + if (request.id) { + this.sendError( + request.id, + -32601, + `Method not found: ${request.method}` + ); + } + } + } catch (error: any) { + if (request.id) { + this.sendError(request.id, -32000, error.message, { + stack: error.stack, + }); + } + } + } + + private async handleListActions(request: JsonRpcRequest) { + if (!request.id) return; // Should be a request + const actions = await this.registry.listResolvableActions(); + const convertedActions: Record = {}; + + Object.keys(actions).forEach((key) => { + const action = actions[key]; + convertedActions[key] = { + key, + name: action.name, + description: action.description, + metadata: action.metadata, + }; + if (action.inputSchema || action.inputJsonSchema) { + convertedActions[key].inputSchema = toJsonSchema({ + schema: action.inputSchema, + jsonSchema: action.inputJsonSchema, + }); + } + if (action.outputSchema || action.outputJsonSchema) { + convertedActions[key].outputSchema = toJsonSchema({ + schema: action.outputSchema, + jsonSchema: action.outputJsonSchema, + }); + } + }); + + this.sendResponse( + request.id, + ReflectionListActionsResponseSchema.parse(convertedActions) + ); + } + + private async handleListValues(request: JsonRpcRequest) { + if (!request.id) return; + const { type } = ReflectionListValuesParamsSchema.parse(request.params); + const values = await this.registry.listValues(type); + this.sendResponse( + request.id, + ReflectionListValuesResponseSchema.parse(values) + ); + } + + private async handleRunAction(request: JsonRpcRequest) { + if (!request.id) return; + + const { key, input, context, telemetryLabels, stream } = + ReflectionRunActionParamsSchema.parse(request.params); + const action = await this.registry.lookupAction(key); + + if (!action) { + this.sendError(request.id, 404, `action ${key} not found`); + return; + } + + const abortController = new AbortController(); + let traceId: string | undefined; + + try { + const onTraceStartCallback = ({ traceId: tid }: { traceId: string }) => { + traceId = tid; + this.activeActions.set(tid, { + abortController, + startTime: new Date(), + }); + // Send early trace ID notification + this.sendNotification( + 'runActionState', + ReflectionRunActionStateParamsSchema.parse({ + requestId: request.id, + state: { traceId: tid }, + }) + ); + }; + + if (stream) { + const callback = (chunk: any) => { + this.sendNotification( + 'streamChunk', + ReflectionStreamChunkParamsSchema.parse({ + requestId: request.id, + chunk, + }) + ); + }; + + const result = await action.run(input, { + context, + onChunk: callback, + telemetryLabels, + onTraceStart: onTraceStartCallback, + abortSignal: abortController.signal, + }); + + await flushTracing(); + + // Send final result + this.sendResponse(request.id, { + result: result.result, + telemetry: { + traceId: result.telemetry.traceId, + }, + }); + } else { + const result = await action.run(input, { + context, + telemetryLabels, + onTraceStart: onTraceStartCallback, + abortSignal: abortController.signal, + }); + await flushTracing(); + + this.sendResponse(request.id, { + result: result.result, + telemetry: { + traceId: result.telemetry.traceId, + }, + }); + } + } catch (err: any) { + const isAbort = + err?.name === 'AbortError' || + (typeof DOMException !== 'undefined' && + err instanceof DOMException && + err.name === 'AbortError'); + + const errorResponse: Status = { + code: isAbort ? StatusCodes.CANCELLED : StatusCodes.INTERNAL, + message: isAbort ? 'Action was cancelled' : err.message, + details: { + stack: err.stack, + }, + }; + if (err.traceId || traceId) { + errorResponse.details.traceId = err.traceId || traceId; + } + + this.sendError(request.id, -32000, errorResponse.message, errorResponse); + } finally { + if (traceId) { + this.activeActions.delete(traceId); + } + } + } + + private handleConfigure(request: JsonRpcRequest) { + const { telemetryServerUrl } = ReflectionConfigureParamsSchema.parse( + request.params + ); + if (telemetryServerUrl && !process.env.GENKIT_TELEMETRY_SERVER) { + setTelemetryServerUrl(telemetryServerUrl); + logger.debug(`Connected to telemetry server on ${telemetryServerUrl}`); + } + } + + private async handleCancelAction(request: JsonRpcRequest) { + if (!request.id) return; + const { traceId } = ReflectionCancelActionParamsSchema.parse( + request.params + ); + const activeAction = this.activeActions.get(traceId); + if (activeAction) { + activeAction.abortController.abort(); + this.activeActions.delete(traceId); + this.sendResponse(request.id, { message: 'Action cancelled' }); + } else { + this.sendError(request.id, 404, 'Action not found or already completed'); + } + } + + private handleSendInputStreamChunk(request: JsonRpcRequest) { + // ReflectionSendInputStreamChunkParamsSchema.parse(request.params); + throw new Error('Not implemented'); + } + + private handleEndInputStream(request: JsonRpcRequest) { + // ReflectionEndInputStreamParamsSchema.parse(request.params); + throw new Error('Not implemented'); + } +} diff --git a/js/core/src/reflection.ts b/js/core/src/reflection.ts index a63959bf13..ec2ca7f726 100644 --- a/js/core/src/reflection.ts +++ b/js/core/src/reflection.ts @@ -92,6 +92,7 @@ export class ReflectionServer { startTime: Date; } >(); + private v2Server: any | null = null; constructor(registry: Registry, options?: ReflectionServerOptions) { this.registry = registry; @@ -135,6 +136,17 @@ export class ReflectionServer { ); return; } + if (process.env.GENKIT_REFLECTION_V2_SERVER) { + const { ReflectionServerV2 } = await import('./reflection-v2.js'); + this.v2Server = new ReflectionServerV2(this.registry, { + configuredEnvs: this.options.configuredEnvs, + name: this.options.name, + url: process.env.GENKIT_REFLECTION_V2_SERVER, + }); + await this.v2Server.start(); + ReflectionServer.RUNNING_SERVERS.push(this); + return; + } const server = express(); @@ -439,6 +451,15 @@ export class ReflectionServer { * Stops the server and removes it from the list of running servers to clean up on exit. */ async stop(): Promise { + if (this.v2Server) { + await this.v2Server.stop(); + const index = ReflectionServer.RUNNING_SERVERS.indexOf(this); + if (index > -1) { + ReflectionServer.RUNNING_SERVERS.splice(index, 1); + } + return; + } + if (!this.server) { return; } diff --git a/js/core/tests/reflection-v2_test.ts b/js/core/tests/reflection-v2_test.ts new file mode 100644 index 0000000000..6b236f9b18 --- /dev/null +++ b/js/core/tests/reflection-v2_test.ts @@ -0,0 +1,377 @@ +/** + * Copyright 2024 Google LLC + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as assert from 'assert'; +import { afterEach, beforeEach, describe, it } from 'node:test'; +import { WebSocketServer } from 'ws'; +import { z } from 'zod'; +import { action } from '../src/action.js'; +import { initNodeFeatures } from '../src/node.js'; +import { ReflectionServerV2 } from '../src/reflection-v2.js'; +import { Registry } from '../src/registry.js'; + +initNodeFeatures(); + +describe('ReflectionServerV2', () => { + let wss: WebSocketServer; + let server: ReflectionServerV2; + let registry: Registry; + let port: number; + let serverWs: any; + + beforeEach(() => { + return new Promise((resolve) => { + wss = new WebSocketServer({ port: 0 }); + wss.on('listening', () => { + port = (wss.address() as any).port; + resolve(); + }); + wss.on('connection', (ws) => { + serverWs = ws; + }); + registry = new Registry(); + }); + }); + + afterEach(async () => { + if (server) { + await server.stop(); + } + if (serverWs) { + serverWs.terminate(); + } + await new Promise((resolve) => { + wss.close(() => resolve()); + }); + }); + + it('should connect to the server and register', async () => { + const connected = new Promise((resolve) => { + wss.on('connection', (ws) => { + ws.on('message', (data) => { + const msg = JSON.parse(data.toString()); + if (msg.method === 'register') { + assert.strictEqual(msg.params.name, 'test-app'); + resolve(); + } + }); + }); + }); + + server = new ReflectionServerV2(registry, { + url: `ws://localhost:${port}`, + name: 'test-app', + }); + await server.start(); + await connected; + }); + + it('should handle listActions', async () => { + // Register a dummy action + const testAction = action( + { + name: 'testAction', + description: 'A test action', + inputSchema: z.object({ foo: z.string() }), + outputSchema: z.object({ bar: z.string() }), + actionType: 'custom', + }, + async (input) => ({ bar: input.foo }) + ); + registry.registerAction('custom', testAction); + + const gotListActions = new Promise((resolve) => { + wss.on('connection', (ws) => { + ws.on('message', (data) => { + const msg = JSON.parse(data.toString()); + if (msg.method === 'register') { + // After registration, request listActions + ws.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'listActions', + id: '123', + }) + ); + } else if (msg.id === '123') { + assert.ok(msg.result['/custom/testAction']); + assert.strictEqual( + msg.result['/custom/testAction'].name, + 'testAction' + ); + resolve(); + } + }); + }); + }); + + server = new ReflectionServerV2(registry, { + url: `ws://localhost:${port}`, + }); + await server.start(); + await gotListActions; + }); + + it('should handle listValues', async () => { + registry.registerValue('prompt', 'my-prompt', { template: 'foo' }); + + const gotListValues = new Promise((resolve) => { + wss.on('connection', (ws) => { + ws.on('message', (data) => { + const msg = JSON.parse(data.toString()); + if (msg.method === 'register') { + ws.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'listValues', + params: { type: 'prompt' }, + id: '124', + }) + ); + } else if (msg.id === '124') { + assert.ok(msg.result['my-prompt']); + assert.strictEqual(msg.result['my-prompt'].template, 'foo'); + resolve(); + } + }); + }); + }); + + server = new ReflectionServerV2(registry, { + url: `ws://localhost:${port}`, + }); + await server.start(); + await gotListValues; + }); + + it('should handle runAction', async () => { + const testAction = action( + { + name: 'testAction', + inputSchema: z.object({ foo: z.string() }), + outputSchema: z.object({ bar: z.string() }), + actionType: 'custom', + }, + async (input) => ({ bar: input.foo }) + ); + registry.registerAction('custom', testAction); + + const actionRun = new Promise((resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error('runAction timeout')), + 2000 + ); + wss.on('connection', (ws) => { + ws.on('message', (data) => { + try { + const msg = JSON.parse(data.toString()); + if (msg.method === 'register') { + ws.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'runAction', + params: { + key: '/custom/testAction', + input: { foo: 'baz' }, + }, + id: '456', + }) + ); + } else if (msg.id === '456') { + if (msg.error) { + reject( + new Error(`runAction error: ${JSON.stringify(msg.error)}`) + ); + return; + } + assert.strictEqual(msg.result.result.bar, 'baz'); + clearTimeout(timeout); + resolve(); + } + } catch (e) { + clearTimeout(timeout); + reject(e); + } + }); + }); + }); + + server = new ReflectionServerV2(registry, { + url: `ws://localhost:${port}`, + }); + await server.start(); + await actionRun; + }); + + it('should handle streaming runAction', async () => { + const streamAction = action( + { + name: 'streamAction', + inputSchema: z.object({ foo: z.string() }), + outputSchema: z.string(), + actionType: 'custom', + }, + async (input, { sendChunk }) => { + sendChunk('chunk1'); + sendChunk('chunk2'); + return 'done'; + } + ); + registry.registerAction('custom', streamAction); + + const chunks: any[] = []; + const actionRun = new Promise((resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error('streamAction timeout')), + 2000 + ); + wss.on('connection', (ws) => { + ws.on('message', (data) => { + try { + const msg = JSON.parse(data.toString()); + if (msg.method === 'register') { + ws.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'runAction', + params: { + key: '/custom/streamAction', + input: { foo: 'baz' }, + stream: true, + }, + id: '789', + }) + ); + } else if (msg.method === 'streamChunk') { + chunks.push(msg.params.chunk); + } else if (msg.id === '789') { + if (msg.error) { + reject( + new Error(`streamAction error: ${JSON.stringify(msg.error)}`) + ); + return; + } + assert.strictEqual(msg.result.result, 'done'); + assert.deepStrictEqual(chunks, ['chunk1', 'chunk2']); + clearTimeout(timeout); + resolve(); + } + } catch (e) { + clearTimeout(timeout); + reject(e); + } + }); + }); + }); + + server = new ReflectionServerV2(registry, { + url: `ws://localhost:${port}`, + }); + await server.start(); + await actionRun; + }); + + it('should handle cancelAction', async () => { + let cancelSignal: AbortSignal | undefined; + const longAction = action( + { + name: 'longAction', + inputSchema: z.any(), + outputSchema: z.any(), + actionType: 'custom', + }, + async (_, { abortSignal }) => { + cancelSignal = abortSignal; + await new Promise((resolve, reject) => { + const timer = setTimeout(resolve, 5000); + if (abortSignal.aborted) { + clearTimeout(timer); + reject(new Error('Action cancelled')); + return; + } + abortSignal.addEventListener('abort', () => { + clearTimeout(timer); + reject(new Error('Action cancelled')); + }); + }); + } + ); + registry.registerAction('custom', longAction); + + const actionCancelled = new Promise((resolve, reject) => { + const timeout = setTimeout( + () => reject(new Error('cancelAction timeout')), + 2000 + ); + wss.on('connection', (ws) => { + ws.on('message', (data) => { + try { + const msg = JSON.parse(data.toString()); + if (msg.method === 'register') { + // Start action + ws.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'runAction', + params: { + key: '/custom/longAction', + input: {}, + }, + id: '999', + }) + ); + } else if (msg.method === 'runActionState') { + // Got traceId, send cancel + const traceId = msg.params.state.traceId; + ws.send( + JSON.stringify({ + jsonrpc: '2.0', + method: 'cancelAction', + params: { traceId }, + id: '1000', + }) + ); + } else if (msg.id === '1000') { + // Cancel response + assert.strictEqual(msg.result.message, 'Action cancelled'); + } else if (msg.id === '999') { + // Run action response (should be error) + if (msg.error) { + // Ensure code indicates cancellation if possible, or just error + // In implementation we send code -32000 and message 'Action was cancelled' + assert.match(msg.error.message, /cancelled/); + assert.ok(cancelSignal?.aborted); + clearTimeout(timeout); + resolve(); + } else { + reject(new Error('Action should have failed')); + } + } + } catch (e) { + clearTimeout(timeout); + reject(e); + } + }); + }); + }); + + server = new ReflectionServerV2(registry, { + url: `ws://localhost:${port}`, + }); + await server.start(); + await actionCancelled; + }); +}); diff --git a/js/pnpm-lock.yaml b/js/pnpm-lock.yaml index c1e6d863a3..6ebbb36b45 100644 --- a/js/pnpm-lock.yaml +++ b/js/pnpm-lock.yaml @@ -144,6 +144,9 @@ importers: json-schema: specifier: ^0.4.0 version: 0.4.0 + ws: + specifier: ^8.18.0 + version: 8.18.3 zod: specifier: ^3.23.8 version: 3.25.67 @@ -157,6 +160,9 @@ importers: '@types/node': specifier: ^20.11.30 version: 20.19.1 + '@types/ws': + specifier: ^8.5.10 + version: 8.18.1 genversion: specifier: ^3.2.0 version: 3.2.0 @@ -2792,11 +2798,27 @@ packages: firebase: optional: true + '@genkit-ai/firebase@1.25.0': + resolution: {integrity: sha512-Z0FbnJHQs8qS0yxG++Dn3CZ7gv+YNaihGaWXoDKy02mNOkeRzHA6UPaWxSTaWkWHYdB0MyOnMGlyqxnWyqVdmg==} + peerDependencies: + '@google-cloud/firestore': ^7.11.0 + firebase: '>=11.5.0' + firebase-admin: '>=12.2' + genkit: ^1.25.0 + peerDependenciesMeta: + firebase: + optional: true + '@genkit-ai/google-cloud@1.16.1': resolution: {integrity: sha512-uujjdGr/sra7iKHApufwkt5jGo7CQcRCJNWPgnSg4g179CjtvtZBGjxmFRVBtKzuF61ktkY6E9JoLz83nWEyAA==} peerDependencies: genkit: ^1.16.1 + '@genkit-ai/google-cloud@1.25.0': + resolution: {integrity: sha512-wHCa8JSTv7MtwzXjUQ9AT5v0kCTJrz0In+ffgAYw1yt8ComAz5o7Ir+xks+sX1vJfN8ptvW0GUa6rsUaXCB3kA==} + peerDependencies: + genkit: ^1.25.0 + '@gerrit0/mini-shiki@1.27.2': resolution: {integrity: sha512-GeWyHz8ao2gBiUW4OJnQDxXQnFgZQwwQk05t/CVVgNBN7/rK8XZ7xY6YhLVv9tH3VppWWmr9DCl3MwemB/i+Og==} @@ -4437,6 +4459,9 @@ packages: '@types/node@20.19.1': resolution: {integrity: sha512-jJD50LtlD2dodAEO653i3YF04NWak6jN3ky+Ri3Em3mGR39/glWiboM/IePaRbgwSfqM1TpGXfAg8ohn/4dTgA==} + '@types/node@20.19.26': + resolution: {integrity: sha512-0l6cjgF0XnihUpndDhk+nyD3exio3iKaYROSgvh/qSevPXax3L8p5DBRFjbvalnwatGgHEQn2R88y2fA3g4irg==} + '@types/node@22.15.32': resolution: {integrity: sha512-3jigKqgSjsH6gYZv2nEsqdXfZqIFGAV36XYYjf9KGZ3PSG+IhLecqPnI310RvjutyMwifE2hhhNEklOUrvx/wA==} @@ -4515,6 +4540,9 @@ packages: '@types/whatwg-url@11.0.5': resolution: {integrity: sha512-coYR071JRaHa+xoEvvYqvnIHaVqaYrLPbsufM9BF63HkwI5Lgmy2QR8Q5K/lYDYo5AK82wOvSOS0UsLTpTG7uQ==} + '@types/ws@8.18.1': + resolution: {integrity: sha512-ThVF6DCVhA8kUGy+aazFQ4kXQ7E1Ty7A3ypFOe0IcJV8O/M511G99AW24irKrW56Wt44yG9+ij8FaqoBGkuBXg==} + '@types/yargs-parser@21.0.3': resolution: {integrity: sha512-I4q9QU9MQv4oEOz4tAHJtNz1cwuLxn2F3xcc2iV5WdqLPpUnj30aUuxt1mAxYTG+oe8CZMV/+6rU4S4gRDzqtQ==} @@ -9792,7 +9820,7 @@ snapshots: dependencies: '@genkit-ai/core': 1.29.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.29.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)) '@opentelemetry/api': 1.9.0 - '@types/node': 20.19.1 + '@types/node': 20.19.26 colorette: 2.0.20 dotprompt: 1.1.1 json5: 2.2.3 @@ -9813,7 +9841,7 @@ snapshots: dependencies: '@genkit-ai/core': 1.29.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) '@opentelemetry/api': 1.9.0 - '@types/node': 20.19.1 + '@types/node': 20.19.26 colorette: 2.0.20 dotprompt: 1.1.1 json5: 2.2.3 @@ -9883,7 +9911,7 @@ snapshots: zod-to-json-schema: 3.24.5(zod@3.25.67) optionalDependencies: '@cfworker/json-schema': 4.1.1 - '@genkit-ai/firebase': 1.16.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) + '@genkit-ai/firebase': 1.25.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) transitivePeerDependencies: - '@google-cloud/firestore' - encoding @@ -9915,9 +9943,22 @@ snapshots: - supports-color optional: true - '@genkit-ai/firebase@1.16.1(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit)': + '@genkit-ai/firebase@1.25.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.28.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1))': dependencies: - '@genkit-ai/google-cloud': 1.16.1(encoding@0.1.13)(genkit@genkit) + '@genkit-ai/google-cloud': 1.25.0(encoding@0.1.13)(genkit@1.28.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)) + '@google-cloud/firestore': 7.11.1(encoding@0.1.13) + firebase-admin: 13.6.0(encoding@0.1.13) + genkit: 1.28.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1) + optionalDependencies: + firebase: 11.9.1 + transitivePeerDependencies: + - encoding + - supports-color + optional: true + + '@genkit-ai/firebase@1.25.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit)': + dependencies: + '@genkit-ai/google-cloud': 1.25.0(encoding@0.1.13)(genkit@genkit) '@google-cloud/firestore': 7.11.1(encoding@0.1.13) firebase-admin: 13.6.0(encoding@0.1.13) genkit: link:genkit @@ -9953,7 +9994,32 @@ snapshots: - supports-color optional: true - '@genkit-ai/google-cloud@1.16.1(encoding@0.1.13)(genkit@genkit)': + '@genkit-ai/google-cloud@1.25.0(encoding@0.1.13)(genkit@1.28.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1))': + dependencies: + '@google-cloud/logging-winston': 6.0.1(encoding@0.1.13)(winston@3.17.0) + '@google-cloud/opentelemetry-cloud-monitoring-exporter': 0.19.0(@opentelemetry/api@1.9.0)(@opentelemetry/core@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/sdk-metrics@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) + '@google-cloud/opentelemetry-cloud-trace-exporter': 2.4.1(@opentelemetry/api@1.9.0)(@opentelemetry/core@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/sdk-trace-base@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) + '@google-cloud/opentelemetry-resource-util': 2.4.0(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) + '@opentelemetry/api': 1.9.0 + '@opentelemetry/auto-instrumentations-node': 0.49.2(@opentelemetry/api@1.9.0)(encoding@0.1.13) + '@opentelemetry/core': 1.25.1(@opentelemetry/api@1.9.0) + '@opentelemetry/instrumentation': 0.52.1(@opentelemetry/api@1.9.0) + '@opentelemetry/instrumentation-pino': 0.41.0(@opentelemetry/api@1.9.0) + '@opentelemetry/instrumentation-winston': 0.39.0(@opentelemetry/api@1.9.0) + '@opentelemetry/resources': 1.25.1(@opentelemetry/api@1.9.0) + '@opentelemetry/sdk-metrics': 1.25.1(@opentelemetry/api@1.9.0) + '@opentelemetry/sdk-node': 0.52.1(@opentelemetry/api@1.9.0) + '@opentelemetry/sdk-trace-base': 1.25.1(@opentelemetry/api@1.9.0) + genkit: 1.28.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1) + google-auth-library: 9.15.1(encoding@0.1.13) + node-fetch: 3.3.2 + winston: 3.17.0 + transitivePeerDependencies: + - encoding + - supports-color + optional: true + + '@genkit-ai/google-cloud@1.25.0(encoding@0.1.13)(genkit@genkit)': dependencies: '@google-cloud/logging-winston': 6.0.1(encoding@0.1.13)(winston@3.17.0) '@google-cloud/opentelemetry-cloud-monitoring-exporter': 0.19.0(@opentelemetry/api@1.9.0)(@opentelemetry/core@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/sdk-metrics@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) @@ -11700,7 +11766,7 @@ snapshots: '@types/graceful-fs@4.1.9': dependencies: - '@types/node': 20.19.1 + '@types/node': 20.19.26 '@types/handlebars@4.1.0': dependencies: @@ -11774,6 +11840,10 @@ snapshots: dependencies: undici-types: 6.21.0 + '@types/node@20.19.26': + dependencies: + undici-types: 6.21.0 + '@types/node@22.15.32': dependencies: undici-types: 6.21.0 @@ -11858,6 +11928,10 @@ snapshots: dependencies: '@types/webidl-conversions': 7.0.3 + '@types/ws@8.18.1': + dependencies: + '@types/node': 20.19.26 + '@types/yargs-parser@21.0.3': {} '@types/yargs@17.0.33': @@ -14371,7 +14445,7 @@ snapshots: '@jest/expect': 29.7.0 '@jest/test-result': 29.7.0 '@jest/types': 29.6.3 - '@types/node': 20.19.1 + '@types/node': 20.19.26 chalk: 4.1.2 co: 4.6.0 dedent: 1.5.3 @@ -14515,7 +14589,7 @@ snapshots: '@jest/environment': 29.7.0 '@jest/fake-timers': 29.7.0 '@jest/types': 29.6.3 - '@types/node': 20.19.1 + '@types/node': 20.19.26 jest-mock: 29.7.0 jest-util: 29.7.0 @@ -14701,7 +14775,7 @@ snapshots: jest-worker@29.7.0: dependencies: - '@types/node': 20.19.1 + '@types/node': 20.19.26 jest-util: 29.7.0 merge-stream: 2.0.0 supports-color: 8.1.1 @@ -15940,7 +16014,7 @@ snapshots: '@protobufjs/path': 1.1.2 '@protobufjs/pool': 1.1.0 '@protobufjs/utf8': 1.1.0 - '@types/node': 20.19.1 + '@types/node': 20.19.26 long: 5.3.2 proxy-addr@2.0.7: From d4f9bf99ecf9a38dd7cba6127582b559ba85beca Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Thu, 12 Mar 2026 14:56:15 -0400 Subject: [PATCH 02/10] updated py model --- py/packages/genkit/src/genkit/core/typing.py | 143 ++++++++++++++++++- py/uv.lock | 24 +--- 2 files changed, 142 insertions(+), 25 deletions(-) diff --git a/py/packages/genkit/src/genkit/core/typing.py b/py/packages/genkit/src/genkit/core/typing.py index f60ed30169..477416bbd2 100644 --- a/py/packages/genkit/src/genkit/core/typing.py +++ b/py/packages/genkit/src/genkit/core/typing.py @@ -412,6 +412,104 @@ class RankedDocumentMetadata(BaseModel): score: float +class ReflectionCancelActionParams(BaseModel): + """Model for reflectioncancelactionparams data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + trace_id: str = Field(...) + + +class ReflectionCancelActionResponse(BaseModel): + """Model for reflectioncancelactionresponse data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + message: str + + +class ReflectionConfigureParams(BaseModel): + """Model for reflectionconfigureparams data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + telemetry_server_url: str | None = Field(default=None) + + +class ReflectionEndInputStreamParams(BaseModel): + """Model for reflectionendinputstreamparams data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + request_id: str = Field(...) + + +class ReflectionListValuesParams(BaseModel): + """Model for reflectionlistvaluesparams data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + type: str + + +class ReflectionListValuesResponse(RootModel[dict[str, Any]]): + """Root model for reflectionlistvaluesresponse.""" + + root: dict[str, Any] + + +class ReflectionRegisterParams(BaseModel): + """Model for reflectionregisterparams data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + id: str + pid: float + name: str | None = None + genkit_version: str | None = Field(default=None) + reflection_api_spec_version: float | None = Field(default=None) + + +class ReflectionRunActionParams(BaseModel): + """Model for reflectionrunactionparams data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + runtime_id: str | None = Field( + default=None, description='ID of the Genkit runtime to run the action on. Typically $pid-$port.' + ) + key: str = Field(..., description='Action key that consists of the action type and ID.') + input: Any | None = Field(default=None, description='An input with the type that this action expects.') + context: Any | None = Field(default=None, description='Additional runtime context data (ex. auth context data).') + telemetry_labels: dict[str, str] | None = Field(default=None, description='Labels to be applied to telemetry data.') + stream: bool | None = None + stream_input: bool | None = Field(default=None) + + +class State(BaseModel): + """Model for state data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + trace_id: str | None = Field(default=None) + + +class ReflectionRunActionStateParams(BaseModel): + """Model for reflectionrunactionstateparams data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + request_id: str = Field(...) + state: State | None = None + + +class ReflectionSendInputStreamChunkParams(BaseModel): + """Model for reflectionsendinputstreamchunkparams data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + request_id: str = Field(...) + chunk: Any | None = None + + +class ReflectionStreamChunkParams(BaseModel): + """Model for reflectionstreamchunkparams data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + request_id: str = Field(...) + chunk: Any | None = None + + class CommonRetrieverOptions(BaseModel): """Model for commonretrieveroptions data.""" @@ -457,8 +555,8 @@ class SameProcessAsParentSpan(BaseModel): value: bool -class State(StrEnum): - """State data type class.""" +class State1(StrEnum): + """State1 data type class.""" SUCCESS = 'success' ERROR = 'error' @@ -469,7 +567,7 @@ class SpanMetadata(BaseModel): model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str - state: State | None = None + state: State1 | None = None input: Any | None = None output: Any | None = None is_root: bool | None = Field(default=None) @@ -510,6 +608,20 @@ class TraceMetadata(BaseModel): timestamp: float +class InputJsonSchema(RootModel[dict[str, Any] | None]): + """Root model for inputjsonschema.""" + + root: dict[str, Any] | None = Field( + ..., description='A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object.' + ) + + +class Field0(RootModel[str | float | int | bool | dict[str, Any]]): + """Root model for field0.""" + + root: str | float | int | bool | dict[str, Any] + + class Context(RootModel[list[Any]]): """Root model for context.""" @@ -801,6 +913,31 @@ class ToolResponsePart(BaseModel): resource: Resource | None = None +class ActionMetadata(BaseModel): + """Model for actionmetadata data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + action_type: str | None = Field(default=None) + name: str + description: str | None = None + input_schema: Any | None = Field(default=None) + input_json_schema: dict[str, Any] | None = Field( + default=None, description='A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object.' + ) + output_schema: Any | None = Field(default=None) + output_json_schema: InputJsonSchema | None = Field( + default=None, description='A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object.' + ) + stream_schema: Any | None = Field(default=None) + metadata: dict[str, str | float | int | bool | dict[str, Any] | list[Field0]] | None = None + + +class ReflectionListActionsResponse(RootModel[dict[str, ActionMetadata]]): + """Root model for reflectionlistactionsresponse.""" + + root: dict[str, ActionMetadata] + + class Link(BaseModel): """Model for link data.""" diff --git a/py/uv.lock b/py/uv.lock index 65897e7b2e..e9f16a6baf 100644 --- a/py/uv.lock +++ b/py/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 3 +revision = 2 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.14'", @@ -53,7 +53,6 @@ members = [ "provider-vertex-ai-vector-search-bigquery", "provider-vertex-ai-vector-search-firestore", "web-fastapi-bugbot", - "web-fastapi-minimal-devui", "web-flask-hello", ] @@ -8029,7 +8028,7 @@ wheels = [ [[package]] name = "web-fastapi-bugbot" version = "0.2.0" -source = { editable = "samples/web-fastapi-bugbot" } +source = { virtual = "samples/web-fastapi-bugbot" } dependencies = [ { name = "genkit" }, { name = "genkit-plugin-fastapi" }, @@ -8054,25 +8053,6 @@ requires-dist = [ ] provides-extras = ["dev"] -[[package]] -name = "web-fastapi-minimal-devui" -version = "0.1.0" -source = { editable = "samples/web-fastapi-minimal-devui" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-compat-oai" }, - { name = "genkit-plugin-fastapi" }, - { name = "uvicorn", extra = ["standard"] }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-compat-oai", editable = "plugins/compat-oai" }, - { name = "genkit-plugin-fastapi", editable = "plugins/fastapi" }, - { name = "uvicorn", extras = ["standard"], specifier = ">=0.34.0" }, -] - [[package]] name = "web-flask-hello" version = "0.2.0" From 06d2365ca91fd71d52a220ba3a38f8275d2eac72 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Thu, 12 Mar 2026 15:51:15 -0400 Subject: [PATCH 03/10] feat: Add optional `key` to action metadata, update schemas to reflect this and allow additional object properties, and ensure key is set during action listing. --- genkit-tools/common/src/manager/manager-v2.ts | 14 +++++++++++++- genkit-tools/common/src/types/action.ts | 9 ++++++++- genkit-tools/genkit-schema.json | 5 ++++- go/ai/gen.go | 1 + js/core/src/action.ts | 1 + js/core/src/reflection-v2.ts | 5 ++++- 6 files changed, 31 insertions(+), 4 deletions(-) diff --git a/genkit-tools/common/src/manager/manager-v2.ts b/genkit-tools/common/src/manager/manager-v2.ts index 286cbf0ead..7874ced572 100644 --- a/genkit-tools/common/src/manager/manager-v2.ts +++ b/genkit-tools/common/src/manager/manager-v2.ts @@ -408,7 +408,19 @@ export class RuntimeManagerV2 extends BaseRuntimeManager { : 'No runtimes found. Make sure your app is running using the `start_runtime` MCP tool or the CLI: `genkit start -- ...`. See getting started documentation.' ); } - return this.sendRequest(runtimeId, 'listActions'); + const result: Record = {}; + const response = ReflectionListActionsResponseSchema.parse( + await this.sendRequest(runtimeId, 'listActions') + ); + // make sure key is set on the action metadata + for (const key of Object.keys(response)) { + const action = response[key]; + if (!action.key) { + action.key = key; + } + result[key] = action as Action; + } + return result; } async listValues( diff --git a/genkit-tools/common/src/types/action.ts b/genkit-tools/common/src/types/action.ts index f9752d84a5..9cc2c59862 100644 --- a/genkit-tools/common/src/types/action.ts +++ b/genkit-tools/common/src/types/action.ts @@ -23,7 +23,13 @@ extendZodWithOpenApi(z); // There's no way around this without defining a custom type. `z.array()` needs an inner type which runs into the same issue. export const SingularAnySchema = z - .union([z.string(), z.number(), z.bigint(), z.boolean(), z.object({})]) + .union([ + z.string(), + z.number(), + z.bigint(), + z.boolean(), + z.object({}).passthrough(), + ]) .openapi('SingularAny'); export const CustomAnySchema = z @@ -60,6 +66,7 @@ export type Action = z.infer; export const ActionMetadataSchema = z .object({ + key: z.string().optional(), actionType: z.string().optional(), name: z.string(), description: z.string().optional(), diff --git a/genkit-tools/genkit-schema.json b/genkit-tools/genkit-schema.json index 9b8cccfa29..f8a53a07bb 100644 --- a/genkit-tools/genkit-schema.json +++ b/genkit-tools/genkit-schema.json @@ -1463,6 +1463,9 @@ "ActionMetadata": { "type": "object", "properties": { + "key": { + "type": "string" + }, "actionType": { "type": "string" }, @@ -1515,7 +1518,7 @@ { "type": "object", "properties": {}, - "additionalProperties": false + "additionalProperties": true } ] }, diff --git a/go/ai/gen.go b/go/ai/gen.go index d50edc3d36..12d245cac3 100644 --- a/go/ai/gen.go +++ b/go/ai/gen.go @@ -24,6 +24,7 @@ type ActionMetadata struct { // A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object. InputJsonSchema any `json:"inputJsonSchema,omitempty"` InputSchema any `json:"inputSchema,omitempty"` + Key string `json:"key,omitempty"` Metadata map[string]any `json:"metadata,omitempty"` Name string `json:"name,omitempty"` // A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object. diff --git a/js/core/src/action.ts b/js/core/src/action.ts index b1992e3aa6..7c320a5d39 100644 --- a/js/core/src/action.ts +++ b/js/core/src/action.ts @@ -63,6 +63,7 @@ export interface ActionMetadata< } export const ActionMetadataSchema = z.object({ + key: z.string().optional(), actionType: z.string().optional(), name: z.string(), description: z.string().optional(), diff --git a/js/core/src/reflection-v2.ts b/js/core/src/reflection-v2.ts index ad05a841f8..4384ff0744 100644 --- a/js/core/src/reflection-v2.ts +++ b/js/core/src/reflection-v2.ts @@ -15,6 +15,7 @@ */ import WebSocket from 'ws'; +import z from 'zod'; import { StatusCodes, type Status } from './action.js'; import { GENKIT_REFLECTION_API_SPEC_VERSION, GENKIT_VERSION } from './index.js'; import { logger } from './logging.js'; @@ -217,7 +218,9 @@ export class ReflectionServerV2 { private async handleListActions(request: JsonRpcRequest) { if (!request.id) return; // Should be a request const actions = await this.registry.listResolvableActions(); - const convertedActions: Record = {}; + const convertedActions: z.infer< + typeof ReflectionListActionsResponseSchema + > = {}; Object.keys(actions).forEach((key) => { const action = actions[key]; From 94d58eb24f0da8daea48940575242917f3128a27 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Thu, 12 Mar 2026 15:53:06 -0400 Subject: [PATCH 04/10] feat: Add optional `key` field to the `ActionMetadata` model. --- py/packages/genkit/src/genkit/core/typing.py | 1 + 1 file changed, 1 insertion(+) diff --git a/py/packages/genkit/src/genkit/core/typing.py b/py/packages/genkit/src/genkit/core/typing.py index 477416bbd2..de64299cb2 100644 --- a/py/packages/genkit/src/genkit/core/typing.py +++ b/py/packages/genkit/src/genkit/core/typing.py @@ -917,6 +917,7 @@ class ActionMetadata(BaseModel): """Model for actionmetadata data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + key: str | None = None action_type: str | None = Field(default=None) name: str description: str | None = None From b63300a31e4a169fde38c6bd99c2f18ec53fea4b Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Thu, 19 Mar 2026 16:38:53 -0400 Subject: [PATCH 05/10] refactor: narrow JSON-RPC message ID type from `number | string` to `string`. --- genkit-tools/common/src/manager/manager-v2.ts | 4 ++-- js/core/src/reflection-v2.ts | 13 ++++--------- 2 files changed, 6 insertions(+), 11 deletions(-) diff --git a/genkit-tools/common/src/manager/manager-v2.ts b/genkit-tools/common/src/manager/manager-v2.ts index 7874ced572..51938f2c33 100644 --- a/genkit-tools/common/src/manager/manager-v2.ts +++ b/genkit-tools/common/src/manager/manager-v2.ts @@ -52,7 +52,7 @@ interface JsonRpcRequest { jsonrpc: '2.0'; method: string; params?: any; - id?: number | string; + id?: string; } interface JsonRpcResponse { @@ -63,7 +63,7 @@ interface JsonRpcResponse { message: string; data?: any; }; - id: number | string; + id: string; } type JsonRpcMessage = JsonRpcRequest | JsonRpcResponse; diff --git a/js/core/src/reflection-v2.ts b/js/core/src/reflection-v2.ts index 4384ff0744..0155684287 100644 --- a/js/core/src/reflection-v2.ts +++ b/js/core/src/reflection-v2.ts @@ -40,7 +40,7 @@ interface JsonRpcRequest { jsonrpc: '2.0'; method: string; params?: any; - id?: number | string; + id?: string; } interface JsonRpcResponse { @@ -51,7 +51,7 @@ interface JsonRpcResponse { message: string; data?: any; }; - id: number | string; + id: string; } type JsonRpcMessage = JsonRpcRequest | JsonRpcResponse; @@ -128,7 +128,7 @@ export class ReflectionServerV2 { } } - private sendResponse(id: number | string, result: any) { + private sendResponse(id: string, result: any) { this.send({ jsonrpc: '2.0', result, @@ -136,12 +136,7 @@ export class ReflectionServerV2 { }); } - private sendError( - id: number | string, - code: number, - message: string, - data?: any - ) { + private sendError(id: string, code: number, message: string, data?: any) { this.send({ jsonrpc: '2.0', error: { code, message, data }, From 1545f93d1fe75ef2f40ed4926e2d408c394d933c Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Thu, 19 Mar 2026 16:42:46 -0400 Subject: [PATCH 06/10] chore: Update pnpm lock file --- js/pnpm-lock.yaml | 213 +++++++++------------------------------------- 1 file changed, 38 insertions(+), 175 deletions(-) diff --git a/js/pnpm-lock.yaml b/js/pnpm-lock.yaml index fe1769dd8c..5b761b6bdc 100644 --- a/js/pnpm-lock.yaml +++ b/js/pnpm-lock.yaml @@ -146,7 +146,7 @@ importers: version: 0.4.0 ws: specifier: ^8.18.0 - version: 8.18.3 + version: 8.19.0 zod: specifier: ^3.23.8 version: 3.25.76 @@ -159,7 +159,7 @@ importers: version: 4.17.25 '@types/node': specifier: ^20.11.30 - version: 20.19.1 + version: 20.19.37 '@types/ws': specifier: ^8.5.10 version: 8.18.1 @@ -187,7 +187,7 @@ importers: version: 4.1.1 '@genkit-ai/firebase': specifier: ^1.16.1 - version: 1.29.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) + version: 1.29.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) doc-snippets: dependencies: @@ -243,7 +243,7 @@ importers: version: 4.17.25 '@types/node': specifier: ^22.15.3 - version: 22.19.15 + version: 22.15.32 '@types/uuid': specifier: ^9.0.6 version: 9.0.8 @@ -2905,11 +2905,11 @@ packages: resolution: {integrity: sha512-Lm/ZLhDZcBECta3TmCQSngiQykFdfw+QtI1/GYMsZd4l3nG+P8WLB16XuS7WaBGLQ+9E+cOcWQsth9cayuGt8g==} engines: {node: ^20.17.0 || >=22.9.0} - '@genkit-ai/ai@1.30.0-rc.0': - resolution: {integrity: sha512-TG+qpnFrTVml9iNdQBCVHCzRy37uIx4qmF8GA1AA511MiEF/klIjdP0KX60s10Lj9DsYVCDkaELH1BtsQKH2cA==} + '@genkit-ai/ai@1.30.1': + resolution: {integrity: sha512-uaf9H8BSUHGBRg0THg6tDLCEEz+FRHTLY4RjOqhVvg0gqZOOaZuccKF3yOb6AU8FptvQE12p7wgDf5yPNBIXGA==} - '@genkit-ai/core@1.30.0-rc.0': - resolution: {integrity: sha512-DGq1paDS0sCfM2gVrPnSEkCgcyIhnhptH+qnr361cWdkPaAVCkPgWJeKu87ayNuoo/kSnS3oshSdJhSnC1uKBA==} + '@genkit-ai/core@1.30.1': + resolution: {integrity: sha512-tW+FjvVpX8cdslHLM85PBe83M5xCsvMazRpN9+XJGalqymtUVc/L3xqAKCL7iGEVveqcJTsjBFfgSBwee80JLg==} '@genkit-ai/express@1.29.0': resolution: {integrity: sha512-+Mf8e45RbAlvns3a3hvIDpUAjNyWTIu1cy7vV86+6NvpPqPhkEaPkptgIX7KIXsPBEwiwsEM4SXx80M/TRbmTg==} @@ -2928,29 +2928,13 @@ packages: firebase: optional: true - '@genkit-ai/firebase@1.25.0': - resolution: {integrity: sha512-Z0FbnJHQs8qS0yxG++Dn3CZ7gv+YNaihGaWXoDKy02mNOkeRzHA6UPaWxSTaWkWHYdB0MyOnMGlyqxnWyqVdmg==} - peerDependencies: - '@google-cloud/firestore': ^7.11.0 - firebase: '>=11.5.0' - firebase-admin: '>=12.2' - genkit: ^1.25.0 - peerDependenciesMeta: - firebase: - optional: true - - '@genkit-ai/google-cloud@1.16.1': - resolution: {integrity: sha512-uujjdGr/sra7iKHApufwkt5jGo7CQcRCJNWPgnSg4g179CjtvtZBGjxmFRVBtKzuF61ktkY6E9JoLz83nWEyAA==} + '@genkit-ai/google-cloud@1.30.1': + resolution: {integrity: sha512-Rx5dh/qSRMZhE1GW258Fy03OhgZP8OD60OWxHIxyHf6LBHzEcGo+YVb3Q0Pe4K37NCdm6wOlSMKGSUCdBDASQw==} peerDependencies: - genkit: ^1.29.0 + genkit: ^1.30.1 - '@genkit-ai/google-cloud@1.25.0': - resolution: {integrity: sha512-wHCa8JSTv7MtwzXjUQ9AT5v0kCTJrz0In+ffgAYw1yt8ComAz5o7Ir+xks+sX1vJfN8ptvW0GUa6rsUaXCB3kA==} - peerDependencies: - genkit: ^1.25.0 - - '@gerrit0/mini-shiki@1.27.2': - resolution: {integrity: sha512-GeWyHz8ao2gBiUW4OJnQDxXQnFgZQwwQk05t/CVVgNBN7/rK8XZ7xY6YhLVv9tH3VppWWmr9DCl3MwemB/i+Og==} + '@gerrit0/mini-shiki@3.23.0': + resolution: {integrity: sha512-bEMORlG0cqdjVyCEuU0cDQbORWX+kYCeo0kV1lbxF5bt4r7SID2l9bqsxJEM0zndaxpOUT7riCyIVEuqq/Ynxg==} '@google-cloud/aiplatform@3.35.0': resolution: {integrity: sha512-Eo+ckr1KbTxAOew9P+MeeR0aQXeW5PeOzrSM1JyGny/SGKejwX/RcGWSFpeapnlegTfI9N9xJeUeo3M+XBOeFg==} @@ -4576,12 +4560,6 @@ packages: '@types/node@20.19.37': resolution: {integrity: sha512-8kzdPJ3FsNsVIurqBs7oodNnCEVbni9yUEkaHbgptDACOPW04jimGagZ51E6+lXUwJjgnBw+hyko/lkFWCldqw==} - '@types/node@20.19.1': - resolution: {integrity: sha512-jJD50LtlD2dodAEO653i3YF04NWak6jN3ky+Ri3Em3mGR39/glWiboM/IePaRbgwSfqM1TpGXfAg8ohn/4dTgA==} - - '@types/node@20.19.26': - resolution: {integrity: sha512-0l6cjgF0XnihUpndDhk+nyD3exio3iKaYROSgvh/qSevPXax3L8p5DBRFjbvalnwatGgHEQn2R88y2fA3g4irg==} - '@types/node@22.15.32': resolution: {integrity: sha512-3jigKqgSjsH6gYZv2nEsqdXfZqIFGAV36XYYjf9KGZ3PSG+IhLecqPnI310RvjutyMwifE2hhhNEklOUrvx/wA==} @@ -5950,8 +5928,8 @@ packages: resolution: {integrity: sha512-zV/5HKTfCeKWnxG0Dmrw51hEWFGfcF2xiXqcA3+J90WDuP0SvoiSO5ORvcBsifmx/FoIjgQN3oNOGaQ5PhLFkg==} engines: {node: '>=18'} - genkit@1.30.0-rc.0: - resolution: {integrity: sha512-ras9lR4DgF+5A1G9pV/eiJt44yamLgGPzQEv/P/FOjo9JvGJCdMd9ze0wPakAlExLbocQ5DZgdU3kxwI3koRhw==} + genkit@1.30.1: + resolution: {integrity: sha512-+v+IuE0SxoCC4AYMAuPi0YkEAT/Tm/oJhYECB+xFieebaxsdBUJ4FEp2u3HrDSyWv1Gj28unh8Og332ZdFJGRQ==} genkitx-openai@0.30.0: resolution: {integrity: sha512-6Dt3o6f+cKrZ1njSzYj2rm0U8SE5nQDwpz5VbFuLjaloE5BBj2c9OSlPGegdcz4PN5wBocGKiNsaKn0t5PpJdA==} @@ -9749,11 +9727,11 @@ snapshots: retry: 0.13.1 optional: true - '@genkit-ai/ai@1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0))': + '@genkit-ai/ai@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0))': dependencies: - '@genkit-ai/core': 1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) + '@genkit-ai/core': 1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) '@opentelemetry/api': 1.9.0 - '@types/node': 20.19.26 + '@types/node': 20.19.37 colorette: 2.0.20 dotprompt: 1.1.2 json5: 2.2.3 @@ -9770,27 +9748,7 @@ snapshots: - supports-color optional: true - '@genkit-ai/ai@1.29.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit)': - dependencies: - '@genkit-ai/core': 1.29.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) - '@opentelemetry/api': 1.9.0 - '@types/node': 20.19.26 - colorette: 2.0.20 - dotprompt: 1.1.1 - json5: 2.2.3 - node-fetch: 3.3.2 - partial-json: 0.1.7 - uri-templates: 0.2.0 - uuid: 10.0.0 - transitivePeerDependencies: - - '@google-cloud/firestore' - - encoding - - firebase - - firebase-admin - - genkit - - supports-color - - '@genkit-ai/core@1.29.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.29.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1))': + '@genkit-ai/core@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0))': dependencies: '@opentelemetry/api': 1.9.0 '@opentelemetry/context-async-hooks': 1.25.1(@opentelemetry/api@1.9.0) @@ -9812,7 +9770,7 @@ snapshots: zod-to-json-schema: 3.25.1(zod@3.25.76) optionalDependencies: '@cfworker/json-schema': 4.1.1 - '@genkit-ai/firebase': 1.29.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) + '@genkit-ai/firebase': 1.29.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) transitivePeerDependencies: - '@google-cloud/firestore' - encoding @@ -9826,39 +9784,17 @@ snapshots: dependencies: body-parser: 1.20.4 cors: 2.8.6 - dotprompt: 1.1.1 - express: 4.21.2 - get-port: 5.1.1 - json-schema: 0.4.0 - zod: 3.25.67 - zod-to-json-schema: 3.24.5(zod@3.25.67) - optionalDependencies: - '@cfworker/json-schema': 4.1.1 - '@genkit-ai/firebase': 1.25.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) - transitivePeerDependencies: - - '@google-cloud/firestore' - - encoding - - firebase - - firebase-admin - - genkit - - supports-color - - '@genkit-ai/express@1.12.0(@genkit-ai/core@1.29.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit))(express@5.1.0)(genkit@genkit)': - dependencies: - '@genkit-ai/core': 1.29.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit) - body-parser: 1.20.3 - cors: 2.8.5 - express: 5.1.0 + express: 5.2.1 genkit: link:genkit transitivePeerDependencies: - supports-color - '@genkit-ai/firebase@1.29.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0))': + '@genkit-ai/firebase@1.29.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0))': dependencies: - '@genkit-ai/google-cloud': 1.29.0(encoding@0.1.13)(genkit@1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) + '@genkit-ai/google-cloud': 1.30.1(encoding@0.1.13)(genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) '@google-cloud/firestore': 7.11.6(encoding@0.1.13) firebase-admin: 13.7.0(encoding@0.1.13) - genkit: 1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0) + genkit: 1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0) optionalDependencies: firebase: 11.10.0 transitivePeerDependencies: @@ -9866,35 +9802,10 @@ snapshots: - supports-color optional: true - '@genkit-ai/firebase@1.25.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@1.28.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1))': + '@genkit-ai/google-cloud@1.30.1(encoding@0.1.13)(genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0))': dependencies: - '@genkit-ai/google-cloud': 1.25.0(encoding@0.1.13)(genkit@1.28.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)) - '@google-cloud/firestore': 7.11.1(encoding@0.1.13) - firebase-admin: 13.6.0(encoding@0.1.13) - genkit: 1.28.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1) - optionalDependencies: - firebase: 11.9.1 - transitivePeerDependencies: - - encoding - - supports-color - optional: true - - '@genkit-ai/firebase@1.25.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1)(genkit@genkit)': - dependencies: - '@genkit-ai/google-cloud': 1.25.0(encoding@0.1.13)(genkit@genkit) - '@google-cloud/firestore': 7.11.1(encoding@0.1.13) - firebase-admin: 13.6.0(encoding@0.1.13) - genkit: link:genkit - optionalDependencies: - firebase: 11.9.1 - transitivePeerDependencies: - - encoding - - supports-color - optional: true - - '@genkit-ai/google-cloud@1.16.1(encoding@0.1.13)(genkit@1.29.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1))': - dependencies: - '@google-cloud/logging-winston': 6.0.1(encoding@0.1.13)(winston@3.17.0) + '@google-cloud/logging-winston': 6.0.1(encoding@0.1.13)(winston@3.19.0) + '@google-cloud/modelarmor': 0.4.1 '@google-cloud/opentelemetry-cloud-monitoring-exporter': 0.19.0(@opentelemetry/api@1.9.0)(@opentelemetry/core@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/sdk-metrics@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) '@google-cloud/opentelemetry-cloud-trace-exporter': 2.4.1(@opentelemetry/api@1.9.0)(@opentelemetry/core@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/sdk-trace-base@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) '@google-cloud/opentelemetry-resource-util': 2.4.0(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) @@ -9908,7 +9819,7 @@ snapshots: '@opentelemetry/sdk-metrics': 1.25.1(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-node': 0.52.1(@opentelemetry/api@1.9.0) '@opentelemetry/sdk-trace-base': 1.25.1(@opentelemetry/api@1.9.0) - genkit: 1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0) + genkit: 1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0) google-auth-library: 9.15.1(encoding@0.1.13) node-fetch: 3.3.2 winston: 3.19.0 @@ -9917,32 +9828,7 @@ snapshots: - supports-color optional: true - '@genkit-ai/google-cloud@1.25.0(encoding@0.1.13)(genkit@1.28.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1))': - dependencies: - '@google-cloud/logging-winston': 6.0.1(encoding@0.1.13)(winston@3.17.0) - '@google-cloud/opentelemetry-cloud-monitoring-exporter': 0.19.0(@opentelemetry/api@1.9.0)(@opentelemetry/core@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/sdk-metrics@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) - '@google-cloud/opentelemetry-cloud-trace-exporter': 2.4.1(@opentelemetry/api@1.9.0)(@opentelemetry/core@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(@opentelemetry/sdk-trace-base@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) - '@google-cloud/opentelemetry-resource-util': 2.4.0(@opentelemetry/resources@1.25.1(@opentelemetry/api@1.9.0))(encoding@0.1.13) - '@opentelemetry/api': 1.9.0 - '@opentelemetry/auto-instrumentations-node': 0.49.2(@opentelemetry/api@1.9.0)(encoding@0.1.13) - '@opentelemetry/core': 1.25.1(@opentelemetry/api@1.9.0) - '@opentelemetry/instrumentation': 0.52.1(@opentelemetry/api@1.9.0) - '@opentelemetry/instrumentation-pino': 0.41.0(@opentelemetry/api@1.9.0) - '@opentelemetry/instrumentation-winston': 0.39.0(@opentelemetry/api@1.9.0) - '@opentelemetry/resources': 1.25.1(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-metrics': 1.25.1(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-node': 0.52.1(@opentelemetry/api@1.9.0) - '@opentelemetry/sdk-trace-base': 1.25.1(@opentelemetry/api@1.9.0) - genkit: 1.28.0(@google-cloud/firestore@7.11.1(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.6.0(encoding@0.1.13))(firebase@11.9.1) - google-auth-library: 9.15.1(encoding@0.1.13) - node-fetch: 3.3.2 - winston: 3.17.0 - transitivePeerDependencies: - - encoding - - supports-color - optional: true - - '@genkit-ai/google-cloud@1.25.0(encoding@0.1.13)(genkit@genkit)': + '@gerrit0/mini-shiki@3.23.0': dependencies: '@shikijs/engine-oniguruma': 3.23.0 '@shikijs/langs': 3.23.0 @@ -11634,11 +11520,7 @@ snapshots: '@types/graceful-fs@4.1.9': dependencies: - '@types/node': 20.19.26 - - '@types/handlebars@4.1.0': - dependencies: - handlebars: 4.7.8 + '@types/node': 20.19.37 '@types/hast@3.0.4': dependencies: @@ -11697,10 +11579,6 @@ snapshots: dependencies: undici-types: 6.21.0 - '@types/node@20.19.26': - dependencies: - undici-types: 6.21.0 - '@types/node@22.15.32': dependencies: undici-types: 6.21.0 @@ -11791,7 +11669,7 @@ snapshots: '@types/ws@8.18.1': dependencies: - '@types/node': 20.19.26 + '@types/node': 20.19.37 '@types/yargs-parser@21.0.3': {} @@ -13430,10 +13308,10 @@ snapshots: transitivePeerDependencies: - supports-color - genkit@1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0): + genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0): dependencies: - '@genkit-ai/ai': 1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) - '@genkit-ai/core': 1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.0-rc.0(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) + '@genkit-ai/ai': 1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) + '@genkit-ai/core': 1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)(genkit@1.30.1(@google-cloud/firestore@7.11.6(encoding@0.1.13))(encoding@0.1.13)(firebase-admin@13.7.0(encoding@0.1.13))(firebase@11.10.0)) uuid: 10.0.0 transitivePeerDependencies: - '@google-cloud/firestore' @@ -14108,7 +13986,7 @@ snapshots: '@jest/expect': 29.7.0 '@jest/test-result': 29.7.0 '@jest/types': 29.6.3 - '@types/node': 20.19.26 + '@types/node': 20.19.37 chalk: 4.1.2 co: 4.6.0 dedent: 1.7.2 @@ -14201,7 +14079,7 @@ snapshots: '@jest/environment': 29.7.0 '@jest/fake-timers': 29.7.0 '@jest/types': 29.6.3 - '@types/node': 20.19.26 + '@types/node': 20.19.37 jest-mock: 29.7.0 jest-util: 29.7.0 @@ -14387,7 +14265,7 @@ snapshots: jest-worker@29.7.0: dependencies: - '@types/node': 20.19.26 + '@types/node': 20.19.37 jest-util: 29.7.0 merge-stream: 2.0.0 supports-color: 8.1.1 @@ -15547,22 +15425,7 @@ snapshots: '@protobufjs/path': 1.1.2 '@protobufjs/pool': 1.1.0 '@protobufjs/utf8': 1.1.0 - '@types/node': 20.19.1 - long: 5.2.3 - - protobufjs@7.5.3: - dependencies: - '@protobufjs/aspromise': 1.1.2 - '@protobufjs/base64': 1.1.2 - '@protobufjs/codegen': 2.0.4 - '@protobufjs/eventemitter': 1.1.0 - '@protobufjs/fetch': 1.1.0 - '@protobufjs/float': 1.0.2 - '@protobufjs/inquire': 1.1.0 - '@protobufjs/path': 1.1.2 - '@protobufjs/pool': 1.1.0 - '@protobufjs/utf8': 1.1.0 - '@types/node': 20.19.26 + '@types/node': 20.19.37 long: 5.3.2 proxy-addr@2.0.7: From dcc0e817edaf401a1db05e6592be9706b4c2bd8b Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Thu, 19 Mar 2026 16:48:46 -0400 Subject: [PATCH 07/10] regen --- genkit-tools/pnpm-lock.yaml | 19 +- .../genkit/src/genkit/_core/_typing.py | 282 +++++++----------- py/uv.lock | 120 +------- 3 files changed, 136 insertions(+), 285 deletions(-) diff --git a/genkit-tools/pnpm-lock.yaml b/genkit-tools/pnpm-lock.yaml index 0c524cf68f..4aedd39f8e 100644 --- a/genkit-tools/pnpm-lock.yaml +++ b/genkit-tools/pnpm-lock.yaml @@ -144,6 +144,9 @@ importers: cors: specifier: ^2.8.5 version: 2.8.6 + events: + specifier: ^3.3.0 + version: 3.3.0 express: specifier: ^4.21.0 version: 4.22.1 @@ -174,6 +177,9 @@ importers: winston: specifier: ^3.11.0 version: 3.19.0 + ws: + specifier: ^8.18.3 + version: 8.18.3 yaml: specifier: ^2.4.1 version: 2.8.2 @@ -1157,6 +1163,9 @@ packages: '@types/cors@2.8.19': resolution: {integrity: sha512-mFNylyeyqN93lfe/9CSxOGREz8cpzAhH+E93xJ4xWQf62V8sQ/24reV2nyzUWM6H6Xji+GGHpkbLe7pVoUEskg==} + '@types/events@3.0.3': + resolution: {integrity: sha512-trOc4AAUThEz9hapPtSd7wf5tiQKvTtu5b371UxXdTuqzIh0ArcRspRP0i0Viu+LXstIQ1z96t1nsPxT9ol01g==} + '@types/express-serve-static-core@4.19.8': resolution: {integrity: sha512-02S5fmqeoKzVZCHPZid4b8JH2eM5HzQLZWN2FohQEy/0eXTq8VXZfSN6Pcr3F6N9R/vNrj7cpgbhjie6m/1tCA==} @@ -1844,6 +1853,10 @@ packages: resolution: {integrity: sha512-i/2XbnSz/uxRCU6+NdVJgKWDTM427+MqYbkQzD321DuCQJUqOuJKIA0IM2+W2xtYHdKOmZ4dR6fExsd4SXL+WQ==} engines: {node: '>=6'} + events@3.3.0: + resolution: {integrity: sha512-mQw+2fkQbALzQ7V0MY0IqdnXNOeTtP4r0lN9z7AAawCXgqea7bDii20AYrIBrFd/Hx0M2Ocz6S111CaFkUcb0Q==} + engines: {node: '>=0.8.x'} + eventsource-parser@3.0.6: resolution: {integrity: sha512-Vo1ab+QXPzZ4tCa8SwIHJFaSzy4R6SHf7BY79rFBDf0idraZWAkYrDjDj8uWaSm3S2TK+hJ7/t1CEmZ7jXw+pg==} engines: {node: '>=18.0.0'} @@ -4553,6 +4566,8 @@ snapshots: dependencies: '@types/node': 20.19.37 + '@types/events@3.0.3': {} + '@types/express-serve-static-core@4.19.8': dependencies: '@types/node': 20.19.37 @@ -4651,7 +4666,7 @@ snapshots: '@types/ws@8.18.1': dependencies: - '@types/node': 20.19.1 + '@types/node': 20.19.37 '@types/yargs-parser@21.0.3': {} @@ -5342,6 +5357,8 @@ snapshots: event-target-shim@5.0.1: {} + events@3.3.0: {} + eventsource-parser@3.0.6: {} eventsource@3.0.7: diff --git a/py/packages/genkit/src/genkit/_core/_typing.py b/py/packages/genkit/src/genkit/_core/_typing.py index 6cbb35580f..b27054f687 100644 --- a/py/packages/genkit/src/genkit/_core/_typing.py +++ b/py/packages/genkit/src/genkit/_core/_typing.py @@ -29,9 +29,7 @@ from genkit._core._base import GenkitModel -warnings.filterwarnings( - 'ignore', message='Field name "schema" in "OutputConfig" shadows an attribute in parent', category=UserWarning -) +warnings.filterwarnings('ignore', message='Field name "schema" in "OutputConfig" shadows an attribute in parent', category=UserWarning) class EvalStatusEnum(StrEnum): @@ -41,7 +39,6 @@ class EvalStatusEnum(StrEnum): PASS_ = 'PASS' FAIL = 'FAIL' - class FinishReason(StrEnum): """FinishReason data type class.""" @@ -52,7 +49,6 @@ class FinishReason(StrEnum): OTHER = 'other' UNKNOWN = 'unknown' - class Role(StrEnum): """Role data type class.""" @@ -61,58 +57,43 @@ class Role(StrEnum): MODEL = 'model' TOOL = 'tool' - class Schema(GenkitModel): """Model for schema data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - class ConfigSchema(GenkitModel): """Model for configschema data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - Metadata = dict[str, Any] # type alias for flexible metadata Custom = dict[str, Any] # type alias for flexible custom data - class DocumentData(GenkitModel): """Model for documentdata data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) content: list[DocumentPart] = Field(...) metadata: Metadata | None = None - class EmbedRequest(GenkitModel): """Model for embedrequest data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) input: list[DocumentData] = Field(...) options: Any | None = Field(default=None) - class EmbedResponse(GenkitModel): """Model for embedresponse data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) embeddings: list[Embedding] = Field(...) - class Embedding(GenkitModel): """Model for embedding data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) embedding: list[float] = Field(...) metadata: Metadata | None = None - class BaseDataPoint(GenkitModel): """Model for basedatapoint data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) input: Any | None = Field(default=None) output: Any | None = Field(default=None) @@ -121,10 +102,8 @@ class BaseDataPoint(GenkitModel): test_case_id: str | None = None trace_ids: list[str] | None = None - class BaseEvalDataPoint(GenkitModel): """Model for baseevaldatapoint data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) input: Any | None = Field(default=None) output: Any | None = Field(default=None) @@ -133,10 +112,8 @@ class BaseEvalDataPoint(GenkitModel): test_case_id: str = Field(...) trace_ids: list[str] | None = None - class EvalFnResponse(GenkitModel): """Model for evalfnresponse data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) sample_index: float | None = None test_case_id: str = Field(...) @@ -144,19 +121,15 @@ class EvalFnResponse(GenkitModel): span_id: str | None = None evaluation: Score = Field(...) - class EvalRequest(GenkitModel): """Model for evalrequest data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) dataset: list[BaseDataPoint] = Field(...) eval_run_id: str = Field(...) options: Any | None = Field(default=None) - class Score(GenkitModel): """Model for score data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) id: str | None = None score: bool | float | str | None = Field(default=None) @@ -164,29 +137,23 @@ class Score(GenkitModel): error: str | None = None details: Details | None = None - class GenkitError(GenkitModel): """Model for genkiterror data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) message: str = Field(...) stack: str | None = None details: Any | None = Field(default=None) data: Data | None = None - class CandidateError(GenkitModel): """Model for candidateerror data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) index: float = Field(...) code: Literal['blocked', 'other', 'unknown'] = Field(...) message: str | None = None - class Candidate(GenkitModel): """Model for candidate data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) index: float = Field(...) message: MessageData = Field(...) @@ -195,10 +162,8 @@ class Candidate(GenkitModel): finish_message: str | None = None custom: Any | None = Field(default=None) - class CustomPart(GenkitModel): """Model for custompart data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -210,10 +175,8 @@ class CustomPart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) - class DataPart(GenkitModel): """Model for datapart data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -225,10 +188,8 @@ class DataPart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) - class GenerateActionOptions(GenkitModel): """Model for generateactionoptions data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) model: str | None = None docs: list[DocumentData] | None = None @@ -244,10 +205,8 @@ class GenerateActionOptions(GenkitModel): step_name: str | None = None use: list[MiddlewareRef] | None = None - class GenerateActionOutputConfig(GenkitModel): """Model for generateactionoutputconfig data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) format: str | None = None content_type: str | None = None @@ -257,10 +216,8 @@ class GenerateActionOutputConfig(GenkitModel): # Store Pydantic type for runtime validation (excluded from JSON) schema_type: Any = Field(default=None, exclude=True) - class GenerateResponseChunk(GenkitModel): """Model for generateresponsechunk data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) role: Role | None = None index: float | None = None @@ -268,10 +225,8 @@ class GenerateResponseChunk(GenkitModel): custom: Any | None = Field(default=None) aggregated: bool | None = None - class GenerationCommonConfig(GenkitModel): """Model for generationcommonconfig data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='allow', populate_by_name=True) version: str | None = None temperature: float | None = None @@ -281,10 +236,8 @@ class GenerationCommonConfig(GenkitModel): stop_sequences: list[str] | None = None api_key: str | None = None - class GenerationUsage(GenkitModel): """Model for generationusage data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) input_tokens: float | None = None output_tokens: float | None = None @@ -301,10 +254,8 @@ class GenerationUsage(GenkitModel): thoughts_tokens: float | None = None cached_content_tokens: float | None = None - class MediaPart(GenkitModel): """Model for mediapart data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Media = Field(...) @@ -316,37 +267,29 @@ class MediaPart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) - class MessageData(GenkitModel): """Model for messagedata data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) role: Role | str = Field(...) content: list[Part] = Field(...) metadata: Metadata | None = None - class MiddlewareDesc(GenkitModel): """Model for middlewaredesc data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) description: str | None = None config_schema: Any | ConfigSchema | None = Field(default=None) metadata: Any | Metadata | None = Field(default=None) - class MiddlewareRef(GenkitModel): """Model for middlewareref data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) config: Any | None = Field(default=None) - class ModelInfo(GenkitModel): """Model for modelinfo data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) versions: list[str] | None = None label: str | None = None @@ -354,10 +297,8 @@ class ModelInfo(GenkitModel): supports: Supports | None = None stage: Stage | None = None - class ModelReference(GenkitModel): """Model for modelreference data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) config_schema: Any | None = Field(default=None) @@ -365,10 +306,8 @@ class ModelReference(GenkitModel): version: str | None = None config: Any | None = Field(default=None) - class ModelResponseChunk(GenkitModel): """Model for modelresponsechunk data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) role: Any | None = Field(default=None) index: float | None = None @@ -376,19 +315,15 @@ class ModelResponseChunk(GenkitModel): custom: Any | None = Field(default=None) aggregated: bool | None = None - class MultipartToolResponse(GenkitModel): """Model for multiparttoolresponse data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) output: Any | None = Field(default=None) content: list[Part] | None = None metadata: Metadata | None = None - class Operation(GenkitModel): """Model for operation data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) action: str | None = None id: str = Field(...) @@ -397,22 +332,16 @@ class Operation(GenkitModel): error: Error | None = None metadata: Metadata | None = None - class OutputConfig(GenkitModel): """Model for outputconfig data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict( - alias_generator=to_camel, extra='forbid', populate_by_name=True, protected_namespaces=() - ) + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True, protected_namespaces=()) format: str | None = None schema_: dict[str, Any] | None = None constrained: bool | None = None content_type: str | None = None - class ReasoningPart(GenkitModel): """Model for reasoningpart data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -424,10 +353,8 @@ class ReasoningPart(GenkitModel): reasoning: str = Field(...) resource: Any | None = Field(default=None) - class ResourcePart(GenkitModel): """Model for resourcepart data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -439,10 +366,8 @@ class ResourcePart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Resource = Field(...) - class TextPart(GenkitModel): """Model for textpart data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: str = Field(...) media: Any | None = Field(default=None) @@ -454,25 +379,17 @@ class TextPart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) - class ToolDefinition(GenkitModel): """Model for tooldefinition data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) description: str = Field(...) - input_schema: Any | dict[str, Any] | None = Field( - default=None, description='Valid JSON Schema representing the input of the tool.' - ) - output_schema: Any | dict[str, Any] | None = Field( - default=None, description='Valid JSON Schema describing the output of the tool.' - ) + input_schema: Any | dict[str, Any] | None = Field(default=None, description='Valid JSON Schema representing the input of the tool.') + output_schema: Any | dict[str, Any] | None = Field(default=None, description='Valid JSON Schema describing the output of the tool.') metadata: Metadata | None = None - class ToolRequestPart(GenkitModel): """Model for toolrequestpart data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -484,10 +401,8 @@ class ToolRequestPart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) - class ToolResponsePart(GenkitModel): """Model for toolresponsepart data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -499,78 +414,145 @@ class ToolResponsePart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) - class Media(GenkitModel): """Model for media data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) content_type: str | None = None url: str = Field(...) - class ToolRequest(GenkitModel): """Model for toolrequest data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) ref: str | None = None name: str = Field(...) input: Any | None = Field(default=None) partial: bool | None = None - class ToolResponse(GenkitModel): """Model for toolresponse data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) ref: str | None = None name: str = Field(...) output: Any | None = Field(default=None) content: list[Any] | None = None +class ActionMetadata(GenkitModel): + """Model for actionmetadata data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + key: str | None = None + action_type: str | None = None + name: str = Field(...) + description: str | None = None + input_schema: Any | None = Field(default=None) + input_json_schema: Any | dict[str, Any] | None = Field(default=None, description='A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object.') + output_schema: Any | None = Field(default=None) + output_json_schema: Any | None = Field(default=None, description='A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object.') + stream_schema: Any | None = Field(default=None) + metadata: Metadata | None = None + +class ReflectionCancelActionParams(GenkitModel): + """Model for reflectioncancelactionparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + trace_id: str = Field(...) + +class ReflectionCancelActionResponse(GenkitModel): + """Model for reflectioncancelactionresponse data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + message: str = Field(...) + +class ReflectionConfigureParams(GenkitModel): + """Model for reflectionconfigureparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + telemetry_server_url: str | None = None + +class ReflectionEndInputStreamParams(GenkitModel): + """Model for reflectionendinputstreamparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + request_id: str = Field(...) + +class ReflectionListActionsResponse(GenkitModel): + """Model for reflectionlistactionsresponse data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + +class ReflectionListValuesParams(GenkitModel): + """Model for reflectionlistvaluesparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + type: str = Field(...) + +class ReflectionListValuesResponse(GenkitModel): + """Model for reflectionlistvaluesresponse data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + +class ReflectionRegisterParams(GenkitModel): + """Model for reflectionregisterparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + id: str = Field(...) + pid: float = Field(...) + name: str | None = None + genkit_version: str | None = None + reflection_api_spec_version: float | None = None + +class ReflectionRunActionParams(GenkitModel): + """Model for reflectionrunactionparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + runtime_id: str | None = None + key: str = Field(..., description='Action key that consists of the action type and ID.') + input: Any | None = Field(default=None, description='An input with the type that this action expects.') + context: Any | None = Field(default=None, description='Additional runtime context data (ex. auth context data).') + telemetry_labels: TelemetryLabels | None = None + stream: bool | None = None + stream_input: bool | None = None + +class ReflectionRunActionStateParams(GenkitModel): + """Model for reflectionrunactionstateparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + request_id: str = Field(...) + state: State | None = None + +class ReflectionSendInputStreamChunkParams(GenkitModel): + """Model for reflectionsendinputstreamchunkparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + request_id: str = Field(...) + chunk: Any | None = Field(default=None) + +class ReflectionStreamChunkParams(GenkitModel): + """Model for reflectionstreamchunkparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + request_id: str = Field(...) + chunk: Any | None = Field(default=None) class InstrumentationLibrary(GenkitModel): """Model for instrumentationlibrary data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) version: str | None = None schema_url: str | None = None - class Link(GenkitModel): """Model for link data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) context: SpanContext | None = None attributes: Attributes | None = None dropped_attributes_count: float | None = None - class PathMetadata(GenkitModel): """Model for pathmetadata data.""" - - model_config: ClassVar[ConfigDict] = ConfigDict( - alias_generator=to_camel, extra='forbid', populate_by_name=True, frozen=True - ) + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True, frozen=True) path: str = Field(...) status: str = Field(...) error: str | None = None latency: float = Field(...) - class SpanContext(GenkitModel): """Model for spancontext data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) span_id: str = Field(...) is_remote: bool | None = None trace_flags: float = Field(...) - class SpanData(GenkitModel): """Model for spandata data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) span_id: str = Field(...) trace_id: str = Field(...) @@ -587,19 +569,15 @@ class SpanData(GenkitModel): time_events: TimeEvents | None = None truncated: bool | None = None - class SpanEndEvent(GenkitModel): """Model for spanendevent data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) span: SpanData = Field(...) type: str = Field(...) - class SpanMetadata(GenkitModel): """Model for spanmetadata data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) state: Literal['success', 'error'] | None = None @@ -609,43 +587,33 @@ class SpanMetadata(GenkitModel): metadata: Metadata | None = None path: str | None = None - class SpanStartEvent(GenkitModel): """Model for spanstartevent data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) span: SpanData = Field(...) type: str = Field(...) - class SpanStatus(GenkitModel): """Model for spanstatus data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) code: float = Field(...) message: str | None = None - class SpantEventBase(GenkitModel): """Model for spanteventbase data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) span: SpanData = Field(...) - class TimeEvent(GenkitModel): """Model for timeevent data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) time: float = Field(...) annotation: Annotation = Field(...) - class TraceData(GenkitModel): """Model for tracedata data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) display_name: str | None = None @@ -653,51 +621,39 @@ class TraceData(GenkitModel): end_time: float | None = None spans: Spans = Field(...) - class TraceMetadata(GenkitModel): """Model for tracemetadata data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) feature_name: str | None = None paths: list[PathMetadata] | None = None timestamp: float = Field(...) - class Details(GenkitModel): """Model for details data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='allow', populate_by_name=True) reasoning: str | None = None - class Data(GenkitModel): """Model for data data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) genkit_error_message: str | None = None genkit_error_details: GenkitErrorDetails | None = None - class GenkitErrorDetails(GenkitModel): """Model for genkiterrordetails data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) stack: str | None = None trace_id: str = Field(...) - class Resume(GenkitModel): """Model for resume data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) respond: list[ToolResponsePart] | None = None restart: list[ToolRequestPart] | None = None metadata: Metadata | None = None - class Supports(GenkitModel): """Model for supports data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) multiturn: bool | None = None media: bool | None = None @@ -710,108 +666,90 @@ class Supports(GenkitModel): tool_choice: bool | None = None long_running: bool | None = None - class Error(GenkitModel): """Model for error data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='allow', populate_by_name=True) message: str = Field(...) - class Resource(GenkitModel): """Model for resource data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) uri: str = Field(...) +class TelemetryLabels(GenkitModel): + """Model for telemetrylabels data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + +class State(GenkitModel): + """Model for state data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + trace_id: str | None = None class Attributes(GenkitModel): """Model for attributes data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - class SameProcessAsParentSpan(GenkitModel): """Model for sameprocessasparentspan data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) value: bool = Field(...) - class TimeEvents(GenkitModel): """Model for timeevents data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) time_event: list[TimeEvent] | None = None - class Annotation(GenkitModel): """Model for annotation data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) attributes: Attributes = Field(...) description: str = Field(...) - class Spans(GenkitModel): """Model for spans data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) - class DocumentPart(RootModel[TextPart | MediaPart]): """Root model for DocumentPart union (Part(root=X), DocumentPart(root=X)).""" - -class Part( - RootModel[ - TextPart | MediaPart | ToolRequestPart | ToolResponsePart | DataPart | CustomPart | ReasoningPart | ResourcePart - ] -): +class Part(RootModel[TextPart | MediaPart | ToolRequestPart | ToolResponsePart | DataPart | CustomPart | ReasoningPart | ResourcePart]): """Root model for Part union (Part(root=X), DocumentPart(root=X)).""" - TraceEvent = SpanStartEvent | SpanEndEvent - class EvalResponse(RootModel[list[EvalFnResponse]]): """Root model for evalresponse.""" - root: list[EvalFnResponse] class Constrained(StrEnum): """Constrained generation support (none, all, no-tools).""" - NONE = 'none' - ALL = 'all' - NO_TOOLS = 'no-tools' - + NONE = "none" + ALL = "all" + NO_TOOLS = "no-tools" class Stage(StrEnum): """Model stage (featured, stable, unstable, legacy, deprecated).""" - FEATURED = 'featured' - STABLE = 'stable' - UNSTABLE = 'unstable' - LEGACY = 'legacy' - DEPRECATED = 'deprecated' - + FEATURED = "featured" + STABLE = "stable" + UNSTABLE = "unstable" + LEGACY = "legacy" + DEPRECATED = "deprecated" class ToolChoice(StrEnum): """Tool choice for generation (auto, required, none).""" - AUTO = 'auto' - REQUIRED = 'required' - NONE = 'none' - + AUTO = "auto" + REQUIRED = "required" + NONE = "none" class MediaModel(RootModel[Any]): """Wrapper for media content (flexible structure).""" - class Text(RootModel[str]): """Plain text content.""" - Resource1 = Resource # alias for Resource (resource with uri) + diff --git a/py/uv.lock b/py/uv.lock index c0913dd934..3094029a6b 100644 --- a/py/uv.lock +++ b/py/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.14'", @@ -28,59 +28,13 @@ members = [ "genkit-plugin-ollama", "genkit-plugin-vertex-ai", "genkit-workspace", - "provider-amazon-bedrock-hello", - "provider-anthropic-hello", - "provider-compat-oai-hello", - "provider-firestore-retriever", - "provider-google-genai-code-execution", - "provider-google-genai-context-caching", - "provider-google-genai-hello", - "provider-google-genai-media-models-demo", - "provider-google-genai-vertexai-hello", - "provider-google-genai-vertexai-image", - "provider-ollama-hello", - "provider-vertex-ai-model-garden", - "provider-vertex-ai-rerank-eval", - "provider-vertex-ai-vector-search-bigquery", - "provider-vertex-ai-vector-search-firestore", - "web-fastapi-bugbot", - "web-flask-hello", -] - -[[package]] -name = "aioboto3" -version = "15.5.0" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiobotocore", extra = ["boto3"] }, - { name = "aiofiles" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/a2/01/92e9ab00f36e2899315f49eefcd5b4685fbb19016c7f19a9edf06da80bb0/aioboto3-15.5.0.tar.gz", hash = "sha256:ea8d8787d315594842fbfcf2c4dce3bac2ad61be275bc8584b2ce9a3402a6979", size = 255069, upload-time = "2025-10-30T13:37:16.122Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/e5/3e/e8f5b665bca646d43b916763c901e00a07e40f7746c9128bdc912a089424/aioboto3-15.5.0-py3-none-any.whl", hash = "sha256:cc880c4d6a8481dd7e05da89f41c384dbd841454fc1998ae25ca9c39201437a6", size = 35913, upload-time = "2025-10-30T13:37:14.549Z" }, -] - -[[package]] -name = "aiobotocore" -version = "2.25.1" -source = { registry = "https://pypi.org/simple" } -dependencies = [ - { name = "aiohttp" }, - { name = "aioitertools" }, - { name = "botocore" }, - { name = "jmespath" }, - { name = "multidict" }, - { name = "python-dateutil" }, - { name = "wrapt" }, -] -sdist = { url = "https://files.pythonhosted.org/packages/62/94/2e4ec48cf1abb89971cb2612d86f979a6240520f0a659b53a43116d344dc/aiobotocore-2.25.1.tar.gz", hash = "sha256:ea9be739bfd7ece8864f072ec99bb9ed5c7e78ebb2b0b15f29781fbe02daedbc", size = 120560, upload-time = "2025-10-28T22:33:21.787Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/95/2a/d275ec4ce5cd0096665043995a7d76f5d0524853c76a3d04656de49f8808/aiobotocore-2.25.1-py3-none-any.whl", hash = "sha256:eb6daebe3cbef5b39a0bb2a97cffbe9c7cb46b2fcc399ad141f369f3c2134b1f", size = 86039, upload-time = "2025-10-28T22:33:19.949Z" }, -] - -[package.optional-dependencies] -boto3 = [ - { name = "boto3" }, + "google-genai-media", + "middleware", + "output-formats", + "prompts", + "tool-interrupts", + "tracing", + "vertexai-imagen", ] overrides = [{ name = "werkzeug", specifier = ">=3.1.6" }] @@ -6837,64 +6791,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/68/5a/199c59e0a824a3db2b89c5d2dade7ab5f9624dbf6448dc291b46d5ec94d3/wcwidth-0.6.0-py3-none-any.whl", hash = "sha256:1a3a1e510b553315f8e146c54764f4fb6264ffad731b3d78088cdb1478ffbdad", size = 94189, upload-time = "2026-02-06T19:19:39.646Z" }, ] -[[package]] -name = "web-fastapi-bugbot" -version = "0.2.0" -source = { virtual = "samples/web-fastapi-bugbot" } -dependencies = [ - { name = "genkit" }, - { name = "genkit-plugin-fastapi" }, - { name = "genkit-plugin-google-genai" }, - { name = "python-dotenv" }, - { name = "uvicorn", extra = ["standard"] }, -] - -[package.optional-dependencies] -dev = [ - { name = "watchdog" }, -] - -[package.metadata] -requires-dist = [ - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-fastapi", editable = "plugins/fastapi" }, - { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, - { name = "python-dotenv", specifier = ">=1.0.0" }, - { name = "uvicorn", extras = ["standard"], specifier = ">=0.34.0" }, - { name = "watchdog", marker = "extra == 'dev'", specifier = ">=6.0.0" }, -] -provides-extras = ["dev"] - -[[package]] -name = "web-flask-hello" -version = "0.2.0" -source = { editable = "samples/web-flask-hello" } -dependencies = [ - { name = "flask" }, - { name = "genkit" }, - { name = "genkit-plugin-flask" }, - { name = "genkit-plugin-google-genai" }, - { name = "rich" }, - { name = "uvloop" }, -] - -[package.optional-dependencies] -dev = [ - { name = "watchdog" }, -] - -[package.metadata] -requires-dist = [ - { name = "flask" }, - { name = "genkit", editable = "packages/genkit" }, - { name = "genkit-plugin-flask", editable = "plugins/flask" }, - { name = "genkit-plugin-google-genai", editable = "plugins/google-genai" }, - { name = "rich", specifier = ">=13.0.0" }, - { name = "uvloop", specifier = ">=0.21.0" }, - { name = "watchdog", marker = "extra == 'dev'", specifier = ">=6.0.0" }, -] -provides-extras = ["dev"] - [[package]] name = "webcolors" version = "25.10.0" From 3321ac6c4b4cf3bc05ae0faaa2e621cba5daf782 Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Thu, 19 Mar 2026 16:49:49 -0400 Subject: [PATCH 08/10] regen --- .../genkit/src/genkit/_core/_typing.py | 230 ++++++++++++++++-- 1 file changed, 210 insertions(+), 20 deletions(-) diff --git a/py/packages/genkit/src/genkit/_core/_typing.py b/py/packages/genkit/src/genkit/_core/_typing.py index b27054f687..7eb9c13cd2 100644 --- a/py/packages/genkit/src/genkit/_core/_typing.py +++ b/py/packages/genkit/src/genkit/_core/_typing.py @@ -29,7 +29,9 @@ from genkit._core._base import GenkitModel -warnings.filterwarnings('ignore', message='Field name "schema" in "OutputConfig" shadows an attribute in parent', category=UserWarning) +warnings.filterwarnings( + 'ignore', message='Field name "schema" in "OutputConfig" shadows an attribute in parent', category=UserWarning +) class EvalStatusEnum(StrEnum): @@ -39,6 +41,7 @@ class EvalStatusEnum(StrEnum): PASS_ = 'PASS' FAIL = 'FAIL' + class FinishReason(StrEnum): """FinishReason data type class.""" @@ -49,6 +52,7 @@ class FinishReason(StrEnum): OTHER = 'other' UNKNOWN = 'unknown' + class Role(StrEnum): """Role data type class.""" @@ -57,43 +61,58 @@ class Role(StrEnum): MODEL = 'model' TOOL = 'tool' + class Schema(GenkitModel): """Model for schema data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + class ConfigSchema(GenkitModel): """Model for configschema data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + Metadata = dict[str, Any] # type alias for flexible metadata Custom = dict[str, Any] # type alias for flexible custom data + class DocumentData(GenkitModel): """Model for documentdata data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) content: list[DocumentPart] = Field(...) metadata: Metadata | None = None + class EmbedRequest(GenkitModel): """Model for embedrequest data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) input: list[DocumentData] = Field(...) options: Any | None = Field(default=None) + class EmbedResponse(GenkitModel): """Model for embedresponse data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) embeddings: list[Embedding] = Field(...) + class Embedding(GenkitModel): """Model for embedding data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) embedding: list[float] = Field(...) metadata: Metadata | None = None + class BaseDataPoint(GenkitModel): """Model for basedatapoint data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) input: Any | None = Field(default=None) output: Any | None = Field(default=None) @@ -102,8 +121,10 @@ class BaseDataPoint(GenkitModel): test_case_id: str | None = None trace_ids: list[str] | None = None + class BaseEvalDataPoint(GenkitModel): """Model for baseevaldatapoint data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) input: Any | None = Field(default=None) output: Any | None = Field(default=None) @@ -112,8 +133,10 @@ class BaseEvalDataPoint(GenkitModel): test_case_id: str = Field(...) trace_ids: list[str] | None = None + class EvalFnResponse(GenkitModel): """Model for evalfnresponse data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) sample_index: float | None = None test_case_id: str = Field(...) @@ -121,15 +144,19 @@ class EvalFnResponse(GenkitModel): span_id: str | None = None evaluation: Score = Field(...) + class EvalRequest(GenkitModel): """Model for evalrequest data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) dataset: list[BaseDataPoint] = Field(...) eval_run_id: str = Field(...) options: Any | None = Field(default=None) + class Score(GenkitModel): """Model for score data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) id: str | None = None score: bool | float | str | None = Field(default=None) @@ -137,23 +164,29 @@ class Score(GenkitModel): error: str | None = None details: Details | None = None + class GenkitError(GenkitModel): """Model for genkiterror data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) message: str = Field(...) stack: str | None = None details: Any | None = Field(default=None) data: Data | None = None + class CandidateError(GenkitModel): """Model for candidateerror data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) index: float = Field(...) code: Literal['blocked', 'other', 'unknown'] = Field(...) message: str | None = None + class Candidate(GenkitModel): """Model for candidate data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) index: float = Field(...) message: MessageData = Field(...) @@ -162,8 +195,10 @@ class Candidate(GenkitModel): finish_message: str | None = None custom: Any | None = Field(default=None) + class CustomPart(GenkitModel): """Model for custompart data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -175,8 +210,10 @@ class CustomPart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) + class DataPart(GenkitModel): """Model for datapart data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -188,8 +225,10 @@ class DataPart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) + class GenerateActionOptions(GenkitModel): """Model for generateactionoptions data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) model: str | None = None docs: list[DocumentData] | None = None @@ -205,8 +244,10 @@ class GenerateActionOptions(GenkitModel): step_name: str | None = None use: list[MiddlewareRef] | None = None + class GenerateActionOutputConfig(GenkitModel): """Model for generateactionoutputconfig data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) format: str | None = None content_type: str | None = None @@ -216,8 +257,10 @@ class GenerateActionOutputConfig(GenkitModel): # Store Pydantic type for runtime validation (excluded from JSON) schema_type: Any = Field(default=None, exclude=True) + class GenerateResponseChunk(GenkitModel): """Model for generateresponsechunk data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) role: Role | None = None index: float | None = None @@ -225,8 +268,10 @@ class GenerateResponseChunk(GenkitModel): custom: Any | None = Field(default=None) aggregated: bool | None = None + class GenerationCommonConfig(GenkitModel): """Model for generationcommonconfig data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='allow', populate_by_name=True) version: str | None = None temperature: float | None = None @@ -236,8 +281,10 @@ class GenerationCommonConfig(GenkitModel): stop_sequences: list[str] | None = None api_key: str | None = None + class GenerationUsage(GenkitModel): """Model for generationusage data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) input_tokens: float | None = None output_tokens: float | None = None @@ -254,8 +301,10 @@ class GenerationUsage(GenkitModel): thoughts_tokens: float | None = None cached_content_tokens: float | None = None + class MediaPart(GenkitModel): """Model for mediapart data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Media = Field(...) @@ -267,29 +316,37 @@ class MediaPart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) + class MessageData(GenkitModel): """Model for messagedata data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) role: Role | str = Field(...) content: list[Part] = Field(...) metadata: Metadata | None = None + class MiddlewareDesc(GenkitModel): """Model for middlewaredesc data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) description: str | None = None config_schema: Any | ConfigSchema | None = Field(default=None) metadata: Any | Metadata | None = Field(default=None) + class MiddlewareRef(GenkitModel): """Model for middlewareref data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) config: Any | None = Field(default=None) + class ModelInfo(GenkitModel): """Model for modelinfo data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) versions: list[str] | None = None label: str | None = None @@ -297,8 +354,10 @@ class ModelInfo(GenkitModel): supports: Supports | None = None stage: Stage | None = None + class ModelReference(GenkitModel): """Model for modelreference data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) config_schema: Any | None = Field(default=None) @@ -306,8 +365,10 @@ class ModelReference(GenkitModel): version: str | None = None config: Any | None = Field(default=None) + class ModelResponseChunk(GenkitModel): """Model for modelresponsechunk data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) role: Any | None = Field(default=None) index: float | None = None @@ -315,15 +376,19 @@ class ModelResponseChunk(GenkitModel): custom: Any | None = Field(default=None) aggregated: bool | None = None + class MultipartToolResponse(GenkitModel): """Model for multiparttoolresponse data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) output: Any | None = Field(default=None) content: list[Part] | None = None metadata: Metadata | None = None + class Operation(GenkitModel): """Model for operation data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) action: str | None = None id: str = Field(...) @@ -332,16 +397,22 @@ class Operation(GenkitModel): error: Error | None = None metadata: Metadata | None = None + class OutputConfig(GenkitModel): """Model for outputconfig data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True, protected_namespaces=()) + + model_config: ClassVar[ConfigDict] = ConfigDict( + alias_generator=to_camel, extra='forbid', populate_by_name=True, protected_namespaces=() + ) format: str | None = None schema_: dict[str, Any] | None = None constrained: bool | None = None content_type: str | None = None + class ReasoningPart(GenkitModel): """Model for reasoningpart data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -353,8 +424,10 @@ class ReasoningPart(GenkitModel): reasoning: str = Field(...) resource: Any | None = Field(default=None) + class ResourcePart(GenkitModel): """Model for resourcepart data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -366,8 +439,10 @@ class ResourcePart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Resource = Field(...) + class TextPart(GenkitModel): """Model for textpart data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: str = Field(...) media: Any | None = Field(default=None) @@ -379,17 +454,25 @@ class TextPart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) + class ToolDefinition(GenkitModel): """Model for tooldefinition data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) description: str = Field(...) - input_schema: Any | dict[str, Any] | None = Field(default=None, description='Valid JSON Schema representing the input of the tool.') - output_schema: Any | dict[str, Any] | None = Field(default=None, description='Valid JSON Schema describing the output of the tool.') + input_schema: Any | dict[str, Any] | None = Field( + default=None, description='Valid JSON Schema representing the input of the tool.' + ) + output_schema: Any | dict[str, Any] | None = Field( + default=None, description='Valid JSON Schema describing the output of the tool.' + ) metadata: Metadata | None = None + class ToolRequestPart(GenkitModel): """Model for toolrequestpart data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -401,8 +484,10 @@ class ToolRequestPart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) + class ToolResponsePart(GenkitModel): """Model for toolresponsepart data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) text: Any | None = Field(default=None) media: Any | None = Field(default=None) @@ -414,77 +499,105 @@ class ToolResponsePart(GenkitModel): reasoning: Any | None = Field(default=None) resource: Any | None = Field(default=None) + class Media(GenkitModel): """Model for media data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) content_type: str | None = None url: str = Field(...) + class ToolRequest(GenkitModel): """Model for toolrequest data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) ref: str | None = None name: str = Field(...) input: Any | None = Field(default=None) partial: bool | None = None + class ToolResponse(GenkitModel): """Model for toolresponse data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) ref: str | None = None name: str = Field(...) output: Any | None = Field(default=None) content: list[Any] | None = None + class ActionMetadata(GenkitModel): """Model for actionmetadata data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) key: str | None = None action_type: str | None = None name: str = Field(...) description: str | None = None input_schema: Any | None = Field(default=None) - input_json_schema: Any | dict[str, Any] | None = Field(default=None, description='A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object.') + input_json_schema: Any | dict[str, Any] | None = Field( + default=None, description='A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object.' + ) output_schema: Any | None = Field(default=None) - output_json_schema: Any | None = Field(default=None, description='A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object.') + output_json_schema: Any | None = Field( + default=None, description='A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object.' + ) stream_schema: Any | None = Field(default=None) metadata: Metadata | None = None + class ReflectionCancelActionParams(GenkitModel): """Model for reflectioncancelactionparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) + class ReflectionCancelActionResponse(GenkitModel): """Model for reflectioncancelactionresponse data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) message: str = Field(...) + class ReflectionConfigureParams(GenkitModel): """Model for reflectionconfigureparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) telemetry_server_url: str | None = None + class ReflectionEndInputStreamParams(GenkitModel): """Model for reflectionendinputstreamparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) request_id: str = Field(...) + class ReflectionListActionsResponse(GenkitModel): """Model for reflectionlistactionsresponse data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + class ReflectionListValuesParams(GenkitModel): """Model for reflectionlistvaluesparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) type: str = Field(...) + class ReflectionListValuesResponse(GenkitModel): """Model for reflectionlistvaluesresponse data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + class ReflectionRegisterParams(GenkitModel): """Model for reflectionregisterparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) id: str = Field(...) pid: float = Field(...) @@ -492,8 +605,10 @@ class ReflectionRegisterParams(GenkitModel): genkit_version: str | None = None reflection_api_spec_version: float | None = None + class ReflectionRunActionParams(GenkitModel): """Model for reflectionrunactionparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) runtime_id: str | None = None key: str = Field(..., description='Action key that consists of the action type and ID.') @@ -503,56 +618,74 @@ class ReflectionRunActionParams(GenkitModel): stream: bool | None = None stream_input: bool | None = None + class ReflectionRunActionStateParams(GenkitModel): """Model for reflectionrunactionstateparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) request_id: str = Field(...) state: State | None = None + class ReflectionSendInputStreamChunkParams(GenkitModel): """Model for reflectionsendinputstreamchunkparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) request_id: str = Field(...) chunk: Any | None = Field(default=None) + class ReflectionStreamChunkParams(GenkitModel): """Model for reflectionstreamchunkparams data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) request_id: str = Field(...) chunk: Any | None = Field(default=None) + class InstrumentationLibrary(GenkitModel): """Model for instrumentationlibrary data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) version: str | None = None schema_url: str | None = None + class Link(GenkitModel): """Model for link data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) context: SpanContext | None = None attributes: Attributes | None = None dropped_attributes_count: float | None = None + class PathMetadata(GenkitModel): """Model for pathmetadata data.""" - model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True, frozen=True) + + model_config: ClassVar[ConfigDict] = ConfigDict( + alias_generator=to_camel, extra='forbid', populate_by_name=True, frozen=True + ) path: str = Field(...) status: str = Field(...) error: str | None = None latency: float = Field(...) + class SpanContext(GenkitModel): """Model for spancontext data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) span_id: str = Field(...) is_remote: bool | None = None trace_flags: float = Field(...) + class SpanData(GenkitModel): """Model for spandata data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) span_id: str = Field(...) trace_id: str = Field(...) @@ -569,15 +702,19 @@ class SpanData(GenkitModel): time_events: TimeEvents | None = None truncated: bool | None = None + class SpanEndEvent(GenkitModel): """Model for spanendevent data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) span: SpanData = Field(...) type: str = Field(...) + class SpanMetadata(GenkitModel): """Model for spanmetadata data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) name: str = Field(...) state: Literal['success', 'error'] | None = None @@ -587,33 +724,43 @@ class SpanMetadata(GenkitModel): metadata: Metadata | None = None path: str | None = None + class SpanStartEvent(GenkitModel): """Model for spanstartevent data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) span: SpanData = Field(...) type: str = Field(...) + class SpanStatus(GenkitModel): """Model for spanstatus data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) code: float = Field(...) message: str | None = None + class SpantEventBase(GenkitModel): """Model for spanteventbase data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) span: SpanData = Field(...) + class TimeEvent(GenkitModel): """Model for timeevent data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) time: float = Field(...) annotation: Annotation = Field(...) + class TraceData(GenkitModel): """Model for tracedata data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str = Field(...) display_name: str | None = None @@ -621,39 +768,51 @@ class TraceData(GenkitModel): end_time: float | None = None spans: Spans = Field(...) + class TraceMetadata(GenkitModel): """Model for tracemetadata data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) feature_name: str | None = None paths: list[PathMetadata] | None = None timestamp: float = Field(...) + class Details(GenkitModel): """Model for details data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='allow', populate_by_name=True) reasoning: str | None = None + class Data(GenkitModel): """Model for data data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) genkit_error_message: str | None = None genkit_error_details: GenkitErrorDetails | None = None + class GenkitErrorDetails(GenkitModel): """Model for genkiterrordetails data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) stack: str | None = None trace_id: str = Field(...) + class Resume(GenkitModel): """Model for resume data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) respond: list[ToolResponsePart] | None = None restart: list[ToolRequestPart] | None = None metadata: Metadata | None = None + class Supports(GenkitModel): """Model for supports data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) multiturn: bool | None = None media: bool | None = None @@ -666,90 +825,121 @@ class Supports(GenkitModel): tool_choice: bool | None = None long_running: bool | None = None + class Error(GenkitModel): """Model for error data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='allow', populate_by_name=True) message: str = Field(...) + class Resource(GenkitModel): """Model for resource data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) uri: str = Field(...) + class TelemetryLabels(GenkitModel): """Model for telemetrylabels data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + class State(GenkitModel): """Model for state data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) trace_id: str | None = None + class Attributes(GenkitModel): """Model for attributes data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + class SameProcessAsParentSpan(GenkitModel): """Model for sameprocessasparentspan data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) value: bool = Field(...) + class TimeEvents(GenkitModel): """Model for timeevents data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) time_event: list[TimeEvent] | None = None + class Annotation(GenkitModel): """Model for annotation data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) attributes: Attributes = Field(...) description: str = Field(...) + class Spans(GenkitModel): """Model for spans data.""" + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + class DocumentPart(RootModel[TextPart | MediaPart]): """Root model for DocumentPart union (Part(root=X), DocumentPart(root=X)).""" -class Part(RootModel[TextPart | MediaPart | ToolRequestPart | ToolResponsePart | DataPart | CustomPart | ReasoningPart | ResourcePart]): + +class Part( + RootModel[ + TextPart | MediaPart | ToolRequestPart | ToolResponsePart | DataPart | CustomPart | ReasoningPart | ResourcePart + ] +): """Root model for Part union (Part(root=X), DocumentPart(root=X)).""" + TraceEvent = SpanStartEvent | SpanEndEvent + class EvalResponse(RootModel[list[EvalFnResponse]]): """Root model for evalresponse.""" + root: list[EvalFnResponse] class Constrained(StrEnum): """Constrained generation support (none, all, no-tools).""" - NONE = "none" - ALL = "all" - NO_TOOLS = "no-tools" + NONE = 'none' + ALL = 'all' + NO_TOOLS = 'no-tools' + class Stage(StrEnum): """Model stage (featured, stable, unstable, legacy, deprecated).""" - FEATURED = "featured" - STABLE = "stable" - UNSTABLE = "unstable" - LEGACY = "legacy" - DEPRECATED = "deprecated" + FEATURED = 'featured' + STABLE = 'stable' + UNSTABLE = 'unstable' + LEGACY = 'legacy' + DEPRECATED = 'deprecated' + class ToolChoice(StrEnum): """Tool choice for generation (auto, required, none).""" - AUTO = "auto" - REQUIRED = "required" - NONE = "none" + AUTO = 'auto' + REQUIRED = 'required' + NONE = 'none' + class MediaModel(RootModel[Any]): """Wrapper for media content (flexible structure).""" + class Text(RootModel[str]): """Plain text content.""" -Resource1 = Resource # alias for Resource (resource with uri) +Resource1 = Resource # alias for Resource (resource with uri) From 4b9599eaf864797a687b8695e6cc8bfd890984ed Mon Sep 17 00:00:00 2001 From: Pavel Jbanov Date: Thu, 19 Mar 2026 17:07:19 -0400 Subject: [PATCH 09/10] refactor: update ReflectionListActionsResponse to encapsulate action metadata within an 'actions' field. --- docs/reflection-v2-protocol.md | 2 +- genkit-tools/common/src/manager/manager-v2.ts | 4 ++-- genkit-tools/common/src/types/reflection.ts | 7 +++---- genkit-tools/genkit-schema.json | 15 ++++++++++++--- go/ai/gen.go | 4 +++- js/core/src/reflection-types.ts | 7 +++---- js/core/src/reflection-v2.ts | 6 ++---- py/packages/genkit/src/genkit/_core/_typing.py | 7 +++++++ 8 files changed, 33 insertions(+), 19 deletions(-) diff --git a/docs/reflection-v2-protocol.md b/docs/reflection-v2-protocol.md index 6c9d7cd168..3d72ea04a4 100644 --- a/docs/reflection-v2-protocol.md +++ b/docs/reflection-v2-protocol.md @@ -157,7 +157,7 @@ JSON-RPC 2.0 does not natively support streaming. We extend it by using Notifica **Result:** | Type | Description | | :--- | :--- | -| `Record` | Map of action keys to Action definitions. (Same schema as V1 `/api/actions`) | +| `{ actions: Record }` | Object with an `actions` field containing a map of action keys to Action definitions. | ### 4. List Values **Direction:** Manager -> Runtime diff --git a/genkit-tools/common/src/manager/manager-v2.ts b/genkit-tools/common/src/manager/manager-v2.ts index 51938f2c33..2c9149920f 100644 --- a/genkit-tools/common/src/manager/manager-v2.ts +++ b/genkit-tools/common/src/manager/manager-v2.ts @@ -413,8 +413,8 @@ export class RuntimeManagerV2 extends BaseRuntimeManager { await this.sendRequest(runtimeId, 'listActions') ); // make sure key is set on the action metadata - for (const key of Object.keys(response)) { - const action = response[key]; + for (const key of Object.keys(response.actions)) { + const action = response.actions[key]; if (!action.key) { action.key = key; } diff --git a/genkit-tools/common/src/types/reflection.ts b/genkit-tools/common/src/types/reflection.ts index 9d46f4a4a1..809931e419 100644 --- a/genkit-tools/common/src/types/reflection.ts +++ b/genkit-tools/common/src/types/reflection.ts @@ -98,10 +98,9 @@ export const ReflectionEndInputStreamParamsSchema = z.object({ /** * ReflectionListActionsResponseSchema is the result for the 'listActions' method. */ -export const ReflectionListActionsResponseSchema = z.record( - z.string(), - ActionMetadataSchema -); +export const ReflectionListActionsResponseSchema = z.object({ + actions: z.record(z.string(), ActionMetadataSchema), +}); /** * ReflectionListValuesResponseSchema is the result for the 'listValues' method. diff --git a/genkit-tools/genkit-schema.json b/genkit-tools/genkit-schema.json index f8a53a07bb..4c25767f8b 100644 --- a/genkit-tools/genkit-schema.json +++ b/genkit-tools/genkit-schema.json @@ -1584,9 +1584,18 @@ }, "ReflectionListActionsResponse": { "type": "object", - "additionalProperties": { - "$ref": "#/$defs/ActionMetadata" - } + "properties": { + "actions": { + "type": "object", + "additionalProperties": { + "$ref": "#/$defs/ActionMetadata" + } + } + }, + "required": [ + "actions" + ], + "additionalProperties": false }, "ReflectionListValuesParams": { "type": "object", diff --git a/go/ai/gen.go b/go/ai/gen.go index 9ba35615ea..a812df9b81 100644 --- a/go/ai/gen.go +++ b/go/ai/gen.go @@ -456,7 +456,9 @@ type ReflectionEndInputStreamParams struct { RequestID string `json:"requestId,omitempty"` } -type ReflectionListActionsResponse map[string]*ActionMetadata +type ReflectionListActionsResponse struct { + Actions map[string]*ActionMetadata `json:"actions,omitempty"` +} type ReflectionListValuesParams struct { Type string `json:"type,omitempty"` diff --git a/js/core/src/reflection-types.ts b/js/core/src/reflection-types.ts index fea0fd5075..eb764378d4 100644 --- a/js/core/src/reflection-types.ts +++ b/js/core/src/reflection-types.ts @@ -102,10 +102,9 @@ export const ReflectionEndInputStreamParamsSchema = z.object({ /** * ReflectionListActionsResponseSchema is the result for the 'listActions' method. */ -export const ReflectionListActionsResponseSchema = z.record( - z.string(), - ActionMetadataSchema -); +export const ReflectionListActionsResponseSchema = z.object({ + actions: z.record(z.string(), ActionMetadataSchema), +}); /** * ReflectionListValuesResponseSchema is the result for the 'listValues' method. diff --git a/js/core/src/reflection-v2.ts b/js/core/src/reflection-v2.ts index 0155684287..66770774de 100644 --- a/js/core/src/reflection-v2.ts +++ b/js/core/src/reflection-v2.ts @@ -213,9 +213,7 @@ export class ReflectionServerV2 { private async handleListActions(request: JsonRpcRequest) { if (!request.id) return; // Should be a request const actions = await this.registry.listResolvableActions(); - const convertedActions: z.infer< - typeof ReflectionListActionsResponseSchema - > = {}; + const convertedActions: Record = {}; Object.keys(actions).forEach((key) => { const action = actions[key]; @@ -241,7 +239,7 @@ export class ReflectionServerV2 { this.sendResponse( request.id, - ReflectionListActionsResponseSchema.parse(convertedActions) + ReflectionListActionsResponseSchema.parse({ actions: convertedActions }) ); } diff --git a/py/packages/genkit/src/genkit/_core/_typing.py b/py/packages/genkit/src/genkit/_core/_typing.py index 7eb9c13cd2..9359522d86 100644 --- a/py/packages/genkit/src/genkit/_core/_typing.py +++ b/py/packages/genkit/src/genkit/_core/_typing.py @@ -580,6 +580,7 @@ class ReflectionListActionsResponse(GenkitModel): """Model for reflectionlistactionsresponse data.""" model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + actions: Actions = Field(...) class ReflectionListValuesParams(GenkitModel): @@ -840,6 +841,12 @@ class Resource(GenkitModel): uri: str = Field(...) +class Actions(GenkitModel): + """Model for actions data.""" + + model_config: ClassVar[ConfigDict] = ConfigDict(alias_generator=to_camel, extra='forbid', populate_by_name=True) + + class TelemetryLabels(GenkitModel): """Model for telemetrylabels data.""" From c41e59fe91fb735a83d97e2577c154e8d86246b4 Mon Sep 17 00:00:00 2001 From: Alex Pascal Date: Fri, 17 Apr 2026 10:35:07 -0700 Subject: [PATCH 10/10] feat(go): implement reflection API v2 over WebSocket --- go/ai/gen.go | 76 +-- go/core/schemas.config | 131 +++++ go/genkit/gen.go | 122 +++++ go/genkit/genkit.go | 22 +- go/genkit/reflection.go | 20 +- go/genkit/reflection_v2.go | 415 +++++++++++++++ go/genkit/reflection_v2_test.go | 496 ++++++++++++++++++ go/go.mod | 1 + go/go.sum | 2 + .../cmd/jsonschemagen/jsonschemagen.go | 11 +- 10 files changed, 1207 insertions(+), 89 deletions(-) create mode 100644 go/genkit/gen.go create mode 100644 go/genkit/reflection_v2.go create mode 100644 go/genkit/reflection_v2_test.go diff --git a/go/ai/gen.go b/go/ai/gen.go index a812df9b81..a6cc34354b 100644 --- a/go/ai/gen.go +++ b/go/ai/gen.go @@ -28,9 +28,9 @@ type ActionMetadata struct { Metadata map[string]any `json:"metadata,omitempty"` Name string `json:"name,omitempty"` // A JSON Schema Draft 7 (http://json-schema.org/draft-07/schema) object. - OutputJsonSchema *ActionMetadata `json:"outputJsonSchema,omitempty"` - OutputSchema any `json:"outputSchema,omitempty"` - StreamSchema any `json:"streamSchema,omitempty"` + OutputJsonSchema any `json:"outputJsonSchema,omitempty"` + OutputSchema any `json:"outputSchema,omitempty"` + StreamSchema any `json:"streamSchema,omitempty"` } type customPart struct { @@ -151,7 +151,7 @@ type GenerateActionOutputConfig struct { // GenerationCommonConfig holds configuration parameters for model generation requests. type GenerationCommonConfig struct { // API Key to use for the model call, overrides API key provided in plugin config. - ApiKey string `json:"apiKey,omitempty"` + APIKey string `json:"apiKey,omitempty"` // MaxOutputTokens limits the maximum number of tokens generated in the response. MaxOutputTokens int `json:"maxOutputTokens,omitempty"` // StopSequences specifies sequences that will cause generation to stop when encountered. @@ -440,74 +440,6 @@ type reasoningPart struct { Reasoning string `json:"reasoning,omitempty"` } -type ReflectionCancelActionParams struct { - TraceID string `json:"traceId,omitempty"` -} - -type ReflectionCancelActionResponse struct { - Message string `json:"message,omitempty"` -} - -type ReflectionConfigureParams struct { - TelemetryServerUrl string `json:"telemetryServerUrl,omitempty"` -} - -type ReflectionEndInputStreamParams struct { - RequestID string `json:"requestId,omitempty"` -} - -type ReflectionListActionsResponse struct { - Actions map[string]*ActionMetadata `json:"actions,omitempty"` -} - -type ReflectionListValuesParams struct { - Type string `json:"type,omitempty"` -} - -type ReflectionListValuesResponse map[string]any - -type ReflectionRegisterParams struct { - GenkitVersion string `json:"genkitVersion,omitempty"` - Id string `json:"id,omitempty"` - Name string `json:"name,omitempty"` - Pid float64 `json:"pid,omitempty"` - ReflectionApiSpecVersion float64 `json:"reflectionApiSpecVersion,omitempty"` -} - -type ReflectionRunActionParams struct { - // Additional runtime context data (ex. auth context data). - Context any `json:"context,omitempty"` - // An input with the type that this action expects. - Input any `json:"input,omitempty"` - // Action key that consists of the action type and ID. - Key string `json:"key,omitempty"` - // ID of the Genkit runtime to run the action on. Typically $pid-$port. - RuntimeID string `json:"runtimeId,omitempty"` - Stream bool `json:"stream,omitempty"` - StreamInput bool `json:"streamInput,omitempty"` - // Labels to be applied to telemetry data. - TelemetryLabels map[string]string `json:"telemetryLabels,omitempty"` -} - -type ReflectionRunActionStateParams struct { - RequestID string `json:"requestId,omitempty"` - State *ReflectionRunActionStateParamsState `json:"state,omitempty"` -} - -type ReflectionRunActionStateParamsState struct { - TraceID string `json:"traceId,omitempty"` -} - -type ReflectionSendInputStreamChunkParams struct { - Chunk any `json:"chunk,omitempty"` - RequestID string `json:"requestId,omitempty"` -} - -type ReflectionStreamChunkParams struct { - Chunk any `json:"chunk,omitempty"` - RequestID string `json:"requestId,omitempty"` -} - // RerankerRequest represents a request to rerank documents based on relevance. type RerankerRequest struct { // Documents is the array of documents to rerank. diff --git a/go/core/schemas.config b/go/core/schemas.config index 1beb6f139e..c2e02d2440 100644 --- a/go/core/schemas.config +++ b/go/core/schemas.config @@ -1111,3 +1111,134 @@ Embedding.embedding type []float32 GenkitError omit GenkitErrorData omit GenkitErrorDataGenkitErrorDetails omit + +# Reflection V2 types (generated into genkit package) +genkit import encoding/json + +ReflectionRegisterParams pkg genkit +ReflectionRegisterParams doc +ReflectionRegisterParams is the payload for the "register" notification +sent by the runtime to the CLI manager on connection. +. +ReflectionRegisterParams.id name ID +ReflectionRegisterParams.id doc +Unique runtime identifier. +. +ReflectionRegisterParams.pid type int +ReflectionRegisterParams.pid name PID +ReflectionRegisterParams.pid doc +Process ID of the runtime. +. +ReflectionRegisterParams.name doc +Application name (optional, defaults to the runtime ID). +. +ReflectionRegisterParams.genkitVersion doc +Genkit library version, e.g. "go/1.4.0". +. +ReflectionRegisterParams.reflectionApiSpecVersion type int +ReflectionRegisterParams.reflectionApiSpecVersion doc +Protocol version for payload compatibility checking. +. + +ReflectionRunActionParams pkg genkit +ReflectionRunActionParams doc +ReflectionRunActionParams is the payload for the "runAction" request +sent by the CLI manager to execute an action on the runtime. +. +ReflectionRunActionParams.input type json.RawMessage +ReflectionRunActionParams.context type json.RawMessage +ReflectionRunActionParams.telemetryLabels type json.RawMessage +ReflectionRunActionParams.stream doc +Whether to stream results via streamChunk notifications. +. +ReflectionRunActionParams.streamInput doc +Whether to stream input to the action (for bidirectional actions). +. + +ReflectionRunActionStateParams pkg genkit +ReflectionRunActionStateParams doc +ReflectionRunActionStateParams is the payload for the "runActionState" +notification sent by the runtime to provide early status updates +(e.g. trace ID) before the action result is ready. +. +ReflectionRunActionStateParams.requestId name RequestID +ReflectionRunActionStateParams.requestId doc +ID of the JSON-RPC request this notification relates to. +. +ReflectionRunActionStateParamsState pkg genkit +ReflectionRunActionStateParamsState doc +ReflectionRunActionStateParamsState holds the state data in a +runActionState notification. +. +ReflectionRunActionStateParamsState.traceId name TraceID +ReflectionRunActionStateParamsState.traceId doc +Trace ID for the action execution. +. + +ReflectionStreamChunkParams pkg genkit +ReflectionStreamChunkParams doc +ReflectionStreamChunkParams is the payload for the "streamChunk" +notification sent by the runtime during a streaming runAction request. +. +ReflectionStreamChunkParams.requestId name RequestID +ReflectionStreamChunkParams.requestId doc +ID of the JSON-RPC request this chunk belongs to. +. +ReflectionStreamChunkParams.chunk doc +The streamed data chunk. +. + +ReflectionConfigureParams pkg genkit +ReflectionConfigureParams doc +ReflectionConfigureParams is the payload for the "configure" notification +sent by the CLI manager to push configuration to the runtime. +. +ReflectionConfigureParams.telemetryServerUrl name TelemetryServerURL +ReflectionConfigureParams.telemetryServerUrl doc +URL of the telemetry server to send traces to (optional). +. + +ReflectionCancelActionParams pkg genkit +ReflectionCancelActionParams doc +ReflectionCancelActionParams is the payload for the "cancelAction" request +sent by the CLI manager to cancel a running action. +. +ReflectionCancelActionParams.traceId name TraceID +ReflectionCancelActionParams.traceId doc +Trace ID of the action to cancel. +. + +ReflectionCancelActionResponse pkg genkit +ReflectionCancelActionResponse doc +ReflectionCancelActionResponse is the result of a successful "cancelAction" request. +. + +ReflectionListValuesParams pkg genkit +ReflectionListValuesParams doc +ReflectionListValuesParams is the payload for the "listValues" request. +. +ReflectionListValuesParams.type doc +The type of values to list (e.g. "model", "prompt", "schema"). +. + +ReflectionListValuesResponse pkg genkit +ReflectionListValuesResponse doc +ReflectionListValuesResponse is the result of a "listValues" request, +mapping value names to their definitions. +. + +ReflectionListActionsResponse omit + +ReflectionSendInputStreamChunkParams pkg genkit +ReflectionSendInputStreamChunkParams doc +ReflectionSendInputStreamChunkParams is the payload for the +"sendInputStreamChunk" notification (bidirectional streaming, not yet implemented). +. +ReflectionSendInputStreamChunkParams.requestId name RequestID + +ReflectionEndInputStreamParams pkg genkit +ReflectionEndInputStreamParams doc +ReflectionEndInputStreamParams is the payload for the "endInputStream" +notification (bidirectional streaming, not yet implemented). +. +ReflectionEndInputStreamParams.requestId name RequestID diff --git a/go/genkit/gen.go b/go/genkit/gen.go new file mode 100644 index 0000000000..5e9354ada3 --- /dev/null +++ b/go/genkit/gen.go @@ -0,0 +1,122 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +// This file was generated by jsonschemagen. DO NOT EDIT. + +package genkit + +import "encoding/json" + +// ReflectionCancelActionParams is the payload for the "cancelAction" request +// sent by the CLI manager to cancel a running action. +type ReflectionCancelActionParams struct { + // Trace ID of the action to cancel. + TraceID string `json:"traceId,omitempty"` +} + +// ReflectionCancelActionResponse is the result of a successful "cancelAction" request. +type ReflectionCancelActionResponse struct { + Message string `json:"message,omitempty"` +} + +// ReflectionConfigureParams is the payload for the "configure" notification +// sent by the CLI manager to push configuration to the runtime. +type ReflectionConfigureParams struct { + // URL of the telemetry server to send traces to (optional). + TelemetryServerURL string `json:"telemetryServerUrl,omitempty"` +} + +// ReflectionEndInputStreamParams is the payload for the "endInputStream" +// notification (bidirectional streaming, not yet implemented). +type ReflectionEndInputStreamParams struct { + RequestID string `json:"requestId,omitempty"` +} + +// ReflectionListValuesParams is the payload for the "listValues" request. +type ReflectionListValuesParams struct { + // The type of values to list (e.g. "model", "prompt", "schema"). + Type string `json:"type,omitempty"` +} + +// ReflectionListValuesResponse is the result of a "listValues" request, +// mapping value names to their definitions. +type ReflectionListValuesResponse map[string]any + +// ReflectionRegisterParams is the payload for the "register" notification +// sent by the runtime to the CLI manager on connection. +type ReflectionRegisterParams struct { + // Genkit library version, e.g. "go/1.4.0". + GenkitVersion string `json:"genkitVersion,omitempty"` + // Unique runtime identifier. + ID string `json:"id,omitempty"` + // Application name (optional, defaults to the runtime ID). + Name string `json:"name,omitempty"` + // Process ID of the runtime. + PID int `json:"pid,omitempty"` + // Protocol version for payload compatibility checking. + ReflectionApiSpecVersion int `json:"reflectionApiSpecVersion,omitempty"` +} + +// ReflectionRunActionParams is the payload for the "runAction" request +// sent by the CLI manager to execute an action on the runtime. +type ReflectionRunActionParams struct { + // Additional runtime context data (ex. auth context data). + Context json.RawMessage `json:"context,omitempty"` + // An input with the type that this action expects. + Input json.RawMessage `json:"input,omitempty"` + // Action key that consists of the action type and ID. + Key string `json:"key,omitempty"` + // ID of the Genkit runtime to run the action on. Typically $pid-$port. + RuntimeID string `json:"runtimeId,omitempty"` + // Whether to stream results via streamChunk notifications. + Stream bool `json:"stream,omitempty"` + // Whether to stream input to the action (for bidirectional actions). + StreamInput bool `json:"streamInput,omitempty"` + // Labels to be applied to telemetry data. + TelemetryLabels json.RawMessage `json:"telemetryLabels,omitempty"` +} + +// ReflectionRunActionStateParams is the payload for the "runActionState" +// notification sent by the runtime to provide early status updates +// (e.g. trace ID) before the action result is ready. +type ReflectionRunActionStateParams struct { + // ID of the JSON-RPC request this notification relates to. + RequestID string `json:"requestId,omitempty"` + State *ReflectionRunActionStateParamsState `json:"state,omitempty"` +} + +// ReflectionRunActionStateParamsState holds the state data in a +// runActionState notification. +type ReflectionRunActionStateParamsState struct { + // Trace ID for the action execution. + TraceID string `json:"traceId,omitempty"` +} + +// ReflectionSendInputStreamChunkParams is the payload for the +// "sendInputStreamChunk" notification (bidirectional streaming, not yet implemented). +type ReflectionSendInputStreamChunkParams struct { + Chunk any `json:"chunk,omitempty"` + RequestID string `json:"requestId,omitempty"` +} + +// ReflectionStreamChunkParams is the payload for the "streamChunk" +// notification sent by the runtime during a streaming runAction request. +type ReflectionStreamChunkParams struct { + // The streamed data chunk. + Chunk any `json:"chunk,omitempty"` + // ID of the JSON-RPC request this chunk belongs to. + RequestID string `json:"requestId,omitempty"` +} diff --git a/go/genkit/genkit.go b/go/genkit/genkit.go index 476908a287..ad24e787b2 100644 --- a/go/genkit/genkit.go +++ b/go/genkit/genkit.go @@ -242,14 +242,20 @@ func Init(ctx context.Context, opts ...GenkitOption) *Genkit { errCh := make(chan error, 1) serverStartCh := make(chan struct{}) - go func() { - if s := startReflectionServer(ctx, g, errCh, serverStartCh); s == nil { - return - } - if err := <-errCh; err != nil { - slog.Error("reflection server error", "err", err) - } - }() + if v2URL := os.Getenv("GENKIT_REFLECTION_V2_SERVER"); v2URL != "" { + // V2: connect to the CLI's WebSocket server. + go startReflectionServerV2(ctx, g, reflectionServerV2Options{URL: v2URL}, errCh, serverStartCh) + } else { + // V1: start an HTTP reflection server. + go func() { + if s := startReflectionServer(ctx, g, errCh, serverStartCh); s == nil { + return + } + if err := <-errCh; err != nil { + slog.Error("reflection server error", "err", err) + } + }() + } select { case err := <-errCh: diff --git a/go/genkit/reflection.go b/go/genkit/reflection.go index 1bd675f75a..fb57755560 100644 --- a/go/genkit/reflection.go +++ b/go/genkit/reflection.go @@ -415,7 +415,7 @@ func handleRunAction(g *Genkit, activeActions *activeActionsMap) func(w http.Res } } - var contextMap core.ActionContext = nil + contextMap := core.ActionContext{} if body.Context != nil { json.Unmarshal(body.Context, &contextMap) } @@ -557,6 +557,15 @@ func handleCancelAction(activeActions *activeActionsMap) func(w http.ResponseWri } } +// configureTelemetry sets up the telemetry client if not already configured via env var. +// Shared between V1 and V2 reflection servers. +func configureTelemetry(url string) { + if os.Getenv("GENKIT_TELEMETRY_SERVER") == "" && url != "" { + tracing.WriteTelemetryImmediate(tracing.NewHTTPTelemetryClient(url)) + slog.Debug("connected to telemetry server", "url", url) + } +} + // handleNotify configures the telemetry server URL from the request. func handleNotify() func(w http.ResponseWriter, r *http.Request) error { return func(w http.ResponseWriter, r *http.Request) error { @@ -570,10 +579,7 @@ func handleNotify() func(w http.ResponseWriter, r *http.Request) error { return core.NewError(core.INVALID_ARGUMENT, err.Error()) } - if os.Getenv("GENKIT_TELEMETRY_SERVER") == "" && body.TelemetryServerURL != "" { - tracing.WriteTelemetryImmediate(tracing.NewHTTPTelemetryClient(body.TelemetryServerURL)) - slog.Debug("connected to telemetry server", "url", body.TelemetryServerURL) - } + configureTelemetry(body.TelemetryServerURL) if body.ReflectionApiSpecVersion != internal.GENKIT_REFLECTION_API_SPEC_VERSION { slog.Error("Genkit CLI version is not compatible with runtime library. Please use `genkit-cli` version compatible with runtime library version.") @@ -662,9 +668,7 @@ func runAction(ctx context.Context, g *Genkit, key string, input json.RawMessage if action == nil { return nil, core.NewError(core.NOT_FOUND, "action %q not found", key) } - if runtimeContext != nil { - ctx = core.WithActionContext(ctx, runtimeContext) - } + ctx = core.WithActionContext(ctx, runtimeContext) // Parse telemetry attributes if provided var telemetryAttributes map[string]string diff --git a/go/genkit/reflection_v2.go b/go/genkit/reflection_v2.go new file mode 100644 index 0000000000..20b7718521 --- /dev/null +++ b/go/genkit/reflection_v2.go @@ -0,0 +1,415 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package genkit + +import ( + "context" + "encoding/json" + "errors" + "fmt" + "log/slog" + "os" + "strconv" + "sync" + "time" + + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" + "github.com/firebase/genkit/go/core" + "github.com/firebase/genkit/go/core/tracing" + "github.com/firebase/genkit/go/internal" +) + +// JSON-RPC 2.0 error codes. +const ( + jsonRPCMethodNotFound = -32601 + jsonRPCInvalidParams = -32602 + jsonRPCServerError = -32000 +) + +// jsonRPCRequest represents an incoming JSON-RPC 2.0 request from the manager. +type jsonRPCRequest struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params json.RawMessage `json:"params,omitempty"` + ID string `json:"id,omitempty"` +} + +// jsonRPCResponse is an outgoing JSON-RPC 2.0 response. +type jsonRPCResponse struct { + JSONRPC string `json:"jsonrpc"` + Result any `json:"result,omitempty"` + Error *jsonRPCError `json:"error,omitempty"` + ID string `json:"id"` +} + +// jsonRPCNotification is an outgoing JSON-RPC 2.0 notification (no ID). +type jsonRPCNotification struct { + JSONRPC string `json:"jsonrpc"` + Method string `json:"method"` + Params any `json:"params,omitempty"` +} + +// jsonRPCError is the error object in a JSON-RPC 2.0 error response. +type jsonRPCError struct { + Code int `json:"code"` + Message string `json:"message"` + Data any `json:"data,omitempty"` +} + +// reflectionRunActionResponse is the success payload for a runAction request. +// Not in the generated schema because only the runtime produces it. +type reflectionRunActionResponse struct { + Result json.RawMessage `json:"result"` + Telemetry telemetry `json:"telemetry"` +} + +// reflectionServerV2 is a WebSocket client that connects to the CLI's +// reflection manager and handles JSON-RPC 2.0 requests. +type reflectionServerV2 struct { + g *Genkit + activeActions *activeActionsMap + + ctx context.Context + writeMu sync.Mutex // Serializes writes to conn. + conn *websocket.Conn +} + +// reflectionServerV2Options configures the V2 reflection client. +type reflectionServerV2Options struct { + Name string // App name (optional, defaults to the runtime ID). + URL string // WebSocket URL of the CLI manager. +} + +// startReflectionServerV2 connects to the CLI's WebSocket server, registers +// this runtime, and spawns a goroutine to handle incoming reflection requests. +// Returns once registration has been sent (or an error has been reported). +func startReflectionServerV2(ctx context.Context, g *Genkit, opts reflectionServerV2Options, errCh chan<- error, serverStartCh chan<- struct{}) *reflectionServerV2 { + if g == nil { + errCh <- fmt.Errorf("nil Genkit provided") + return nil + } + + conn, _, err := websocket.Dial(ctx, opts.URL, nil) + if err != nil { + errCh <- fmt.Errorf("failed to connect to reflection V2 server at %s: %w", opts.URL, err) + return nil + } + + s := &reflectionServerV2{ + g: g, + activeActions: newActiveActionsMap(), + ctx: ctx, + conn: conn, + } + + slog.Debug("reflection V2: connected", "url", opts.URL) + + if err := s.register(opts.Name); err != nil { + conn.Close(websocket.StatusInternalError, "registration failed") + errCh <- fmt.Errorf("failed to register with reflection V2 server: %w", err) + return nil + } + + close(serverStartCh) + + go func() { + defer conn.Close(websocket.StatusNormalClosure, "shutting down") + s.readLoop(ctx) + }() + + return s +} + +// register sends a registration notification to the manager. +func (s *reflectionServerV2) register(name string) error { + runtimeID := os.Getenv("GENKIT_RUNTIME_ID") + if runtimeID == "" { + runtimeID = strconv.Itoa(os.Getpid()) + } + if name == "" { + name = runtimeID + } + + return s.sendNotification("register", &ReflectionRegisterParams{ + ID: runtimeID, + PID: os.Getpid(), + Name: name, + GenkitVersion: "go/" + internal.Version, + ReflectionApiSpecVersion: internal.GENKIT_REFLECTION_API_SPEC_VERSION, + }) +} + +// readLoop reads and dispatches JSON-RPC messages until the context is +// cancelled or the connection is closed. +func (s *reflectionServerV2) readLoop(ctx context.Context) { + for { + var req jsonRPCRequest + if err := wsjson.Read(ctx, s.conn, &req); err != nil { + if ctx.Err() == nil && websocket.CloseStatus(err) == -1 { + slog.Error("reflection V2: failed to read message", "err", err) + } + return + } + + if req.JSONRPC != "2.0" || req.Method == "" { + continue + } + + go s.handleRequest(ctx, &req) + } +} + +// handleRequest dispatches a JSON-RPC request to the appropriate handler. +// Each handler is responsible for sending its own response (or none, for +// notifications). Unknown methods with a request ID return "method not found"; +// unknown notifications are logged and ignored. +func (s *reflectionServerV2) handleRequest(ctx context.Context, req *jsonRPCRequest) { + switch req.Method { + case "listActions": + s.handleListActions(ctx, req) + case "listValues": + s.handleListValues(req) + case "runAction": + s.handleRunAction(ctx, req) + case "cancelAction": + s.handleCancelAction(req) + case "configure": + s.handleConfigure(req) + case "sendInputStreamChunk", "endInputStream": + // Bidirectional input streaming is not yet implemented. + slog.Debug("reflection V2: method not implemented", "method", req.Method) + default: + if req.ID != "" { + s.sendErrorResponse(req.ID, jsonRPCMethodNotFound, "method not found: "+req.Method, nil) + } else { + slog.Debug("reflection V2: unknown notification", "method", req.Method) + } + } +} + +// handleListActions responds with all registered and resolvable actions. +func (s *reflectionServerV2) handleListActions(ctx context.Context, req *jsonRPCRequest) { + if req.ID == "" { + return + } + ads := listResolvableActions(ctx, s.g) + actionsMap := make(map[string]any, len(ads)) + for _, d := range ads { + actionsMap[d.Key] = d + } + s.sendResponse(req.ID, struct { + Actions map[string]any `json:"actions"` + }{Actions: actionsMap}) +} + +// handleListValues responds with registered values. The "type" param is +// parsed for protocol compliance but the Go registry does not currently +// support filtering by type, so all values are returned. +func (s *reflectionServerV2) handleListValues(req *jsonRPCRequest) { + if req.ID == "" { + return + } + + var params ReflectionListValuesParams + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + s.sendErrorResponse(req.ID, jsonRPCInvalidParams, "invalid params: "+err.Error(), nil) + return + } + + s.sendResponse(req.ID, ReflectionListValuesResponse(s.g.reg.ListValues())) +} + +// handleRunAction executes an action and sends the result (with optional streaming). +func (s *reflectionServerV2) handleRunAction(ctx context.Context, req *jsonRPCRequest) { + if req.ID == "" { + return + } + + var params ReflectionRunActionParams + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + s.sendErrorResponse(req.ID, jsonRPCInvalidParams, "invalid params: "+err.Error(), nil) + return + } + + slog.Debug("reflection V2: running action", "key", params.Key, "stream", params.Stream) + + actionCtx, cancel := context.WithCancel(ctx) + defer cancel() + + // Capture the trace ID asynchronously so we can both track the action for + // cancellation and include the trace ID in any error we send back. + var traceIDMu sync.Mutex + var traceID string + + telemetryCb := func(tid, _ string) { + traceIDMu.Lock() + traceID = tid + traceIDMu.Unlock() + + s.activeActions.Set(tid, &activeAction{ + cancel: cancel, + startTime: time.Now(), + traceID: tid, + }) + + s.sendNotification("runActionState", &ReflectionRunActionStateParams{ + RequestID: req.ID, + State: &ReflectionRunActionStateParamsState{TraceID: tid}, + }) + } + + var streamCb streamingCallback[json.RawMessage] + if params.Stream { + streamCb = func(_ context.Context, chunk json.RawMessage) error { + return s.sendNotification("streamChunk", &ReflectionStreamChunkParams{ + RequestID: req.ID, + Chunk: chunk, + }) + } + } + + contextMap := core.ActionContext{} + if params.Context != nil { + if err := json.Unmarshal(params.Context, &contextMap); err != nil { + s.sendErrorResponse(req.ID, jsonRPCInvalidParams, "invalid context: "+err.Error(), nil) + return + } + } + + actionCtx = tracing.WithTelemetryCallback(actionCtx, telemetryCb) + resp, err := runAction(actionCtx, s.g, params.Key, params.Input, params.TelemetryLabels, streamCb, contextMap) + + traceIDMu.Lock() + capturedTraceID := traceID + traceIDMu.Unlock() + if capturedTraceID != "" { + s.activeActions.Delete(capturedTraceID) + } + + if err != nil { + s.sendRunActionError(req.ID, err, capturedTraceID) + return + } + + s.sendResponse(req.ID, &reflectionRunActionResponse{ + Result: resp.Result, + Telemetry: telemetry{TraceID: resp.Telemetry.TraceID}, + }) +} + +// sendRunActionError maps a runAction error to a JSON-RPC error response +// with a Status-shaped data field matching the JS implementation. +func (s *reflectionServerV2) sendRunActionError(id string, err error, traceID string) { + code := core.INTERNAL + msg := err.Error() + if errors.Is(err, context.Canceled) { + code = core.CANCELLED + msg = "Action was cancelled" + } + + details := map[string]any{} + if traceID != "" { + details["traceId"] = traceID + } + var ge *core.GenkitError + if errors.As(err, &ge) && ge.Details != nil { + if stack, ok := ge.Details["stack"].(string); ok { + details["stack"] = stack + } + } + + data := map[string]any{ + "code": core.StatusNameToCode[code], + "message": msg, + } + if len(details) > 0 { + data["details"] = details + } + + s.sendErrorResponse(id, jsonRPCServerError, msg, data) +} + +// handleConfigure processes a configuration notification from the manager. +func (s *reflectionServerV2) handleConfigure(req *jsonRPCRequest) { + var params ReflectionConfigureParams + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + slog.Error("reflection V2: invalid configure params", "err", err) + return + } + configureTelemetry(params.TelemetryServerURL) +} + +// handleCancelAction cancels an in-flight action by trace ID. +func (s *reflectionServerV2) handleCancelAction(req *jsonRPCRequest) { + if req.ID == "" { + return + } + + var params ReflectionCancelActionParams + if err := json.Unmarshal(req.Params, ¶ms); err != nil { + s.sendErrorResponse(req.ID, jsonRPCInvalidParams, "invalid params: "+err.Error(), nil) + return + } + if params.TraceID == "" { + s.sendErrorResponse(req.ID, jsonRPCInvalidParams, "traceId is required", nil) + return + } + + action, ok := s.activeActions.Get(params.TraceID) + if !ok { + s.sendErrorResponse(req.ID, jsonRPCServerError, "Action not found or already completed", nil) + return + } + + action.cancel() + s.activeActions.Delete(params.TraceID) + s.sendResponse(req.ID, &ReflectionCancelActionResponse{Message: "Action cancelled"}) +} + +// sendResponse sends a JSON-RPC success response. Send errors are logged but +// not returned: the read loop will pick up on a broken connection on its +// next read. +func (s *reflectionServerV2) sendResponse(id string, result any) { + if err := s.send(&jsonRPCResponse{JSONRPC: "2.0", Result: result, ID: id}); err != nil { + slog.Error("reflection V2: failed to send response", "err", err, "id", id) + } +} + +// sendErrorResponse sends a JSON-RPC error response. +func (s *reflectionServerV2) sendErrorResponse(id string, code int, message string, data any) { + if err := s.send(&jsonRPCResponse{ + JSONRPC: "2.0", + Error: &jsonRPCError{Code: code, Message: message, Data: data}, + ID: id, + }); err != nil { + slog.Error("reflection V2: failed to send error response", "err", err, "id", id) + } +} + +// sendNotification sends a JSON-RPC notification (no ID, no response expected). +func (s *reflectionServerV2) sendNotification(method string, params any) error { + return s.send(&jsonRPCNotification{JSONRPC: "2.0", Method: method, Params: params}) +} + +// send writes a JSON message to the WebSocket connection. +// It is safe for concurrent use. +func (s *reflectionServerV2) send(msg any) error { + s.writeMu.Lock() + defer s.writeMu.Unlock() + return wsjson.Write(s.ctx, s.conn, msg) +} diff --git a/go/genkit/reflection_v2_test.go b/go/genkit/reflection_v2_test.go new file mode 100644 index 0000000000..19a2b7c49c --- /dev/null +++ b/go/genkit/reflection_v2_test.go @@ -0,0 +1,496 @@ +// Copyright 2025 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. +// +// SPDX-License-Identifier: Apache-2.0 + +package genkit + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/coder/websocket" + "github.com/coder/websocket/wsjson" + "github.com/firebase/genkit/go/core" + "github.com/firebase/genkit/go/core/api" + "github.com/firebase/genkit/go/core/tracing" +) + +// fakeManager is a test double for the CLI's reflection V2 manager. It accepts +// one WebSocket connection, records inbound messages, and lets tests drive the +// runtime by sending JSON-RPC requests / reading responses. +type fakeManager struct { + server *httptest.Server + url string + + mu sync.Mutex + conn *websocket.Conn + connCh chan *websocket.Conn +} + +func newFakeManager(t *testing.T) *fakeManager { + t.Helper() + m := &fakeManager{connCh: make(chan *websocket.Conn, 1)} + + m.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c, err := websocket.Accept(w, r, nil) + if err != nil { + t.Errorf("accept: %v", err) + return + } + m.mu.Lock() + m.conn = c + m.mu.Unlock() + m.connCh <- c + // Block until test closes the connection so the handler doesn't exit. + <-r.Context().Done() + })) + m.url = "ws" + strings.TrimPrefix(m.server.URL, "http") + return m +} + +func (m *fakeManager) close() { + m.mu.Lock() + if m.conn != nil { + m.conn.Close(websocket.StatusNormalClosure, "") + } + m.mu.Unlock() + m.server.Close() +} + +// waitForConnection blocks until the runtime has connected. +func (m *fakeManager) waitForConnection(t *testing.T) *websocket.Conn { + t.Helper() + select { + case c := <-m.connCh: + return c + case <-time.After(2 * time.Second): + t.Fatal("timed out waiting for runtime to connect") + return nil + } +} + +// read reads the next JSON-RPC message from the runtime. +func (m *fakeManager) read(t *testing.T, ctx context.Context, conn *websocket.Conn) map[string]any { + t.Helper() + var msg map[string]any + readCtx, cancel := context.WithTimeout(ctx, 2*time.Second) + defer cancel() + if err := wsjson.Read(readCtx, conn, &msg); err != nil { + t.Fatalf("read: %v", err) + } + return msg +} + +// write sends a JSON-RPC message to the runtime. +func (m *fakeManager) write(t *testing.T, ctx context.Context, conn *websocket.Conn, msg any) { + t.Helper() + if err := wsjson.Write(ctx, conn, msg); err != nil { + t.Fatalf("write: %v", err) + } +} + +// startRuntime starts a reflection V2 client connected to the fake manager +// and waits for registration. Returns a teardown function. +func startRuntime(t *testing.T, g *Genkit, m *fakeManager) (context.Context, func()) { + t.Helper() + tracing.WriteTelemetryImmediate(tracing.NewTestOnlyTelemetryClient()) + + ctx, cancel := context.WithCancel(context.Background()) + errCh := make(chan error, 1) + startedCh := make(chan struct{}) + + go startReflectionServerV2(ctx, g, reflectionServerV2Options{URL: m.url, Name: "test-app"}, errCh, startedCh) + + select { + case err := <-errCh: + cancel() + t.Fatalf("runtime failed to start: %v", err) + case <-startedCh: + case <-time.After(2 * time.Second): + cancel() + t.Fatal("timed out waiting for runtime startup") + } + + return ctx, cancel +} + +func TestReflectionServerV2_Register(t *testing.T) { + m := newFakeManager(t) + defer m.close() + + g := Init(context.Background()) + _, cancel := startRuntime(t, g, m) + defer cancel() + + conn := m.waitForConnection(t) + msg := m.read(t, context.Background(), conn) + + if msg["method"] != "register" { + t.Fatalf("first message method = %q, want register", msg["method"]) + } + params, ok := msg["params"].(map[string]any) + if !ok { + t.Fatalf("params is not object: %v", msg["params"]) + } + if params["name"] != "test-app" { + t.Errorf("name = %q, want test-app", params["name"]) + } + if params["id"] == "" || params["id"] == nil { + t.Error("id should be set") + } + if _, ok := params["pid"].(float64); !ok { + t.Errorf("pid should be a number, got %T", params["pid"]) + } + if !strings.HasPrefix(params["genkitVersion"].(string), "go/") { + t.Errorf("genkitVersion = %q, want prefix go/", params["genkitVersion"]) + } + if _, ok := params["reflectionApiSpecVersion"].(float64); !ok { + t.Errorf("reflectionApiSpecVersion should be a number, got %T", params["reflectionApiSpecVersion"]) + } + // register is a notification; it must not have an id field. + if _, hasID := msg["id"]; hasID { + t.Error("register notification should not have id") + } +} + +func TestReflectionServerV2_ListActions(t *testing.T) { + m := newFakeManager(t) + defer m.close() + + g := Init(context.Background()) + core.DefineAction(g.reg, "test/inc", api.ActionTypeCustom, nil, nil, inc) + core.DefineAction(g.reg, "test/dec", api.ActionTypeCustom, nil, nil, dec) + + ctx, cancel := startRuntime(t, g, m) + defer cancel() + + conn := m.waitForConnection(t) + if got := m.read(t, ctx, conn)["method"]; got != "register" { + t.Fatalf("expected register, got %v", got) + } + + m.write(t, ctx, conn, map[string]any{ + "jsonrpc": "2.0", + "method": "listActions", + "id": "1", + }) + + resp := m.read(t, ctx, conn) + if resp["id"] != "1" { + t.Fatalf("id = %v, want 1", resp["id"]) + } + result, ok := resp["result"].(map[string]any) + if !ok { + t.Fatalf("result is not object: %v", resp["result"]) + } + actions, ok := result["actions"].(map[string]any) + if !ok { + t.Fatalf("actions is not object: %v", result["actions"]) + } + for _, key := range []string{"/custom/test/inc", "/custom/test/dec"} { + if _, ok := actions[key]; !ok { + t.Errorf("action %q missing from response", key) + } + } +} + +func TestReflectionServerV2_ListValues(t *testing.T) { + m := newFakeManager(t) + defer m.close() + + g := Init(context.Background()) + g.reg.RegisterValue("prompts/my-prompt", "hello") + + ctx, cancel := startRuntime(t, g, m) + defer cancel() + + conn := m.waitForConnection(t) + m.read(t, ctx, conn) // drain register + + m.write(t, ctx, conn, map[string]any{ + "jsonrpc": "2.0", + "method": "listValues", + "params": map[string]any{"type": "prompt"}, + "id": "2", + }) + + resp := m.read(t, ctx, conn) + if resp["id"] != "2" { + t.Fatalf("id = %v, want 2", resp["id"]) + } + result, ok := resp["result"].(map[string]any) + if !ok { + t.Fatalf("result is not object: %v", resp["result"]) + } + if result["prompts/my-prompt"] != "hello" { + t.Errorf("value = %v, want hello", result["prompts/my-prompt"]) + } +} + +func TestReflectionServerV2_RunAction(t *testing.T) { + m := newFakeManager(t) + defer m.close() + + g := Init(context.Background()) + core.DefineAction(g.reg, "test/inc", api.ActionTypeCustom, nil, nil, inc) + + ctx, cancel := startRuntime(t, g, m) + defer cancel() + + conn := m.waitForConnection(t) + m.read(t, ctx, conn) // drain register + + m.write(t, ctx, conn, map[string]any{ + "jsonrpc": "2.0", + "method": "runAction", + "params": map[string]any{ + "key": "/custom/test/inc", + "input": 3, + }, + "id": "3", + }) + + // Drain any runActionState notifications, then expect the final response. + var resp map[string]any + for { + msg := m.read(t, ctx, conn) + if msg["method"] == "runActionState" { + continue + } + resp = msg + break + } + if resp["id"] != "3" { + t.Fatalf("id = %v, want 3", resp["id"]) + } + if resp["error"] != nil { + t.Fatalf("unexpected error: %v", resp["error"]) + } + result, ok := resp["result"].(map[string]any) + if !ok { + t.Fatalf("result is not object: %v", resp["result"]) + } + // action result is json.RawMessage, which unmarshals to a number here. + if got := result["result"]; got != float64(4) { + t.Errorf("result = %v, want 4", got) + } + telemetry, ok := result["telemetry"].(map[string]any) + if !ok || telemetry["traceId"] == "" { + t.Errorf("expected non-empty traceId, got %v", result["telemetry"]) + } +} + +func TestReflectionServerV2_StreamingRunAction(t *testing.T) { + m := newFakeManager(t) + defer m.close() + + g := Init(context.Background()) + streamInc := func(_ context.Context, x int, cb streamingCallback[json.RawMessage]) (int, error) { + for i := range x { + msg, _ := json.Marshal(i) + if err := cb(context.Background(), msg); err != nil { + return 0, err + } + } + return x, nil + } + core.DefineStreamingAction(g.reg, "test/streaming", api.ActionTypeCustom, nil, nil, streamInc) + + ctx, cancel := startRuntime(t, g, m) + defer cancel() + + conn := m.waitForConnection(t) + m.read(t, ctx, conn) // drain register + + m.write(t, ctx, conn, map[string]any{ + "jsonrpc": "2.0", + "method": "runAction", + "params": map[string]any{ + "key": "/custom/test/streaming", + "input": 3, + "stream": true, + }, + "id": "4", + }) + + var chunks []float64 + var final map[string]any + for { + msg := m.read(t, ctx, conn) + switch msg["method"] { + case "streamChunk": + params := msg["params"].(map[string]any) + if params["requestId"] != "4" { + t.Errorf("streamChunk requestId = %v, want 4", params["requestId"]) + } + chunks = append(chunks, params["chunk"].(float64)) + continue + case "runActionState": + continue + } + final = msg + break + } + if len(chunks) != 3 { + t.Errorf("got %d chunks, want 3", len(chunks)) + } + for i, c := range chunks { + if c != float64(i) { + t.Errorf("chunk[%d] = %v, want %d", i, c, i) + } + } + result := final["result"].(map[string]any) + if result["result"] != float64(3) { + t.Errorf("final result = %v, want 3", result["result"]) + } +} + +func TestReflectionServerV2_RunActionNotFound(t *testing.T) { + m := newFakeManager(t) + defer m.close() + + g := Init(context.Background()) + ctx, cancel := startRuntime(t, g, m) + defer cancel() + + conn := m.waitForConnection(t) + m.read(t, ctx, conn) // drain register + + m.write(t, ctx, conn, map[string]any{ + "jsonrpc": "2.0", + "method": "runAction", + "params": map[string]any{"key": "/custom/does-not-exist", "input": nil}, + "id": "5", + }) + + resp := m.read(t, ctx, conn) + errObj, ok := resp["error"].(map[string]any) + if !ok { + t.Fatalf("expected error object, got %v", resp) + } + if code := errObj["code"].(float64); code != float64(jsonRPCServerError) { + t.Errorf("code = %v, want %d", code, jsonRPCServerError) + } + data, ok := errObj["data"].(map[string]any) + if !ok { + t.Fatalf("expected error data, got %v", errObj["data"]) + } + if data["code"] == nil { + t.Error("data.code missing") + } + if data["message"] == nil { + t.Error("data.message missing") + } +} + +func TestReflectionServerV2_CancelAction(t *testing.T) { + m := newFakeManager(t) + defer m.close() + + g := Init(context.Background()) + started := make(chan struct{}) + core.DefineAction(g.reg, "test/slow", api.ActionTypeCustom, nil, nil, + func(ctx context.Context, _ any) (any, error) { + close(started) + <-ctx.Done() + return nil, ctx.Err() + }) + + ctx, cancel := startRuntime(t, g, m) + defer cancel() + + conn := m.waitForConnection(t) + m.read(t, ctx, conn) // drain register + + m.write(t, ctx, conn, map[string]any{ + "jsonrpc": "2.0", + "method": "runAction", + "params": map[string]any{"key": "/custom/test/slow", "input": nil}, + "id": "6", + }) + + // Wait for action to start, then for the runActionState notification with trace ID. + <-started + var traceID string + for traceID == "" { + msg := m.read(t, ctx, conn) + if msg["method"] == "runActionState" { + state := msg["params"].(map[string]any)["state"].(map[string]any) + traceID = state["traceId"].(string) + } + } + + m.write(t, ctx, conn, map[string]any{ + "jsonrpc": "2.0", + "method": "cancelAction", + "params": map[string]any{"traceId": traceID}, + "id": "7", + }) + + // Expect both the cancelAction response and the runAction error response. + var sawCancel, sawRunErr bool + for !sawCancel || !sawRunErr { + msg := m.read(t, ctx, conn) + switch msg["id"] { + case "7": + if result, ok := msg["result"].(map[string]any); !ok || result["message"] != "Action cancelled" { + t.Errorf("cancel response = %v", msg) + } + sawCancel = true + case "6": + errObj, ok := msg["error"].(map[string]any) + if !ok { + t.Fatalf("expected runAction error, got %v", msg) + } + if !strings.Contains(errObj["message"].(string), "cancel") { + t.Errorf("error message = %q, want contains 'cancel'", errObj["message"]) + } + sawRunErr = true + } + } +} + +func TestReflectionServerV2_MethodNotFound(t *testing.T) { + m := newFakeManager(t) + defer m.close() + + g := Init(context.Background()) + ctx, cancel := startRuntime(t, g, m) + defer cancel() + + conn := m.waitForConnection(t) + m.read(t, ctx, conn) // drain register + + m.write(t, ctx, conn, map[string]any{ + "jsonrpc": "2.0", + "method": "unknownMethod", + "id": "8", + }) + + resp := m.read(t, ctx, conn) + errObj, ok := resp["error"].(map[string]any) + if !ok { + t.Fatalf("expected error, got %v", resp) + } + if code := errObj["code"].(float64); code != float64(jsonRPCMethodNotFound) { + t.Errorf("code = %v, want %d", code, jsonRPCMethodNotFound) + } +} diff --git a/go/go.mod b/go/go.mod index 86d45da02f..b6a1f4fd3b 100644 --- a/go/go.mod +++ b/go/go.mod @@ -19,6 +19,7 @@ require ( github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/trace v1.27.0 github.com/anthropics/anthropic-sdk-go v1.23.0 github.com/blues/jsonata-go v1.5.4 + github.com/coder/websocket v1.8.14 github.com/goccy/go-yaml v1.17.1 github.com/google/dotprompt/go v0.0.0-20251014011017-8d056e027254 github.com/google/go-cmp v0.7.0 diff --git a/go/go.sum b/go/go.sum index 43edb63227..c17d1ffbad 100644 --- a/go/go.sum +++ b/go/go.sum @@ -75,6 +75,8 @@ github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDk github.com/cncf/udpa/go v0.0.0-20191209042840-269d4d468f6f/go.mod h1:M8M6+tZqaGXZJjfX53e64911xZQV5JYwmTeXPW+k8Sc= github.com/cncf/xds/go v0.0.0-20250326154945-ae57f3c0d45f h1:C5bqEmzEPLsHm9Mv73lSE9e9bKV23aB1vxOsmZrkl3k= github.com/cncf/xds/go v0.0.0-20250326154945-ae57f3c0d45f/go.mod h1:W+zGtBO5Y1IgJhy4+A9GOqVhqLpfZi+vwmdNXUehLA8= +github.com/coder/websocket v1.8.14 h1:9L0p0iKiNOibykf283eHkKUHHrpG7f65OE3BhhO7v9g= +github.com/coder/websocket v1.8.14/go.mod h1:NX3SzP+inril6yawo5CQXx8+fk145lPDC6pumgx0mVg= github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ33E= github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= diff --git a/go/internal/cmd/jsonschemagen/jsonschemagen.go b/go/internal/cmd/jsonschemagen/jsonschemagen.go index 03e11e00de..3759d26ebd 100644 --- a/go/internal/cmd/jsonschemagen/jsonschemagen.go +++ b/go/internal/cmd/jsonschemagen/jsonschemagen.go @@ -442,7 +442,11 @@ func (g *generator) generateStruct(name string, s *Schema, tcfg *itemConfig) err if skipOmitEmpty(goName, field) { jsonTag = fmt.Sprintf(`json:"%s"`, field) } - g.pr(fmt.Sprintf(" %s %s `%s`\n", adjustIdentifier(field), typeExpr, jsonTag)) + fieldName := fcfg.name + if fieldName == "" { + fieldName = adjustIdentifier(field) + } + g.pr(fmt.Sprintf(" %s %s `%s`\n", fieldName, typeExpr, jsonTag)) } for _, f := range tcfg.fields { g.pr(fmt.Sprintf(" %s %s\n", f.name, f.typeExpr)) @@ -510,6 +514,11 @@ func (g *generator) typeExpr(s *Schema) (string, error) { if err != nil { return "", err } + // Nested refs (e.g. "#/$defs/Foo/properties/bar") don't correspond to a + // generated Go type; inline the resolved sub-schema's type instead. + if strings.Count(s.Ref, "/") > 2 { + return g.typeExpr(s2) + } ic := g.cfg.configFor(name) if s2 == nil { // If there is no schema, perhaps there is a config value.