agentcap-cli 0.1.0__py3-none-any.whl

agentcap/devtools/fake_upstream.py

@@ -0,0 +1,412 @@
+"""A tiny local fake upstream for offline AgentCap experiments.
+
+Run it with::
+
+    python -m agentcap.devtools.fake_upstream --port 9000
+
+and point AgentCap at it::
+
+    AGENTCAP_GEMINI_UPSTREAM=http://127.0.0.1:9000 \
+        agentcap run --cap 0.05 -- python my_agent.py
+
+It imitates just enough of the OpenAI and Gemini wire formats to exercise
+metering, streaming and output clamping:
+
+* ``POST /v1/chat/completions`` (JSON, or SSE with a final usage chunk when
+  ``stream_options.include_usage`` is set)
+* ``POST /v1/responses`` (JSON with ``usage``)
+* ``POST /v1beta/models/<model>:generateContent`` (JSON with ``usageMetadata``)
+* ``POST /v1beta/models/<model>:streamGenerateContent`` (SSE when ``alt=sse``,
+  otherwise a JSON array; final ``usageMetadata``)
+
+Prompt usage is realistic rather than fixed: ``max(1, ceil(text characters /
+4))`` where the text characters are the concatenated ``content``, ``input``,
+``prompt``, ``text``, ``system`` and tool-description strings in the request.
+Output usage is ``min(requested output limit, 2000)`` tokens, so a clamped
+request visibly produces a different bill.  Every response carries
+``x-fake-output-limit`` with the limit that was actually received (``none``
+when the request had no output limit).
+
+Request bodies and headers are never logged or stored.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import math
+from typing import Any, AsyncIterator, Mapping
+
+import uvicorn
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import JSONResponse, Response, StreamingResponse
+from starlette.routing import Route
+
+__all__ = ["create_app", "main"]
+
+PROMPT_TOKENS_PER_CHAR = 4
+MAX_OUTPUT_TOKENS = 2000
+LIMIT_HEADER = "x-fake-output-limit"
+SSE_MEDIA_TYPE = "text/event-stream"
+
+# Keys whose string values are prompt text.  Structural fields (``model``,
+# ``role``, ``type`` ...) are deliberately excluded: they are not billed as
+# prompt text by the real APIs.
+_TEXT_KEYS = frozenset(
+    {
+        "text",
+        "content",
+        "input",
+        "prompt",
+        "system",
+        "system_instruction",
+        "systemInstruction",
+        "description",
+    }
+)
+
+
+def _text_chars(value: Any, *, is_text: bool = False) -> int:
+    """Count characters of prompt text anywhere in a request body."""
+
+    if isinstance(value, str):
+        return len(value) if is_text else 0
+    if isinstance(value, Mapping):
+        total = 0
+        for key, child in value.items():
+            child_is_text = isinstance(key, str) and key in _TEXT_KEYS
+            total += _text_chars(child, is_text=child_is_text)
+        return total
+    if isinstance(value, (list, tuple)):
+        # A list under a text key may be a list of strings (embeddings input)
+        # or a list of message/part objects (whose own text keys are counted).
+        return sum(
+            _text_chars(child, is_text=is_text and isinstance(child, str))
+            for child in value
+        )
+    return 0
+
+
+def _prompt_tokens(payload: Mapping[str, Any]) -> int:
+    return max(1, math.ceil(_text_chars(payload) / PROMPT_TOKENS_PER_CHAR))
+
+
+def _positive_int(value: Any) -> int | None:
+    if isinstance(value, bool) or value is None:
+        return None
+    if isinstance(value, int):
+        return value if value > 0 else None
+    if isinstance(value, float) and value.is_integer() and value > 0:
+        return int(value)
+    if isinstance(value, str) and value.strip().isdigit():
+        number = int(value.strip())
+        return number if number > 0 else None
+    return None
+
+
+def _mapping(value: Any) -> Mapping[str, Any]:
+    return value if isinstance(value, Mapping) else {}
+
+
+def _openai_chat_limit(payload: Mapping[str, Any]) -> int | None:
+    return _positive_int(
+        payload.get("max_completion_tokens", payload.get("max_tokens"))
+    )
+
+
+def _openai_responses_limit(payload: Mapping[str, Any]) -> int | None:
+    return _positive_int(payload.get("max_output_tokens"))
+
+
+def _gemini_limit(payload: Mapping[str, Any]) -> int | None:
+    config = _mapping(payload.get("generationConfig"))
+    if not config:
+        config = _mapping(payload.get("generation_config"))
+    return _positive_int(
+        config.get("maxOutputTokens", config.get("max_output_tokens"))
+    )
+
+
+def _completion_tokens(limit: int | None) -> int:
+    return min(limit, MAX_OUTPUT_TOKENS) if limit else MAX_OUTPUT_TOKENS
+
+
+def _limit_header(limit: int | None) -> dict[str, str]:
+    return {LIMIT_HEADER: str(limit) if limit is not None else "none"}
+
+
+def _usage(limit: int | None, prompt_tokens: int) -> dict[str, int]:
+    output = _completion_tokens(limit)
+    return {
+        "prompt_tokens": prompt_tokens,
+        "completion_tokens": output,
+        "total_tokens": prompt_tokens + output,
+    }
+
+
+def _gemini_usage(limit: int | None, prompt_tokens: int) -> dict[str, int]:
+    output = _completion_tokens(limit)
+    return {
+        "promptTokenCount": prompt_tokens,
+        "candidatesTokenCount": output,
+        "totalTokenCount": prompt_tokens + output,
+    }
+
+
+async def _json_body(request: Request) -> Mapping[str, Any] | None:
+    try:
+        payload = json.loads(await request.body())
+    except (json.JSONDecodeError, UnicodeDecodeError, ValueError):
+        return None
+    return payload if isinstance(payload, Mapping) else None
+
+
+def _invalid_body() -> JSONResponse:
+    return JSONResponse(
+        {"error": {"message": "fake upstream: body is not a JSON object"}},
+        status_code=400,
+    )
+
+
+def _not_found(what: str) -> JSONResponse:
+    return JSONResponse(
+        {"error": {"message": f"fake upstream: unsupported {what}"}},
+        status_code=404,
+    )
+
+
+# ------------------------------------------------------------------- OpenAI
+async def openai_chat(request: Request) -> Response:
+    payload = await _json_body(request)
+    if payload is None:
+        return _invalid_body()
+    limit = _openai_chat_limit(payload)
+    model = payload.get("model") if isinstance(payload.get("model"), str) else "fake-model"
+    headers = _limit_header(limit)
+    usage = _usage(limit, _prompt_tokens(payload))
+
+    if payload.get("stream") is not True:
+        body = {
+            "id": "chatcmpl-fake",
+            "object": "chat.completion",
+            "created": 0,
+            "model": model,
+            "choices": [
+                {
+                    "index": 0,
+                    "message": {"role": "assistant", "content": "fake response"},
+                    "finish_reason": "stop",
+                }
+            ],
+            "usage": usage,
+        }
+        return JSONResponse(body, headers=headers)
+
+    options = _mapping(payload.get("stream_options"))
+    include_usage = options.get("include_usage") is True
+
+    async def events() -> AsyncIterator[bytes]:
+        def send(data: Mapping[str, Any]) -> bytes:
+            return f"data: {json.dumps(data)}\n\n".encode("utf-8")
+
+        base = {
+            "id": "chatcmpl-fake",
+            "object": "chat.completion.chunk",
+            "created": 0,
+            "model": model,
+        }
+        yield send(
+            {
+                **base,
+                "choices": [
+                    {
+                        "index": 0,
+                        "delta": {"role": "assistant", "content": "fake "},
+                        "finish_reason": None,
+                    }
+                ],
+            }
+        )
+        yield send(
+            {
+                **base,
+                "choices": [
+                    {
+                        "index": 0,
+                        "delta": {"content": "response"},
+                        "finish_reason": None,
+                    }
+                ],
+            }
+        )
+        yield send(
+            {
+                **base,
+                "choices": [
+                    {"index": 0, "delta": {}, "finish_reason": "stop"}
+                ],
+            }
+        )
+        if include_usage:
+            yield send({**base, "choices": [], "usage": usage})
+        yield b"data: [DONE]\n\n"
+
+    return StreamingResponse(
+        events(),
+        media_type=SSE_MEDIA_TYPE,
+        headers={**headers, "cache-control": "no-cache"},
+    )
+
+
+async def openai_responses(request: Request) -> Response:
+    payload = await _json_body(request)
+    if payload is None:
+        return _invalid_body()
+    limit = _openai_responses_limit(payload)
+    model = payload.get("model") if isinstance(payload.get("model"), str) else "fake-model"
+    usage = _usage(limit, _prompt_tokens(payload))
+    body = {
+        "id": "resp_fake",
+        "object": "response",
+        "created_at": 0,
+        "model": model,
+        "status": "completed",
+        "output": [
+            {
+                "type": "message",
+                "role": "assistant",
+                "content": [{"type": "output_text", "text": "fake response"}],
+            }
+        ],
+        "usage": {
+            "input_tokens": usage["prompt_tokens"],
+            "output_tokens": usage["completion_tokens"],
+            "total_tokens": usage["total_tokens"],
+        },
+    }
+    return JSONResponse(body, headers=_limit_header(limit))
+
+
+# ------------------------------------------------------------------- Gemini
+def _gemini_chunks(limit: int | None, prompt_tokens: int) -> list[dict[str, Any]]:
+    return [
+        {
+            "candidates": [
+                {
+                    "content": {
+                        "role": "model",
+                        "parts": [{"text": "fake "}],
+                    }
+                }
+            ]
+        },
+        {
+            "candidates": [
+                {
+                    "content": {
+                        "role": "model",
+                        "parts": [{"text": "response"}],
+                    },
+                    "finishReason": "STOP",
+                }
+            ],
+            "usageMetadata": _gemini_usage(limit, prompt_tokens),
+        },
+    ]
+
+
+async def gemini_model(request: Request) -> Response:
+    rest = request.path_params.get("rest", "")
+    action = rest.rsplit(":", 1)[-1] if ":" in rest else ""
+    payload = await _json_body(request)
+    if payload is None:
+        return _invalid_body()
+    limit = _gemini_limit(payload)
+    headers = _limit_header(limit)
+    prompt_tokens = _prompt_tokens(payload)
+
+    if action == "generateContent":
+        chunks = _gemini_chunks(limit, prompt_tokens)
+        final = chunks[-1]
+        return JSONResponse(
+            {
+                "candidates": final["candidates"],
+                "usageMetadata": final["usageMetadata"],
+            },
+            headers=headers,
+        )
+
+    if action == "streamGenerateContent":
+        chunks = _gemini_chunks(limit, prompt_tokens)
+        if request.query_params.get("alt") == "sse":
+
+            async def events() -> AsyncIterator[bytes]:
+                for chunk in chunks:
+                    yield f"data: {json.dumps(chunk)}\n\n".encode("utf-8")
+
+            return StreamingResponse(
+                events(),
+                media_type=SSE_MEDIA_TYPE,
+                headers={**headers, "cache-control": "no-cache"},
+            )
+        return JSONResponse(chunks, headers=headers)
+
+    return _not_found("Gemini action")
+
+
+# ------------------------------------------------------------------- routes
+async def healthz(_request: Request) -> Response:
+    return JSONResponse({"status": "ok"})
+
+
+def create_app() -> Starlette:
+    """Build the fake upstream ASGI app (no state, no logging)."""
+
+    return Starlette(
+        routes=[
+            Route("/healthz", healthz, methods=["GET"]),
+            Route(
+                "/v1/chat/completions",
+                openai_chat,
+                methods=["POST"],
+            ),
+            Route("/v1/responses", openai_responses, methods=["POST"]),
+            Route(
+                "/v1beta/models/{rest:path}",
+                gemini_model,
+                methods=["POST"],
+            ),
+        ]
+    )
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="python -m agentcap.devtools.fake_upstream",
+        description="Local fake OpenAI/Gemini upstream for AgentCap testing.",
+    )
+    parser.add_argument(
+        "--port", type=int, default=0, help="port to listen on (default: random)"
+    )
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    args = _build_parser().parse_args(argv)
+    if args.port < 0 or args.port > 65535:
+        print("fake_upstream: --port must be between 0 and 65535")
+        return 2
+    # 127.0.0.1 only: never expose a fake upstream to the network.
+    uvicorn.run(
+        create_app(),
+        host="127.0.0.1",
+        port=args.port,
+        log_level="warning",
+        access_log=False,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

agentcap/estimate.py

@@ -0,0 +1,271 @@
+"""Conservative token estimation and worst-case cost reservation."""
+
+from __future__ import annotations
+
+import json
+import math
+from datetime import date, datetime
+from decimal import Decimal, ROUND_CEILING
+from typing import Any, Mapping
+
+from .pricing import MICRO_USD, ModelPrice, effective_rates
+
+_TEXT_FIELDS = ("messages", "input", "contents", "system", "tools")
+_MEDIA_TOKENS = 2000
+
+_OPENAI_MEDIA_TYPES = {"image_url", "input_image", "file", "input_file"}
+_GEMINI_MEDIA_KEYS = ("inline_data", "inlineData", "file_data", "fileData")
+
+
+def _as_mapping(body: Mapping[str, Any] | dict[str, Any]) -> Mapping[str, Any]:
+    if not isinstance(body, Mapping):
+        raise TypeError("request body must be a mapping")
+    return body
+
+
+def count_media_parts(body: Mapping[str, Any]) -> int:
+    """Count image/file parts (OpenAI and Gemini) anywhere in the body."""
+
+    body = _as_mapping(body)
+
+    def walk(value: Any) -> int:
+        count = 0
+        if isinstance(value, Mapping):
+            part_type = value.get("type")
+            is_openai_media = isinstance(part_type, str) and part_type in _OPENAI_MEDIA_TYPES
+            is_gemini_media = any(key in value for key in _GEMINI_MEDIA_KEYS)
+            # A part may carry both a type marker and a media payload; count once.
+            if is_openai_media or is_gemini_media:
+                count += 1
+            for child in value.values():
+                count += walk(child)
+        elif isinstance(value, (list, tuple)):
+            for child in value:
+                count += walk(child)
+        return count
+
+    return walk(body)
+
+
+def _relevant_text_chars(body: Mapping[str, Any]) -> int:
+    chars = 0
+    for field in _TEXT_FIELDS:
+        if field in body and body[field] is not None:
+            serialized = json.dumps(
+                body[field], ensure_ascii=False, sort_keys=True, separators=(",", ":")
+            )
+            chars += len(serialized)
+    return chars
+
+
+def estimate_input_tokens(body: Mapping[str, Any] | dict[str, Any]) -> int:
+    """Estimate prompt tokens with a deliberately conservative formula.
+
+    ``ceil(characters / 3) * 1.10 + 50`` plus a flat 2000 tokens for every
+    image or file part.  Overestimating is safe: it only widens the
+    reservation, never the spend.
+    """
+
+    body = _as_mapping(body)
+    base = math.ceil(_relevant_text_chars(body) / 3)
+    overhead = math.ceil(base * 1.10) + 50
+    return overhead + count_media_parts(body) * _MEDIA_TOKENS
+
+
+def apply_input_safety_factor(input_tokens: int, factor: Decimal) -> int:
+    """Widen an input-token estimate for a reservation (rounded up).
+
+    The factor only pads the worst-case reservation; settlement always uses
+    the usage the provider reported.  ``ProxyConfig`` requires a factor of at
+    least 1 so the reservation can only widen.
+    """
+
+    if (
+        isinstance(input_tokens, bool)
+        or not isinstance(input_tokens, int)
+        or input_tokens < 0
+    ):
+        raise ValueError("input_tokens must be a non-negative integer")
+    if isinstance(factor, bool) or not isinstance(factor, Decimal):
+        raise TypeError("factor must be a Decimal")
+    if not factor.is_finite() or factor < 0:
+        raise ValueError("factor must be a finite non-negative decimal")
+    return int(
+        (Decimal(input_tokens) * factor).to_integral_value(
+            rounding=ROUND_CEILING
+        )
+    )
+
+
+def cache_write_requested(body: Mapping[str, Any] | dict[str, Any]) -> bool:
+    """True when the request can write to a prompt cache.
+
+    OpenAI-style cache retention is signalled by a top-level ``prompt_cache*``
+    key (e.g. ``prompt_cache_retention``); cache-write capable content parts
+    carry a ``cache_control`` key anywhere in the body.
+    """
+
+    body = _as_mapping(body)
+    for key in body:
+        if isinstance(key, str) and key.startswith("prompt_cache"):
+            return True
+
+    def walk(value: Any) -> bool:
+        if isinstance(value, Mapping):
+            if "cache_control" in value:
+                return True
+            return any(walk(child) for child in value.values())
+        if isinstance(value, (list, tuple)):
+            return any(walk(child) for child in value)
+        return False
+
+    return walk(body)
+
+
+def _positive_int(value: Any) -> int | None:
+    if isinstance(value, bool) or value is None:
+        return None
+    if isinstance(value, int):
+        return value if value > 0 else None
+    if isinstance(value, float) and value.is_integer() and value > 0:
+        return int(value)
+    if isinstance(value, str):
+        stripped = value.strip()
+        if stripped.isdigit() and int(stripped) > 0:
+            return int(stripped)
+    return None
+
+
+def _normalise_provider(provider: str) -> str:
+    normalised = provider.strip().lower()
+    if normalised in {"openai", "open-ai"}:
+        return "openai"
+    if normalised in {"google", "gemini", "google-gemini", "google_genai"}:
+        return "gemini"
+    raise ValueError(f"Unsupported provider for token estimation: {provider!r}")
+
+
+def _openai_max_tokens(body: Mapping[str, Any]) -> int | None:
+    # Chat Completions uses max_completion_tokens (new) or max_tokens (old).
+    # Responses API uses max_output_tokens.  Prefer the most explicit fields.
+    for field in ("max_completion_tokens", "max_output_tokens", "max_tokens"):
+        value = _positive_int(body.get(field))
+        if value is not None:
+            return value
+    return None
+
+
+def _gemini_max_tokens(body: Mapping[str, Any]) -> int | None:
+    candidates: list[Mapping[str, Any]] = []
+    for config_key in ("generationConfig", "generation_config"):
+        value = body.get(config_key)
+        if isinstance(value, Mapping):
+            candidates.append(value)
+    for config in candidates:
+        for field in ("maxOutputTokens", "max_output_tokens"):
+            value = _positive_int(config.get(field))
+            if value is not None:
+                return value
+    return None
+
+
+def _openai_candidates(body: Mapping[str, Any]) -> int:
+    # ``n`` produces n completions; each may use the full output budget.
+    return _positive_int(body.get("n")) or 1
+
+
+def _gemini_candidates(body: Mapping[str, Any]) -> int:
+    for config_key in ("generationConfig", "generation_config"):
+        value = body.get(config_key)
+        if isinstance(value, Mapping):
+            for field in ("candidateCount", "candidate_count"):
+                count = _positive_int(value.get(field))
+                if count is not None:
+                    return count
+    return 1
+
+
+def explicit_output_tokens(provider: str, body: Mapping[str, Any]) -> int | None:
+    """The output limit the client itself asked for, if any."""
+
+    body = _as_mapping(body)
+    if _normalise_provider(provider) == "openai":
+        return _openai_max_tokens(body)
+    return _gemini_max_tokens(body)
+
+
+def candidate_count(provider: str, body: Mapping[str, Any]) -> int:
+    """How many completions/candidates one request can produce (>= 1)."""
+
+    body = _as_mapping(body)
+    if _normalise_provider(provider) == "openai":
+        return _openai_candidates(body)
+    return _gemini_candidates(body)
+
+
+def max_output_tokens(
+    provider: str, body: Mapping[str, Any], model_price: ModelPrice
+) -> int:
+    """Worst-case output tokens for a request, clamped to the model default."""
+
+    body = _as_mapping(body)
+    normalised = _normalise_provider(provider)
+
+    if normalised == "openai":
+        requested = _openai_max_tokens(body)
+        candidates = _openai_candidates(body)
+    else:
+        requested = _gemini_max_tokens(body)
+        candidates = _gemini_candidates(body)
+
+    per_candidate = requested if requested is not None else model_price.max_output_tokens
+    return per_candidate * candidates
+
+
+def reservation_usd(
+    model_price: ModelPrice,
+    input_tokens: int,
+    output_tokens: int,
+    now: datetime | date | None = None,
+    cache_write: bool = False,
+) -> Decimal:
+    """Worst-case USD reservation, rounded up to 6 decimals.
+
+    Cached prompt tokens are ignored here: the reservation assumes no cache
+    hits.  A long-context prompt is reserved at the higher tier.  When
+    ``cache_write`` is true, input is reserved at
+    ``max(input, cache_write)`` (conservative).
+    """
+
+    if isinstance(input_tokens, bool) or not isinstance(input_tokens, int) or input_tokens < 0:
+        raise ValueError("input_tokens must be a non-negative integer")
+    if isinstance(output_tokens, bool) or not isinstance(output_tokens, int) or output_tokens < 0:
+        raise ValueError("output_tokens must be a non-negative integer")
+
+    input_rate, _cached_rate, output_rate = effective_rates(
+        model_price, input_tokens, now=now
+    )
+    if cache_write and model_price.cache_write_per_mtok is not None:
+        input_rate = max(input_rate, model_price.cache_write_per_mtok)
+    total = (
+        Decimal(input_tokens) * input_rate + Decimal(output_tokens) * output_rate
+    ) / Decimal("1000000")
+    return total.quantize(MICRO_USD, rounding=ROUND_CEILING)
+
+
+def estimate_request(
+    provider: str,
+    body: Mapping[str, Any],
+    model_price: ModelPrice,
+    now: datetime | date | None = None,
+    cache_write: bool | None = None,
+) -> tuple[int, int, Decimal]:
+    """Convenience helper: ``(input_tokens, output_tokens, reservation_usd)``."""
+
+    input_tokens = estimate_input_tokens(body)
+    output_tokens = max_output_tokens(provider, body, model_price)
+    if cache_write is None:
+        cache_write = cache_write_requested(body)
+    return input_tokens, output_tokens, reservation_usd(
+        model_price, input_tokens, output_tokens, now=now, cache_write=cache_write
+    )