@copilotkit/runtime 0.0.0-mme-load-agent-state-20250117154700

package/src/lib/integrations/shared.ts

@@ -0,0 +1,110 @@
+import { YogaInitialContext } from "graphql-yoga";
+import { buildSchemaSync } from "type-graphql";
+import { CopilotResolver } from "../../graphql/resolvers/copilot.resolver";
+import { useDeferStream } from "@graphql-yoga/plugin-defer-stream";
+import { CopilotRuntime } from "../runtime/copilot-runtime";
+import { CopilotServiceAdapter } from "../../service-adapters";
+import { CopilotCloudOptions } from "../cloud";
+import { LogLevel, createLogger } from "../../lib/logger";
+import { createYoga } from "graphql-yoga";
+import telemetry from "../telemetry-client";
+import { StateResolver } from "../../graphql/resolvers/state.resolver";
+
+const logger = createLogger();
+
+type AnyPrimitive = string | boolean | number | null;
+export type CopilotRequestContextProperties = Record<
+  string,
+  AnyPrimitive | Record<string, AnyPrimitive>
+>;
+
+export type GraphQLContext = YogaInitialContext & {
+  _copilotkit: CreateCopilotRuntimeServerOptions;
+  properties: CopilotRequestContextProperties;
+  logger: typeof logger;
+};
+
+export interface CreateCopilotRuntimeServerOptions {
+  runtime: CopilotRuntime<any>;
+  serviceAdapter: CopilotServiceAdapter;
+  endpoint: string;
+  baseUrl?: string;
+  cloud?: CopilotCloudOptions;
+  properties?: CopilotRequestContextProperties;
+  logLevel?: LogLevel;
+}
+
+export async function createContext(
+  initialContext: YogaInitialContext,
+  copilotKitContext: CreateCopilotRuntimeServerOptions,
+  contextLogger: typeof logger,
+  properties: CopilotRequestContextProperties = {},
+): Promise<Partial<GraphQLContext>> {
+  logger.debug({ copilotKitContext }, "Creating GraphQL context");
+  const ctx: GraphQLContext = {
+    ...initialContext,
+    _copilotkit: {
+      ...copilotKitContext,
+    },
+    properties: { ...properties },
+    logger: contextLogger,
+  };
+  return ctx;
+}
+
+export function buildSchema(
+  options: {
+    emitSchemaFile?: string;
+  } = {},
+) {
+  logger.debug("Building GraphQL schema...");
+  const schema = buildSchemaSync({
+    resolvers: [CopilotResolver, StateResolver],
+    emitSchemaFile: options.emitSchemaFile,
+  });
+  logger.debug("GraphQL schema built successfully");
+  return schema;
+}
+
+export type CommonConfig = {
+  logging: typeof logger;
+  schema: ReturnType<typeof buildSchema>;
+  plugins: Parameters<typeof createYoga>[0]["plugins"];
+  context: (ctx: YogaInitialContext) => Promise<Partial<GraphQLContext>>;
+};
+
+export function getCommonConfig(options: CreateCopilotRuntimeServerOptions): CommonConfig {
+  const logLevel = (process.env.LOG_LEVEL as LogLevel) || (options.logLevel as LogLevel) || "error";
+  const logger = createLogger({ level: logLevel, component: "getCommonConfig" });
+
+  const contextLogger = createLogger({ level: logLevel });
+
+  if (options.cloud) {
+    telemetry.setCloudConfiguration({
+      publicApiKey: options.cloud.publicApiKey,
+      baseUrl: options.cloud.baseUrl,
+    });
+  }
+
+  if (options.properties?._copilotkit) {
+    telemetry.setGlobalProperties({
+      _copilotkit: {
+        ...(options.properties._copilotkit as Record<string, any>),
+      },
+    });
+  }
+
+  telemetry.setGlobalProperties({
+    runtime: {
+      serviceAdapter: options.serviceAdapter.constructor.name,
+    },
+  });
+
+  return {
+    logging: createLogger({ component: "Yoga GraphQL", level: logLevel }),
+    schema: buildSchema(),
+    plugins: [useDeferStream()],
+    context: (ctx: YogaInitialContext): Promise<Partial<GraphQLContext>> =>
+      createContext(ctx, options, contextLogger, options.properties),
+  };
+}

package/src/lib/logger.ts

@@ -0,0 +1,28 @@
+import createPinoLogger from "pino";
+import pretty from "pino-pretty";
+
+export type LogLevel = "debug" | "info" | "warn" | "error";
+
+export type CopilotRuntimeLogger = ReturnType<typeof createLogger>;
+
+export function createLogger(options?: { level?: LogLevel; component?: string }) {
+  const { level, component } = options || {};
+  const stream = pretty({ colorize: true });
+
+  const logger = createPinoLogger(
+    {
+      level: process.env.LOG_LEVEL || level || "error",
+      redact: {
+        paths: ["pid", "hostname"],
+        remove: true,
+      },
+    },
+    stream,
+  );
+
+  if (component) {
+    return logger.child({ component });
+  } else {
+    return logger;
+  }
+}

package/src/lib/runtime/copilot-runtime.ts

@@ -0,0 +1,571 @@
+/**
+ * <Callout type="info">
+ *   This is the reference for the `CopilotRuntime` class. For more information and example code snippets, please see [Concept: Copilot Runtime](/concepts/copilot-runtime).
+ * </Callout>
+ *
+ * ## Usage
+ *
+ * ```tsx
+ * import { CopilotRuntime } from "@copilotkit/runtime";
+ *
+ * const copilotKit = new CopilotRuntime();
+ * ```
+ */
+
+import { Action, actionParametersToJsonSchema, Parameter, randomId } from "@copilotkit/shared";
+import {
+  CopilotServiceAdapter,
+  ExperimentalEmptyAdapter,
+  RemoteChain,
+  RemoteChainParameters,
+} from "../../service-adapters";
+import { MessageInput } from "../../graphql/inputs/message.input";
+import { ActionInput } from "../../graphql/inputs/action.input";
+import { RuntimeEventSource } from "../../service-adapters/events";
+import { convertGqlInputToMessages } from "../../service-adapters/conversion";
+import { Message } from "../../graphql/types/converted";
+import { ForwardedParametersInput } from "../../graphql/inputs/forwarded-parameters.input";
+import {
+  isLangGraphAgentAction,
+  LangGraphAgentAction,
+  EndpointType,
+  setupRemoteActions,
+  EndpointDefinition,
+  CopilotKitEndpoint,
+  LangGraphPlatformEndpoint,
+} from "./remote-actions";
+import { GraphQLContext } from "../integrations/shared";
+import { AgentSessionInput } from "../../graphql/inputs/agent-session.input";
+import { from } from "rxjs";
+import { AgentStateInput } from "../../graphql/inputs/agent-state.input";
+import { ActionInputAvailability } from "../../graphql/types/enums";
+import { createHeaders } from "./remote-action-constructors";
+import { Agent } from "../../graphql/types/agents-response.type";
+import { ExtensionsInput } from "../../graphql/inputs/extensions.input";
+import { ExtensionsResponse } from "../../graphql/types/extensions-response.type";
+import { LoadAgentStateResponse } from "../../graphql/types/load-agent-state-response.type";
+import { Client as LangGraphClient } from "@langchain/langgraph-sdk";
+import { langchainMessagesToCopilotKit } from "./remote-lg-action";
+
+interface CopilotRuntimeRequest {
+  serviceAdapter: CopilotServiceAdapter;
+  messages: MessageInput[];
+  actions: ActionInput[];
+  agentSession?: AgentSessionInput;
+  agentStates?: AgentStateInput[];
+  outputMessagesPromise: Promise<Message[]>;
+  threadId?: string;
+  runId?: string;
+  publicApiKey?: string;
+  graphqlContext: GraphQLContext;
+  forwardedParameters?: ForwardedParametersInput;
+  url?: string;
+  extensions?: ExtensionsInput;
+}
+
+interface CopilotRuntimeResponse {
+  threadId: string;
+  runId?: string;
+  eventSource: RuntimeEventSource;
+  serverSideActions: Action<any>[];
+  actionInputsWithoutAgents: ActionInput[];
+  extensions?: ExtensionsResponse;
+}
+
+type ActionsConfiguration<T extends Parameter[] | [] = []> =
+  | Action<T>[]
+  | ((ctx: { properties: any; url?: string }) => Action<T>[]);
+
+interface OnBeforeRequestOptions {
+  threadId?: string;
+  runId?: string;
+  inputMessages: Message[];
+  properties: any;
+  url?: string;
+}
+
+type OnBeforeRequestHandler = (options: OnBeforeRequestOptions) => void | Promise<void>;
+
+interface OnAfterRequestOptions {
+  threadId: string;
+  runId?: string;
+  inputMessages: Message[];
+  outputMessages: Message[];
+  properties: any;
+  url?: string;
+}
+
+type OnAfterRequestHandler = (options: OnAfterRequestOptions) => void | Promise<void>;
+
+interface Middleware {
+  /**
+   * A function that is called before the request is processed.
+   */
+  onBeforeRequest?: OnBeforeRequestHandler;
+
+  /**
+   * A function that is called after the request is processed.
+   */
+  onAfterRequest?: OnAfterRequestHandler;
+}
+
+type AgentWithEndpoint = Agent & { endpoint: EndpointDefinition };
+
+export interface CopilotRuntimeConstructorParams<T extends Parameter[] | [] = []> {
+  /**
+   * Middleware to be used by the runtime.
+   *
+   * ```ts
+   * onBeforeRequest: (options: {
+   *   threadId?: string;
+   *   runId?: string;
+   *   inputMessages: Message[];
+   *   properties: any;
+   * }) => void | Promise<void>;
+   * ```
+   *
+   * ```ts
+   * onAfterRequest: (options: {
+   *   threadId?: string;
+   *   runId?: string;
+   *   inputMessages: Message[];
+   *   outputMessages: Message[];
+   *   properties: any;
+   * }) => void | Promise<void>;
+   * ```
+   */
+  middleware?: Middleware;
+
+  /*
+   * A list of server side actions that can be executed.
+   */
+  actions?: ActionsConfiguration<T>;
+
+  /*
+   * Deprecated: Use `remoteEndpoints`.
+   */
+  remoteActions?: CopilotKitEndpoint[];
+
+  /*
+   * A list of remote actions that can be executed.
+   */
+  remoteEndpoints?: EndpointDefinition[];
+
+  /*
+   * An array of LangServer URLs.
+   */
+  langserve?: RemoteChainParameters[];
+}
+
+export class CopilotRuntime<const T extends Parameter[] | [] = []> {
+  public actions: ActionsConfiguration<T>;
+  public remoteEndpointDefinitions: EndpointDefinition[];
+  private langserve: Promise<Action<any>>[] = [];
+  private onBeforeRequest?: OnBeforeRequestHandler;
+  private onAfterRequest?: OnAfterRequestHandler;
+
+  constructor(params?: CopilotRuntimeConstructorParams<T>) {
+    this.actions = params?.actions || [];
+
+    for (const chain of params?.langserve || []) {
+      const remoteChain = new RemoteChain(chain);
+      this.langserve.push(remoteChain.toAction());
+    }
+
+    this.remoteEndpointDefinitions = params?.remoteEndpoints ?? params?.remoteActions ?? [];
+
+    this.onBeforeRequest = params?.middleware?.onBeforeRequest;
+    this.onAfterRequest = params?.middleware?.onAfterRequest;
+  }
+
+  async processRuntimeRequest(request: CopilotRuntimeRequest): Promise<CopilotRuntimeResponse> {
+    const {
+      serviceAdapter,
+      messages: rawMessages,
+      actions: clientSideActionsInput,
+      threadId,
+      runId,
+      outputMessagesPromise,
+      graphqlContext,
+      forwardedParameters,
+      agentSession,
+      url,
+      extensions,
+    } = request;
+
+    const eventSource = new RuntimeEventSource();
+
+    try {
+      if (agentSession) {
+        return await this.processAgentRequest(request);
+      }
+      if (serviceAdapter instanceof ExperimentalEmptyAdapter) {
+        // TODO: use CPK error here
+        throw new Error(`Invalid adapter configuration: ExperimentalEmptyAdapter is only meant to be used with agent lock mode. 
+For non-agent components like useCopilotChatSuggestions, CopilotTextarea, or CopilotTask, 
+please use an LLM adapter instead.`);
+      }
+
+      const messages = rawMessages.filter((message) => !message.agentStateMessage);
+
+      const inputMessages = convertGqlInputToMessages(messages);
+      const serverSideActions = await this.getServerSideActions(request);
+
+      const serverSideActionsInput: ActionInput[] = serverSideActions.map((action) => ({
+        name: action.name,
+        description: action.description,
+        jsonSchema: JSON.stringify(actionParametersToJsonSchema(action.parameters)),
+      }));
+
+      const actionInputs = flattenToolCallsNoDuplicates([
+        ...serverSideActionsInput,
+        ...clientSideActionsInput.filter(
+          // Filter remote actions from CopilotKit core loop
+          (action) => action.available !== ActionInputAvailability.remote,
+        ),
+      ]);
+
+      await this.onBeforeRequest?.({
+        threadId,
+        runId,
+        inputMessages,
+        properties: graphqlContext.properties,
+        url,
+      });
+
+      const result = await serviceAdapter.process({
+        messages: inputMessages,
+        actions: actionInputs,
+        threadId,
+        runId,
+        eventSource,
+        forwardedParameters,
+        extensions,
+      });
+
+      // for backwards compatibility, we deal with the case that no threadId is provided
+      // by the frontend, by using the threadId from the response
+      const nonEmptyThreadId = threadId ?? result.threadId;
+
+      outputMessagesPromise
+        .then((outputMessages) => {
+          this.onAfterRequest?.({
+            threadId: nonEmptyThreadId,
+            runId: result.runId,
+            inputMessages,
+            outputMessages,
+            properties: graphqlContext.properties,
+            url,
+          });
+        })
+        .catch((_error) => {});
+
+      return {
+        threadId: nonEmptyThreadId,
+        runId: result.runId,
+        eventSource,
+        serverSideActions,
+        actionInputsWithoutAgents: actionInputs.filter(
+          (action) =>
+            // TODO-AGENTS: do not exclude ALL server side actions
+            !serverSideActions.find((serverSideAction) => serverSideAction.name == action.name),
+          // !isLangGraphAgentAction(
+          //   serverSideActions.find((serverSideAction) => serverSideAction.name == action.name),
+          // ),
+        ),
+        extensions: result.extensions,
+      };
+    } catch (error) {
+      console.error("Error getting response:", error);
+      eventSource.sendErrorMessageToChat();
+      throw error;
+    }
+  }
+
+  async discoverAgentsFromEndpoints(graphqlContext: GraphQLContext): Promise<AgentWithEndpoint[]> {
+    const headers = createHeaders(null, graphqlContext);
+    const agents = this.remoteEndpointDefinitions.reduce(
+      async (acc: Promise<Agent[]>, endpoint) => {
+        const agents = await acc;
+        if (endpoint.type === EndpointType.LangGraphPlatform) {
+          const client = new LangGraphClient({
+            apiUrl: endpoint.deploymentUrl,
+            apiKey: endpoint.langsmithApiKey,
+          });
+
+          const data: Array<{ assistant_id: string; graph_id: string }> =
+            await client.assistants.search();
+
+          const endpointAgents = (data ?? []).map((entry) => ({
+            name: entry.graph_id,
+            id: entry.assistant_id,
+            description: "",
+            endpoint,
+          }));
+          return [...agents, ...endpointAgents];
+        }
+
+        interface InfoResponse {
+          agents?: Array<{
+            name: string;
+            description: string;
+          }>;
+        }
+
+        const response = await fetch(`${(endpoint as CopilotKitEndpoint).url}/info`, {
+          method: "POST",
+          headers,
+          body: JSON.stringify({ properties: graphqlContext.properties }),
+        });
+        const data: InfoResponse = await response.json();
+        const endpointAgents = (data?.agents ?? []).map((agent) => ({
+          name: agent.name,
+          description: agent.description ?? "",
+          id: randomId(), // Required by Agent type
+          endpoint,
+        }));
+        return [...agents, ...endpointAgents];
+      },
+      Promise.resolve([]),
+    );
+
+    return agents;
+  }
+
+  async loadAgentState(
+    graphqlContext: GraphQLContext,
+    threadId: string,
+    agentName: string,
+  ): Promise<LoadAgentStateResponse> {
+    const agentsWithEndpoints = await this.discoverAgentsFromEndpoints(graphqlContext);
+
+    const agentWithEndpoint = agentsWithEndpoints.find((agent) => agent.name === agentName);
+    if (!agentWithEndpoint) {
+      throw new Error("Agent not found");
+    }
+    const headers = createHeaders(null, graphqlContext);
+
+    if (agentWithEndpoint.endpoint.type === EndpointType.LangGraphPlatform) {
+      const client = new LangGraphClient({
+        apiUrl: agentWithEndpoint.endpoint.deploymentUrl,
+        apiKey: agentWithEndpoint.endpoint.langsmithApiKey,
+      });
+      const state = (await client.threads.getState(threadId)).values as any;
+
+      if (Object.keys(state).length === 0) {
+        return {
+          threadId,
+          threadExists: false,
+          state: JSON.stringify({}),
+          messages: JSON.stringify([]),
+        };
+      } else {
+        console.log(state);
+        const { messages, ...stateWithoutMessages } = state;
+        const copilotkitMessages = langchainMessagesToCopilotKit(messages);
+        return {
+          threadId,
+          threadExists: true,
+          state: JSON.stringify(stateWithoutMessages),
+          messages: JSON.stringify(copilotkitMessages),
+        };
+      }
+    } else if (
+      agentWithEndpoint.endpoint.type === EndpointType.CopilotKit ||
+      !("type" in agentWithEndpoint.endpoint)
+    ) {
+      const response = await fetch(
+        `${(agentWithEndpoint.endpoint as CopilotKitEndpoint).url}/agents/state`,
+        {
+          method: "POST",
+          headers,
+          body: JSON.stringify({
+            properties: graphqlContext.properties,
+            threadId,
+            name: agentName,
+          }),
+        },
+      );
+      const data: LoadAgentStateResponse = await response.json();
+
+      return {
+        ...data,
+        state: JSON.stringify(data.state),
+        messages: JSON.stringify(data.messages),
+      };
+    } else {
+      throw new Error(`Unknown endpoint type: ${(agentWithEndpoint.endpoint as any).type}`);
+    }
+  }
+
+  private async processAgentRequest(
+    request: CopilotRuntimeRequest,
+  ): Promise<CopilotRuntimeResponse> {
+    const {
+      messages: rawMessages,
+      outputMessagesPromise,
+      graphqlContext,
+      agentSession,
+      threadId: threadIdFromRequest,
+    } = request;
+    const { agentName, nodeName } = agentSession;
+
+    // for backwards compatibility, deal with the case when no threadId is provided
+    const threadId = threadIdFromRequest ?? agentSession.threadId;
+
+    const serverSideActions = await this.getServerSideActions(request);
+
+    const messages = convertGqlInputToMessages(rawMessages);
+
+    const agent = serverSideActions.find(
+      (action) => action.name === agentName && isLangGraphAgentAction(action),
+    ) as LangGraphAgentAction;
+
+    if (!agent) {
+      throw new Error(`Agent ${agentName} not found`);
+    }
+
+    const serverSideActionsInput: ActionInput[] = serverSideActions
+      .filter((action) => !isLangGraphAgentAction(action))
+      .map((action) => ({
+        name: action.name,
+        description: action.description,
+        jsonSchema: JSON.stringify(actionParametersToJsonSchema(action.parameters)),
+      }));
+
+    const actionInputsWithoutAgents = flattenToolCallsNoDuplicates([
+      ...serverSideActionsInput,
+      ...request.actions,
+    ]);
+
+    await this.onBeforeRequest?.({
+      threadId,
+      runId: undefined,
+      inputMessages: messages,
+      properties: graphqlContext.properties,
+    });
+    try {
+      const eventSource = new RuntimeEventSource();
+      const stream = await agent.langGraphAgentHandler({
+        name: agentName,
+        threadId,
+        nodeName,
+        actionInputsWithoutAgents,
+      });
+
+      eventSource.stream(async (eventStream$) => {
+        from(stream).subscribe({
+          next: (event) => eventStream$.next(event),
+          error: (err) => {
+            console.error("Error in stream", err);
+            eventStream$.error(err);
+            eventStream$.complete();
+          },
+          complete: () => eventStream$.complete(),
+        });
+      });
+
+      outputMessagesPromise
+        .then((outputMessages) => {
+          this.onAfterRequest?.({
+            threadId,
+            runId: undefined,
+            inputMessages: messages,
+            outputMessages,
+            properties: graphqlContext.properties,
+          });
+        })
+        .catch((_error) => {});
+
+      return {
+        threadId,
+        runId: undefined,
+        eventSource,
+        serverSideActions: [],
+        actionInputsWithoutAgents,
+      };
+    } catch (error) {
+      console.error("Error getting response:", error);
+      throw error;
+    }
+  }
+
+  private async getServerSideActions(request: CopilotRuntimeRequest): Promise<Action<any>[]> {
+    const { messages: rawMessages, graphqlContext, agentStates, url } = request;
+    const inputMessages = convertGqlInputToMessages(rawMessages);
+    const langserveFunctions: Action<any>[] = [];
+
+    for (const chainPromise of this.langserve) {
+      try {
+        const chain = await chainPromise;
+        langserveFunctions.push(chain);
+      } catch (error) {
+        console.error("Error loading langserve chain:", error);
+      }
+    }
+
+    const remoteEndpointDefinitions = this.remoteEndpointDefinitions.map(
+      (endpoint) =>
+        ({
+          ...endpoint,
+          type: resolveEndpointType(endpoint),
+        }) as EndpointDefinition,
+    );
+
+    const remoteActions = await setupRemoteActions({
+      remoteEndpointDefinitions,
+      graphqlContext,
+      messages: inputMessages,
+      agentStates,
+      frontendUrl: url,
+    });
+
+    const configuredActions =
+      typeof this.actions === "function"
+        ? this.actions({ properties: graphqlContext.properties, url })
+        : this.actions;
+
+    return [...configuredActions, ...langserveFunctions, ...remoteActions];
+  }
+}
+
+export function flattenToolCallsNoDuplicates(toolsByPriority: ActionInput[]): ActionInput[] {
+  let allTools: ActionInput[] = [];
+  const allToolNames: string[] = [];
+  for (const tool of toolsByPriority) {
+    if (!allToolNames.includes(tool.name)) {
+      allTools.push(tool);
+      allToolNames.push(tool.name);
+    }
+  }
+  return allTools;
+}
+
+// The two functions below are "factory functions", meant to create the action objects that adhere to the expected interfaces
+export function copilotKitEndpoint(config: Omit<CopilotKitEndpoint, "type">): CopilotKitEndpoint {
+  return {
+    ...config,
+    type: EndpointType.CopilotKit,
+  };
+}
+
+export function langGraphPlatformEndpoint(
+  config: Omit<LangGraphPlatformEndpoint, "type">,
+): LangGraphPlatformEndpoint {
+  return {
+    ...config,
+    type: EndpointType.LangGraphPlatform,
+  };
+}
+
+export function resolveEndpointType(endpoint: EndpointDefinition) {
+  if (!endpoint.type) {
+    if ("langsmithApiKey" in endpoint && "deploymentUrl" in endpoint && "agents" in endpoint) {
+      return EndpointType.LangGraphPlatform;
+    } else {
+      return EndpointType.CopilotKit;
+    }
+  }
+
+  return endpoint.type;
+}