company-ai-platform/tests/test_ai_adapters.py

from typing import Any

import pytest
from fastapi import HTTPException

from app.core.config import Settings
from app.modules.ai_agent import adapters
from app.modules.ai_agent.constants import (
    OPENCLAW_HERMES_PIPELINE,
    AIChatRole,
    AIContextKey,
    AIHttpHeader,
    AIHttpPath,
    AIHttpPayloadKey,
    AIMemoryMode,
    AIResponseKey,
)


class DummyResponse:
    def __init__(self, data: dict[str, Any], status_code: int = 200):
        self._data = data
        self.status_code = status_code
        self.text = str(data)

    def json(self) -> dict[str, Any]:
        return self._data


def chat_response(content: str) -> DummyResponse:
    return DummyResponse(
        {
            AIHttpPayloadKey.CHOICES: [
                {
                    AIHttpPayloadKey.MESSAGE: {
                        AIHttpPayloadKey.CONTENT: content,
                    }
                }
            ]
        }
    )


class DummyClient:
    calls: list[dict[str, Any]] = []

    def __init__(self, timeout: int, trust_env: bool = True):
        self.timeout = timeout
        self.trust_env = trust_env

    def __enter__(self) -> "DummyClient":
        return self

    def __exit__(self, exc_type, exc, traceback) -> None:
        return None

    def post(
        self,
        url: str,
        json: dict[str, Any],
        headers: dict[str, str],
    ) -> DummyResponse:
        self.calls.append({"method": "POST", "url": url, "json": json, "headers": headers})
        if url.endswith(AIHttpPath.CHAT_COMPLETIONS):
            user_content = json[AIHttpPayloadKey.MESSAGES][-1][AIHttpPayloadKey.CONTENT]
            if AIMemoryMode.RECALL in user_content:
                return chat_response("remembered project preference")
            if AIMemoryMode.WRITE in user_content:
                return chat_response("stored")
            if AIContextKey.AGENT_PIPELINE in user_content:
                return chat_response("hermes final answer")
            return chat_response("hermes direct answer")
        if url.endswith(AIHttpPath.TOOLS_INVOKE):
            return DummyResponse(
                {
                    AIResponseKey.OK: True,
                    AIHttpPayloadKey.TOOL: json[AIHttpPayloadKey.TOOL],
                    AIHttpPayloadKey.ACTION: json[AIHttpPayloadKey.ACTION],
                    AIHttpPayloadKey.ARGS: json[AIHttpPayloadKey.ARGS],
                    AIHttpPayloadKey.SESSION_KEY: json[AIHttpPayloadKey.SESSION_KEY],
                }
            )
        raise AssertionError(f"unexpected URL: {url}")

    def get(self, url: str, headers: dict[str, str]) -> DummyResponse:
        self.calls.append({"method": "GET", "url": url, "headers": headers})
        if url.endswith(AIHttpPath.HEALTHZ):
            return DummyResponse({"status": "ok"})
        if url.endswith(AIHttpPath.READYZ):
            return DummyResponse({"status": "ready"})
        if url.endswith(AIHttpPath.HEALTH):
            return DummyResponse({"status": "ok", "platform": "hermes-agent"})
        raise AssertionError(f"unexpected URL: {url}")


def test_hermes_adapter_uses_openai_chat_completions(monkeypatch) -> None:
    DummyClient.calls = []
    monkeypatch.setattr(adapters.httpx, "Client", DummyClient)
    settings = Settings(
        hermes_base_url="http://hermes.local/v1",
        hermes_api_key="hermes-key",
        hermes_model="hermes-agent",
        hermes_session_id="company-ai-platform-test",
    )

    result = adapters.HermesAdapter(settings).ask(
        "summarize project risk",
        {"project_code": "P-001"},
    )

    assert result[AIResponseKey.ANSWER] == "hermes direct answer"
    assert DummyClient.calls[0]["url"] == "http://hermes.local/v1/chat/completions"
    assert DummyClient.calls[0]["headers"][AIHttpHeader.AUTHORIZATION] == "Bearer hermes-key"
    assert (
        DummyClient.calls[0]["headers"][AIHttpHeader.HERMES_SESSION_ID]
        == "company-ai-platform-test"
    )
    assert DummyClient.calls[0]["json"][AIHttpPayloadKey.MODEL] == "hermes-agent"
    assert (
        DummyClient.calls[0]["json"][AIHttpPayloadKey.MESSAGES][0][AIHttpPayloadKey.ROLE]
        == AIChatRole.SYSTEM
    )
    assert "P-001" in DummyClient.calls[0]["json"][AIHttpPayloadKey.MESSAGES][1][
        AIHttpPayloadKey.CONTENT
    ]


def test_openclaw_adapter_uses_gateway_health_and_tool_invoke(monkeypatch) -> None:
    DummyClient.calls = []
    monkeypatch.setattr(adapters.httpx, "Client", DummyClient)
    settings = Settings(
        openclaw_http_url="http://openclaw.local",
        openclaw_gateway_token="gateway-token",
        openclaw_allowed_tools=["sessions_list"],
    )

    health = adapters.OpenClawAdapter(settings).health()
    result = adapters.OpenClawAdapter(settings).ask(
        "list sessions",
        {
            AIContextKey.OPENCLAW_TOOL: "sessions_list",
            AIContextKey.OPENCLAW_ARGS: {},
            AIContextKey.OPENCLAW_SESSION_KEY: "main",
        },
    )

    assert health[AIResponseKey.OK] is True
    assert result[AIResponseKey.RAW][AIResponseKey.DATA][AIHttpPayloadKey.TOOL] == "sessions_list"
    assert [call["url"] for call in DummyClient.calls] == [
        "http://openclaw.local/healthz",
        "http://openclaw.local/readyz",
        "http://openclaw.local/tools/invoke",
    ]
    assert DummyClient.calls[0]["headers"][AIHttpHeader.AUTHORIZATION] == "Bearer gateway-token"
    assert DummyClient.calls[2]["json"][AIHttpPayloadKey.SESSION_KEY] == "main"


def test_openclaw_hermes_adapter_runs_recall_answer_and_remember(monkeypatch) -> None:
    DummyClient.calls = []
    monkeypatch.setattr(adapters.httpx, "Client", DummyClient)
    settings = Settings(
        model_provider="openclaw_hermes",
        openclaw_http_url="http://openclaw.local",
        openclaw_gateway_token="openclaw-key",
        openclaw_allowed_tools=["sessions_list"],
        hermes_base_url="http://hermes.local/v1",
        hermes_api_key="hermes-key",
    )

    result = adapters.OpenClawHermesAdapter(settings).ask(
        "summarize project risk",
        {
            "project_code": "P-001",
            AIContextKey.OPENCLAW_TOOL: "sessions_list",
            AIContextKey.OPENCLAW_ARGS: {},
        },
    )

    assert result[AIResponseKey.ANSWER] == "hermes final answer"
    assert result[AIResponseKey.RAW][AIResponseKey.PIPELINE] == OPENCLAW_HERMES_PIPELINE
    assert [call["url"] for call in DummyClient.calls] == [
        "http://hermes.local/v1/chat/completions",
        "http://openclaw.local/healthz",
        "http://openclaw.local/readyz",
        "http://openclaw.local/tools/invoke",
        "http://hermes.local/v1/chat/completions",
        "http://hermes.local/v1/chat/completions",
    ]
    assert DummyClient.calls[0]["headers"][AIHttpHeader.AUTHORIZATION] == "Bearer hermes-key"
    assert DummyClient.calls[1]["headers"][AIHttpHeader.AUTHORIZATION] == "Bearer openclaw-key"
    assert DummyClient.calls[3]["json"][AIHttpPayloadKey.TOOL] == "sessions_list"
    assert AIContextKey.OPENCLAW in DummyClient.calls[4]["json"][AIHttpPayloadKey.MESSAGES][1][
        AIHttpPayloadKey.CONTENT
    ]


def test_openclaw_adapter_blocks_tools_not_in_allowlist(monkeypatch) -> None:
    DummyClient.calls = []
    monkeypatch.setattr(adapters.httpx, "Client", DummyClient)
    settings = Settings(
        openclaw_http_url="http://openclaw.local",
        openclaw_gateway_token="gateway-token",
        openclaw_allowed_tools=["sessions_list"],
    )

    with pytest.raises(HTTPException) as exc_info:
        adapters.OpenClawAdapter(settings).invoke_tool("filesystem_write")

    assert exc_info.value.status_code == 403
    assert DummyClient.calls == []