lambdagent/tests/test_ollama.py

"""
Tests for Ollama (local / self-hosted) provider support.

Two groups:
  1. Cost accounting — pure unit tests, no Ollama needed. Verify local models
     are billed at $0 and the cost source label reflects the real provider.
  2. Live provider — skipped automatically unless an Ollama server is reachable
     at http://127.0.0.1:11434. Run a local model first, e.g.:
         ollama pull qwen2.5:7b && ollama serve
"""

import urllib.request

import pytest

from agentpaas.engine.sandbox import _compute_cost_usd


# ============================================================
# 1. Cost accounting (no Ollama required)
# ============================================================

class TestOllamaCost:

    def test_ollama_costs_zero(self):
        """Local Ollama models incur no per-token API cost."""
        usage = {"input_tokens": 10_000, "output_tokens": 5_000, "provider": "ollama"}
        assert _compute_cost_usd(usage) == 0.0

    def test_claude_still_billed(self):
        """Regression: non-local providers keep their pricing."""
        usage = {"input_tokens": 1_000_000, "output_tokens": 0, "provider": "claude"}
        assert _compute_cost_usd(usage) == pytest.approx(3.00)

    def test_qwen_still_billed(self):
        usage = {"input_tokens": 1_000_000, "output_tokens": 0, "provider": "qwen"}
        assert _compute_cost_usd(usage) == pytest.approx(5.56)


# ============================================================
# 2. Live provider (skipped unless Ollama is up)
# ============================================================

_OLLAMA_URL = "http://127.0.0.1:11434/api/tags"


def _ollama_models():
    """Return the list of locally available Ollama model names, or [] if down."""
    try:
        with urllib.request.urlopen(_OLLAMA_URL, timeout=3) as resp:
            import json
            data = json.loads(resp.read())
        return [m["name"] for m in data.get("models", [])]
    except Exception:
        return []


_MODELS = _ollama_models()
_HAS_OLLAMA = bool(_MODELS)
# Prefer a small model if present, else use whatever is available.
_MODEL = next((m for m in _MODELS if "7b" in m or "3b" in m), _MODELS[0] if _MODELS else "qwen2.5:7b")

ollama_required = pytest.mark.skipif(
    not _HAS_OLLAMA, reason="No Ollama server reachable at 127.0.0.1:11434"
)


@ollama_required
class TestOllamaLive:

    def test_chat_roundtrip(self):
        from lambdagent.providers import create_provider
        p = create_provider("ollama", model=_MODEL, timeout=120)
        assert p.provider_name == "ollama"
        assert p.base_url.startswith("http://127.0.0.1:11434")
        out = p.chat([{"role": "user", "content": "Reply with exactly: OLLAMA_OK"}])
        assert "OLLAMA_OK" in out

    def test_usage_accumulates(self):
        from lambdagent.providers import create_provider, ChatMessage
        p = create_provider("ollama", model=_MODEL, timeout=120)
        r = p.chat_typed([ChatMessage(role="user", content="Say hi.")], max_tokens=16)
        assert r.text
        assert r.input_tokens > 0
        usage = p.get_usage()
        assert usage["provider"] == "ollama"
        assert usage["input_tokens"] > 0

    def test_full_stack_run_costs_zero(self):
        """A simple agent run on Ollama produces artifacts and bills $0."""
        import asyncio, tempfile, os, json, shutil
        from agentpaas.engine.sandbox import Sandbox

        config = {
            "name": "ollama-smoke", "type": "simple",
            "systemPrompt": "You are concise.",
            "model": {"provider": "ollama", "name": _MODEL,
                      "temperature": 0.0, "maxTokens": 64},
        }
        agent_dir = tempfile.mkdtemp(prefix="ollama_test_")
        try:
            sb = Sandbox(level=0)
            r = asyncio.run(sb.execute(
                config, "Reply with exactly: FULLSTACK_OK",
                timeout=120, agent_dir=agent_dir, run_id="run_olltest",
            ))
            assert r.status == "completed", r.error
            ws = r.workspace_path
            for fn in ("output.json", "cost.json", "manifest.json"):
                assert os.path.isfile(os.path.join(ws, fn))
            cost = json.load(open(os.path.join(ws, "cost.json")))
            assert cost["cost_usd"] == 0.0
        finally:
            shutil.rmtree(agent_dir, ignore_errors=True)