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

langgraph_agent_toolkit/client/client.py

@@ -0,0 +1,344 @@
+import json
+import os
+from collections.abc import AsyncGenerator, Generator
+from typing import Any
+
+import httpx
+
+from langgraph_agent_toolkit.schema import (
+    ChatHistory,
+    ChatHistoryInput,
+    ChatMessage,
+    Feedback,
+    ServiceMetadata,
+    StreamInput,
+    UserInput,
+)
+
+
+class AgentClientError(Exception):
+    pass
+
+
+class AgentClient:
+    """Client for interacting with the agent service."""
+
+    def __init__(
+        self,
+        base_url: str = "http://0.0.0.0",
+        agent: str = None,
+        timeout: float | None = None,
+        get_info: bool = True,
+        verify: bool = False,
+    ) -> None:
+        """Initialize the client.
+
+        Args:
+            base_url (str): The base URL of the agent service.
+            agent (str): The name of the default agent to use.
+            timeout (float, optional): The timeout for requests.
+            get_info (bool, optional): Whether to fetch agent information on init.
+                Default: True
+            verify (bool, optional): Whether to verify the agent information.
+                Default: False
+
+        """
+        self.base_url = base_url
+        self.auth_secret = os.getenv("AUTH_SECRET")
+        self.timeout = timeout
+        self.info: ServiceMetadata | None = None
+        self.agent: str | None = None
+        if get_info:
+            self.retrieve_info()
+        if agent:
+            self.update_agent(agent, verify=verify)
+
+    @property
+    def _headers(self) -> dict[str, str]:
+        headers = {}
+        if self.auth_secret:
+            headers["Authorization"] = f"Bearer {self.auth_secret}"
+        return headers
+
+    def retrieve_info(self) -> None:
+        try:
+            response = httpx.get(
+                f"{self.base_url}/info",
+                headers=self._headers,
+                timeout=self.timeout,
+            )
+            response.raise_for_status()
+        except httpx.HTTPError as e:
+            raise AgentClientError(f"Error getting service info: {e}")
+
+        self.info: ServiceMetadata = ServiceMetadata.model_validate(response.json())
+        if not self.agent or self.agent not in [a.key for a in self.info.agents]:
+            self.agent = self.info.default_agent
+
+    def update_agent(self, agent: str, verify: bool = True) -> None:
+        if verify:
+            if not self.info:
+                self.retrieve_info()
+            agent_keys = [a.key for a in self.info.agents]
+            if agent not in agent_keys:
+                raise AgentClientError(f"Agent {agent} not found in available agents: {', '.join(agent_keys)}")
+        self.agent = agent
+
+    async def ainvoke(
+        self,
+        message: str,
+        model: str | None = None,
+        thread_id: str | None = None,
+        agent_config: dict[str, Any] | None = None,
+    ) -> ChatMessage:
+        """Invoke the agent asynchronously. Only the final message is returned.
+
+        Args:
+            message (str): The message to send to the agent
+            model (str, optional): LLM model to use for the agent
+            thread_id (str, optional): Thread ID for continuing a conversation
+            agent_config (dict[str, Any], optional): Additional configuration to pass through to the agent
+
+        Returns:
+            AnyMessage: The response from the agent
+
+        """
+        if not self.agent:
+            raise AgentClientError("No agent selected. Use update_agent() to select an agent.")
+        request = UserInput(message=message)
+        if thread_id:
+            request.thread_id = thread_id
+        if model:
+            request.model = model
+        if agent_config:
+            request.agent_config = agent_config
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(
+                    f"{self.base_url}/{self.agent}/invoke",
+                    json=request.model_dump(),
+                    headers=self._headers,
+                    timeout=self.timeout,
+                )
+                response.raise_for_status()
+            except httpx.HTTPError as e:
+                raise AgentClientError(f"Error: {e}")
+
+        return ChatMessage.model_validate(response.json())
+
+    def invoke(
+        self,
+        message: str,
+        model: str | None = None,
+        thread_id: str | None = None,
+        agent_config: dict[str, Any] | None = None,
+    ) -> ChatMessage:
+        """Invoke the agent synchronously. Only the final message is returned.
+
+        Args:
+            message (str): The message to send to the agent
+            model (str, optional): LLM model to use for the agent
+            thread_id (str, optional): Thread ID for continuing a conversation
+            agent_config (dict[str, Any], optional): Additional configuration to pass through to the agent
+
+        Returns:
+            ChatMessage: The response from the agent
+
+        """
+        if not self.agent:
+            raise AgentClientError("No agent selected. Use update_agent() to select an agent.")
+        request = UserInput(message=message)
+        if thread_id:
+            request.thread_id = thread_id
+        if model:
+            request.model = model
+        if agent_config:
+            request.agent_config = agent_config
+        try:
+            response = httpx.post(
+                f"{self.base_url}/{self.agent}/invoke",
+                json=request.model_dump(),
+                headers=self._headers,
+                timeout=self.timeout,
+            )
+            response.raise_for_status()
+        except httpx.HTTPError as e:
+            raise AgentClientError(f"Error: {e}")
+
+        return ChatMessage.model_validate(response.json())
+
+    def _parse_stream_line(self, line: str) -> ChatMessage | str | None:
+        line = line.strip()
+        if line.startswith("data: "):
+            data = line[6:]
+            if data == "[DONE]":
+                return None
+            try:
+                parsed = json.loads(data)
+            except Exception as e:
+                raise Exception(f"Error JSON parsing message from server: {e}")
+            match parsed["type"]:
+                case "message":
+                    # Convert the JSON formatted message to an AnyMessage
+                    try:
+                        return ChatMessage.model_validate(parsed["content"])
+                    except Exception as e:
+                        raise Exception(f"Server returned invalid message: {e}")
+                case "token":
+                    # Yield the str token directly
+                    return parsed["content"]
+                case "error":
+                    raise Exception(parsed["content"])
+        return None
+
+    def stream(
+        self,
+        message: str,
+        model: str | None = None,
+        thread_id: str | None = None,
+        agent_config: dict[str, Any] | None = None,
+        stream_tokens: bool = True,
+    ) -> Generator[ChatMessage | str, None, None]:
+        """Stream the agent's response synchronously.
+
+        Each intermediate message of the agent process is yielded as a ChatMessage.
+        If stream_tokens is True (the default value), the response will also yield
+        content tokens from streaming models as they are generated.
+
+        Args:
+            message (str): The message to send to the agent
+            model (str, optional): LLM model to use for the agent
+            thread_id (str, optional): Thread ID for continuing a conversation
+            agent_config (dict[str, Any], optional): Additional configuration to pass through to the agent
+            stream_tokens (bool, optional): Stream tokens as they are generated
+                Default: True
+
+        Returns:
+            Generator[ChatMessage | str, None, None]: The response from the agent
+
+        """
+        if not self.agent:
+            raise AgentClientError("No agent selected. Use update_agent() to select an agent.")
+        request = StreamInput(message=message, stream_tokens=stream_tokens)
+        if thread_id:
+            request.thread_id = thread_id
+        if model:
+            request.model = model
+        if agent_config:
+            request.agent_config = agent_config
+        try:
+            with httpx.stream(
+                "POST",
+                f"{self.base_url}/{self.agent}/stream",
+                json=request.model_dump(),
+                headers=self._headers,
+                timeout=self.timeout,
+            ) as response:
+                response.raise_for_status()
+                for line in response.iter_lines():
+                    if line.strip():
+                        parsed = self._parse_stream_line(line)
+                        if parsed is None:
+                            break
+                        yield parsed
+        except httpx.HTTPError as e:
+            raise AgentClientError(f"Error: {e}")
+
+    async def astream(
+        self,
+        message: str,
+        model: str | None = None,
+        thread_id: str | None = None,
+        agent_config: dict[str, Any] | None = None,
+        stream_tokens: bool = True,
+    ) -> AsyncGenerator[ChatMessage | str, None]:
+        """Stream the agent's response asynchronously.
+
+        Each intermediate message of the agent process is yielded as an AnyMessage.
+        If stream_tokens is True (the default value), the response will also yield
+        content tokens from streaming modelsas they are generated.
+
+        Args:
+            message (str): The message to send to the agent
+            model (str, optional): LLM model to use for the agent
+            thread_id (str, optional): Thread ID for continuing a conversation
+            agent_config (dict[str, Any], optional): Additional configuration to pass through to the agent
+            stream_tokens (bool, optional): Stream tokens as they are generated
+                Default: True
+
+        Returns:
+            AsyncGenerator[ChatMessage | str, None]: The response from the agent
+
+        """
+        if not self.agent:
+            raise AgentClientError("No agent selected. Use update_agent() to select an agent.")
+        request = StreamInput(message=message, stream_tokens=stream_tokens)
+        if thread_id:
+            request.thread_id = thread_id
+        if model:
+            request.model = model
+        if agent_config:
+            request.agent_config = agent_config
+        async with httpx.AsyncClient() as client:
+            try:
+                async with client.stream(
+                    "POST",
+                    f"{self.base_url}/{self.agent}/stream",
+                    json=request.model_dump(),
+                    headers=self._headers,
+                    timeout=self.timeout,
+                ) as response:
+                    response.raise_for_status()
+                    async for line in response.aiter_lines():
+                        if line.strip():
+                            parsed = self._parse_stream_line(line)
+                            if parsed is None:
+                                break
+                            yield parsed
+            except httpx.HTTPError as e:
+                raise AgentClientError(f"Error: {e}")
+
+    async def acreate_feedback(self, run_id: str, key: str, score: float, kwargs: dict[str, Any] = {}) -> None:
+        """Create a feedback record for a run.
+
+        This is a simple wrapper for the LangSmith create_feedback API, so the
+        credentials can be stored and managed in the service rather than the client.
+        See: https://api.smith.langchain.com/redoc#tag/feedback/operation/create_feedback_api_v1_feedback_post
+        """
+        request = Feedback(run_id=run_id, key=key, score=score, kwargs=kwargs)
+        async with httpx.AsyncClient() as client:
+            try:
+                response = await client.post(
+                    f"{self.base_url}/feedback",
+                    json=request.model_dump(),
+                    headers=self._headers,
+                    timeout=self.timeout,
+                )
+                response.raise_for_status()
+                response.json()
+            except httpx.HTTPError as e:
+                raise AgentClientError(f"Error: {e}")
+
+    def get_history(
+        self,
+        thread_id: str,
+    ) -> ChatHistory:
+        """Get chat history.
+
+        Args:
+            thread_id (str, optional): Thread ID for identifying a conversation
+
+        """
+        request = ChatHistoryInput(thread_id=thread_id)
+        try:
+            response = httpx.post(
+                f"{self.base_url}/history",
+                json=request.model_dump(),
+                headers=self._headers,
+                timeout=self.timeout,
+            )
+            response.raise_for_status()
+        except httpx.HTTPError as e:
+            raise AgentClientError(f"Error: {e}")
+
+        return ChatHistory.model_validate(response.json())

langgraph_agent_toolkit/core/__init__.py

@@ -0,0 +1,5 @@
+from langgraph_agent_toolkit.core.models.factory import ModelFactory
+from langgraph_agent_toolkit.core.settings import settings
+
+
+__all__ = ["settings", "ModelFactory"]

langgraph_agent_toolkit/core/memory/base.py

@@ -0,0 +1,33 @@
+from abc import ABC, abstractmethod
+from contextlib import AbstractAsyncContextManager
+from typing import Any, TypeVar
+
+
+T = TypeVar("T", bound=Any)
+
+
+class BaseMemoryBackend(ABC):
+    """Base class for memory backends."""
+
+    @abstractmethod
+    def validate_config(self) -> bool:
+        """Validate that all necessary configuration is set.
+
+        Returns:
+            True if configuration is valid
+
+        Raises:
+            ValueError: If required configuration is missing
+
+        """
+        pass
+
+    @abstractmethod
+    def get_checkpoint_saver(self) -> AbstractAsyncContextManager[T]:
+        """Get the checkpoint saver for the memory backend.
+
+        Returns:
+            A configured checkpoint saver
+
+        """
+        pass

langgraph_agent_toolkit/core/memory/factory.py

@@ -0,0 +1,30 @@
+from langgraph_agent_toolkit.core.memory.base import BaseMemoryBackend
+from langgraph_agent_toolkit.core.memory.postgres import PostgresMemoryBackend
+from langgraph_agent_toolkit.core.memory.sqlite import SQLiteMemoryBackend
+from langgraph_agent_toolkit.core.memory.types import MemoryBackends
+
+
+class MemoryFactory:
+    """Factory for creating memory backend instances."""
+
+    @staticmethod
+    def create(backend: MemoryBackends) -> BaseMemoryBackend:
+        """Create and return a memory backend instance.
+
+        Args:
+            backend: The memory backend to create
+
+        Returns:
+            An instance of the requested memory backend
+
+        Raises:
+            ValueError: If the requested backend is not supported
+
+        """
+        match backend:
+            case MemoryBackends.POSTGRES:
+                return PostgresMemoryBackend()
+            case MemoryBackends.SQLITE:
+                return SQLiteMemoryBackend()
+            case _:
+                raise ValueError(f"Unsupported memory backend: {backend}")

langgraph_agent_toolkit/core/memory/postgres.py

@@ -0,0 +1,76 @@
+from collections.abc import AsyncGenerator
+from contextlib import AbstractAsyncContextManager, asynccontextmanager
+
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from psycopg.rows import dict_row
+from psycopg_pool import AsyncConnectionPool
+
+from langgraph_agent_toolkit.core.memory.base import BaseMemoryBackend
+from langgraph_agent_toolkit.core.settings import settings
+from langgraph_agent_toolkit.helper.logging import logger
+
+
+class PostgresMemoryBackend(BaseMemoryBackend):
+    """PostgreSQL implementation of memory backend."""
+
+    def validate_config(self) -> bool:
+        """Validate that all required PostgreSQL configuration is present."""
+        required_vars = [
+            "POSTGRES_USER",
+            "POSTGRES_PASSWORD",
+            "POSTGRES_HOST",
+            "POSTGRES_PORT",
+            "POSTGRES_DB",
+        ]
+
+        missing = [var for var in required_vars if not getattr(settings, var, None)]
+        if missing:
+            raise ValueError(
+                f"Missing required PostgreSQL configuration: {', '.join(missing)}. "
+                "These environment variables must be set to use PostgreSQL persistence."
+            )
+        return True
+
+    def get_connection_string(self) -> str:
+        """Build and return the PostgreSQL connection string from settings."""
+        return (
+            f"postgresql://{settings.POSTGRES_USER}:"
+            f"{settings.POSTGRES_PASSWORD.get_secret_value()}@"
+            f"{settings.POSTGRES_HOST}:{settings.POSTGRES_PORT}/"
+            f"{settings.POSTGRES_DB}"
+        )
+
+    @asynccontextmanager
+    async def get_saver(self) -> AsyncGenerator[AsyncPostgresSaver, None]:
+        """Asynchronous context manager for acquiring and releasing the PostgreSQL connection pool.
+
+        Yields:
+            AsyncPostgresSaver: The database saver instance
+
+        """
+        conn_string = self.get_connection_string()
+
+        logger.info(
+            f"Creating PostgreSQL connection pool: min_size={settings.POSTGRES_MIN_SIZE}, "
+            f"max_size={settings.POSTGRES_POOL_SIZE}, max_idle={settings.POSTGRES_MAX_IDLE}"
+        )
+
+        # Use AsyncConnectionPool as an async context manager
+        async with AsyncConnectionPool(
+            conn_string,
+            min_size=settings.POSTGRES_MIN_SIZE,
+            max_size=settings.POSTGRES_POOL_SIZE,
+            max_idle=settings.POSTGRES_MAX_IDLE,
+            kwargs=dict(autocommit=True, prepare_threshold=0, row_factory=dict_row),
+        ) as pool:
+            logger.info("PostgreSQL connection pool opened successfully")
+
+            try:
+                yield AsyncPostgresSaver(conn=pool)
+            finally:
+                logger.info("PostgreSQL connection pool will be closed automatically")
+
+    def get_checkpoint_saver(self) -> AbstractAsyncContextManager[AsyncPostgresSaver]:
+        """Initialize and return a PostgreSQL saver instance."""
+        self.validate_config()
+        return self.get_saver()

langgraph_agent_toolkit/core/memory/sqlite.py

@@ -0,0 +1,21 @@
+from contextlib import AbstractAsyncContextManager
+
+from langgraph.checkpoint.sqlite.aio import AsyncSqliteSaver
+
+from langgraph_agent_toolkit.core.memory.base import BaseMemoryBackend
+from langgraph_agent_toolkit.core.settings import settings
+
+
+class SQLiteMemoryBackend(BaseMemoryBackend):
+    """SQLite implementation of memory backend."""
+
+    def validate_config(self) -> bool:
+        """Validate that SQLite configuration is present."""
+        if not getattr(settings, "SQLITE_DB_PATH", None):
+            raise ValueError("Missing SQLITE_DB_PATH configuration. This must be set to use SQLite persistence.")
+        return True
+
+    def get_checkpoint_saver(self) -> AbstractAsyncContextManager[AsyncSqliteSaver]:
+        """Initialize and return a SQLite saver instance."""
+        self.validate_config()
+        return AsyncSqliteSaver.from_conn_string(settings.SQLITE_DB_PATH)

langgraph_agent_toolkit/core/memory/types.py

@@ -0,0 +1,6 @@
+from enum import StrEnum, auto
+
+
+class MemoryBackends(StrEnum):
+    POSTGRES = auto()
+    SQLITE = auto()

langgraph_agent_toolkit/core/models/__init__.py

@@ -0,0 +1,5 @@
+from langgraph_agent_toolkit.core.models.chat_openai import ChatOpenAIPatched
+from langgraph_agent_toolkit.core.models.fake import FakeToolModel
+
+
+__all__ = ["ChatOpenAIPatched", "FakeToolModel"]

langgraph_agent_toolkit/core/models/chat_openai.py

@@ -0,0 +1,21 @@
+from typing import Dict, Optional, Union
+
+import openai
+from langchain_core.outputs import ChatResult
+from langchain_openai import ChatOpenAI
+
+
+class ChatOpenAIPatched(ChatOpenAI):
+    def _create_chat_result(
+        self,
+        response: Union[dict, openai.BaseModel],
+        generation_info: Optional[Dict] = None,
+    ) -> ChatResult:
+        for idx in range(len(response.choices)):
+            resp = response.choices[idx]
+            # fix the role of the message
+            if resp.message.role.startswith("assistant"):
+                resp.message.role = "assistant"
+            response.choices[idx] = resp
+
+        return super()._create_chat_result(response, generation_info)

langgraph_agent_toolkit/core/models/factory.py

@@ -0,0 +1,118 @@
+import warnings
+from functools import cache
+from typing import (
+    Any,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+    cast,
+)
+
+from langchain.chat_models.base import _ConfigurableModel, _init_chat_model_helper
+from langchain_core.language_models import BaseChatModel
+from langchain_core.runnables import RunnableSerializable
+from typing_extensions import TypeAlias
+
+from langgraph_agent_toolkit.core.models import ChatOpenAIPatched, FakeToolModel
+from langgraph_agent_toolkit.core.settings import settings
+from langgraph_agent_toolkit.helper.constants import DEFAULT_OPENAI_COMPATIBLE_MODEL_PARAMS
+from langgraph_agent_toolkit.schema.models import (
+    AllModelEnum,
+    FakeModelName,
+    OpenAICompatibleName,
+)
+
+
+ModelT: TypeAlias = ChatOpenAIPatched | FakeToolModel | RunnableSerializable | _ConfigurableModel
+
+
+class ModelFactory:
+    """Factory for creating model instances."""
+
+    # Map model enum names to their respective API model names
+    _MODEL_TABLE = {
+        OpenAICompatibleName.OPENAI_COMPATIBLE: settings.COMPATIBLE_MODEL,
+        FakeModelName.FAKE: "fake",
+    }
+
+    @staticmethod
+    def __init_chat_model_helper(model: str, *, model_provider: Optional[str] = None, **kwargs: Any) -> BaseChatModel:
+        if model_provider == "openai":
+            return ChatOpenAIPatched(model_name=model, **kwargs)
+        else:
+            return _init_chat_model_helper(model, model_provider=model_provider, **kwargs)
+
+    @staticmethod
+    def init_chat_model(
+        model: Optional[str] = None,
+        *,
+        model_provider: Optional[str] = None,
+        configurable_fields: Optional[Union[Literal["any"], List[str], Tuple[str, ...]]] = None,
+        config_prefix: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Union[BaseChatModel, _ConfigurableModel]:
+        if not model and not configurable_fields:
+            configurable_fields = ("model", "model_provider")
+        config_prefix = config_prefix or ""
+
+        if config_prefix and not configurable_fields:
+            warnings.warn(
+                f"{config_prefix=} has been set but no fields are configurable. Set "
+                f"`configurable_fields=(...)` to specify the model params that are "
+                f"configurable."
+            )
+
+        if not configurable_fields:
+            return ModelFactory.__init_chat_model_helper(cast(str, model), model_provider=model_provider, **kwargs)
+        else:
+            if model:
+                kwargs["model"] = model
+            if model_provider:
+                kwargs["model_provider"] = model_provider
+            return _ConfigurableModel(
+                default_config=kwargs,
+                config_prefix=config_prefix,
+                configurable_fields=configurable_fields,
+            )
+
+    @staticmethod
+    @cache
+    def create(model_name: AllModelEnum) -> ModelT:
+        """Create and return a model instance.
+
+        Args:
+            model_name: The model to create from AllModelEnum
+
+        Returns:
+            An instance of the requested model
+
+        Raises:
+            ValueError: If the requested model is not supported
+
+        """
+        api_model_name = ModelFactory._MODEL_TABLE.get(model_name)
+        if not api_model_name:
+            raise ValueError(f"Unsupported model: {model_name}")
+
+        match model_name:
+            case name if name in OpenAICompatibleName:
+                if not settings.COMPATIBLE_BASE_URL or not settings.COMPATIBLE_MODEL:
+                    raise ValueError("OpenAICompatible base url and endpoint must be configured")
+
+                model = ModelFactory.init_chat_model(
+                    model=settings.COMPATIBLE_MODEL,
+                    model_provider="openai",
+                    configurable_fields=("temperature", "max_tokens", "top_p", "streaming"),
+                    config_prefix="agent",
+                    openai_api_base=settings.COMPATIBLE_BASE_URL,
+                    openai_api_key=settings.COMPATIBLE_API_KEY,
+                    **DEFAULT_OPENAI_COMPATIBLE_MODEL_PARAMS,
+                )
+
+                return model
+            case name if name in FakeModelName:
+                return FakeToolModel(responses=["This is a test response from the fake model."])
+            case _:
+                raise ValueError(f"Unsupported model: {model_name}")