191 lines
6.9 KiB
Python
191 lines
6.9 KiB
Python
"""Shared test seams for the Fase 1 de-risk spikes (throwaway, dev-only).
|
|
|
|
Three reusable pieces every spike leans on:
|
|
|
|
1. **Cost/stop invariant (B4 / D6).** ``Budget`` refuses to start without positive
|
|
token and round caps (``ValueError`` on a bad construction argument), and
|
|
``TokenMeter`` raises ``BudgetExceeded`` the moment a cap is crossed at runtime.
|
|
Two exception types is intentional: ``ValueError`` = bad ctor argument (you never
|
|
even started), ``BudgetExceeded`` = a cap was breached while running.
|
|
|
|
2. **Deterministic fake model.** ``FakeChatClient`` subclasses the GA
|
|
``BaseChatClient`` with scripted, deterministic replies and counts "tokens" by
|
|
word-count — no network, no endpoint. ``fake_agent`` wraps it in a real
|
|
``agent_framework.Agent`` so the orchestration builders get genuine participants.
|
|
The Step 2 builder smoke (in ``tests/spikes/test_harness.py``) proves this client
|
|
can actually drive the GA ``GroupChatBuilder`` / ``ConcurrentBuilder``.
|
|
|
|
3. **Gated live arm.** ``live_local_client_or_skip`` builds an
|
|
``agent_framework.openai.OpenAIChatClient`` against an OpenAI-compatible LOCAL
|
|
endpoint **directly** (the D2 ``LocalBackend`` seam is deliberately left un-wired
|
|
until Fase 2, so ``src/`` stays untouched), or ``pytest.skip``s when the
|
|
``PORTFOLIO_LOCAL_*`` env is unset. No silent egress (D6).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from collections.abc import Sequence
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from agent_framework import (
|
|
Agent,
|
|
BaseChatClient,
|
|
ChatResponse,
|
|
ChatResponseUpdate,
|
|
Message,
|
|
)
|
|
|
|
|
|
class BudgetExceeded(RuntimeError):
|
|
"""Raised the moment a runtime cap (tokens or rounds) is crossed (B4).
|
|
|
|
Carries the breached ``kind`` ("tokens" | "rounds"), the ``limit`` that was
|
|
set, and the ``observed`` value that crossed it — a structured stop event,
|
|
never a silent hang.
|
|
"""
|
|
|
|
def __init__(self, kind: str, limit: int, observed: int) -> None:
|
|
self.kind = kind
|
|
self.limit = limit
|
|
self.observed = observed
|
|
super().__init__(f"budget exceeded: {kind} limit={limit} observed={observed}")
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Budget:
|
|
"""Hard token + round/iteration caps, required at startup (A4 / D6).
|
|
|
|
Refuses to construct without positive caps — fail-fast, never an unbounded loop.
|
|
"""
|
|
|
|
max_tokens: int
|
|
max_rounds: int
|
|
|
|
def __post_init__(self) -> None:
|
|
if self.max_tokens <= 0:
|
|
raise ValueError(f"max_tokens must be positive, got {self.max_tokens}")
|
|
if self.max_rounds <= 0:
|
|
raise ValueError(f"max_rounds must be positive, got {self.max_rounds}")
|
|
|
|
|
|
class TokenMeter:
|
|
"""Accumulates token and round usage against a ``Budget``; raises the moment
|
|
a cap is crossed."""
|
|
|
|
def __init__(self, budget: Budget) -> None:
|
|
self.budget = budget
|
|
self.tokens = 0
|
|
self.rounds = 0
|
|
|
|
def charge(self, tokens: int) -> int:
|
|
"""Add ``tokens`` to the running total; raise ``BudgetExceeded`` if over cap."""
|
|
self.tokens += tokens
|
|
if self.tokens > self.budget.max_tokens:
|
|
raise BudgetExceeded("tokens", self.budget.max_tokens, self.tokens)
|
|
return self.tokens
|
|
|
|
def tick_round(self) -> int:
|
|
"""Increment the round counter; raise ``BudgetExceeded`` if over cap."""
|
|
self.rounds += 1
|
|
if self.rounds > self.budget.max_rounds:
|
|
raise BudgetExceeded("rounds", self.budget.max_rounds, self.rounds)
|
|
return self.rounds
|
|
|
|
|
|
def _word_tokens(text: str) -> int:
|
|
"""Token proxy: word count. Deterministic, endpoint-free."""
|
|
return len(text.split())
|
|
|
|
|
|
class FakeChatClient(BaseChatClient):
|
|
"""A deterministic, network-free ``BaseChatClient`` for driving MAF agents in tests.
|
|
|
|
Returns scripted replies in order; once the script is exhausted it falls back to
|
|
``default_reply``. Counts "tokens" by word-count of each reply it emits, exposing
|
|
``total_tokens`` and ``call_count`` for the spike measurements.
|
|
"""
|
|
|
|
OTEL_PROVIDER_NAME = "fake"
|
|
|
|
def __init__(self, scripted: Sequence[str] | None = None, *, default_reply: str = "ok") -> None:
|
|
super().__init__()
|
|
self._scripted: list[str] = list(scripted or [])
|
|
self._idx = 0
|
|
self._default = default_reply
|
|
self.total_tokens = 0
|
|
self.call_count = 0
|
|
|
|
def _next_reply(self) -> str:
|
|
reply = self._scripted[self._idx] if self._idx < len(self._scripted) else self._default
|
|
self._idx += 1
|
|
self.call_count += 1
|
|
self.total_tokens += _word_tokens(reply)
|
|
return reply
|
|
|
|
def _inner_get_response(
|
|
self,
|
|
*,
|
|
messages: Sequence[Message],
|
|
stream: bool,
|
|
options: Any,
|
|
**kwargs: Any,
|
|
) -> Any:
|
|
# Matches the GA BaseChatClient contract: return a ResponseStream when
|
|
# streaming, otherwise an awaitable resolving to a ChatResponse.
|
|
reply = self._next_reply()
|
|
|
|
if stream:
|
|
|
|
async def _agen() -> Any:
|
|
yield ChatResponseUpdate(
|
|
role="assistant", contents=[{"type": "text", "text": reply}]
|
|
)
|
|
|
|
return self._build_response_stream(_agen())
|
|
|
|
async def _coro() -> ChatResponse:
|
|
return ChatResponse(
|
|
messages=[Message(role="assistant", contents=[reply])],
|
|
response_id="fake",
|
|
)
|
|
|
|
return _coro()
|
|
|
|
|
|
def fake_agent(
|
|
client: BaseChatClient,
|
|
name: str,
|
|
instructions: str = "You are a terse participant. Answer in one short line.",
|
|
) -> Agent:
|
|
"""Build a minimal real ``agent_framework.Agent`` backed by ``client`` so the
|
|
orchestration builders get a genuine participant."""
|
|
return Agent(client, instructions, name=name)
|
|
|
|
|
|
def live_local_client_or_skip() -> Any:
|
|
"""Build an ``OpenAIChatClient`` against the OpenAI-compatible LOCAL endpoint
|
|
(``PORTFOLIO_LOCAL_BASE_URL`` + ``PORTFOLIO_LOCAL_MODEL``), or ``pytest.skip``
|
|
when unset.
|
|
|
|
The D2 ``LocalBackend`` seam is intentionally NOT used here — its live wiring is a
|
|
Fase 2 concern; the throwaway spike builds the client directly so ``src/`` stays
|
|
untouched. No silent egress: without the env vars the live arm simply skips (D6).
|
|
"""
|
|
base_url = os.environ.get("PORTFOLIO_LOCAL_BASE_URL")
|
|
model = os.environ.get("PORTFOLIO_LOCAL_MODEL")
|
|
if not base_url or not model:
|
|
import pytest
|
|
|
|
pytest.skip(
|
|
"LOCAL endpoint not configured "
|
|
"(set PORTFOLIO_LOCAL_BASE_URL and PORTFOLIO_LOCAL_MODEL to run the live arm)"
|
|
)
|
|
|
|
from agent_framework.openai import OpenAIChatClient
|
|
|
|
# Most local OpenAI-compatible servers (Ollama / LM Studio) accept any non-empty
|
|
# key; allow an override but default to a dummy so construction never blocks.
|
|
api_key = os.environ.get("PORTFOLIO_LOCAL_API_KEY", "local")
|
|
return OpenAIChatClient(model=model, api_key=api_key, base_url=base_url)
|