From fea916fc9203f8fe5db36bf92d71d84bfc2dce90 Mon Sep 17 00:00:00 2001 From: Ran Shemtov Date: Mon, 24 Feb 2025 12:04:40 +0100 Subject: [PATCH] feat: support input and output schema of langgraph (#1406) --- CopilotKit/.changeset/calm-carpets-peel.md | 6 ++ .../src/lib/runtime/remote-lg-action.ts | 34 ++++++- .../docs/coagents/shared-state/meta.json | 1 + .../shared-state/state-inputs-outputs.mdx | 98 +++++++++++++++++++ sdk-python/copilotkit/langgraph_agent.py | 46 ++++++++- 5 files changed, 183 insertions(+), 2 deletions(-) create mode 100644 CopilotKit/.changeset/calm-carpets-peel.md create mode 100644 docs/content/docs/coagents/shared-state/state-inputs-outputs.mdx diff --git a/CopilotKit/.changeset/calm-carpets-peel.md b/CopilotKit/.changeset/calm-carpets-peel.md new file mode 100644 index 0000000000..3018a4fe2c --- /dev/null +++ b/CopilotKit/.changeset/calm-carpets-peel.md @@ -0,0 +1,6 @@ +--- +"@copilotkit/runtime": minor +--- + +- feat: support input and output schema of langgraph +- docs: add input output schema docs 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 cc5c9be05a..53e74b7c27 100644 --- a/CopilotKit/packages/runtime/src/lib/runtime/remote-lg-action.ts +++ b/CopilotKit/packages/runtime/src/lib/runtime/remote-lg-action.ts @@ -1,4 +1,4 @@ -import { Client as LangGraphClient } from "@langchain/langgraph-sdk"; +import { AssistantGraph, Client as LangGraphClient, GraphSchema } from "@langchain/langgraph-sdk"; import { createHash } from "node:crypto"; import { isValidUUID, randomUUID } from "@copilotkit/shared"; import { parse as parsePartialJson } from "partial-json"; @@ -204,6 +204,15 @@ async function streamEvents(controller: ReadableStreamDefaultController, args: E 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); + + // 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)), + ); + } let streamingStateExtractor = new StreamingStateExtractor([]); let prevNodeName = null; @@ -330,6 +339,7 @@ async function streamEvents(controller: ReadableStreamDefaultController, args: E state: manuallyEmittedState, running: true, active: true, + schemaKeys, }), ); continue; @@ -380,6 +390,7 @@ async function streamEvents(controller: ReadableStreamDefaultController, args: E state, running: true, active: !exitingNode, + schemaKeys, }), ); } @@ -404,6 +415,7 @@ async function streamEvents(controller: ReadableStreamDefaultController, args: E running: !shouldExit, active: false, includeMessages: true, + schemaKeys, }), ); @@ -427,6 +439,7 @@ function getStateSyncEvent({ running, active, includeMessages = false, + schemaKeys, }: { threadId: string; runId: string; @@ -436,6 +449,7 @@ function getStateSyncEvent({ running: boolean; active: boolean; includeMessages?: boolean; + schemaKeys: { input: string[] | null; output: string[] | null }; }): string { if (!includeMessages) { state = Object.keys(state).reduce((acc, key) => { @@ -451,6 +465,13 @@ 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)), + ); + } + return ( JSON.stringify({ event: LangGraphEventTypes.OnCopilotKitStateSync, @@ -739,3 +760,14 @@ function copilotkitMessagesToLangChain(messages: Message[]): LangGraphPlatformMe return result; } + +function getSchemaKeys(graphSchema: GraphSchema) { + const CONSTANT_KEYS = ["messages", "copilotkit"]; + const inputSchema = Object.keys(graphSchema.input_schema.properties); + const outputSchema = Object.keys(graphSchema.output_schema.properties); + + return { + input: inputSchema && inputSchema.length ? [...inputSchema, ...CONSTANT_KEYS] : null, + output: outputSchema && outputSchema.length ? [...outputSchema, ...CONSTANT_KEYS] : null, + }; +} diff --git a/docs/content/docs/coagents/shared-state/meta.json b/docs/content/docs/coagents/shared-state/meta.json index 139322c59b..35d28d1b5c 100644 --- a/docs/content/docs/coagents/shared-state/meta.json +++ b/docs/content/docs/coagents/shared-state/meta.json @@ -2,6 +2,7 @@ "pages": [ "in-app-agent-read", "in-app-agent-write", + "state-inputs-outputs", "predictive-state-updates" ] } \ No newline at end of file diff --git a/docs/content/docs/coagents/shared-state/state-inputs-outputs.mdx b/docs/content/docs/coagents/shared-state/state-inputs-outputs.mdx new file mode 100644 index 0000000000..08d8612c29 --- /dev/null +++ b/docs/content/docs/coagents/shared-state/state-inputs-outputs.mdx @@ -0,0 +1,98 @@ +--- +title: Agent state inputs and outputs +icon: "lucide/ArrowRightLeft" +description: Decide which state properties are received and returned to the frontend +--- + +## What is this? + +Not all state properties are relevant for frontend-backend sharing. +This guide shows how to ensure only the right portion of state is communicated back and forth. + +This guide is based on [LangGraph's Input/Output Schema feature](https://langchain-ai.github.io/langgraph/how-tos/input_output_schema/) + +## When should I use this? + +Depending on your implementation, some properties are meant to be processed internally, while some others are the way for the UI to communicate user input. +In addition, some state properties contain a lot of information. Syncing them back and forth between the agent and UI can be costly, while it might not have any practical benefit. + +## Implementation + + + Due to LangGraph support, this feature is currently only available for Python agents + + + + + ### Examine our old state + LangGraph is stateful. As you transition between nodes, that state is updated and passed to the next node. For this example, + let's assume that the state our agent should be using, can be described like this: + ```python title="agent-py/sample_agent/agent.py" + from copilotkit import CopilotKitState + from typing import Literal + + class AgentState(CopilotKitState): + question: str + answer: str + resources: List[str] + ``` + + + ### Divide state to Input and Output + Our example case lists several state properties, which with its own purpose: + - The question is being asked by the user, expecting the llm to answer + - The answer is what the LLM returns + - The resources list will be used by the LLM to answer the question, and should not be communicated to the user, or set by them. + + ```python title="agent-py/sample_agent/agent.py" + from copilotkit import CopilotKitState + from typing import Literal + + # divide the state into 3 + class InputState(CopilotKitState): + question: str + + class OutputState(CopilotKitState): + answer: str + + class OverallState(InputState, OutputState): + resources: List[str] + + # ...add the rest of the agent implementation + + # finally, before compiling the graph, we define the 3 state components + builder = StateGraph(OverallState, input=InputState, output=OutputState) + + # add all the different nodes and edges and compile the graph + builder.add_node(answer_node) + builder.add_edge(START, "answer_node") + builder.add_edge("answer_node", END) + graph = builder.compile() + ``` + + + ### Give it a try! + Now that we know which state properties our agent emits, we can inspect the state and expect the following to happen: + - While we are able to provide a question, we will not receive it back from the agent. If we are using it in our UI, we need to remember the UI is the source of truth for it + - Answer will change once it's returned back from the agent + - The UI has no access to resources. + + ```tsx + import { useCoAgent } from "@copilotkit/react-core"; + + type AgentState = { + question: string; + answer: string; + } + + const { state } = useCoAgent({ + name: "sample_agent", + initialState: { + question: "How's is the weather in SF?", + } + }); + + console.log(state) // You can expect seeing "answer" change, while the others are not returned from the agent + ``` + + diff --git a/sdk-python/copilotkit/langgraph_agent.py b/sdk-python/copilotkit/langgraph_agent.py index 5af02c9fe6..c60b127fca 100644 --- a/sdk-python/copilotkit/langgraph_agent.py +++ b/sdk-python/copilotkit/langgraph_agent.py @@ -2,7 +2,7 @@ import uuid import json -from typing import Optional, List, Callable, Any, cast, Union, TypedDict +from typing import Optional, List, Callable, Any, cast, Union, TypedDict, Literal from typing_extensions import NotRequired from langgraph.graph.graph import CompiledGraph @@ -191,6 +191,7 @@ class LangGraphAgent(Agent): active: bool, include_messages: bool = False ): + # First handle messages as before if not include_messages: state = { k: v for k, v in state.items() if k != "messages" @@ -201,6 +202,9 @@ class LangGraphAgent(Agent): "messages": langchain_messages_to_copilotkit(state.get("messages", [])) } + # Filter by schema keys if available + state = self.filter_state_on_schema_keys(state, 'output') + return langchain_dumps({ "event": "on_copilotkit_state_sync", "thread_id": thread_id, @@ -292,6 +296,12 @@ class LangGraphAgent(Agent): # Use provided input or fallback to initial_state 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) + self.output_schema_keys = output_keys + self.input_schema_keys = input_keys + + stream_input = self.filter_state_on_schema_keys(stream_input, 'input') async for event in self.graph.astream_events(stream_input, config, version="v2"): current_node_name = event.get("name") event_type = event.get("event") @@ -457,6 +467,40 @@ class LangGraphAgent(Agent): 'type': 'langgraph' } + def get_schema_keys(self, config): + CONSTANT_KEYS = ['copilotkit', 'messages'] + try: + graph = self.graph.get_graph(config) + end_node = graph.nodes["__end__"] + start_node = graph.nodes["__start__"] + input_schema = start_node.data.schema() + output_schema = end_node.data.schema() + input_schema_keys = list(input_schema["properties"].keys()) + output_schema_keys = list(output_schema["properties"].keys()) + + # We add "copilotkit" and "messages" as they are always sent and received. + for key in CONSTANT_KEYS: + if key not in input_schema_keys: + input_schema_keys.append(key) + if key not in output_schema_keys: + output_schema_keys.append(key) + + return input_schema_keys, output_schema_keys + except Exception: + return None + + def filter_state_on_schema_keys(self, state, schema_type: Literal["input", "output"]): + 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" + } + except Exception: + return state + + class _StreamingStateExtractor: def __init__(self, emit_intermediate_state: List[dict]): self.emit_intermediate_state = emit_intermediate_state