fengqun/swarm_minimal/newapi_agnet.py

"""NewAPI-backed Agnet adapter for the minimal swarm.

The adapter uses the OpenAI-compatible chat completions shape exposed by NewAPI.
Credentials are read from environment variables by examples, not stored in code.
"""

from __future__ import annotations

from dataclasses import dataclass
import json
import os
from typing import Protocol
from urllib.error import HTTPError
from urllib.request import Request, urlopen

from .config import MissingConfigError
from .core import Agent, Task


class HttpClient(Protocol):
    def get_json(
        self,
        url: str,
        headers: dict[str, str],
        timeout: float,
    ) -> dict[str, object]:
        pass

    def post_json(
        self,
        url: str,
        headers: dict[str, str],
        payload: dict[str, object],
        timeout: float,
    ) -> dict[str, object]:
        pass


@dataclass(frozen=True)
class NewApiChannelConfig:
    base_url: str
    api_key: str
    model: str | None = None
    timeout_seconds: float = 30.0

    @classmethod
    def from_env(cls) -> "NewApiChannelConfig":
        base_url = os.environ.get("NEWAPI_BASE_URL") or os.environ.get("NEWAPI_URL")
        api_key = os.environ.get("NEWAPI_API_KEY") or os.environ.get("NEWAPI_KEY")
        if not base_url or not api_key:
            raise MissingConfigError("missing NEWAPI_BASE_URL/NEWAPI_URL or NEWAPI_API_KEY/NEWAPI_KEY")

        return cls(
            base_url=base_url,
            api_key=api_key,
            model=os.environ.get("NEWAPI_MODEL"),
            timeout_seconds=float(os.environ.get("NEWAPI_TIMEOUT_SECONDS", "30")),
        )

    def redacted_summary(self) -> dict[str, object]:
        return {
            "base_url": self.base_url.rstrip("/"),
            "fallback_model": self.model or "<not-set>",
            "api_key": "<redacted>",
            "timeout_seconds": self.timeout_seconds,
        }


class UrlLibHttpClient:
    def get_json(
        self,
        url: str,
        headers: dict[str, str],
        timeout: float,
    ) -> dict[str, object]:
        request = Request(url=url, headers=headers, method="GET")
        try:
            with urlopen(request, timeout=timeout) as response:
                body = response.read().decode("utf-8")
        except HTTPError as exc:
            error_body = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"NewAPI request failed with HTTP {exc.code}: {error_body}") from exc
        return json.loads(body)

    def post_json(
        self,
        url: str,
        headers: dict[str, str],
        payload: dict[str, object],
        timeout: float,
    ) -> dict[str, object]:
        request = Request(
            url=url,
            data=json.dumps(payload).encode("utf-8"),
            headers=headers,
            method="POST",
        )
        try:
            with urlopen(request, timeout=timeout) as response:
                body = response.read().decode("utf-8")
        except HTTPError as exc:
            error_body = exc.read().decode("utf-8", errors="replace")
            raise RuntimeError(f"NewAPI request failed with HTTP {exc.code}: {error_body}") from exc
        return json.loads(body)


class NewApiAgnet:
    """Small Agnet wrapper that can join the minimal swarm as an Agent."""

    def __init__(
        self,
        config: NewApiChannelConfig,
        *,
        agent_id: str = "newapi-agnet",
        capability: str = "verify",
        http_client: HttpClient | None = None,
    ) -> None:
        self.config = config
        self.agent_id = agent_id
        self.capability = capability
        self.http_client = http_client or UrlLibHttpClient()

    def as_agent(self) -> Agent:
        return Agent(id=self.agent_id, capability=self.capability, run=self.run_task)

    def run_task(self, task: Task, shared_state: dict[str, str]) -> tuple[str, float]:
        system_prompt = (
            "You are a minimal Agnet worker inside a swarm. "
            "Return a concise result that can be scored and converged."
        )
        user_prompt = (
            f"Task kind: {task.kind}\n"
            f"Task input: {task.input}\n"
            f"Known shared state keys: {', '.join(sorted(shared_state.keys()))}"
        )
        content = self.chat(system_prompt=system_prompt, user_prompt=user_prompt)
        score = 0.9 if content.strip() else 0.0
        return content, score

    def chat(self, *, system_prompt: str, user_prompt: str) -> str:
        if not self.config.model:
            raise RuntimeError("NewAPI model is required for chat calls")
        url = self.config.base_url.rstrip("/") + "/v1/chat/completions"
        payload = {
            "model": self.config.model,
            "messages": [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": user_prompt},
            ],
            "temperature": 0.2,
        }
        response = self.http_client.post_json(
            url=url,
            headers={
                "Authorization": f"Bearer {self.config.api_key}",
                "Content-Type": "application/json",
            },
            payload=payload,
            timeout=self.config.timeout_seconds,
        )
        return _extract_chat_content(response)


def discover_newapi_models(
    config: NewApiChannelConfig,
    *,
    http_client: HttpClient | None = None,
    candidate_paths: tuple[str, ...] = ("/v1/models", "/models", "/model"),
) -> list[str]:
    """Discover available NewAPI model ids.

    NewAPI deployments are often OpenAI-compatible at ``/v1/models``. Some
    installations expose shorter model-list routes, so the helper tries a small
    fallback list and returns the first valid model list.
    """

    client = http_client or UrlLibHttpClient()
    headers = {"Authorization": f"Bearer {config.api_key}"}
    errors: list[str] = []
    for path in candidate_paths:
        url = config.base_url.rstrip("/") + path
        try:
            response = client.get_json(url=url, headers=headers, timeout=config.timeout_seconds)
            models = _extract_model_ids(response)
        except Exception as exc:
            errors.append(f"{path}: {exc}")
            continue
        if models:
            return models
    raise RuntimeError("could not discover NewAPI models: " + "; ".join(errors))


def select_distinct_models(models: list[str], count: int = 3) -> list[str]:
    unique: list[str] = []
    seen: set[str] = set()
    for model in models:
        if model in seen:
            continue
        seen.add(model)
        unique.append(model)
        if len(unique) == count:
            return unique
    raise RuntimeError(f"need at least {count} distinct models, got {len(unique)}")


def build_model_test_agnets(
    config: NewApiChannelConfig,
    *,
    models: list[str],
    http_client: HttpClient | None = None,
) -> list[Agent]:
    """Build one Agnet per model with distinct capabilities."""

    agents: list[Agent] = []
    for index, model in enumerate(models):
        model_config = NewApiChannelConfig(
            base_url=config.base_url,
            api_key=config.api_key,
            model=model,
            timeout_seconds=config.timeout_seconds,
        )
        agents.append(
            NewApiAgnet(
                model_config,
                agent_id=f"newapi-model-{index + 1}",
                capability=f"model_test_{index + 1}",
                http_client=http_client,
            ).as_agent()
        )
    return agents


def _extract_chat_content(response: dict[str, object]) -> str:
    choices = response.get("choices")
    if not isinstance(choices, list) or not choices:
        raise RuntimeError("NewAPI response does not include choices")
    first = choices[0]
    if not isinstance(first, dict):
        raise RuntimeError("NewAPI response choice is invalid")
    message = first.get("message")
    if not isinstance(message, dict):
        raise RuntimeError("NewAPI response choice does not include message")
    content = message.get("content")
    if not isinstance(content, str):
        raise RuntimeError("NewAPI response message does not include string content")
    return content


def _extract_model_ids(response: object) -> list[str]:
    if isinstance(response, list):
        return [item for item in response if isinstance(item, str)]

    if not isinstance(response, dict):
        return []

    data = response.get("data")
    if isinstance(data, list):
        models: list[str] = []
        for item in data:
            if isinstance(item, str):
                models.append(item)
            elif isinstance(item, dict) and isinstance(item.get("id"), str):
                models.append(item["id"])
        return models

    models_value = response.get("models") or response.get("model")
    if isinstance(models_value, list):
        return [
            item if isinstance(item, str) else item.get("id")
            for item in models_value
            if isinstance(item, str) or (isinstance(item, dict) and isinstance(item.get("id"), str))
        ]

    return []