From a44d44e1f0f50f3c34bfa40c4d974304d1387604 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 9 Jul 2025 10:39:17 -0400 Subject: [PATCH 1/5] enable linting --- Makefile | 2 +- tests/test_swarm.py | 11 +++++++---- uv.lock | 2 +- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/Makefile b/Makefile index f31cbb2..7e51bfa 100644 --- a/Makefile +++ b/Makefile @@ -28,7 +28,7 @@ lint_diff format_diff: PYTHON_FILES=$(shell git diff --relative=langgraph_swarm/ lint lint_diff: [ "$(PYTHON_FILES)" = "" ] || uv run ruff format $(PYTHON_FILES) --diff [ "$(PYTHON_FILES)" = "" ] || uv run ruff check $(PYTHON_FILES) --diff - # [ "$(PYTHON_FILES)" = "" ] || uv run mypy $(PYTHON_FILES) + [ "$(PYTHON_FILES)" = "" ] || uv run mypy $(PYTHON_FILES) format format_diff: [ "$(PYTHON_FILES)" = "" ] || uv run ruff check --fix $(PYTHON_FILES) diff --git a/tests/test_swarm.py b/tests/test_swarm.py index e2dfe80..af46f95 100644 --- a/tests/test_swarm.py +++ b/tests/test_swarm.py @@ -1,7 +1,10 @@ +from typing import Any, Callable, Sequence from langchain_core.callbacks.manager import CallbackManagerForLLMRun from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import AIMessage, BaseMessage from langchain_core.outputs import ChatGeneration, ChatResult +from langchain_core.tools import BaseTool +from langchain_core.runnables.config import RunnableConfig from langgraph.checkpoint.memory import MemorySaver from langgraph.prebuilt import create_react_agent @@ -21,13 +24,13 @@ def _generate( messages: list[BaseMessage], stop: list[str] | None = None, run_manager: CallbackManagerForLLMRun | None = None, - **kwargs, + **kwargs: Any, ) -> ChatResult: generation = ChatGeneration(message=self.responses[self.idx]) self.idx += 1 return ChatResult(generations=[generation]) - def bind_tools(self, tools: list[any]) -> "FakeChatModel": + def bind_tools(self, tools: Sequence[dict[str, Any] | type | Callable[..., Any] | BaseTool], *, tool_choice: str | None = None, **kwargs: Any) -> "FakeChatModel": return self @@ -80,7 +83,7 @@ def test_basic_swarm() -> None: ), ] - model = FakeChatModel(responses=recorded_messages) + model = FakeChatModel(responses=recorded_messages) # type: ignore[arg-type] def add(a: int, b: int) -> int: """Add two numbers.""" @@ -109,7 +112,7 @@ def add(a: int, b: int) -> int: workflow = create_swarm([alice, bob], default_active_agent="Alice") app = workflow.compile(checkpointer=checkpointer) - config = {"configurable": {"thread_id": "1"}} + config: RunnableConfig = {"configurable": {"thread_id": "1"}} turn_1 = app.invoke( {"messages": [{"role": "user", "content": "i'd like to speak to Bob"}]}, config, diff --git a/uv.lock b/uv.lock index eedaba2..6aed801 100644 --- a/uv.lock +++ b/uv.lock @@ -326,7 +326,7 @@ wheels = [ [[package]] name = "langgraph-swarm" -version = "0.0.11" +version = "0.0.12" source = { editable = "." } dependencies = [ { name = "langchain-core" }, From ad94b74f67bb8712ba79291c05d6fc8acc29f987 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 9 Jul 2025 10:39:45 -0400 Subject: [PATCH 2/5] x --- tests/test_swarm.py | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/tests/test_swarm.py b/tests/test_swarm.py index af46f95..b5852bd 100644 --- a/tests/test_swarm.py +++ b/tests/test_swarm.py @@ -1,15 +1,19 @@ -from typing import Any, Callable, Sequence +from collections.abc import Callable, Sequence +from typing import TYPE_CHECKING, Any + from langchain_core.callbacks.manager import CallbackManagerForLLMRun from langchain_core.language_models.chat_models import BaseChatModel from langchain_core.messages import AIMessage, BaseMessage from langchain_core.outputs import ChatGeneration, ChatResult from langchain_core.tools import BaseTool -from langchain_core.runnables.config import RunnableConfig from langgraph.checkpoint.memory import MemorySaver from langgraph.prebuilt import create_react_agent from langgraph_swarm import create_handoff_tool, create_swarm +if TYPE_CHECKING: + from langchain_core.runnables.config import RunnableConfig + class FakeChatModel(BaseChatModel): idx: int = 0 @@ -30,7 +34,13 @@ def _generate( self.idx += 1 return ChatResult(generations=[generation]) - def bind_tools(self, tools: Sequence[dict[str, Any] | type | Callable[..., Any] | BaseTool], *, tool_choice: str | None = None, **kwargs: Any) -> "FakeChatModel": + def bind_tools( + self, + tools: Sequence[dict[str, Any] | type | Callable[..., Any] | BaseTool], + *, + tool_choice: str | None = None, + **kwargs: Any, + ) -> "FakeChatModel": return self From 7ab121a33f346da72521b43f81c3aef0f76487fc Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 9 Jul 2025 14:11:29 -0400 Subject: [PATCH 3/5] x --- langgraph_swarm/__init__.py | 7 ++++++- langgraph_swarm/handoff.py | 14 ++++++++++---- langgraph_swarm/swarm.py | 13 ++++++++----- pyproject.toml | 2 +- 4 files changed, 25 insertions(+), 11 deletions(-) diff --git a/langgraph_swarm/__init__.py b/langgraph_swarm/__init__.py index d6f30ac..b3f90c9 100644 --- a/langgraph_swarm/__init__.py +++ b/langgraph_swarm/__init__.py @@ -1,4 +1,9 @@ from langgraph_swarm.handoff import create_handoff_tool from langgraph_swarm.swarm import SwarmState, add_active_agent_router, create_swarm -__all__ = ["SwarmState", "add_active_agent_router", "create_handoff_tool", "create_swarm"] +__all__ = [ + "SwarmState", + "add_active_agent_router", + "create_handoff_tool", + "create_swarm", +] diff --git a/langgraph_swarm/handoff.py b/langgraph_swarm/handoff.py index 08d1ec0..cccf9f1 100644 --- a/langgraph_swarm/handoff.py +++ b/langgraph_swarm/handoff.py @@ -47,7 +47,7 @@ def create_handoff_tool( def handoff_to_agent( state: Annotated[dict, InjectedState], tool_call_id: Annotated[str, InjectedToolCallId], - ): + ) -> Command: tool_message = ToolMessage( content=f"Successfully transferred to {agent_name}", name=name, @@ -56,14 +56,19 @@ def handoff_to_agent( return Command( goto=agent_name, graph=Command.PARENT, - update={"messages": state["messages"] + [tool_message], "active_agent": agent_name}, + update={ + "messages": state["messages"] + [tool_message], + "active_agent": agent_name, + }, ) handoff_to_agent.metadata = {METADATA_KEY_HANDOFF_DESTINATION: agent_name} return handoff_to_agent -def get_handoff_destinations(agent: CompiledStateGraph, tool_node_name: str = "tools") -> list[str]: +def get_handoff_destinations( + agent: CompiledStateGraph, tool_node_name: str = "tools" +) -> list[str]: """Get a list of destinations from agent's handoff tools.""" nodes = agent.get_graph().nodes if tool_node_name not in nodes: @@ -77,5 +82,6 @@ def get_handoff_destinations(agent: CompiledStateGraph, tool_node_name: str = "t return [ tool.metadata[METADATA_KEY_HANDOFF_DESTINATION] for tool in tools - if tool.metadata is not None and METADATA_KEY_HANDOFF_DESTINATION in tool.metadata + if tool.metadata is not None + and METADATA_KEY_HANDOFF_DESTINATION in tool.metadata ] diff --git a/langgraph_swarm/swarm.py b/langgraph_swarm/swarm.py index 16bd815..4c1ab86 100644 --- a/langgraph_swarm/swarm.py +++ b/langgraph_swarm/swarm.py @@ -1,4 +1,4 @@ -from typing import Literal, Optional, Union, get_args, get_origin +from typing import Literal, Optional, Union, cast, get_args, get_origin from langgraph.graph import START, MessagesState, StateGraph from langgraph.pregel import Pregel @@ -30,7 +30,8 @@ def _update_state_schema_agent_names( # Check if the annotation is str or Optional[str] is_str_type = active_agent_annotation is str is_optional_str = ( - get_origin(active_agent_annotation) is Union and get_args(active_agent_annotation)[0] is str + get_origin(active_agent_annotation) is Union + and get_args(active_agent_annotation)[0] is str ) # We only update if the 'active_agent' is a str or Optional[str] @@ -135,8 +136,8 @@ def add(a: int, b: int) -> int: msg, ) - def route_to_active_agent(state: dict): - return state.get("active_agent", default_active_agent) + def route_to_active_agent(state: dict) -> str: + return cast(str, state.get("active_agent", default_active_agent)) builder.add_conditional_edges(START, route_to_active_agent, path_map=route_to) return builder @@ -223,7 +224,9 @@ def add(a: int, b: int) -> int: for agent in agents: builder.add_node( agent.name, - agent, + # We need to update the type signatures in add_node to match + # the fact that more flexible Pregel objects are allowed. + agent, # type: ignore[arg-type] destinations=tuple(get_handoff_destinations(agent)), ) diff --git a/pyproject.toml b/pyproject.toml index 123645e..a29d1d0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ python_files = ["test_*.py"] python_functions = ["test_*"] [tool.ruff] -line-length = 100 +line-length = 88 target-version = "py310" [tool.ruff.lint] From f04c8ea220a38ca20b7d5bbb9a6b65df66a5167a Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 9 Jul 2025 14:32:23 -0400 Subject: [PATCH 4/5] x --- langgraph_swarm/swarm.py | 5 ++++- tests/test_swarm.py | 8 ++++++-- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/langgraph_swarm/swarm.py b/langgraph_swarm/swarm.py index 4c1ab86..33c1c94 100644 --- a/langgraph_swarm/swarm.py +++ b/langgraph_swarm/swarm.py @@ -227,7 +227,10 @@ def add(a: int, b: int) -> int: # We need to update the type signatures in add_node to match # the fact that more flexible Pregel objects are allowed. agent, # type: ignore[arg-type] - destinations=tuple(get_handoff_destinations(agent)), + destinations=tuple( + # Need to update implementation to support Pregel objects + get_handoff_destinations(agent) # type: ignore[arg-type] + ), ) return builder diff --git a/tests/test_swarm.py b/tests/test_swarm.py index b5852bd..7d324e5 100644 --- a/tests/test_swarm.py +++ b/tests/test_swarm.py @@ -124,7 +124,9 @@ def add(a: int, b: int) -> int: config: RunnableConfig = {"configurable": {"thread_id": "1"}} turn_1 = app.invoke( - {"messages": [{"role": "user", "content": "i'd like to speak to Bob"}]}, + { # type: ignore[arg-type] + "messages": [{"role": "user", "content": "i'd like to speak to Bob"}] + }, config, ) @@ -135,7 +137,9 @@ def add(a: int, b: int) -> int: assert turn_1["active_agent"] == "Bob" turn_2 = app.invoke( - {"messages": [{"role": "user", "content": "what's 5 + 7?"}]}, + { # type: ignore[arg-type] + "messages": [{"role": "user", "content": "what's 5 + 7?"}] + }, config, ) From 74cb5c925663ae5b82bdaabf5699282d97428726 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Wed, 9 Jul 2025 14:34:08 -0400 Subject: [PATCH 5/5] x --- langgraph_swarm/swarm.py | 4 ++-- tests/test_import.py | 5 ----- 2 files changed, 2 insertions(+), 7 deletions(-) diff --git a/langgraph_swarm/swarm.py b/langgraph_swarm/swarm.py index 33c1c94..eb8e3c6 100644 --- a/langgraph_swarm/swarm.py +++ b/langgraph_swarm/swarm.py @@ -49,7 +49,7 @@ def _update_state_schema_agent_names( # If it was Optional[str], make it Optional[Literal[...]] if is_optional_str: - updated_schema.__annotations__["active_agent"] = Optional[literal_type] + updated_schema.__annotations__["active_agent"] = Optional[literal_type] # noqa: UP045 else: updated_schema.__annotations__["active_agent"] = literal_type @@ -137,7 +137,7 @@ def add(a: int, b: int) -> int: ) def route_to_active_agent(state: dict) -> str: - return cast(str, state.get("active_agent", default_active_agent)) + return cast("str", state.get("active_agent", default_active_agent)) builder.add_conditional_edges(START, route_to_active_agent, path_map=route_to) return builder diff --git a/tests/test_import.py b/tests/test_import.py index 2c6665e..68c05ec 100644 --- a/tests/test_import.py +++ b/tests/test_import.py @@ -1,7 +1,2 @@ def test_import() -> None: """Test that the code can be imported.""" - from langgraph_swarm import ( # noqa: F401 - add_active_agent_router, - create_handoff_tool, - create_swarm, - )