feat: support input and output schema of langgraph (CopilotKit#1406)

Author: ranst91

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

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,
+  };
+}

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"
     ]
 }

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
+
+<Callout>
+    Due to LangGraph support, this feature is currently only available for Python agents
+</Callout>
+
+<Steps>
+  <Step>
+    ### 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]
+    ```
+  </Step>
+  <Step>
+    ### 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()
+      ```
+  </Step>
+  <Step>
+    ### 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<AgentState>({
+      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
+    ```
+  </Step>
+</Steps>

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 @@ def _emit_state_sync_event(
             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 @@ def _emit_state_sync_event(
                 "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 @@ async def _stream_events( # pylint: disable=too-many-locals
         # 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 @@ def dict_repr(self):
             '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