feat(langgraph-py): several fixes

Author: ranst91

typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/__init__.py

@@ -14,6 +14,7 @@
     LangGraphPlatformMessage,
     PredictStateTool
 )
+from .endpoint import add_langgraph_fastapi_endpoint
 
 __all__ = [
     "LangGraphAgent",
@@ -29,5 +30,6 @@
     "LangGraphPlatformResultMessage",
     "LangGraphPlatformActionExecutionMessage",
     "LangGraphPlatformMessage",
-    "PredictStateTool"
+    "PredictStateTool",
+    "add_langgraph_fastapi_endpoint"
 ]

typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/agent.py

@@ -2,9 +2,7 @@
 import json
 from typing import Optional, List, Any, Union, AsyncGenerator, Generator
 
-from fastapi.responses import StreamingResponse
-
-from langgraph.graph.graph import CompiledGraph
+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
@@ -78,13 +76,14 @@
 ]
 
 class LangGraphAgent:
-    def __init__(self, *, name: str, graph: CompiledGraph, description: Optional[str] = None, config:  Union[Optional[RunnableConfig], dict] = None):
+    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
@@ -352,7 +351,6 @@ def set_message_in_progress(self, run_id: str, data: MessageInProgress):
         }
 
     def get_schema_keys(self, config) -> SchemaKeys:
-        CONSTANT_KEYS = ['messages']
         try:
             input_schema = self.graph.get_input_jsonschema(config)
             output_schema = self.graph.get_output_jsonschema(config)
@@ -363,14 +361,14 @@ def get_schema_keys(self, config) -> SchemaKeys:
             config_schema_keys = list(config_schema["properties"].keys()) if "properties" in config_schema else []
 
             return {
-                "input": [*input_schema_keys, *CONSTANT_KEYS],
-                "output": [*output_schema_keys, *CONSTANT_KEYS],
+                "input": [*input_schema_keys, *self.constant_schema_keys],
+                "output": [*output_schema_keys, *self.constant_schema_keys],
                 "config": config_schema_keys,
             }
         except Exception:
             return {
-                "input": CONSTANT_KEYS,
-                "output": CONSTANT_KEYS,
+                "input": self.constant_schema_keys,
+                "output": self.constant_schema_keys,
                 "config": [],
             }
 

typescript-sdk/integrations/langgraph/python/.env.example renamed to typescript-sdk/integrations/langgraph/python/ag_ui_langgraph/examples/.env.example

@@ -5,14 +5,13 @@
 from dotenv import load_dotenv
 load_dotenv()
 
-from .endpoint import add_langgraph_fastapi_endpoint
-from .agent import LangGraphAgent
-from .examples.human_in_the_loop import human_in_the_loop_graph
-from .examples.predictive_state_updates import predictive_state_updates_graph
-from .examples.shared_state import shared_state_graph
-from .examples.tool_based_generative_ui import tool_based_generative_ui_graph
-from .examples.agentic_chat import agentic_chat_graph
-from .examples.agentic_generative_ui import graph
+from ag_ui_langgraph import LangGraphAgent, add_langgraph_fastapi_endpoint
+from .human_in_the_loop import human_in_the_loop_graph
+from .predictive_state_updates import predictive_state_updates_graph
+from .shared_state import shared_state_graph
+from .tool_based_generative_ui import tool_based_generative_ui_graph
+from .agentic_chat import agentic_chat_graph
+from .agentic_generative_ui import graph
 
 app = FastAPI(title="LangGraph Dojo Example Server")
 
@@ -90,7 +89,7 @@ def main():
     """Run the uvicorn server."""
     port = int(os.getenv("PORT", "8000"))
     uvicorn.run(
-        "ag_ui_langgraph.dojo:app",
+        "agents.dojo:app",
         host="0.0.0.0",
         port=port,
         reload=True