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