From 5a0279892a307b45871b4c4ca210b85e22bdee6e Mon Sep 17 00:00:00 2001 From: Nadeesha Cabral Date: Sun, 15 Dec 2024 09:13:41 +1100 Subject: [PATCH] update --- control-plane/src/modules/machines.test.ts | 17 +++++++------- control-plane/src/modules/management.ts | 23 ------------------- .../src/modules/observability/hyperdx.ts | 2 +- control-plane/src/modules/prompt-templates.ts | 2 +- control-plane/src/modules/router.ts | 21 ++++++++++++++--- .../src/modules/service-definitions.ts | 21 +++++++++++------ .../src/modules/workflows/agent/run.ts | 17 ++++++++------ 7 files changed, 52 insertions(+), 51 deletions(-) diff --git a/control-plane/src/modules/machines.test.ts b/control-plane/src/modules/machines.test.ts index f5f5eb95..9a0d6f1e 100644 --- a/control-plane/src/modules/machines.test.ts +++ b/control-plane/src/modules/machines.test.ts @@ -1,11 +1,10 @@ import { upsertMachine } from "./machines"; -import { - createCluster, - getClusterMachines, - getClusterServices, -} from "./management"; +import { createCluster, getClusterMachines } from "./management"; import * as redis from "./redis"; -import { upsertServiceDefinition } from "./service-definitions"; +import { + getServiceDefinitions, + upsertServiceDefinition, +} from "./service-definitions"; describe("machines", () => { beforeAll(async () => { @@ -60,10 +59,10 @@ describe("machines", () => { ); expect( - await getClusterServices({ + await getServiceDefinitions({ clusterId: id, - }), - ).toHaveLength(2); + }).then((services) => services.map((s) => s.service)), + ).toEqual(services); }); }); }); diff --git a/control-plane/src/modules/management.ts b/control-plane/src/modules/management.ts index 94ca95aa..ca666ca5 100644 --- a/control-plane/src/modules/management.ts +++ b/control-plane/src/modules/management.ts @@ -187,29 +187,6 @@ export const getClusterDetails = async ({ }; }; -export const getClusterServices = async ({ - clusterId, -}: { - clusterId: string; -}) => { - const services = await data.db - .select({ - definition: data.services.definition, - timestamp: data.services.timestamp, - }) - .from(data.services) - .where(eq(data.services.cluster_id, clusterId)); - - const serviceDefinitions = storedServiceDefinitionSchema.parse( - services.map((s) => ({ - ...(s.definition as object), - timestamp: s.timestamp, - })), - ); - - return serviceDefinitions; -}; - export const getClusterMachines = async ({ clusterId, }: { diff --git a/control-plane/src/modules/observability/hyperdx.ts b/control-plane/src/modules/observability/hyperdx.ts index f827a1c2..613c3091 100644 --- a/control-plane/src/modules/observability/hyperdx.ts +++ b/control-plane/src/modules/observability/hyperdx.ts @@ -4,7 +4,7 @@ import * as HyperDX from "@hyperdx/node-opentelemetry"; export const hdx = env.HYPERDX_API_KEY ? HyperDX : null; -if (!hdx) { +if (!hdx && env.NODE_ENV === "production") { // eslint-disable-next-line no-console console.log("HyperDX is not configured"); } diff --git a/control-plane/src/modules/prompt-templates.ts b/control-plane/src/modules/prompt-templates.ts index 368ccb88..dc852ede 100644 --- a/control-plane/src/modules/prompt-templates.ts +++ b/control-plane/src/modules/prompt-templates.ts @@ -87,7 +87,7 @@ export async function upsertRunConfig({ name, initial_prompt: initialPrompt, system_prompt: systemPrompt, - attached_functions: attachedFunctions, + attached_functions: attachedFunctions ?? [], result_schema: resultSchema, input_schema: inputSchema, }) diff --git a/control-plane/src/modules/router.ts b/control-plane/src/modules/router.ts index 030d82c6..8ad72ebb 100644 --- a/control-plane/src/modules/router.ts +++ b/control-plane/src/modules/router.ts @@ -44,7 +44,10 @@ import { } from "./knowledge/knowledgebase"; import { callsRouter } from "./calls/router"; import { buildModel } from "./models"; -import { deserializeFunctionSchema } from "./service-definitions"; +import { + deserializeFunctionSchema, + getServiceDefinitions, +} from "./service-definitions"; import { integrationsRouter } from "./integrations/router"; import { getServerStats } from "./server-stats"; @@ -439,13 +442,25 @@ export const router = initServer().router(contract, { const user = request.request.getAuth(); await user.canAccess({ cluster: { clusterId } }); - const services = await management.getClusterServices({ + const services = await getServiceDefinitions({ clusterId, }); + const transformedServices = services.map((service) => ({ + name: service.service, + timestamp: service.timestamp ?? new Date(), + description: service.definition.description, + functions: service.definition.functions?.map((fn) => ({ + name: fn.name, + description: fn.description, + schema: fn.schema, + config: fn.config, + })), + })); + return { status: 200, - body: services, + body: transformedServices, }; }, getBlobData: async (request) => { diff --git a/control-plane/src/modules/service-definitions.ts b/control-plane/src/modules/service-definitions.ts index 405fbedb..fcdd7430 100644 --- a/control-plane/src/modules/service-definitions.ts +++ b/control-plane/src/modules/service-definitions.ts @@ -43,7 +43,6 @@ export const storedServiceDefinitionSchema = z.array( z.object({ name: z.string(), description: z.string().optional(), - timestamp: z.date(), functions: z .array( z.object({ @@ -194,10 +193,18 @@ export const getServiceDefinition = async ({ export const getServiceDefinitions = async (owner: { clusterId: string; -}): Promise => { +}): Promise< + { + service: string; + definition: ServiceDefinition; + timestamp: Date | null; + }[] +> => { const serviceDefinitions = await data.db .select({ + service: data.services.service, definition: data.services.definition, + timestamp: data.services.timestamp, }) .from(data.services) .where(eq(data.services.cluster_id, owner.clusterId)); @@ -210,11 +217,11 @@ export const getServiceDefinitions = async (owner: { return []; } - const retrieved = parseServiceDefinition( - serviceDefinitions.map((d) => d.definition), - ); - - return retrieved; + return serviceDefinitions.map((r) => ({ + service: r.service, + definition: parseServiceDefinition([r.definition])[0], + timestamp: r.timestamp, + })); }; export const parseServiceDefinition = ( diff --git a/control-plane/src/modules/workflows/agent/run.ts b/control-plane/src/modules/workflows/agent/run.ts index 8d3e2db2..f2ee89bf 100644 --- a/control-plane/src/modules/workflows/agent/run.ts +++ b/control-plane/src/modules/workflows/agent/run.ts @@ -36,7 +36,10 @@ import { events } from "../../observability/events"; /** * Run a workflow from the most recent saved state **/ -export const processRun = async (run: Run, metadata?: Record) => { +export const processRun = async ( + run: Run, + metadata?: Record, +) => { logger.info("Running workflow"); // Parallelize fetching additional context and service definitions @@ -64,7 +67,7 @@ export const processRun = async (run: Run, metadata?: Record) => allAvailableTools.push(...attachedFunctions); serviceDefinitions.flatMap((service) => - (service.functions ?? []).forEach((f) => { + (service.definition.functions ?? []).forEach((f) => { // Do not attach additional tools if `attachedFunctions` is provided if (attachedFunctions.length > 0 || f.config?.private) { return; @@ -72,7 +75,7 @@ export const processRun = async (run: Run, metadata?: Record) => allAvailableTools.push( serviceFunctionEmbeddingId({ - serviceName: service.name, + serviceName: service.definition.name, functionName: f.name, }), ); @@ -167,8 +170,8 @@ export const processRun = async (run: Run, metadata?: Record) => for (const message of state.messages.filter((m) => !m.persisted)) { await Promise.all([ insertRunMessage(message), - notifyNewMessage({message, metadata}), - ]) + notifyNewMessage({ message, metadata }), + ]); message.persisted = true; } }, @@ -447,7 +450,7 @@ export const buildMockTools = async (workflow: Run) => { } const serviceDefinition = serviceDefinitions.find( - (sd) => sd.name === serviceName, + (sd) => sd.service === serviceName, ); if (!serviceDefinition) { @@ -461,7 +464,7 @@ export const buildMockTools = async (workflow: Run) => { continue; } - const functionDefinition = serviceDefinition.functions?.find( + const functionDefinition = serviceDefinition.definition.functions?.find( (f) => f.name === functionName, );