diff --git a/CopilotKit/.changeset/metal-actors-double.md b/CopilotKit/.changeset/metal-actors-double.md new file mode 100644 index 0000000000..82ee0f84a5 --- /dev/null +++ b/CopilotKit/.changeset/metal-actors-double.md @@ -0,0 +1,5 @@ +--- +"@copilotkit/runtime": patch +--- + +- feat: add support for langgraph config schema diff --git a/CopilotKit/packages/runtime/src/lib/runtime/remote-lg-action.ts b/CopilotKit/packages/runtime/src/lib/runtime/remote-lg-action.ts index 681a49e4ef..fb6c9a5a90 100644 --- a/CopilotKit/packages/runtime/src/lib/runtime/remote-lg-action.ts +++ b/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, schemaKeys: string[]) { + return Object.fromEntries(Object.entries(obj).filter(([key]) => schemaKeys.includes(key))); +} diff --git a/docs/content/docs/coagents/advanced/adding-runtime-configuration.mdx b/docs/content/docs/coagents/advanced/adding-runtime-configuration.mdx index caf7f4fdfb..2327e61b4f 100644 --- a/docs/content/docs/coagents/advanced/adding-runtime-configuration.mdx +++ b/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 + + ### 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] + + + ```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) + ``` + + + ```typescript + import { Annotation } from "@langchain/langgraph"; + + // define which properties will be allowed in the configuration + export const ConfigSchemaAnnotation = Annotation.Root({ + authToken: Annotation + }) + + // ...add all necessary graph nodes + + // when defining the state graph, apply the config schema + const workflow = new StateGraph(AgentStateAnnotation, ConfigSchemaAnnotation) + ``` + + + diff --git a/sdk-python/copilotkit/langgraph_agent.py b/sdk-python/copilotkit/langgraph_agent.py index 0b4671122e..8b93f6f3d4 100644 --- a/sdk-python/copilotkit/langgraph_agent.py +++ b/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 diff --git a/sdk-python/copilotkit/utils.py b/sdk-python/copilotkit/utils.py new file mode 100644 index 0000000000..e123e558c7 --- /dev/null +++ b/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 \ No newline at end of file