langgraph-agent-common 2.0.0__py3-none-any.whl

langgraph_a2a_api/fastapi_application.py

@@ -0,0 +1,152 @@
+import json
+from pathlib import Path
+from typing import Any
+
+from fastapi import FastAPI
+from fastapi.openapi.utils import get_openapi
+from fastapi.responses import FileResponse
+from fastapi.staticfiles import StaticFiles
+from starlette.requests import Request
+from starlette.responses import Response
+from starlette.routing import BaseRoute, Mount
+
+from a2a.server.request_handlers import DefaultRequestHandler
+from a2a.server.routes.agent_card_routes import create_agent_card_routes
+from a2a.server.routes.jsonrpc_routes import create_jsonrpc_routes
+from a2a.server.routes.rest_routes import create_rest_routes
+from a2a.types import AgentCard
+from a2a.utils import constants
+
+_OPENAPI_SCHEMA_PATH = Path(__file__).parent / "resources" / "openapi_schema.json"
+_STATIC_PATH = Path(__file__).parent / "resources" / "static"
+
+
+def _load_openapi_schema() -> dict:
+    with open(_OPENAPI_SCHEMA_PATH) as f:
+        return json.load(f)
+
+
+def _create_a2a_routes(agent_card: AgentCard, request_handler: DefaultRequestHandler) -> list[BaseRoute]:
+    return [
+        *create_agent_card_routes(agent_card),
+        *create_jsonrpc_routes(request_handler, rpc_url='/'),
+        *create_rest_routes(request_handler),
+    ]
+
+
+async def _schema_placeholder() -> Response:
+    return Response(status_code=501)
+
+
+def _derive_route_tag(path: str) -> str:
+    if path == '/':
+        return 'jsonrpc'
+    if path == '/.well-known/agent-card.json':
+        return 'agent-card'
+    return 'rest'
+
+
+def _derive_route_summary(route: BaseRoute) -> str:
+    endpoint = getattr(route, 'endpoint', None)
+    path = getattr(route, 'path', 'A2A route')
+    if endpoint is None:
+        return f'Handle {path}'
+    doc = getattr(endpoint, '__doc__', None)
+    if doc:
+        return doc.strip().splitlines()[0]
+    name = getattr(route, 'name', None) or getattr(endpoint, '__name__', None)
+    if name:
+        return name.replace('_', ' ').strip().capitalize()
+    return f'Handle {path}'
+
+
+def _register_routes(app: FastAPI, routes: list[BaseRoute]):
+    for route in routes:
+        app.router.routes.append(route)
+        if isinstance(route, Mount):
+            continue
+        methods = sorted(method for method in (getattr(route, 'methods', None) or []) if method != 'HEAD')
+        if not methods:
+            continue
+        app.add_api_route(
+            path=route.path,
+            endpoint=_schema_placeholder,
+            methods=methods,
+            include_in_schema=True,
+            tags=[_derive_route_tag(route.path)],
+            summary=_derive_route_summary(route),
+            name=f"openapi:{','.join(methods)}:{route.path}",
+        )
+
+
+def _add_common_routes(app: FastAPI):
+    if not _STATIC_PATH.exists():
+        return
+
+    app.mount('/static', StaticFiles(directory=str(_STATIC_PATH)), name='static')
+
+    @app.get('/chat')
+    async def chat_ui():
+        return FileResponse(str(_STATIC_PATH / 'index.html'))
+
+
+def build_fastapi_application(agent_card: AgentCard, request_handler: DefaultRequestHandler, lifespan: Any, root_path: str = "") -> FastAPI:
+    app = FastAPI(
+        lifespan=lifespan,
+        title=agent_card.name or 'A2A Agent API',
+        description=agent_card.description or 'A2A Protocol API with JSON-RPC and REST endpoints',
+        version=agent_card.version or '1.0.0',
+        docs_url='/docs',
+        redoc_url='/redoc',
+        openapi_url='/openapi.json',
+        root_path=root_path,
+    )
+
+    routes = _create_a2a_routes(agent_card, request_handler)
+
+    @app.middleware('http')
+    async def ensure_a2a_version_header(request: Request, call_next):
+        if constants.VERSION_HEADER.encode() not in request.headers.raw:
+            request.scope['headers'] = [
+                *request.scope['headers'],
+                (constants.VERSION_HEADER.lower().encode(), constants.PROTOCOL_VERSION_1_0.encode()),
+            ]
+        return await call_next(request)
+
+    _add_common_routes(app)
+    _register_routes(app, routes)
+    _customize_openapi(app)
+    return app
+
+
+def _customize_openapi(app: FastAPI):
+    external_schema = _load_openapi_schema()
+
+    def openapi_schema():
+        if app.openapi_schema:
+            return app.openapi_schema
+
+        schema = get_openapi(
+            title=app.title,
+            version=app.version,
+            openapi_version=app.openapi_version,
+            description=app.description,
+            routes=app.routes,
+        )
+
+        external_schemas = external_schema.get('components', {}).get('schemas', {})
+        if external_schemas:
+            schema.setdefault('components', {}).setdefault('schemas', {}).update(external_schemas)
+
+        external_paths = external_schema.get('paths', {})
+        app_paths = schema.get('paths', {})
+        for path, methods in external_paths.items():
+            if path in app_paths:
+                for method, config in methods.items():
+                    if method in app_paths[path]:
+                        app_paths[path][method].update(config)
+
+        app.openapi_schema = schema
+        return app.openapi_schema
+
+    app.openapi = openapi_schema

langgraph_a2a_api/langgraph/checkpointer.py

@@ -0,0 +1,102 @@
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.checkpoint.memory import InMemorySaver
+from langgraph.checkpoint.base import BaseCheckpointSaver
+
+from contextlib import asynccontextmanager
+from collections.abc import AsyncIterator
+
+from psycopg.rows import DictRow
+from psycopg import AsyncCursor
+from psycopg_pool import AsyncConnectionPool
+
+import re
+import os
+import textwrap
+
+class AsyncPostgresSaverWrapper(AsyncPostgresSaver):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.checkpoint_table_prefix = None
+        self.GET_EXPIRED_THREADS_SQL = textwrap.dedent("""
+        SELECT thread_id
+        FROM checkpoints
+        GROUP BY thread_id
+        HAVING MAX((checkpoint ->> 'ts')::timestamptz) < (NOW() - (%s || ' MINUTES')::INTERVAL);
+        """)
+
+    def set_checkpoint_table_prefix(self, checkpoint_table_prefix: str):
+        self.checkpoint_table_prefix = checkpoint_table_prefix
+
+    @asynccontextmanager
+    async def _cursor(self, *, pipeline: bool = False) -> AsyncIterator[AsyncCursor[DictRow]]:
+        async with super()._cursor(pipeline = pipeline) as cur:
+            yield AsyncCursorWrapper(cur, self.checkpoint_table_prefix)
+
+    async def sweep_threads(self, ttl: int) -> int:
+        async with self._cursor() as cur:
+            await cur.execute(self.GET_EXPIRED_THREADS_SQL, (ttl,), binary=True)
+            values = await cur.fetchmany(1024)
+        for value in values:
+            await self.adelete_thread(value["thread_id"])
+        return len(values)
+
+class AsyncCursorWrapper():
+    def __init__(self, cursor: AsyncCursor, checkpoint_table_prefix: str):
+        self._cursor = cursor
+        self.checkpoint_table_prefix = checkpoint_table_prefix
+
+    def _add_checkpoint_table_prefix(self, query: str) -> str:
+        checkpoint_migrations_regex = r"(?<=[ ])checkpoint_migrations(?=[_ .(]|\n)"
+        checkpoints_regex = r"(?<=[ ])checkpoints(?=[_ .(]|\n)"
+        checkpoint_blobs_regex = r"(?<=[ ])checkpoint_blobs(?=[_ .(]|\n)"
+        checkpoint_writes_regex = r"(?<=[ ])checkpoint_writes(?=[_ .(]|\n)"
+        checkpoint_index_regex = r"(?<=\W)CREATE INDEX CONCURRENTLY(?=[ ])"
+
+        checkpoint_migrations_sub = re.sub(checkpoint_migrations_regex, f"{self.checkpoint_table_prefix}__checkpoint_migrations", query)
+        checkpoints_sub = re.sub(checkpoints_regex, f"{self.checkpoint_table_prefix}__checkpoints", checkpoint_migrations_sub)
+        checkpoint_blobs_sub = re.sub(checkpoint_blobs_regex, f"{self.checkpoint_table_prefix}__checkpoint_blobs", checkpoints_sub)
+        checkpoint_writes_sub = re.sub(checkpoint_writes_regex, f"{self.checkpoint_table_prefix}__checkpoint_writes", checkpoint_blobs_sub)
+        checkpoint_index_sub = re.sub(checkpoint_index_regex, "CREATE INDEX", checkpoint_writes_sub)
+        return checkpoint_index_sub
+
+
+    def execute(self, *args, **kwargs):
+        modified_args = (self._add_checkpoint_table_prefix(args[0]),) + args[1:]
+        return self._cursor.execute(*modified_args, **kwargs)
+
+    def executemany(self, *args, **kwargs):
+        modified_args = (self._add_checkpoint_table_prefix(args[0]),) + args[1:]
+        return self._cursor.executemany(*modified_args, **kwargs)
+
+    def __getattr__(self, name):
+        return getattr(self._cursor, name)
+
+@asynccontextmanager
+async def get_async_postgres_checkpointer(checkpoint_table_prefix: str) -> AsyncIterator[AsyncPostgresSaver]:
+    postgres_uri = os.environ["POSTGRES_URI"]
+    async with AsyncConnectionPool(postgres_uri, min_size=1, max_size=8, max_idle=600, timeout=30) as pool:
+        checkpointer = AsyncPostgresSaverWrapper(conn=pool)
+        checkpointer.set_checkpoint_table_prefix(checkpoint_table_prefix)
+        await checkpointer.setup()
+        yield checkpointer
+
+
+class InMemorySaverWrapper(InMemorySaver):
+    async def sweep_threads(self, ttl: int) -> int:
+        return 0
+
+
+@asynccontextmanager
+async def get_inmemory_checkpointer(checkpoint_table_prefix: str) -> AsyncIterator[BaseCheckpointSaver]:
+    yield InMemorySaverWrapper()
+
+
+@asynccontextmanager
+async def get_checkpointer(checkpoint_table_prefix: str) -> AsyncIterator[BaseCheckpointSaver]:
+    if os.getenv("POSTGRES_URI"):
+        async with get_async_postgres_checkpointer(checkpoint_table_prefix) as cp:
+            yield cp
+    else:
+        async with get_inmemory_checkpointer(checkpoint_table_prefix) as cp:
+            yield cp
+

langgraph_a2a_api/langgraph/graph.py

@@ -0,0 +1,49 @@
+import os
+import importlib
+from typing import NamedTuple
+from random import choice
+import sys
+from contextlib import asynccontextmanager
+from collections.abc import AsyncIterator
+import inflection
+
+from langgraph.graph.state import CompiledStateGraph
+from langgraph.types import Checkpointer
+
+from langgraph.graph.message import MessagesState
+from a2a.types import AgentSkill
+
+
+class GraphSpec(NamedTuple):
+    name: str
+    description: str
+    version: str
+    skills: list[AgentSkill]
+
+
+def get_graph() -> (CompiledStateGraph, GraphSpec):
+    path_and_var = os.environ["GRAPH_PATH_AND_VAR"]
+    (path, var) = path_and_var.rsplit(":", maxsplit = 1)
+    module_name = "".join(choice("abcdefghijklmnopqrstuvwxyz") for _ in range(24))
+    module_spec = importlib.util.spec_from_file_location(module_name, path)
+    if module_spec is None:
+        raise ValueError(f"Could not find python file: {path} for graph: {var}")
+    module = importlib.util.module_from_spec(module_spec)
+    sys.modules[module_name] = module
+    module_spec.loader.exec_module(module)
+
+    graph = module.__dict__[var]
+    if not graph.get_name():
+        raise Exception(f"Langgraph graph must include graph name")
+    if not issubclass(type(graph.InputType), type(MessagesState)):
+        raise Exception(f"Langgraph state must be a subclass of {MessagesState}.")
+    description = graph.config.get("metadata", {}).get("description", "")
+    version = graph.config.get("metadata", {}).get("version", "1.0.0")
+    graph_spec = GraphSpec(
+        name = inflection.underscore(graph.get_name()),
+        description = description,
+        version = version,
+        skills = graph.config.get("metadata", {}).get("skills", []),
+    )
+    return graph, graph_spec
+

langgraph_a2a_api/langgraph/store.py

@@ -0,0 +1,64 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
+
+import os
+
+from psycopg.rows import dict_row
+from psycopg_pool import AsyncConnectionPool
+
+from langgraph.store.base import BaseStore
+from langgraph.store.postgres.aio import AsyncPostgresStore
+from langgraph.store.memory import InMemoryStore
+
+
+class AsyncPostgresStoreWrapper(AsyncPostgresStore):
+    async def sweep(self, ttl: int) -> int:
+        async with self._cursor() as cur:
+            await cur.execute(
+                """
+                DELETE FROM store
+                WHERE updated_at < (NOW() - (%s || ' MINUTES')::INTERVAL)
+                """,
+                (ttl,),
+            )
+            return cur.rowcount
+
+
+@asynccontextmanager
+async def get_async_postgres_store() -> AsyncIterator[BaseStore]:
+    postgres_uri = os.environ["POSTGRES_URI"]
+    async with AsyncConnectionPool(
+        postgres_uri,
+        min_size = 1,
+        max_size = 8,
+        max_idle = 600,
+        timeout = 30,
+        kwargs = {
+            "autocommit": True,
+            "prepare_threshold": 0,
+            "row_factory": dict_row,
+        },
+    ) as pool:
+        store = AsyncPostgresStoreWrapper(conn = pool)
+        await store.setup()
+        yield store
+
+
+class InMemoryStoreWrapper(InMemoryStore):
+    async def sweep(self, ttl: int) -> int:
+        return 0
+
+
+@asynccontextmanager
+async def get_inmemory_store() -> AsyncIterator[BaseStore]:
+    yield InMemoryStoreWrapper()
+
+
+@asynccontextmanager
+async def get_store() -> AsyncIterator[BaseStore]:
+    if os.getenv("POSTGRES_URI"):
+        async with get_async_postgres_store() as s:
+            yield s
+    else:
+        async with get_inmemory_store() as s:
+            yield s

langgraph_a2a_api/langgraph/validator.py

@@ -0,0 +1,42 @@
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict, ValidationError
+
+
+class HitlAction(BaseModel):
+    name: str
+    args: dict[str, Any]
+
+
+class ApproveDecision(BaseModel):
+    type: Literal["approve"]
+
+
+class EditDecision(BaseModel):
+    type: Literal["edit"]
+    edited_action: HitlAction
+
+
+class RejectDecision(BaseModel):
+    type: Literal["reject"]
+    message: str | None = None
+
+
+class RespondDecision(BaseModel):
+    type: Literal["respond"]
+    message: str
+
+
+class HITLResponseModel(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+    hitl_id: str
+    decision: ApproveDecision | EditDecision | RejectDecision | RespondDecision
+
+
+def validate_HITLResponse(payload: Any) -> bool:
+    try:
+        HITLResponseModel.model_validate(payload)
+    except ValidationError:
+        return False
+    return True

langgraph_a2a_api/resources/openapi_schema.json

@@ -0,0 +1,277 @@
+{
+  "components": {
+    "schemas": {
+      "Part": {
+        "type": "object",
+        "properties": {
+          "text": { "type": "string", "description": "Text content" },
+          "raw": { "type": "string", "format": "byte", "description": "Base64 encoded raw bytes" },
+          "url": { "type": "string", "format": "uri", "description": "URL to external content" },
+          "data": { "description": "Structured data (JSON object/value)" },
+          "metadata": { "type": "object", "description": "Optional metadata" },
+          "filename": { "type": "string", "description": "File name" },
+          "mediaType": { "type": "string", "description": "MIME type" }
+        }
+      },
+      "Message": {
+        "type": "object",
+        "required": ["messageId", "role", "parts"],
+        "properties": {
+          "messageId": { "type": "string", "description": "Unique message ID (UUID)" },
+          "contextId": { "type": "string", "description": "Context ID for grouping related interactions" },
+          "taskId": { "type": "string", "description": "Task ID this message belongs to" },
+          "role": {
+            "type": "string",
+            "enum": ["ROLE_USER", "ROLE_AGENT"],
+            "description": "Message sender role"
+          },
+          "parts": {
+            "type": "array",
+            "items": { "$ref": "#/components/schemas/Part" },
+            "description": "Content parts of the message"
+          },
+          "metadata": { "type": "object", "description": "Optional metadata for extensions" },
+          "extensions": {
+            "type": "array",
+            "items": { "type": "string" },
+            "description": "Extension URIs relevant to this message"
+          },
+          "referenceTaskIds": {
+            "type": "array",
+            "items": { "type": "string" },
+            "description": "Related task IDs for context"
+          }
+        }
+      },
+      "SendMessageConfiguration": {
+        "type": "object",
+        "properties": {
+          "acceptedOutputModes": {
+            "type": "array",
+            "items": { "type": "string" },
+            "description": "Accepted output MIME types"
+          },
+          "taskPushNotificationConfig": { "$ref": "#/components/schemas/TaskPushNotificationConfig" },
+          "historyLength": { "type": "integer", "description": "Number of recent history messages to return" },
+          "returnImmediately": { "type": "boolean", "description": "Return immediately without waiting" }
+        }
+      },
+      "SendMessageRequest": {
+        "type": "object",
+        "required": ["message"],
+        "properties": {
+          "tenant": { "type": "string", "description": "Tenant identifier" },
+          "message": { "$ref": "#/components/schemas/Message" },
+          "configuration": { "$ref": "#/components/schemas/SendMessageConfiguration" },
+          "metadata": { "type": "object", "description": "Optional metadata" }
+        }
+      },
+      "TaskStatus": {
+        "type": "object",
+        "required": ["state"],
+        "properties": {
+          "state": {
+            "type": "string",
+            "enum": [
+              "TASK_STATE_UNSPECIFIED",
+              "TASK_STATE_SUBMITTED",
+              "TASK_STATE_WORKING",
+              "TASK_STATE_COMPLETED",
+              "TASK_STATE_FAILED",
+              "TASK_STATE_CANCELED",
+              "TASK_STATE_INPUT_REQUIRED",
+              "TASK_STATE_REJECTED",
+              "TASK_STATE_AUTH_REQUIRED"
+            ],
+            "description": "Current task state"
+          },
+          "message": { "$ref": "#/components/schemas/Message" },
+          "timestamp": { "type": "string", "format": "date-time", "description": "ISO 8601 timestamp" }
+        }
+      },
+      "Artifact": {
+        "type": "object",
+        "required": ["artifactId", "parts"],
+        "properties": {
+          "artifactId": { "type": "string", "description": "Unique artifact ID" },
+          "name": { "type": "string", "description": "Artifact name" },
+          "description": { "type": "string", "description": "Artifact description" },
+          "parts": {
+            "type": "array",
+            "items": { "$ref": "#/components/schemas/Part" },
+            "description": "Content parts"
+          },
+          "metadata": { "type": "object", "description": "Optional metadata" },
+          "extensions": {
+            "type": "array",
+            "items": { "type": "string" },
+            "description": "Extension URIs"
+          }
+        }
+      },
+      "Task": {
+        "type": "object",
+        "required": ["id", "contextId", "status"],
+        "properties": {
+          "id": { "type": "string", "description": "Unique task ID" },
+          "contextId": { "type": "string", "description": "Context ID for this task" },
+          "status": { "$ref": "#/components/schemas/TaskStatus" },
+          "artifacts": {
+            "type": "array",
+            "items": { "$ref": "#/components/schemas/Artifact" },
+            "description": "Generated artifacts"
+          },
+          "history": {
+            "type": "array",
+            "items": { "$ref": "#/components/schemas/Message" },
+            "description": "Conversation history"
+          },
+          "metadata": { "type": "object", "description": "Optional metadata" }
+        }
+      },
+      "SendMessageResponse": {
+        "type": "object",
+        "properties": {
+          "task": { "$ref": "#/components/schemas/Task" },
+          "message": { "$ref": "#/components/schemas/Message" }
+        }
+      },
+      "TaskPushNotificationConfig": {
+        "type": "object",
+        "required": ["id", "taskId", "url"],
+        "properties": {
+          "tenant": { "type": "string" },
+          "id": { "type": "string", "description": "Config ID" },
+          "taskId": { "type": "string", "description": "Task ID" },
+          "url": { "type": "string", "format": "uri", "description": "Callback URL" },
+          "token": { "type": "string", "description": "Auth token" },
+          "authentication": { "$ref": "#/components/schemas/AuthenticationInfo" }
+        }
+      },
+      "AuthenticationInfo": {
+        "type": "object",
+        "properties": {
+          "scheme": { "type": "string", "description": "Auth scheme" },
+          "credentials": { "type": "string", "description": "Auth credentials" }
+        }
+      },
+      "JSONRPCRequest": {
+        "type": "object",
+        "required": ["jsonrpc", "method"],
+        "properties": {
+          "jsonrpc": { "type": "string", "default": "2.0" },
+          "id": { "type": "string", "description": "Request ID" },
+          "method": {
+            "type": "string",
+            "enum": [
+              "SendMessage",
+              "SendStreamingMessage",
+              "GetTask",
+              "ListTasks",
+              "CancelTask",
+              "CreateTaskPushNotificationConfig",
+              "GetTaskPushNotificationConfig",
+              "ListTaskPushNotificationConfigs",
+              "DeleteTaskPushNotificationConfig",
+              "SubscribeToTask",
+              "GetExtendedAgentCard"
+            ]
+          },
+          "params": { "type": "object" }
+        }
+      }
+    }
+  },
+  "paths": {
+    "/message:send": {
+      "post": {
+        "requestBody": {
+          "required": true,
+          "content": {
+            "application/json": {
+              "schema": { "$ref": "#/components/schemas/SendMessageRequest" }
+            }
+          }
+        },
+        "responses": {
+          "200": {
+            "description": "Successful response",
+            "content": {
+              "application/json": {
+                "schema": { "$ref": "#/components/schemas/SendMessageResponse" }
+              }
+            }
+          }
+        }
+      }
+    },
+    "/message:stream": {
+      "post": {
+        "requestBody": {
+          "required": true,
+          "content": {
+            "application/json": {
+              "schema": { "$ref": "#/components/schemas/SendMessageRequest" }
+            }
+          }
+        },
+        "responses": {
+          "200": {
+            "description": "SSE stream of events",
+            "content": {
+              "text/event-stream": {
+                "schema": { "type": "string" }
+              }
+            }
+          }
+        }
+      }
+    },
+    "/tasks/{id}": {
+      "get": {
+        "parameters": [
+          {
+            "name": "id",
+            "in": "path",
+            "required": true,
+            "schema": { "type": "string" },
+            "description": "Task ID"
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Task details",
+            "content": {
+              "application/json": {
+                "schema": { "$ref": "#/components/schemas/Task" }
+              }
+            }
+          }
+        }
+      }
+    },
+    "/tasks/{id}:cancel": {
+      "post": {
+        "parameters": [
+          {
+            "name": "id",
+            "in": "path",
+            "required": true,
+            "schema": { "type": "string" },
+            "description": "Task ID to cancel"
+          }
+        ],
+        "responses": {
+          "200": {
+            "description": "Canceled task",
+            "content": {
+              "application/json": {
+                "schema": { "$ref": "#/components/schemas/Task" }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}