diff --git a/README.md b/README.md index 7efd41a..81bc542 100644 --- a/README.md +++ b/README.md @@ -4,7 +4,7 @@ [![Python](https://img.shields.io/pypi/pyversions/blockrun-litellm.svg)](https://pypi.org/project/blockrun-litellm/) [![License](https://img.shields.io/badge/license-MIT-blue.svg)](LICENSE) -LiteLLM adapter for [BlockRun](https://blockrun.ai) β€” call x402-paid AI models through [LiteLLM](https://github.com/BerriAI/litellm) with zero changes to your existing code. **Base and Solana chains supported.** +LiteLLM adapter for [BlockRun](https://blockrun.ai) β€” call AI models with account API keys or x402 wallet payments through [LiteLLM](https://github.com/BerriAI/litellm) with zero changes to your existing code. **Solana and Base wallets supported.** πŸ“š **Full docs in [`docs/`](docs/)** β€” bilingual (English + δΈ­ζ–‡): - [`CUSTOMER-ONBOARDING`](docs/CUSTOMER-ONBOARDING.md) / [`δΈ­ζ–‡`](docs/CUSTOMER-ONBOARDING.zh.md) β€” 5-minute walkthrough, both modes @@ -14,7 +14,7 @@ LiteLLM adapter for [BlockRun](https://blockrun.ai) β€” call x402-paid AI models - [Chat Completions API](https://blockrun.ai/docs/api-reference/chat-completions) - [Models & pricing](https://blockrun.ai/docs/api-reference/models) -> **TL;DR** β€” BlockRun's `/v1/chat/completions` is already OpenAI-compatible at the protocol level. The only thing that differs is *authentication*: BlockRun uses per-request x402 wallet signatures (non-custodial USDC micropayments on Base / Solana), not a Bearer API key. This package bridges that gap. +> **Authentication:** use `BLOCKRUN_API_KEY` for prepaid account credits, or an x402 wallet for USDC payments on Solana or Base. Both work with the custom provider and sidecar. [中文文摣见底部 / Chinese docs at the bottom](#δΈ­ζ–‡ζ–‡ζ‘£) @@ -27,7 +27,7 @@ LiteLLM adapter for [BlockRun](https://blockrun.ai) β€” call x402-paid AI models | **1. Custom provider** (in-process) | Apps using the LiteLLM **Python library** | `litellm.completion(model="blockrun/openai/gpt-5.5", ...)` | | **2. Local proxy** (sidecar) | Apps using the LiteLLM **Proxy Server** (or any OpenAI client) | `api_base="http://localhost:4001/v1"` | -Both modes share the same underlying wallet/signing flow (via the [`blockrun-llm`](https://github.com/BlockRunAI/blockrun-llm) SDK), so they behave identically. Pick whichever fits your deployment. +Both modes share account authentication or wallet signing (via the [`blockrun-llm`](https://github.com/BlockRunAI/blockrun-llm) SDK), so they behave identically. Pick whichever fits your deployment. ### Verified end-to-end against the live BlockRun gateway @@ -53,16 +53,48 @@ $ curl -sS http://127.0.0.1:4001/v1/chat/completions \ --- +## Account API quick start + +[Register or sign in](https://user.blockrun.ai), [create an API key](https://user.blockrun.ai/dashboard/keys), and [add credits](https://user.blockrun.ai/dashboard/credits). + +This review branch requires the Python SDK implementation in [SDK PR #58](https://github.com/BlockRunAI/blockrun-llm/pull/58). Until that dependency is released, install the exact preview revision from this checkout: + +```bash +pip install -r requirements-api-preview.txt +pip install -e '.[proxy]' +export BLOCKRUN_API_KEY=brk_live_YOUR_KEY +blockrun-litellm-proxy --port 4001 +``` + +Do not publish this integration before the SDK release is available and its minimum version is reflected in `pyproject.toml`. The existing PyPI installation instructions below describe released wallet support. + +```python +import litellm +from blockrun_litellm import register +register() +response = litellm.completion( + model="blockrun/openai/gpt-4o-mini", + messages=[{"role": "user", "content": "Hello"}], +) +print(response.choices[0].message.content) +``` + +`BLOCKRUN_API_KEY` selects account mode without loading or creating a wallet. Per-call `api_key="brk_live_..."` is also supported; an explicit wallet `private_key` selects wallet mode. Supplying both explicitly is rejected. `BLOCKRUN_API_BASE_URL` defaults to `https://api.blockrun.ai` (an optional `/v1` suffix is accepted); a leftover wallet `BLOCKRUN_API_URL` does not override account mode. Native OpenAI chat, Anthropic Messages and Responses are forwarded by the sidecar; media uses authenticated SDK requests and polling. Account errors 401/402/429 preserve their status and retry information without x402 payment retries. + +The account portal is authoritative for credit balance and charges. Account calls are marked `api-key` in audit logs and are not reported as wallet settlements or a zero-dollar charge. Any LiteLLM token-based cost is an estimate. `BLOCKRUN_PROXY_TOKEN` separately protects access to your local sidecar. + +Live account acceptance covers chat JSON/SSE, Messages and Responses JSON. Responses SSE and video completion still depend on [Enterprise PR #10](https://github.com/BlockRunAI/enterprise/pull/10) and its deployment/configuration. Native Gemini `/v1beta` account availability is not verified; call Google models through chat completions when using account keys. Media authentication/polling is tested with deterministic local responses; this is not a claim that every paid media model was exercised live. + ## Install ```bash -# Base chain only β€” minimal +# Account API or Base wallet β€” minimal pip install blockrun-litellm -# Base chain + local OpenAI-compatible proxy (FastAPI/uvicorn) +# Account API or Base wallet + local proxy pip install 'blockrun-litellm[proxy]' -# Base + Solana (adds the x402 SVM toolchain) +# Solana wallet + proxy (adds the x402 SVM toolchain) pip install 'blockrun-litellm[proxy,solana]' ``` @@ -72,8 +104,8 @@ Requires Python β‰₯ 3.9. | Chain | Gateway URL | Wallet env var | Status | |---|---|---|---| -| Base (USDC) | `https://blockrun.ai/api` *(default)* | `BLOCKRUN_WALLET_KEY` | sync + async, streaming | -| Solana (USDC) | `https://sol.blockrun.ai/api` | `SOLANA_WALLET_KEY` | sync + async, streaming on both (since 0.3.1) | +| Solana (USDC) | `https://sol.blockrun.ai/api` *(new wallet default)* | `SOLANA_WALLET_KEY` | sync + async, streaming | +| Base (USDC) | `https://blockrun.ai/api` | `BLOCKRUN_WALLET_KEY` | sync + async, streaming | To route on Solana, pass `api_base="https://sol.blockrun.ai/api"` plus `api_key=` to `litellm.completion(...)` β€” the adapter detects the chain from the URL and uses the right SDK client. @@ -678,7 +710,7 @@ Yes, as of v0.2.0. Pass `stream=True` and the adapter routes through `blockrun-l On your machine only β€” `BLOCKRUN_WALLET_KEY` env var, or `~/.blockrun/.session` if you used `setup_agent_wallet()`. The proxy and provider both read from those sources via `blockrun-llm`. Only EIP-712 signatures are transmitted. **Q: How do I switch between Base and Solana?** -Today this adapter wires to BlockRun's Base gateway (USDC on Base). Solana support tracks the `blockrun-llm` `SolanaLLMClient` and will be added in a follow-up release. +Set `BLOCKRUN_CHAIN=solana` (the new-wallet default) or `BLOCKRUN_CHAIN=base`, or set `BLOCKRUN_API_URL` explicitly. Existing Base-only wallets and saved chain selections are preserved. Account API mode does not require a chain. **Q: Can I run the proxy in Docker / k8s?** Yes β€” it's a vanilla FastAPI app. Pass the wallet key via secret (env var), bind to `0.0.0.0` only inside a private network, and set `BLOCKRUN_PROXY_TOKEN` for an additional auth layer. diff --git a/blockrun_litellm/_adapter.py b/blockrun_litellm/_adapter.py index 3ce9397..0cdb1f0 100644 --- a/blockrun_litellm/_adapter.py +++ b/blockrun_litellm/_adapter.py @@ -28,9 +28,10 @@ import logging import os import threading -from typing import Any, AsyncIterator, Dict, Iterator, List, Optional, Union +from typing import Any, AsyncIterator, Dict, Iterator, List, Optional from blockrun_llm import AsyncLLMClient, ImageClient, LLMClient +from ._auth import account_key, account_auth, cache_key, wallet_url from blockrun_llm.types import APIError, ChatCompletionChunk, PaymentError try: @@ -149,76 +150,50 @@ def _wallet_env_var(api_url: Optional[str]) -> str: return "SOLANA_WALLET_KEY" if _is_solana_url(api_url) else "BLOCKRUN_WALLET_KEY" -def _client_key(api_url: Optional[str], private_key: Optional[str]) -> str: - chain = "solana" if _is_solana_url(api_url) else "base" - fallback_env = os.environ.get(_wallet_env_var(api_url), "") - return f"{chain}::{api_url or ''}::{private_key or fallback_env}" +def _client_key(api_url, private_key, api_key=None): + return cache_key(api_url, private_key, api_key) -def get_sync_client( - api_url: Optional[str] = None, - private_key: Optional[str] = None, -) -> Union[LLMClient, "SolanaLLMClient"]: # type: ignore[name-defined] - """Return a cached sync client for the given creds/url. - - Routes to :class:`SolanaLLMClient` when ``api_url`` points at - ``sol.blockrun.ai``, otherwise :class:`LLMClient` (Base). - """ - key = _client_key(api_url, private_key) +def get_sync_client(api_url=None, private_key=None, api_key=None): + key = _client_key(api_url, private_key, api_key) with _lock: - client = _sync_clients.get(key) - if client is None: - if _is_solana_url(api_url): - if not _HAS_SOLANA: - raise ImportError( - "Solana support requires the solana extra. " - "Install with: pip install 'blockrun-llm[solana]'" - ) - # SolanaLLMClient also reads from SOLANA_WALLET_KEY if no - # explicit key was passed. - client = SolanaLLMClient( - private_key=private_key, - api_url=api_url or SOLANA_API_URL, + if key not in _sync_clients: + auth = account_auth(api_key, private_key, api_url) + if auth: + client = LLMClient( + api_key=account_key(api_key, private_key), + api_url=auth.api_url, timeout=_CHAT_TIMEOUT, ) else: - client = LLMClient(private_key=private_key, api_url=api_url, timeout=_CHAT_TIMEOUT) + url = wallet_url(api_url, private_key) + cls = SolanaLLMClient if _is_solana_url(url) else LLMClient + if cls is None: + raise ImportError("Install blockrun-litellm[solana] for Solana wallets") + client = cls(private_key=private_key, api_url=url, timeout=_CHAT_TIMEOUT) _sync_clients[key] = client - return client + return _sync_clients[key] -def get_async_client( - api_url: Optional[str] = None, - private_key: Optional[str] = None, -) -> Union[AsyncLLMClient, "AsyncSolanaLLMClient"]: # type: ignore[name-defined] - """Return a cached async client for the given creds/url. - - Routes to :class:`AsyncSolanaLLMClient` when ``api_url`` points at - ``sol.blockrun.ai``, otherwise :class:`AsyncLLMClient` (Base). - Requires ``blockrun-llm>=0.22.0`` for the async Solana client. - """ - is_solana = _is_solana_url(api_url) - key = _client_key(api_url, private_key) +def get_async_client(api_url=None, private_key=None, api_key=None): + key = _client_key(api_url, private_key, api_key) with _lock: - client = _async_clients.get(key) - if client is None: - if is_solana: - if not _HAS_SOLANA or AsyncSolanaLLMClient is None: - raise ImportError( - "Solana support requires the solana extra. " - "Install with: pip install 'blockrun-litellm[solana]'" - ) - client = AsyncSolanaLLMClient( - private_key=private_key, - api_url=api_url or SOLANA_API_URL, + if key not in _async_clients: + auth = account_auth(api_key, private_key, api_url) + if auth: + client = AsyncLLMClient( + api_key=account_key(api_key, private_key), + api_url=auth.api_url, timeout=_CHAT_TIMEOUT, ) else: - client = AsyncLLMClient( - private_key=private_key, api_url=api_url, timeout=_CHAT_TIMEOUT - ) + url = wallet_url(api_url, private_key) + cls = AsyncSolanaLLMClient if _is_solana_url(url) else AsyncLLMClient + if cls is None: + raise ImportError("Install blockrun-litellm[solana] for Solana wallets") + client = cls(private_key=private_key, api_url=url, timeout=_CHAT_TIMEOUT) _async_clients[key] = client - return client + return _async_clients[key] # --------------------------------------------------------------------------- @@ -280,6 +255,8 @@ def _strip_real_cost(payload: Dict[str, Any], client: Any) -> Dict[str, Any]: a ``{cost_usd, settlement}`` meta dict (cost may be ``None`` if unavailable).""" cost = payload.pop("cost_usd", None) settlement = payload.pop("settlement", None) + if getattr(client, "auth_mode", None) == "api-key": + return {"auth_mode": "api-key", "cost_usd": None, "settlement": None} if cost is None: cost = getattr(client, "_last_call_cost", None) return {"cost_usd": cost, "settlement": settlement} @@ -296,6 +273,7 @@ def chat_completion_sync( *, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, **openai_kwargs: Any, ) -> Dict[str, Any]: """ @@ -311,7 +289,11 @@ def chat_completion_sync( openai_kwargs.pop("stream", None) is_solana = _is_solana_url(api_url) kwargs = _filter_kwargs(openai_kwargs, is_solana=is_solana) - client = get_sync_client(api_url=api_url, private_key=private_key) + client = get_sync_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) response = client.chat_completion(model=model, messages=messages, **kwargs) payload = response.model_dump(exclude_none=True) payload[_BLOCKRUN_META_KEY] = _strip_real_cost(payload, client) @@ -324,6 +306,7 @@ async def chat_completion_async( *, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, **openai_kwargs: Any, ) -> Dict[str, Any]: """Async variant of :func:`chat_completion_sync`. @@ -334,7 +317,11 @@ async def chat_completion_async( openai_kwargs.pop("stream", None) is_solana = _is_solana_url(api_url) kwargs = _filter_kwargs(openai_kwargs, is_solana=is_solana) - client = get_async_client(api_url=api_url, private_key=private_key) + client = get_async_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) response = await client.chat_completion(model=model, messages=messages, **kwargs) payload = response.model_dump(exclude_none=True) payload[_BLOCKRUN_META_KEY] = _strip_real_cost(payload, client) @@ -352,6 +339,7 @@ def chat_completion_stream_sync( *, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, **openai_kwargs: Any, ) -> Iterator[ChatCompletionChunk]: """ @@ -365,7 +353,11 @@ def chat_completion_stream_sync( openai_kwargs.pop("stream", None) is_solana = _is_solana_url(api_url) kwargs = _filter_kwargs(openai_kwargs, is_solana=is_solana) - client = get_sync_client(api_url=api_url, private_key=private_key) + client = get_sync_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) yield from client.chat_completion_stream(model=model, messages=messages, **kwargs) @@ -375,6 +367,7 @@ async def chat_completion_stream_async( *, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, **openai_kwargs: Any, ) -> AsyncIterator[ChatCompletionChunk]: """Async variant of :func:`chat_completion_stream_sync`. @@ -385,7 +378,11 @@ async def chat_completion_stream_async( openai_kwargs.pop("stream", None) is_solana = _is_solana_url(api_url) kwargs = _filter_kwargs(openai_kwargs, is_solana=is_solana) - client = get_async_client(api_url=api_url, private_key=private_key) + client = get_async_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) async for chunk in client.chat_completion_stream(model=model, messages=messages, **kwargs): yield chunk @@ -413,51 +410,29 @@ def _solana_image_timeout() -> float: return _DEFAULT_SOLANA_IMAGE_TIMEOUT_S -def get_image_client( - api_url: Optional[str] = None, - private_key: Optional[str] = None, -) -> Union[ImageClient, "SolanaLLMClient"]: # type: ignore[name-defined] - """Return a cached image client for the given creds/url. - - Routes to :class:`SolanaLLMClient` when ``api_url`` points at - ``sol.blockrun.ai``, otherwise :class:`ImageClient` (Base). - - The Solana branch is required because ``ImageClient`` only signs EIP-712 - over the EVM ``Account`` β€” sending those payments to the Solana gateway - fails at x402 settlement (``transaction_simulation_failed``). - ``SolanaLLMClient`` exposes ``.image()`` / ``.image_edit()`` that hit the - same ``/v1/images/*`` endpoints with SVM-scheme x402 payments. - """ - key = _client_key(api_url, private_key) +def get_image_client(api_url=None, private_key=None, api_key=None): + key = _client_key(api_url, private_key, api_key) with _lock: - client = _image_clients.get(key) - if client is None: - if _is_solana_url(api_url): - if not _HAS_SOLANA or SolanaLLMClient is None: - raise ImportError( - "Solana support requires the solana extra. " - "Install with: pip install 'blockrun-litellm[solana]'" - ) + if key not in _image_clients: + auth = account_auth(api_key, private_key, api_url) + if auth: + client = ImageClient( + api_key=account_key(api_key, private_key), api_url=auth.api_url + ) + elif _is_solana_url(wallet_url(api_url, private_key)): + if SolanaLLMClient is None: + raise ImportError("Install blockrun-litellm[solana]") client = SolanaLLMClient( private_key=private_key, - api_url=api_url or SOLANA_API_URL, - # Raise the per-image-request timeout ceiling. The SDK caps - # each image POST at ``image_timeout`` (SolanaLLMClient - # default 200s); slow models such as ``openai/gpt-image-2`` - # can exceed that on the synchronous Solana path, so the - # sidecar would otherwise throw ``httpx.ReadTimeout`` mid- - # generation. NOTE: the general ``timeout=`` kwarg is the - # chat baseline and is overridden per-request for images - # (``_request_image_with_payment`` passes ``image_timeout``), - # so ``image_timeout=`` is the knob that actually governs - # image calls. Tunable via BLOCKRUN_SOLANA_IMAGE_TIMEOUT - # for ops without a redeploy. + api_url=wallet_url(api_url, private_key), image_timeout=_solana_image_timeout(), ) else: - client = ImageClient(private_key=private_key, api_url=api_url) + client = ImageClient( + private_key=private_key, api_url=wallet_url(api_url, private_key) + ) _image_clients[key] = client - return client + return _image_clients[key] def _is_solana_image_client(client: Any) -> bool: @@ -546,11 +521,14 @@ def image_generation_sync( quality: Optional[str] = None, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, ) -> Dict[str, Any]: - client = get_image_client(api_url=api_url, private_key=private_key) - response = _invoke_image_generate( - client, prompt, model=model, size=size, n=n, quality=quality + client = get_image_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), ) + response = _invoke_image_generate(client, prompt, model=model, size=size, n=n, quality=quality) return response.model_dump(exclude_none=True) @@ -563,8 +541,13 @@ async def image_generation_async( quality: Optional[str] = None, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, ) -> Dict[str, Any]: - client = get_image_client(api_url=api_url, private_key=private_key) + client = get_image_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) loop = asyncio.get_event_loop() response = await loop.run_in_executor( _image_executor, @@ -586,8 +569,13 @@ def image_edit_sync( quality: Optional[str] = None, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, ) -> Dict[str, Any]: - client = get_image_client(api_url=api_url, private_key=private_key) + client = get_image_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) response = _invoke_image_edit( client, prompt, @@ -612,8 +600,13 @@ async def image_edit_async( quality: Optional[str] = None, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, ) -> Dict[str, Any]: - client = get_image_client(api_url=api_url, private_key=private_key) + client = get_image_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) loop = asyncio.get_event_loop() response = await loop.run_in_executor( _image_executor, @@ -665,21 +658,22 @@ async def image_edit_async( _BASE_MEDIA_CLASSES = {"video": "VideoClient", "music": "MusicClient", "speech": "SpeechClient"} -def _get_media_client(medium: str, api_url: Optional[str], private_key: Optional[str]) -> Any: - """Dedicated Base client for ``medium``, or the unified SolanaLLMClient - (which get_image_client already builds + caches) when the URL is Solana.""" - if _is_solana_url(api_url): - return get_image_client(api_url=api_url, private_key=private_key) +def _get_media_client(medium, api_url, private_key, api_key=None): import blockrun_llm - base_cls = getattr(blockrun_llm, _BASE_MEDIA_CLASSES[medium]) - key = f"{base_cls.__name__}::{_client_key(api_url, private_key)}" + auth = account_auth(api_key, private_key, api_url) + if not auth and _is_solana_url(wallet_url(api_url, private_key)): + return get_image_client(api_url=api_url, private_key=private_key) + cls = getattr(blockrun_llm, _BASE_MEDIA_CLASSES[medium]) + key = cls.__name__ + "::" + _client_key(api_url, private_key, api_key) with _lock: - client = _media_clients.get(key) - if client is None: - client = base_cls(private_key=private_key, api_url=api_url) - _media_clients[key] = client - return client + if key not in _media_clients: + _media_clients[key] = ( + cls(api_key=account_key(api_key, private_key), api_url=auth.api_url) + if auth + else cls(private_key=private_key, api_url=wallet_url(api_url, private_key)) + ) + return _media_clients[key] def _is_solana_client(client: Any) -> bool: @@ -700,20 +694,26 @@ def _solana_media_method(client: Any, method: str) -> Any: return fn -def get_video_client(api_url: Optional[str] = None, private_key: Optional[str] = None) -> Any: +def get_video_client( + api_url: Optional[str] = None, private_key: Optional[str] = None, api_key: Optional[str] = None +) -> Any: """VideoClient (Base) or the unified SolanaLLMClient (Solana).""" - return _get_media_client("video", api_url, private_key) + return _get_media_client("video", api_url, private_key, api_key) -def get_music_client(api_url: Optional[str] = None, private_key: Optional[str] = None) -> Any: +def get_music_client( + api_url: Optional[str] = None, private_key: Optional[str] = None, api_key: Optional[str] = None +) -> Any: """MusicClient (Base) or the unified SolanaLLMClient (Solana).""" - return _get_media_client("music", api_url, private_key) + return _get_media_client("music", api_url, private_key, api_key) -def get_speech_client(api_url: Optional[str] = None, private_key: Optional[str] = None) -> Any: +def get_speech_client( + api_url: Optional[str] = None, private_key: Optional[str] = None, api_key: Optional[str] = None +) -> Any: """SpeechClient (Base) or the unified SolanaLLMClient (Solana). Serves both TTS (speech) and sound-effects.""" - return _get_media_client("speech", api_url, private_key) + return _get_media_client("speech", api_url, private_key, api_key) async def _run_media( @@ -764,6 +764,7 @@ async def video_generation_async( model: Optional[str] = None, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, **params: Any, ) -> Dict[str, Any]: """Generate a video. Extra kwargs (see :data:`VIDEO_PARAM_KEYS`) forward to @@ -771,7 +772,11 @@ async def video_generation_async( arg). Client-supplied ``budget_seconds``/``timeout`` are clamped to the server cap so a request body can't pin a worker thread indefinitely; a malformed (non-numeric) value raises ValueError β†’ HTTP 400 at the proxy.""" - client = get_video_client(api_url=api_url, private_key=private_key) + client = get_video_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) model = _canonical_video_model(model) params = {k: v for k, v in params.items() if v is not None} for knob in ("budget_seconds", "timeout"): @@ -802,10 +807,15 @@ async def music_generation_async( lyrics: Optional[str] = None, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, ) -> Dict[str, Any]: """Generate a music track. Raises ValueError (β†’ HTTP 400 at the proxy) when ``lyrics`` is combined with ``instrumental=True`` β€” the SDK rejects that.""" - client = get_music_client(api_url=api_url, private_key=private_key) + client = get_music_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) # Same call shape on both chains; only the method name differs. media_fn = ( _solana_media_method(client, "music") if _is_solana_client(client) else client.generate @@ -827,9 +837,14 @@ async def speech_generation_async( speed: Optional[float] = None, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, ) -> Dict[str, Any]: """Synthesize speech (TTS).""" - client = get_speech_client(api_url=api_url, private_key=private_key) + client = get_speech_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) kw = {"model": model, "voice": voice, "response_format": response_format, "speed": speed} kw = {k: v for k, v in kw.items() if v is not None} media_fn = ( @@ -848,9 +863,14 @@ async def sound_effect_async( response_format: Optional[str] = None, api_url: Optional[str] = None, private_key: Optional[str] = None, + api_key: Optional[str] = None, ) -> Dict[str, Any]: """Generate a cinematic sound effect.""" - client = get_speech_client(api_url=api_url, private_key=private_key) + client = get_speech_client( + api_url=api_url, + private_key=private_key, + **({"api_key": api_key} if api_key is not None else {}), + ) kw = { "model": model, "duration_seconds": duration_seconds, diff --git a/blockrun_litellm/_auth.py b/blockrun_litellm/_auth.py new file mode 100644 index 0000000..61d8d30 --- /dev/null +++ b/blockrun_litellm/_auth.py @@ -0,0 +1,99 @@ +"""Resolve account vs wallet credentials without touching process-wide auth state.""" + +from __future__ import annotations +import hashlib +import os +from pathlib import Path +from typing import Optional + + +def account_key(api_key: Optional[str] = None, private_key: Optional[str] = None) -> Optional[str]: + if api_key is not None and private_key is not None: + raise ValueError("Pass either api_key or private_key, not both") + key = ( + api_key + if api_key is not None + else (os.getenv("BLOCKRUN_API_KEY") if private_key is None else None) + ) + if key is not None and ( + not key.startswith("brk_live_") or len(key) <= 9 or any(c.isspace() for c in key) + ): + raise ValueError( + "Invalid BlockRun API key; create one at https://user.blockrun.ai/dashboard/keys" + ) + return key + + +def account_auth(api_key=None, private_key=None, api_url=None): + key = account_key(api_key, private_key) + if key is None: + return None + try: + from blockrun_llm.api_key import resolve_api_auth + except ImportError as exc: + raise RuntimeError( + "Account API mode requires the SDK from BlockRunAI/blockrun-llm#58; use requirements-api-preview.txt until its release" + ) from exc + return resolve_api_auth(key, None, api_url or os.getenv("BLOCKRUN_API_BASE_URL")) + + +def wallet_url(api_url=None, private_key=None): + explicit = api_url or os.getenv("BLOCKRUN_API_URL") + if explicit: + return explicit.rstrip("/") + chain = os.getenv("BLOCKRUN_CHAIN") + home = Path.home() / ".blockrun" + if private_key: + chain = "base" if private_key.startswith("0x") or len(private_key) == 64 else "solana" + if not chain: + for name in ("payment-chain", ".chain"): + path = home / name + if path.exists(): + chain = path.read_text().strip() + if chain: + break + if not chain: + base = ( + os.getenv("BLOCKRUN_WALLET_KEY") + or os.getenv("BASE_CHAIN_WALLET_KEY") + or (home / ".session").exists() + ) + sol = os.getenv("SOLANA_WALLET_KEY") or (home / ".solana-session").exists() + chain = "base" if base and not sol else "solana" + if chain not in ("base", "solana"): + raise ValueError("BLOCKRUN_CHAIN must be solana or base") + return "https://sol.blockrun.ai/api" if chain == "solana" else "https://blockrun.ai/api" + + +def cache_key(api_url=None, private_key=None, api_key=None): + key = account_key(api_key, private_key) + if key is not None: + url = ( + (api_url or os.getenv("BLOCKRUN_API_BASE_URL") or "https://api.blockrun.ai") + .rstrip("/") + .removesuffix("/v1") + ) + mode = "api-key" + else: + url = wallet_url(api_url, private_key) + mode = "wallet" + key = ( + private_key + or os.getenv("SOLANA_WALLET_KEY" if "sol.blockrun.ai" in url else "BLOCKRUN_WALLET_KEY") + or os.getenv("BASE_CHAIN_WALLET_KEY") + or "" + ) + return mode + "::" + url + "::" + hashlib.sha256(key.encode()).hexdigest() + + +def provider_credentials(api_key, kwargs): + private = kwargs.get("private_key") or (kwargs.get("optional_params") or {}).get("private_key") + if api_key is not None: + if private is not None: + raise ValueError("Pass either api_key or private_key, not both") + if api_key.startswith("brk_"): + account_key(api_key) + return {"api_key": api_key} + # Backward compatibility for the old api_key-as-wallet interface. + return {"private_key": api_key} + return {"private_key": private} diff --git a/blockrun_litellm/logger.py b/blockrun_litellm/logger.py index b29e30b..35c5590 100644 --- a/blockrun_litellm/logger.py +++ b/blockrun_litellm/logger.py @@ -86,11 +86,7 @@ def _resolve_path(path: Optional[str | Path] = None) -> Path: """Pick the log destination β€” explicit arg > env var > default.""" - resolved = Path( - path - or os.environ.get("BLOCKRUN_LITELLM_LOG") - or DEFAULT_LOG_PATH - ) + resolved = Path(path or os.environ.get("BLOCKRUN_LITELLM_LOG") or DEFAULT_LOG_PATH) resolved.parent.mkdir(parents=True, exist_ok=True) return resolved @@ -226,23 +222,24 @@ def _build_entry( "stream": stream, "latency_ms": _latency_ms(start_time, end_time), "request_id": ( - kwargs.get("litellm_call_id") - or (kwargs.get("metadata") or {}).get("litellm_call_id") + kwargs.get("litellm_call_id") or (kwargs.get("metadata") or {}).get("litellm_call_id") ), } if failure is not None: - entry.update({ - "status": "failure", - "completion": None, - "usage": None, - "cost_usd": None, - "cost_source": None, - "estimated_cost_usd": None, - "settlement": None, - "error_type": type(failure).__name__, - "error_message": str(failure), - }) + entry.update( + { + "status": "failure", + "completion": None, + "usage": None, + "cost_usd": None, + "cost_source": None, + "estimated_cost_usd": None, + "settlement": None, + "error_type": type(failure).__name__, + "error_message": str(failure), + } + ) return entry usage = _extract_usage(response_obj) @@ -254,23 +251,29 @@ def _build_entry( return None estimate = _extract_cost(response_obj, kwargs) real = _extract_real_cost(response_obj) + hidden = getattr(response_obj, "_hidden_params", {}) or {} + account_mode = hidden.get("blockrun_auth_mode") == "api-key" + if account_mode: + real = {"cost_usd": None, "settlement": None} if real["cost_usd"] is not None: cost_usd = real["cost_usd"] cost_source = "blockrun_x402" else: cost_usd = estimate - cost_source = "litellm_estimate" - entry.update({ - "status": "success", - "completion": completion, - "usage": usage, - # Real wallet deduction when known (x402), else LiteLLM's estimate. - "cost_usd": cost_usd, - "cost_source": cost_source, - # Keep LiteLLM's tokenΓ—list-price estimate alongside for comparison. - "estimated_cost_usd": estimate, - "settlement": real["settlement"], - }) + cost_source = "account_estimate" if account_mode else "litellm_estimate" + entry.update( + { + "status": "success", + "completion": completion, + "usage": usage, + # Real wallet deduction when known (x402), else LiteLLM's estimate. + "cost_usd": cost_usd, + "cost_source": cost_source, + # Keep LiteLLM's tokenΓ—list-price estimate alongside for comparison. + "estimated_cost_usd": estimate, + "settlement": real["settlement"], + } + ) return entry @@ -285,6 +288,7 @@ def log_proxy_call( latency_ms: Optional[float], request_id: Optional[str] = None, settlement_status: Optional[str] = None, + auth_mode: Optional[str] = None, ) -> None: """Append a JSONL audit row for a raw FastAPI sidecar passthrough call. @@ -328,6 +332,11 @@ def log_proxy_call( "settlement": settlement, "request_id": request_id, } + if auth_mode == "api-key": + entry.update( + auth_mode="api-key", cost_usd=None, cost_source="account_portal", settlement=None + ) + settlement_status = None # Omitted when there's nothing to flag, so existing rows keep their shape # and anything parsing this file sees the key only when it means something. if settlement_status is not None: @@ -366,8 +375,10 @@ def log_success_event(self, kwargs, response_obj, start_time, end_time): def log_failure_event(self, kwargs, response_obj, start_time, end_time): # LiteLLM passes the exception object as the second positional arg in # the failure hook (it's called ``response_obj`` for legacy reasons). - exc = response_obj if isinstance(response_obj, BaseException) else ( - kwargs.get("exception") if isinstance(kwargs, dict) else None + exc = ( + response_obj + if isinstance(response_obj, BaseException) + else (kwargs.get("exception") if isinstance(kwargs, dict) else None) ) entry = _build_entry( kwargs, None, start_time, end_time, failure=exc or Exception("unknown") diff --git a/blockrun_litellm/provider.py b/blockrun_litellm/provider.py index 7a8135d..cfdad1b 100644 --- a/blockrun_litellm/provider.py +++ b/blockrun_litellm/provider.py @@ -45,6 +45,7 @@ from blockrun_llm.types import ChatCompletionChunk, ChatUsage from blockrun_litellm import _adapter +from ._auth import provider_credentials, account_key # Provider name surfaced to LiteLLM exception classes so the router knows @@ -81,10 +82,23 @@ def _translate_to_litellm(exc: Exception, model: str) -> Optional[Exception]: ) if isinstance(exc, BlockRunAPIError): status = getattr(exc, "status_code", 0) - if status == 429: - return litellm.RateLimitError( - message=str(exc), model=model, llm_provider=_LITELLM_PROVIDER - ) + if status in (401, 402, 429): + if status == 401: + error = litellm.AuthenticationError( + message=str(exc), model=model, llm_provider=_LITELLM_PROVIDER + ) + elif status == 429: + error = litellm.RateLimitError( + message=str(exc), model=model, llm_provider=_LITELLM_PROVIDER + ) + else: + error = litellm.APIError( + status_code=402, message=str(exc), model=model, llm_provider=_LITELLM_PROVIDER + ) + error.retry_after = getattr(exc, "retry_after", None) + if error.retry_after and getattr(error, "response", None) is not None: + error.response.headers["Retry-After"] = error.retry_after + return error if status == 500: return litellm.InternalServerError( message=str(exc), model=model, llm_provider=_LITELLM_PROVIDER @@ -103,7 +117,7 @@ def _translate_to_litellm(exc: Exception, model: str) -> Optional[Exception]: # LiteLLM passes the provider-stripped model name *and* an "optional_params" # dict containing the OpenAI-style params (temperature, max_tokens, ...). # It also forwards ``api_base`` / ``api_key`` from the call site, which we -# repurpose: ``api_base`` β†’ BlockRun ``api_url``, ``api_key`` β†’ wallet key. +# map to the gateway URL and account key (legacy wallet keys remain accepted). _OPTIONAL_KEYS = ( "max_tokens", "temperature", @@ -169,6 +183,10 @@ def _attach_real_cost(response: litellm.ModelResponse, meta: Optional[Dict[str, """ if not meta: return + if meta.get("auth_mode") == "api-key": + response._hidden_params["blockrun_auth_mode"] = "api-key" + response._hidden_params["blockrun_cost_source"] = "account_portal" + return cost = meta.get("cost_usd") if cost is None: return @@ -474,7 +492,7 @@ def completion( model=model, messages=messages, api_url=api_base, - private_key=api_key, + **provider_credentials(api_key, kwargs), **openai_kwargs, ) except Exception as exc: @@ -498,7 +516,7 @@ async def acompletion( model=model, messages=messages, api_url=api_base, - private_key=api_key, + **provider_credentials(api_key, kwargs), **openai_kwargs, ) except Exception as exc: @@ -535,15 +553,23 @@ def streaming( model=model, messages=messages, api_url=api_base, - private_key=api_key, + **provider_credentials(api_key, kwargs), **openai_kwargs, ): # Per-call x402 charge the SDK attaches to each chunk (race-free, # vs the shared client._last_call_cost). ``None`` on older SDKs # that don't attach it -> estimate fallback, no injection. - cost = getattr(chunk, "cost_usd", None) + creds = provider_credentials(api_key, kwargs) + account_mode = ( + account_key(creds.get("api_key"), creds.get("private_key")) is not None + ) + cost = None if account_mode else getattr(chunk, "cost_usd", None) for gchunk in _iter_stream_chunks(chunk, stream_state): - if cost is not None: + if account_mode: + gchunk.setdefault("_hidden_params", {}).update( + blockrun_auth_mode="api-key", blockrun_cost_source="account_portal" + ) + elif cost is not None: _inject_real_cost(gchunk, cost) yield gchunk except Exception as exc: @@ -569,12 +595,20 @@ async def astreaming( model=model, messages=messages, api_url=api_base, - private_key=api_key, + **provider_credentials(api_key, kwargs), **openai_kwargs, ): - cost = getattr(chunk, "cost_usd", None) + creds = provider_credentials(api_key, kwargs) + account_mode = ( + account_key(creds.get("api_key"), creds.get("private_key")) is not None + ) + cost = None if account_mode else getattr(chunk, "cost_usd", None) for gchunk in _iter_stream_chunks(chunk, stream_state): - if cost is not None: + if account_mode: + gchunk.setdefault("_hidden_params", {}).update( + blockrun_auth_mode="api-key", blockrun_cost_source="account_portal" + ) + elif cost is not None: _inject_real_cost(gchunk, cost) yield gchunk except Exception as exc: diff --git a/blockrun_litellm/proxy.py b/blockrun_litellm/proxy.py index 7d9e879..ce8ac92 100644 --- a/blockrun_litellm/proxy.py +++ b/blockrun_litellm/proxy.py @@ -73,6 +73,7 @@ from blockrun_llm.tx_log import decode_settlement_header from blockrun_litellm import _adapter +from ._auth import account_key, account_auth, wallet_url, cache_key from blockrun_litellm import logger as _logger # Optional β€” present when blockrun-llm[solana] is installed alongside solana-py. @@ -326,17 +327,35 @@ def close(self) -> None: self._base.close() +def _account_error_response(exc: APIError) -> JSONResponse: + headers = {} + retry = getattr(exc, "retry_after", None) + if retry: + headers["Retry-After"] = retry + detail = getattr(exc, "response", None) or {"message": str(exc)} + return JSONResponse( + status_code=exc.status_code or 502, content={"error": detail}, headers=headers + ) + + def _resolve_api_url() -> str: - return (os.environ.get("BLOCKRUN_API_URL") or _adapter.BASE_API_URL).rstrip("/") + auth = account_auth() + return auth.api_url if auth else wallet_url() def _messages_client(api_url: str) -> httpx.Client: """Cached httpx client whose transport signs x402 for any path on the chain implied by ``api_url`` (Base via EIP-712, Solana via SVM).""" - existing = _messages_http_clients.get(api_url) + auth = account_auth(api_url=api_url) + cache_id = cache_key(api_url) if auth else api_url + existing = _messages_http_clients.get(cache_id) if existing is not None: return existing + if auth: + client = httpx.Client(auth=auth, timeout=_adapter._CHAT_TIMEOUT, follow_redirects=False) + _messages_http_clients[cache_id] = client + return client if _adapter._is_solana_url(api_url): from blockrun_llm.solana_wallet import load_solana_wallet @@ -364,7 +383,7 @@ def _messages_client(api_url: str) -> httpx.Client: client = httpx.Client( transport=_SignedAmountTransport(transport), timeout=_adapter._CHAT_TIMEOUT ) - _messages_http_clients[api_url] = client + _messages_http_clients[cache_id] = client return client @@ -569,6 +588,7 @@ async def _forward_passthrough( qs = request.url.query if forward_query else "" target = f"{api_url}{path}" + (f"?{qs}" if qs else "") + auth_mode = "api-key" if account_key() is not None else "wallet" _t0 = time.monotonic() model = model_override or _body_model(raw) req_id = request.headers.get("x-request-id") @@ -590,6 +610,19 @@ def _latency_ms() -> float: async with _get_semaphore(): try: resp = await run_in_threadpool(_open_upstream_stream, client, target, raw, headers) + except APIError as exc: + _logger.log_proxy_call( + auth_mode=auth_mode, + model=model, + path=path, + stream=wants_stream, + http_status=exc.status_code, + cost_usd=None, + settlement=None, + latency_ms=_latency_ms(), + request_id=req_id, + ) + return _account_error_response(exc) except Exception as exc: # noqa: BLE001 # A Solana JSON-RPC fault during x402 signing (e.g. getAccountInfo # timeout) surfaces here, BEFORE any upstream status exists. Map it @@ -611,6 +644,7 @@ def _latency_ms() -> float: body = await run_in_threadpool(resp.read) await run_in_threadpool(resp.close) _logger.log_proxy_call( + auth_mode=auth_mode, model=model, path=path, stream=True, @@ -633,6 +667,7 @@ def _latency_ms() -> float: ) _logger.log_proxy_call( + auth_mode=auth_mode, model=model, path=path, stream=True, @@ -676,12 +711,26 @@ def _post(): async with _get_semaphore(): try: status, ctype, content, cost, settlement = await run_in_threadpool(_post) + except APIError as exc: + _logger.log_proxy_call( + auth_mode=auth_mode, + model=model, + path=path, + stream=False, + http_status=exc.status_code, + cost_usd=None, + settlement=None, + latency_ms=_latency_ms(), + request_id=req_id, + ) + return _account_error_response(exc) except Exception as exc: # noqa: BLE001 if _is_solana_rpc_exc(exc): log.warning("solana rpc error during payment signing: %s", _solana_rpc_msg(exc)) return JSONResponse(status_code=503, content={"error": _solana_rpc_msg(exc)}) raise _logger.log_proxy_call( + auth_mode=auth_mode, model=model, path=path, stream=False, @@ -939,8 +988,10 @@ async def _media_endpoint( the ValueError arm they'd answer 400 β€” blaming the caller for a call they paid for. It is upstream's fault: 502, flagged, and logged. """ + auth_mode = "api-key" if account_key() is not None else "wallet" _t0 = time.monotonic() req_id = uuid.uuid4().hex + error_headers = {} result: Optional[Dict[str, Any]] = None parse_failed_after_settlement = False reached_gateway = True @@ -979,6 +1030,8 @@ async def _media_endpoint( except APIError as exc: status = exc.status_code if 400 <= getattr(exc, "status_code", 0) < 600 else 502 payload = {"error": str(exc)} + if getattr(exc, "retry_after", None) is not None: + error_headers["Retry-After"] = str(exc.retry_after) except Exception as exc: # noqa: BLE001 - a missing row is worse than a broad catch # Transport errors (httpx.ReadTimeout on a 10-minute image call, # connection resets) escape the SDK unwrapped. If one lands after the @@ -994,6 +1047,7 @@ async def _media_endpoint( } settlement = _media_settlement(result) _logger.log_proxy_call( + auth_mode=auth_mode, model=model, path=path, stream=False, @@ -1011,6 +1065,7 @@ async def _media_endpoint( ), ) headers = _cost_response_headers(None, settlement) + headers.update(error_headers) if warning: headers[_WARNING_HEADER] = warning return JSONResponse(status_code=status, content=payload, headers=headers) @@ -1248,9 +1303,7 @@ async def image_edits(request: Request) -> Any: if not values: raise HTTPException(400, "`image` is required") if len(values) > _MAX_IMAGE_PARTS: - raise HTTPException( - 400, f"at most {_MAX_IMAGE_PARTS} image parts (got {len(values)})" - ) + raise HTTPException(400, f"at most {_MAX_IMAGE_PARTS} image parts (got {len(values)})") images = [await _image_form_value(value, "image") for value in values] image: Any = images[0] if len(images) == 1 else images mask_value = form.get("mask") @@ -1402,7 +1455,7 @@ def _openai_video_kwargs(body: Dict[str, Any]) -> Dict[str, Any]: try: kwargs["duration_seconds"] = int(float(seconds)) except (TypeError, ValueError): - raise HTTPException(400, "`seconds` must be numeric (e.g. \"8\")") + raise HTTPException(400, '`seconds` must be numeric (e.g. "8")') if body.get("size") is not None: kwargs.update(_map_openai_video_size(body["size"])) kwargs.update({k: body[k] for k in _adapter.VIDEO_PARAM_KEYS if body.get(k) is not None}) @@ -1448,6 +1501,7 @@ async def _run_video_job(job: Dict[str, Any], prompt: str, kwargs: Dict[str, Any """Drive the blocking SDK submit+poll for one video job and record the outcome on the job dict. Errors are folded into the OpenAI ``error`` shape so a poller sees status=failed instead of a hung queue.""" + auth_mode = "api-key" if account_key() is not None else "wallet" _t0 = time.monotonic() status = 200 result: Optional[Dict[str, Any]] = None @@ -1497,6 +1551,7 @@ async def _run_video_job(job: Dict[str, Any], prompt: str, kwargs: Dict[str, Any job["error"] = {"code": "server_error", "message": str(exc)} settlement = _video_job_settlement(job) _logger.log_proxy_call( + auth_mode=auth_mode, model=job["model"], path="/v1/videos", stream=False, @@ -1560,9 +1615,12 @@ async def openai_videos_create(request: Request) -> Any: def _get_video_job_or_404(video_id: str) -> Dict[str, Any]: job = _video_jobs.get(video_id) if job is None: - raise HTTPException(404, f"video job '{video_id}' not found (jobs expire after " - f"{int(_VIDEO_JOB_TTL_S)}s and live on the sidecar instance " - "that accepted the create)") + raise HTTPException( + 404, + f"video job '{video_id}' not found (jobs expire after " + f"{int(_VIDEO_JOB_TTL_S)}s and live on the sidecar instance " + "that accepted the create)", + ) return job @@ -1780,9 +1838,7 @@ def _usage_to_responses(usage: Dict[str, Any]) -> Dict[str, Any]: "output_tokens": _int_or_zero(usage.get("completion_tokens")), "total_tokens": _int_or_zero(usage.get("total_tokens")), "input_tokens_details": {"cached_tokens": cached}, - "output_tokens_details": { - "reasoning_tokens": _int_or_zero(ctd.get("reasoning_tokens")) - }, + "output_tokens_details": {"reasoning_tokens": _int_or_zero(ctd.get("reasoning_tokens"))}, } @@ -1987,7 +2043,11 @@ def base( @app.post("/v1/responses", dependencies=[Depends(_require_token)]) async def responses(request: Request) -> Any: - """OpenAI Responses API bridge β†’ BlockRun Chat Completions.""" + """Native account Responses, or the legacy wallet Chat Completions bridge.""" + if account_key() is not None: + return await _forward_passthrough( + request, "/v1/responses", _openai_fwd_headers(request), allow_stream=True + ) try: body = await request.json() except Exception: @@ -2052,13 +2112,15 @@ def main() -> None: args = parser.parse_args() if args.api_url: - os.environ["BLOCKRUN_API_URL"] = args.api_url + os.environ["BLOCKRUN_API_BASE_URL" if account_key() is not None else "BLOCKRUN_API_URL"] = ( + args.api_url + ) # Fail fast if no wallet β€” better than waiting for first request. try: _adapter.get_sync_client() except ValueError as exc: - parser.exit(2, f"\nWallet not configured:\n {exc}\n") + parser.exit(2, f"\nAuthentication not configured:\n {exc}\n") import uvicorn diff --git a/requirements-api-preview.txt b/requirements-api-preview.txt new file mode 100644 index 0000000..38c6ccf --- /dev/null +++ b/requirements-api-preview.txt @@ -0,0 +1,2 @@ +# Review-only SDK dependency. Replace with the canonical release before publishing. +blockrun-llm @ git+https://github.com/KillerQueen-Z/blockrun-llm.git@28052aa4d5f7a54ab23f61375dfc222d22110300 diff --git a/tests/conftest.py b/tests/conftest.py index 6b17c33..db2fcde 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -60,6 +60,8 @@ def _no_wallet_required(monkeypatch: pytest.MonkeyPatch) -> None: """Make sure tests run without `BLOCKRUN_WALLET_KEY` set.""" monkeypatch.delenv("BLOCKRUN_WALLET_KEY", raising=False) monkeypatch.delenv("BASE_CHAIN_WALLET_KEY", raising=False) + monkeypatch.delenv("BLOCKRUN_API_KEY", raising=False) + monkeypatch.setenv("BLOCKRUN_API_URL", "https://blockrun.ai/api") @pytest.fixture diff --git a/tests/test_api_key.py b/tests/test_api_key.py new file mode 100644 index 0000000..55ca5cb --- /dev/null +++ b/tests/test_api_key.py @@ -0,0 +1,248 @@ +import asyncio +import json +import httpx +import pytest +from fastapi.testclient import TestClient +from blockrun_litellm import _adapter, proxy +from blockrun_litellm._auth import account_key, cache_key, wallet_url +from blockrun_litellm.provider import BlockRunLLM + +KEY = "brk_live_account_test" +CHAT = { + "id": "chat-test", + "model": "openai/gpt-4o-mini", + "object": "chat.completion", + "created": 1, + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, +} + + +@pytest.fixture(autouse=True) +def account_fixture(monkeypatch, tmp_path): + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setenv("BLOCKRUN_API_KEY", KEY) + monkeypatch.delenv("BLOCKRUN_API_BASE_URL", raising=False) + monkeypatch.delenv("BLOCKRUN_PROXY_TOKEN", raising=False) + # A leftover wallet URL must not redirect account traffic onto an x402 gateway. + monkeypatch.setenv("BLOCKRUN_API_URL", "https://sol.blockrun.ai/api") + for d in ( + _adapter._sync_clients, + _adapter._async_clients, + _adapter._image_clients, + _adapter._media_clients, + proxy._messages_http_clients, + ): + d.clear() + yield + for d in ( + _adapter._sync_clients, + _adapter._async_clients, + _adapter._image_clients, + _adapter._media_clients, + proxy._messages_http_clients, + ): + d.clear() + + +def mock_http(monkeypatch, respond): + calls = [] + + def handle(self, r): + calls.append(r) + return respond(r) + + async def ahandle(self, r): + return handle(self, r) + + monkeypatch.setattr(httpx.HTTPTransport, "handle_request", handle) + monkeypatch.setattr(httpx.AsyncHTTPTransport, "handle_async_request", ahandle) + return calls + + +def test_provider_sync_async_uses_account_key_and_no_wallet_cost(monkeypatch): + calls = mock_http(monkeypatch, lambda r: httpx.Response(200, json=CHAT, request=r)) + handler = BlockRunLLM() + r = handler.completion( + model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], api_key=KEY + ) + assert r.choices[0].message.content == "OK" + assert r._hidden_params["blockrun_auth_mode"] == "api-key" + assert "blockrun_cost_usd" not in r._hidden_params + r = asyncio.run( + handler.acompletion( + model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}], api_key=KEY + ) + ) + assert r.choices[0].message.content == "OK" + assert len(calls) == 2 + for c in calls: + assert str(c.url) == "https://api.blockrun.ai/v1/chat/completions" + assert c.headers["authorization"] == "Bearer " + KEY + assert not any("payment" in h for h in c.headers) + + +@pytest.mark.parametrize("status", [401, 402, 429]) +def test_provider_preserves_account_errors(monkeypatch, status): + calls = mock_http( + monkeypatch, + lambda r: httpx.Response( + status, + json={"error": {"message": KEY, "code": "quota"}}, + headers={"retry-after": "12"}, + request=r, + ), + ) + with pytest.raises(Exception) as err: + BlockRunLLM().completion( + model="openai/gpt-4o-mini", messages=[{"role": "user", "content": "hi"}] + ) + assert err.value.status_code == status + assert err.value.retry_after == "12" + assert KEY not in str(err.value) + assert len(calls) == 1 + + +@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/messages", "/v1/responses"]) +@pytest.mark.parametrize("stream", [False, True]) +def test_sidecar_native_protocol_and_headers(monkeypatch, path, stream): + payload = { + "model": "openai/gpt-4o-mini", + "input": "hi", + "messages": [{"role": "user", "content": "hi"}], + "stream": stream, + "tools": [{"type": "function", "name": "lookup", "parameters": {"type": "object"}}], + } + wire = ( + b'event: response.completed\ndata: {"type":"response.completed","response":{"output":[]}}\n\n' + if stream + else json.dumps( + { + "id": "native-id", + "output": [], + "content": [{"type": "text", "text": "OK"}], + "choices": [], + } + ).encode() + ) + calls = mock_http( + monkeypatch, + lambda r: httpx.Response( + 200, + stream=httpx.ByteStream(wire), + headers={"content-type": "text/event-stream" if stream else "application/json"}, + request=r, + ), + ) + with TestClient(proxy.app) as client: + r = client.post( + path, json=payload, headers={"x-api-key": "placeholder", "payment-signature": "remove"} + ) + assert r.status_code == 200 + assert r.content == wire + assert len(calls) == 1 + assert str(calls[0].url) == "https://api.blockrun.ai" + path + assert json.loads(calls[0].content) == payload + assert calls[0].headers["authorization"] == "Bearer " + KEY + assert "x-api-key" not in calls[0].headers + + +@pytest.mark.parametrize("status", [401, 402, 429]) +@pytest.mark.parametrize("stream", [False, True]) +def test_sidecar_error_code_retry_after_no_x402(monkeypatch, status, stream): + calls = mock_http( + monkeypatch, + lambda r: httpx.Response( + status, + json={"error": {"message": KEY, "code": "quota"}}, + headers={"retry-after": "12"}, + request=r, + ), + ) + with TestClient(proxy.app) as client: + r = client.post("/v1/chat/completions", json={"model": "m", "stream": stream}) + assert r.status_code == status + assert r.headers["retry-after"] == "12" + assert KEY not in r.text + assert len(calls) == 1 + + +def test_cache_uses_key_fingerprint_and_rotation(monkeypatch): + a = _adapter.get_sync_client() + key1 = cache_key() + assert KEY not in key1 + monkeypatch.setenv("BLOCKRUN_API_KEY", "brk_live_second_account") + b = _adapter.get_sync_client() + assert a is not b + assert key1 != cache_key() + assert proxy._resolve_api_url() == "https://api.blockrun.ai" + + +def test_media_clients_are_account_mode(): + for f in ( + _adapter.get_image_client, + _adapter.get_video_client, + _adapter.get_music_client, + _adapter.get_speech_client, + ): + assert f().auth_mode == "api-key" + + +@pytest.mark.asyncio +async def test_account_image_202_polling(monkeypatch): + def response(r): + if r.method == "POST": + return httpx.Response( + 202, + json={"status": "queued", "poll_url": "/api/v1/images/generations/test"}, + request=r, + ) + return httpx.Response( + 200, + json={ + "status": "completed", + "created": 1, + "data": [{"url": "https://cdn.example/test.png"}], + }, + request=r, + ) + + calls = mock_http(monkeypatch, response) + result = await _adapter.image_generation_async("cat", api_key=KEY) + assert result["data"] + assert [r.method for r in calls] == ["POST", "GET"] + assert calls[1].url.path == "/v1/images/generations/test" + + +def test_logger_marks_account_cost_as_unavailable_not_x402(monkeypatch, tmp_path): + from blockrun_litellm import logger + + path = tmp_path / "audit.jsonl" + monkeypatch.setenv("BLOCKRUN_LITELLM_LOG", str(path)) + logger.log_proxy_call( + model="m", + path="/v1/messages", + stream=True, + http_status=200, + cost_usd=99, + settlement={"tx_hash": "mock"}, + latency_ms=1, + auth_mode="api-key", + ) + row = json.loads(path.read_text()) + assert row["cost_source"] == "account_portal" + assert row["cost_usd"] is None + assert row["settlement"] is None + + +def test_wallet_selection_preserves_explicit_base_and_prefers_new_solana(monkeypatch, tmp_path): + monkeypatch.delenv("BLOCKRUN_API_URL") + monkeypatch.delenv("BLOCKRUN_API_KEY") + monkeypatch.delenv("BLOCKRUN_CHAIN", raising=False) + monkeypatch.delenv("SOLANA_WALLET_KEY", raising=False) + assert wallet_url() == "https://sol.blockrun.ai/api" + assert wallet_url(private_key="0x" + "1" * 64) == "https://blockrun.ai/api" + with pytest.raises(ValueError): + account_key("", "0x" + "1" * 64)