diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index ad92982..f0ca04f 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -44,6 +44,34 @@ jobs: pytest tests/unit --ignore=tests/unit/test_solana_client.py --ignore=tests/unit/test_solana_wallet.py -k "not SolanaX402" fi + # llama-index-llms-blockrun, published from integrations/ as its own PyPI + # package. Installed against THIS checkout of the SDK, so an SDK change that + # breaks the adapter fails here rather than after both are released. The + # lowest leg pins the floor of every declared range: a floor that is never + # installed is a floor nobody has checked (0.6.5 of the OpenAI serializer + # raises on a ThinkingBlock that 0.8 skips). + llama-index: + runs-on: ubuntu-latest + strategy: + matrix: + include: + - python-version: '3.10' + pins: 'llama-index-core==0.14.5 llama-index-llms-openai==0.6.5' + - python-version: '3.12' + pins: '' + defaults: + run: + working-directory: integrations/llama-index-llms-blockrun + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: ${{ matrix.python-version }} + - name: Install + run: pip install -e "../..[solana]" -e ".[dev]" ${{ matrix.pins }} + - name: Run unit tests + run: pytest + # Fails when a number in the docs disagrees with brand-numbers.json. # Offline by design: it reads the committed snapshot and never fetches, so a # blockrun.ai deploy in progress cannot fail this repo's CI. Its own job diff --git a/.github/workflows/publish-llama-index.yml b/.github/workflows/publish-llama-index.yml new file mode 100644 index 0000000..0137b05 --- /dev/null +++ b/.github/workflows/publish-llama-index.yml @@ -0,0 +1,66 @@ +name: Publish llama-index-llms-blockrun to PyPI + +# Released separately from the SDK: cut a GitHub release tagged +# llama-index-llms-blockrun-v (the version in +# integrations/llama-index-llms-blockrun/pyproject.toml). publish.yml skips +# these tags. +on: + release: + types: [published] + +jobs: + build: + if: ${{ startsWith(github.event.release.tag_name, 'llama-index-llms-blockrun-v') }} + runs-on: ubuntu-latest + defaults: + run: + working-directory: integrations/llama-index-llms-blockrun + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: '3.12' + + - name: Install package and test deps + run: pip install -e "../..[solana]" -e ".[dev]" build + + # PyPI never lets a published file be replaced, so a tag that names the + # wrong version is refused here, before anything is uploaded. + - name: Verify the release tag matches the package version + run: | + tag="${{ github.event.release.tag_name }}" + declared="llama-index-llms-blockrun-v$(python -c 'import tomllib; print(tomllib.load(open("pyproject.toml", "rb"))["project"]["version"])')" + if [ "$tag" != "$declared" ]; then + echo "Release tag $tag does not match pyproject version ($declared)." >&2 + exit 1 + fi + echo "Release tag $tag matches pyproject version." + + - name: Verify the release is publishable + run: pytest -q + + - name: Build package + run: python -m build + + - name: Upload artifacts + uses: actions/upload-artifact@v4 + with: + name: dist-llama-index + path: integrations/llama-index-llms-blockrun/dist/ + + publish: + needs: build + runs-on: ubuntu-latest + environment: pypi + permissions: + id-token: write + steps: + - name: Download artifacts + uses: actions/download-artifact@v4 + with: + name: dist-llama-index + path: dist/ + + - name: Publish to PyPI + uses: pypa/gh-action-pypi-publish@release/v1 diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 5ee5c8e..dd417d4 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -6,6 +6,9 @@ on: jobs: build: + # llama-index-llms-blockrun-v* releases publish a different package, from + # publish-llama-index.yml; this job would only fail their VERSION check. + if: ${{ !startsWith(github.event.release.tag_name, 'llama-index-llms-blockrun-v') }} runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 diff --git a/README.md b/README.md index 3177d0a..ed3753e 100644 --- a/README.md +++ b/README.md @@ -169,6 +169,24 @@ The 402 from that host names `eip155:5042`, and the SDK signs the EIP-3009 autho 2. Fund it with USDC on Arc (Arc's native token, shown as the ERC-20 at `0x3600…0000`) 3. `api_url="https://arc.blockrun.ai/api"` — payments are automatic via x402 +## LlamaIndex + +[`llama-index-llms-blockrun`](integrations/llama-index-llms-blockrun) makes any +BlockRun chat model a LlamaIndex LLM — chat, streaming, async, tool calling and +agents — paying through this SDK on the chain you pick: + +```bash +pip install llama-index-llms-blockrun # pay on Base +pip install "llama-index-llms-blockrun[solana]" # pay on Solana +``` + +```python +from llama_index.llms.blockrun import BlockRun + +llm = BlockRun(model="openai/gpt-5.5") # USDC on Base +llm = BlockRun(model="openai/gpt-5.5", chain="solana") # USDC on Solana +``` + ## Smart Routing (Router Core) Let the SDK automatically pick the cheapest capable model for each request: diff --git a/integrations/llama-index-llms-blockrun/LICENSE b/integrations/llama-index-llms-blockrun/LICENSE new file mode 100644 index 0000000..7812798 --- /dev/null +++ b/integrations/llama-index-llms-blockrun/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 BlockRun + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/integrations/llama-index-llms-blockrun/README.md b/integrations/llama-index-llms-blockrun/README.md new file mode 100644 index 0000000..b14d90a --- /dev/null +++ b/integrations/llama-index-llms-blockrun/README.md @@ -0,0 +1,128 @@ +# LlamaIndex LLM Integration: BlockRun + +[BlockRun](https://blockrun.ai) is one OpenAI-compatible gateway to models from +OpenAI, Anthropic, Google, DeepSeek, xAI, NVIDIA and others, where every call +pays for itself in USDC over [x402](https://x402.org). There is no API key and +no account: you pick the chain the wallet pays from, Base or Solana. + +## Why not `OpenAILike`? + +`OpenAILike(api_base="https://blockrun.ai/api/v1", api_key="fake")` reaches the +gateway and stops at the first paid call. The gateway answers an unpaid request +with HTTP 402 and a price, and nothing in the OpenAI client knows how to sign a +payment for it. This package hands that step to the +[`blockrun-llm`](https://pypi.org/project/blockrun-llm/) SDK, which signs the +payment locally and retries the request with the signature attached. + +## Install + +```bash +pip install llama-index-llms-blockrun # pay on Base +pip install "llama-index-llms-blockrun[solana]" # pay on Solana +``` + +## Pick a chain + +| `chain=` | Pays with | Wallet, in order of precedence | +| ------------------ | -------------- | ---------------------------------------------------------------------------------------- | +| `"base"` (default) | USDC on Base | `private_key=` (EVM hex key), `BLOCKRUN_WALLET_KEY`, `~/.blockrun/.session` | +| `"solana"` | USDC on Solana | `private_key=` (base58 secret key), `SOLANA_WALLET_KEY`, a Solana wallet file on disk | + +The key never leaves your machine; it only signs the payment for each call. +`llm.wallet_address` prints the address to fund. + +```python +from llama_index.llms.blockrun import BlockRun + +llm = BlockRun(model="openai/gpt-5.5") # Base +llm = BlockRun(model="openai/gpt-5.5", chain="solana") # Solana + +print(llm.wallet_address) +``` + +## Usage + +```python +from llama_index.core.llms import ChatMessage +from llama_index.llms.blockrun import BlockRun + +llm = BlockRun(model="anthropic/claude-sonnet-4.6") + +print(llm.complete("Paris is the capital of")) + +response = llm.chat( + [ + ChatMessage(role="system", content="Answer in one sentence."), + ChatMessage(role="user", content="What is x402?"), + ] +) +print(response.message.content) +print(response.additional_kwargs.get("cost_usd")) # USD charged, where the SDK reports it (Base) + +for chunk in llm.stream_complete("Write a haiku about USDC."): + print(chunk.delta, end="", flush=True) +``` + +Async works the same way: `achat`, `acomplete`, `astream_chat`, `astream_complete`. + +### As the default LLM + +```python +from llama_index.core import Settings + +Settings.llm = BlockRun(model="openai/gpt-5.5") +``` + +### Tools and agents + +`BlockRun` is a `FunctionCallingLLM`, so tool calling, `predict_and_call`, +structured outputs and `FunctionAgent` work as they do with OpenAI: + +```python +from llama_index.core.agent.workflow import FunctionAgent + + +def multiply(a: float, b: float) -> float: + """Multiply two numbers.""" + return a * b + + +agent = FunctionAgent(tools=[multiply], llm=BlockRun(model="openai/gpt-5.5")) +print(await agent.run("What is 1234 * 4567?")) +``` + +## Spending + +Each call is quoted by the gateway before anything is signed. Cap what a single +call may cost, and a quote above it is refused before it is paid: + +```python +llm = BlockRun(model="openai/gpt-5.5", max_cost_per_call=0.05) +``` + +A refused quote raises `blockrun_llm.SpendLimitError`; an empty wallet raises +`blockrun_llm.PaymentError`. Prices per model are listed at +https://blockrun.ai/api/v1/models. + +## Options + +| Parameter | Default | Notes | +| --------------------------- | ---------------- | ---------------------------------------------------------------------------------------- | +| `model` | required | Any chat model id from https://blockrun.ai/api/v1/models | +| `chain` | `"base"` | `"base"` or `"solana"` | +| `private_key` | from environment | Never serialized with the LLM | +| `temperature` | unset | Unset lets each model use its own default | +| `max_tokens` | unset | The SDK sends 1024 when unset | +| `context_window` | from catalog | Read once from `/v1/models`; set it to skip that request | +| `max_cost_per_call` | unset | USD ceiling per call | +| `api_url` | chain's mainnet | `https://testnet.blockrun.ai/api` for Base Sepolia | +| `timeout` | 600 s | Reasoning models can think for minutes | +| `additional_kwargs` | `{}` | Sent with every request, e.g. `top_p`, `stop`, `response_format` | +| `is_function_calling_model` | `True` | Set `False` for a model without tool support | + +## Development + +```bash +pip install -e "../..[dev]" -e ".[dev]" +pytest +``` diff --git a/integrations/llama-index-llms-blockrun/llama_index/llms/blockrun/__init__.py b/integrations/llama-index-llms-blockrun/llama_index/llms/blockrun/__init__.py new file mode 100644 index 0000000..3036ed9 --- /dev/null +++ b/integrations/llama-index-llms-blockrun/llama_index/llms/blockrun/__init__.py @@ -0,0 +1,3 @@ +from llama_index.llms.blockrun.base import BlockRun + +__all__ = ["BlockRun"] diff --git a/integrations/llama-index-llms-blockrun/llama_index/llms/blockrun/base.py b/integrations/llama-index-llms-blockrun/llama_index/llms/blockrun/base.py new file mode 100644 index 0000000..36db0fe --- /dev/null +++ b/integrations/llama-index-llms-blockrun/llama_index/llms/blockrun/base.py @@ -0,0 +1,631 @@ +"""BlockRun LLM for LlamaIndex. + +BlockRun is an OpenAI-compatible gateway where every call pays for itself in +USDC over x402. There is no API key: the gateway answers an unpaid request with +HTTP 402 and a price, the client signs a payment for exactly that amount with a +local wallet, and the request is retried with the signature attached. + +That signing step is why this is a package rather than an ``OpenAILike`` +snippet. ``OpenAILike(api_base=..., api_key="fake")`` reaches the gateway and +then stops at the first 402, because nothing in the OpenAI client knows how to +pay. Payment here is delegated to the ``blockrun-llm`` SDK, which owns the +signing, the per-chain differences, and the retry rules that keep a request +from being paid twice. +""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import AsyncIterator, Iterable, Sequence +from typing import TYPE_CHECKING, Any, Literal + +from llama_index.core.base.llms.generic_utils import ( + achat_to_completion_decorator, + astream_chat_to_completion_decorator, + chat_to_completion_decorator, + stream_chat_to_completion_decorator, +) +from llama_index.core.base.llms.types import ( + ChatMessage, + ChatResponse, + ChatResponseAsyncGen, + ChatResponseGen, + CompletionResponse, + CompletionResponseAsyncGen, + CompletionResponseGen, + ContentBlock, + LLMMetadata, + MessageRole, + TextBlock, + ThinkingBlock, + ToolCallBlock, +) +from llama_index.core.bridge.pydantic import Field, PrivateAttr +from llama_index.core.callbacks import CallbackManager +from llama_index.core.constants import DEFAULT_CONTEXT_WINDOW +from llama_index.core.llms.callbacks import llm_chat_callback, llm_completion_callback +from llama_index.core.llms.function_calling import FunctionCallingLLM +from llama_index.core.llms.llm import ToolSelection +from llama_index.core.llms.utils import parse_partial_json +from llama_index.core.types import BaseOutputParser, PydanticProgramMode +from llama_index.llms.openai.utils import resolve_tool_choice, to_openai_message_dicts + +if TYPE_CHECKING: + from llama_index.core.tools.types import BaseTool + +logger = logging.getLogger(__name__) + +Chain = Literal["base", "solana"] + +# What the SDK sends when max_tokens is left unset, on both chains. Reported +# as num_output so LlamaIndex budgets the prompt against what will be generated. +DEFAULT_MAX_TOKENS = 1024 + +# The parameters both chains' clients accept by name. The Base client also +# forwards any other keyword into the request body; the Solana client takes +# only these, so an unknown key there would be a TypeError from deep inside +# the SDK. It is refused up front instead, naming the key. +_SHARED_REQUEST_PARAMS = frozenset( + { + "max_tokens", + "temperature", + "top_p", + "search", + "search_parameters", + "tools", + "tool_choice", + "response_format", + "stop", + "fallback_models", + } +) + + +class BlockRun(FunctionCallingLLM): + """BlockRun LLM: one OpenAI-compatible gateway, paid per call in USDC over x402. + + Every BlockRun chat model is reachable through this class, addressed by its + catalog id (``openai/gpt-5.5``, ``anthropic/claude-sonnet-4.6``, + ``deepseek/deepseek-chat`` ...). The full list is at + https://blockrun.ai/api/v1/models. + + There is no API key. Requests are paid from a wallet on the chain you pick: + + * ``chain="base"`` (default): USDC on Base. The wallet is an EVM private + key, read from ``private_key``, then ``BLOCKRUN_WALLET_KEY``, then the + wallet file ``~/.blockrun/.session``. + * ``chain="solana"``: USDC on Solana. The wallet is a base58 secret key, + read from ``private_key``, then ``SOLANA_WALLET_KEY``, then a Solana + wallet file on disk. Needs ``pip install "llama-index-llms-blockrun[solana]"``. + + The key is used only to sign payments locally; it is never sent anywhere. + + Install: + ``pip install llama-index-llms-blockrun`` + + Examples: + .. code-block:: python + + from llama_index.llms.blockrun import BlockRun + + # Pays in USDC on Base (the default chain). + llm = BlockRun(model="openai/gpt-5.5") + print(llm.complete("Paris is the capital of")) + + # Pays in USDC on Solana. + llm = BlockRun(model="openai/gpt-5.5", chain="solana") + + # Refuse any single call quoted above $0.05 before signing it. + llm = BlockRun(model="openai/gpt-5.5", max_cost_per_call=0.05) + + """ + + model: str = Field(description="BlockRun model id, e.g. 'openai/gpt-5.5'.") + chain: Chain = Field( + default="base", + description="Chain the wallet pays from: 'base' (USDC on Base) or 'solana' (USDC on Solana).", + ) + temperature: float | None = Field( + default=None, + description=( + "Sampling temperature. Unset sends none, so each model uses its own " + "default; several reasoning models accept no other value." + ), + ge=0.0, + le=2.0, + ) + max_tokens: int | None = Field( + default=None, + description="Maximum tokens to generate. Unset sends the SDK default of 1024.", + gt=0, + ) + context_window: int | None = Field( + default=None, + description=( + "Context window in tokens. Unset reads the model's own window from the " + "gateway's /v1/models catalog the first time it is needed." + ), + gt=0, + ) + is_function_calling_model: bool = Field( + default=True, + description="Whether the model supports OpenAI-style tool calling.", + ) + max_cost_per_call: float | None = Field( + default=None, + description=( + "Refuse to sign any single payment quoted above this many USD. The call " + "fails before anything is sent or paid. Unset signs every quote." + ), + gt=0, + ) + api_url: str | None = Field( + default=None, + description=( + "Gateway URL override, e.g. https://testnet.blockrun.ai/api for Base " + "Sepolia. Unset uses the chain's mainnet gateway." + ), + ) + timeout: float | None = Field( + default=None, + description="HTTP timeout in seconds. Unset uses the SDK default (600).", + gt=0, + ) + additional_kwargs: dict[str, Any] = Field( + default_factory=dict, + description="Extra request parameters sent with every call, e.g. top_p or stop.", + ) + + # Never a field: fields are serialized by to_dict()/to_json() and by any + # callback that logs the LLM, and this is a wallet key. + _private_key: str | None = PrivateAttr(default=None) + _client: Any = PrivateAttr(default=None) + _aclient: Any = PrivateAttr(default=None) + _aclient_loop: asyncio.AbstractEventLoop | None = PrivateAttr(default=None) + _catalog_entry: dict[str, Any] | None = PrivateAttr(default=None) + _catalog_loaded: bool = PrivateAttr(default=False) + + def __init__( + self, + model: str, + chain: Chain = "base", + private_key: str | None = None, + temperature: float | None = None, + max_tokens: int | None = None, + context_window: int | None = None, + is_function_calling_model: bool = True, + max_cost_per_call: float | None = None, + api_url: str | None = None, + timeout: float | None = None, + additional_kwargs: dict[str, Any] | None = None, + callback_manager: CallbackManager | None = None, + system_prompt: str | None = None, + output_parser: BaseOutputParser | None = None, + pydantic_program_mode: PydanticProgramMode = PydanticProgramMode.DEFAULT, + **kwargs: Any, + ) -> None: + if chain not in ("base", "solana"): + raise ValueError(f"chain must be 'base' or 'solana', got {chain!r}") + super().__init__( + model=model, + chain=chain, + temperature=temperature, + max_tokens=max_tokens, + context_window=context_window, + is_function_calling_model=is_function_calling_model, + max_cost_per_call=max_cost_per_call, + api_url=api_url, + timeout=timeout, + additional_kwargs=additional_kwargs or {}, + callback_manager=callback_manager, + system_prompt=system_prompt, + output_parser=output_parser, + pydantic_program_mode=pydantic_program_mode, + **kwargs, + ) + self._private_key = private_key + if chain == "solana": + # Checked before the request instead of inside it: on Solana an + # unknown parameter is a TypeError from the SDK, not a 400. + _check_solana_params(self.additional_kwargs) + + @classmethod + def class_name(cls) -> str: + return "BlockRun_LLM" + + # ------------------------------------------------------------------ + # Clients + # ------------------------------------------------------------------ + + def _client_kwargs(self) -> dict[str, Any]: + # Only what the caller set. The Solana client's api_url default is a + # string, not None, so passing None through would replace its gateway. + kwargs: dict[str, Any] = {"private_key": self._private_key} + if self.api_url is not None: + kwargs["api_url"] = self.api_url + if self.timeout is not None: + kwargs["timeout"] = self.timeout + if self.max_cost_per_call is not None: + kwargs["max_cost_per_call"] = self.max_cost_per_call + return kwargs + + def _get_client(self) -> Any: + if self._client is None: + if self.chain == "solana": + from blockrun_llm import SolanaLLMClient as client_cls + else: + from blockrun_llm import LLMClient as client_cls + self._client = client_cls(**self._client_kwargs()) + return self._client + + def _get_aclient(self) -> Any: + # An httpx.AsyncClient is bound to the loop it first ran on. Each + # asyncio.run() is a new loop, so a client cached across them fails + # with "Event loop is closed". Keep one per running loop. + loop = asyncio.get_running_loop() + if self._aclient is None or self._aclient_loop is not loop: + if self.chain == "solana": + from blockrun_llm import AsyncSolanaLLMClient as aclient_cls + else: + from blockrun_llm import AsyncLLMClient as aclient_cls + self._aclient = aclient_cls(**self._client_kwargs()) + self._aclient_loop = loop + return self._aclient + + @property + def wallet_address(self) -> str: + """Address of the wallet paying for calls (EVM on Base, base58 on Solana).""" + return str(self._get_client().get_wallet_address()) + + # ------------------------------------------------------------------ + # Metadata + # ------------------------------------------------------------------ + + def _lookup_catalog_entry(self) -> dict[str, Any] | None: + # One free GET /v1/models per instance. A failure is cached too, so an + # unreachable catalog costs one warning rather than a request every + # time LlamaIndex reads metadata; set context_window to skip this. + if not self._catalog_loaded: + self._catalog_loaded = True + try: + models = self._get_client().list_models() + except Exception as exc: + logger.warning( + "BlockRun: could not read the model catalog (%s); using a " + "context window of %d. Pass context_window= to set it.", + exc, + DEFAULT_CONTEXT_WINDOW, + ) + models = [] + self._catalog_entry = next((m for m in models if m.get("id") == self.model), None) + return self._catalog_entry + + @property + def metadata(self) -> LLMMetadata: + context_window = self.context_window + if context_window is None: + entry = self._lookup_catalog_entry() + context_window = (entry or {}).get("context_window") or DEFAULT_CONTEXT_WINDOW + return LLMMetadata( + context_window=context_window, + num_output=self.max_tokens or DEFAULT_MAX_TOKENS, + is_chat_model=True, + is_function_calling_model=self.is_function_calling_model, + model_name=self.model, + ) + + # ------------------------------------------------------------------ + # Request building + # ------------------------------------------------------------------ + + def _request_kwargs(self, **kwargs: Any) -> dict[str, Any]: + merged = { + "temperature": self.temperature, + "max_tokens": self.max_tokens, + **self.additional_kwargs, + **kwargs, + } + # None means "not set" throughout; the SDK applies its own defaults. + merged = {k: v for k, v in merged.items() if v is not None} + if self.chain == "solana": + _check_solana_params(merged) + return merged + + def _message_dicts(self, messages: Sequence[ChatMessage]) -> list[dict[str, Any]]: + # A reasoning trace is output, never input: the chat schema has no + # field for it. Dropped here rather than left to the serializer: + # llama-index-llms-openai 0.8 skips a ThinkingBlock, but 0.6.5, inside + # the supported range, raises on it. + cleaned = [ + ( + message.model_copy( + update={ + "blocks": [b for b in message.blocks if not isinstance(b, ThinkingBlock)] + } + ) + if any(isinstance(b, ThinkingBlock) for b in message.blocks) + else message + ) + for message in messages + ] + return list(to_openai_message_dicts(cleaned)) # type: ignore[arg-type] + + # ------------------------------------------------------------------ + # Chat + # ------------------------------------------------------------------ + + @llm_chat_callback() + def chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse: + response = self._get_client().chat_completion( + self.model, self._message_dicts(messages), **self._request_kwargs(**kwargs) + ) + return _to_chat_response(response) + + @llm_chat_callback() + def stream_chat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponseGen: + chunks = self._get_client().chat_completion_stream( + self.model, self._message_dicts(messages), **self._request_kwargs(**kwargs) + ) + + def gen() -> ChatResponseGen: + accumulator = _StreamAccumulator() + for chunk in chunks: + yield accumulator.step(chunk) + + return gen() + + @llm_chat_callback() + async def achat(self, messages: Sequence[ChatMessage], **kwargs: Any) -> ChatResponse: + response = await self._get_aclient().chat_completion( + self.model, self._message_dicts(messages), **self._request_kwargs(**kwargs) + ) + return _to_chat_response(response) + + @llm_chat_callback() + async def astream_chat( + self, messages: Sequence[ChatMessage], **kwargs: Any + ) -> ChatResponseAsyncGen: + chunks: AsyncIterator[Any] = self._get_aclient().chat_completion_stream( + self.model, self._message_dicts(messages), **self._request_kwargs(**kwargs) + ) + + async def gen() -> ChatResponseAsyncGen: + accumulator = _StreamAccumulator() + async for chunk in chunks: + yield accumulator.step(chunk) + + return gen() + + # ------------------------------------------------------------------ + # Completion (a single user turn through the chat endpoint) + # ------------------------------------------------------------------ + + @llm_completion_callback() + def complete(self, prompt: str, formatted: bool = False, **kwargs: Any) -> CompletionResponse: + return chat_to_completion_decorator(self.chat)(prompt, **kwargs) + + @llm_completion_callback() + def stream_complete( + self, prompt: str, formatted: bool = False, **kwargs: Any + ) -> CompletionResponseGen: + return stream_chat_to_completion_decorator(self.stream_chat)(prompt, **kwargs) + + @llm_completion_callback() + async def acomplete( + self, prompt: str, formatted: bool = False, **kwargs: Any + ) -> CompletionResponse: + return await achat_to_completion_decorator(self.achat)(prompt, **kwargs) + + @llm_completion_callback() + async def astream_complete( + self, prompt: str, formatted: bool = False, **kwargs: Any + ) -> CompletionResponseAsyncGen: + return await astream_chat_to_completion_decorator(self.astream_chat)(prompt, **kwargs) + + # ------------------------------------------------------------------ + # Tool calling + # ------------------------------------------------------------------ + + def _prepare_chat_with_tools( + self, + tools: Sequence[BaseTool], + user_msg: str | ChatMessage | None = None, + chat_history: list[ChatMessage] | None = None, + verbose: bool = False, + allow_parallel_tool_calls: bool = False, + tool_required: bool = False, + tool_choice: str | dict[str, Any] | None = None, + **kwargs: Any, + ) -> dict[str, Any]: + tool_specs = [tool.metadata.to_openai_tool(skip_length_check=True) for tool in tools] + + # A copy: appending to the caller's chat_history would leave the user + # turn in their memory a second time when they record it themselves. + messages = list(chat_history or []) + if isinstance(user_msg, str): + user_msg = ChatMessage(role=MessageRole.USER, content=user_msg) + if user_msg is not None: + messages.append(user_msg) + + # parallel_tool_calls is not sent: the Solana client does not accept + # it, and not every upstream behind the gateway honors it. A single + # call is enforced on the response instead, in + # _validate_chat_with_tools_response. + return { + "messages": messages, + "tools": tool_specs or None, + "tool_choice": ( + resolve_tool_choice(tool_choice, tool_required) if tool_specs else None + ), + **kwargs, + } + + def _validate_chat_with_tools_response( + self, + response: ChatResponse, + tools: Sequence[BaseTool], + allow_parallel_tool_calls: bool = False, + **kwargs: Any, + ) -> ChatResponse: + if not allow_parallel_tool_calls: + _force_single_tool_call(response) + return response + + def get_tool_calls_from_response( + self, + response: ChatResponse, + error_on_no_tool_call: bool = True, + **kwargs: Any, + ) -> list[ToolSelection]: + tool_calls = [b for b in response.message.blocks if isinstance(b, ToolCallBlock)] + if not tool_calls: + if error_on_no_tool_call: + raise ValueError("Expected at least one tool call, but got 0 tool calls.") + return [] + + selections = [] + for call in tool_calls: + arguments = call.tool_kwargs + if isinstance(arguments, str): + # Partial JSON too: a streamed call can be read mid-stream. + try: + arguments = parse_partial_json(arguments) if arguments.strip() else {} + except (ValueError, TypeError): + arguments = {} + selections.append( + ToolSelection( + tool_id=call.tool_call_id or "", + tool_name=call.tool_name, + tool_kwargs=arguments if isinstance(arguments, dict) else {}, + ) + ) + return selections + + +# ---------------------------------------------------------------------- +# Helpers +# ---------------------------------------------------------------------- + + +def _check_solana_params(params: Iterable[str]) -> None: + unsupported = sorted(set(params) - _SHARED_REQUEST_PARAMS) + if unsupported: + raise ValueError( + f"chain='solana' does not support the request parameter(s) {unsupported}. " + f"Supported: {sorted(_SHARED_REQUEST_PARAMS)}." + ) + + +def _response_kwargs( + usage: Any, cost_usd: float | None = None, citations: list[str] | None = None +) -> dict[str, Any]: + out: dict[str, Any] = {} + if usage is not None: + out["prompt_tokens"] = usage.prompt_tokens + out["completion_tokens"] = usage.completion_tokens + out["total_tokens"] = usage.total_tokens + if cost_usd is not None: + out["cost_usd"] = cost_usd + if citations: + out["citations"] = citations + return out + + +def _to_chat_response(response: Any) -> ChatResponse: + message = response.choices[0].message + blocks: list[ContentBlock] = [] + reasoning = message.reasoning_content or message.thinking + if reasoning: + blocks.append(ThinkingBlock(content=reasoning)) + blocks.append(TextBlock(text=message.content or "")) + for call in message.tool_calls or []: + blocks.append( + ToolCallBlock( + tool_call_id=call.id, + tool_name=call.function.name, + tool_kwargs=call.function.arguments or {}, + ) + ) + return ChatResponse( + message=ChatMessage(role=MessageRole.ASSISTANT, blocks=blocks), + raw=response, + additional_kwargs=_response_kwargs( + response.usage, getattr(response, "cost_usd", None), response.citations + ), + ) + + +def _force_single_tool_call(response: ChatResponse) -> None: + tool_calls = [b for b in response.message.blocks if isinstance(b, ToolCallBlock)] + if len(tool_calls) > 1: + response.message.blocks = [ + b for b in response.message.blocks if not isinstance(b, ToolCallBlock) + ] + [tool_calls[0]] + + +class _StreamAccumulator: + """Folds SSE chunks into the running message each streamed ChatResponse carries.""" + + def __init__(self) -> None: + self.content = "" + self.reasoning = "" + # Keyed by the delta's tool-call index; dicts keep first-seen order. + self.tool_calls: dict[int, dict[str, str]] = {} + + def step(self, chunk: Any) -> ChatResponse: + delta = chunk.choices[0].delta if chunk.choices else None + content_delta = (delta.content if delta else None) or "" + reasoning_delta = ((delta.reasoning_content or delta.thinking) if delta else None) or "" + self.content += content_delta + self.reasoning += reasoning_delta + for call_delta in (delta.tool_calls if delta else None) or []: + self._merge_tool_call(call_delta) + + blocks: list[ContentBlock] = [] + if self.reasoning: + blocks.append(ThinkingBlock(content=self.reasoning)) + blocks.append(TextBlock(text=self.content)) + for call in self.tool_calls.values(): + blocks.append( + ToolCallBlock( + tool_call_id=call["id"] or None, + tool_name=call["name"], + tool_kwargs=call["arguments"] or {}, + ) + ) + + # The Base SDK attaches the real x402 charge to every chunk; read it + # with getattr because it is an extra attribute, not a declared field. + additional = _response_kwargs( + chunk.usage, getattr(chunk, "cost_usd", None), chunk.citations + ) + if reasoning_delta: + additional["thinking_delta"] = reasoning_delta + return ChatResponse( + message=ChatMessage(role=MessageRole.ASSISTANT, blocks=blocks), + delta=content_delta, + raw=chunk, + additional_kwargs=additional, + ) + + def _merge_tool_call(self, call_delta: Any) -> None: + if call_delta.index is not None: + index = call_delta.index + elif self.tool_calls and ( + not call_delta.id or call_delta.id == self.tool_calls[self._last_index()]["id"] + ): + # No index: a delta without a new id continues the open call. + index = self._last_index() + else: + index = max(self.tool_calls) + 1 if self.tool_calls else 0 + + call = self.tool_calls.setdefault(index, {"id": "", "name": "", "arguments": ""}) + if call_delta.id: + call["id"] = call_delta.id + if call_delta.function is not None: + call["name"] += call_delta.function.name or "" + call["arguments"] += call_delta.function.arguments or "" + + def _last_index(self) -> int: + return next(reversed(self.tool_calls)) diff --git a/integrations/llama-index-llms-blockrun/llama_index/llms/blockrun/py.typed b/integrations/llama-index-llms-blockrun/llama_index/llms/blockrun/py.typed new file mode 100644 index 0000000..e69de29 diff --git a/integrations/llama-index-llms-blockrun/pyproject.toml b/integrations/llama-index-llms-blockrun/pyproject.toml new file mode 100644 index 0000000..d405a91 --- /dev/null +++ b/integrations/llama-index-llms-blockrun/pyproject.toml @@ -0,0 +1,45 @@ +[build-system] +requires = ["hatchling"] +build-backend = "hatchling.build" + +[project] +name = "llama-index-llms-blockrun" +version = "0.1.0" +description = "LlamaIndex LLM for BlockRun: one OpenAI-compatible gateway, paid per call in USDC over x402 on Base or Solana. No API key." +authors = [{ name = "BlockRun", email = "hello@blockrun.ai" }] +requires-python = ">=3.10,<4.0" +readme = "README.md" +license = "MIT" +keywords = ["llama-index", "llm", "x402", "usdc", "base", "solana", "blockrun"] +dependencies = [ + "llama-index-core>=0.14.5,<0.15", + # Only for to_openai_message_dicts / resolve_tool_choice: the gateway speaks + # the OpenAI chat schema, and that is LlamaIndex's own serializer for it. + "llama-index-llms-openai>=0.6.5,<0.9", + "blockrun-llm>=1.17.0,<2", +] + +[project.optional-dependencies] +solana = ["blockrun-llm[solana]>=1.17.0,<2"] +dev = ["pytest>=7.0.0", "pytest-asyncio>=0.21.0"] + +[project.urls] +Homepage = "https://blockrun.ai" +Repository = "https://github.com/BlockRunAI/blockrun-llm/tree/main/integrations/llama-index-llms-blockrun" + +[tool.hatch.build.targets.sdist] +include = ["llama_index/", "tests/", "README.md", "LICENSE"] + +[tool.hatch.build.targets.wheel] +include = ["llama_index/"] + +[tool.llamahub] +contains_example = false +import_path = "llama_index.llms.blockrun" + +[tool.llamahub.class_authors] +BlockRun = "BlockRunAI" + +[tool.pytest.ini_options] +asyncio_mode = "auto" +testpaths = ["tests"] diff --git a/integrations/llama-index-llms-blockrun/tests/test_llms_blockrun.py b/integrations/llama-index-llms-blockrun/tests/test_llms_blockrun.py new file mode 100644 index 0000000..4c3b00d --- /dev/null +++ b/integrations/llama-index-llms-blockrun/tests/test_llms_blockrun.py @@ -0,0 +1,634 @@ +"""Unit tests for the BlockRun LlamaIndex LLM. + +No network and no wallet: the SDK clients are replaced with fakes that record +what they were asked and answer with real blockrun_llm response types, so the +conversion code is exercised against the shapes the SDK actually returns. +""" + +from __future__ import annotations + +import asyncio +import json +from typing import Any + +import pytest +from llama_index.core.base.llms.types import ( + ChatMessage, + MessageRole, + ThinkingBlock, + ToolCallBlock, +) +from llama_index.core.constants import DEFAULT_CONTEXT_WINDOW +from llama_index.core.llms.function_calling import FunctionCallingLLM +from llama_index.core.tools import FunctionTool +from llama_index.llms.blockrun import BlockRun + +import blockrun_llm +from blockrun_llm.types import ChatCompletionChunk +from blockrun_llm.types import ChatResponse as SdkChatResponse + +MODEL = "openai/gpt-5.5" + + +# ---------------------------------------------------------------------- +# Fakes +# ---------------------------------------------------------------------- + + +def sdk_response( + content: str | None = "Hello!", + tool_calls: list[dict[str, Any]] | None = None, + reasoning: str | None = None, + cost_usd: float | None = 0.0021, + citations: list[str] | None = None, +) -> SdkChatResponse: + message: dict[str, Any] = {"role": "assistant", "content": content} + if tool_calls is not None: + message["tool_calls"] = tool_calls + if reasoning is not None: + message["reasoning_content"] = reasoning + return SdkChatResponse.model_validate( + { + "id": "chatcmpl-1", + "created": 1, + "model": MODEL, + "choices": [{"index": 0, "message": message, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 11, "completion_tokens": 7, "total_tokens": 18}, + "citations": citations, + "cost_usd": cost_usd, + } + ) + + +def chunk( + content: str | None = None, + tool_calls: list[dict[str, Any]] | None = None, + reasoning: str | None = None, + usage: dict[str, int] | None = None, + no_choices: bool = False, +) -> ChatCompletionChunk: + delta: dict[str, Any] = {} + if content is not None: + delta["content"] = content + if tool_calls is not None: + delta["tool_calls"] = tool_calls + if reasoning is not None: + delta["reasoning_content"] = reasoning + return ChatCompletionChunk.model_validate( + { + "id": "chatcmpl-1", + "created": 1, + "model": MODEL, + "choices": [] if no_choices else [{"index": 0, "delta": delta}], + "usage": usage, + } + ) + + +def tool_call(call_id: str, name: str, arguments: str) -> dict[str, Any]: + return {"id": call_id, "type": "function", "function": {"name": name, "arguments": arguments}} + + +class FakeClient: + def __init__( + self, + response: SdkChatResponse | None = None, + chunks: list[ChatCompletionChunk] | None = None, + models: list[dict[str, Any]] | None = None, + models_error: Exception | None = None, + ) -> None: + self.response = response or sdk_response() + self.chunks = chunks or [] + self.models = models or [] + self.models_error = models_error + self.calls: list[dict[str, Any]] = [] + self.list_models_calls = 0 + + def chat_completion(self, model: str, messages: list[dict[str, Any]], **kwargs: Any) -> Any: + self.calls.append({"model": model, "messages": messages, **kwargs}) + return self.response + + def chat_completion_stream( + self, model: str, messages: list[dict[str, Any]], **kwargs: Any + ) -> Any: + self.calls.append({"model": model, "messages": messages, "stream": True, **kwargs}) + yield from self.chunks + + def list_models(self) -> list[dict[str, Any]]: + self.list_models_calls += 1 + if self.models_error is not None: + raise self.models_error + return self.models + + def get_wallet_address(self) -> str: + return "0xabc" + + +class FakeAsyncClient(FakeClient): + async def chat_completion( # type: ignore[override] + self, model: str, messages: list[dict[str, Any]], **kwargs: Any + ) -> Any: + self.calls.append({"model": model, "messages": messages, **kwargs}) + return self.response + + async def chat_completion_stream( # type: ignore[override] + self, model: str, messages: list[dict[str, Any]], **kwargs: Any + ) -> Any: + self.calls.append({"model": model, "messages": messages, "stream": True, **kwargs}) + for c in self.chunks: + yield c + + +def make_llm(client: FakeClient | None = None, **kwargs: Any) -> tuple[BlockRun, FakeClient]: + llm = BlockRun(model=MODEL, **kwargs) + client = client or FakeClient() + llm._client = client + return llm, client + + +def make_async_llm(client: FakeAsyncClient, **kwargs: Any) -> BlockRun: + llm = BlockRun(model=MODEL, **kwargs) + # _get_aclient builds a new client whenever the running loop changes, so + # pin the fake to whichever loop the test runs on. + llm._get_aclient = lambda: client # type: ignore[method-assign] + return llm + + +# ---------------------------------------------------------------------- +# Construction +# ---------------------------------------------------------------------- + + +class TestConstruction: + def test_is_a_function_calling_llm(self) -> None: + assert issubclass(BlockRun, FunctionCallingLLM) + assert BlockRun.class_name() == "BlockRun_LLM" + + def test_defaults_to_base(self) -> None: + assert BlockRun(model=MODEL).chain == "base" + + def test_rejects_unknown_chain(self) -> None: + with pytest.raises(ValueError, match="'base' or 'solana'"): + BlockRun(model=MODEL, chain="ethereum") # type: ignore[arg-type] + + def test_private_key_is_never_serialized(self) -> None: + secret = "0x" + "ab" * 32 + llm = BlockRun(model=MODEL, private_key=secret) + assert secret not in json.dumps(llm.to_dict()) + assert secret not in llm.model_dump_json() + assert secret not in repr(llm) + + def test_solana_refuses_unsupported_additional_kwargs_at_construction(self) -> None: + with pytest.raises(ValueError, match="presence_penalty"): + BlockRun(model=MODEL, chain="solana", additional_kwargs={"presence_penalty": 1}) + + +class TestClientSelection: + @pytest.fixture + def recorded(self, monkeypatch: pytest.MonkeyPatch) -> dict[str, Any]: + seen: dict[str, Any] = {} + + def recorder(name: str) -> type: + class Recorder: + def __init__(self, **kwargs: Any) -> None: + seen[name] = kwargs + + return Recorder + + for name in ("LLMClient", "AsyncLLMClient", "SolanaLLMClient", "AsyncSolanaLLMClient"): + monkeypatch.setattr(blockrun_llm, name, recorder(name)) + return seen + + def test_base_uses_the_evm_client(self, recorded: dict[str, Any]) -> None: + BlockRun(model=MODEL, private_key="k")._get_client() + assert recorded == {"LLMClient": {"private_key": "k"}} + + def test_solana_uses_the_solana_client(self, recorded: dict[str, Any]) -> None: + BlockRun(model=MODEL, chain="solana", private_key="k")._get_client() + assert recorded == {"SolanaLLMClient": {"private_key": "k"}} + + def test_unset_options_are_not_passed(self, recorded: dict[str, Any]) -> None: + # SolanaLLMClient's api_url default is a URL string; a None passed + # through would replace the gateway it talks to. + BlockRun(model=MODEL, chain="solana")._get_client() + assert recorded["SolanaLLMClient"] == {"private_key": None} + + def test_set_options_are_passed(self, recorded: dict[str, Any]) -> None: + BlockRun( + model=MODEL, + api_url="https://testnet.blockrun.ai/api", + timeout=30, + max_cost_per_call=0.05, + )._get_client() + assert recorded["LLMClient"] == { + "private_key": None, + "api_url": "https://testnet.blockrun.ai/api", + "timeout": 30, + "max_cost_per_call": 0.05, + } + + def test_client_is_reused(self, recorded: dict[str, Any]) -> None: + llm = BlockRun(model=MODEL) + assert llm._get_client() is llm._get_client() + + @pytest.mark.parametrize( + ("chain", "name"), [("base", "AsyncLLMClient"), ("solana", "AsyncSolanaLLMClient")] + ) + def test_async_client_per_event_loop( + self, recorded: dict[str, Any], chain: str, name: str + ) -> None: + llm = BlockRun(model=MODEL, chain=chain) # type: ignore[arg-type] + + async def get_twice() -> tuple[Any, Any]: + return llm._get_aclient(), llm._get_aclient() + + first_a, first_b = asyncio.run(get_twice()) + second, _ = asyncio.run(get_twice()) + assert name in recorded + assert first_a is first_b # reused within a loop + assert second is not first_a # rebuilt on a new loop + + +# ---------------------------------------------------------------------- +# Request parameters +# ---------------------------------------------------------------------- + + +class TestRequestParameters: + def test_unset_parameters_are_not_sent(self) -> None: + llm, client = make_llm() + llm.chat([ChatMessage(role="user", content="hi")]) + assert client.calls[0] == { + "model": MODEL, + "messages": [{"role": "user", "content": "hi"}], + } + + def test_set_parameters_and_additional_kwargs_are_sent(self) -> None: + llm, client = make_llm(temperature=0.2, max_tokens=50, additional_kwargs={"top_p": 0.9}) + llm.chat([ChatMessage(role="user", content="hi")], stop=["\n"]) + call = client.calls[0] + assert call["temperature"] == 0.2 + assert call["max_tokens"] == 50 + assert call["top_p"] == 0.9 + assert call["stop"] == ["\n"] + + def test_call_kwargs_override_instance_defaults(self) -> None: + llm, client = make_llm(temperature=0.2) + llm.chat([ChatMessage(role="user", content="hi")], temperature=0.7) + assert client.calls[0]["temperature"] == 0.7 + + def test_base_forwards_other_parameters(self) -> None: + llm, client = make_llm() + llm.chat([ChatMessage(role="user", content="hi")], reasoning_effort="low") + assert client.calls[0]["reasoning_effort"] == "low" + + def test_solana_refuses_other_parameters_before_sending(self) -> None: + llm, client = make_llm(chain="solana") + with pytest.raises(ValueError, match="reasoning_effort"): + llm.chat([ChatMessage(role="user", content="hi")], reasoning_effort="low") + assert client.calls == [] + + def test_system_prompt_is_sent_as_a_system_message(self) -> None: + llm, client = make_llm() + llm.chat( + [ + ChatMessage(role="system", content="Be terse."), + ChatMessage(role="user", content="hi"), + ] + ) + assert client.calls[0]["messages"][0] == {"role": "system", "content": "Be terse."} + + +# ---------------------------------------------------------------------- +# Responses +# ---------------------------------------------------------------------- + + +class TestChat: + def test_text_usage_and_cost(self) -> None: + llm, _ = make_llm(FakeClient(sdk_response("Paris.", citations=["https://x"]))) + response = llm.chat([ChatMessage(role="user", content="Capital of France?")]) + assert response.message.role == MessageRole.ASSISTANT + assert response.message.content == "Paris." + assert response.additional_kwargs == { + "prompt_tokens": 11, + "completion_tokens": 7, + "total_tokens": 18, + "cost_usd": 0.0021, + "citations": ["https://x"], + } + assert isinstance(response.raw, SdkChatResponse) + + def test_reasoning_becomes_a_thinking_block(self) -> None: + llm, _ = make_llm(FakeClient(sdk_response("4", reasoning="2+2 is 4"))) + response = llm.chat([ChatMessage(role="user", content="2+2?")]) + thinking = [b for b in response.message.blocks if isinstance(b, ThinkingBlock)] + assert [b.content for b in thinking] == ["2+2 is 4"] + assert response.message.content == "4" + + def test_complete(self) -> None: + llm, client = make_llm(FakeClient(sdk_response("Paris."))) + assert llm.complete("Capital of France?").text == "Paris." + assert client.calls[0]["messages"] == [{"role": "user", "content": "Capital of France?"}] + + def test_payment_errors_propagate(self) -> None: + class Refusing(FakeClient): + def chat_completion(self, *args: Any, **kwargs: Any) -> Any: + raise blockrun_llm.SpendLimitError( + "quote $0.40 exceeds max_cost_per_call $0.05", + quoted_usd=0.40, + limit_usd=0.05, + scope="call", + ) + + llm, _ = make_llm(Refusing()) + with pytest.raises(blockrun_llm.SpendLimitError): + llm.chat([ChatMessage(role="user", content="hi")]) + + +class TestAsync: + async def test_achat(self) -> None: + client = FakeAsyncClient(sdk_response("async!")) + llm = make_async_llm(client, temperature=0.3) + response = await llm.achat([ChatMessage(role="user", content="hi")]) + assert response.message.content == "async!" + assert client.calls[0]["temperature"] == 0.3 + + async def test_acomplete(self) -> None: + llm = make_async_llm(FakeAsyncClient(sdk_response("done"))) + assert (await llm.acomplete("go")).text == "done" + + async def test_astream_chat(self) -> None: + client = FakeAsyncClient(chunks=[chunk("Hel"), chunk("lo")]) + llm = make_async_llm(client) + stream = await llm.astream_chat([ChatMessage(role="user", content="hi")]) + deltas = [r.delta async for r in stream] + assert deltas == ["Hel", "lo"] + + async def test_astream_complete(self) -> None: + llm = make_async_llm(FakeAsyncClient(chunks=[chunk("a"), chunk("b")])) + stream = await llm.astream_complete("go") + texts = [r.text async for r in stream] + assert texts == ["a", "ab"] + + +# ---------------------------------------------------------------------- +# Streaming +# ---------------------------------------------------------------------- + + +class TestStreaming: + def test_content_accumulates_and_usage_arrives_on_the_last_chunk(self) -> None: + usage = {"prompt_tokens": 3, "completion_tokens": 2, "total_tokens": 5} + llm, client = make_llm( + FakeClient(chunks=[chunk("Hel"), chunk("lo"), chunk(usage=usage, no_choices=True)]) + ) + responses = list(llm.stream_chat([ChatMessage(role="user", content="hi")])) + assert [r.delta for r in responses] == ["Hel", "lo", ""] + assert [r.message.content for r in responses] == ["Hel", "Hello", "Hello"] + assert responses[-1].additional_kwargs == usage + assert client.calls[0]["stream"] is True + + def test_charge_attached_to_chunks_is_surfaced(self) -> None: + paid = chunk("hi") + paid.cost_usd = 0.0021 # type: ignore[attr-defined] # as the Base SDK does + llm, _ = make_llm(FakeClient(chunks=[paid])) + (response,) = list(llm.stream_chat([ChatMessage(role="user", content="hi")])) + assert response.additional_kwargs["cost_usd"] == 0.0021 + + def test_stream_complete(self) -> None: + llm, _ = make_llm(FakeClient(chunks=[chunk("a"), chunk("b")])) + assert [r.text for r in llm.stream_complete("go")] == ["a", "ab"] + + def test_reasoning_streams_into_a_thinking_block(self) -> None: + llm, _ = make_llm(FakeClient(chunks=[chunk(reasoning="think"), chunk("answer")])) + responses = list(llm.stream_chat([ChatMessage(role="user", content="hi")])) + assert responses[0].additional_kwargs["thinking_delta"] == "think" + thinking = [b for b in responses[-1].message.blocks if isinstance(b, ThinkingBlock)] + assert [b.content for b in thinking] == ["think"] + assert responses[-1].message.content == "answer" + + def test_tool_call_arguments_accumulate_across_chunks(self) -> None: + chunks = [ + chunk(tool_calls=[{"index": 0, "id": "call_1", "function": {"name": "add"}}]), + chunk(tool_calls=[{"index": 0, "function": {"arguments": '{"a": 1, '}}]), + chunk(tool_calls=[{"index": 0, "function": {"arguments": '"b": 2}'}}]), + ] + llm, _ = make_llm(FakeClient(chunks=chunks)) + last = list(llm.stream_chat([ChatMessage(role="user", content="1+2")]))[-1] + (selection,) = llm.get_tool_calls_from_response(last) + assert (selection.tool_id, selection.tool_name) == ("call_1", "add") + assert selection.tool_kwargs == {"a": 1, "b": 2} + + def test_parallel_tool_calls_stay_separate(self) -> None: + chunks = [ + chunk( + tool_calls=[ + {"index": 0, "id": "c0", "function": {"name": "f", "arguments": '{"x":'}}, + {"index": 1, "id": "c1", "function": {"name": "g", "arguments": '{"y":'}}, + ] + ), + chunk(tool_calls=[{"index": 1, "function": {"arguments": "2}"}}]), + chunk(tool_calls=[{"index": 0, "function": {"arguments": "1}"}}]), + ] + llm, _ = make_llm(FakeClient(chunks=chunks)) + last = list(llm.stream_chat([ChatMessage(role="user", content="go")]))[-1] + selections = llm.get_tool_calls_from_response(last) + assert [(s.tool_id, s.tool_name, s.tool_kwargs) for s in selections] == [ + ("c0", "f", {"x": 1}), + ("c1", "g", {"y": 2}), + ] + + def test_tool_call_deltas_without_an_index(self) -> None: + # Some upstreams omit `index`. A delta with no new id continues the + # open call; a new id starts the next one. + chunks = [ + chunk(tool_calls=[{"id": "a", "function": {"name": "f", "arguments": "{"}}]), + chunk(tool_calls=[{"function": {"arguments": "}"}}]), + chunk(tool_calls=[{"id": "b", "function": {"name": "g", "arguments": "{}"}}]), + ] + llm, _ = make_llm(FakeClient(chunks=chunks)) + last = list(llm.stream_chat([ChatMessage(role="user", content="go")]))[-1] + selections = llm.get_tool_calls_from_response(last) + assert [(s.tool_id, s.tool_name, s.tool_kwargs) for s in selections] == [ + ("a", "f", {}), + ("b", "g", {}), + ] + + def test_a_mid_stream_tool_call_reads_as_partial_json(self) -> None: + chunks = [chunk(tool_calls=[{"index": 0, "id": "c", "function": {"name": "f"}}])] + chunks.append(chunk(tool_calls=[{"index": 0, "function": {"arguments": '{"q": "x", "n'}}])) + llm, _ = make_llm(FakeClient(chunks=chunks)) + last = list(llm.stream_chat([ChatMessage(role="user", content="go")]))[-1] + (selection,) = llm.get_tool_calls_from_response(last) + assert selection.tool_kwargs == {"q": "x"} + + +# ---------------------------------------------------------------------- +# Tool calling +# ---------------------------------------------------------------------- + + +def add(a: int, b: int) -> int: + """Add two integers.""" + return a + b + + +def multiply(a: int, b: int) -> int: + """Multiply two integers.""" + return a * b + + +ADD = FunctionTool.from_defaults(fn=add) +MULTIPLY = FunctionTool.from_defaults(fn=multiply) + + +class TestToolCalling: + def test_chat_with_tools_sends_openai_tool_specs(self) -> None: + llm, client = make_llm( + FakeClient(sdk_response(None, tool_calls=[tool_call("c1", "add", '{"a":1,"b":2}')])) + ) + response = llm.chat_with_tools([ADD], user_msg="What is 1 + 2?") + call = client.calls[0] + assert call["tool_choice"] == "auto" + assert [t["function"]["name"] for t in call["tools"]] == ["add"] + assert call["tools"][0]["type"] == "function" + assert "parallel_tool_calls" not in call + (selection,) = llm.get_tool_calls_from_response(response) + assert (selection.tool_id, selection.tool_name) == ("c1", "add") + assert selection.tool_kwargs == {"a": 1, "b": 2} + + def test_tool_required_and_named_tool_choice(self) -> None: + llm, client = make_llm() + llm.chat_with_tools([ADD], user_msg="hi", tool_required=True) + llm.chat_with_tools([ADD, MULTIPLY], user_msg="hi", tool_choice="multiply") + assert client.calls[0]["tool_choice"] == "required" + assert client.calls[1]["tool_choice"] == { + "type": "function", + "function": {"name": "multiply"}, + } + + def test_chat_history_is_not_mutated(self) -> None: + llm, client = make_llm() + history = [ChatMessage(role="user", content="earlier")] + llm.chat_with_tools([ADD], user_msg="now", chat_history=history) + assert len(history) == 1 + assert [m["content"] for m in client.calls[0]["messages"]] == ["earlier", "now"] + + def test_single_tool_call_unless_parallel_allowed(self) -> None: + calls = [tool_call("c1", "add", "{}"), tool_call("c2", "multiply", "{}")] + llm, _ = make_llm(FakeClient(sdk_response(None, tool_calls=calls))) + single = llm.chat_with_tools([ADD, MULTIPLY], user_msg="go") + parallel = llm.chat_with_tools( + [ADD, MULTIPLY], user_msg="go", allow_parallel_tool_calls=True + ) + assert [s.tool_id for s in llm.get_tool_calls_from_response(single)] == ["c1"] + assert [s.tool_id for s in llm.get_tool_calls_from_response(parallel)] == ["c1", "c2"] + + def test_no_tool_call(self) -> None: + llm, _ = make_llm() + response = llm.chat([ChatMessage(role="user", content="hi")]) + with pytest.raises(ValueError, match="at least one tool call"): + llm.get_tool_calls_from_response(response) + assert llm.get_tool_calls_from_response(response, error_on_no_tool_call=False) == [] + + def test_tool_turn_round_trips_into_the_next_request(self) -> None: + # The assistant turn that called a tool, and the tool's result, must be + # sent back in the OpenAI shape or the upstream rejects the follow-up. + llm, client = make_llm( + FakeClient(sdk_response(None, tool_calls=[tool_call("c1", "add", '{"a":1,"b":2}')])) + ) + first = llm.chat_with_tools([ADD], user_msg="1+2?") + tool_result = ChatMessage( + role=MessageRole.TOOL, + content="3", + additional_kwargs={"tool_call_id": "c1", "name": "add"}, + ) + llm.chat([ChatMessage(role="user", content="1+2?"), first.message, tool_result]) + _, assistant, tool = client.calls[1]["messages"] + assert assistant["role"] == "assistant" + assert assistant["content"] is None + assert assistant["tool_calls"] == [ + { + "type": "function", + "function": {"name": "add", "arguments": '{"a":1,"b":2}'}, + "id": "c1", + } + ] + assert tool["role"] == "tool" + assert tool["tool_call_id"] == "c1" + assert tool["content"] == "3" + + def test_thinking_is_not_sent_back(self) -> None: + llm, client = make_llm(FakeClient(sdk_response("4", reasoning="secret chain"))) + first = llm.chat([ChatMessage(role="user", content="2+2?")]) + llm.chat([ChatMessage(role="user", content="2+2?"), first.message]) + assert "secret chain" not in json.dumps(client.calls[1]["messages"]) + + def test_tool_call_block_kwargs_are_kept_as_sent(self) -> None: + llm, _ = make_llm(FakeClient(sdk_response(None, tool_calls=[tool_call("c", "f", "{}")]))) + response = llm.chat([ChatMessage(role="user", content="go")]) + (block,) = [b for b in response.message.blocks if isinstance(b, ToolCallBlock)] + assert block.tool_kwargs == "{}" + + def test_malformed_arguments_become_empty_kwargs(self) -> None: + llm, _ = make_llm( + FakeClient(sdk_response(None, tool_calls=[tool_call("c", "f", "not json")])) + ) + response = llm.chat([ChatMessage(role="user", content="go")]) + (selection,) = llm.get_tool_calls_from_response(response) + assert selection.tool_kwargs == {} + + def test_predict_and_call_runs_the_tool(self) -> None: + llm, _ = make_llm( + FakeClient(sdk_response(None, tool_calls=[tool_call("c1", "add", '{"a":2,"b":3}')])) + ) + output = llm.predict_and_call([ADD], user_msg="2+3?") + assert output.response == "5" + + +# ---------------------------------------------------------------------- +# Metadata +# ---------------------------------------------------------------------- + + +class TestMetadata: + def test_context_window_comes_from_the_catalog_once(self) -> None: + client = FakeClient(models=[{"id": "other"}, {"id": MODEL, "context_window": 1_050_000}]) + llm, _ = make_llm(client) + assert llm.metadata.context_window == 1_050_000 + assert llm.metadata.context_window == 1_050_000 + assert client.list_models_calls == 1 + + def test_explicit_context_window_skips_the_catalog(self) -> None: + llm, client = make_llm(context_window=8000) + assert llm.metadata.context_window == 8000 + assert client.list_models_calls == 0 + + def test_unknown_model_falls_back(self) -> None: + llm, _ = make_llm(FakeClient(models=[{"id": "other", "context_window": 5}])) + assert llm.metadata.context_window == DEFAULT_CONTEXT_WINDOW + + def test_unreachable_catalog_falls_back_once(self, caplog: pytest.LogCaptureFixture) -> None: + client = FakeClient(models_error=RuntimeError("offline")) + llm, _ = make_llm(client) + assert llm.metadata.context_window == DEFAULT_CONTEXT_WINDOW + assert llm.metadata.context_window == DEFAULT_CONTEXT_WINDOW + assert client.list_models_calls == 1 + assert "context_window" in caplog.text + + def test_other_fields(self) -> None: + llm, _ = make_llm(context_window=8000, max_tokens=300, is_function_calling_model=False) + meta = llm.metadata + assert meta.model_name == MODEL + assert meta.num_output == 300 + assert meta.is_chat_model is True + assert meta.is_function_calling_model is False + + def test_num_output_defaults_to_the_sdk_default(self) -> None: + llm, _ = make_llm(context_window=8000) + assert llm.metadata.num_output == 1024 + + def test_wallet_address(self) -> None: + llm, _ = make_llm() + assert llm.wallet_address == "0xabc" diff --git a/pyproject.toml b/pyproject.toml index 9704fd0..cad3830 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -68,6 +68,8 @@ exclude = [ ".claude/", ".ruff_cache/", "*.png", + # Separate packages, published from their own pyproject.toml. + "integrations/", ] [tool.black]