agent-harness-spec/template/app/agent.py

"""<AgentName> — Starlette entry point.

This module is the canonical agent harness. It implements the cross-cutting
plumbing every agent on the platform must ship:

    - GET /health                       liveness/readiness probe
    - GET /.well-known/agent.json       A2A discovery manifest
    - lifespan-driven self-registration with the agent-gateway
    - OTel log bridge wired at import time
    - LangChain instrumentation wired before any LLM client construction

Replace the example `/echo` route with your domain endpoints.
"""
from __future__ import annotations

import asyncio
import logging
import os
import time
from contextlib import asynccontextmanager

import httpx
from dotenv import load_dotenv
from starlette.applications import Starlette
from starlette.requests import Request
from starlette.responses import JSONResponse
from starlette.routing import Route

try:
    from openinference.instrumentation.langchain import LangChainInstrumentor
except ImportError:
    LangChainInstrumentor = None  # type: ignore[assignment]

from .logging_setup import configure_otlp_log_handler
from .metrics import echo_duration, echoes
from .skills import AGENT_CONFIG, AGENT_SKILLS

# ─── One-shot startup wiring ────────────────────────────────────────────
load_dotenv(".env.local")

LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
configure_otlp_log_handler(getattr(logging, LOG_LEVEL.upper(), logging.INFO))
logger = logging.getLogger(__name__)

if LangChainInstrumentor is not None:
    try:
        LangChainInstrumentor().instrument()
    except Exception as exc:  # pragma: no cover — instrumentation failure is non-fatal
        logger.warning("LangChainInstrumentor failed to attach: %s", exc)

# ─── Configuration ──────────────────────────────────────────────────────
REGISTRY_URL = os.getenv("REGISTRY_URL", "").strip()
AGENT_SELF_URL = os.getenv(
    "AGENT_SELF_URL",
    "http://<agent-name>.agents.svc.cluster.local",
).strip()
PORT = int(os.getenv("PORT", "8000"))


# ─── Skill serialisation (handles pydantic + protobuf AgentSkill) ───────
def _skill_dict(s) -> dict:
    """Serialise an A2A skill to a JSON-safe dict.

    a2a-sdk exposes AgentSkill as a protobuf Message in current releases, so
    `s.model_dump()` raises AttributeError. Handle pydantic v2, pydantic v1,
    protobuf, and fall back to attribute scraping.
    """
    if hasattr(s, "model_dump"):
        return s.model_dump()
    if hasattr(s, "DESCRIPTOR"):
        try:
            from google.protobuf.json_format import MessageToDict
            return MessageToDict(s, preserving_proto_field_name=True)
        except Exception:  # pragma: no cover
            pass
    if hasattr(s, "dict") and callable(getattr(s, "dict")):
        try:
            return s.dict()
        except Exception:  # pragma: no cover
            pass
    return {
        "id": getattr(s, "id", None),
        "name": getattr(s, "name", None),
        "description": getattr(s, "description", ""),
        "tags": list(getattr(s, "tags", []) or []),
        "examples": list(getattr(s, "examples", []) or []),
    }


# ─── Route handlers ─────────────────────────────────────────────────────
async def health_check(request: Request) -> JSONResponse:
    return JSONResponse({
        "status": "healthy",
        "service": AGENT_CONFIG["name"],
        "version": AGENT_CONFIG["version"],
    })


async def agent_manifest(request: Request) -> JSONResponse:
    """GET /.well-known/agent.json — discovery manifest consumed by agent-gateway."""
    return JSONResponse({
        "name": AGENT_CONFIG["name"],
        "version": AGENT_CONFIG["version"],
        "description": AGENT_CONFIG["description"],
        "url": "/",
        "skills": [_skill_dict(s) for s in AGENT_SKILLS],
        "capabilities": AGENT_CONFIG["capabilities"],
    })


async def echo(request: Request) -> JSONResponse:
    """POST /echo — example endpoint. Replace with your domain logic."""
    started = time.perf_counter()
    try:
        body = await request.json()
    except Exception:
        echoes.add(1, {"outcome": "bad_request"})
        return JSONResponse({"error": "invalid JSON body"}, status_code=400)

    try:
        response = {"received": body}
        echoes.add(1, {"outcome": "success"})
        return JSONResponse(response)
    except Exception as exc:
        echoes.add(1, {"outcome": "error"})
        logger.exception("echo failed: %s", exc)
        return JSONResponse({"error": str(exc)}, status_code=500)
    finally:
        echo_duration.record(time.perf_counter() - started)


# ─── Self-registration with the agent-gateway ───────────────────────────
async def _register_with_registry() -> None:
    """POST our endpoint to the registry; it will GET /.well-known/agent.json back."""
    if not REGISTRY_URL:
        logger.info("REGISTRY_URL not set — skipping self-registration")
        return

    await asyncio.sleep(2)  # let uvicorn finish binding before the gateway calls back

    url = f"{REGISTRY_URL.rstrip('/')}/agents/register-url"
    payload = {"endpoint": AGENT_SELF_URL}

    for attempt in range(3):
        try:
            async with httpx.AsyncClient(timeout=10.0) as client:
                resp = await client.post(url, json=payload)
            if resp.status_code in (200, 201):
                logger.info("Self-registered with agent-registry at %s", REGISTRY_URL)
                return
            logger.warning(
                "Registry registration attempt %d: HTTP %d — %s",
                attempt + 1, resp.status_code, resp.text[:200],
            )
        except Exception as exc:
            logger.warning("Registry registration attempt %d failed: %s", attempt + 1, exc)
        await asyncio.sleep(5)

    logger.error("Failed to self-register with agent-registry after 3 attempts")


@asynccontextmanager
async def lifespan(app: Starlette):
    register_task = asyncio.create_task(_register_with_registry())
    try:
        yield
    finally:
        register_task.cancel()


# ─── Application ────────────────────────────────────────────────────────
app = Starlette(
    routes=[
        Route("/health", health_check, methods=["GET"]),
        Route("/.well-known/agent.json", agent_manifest, methods=["GET"]),
        Route("/echo", echo, methods=["POST"]),
    ],
    lifespan=lifespan,
)


if __name__ == "__main__":
    import uvicorn

    uvicorn.run("app.agent:app", host="0.0.0.0", port=PORT, log_level=LOG_LEVEL.lower())