Apply stricter linting (#89)

Apply stricter linting across the codebase

Author: eyurtsev

.gitignore

@@ -39,3 +39,4 @@ MANIFEST
 dmypy.json
 
 .langgraph_api
+.idea

examples/customer_support/src/agent/customer_support.ipynb

@@ -128,7 +128,7 @@
     "    hotel_id: str,\n",
     "    config: RunnableConfig,\n",
     ") -> str:\n",
-    "    \"\"\"Book a hotel\"\"\"\n",
+    "    \"\"\"Book a hotel.\"\"\"\n",
     "    user_id = config[\"configurable\"].get(\"user_id\")\n",
     "    hotel = [hotel for hotel in HOTELS if hotel[\"id\"] == hotel_id][0]\n",
     "    RESERVATIONS[user_id][\"hotel_info\"] = hotel\n",
@@ -239,7 +239,6 @@
    "source": [
     "def print_stream(stream):\n",
     "    for ns, update in stream:\n",
-    "        print(f\"Namespace '{ns}'\")\n",
     "        for node, node_updates in update.items():\n",
     "            if node_updates is None:\n",
     "                continue\n",
@@ -252,21 +251,16 @@
     "                raise ValueError(node_updates)\n",
     "\n",
     "            for node_updates in node_updates_list:\n",
-    "                print(f\"Update from node '{node}'\")\n",
     "                if isinstance(node_updates, tuple):\n",
-    "                    print(node_updates)\n",
     "                    continue\n",
     "                messages_key = next(\n",
     "                    (k for k in node_updates.keys() if \"messages\" in k), None\n",
     "                )\n",
     "                if messages_key is not None:\n",
     "                    node_updates[messages_key][-1].pretty_print()\n",
     "                else:\n",
-    "                    print(node_updates)\n",
-    "\n",
-    "        print(\"\\n\\n\")\n",
-    "\n",
-    "    print(\"\\n===\\n\")"
+    "                    pass\n",
+    "\n"
    ]
   },
   {

examples/customer_support/src/agent/customer_support.py

@@ -75,7 +75,7 @@ def book_hotel(
     hotel_id: str,
     config: RunnableConfig,
 ) -> str:
-    """Book a hotel"""
+    """Book a hotel."""
     user_id = config["configurable"].get("user_id")
     hotel = [hotel for hotel in HOTELS if hotel["id"] == hotel_id][0]
     RESERVATIONS[user_id]["hotel_info"] = hotel

examples/research/src/agent/utils.py

@@ -6,7 +6,6 @@
 
 def print_stream(stream):
     for ns, update in stream:
-        print(f"Namespace '{ns}'")
         for node, node_updates in update.items():
             if node_updates is None:
                 continue
@@ -19,21 +18,17 @@ def print_stream(stream):
                 raise ValueError(node_updates)
 
             for node_updates in node_updates_list:
-                print(f"Update from node '{node}'")
                 if isinstance(node_updates, tuple):
-                    print(node_updates)
                     continue
                 messages_key = next(
                     (k for k in node_updates.keys() if "messages" in k), None
                 )
                 if messages_key is not None:
                     node_updates[messages_key][-1].pretty_print()
                 else:
-                    print(node_updates)
+                    pass
 
-        print("\n\n")
 
-    print("\n===\n")
 
 
 def fetch_doc(url: str) -> str:

langgraph_swarm/__init__.py

@@ -1,4 +1,4 @@
 from langgraph_swarm.handoff import create_handoff_tool
 from langgraph_swarm.swarm import SwarmState, add_active_agent_router, create_swarm
 
-__all__ = ["create_swarm", "add_active_agent_router", "create_handoff_tool", "SwarmState"]
+__all__ = ["SwarmState", "add_active_agent_router", "create_handoff_tool", "create_swarm"]

langgraph_swarm/handoff.py

@@ -1,11 +1,11 @@
 import re
+from typing import Annotated
 
 from langchain_core.messages import ToolMessage
 from langchain_core.tools import BaseTool, InjectedToolCallId, tool
 from langgraph.graph.state import CompiledStateGraph
 from langgraph.prebuilt import InjectedState, ToolNode
 from langgraph.types import Command
-from typing_extensions import Annotated
 
 WHITESPACE_RE = re.compile(r"\s+")
 METADATA_KEY_HANDOFF_DESTINATION = "__handoff_destination"
@@ -17,7 +17,10 @@ def _normalize_agent_name(agent_name: str) -> str:
 
 
 def create_handoff_tool(
-    *, agent_name: str, name: str | None = None, description: str | None = None
+    *,
+    agent_name: str,
+    name: str | None = None,
+    description: str | None = None,
 ) -> BaseTool:
     """Create a tool that can handoff control to the requested agent.
 
@@ -32,6 +35,7 @@ def create_handoff_tool(
             If not provided, the tool name will be `transfer_to_<agent_name>`.
         description: Optional description for the handoff tool.
             If not provided, the tool description will be `Ask agent <agent_name> for help`.
+
     """
     if name is None:
         name = f"transfer_to_{_normalize_agent_name(agent_name)}"

langgraph_swarm/swarm.py

@@ -1,6 +1,8 @@
+from typing import Literal, Optional, Union, get_args, get_origin
+
 from langgraph.graph import START, MessagesState, StateGraph
 from langgraph.pregel import Pregel
-from typing_extensions import Any, Literal, Optional, Type, TypeVar, Union, get_args, get_origin
+from typing_extensions import Any, TypeVar
 
 from langgraph_swarm.handoff import get_handoff_destinations
 
@@ -11,18 +13,18 @@ class SwarmState(MessagesState):
     # NOTE: this state field is optional and is not expected to be provided by the user.
     # If a user does provide it, the graph will start from the specified active agent.
     # If active agent is typed as a `str`, we turn it into enum of all active agent names.
-    active_agent: Optional[str]
+    active_agent: str | None
 
 
 StateSchema = TypeVar("StateSchema", bound=SwarmState)
-StateSchemaType = Type[StateSchema]
+StateSchemaType = type[StateSchema]
 
 
 def _update_state_schema_agent_names(
-    state_schema: StateSchemaType, agent_names: list[str]
+    state_schema: StateSchemaType,
+    agent_names: list[str],
 ) -> StateSchemaType:
     """Update the state schema to use Literal with agent names for 'active_agent'."""
-
     active_agent_annotation = state_schema.__annotations__["active_agent"]
 
     # Check if the annotation is str or Optional[str]
@@ -120,14 +122,17 @@ def add(a: int, b: int) -> int:
             config,
         )
         ```
+
     """
     channels = builder.schemas[builder.state_schema]
     if "active_agent" not in channels:
-        raise ValueError("Missing required key 'active_agent' in in builder's state_schema")
+        msg = "Missing required key 'active_agent' in in builder's state_schema"
+        raise ValueError(msg)
 
     if default_active_agent not in route_to:
+        msg = f"Default active agent '{default_active_agent}' not found in routes {route_to}"
         raise ValueError(
-            f"Default active agent '{default_active_agent}' not found in routes {route_to}"
+            msg,
         )
 
     def route_to_active_agent(state: dict):
@@ -142,7 +147,7 @@ def create_swarm(
     *,
     default_active_agent: str,
     state_schema: StateSchemaType = SwarmState,
-    config_schema: Type[Any] | None = None,
+    config_schema: type[Any] | None = None,
 ) -> StateGraph:
     """Create a multi-agent swarm.
 
@@ -200,10 +205,12 @@ def add(a: int, b: int) -> int:
             config,
         )
         ```
+
     """
     active_agent_annotation = state_schema.__annotations__.get("active_agent")
     if active_agent_annotation is None:
-        raise ValueError("Missing required key 'active_agent' in state_schema")
+        msg = "Missing required key 'active_agent' in state_schema"
+        raise ValueError(msg)
 
     agent_names = [agent.name for agent in agents]
     state_schema = _update_state_schema_agent_names(state_schema, agent_names)

pyproject.toml

@@ -42,14 +42,17 @@ target-version = "py310"
 
 [tool.ruff.lint]
 select = [
-    "E",  # pycodestyle errors
-    "W",  # pycodestyle warnings
-    "F",  # pyflakes
-    "I",  # isort
-    "B",  # flake8-bugbear
+  "ALL",
 ]
 ignore = [
-  "E501" # line-length
+  "E501", # line-length
+  "D100",
+  "D104",
+  "D203",
+  "D213",
+  "UP007",
+  "COM812",
+  "ANN",
 ]
 
 
@@ -59,3 +62,24 @@ warn_return_any = true
 warn_unused_configs = true
 disallow_untyped_defs = true
 check_untyped_defs = true
+
+
+[tool.ruff.lint.extend-per-file-ignores]
+"tests/**/*.py" = [
+    # at least this three should be fine in tests:
+    "S101", # asserts allowed in tests...
+    "ARG", # Unused function args -> fixtures nevertheless are functionally relevant...
+    "FBT", # Don't care about booleans as positional arguments in tests, e.g. via @pytest.mark.parametrize()
+    "D104",
+    # The below are debateable
+    "PLR2004", # Magic value used in comparison, ...
+    "S311", # Standard pseudo-random generators are not suitable for cryptographic purposes
+    "D100", # Missing docstring in public module
+    "D101", # Missing docstring in public class
+    "D102", # Missing docstring in public method
+    "D103", # Missing docstring in public function
+    "D104", # Missing docstring in public package
+    "D105", # Missing docstring in magic method
+    "D415", # First line should end with a period, question mark, or exclamation point
+]
+

tests/test_import.py

@@ -1,5 +1,5 @@
 def test_import() -> None:
-    """Test that the code can be imported"""
+    """Test that the code can be imported."""
     from langgraph_swarm import (  # noqa: F401
         add_active_agent_router,
         create_handoff_tool,

tests/test_swarm.py

@@ -1,5 +1,3 @@
-from typing import Optional
-
 from langchain_core.callbacks.manager import CallbackManagerForLLMRun
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import AIMessage, BaseMessage
@@ -21,8 +19,8 @@ def _llm_type(self) -> str:
     def _generate(
         self,
         messages: list[BaseMessage],
-        stop: Optional[list[str]] = None,
-        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        stop: list[str] | None = None,
+        run_manager: CallbackManagerForLLMRun | None = None,
         **kwargs,
     ) -> ChatResult:
         generation = ChatGeneration(message=self.responses[self.idx])
@@ -44,7 +42,7 @@ def test_basic_swarm() -> None:
                     "name": "transfer_to_bob",
                     "args": {},
                     "id": "call_1LlFyjm6iIhDjdn7juWuPYr4",
-                }
+                },
             ],
         ),
         AIMessage(
@@ -59,7 +57,7 @@ def test_basic_swarm() -> None:
                     "name": "transfer_to_alice",
                     "args": {},
                     "id": "call_T6pNmo2jTfZEK3a9avQ14f8Q",
-                }
+                },
             ],
         ),
         AIMessage(
@@ -73,7 +71,7 @@ def test_basic_swarm() -> None:
                         "b": 7,
                     },
                     "id": "call_4kLYO1amR2NfhAxfECkALCr1",
-                }
+                },
             ],
         ),
         AIMessage(
@@ -85,7 +83,7 @@ def test_basic_swarm() -> None:
     model = FakeChatModel(responses=recorded_messages)
 
     def add(a: int, b: int) -> int:
-        """Add two numbers"""
+        """Add two numbers."""
         return a + b
 
     alice = create_react_agent(
@@ -99,8 +97,9 @@ def add(a: int, b: int) -> int:
         model,
         [
             create_handoff_tool(
-                agent_name="Alice", description="Transfer to Alice, she can help with math"
-            )
+                agent_name="Alice",
+                description="Transfer to Alice, she can help with math",
+            ),
         ],
         prompt="You are Bob, you speak like a pirate.",
         name="Bob",