ag-ui-langgraph 0.0.1__tar.gz

ag_ui_langgraph-0.0.1/PKG-INFO

@@ -0,0 +1,21 @@
+Metadata-Version: 2.3
+Name: ag-ui-langgraph
+Version: 0.0.1
+Summary: 
+Author: Ran Shem Tov
+Author-email: ran@copilotkit.ai
+Requires-Python: >=3.10,<3.14
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Provides-Extra: fastapi
+Requires-Dist: ag-ui-protocol (==0.1.7)
+Requires-Dist: fastapi (>=0.115.12,<0.116.0) ; extra == "fastapi"
+Requires-Dist: langchain (>=0.3.0)
+Requires-Dist: langchain-core (>=0.3.0)
+Requires-Dist: langgraph (>=0.3.25,<=0.5.0)
+Description-Content-Type: text/markdown
+
+

ag_ui_langgraph-0.0.1/ag_ui_langgraph/__init__.py

@@ -0,0 +1,35 @@
+from .agent import LangGraphAgent
+from .types import (
+    LangGraphEventTypes,
+    CustomEventNames,
+    State,
+    SchemaKeys,
+    MessageInProgress,
+    RunMetadata,
+    MessagesInProgressRecord,
+    ToolCall,
+    BaseLangGraphPlatformMessage,
+    LangGraphPlatformResultMessage,
+    LangGraphPlatformActionExecutionMessage,
+    LangGraphPlatformMessage,
+    PredictStateTool
+)
+from .endpoint import add_langgraph_fastapi_endpoint
+
+__all__ = [
+    "LangGraphAgent",
+    "LangGraphEventTypes",
+    "CustomEventNames",
+    "State",
+    "SchemaKeys",
+    "MessageInProgress",
+    "RunMetadata",
+    "MessagesInProgressRecord",
+    "ToolCall",
+    "BaseLangGraphPlatformMessage",
+    "LangGraphPlatformResultMessage",
+    "LangGraphPlatformActionExecutionMessage",
+    "LangGraphPlatformMessage",
+    "PredictStateTool",
+    "add_langgraph_fastapi_endpoint"
+]

ag_ui_langgraph-0.0.1/ag_ui_langgraph/agent.py

@@ -0,0 +1,682 @@
+import uuid
+import json
+from typing import Optional, List, Any, Union, AsyncGenerator, Generator
+
+from langgraph.graph.state import CompiledStateGraph
+from langchain.schema import BaseMessage, SystemMessage
+from langchain_core.runnables import RunnableConfig, ensure_config
+from langchain_core.messages import HumanMessage
+from langgraph.types import Command
+
+from .types import (
+    State,
+    LangGraphPlatformMessage,
+    MessagesInProgressRecord,
+    SchemaKeys,
+    MessageInProgress,
+    RunMetadata,
+    LangGraphEventTypes,
+    CustomEventNames,
+    LangGraphReasoning
+)
+from .utils import (
+    agui_messages_to_langchain,
+    DEFAULT_SCHEMA_KEYS,
+    filter_object_by_schema_keys,
+    get_stream_payload_input,
+    langchain_messages_to_agui,
+    resolve_reasoning_content,
+    resolve_message_content,
+    camel_to_snake
+)
+
+from ag_ui.core import (
+    EventType,
+    CustomEvent,
+    MessagesSnapshotEvent,
+    RawEvent,
+    RunAgentInput,
+    RunErrorEvent,
+    RunFinishedEvent,
+    RunStartedEvent,
+    StateDeltaEvent,
+    StateSnapshotEvent,
+    StepFinishedEvent,
+    StepStartedEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    TextMessageStartEvent,
+    ToolCallArgsEvent,
+    ToolCallEndEvent,
+    ToolCallStartEvent,
+    ThinkingTextMessageStartEvent,
+    ThinkingTextMessageContentEvent,
+    ThinkingTextMessageEndEvent,
+    ThinkingStartEvent,
+    ThinkingEndEvent,
+)
+from ag_ui.encoder import EventEncoder
+
+ProcessedEvents = Union[
+    TextMessageStartEvent,
+    TextMessageContentEvent,
+    TextMessageEndEvent,
+    ToolCallStartEvent,
+    ToolCallArgsEvent,
+    ToolCallEndEvent,
+    StateSnapshotEvent,
+    StateDeltaEvent,
+    MessagesSnapshotEvent,
+    RawEvent,
+    CustomEvent,
+    RunStartedEvent,
+    RunFinishedEvent,
+    RunErrorEvent,
+    StepStartedEvent,
+    StepFinishedEvent,
+]
+
+class LangGraphAgent:
+    def __init__(self, *, name: str, graph: CompiledStateGraph, description: Optional[str] = None, config:  Union[Optional[RunnableConfig], dict] = None):
+        self.name = name
+        self.description = description
+        self.graph = graph
+        self.config = config or {}
+        self.messages_in_process: MessagesInProgressRecord = {}
+        self.active_run: Optional[RunMetadata] = None
+        self.constant_schema_keys = ['messages', 'tools']
+
+    def _dispatch_event(self, event: ProcessedEvents) -> str:
+        return event  # Fallback if no encoder
+
+    async def run(self, input: RunAgentInput) -> AsyncGenerator[str, None]:
+        forwarded_props = {}
+        if hasattr(input, "forwarded_props") and input.forwarded_props:
+            forwarded_props = {
+                camel_to_snake(k): v for k, v in input.forwarded_props.items()
+            }
+        async for event_str in self._handle_stream_events(input.copy(update={"forwarded_props": forwarded_props})):
+            yield event_str
+
+    async def _handle_stream_events(self, input: RunAgentInput) -> AsyncGenerator[str, None]:
+        thread_id = input.thread_id or str(uuid.uuid4())
+        self.active_run = {
+            "id": input.run_id,
+            "thread_id": thread_id,
+            "thinking_process": None,
+        }
+
+        messages = input.messages or []
+        forwarded_props = input.forwarded_props
+        node_name_input = forwarded_props.get('node_name', None) if forwarded_props else None
+
+        self.active_run["manually_emitted_state"] = None
+        self.active_run["node_name"] = node_name_input
+        if self.active_run["node_name"] == "__end__":
+            self.active_run["node_name"] = None
+
+        config = ensure_config(self.config.copy() if self.config else {})
+        config["configurable"] = {**(config.get('configurable', {})), "thread_id": thread_id}
+
+        agent_state = await self.graph.aget_state(config)
+        self.active_run["mode"] = "continue" if thread_id and self.active_run.get("node_name") != "__end__" and self.active_run.get("node_name") else "start"
+        prepared_stream_response = await self.prepare_stream(input=input, agent_state=agent_state, config=config)
+
+        yield self._dispatch_event(
+            RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=self.active_run["id"])
+        )
+
+        langchain_messages = agui_messages_to_langchain(messages)
+        non_system_messages = [msg for msg in langchain_messages if not isinstance(msg, SystemMessage)]
+
+        if len(agent_state.values.get("messages", [])) > len(non_system_messages):
+            # Find the last user message by working backwards from the last message
+            last_user_message = None
+            for i in range(len(langchain_messages) - 1, -1, -1):
+                if isinstance(langchain_messages[i], HumanMessage):
+                    last_user_message = langchain_messages[i]
+                    break
+
+            if last_user_message:
+                prepared_stream_response = await self.prepare_regenerate_stream(
+                    input=input,
+                    message_checkpoint=last_user_message,
+                    config=config
+                )
+
+        state = prepared_stream_response["state"]
+        stream = prepared_stream_response["stream"]
+        config = prepared_stream_response["config"]
+        events_to_dispatch = prepared_stream_response.get('events_to_dispatch', None)
+
+        if events_to_dispatch is not None and len(events_to_dispatch) > 0:
+            for event in events_to_dispatch:
+                yield self._dispatch_event(event)
+            return
+
+        should_exit = False
+        current_graph_state = state
+        async for event in stream:
+            if event["event"] == "error":
+                yield self._dispatch_event(
+                    RunErrorEvent(type=EventType.RUN_ERROR, message=event["data"]["message"], raw_event=event)
+                )
+                break
+
+            current_node_name = event.get("metadata", {}).get("langgraph_node")
+            event_type = event.get("event")
+            self.active_run["id"] = event.get("run_id")
+            exiting_node = False
+
+            if event_type == "on_chain_end" and isinstance(
+                    event.get("data", {}).get("output"), dict
+            ):
+                current_graph_state.update(event["data"]["output"])
+                exiting_node = self.active_run["node_name"] == current_node_name
+
+            should_exit = should_exit or (
+                    event_type == "on_custom_event" and
+                    event["name"] == "exit"
+                )
+
+            if current_node_name and current_node_name != self.active_run.get("node_name"):
+                if self.active_run["node_name"] and self.active_run["node_name"] != node_name_input:
+                    yield self._dispatch_event(
+                        StepFinishedEvent(type=EventType.STEP_FINISHED, step_name=self.active_run["node_name"])
+                    )
+                    self.active_run["node_name"] = None
+
+                yield self._dispatch_event(
+                    StepStartedEvent(type=EventType.STEP_STARTED, step_name=current_node_name)
+                )
+                self.active_run["node_name"] = current_node_name
+
+            updated_state = self.active_run.get("manually_emitted_state") or current_graph_state
+            has_state_diff = updated_state != state
+            if exiting_node or (has_state_diff and not self.get_message_in_progress(self.active_run["id"])):
+                state = updated_state
+                self.active_run["prev_node_name"] = self.active_run["node_name"]
+                current_graph_state.update(updated_state)
+                yield self._dispatch_event(
+                    StateSnapshotEvent(
+                        type=EventType.STATE_SNAPSHOT,
+                        snapshot=self.get_state_snapshot(state),
+                        raw_event=event,
+                    )
+                )
+
+            yield self._dispatch_event(
+                RawEvent(type=EventType.RAW, event=event)
+            )
+
+            async for single_event in self._handle_single_event(event, state):
+                yield single_event
+
+        state = await self.graph.aget_state(config)
+
+        tasks = state.tasks if len(state.tasks) > 0 else None
+        interrupts = tasks[0].interrupts if tasks else []
+
+        writes = state.metadata.get("writes", {}) or {}
+        node_name = self.active_run["node_name"] if interrupts else next(iter(writes), None)
+        next_nodes = state.next or ()
+        is_end_node = len(next_nodes) == 0 and not interrupts
+
+        node_name = "__end__" if is_end_node else node_name
+
+        for interrupt in interrupts:
+            yield self._dispatch_event(
+                CustomEvent(
+                    type=EventType.CUSTOM,
+                    name=LangGraphEventTypes.OnInterrupt.value,
+                    value=json.dumps(interrupt.value) if not isinstance(interrupt.value, str) else interrupt.value,
+                    raw_event=interrupt,
+                )
+            )
+
+        if self.active_run.get("node_name") != node_name:
+            yield self._dispatch_event(
+                StepFinishedEvent(type=EventType.STEP_FINISHED, step_name=self.active_run["node_name"])
+            )
+            self.active_run["node_name"] = node_name
+            yield self._dispatch_event(
+                StepStartedEvent(type=EventType.STEP_STARTED, step_name=self.active_run["node_name"])
+            )
+
+        state_values = state.values if state.values else state
+        yield self._dispatch_event(
+            StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=self.get_state_snapshot(state_values))
+        )
+
+        yield self._dispatch_event(
+            MessagesSnapshotEvent(
+                type=EventType.MESSAGES_SNAPSHOT,
+                messages=langchain_messages_to_agui(state_values.get("messages", [])),
+            )
+        )
+
+        yield self._dispatch_event(
+            StepFinishedEvent(type=EventType.STEP_FINISHED, step_name=self.active_run["node_name"])
+        )
+        self.active_run["node_name"] = None
+
+        yield self._dispatch_event(
+            RunFinishedEvent(type=EventType.RUN_FINISHED, thread_id=thread_id, run_id=self.active_run["id"])
+        )
+        self.active_run = None
+
+
+    async def prepare_stream(self, input: RunAgentInput, agent_state: State, config: RunnableConfig):
+        state_input = input.state or {}
+        messages = input.messages or []
+        tools = input.tools or []
+        forwarded_props = input.forwarded_props or {}
+        thread_id = input.thread_id
+
+        state_input["messages"] = agent_state.values.get("messages", [])
+        self.active_run["current_graph_state"] = agent_state.values
+        langchain_messages = agui_messages_to_langchain(messages)
+        state = self.langgraph_default_merge_state(state_input, langchain_messages, tools)
+        self.active_run["current_graph_state"].update(state)
+        config["configurable"]["thread_id"] = thread_id
+        interrupts = agent_state.tasks[0].interrupts if agent_state.tasks and len(agent_state.tasks) > 0 else []
+        has_active_interrupts = len(interrupts) > 0
+        resume_input = forwarded_props.get('command', {}).get('resume', None)
+
+        events_to_dispatch = []
+        if has_active_interrupts and not resume_input:
+            events_to_dispatch.append(
+                RunStartedEvent(type=EventType.RUN_STARTED, thread_id=thread_id, run_id=self.active_run["id"])
+            )
+
+            for interrupt in interrupts:
+                events_to_dispatch.append(
+                    CustomEvent(
+                        type=EventType.CUSTOM,
+                        name=LangGraphEventTypes.OnInterrupt.value,
+                        value=json.dumps(interrupt.value) if not isinstance(interrupt.value, str) else interrupt.value,
+                        raw_event=interrupt,
+                    )
+                )
+
+            events_to_dispatch.append(
+                RunFinishedEvent(type=EventType.RUN_FINISHED, thread_id=thread_id, run_id=self.active_run["id"])
+            )
+            return {
+                "stream": None,
+                "state": None,
+                "config": None,
+                "events_to_dispatch": events_to_dispatch,
+            }
+
+        if self.active_run["mode"] == "continue":
+            await self.graph.aupdate_state(config, state, as_node=self.active_run.get("node_name"))
+
+        self.active_run["schema_keys"] = self.get_schema_keys(config)
+
+        if resume_input:
+            stream_input = Command(resume=resume_input)
+        else:
+            payload_input = get_stream_payload_input(
+                mode=self.active_run["mode"],
+                state=state,
+                schema_keys=self.active_run["schema_keys"],
+            )
+            stream_input = {**forwarded_props, **payload_input} if payload_input else None
+
+        return {
+            "stream": self.graph.astream_events(stream_input, config, version="v2"),
+            "state": state,
+            "config": config
+        }
+
+    async def prepare_regenerate_stream( # pylint: disable=too-many-arguments
+            self,
+            input: RunAgentInput,
+            message_checkpoint: HumanMessage,
+            config: RunnableConfig
+    ):
+        tools = input.tools or []
+        thread_id = input.thread_id
+
+        time_travel_checkpoint = await self.get_checkpoint_before_message(message_checkpoint.id, thread_id)
+        if time_travel_checkpoint is None:
+            return None
+
+        fork = await self.graph.aupdate_state(
+            time_travel_checkpoint.config,
+            time_travel_checkpoint.values,
+            as_node=time_travel_checkpoint.next[0] if time_travel_checkpoint.next else "__start__"
+        )
+
+        stream_input = self.langgraph_default_merge_state(time_travel_checkpoint.values, [message_checkpoint], tools)
+        stream = self.graph.astream_events(stream_input, fork, version="v2")
+
+        return {
+            "stream": stream,
+            "state": time_travel_checkpoint.values,
+            "config": config
+        }
+
+    def get_message_in_progress(self, run_id: str) -> Optional[MessageInProgress]:
+        return self.messages_in_process.get(run_id)
+
+    def set_message_in_progress(self, run_id: str, data: MessageInProgress):
+        current_message_in_progress = self.messages_in_process.get(run_id, {})
+        self.messages_in_process[run_id] = {
+            **current_message_in_progress,
+            **data,
+        }
+
+    def get_schema_keys(self, config) -> SchemaKeys:
+        try:
+            input_schema = self.graph.get_input_jsonschema(config)
+            output_schema = self.graph.get_output_jsonschema(config)
+            config_schema = self.graph.config_schema().schema()
+
+            input_schema_keys = list(input_schema["properties"].keys()) if "properties" in input_schema else []
+            output_schema_keys = list(output_schema["properties"].keys()) if "properties" in output_schema else []
+            config_schema_keys = list(config_schema["properties"].keys()) if "properties" in config_schema else []
+
+            return {
+                "input": [*input_schema_keys, *self.constant_schema_keys],
+                "output": [*output_schema_keys, *self.constant_schema_keys],
+                "config": config_schema_keys,
+            }
+        except Exception:
+            return {
+                "input": self.constant_schema_keys,
+                "output": self.constant_schema_keys,
+                "config": [],
+            }
+
+    def langgraph_default_merge_state(self, state: State, messages: List[BaseMessage], tools: Any) -> State:
+        if messages and isinstance(messages[0], SystemMessage):
+            messages = messages[1:]
+
+        existing_messages: List[LangGraphPlatformMessage] = state.get("messages", [])
+        existing_message_ids = {msg.id for msg in existing_messages}
+
+        new_messages = [msg for msg in messages if msg.id not in existing_message_ids]
+
+        tools_as_dicts = []
+        if tools:
+            for tool in tools:
+                if hasattr(tool, "model_dump"):
+                    tools_as_dicts.append(tool.model_dump())
+                elif hasattr(tool, "dict"):
+                    tools_as_dicts.append(tool.dict())
+                else:
+                    tools_as_dicts.append(tool)
+
+        return {
+            **state,
+            "messages": new_messages,
+            "tools": [*state.get("tools", []), *tools_as_dicts],
+        }
+
+    def get_state_snapshot(self, state: State) -> State:
+        schema_keys = self.active_run["schema_keys"]
+        if schema_keys and schema_keys.get("output"):
+            state = filter_object_by_schema_keys(state, [*DEFAULT_SCHEMA_KEYS, *schema_keys["output"]])
+        return state
+
+    async def _handle_single_event(self, event: Any, state: State) -> AsyncGenerator[str, None]:
+        event_type = event.get("event")
+        if event_type == LangGraphEventTypes.OnChatModelStream:
+            should_emit_messages = event["metadata"].get("emit-messages", True)
+            should_emit_tool_calls = event["metadata"].get("emit-tool-calls", True)
+
+            if event["data"]["chunk"].response_metadata.get('finish_reason', None):
+                return
+
+            current_stream = self.get_message_in_progress(self.active_run["id"])
+            has_current_stream = bool(current_stream and current_stream.get("id"))
+            tool_call_data = event["data"]["chunk"].tool_call_chunks[0] if event["data"]["chunk"].tool_call_chunks else None
+            predict_state_metadata = event["metadata"].get("predict_state", [])
+            tool_call_used_to_predict_state = False
+            if tool_call_data and tool_call_data.get("name") and predict_state_metadata:
+                tool_call_used_to_predict_state = any(
+                    predict_tool.get("tool") == tool_call_data["name"]
+                    for predict_tool in predict_state_metadata
+                )
+
+            is_tool_call_start_event = not has_current_stream and tool_call_data and tool_call_data.get("name")
+            is_tool_call_args_event = has_current_stream and current_stream.get("tool_call_id") and tool_call_data and tool_call_data.get("args")
+            is_tool_call_end_event = has_current_stream and current_stream.get("tool_call_id") and not tool_call_data
+
+            reasoning_data = resolve_reasoning_content(event["data"]["chunk"]) if event["data"]["chunk"] else None
+            message_content = resolve_message_content(event["data"]["chunk"].content) if event["data"]["chunk"] and event["data"]["chunk"].content else None
+            is_message_content_event = tool_call_data is None and message_content
+            is_message_end_event = has_current_stream and not current_stream.get("tool_call_id") and not is_message_content_event
+
+            if reasoning_data:
+                self.handle_thinking_event(reasoning_data)
+                return
+
+            if reasoning_data is None and self.active_run.get('thinking_process', None) is not None:
+                yield self._dispatch_event(
+                    ThinkingTextMessageEndEvent(
+                        type=EventType.THINKING_TEXT_MESSAGE_END,
+                    )
+                )
+                yield self._dispatch_event(
+                    ThinkingEndEvent(
+                        type=EventType.THINKING_END,
+                    )
+                )
+                self.active_run["thinking_process"] = None
+
+            if tool_call_used_to_predict_state:
+                yield self._dispatch_event(
+                    CustomEvent(
+                        type=EventType.CUSTOM,
+                        name="PredictState",
+                        value=predict_state_metadata,
+                        raw_event=event
+                    )
+                )
+
+            if is_tool_call_end_event:
+                yield self._dispatch_event(
+                    ToolCallEndEvent(type=EventType.TOOL_CALL_END, tool_call_id=current_stream["tool_call_id"], raw_event=event)
+                )
+                self.messages_in_process[self.active_run["id"]] = None
+                return
+
+
+            if is_message_end_event:
+                yield self._dispatch_event(
+                    TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=current_stream["id"], raw_event=event)
+                )
+                self.messages_in_process[self.active_run["id"]] = None
+                return
+
+            if is_tool_call_start_event and should_emit_tool_calls:
+                yield self._dispatch_event(
+                    ToolCallStartEvent(
+                        type=EventType.TOOL_CALL_START,
+                        tool_call_id=tool_call_data["id"],
+                        tool_call_name=tool_call_data["name"],
+                        parent_message_id=event["data"]["chunk"].id,
+                        raw_event=event,
+                    )
+                )
+                self.set_message_in_progress(
+                    self.active_run["id"],
+                    MessageInProgress(id=event["data"]["chunk"].id, tool_call_id=tool_call_data["id"], tool_call_name=tool_call_data["name"])
+                )
+                return
+
+            if is_tool_call_args_event and should_emit_tool_calls:
+                yield self._dispatch_event(
+                    ToolCallArgsEvent(
+                        type=EventType.TOOL_CALL_ARGS,
+                        tool_call_id=current_stream["tool_call_id"],
+                        delta=tool_call_data["args"],
+                        raw_event=event
+                    )
+                )
+                return
+
+            if is_message_content_event and should_emit_messages:
+                if bool(current_stream and current_stream.get("id")) == False:
+                    yield self._dispatch_event(
+                        TextMessageStartEvent(
+                            type=EventType.TEXT_MESSAGE_START,
+                            role="assistant",
+                            message_id=event["data"]["chunk"].id,
+                            raw_event=event,
+                        )
+                    )
+                    self.set_message_in_progress(
+                        self.active_run["id"],
+                        MessageInProgress(
+                            id=event["data"]["chunk"].id,
+                            tool_call_id=None,
+                            tool_call_name=None
+                        )
+                    )
+                    current_stream = self.get_message_in_progress(self.active_run["id"])
+
+                yield self._dispatch_event(
+                    TextMessageContentEvent(
+                        type=EventType.TEXT_MESSAGE_CONTENT,
+                        message_id=current_stream["id"],
+                        delta=event["data"]["chunk"].content,
+                        raw_event=event,
+                    )
+                )
+                return
+
+        elif event_type == LangGraphEventTypes.OnChatModelEnd:
+            if self.get_message_in_progress(self.active_run["id"]) and self.get_message_in_progress(self.active_run["id"]).get("tool_call_id"):
+                resolved = self._dispatch_event(
+                    ToolCallEndEvent(type=EventType.TOOL_CALL_END, tool_call_id=self.get_message_in_progress(self.active_run["id"])["tool_call_id"], raw_event=event)
+                )
+                if resolved:
+                    self.messages_in_process[self.active_run["id"]] = None
+                yield resolved
+            elif self.get_message_in_progress(self.active_run["id"]) and self.get_message_in_progress(self.active_run["id"]).get("id"):
+                resolved = self._dispatch_event(
+                    TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=self.get_message_in_progress(self.active_run["id"])["id"], raw_event=event)
+                )
+                if resolved:
+                    self.messages_in_process[self.active_run["id"]] = None
+                yield resolved
+
+        elif event_type == LangGraphEventTypes.OnCustomEvent:
+            if event["name"] == CustomEventNames.ManuallyEmitMessage:
+                yield self._dispatch_event(
+                    TextMessageStartEvent(type=EventType.TEXT_MESSAGE_START, role="assistant", message_id=event["data"]["message_id"], raw_event=event)
+                )
+                yield self._dispatch_event(
+                    TextMessageContentEvent(
+                        type=EventType.TEXT_MESSAGE_CONTENT,
+                        message_id=event["data"]["message_id"],
+                        delta=event["data"]["message"],
+                        raw_event=event,
+                    )
+                )
+                yield self._dispatch_event(
+                    TextMessageEndEvent(type=EventType.TEXT_MESSAGE_END, message_id=event["data"]["message_id"], raw_event=event)
+                )
+
+            elif event["name"] == CustomEventNames.ManuallyEmitToolCall:
+                yield self._dispatch_event(
+                    ToolCallStartEvent(
+                        type=EventType.TOOL_CALL_START,
+                        tool_call_id=event["data"]["id"],
+                        tool_call_name=event["data"]["name"],
+                        parent_message_id=event["data"]["id"],
+                        raw_event=event,
+                    )
+                )
+                yield self._dispatch_event(
+                    ToolCallArgsEvent(type=EventType.TOOL_CALL_ARGS, tool_call_id=event["data"]["id"], delta=event["data"]["args"], raw_event=event)
+                )
+                yield self._dispatch_event(
+                    ToolCallEndEvent(type=EventType.TOOL_CALL_END, tool_call_id=event["data"]["id"], raw_event=event)
+                )
+
+            elif event["name"] == CustomEventNames.ManuallyEmitState:
+                self.active_run["manually_emitted_state"] = event["data"]
+                yield self._dispatch_event(
+                    StateSnapshotEvent(type=EventType.STATE_SNAPSHOT, snapshot=self.get_state_snapshot(state), raw_event=event)
+                )
+            
+            yield self._dispatch_event(
+                CustomEvent(type=EventType.CUSTOM, name=event["name"], value=event["data"], raw_event=event)
+            )
+
+    def handle_thinking_event(self, reasoning_data: LangGraphReasoning) -> Generator[str, Any, str | None]:
+        if not reasoning_data or "type" not in reasoning_data or "text" not in reasoning_data:
+            return ""
+
+        thinking_step_index = reasoning_data.get("index")
+
+        if (self.active_run.get("thinking_process") and
+                self.active_run["thinking_process"].get("index") and
+                self.active_run["thinking_process"]["index"] != thinking_step_index):
+
+            if self.active_run["thinking_process"].get("type"):
+                yield self._dispatch_event(
+                    ThinkingTextMessageEndEvent(
+                        type=EventType.THINKING_TEXT_MESSAGE_END,
+                    )
+                )
+            yield self._dispatch_event(
+                ThinkingEndEvent(
+                    type=EventType.THINKING_END,
+                )
+            )
+            self.active_run["thinking_process"] = None
+
+        if not self.active_run.get("thinking_process"):
+            yield self._dispatch_event(
+                ThinkingStartEvent(
+                    type=EventType.THINKING_START,
+                )
+            )
+            self.active_run["thinking_process"] = {
+                "index": thinking_step_index
+            }
+
+        if self.active_run["thinking_process"].get("type") != reasoning_data["type"]:
+            yield self._dispatch_event(
+                ThinkingTextMessageStartEvent(
+                    type=EventType.THINKING_TEXT_MESSAGE_START,
+                )
+            )
+            self.active_run["thinking_process"]["type"] = reasoning_data["type"]
+
+        if self.active_run["thinking_process"].get("type"):
+            yield self._dispatch_event(
+                ThinkingTextMessageContentEvent(
+                    type=EventType.THINKING_TEXT_MESSAGE_CONTENT,
+                    delta=reasoning_data["text"]
+                )
+            )
+
+    async def get_checkpoint_before_message(self, message_id: str, thread_id: str):
+        if not thread_id:
+            raise ValueError("Missing thread_id in config")
+
+        history_list = []
+        async for snapshot in self.graph.aget_state_history({"configurable": {"thread_id": thread_id}}):
+            history_list.append(snapshot)
+
+        history_list.reverse()
+        for idx, snapshot in enumerate(history_list):
+            messages = snapshot.values.get("messages", [])
+            if any(getattr(m, "id", None) == message_id for m in messages):
+                if idx == 0:
+                    # No snapshot before this
+                    # Return synthetic "empty before" version
+                    empty_snapshot = snapshot
+                    empty_snapshot.values["messages"] = []
+                    return empty_snapshot
+                return history_list[idx - 1]  # return one snapshot *before* the one that includes the message
+
+        raise ValueError("Message ID not found in history")
+

ag_ui_langgraph-0.0.1/ag_ui_langgraph/endpoint.py

@@ -0,0 +1,27 @@
+from fastapi import FastAPI, HTTPException, Request
+from fastapi.responses import StreamingResponse
+
+from ag_ui.core.types import RunAgentInput
+from ag_ui.encoder import EventEncoder
+
+from .agent import LangGraphAgent
+
+def add_langgraph_fastapi_endpoint(app: FastAPI, agent: LangGraphAgent, path: str = "/"):
+    """Adds an endpoint to the FastAPI app."""
+
+    @app.post(path)
+    async def langgraph_agent_endpoint(input_data: RunAgentInput, request: Request):
+        # Get the accept header from the request
+        accept_header = request.headers.get("accept")
+
+        # Create an event encoder to properly format SSE events
+        encoder = EventEncoder(accept=accept_header)
+
+        async def event_generator():
+            async for event in agent.run(input_data):
+                yield encoder.encode(event)
+
+        return StreamingResponse(
+            event_generator(),
+            media_type=encoder.get_content_type()
+        )

ag_ui_langgraph-0.0.1/ag_ui_langgraph/types.py

@@ -0,0 +1,91 @@
+from typing import TypedDict, Optional, List, Any, Dict, Union, Literal
+from typing_extensions import NotRequired
+from enum import Enum
+
+class LangGraphEventTypes(str, Enum):
+    OnChainStart = "on_chain_start"
+    OnChainStream = "on_chain_stream"
+    OnChainEnd = "on_chain_end"
+    OnChatModelStart = "on_chat_model_start"
+    OnChatModelStream = "on_chat_model_stream"
+    OnChatModelEnd = "on_chat_model_end"
+    OnToolStart = "on_tool_start"
+    OnToolEnd = "on_tool_end"
+    OnCustomEvent = "on_custom_event"
+    OnInterrupt = "on_interrupt"
+
+class CustomEventNames(str, Enum):
+    ManuallyEmitMessage = "manually_emit_message"
+    ManuallyEmitToolCall = "manually_emit_tool_call"
+    ManuallyEmitState = "manually_emit_state"
+    Exit = "exit"
+
+State = Dict[str, Any]
+
+SchemaKeys = TypedDict("SchemaKeys", {
+    "input": NotRequired[Optional[List[str]]],
+    "output": NotRequired[Optional[List[str]]],
+    "config": NotRequired[Optional[List[str]]]
+})
+
+ThinkingProcess = TypedDict("ThinkingProcess", {
+    "index": int,
+    "type": NotRequired[Optional[Literal['text']]],
+})
+
+MessageInProgress = TypedDict("MessageInProgress", {
+    "id": str,
+    "tool_call_id": NotRequired[Optional[str]],
+    "tool_call_name": NotRequired[Optional[str]]
+})
+
+RunMetadata = TypedDict("RunMetadata", {
+    "id": str,
+    "schema_keys": NotRequired[Optional[SchemaKeys]],
+    "node_name": NotRequired[Optional[str]],
+    "prev_node_name": NotRequired[Optional[str]],
+    "exiting_node": NotRequired[bool],
+    "manually_emitted_state": NotRequired[Optional[State]],
+    "thread_id": NotRequired[Optional[ThinkingProcess]],
+    "thinking_process": NotRequired[Optional[str]]
+})
+
+MessagesInProgressRecord = Dict[str, Optional[MessageInProgress]]
+
+ToolCall = TypedDict("ToolCall", {
+    "id": str,
+    "name": str,
+    "args": Dict[str, Any]
+})
+
+class BaseLangGraphPlatformMessage(TypedDict):
+    content: str
+    role: str
+    additional_kwargs: NotRequired[Dict[str, Any]]
+    type: str
+    id: str
+
+class LangGraphPlatformResultMessage(BaseLangGraphPlatformMessage):
+    tool_call_id: str
+    name: str
+
+class LangGraphPlatformActionExecutionMessage(BaseLangGraphPlatformMessage):
+    tool_calls: List[ToolCall]
+
+LangGraphPlatformMessage = Union[
+    LangGraphPlatformActionExecutionMessage,
+    LangGraphPlatformResultMessage,
+    BaseLangGraphPlatformMessage,
+]
+
+PredictStateTool = TypedDict("PredictStateTool", {
+    "tool": str,
+    "state_key": str,
+    "tool_argument": str
+})
+
+LangGraphReasoning = TypedDict("LangGraphReasoning", {
+    "type": str,
+    "text": str,
+    "index": int
+})

ag_ui_langgraph-0.0.1/ag_ui_langgraph/utils.py

@@ -0,0 +1,179 @@
+import json
+import re
+from typing import List, Any, Dict, Union
+
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage, ToolMessage
+from ag_ui.core import (
+    Message as AGUIMessage,
+    UserMessage as AGUIUserMessage,
+    AssistantMessage as AGUIAssistantMessage,
+    SystemMessage as AGUISystemMessage,
+    ToolMessage as AGUIToolMessage,
+    ToolCall as AGUIToolCall,
+    FunctionCall as AGUIFunctionCall,
+)
+from .types import State, SchemaKeys, LangGraphReasoning
+
+DEFAULT_SCHEMA_KEYS = ["tools"]
+
+def filter_object_by_schema_keys(obj: Dict[str, Any], schema_keys: List[str]) -> Dict[str, Any]:
+    if not obj:
+        return {}
+    return {k: v for k, v in obj.items() if k in schema_keys}
+
+def get_stream_payload_input(
+    *,
+    mode: str,
+    state: State,
+    schema_keys: SchemaKeys,
+) -> Union[State, None]:
+    input_payload = state if mode == "start" else None
+    if input_payload and schema_keys and schema_keys.get("input"):
+        input_payload = filter_object_by_schema_keys(input_payload, [*DEFAULT_SCHEMA_KEYS, *schema_keys["input"]])
+    return input_payload
+
+def stringify_if_needed(item: Any) -> str:
+    if item is None:
+        return ''
+    if isinstance(item, str):
+        return item
+    return json.dumps(item)
+
+def langchain_messages_to_agui(messages: List[BaseMessage]) -> List[AGUIMessage]:
+    agui_messages: List[AGUIMessage] = []
+    for message in messages:
+        if isinstance(message, HumanMessage):
+            agui_messages.append(AGUIUserMessage(
+                id=str(message.id),
+                role="user",
+                content=stringify_if_needed(resolve_message_content(message.content)),
+                name=message.name,
+            ))
+        elif isinstance(message, AIMessage):
+            tool_calls = None
+            if message.tool_calls:
+                tool_calls = [
+                    AGUIToolCall(
+                        id=str(tc["id"]),
+                        type="function",
+                        function=AGUIFunctionCall(
+                            name=tc["name"],
+                            arguments=json.dumps(tc.get("args", {})),
+                        ),
+                    )
+                    for tc in message.tool_calls
+                ]
+
+            agui_messages.append(AGUIAssistantMessage(
+                id=str(message.id),
+                role="assistant",
+                content=stringify_if_needed(resolve_message_content(message.content)),
+                tool_calls=tool_calls,
+                name=message.name,
+            ))
+        elif isinstance(message, SystemMessage):
+            agui_messages.append(AGUISystemMessage(
+                id=str(message.id),
+                role="system",
+                content=stringify_if_needed(resolve_message_content(message.content)),
+                name=message.name,
+            ))
+        elif isinstance(message, ToolMessage):
+            agui_messages.append(AGUIToolMessage(
+                id=str(message.id),
+                role="tool",
+                content=stringify_if_needed(resolve_message_content(message.content)),
+                tool_call_id=message.tool_call_id,
+            ))
+        else:
+            raise TypeError(f"Unsupported message type: {type(message)}")
+    return agui_messages
+
+def agui_messages_to_langchain(messages: List[AGUIMessage]) -> List[BaseMessage]:
+    langchain_messages = []
+    for message in messages:
+        role = message.role
+        if role == "user":
+            langchain_messages.append(HumanMessage(
+                id=message.id,
+                content=message.content,
+                name=message.name,
+            ))
+        elif role == "assistant":
+            tool_calls = []
+            if hasattr(message, "tool_calls") and message.tool_calls:
+                for tc in message.tool_calls:
+                    tool_calls.append({
+                        "id": tc.id,
+                        "name": tc.function.name,
+                        "args": json.loads(tc.function.arguments) if hasattr(tc, "function") and tc.function.arguments else {},
+                        "type": "tool_call",
+                    })
+            langchain_messages.append(AIMessage(
+                id=message.id,
+                content=message.content or "",
+                tool_calls=tool_calls,
+                name=message.name,
+            ))
+        elif role == "system":
+            langchain_messages.append(SystemMessage(
+                id=message.id,
+                content=message.content,
+                name=message.name,
+            ))
+        elif role == "tool":
+            langchain_messages.append(ToolMessage(
+                id=message.id,
+                content=message.content,
+                tool_call_id=message.tool_call_id,
+            ))
+        else:
+            raise ValueError(f"Unsupported message role: {role}")
+    return langchain_messages
+
+def resolve_reasoning_content(chunk: Any) -> LangGraphReasoning | None:
+    content = chunk.content
+    if not content:
+        return None
+
+    # Anthropic reasoning response
+    if isinstance(content, list) and content and content[0]:
+        if not content[0].get("thinking"):
+            return None
+        return LangGraphReasoning(
+            text=content[0]["thinking"],
+            type="text",
+            index=content[0].get("index", 0)
+        )
+
+    # OpenAI reasoning response
+    if hasattr(chunk, "additional_kwargs"):
+        reasoning = chunk.additional_kwargs.get("reasoning", {})
+        summary = reasoning.get("summary", [])
+        if summary:
+            data = summary[0]
+            if not data or not data.get("text"):
+                return None
+            return LangGraphReasoning(
+                type="text",
+                text=data["text"],
+                index=data.get("index", 0)
+            )
+
+    return None
+
+def resolve_message_content(content: Any) -> str | None:
+    if not content:
+        return None
+
+    if isinstance(content, str):
+        return content
+
+    if isinstance(content, list) and content:
+        content_text = next((c.get("text") for c in content if isinstance(c, dict) and c.get("type") == "text"), None)
+        return content_text
+
+    return None
+
+def camel_to_snake(name):
+    return re.sub(r'(?<!^)(?=[A-Z])', '_', name).lower()

ag_ui_langgraph-0.0.1/pyproject.toml

@@ -0,0 +1,27 @@
+[tool.poetry]
+name = "ag-ui-langgraph"
+version = "0.0.1"
+description = ""
+authors = ["Ran Shem Tov <ran@copilotkit.ai>"]
+readme = "README.md"
+exclude = [
+    "ag_ui_langgraph/examples/**",
+]
+
+[tool.poetry.dependencies]
+python = "<3.14,>=3.10"
+ag-ui-protocol = "==0.1.7"
+fastapi = { version = "^0.115.12", optional = true }
+langchain = ">=0.3.0"
+langchain-core = ">=0.3.0"
+langgraph = ">=0.3.25,<=0.5.0"
+
+[tool.poetry.extras]
+fastapi = ["fastapi"]
+
+[build-system]
+requires = ["poetry-core"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.poetry.scripts]
+dev = "ag_ui_langgraph.dojo:main"