langgraph-agent-toolkit 0.1.0__py3-none-any.whl

langgraph_agent_toolkit/__init__.py

@@ -0,0 +1,7 @@
+try:
+    from importlib.metadata import version
+
+    __version__ = version("langgraph_agent_toolkit")
+except Exception:
+    # Fallback for local development
+    __version__ = "0.0.0.dev0"

langgraph_agent_toolkit/agents/agent.py

@@ -0,0 +1,14 @@
+from dataclasses import dataclass
+
+from langgraph.func import Pregel
+from langgraph.graph.state import CompiledStateGraph
+
+from langgraph_agent_toolkit.core.observability.base import BaseObservabilityPlatform
+
+
+@dataclass
+class Agent:
+    name: str
+    description: str
+    graph: CompiledStateGraph | Pregel
+    observability: BaseObservabilityPlatform | None = None

langgraph_agent_toolkit/agents/agent_executor.py

@@ -0,0 +1,415 @@
+import asyncio
+import functools
+import importlib
+import os
+from pathlib import Path
+from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple, TypeVar
+from uuid import UUID, uuid4
+
+import joblib
+from langchain_core.messages import AIMessage, AIMessageChunk, HumanMessage, ToolMessage
+from langchain_core.runnables import RunnableConfig
+from langgraph.errors import GraphRecursionError
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.pregel import Pregel
+from langgraph.types import Command, Interrupt
+
+from langgraph_agent_toolkit.agents.agent import Agent
+from langgraph_agent_toolkit.helper.constants import DEFAULT_AGENT, DEFAULT_RECURSION_LIMIT
+from langgraph_agent_toolkit.helper.logging import logger
+from langgraph_agent_toolkit.helper.utils import (
+    convert_message_content_to_string,
+    langchain_to_chat_message,
+    remove_tool_calls,
+)
+from langgraph_agent_toolkit.schema import AgentInfo, ChatMessage
+
+
+# Type variable for the decorator
+T = TypeVar("T")
+
+
+class AgentExecutor:
+    """Handles the loading, execution and saving logic for different LangGraph agents."""
+
+    def __init__(self, *args):
+        """Initialize the AgentExecutor by importing agents.
+
+        Args:
+            *args: Variable length strings specifying the agents to import,
+                  e.g., "langgraph_agent_toolkit.agents.blueprints.react.agent:react_agent".
+
+        Raises:
+            ValueError: If no agents are provided.
+
+        """
+        self.agents: Dict[str, Agent] = {}
+
+        # Check if args is empty and raise an error
+        if not args:
+            raise ValueError("At least one agent must be provided to AgentExecutor.")
+
+        # Load agents from import strings
+        self.load_agents_from_imports(args)
+        self._validate_default_agent_loaded()
+
+    def load_agents_from_imports(self, args: tuple) -> None:
+        """Dynamically imports agents based on the provided import strings."""
+        for import_str in args:
+            try:
+                module_path, object_name = import_str.split(":")
+                module = importlib.import_module(module_path)
+                agent_obj = getattr(module, object_name)
+
+                # Check if it's a raw graph or already an Agent instance
+                if isinstance(agent_obj, (CompiledStateGraph, Pregel)):
+                    agent = Agent(name=object_name, description=f"Dynamically loaded {object_name}", graph=agent_obj)
+                    self.agents[agent.name] = agent
+                elif isinstance(agent_obj, Agent):
+                    self.agents[agent_obj.name] = agent_obj
+                else:
+                    print(f"Warning: Object '{object_name}' is neither a graph nor an Agent instance")
+            except (ImportError, AttributeError, ValueError) as e:
+                print(f"Error loading agent from '{import_str}': {e}")
+
+    def _validate_default_agent_loaded(self) -> None:
+        """Validate that the default agent is loaded.
+
+        Raises:
+            ValueError: If the default agent is not loaded
+
+        """
+        if not self.agents or DEFAULT_AGENT not in self.agents:
+            raise ValueError(
+                f"Default agent '{DEFAULT_AGENT}' was not imported. Make sure to include it in your agent imports."
+            )
+
+    def get_agent(self, agent_id: str) -> Agent:
+        """Get an agent by its ID.
+
+        Args:
+            agent_id: The ID of the agent to retrieve
+
+        Returns:
+            The requested Agent instance
+
+        Raises:
+            KeyError: If the agent_id is not found
+
+        """
+        if agent_id not in self.agents:
+            raise KeyError(f"Agent '{agent_id}' not found")
+        return self.agents[agent_id]
+
+    def get_all_agent_info(self) -> list[AgentInfo]:
+        """Get information about all available agents.
+
+        Returns:
+            A list of AgentInfo objects containing agent IDs and descriptions
+
+        """
+        return [AgentInfo(key=agent_id, description=agent.description) for agent_id, agent in self.agents.items()]
+
+    def add_agent(self, agent_id: str, agent: Agent) -> None:
+        """Add a new agent to the executor.
+
+        Args:
+            agent_id: The ID to assign to the agent
+            agent: The Agent instance to add
+
+        """
+        self.agents[agent_id] = agent
+
+    @staticmethod
+    def handle_agent_errors(func: Callable[..., T]) -> Callable[..., T]:
+        """Handle errors occurring during agent execution.
+
+        Specifically handles GraphRecursionError and other exceptions.
+
+        Args:
+            func: The function to decorate
+
+        Returns:
+            The decorated function
+
+        """
+
+        def _handle_error(e: Exception):
+            """Handle and re-raise errors with logging."""
+            if isinstance(e, GraphRecursionError):
+                logger.error(f"GraphRecursionError occurred: {e}")
+            else:
+                logger.error(f"Error during agent execution: {e}")
+            raise e
+
+        @functools.wraps(func)
+        async def async_wrapper(self, *args, **kwargs):
+            try:
+                return await func(self, *args, **kwargs)
+            except Exception as e:
+                return _handle_error(e)
+
+        @functools.wraps(func)
+        def sync_wrapper(self, *args, **kwargs):
+            try:
+                return func(self, *args, **kwargs)
+            except Exception as e:
+                return _handle_error(e)
+
+        if asyncio.iscoroutinefunction(func):
+            return async_wrapper
+        else:
+            return sync_wrapper
+
+    async def _setup_agent_execution(
+        self,
+        agent_id: str,
+        message: str,
+        thread_id: Optional[str] = None,
+        model: Optional[str] = None,
+        agent_config: Optional[Dict[str, Any]] = None,
+        recursion_limit: Optional[int] = None,
+    ) -> Tuple[Agent, Any, Any, UUID]:
+        """Apply common setup for agent execution that both invoke and stream methods share.
+
+        Args:
+            agent_id: ID of the agent to invoke
+            message: User message to send to the agent
+            thread_id: Optional thread ID for conversation history
+            model: Optional model name to override the default
+            agent_config: Optional additional configuration for the agent
+            recursion_limit: Optional recursion limit for the agent
+
+        Returns:
+            Tuple containing:
+                - agent: The Agent instance
+                - input_data: The properly formatted input for the agent
+                - config: The RunnableConfig for the agent
+                - run_id: The UUID for this run
+
+        """
+        agent = self.get_agent(agent_id)
+        agent_graph = agent.graph
+
+        run_id = uuid4()
+        thread_id = thread_id or str(uuid4())
+
+        recursion_limit = recursion_limit or DEFAULT_RECURSION_LIMIT
+
+        configurable = {
+            "thread_id": thread_id,
+        }
+
+        if model:
+            configurable["model"] = model
+
+        if agent_config:
+            configurable.update(agent_config)
+
+        callback = agent.observability.get_callback_handler(session_id=thread_id)
+
+        config = RunnableConfig(
+            configurable=configurable,
+            run_id=run_id,
+            callbacks=[callback] if callback else None,
+            recursion_limit=recursion_limit,
+        )
+
+        # Check if there are any interrupts that need to be resumed
+        state = await agent_graph.aget_state(config=config)
+        interrupted_tasks = [task for task in state.tasks if hasattr(task, "interrupts") and task.interrupts]
+
+        if interrupted_tasks:
+            # User input is a response to resume agent execution from interrupt
+            input_data = Command(resume=message)
+        else:
+            input_data = {"messages": [HumanMessage(content=message)]}
+
+        return agent, input_data, config, run_id
+
+    @handle_agent_errors
+    async def invoke(
+        self,
+        agent_id: str,
+        message: str,
+        thread_id: Optional[str] = None,
+        model: Optional[str] = None,
+        agent_config: Optional[Dict[str, Any]] = None,
+        recursion_limit: Optional[int] = None,
+    ) -> ChatMessage:
+        """Invoke an agent with a message and return the response.
+
+        Args:
+            agent_id: ID of the agent to invoke
+            message: User message to send to the agent
+            thread_id: Optional thread ID for conversation history
+            model: Optional model name to override the default
+            agent_config: Optional additional configuration for the agent
+            recursion_limit: Optional recursion limit for the agent
+
+        Returns:
+            ChatMessage: The agent's response
+
+        """
+        agent, input_data, config, run_id = await self._setup_agent_execution(
+            agent_id=agent_id,
+            message=message,
+            thread_id=thread_id,
+            model=model,
+            agent_config=agent_config,
+            recursion_limit=recursion_limit,
+        )
+
+        # Invoke the agent
+        response_events = await agent.graph.ainvoke(input=input_data, config=config, stream_mode=["updates", "values"])
+
+        response_type, response = response_events[-1]
+        if response_type == "values":
+            # Normal response, the agent completed successfully
+            output = langchain_to_chat_message(response["messages"][-1])
+        elif response_type == "updates" and "__interrupt__" in response:
+            # The last thing to occur was an interrupt
+            # Return the value of the first interrupt as an AIMessage
+            output = langchain_to_chat_message(AIMessage(content=response["__interrupt__"][0].value))
+        else:
+            raise ValueError(f"Unexpected response type: {response_type}")
+
+        output.run_id = str(run_id)
+        return output
+
+    @handle_agent_errors
+    async def stream(
+        self,
+        agent_id: str,
+        message: str,
+        thread_id: Optional[str] = None,
+        model: Optional[str] = None,
+        stream_tokens: bool = True,
+        agent_config: Optional[Dict[str, Any]] = None,
+        recursion_limit: Optional[int] = None,
+    ) -> AsyncGenerator[str | ChatMessage, None]:
+        """Stream an agent's response to a message, yielding either tokens or messages.
+
+        Args:
+            agent_id: ID of the agent to invoke
+            message: User message to send to the agent
+            thread_id: Optional thread ID for conversation history
+            model: Optional model name to override the default
+            stream_tokens: Whether to stream individual tokens
+            agent_config: Optional additional configuration for the agent
+            recursion_limit: Optional recursion limit for the agent
+
+        Yields:
+            Either ChatMessage objects for full messages or strings for token chunks
+
+        """
+        agent, input_data, config, run_id = await self._setup_agent_execution(
+            agent_id=agent_id,
+            message=message,
+            thread_id=thread_id,
+            model=model,
+            agent_config=agent_config,
+            recursion_limit=recursion_limit,
+        )
+
+        # Stream from the agent with appropriate modes
+        stream_mode = ["updates", "messages", "custom"] if stream_tokens else ["updates"]
+
+        async for stream_event in agent.graph.astream(input=input_data, config=config, stream_mode=stream_mode):
+            if not isinstance(stream_event, tuple):
+                continue
+
+            stream_mode, event = stream_event
+            new_messages = []
+
+            if stream_mode == "updates":
+                for node, updates in event.items():
+                    # A simple approach to handle agent interrupts.
+                    # In a more sophisticated implementation, we could add
+                    # some structured ChatMessage type to return the interrupt value.
+                    if node == "__interrupt__":
+                        interrupt: Interrupt
+                        for interrupt in updates:
+                            new_messages.append(AIMessage(content=interrupt.value))
+                        continue
+
+                    update_messages = updates.get("messages", [])
+
+                    # Special case for supervisor agent
+                    if node == "supervisor":
+                        # Get only the last AIMessage since supervisor includes all previous messages
+                        ai_messages = [msg for msg in update_messages if isinstance(msg, AIMessage)]
+                        if ai_messages:
+                            update_messages = [ai_messages[-1]]
+
+                    # Special case for expert agents
+                    if node in ("research_expert", "math_expert"):
+                        # Convert to ToolMessage so it displays in the UI as a tool response
+                        if update_messages:
+                            msg = ToolMessage(
+                                content=update_messages[0].content,
+                                name=node,
+                                tool_call_id="",
+                            )
+                            update_messages = [msg]
+                    new_messages.extend(update_messages)
+
+            elif stream_mode == "custom":
+                new_messages = [event]
+
+            elif stream_mode == "messages" and stream_tokens:
+                msg, metadata = event
+                if "skip_stream" in metadata.get("tags", []):
+                    continue
+                # Skip non-LLM nodes that might send messages
+                if not isinstance(msg, AIMessageChunk):
+                    continue
+                content = remove_tool_calls(msg.content)
+                if content:
+                    # Empty content in OpenAI context usually means the model is asking for a tool to be invoked
+                    yield convert_message_content_to_string(content)
+
+            # Process and yield all collected messages
+            for msg in new_messages:
+                try:
+                    chat_message = langchain_to_chat_message(msg)
+                    chat_message.run_id = str(run_id)
+                    # Skip the input message if it's repeated by LangGraph
+                    if chat_message.type == "human" and chat_message.content == message:
+                        continue
+                    yield chat_message
+                except Exception as e:
+                    logger.error(f"Error parsing message: {e}")
+                    continue
+
+    def save(self, path: str, agent_ids: Optional[List[str]] = None) -> None:
+        """Save agents to disk using joblib.
+
+        Args:
+            path: Directory path where to save agents
+            agent_ids: List of agent IDs to save. If None, saves all agents.
+
+        """
+        _path = Path(path)
+        _path.mkdir(exist_ok=True, parents=True)
+
+        agents_to_save = self.agents
+        if agent_ids:
+            agents_to_save = {k: v for k, v in self.agents.items() if k in agent_ids}
+
+        for agent_id, agent in agents_to_save.items():
+            joblib.dump(agent, _path / f"{agent_id}.joblib")
+
+    def load_saved_agents(self, path: str) -> None:
+        """Load saved agents from disk using joblib.
+
+        Args:
+            path: Directory path from which to load agents
+
+        """
+        for filename in os.listdir(path):
+            if filename.endswith(".joblib"):
+                agent = joblib.load(os.path.join(path, filename))
+                self.agents[agent.name] = agent
+
+        self._validate_default_agent_loaded()

langgraph_agent_toolkit/agents/blueprints/bg_task_agent/agent.py

@@ -0,0 +1,69 @@
+import asyncio
+
+from langchain_core.language_models.chat_models import BaseChatModel
+from langchain_core.messages import AIMessage
+from langchain_core.runnables import RunnableConfig, RunnableLambda, RunnableSerializable
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph import END, MessagesState, StateGraph
+from langgraph.types import StreamWriter
+
+from langgraph_agent_toolkit.agents.agent import Agent
+from langgraph_agent_toolkit.agents.blueprints.bg_task_agent.task import Task
+from langgraph_agent_toolkit.core import settings
+from langgraph_agent_toolkit.core.models.factory import ModelFactory
+
+
+class AgentState(MessagesState, total=False):
+    """`total=False` is PEP589 specs.
+
+    documentation: https://typing.readthedocs.io/en/latest/spec/typeddict.html#totality
+    """
+
+
+def wrap_model(model: BaseChatModel) -> RunnableSerializable[AgentState, AIMessage]:
+    preprocessor = RunnableLambda(
+        lambda state: state["messages"],
+        name="StateModifier",
+    )
+    return preprocessor | model
+
+
+async def acall_model(state: AgentState, config: RunnableConfig) -> AgentState:
+    m = ModelFactory.create(config["configurable"].get("model", settings.DEFAULT_MODEL))
+    model_runnable = wrap_model(m)
+    response = await model_runnable.ainvoke(state, config)
+
+    # We return a list, because this will get added to the existing list
+    return {"messages": [response]}
+
+
+async def bg_task(state: AgentState, writer: StreamWriter) -> AgentState:
+    task1 = Task("Simple task 1...", writer)
+    task2 = Task("Simple task 2...", writer)
+
+    task1.start()
+    await asyncio.sleep(2)
+    task2.start()
+    await asyncio.sleep(2)
+    task1.write_data(data={"status": "Still running..."})
+    await asyncio.sleep(2)
+    task2.finish(result="error", data={"output": 42})
+    await asyncio.sleep(2)
+    task1.finish(result="success", data={"output": 42})
+    return {"messages": []}
+
+
+# Define the graph
+agent = StateGraph(AgentState)
+agent.add_node("model", acall_model)
+agent.add_node("bg_task", bg_task)
+agent.set_entry_point("bg_task")
+
+agent.add_edge("bg_task", "model")
+agent.add_edge("model", END)
+
+bg_task_agent = Agent(
+    name="bg-task-agent",
+    description="A background task agent.",
+    graph=agent.compile(checkpointer=MemorySaver()),
+)

langgraph_agent_toolkit/agents/blueprints/bg_task_agent/task.py

@@ -0,0 +1,52 @@
+from typing import Any, Dict, Literal, Optional
+from uuid import uuid4
+
+from langchain_core.messages import BaseMessage
+from langgraph.types import StreamWriter
+
+from langgraph_agent_toolkit.agents.blueprints.bg_task_agent.utils import CustomData
+from langgraph_agent_toolkit.schema.task_data import TaskData
+
+
+class Task:
+    def __init__(self, task_name: str, writer: StreamWriter | None = None) -> None:
+        self.name = task_name
+        self.id = str(uuid4())
+        self.state: Literal["new", "running", "complete"] = "new"
+        self.result: Literal["success", "error"] | None = None
+        self.writer = writer
+
+    def _generate_and_dispatch_message(self, writer: StreamWriter, data: dict):
+        writer = writer or self.writer
+        task_data = TaskData(name=self.name, run_id=self.id, state=self.state, data=data)
+        if self.result:
+            task_data.result = self.result
+        task_custom_data = CustomData(
+            type=self.name,
+            data=task_data.model_dump(),
+        )
+        task_custom_data.dispatch(writer)
+        return task_custom_data.to_langchain()
+
+    def start(self, writer: StreamWriter | None = None, data: dict = {}) -> BaseMessage:
+        self.state = "new"
+        task_message = self._generate_and_dispatch_message(writer, data)
+        return task_message
+
+    def write_data(self, writer: StreamWriter | None = None, data: dict = {}) -> BaseMessage:
+        if self.state == "complete":
+            raise ValueError("Only incomplete tasks can output data.")
+        self.state = "running"
+        task_message = self._generate_and_dispatch_message(writer, data)
+        return task_message
+
+    def finish(
+        self,
+        result: Literal["success", "error"],
+        writer: StreamWriter | None = None,
+        data: Optional[Dict[str, Any]] = None,
+    ) -> BaseMessage:
+        self.state = "complete"
+        self.result = result
+        task_message = self._generate_and_dispatch_message(writer, data or {})
+        return task_message

langgraph_agent_toolkit/agents/blueprints/bg_task_agent/utils.py

@@ -0,0 +1,17 @@
+from typing import Any
+
+from langchain_core.messages import ChatMessage
+from langgraph.types import StreamWriter
+from pydantic import BaseModel, Field
+
+
+class CustomData(BaseModel):
+    """Custom data being sent by an agent."""
+
+    data: dict[str, Any] = Field(description="The custom data")
+
+    def to_langchain(self) -> ChatMessage:
+        return ChatMessage(content=[self.data], role="custom")
+
+    def dispatch(self, writer: StreamWriter) -> None:
+        writer(self.to_langchain())

langgraph_agent_toolkit/agents/blueprints/chatbot/agent.py

@@ -0,0 +1,34 @@
+from langchain_core.messages import BaseMessage
+from langchain_core.runnables import RunnableConfig
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.func import entrypoint
+from langgraph.graph import add_messages
+
+from langgraph_agent_toolkit.agents.agent import Agent
+from langgraph_agent_toolkit.core import settings
+from langgraph_agent_toolkit.core.models.factory import ModelFactory
+
+
+@entrypoint(
+    # checkpointer=MemorySaver(),  # Uncomment if you want to save the state of the agent
+)
+async def chatbot(
+    inputs: dict[str, list[BaseMessage]],
+    *,
+    previous: dict[str, list[BaseMessage]],
+    config: RunnableConfig,
+):
+    messages = inputs["messages"]
+    if previous:
+        messages = add_messages(previous["messages"], messages)
+
+    model = ModelFactory.create(config["configurable"].get("model", settings.DEFAULT_MODEL))
+    response = await model.ainvoke(messages)
+    return entrypoint.final(value={"messages": [response]}, save={"messages": add_messages(messages, response)})
+
+
+chatbot_agent = Agent(
+    name="chatbot-agent",
+    description="A simple chatbot.",
+    graph=chatbot,
+)

langgraph_agent_toolkit/agents/blueprints/command_agent/agent.py

@@ -0,0 +1,54 @@
+import random
+from typing import Literal
+
+from langchain_core.messages import AIMessage
+from langgraph.graph import START, MessagesState, StateGraph
+from langgraph.types import Command
+
+from langgraph_agent_toolkit.agents.agent import Agent
+
+
+class AgentState(MessagesState, total=False):
+    pass
+
+
+def node_a(state: AgentState) -> Command[Literal["node_b", "node_c"]]:
+    print("Called A")
+    value = random.choice(["a", "b"])
+    # this is a replacement for a conditional edge function
+    if value == "a":
+        goto = "node_b"
+    else:
+        goto = "node_c"
+
+    # note how Command allows you to BOTH update the graph state AND route to the next node
+    return Command(
+        # this is the state update
+        update={"messages": [AIMessage(content=f"Hello {value}")]},
+        # this is a replacement for an edge
+        goto=goto,
+    )
+
+
+def node_b(state: AgentState):
+    print("Called B")
+    return {"messages": [AIMessage(content="Hello B")]}
+
+
+def node_c(state: AgentState):
+    print("Called C")
+    return {"messages": [AIMessage(content="Hello C")]}
+
+
+builder = StateGraph(AgentState)
+builder.add_edge(START, "node_a")
+builder.add_node(node_a)
+builder.add_node(node_b)
+builder.add_node(node_c)
+# NOTE: there are no edges between nodes A, B and C!
+
+command_agent = Agent(
+    name="command-agent",
+    description="A command agent.",
+    graph=builder.compile(),
+)