feat: add support for langgraph config schema (#1596)

CopilotKit/.changeset/metal-actors-double.md

@@ -0,0 +1,5 @@
+---
+"@copilotkit/runtime": patch
+---
+
+- feat: add support for langgraph config schema

CopilotKit/packages/runtime/src/lib/runtime/remote-lg-action.ts

@@ -68,7 +68,11 @@ type LangGraphPlatformMessage =
   | LangGraphPlatformResultMessage
   | BaseLangGraphPlatformMessage;
 
-type SchemaKeys = { input: string[] | null; output: string[] | null } | null;
+type SchemaKeys = {
+  input: string[] | null;
+  output: string[] | null;
+  config: string[] | null;
+} | null;
 
 let activeInterruptEvent = false;
 
@@ -202,17 +206,19 @@ async function streamEvents(controller: ReadableStreamDefaultController, args: E
   }
   const assistantId = retrievedAssistant.assistant_id;
 
-  if (configurable) {
-    await client.assistants.update(assistantId, { config: { configurable } });
-  }
   const graphInfo = await client.assistants.getGraph(assistantId);
   const graphSchema = await client.assistants.getSchemas(assistantId);
   const schemaKeys = getSchemaKeys(graphSchema);
+  if (configurable) {
+    const filteredConfigurable = schemaKeys?.config
+      ? filterObjectBySchemaKeys(configurable, schemaKeys?.config)
+      : configurable;
+    await client.assistants.update(assistantId, { config: { configurable: filteredConfigurable } });
+  }
+
   // Do not input keys that are not part of the input schema
   if (payload.input && schemaKeys?.input) {
-    payload.input = Object.fromEntries(
-      Object.entries(payload.input).filter(([key]) => schemaKeys.input.includes(key)),
-    );
+    payload.input = filterObjectBySchemaKeys(payload.input, schemaKeys.input);
   }
 
   let streamingStateExtractor = new StreamingStateExtractor([]);
@@ -468,9 +474,7 @@ function getStateSyncEvent({
 
   // Do not emit state keys that are not part of the output schema
   if (schemaKeys?.output) {
-    state = Object.fromEntries(
-      Object.entries(state).filter(([key]) => schemaKeys.output.includes(key)),
-    );
+    state = filterObjectBySchemaKeys(state, schemaKeys.output);
   }
 
   return (
@@ -764,8 +768,12 @@ function copilotkitMessagesToLangChain(messages: Message[]): LangGraphPlatformMe
 
 function getSchemaKeys(graphSchema: GraphSchema): SchemaKeys {
   const CONSTANT_KEYS = ["messages", "copilotkit"];
+  let configSchema = null;
+  if (graphSchema.config_schema.properties) {
+    configSchema = Object.keys(graphSchema.config_schema.properties);
+  }
   if (!graphSchema.input_schema.properties || !graphSchema.output_schema.properties) {
-    return null;
+    return configSchema;
   }
   const inputSchema = Object.keys(graphSchema.input_schema.properties);
   const outputSchema = Object.keys(graphSchema.output_schema.properties);
@@ -773,5 +781,10 @@ function getSchemaKeys(graphSchema: GraphSchema): SchemaKeys {
   return {
     input: inputSchema && inputSchema.length ? [...inputSchema, ...CONSTANT_KEYS] : null,
     output: outputSchema && outputSchema.length ? [...outputSchema, ...CONSTANT_KEYS] : null,
+    config: configSchema,
   };
 }
+
+function filterObjectBySchemaKeys(obj: Record<string, any>, schemaKeys: string[]) {
+  return Object.fromEntries(Object.entries(obj).filter(([key]) => schemaKeys.includes(key)));
+}

docs/content/docs/coagents/advanced/adding-runtime-configuration.mdx

@@ -68,4 +68,42 @@ Now you can simply pull the values from the provided config argument in any agen
     </Tab>
 </Tabs>
 </Step>
+<Step>
+    ### Optional: Define configurables schema
+    If you'd like, you can define a schema to indicate which configurables you wish to receive.
+    Any item passed to "configurables" which is not included in the schema, will be filtered out.
+
+    You can read more about this (here)[https://langchain-ai.github.io/langgraph/how-tos/configuration/#define-graph]
+    <Tabs groupId="language" items={['Python', 'TypeScript']} default="Python">
+        <Tab value="Python">
+            ```python
+            from typing import TypedDict
+
+            # define which properties will be allowed in the configuration
+            class ConfigSchema(TypedDict):
+              authToken: str
+
+            # ...add all necessary graph nodes
+
+            # when defining the state graph, apply the config schema
+            workflow = StateGraph(AgentState, config_schema=ConfigSchema)
+            ```
+        </Tab>
+        <Tab value="TypeScript">
+            ```typescript
+            import { Annotation } from "@langchain/langgraph";
+
+            // define which properties will be allowed in the configuration
+            export const ConfigSchemaAnnotation = Annotation.Root({
+              authToken: Annotation<string>
+            })
+
+            // ...add all necessary graph nodes
+
+            // when defining the state graph, apply the config schema
+            const workflow = new StateGraph(AgentStateAnnotation, ConfigSchemaAnnotation)
+            ```
+        </Tab>
+    </Tabs>
+</Step>
 </Steps>

sdk-python/copilotkit/langgraph_agent.py

@@ -14,6 +14,7 @@ from langchain_core.runnables import RunnableConfig, ensure_config
 from partialjson.json_parser import JSONParser
 
 from .types import Message, MetaEvent
+from .utils import filter_by_schema_keys
 from .langgraph import copilotkit_messages_to_langchain, langchain_messages_to_copilotkit
 from .action import ActionDict
 from .agent import Agent
@@ -299,11 +300,12 @@ class LangGraphAgent(Agent):
         stream_input = resume_input if resume_input else initial_state
 
         # Get the output and input schema keys the user has allowed for this graph
-        input_keys, output_keys = self.get_schema_keys(config)
+        input_keys, output_keys, config_keys = self.get_schema_keys(config)
         self.output_schema_keys = output_keys
         self.input_schema_keys = input_keys
 
         stream_input = self.filter_state_on_schema_keys(stream_input, 'input')
+        config["configurable"] = filter_by_schema_keys(config["configurable"], config_keys)
         async for event in self.graph.astream_events(stream_input, config, version="v2"):
             current_node_name = event.get("name")
             event_type = event.get("event")
@@ -478,12 +480,24 @@ class LangGraphAgent(Agent):
 
     def get_schema_keys(self, config):
         CONSTANT_KEYS = ['copilotkit', 'messages']
+        CONSTANT_CONFIG_KEYS = ['checkpoint_id', 'checkpoint_ns', 'thread_id']
         try:
             input_schema = self.graph.get_input_jsonschema(config)
             output_schema = self.graph.get_output_jsonschema(config)
             input_schema_keys = list(input_schema["properties"].keys())
             output_schema_keys = list(output_schema["properties"].keys())
 
+            try:
+                schema_dict = self.graph.config_schema().schema()
+                configurable_schema = schema_dict["$defs"]["Configurable"]
+                config_schema_keys = list(configurable_schema["properties"].keys())
+
+                # If only constant keys are present, it means no schema was passed, we allow everything
+                if set(config_schema_keys) == set(CONSTANT_CONFIG_KEYS):
+                    config_schema_keys = None
+            except:
+                config_schema_keys = None
+
             # We add "copilotkit" and "messages" as they are always sent and received.
             for key in CONSTANT_KEYS:
                 if key not in input_schema_keys:
@@ -491,7 +505,7 @@ class LangGraphAgent(Agent):
                 if key not in output_schema_keys:
                     output_schema_keys.append(key)
 
-            return input_schema_keys, output_schema_keys
+            return input_schema_keys, output_schema_keys, config_schema_keys
         except Exception:
             return None
 
@@ -499,10 +513,7 @@ class LangGraphAgent(Agent):
         try:
             schema_keys_name = f"{schema_type}_schema_keys"
             if hasattr(self, schema_keys_name) and getattr(self, schema_keys_name):
-                return {
-                    k: v for k, v in state.items()
-                    if k in getattr(self, schema_keys_name) or k == "messages"
-                }
+                return filter_by_schema_keys(state, getattr(self, schema_keys_name))
         except Exception:
             return state
 

sdk-python/copilotkit/utils.py

@@ -0,0 +1,8 @@
+def filter_by_schema_keys(obj, schema):
+    try:
+        return {
+            k: v for k, v in obj.items()
+            if k in schema or k == "messages"
+        }
+    except Exception:
+        return obj