relax agent types (#62)

Author: vbarda

.gitignore

@@ -36,4 +36,6 @@ MANIFEST
 # mypy
 .mypy_cache/
 .dmypy.json
-dmypy.json
+dmypy.json
+
+.langgraph_api

Makefile

@@ -22,8 +22,8 @@ test_watch:
 ######################
 
 # Define a variable for Python and notebook files.
-lint format: PYTHON_FILES=.
-lint_diff format_diff: PYTHON_FILES=$(shell git diff --relative=. --name-only --diff-filter=d master | grep -E '\.py$$|\.ipynb$$')
+lint format: PYTHON_FILES=langgraph_swarm/ tests/
+lint_diff format_diff: PYTHON_FILES=$(shell git diff --relative=langgraph_swarm/ tests/ --name-only --diff-filter=d master | grep -E '\.py$$|\.ipynb$$')
 
 lint lint_diff:
 	[ "$(PYTHON_FILES)" = "" ] ||	uv run ruff format $(PYTHON_FILES) --diff

langgraph_swarm/swarm.py

@@ -1,5 +1,5 @@
 from langgraph.graph import START, MessagesState, StateGraph
-from langgraph.graph.state import CompiledStateGraph
+from langgraph.pregel import Pregel
 from typing_extensions import Any, Literal, Optional, Type, TypeVar, Union, get_args, get_origin
 
 from langgraph_swarm.handoff import get_handoff_destinations
@@ -86,7 +86,7 @@ def route_to_active_agent(state: dict):
 
 
 def create_swarm(
-    agents: list[CompiledStateGraph],
+    agents: list[Pregel],
     *,
     default_active_agent: str,
     state_schema: StateSchemaType = SwarmState,