agent-platform/services/team-service/app/infrastructure/agent_client.py

import httpx

from core_domain import AgentRunContract
from core_shared import JSONValue


class AgentServiceClientError(Exception):
    pass


class AgentServiceClient:
    def __init__(self, base_url: str, timeout_seconds: float = 30.0) -> None:
        self.base_url = base_url.rstrip("/")
        self.timeout_seconds = timeout_seconds

    def create_agent_run(
        self,
        *,
        tenant_id: str,
        agent_id: str,
        agent_version_id: str | None,
        session_id: str | None,
        input_text: str | None,
        input_json: dict[str, JSONValue] | None,
    ) -> AgentRunContract:
        payload: dict[str, JSONValue] = {
            "tenant_id": tenant_id,
            "agent_id": agent_id,
        }
        if agent_version_id is not None:
            payload["agent_version_id"] = agent_version_id
        if session_id is not None:
            payload["session_id"] = session_id
        if input_text is not None:
            payload["input_text"] = input_text
        if input_json is not None:
            payload["input_json"] = input_json

        try:
            with httpx.Client(timeout=self.timeout_seconds) as client:
                response = client.post(f"{self.base_url}/agents/runs", json=payload)
                response.raise_for_status()
                return AgentRunContract.model_validate(response.json())
        except httpx.HTTPError as exc:
            raise AgentServiceClientError(f"agent-service create run failed: {exc}") from exc

    def execute_agent_run(
        self,
        *,
        tenant_id: str,
        agent_run_id: str,
        worker_key: str | None,
        dry_run: bool,
    ) -> AgentRunContract:
        payload: dict[str, JSONValue] = {
            "tenant_id": tenant_id,
            "dry_run": dry_run,
        }
        if worker_key is not None:
            payload["worker_key"] = worker_key

        try:
            with httpx.Client(timeout=self.timeout_seconds) as client:
                response = client.post(
                    f"{self.base_url}/agents/runs/{agent_run_id}/execute",
                    json=payload,
                )
                response.raise_for_status()
                response_payload = response.json()
                run_payload = response_payload.get("run")
                if not isinstance(run_payload, dict):
                    raise AgentServiceClientError("agent-service execute response missing run")
                return AgentRunContract.model_validate(run_payload)
        except httpx.HTTPError as exc:
            raise AgentServiceClientError(f"agent-service execute run failed: {exc}") from exc