diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index 9f10df0b4..bfa9792ee 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -151,8 +151,8 @@ class StateGraph(Graph): self.input = input self.output = output self._add_schema(state_schema) - self._add_schema(input) - self._add_schema(output) + self._add_schema(input, allow_managed=False) + self._add_schema(output, allow_managed=False) self.config_schema = config_schema self.waiting_edges: set[tuple[tuple[str, ...], str]] = set() @@ -162,10 +162,17 @@ class StateGraph(Graph): (start, end) for starts, end in self.waiting_edges for start in starts } - def _add_schema(self, schema: Type[Any]) -> None: + def _add_schema(self, schema: Type[Any], /, allow_managed: bool = True) -> None: if schema not in self.schemas: _warn_invalid_state_schema(schema) channels, managed = _get_channels(schema) + if managed and not allow_managed: + names = ", ".join(managed) + schema_name = getattr(schema, "__name__", "") + raise ValueError( + f"Invalid managed channels detected in {schema_name}: {names}." + " Managed channels are not permitted in Input/Output schema." + ) self.schemas[schema] = {**channels, **managed} for key, channel in channels.items(): if key in self.channels: diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 832266064..d2a31690f 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -208,8 +208,61 @@ def test_state_schema_default_values(kw_only_: bool): "val11", } - assert set(json_schema.get("required", set())) == expected_required - assert ( - set(json_schema["properties"].keys()) - == expected_required | expected_optional - ) + assert set(json_schema.get("required", set())) == expected_required + assert ( + set(json_schema["properties"].keys()) == expected_required | expected_optional + ) + + +def test_raises_invalid_managed(): + class BadInputState(TypedDict): + some_thing: str + some_input_channel: Annotated[str, SharedValue.on("assistant_id")] + + class InputState(TypedDict): + some_thing: str + some_input_channel: str + + class BadOutputState(TypedDict): + some_thing: str + some_output_channel: Annotated[str, SharedValue.on("assistant_id")] + + class OutputState(TypedDict): + some_thing: str + some_output_channel: str + + class State(TypedDict): + some_thing: str + some_channel: Annotated[str, SharedValue.on("assistant_id")] + + # All OK + StateGraph(State, input=InputState, output=OutputState) + StateGraph(State) + StateGraph(State, input=State, output=State) + StateGraph(State, input=InputState) + StateGraph(State, input=InputState) + + bad_input_examples = [ + (State, BadInputState, OutputState), + (State, BadInputState, BadOutputState), + (State, BadInputState, State), + (State, BadInputState, None), + ] + for _state, _inp, _outp in bad_input_examples: + with pytest.raises( + ValueError, + match="Invalid managed channels detected in BadInputState: some_input_channel. Managed channels are not permitted in Input/Output schema.", + ): + StateGraph(_state, input=_inp, output=_outp) + bad_output_examples = [ + (State, InputState, BadOutputState), + (None, InputState, BadOutputState), + (None, State, BadOutputState), + (State, None, BadOutputState), + ] + for _state, _inp, _outp in bad_output_examples: + with pytest.raises( + ValueError, + match="Invalid managed channels detected in BadOutputState: some_output_channel. Managed channels are not permitted in Input/Output schema.", + ): + StateGraph(_state, input=_inp, output=_outp) diff --git a/libs/sdk-js/src/client.mts b/libs/sdk-js/src/client.mts index 459e9e6b1..9081749af 100644 --- a/libs/sdk-js/src/client.mts +++ b/libs/sdk-js/src/client.mts @@ -654,6 +654,29 @@ export class RunsClient extends BaseClient { }); } + /** + * Create a batch of stateless background runs. + * + * @param payloads An array of payloads for creating runs. + * @returns An array of created runs. + */ + async createBatch( + payloads: (RunsCreatePayload & { assistantId: string })[], + ): Promise { + const filteredPayloads = payloads + .map((payload) => ({ ...payload, assistant_id: payload.assistantId })) + .map((payload) => { + return Object.fromEntries( + Object.entries(payload).filter(([_, v]) => v !== undefined), + ); + }); + + return this.fetch("/runs/batch", { + method: "POST", + json: filteredPayloads, + }); + } + async wait( threadId: null, assistantId: string, @@ -775,6 +798,71 @@ export class RunsClient extends BaseClient { return this.fetch(`/threads/${threadId}/runs/${runId}/join`); } + /** + * Stream output from a run in real-time, until the run is done. + * Output is not buffered, so any output produced before this call will + * not be received here. + * + * @param threadId The ID of the thread. + * @param runId The ID of the run. + * @param signal An optional abort signal. + * @returns An async generator yielding stream parts. + */ + async *joinStream( + threadId: string, + runId: string, + signal?: AbortSignal, + ): AsyncGenerator<{ event: StreamEvent; data: any }> { + const response = await this.asyncCaller.fetch( + ...this.prepareFetchOptions(`/threads/${threadId}/runs/${runId}/stream`, { + method: "GET", + signal, + }), + ); + + let parser: EventSourceParser; + let onEndEvent: () => void; + const textDecoder = new TextDecoder(); + + const stream: ReadableStream<{ event: string; data: any }> = ( + response.body || new ReadableStream({ start: (ctrl) => ctrl.close() }) + ).pipeThrough( + new TransformStream({ + async start(ctrl) { + parser = createParser((event) => { + if ( + (signal && signal.aborted) || + (event.type === "event" && event.data === "[DONE]") + ) { + ctrl.terminate(); + return; + } + + if ("data" in event) { + ctrl.enqueue({ + event: event.event ?? "message", + data: JSON.parse(event.data), + }); + } + }); + onEndEvent = () => { + ctrl.enqueue({ event: "end", data: undefined }); + }; + }, + async transform(chunk) { + const payload = textDecoder.decode(chunk); + parser.feed(payload); + + // eventsource-parser will ignore events + // that are not terminated by a newline + if (payload.trim() === "event: end") onEndEvent(); + }, + }), + ); + + yield* IterableReadableStream.fromReadableStream(stream); + } + /** * Delete a run. * diff --git a/libs/sdk-py/langgraph_sdk/client.py b/libs/sdk-py/langgraph_sdk/client.py index 5d6779c5d..1a7766a92 100644 --- a/libs/sdk-py/langgraph_sdk/client.py +++ b/libs/sdk-py/langgraph_sdk/client.py @@ -1234,7 +1234,7 @@ class RunsClient: return await self.http.post("/runs", json=payload) async def create_batch(self, payloads: list[RunCreate]) -> list[Run]: - """Create a batch of background runs.""" + """Create a batch of stateless background runs.""" def filter_payload(payload: RunCreate): return {k: v for k, v in payload.items() if v is not None} @@ -1484,7 +1484,7 @@ class RunsClient: Example Usage: - await client.runs.join( + await client.runs.join_stream( thread_id="thread_id_to_join", run_id="run_id_to_join" )