@copilotkit/runtime 0.37.0 → 0.38.0-mme-alpha.0

package/src/graphql/resolvers/copilot.resolver.ts

@@ -1,5 +1,16 @@
 import { Arg, Ctx, Mutation, Query, Resolver } from "type-graphql";
-import { Subject, firstValueFrom, shareReplay, skipWhile, takeWhile } from "rxjs";
+import {
+  ReplaySubject,
+  Subject,
+  Subscription,
+  finalize,
+  firstValueFrom,
+  shareReplay,
+  skipWhile,
+  take,
+  takeWhile,
+  tap,
+} from "rxjs";
 import { GenerateCopilotResponseInput } from "../inputs/generate-copilot-response.input";
 import { CopilotResponse } from "../types/copilot-response.type";
 import { MessageRole } from "../types/enums";
@@ -7,9 +18,82 @@ import { Repeater } from "graphql-yoga";
 import type { CopilotRequestContextProperties, GraphQLContext } from "../../lib/integrations";
 import { nanoid } from "nanoid";
 import { RuntimeEvent, RuntimeEventTypes } from "../../service-adapters/events";
-import { MessageStatusUnion, SuccessMessageStatus } from "../types/message-status.type";
+import {
+  FailedMessageStatus,
+  MessageStatusUnion,
+  SuccessMessageStatus,
+} from "../types/message-status.type";
 import { ResponseStatusUnion, SuccessResponseStatus } from "../types/response-status.type";
 import { GraphQLJSONObject } from "graphql-scalars";
+import { plainToInstance } from "class-transformer";
+import { GuardrailsResult, GuardrailsResultStatus } from "../types/guardrails-result.type";
+import { GraphQLError } from "graphql";
+import {
+  GuardrailsValidationFailureResponse,
+  MessageStreamInterruptedResponse,
+  UnknownErrorResponse,
+} from "../../utils";
+import { CopilotRuntimeLogger } from "../../lib/logger";
+
+const invokeGuardrails = async ({
+  baseUrl,
+  copilotCloudPublicApiKey,
+  data,
+  logger,
+  onResult,
+}: {
+  baseUrl: string;
+  copilotCloudPublicApiKey: string;
+  data: GenerateCopilotResponseInput;
+  logger: CopilotRuntimeLogger;
+  onResult: (result: GuardrailsResult) => void;
+}) => {
+  if (
+    data.messages.length &&
+    data.messages[data.messages.length - 1].textMessage?.role === MessageRole.user
+  ) {
+    const messages = data.messages
+      .filter(
+        (m) =>
+          m.textMessage !== undefined &&
+          (m.textMessage.role === MessageRole.user || m.textMessage.role === MessageRole.assistant),
+      )
+      .map((m) => ({
+        role: m.textMessage!.role,
+        content: m.textMessage.content,
+      }));
+
+    const lastMessage = messages[messages.length - 1];
+    const restOfMessages = messages.slice(0, -1);
+
+    const body = {
+      input: lastMessage.content,
+      validTopics: data.cloud.guardrails.inputValidationRules.allowList,
+      invalidTopics: data.cloud.guardrails.inputValidationRules.denyList,
+      messages: restOfMessages,
+    };
+
+    const guardrailsResult = await fetch(`${baseUrl}/guardrails/validate`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "X-CopilotCloud-Public-API-Key": copilotCloudPublicApiKey,
+      },
+      body: JSON.stringify(body),
+    });
+
+    const responseBody = await guardrailsResult.json();
+
+    if (!guardrailsResult.ok) {
+      logger.error({ error: responseBody }, "Failed to invoke guardrails");
+      throw new GraphQLError("Failed to invoke guardrails");
+    }
+
+    onResult(responseBody);
+  } else {
+    onResult({ status: GuardrailsResultStatus.ALLOWED });
+  }
+};
 
 @Resolver(() => CopilotResponse)
 export class CopilotResolver {
@@ -25,17 +109,53 @@ export class CopilotResolver {
     @Arg("properties", () => GraphQLJSONObject, { nullable: true })
     properties?: CopilotRequestContextProperties,
   ) {
-    if (properties) {
-      ctx._copilotkit.properties = { ...ctx._copilotkit.properties, ...properties };
-    }
+    let logger = ctx.logger.child({ component: "CopilotResolver.generateCopilotResponse" });
+    logger.debug({ data }, "Generating Copilot response");
+
     const copilotRuntime = ctx._copilotkit.runtime;
     const serviceAdapter = ctx._copilotkit.serviceAdapter;
-    const responseStatus = new Subject<typeof ResponseStatusUnion>();
 
+    if (properties) {
+      logger.debug("Properties provided, merging with context properties");
+      ctx.properties = { ...ctx.properties, ...properties };
+    }
+
+    let copilotCloudPublicApiKey: string | null = null;
+    let copilotCloudBaseUrl: string;
+
+    if (data.cloud) {
+      logger = logger.child({ cloud: true });
+      logger.debug("Cloud configuration provided, checking for public API key in headers");
+      const key = ctx.request.headers.get("x-copilotcloud-public-api-key");
+      if (key) {
+        logger.debug("Public API key found in headers");
+        copilotCloudPublicApiKey = key;
+      } else {
+        logger.error("Public API key not found in headers");
+        throw new GraphQLError("X-CopilotCloud-Public-API-Key header is required");
+      }
+
+      if (process.env.COPILOT_CLOUD_BASE_URL) {
+        copilotCloudBaseUrl = process.env.COPILOT_CLOUD_BASE_URL;
+      } else if (ctx._copilotkit.baseUrl) {
+        copilotCloudBaseUrl = ctx._copilotkit.baseUrl;
+      } else {
+        copilotCloudBaseUrl = "https://api.cloud.copilotkit.ai";
+      }
+
+      logger = logger.child({ copilotCloudBaseUrl });
+    }
+    logger.debug("Setting up subjects");
+    const responseStatus$ = new ReplaySubject<typeof ResponseStatusUnion>();
+    const interruptStreaming$ = new ReplaySubject<{ reason: string; messageId?: string }>();
+    const guardrailsResult$ = new ReplaySubject<GuardrailsResult>();
+
+    logger.debug("Processing");
     const {
       eventSource,
       threadId = nanoid(),
       runId,
+      actions,
     } = await copilotRuntime.process({
       serviceAdapter,
       messages: data.messages,
@@ -45,18 +165,60 @@ export class CopilotResolver {
       publicApiKey: undefined,
     });
 
+    logger.debug("Event source created, creating response");
+
     const response = {
       threadId,
       runId,
-      status: firstValueFrom(responseStatus),
+      status: firstValueFrom(responseStatus$),
       messages: new Repeater(async (pushMessage, stopStreamingMessages) => {
+        logger.debug("Messages repeater created");
+
+        if (data.cloud?.guardrails) {
+          logger = logger.child({ guardrails: true });
+          logger.debug("Guardrails is enabled, validating input");
+
+          invokeGuardrails({
+            baseUrl: copilotCloudBaseUrl,
+            copilotCloudPublicApiKey,
+            data,
+            logger,
+            onResult: (result) => {
+              logger.debug({ status: result.status }, "Guardrails validation done");
+              guardrailsResult$.next(result);
+              if (result.status === "denied") {
+                responseStatus$.next(
+                  new GuardrailsValidationFailureResponse({ guardrailsReason: result.reason }),
+                );
+                interruptStreaming$.next({
+                  reason: `Interrupted due to Guardrails validation failure. Reason: ${result.reason}`,
+                });
+              }
+            },
+          });
+        }
+
+        let eventStreamSubscription: Subscription;
+
         // run and process the event stream
-        const eventStream = eventSource.process(copilotRuntime.actions).pipe(
-          // shareReplay() ensures that later subscribers will see the whole stream instead of
-          // just the events that were emitted after the subscriber was added.
-          shareReplay(),
-        );
-        eventStream.subscribe({
+        const eventStream = eventSource
+          .process({
+            serversideActions: copilotRuntime.actions,
+            guardrailsResult$: data.cloud?.guardrails ? guardrailsResult$ : null,
+          })
+          .pipe(
+            // shareReplay() ensures that later subscribers will see the whole stream instead of
+            // just the events that were emitted after the subscriber was added.
+            shareReplay(),
+            finalize(() => {
+              logger.debug("Event stream finalized, stopping streaming messages");
+              stopStreamingMessages();
+            }),
+          );
+
+        logger.debug("Event stream created, subscribing to event stream");
+
+        eventStreamSubscription = eventStream.subscribe({
           next: async (event) => {
             switch (event.type) {
               ////////////////////////////////
@@ -74,21 +236,61 @@ export class CopilotResolver {
                 // signal when we are done streaming
                 const streamingTextStatus = new Subject<typeof MessageStatusUnion>();
 
+                const messageId = nanoid();
+
                 // push the new message
                 pushMessage({
-                  id: nanoid(),
+                  id: messageId,
                   status: firstValueFrom(streamingTextStatus),
                   createdAt: new Date(),
                   role: MessageRole.assistant,
                   content: new Repeater(async (pushTextChunk, stopStreamingText) => {
-                    // push the message content
-                    await textMessageContentStream.forEach(async (e: RuntimeEvent) => {
-                      if (e.type == RuntimeEventTypes.TextMessageContent) {
-                        await pushTextChunk(e.content);
-                      }
+                    logger.debug("Text message content repeater created");
+
+                    let textSubscription: Subscription;
+
+                    interruptStreaming$
+                      .pipe(
+                        shareReplay(),
+                        take(1),
+                        tap(({ reason, messageId }) => {
+                          logger.debug({ reason, messageId }, "Text streaming interrupted");
+
+                          streamingTextStatus.next(
+                            plainToInstance(FailedMessageStatus, { reason }),
+                          );
+
+                          responseStatus$.next(new MessageStreamInterruptedResponse({ messageId }));
+                          stopStreamingText();
+                          textSubscription.unsubscribe();
+                        }),
+                      )
+                      .subscribe();
+
+                    logger.debug("Subscribing to text message content stream");
+
+                    textSubscription = textMessageContentStream.subscribe({
+                      next: async (e: RuntimeEvent) => {
+                        if (e.type == RuntimeEventTypes.TextMessageContent) {
+                          await pushTextChunk(e.content);
+                        }
+                      },
+                      error: (err) => {
+                        logger.error({ err }, "Error in text message content stream");
+                        interruptStreaming$.next({
+                          reason: "Error streaming message content",
+                          messageId,
+                        });
+                        stopStreamingText();
+                        textSubscription.unsubscribe();
+                      },
+                      complete: () => {
+                        logger.debug("Text message content stream completed");
+                        streamingTextStatus.next(new SuccessMessageStatus());
+                        stopStreamingText();
+                        textSubscription.unsubscribe();
+                      },
                     });
-                    stopStreamingText();
-                    streamingTextStatus.next(new SuccessMessageStatus());
                   }),
                 });
                 break;
@@ -96,6 +298,7 @@ export class CopilotResolver {
               // ActionExecutionStart
               ////////////////////////////////
               case RuntimeEventTypes.ActionExecutionStart:
+                logger.debug("Action execution start event received");
                 const actionExecutionArgumentStream = eventStream.pipe(
                   skipWhile((e) => e !== event),
                   takeWhile((e) => e.type != RuntimeEventTypes.ActionExecutionEnd),
@@ -108,13 +311,33 @@ export class CopilotResolver {
                   name: event.actionName,
                   scope: event.scope!,
                   arguments: new Repeater(async (pushArgumentsChunk, stopStreamingArguments) => {
-                    await actionExecutionArgumentStream.forEach(async (e: RuntimeEvent) => {
-                      if (e.type == RuntimeEventTypes.ActionExecutionArgs) {
-                        await pushArgumentsChunk(e.args);
-                      }
+                    logger.debug("Action execution argument stream created");
+
+                    let actionExecutionArgumentSubscription: Subscription;
+                    actionExecutionArgumentSubscription = actionExecutionArgumentStream.subscribe({
+                      next: async (e: RuntimeEvent) => {
+                        if (e.type == RuntimeEventTypes.ActionExecutionArgs) {
+                          await pushArgumentsChunk(e.args);
+                        }
+                      },
+                      error: (err) => {
+                        logger.error({ err }, "Error in action execution argument stream");
+                        streamingArgumentsStatus.next(
+                          plainToInstance(FailedMessageStatus, {
+                            reason:
+                              "An unknown error has occurred in the action execution argument stream",
+                          }),
+                        );
+                        stopStreamingArguments();
+                        actionExecutionArgumentSubscription.unsubscribe();
+                      },
+                      complete: () => {
+                        logger.debug("Action execution argument stream completed");
+                        streamingArgumentsStatus.next(new SuccessMessageStatus());
+                        stopStreamingArguments();
+                        actionExecutionArgumentSubscription.unsubscribe();
+                      },
                     });
-                    stopStreamingArguments();
-                    streamingArgumentsStatus.next(new SuccessMessageStatus());
                   }),
                 });
                 break;
@@ -122,6 +345,7 @@ export class CopilotResolver {
               // ActionExecutionResult
               ////////////////////////////////
               case RuntimeEventTypes.ActionExecutionResult:
+                logger.debug("Action execution result event received");
                 pushMessage({
                   id: nanoid(),
                   status: new SuccessMessageStatus(),
@@ -133,9 +357,24 @@ export class CopilotResolver {
                 break;
             }
           },
-          error: (err) => console.error("Error in event source", err),
-          complete: () => {
-            responseStatus.next(new SuccessResponseStatus());
+          error: (err) => {
+            logger.error({ err }, "Error in event stream");
+            responseStatus$.next(
+              new UnknownErrorResponse({
+                description: `An unknown error has occurred in the event stream`,
+              }),
+            );
+            eventStreamSubscription.unsubscribe();
+            stopStreamingMessages();
+          },
+          complete: async () => {
+            logger.debug("Event stream completed");
+            if (data.cloud?.guardrails) {
+              logger.debug("Guardrails is enabled, waiting for guardrails result");
+              await firstValueFrom(guardrailsResult$);
+            }
+            responseStatus$.next(new SuccessResponseStatus());
+            eventStreamSubscription.unsubscribe();
             stopStreamingMessages();
           },
         });

package/src/graphql/types/response-status.type.ts

@@ -1,3 +1,4 @@
+import { GraphQLJSON } from "graphql-scalars";
 import { Field, InterfaceType, ObjectType, createUnionType, registerEnumType } from "type-graphql";
 
 export enum ResponseStatusCode {
@@ -38,12 +39,25 @@ export class SuccessResponseStatus extends BaseResponseStatus {
   code: ResponseStatusCode = ResponseStatusCode.Success;
 }
 
+export enum FailedResponseStatusReason {
+  GUARDRAILS_VALIDATION_FAILED = "GUARDRAILS_VALIDATION_FAILED",
+  MESSAGE_STREAM_INTERRUPTED = "MESSAGE_STREAM_INTERRUPTED",
+  UNKNOWN_ERROR = "UNKNOWN_ERROR",
+}
+
+registerEnumType(FailedResponseStatusReason, {
+  name: "FailedResponseStatusReason",
+});
+
 @ObjectType({ implements: BaseResponseStatus })
 export class FailedResponseStatus extends BaseResponseStatus {
   code: ResponseStatusCode = ResponseStatusCode.Failed;
 
-  @Field(() => String)
-  reason: string;
+  @Field(() => FailedResponseStatusReason)
+  reason: FailedResponseStatusReason;
+
+  @Field(() => GraphQLJSON, { nullable: true })
+  details?: Record<string, any> = null;
 }
 
 export const ResponseStatusUnion = createUnionType({

package/src/index.ts

@@ -1,2 +1,3 @@
 import "reflect-metadata";
 export * from "./lib";
+export * from "./utils";

package/src/lib/cloud/index.ts

@@ -0,0 +1,4 @@
+export interface CopilotCloudOptions {
+  baseUrl?: string;
+  publicApiKey?: string;
+}

package/src/lib/copilot-runtime.ts

@@ -99,7 +99,7 @@
  */
 
 import { Action, actionParametersToJsonSchema, Parameter } from "@copilotkit/shared";
-import { RemoteChain, CopilotServiceAdapter } from "../service-adapters";
+import { RemoteChain, RemoteChainParameters, CopilotServiceAdapter } from "../service-adapters";
 import { CopilotCloud, RemoteCopilotCloud } from "./copilot-cloud";
 import { MessageInput } from "../graphql/inputs/message.input";
 import { ActionInput } from "../graphql/inputs/action.input";
@@ -119,6 +119,7 @@ interface CopilotRuntimeResponse {
   threadId?: string;
   runId?: string;
   eventSource: RuntimeEventSource;
+  actions: Action<any>[];
 }
 
 export interface CopilotRuntimeConstructorParams<T extends Parameter[] | [] = []> {
@@ -130,7 +131,7 @@ export interface CopilotRuntimeConstructorParams<T extends Parameter[] | [] = []
   /*
    * An array of LangServer URLs.
    */
-  langserve?: RemoteChain[];
+  langserve?: RemoteChainParameters[];
 
   debug?: boolean;
   copilotCloud?: CopilotCloud;
@@ -170,9 +171,6 @@ export class CopilotRuntime<const T extends Parameter[] | [] = []> {
     runId,
     publicApiKey,
   }: CopilotRuntimeRequest): Promise<CopilotRuntimeResponse> {
-    // TODO-PROTOCOL: cloud configuration
-    // const cloud: CopilotCloudConfig = forwardedProps.cloud;
-
     const langserveFunctions: Action<any>[] = [];
 
     for (const chainPromise of this.langserve) {
@@ -184,15 +182,15 @@ export class CopilotRuntime<const T extends Parameter[] | [] = []> {
       }
     }
 
-    const serverSideActionsInput: ActionInput[] = [...this.actions, ...langserveFunctions].map(
-      (action) => ({
-        name: action.name,
-        description: action.description,
-        jsonSchema: JSON.stringify(actionParametersToJsonSchema(action.parameters)),
-      }),
-    );
+    const actions = [...this.actions, ...langserveFunctions];
+
+    const serverSideActionsInput: ActionInput[] = actions.map((action) => ({
+      name: action.name,
+      description: action.description,
+      jsonSchema: JSON.stringify(actionParametersToJsonSchema(action.parameters)),
+    }));
 
-    const actions = flattenToolCallsNoDuplicates([
+    const actionInputs = flattenToolCallsNoDuplicates([
       ...serverSideActionsInput,
       ...clientSideActionsInput,
     ]);
@@ -202,39 +200,17 @@ export class CopilotRuntime<const T extends Parameter[] | [] = []> {
       // TODO-PROTOCOL: type this and support function calls
       const result = await serviceAdapter.process({
         messages: convertGqlInputToMessages(messages),
-        actions,
+        actions: actionInputs,
         threadId,
         runId,
         eventSource,
       });
 
-      // TODO-PROTOCOL add guardrails
-      //
-      // if (publicApiKey !== undefined) {
-      //   // wait for the cloud log chat to finish before streaming back the response
-      //   try {
-      //     const checkGuardrailsInputResult = await this.copilotCloud.checkGuardrailsInput({
-      //       cloud,
-      //       publicApiKey,
-      //       messages: forwardedProps.messages || [],
-      //     });
-
-      //     if (checkGuardrailsInputResult.status === "denied") {
-      //       // the chat was denied. instead of streaming back the response,
-      //       // we let the client know...
-      //       return {
-      //         stream: new SingleChunkReadableStream(checkGuardrailsInputResult.reason),
-      //         // headers: result.headers,
-      //       };
-      //     }
-      //   } catch (error) {
-      //     console.error("Error checking guardrails:", error);
-      //   }
-      // }
       return {
         threadId: result.threadId,
         runId: result.runId,
         eventSource,
+        actions,
       };
     } catch (error) {
       console.error("Error getting response:", error);

package/src/lib/integrations/nextjs/app-router.ts

@@ -1,29 +1,21 @@
 import { createYoga } from "graphql-yoga";
-import { getCommonConfig } from "../shared";
-import { CopilotRuntime } from "../../copilot-runtime";
-import { CopilotServiceAdapter } from "../../../service-adapters";
+import { CreateCopilotRuntimeServerOptions, getCommonConfig } from "../shared";
 
-export function copilotRuntimeNextJSAppRouterEndpoint({
-  runtime,
-  endpoint,
-  serviceAdapter,
-}: {
-  runtime: CopilotRuntime;
-  serviceAdapter: CopilotServiceAdapter;
-  endpoint: string;
-}) {
-  const commonConfig = getCommonConfig({ runtime, serviceAdapter });
+export function copilotRuntimeNextJSAppRouterEndpoint(options: CreateCopilotRuntimeServerOptions) {
+  const commonConfig = getCommonConfig(options);
+  const logger = commonConfig.logging;
+  logger.debug("Creating NextJS App Router endpoint");
 
   const yoga = createYoga({
     ...commonConfig,
-    graphqlEndpoint: endpoint,
+    graphqlEndpoint: options.endpoint,
     fetchAPI: { Response: globalThis.Response },
   });
 
   return {
     handleRequest: yoga,
-    GET: yoga,
-    POST: yoga,
-    OPTIONS: yoga,
+    GET: yoga as any,
+    POST: yoga as any,
+    OPTIONS: yoga as any,
   };
 }

package/src/lib/integrations/nextjs/pages-router.ts

@@ -1,7 +1,5 @@
 import { YogaServerInstance, createYoga } from "graphql-yoga";
-import { GraphQLContext, getCommonConfig } from "../shared";
-import { CopilotRuntime } from "../../copilot-runtime";
-import { CopilotServiceAdapter } from "../../../service-adapters";
+import { CreateCopilotRuntimeServerOptions, GraphQLContext, getCommonConfig } from "../shared";
 
 export const config = {
   api: {
@@ -11,25 +9,21 @@ export const config = {
 
 export type CopilotRuntimeServerInstance<T> = YogaServerInstance<T, Partial<GraphQLContext>>;
 
-// Theis import is needed to fix the type error
+// This import is needed to fix the type error
 // Fix is currently in TypeScript 5.5 beta, waiting for stable version
 // https://github.com/microsoft/TypeScript/issues/42873#issuecomment-2066874644
 export type {} from "@whatwg-node/server";
 
-export function copilotRuntimeNextJSPagesRouterEndpoint({
-  runtime,
-  endpoint,
-  serviceAdapter,
-}: {
-  runtime: CopilotRuntime;
-  serviceAdapter: CopilotServiceAdapter;
-  endpoint: string;
-}): CopilotRuntimeServerInstance<GraphQLContext> {
-  const commonConfig = getCommonConfig({ runtime, serviceAdapter });
+export function copilotRuntimeNextJSPagesRouterEndpoint(
+  options: CreateCopilotRuntimeServerOptions,
+): CopilotRuntimeServerInstance<GraphQLContext> {
+  const commonConfig = getCommonConfig(options);
+  const logger = commonConfig.logging;
+  logger.debug("Creating NextJS Pages Router endpoint");
 
   const yoga = createYoga({
     ...commonConfig,
-    graphqlEndpoint: endpoint,
+    graphqlEndpoint: options.endpoint,
   });
 
   return yoga;

package/src/lib/integrations/node-http/index.ts

@@ -1,22 +1,14 @@
 import { createYoga } from "graphql-yoga";
-import { getCommonConfig } from "../shared";
-import { CopilotRuntime } from "../../copilot-runtime";
-import { CopilotServiceAdapter } from "../../../service-adapters";
+import { CreateCopilotRuntimeServerOptions, getCommonConfig } from "../shared";
 
-export function copilotRuntimeNodeHttpEndpoint({
-  runtime,
-  endpoint,
-  serviceAdapter,
-}: {
-  runtime: CopilotRuntime;
-  serviceAdapter: CopilotServiceAdapter;
-  endpoint: string;
-}) {
-  const commonConfig = getCommonConfig({ runtime, serviceAdapter });
+export function copilotRuntimeNodeHttpEndpoint(options: CreateCopilotRuntimeServerOptions) {
+  const commonConfig = getCommonConfig(options);
+  const logger = commonConfig.logging;
+  logger.debug("Creating Node HTTP endpoint");
 
   const yoga = createYoga({
     ...commonConfig,
-    graphqlEndpoint: endpoint,
+    graphqlEndpoint: options.endpoint,
   });
 
   return yoga;