mirror of
https://github.com/Routstr/routstr-core.git
synced 2026-07-30 23:36:15 +00:00
Compare commits
2
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Shroominic
|
1b7e0b7d22 | ||
|
Shroominic
|
1c5e709b0a |
@@ -14,7 +14,6 @@ UPSTREAM_API_KEY=your-upstream-api-key
|
||||
# HTTP_URL=https://api.mynode.com
|
||||
# ONION_URL=http://mynode.onion (auto fetched from compose)
|
||||
# RELAYS="wss://relay.damus.io,wss://relay.nostr.band,wss://eden.nostr.land,wss://relay.routstr.com"
|
||||
# ENABLE_ANALYTICS_SHARING=true
|
||||
# CASHU_MINTS="https://mint.minibits.cash/Bitcoin,https://mint.cubabitcoin.org,https://ecashmint.otrta.me"
|
||||
# RECEIVE_LN_ADDRESS=
|
||||
|
||||
|
||||
@@ -14,14 +14,6 @@ jobs:
|
||||
steps:
|
||||
- name: Checkout code
|
||||
uses: actions/checkout@v3
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Resolve git metadata
|
||||
id: gitmeta
|
||||
run: |
|
||||
echo "sha=$(git rev-parse --short=7 HEAD)" >> "$GITHUB_OUTPUT"
|
||||
echo "tag=$(git describe --tags --exact-match HEAD 2>/dev/null || true)" >> "$GITHUB_OUTPUT"
|
||||
|
||||
- name: Set up Docker Buildx
|
||||
uses: docker/setup-buildx-action@v2
|
||||
@@ -37,11 +29,7 @@ jobs:
|
||||
uses: docker/build-push-action@v4
|
||||
with:
|
||||
context: .
|
||||
file: Dockerfile.full
|
||||
push: true
|
||||
build-args: |
|
||||
GIT_COMMIT=${{ steps.gitmeta.outputs.sha }}
|
||||
GIT_TAG=${{ steps.gitmeta.outputs.tag }}
|
||||
tags: |
|
||||
ghcr.io/routstr/proxy:latest
|
||||
ghcr.io/routstr/core:latest
|
||||
|
||||
@@ -67,7 +67,7 @@ jobs:
|
||||
- name: Setup Node.js
|
||||
uses: actions/setup-node@v4
|
||||
with:
|
||||
node-version: "20"
|
||||
node-version: "18"
|
||||
cache: "pnpm"
|
||||
cache-dependency-path: ui/pnpm-lock.yaml
|
||||
|
||||
|
||||
@@ -21,10 +21,6 @@ WORKDIR /app
|
||||
|
||||
COPY . .
|
||||
|
||||
ARG GIT_COMMIT=""
|
||||
ARG GIT_TAG=""
|
||||
ENV GIT_COMMIT=${GIT_COMMIT}
|
||||
ENV GIT_TAG=${GIT_TAG}
|
||||
ENV PORT=8000
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
|
||||
@@ -1,54 +0,0 @@
|
||||
# Multi-stage Dockerfile for Routstr (includes UI build)
|
||||
# Stage 1: Build the UI
|
||||
FROM node:23-alpine AS ui-builder
|
||||
WORKDIR /app/ui
|
||||
|
||||
# Install pnpm
|
||||
RUN corepack enable pnpm && corepack prepare pnpm@10.15.0 --activate
|
||||
|
||||
# Copy UI source
|
||||
COPY ui/package.json ui/pnpm-lock.yaml* ./
|
||||
RUN pnpm install --frozen-lockfile
|
||||
|
||||
COPY ui/ ./
|
||||
ENV NEXT_TELEMETRY_DISABLED=1
|
||||
# Next.js build produces a static export in 'out' directory
|
||||
RUN pnpm run build
|
||||
|
||||
# Stage 2: Build the Routstr Node
|
||||
FROM ghcr.io/astral-sh/uv:python3.11-alpine AS runner
|
||||
|
||||
# Install system dependencies
|
||||
RUN apk add --no-cache \
|
||||
pkgconf \
|
||||
build-base \
|
||||
automake \
|
||||
autoconf \
|
||||
libtool \
|
||||
m4 \
|
||||
perl \
|
||||
git
|
||||
|
||||
WORKDIR /app
|
||||
|
||||
# Copy the rest of the application (required for uv sync to find the package)
|
||||
COPY . .
|
||||
|
||||
# Install dependencies including the specific secp256k1 branch
|
||||
RUN uv add git+https://github.com/saschanaz/secp256k1-py.git#branch=upgrade060
|
||||
RUN uv sync --no-dev
|
||||
|
||||
# Copy the built UI from the ui-builder stage
|
||||
COPY --from=ui-builder /app/ui/out ./ui_out
|
||||
|
||||
ARG GIT_COMMIT=""
|
||||
ARG GIT_TAG=""
|
||||
ENV GIT_COMMIT=${GIT_COMMIT}
|
||||
ENV GIT_TAG=${GIT_TAG}
|
||||
ENV PORT=8000
|
||||
ENV PYTHONUNBUFFERED=1
|
||||
|
||||
EXPOSE 8000
|
||||
|
||||
# Run the application
|
||||
CMD ["/app/.venv/bin/fastapi", "run", "routstr", "--host", "0.0.0.0"]
|
||||
@@ -51,26 +51,14 @@ curl https://api.routstr.com/v1/chat/completions \
|
||||
|
||||
## Quick Start (Docker)
|
||||
|
||||
If you are a node runner, start a Routstr Core instance using Docker Compose:
|
||||
If you are a node runner, start a Routstr Core instance and configure upstream access in the dashboard.
|
||||
|
||||
1. **Prepare your `.env`**:
|
||||
```bash
|
||||
ADMIN_PASSWORD=mysecretpassword
|
||||
NAME="My AI Node"
|
||||
DESCRIPTION="Fast access to models"
|
||||
NSEC=yournsec
|
||||
RECEIVE_LN_ADDRESS=yourname@wallet.com
|
||||
```
|
||||
|
||||
2. **Start the services**:
|
||||
```bash
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
3. **Configure**:
|
||||
Open [http://localhost:8000/admin/](http://localhost:8000/admin/) to connect your AI providers and set pricing.
|
||||
|
||||
For full instructions, see the **[Provider Quick Start Guide](https://docs.routstr.com/provider/quickstart/)**.
|
||||
```bash
|
||||
docker run -d \
|
||||
--name routstr-proxy \
|
||||
-p 8000:8000 \
|
||||
ghcr.io/routstr/proxy:latest
|
||||
```
|
||||
|
||||
## Development
|
||||
|
||||
|
||||
+1
-1
@@ -6,7 +6,7 @@ services:
|
||||
context: ./ui
|
||||
dockerfile: Dockerfile.build
|
||||
args:
|
||||
# NEXT_PUBLIC_API_URL: ${NEXT_PUBLIC_API_URL:-http://127.0.0.1:8000}
|
||||
NEXT_PUBLIC_API_URL: ${NEXT_PUBLIC_API_URL:-http://127.0.0.1:8000}
|
||||
NEXT_PUBLIC_ADMIN_API_KEY: ${NEXT_PUBLIC_ADMIN_API_KEY:-}
|
||||
volumes:
|
||||
- ./ui_out:/output
|
||||
|
||||
@@ -360,42 +360,6 @@ POST /v1/wallet/create
|
||||
}
|
||||
```
|
||||
|
||||
### Get Key Information
|
||||
|
||||
Get current balance, consumption data, and child keys for an API key.
|
||||
|
||||
```http
|
||||
GET /v1/balance/info
|
||||
Authorization: Bearer sk-...
|
||||
```
|
||||
|
||||
**Response:**
|
||||
|
||||
```json
|
||||
{
|
||||
"api_key": "sk-abc...",
|
||||
"balance": 8500000,
|
||||
"reserved": 0,
|
||||
"is_child": false,
|
||||
"parent_key": null,
|
||||
"total_requests": 42,
|
||||
"total_spent": 1500000,
|
||||
"balance_limit": null,
|
||||
"balance_limit_reset": null,
|
||||
"validity_date": null,
|
||||
"child_keys": [
|
||||
{
|
||||
"api_key": "sk-child1...",
|
||||
"total_requests": 10,
|
||||
"total_spent": 500000,
|
||||
"balance_limit": 1000000,
|
||||
"balance_limit_reset": "daily",
|
||||
"validity_date": 1738000000
|
||||
}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
### Check Balance
|
||||
|
||||
Get current wallet balance.
|
||||
|
||||
@@ -6,32 +6,6 @@ For automated deployments, you can optionally pre-configure settings via environ
|
||||
|
||||
---
|
||||
|
||||
## Initial Setup (.env file)
|
||||
|
||||
Before running your node, you should create a `.env` file in the project root. This file is used to bootstrap the initial configuration and store sensitive secrets.
|
||||
|
||||
### Example .env
|
||||
|
||||
```bash
|
||||
ADMIN_PASSWORD=your-secure-password
|
||||
|
||||
# Node Identity
|
||||
NAME="My AI Node"
|
||||
DESCRIPTION="Fast access to models"
|
||||
|
||||
# Lightning Payouts
|
||||
RECEIVE_LN_ADDRESS=yourname@wallet.com
|
||||
```
|
||||
|
||||
### Setting the UI Password
|
||||
|
||||
There are two ways to set or change your Admin Dashboard password:
|
||||
|
||||
1. **Via Environment Variable**: Set `ADMIN_PASSWORD` in your `.env` file before starting the container. This will be the password used for the first login.
|
||||
2. **Via Dashboard**: Once logged in, go to **Settings** → **Security** to update your password. Dashboard settings override the `.env` file once saved.
|
||||
|
||||
---
|
||||
|
||||
## Admin Dashboard (Primary)
|
||||
|
||||
Access the dashboard at `/admin/` on your node.
|
||||
@@ -40,29 +14,29 @@ Access the dashboard at `/admin/` on your node.
|
||||
|
||||
Connect to your AI provider(s):
|
||||
|
||||
| Setting | Description |
|
||||
| ---------------- | ------------------------------------------------ |
|
||||
| Setting | Description |
|
||||
|---------|-------------|
|
||||
| **Upstream URL** | API endpoint (e.g., `https://api.openai.com/v1`) |
|
||||
| **API Key** | Your provider's API key |
|
||||
| **API Key** | Your provider's API key |
|
||||
|
||||
### Node Identity
|
||||
|
||||
How your node appears to clients:
|
||||
|
||||
| Setting | Description |
|
||||
| --------------- | -------------------------------------- |
|
||||
| **Name** | Display name (e.g., "Fast GPT-4 Node") |
|
||||
| **Description** | Brief description of your service |
|
||||
| Setting | Description |
|
||||
|---------|-------------|
|
||||
| **Name** | Display name (e.g., "Fast GPT-4 Node") |
|
||||
| **Description** | Brief description of your service |
|
||||
|
||||
### Pricing
|
||||
|
||||
Control your profit margins:
|
||||
|
||||
| Setting | Description | Default |
|
||||
| ----------------- | ------------------------------------------ | ------------ |
|
||||
| **Fixed Pricing** | Charge flat rate per request vs. per-token | Off |
|
||||
| **Exchange Fee** | Buffer for BTC volatility | 1.005 (0.5%) |
|
||||
| **Upstream Fee** | Your profit markup | 1.10 (10%) |
|
||||
| Setting | Description | Default |
|
||||
|---------|-------------|---------|
|
||||
| **Fixed Pricing** | Charge flat rate per request vs. per-token | Off |
|
||||
| **Exchange Fee** | Buffer for BTC volatility | 1.005 (0.5%) |
|
||||
| **Upstream Fee** | Your profit markup | 1.10 (10%) |
|
||||
|
||||
See [Pricing](pricing.md) for detailed strategies.
|
||||
|
||||
@@ -70,40 +44,33 @@ See [Pricing](pricing.md) for detailed strategies.
|
||||
|
||||
Which mints to accept payments from:
|
||||
|
||||
| Setting | Description |
|
||||
| --------- | ------------------------------- |
|
||||
| Setting | Description |
|
||||
|---------|-------------|
|
||||
| **Mints** | List of trusted Cashu mint URLs |
|
||||
|
||||
### Lightning Withdrawals
|
||||
|
||||
Automatic profit withdrawal:
|
||||
|
||||
| Setting | Description | Default |
|
||||
| ------------------------------------ | ---------------------------------------------------------------------------------------------------------- | ------- |
|
||||
| **Lightning Address** | Your LN address for withdrawals | — |
|
||||
| **Minimum Payout (sat)** | Min available balance (in sats) before profit is paid out. Applies to both `sat` and `msat` mints (auto-converted). | `210` |
|
||||
| **Payout Interval (seconds)** | How often the payout loop wakes up and checks balances | `900` |
|
||||
|
||||
All payout amounts must be positive. Set the minimums above your wallet's
|
||||
minimum-invoice constraint (typically 1 sat) and high enough to amortise
|
||||
routing fees.
|
||||
| Setting | Description |
|
||||
|---------|-------------|
|
||||
| **Lightning Address** | Your LN address for withdrawals |
|
||||
|
||||
### Security
|
||||
|
||||
| Setting | Description |
|
||||
| ------------------ | ----------------------------- |
|
||||
| Setting | Description |
|
||||
|---------|-------------|
|
||||
| **Admin Password** | Password for dashboard access |
|
||||
|
||||
### Nostr Discovery
|
||||
|
||||
Announce your node on the network:
|
||||
|
||||
| Setting | Description |
|
||||
| ---------- | ------------------------------------ |
|
||||
| **Npub** | Your Nostr public key |
|
||||
| **Nsec** | Your Nostr private key (for signing) |
|
||||
| **Relays** | Relays to publish announcements |
|
||||
| **Share Analytics** | Publish aggregate usage stats to Nostr |
|
||||
| Setting | Description |
|
||||
|---------|-------------|
|
||||
| **Npub** | Your Nostr public key |
|
||||
| **Nsec** | Your Nostr private key (for signing) |
|
||||
| **Relays** | Relays to publish announcements |
|
||||
|
||||
See [Discovery](discovery.md) for details.
|
||||
|
||||
@@ -119,24 +86,21 @@ Use environment variables for:
|
||||
|
||||
### All Variables
|
||||
|
||||
| Variable | Description | Default |
|
||||
| -------------------- | --------------------------------- | ------------------------------------ |
|
||||
| `UPSTREAM_BASE_URL` | Upstream API endpoint | — |
|
||||
| `UPSTREAM_API_KEY` | Upstream API key | — |
|
||||
| `ADMIN_PASSWORD` | Dashboard password | (none) |
|
||||
| `DATABASE_URL` | Database connection string | `sqlite+aiosqlite:///keys.db` |
|
||||
| `NAME` | Node display name | `ARoutstrNode` |
|
||||
| `DESCRIPTION` | Node description | `A Routstr Node` |
|
||||
| `NPUB` | Nostr public key (bech32) | — |
|
||||
| `NSEC` | Nostr private key | — |
|
||||
| `ENABLE_ANALYTICS_SHARING` | Enable usage analytics sharing to Nostr | `true` |
|
||||
| `CASHU_MINTS` | Comma-separated mint URLs | `https://mint.minibits.cash/Bitcoin` |
|
||||
| `RECEIVE_LN_ADDRESS` | Lightning address for withdrawals | — |
|
||||
| `MIN_PAYOUT_SAT` | Min payout balance in sats (applies to all mints) | `210` |
|
||||
| `PAYOUT_INTERVAL_SECONDS` | Payout loop interval (seconds) | `900` |
|
||||
| `TOR_PROXY_URL` | SOCKS5 proxy for Tor | `socks5://127.0.0.1:9050` |
|
||||
| `CORS_ORIGINS` | Allowed CORS origins | `*` |
|
||||
| `RELAYS` | Nostr relays (comma-separated) | (default set) |
|
||||
| Variable | Description | Default |
|
||||
|----------|-------------|---------|
|
||||
| `UPSTREAM_BASE_URL` | Upstream API endpoint | — |
|
||||
| `UPSTREAM_API_KEY` | Upstream API key | — |
|
||||
| `ADMIN_PASSWORD` | Dashboard password | (none) |
|
||||
| `DATABASE_URL` | Database connection string | `sqlite+aiosqlite:///keys.db` |
|
||||
| `NAME` | Node display name | `ARoutstrNode` |
|
||||
| `DESCRIPTION` | Node description | `A Routstr Node` |
|
||||
| `NPUB` | Nostr public key (bech32) | — |
|
||||
| `NSEC` | Nostr private key | — |
|
||||
| `CASHU_MINTS` | Comma-separated mint URLs | `https://mint.minibits.cash/Bitcoin` |
|
||||
| `RECEIVE_LN_ADDRESS` | Lightning address for withdrawals | — |
|
||||
| `TOR_PROXY_URL` | SOCKS5 proxy for Tor | `socks5://127.0.0.1:9050` |
|
||||
| `CORS_ORIGINS` | Allowed CORS origins | `*` |
|
||||
| `RELAYS` | Nostr relays (comma-separated) | (default set) |
|
||||
|
||||
### Priority
|
||||
|
||||
|
||||
@@ -152,7 +152,6 @@ Manage which mints you accept payments from:
|
||||
|-------|-------------|
|
||||
| **Nsec** | Private key for signing announcements |
|
||||
| **Relays** | Where to publish your node advertisement |
|
||||
| **Share Analytics** | Toggle publishing aggregate usage stats to Nostr |
|
||||
|
||||
### Security
|
||||
|
||||
|
||||
+22
-67
@@ -2,71 +2,34 @@
|
||||
|
||||
Production deployment guide for Routstr Provider nodes.
|
||||
|
||||
## All-in-One Docker Image (Preferred)
|
||||
|
||||
The easiest way to deploy Routstr is using the all-in-one Docker image from Docker Hub, which includes both the FastAPI backend and the Next.js admin dashboard in a single container.
|
||||
|
||||
### Quick Start
|
||||
|
||||
```bash
|
||||
docker run -d \
|
||||
--name routstr \
|
||||
-p 8000:8000 \
|
||||
-v routstr-data:/app/data \
|
||||
-e DATABASE_URL="sqlite:////app/data/routstr.db" \
|
||||
9qeklajc/routstr:latest
|
||||
```
|
||||
|
||||
Access your node:
|
||||
- **API & Admin Dashboard**: http://localhost:8000
|
||||
|
||||
### Docker Compose Setup
|
||||
|
||||
Create `docker-compose.yml`:
|
||||
|
||||
```yaml
|
||||
version: '3.8'
|
||||
|
||||
services:
|
||||
routstr:
|
||||
image: 9qeklajc/routstr:latest
|
||||
container_name: routstr
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "8000:8000"
|
||||
volumes:
|
||||
- routstr-data:/app/data
|
||||
environment:
|
||||
DATABASE_URL: "sqlite:////app/data/routstr.db"
|
||||
ADMIN_KEY: "your-secure-admin-key"
|
||||
LOG_LEVEL: "info"
|
||||
|
||||
volumes:
|
||||
routstr-data:
|
||||
```
|
||||
|
||||
Start it:
|
||||
|
||||
```bash
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Docker Compose (Recommended)
|
||||
|
||||
For production, use Docker Compose with persistent storage and optional Tor support.
|
||||
|
||||
Use the included `compose.yml` for a flexible setup that handles both the UI and the node execution. This is useful for development or when you want to manage Tor as a separate service.
|
||||
### Basic Setup
|
||||
|
||||
Create a `compose.yml`:
|
||||
|
||||
```yaml
|
||||
services:
|
||||
routstr:
|
||||
image: ghcr.io/routstr/proxy:latest
|
||||
container_name: routstr
|
||||
restart: unless-stopped
|
||||
ports:
|
||||
- "8000:8000"
|
||||
volumes:
|
||||
- ./data:/app/data
|
||||
- ./logs:/app/logs
|
||||
```
|
||||
|
||||
Start the node:
|
||||
|
||||
```bash
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
This will:
|
||||
1. **Build the UI**: Compiles the frontend and copies it to a shared volume.
|
||||
2. **Start Routstr**: Runs the Python node, mounting the built UI.
|
||||
3. **Start Tor**: Provides anonymous access via a `.onion` address.
|
||||
Then configure everything via the [Admin Dashboard](http://localhost:8000/admin/).
|
||||
|
||||
---
|
||||
|
||||
@@ -226,16 +189,8 @@ docker compose up -d
|
||||
|
||||
## Building from Source
|
||||
|
||||
### Using Docker Compose
|
||||
The easiest way to build everything from source:
|
||||
|
||||
```bash
|
||||
docker compose build
|
||||
```
|
||||
|
||||
### Individual Components
|
||||
If you prefer building the node only (requires manual UI build first):
|
||||
|
||||
```bash
|
||||
docker build -t routstr-node .
|
||||
git clone https://github.com/routstr/routstr-core.git
|
||||
cd routstr-core
|
||||
docker build -t routstr-local .
|
||||
```
|
||||
|
||||
+10
-35
@@ -13,7 +13,7 @@ A **Routstr Provider Node** acts as a gateway that:
|
||||
You bring the API keys, Routstr handles the billing, payments, and client management.
|
||||
|
||||
!!! tip "Future: Node-to-Node Routing"
|
||||
In future versions, you'll be able to run a node that connects to other Routstr nodes—eliminating the need to configure upstream providers yourself. For now, you'll need your own API credentials.
|
||||
In future versions, you'll be able to run a node that connects to other Routstr nodes—eliminating the need to configure upstream providers yourself. For now, you'll need your own API credentials.
|
||||
|
||||
---
|
||||
|
||||
@@ -24,30 +24,14 @@ In future versions, you'll be able to run a node that connects to other Routstr
|
||||
|
||||
---
|
||||
|
||||
## 1. Prepare Configuration
|
||||
|
||||
Create a `.env` file in the root of the project to store your secrets:
|
||||
## 1. Start the Node
|
||||
|
||||
```bash
|
||||
# Initial Admin Password
|
||||
ADMIN_PASSWORD=mysecretpassword
|
||||
|
||||
# Node Identity
|
||||
NAME="My AI Node"
|
||||
DESCRIPTION="Fast access to models"
|
||||
NSEC=yournsec
|
||||
|
||||
# Lightning Payouts
|
||||
RECEIVE_LN_ADDRESS=yourname@wallet.com
|
||||
|
||||
```
|
||||
|
||||
## 2. Start the Node
|
||||
|
||||
The recommended way to run Routstr is using Docker Compose, which handles the node, the UI, and optional services like Tor.
|
||||
|
||||
```bash
|
||||
docker compose up -d
|
||||
docker run -d \
|
||||
--name routstr \
|
||||
-p 8000:8000 \
|
||||
-v routstr-data:/app/data \
|
||||
ghcr.io/routstr/proxy:latest
|
||||
```
|
||||
|
||||
Verify it's running:
|
||||
@@ -56,23 +40,14 @@ Verify it's running:
|
||||
curl http://localhost:8000/v1/info
|
||||
```
|
||||
|
||||
### Build from Source (Optional)
|
||||
|
||||
If you've cloned the repository and want to build the images yourself:
|
||||
|
||||
```bash
|
||||
docker compose build
|
||||
docker compose up -d
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## 3. Configure via Dashboard
|
||||
## 2. Configure via Dashboard
|
||||
|
||||
Open the **Admin Dashboard** at [http://localhost:8000/admin/](http://localhost:8000/admin/).
|
||||
|
||||
!!! note "Login"
|
||||
Use the `ADMIN_PASSWORD` you defined in your `.env` file to log in. If you didn't set one, the dashboard will prompt you to set one on first visit.
|
||||
!!! note "Default Access"
|
||||
The dashboard has no password by default. Set one immediately in Settings for production use.
|
||||
|
||||
### Connect Your AI Providers
|
||||
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
|
||||
|
||||
@@ -1,32 +0,0 @@
|
||||
"""add routstr_fees table
|
||||
|
||||
Revision ID: 02650cd6f028
|
||||
Revises: c3d4e5f6a7b8
|
||||
Create Date: 2026-04-24 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "02650cd6f028"
|
||||
down_revision = "c3d4e5f6a7b8"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"routstr_fees",
|
||||
sa.Column("id", sa.Integer(), nullable=False),
|
||||
sa.Column("accumulated_msats", sa.Integer(), nullable=False, server_default="0"),
|
||||
sa.Column("total_paid_msats", sa.Integer(), nullable=False, server_default="0"),
|
||||
sa.Column("last_paid_at", sa.Integer(), nullable=True),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
# Seed with a single row
|
||||
op.execute("INSERT INTO routstr_fees (id, accumulated_msats, total_paid_msats) VALUES (1, 0, 0)")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_table("routstr_fees")
|
||||
@@ -1,37 +0,0 @@
|
||||
"""add key management and reset fields to api_keys
|
||||
|
||||
Revision ID: 06f81c0fc88d
|
||||
Revises: c2d3e4f5a6b7
|
||||
Create Date: 2026-02-04 22:44:03.311983
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "06f81c0fc88d"
|
||||
down_revision = "c2d3e4f5a6b7"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column("api_keys", sa.Column("balance_limit", sa.Integer(), nullable=True))
|
||||
op.add_column(
|
||||
"api_keys",
|
||||
sa.Column(
|
||||
"balance_limit_reset", sqlmodel.sql.sqltypes.AutoString(), nullable=True
|
||||
),
|
||||
)
|
||||
op.add_column(
|
||||
"api_keys", sa.Column("balance_limit_reset_date", sa.Integer(), nullable=True)
|
||||
)
|
||||
op.add_column("api_keys", sa.Column("validity_date", sa.Integer(), nullable=True))
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("api_keys", "validity_date")
|
||||
op.drop_column("api_keys", "balance_limit_reset_date")
|
||||
op.drop_column("api_keys", "balance_limit_reset")
|
||||
op.drop_column("api_keys", "balance_limit")
|
||||
@@ -1,32 +0,0 @@
|
||||
"""add provider_settings to upstream_providers
|
||||
|
||||
Revision ID: 614c0a740e68
|
||||
Revises: 06f81c0fc88d
|
||||
Create Date: 2026-02-13 22:36:53.608737
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "614c0a740e68"
|
||||
down_revision = "06f81c0fc88d"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
# Check if column exists before adding it
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
columns = [c["name"] for c in inspector.get_columns("upstream_providers")]
|
||||
|
||||
if "provider_settings" not in columns:
|
||||
op.add_column(
|
||||
"upstream_providers",
|
||||
sa.Column("provider_settings", sa.Text(), nullable=True),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("upstream_providers", "provider_settings")
|
||||
@@ -1,42 +0,0 @@
|
||||
"""add cashu_transactions table
|
||||
|
||||
Revision ID: a776ca70e5fe
|
||||
Revises: 614c0a740e68
|
||||
Create Date: 2026-03-11 22:00:01.554762
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "a776ca70e5fe"
|
||||
down_revision = "614c0a740e68"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"cashu_transactions",
|
||||
sa.Column("id", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
|
||||
sa.Column("token", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
|
||||
sa.Column("amount", sa.Integer(), nullable=False),
|
||||
sa.Column("unit", sqlmodel.sql.sqltypes.AutoString(), nullable=False),
|
||||
sa.Column("mint_url", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
|
||||
sa.Column(
|
||||
"type",
|
||||
sqlmodel.sql.sqltypes.AutoString(),
|
||||
nullable=False,
|
||||
server_default="out",
|
||||
),
|
||||
sa.Column("request_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True),
|
||||
sa.Column("created_at", sa.Integer(), nullable=False),
|
||||
sa.Column("collected", sa.Boolean(), nullable=False),
|
||||
sa.Column("swept", sa.Boolean(), nullable=False),
|
||||
sa.PrimaryKeyConstraint("id"),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_table("cashu_transactions")
|
||||
@@ -1,33 +0,0 @@
|
||||
"""add forwarded_model_id to models
|
||||
|
||||
Revision ID: b1c2d3e4f5a6
|
||||
Revises: a776ca70e5fe
|
||||
Create Date: 2026-04-05 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "b1c2d3e4f5a6"
|
||||
down_revision = "a776ca70e5fe"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.add_column(
|
||||
"models",
|
||||
sa.Column(
|
||||
"forwarded_model_id",
|
||||
sqlmodel.sql.sqltypes.AutoString(),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
# Backfill: set forwarded_model_id = id for all existing rows
|
||||
op.execute("UPDATE models SET forwarded_model_id = id WHERE forwarded_model_id IS NULL")
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("models", "forwarded_model_id")
|
||||
@@ -1,36 +0,0 @@
|
||||
"""add source to cashu_transactions
|
||||
|
||||
Revision ID: c3d4e5f6a7b8
|
||||
Revises: b1c2d3e4f5a6
|
||||
Create Date: 2026-04-10 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "c3d4e5f6a7b8"
|
||||
down_revision = "b1c2d3e4f5a6"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
columns = [col["name"] for col in inspector.get_columns("cashu_transactions")]
|
||||
if "source" not in columns:
|
||||
op.add_column(
|
||||
"cashu_transactions",
|
||||
sa.Column(
|
||||
"source",
|
||||
sqlmodel.sql.sqltypes.AutoString(),
|
||||
nullable=False,
|
||||
server_default="x-cashu",
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_column("cashu_transactions", "source")
|
||||
@@ -1,34 +0,0 @@
|
||||
"""add cli_tokens table
|
||||
|
||||
Revision ID: cli_tokens_001
|
||||
Revises: e8f9a0b1c2d3
|
||||
Create Date: 2026-04-25 00:00:00.000000
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import sqlalchemy as sa
|
||||
from alembic import op
|
||||
|
||||
revision = "cli_tokens_001"
|
||||
down_revision = "e8f9a0b1c2d3"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
op.create_table(
|
||||
"cli_tokens",
|
||||
sa.Column("id", sa.String(), primary_key=True, nullable=False),
|
||||
sa.Column("token", sa.String(), nullable=False, unique=True),
|
||||
sa.Column("name", sa.String(), nullable=False),
|
||||
sa.Column("created_at", sa.Integer(), nullable=False),
|
||||
sa.Column("last_used_at", sa.Integer(), nullable=True),
|
||||
sa.Column("expires_at", sa.Integer(), nullable=True),
|
||||
)
|
||||
op.create_index("ix_cli_tokens_token", "cli_tokens", ["token"], unique=True)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_cli_tokens_token", table_name="cli_tokens")
|
||||
op.drop_table("cli_tokens")
|
||||
@@ -1,46 +0,0 @@
|
||||
"""add api key link to cashu_transactions
|
||||
|
||||
Revision ID: d4e5f6a7b8c9
|
||||
Revises: c3d4e5f6a7b8
|
||||
Create Date: 2026-04-20 00:00:00.000000
|
||||
"""
|
||||
|
||||
import sqlalchemy as sa
|
||||
import sqlmodel
|
||||
from alembic import op
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "d4e5f6a7b8c9"
|
||||
down_revision = "c3d4e5f6a7b8"
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
conn = op.get_bind()
|
||||
inspector = sa.inspect(conn)
|
||||
columns = [col["name"] for col in inspector.get_columns("cashu_transactions")]
|
||||
indexes = {index["name"] for index in inspector.get_indexes("cashu_transactions")}
|
||||
|
||||
if "api_key_hashed_key" not in columns:
|
||||
op.add_column(
|
||||
"cashu_transactions",
|
||||
sa.Column(
|
||||
"api_key_hashed_key",
|
||||
sqlmodel.sql.sqltypes.AutoString(),
|
||||
nullable=True,
|
||||
),
|
||||
)
|
||||
|
||||
if "ix_cashu_transactions_api_key_hashed_key" not in indexes:
|
||||
op.create_index(
|
||||
"ix_cashu_transactions_api_key_hashed_key",
|
||||
"cashu_transactions",
|
||||
["api_key_hashed_key"],
|
||||
unique=False,
|
||||
)
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
op.drop_index("ix_cashu_transactions_api_key_hashed_key", table_name="cashu_transactions")
|
||||
op.drop_column("cashu_transactions", "api_key_hashed_key")
|
||||
@@ -1,20 +0,0 @@
|
||||
"""merge heads: routstr_fees + api_key_to_cashu_transactions
|
||||
|
||||
Revision ID: e8f9a0b1c2d3
|
||||
Revises: 02650cd6f028, d4e5f6a7b8c9
|
||||
Create Date: 2026-04-24 00:00:00.000000
|
||||
"""
|
||||
|
||||
# revision identifiers, used by Alembic.
|
||||
revision = "e8f9a0b1c2d3"
|
||||
down_revision = ("02650cd6f028", "d4e5f6a7b8c9")
|
||||
branch_labels = None
|
||||
depends_on = None
|
||||
|
||||
|
||||
def upgrade() -> None:
|
||||
pass
|
||||
|
||||
|
||||
def downgrade() -> None:
|
||||
pass
|
||||
@@ -1,355 +0,0 @@
|
||||
# Plan: Track unsupported incoming mints in `other_mints` and include them in balances
|
||||
|
||||
## Goal
|
||||
|
||||
When a Cashu token arrives from a mint that is **not** in `settings.cashu_mints`, we currently create/use a wallet for that mint and swap value into the primary mint. Any change/surplus left behind after `melt()` remains in the foreign `token_wallet`, but that balance is not surfaced by `fetch_all_balances()` because it only iterates over configured mints.
|
||||
|
||||
This plan adds persistent tracking for those foreign mints in a new database table called `other_mints`, and updates balance reporting to include them.
|
||||
|
||||
---
|
||||
|
||||
## Current behavior
|
||||
|
||||
### Incoming unsupported mint flow
|
||||
|
||||
In `routstr/wallet.py`:
|
||||
|
||||
- `recieve_token()` deserializes the token
|
||||
- if `token_obj.mint not in settings.cashu_mints`, it calls `swap_to_primary_mint(token_obj, wallet)`
|
||||
- `swap_to_primary_mint()` calls `token_wallet.melt(...)` to pay the primary mint invoice
|
||||
|
||||
### Important detail: change is retained, not discarded
|
||||
|
||||
The underlying Cashu wallet library keeps any melt change:
|
||||
|
||||
- `Wallet.melt()` constructs blank outputs for change
|
||||
- when the melt succeeds, returned change is reconstructed into proofs
|
||||
- those proofs are appended to `self.proofs` and stored in the wallet DB
|
||||
|
||||
So surplus from unsupported mints is **not discarded**, but it may become invisible operationally.
|
||||
|
||||
### Visibility problem
|
||||
|
||||
`fetch_all_balances()` currently only loops over:
|
||||
|
||||
- `settings.cashu_mints`
|
||||
- units `sat` and `msat`
|
||||
|
||||
This means balances left on unsupported mints are not shown in admin balance reporting.
|
||||
|
||||
---
|
||||
|
||||
## Proposed design
|
||||
|
||||
## 1. Add a new DB table: `other_mints`
|
||||
|
||||
Add a small table in `routstr/core/db.py` to persist unsupported mints we have seen in incoming tokens.
|
||||
|
||||
Suggested schema:
|
||||
|
||||
- `mint_url: str` primary key
|
||||
- `created_at: int`
|
||||
- `last_seen_at: int`
|
||||
|
||||
Minimal model:
|
||||
|
||||
```python
|
||||
class OtherMint(SQLModel, table=True):
|
||||
__tablename__ = "other_mints"
|
||||
|
||||
mint_url: str = Field(primary_key=True)
|
||||
created_at: int = Field(default_factory=lambda: int(time.time()))
|
||||
last_seen_at: int = Field(default_factory=lambda: int(time.time()))
|
||||
```
|
||||
|
||||
Why minimal:
|
||||
|
||||
- the only required function is mint discovery/tracking
|
||||
- unit handling can remain dynamic via existing balance queries over `sat` and `msat`
|
||||
|
||||
---
|
||||
|
||||
## 2. Add DB helpers for `other_mints`
|
||||
|
||||
In `routstr/core/db.py`, add helper functions:
|
||||
|
||||
### `register_other_mint(mint_url: str) -> None`
|
||||
|
||||
Behavior:
|
||||
|
||||
- if the mint is not present, insert it
|
||||
- if it already exists, update `last_seen_at`
|
||||
|
||||
### `list_other_mints(session) -> list[str]`
|
||||
|
||||
Behavior:
|
||||
|
||||
- return all tracked unsupported mint URLs
|
||||
|
||||
Optional later:
|
||||
|
||||
- `delete_other_mint(...)`
|
||||
- admin cleanup helpers
|
||||
|
||||
---
|
||||
|
||||
## 3. Register unsupported mints during token receipt
|
||||
|
||||
Update `recieve_token()` in `routstr/wallet.py`.
|
||||
|
||||
Current logic:
|
||||
|
||||
```python
|
||||
if token_obj.mint not in settings.cashu_mints:
|
||||
return await swap_to_primary_mint(token_obj, wallet)
|
||||
```
|
||||
|
||||
Planned logic:
|
||||
|
||||
```python
|
||||
if token_obj.mint not in settings.cashu_mints:
|
||||
await db.register_other_mint(token_obj.mint)
|
||||
return await swap_to_primary_mint(token_obj, wallet)
|
||||
```
|
||||
|
||||
Why here:
|
||||
|
||||
- this is the earliest reliable point where we know the mint came in via an actual token
|
||||
- this is exactly the path that can leave foreign-mint change behind
|
||||
- it avoids needing to infer unsupported mints later from wallet internals
|
||||
|
||||
---
|
||||
|
||||
## 4. Update `fetch_all_balances()` to include `other_mints`
|
||||
|
||||
Current behavior only includes configured mints.
|
||||
|
||||
Planned behavior:
|
||||
|
||||
- load tracked unsupported mints from DB
|
||||
- combine them with `settings.cashu_mints`
|
||||
- dedupe while preserving order
|
||||
- fetch balances for all tracked mints across requested units
|
||||
|
||||
Conceptual flow:
|
||||
|
||||
```python
|
||||
tracked_mints = dedupe(settings.cashu_mints + other_mints_from_db)
|
||||
```
|
||||
|
||||
Then existing per-mint/per-unit balance logic can remain mostly unchanged.
|
||||
|
||||
This ensures that retained change on unsupported mints becomes visible in admin balance reporting.
|
||||
|
||||
---
|
||||
|
||||
## 5. Add a balance source marker
|
||||
|
||||
Extend `BalanceDetail` in `routstr/wallet.py` to identify whether a balance row comes from a configured mint or an `other_mints` entry.
|
||||
|
||||
Suggested field:
|
||||
|
||||
- `source: str` with values:
|
||||
- `"configured"`
|
||||
- `"other"`
|
||||
|
||||
Updated shape:
|
||||
|
||||
```python
|
||||
class BalanceDetail(TypedDict, total=False):
|
||||
mint_url: str
|
||||
unit: str
|
||||
source: str
|
||||
wallet_balance: int
|
||||
user_balance: int
|
||||
owner_balance: int
|
||||
error: str
|
||||
```
|
||||
|
||||
Why this helps:
|
||||
|
||||
- admin can distinguish normal configured wallet balances from foreign/unsupported balances
|
||||
- avoids confusion if unexpected mint URLs show up in the balances API/UI
|
||||
|
||||
---
|
||||
|
||||
## 6. Admin/API impact
|
||||
|
||||
Backend impact is minimal because `/admin/api/balances` already returns `fetch_all_balances()` output.
|
||||
|
||||
Effects:
|
||||
|
||||
- supported mints continue to show as before
|
||||
- tracked unsupported mints will also appear
|
||||
- UI can optionally display the new `source` field
|
||||
|
||||
No API contract break is expected if the frontend ignores unknown fields.
|
||||
|
||||
---
|
||||
|
||||
## 7. Payout behavior: do not change in phase 1
|
||||
|
||||
`periodic_payout()` currently only iterates over `settings.cashu_mints`.
|
||||
|
||||
Recommendation for this change:
|
||||
|
||||
- **do not** expand `periodic_payout()` to include `other_mints` yet
|
||||
- only improve visibility through balance reporting
|
||||
|
||||
Reason:
|
||||
|
||||
- automatic payout from unsupported/foreign mints may be operationally undesirable
|
||||
- visibility should come first, automation second
|
||||
|
||||
Possible future phase:
|
||||
|
||||
- add optional sweeping/payout support for `other_mints`
|
||||
- or provide an admin-triggered withdrawal/sweep flow
|
||||
|
||||
---
|
||||
|
||||
## 8. Logging improvements (optional)
|
||||
|
||||
Optional follow-up improvement in `swap_to_primary_mint()`:
|
||||
|
||||
- capture the return value from `token_wallet.melt(...)`
|
||||
- if feasible, log any reported change amount
|
||||
- otherwise, rely on wallet balance reporting to surface residual amounts
|
||||
|
||||
This is useful but not required for the first implementation.
|
||||
|
||||
---
|
||||
|
||||
## Files to change
|
||||
|
||||
### `routstr/core/db.py`
|
||||
|
||||
Add:
|
||||
|
||||
- `OtherMint` SQLModel
|
||||
- `register_other_mint()`
|
||||
- `list_other_mints()`
|
||||
|
||||
### `migrations/versions/<new_revision>_add_other_mints_table.py`
|
||||
|
||||
Create migration to add the `other_mints` table.
|
||||
|
||||
### `routstr/wallet.py`
|
||||
|
||||
Update:
|
||||
|
||||
- `recieve_token()` to register unsupported mints
|
||||
- `BalanceDetail` to include `source`
|
||||
- `fetch_all_balances()` to include both configured and tracked unsupported mints
|
||||
|
||||
### `routstr/core/admin.py`
|
||||
|
||||
Likely no backend changes required unless a dedicated `other_mints` API is desired.
|
||||
|
||||
---
|
||||
|
||||
## Behavior rules
|
||||
|
||||
### Register a mint when
|
||||
|
||||
- an incoming token is processed
|
||||
- the token mint is not in `settings.cashu_mints`
|
||||
|
||||
### Do not remove automatically when
|
||||
|
||||
- balance reaches zero
|
||||
|
||||
Reason:
|
||||
|
||||
- historical visibility is useful
|
||||
- avoids flapping entries in the admin balance list
|
||||
- mint may receive additional unsupported tokens later
|
||||
|
||||
Potential future enhancement:
|
||||
|
||||
- admin endpoint to prune zero-balance `other_mints`
|
||||
|
||||
---
|
||||
|
||||
## Edge cases
|
||||
|
||||
### A mint later becomes configured
|
||||
|
||||
If a mint in `other_mints` is later added to `settings.cashu_mints`:
|
||||
|
||||
- deduplication prevents duplicate balance rows
|
||||
- `source` should resolve to `configured`
|
||||
|
||||
### Unsupported mint with zero balance
|
||||
|
||||
A tracked unsupported mint may show zero balances.
|
||||
|
||||
Initial recommendation:
|
||||
|
||||
- allow it to appear
|
||||
- consider later filtering zero-balance `other` rows if the UI becomes noisy
|
||||
|
||||
### Units
|
||||
|
||||
Balance fetching can continue to query both `sat` and `msat` for each tracked mint.
|
||||
|
||||
If a mint has no proofs in one unit, current error/zero handling can continue to apply.
|
||||
|
||||
---
|
||||
|
||||
## Test plan
|
||||
|
||||
### DB tests
|
||||
|
||||
- registering a new unsupported mint inserts a row
|
||||
- registering the same mint again updates `last_seen_at` without duplication
|
||||
- listing other mints returns expected mint URLs
|
||||
|
||||
### Wallet tests
|
||||
|
||||
#### `recieve_token()`
|
||||
|
||||
- when mint is unsupported, `db.register_other_mint()` is called before swap
|
||||
- when mint is configured, `db.register_other_mint()` is not called
|
||||
|
||||
#### `fetch_all_balances()`
|
||||
|
||||
- includes configured mints
|
||||
- includes `other_mints` from DB
|
||||
- dedupes if a mint exists in both configured and other lists
|
||||
- sets `source` correctly
|
||||
|
||||
### Regression tests
|
||||
|
||||
- existing trusted mint balance reporting remains unchanged
|
||||
- `/admin/api/balances` continues to work
|
||||
|
||||
---
|
||||
|
||||
## Recommended implementation order
|
||||
|
||||
1. Add `OtherMint` model to `routstr/core/db.py`
|
||||
2. Add Alembic migration for `other_mints`
|
||||
3. Add `register_other_mint()` and `list_other_mints()` helpers
|
||||
4. Update `recieve_token()` to register unsupported mints
|
||||
5. Update `fetch_all_balances()` to union configured + tracked other mints
|
||||
6. Add `source` to `BalanceDetail`
|
||||
7. Add/adjust tests
|
||||
|
||||
---
|
||||
|
||||
## Summary
|
||||
|
||||
This change solves an operational visibility problem:
|
||||
|
||||
- unsupported incoming mints can leave retained change in foreign wallets
|
||||
- those funds are currently preserved but not surfaced in balance reporting
|
||||
- introducing `other_mints` makes those mints discoverable and auditable
|
||||
- expanding `fetch_all_balances()` ensures their balances are visible in admin tooling
|
||||
|
||||
Recommended scope for the first pass:
|
||||
|
||||
- track unsupported mints in DB
|
||||
- include them in balance reporting
|
||||
- mark them as `source="other"`
|
||||
- do not yet change payout/sweeping behavior
|
||||
+2
-3
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "routstr"
|
||||
version = "0.4.3"
|
||||
version = "0.3.0"
|
||||
description = "Payment proxy for your LLM endpoint using cashu and nostr."
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.11"
|
||||
@@ -13,7 +13,7 @@ dependencies = [
|
||||
"greenlet>=3.2.1",
|
||||
"alembic>=1.13",
|
||||
"python-json-logger>=2.0.0",
|
||||
"cashu>=0.20",
|
||||
"cashu",
|
||||
"secp256k1",
|
||||
"marshmallow>=3.13,<4.0",
|
||||
"websockets>=12.0",
|
||||
@@ -21,7 +21,6 @@ dependencies = [
|
||||
"mdurl==0.1.2",
|
||||
"pillow>=10",
|
||||
"openai>=1.98.0",
|
||||
"litellm>=1.55.0",
|
||||
]
|
||||
|
||||
[dependency-groups]
|
||||
|
||||
+4
-144
@@ -118,32 +118,12 @@ def create_model_mappings(
|
||||
|
||||
candidates: dict[str, list[tuple["Model", "BaseUpstreamProvider"]]] = {}
|
||||
unique_models: dict[str, "Model"] = {}
|
||||
seen_model_provider: set[tuple[str, str]] = set()
|
||||
|
||||
providers_by_db_id: dict[int, "BaseUpstreamProvider"] = {}
|
||||
for upstream in upstreams:
|
||||
db_id = getattr(upstream, "db_id", None)
|
||||
if isinstance(db_id, int):
|
||||
providers_by_db_id[db_id] = upstream
|
||||
|
||||
# Group upstreams by URL and keep only the one with the lowest fee for each URL
|
||||
upstreams_by_url: dict[str, list["BaseUpstreamProvider"]] = {}
|
||||
for upstream in upstreams:
|
||||
url = getattr(upstream, "base_url", "")
|
||||
if url not in upstreams_by_url:
|
||||
upstreams_by_url[url] = []
|
||||
upstreams_by_url[url].append(upstream)
|
||||
|
||||
filtered_upstreams: list["BaseUpstreamProvider"] = []
|
||||
for providers in upstreams_by_url.values():
|
||||
best_provider = min(providers, key=lambda p: p.provider_fee)
|
||||
filtered_upstreams.append(best_provider)
|
||||
|
||||
# Separate OpenRouter from other providers
|
||||
openrouter: "BaseUpstreamProvider" | None = None
|
||||
other_upstreams: list["BaseUpstreamProvider"] = []
|
||||
|
||||
for upstream in filtered_upstreams:
|
||||
for upstream in upstreams:
|
||||
base_url = getattr(upstream, "base_url", "")
|
||||
if base_url == "https://openrouter.ai/api/v1":
|
||||
openrouter = upstream
|
||||
@@ -154,16 +134,6 @@ def create_model_mappings(
|
||||
"""Get base model ID by removing provider prefix."""
|
||||
return model_id.split("/", 1)[1] if "/" in model_id else model_id
|
||||
|
||||
def get_provider_identity(upstream: "BaseUpstreamProvider") -> str:
|
||||
"""Get a stable provider identity used for deduplication."""
|
||||
db_id = getattr(upstream, "db_id", None)
|
||||
if isinstance(db_id, int):
|
||||
return f"db:{db_id}"
|
||||
|
||||
provider_type = str(getattr(upstream, "provider_type", "") or "").lower()
|
||||
base_url = str(getattr(upstream, "base_url", "") or "").lower()
|
||||
return f"{provider_type}|{base_url}"
|
||||
|
||||
def _add_candidate(
|
||||
alias: str, model: "Model", provider: "BaseUpstreamProvider"
|
||||
) -> None:
|
||||
@@ -178,7 +148,6 @@ def create_model_mappings(
|
||||
) -> None:
|
||||
"""Process all models from a given provider."""
|
||||
upstream_prefix = getattr(upstream, "upstream_name", None)
|
||||
provider_key = get_provider_identity(upstream)
|
||||
|
||||
for model in upstream.get_cached_models():
|
||||
if not model.enabled or model.id in disabled_model_ids:
|
||||
@@ -195,15 +164,14 @@ def create_model_mappings(
|
||||
|
||||
# Add to unique models
|
||||
base_id = get_base_model_id(model_to_use.id)
|
||||
unique_key = model_to_use.forwarded_model_id or base_id
|
||||
if not is_openrouter or unique_key not in unique_models:
|
||||
if not is_openrouter or base_id not in unique_models:
|
||||
unique_model = model_to_use.copy(
|
||||
update={
|
||||
"id": base_id,
|
||||
"upstream_provider_id": upstream.provider_type,
|
||||
}
|
||||
)
|
||||
unique_models[unique_key] = unique_model
|
||||
unique_models[base_id] = unique_model
|
||||
|
||||
# Get all aliases for this model
|
||||
aliases = resolve_model_alias(
|
||||
@@ -218,14 +186,9 @@ def create_model_mappings(
|
||||
if prefixed_id not in aliases:
|
||||
aliases.append(prefixed_id)
|
||||
|
||||
# Register forwarded_model_id as a routable alias
|
||||
if model_to_use.forwarded_model_id and model_to_use.forwarded_model_id not in aliases:
|
||||
aliases.append(model_to_use.forwarded_model_id)
|
||||
|
||||
# Try to set each alias
|
||||
for alias in aliases:
|
||||
_add_candidate(alias, model_to_use, upstream)
|
||||
seen_model_provider.add((model_to_use.id.lower(), provider_key))
|
||||
|
||||
# Process non-OpenRouter providers first
|
||||
for upstream in other_upstreams:
|
||||
@@ -235,115 +198,12 @@ def create_model_mappings(
|
||||
if openrouter:
|
||||
process_provider_models(openrouter, is_openrouter=True)
|
||||
|
||||
# Include enabled DB overrides even when provider discovery misses models.
|
||||
# This is important for deployment-based providers like Azure.
|
||||
for model_id, override_data in overrides_by_id.items():
|
||||
if model_id in disabled_model_ids:
|
||||
continue
|
||||
override_row, provider_fee = override_data
|
||||
upstream_provider_id = getattr(override_row, "upstream_provider_id", None)
|
||||
if not isinstance(upstream_provider_id, int):
|
||||
continue
|
||||
|
||||
upstream_for_override = providers_by_db_id.get(upstream_provider_id)
|
||||
if upstream_for_override is None:
|
||||
continue
|
||||
|
||||
provider_key = get_provider_identity(upstream_for_override)
|
||||
dedupe_key = (model_id.lower(), provider_key)
|
||||
if dedupe_key in seen_model_provider:
|
||||
continue
|
||||
|
||||
try:
|
||||
model_to_use = _row_to_model(
|
||||
override_row, apply_provider_fee=True, provider_fee=provider_fee
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Skipping invalid model override while building model mappings",
|
||||
extra={
|
||||
"model_id": model_id,
|
||||
"upstream_provider_id": upstream_provider_id,
|
||||
"error": str(exc),
|
||||
"error_type": type(exc).__name__,
|
||||
},
|
||||
)
|
||||
continue
|
||||
if not model_to_use.enabled:
|
||||
continue
|
||||
|
||||
base_id = get_base_model_id(model_to_use.id)
|
||||
unique_key = model_to_use.forwarded_model_id or base_id
|
||||
is_openrouter = (
|
||||
getattr(upstream_for_override, "base_url", "")
|
||||
== "https://openrouter.ai/api/v1"
|
||||
)
|
||||
if not is_openrouter or unique_key not in unique_models:
|
||||
unique_model = model_to_use.copy(
|
||||
update={
|
||||
"id": base_id,
|
||||
"upstream_provider_id": upstream_for_override.provider_type,
|
||||
}
|
||||
)
|
||||
unique_models[unique_key] = unique_model
|
||||
|
||||
try:
|
||||
aliases = resolve_model_alias(
|
||||
model_to_use.id,
|
||||
model_to_use.canonical_slug,
|
||||
alias_ids=model_to_use.alias_ids,
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"Skipping model aliases for invalid override model",
|
||||
extra={
|
||||
"model_id": model_id,
|
||||
"upstream_provider_id": upstream_provider_id,
|
||||
"error": str(exc),
|
||||
"error_type": type(exc).__name__,
|
||||
},
|
||||
)
|
||||
continue
|
||||
|
||||
upstream_prefix = getattr(upstream_for_override, "upstream_name", None)
|
||||
if upstream_prefix and "/" not in model_to_use.id:
|
||||
prefixed_id = f"{upstream_prefix}/{model_to_use.id}"
|
||||
if prefixed_id not in aliases:
|
||||
aliases.append(prefixed_id)
|
||||
|
||||
# Register forwarded_model_id as a routable alias
|
||||
if model_to_use.forwarded_model_id and model_to_use.forwarded_model_id not in aliases:
|
||||
aliases.append(model_to_use.forwarded_model_id)
|
||||
|
||||
for alias in aliases:
|
||||
_add_candidate(alias, model_to_use, upstream_for_override)
|
||||
seen_model_provider.add(dedupe_key)
|
||||
|
||||
# Sort candidates and build final maps
|
||||
model_instances: dict[str, "Model"] = {}
|
||||
provider_map: dict[str, list["BaseUpstreamProvider"]] = {}
|
||||
|
||||
def alias_priority(model: "Model", alias: str) -> int:
|
||||
"""Rank how strong the mapping of alias->model is.
|
||||
|
||||
forwarded_model_id is the most specific identifier (set per-provider
|
||||
instance), so a match there should beat a model_id match. This way,
|
||||
when multiple providers have the same model_id but different
|
||||
forwarded_model_ids, the one whose forwarded_model_id equals the
|
||||
requested alias wins.
|
||||
"""
|
||||
if (
|
||||
model.forwarded_model_id
|
||||
and model.forwarded_model_id.lower() == alias
|
||||
):
|
||||
return 5
|
||||
|
||||
if (
|
||||
model.id
|
||||
and model.id.lower() == alias
|
||||
):
|
||||
return 4
|
||||
|
||||
"""Rank how strong the mapping of alias->model is."""
|
||||
model_base = get_base_model_id(model.id)
|
||||
if model_base == alias:
|
||||
return 3
|
||||
|
||||
+80
-523
@@ -1,18 +1,13 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import math
|
||||
import random
|
||||
import time
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import HTTPException
|
||||
from sqlalchemy import case
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlmodel import col, select, update
|
||||
from sqlmodel import col, update
|
||||
|
||||
from .core import get_logger
|
||||
from .core.db import ApiKey, AsyncSession, accumulate_routstr_fee
|
||||
from .core.db import ApiKey, AsyncSession
|
||||
from .core.settings import settings
|
||||
from .payment.cost_calculation import (
|
||||
CostData,
|
||||
@@ -23,73 +18,22 @@ from .payment.cost_calculation import (
|
||||
from .wallet import credit_balance, deserialize_token_from_string
|
||||
|
||||
logger = get_logger(__name__)
|
||||
payments_logger = get_logger("routstr.payments")
|
||||
|
||||
# Routstr platform fee constants
|
||||
ROUTSTR_FEE_PERCENT: float = 2.1
|
||||
ROUTSTR_LN_ADDRESS: str = "npub130mznv74rxs032peqym6g3wqavh472623mt3z5w73xq9r6qqdufs7ql29s@npub.cash"
|
||||
ROUTSTR_FEE_PAYOUT_INTERVAL_SECONDS: int = 900
|
||||
ROUTSTR_FEE_DEFAULT_PAYOUT: int = 200
|
||||
|
||||
# TODO: implement prepaid api key (not like it was before)
|
||||
# PREPAID_API_KEY = os.environ.get("PREPAID_API_KEY", None)
|
||||
# PREPAID_BALANCE = int(os.environ.get("PREPAID_BALANCE", "0")) * 1000 # Convert to msats
|
||||
|
||||
|
||||
async def check_and_reset_limit(key: ApiKey, session: AsyncSession) -> bool:
|
||||
"""Checks if a key's balance limit should be reset based on its policy."""
|
||||
if key.balance_limit is not None and key.balance_limit_reset:
|
||||
now = int(time.time())
|
||||
reset_date = key.balance_limit_reset_date or 0
|
||||
should_reset = False
|
||||
|
||||
if key.balance_limit_reset == "daily":
|
||||
if (
|
||||
datetime.fromtimestamp(now).date()
|
||||
> datetime.fromtimestamp(reset_date).date()
|
||||
):
|
||||
should_reset = True
|
||||
elif key.balance_limit_reset == "weekly":
|
||||
if (
|
||||
datetime.fromtimestamp(now).isocalendar()[:2]
|
||||
> datetime.fromtimestamp(reset_date).isocalendar()[:2]
|
||||
):
|
||||
should_reset = True
|
||||
elif key.balance_limit_reset == "monthly":
|
||||
dt_now = datetime.fromtimestamp(now)
|
||||
dt_reset = datetime.fromtimestamp(reset_date)
|
||||
if dt_now.year > dt_reset.year or dt_now.month > dt_reset.month:
|
||||
should_reset = True
|
||||
|
||||
if should_reset:
|
||||
logger.info(
|
||||
"Resetting balance limit for key",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"policy": key.balance_limit_reset,
|
||||
"old_spent": key.total_spent,
|
||||
},
|
||||
)
|
||||
key.total_spent = 0
|
||||
key.balance_limit_reset_date = now
|
||||
session.add(key)
|
||||
await session.flush()
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
async def validate_bearer_key(
|
||||
bearer_key: str,
|
||||
session: AsyncSession,
|
||||
refund_address: Optional[str] = None,
|
||||
key_expiry_time: Optional[int] = None,
|
||||
min_cost: int = 0,
|
||||
) -> ApiKey:
|
||||
"""
|
||||
Validates the provided API key using SQLModel.
|
||||
If it's a cashu key, it redeems it and stores its hash and balance.
|
||||
Otherwise checks if the hash of the key exists.
|
||||
Includes a balance check against min_cost for limited keys.
|
||||
"""
|
||||
logger.debug(
|
||||
"Starting bearer key validation",
|
||||
@@ -99,7 +43,6 @@ async def validate_bearer_key(
|
||||
else bearer_key,
|
||||
"has_refund_address": bool(refund_address),
|
||||
"has_expiry_time": bool(key_expiry_time),
|
||||
"min_cost": min_cost,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -154,50 +97,6 @@ async def validate_bearer_key(
|
||||
},
|
||||
)
|
||||
|
||||
# Check and reset limit if needed
|
||||
await check_and_reset_limit(existing_key, session)
|
||||
|
||||
# Early check: Billing balance check (Parent balance)
|
||||
billing_key = await get_billing_key(existing_key, session)
|
||||
if min_cost > 0 and billing_key.total_balance < min_cost:
|
||||
logger.warning(
|
||||
"Insufficient billing balance during validation",
|
||||
extra={
|
||||
"key_hash": existing_key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"balance": billing_key.total_balance,
|
||||
"required": min_cost,
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Insufficient balance: {min_cost} mSats required for this model. {billing_key.total_balance} available.",
|
||||
"type": "insufficient_quota",
|
||||
"code": "insufficient_balance",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
# Early check: Spending limit check (Child key limit)
|
||||
if (
|
||||
min_cost > 0
|
||||
and existing_key.balance_limit is not None
|
||||
and existing_key.total_spent + existing_key.reserved_balance + min_cost
|
||||
> existing_key.balance_limit
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Balance limit exceeded: {existing_key.balance_limit} mSats limit. {existing_key.total_spent} already spent ({existing_key.reserved_balance} reserved), {min_cost} minimum required for this model.",
|
||||
"type": "insufficient_quota",
|
||||
"code": "balance_limit_exceeded",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
return existing_key
|
||||
else:
|
||||
logger.warning(
|
||||
@@ -253,19 +152,6 @@ async def validate_bearer_key(
|
||||
},
|
||||
)
|
||||
|
||||
# Early check: Billing balance check
|
||||
if min_cost > 0 and existing_key.total_balance < min_cost:
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Insufficient balance: {min_cost} mSats required for this model. {existing_key.total_balance} available.",
|
||||
"type": "insufficient_quota",
|
||||
"code": "insufficient_balance",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
return existing_key
|
||||
|
||||
logger.info(
|
||||
@@ -317,13 +203,13 @@ async def validate_bearer_key(
|
||||
extra={"key_hash": hashed_key[:8] + "..."},
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
logger.info(
|
||||
"AUTH: About to call credit_balance",
|
||||
extra={"token_preview": bearer_key[:50]},
|
||||
)
|
||||
try:
|
||||
msats = await credit_balance(bearer_key, new_key, session)
|
||||
logger.debug(
|
||||
logger.info(
|
||||
"AUTH: credit_balance returned successfully", extra={"msats": msats}
|
||||
)
|
||||
except Exception as credit_error:
|
||||
@@ -356,8 +242,6 @@ async def validate_bearer_key(
|
||||
)
|
||||
|
||||
return new_key
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Cashu token redemption failed",
|
||||
@@ -380,12 +264,12 @@ async def validate_bearer_key(
|
||||
},
|
||||
)
|
||||
|
||||
key_preview = bearer_key[:10] + "..." if len(bearer_key) > 10 else bearer_key
|
||||
logger.error(
|
||||
f"Invalid API key format: preview={key_preview!r} length={len(bearer_key)} "
|
||||
f"(expected 'sk-...' or 'cashu...' token)",
|
||||
"Invalid API key format",
|
||||
extra={
|
||||
"key_preview": key_preview,
|
||||
"key_preview": bearer_key[:10] + "..."
|
||||
if len(bearer_key) > 10
|
||||
else bearer_key,
|
||||
"key_length": len(bearer_key),
|
||||
},
|
||||
)
|
||||
@@ -394,7 +278,7 @@ async def validate_bearer_key(
|
||||
status_code=401,
|
||||
detail={
|
||||
"error": {
|
||||
"message": "Invalid API key format. Expected an 'sk-...' API key or a 'cashu...' token.",
|
||||
"message": "Invalid API key",
|
||||
"type": "invalid_request_error",
|
||||
"code": "invalid_api_key",
|
||||
}
|
||||
@@ -427,8 +311,6 @@ async def pay_for_request(
|
||||
key: ApiKey, cost_per_request: int, session: AsyncSession
|
||||
) -> int:
|
||||
"""Process payment for a request."""
|
||||
# Ensure cost_per_request is at least the minimum allowed request cost
|
||||
cost_per_request = max(cost_per_request, settings.min_request_msat)
|
||||
|
||||
billing_key = await get_billing_key(key, session)
|
||||
|
||||
@@ -467,57 +349,6 @@ async def pay_for_request(
|
||||
},
|
||||
)
|
||||
|
||||
# Check validity date
|
||||
if key.validity_date is not None:
|
||||
if time.time() > key.validity_date:
|
||||
logger.warning(
|
||||
"Key validity date expired",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"validity_date": key.validity_date,
|
||||
"current_time": time.time(),
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": {
|
||||
"message": "API key has expired (validity date reached).",
|
||||
"type": "invalid_request_error",
|
||||
"code": "key_expired",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
# Check balance limit for child keys (or any key with a limit)
|
||||
if key.balance_limit is not None:
|
||||
await check_and_reset_limit(key, session)
|
||||
|
||||
if (
|
||||
key.total_spent + key.reserved_balance + cost_per_request
|
||||
> key.balance_limit
|
||||
):
|
||||
logger.warning(
|
||||
"Balance limit exceeded",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"total_spent": key.total_spent,
|
||||
"reserved": key.reserved_balance,
|
||||
"balance_limit": key.balance_limit,
|
||||
"required": cost_per_request,
|
||||
},
|
||||
)
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"Balance limit exceeded: {key.balance_limit} mSats limit. {key.total_spent} already spent ({key.reserved_balance} reserved), {cost_per_request} required for this request.",
|
||||
"type": "insufficient_quota",
|
||||
"code": "balance_limit_exceeded",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"Charging base cost for request",
|
||||
extra={
|
||||
@@ -540,15 +371,12 @@ async def pay_for_request(
|
||||
)
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also increment total_requests and reserved_balance on the child key if it's different
|
||||
# Also increment total_requests on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
total_requests=col(ApiKey.total_requests) + 1,
|
||||
reserved_balance=col(ApiKey.reserved_balance) + cost_per_request,
|
||||
)
|
||||
.values(total_requests=col(ApiKey.total_requests) + 1)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
@@ -592,33 +420,18 @@ async def pay_for_request(
|
||||
"total_requests": billing_key.total_requests,
|
||||
},
|
||||
)
|
||||
payments_logger.info(
|
||||
"RESERVE",
|
||||
extra={
|
||||
"event": "reserve",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"cost_reserved": cost_per_request,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
},
|
||||
)
|
||||
|
||||
return cost_per_request
|
||||
|
||||
|
||||
async def revert_pay_for_request(
|
||||
key: ApiKey, session: AsyncSession, cost_per_request: int
|
||||
) -> bool:
|
||||
"""Revert a previously reserved payment. Returns True if revert succeeded,
|
||||
False if the reservation was already released (prevents negative reserved_balance)."""
|
||||
) -> None:
|
||||
billing_key = await get_billing_key(key, session)
|
||||
|
||||
stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= cost_per_request)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
|
||||
total_requests=col(ApiKey.total_requests) - 1,
|
||||
@@ -627,23 +440,19 @@ async def revert_pay_for_request(
|
||||
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also decrement total_requests and reserved_balance on the child key if it's different
|
||||
# Also decrement total_requests on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= cost_per_request)
|
||||
.values(
|
||||
total_requests=col(ApiKey.total_requests) - 1,
|
||||
reserved_balance=col(ApiKey.reserved_balance) - cost_per_request,
|
||||
)
|
||||
.values(total_requests=col(ApiKey.total_requests) - 1)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
if result.rowcount == 0:
|
||||
logger.warning(
|
||||
"Revert skipped - reservation already released (no-op to prevent negative reserved_balance)",
|
||||
logger.error(
|
||||
"Failed to revert payment - insufficient reserved balance",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
@@ -651,22 +460,19 @@ async def revert_pay_for_request(
|
||||
"current_reserved_balance": billing_key.reserved_balance,
|
||||
},
|
||||
)
|
||||
return False
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
detail={
|
||||
"error": {
|
||||
"message": f"failed to revert request payment: {cost_per_request} mSats required. {billing_key.balance} available.",
|
||||
"type": "payment_error",
|
||||
"code": "payment_error",
|
||||
}
|
||||
},
|
||||
)
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
payments_logger.info(
|
||||
"REVERT",
|
||||
extra={
|
||||
"event": "revert",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"cost_reverted": cost_per_request,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
},
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
async def adjust_payment_for_tokens(
|
||||
@@ -698,45 +504,20 @@ async def adjust_payment_for_tokens(
|
||||
release_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost
|
||||
)
|
||||
)
|
||||
result = await session.exec(release_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also release on child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_release_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost
|
||||
)
|
||||
)
|
||||
await session.exec(child_release_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.exec(release_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
if result.rowcount == 0: # type: ignore[union-attr]
|
||||
logger.warning(
|
||||
"Release reservation skipped - already released (no-op to prevent negative reserved_balance)",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
},
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Released reservation without charging (fallback)",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
},
|
||||
)
|
||||
logger.warning(
|
||||
"Released reservation without charging (fallback)",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to release reservation in fallback",
|
||||
@@ -747,17 +528,6 @@ async def adjust_payment_for_tokens(
|
||||
},
|
||||
)
|
||||
|
||||
async def _accumulate_fee(total_cost_msats: int) -> None:
|
||||
if total_cost_msats > 0 and ROUTSTR_FEE_PERCENT > 0:
|
||||
fee_msats = math.ceil(total_cost_msats * ROUTSTR_FEE_PERCENT / 100)
|
||||
try:
|
||||
await accumulate_routstr_fee(session, fee_msats)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Failed to accumulate Routstr fee",
|
||||
extra={"error": str(e), "fee_msats": fee_msats},
|
||||
)
|
||||
|
||||
match await calculate_cost(response_data, deducted_max_cost, session):
|
||||
case MaxCostData() as cost:
|
||||
logger.debug(
|
||||
@@ -770,52 +540,23 @@ async def adjust_payment_for_tokens(
|
||||
},
|
||||
)
|
||||
# Finalize by releasing reservation and charging max cost
|
||||
if billing_key.reserved_balance < deducted_max_cost:
|
||||
logger.error(
|
||||
"reserved_balance below deducted_max_cost before MaxCost finalization — clamping to 0",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"total_cost_msats": cost.total_msats,
|
||||
"balance": billing_key.balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=safe_reserved,
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
|
||||
balance=col(ApiKey.balance) - cost.total_msats,
|
||||
total_spent=col(ApiKey.total_spent) + cost.total_msats,
|
||||
)
|
||||
)
|
||||
result = await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent and reserved_balance on the child key if it's different
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
total_spent=col(ApiKey.total_spent) + cost.total_msats,
|
||||
reserved_balance=child_safe_reserved,
|
||||
)
|
||||
.values(total_spent=col(ApiKey.total_spent) + cost.total_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
@@ -843,30 +584,10 @@ async def adjust_payment_for_tokens(
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"charged_amount": cost.total_msats,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
"new_balance": billing_key.balance,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await _accumulate_fee(cost.total_msats)
|
||||
payments_logger.info(
|
||||
"FINALIZE",
|
||||
extra={
|
||||
"event": "finalize",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"cost_reserved": deducted_max_cost,
|
||||
"cost_charged": cost.total_msats,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"finalize_type": "max_cost",
|
||||
},
|
||||
)
|
||||
return cost.dict()
|
||||
|
||||
case CostData() as cost:
|
||||
@@ -886,8 +607,6 @@ async def adjust_payment_for_tokens(
|
||||
"cost_difference": cost_difference,
|
||||
"input_msats": cost.input_msats,
|
||||
"output_msats": cost.output_msats,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -900,52 +619,24 @@ async def adjust_payment_for_tokens(
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
if billing_key.reserved_balance < deducted_max_cost:
|
||||
logger.error(
|
||||
"reserved_balance below deducted_max_cost on exact-cost finalization — clamping to 0",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"total_cost_msats": total_cost_msats,
|
||||
"balance": billing_key.balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
exact_safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=exact_safe_reserved,
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
balance=col(ApiKey.balance) - total_cost_msats,
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
)
|
||||
)
|
||||
await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent and reserved_balance on the child key if it's different
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_exact_safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
reserved_balance=child_exact_safe_reserved,
|
||||
)
|
||||
.values(total_spent=col(ApiKey.total_spent) + total_cost_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
@@ -953,67 +644,52 @@ async def adjust_payment_for_tokens(
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
await _accumulate_fee(total_cost_msats)
|
||||
payments_logger.info(
|
||||
"FINALIZE",
|
||||
extra={
|
||||
"event": "finalize",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"cost_reserved": deducted_max_cost,
|
||||
"cost_charged": total_cost_msats,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"finalize_type": "exact",
|
||||
},
|
||||
)
|
||||
return cost.dict()
|
||||
|
||||
# actual cost exceeded discounted reservation (due to tolerance_percentage)
|
||||
# this should never happen why do we handle this???
|
||||
if cost_difference > 0:
|
||||
# Always release the reservation and charge min(actual_cost, balance).
|
||||
# Using a CASE expression makes this a single atomic UPDATE — no
|
||||
# multi-level fallback needed and balance can never go negative.
|
||||
chargeable = case(
|
||||
(col(ApiKey.balance) >= total_cost_msats, total_cost_msats),
|
||||
else_=col(ApiKey.balance),
|
||||
# Need to charge more than reserved, finalize by releasing reservation and charging total
|
||||
logger.info(
|
||||
"Additional charge required for token usage",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"additional_charge": cost_difference,
|
||||
"current_balance": billing_key.balance,
|
||||
"sufficient_balance": billing_key.balance >= cost_difference,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
finalize_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
|
||||
balance=col(ApiKey.balance) - chargeable,
|
||||
total_spent=col(ApiKey.total_spent) + chargeable,
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
balance=col(ApiKey.balance) - total_cost_msats,
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
)
|
||||
)
|
||||
result = await session.exec(finalize_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.reserved_balance) >= deducted_max_cost)
|
||||
.values(
|
||||
reserved_balance=col(ApiKey.reserved_balance) - deducted_max_cost,
|
||||
total_spent=col(ApiKey.total_spent) + min(billing_key.balance, total_cost_msats),
|
||||
)
|
||||
.values(total_spent=col(ApiKey.total_spent) + total_cost_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
await session.commit()
|
||||
|
||||
if result.rowcount:
|
||||
cost.total_msats = total_cost_msats
|
||||
await session.refresh(billing_key)
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
await session.refresh(key)
|
||||
cost.total_msats = total_cost_msats
|
||||
|
||||
logger.info(
|
||||
"Finalized payment with additional charge",
|
||||
extra={
|
||||
@@ -1024,29 +700,9 @@ async def adjust_payment_for_tokens(
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await _accumulate_fee(total_cost_msats)
|
||||
payments_logger.info(
|
||||
"FINALIZE",
|
||||
extra={
|
||||
"event": "finalize",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"cost_reserved": deducted_max_cost,
|
||||
"cost_charged": total_cost_msats,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"finalize_type": "overrun",
|
||||
},
|
||||
)
|
||||
else:
|
||||
# Guard fired: reservation was already released by a concurrent
|
||||
# finalization for this key. Nothing left to do.
|
||||
logger.warning(
|
||||
"Finalization skipped - reservation already released",
|
||||
"Failed to finalize additional charge - releasing reservation",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
@@ -1054,6 +710,7 @@ async def adjust_payment_for_tokens(
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await release_reservation_only()
|
||||
else:
|
||||
# Refund some of the base cost
|
||||
refund = abs(cost_difference)
|
||||
@@ -1068,53 +725,24 @@ async def adjust_payment_for_tokens(
|
||||
},
|
||||
)
|
||||
|
||||
if billing_key.reserved_balance < deducted_max_cost:
|
||||
logger.error(
|
||||
"reserved_balance below deducted_max_cost on refund finalization — clamping to 0",
|
||||
extra={
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"deducted_max_cost": deducted_max_cost,
|
||||
"total_cost_msats": total_cost_msats,
|
||||
"refund_amount": refund,
|
||||
"balance": billing_key.balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
|
||||
refund_safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
|
||||
refund_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == billing_key.hashed_key)
|
||||
.values(
|
||||
reserved_balance=refund_safe_reserved,
|
||||
reserved_balance=col(ApiKey.reserved_balance)
|
||||
- deducted_max_cost,
|
||||
balance=col(ApiKey.balance) - total_cost_msats,
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
)
|
||||
)
|
||||
result = await session.exec(refund_stmt) # type: ignore[call-overload]
|
||||
|
||||
# Also update total_spent and reserved_balance on the child key if it's different
|
||||
# Also update total_spent on the child key if it's different
|
||||
if billing_key.hashed_key != key.hashed_key:
|
||||
child_refund_safe_reserved = case(
|
||||
(col(ApiKey.reserved_balance) >= deducted_max_cost,
|
||||
col(ApiKey.reserved_balance) - deducted_max_cost),
|
||||
else_=0,
|
||||
)
|
||||
child_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.values(
|
||||
total_spent=col(ApiKey.total_spent) + total_cost_msats,
|
||||
reserved_balance=child_refund_safe_reserved,
|
||||
)
|
||||
.values(total_spent=col(ApiKey.total_spent) + total_cost_msats)
|
||||
)
|
||||
await session.exec(child_stmt) # type: ignore[call-overload]
|
||||
|
||||
@@ -1150,25 +778,6 @@ async def adjust_payment_for_tokens(
|
||||
"model": model,
|
||||
},
|
||||
)
|
||||
await _accumulate_fee(total_cost_msats)
|
||||
payments_logger.info(
|
||||
"FINALIZE",
|
||||
extra={
|
||||
"event": "finalize",
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
"billing_key_hash": billing_key.hashed_key[:8] + "...",
|
||||
"model": model,
|
||||
"cost_reserved": deducted_max_cost,
|
||||
"cost_charged": total_cost_msats,
|
||||
"refunded": refund,
|
||||
"input_tokens": cost.input_tokens,
|
||||
"output_tokens": cost.output_tokens,
|
||||
"balance": billing_key.balance,
|
||||
"reserved_balance": billing_key.reserved_balance,
|
||||
"total_spent": billing_key.total_spent,
|
||||
"finalize_type": "refund",
|
||||
},
|
||||
)
|
||||
|
||||
return cost.dict()
|
||||
|
||||
@@ -1194,67 +803,15 @@ async def adjust_payment_for_tokens(
|
||||
}
|
||||
},
|
||||
)
|
||||
# All calculate_cost variants are handled above.
|
||||
raise AssertionError("Unreachable: unhandled calculate_cost result")
|
||||
|
||||
|
||||
async def periodic_key_reset() -> None:
|
||||
"""Background task to reset key limits based on their policy."""
|
||||
from .core.db import create_session
|
||||
|
||||
while True:
|
||||
try:
|
||||
interval = 3600 # Run every hour
|
||||
jitter = 300
|
||||
await asyncio.sleep(interval + random.uniform(0, jitter))
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
|
||||
try:
|
||||
async with create_session() as session:
|
||||
# Find all keys that have a reset policy
|
||||
stmt = select(ApiKey).where(ApiKey.balance_limit_reset.is_not(None)) # type: ignore
|
||||
keys = (await session.exec(stmt)).all()
|
||||
|
||||
now = int(time.time())
|
||||
updated_count = 0
|
||||
|
||||
for key in keys:
|
||||
reset_date = key.balance_limit_reset_date or 0
|
||||
should_reset = False
|
||||
|
||||
if key.balance_limit_reset == "daily":
|
||||
if (
|
||||
datetime.fromtimestamp(now).date()
|
||||
> datetime.fromtimestamp(reset_date).date()
|
||||
):
|
||||
should_reset = True
|
||||
elif key.balance_limit_reset == "weekly":
|
||||
if (
|
||||
datetime.fromtimestamp(now).isocalendar()[:2]
|
||||
> datetime.fromtimestamp(reset_date).isocalendar()[:2]
|
||||
):
|
||||
should_reset = True
|
||||
elif key.balance_limit_reset == "monthly":
|
||||
dt_now = datetime.fromtimestamp(now)
|
||||
dt_reset = datetime.fromtimestamp(reset_date)
|
||||
if dt_now.year > dt_reset.year or dt_now.month > dt_reset.month:
|
||||
should_reset = True
|
||||
|
||||
if should_reset:
|
||||
key.total_spent = 0
|
||||
key.balance_limit_reset_date = now
|
||||
session.add(key)
|
||||
updated_count += 1
|
||||
|
||||
if updated_count > 0:
|
||||
await session.commit()
|
||||
logger.info(
|
||||
"Periodic key reset complete",
|
||||
extra={"keys_reset": updated_count},
|
||||
)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(f"Error in periodic_key_reset: {e}")
|
||||
# Fallback: should not reach here, but release reservation just in case
|
||||
logger.error(
|
||||
"Unexpected fallback in adjust_payment_for_tokens - releasing reservation",
|
||||
extra={"key_hash": key.hashed_key[:8] + "...", "model": model},
|
||||
)
|
||||
await release_reservation_only()
|
||||
return {
|
||||
"base_msats": deducted_max_cost,
|
||||
"input_msats": 0,
|
||||
"output_msats": 0,
|
||||
"total_msats": deducted_max_cost,
|
||||
}
|
||||
|
||||
+23
-300
@@ -1,22 +1,13 @@
|
||||
import asyncio
|
||||
import hashlib
|
||||
import time
|
||||
from time import monotonic
|
||||
from typing import Annotated, NoReturn
|
||||
|
||||
from fastapi import APIRouter, Depends, Header, HTTPException
|
||||
from fastapi.responses import JSONResponse
|
||||
from pydantic import BaseModel
|
||||
from sqlmodel import col, select, update
|
||||
|
||||
from .auth import get_billing_key, validate_bearer_key
|
||||
from .core.db import (
|
||||
ApiKey,
|
||||
AsyncSession,
|
||||
CashuTransaction,
|
||||
get_session,
|
||||
store_cashu_transaction,
|
||||
)
|
||||
from .auth import validate_bearer_key
|
||||
from .core.db import ApiKey, AsyncSession, get_session
|
||||
from .core.logging import get_logger
|
||||
from .core.settings import settings
|
||||
from .lightning import lightning_router
|
||||
@@ -42,8 +33,10 @@ async def get_key_from_header(
|
||||
|
||||
|
||||
async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
|
||||
from .auth import get_billing_key
|
||||
|
||||
billing_key = await get_billing_key(key, session)
|
||||
info = {
|
||||
return {
|
||||
"api_key": "sk-" + key.hashed_key,
|
||||
"balance": billing_key.balance,
|
||||
"reserved": billing_key.reserved_balance,
|
||||
@@ -51,31 +44,8 @@ async def get_balance_info(key: ApiKey, session: AsyncSession) -> dict:
|
||||
"parent_key": "sk-" + key.parent_key_hash if key.parent_key_hash else None,
|
||||
"total_requests": key.total_requests,
|
||||
"total_spent": key.total_spent,
|
||||
"balance_limit": key.balance_limit,
|
||||
"balance_limit_reset": key.balance_limit_reset,
|
||||
"validity_date": key.validity_date,
|
||||
}
|
||||
|
||||
if not key.parent_key_hash:
|
||||
# Fetch child keys if this is a parent key
|
||||
statement = select(ApiKey).where(ApiKey.parent_key_hash == key.hashed_key)
|
||||
results = await session.exec(statement)
|
||||
child_keys = results.all()
|
||||
if child_keys:
|
||||
info["child_keys"] = [
|
||||
{
|
||||
"api_key": "sk-" + ck.hashed_key,
|
||||
"total_requests": ck.total_requests,
|
||||
"total_spent": ck.total_spent,
|
||||
"balance_limit": ck.balance_limit,
|
||||
"balance_limit_reset": ck.balance_limit_reset,
|
||||
"validity_date": ck.validity_date,
|
||||
}
|
||||
for ck in child_keys
|
||||
]
|
||||
|
||||
return info
|
||||
|
||||
|
||||
# TODO: remove this endpoint when frontend is updated
|
||||
@router.get("/", include_in_schema=False)
|
||||
@@ -100,24 +70,9 @@ async def account_info(
|
||||
|
||||
@router.get("/create")
|
||||
async def create_balance(
|
||||
initial_balance_token: str,
|
||||
balance_limit: int | None = None,
|
||||
balance_limit_reset: str | None = None,
|
||||
validity_date: int | None = None,
|
||||
session: AsyncSession = Depends(get_session),
|
||||
initial_balance_token: str, session: AsyncSession = Depends(get_session)
|
||||
) -> dict:
|
||||
key = await validate_bearer_key(initial_balance_token, session)
|
||||
|
||||
if balance_limit is not None or balance_limit_reset or validity_date:
|
||||
key.balance_limit = balance_limit
|
||||
key.balance_limit_reset = balance_limit_reset
|
||||
key.validity_date = validity_date
|
||||
if balance_limit_reset:
|
||||
key.balance_limit_reset_date = int(time.time())
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
await session.refresh(key)
|
||||
|
||||
return {
|
||||
"api_key": "sk-" + key.hashed_key,
|
||||
"balance": key.balance,
|
||||
@@ -143,6 +98,8 @@ async def topup_wallet_endpoint(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict[str, int]:
|
||||
from .auth import get_billing_key
|
||||
|
||||
billing_key = await get_billing_key(key, session)
|
||||
|
||||
if topup_request is not None:
|
||||
@@ -161,23 +118,9 @@ async def topup_wallet_endpoint(
|
||||
raise HTTPException(status_code=400, detail="Token already spent")
|
||||
elif "invalid" in error_msg.lower() or "decode" in error_msg.lower():
|
||||
raise HTTPException(status_code=400, detail="Invalid token format")
|
||||
elif "insufficient" in error_msg.lower() or "melt fee" in error_msg.lower():
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Token value is too small to cover swap fees. {error_msg}",
|
||||
)
|
||||
elif "failed to melt" in error_msg.lower():
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail=f"Failed to swap foreign mint token. {error_msg}",
|
||||
)
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail=f"Failed to redeem token: {error_msg}")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"topup_wallet_endpoint: unhandled error",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
raise HTTPException(status_code=400, detail="Failed to redeem token")
|
||||
except Exception:
|
||||
raise HTTPException(status_code=500, detail="Internal server error")
|
||||
return {"msats": amount_msats}
|
||||
|
||||
@@ -211,82 +154,23 @@ async def _refund_cache_set(authorization: str, value: dict[str, str]) -> None:
|
||||
_refund_cache[key] = (expiry, value)
|
||||
|
||||
|
||||
async def _restore_balance(
|
||||
session: AsyncSession, hashed_key: str, balance: int, reserved_balance: int
|
||||
) -> None:
|
||||
"""Restore balance after a failed refund mint attempt."""
|
||||
restore_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == hashed_key)
|
||||
.values(
|
||||
balance=col(ApiKey.balance) + balance,
|
||||
reserved_balance=col(ApiKey.reserved_balance) + reserved_balance,
|
||||
)
|
||||
)
|
||||
await session.exec(restore_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
logger.info(
|
||||
"refund_wallet_endpoint: balance restored after mint failure",
|
||||
extra={"hashed_key": hashed_key, "restored_balance": balance},
|
||||
)
|
||||
|
||||
|
||||
@router.post("/refund", response_model=None)
|
||||
@router.post("/refund")
|
||||
async def refund_wallet_endpoint(
|
||||
authorization: Annotated[str | None, Header()] = None,
|
||||
x_cashu: Annotated[str | None, Header()] = None,
|
||||
authorization: Annotated[str, Header(...)],
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> JSONResponse | dict[str, str]:
|
||||
if x_cashu:
|
||||
# Find the "in" transaction by the original payment token
|
||||
in_tx_result = await session.exec(
|
||||
select(CashuTransaction).where(
|
||||
CashuTransaction.token == x_cashu,
|
||||
CashuTransaction.type == "in",
|
||||
)
|
||||
)
|
||||
in_tx = in_tx_result.first()
|
||||
if in_tx is None:
|
||||
raise HTTPException(status_code=404, detail="Refund not found")
|
||||
|
||||
# Use the request_id to find the associated "out" (refund) transaction
|
||||
if in_tx.request_id is None:
|
||||
raise HTTPException(status_code=404, detail="Refund not found")
|
||||
|
||||
out_tx_result = await session.exec(
|
||||
select(CashuTransaction).where(
|
||||
CashuTransaction.request_id == in_tx.request_id,
|
||||
CashuTransaction.type == "out",
|
||||
)
|
||||
)
|
||||
out_tx = out_tx_result.first()
|
||||
if out_tx is None:
|
||||
raise HTTPException(status_code=404, detail="Refund not found")
|
||||
if out_tx.swept:
|
||||
raise HTTPException(status_code=410, detail="Refund has been swept")
|
||||
|
||||
out_tx.collected = True
|
||||
session.add(out_tx)
|
||||
await session.commit()
|
||||
body: dict[str, str] = {"token": out_tx.token}
|
||||
if out_tx.unit == "sat":
|
||||
body["sats"] = str(out_tx.amount)
|
||||
else:
|
||||
body["msats"] = str(out_tx.amount)
|
||||
return JSONResponse(content=body, headers={"X-Cashu": out_tx.token})
|
||||
|
||||
if authorization is None or not authorization.startswith("Bearer "):
|
||||
) -> dict[str, str]:
|
||||
if not authorization.startswith("Bearer "):
|
||||
raise HTTPException(
|
||||
status_code=401,
|
||||
detail="Invalid authorization. Use 'Bearer <cashu-token>' or 'Bearer <api-key>'",
|
||||
)
|
||||
|
||||
bearer_value: str = authorization[7:]
|
||||
key: ApiKey = await validate_bearer_key(bearer_value, session)
|
||||
|
||||
if key.total_balance <= 0:
|
||||
if cached := await _refund_cache_get(bearer_value):
|
||||
return cached
|
||||
if cached := await _refund_cache_get(bearer_value):
|
||||
return cached
|
||||
|
||||
key: ApiKey = await validate_bearer_key(bearer_value, session)
|
||||
|
||||
if key.parent_key_hash:
|
||||
raise HTTPException(
|
||||
@@ -294,12 +178,6 @@ async def refund_wallet_endpoint(
|
||||
detail="Cannot refund child key. Please refund the parent key instead.",
|
||||
)
|
||||
|
||||
if key.reserved_balance > 0:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Cannot refund key. There are ongoing requests for this api key.",
|
||||
)
|
||||
|
||||
remaining_balance_msats: int = key.total_balance
|
||||
|
||||
if key.refund_currency == "sat":
|
||||
@@ -312,31 +190,7 @@ async def refund_wallet_endpoint(
|
||||
elif remaining_balance <= 0:
|
||||
raise HTTPException(status_code=400, detail="No balance to refund")
|
||||
|
||||
# Capture values before debit — the session may refresh key after commit
|
||||
pre_debit_balance = key.balance
|
||||
pre_debit_reserved = key.reserved_balance
|
||||
|
||||
# --- DEBIT FIRST: atomically zero the balance before minting tokens ---
|
||||
# This prevents the race where a concurrent topup/spend happens between
|
||||
# reading the balance and minting the refund token (double-spend).
|
||||
debit_stmt = (
|
||||
update(ApiKey)
|
||||
.where(col(ApiKey.hashed_key) == key.hashed_key)
|
||||
.where(col(ApiKey.balance) == pre_debit_balance)
|
||||
.where(col(ApiKey.reserved_balance) == pre_debit_reserved)
|
||||
.values(balance=0, reserved_balance=0)
|
||||
)
|
||||
debit_result = await session.exec(debit_stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
|
||||
if debit_result.rowcount == 0:
|
||||
# Balance changed between read and debit — another request is active
|
||||
raise HTTPException(
|
||||
status_code=409,
|
||||
detail="Balance changed concurrently. Please retry the refund.",
|
||||
)
|
||||
|
||||
# --- MINT: balance is locked at zero, safe to create the refund token ---
|
||||
# Perform refund operation first, before modifying balance
|
||||
try:
|
||||
if key.refund_address:
|
||||
from .core.settings import settings as global_settings
|
||||
@@ -360,24 +214,11 @@ async def refund_wallet_endpoint(
|
||||
else:
|
||||
result["msats"] = str(remaining_balance_msats)
|
||||
|
||||
if "token" in result:
|
||||
logger.info(
|
||||
"refund_wallet_endpoint: cashu token issued",
|
||||
extra={
|
||||
"path": "/v1/wallet/refund",
|
||||
"token": result["token"],
|
||||
"amount": remaining_balance,
|
||||
"currency": key.refund_currency or "sat",
|
||||
},
|
||||
)
|
||||
|
||||
except HTTPException:
|
||||
# Minting failed — restore the debited balance
|
||||
await _restore_balance(session, key.hashed_key, pre_debit_balance, pre_debit_reserved)
|
||||
# Re-raise HTTP exceptions (like 400 for balance too small)
|
||||
raise
|
||||
except Exception as e:
|
||||
# Minting failed — restore the debited balance
|
||||
await _restore_balance(session, key.hashed_key, pre_debit_balance, pre_debit_reserved)
|
||||
# If refund fails, don't modify the database
|
||||
error_msg = str(e)
|
||||
if (
|
||||
"mint" in error_msg.lower()
|
||||
@@ -391,67 +232,12 @@ async def refund_wallet_endpoint(
|
||||
|
||||
await _refund_cache_set(bearer_value, result)
|
||||
|
||||
if "token" in result:
|
||||
try:
|
||||
await store_cashu_transaction(
|
||||
token=result["token"],
|
||||
amount=remaining_balance,
|
||||
unit=key.refund_currency or "sat",
|
||||
mint_url=key.refund_mint_url,
|
||||
typ="out",
|
||||
collected=False,
|
||||
source="apikey",
|
||||
api_key_hashed_key=key.hashed_key,
|
||||
)
|
||||
except Exception:
|
||||
pass # store_cashu_transaction already logs
|
||||
|
||||
logger.info(
|
||||
"refund_wallet_endpoint: refund successful",
|
||||
extra={
|
||||
"refunded_msats": remaining_balance_msats,
|
||||
"previous_reserved_balance": key.reserved_balance,
|
||||
},
|
||||
)
|
||||
await session.delete(key)
|
||||
await session.commit()
|
||||
|
||||
return result
|
||||
|
||||
|
||||
@router.get("/history")
|
||||
async def wallet_history(
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict[str, list[dict[str, str | int | bool | None]]]:
|
||||
if key.parent_key_hash:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail="Cannot view child key history. Please use the parent key instead.",
|
||||
)
|
||||
|
||||
result = await session.exec(
|
||||
select(CashuTransaction)
|
||||
.where(CashuTransaction.api_key_hashed_key == key.hashed_key)
|
||||
.order_by(col(CashuTransaction.created_at).desc())
|
||||
)
|
||||
transactions = result.all()
|
||||
return {
|
||||
"transactions": [
|
||||
{
|
||||
"id": tx.id,
|
||||
"type": tx.type,
|
||||
"source": tx.source,
|
||||
"amount": tx.amount,
|
||||
"unit": tx.unit,
|
||||
"mint_url": tx.mint_url,
|
||||
"created_at": tx.created_at,
|
||||
"collected": tx.collected,
|
||||
"swept": tx.swept,
|
||||
}
|
||||
for tx in transactions
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@router.post("/donate")
|
||||
async def donate(token: str, ref: str | None = None) -> str:
|
||||
try:
|
||||
@@ -467,9 +253,6 @@ async def donate(token: str, ref: str | None = None) -> str:
|
||||
|
||||
class ChildKeyRequest(BaseModel):
|
||||
count: int
|
||||
balance_limit: int | None = None
|
||||
balance_limit_reset: str | None = None
|
||||
validity_date: int | None = None
|
||||
|
||||
|
||||
@router.post("/child-key")
|
||||
@@ -519,12 +302,6 @@ async def create_child_key(
|
||||
hashed_key=new_key_hash,
|
||||
balance=0,
|
||||
parent_key_hash=key.hashed_key,
|
||||
balance_limit=payload.balance_limit,
|
||||
balance_limit_reset=payload.balance_limit_reset,
|
||||
balance_limit_reset_date=int(time.time())
|
||||
if payload.balance_limit_reset
|
||||
else None,
|
||||
validity_date=payload.validity_date,
|
||||
)
|
||||
session.add(child_key)
|
||||
new_keys.append("sk-" + new_key_hash)
|
||||
@@ -543,60 +320,6 @@ async def create_child_key(
|
||||
return response_data
|
||||
|
||||
|
||||
class ChildKeyResetRequest(BaseModel):
|
||||
child_key: str
|
||||
|
||||
|
||||
@router.post("/child-key/reset")
|
||||
async def reset_child_key_spent(
|
||||
payload: ChildKeyResetRequest,
|
||||
key: ApiKey = Depends(get_key_from_header),
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict:
|
||||
"""Resets the total_spent of a child key. Must be called by the parent."""
|
||||
child_key_raw = payload.child_key
|
||||
if child_key_raw.startswith("sk-"):
|
||||
child_key_raw = child_key_raw[3:]
|
||||
|
||||
child_key = await session.get(ApiKey, child_key_raw)
|
||||
if not child_key:
|
||||
raise HTTPException(status_code=404, detail="Child key not found.")
|
||||
|
||||
if child_key.parent_key_hash != key.hashed_key:
|
||||
raise HTTPException(
|
||||
status_code=403, detail="Unauthorized. You are not the parent of this key."
|
||||
)
|
||||
|
||||
child_key.total_spent = 0
|
||||
if child_key.balance_limit_reset:
|
||||
child_key.balance_limit_reset_date = int(time.time())
|
||||
session.add(child_key)
|
||||
await session.commit()
|
||||
|
||||
return {"success": True, "message": "Child key balance reset successfully."}
|
||||
|
||||
|
||||
@router.get("/cashu-refund/{payment_token_hash}")
|
||||
async def get_cashu_refund(
|
||||
payment_token_hash: str,
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict:
|
||||
"""Retrieve a stored Cashu refund token by the hash of the original payment token."""
|
||||
result = await session.get(CashuTransaction, payment_token_hash)
|
||||
if result is None:
|
||||
raise HTTPException(status_code=404, detail="Refund not found")
|
||||
if result.swept:
|
||||
raise HTTPException(status_code=410, detail="Refund has been swept")
|
||||
result.collected = True
|
||||
session.add(result)
|
||||
await session.commit()
|
||||
return {
|
||||
"refund_token": result.token,
|
||||
"amount": result.amount,
|
||||
"unit": result.unit,
|
||||
}
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/{path:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE"],
|
||||
|
||||
+2242
-627
File diff suppressed because it is too large
Load Diff
+8
-239
@@ -1,18 +1,13 @@
|
||||
import os
|
||||
import pathlib
|
||||
import sqlite3
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from alembic import command
|
||||
from alembic.config import Config
|
||||
from alembic.util.exc import CommandError
|
||||
from sqlalchemy import UniqueConstraint
|
||||
from sqlalchemy.exc import OperationalError
|
||||
from sqlalchemy.ext.asyncio.engine import create_async_engine
|
||||
from sqlmodel import Field, Relationship, SQLModel, col, func, select, update
|
||||
from sqlmodel import Field, Relationship, SQLModel, func, select, update
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from .logging import get_logger
|
||||
@@ -56,22 +51,6 @@ class ApiKey(SQLModel, table=True): # type: ignore
|
||||
parent_key_hash: str | None = Field(
|
||||
default=None, foreign_key="api_keys.hashed_key", index=True
|
||||
)
|
||||
balance_limit: int | None = Field(
|
||||
default=None,
|
||||
description="Max spendable balance in msats for this key (mostly for child keys)",
|
||||
)
|
||||
balance_limit_reset: str | None = Field(
|
||||
default=None,
|
||||
description="Reset policy for balance limit (manual, daily, monthly, etc.)",
|
||||
)
|
||||
balance_limit_reset_date: int | None = Field(
|
||||
default=None,
|
||||
description="Unix timestamp of the last time the balance limit was reset",
|
||||
)
|
||||
validity_date: int | None = Field(
|
||||
default=None,
|
||||
description="Unix timestamp after which the key is no longer valid",
|
||||
)
|
||||
|
||||
@property
|
||||
def total_balance(self) -> int:
|
||||
@@ -79,10 +58,11 @@ class ApiKey(SQLModel, table=True): # type: ignore
|
||||
|
||||
|
||||
async def reset_all_reserved_balances(session: AsyncSession) -> None:
|
||||
logger.info("Resetting all reserved balances to 0")
|
||||
stmt = update(ApiKey).values(reserved_balance=0)
|
||||
await session.exec(stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
logger.info("Reset reserved balances on startup")
|
||||
logger.info("Reserved balances reset successfully")
|
||||
|
||||
|
||||
class ModelRow(SQLModel, table=True): # type: ignore
|
||||
@@ -105,10 +85,6 @@ class ModelRow(SQLModel, table=True): # type: ignore
|
||||
default=None, description="JSON array of model alias IDs"
|
||||
)
|
||||
enabled: bool = Field(default=True, description="Whether this model is enabled")
|
||||
forwarded_model_id: str | None = Field(
|
||||
default=None,
|
||||
description="Model ID to use when forwarding requests to upstream provider. Defaults to id if not set.",
|
||||
)
|
||||
upstream_provider: "UpstreamProviderRow" = Relationship(back_populates="models")
|
||||
|
||||
|
||||
@@ -134,79 +110,10 @@ class LightningInvoice(SQLModel, table=True): # type: ignore
|
||||
paid_at: int | None = Field(default=None, description="Unix timestamp when paid")
|
||||
|
||||
|
||||
class CashuTransaction(SQLModel, table=True): # type: ignore
|
||||
__tablename__ = "cashu_transactions"
|
||||
|
||||
id: str = Field(
|
||||
primary_key=True,
|
||||
default_factory=lambda: uuid.uuid4().hex,
|
||||
description="Unique transaction identifier",
|
||||
)
|
||||
token: str = Field(description="Serialized Cashu token")
|
||||
amount: int = Field(description="Amount in the token's unit")
|
||||
unit: str = Field(description="Token unit (sat or msat)")
|
||||
mint_url: str | None = Field(default=None, description="Mint URL for the token")
|
||||
type: str = Field(default="out", description="Transaction type: in or out")
|
||||
request_id: str | None = Field(default=None, description="Associated request ID")
|
||||
created_at: int = Field(
|
||||
default_factory=lambda: int(time.time()),
|
||||
description="Unix timestamp",
|
||||
)
|
||||
collected: bool = Field(default=False)
|
||||
swept: bool = Field(default=False)
|
||||
source: str = Field(
|
||||
default="x-cashu",
|
||||
description="Payment source: x-cashu or apikey",
|
||||
)
|
||||
api_key_hashed_key: str | None = Field(
|
||||
default=None,
|
||||
foreign_key="api_keys.hashed_key",
|
||||
index=True,
|
||||
description="Associated API key hash for wallet history",
|
||||
)
|
||||
|
||||
|
||||
async def store_cashu_transaction(
|
||||
token: str,
|
||||
amount: int,
|
||||
unit: str,
|
||||
mint_url: str | None = None,
|
||||
typ: str = "out",
|
||||
request_id: str | None = None,
|
||||
collected: bool = False,
|
||||
created_at: int | None = None,
|
||||
source: str = "x-cashu",
|
||||
api_key_hashed_key: str | None = None,
|
||||
) -> None:
|
||||
try:
|
||||
async with create_session() as session:
|
||||
tx = CashuTransaction(
|
||||
token=token,
|
||||
amount=amount,
|
||||
unit=unit,
|
||||
mint_url=mint_url,
|
||||
type=typ,
|
||||
request_id=request_id,
|
||||
collected=collected,
|
||||
created_at=created_at or int(time.time()),
|
||||
source=source,
|
||||
api_key_hashed_key=api_key_hashed_key,
|
||||
)
|
||||
session.add(tx)
|
||||
await session.commit()
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
f"Failed to store cashu transaction: {e} (type={typ})",
|
||||
extra={"error": str(e), "type": typ},
|
||||
)
|
||||
|
||||
|
||||
class UpstreamProviderRow(SQLModel, table=True): # type: ignore
|
||||
__tablename__ = "upstream_providers"
|
||||
__table_args__ = (
|
||||
UniqueConstraint(
|
||||
"base_url", "api_key", name="uq_upstream_providers_base_url_api_key"
|
||||
),
|
||||
UniqueConstraint("base_url", "api_key", name="uq_upstream_providers_base_url_api_key"),
|
||||
)
|
||||
id: int | None = Field(default=None, primary_key=True)
|
||||
provider_type: str = Field(
|
||||
@@ -221,75 +128,12 @@ class UpstreamProviderRow(SQLModel, table=True): # type: ignore
|
||||
provider_fee: float = Field(
|
||||
default=1.01, description="Provider fee multiplier (default 1%)"
|
||||
)
|
||||
provider_settings: str | None = Field(
|
||||
default=None, description="JSON string for provider-specific settings"
|
||||
)
|
||||
models: list["ModelRow"] = Relationship(
|
||||
back_populates="upstream_provider",
|
||||
sa_relationship_kwargs={"cascade": "all, delete-orphan"},
|
||||
)
|
||||
|
||||
|
||||
class RoutstrFee(SQLModel, table=True): # type: ignore
|
||||
__tablename__ = "routstr_fees"
|
||||
id: int = Field(default=1, primary_key=True)
|
||||
accumulated_msats: int = Field(default=0)
|
||||
total_paid_msats: int = Field(default=0)
|
||||
last_paid_at: int | None = Field(default=None)
|
||||
|
||||
|
||||
class CliToken(SQLModel, table=True): # type: ignore
|
||||
"""Long-lived authorization token for CLI/agent use against admin endpoints."""
|
||||
|
||||
__tablename__ = "cli_tokens"
|
||||
id: str = Field(
|
||||
primary_key=True, default_factory=lambda: uuid.uuid4().hex
|
||||
)
|
||||
token: str = Field(unique=True, index=True, description="Bearer token value")
|
||||
name: str = Field(description="Human-readable label for this token")
|
||||
created_at: int = Field(default_factory=lambda: int(time.time()))
|
||||
last_used_at: int | None = Field(default=None)
|
||||
expires_at: int | None = Field(
|
||||
default=None, description="Optional expiry unix timestamp; null = never expires"
|
||||
)
|
||||
|
||||
|
||||
async def accumulate_routstr_fee(session: AsyncSession, amount_msats: int) -> None:
|
||||
stmt = (
|
||||
update(RoutstrFee)
|
||||
.where(col(RoutstrFee.id) == 1)
|
||||
.values(accumulated_msats=RoutstrFee.accumulated_msats + amount_msats)
|
||||
)
|
||||
result = await session.exec(stmt) # type: ignore[call-overload]
|
||||
if result.rowcount == 0:
|
||||
session.add(RoutstrFee(id=1, accumulated_msats=amount_msats))
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def get_routstr_fee(session: AsyncSession) -> RoutstrFee:
|
||||
fee = await session.get(RoutstrFee, 1)
|
||||
if fee is None:
|
||||
fee = RoutstrFee(id=1, accumulated_msats=0, total_paid_msats=0)
|
||||
session.add(fee)
|
||||
await session.commit()
|
||||
await session.refresh(fee)
|
||||
return fee
|
||||
|
||||
|
||||
async def reset_routstr_fee(session: AsyncSession, paid_msats: int) -> None:
|
||||
stmt = (
|
||||
update(RoutstrFee)
|
||||
.where(col(RoutstrFee.id) == 1)
|
||||
.values(
|
||||
accumulated_msats=RoutstrFee.accumulated_msats - paid_msats,
|
||||
total_paid_msats=RoutstrFee.total_paid_msats + paid_msats,
|
||||
last_paid_at=int(time.time()),
|
||||
)
|
||||
)
|
||||
await session.exec(stmt) # type: ignore[call-overload]
|
||||
await session.commit()
|
||||
|
||||
|
||||
async def balances_for_mint_and_unit(
|
||||
db_session: AsyncSession, mint_url: str, unit: str
|
||||
) -> int:
|
||||
@@ -319,64 +163,11 @@ async def create_session() -> AsyncGenerator[AsyncSession, None]:
|
||||
yield session
|
||||
|
||||
|
||||
def fix_cashu_migrations() -> None:
|
||||
"""
|
||||
Fixes Cashu wallet migrations that are not idempotent.
|
||||
This specifically addresses the 'duplicate column name: public_keys' error
|
||||
in the keysets table of Cashu's internal SQLite databases.
|
||||
"""
|
||||
project_root = pathlib.Path(__file__).resolve().parents[2]
|
||||
wallet_dir = project_root / ".wallet"
|
||||
|
||||
if not wallet_dir.exists() or not wallet_dir.is_dir():
|
||||
return
|
||||
|
||||
logger.info("Checking Cashu wallet databases for migration idempotency")
|
||||
|
||||
for db_file in wallet_dir.glob("*.sqlite3"):
|
||||
try:
|
||||
conn = sqlite3.connect(db_file)
|
||||
cursor = conn.cursor()
|
||||
|
||||
# Check if keysets table exists
|
||||
cursor.execute(
|
||||
"SELECT name FROM sqlite_master WHERE type='table' AND name='keysets'"
|
||||
)
|
||||
if not cursor.fetchone():
|
||||
conn.close()
|
||||
continue
|
||||
|
||||
# Check if public_keys column exists
|
||||
cursor.execute("PRAGMA table_info(keysets)")
|
||||
columns = [info[1] for info in cursor.fetchall()]
|
||||
|
||||
if "public_keys" not in columns:
|
||||
logger.info(f"Adding missing public_keys column to {db_file.name}")
|
||||
cursor.execute("ALTER TABLE keysets ADD COLUMN public_keys TEXT")
|
||||
conn.commit()
|
||||
|
||||
conn.close()
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not check/fix Cashu database {db_file}: {e}")
|
||||
|
||||
|
||||
def _clear_alembic_version() -> None:
|
||||
"""Clear the alembic_version table so stamp/upgrade can proceed."""
|
||||
sync_url = DATABASE_URL.replace("+aiosqlite", "")
|
||||
from sqlalchemy import create_engine, text
|
||||
|
||||
eng = create_engine(sync_url)
|
||||
with eng.begin() as conn:
|
||||
conn.execute(text("DELETE FROM alembic_version"))
|
||||
eng.dispose()
|
||||
|
||||
|
||||
def run_migrations() -> None:
|
||||
"""Run Alembic migrations programmatically."""
|
||||
try:
|
||||
# Run Cashu migration fix first
|
||||
fix_cashu_migrations()
|
||||
import pathlib
|
||||
|
||||
try:
|
||||
# Get the path to the alembic.ini file
|
||||
project_root = pathlib.Path(__file__).resolve().parents[2]
|
||||
alembic_ini_path = project_root / "alembic.ini"
|
||||
@@ -392,30 +183,8 @@ def run_migrations() -> None:
|
||||
# Set the database URL in the config
|
||||
alembic_cfg.set_main_option("sqlalchemy.url", DATABASE_URL)
|
||||
|
||||
try:
|
||||
command.upgrade(alembic_cfg, "head")
|
||||
except CommandError as e:
|
||||
if "Can't locate revision" in str(e):
|
||||
logger.warning(
|
||||
"Database stamped with unknown revision (likely from another branch). "
|
||||
"Re-stamping to current head.",
|
||||
extra={"error": str(e)},
|
||||
)
|
||||
_clear_alembic_version()
|
||||
command.stamp(alembic_cfg, "head")
|
||||
else:
|
||||
raise
|
||||
except OperationalError as e:
|
||||
if "duplicate column name" in str(e).lower():
|
||||
logger.warning(
|
||||
"Migration hit a column that already exists (likely added via "
|
||||
"create_all on another branch). Stamping to current head.",
|
||||
extra={"error": str(e)},
|
||||
)
|
||||
_clear_alembic_version()
|
||||
command.stamp(alembic_cfg, "head")
|
||||
else:
|
||||
raise
|
||||
# Run migrations to the latest revision
|
||||
command.upgrade(alembic_cfg, "head")
|
||||
|
||||
logger.info("Database migrations completed successfully")
|
||||
|
||||
|
||||
@@ -22,20 +22,16 @@ async def http_exception_handler(request: Request, exc: Exception) -> JSONRespon
|
||||
# Get status code and detail - works for both FastAPI and Starlette HTTPException
|
||||
status_code = getattr(exc, "status_code", 500)
|
||||
detail = getattr(exc, "detail", str(exc))
|
||||
path = request.url.path
|
||||
|
||||
# 4xx is client behaviour; the uvicorn access log already records it.
|
||||
# Only 5xx warrants a server-side warning/error log here.
|
||||
if status_code >= 500:
|
||||
logger.error(
|
||||
f"HTTP {status_code} on {path}: {detail}",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"status_code": status_code,
|
||||
"detail": detail,
|
||||
"path": path,
|
||||
},
|
||||
)
|
||||
logger.warning(
|
||||
"HTTP exception",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"status_code": status_code,
|
||||
"detail": detail,
|
||||
"path": request.url.path,
|
||||
},
|
||||
)
|
||||
|
||||
return JSONResponse(
|
||||
status_code=status_code,
|
||||
|
||||
+221
-830
File diff suppressed because it is too large
Load Diff
+20
-47
@@ -3,61 +3,50 @@ Logging configuration for Routstr.
|
||||
|
||||
CRITICAL LOG MESSAGES FOR USAGE STATISTICS:
|
||||
===========================================
|
||||
The following log messages are parsed by the usage tracking system
|
||||
(routstr/core/usage_analytics_store.py and routstr/core/log_manager.py).
|
||||
The following log messages are parsed by the usage tracking system (routstr/core/admin.py).
|
||||
DO NOT modify or remove these messages without updating the usage tracking logic:
|
||||
|
||||
1. "Received proxy request" (INFO) - routstr/proxy.py
|
||||
- Used to count total incoming requests
|
||||
- Includes model information in context
|
||||
|
||||
2. "Calculated token-based cost" (INFO) - routstr/auth.py
|
||||
2. "Payment adjustment completed for streaming" (INFO) - routstr/upstream/base.py
|
||||
"Payment adjustment completed for non-streaming" (INFO) - routstr/upstream/base.py
|
||||
- Used to track successful completions and revenue
|
||||
- The 'token_cost', 'model', 'input_tokens', and 'output_tokens' fields are extracted for dashboard metrics
|
||||
- The 'cost_data.total_msats' field is extracted for revenue calculation
|
||||
- Must include 'cost_data' in extra dict
|
||||
|
||||
3. "Max cost payment finalized" (INFO) - routstr/auth.py
|
||||
- Used as the successful completion fallback when token usage is unavailable
|
||||
- The 'charged_amount', 'model', 'input_tokens', and 'output_tokens' fields are extracted for dashboard metrics
|
||||
|
||||
4. "Payment processed successfully" (INFO) - routstr/auth.py
|
||||
3. "Payment processed successfully" (INFO) - routstr/auth.py
|
||||
- Used to count successful payment processing events
|
||||
- Tracks payment-related metrics
|
||||
|
||||
5. "Upstream request failed, revert payment" (WARNING) - routstr/proxy.py
|
||||
4. "Upstream request failed, revert payment" (WARNING) - routstr/proxy.py
|
||||
- Used to track failed requests and refunds
|
||||
- The 'max_cost_for_model' field is extracted for refund calculation
|
||||
- Must include 'max_cost_for_model' in extra dict
|
||||
|
||||
6. Any ERROR level logs with "upstream" in the message
|
||||
5. Any ERROR level logs with "upstream" in the message
|
||||
- Used to count upstream provider errors
|
||||
- Helps identify service reliability issues
|
||||
|
||||
If you need to modify these messages, ensure you also update the parsing logic in:
|
||||
- routstr/core/usage_analytics_store.py
|
||||
- routstr/core/log_manager.py
|
||||
- routstr/core/admin.py:_aggregate_metrics_by_time()
|
||||
- routstr/core/admin.py:_get_summary_stats()
|
||||
- routstr/core/admin.py:get_revenue_by_model()
|
||||
"""
|
||||
|
||||
import logging.config
|
||||
import logging.handlers
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import tomllib
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from pythonjsonlogger import jsonlogger
|
||||
from rich.console import Console
|
||||
from rich.logging import RichHandler
|
||||
|
||||
# Only use RichHandler when stdout is a real TTY. In non-TTY contexts
|
||||
# (docker logs, pipes, CI) Rich pads every line to width and wraps long
|
||||
# records, producing visually-empty trailing whitespace and split records.
|
||||
# A plain StreamHandler avoids both problems.
|
||||
_stdout_is_tty = sys.stdout.isatty()
|
||||
_console = Console(soft_wrap=True) if _stdout_is_tty else None
|
||||
|
||||
# Define custom TRACE level
|
||||
TRACE_LEVEL = 5
|
||||
logging.addLevelName(TRACE_LEVEL, "TRACE")
|
||||
@@ -270,26 +259,6 @@ def setup_logging() -> None:
|
||||
if console_enabled:
|
||||
handlers.append("console")
|
||||
|
||||
if _stdout_is_tty:
|
||||
console_handler: dict[str, Any] = {
|
||||
"()": RichHandler,
|
||||
"level": log_level,
|
||||
"show_time": False,
|
||||
"show_path": False,
|
||||
"rich_tracebacks": True,
|
||||
"markup": True,
|
||||
"console": _console,
|
||||
"filters": ["request_id_filter", "security_filter"],
|
||||
}
|
||||
else:
|
||||
console_handler = {
|
||||
"class": "logging.StreamHandler",
|
||||
"level": log_level,
|
||||
"formatter": "plain",
|
||||
"stream": "ext://sys.stdout",
|
||||
"filters": ["request_id_filter", "security_filter"],
|
||||
}
|
||||
|
||||
LOGGING_CONFIG = {
|
||||
"version": 1,
|
||||
"disable_existing_loggers": False,
|
||||
@@ -299,10 +268,6 @@ def setup_logging() -> None:
|
||||
"format": "%(asctime)s %(name)s %(levelname)s %(message)s %(pathname)s %(lineno)d %(version)s %(request_id)s",
|
||||
"datefmt": "%Y-%m-%d %H:%M:%S",
|
||||
},
|
||||
"plain": {
|
||||
"format": "%(asctime)s %(levelname)-7s %(name)s %(message)s",
|
||||
"datefmt": "%Y-%m-%d %H:%M:%S",
|
||||
},
|
||||
},
|
||||
"filters": {
|
||||
"version_filter": {"()": VersionFilter},
|
||||
@@ -310,7 +275,15 @@ def setup_logging() -> None:
|
||||
"security_filter": {"()": SecurityFilter},
|
||||
},
|
||||
"handlers": {
|
||||
"console": console_handler,
|
||||
"console": {
|
||||
"()": RichHandler,
|
||||
"level": log_level,
|
||||
"show_time": False,
|
||||
"show_path": False,
|
||||
"rich_tracebacks": True,
|
||||
"markup": True,
|
||||
"filters": ["request_id_filter", "security_filter"],
|
||||
},
|
||||
"file": {
|
||||
"()": DailyRotatingFileHandler,
|
||||
"level": log_level,
|
||||
|
||||
+14
-305
@@ -1,219 +1,30 @@
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
from pathlib import Path
|
||||
from typing import AsyncGenerator
|
||||
import os
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.middleware.cors import CORSMiddleware
|
||||
from fastapi.responses import FileResponse, RedirectResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
from fastapi.responses import RedirectResponse
|
||||
from starlette.exceptions import HTTPException
|
||||
from starlette.responses import Response as StarletteResponse
|
||||
from starlette.types import Scope
|
||||
|
||||
from ..auth import periodic_key_reset
|
||||
from ..balance import balance_router, deprecated_wallet_router
|
||||
from ..nostr import (
|
||||
announce_provider,
|
||||
providers_cache_refresher,
|
||||
publish_usage_analytics,
|
||||
)
|
||||
from ..nostr.discovery import providers_router
|
||||
from ..payment.models import models_router, update_sats_pricing
|
||||
from ..payment.price import update_prices_periodically
|
||||
from ..proxy import initialize_upstreams, proxy_router, refresh_model_maps_periodically
|
||||
from ..upstream.auto_topup import periodic_auto_topup
|
||||
from ..upstream.litellm_routing import configure_litellm
|
||||
from ..wallet import periodic_payout, periodic_refund_sweep, periodic_routstr_fee_payout
|
||||
from ..discovery import providers_router
|
||||
from ..payment.models import models_router
|
||||
from ..proxy import proxy_router
|
||||
from .admin import admin_router
|
||||
from .db import create_session, init_db, run_migrations
|
||||
from .exceptions import general_exception_handler, http_exception_handler
|
||||
from .logging import get_logger, setup_logging
|
||||
from .middleware import LoggingMiddleware
|
||||
from .settings import SettingsService
|
||||
from .settings import settings as global_settings
|
||||
from .version import __version__
|
||||
from .tasks import lifespan
|
||||
from .ui import setup_ui
|
||||
|
||||
# Initialize logging first
|
||||
setup_logging()
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(_: FastAPI) -> AsyncGenerator[None, None]:
|
||||
logger.info("Application startup initiated", extra={"version": __version__})
|
||||
|
||||
btc_price_task = None
|
||||
pricing_task = None
|
||||
payout_task = None
|
||||
nip91_task = None
|
||||
analytics_task = None
|
||||
providers_task = None
|
||||
models_refresh_task = None
|
||||
model_maps_refresh_task = None
|
||||
key_reset_task = None
|
||||
auto_topup_task = None
|
||||
refund_sweep_task = None
|
||||
routstr_fee_task = None
|
||||
|
||||
try:
|
||||
# Apply litellm-wide settings (drop_params, chat-completions URL,
|
||||
# debug logging) before any upstream provider dispatches a request.
|
||||
configure_litellm()
|
||||
|
||||
# Run database migrations on startup
|
||||
run_migrations()
|
||||
|
||||
# Initialize database connection pools
|
||||
# This creates any tables that might not be tracked by migrations yet
|
||||
await init_db()
|
||||
|
||||
# Initialize application settings (env -> computed -> DB precedence)
|
||||
async with create_session() as session:
|
||||
s = await SettingsService.initialize(session)
|
||||
if s.reset_reserved_balance_on_startup:
|
||||
from .db import reset_all_reserved_balances
|
||||
|
||||
await reset_all_reserved_balances(session)
|
||||
|
||||
if not s.admin_password:
|
||||
logger.warning(
|
||||
f"Admin password is not set. Visit {s.http_url or 'http://localhost:8000'}/admin to set the password."
|
||||
)
|
||||
|
||||
# Apply app metadata from settings
|
||||
try:
|
||||
app.title = s.name
|
||||
app.description = s.description
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# await ensure_models_bootstrapped()
|
||||
|
||||
from ..payment.price import _update_prices
|
||||
from ..proxy import get_upstreams
|
||||
from ..upstream.helpers import refresh_upstreams_models_periodically
|
||||
|
||||
_update_prices_task = asyncio.create_task(_update_prices())
|
||||
_initialize_upstreams_task = asyncio.create_task(initialize_upstreams())
|
||||
|
||||
# ensure both setup tasks complete
|
||||
await asyncio.gather(
|
||||
_update_prices_task, _initialize_upstreams_task, return_exceptions=True
|
||||
)
|
||||
|
||||
btc_price_task = asyncio.create_task(update_prices_periodically())
|
||||
pricing_task = asyncio.create_task(update_sats_pricing())
|
||||
if global_settings.models_refresh_interval_seconds > 0:
|
||||
# Pass the accessor (not its current value) so the loop sees providers
|
||||
# added/changed via reinitialize_upstreams() instead of staying pinned
|
||||
# to the startup snapshot.
|
||||
models_refresh_task = asyncio.create_task(
|
||||
refresh_upstreams_models_periodically(get_upstreams)
|
||||
)
|
||||
model_maps_refresh_task = asyncio.create_task(refresh_model_maps_periodically())
|
||||
payout_task = asyncio.create_task(periodic_payout())
|
||||
if global_settings.nsec:
|
||||
nip91_task = asyncio.create_task(announce_provider())
|
||||
analytics_task = asyncio.create_task(publish_usage_analytics())
|
||||
if global_settings.providers_refresh_interval_seconds > 0:
|
||||
providers_task = asyncio.create_task(providers_cache_refresher())
|
||||
key_reset_task = asyncio.create_task(periodic_key_reset())
|
||||
auto_topup_task = asyncio.create_task(periodic_auto_topup())
|
||||
refund_sweep_task = asyncio.create_task(periodic_refund_sweep())
|
||||
routstr_fee_task = asyncio.create_task(periodic_routstr_fee_payout())
|
||||
|
||||
yield
|
||||
|
||||
except asyncio.CancelledError:
|
||||
# Expected during shutdown
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Application startup failed",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
logger.info("Application shutdown initiated")
|
||||
|
||||
if btc_price_task is not None:
|
||||
btc_price_task.cancel()
|
||||
if pricing_task is not None:
|
||||
pricing_task.cancel()
|
||||
if payout_task is not None:
|
||||
payout_task.cancel()
|
||||
if nip91_task is not None:
|
||||
nip91_task.cancel()
|
||||
if analytics_task is not None:
|
||||
analytics_task.cancel()
|
||||
if providers_task is not None:
|
||||
providers_task.cancel()
|
||||
if models_refresh_task is not None:
|
||||
models_refresh_task.cancel()
|
||||
if model_maps_refresh_task is not None:
|
||||
model_maps_refresh_task.cancel()
|
||||
if key_reset_task is not None:
|
||||
key_reset_task.cancel()
|
||||
if auto_topup_task is not None:
|
||||
auto_topup_task.cancel()
|
||||
if refund_sweep_task is not None:
|
||||
refund_sweep_task.cancel()
|
||||
if routstr_fee_task is not None:
|
||||
routstr_fee_task.cancel()
|
||||
|
||||
try:
|
||||
tasks_to_wait = []
|
||||
if btc_price_task is not None:
|
||||
tasks_to_wait.append(btc_price_task)
|
||||
if pricing_task is not None:
|
||||
tasks_to_wait.append(pricing_task)
|
||||
if payout_task is not None:
|
||||
tasks_to_wait.append(payout_task)
|
||||
if nip91_task is not None:
|
||||
tasks_to_wait.append(nip91_task)
|
||||
if analytics_task is not None:
|
||||
tasks_to_wait.append(analytics_task)
|
||||
if providers_task is not None:
|
||||
tasks_to_wait.append(providers_task)
|
||||
if models_refresh_task is not None:
|
||||
tasks_to_wait.append(models_refresh_task)
|
||||
if model_maps_refresh_task is not None:
|
||||
tasks_to_wait.append(model_maps_refresh_task)
|
||||
if key_reset_task is not None:
|
||||
tasks_to_wait.append(key_reset_task)
|
||||
if auto_topup_task is not None:
|
||||
tasks_to_wait.append(auto_topup_task)
|
||||
if refund_sweep_task is not None:
|
||||
tasks_to_wait.append(refund_sweep_task)
|
||||
if routstr_fee_task is not None:
|
||||
tasks_to_wait.append(routstr_fee_task)
|
||||
|
||||
if tasks_to_wait:
|
||||
await asyncio.gather(*tasks_to_wait, return_exceptions=True)
|
||||
logger.info("Background tasks stopped successfully")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error stopping background tasks",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
|
||||
|
||||
class _ImmutableStaticFiles(StaticFiles):
|
||||
"""Static files with long Cache-Control for content-hashed Next.js assets.
|
||||
|
||||
Files under `/_next/static/` are emitted with content hashes in their
|
||||
filenames and never mutate, so we serve them with a one-year immutable
|
||||
cache header so browsers and CDNs stop revalidating on every reload.
|
||||
"""
|
||||
|
||||
async def get_response(self, path: str, scope: Scope) -> StarletteResponse:
|
||||
response = await super().get_response(path, scope)
|
||||
if response.status_code == 200:
|
||||
response.headers["Cache-Control"] = (
|
||||
"public, max-age=31536000, immutable"
|
||||
)
|
||||
return response
|
||||
if os.getenv("VERSION_SUFFIX") is not None:
|
||||
__version__ = f"0.3.0-{os.getenv('VERSION_SUFFIX')}"
|
||||
else:
|
||||
__version__ = "0.3.0"
|
||||
|
||||
|
||||
app = FastAPI(version=__version__, lifespan=lifespan)
|
||||
@@ -225,7 +36,7 @@ app.add_middleware(
|
||||
allow_credentials=True,
|
||||
allow_methods=["*"],
|
||||
allow_headers=["*"],
|
||||
expose_headers=["x-routstr-request-id", "x-cashu"],
|
||||
expose_headers=["x-routstr-request-id"],
|
||||
)
|
||||
|
||||
# Add logging middleware
|
||||
@@ -255,110 +66,8 @@ async def providers() -> RedirectResponse:
|
||||
return RedirectResponse("/v1/providers/")
|
||||
|
||||
|
||||
UI_DIST_PATH = Path(__file__).parent.parent.parent / "ui_out"
|
||||
|
||||
if UI_DIST_PATH.exists() and UI_DIST_PATH.is_dir():
|
||||
logger.info(f"Serving static UI from {UI_DIST_PATH}")
|
||||
|
||||
app.mount(
|
||||
"/_next",
|
||||
_ImmutableStaticFiles(directory=UI_DIST_PATH / "_next", check_dir=True),
|
||||
name="next-static",
|
||||
)
|
||||
|
||||
@app.get("/", include_in_schema=False)
|
||||
async def serve_root_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "index.html")
|
||||
|
||||
# Serve the App Router RSC payload for the home page.
|
||||
@app.get("/index.txt", include_in_schema=False)
|
||||
async def serve_root_rsc() -> FileResponse:
|
||||
return FileResponse(
|
||||
UI_DIST_PATH / "index.txt", media_type="text/x-component"
|
||||
)
|
||||
|
||||
# Next.js is built with `trailingSlash: true`, so all UI page URLs end
|
||||
# with a slash (e.g. `/login/`). The proxy router catches `/{path:path}`
|
||||
# before FastAPI's `redirect_slashes` logic can normalize the URL, so we
|
||||
# must register both the with-slash and without-slash variants here.
|
||||
UI_PAGES = (
|
||||
"dashboard",
|
||||
"login",
|
||||
"model",
|
||||
"providers",
|
||||
"settings",
|
||||
"transactions",
|
||||
"balances",
|
||||
"logs",
|
||||
"usage",
|
||||
"unauthorized",
|
||||
)
|
||||
|
||||
def _register_ui_page(name: str) -> None:
|
||||
page_dir = UI_DIST_PATH / name
|
||||
index_html = page_dir / "index.html"
|
||||
index_txt = page_dir / "index.txt"
|
||||
|
||||
async def serve_page() -> FileResponse:
|
||||
return FileResponse(index_html)
|
||||
|
||||
async def serve_page_rsc() -> FileResponse:
|
||||
return FileResponse(index_txt, media_type="text/x-component")
|
||||
|
||||
app.add_api_route(
|
||||
f"/{name}",
|
||||
serve_page,
|
||||
methods=["GET"],
|
||||
include_in_schema=False,
|
||||
name=f"serve_{name}_ui",
|
||||
)
|
||||
app.add_api_route(
|
||||
f"/{name}/",
|
||||
serve_page,
|
||||
methods=["GET"],
|
||||
include_in_schema=False,
|
||||
name=f"serve_{name}_ui_slash",
|
||||
)
|
||||
app.add_api_route(
|
||||
f"/{name}/index.txt",
|
||||
serve_page_rsc,
|
||||
methods=["GET"],
|
||||
include_in_schema=False,
|
||||
name=f"serve_{name}_rsc",
|
||||
)
|
||||
|
||||
for _page in UI_PAGES:
|
||||
_register_ui_page(_page)
|
||||
|
||||
@app.get("/admin")
|
||||
async def admin_redirect() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "index.html")
|
||||
|
||||
@app.get("/favicon.ico", include_in_schema=False)
|
||||
async def serve_favicon() -> FileResponse:
|
||||
icon_path = UI_DIST_PATH / "icon.ico"
|
||||
if icon_path.exists():
|
||||
return FileResponse(icon_path)
|
||||
return FileResponse(UI_DIST_PATH / "favicon.ico")
|
||||
|
||||
@app.get("/icon.ico", include_in_schema=False)
|
||||
async def serve_icon() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "icon.ico")
|
||||
|
||||
else:
|
||||
logger.warning(
|
||||
f"UI dist directory not found at {UI_DIST_PATH}, skipping static file serving"
|
||||
)
|
||||
|
||||
@app.get("/", include_in_schema=False)
|
||||
async def root_fallback() -> dict:
|
||||
return {
|
||||
"name": global_settings.name,
|
||||
"description": global_settings.description,
|
||||
"version": __version__,
|
||||
"status": "running",
|
||||
"ui": "not available",
|
||||
}
|
||||
# Setup UI serving
|
||||
setup_ui(app)
|
||||
|
||||
|
||||
app.include_router(models_router)
|
||||
|
||||
+71
-69
@@ -14,54 +14,8 @@ logger = get_logger(__name__)
|
||||
request_id_context: ContextVar[str | None] = ContextVar("request_id")
|
||||
|
||||
|
||||
# Methods that are never logged: HEAD requests are health probes from
|
||||
# monitoring/load balancers, OPTIONS are CORS preflights — both are framework
|
||||
# chatter, not user-meaningful events.
|
||||
_SKIP_LOG_METHODS: frozenset[str] = frozenset({"HEAD", "OPTIONS"})
|
||||
|
||||
# Path prefixes to skip. Includes Next.js static chunks and the admin
|
||||
# dashboard's internal polling API (/admin/api/*) which the UI hits on a timer
|
||||
# to refresh balances, logs, providers, etc. — high volume, low diagnostic
|
||||
# value. Mutating admin actions are recorded separately in the audit log.
|
||||
_SKIP_LOG_PREFIXES: tuple[str, ...] = (
|
||||
"/_next/",
|
||||
"/admin/api/",
|
||||
)
|
||||
|
||||
# Exact paths to skip. RSC payload prefetches (`*/index.txt`) fire automatically
|
||||
# as the user hovers near `<Link>`s, and `/v1/wallet/info` is polled by the UI.
|
||||
_SKIP_LOG_EXACT: frozenset[str] = frozenset(
|
||||
{
|
||||
"/favicon.ico",
|
||||
"/icon.ico",
|
||||
"/v1/wallet/info",
|
||||
"/index.txt",
|
||||
"/login/index.txt",
|
||||
"/model/index.txt",
|
||||
"/providers/index.txt",
|
||||
"/settings/index.txt",
|
||||
"/transactions/index.txt",
|
||||
"/balances/index.txt",
|
||||
"/logs/index.txt",
|
||||
"/usage/index.txt",
|
||||
"/unauthorized/index.txt",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _should_log(method: str, path: str) -> bool:
|
||||
if method in _SKIP_LOG_METHODS:
|
||||
return False
|
||||
if path in _SKIP_LOG_EXACT:
|
||||
return False
|
||||
return not any(path.startswith(prefix) for prefix in _SKIP_LOG_PREFIXES)
|
||||
|
||||
|
||||
class LoggingMiddleware(BaseHTTPMiddleware):
|
||||
"""Middleware to log proxy interactions and page navigation.
|
||||
|
||||
Skips logging for static assets and Next.js chunks to avoid noise.
|
||||
"""
|
||||
"""Middleware to log detailed request and response information."""
|
||||
|
||||
async def dispatch(self, request: Request, call_next: Callable) -> Response:
|
||||
# Generate request ID
|
||||
@@ -71,20 +25,62 @@ class LoggingMiddleware(BaseHTTPMiddleware):
|
||||
# Set request ID in context for logging
|
||||
token = request_id_context.set(request_id)
|
||||
|
||||
path = request.url.path
|
||||
should_log = _should_log(request.method, path)
|
||||
|
||||
# Start timing
|
||||
start_time = time.time()
|
||||
|
||||
if should_log:
|
||||
logger.info(
|
||||
"Incoming request",
|
||||
# Log request details
|
||||
request_body = None
|
||||
if request.method in ["POST", "PUT", "PATCH"]:
|
||||
try:
|
||||
# Only read body for non-streaming requests
|
||||
if hasattr(request, "_body"):
|
||||
request_body = await request.body()
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Extract request info
|
||||
client_host = None
|
||||
if request.client:
|
||||
client_host = request.client.host
|
||||
|
||||
# Log incoming request
|
||||
logger.info(
|
||||
"Incoming request",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"method": request.method,
|
||||
"path": request.url.path,
|
||||
"query_params": dict(request.query_params),
|
||||
"client_host": client_host,
|
||||
"headers": {
|
||||
k: v
|
||||
for k, v in request.headers.items()
|
||||
if k.lower()
|
||||
not in [
|
||||
"authorization",
|
||||
"x-cashu",
|
||||
"cookie",
|
||||
"cf-connecting-ip",
|
||||
"cf-ipcountry",
|
||||
"x-forwarded-for",
|
||||
"x-real-ip",
|
||||
]
|
||||
},
|
||||
"body_size": len(request_body) if request_body else 0,
|
||||
},
|
||||
)
|
||||
|
||||
# Log at TRACE level for full body (security filter will redact sensitive data)
|
||||
if request_body and hasattr(logger, "exception"):
|
||||
logger.exception(
|
||||
"Request body",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"method": request.method,
|
||||
"path": path,
|
||||
"query_params": dict(request.query_params),
|
||||
"path": request.url.path,
|
||||
"body": request_body.decode("utf-8", errors="ignore")[
|
||||
:1000
|
||||
], # Limit size
|
||||
},
|
||||
)
|
||||
|
||||
@@ -92,33 +88,39 @@ class LoggingMiddleware(BaseHTTPMiddleware):
|
||||
try:
|
||||
response = await call_next(request)
|
||||
|
||||
if should_log:
|
||||
duration = time.time() - start_time
|
||||
logger.info(
|
||||
"Request completed",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"method": request.method,
|
||||
"path": path,
|
||||
"status_code": response.status_code,
|
||||
"duration_ms": round(duration * 1000, 2),
|
||||
},
|
||||
)
|
||||
# Calculate duration
|
||||
duration = time.time() - start_time
|
||||
|
||||
# Log response
|
||||
logger.info(
|
||||
"Request completed",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"method": request.method,
|
||||
"path": request.url.path,
|
||||
"status_code": response.status_code,
|
||||
"duration_ms": round(duration * 1000, 2),
|
||||
"client_host": client_host,
|
||||
},
|
||||
)
|
||||
if hasattr(response, "headers"):
|
||||
response.headers["x-routstr-request-id"] = request_id
|
||||
|
||||
return response
|
||||
|
||||
except Exception as e:
|
||||
# Always log failures, even for skipped paths, so we don't lose errors.
|
||||
# Calculate duration
|
||||
duration = time.time() - start_time
|
||||
|
||||
# Log error
|
||||
logger.error(
|
||||
"Request failed",
|
||||
extra={
|
||||
"request_id": request_id,
|
||||
"method": request.method,
|
||||
"path": path,
|
||||
"path": request.url.path,
|
||||
"duration_ms": round(duration * 1000, 2),
|
||||
"client_host": client_host,
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
},
|
||||
|
||||
@@ -41,15 +41,6 @@ class Settings(BaseSettings):
|
||||
primary_mint: str = Field(default="", env="PRIMARY_MINT_URL")
|
||||
primary_mint_unit: str = Field(default="sat", env="PRIMARY_MINT_UNIT")
|
||||
|
||||
# Lightning payout configuration
|
||||
# Minimum available balance (in satoshis) before profit is paid out over
|
||||
# Lightning
|
||||
min_payout_sat: int = Field(default=210, gt=0, env="MIN_PAYOUT_SAT")
|
||||
# Interval (seconds) between periodic payout attempts. Must be positive.
|
||||
payout_interval_seconds: int = Field(
|
||||
default=900, gt=0, env="PAYOUT_INTERVAL_SECONDS"
|
||||
)
|
||||
|
||||
# Pricing
|
||||
# Default behavior: derive pricing from MODELS
|
||||
# If fixed_pricing is True -> use fixed_cost_per_request and ignore tokens
|
||||
@@ -61,7 +52,7 @@ class Settings(BaseSettings):
|
||||
exchange_fee: float = Field(default=1.005, env="EXCHANGE_FEE")
|
||||
upstream_provider_fee: float = Field(default=1.05, env="UPSTREAM_PROVIDER_FEE")
|
||||
tolerance_percentage: float = Field(default=1.0, env="TOLERANCE_PERCENTAGE")
|
||||
child_key_cost: int = Field(default=0, env="CHILD_KEY_COST")
|
||||
child_key_cost: int = Field(default=1000, env="CHILD_KEY_COST")
|
||||
# Minimum per-request charge in millisatoshis when model pricing is free/zero
|
||||
min_request_msat: int = Field(default=1, env="MIN_REQUEST_MSAT")
|
||||
reset_reserved_balance_on_startup: bool = Field(
|
||||
@@ -83,7 +74,6 @@ class Settings(BaseSettings):
|
||||
enable_pricing_refresh: bool = Field(default=True, env="ENABLE_PRICING_REFRESH")
|
||||
enable_models_refresh: bool = Field(default=True, env="ENABLE_MODELS_REFRESH")
|
||||
refund_cache_ttl_seconds: int = Field(default=3600, env="REFUND_CACHE_TTL_SECONDS")
|
||||
refund_sweep_ttl_seconds: int = Field(default=604800, env="REFUND_SWEEP_TTL_SECONDS")
|
||||
|
||||
# Logging
|
||||
log_level: str = Field(default="INFO", env="LOG_LEVEL")
|
||||
@@ -102,20 +92,6 @@ class Settings(BaseSettings):
|
||||
|
||||
# Discovery
|
||||
relays: list[str] = Field(default_factory=list, env="RELAYS")
|
||||
enable_analytics_sharing: bool = Field(
|
||||
default=True, env="ENABLE_ANALYTICS_SHARING"
|
||||
)
|
||||
|
||||
def _normalize_settings_data(data: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Discard unknown keys from persisted settings."""
|
||||
normalized: dict[str, Any] = {}
|
||||
known_fields = Settings.__fields__
|
||||
|
||||
for key, value in data.items():
|
||||
if key in known_fields:
|
||||
normalized[key] = value
|
||||
|
||||
return normalized
|
||||
|
||||
|
||||
def _compute_primary_mint(cashu_mints: list[str]) -> str:
|
||||
@@ -167,7 +143,7 @@ def resolve_bootstrap() -> Settings:
|
||||
pass
|
||||
if not base.onion_url:
|
||||
try:
|
||||
from ..nostr.listing import discover_onion_url_from_tor # type: ignore
|
||||
from ..nip91 import discover_onion_url_from_tor # type: ignore
|
||||
|
||||
discovered = discover_onion_url_from_tor()
|
||||
if discovered:
|
||||
@@ -255,21 +231,16 @@ class SettingsService:
|
||||
|
||||
db_id, db_data, _updated_at = row
|
||||
try:
|
||||
db_json_raw = (
|
||||
db_json = (
|
||||
json.loads(db_data) if isinstance(db_data, str) else dict(db_data)
|
||||
)
|
||||
if not isinstance(db_json_raw, dict):
|
||||
db_json_raw = {}
|
||||
except Exception:
|
||||
db_json_raw = {}
|
||||
db_json = _normalize_settings_data(db_json_raw)
|
||||
db_json = {}
|
||||
|
||||
valid_fields = set(env_resolved.dict().keys())
|
||||
merged_dict: dict[str, Any] = dict(env_resolved.dict())
|
||||
merged_dict.update(
|
||||
{k: v for k, v in db_json.items() if v not in (None, "", [], {}) and k in valid_fields}
|
||||
{k: v for k, v in db_json.items() if v not in (None, "", [], {})}
|
||||
)
|
||||
merged_dict = Settings(**merged_dict).dict()
|
||||
|
||||
# Ensure primary_mint is consistent with cashu_mints if not explicitly set
|
||||
if not merged_dict.get("primary_mint"):
|
||||
@@ -277,7 +248,7 @@ class SettingsService:
|
||||
merged_dict.get("cashu_mints", [])
|
||||
)
|
||||
|
||||
if db_json_raw != merged_dict:
|
||||
if any(k not in db_json for k in merged_dict.keys()):
|
||||
await db_session.exec( # type: ignore
|
||||
text(
|
||||
"UPDATE settings SET data = :data, updated_at = :updated_at WHERE id = 1"
|
||||
@@ -300,7 +271,7 @@ class SettingsService:
|
||||
) -> Settings:
|
||||
async with cls._lock:
|
||||
current = cls.get()
|
||||
candidate_dict = {**current.dict(), **_normalize_settings_data(partial)}
|
||||
candidate_dict = {**current.dict(), **partial}
|
||||
candidate = Settings(**candidate_dict)
|
||||
from sqlmodel import text
|
||||
|
||||
@@ -334,10 +305,8 @@ class SettingsService:
|
||||
raise RuntimeError("Settings row missing")
|
||||
(data_str,) = row
|
||||
data = json.loads(data_str) if isinstance(data_str, str) else dict(data_str)
|
||||
valid_fields = set(settings.dict().keys())
|
||||
# Update in-place
|
||||
for k, v in data.items():
|
||||
if k in valid_fields:
|
||||
setattr(settings, k, v)
|
||||
setattr(settings, k, v)
|
||||
cls._current = settings
|
||||
return settings
|
||||
|
||||
@@ -0,0 +1,143 @@
|
||||
import asyncio
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import AsyncGenerator
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
||||
# when other refactor is merged:
|
||||
# from ..nostr import announce_provider, providers_cache_refresher
|
||||
from ..discovery import providers_cache_refresher
|
||||
from ..nip91 import announce_provider
|
||||
from ..payment.models import update_sats_pricing
|
||||
from ..payment.price import update_prices_periodically
|
||||
from ..proxy import initialize_upstreams, refresh_model_maps_periodically
|
||||
from ..wallet import periodic_payout
|
||||
from .db import create_session, init_db, run_migrations
|
||||
from .logging import get_logger
|
||||
from .settings import SettingsService
|
||||
from .settings import settings as global_settings
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
|
||||
logger.info("Application startup initiated", extra={"version": app.version})
|
||||
|
||||
btc_price_task = None
|
||||
pricing_task = None
|
||||
payout_task = None
|
||||
nip91_task = None
|
||||
providers_task = None
|
||||
models_refresh_task = None
|
||||
model_maps_refresh_task = None
|
||||
|
||||
try:
|
||||
# Run database migrations on startup
|
||||
run_migrations()
|
||||
|
||||
# Initialize database connection pools
|
||||
# This creates any tables that might not be tracked by migrations yet
|
||||
await init_db()
|
||||
|
||||
# Initialize application settings (env -> computed -> DB precedence)
|
||||
async with create_session() as session:
|
||||
s = await SettingsService.initialize(session)
|
||||
if s.reset_reserved_balance_on_startup:
|
||||
from .db import reset_all_reserved_balances
|
||||
|
||||
await reset_all_reserved_balances(session)
|
||||
|
||||
if not s.admin_password:
|
||||
logger.warning(
|
||||
f"Admin password is not set. Visit {s.http_url or 'http://localhost:8000'}/admin to set the password."
|
||||
)
|
||||
|
||||
# Apply app metadata from settings
|
||||
try:
|
||||
app.title = s.name
|
||||
app.description = s.description
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# await ensure_models_bootstrapped()
|
||||
|
||||
from ..payment.price import _update_prices
|
||||
from ..proxy import get_upstreams
|
||||
from ..upstream.helpers import refresh_upstreams_models_periodically
|
||||
|
||||
_update_prices_task = asyncio.create_task(_update_prices())
|
||||
_initialize_upstreams_task = asyncio.create_task(initialize_upstreams())
|
||||
|
||||
# ensure both setup tasks complete
|
||||
await asyncio.gather(
|
||||
_update_prices_task, _initialize_upstreams_task, return_exceptions=True
|
||||
)
|
||||
|
||||
btc_price_task = asyncio.create_task(update_prices_periodically())
|
||||
pricing_task = asyncio.create_task(update_sats_pricing())
|
||||
if global_settings.models_refresh_interval_seconds > 0:
|
||||
models_refresh_task = asyncio.create_task(
|
||||
refresh_upstreams_models_periodically(get_upstreams())
|
||||
)
|
||||
model_maps_refresh_task = asyncio.create_task(refresh_model_maps_periodically())
|
||||
payout_task = asyncio.create_task(periodic_payout())
|
||||
if global_settings.nsec:
|
||||
nip91_task = asyncio.create_task(announce_provider())
|
||||
if global_settings.providers_refresh_interval_seconds > 0:
|
||||
providers_task = asyncio.create_task(providers_cache_refresher())
|
||||
|
||||
yield
|
||||
|
||||
except asyncio.CancelledError:
|
||||
# Expected during shutdown
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Application startup failed",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
logger.info("Application shutdown initiated")
|
||||
|
||||
if btc_price_task is not None:
|
||||
btc_price_task.cancel()
|
||||
if pricing_task is not None:
|
||||
pricing_task.cancel()
|
||||
if payout_task is not None:
|
||||
payout_task.cancel()
|
||||
if nip91_task is not None:
|
||||
nip91_task.cancel()
|
||||
if providers_task is not None:
|
||||
providers_task.cancel()
|
||||
if models_refresh_task is not None:
|
||||
models_refresh_task.cancel()
|
||||
if model_maps_refresh_task is not None:
|
||||
model_maps_refresh_task.cancel()
|
||||
|
||||
try:
|
||||
tasks_to_wait = []
|
||||
if btc_price_task is not None:
|
||||
tasks_to_wait.append(btc_price_task)
|
||||
if pricing_task is not None:
|
||||
tasks_to_wait.append(pricing_task)
|
||||
if payout_task is not None:
|
||||
tasks_to_wait.append(payout_task)
|
||||
if nip91_task is not None:
|
||||
tasks_to_wait.append(nip91_task)
|
||||
if providers_task is not None:
|
||||
tasks_to_wait.append(providers_task)
|
||||
if models_refresh_task is not None:
|
||||
tasks_to_wait.append(models_refresh_task)
|
||||
if model_maps_refresh_task is not None:
|
||||
tasks_to_wait.append(model_maps_refresh_task)
|
||||
|
||||
if tasks_to_wait:
|
||||
await asyncio.gather(*tasks_to_wait, return_exceptions=True)
|
||||
logger.info("Background tasks stopped successfully")
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error stopping background tasks",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
@@ -0,0 +1,152 @@
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.responses import FileResponse, RedirectResponse
|
||||
from fastapi.staticfiles import StaticFiles
|
||||
|
||||
from .logging import get_logger
|
||||
from .settings import settings as global_settings
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def setup_ui(app: FastAPI) -> None:
|
||||
UI_DIST_PATH = Path(__file__).parent.parent.parent / "ui_out"
|
||||
|
||||
if UI_DIST_PATH.exists() and UI_DIST_PATH.is_dir():
|
||||
logger.info(f"Serving static UI from {UI_DIST_PATH}")
|
||||
|
||||
app.mount(
|
||||
"/_next",
|
||||
StaticFiles(directory=UI_DIST_PATH / "_next", check_dir=True),
|
||||
name="next-static",
|
||||
)
|
||||
|
||||
@app.get("/", include_in_schema=False)
|
||||
async def serve_root_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "index.html")
|
||||
|
||||
# Add explicit route for /index.txt to redirect to /
|
||||
@app.get("/index.txt", include_in_schema=False)
|
||||
async def redirect_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/")
|
||||
|
||||
@app.get("/admin")
|
||||
async def admin_redirect() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "index.html")
|
||||
|
||||
@app.get("/dashboard", include_in_schema=False)
|
||||
async def serve_dashboard_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "index.html")
|
||||
|
||||
@app.get("/login", include_in_schema=False)
|
||||
async def serve_login_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "login" / "index.html")
|
||||
|
||||
# Add explicit route for /login/index.txt to redirect to /login
|
||||
@app.get("/login/index.txt", include_in_schema=False)
|
||||
async def redirect_login_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/login")
|
||||
|
||||
@app.get("/model", include_in_schema=False)
|
||||
async def serve_models_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "model" / "index.html")
|
||||
|
||||
# Add explicit route for /model/index.txt to redirect to /model
|
||||
@app.get("/model/index.txt", include_in_schema=False)
|
||||
async def redirect_model_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/model")
|
||||
|
||||
@app.get("/providers", include_in_schema=False)
|
||||
async def serve_providers_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "providers" / "index.html")
|
||||
|
||||
# Add explicit route for /providers/index.txt to redirect to /providers
|
||||
@app.get("/providers/index.txt", include_in_schema=False)
|
||||
async def redirect_providers_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/providers")
|
||||
|
||||
@app.get("/settings", include_in_schema=False)
|
||||
async def serve_settings_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "settings" / "index.html")
|
||||
|
||||
# Add explicit route for /settings/index.txt to redirect to /settings
|
||||
@app.get("/settings/index.txt", include_in_schema=False)
|
||||
async def redirect_settings_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/settings")
|
||||
|
||||
@app.get("/transactions", include_in_schema=False)
|
||||
async def serve_transactions_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "transactions" / "index.html")
|
||||
|
||||
# Add explicit route for /transactions/index.txt to redirect to /transactions
|
||||
@app.get("/transactions/index.txt", include_in_schema=False)
|
||||
async def redirect_transactions_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/transactions")
|
||||
|
||||
@app.get("/balances", include_in_schema=False)
|
||||
async def serve_balances_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "balances" / "index.html")
|
||||
|
||||
# Add explicit route for /balances/index.txt to redirect to /balances
|
||||
@app.get("/balances/index.txt", include_in_schema=False)
|
||||
async def redirect_balances_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/balances")
|
||||
|
||||
@app.get("/logs", include_in_schema=False)
|
||||
async def serve_logs_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "logs" / "index.html")
|
||||
|
||||
# Add explicit route for /logs/index.txt to redirect to /logs
|
||||
@app.get("/logs/index.txt", include_in_schema=False)
|
||||
async def redirect_logs_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/logs")
|
||||
|
||||
@app.get("/usage", include_in_schema=False)
|
||||
async def serve_usage_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "usage" / "index.html")
|
||||
|
||||
# Add explicit route for /usage/index.txt to redirect to /usage
|
||||
@app.get("/usage/index.txt", include_in_schema=False)
|
||||
async def redirect_usage_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/usage")
|
||||
|
||||
@app.get("/unauthorized", include_in_schema=False)
|
||||
async def serve_unauthorized_ui() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "unauthorized" / "index.html")
|
||||
|
||||
# Add explicit route for /unauthorized/index.txt to redirect to /unauthorized
|
||||
@app.get("/unauthorized/index.txt", include_in_schema=False)
|
||||
async def redirect_unauthorized_index_txt() -> RedirectResponse:
|
||||
return RedirectResponse("/unauthorized")
|
||||
|
||||
@app.get("/favicon.ico", include_in_schema=False)
|
||||
async def serve_favicon() -> FileResponse:
|
||||
icon_path = UI_DIST_PATH / "icon.ico"
|
||||
if icon_path.exists():
|
||||
return FileResponse(icon_path)
|
||||
return FileResponse(UI_DIST_PATH / "favicon.ico")
|
||||
|
||||
@app.get("/icon.ico", include_in_schema=False)
|
||||
async def serve_icon() -> FileResponse:
|
||||
return FileResponse(UI_DIST_PATH / "icon.ico")
|
||||
|
||||
app.mount(
|
||||
"/static",
|
||||
StaticFiles(directory=UI_DIST_PATH, check_dir=True),
|
||||
name="ui-static",
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"UI dist directory not found at {UI_DIST_PATH}, skipping static file serving"
|
||||
)
|
||||
|
||||
@app.get("/", include_in_schema=False)
|
||||
async def root_fallback() -> dict:
|
||||
return {
|
||||
"name": global_settings.name,
|
||||
"description": global_settings.description,
|
||||
"version": app.version,
|
||||
"status": "running",
|
||||
"ui": "not available",
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,86 +0,0 @@
|
||||
"""Application version resolution.
|
||||
|
||||
Priority order:
|
||||
|
||||
1. ``VERSION_SUFFIX`` env var (manual override; preserves prior behaviour).
|
||||
2. Bare base version when HEAD is on the matching release tag (detected via
|
||||
``GIT_TAG`` env or ``git describe --tags --exact-match HEAD``).
|
||||
3. ``GIT_COMMIT`` env var (build-time injection) -> ``<base>+g<sha>``.
|
||||
4. Local ``.git`` lookup (source checkouts) -> ``<base>+g<sha>``.
|
||||
5. Fallback: bare base version.
|
||||
|
||||
The ``+g<sha>`` form is PEP 440 local-version syntax so the result remains a
|
||||
valid package version.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
from functools import lru_cache
|
||||
from pathlib import Path
|
||||
|
||||
BASE_VERSION = "0.4.3"
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[2]
|
||||
_GIT_TIMEOUT_SECONDS = 2.0
|
||||
|
||||
|
||||
def _run_git(*args: str) -> str | None:
|
||||
try:
|
||||
result = subprocess.run( # noqa: S603 - fixed argv, no shell
|
||||
["git", *args],
|
||||
cwd=_REPO_ROOT,
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
timeout=_GIT_TIMEOUT_SECONDS,
|
||||
)
|
||||
except (FileNotFoundError, subprocess.SubprocessError, OSError):
|
||||
return None
|
||||
if result.returncode != 0:
|
||||
return None
|
||||
return result.stdout.strip() or None
|
||||
|
||||
|
||||
def _git_short_sha() -> str | None:
|
||||
sha = os.getenv("GIT_COMMIT", "").strip()
|
||||
if sha:
|
||||
return sha[:7]
|
||||
return _run_git("rev-parse", "--short=7", "HEAD")
|
||||
|
||||
|
||||
def _on_tagged_release() -> bool:
|
||||
tag = os.getenv("GIT_TAG", "").strip()
|
||||
if tag:
|
||||
return tag.lstrip("v") == BASE_VERSION
|
||||
described = _run_git("describe", "--tags", "--exact-match", "HEAD")
|
||||
if described and described.lstrip("v") == BASE_VERSION:
|
||||
return True
|
||||
current_branch = _run_git("rev-parse", "--abbrev-ref", "HEAD")
|
||||
if current_branch == "main":
|
||||
branch_tag = _run_git("describe", "--tags", "--abbrev=0", "HEAD")
|
||||
if branch_tag and branch_tag.lstrip("v") == BASE_VERSION:
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
@lru_cache(maxsize=1)
|
||||
def get_version() -> str:
|
||||
suffix = os.getenv("VERSION_SUFFIX")
|
||||
if suffix is not None:
|
||||
return f"{BASE_VERSION}-{suffix}"
|
||||
|
||||
if _on_tagged_release():
|
||||
return BASE_VERSION
|
||||
|
||||
sha = _git_short_sha()
|
||||
if not sha:
|
||||
return BASE_VERSION
|
||||
|
||||
return f"{BASE_VERSION}+g{sha}"
|
||||
|
||||
|
||||
__version__ = get_version()
|
||||
|
||||
__all__ = ["BASE_VERSION", "__version__", "get_version"]
|
||||
@@ -8,8 +8,8 @@ import httpx
|
||||
import websockets
|
||||
from fastapi import APIRouter, HTTPException
|
||||
|
||||
from ..core.logging import get_logger
|
||||
from ..core.settings import settings
|
||||
from .core.logging import get_logger
|
||||
from .core.settings import settings
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
@@ -72,6 +72,8 @@ async def query_nostr_relay_for_providers(
|
||||
elif data[0] == "NOTICE":
|
||||
try:
|
||||
msg = str(data[1])
|
||||
if len(msg) > 200:
|
||||
msg = msg[:200] + "..."
|
||||
logger.debug(f"Relay notice: {msg}")
|
||||
except Exception:
|
||||
logger.debug("Relay notice received")
|
||||
@@ -23,9 +23,6 @@ class InvoiceCreateRequest(BaseModel):
|
||||
api_key: str | None = Field(
|
||||
default=None, description="Required for topup operations"
|
||||
)
|
||||
balance_limit: int | None = Field(default=None)
|
||||
balance_limit_reset: str | None = Field(default=None)
|
||||
validity_date: int | None = Field(default=None)
|
||||
|
||||
|
||||
class InvoiceCreateResponse(BaseModel):
|
||||
@@ -97,9 +94,6 @@ async def create_invoice(
|
||||
status="pending",
|
||||
api_key_hash=request.api_key[3:] if request.api_key else None,
|
||||
purpose=request.purpose,
|
||||
balance_limit=request.balance_limit,
|
||||
balance_limit_reset=request.balance_limit_reset,
|
||||
validity_date=request.validity_date,
|
||||
expires_at=expires_at,
|
||||
)
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Listing: Routstr Provider Discoverability Implementation
|
||||
NIP-91: Routstr Provider Discoverability Implementation
|
||||
Automatically announces this Routstr proxy instance to Nostr relays.
|
||||
"""
|
||||
|
||||
@@ -18,15 +18,15 @@ from nostr.key import PrivateKey
|
||||
from nostr.message_type import ClientMessageType
|
||||
from nostr.relay_manager import RelayManager
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.settings import settings
|
||||
from .core import get_logger
|
||||
from .core.settings import settings
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
def get_app_version() -> str | None:
|
||||
try:
|
||||
from ..core.version import __version__ as imported_version
|
||||
from .core.main import __version__ as imported_version
|
||||
|
||||
return imported_version
|
||||
except Exception:
|
||||
@@ -71,7 +71,7 @@ def nsec_to_keypair(nsec: str) -> tuple[str, str] | None:
|
||||
return None
|
||||
|
||||
|
||||
def create_listing_event(
|
||||
def create_nip91_event(
|
||||
private_key_hex: str,
|
||||
provider_id: str,
|
||||
endpoint_urls: list[str],
|
||||
@@ -80,7 +80,7 @@ def create_listing_event(
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> dict[str, Any]:
|
||||
"""
|
||||
Create a listing provider announcement event (kind:38421).
|
||||
Create a NIP-91 compliant provider announcement event (kind:38421).
|
||||
|
||||
Args:
|
||||
private_key_hex: 32-byte hex private key for signing
|
||||
@@ -164,14 +164,14 @@ def events_semantically_equal(a: dict[str, Any], b: dict[str, Any]) -> bool:
|
||||
return True
|
||||
|
||||
|
||||
async def query_listing_events(
|
||||
async def query_nip91_events(
|
||||
relay_url: str,
|
||||
pubkey: str,
|
||||
provider_id: str | None = None,
|
||||
timeout: int = 30,
|
||||
) -> tuple[list[dict[str, Any]], bool]:
|
||||
"""
|
||||
Query a Nostr relay for listing provider announcements (kind:38421) via nostr library.
|
||||
Query a Nostr relay for NIP-91 provider announcements (kind:38421) via nostr library.
|
||||
|
||||
Returns a tuple of (events, ok) where ok indicates whether the relay interaction
|
||||
succeeded without transport-level errors.
|
||||
@@ -188,7 +188,7 @@ async def query_listing_events(
|
||||
|
||||
flt = Filter(kinds=[38421], authors=[pubkey], limit=10)
|
||||
filters = Filters([flt])
|
||||
sub_id = f"routstr_listing_{int(time.time())}"
|
||||
sub_id = f"nip91_{int(time.time())}"
|
||||
rm.add_subscription(sub_id, filters)
|
||||
req: list[Any] = [ClientMessageType.REQUEST, sub_id]
|
||||
req.extend(filters.to_json_array())
|
||||
@@ -294,7 +294,7 @@ async def _determine_provider_id(public_key_hex: str, relay_urls: list[str]) ->
|
||||
|
||||
async def query_single_relay(relay_url: str) -> list[dict[str, Any]]:
|
||||
try:
|
||||
events, _ok = await query_listing_events(relay_url, public_key_hex, None)
|
||||
events, _ok = await query_nip91_events(relay_url, public_key_hex, None)
|
||||
return events
|
||||
except Exception:
|
||||
return []
|
||||
@@ -330,7 +330,7 @@ async def publish_to_relay(
|
||||
timeout: int = 30,
|
||||
) -> bool:
|
||||
"""
|
||||
Publish a listing event to a nostr relay via nostr library.
|
||||
Publish a NIP-91 event to a nostr relay via nostr library.
|
||||
"""
|
||||
|
||||
def _sync_publish() -> bool:
|
||||
@@ -341,7 +341,7 @@ async def publish_to_relay(
|
||||
time.sleep(1.0)
|
||||
# Publish the event as-is via publish_message to preserve signature
|
||||
rm.publish_message(json.dumps(["EVENT", event]))
|
||||
logger.debug(f"Sent listing event {event.get('id', '')} to {relay_url}")
|
||||
logger.debug(f"Sent NIP-91 event {event.get('id', '')} to {relay_url}")
|
||||
time.sleep(1.0)
|
||||
return True
|
||||
except Exception as e:
|
||||
@@ -364,13 +364,13 @@ async def announce_provider() -> None:
|
||||
# Check for NSEC in environment (use NSEC only)
|
||||
nsec = settings.nsec
|
||||
if not nsec:
|
||||
logger.info("Nostr private key not found (NSEC), skipping listing announcement")
|
||||
logger.info("Nostr private key not found (NSEC), skipping NIP-91 announcement")
|
||||
return
|
||||
|
||||
# Convert NSEC to keypair
|
||||
keypair = nsec_to_keypair(nsec)
|
||||
if not keypair:
|
||||
logger.error("Failed to parse NSEC, skipping listing announcement")
|
||||
logger.error("Failed to parse NSEC, skipping NIP-91 announcement")
|
||||
return
|
||||
|
||||
private_key_hex, public_key_hex = keypair
|
||||
@@ -409,7 +409,7 @@ async def announce_provider() -> None:
|
||||
|
||||
if not endpoint_urls:
|
||||
logger.warning(
|
||||
"No valid endpoints configured (HTTP_URL/ONION_URL). Skipping listing publish."
|
||||
"No valid endpoints configured (HTTP_URL/ONION_URL). Skipping NIP-91 publish."
|
||||
)
|
||||
return
|
||||
|
||||
@@ -434,7 +434,7 @@ async def announce_provider() -> None:
|
||||
|
||||
# Create the candidate event that we would publish
|
||||
version_str = get_app_version()
|
||||
candidate_event = create_listing_event(
|
||||
candidate_event = create_nip91_event(
|
||||
private_key_hex=private_key_hex,
|
||||
provider_id=provider_id,
|
||||
endpoint_urls=endpoint_urls,
|
||||
@@ -474,7 +474,7 @@ async def announce_provider() -> None:
|
||||
if _should_skip(relay_url):
|
||||
logger.debug(f"Skipping {relay_url} due to backoff")
|
||||
continue
|
||||
events, ok = await query_listing_events(relay_url, public_key_hex, provider_id)
|
||||
events, ok = await query_nip91_events(relay_url, public_key_hex, provider_id)
|
||||
if ok:
|
||||
_register_success(relay_url)
|
||||
existing_events.extend(events)
|
||||
@@ -489,7 +489,7 @@ async def announce_provider() -> None:
|
||||
|
||||
if not all_match:
|
||||
logger.debug(
|
||||
"No matching listing announcement found or differences detected; publishing update"
|
||||
"No matching NIP-91 announcement found or differences detected; publishing update"
|
||||
)
|
||||
success_count = 0
|
||||
for relay_url in relay_urls:
|
||||
@@ -502,11 +502,11 @@ async def announce_provider() -> None:
|
||||
else:
|
||||
_register_failure(relay_url)
|
||||
logger.info(
|
||||
f"Published listing announcement to {success_count}/{len(relay_urls)} relays"
|
||||
f"Published NIP-91 announcement to {success_count}/{len(relay_urls)} relays"
|
||||
)
|
||||
else:
|
||||
logger.debug(
|
||||
"Matching listing announcement already present; skipping publish on startup"
|
||||
"Matching NIP-91 announcement already present; skipping publish on startup"
|
||||
)
|
||||
|
||||
# Re-announce periodically (every 24 hours)
|
||||
@@ -518,7 +518,7 @@ async def announce_provider() -> None:
|
||||
|
||||
# Build fresh candidate event for comparison
|
||||
version_str = get_app_version()
|
||||
candidate_event = create_listing_event(
|
||||
candidate_event = create_nip91_event(
|
||||
private_key_hex=private_key_hex,
|
||||
provider_id=provider_id,
|
||||
endpoint_urls=endpoint_urls,
|
||||
@@ -533,7 +533,7 @@ async def announce_provider() -> None:
|
||||
if _should_skip(relay_url):
|
||||
logger.debug(f"Skipping {relay_url} due to backoff")
|
||||
continue
|
||||
events, ok = await query_listing_events(
|
||||
events, ok = await query_nip91_events(
|
||||
relay_url, public_key_hex, provider_id
|
||||
)
|
||||
if ok:
|
||||
@@ -549,7 +549,7 @@ async def announce_provider() -> None:
|
||||
|
||||
if all_match:
|
||||
logger.debug(
|
||||
"Matching listing announcement already present; skipping periodic re-announce"
|
||||
"Matching NIP-91 announcement already present; skipping periodic re-announce"
|
||||
)
|
||||
continue
|
||||
|
||||
@@ -567,8 +567,8 @@ async def announce_provider() -> None:
|
||||
_register_failure(relay_url)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Listing announcement task cancelled")
|
||||
logger.info("NIP-91 announcement task cancelled")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.debug(f"Error in listing announcement loop: {type(e).__name__}")
|
||||
logger.debug(f"Error in NIP-91 announcement loop: {type(e).__name__}")
|
||||
# Continue running despite errors
|
||||
@@ -1,5 +0,0 @@
|
||||
from .analytics import publish_usage_analytics
|
||||
from .discovery import providers_cache_refresher
|
||||
from .listing import announce_provider
|
||||
|
||||
__all__ = ["providers_cache_refresher", "announce_provider", "publish_usage_analytics"]
|
||||
@@ -1,419 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Nostr usage analytics publisher.
|
||||
Publishes a single replaceable analytics snapshot for each provider.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import hashlib
|
||||
import json
|
||||
import time
|
||||
from typing import Any
|
||||
|
||||
from nostr.event import Event
|
||||
from nostr.key import PrivateKey
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.log_manager import log_manager
|
||||
from ..core.settings import settings
|
||||
from .listing import nsec_to_keypair, publish_to_relay
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
ANALYTICS_KIND = 38422
|
||||
ANALYTICS_SCHEMA = "routstr.analytics.snapshot.v1"
|
||||
DEFAULT_RELAYS = [
|
||||
"wss://relay.nostr.band",
|
||||
"wss://relay.damus.io",
|
||||
"wss://relay.routstr.com",
|
||||
"wss://nos.lol",
|
||||
]
|
||||
PUBLISH_INTERVAL_SECONDS = 15 * 60
|
||||
DISABLED_POLL_SECONDS = 60
|
||||
DASHBOARD_WINDOW_HOURS = 24
|
||||
DASHBOARD_INTERVAL_MINUTES = 60
|
||||
MODEL_LIMIT = 20
|
||||
WINDOW_DEFINITIONS: tuple[tuple[str, int, int], ...] = (
|
||||
("24h", 24, 60),
|
||||
("7d", 7 * 24, 6 * 60),
|
||||
("30d", 30 * 24, 24 * 60),
|
||||
("3m", 90 * 24, 24 * 60),
|
||||
("1y", 365 * 24, 7 * 24 * 60),
|
||||
)
|
||||
|
||||
|
||||
def _event_to_dict(ev: Event) -> dict[str, Any]:
|
||||
return {
|
||||
"id": ev.id,
|
||||
"pubkey": ev.public_key,
|
||||
"created_at": ev.created_at,
|
||||
"kind": int(ev.kind) if not isinstance(ev.kind, int) else ev.kind,
|
||||
"tags": ev.tags,
|
||||
"content": ev.content,
|
||||
"sig": ev.signature,
|
||||
}
|
||||
|
||||
|
||||
def _resolve_provider_id(public_key_hex: str) -> str:
|
||||
explicit_provider_id = (settings.provider_id or "").strip()
|
||||
if explicit_provider_id:
|
||||
return explicit_provider_id
|
||||
return public_key_hex[:12]
|
||||
|
||||
|
||||
def _resolve_endpoint_urls() -> list[str]:
|
||||
urls: list[str] = []
|
||||
http_url = (settings.http_url or "").strip()
|
||||
onion_url = (settings.onion_url or "").strip()
|
||||
|
||||
if http_url and http_url != "http://localhost:8000":
|
||||
urls.append(http_url)
|
||||
|
||||
if onion_url:
|
||||
if onion_url.endswith(".onion") and not (
|
||||
onion_url.startswith("http://") or onion_url.startswith("https://")
|
||||
):
|
||||
onion_url = f"http://{onion_url}"
|
||||
urls.append(onion_url)
|
||||
|
||||
return urls
|
||||
|
||||
|
||||
def _resolve_relays() -> list[str]:
|
||||
configured = [url.strip() for url in settings.relays if url.strip()]
|
||||
return configured if configured else list(DEFAULT_RELAYS)
|
||||
|
||||
|
||||
def _to_int(value: Any) -> int:
|
||||
if isinstance(value, bool):
|
||||
return int(value)
|
||||
if isinstance(value, int):
|
||||
return value
|
||||
if isinstance(value, float):
|
||||
return int(value)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return int(float(value))
|
||||
except ValueError:
|
||||
return 0
|
||||
return 0
|
||||
|
||||
|
||||
def _to_float(value: Any) -> float:
|
||||
if isinstance(value, bool):
|
||||
return float(int(value))
|
||||
if isinstance(value, (int, float)):
|
||||
return float(value)
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return float(value)
|
||||
except ValueError:
|
||||
return 0.0
|
||||
return 0.0
|
||||
|
||||
|
||||
def _aggregate_top_model_usage(
|
||||
model_usage_mix: dict[str, Any],
|
||||
) -> tuple[list[dict[str, Any]], dict[str, Any]]:
|
||||
top_models_raw = model_usage_mix.get("top_models", [])
|
||||
mix_metrics_raw = model_usage_mix.get("metrics", [])
|
||||
|
||||
top_models = [model for model in top_models_raw if isinstance(model, str)]
|
||||
metrics = [row for row in mix_metrics_raw if isinstance(row, dict)]
|
||||
|
||||
model_totals: dict[str, dict[str, float | int]] = {
|
||||
model: {
|
||||
"successful_requests": 0,
|
||||
"revenue_msats": 0.0,
|
||||
"total_tokens": 0,
|
||||
}
|
||||
for model in top_models
|
||||
}
|
||||
others = {
|
||||
"successful_requests": 0,
|
||||
"revenue_msats": 0.0,
|
||||
"total_tokens": 0,
|
||||
}
|
||||
|
||||
for metric in metrics:
|
||||
model_counts = metric.get("model_counts", {})
|
||||
model_revenue = metric.get("model_revenue_msats", {})
|
||||
model_tokens = metric.get("model_tokens", {})
|
||||
|
||||
if isinstance(model_counts, dict):
|
||||
for model, count in model_counts.items():
|
||||
if model in model_totals:
|
||||
model_totals[model]["successful_requests"] += _to_int(count)
|
||||
|
||||
if isinstance(model_revenue, dict):
|
||||
for model, amount in model_revenue.items():
|
||||
if model in model_totals:
|
||||
model_totals[model]["revenue_msats"] += _to_float(amount)
|
||||
|
||||
if isinstance(model_tokens, dict):
|
||||
for model, token_count in model_tokens.items():
|
||||
if model in model_totals:
|
||||
model_totals[model]["total_tokens"] += _to_int(token_count)
|
||||
|
||||
others["successful_requests"] += _to_int(metric.get("others", 0))
|
||||
others["revenue_msats"] += _to_float(metric.get("others_revenue_msats", 0.0))
|
||||
others["total_tokens"] += _to_int(metric.get("others_tokens", 0))
|
||||
|
||||
model_rows = [
|
||||
{
|
||||
"model": model,
|
||||
"successful_requests": int(values["successful_requests"]),
|
||||
"revenue_msats": float(values["revenue_msats"]),
|
||||
"total_tokens": int(values["total_tokens"]),
|
||||
}
|
||||
for model, values in model_totals.items()
|
||||
]
|
||||
model_rows.sort(
|
||||
key=lambda row: _to_int(row.get("successful_requests", 0)),
|
||||
reverse=True,
|
||||
)
|
||||
|
||||
return model_rows, others
|
||||
|
||||
|
||||
def _build_summary_payload(summary: dict[str, Any]) -> dict[str, Any]:
|
||||
return {
|
||||
"total_requests": _to_int(summary.get("total_requests", 0)),
|
||||
"successful_chat_completions": _to_int(
|
||||
summary.get("successful_chat_completions", 0)
|
||||
),
|
||||
"failed_requests": _to_int(summary.get("failed_requests", 0)),
|
||||
"success_rate": _to_float(summary.get("success_rate", 0.0)),
|
||||
"unique_models_count": _to_int(summary.get("unique_models_count", 0)),
|
||||
"input_tokens": _to_int(summary.get("input_tokens", 0)),
|
||||
"output_tokens": _to_int(summary.get("output_tokens", 0)),
|
||||
"total_tokens": _to_int(summary.get("total_tokens", 0)),
|
||||
"revenue_msats": _to_float(summary.get("revenue_msats", 0.0)),
|
||||
"refunds_msats": _to_float(summary.get("refunds_msats", 0.0)),
|
||||
"net_revenue_msats": _to_float(summary.get("net_revenue_msats", 0.0)),
|
||||
"revenue_sats": _to_float(summary.get("revenue_sats", 0.0)),
|
||||
"refunds_sats": _to_float(summary.get("refunds_sats", 0.0)),
|
||||
"net_revenue_sats": _to_float(summary.get("net_revenue_sats", 0.0)),
|
||||
}
|
||||
|
||||
|
||||
def _build_window_payload(
|
||||
*,
|
||||
hours: int,
|
||||
interval_minutes: int,
|
||||
model_limit: int,
|
||||
) -> dict[str, Any]:
|
||||
dashboard = log_manager.get_usage_dashboard(
|
||||
interval=interval_minutes,
|
||||
hours=hours,
|
||||
error_limit=1,
|
||||
model_limit=model_limit,
|
||||
)
|
||||
|
||||
summary = dashboard.get("summary", {})
|
||||
model_usage_mix = dashboard.get("model_usage_mix", {})
|
||||
|
||||
summary_payload = _build_summary_payload(summary if isinstance(summary, dict) else {})
|
||||
usage_mix_payload = model_usage_mix if isinstance(model_usage_mix, dict) else {}
|
||||
top_model_usage, others_usage = _aggregate_top_model_usage(usage_mix_payload)
|
||||
|
||||
return {
|
||||
"window_hours": hours,
|
||||
"interval_minutes": interval_minutes,
|
||||
"summary": summary_payload,
|
||||
"model_usage_mix": usage_mix_payload,
|
||||
"top_model_usage": top_model_usage,
|
||||
"others_usage": others_usage,
|
||||
}
|
||||
|
||||
|
||||
def build_stats_snapshot_payload(
|
||||
provider_id: str,
|
||||
*,
|
||||
public_key_hex: str,
|
||||
generated_at: int,
|
||||
window_hours: int = DASHBOARD_WINDOW_HOURS,
|
||||
interval_minutes: int = DASHBOARD_INTERVAL_MINUTES,
|
||||
model_limit: int = MODEL_LIMIT,
|
||||
) -> dict[str, Any]:
|
||||
_ = (window_hours, interval_minutes)
|
||||
windows: dict[str, dict[str, Any]] = {}
|
||||
for key, hours, window_interval_minutes in WINDOW_DEFINITIONS:
|
||||
windows[key] = _build_window_payload(
|
||||
hours=hours,
|
||||
interval_minutes=window_interval_minutes,
|
||||
model_limit=model_limit,
|
||||
)
|
||||
|
||||
primary_window = windows.get("24h", {})
|
||||
summary_payload = (
|
||||
primary_window.get("summary", {})
|
||||
if isinstance(primary_window.get("summary", {}), dict)
|
||||
else {}
|
||||
)
|
||||
usage_mix_payload = (
|
||||
primary_window.get("model_usage_mix", {})
|
||||
if isinstance(primary_window.get("model_usage_mix", {}), dict)
|
||||
else {}
|
||||
)
|
||||
top_model_usage = (
|
||||
primary_window.get("top_model_usage", [])
|
||||
if isinstance(primary_window.get("top_model_usage", []), list)
|
||||
else []
|
||||
)
|
||||
others_usage = (
|
||||
primary_window.get("others_usage", {})
|
||||
if isinstance(primary_window.get("others_usage", {}), dict)
|
||||
else {}
|
||||
)
|
||||
|
||||
return {
|
||||
"schema": ANALYTICS_SCHEMA,
|
||||
"generated_at": generated_at,
|
||||
"provider_id": provider_id,
|
||||
"pubkey": public_key_hex,
|
||||
"npub": settings.npub or "",
|
||||
"endpoint_urls": _resolve_endpoint_urls(),
|
||||
"window_hours": DASHBOARD_WINDOW_HOURS,
|
||||
"interval_minutes": DASHBOARD_INTERVAL_MINUTES,
|
||||
"summary": summary_payload,
|
||||
"model_usage_mix": usage_mix_payload,
|
||||
"top_model_usage": top_model_usage,
|
||||
"others_usage": others_usage,
|
||||
"windows": windows,
|
||||
}
|
||||
|
||||
|
||||
def create_stats_snapshot_event(
|
||||
private_key_hex: str,
|
||||
provider_id: str,
|
||||
payload_json: str,
|
||||
*,
|
||||
d_tag: str,
|
||||
) -> dict[str, Any]:
|
||||
private_key = PrivateKey(bytes.fromhex(private_key_hex))
|
||||
tags = [
|
||||
["d", d_tag],
|
||||
["provider", provider_id],
|
||||
["schema", ANALYTICS_SCHEMA],
|
||||
]
|
||||
|
||||
event = Event(
|
||||
public_key=private_key.public_key.hex(),
|
||||
content=payload_json,
|
||||
kind=ANALYTICS_KIND,
|
||||
tags=tags,
|
||||
)
|
||||
private_key.sign_event(event)
|
||||
return _event_to_dict(event)
|
||||
|
||||
|
||||
def _fingerprint_payload(payload: dict[str, Any]) -> str:
|
||||
normalized = dict(payload)
|
||||
# Ignore generated timestamp for semantic dedupe.
|
||||
normalized.pop("generated_at", None)
|
||||
payload_json = json.dumps(normalized, separators=(",", ":"), sort_keys=True)
|
||||
return hashlib.sha256(payload_json.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
async def publish_usage_analytics() -> None:
|
||||
last_payload_hash: str | None = None
|
||||
|
||||
parsed_nsec: str | None = None
|
||||
private_key_hex: str | None = None
|
||||
public_key_hex: str | None = None
|
||||
provider_id: str | None = None
|
||||
warned_missing_nsec = False
|
||||
|
||||
logger.info("Usage analytics sharing task started")
|
||||
|
||||
while True:
|
||||
try:
|
||||
if not settings.enable_analytics_sharing:
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
continue
|
||||
|
||||
nsec = (settings.nsec or "").strip()
|
||||
if not nsec:
|
||||
if not warned_missing_nsec:
|
||||
logger.info("NSEC is not configured; skipping analytics sharing to Nostr")
|
||||
warned_missing_nsec = True
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
continue
|
||||
|
||||
warned_missing_nsec = False
|
||||
if nsec != parsed_nsec or private_key_hex is None or public_key_hex is None:
|
||||
keypair = nsec_to_keypair(nsec)
|
||||
if not keypair:
|
||||
logger.error("Invalid NSEC; analytics sharing is paused")
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
continue
|
||||
private_key_hex, public_key_hex = keypair
|
||||
parsed_nsec = nsec
|
||||
provider_id = _resolve_provider_id(public_key_hex)
|
||||
last_payload_hash = None
|
||||
|
||||
if private_key_hex is None or public_key_hex is None:
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
continue
|
||||
|
||||
relay_urls = _resolve_relays()
|
||||
if not relay_urls:
|
||||
logger.warning("No Nostr relays configured; analytics sharing skipped")
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
continue
|
||||
|
||||
resolved_provider_id = provider_id or _resolve_provider_id(public_key_hex)
|
||||
now_ts = int(time.time())
|
||||
payload = build_stats_snapshot_payload(
|
||||
resolved_provider_id,
|
||||
public_key_hex=public_key_hex,
|
||||
generated_at=now_ts,
|
||||
)
|
||||
|
||||
payload_hash = _fingerprint_payload(payload)
|
||||
if last_payload_hash == payload_hash:
|
||||
await asyncio.sleep(PUBLISH_INTERVAL_SECONDS)
|
||||
continue
|
||||
|
||||
payload_json = json.dumps(payload, separators=(",", ":"), sort_keys=True)
|
||||
d_tag = f"{resolved_provider_id}:stats"
|
||||
event = create_stats_snapshot_event(
|
||||
private_key_hex,
|
||||
resolved_provider_id,
|
||||
payload_json,
|
||||
d_tag=d_tag,
|
||||
)
|
||||
|
||||
success_count = 0
|
||||
for relay_url in relay_urls:
|
||||
if await publish_to_relay(relay_url, event):
|
||||
success_count += 1
|
||||
|
||||
if success_count > 0:
|
||||
last_payload_hash = payload_hash
|
||||
|
||||
logger.info(
|
||||
"Published analytics snapshot (success=%s/%s provider=%s)",
|
||||
success_count,
|
||||
len(relay_urls),
|
||||
resolved_provider_id,
|
||||
extra={
|
||||
"relay_success_count": success_count,
|
||||
"relay_total": len(relay_urls),
|
||||
"provider_id": resolved_provider_id,
|
||||
},
|
||||
)
|
||||
await asyncio.sleep(PUBLISH_INTERVAL_SECONDS)
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info("Usage analytics sharing task cancelled")
|
||||
break
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Usage analytics sharing error",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
await asyncio.sleep(DISABLED_POLL_SECONDS)
|
||||
+116
-328
@@ -16,12 +16,6 @@ class CostData(BaseModel):
|
||||
output_msats: int
|
||||
total_msats: int
|
||||
total_usd: float = 0.0
|
||||
input_tokens: int = 0
|
||||
output_tokens: int = 0
|
||||
cache_read_input_tokens: int = 0
|
||||
cache_creation_input_tokens: int = 0
|
||||
cache_read_msats: int = 0
|
||||
cache_creation_msats: int = 0
|
||||
|
||||
|
||||
class MaxCostData(CostData):
|
||||
@@ -33,10 +27,11 @@ class CostDataError(BaseModel):
|
||||
code: str
|
||||
|
||||
|
||||
async def calculate_cost(
|
||||
async def calculate_cost( # todo: can be sync
|
||||
response_data: dict, max_cost: int, session: AsyncSession
|
||||
) -> CostData | MaxCostData | CostDataError:
|
||||
"""Calculate the cost of an API request based on token usage.
|
||||
"""
|
||||
Calculate the cost of an API request based on token usage.
|
||||
|
||||
Args:
|
||||
response_data: Response data containing usage information
|
||||
@@ -54,20 +49,12 @@ async def calculate_cost(
|
||||
},
|
||||
)
|
||||
|
||||
# Check for usage data
|
||||
if "usage" not in response_data or response_data["usage"] is None:
|
||||
logger.warning(
|
||||
"No usage data in response — billing at MaxCostData with zero "
|
||||
"tokens. Dashboard will show this request as `(0+0)`. Most "
|
||||
"common cause: upstream stream did not include a final usage "
|
||||
"chunk (OpenAI-compat backends require "
|
||||
"`stream_options.include_usage=true`).",
|
||||
"No usage data in response, using base cost only",
|
||||
extra={
|
||||
"max_cost_msats": max_cost,
|
||||
"model": response_data.get("model", "unknown"),
|
||||
"response_keys": sorted(response_data.keys())
|
||||
if isinstance(response_data, dict)
|
||||
else None,
|
||||
},
|
||||
)
|
||||
return MaxCostData(
|
||||
@@ -76,60 +63,47 @@ async def calculate_cost(
|
||||
output_msats=0,
|
||||
total_msats=0,
|
||||
total_usd=0.0,
|
||||
input_tokens=0,
|
||||
output_tokens=0,
|
||||
cache_read_input_tokens=0,
|
||||
cache_creation_input_tokens=0,
|
||||
cache_read_msats=0,
|
||||
cache_creation_msats=0,
|
||||
)
|
||||
|
||||
usage_data = response_data["usage"]
|
||||
|
||||
# Extract token counts
|
||||
input_tokens = _extract_token_pair(usage_data, "prompt_tokens", "input_tokens")
|
||||
output_tokens = _extract_token_pair(usage_data, "completion_tokens", "output_tokens")
|
||||
usd_cost = 0.0
|
||||
|
||||
# Extract cache tokens (handles OpenAI vs Anthropic formats)
|
||||
cache_read_tokens, cache_creation_tokens, input_tokens = _extract_cache_tokens(
|
||||
usage_data, input_tokens
|
||||
)
|
||||
# Prioritize cost_details.upstream_inference_cost
|
||||
if "cost_details" in usage_data:
|
||||
usd_cost = float(
|
||||
usage_data["cost_details"].get("upstream_inference_cost", 0) or 0
|
||||
)
|
||||
|
||||
# Fallback to cost field if upstream_inference_cost is 0
|
||||
if usd_cost == 0 and "cost" in usage_data:
|
||||
try:
|
||||
usd_cost = float(usage_data.get("cost", 0) or 0)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
# Try USD cost first
|
||||
usd_cost = _resolve_usd_cost(usage_data, response_data)
|
||||
if usd_cost > 0:
|
||||
if input_tokens == 0 and output_tokens == 0:
|
||||
logger.warning(
|
||||
"Upstream reported a USD cost but no token counts — "
|
||||
"billing the USD-derived cost while the dashboard will "
|
||||
"show this request as `(0+0)` tokens. Check that the "
|
||||
"upstream actually emits `usage.input_tokens` and "
|
||||
"`usage.output_tokens` (OpenAI-compat streams require "
|
||||
"`stream_options.include_usage=true`).",
|
||||
try:
|
||||
sats_per_usd = 1.0 / sats_usd_price()
|
||||
cost_in_sats = usd_cost * sats_per_usd
|
||||
cost_in_msats = math.ceil(cost_in_sats * 1000)
|
||||
|
||||
logger.info(
|
||||
"Using cost from usage data/details",
|
||||
extra={
|
||||
"model": response_data.get("model", "unknown"),
|
||||
"usd_cost": usd_cost,
|
||||
"usage_keys": sorted(usage_data.keys())
|
||||
if isinstance(usage_data, dict)
|
||||
else None,
|
||||
"cost_in_sats": cost_in_sats,
|
||||
"cost_in_msats": cost_in_msats,
|
||||
"model": response_data.get("model", "unknown"),
|
||||
},
|
||||
)
|
||||
try:
|
||||
input_usd = _coerce_usd(
|
||||
usage_data.get("cost_details", {}).get("input_cost", 0)
|
||||
)
|
||||
output_usd = _coerce_usd(
|
||||
usage_data.get("cost_details", {}).get("output_cost", 0)
|
||||
)
|
||||
return _calculate_from_usd_cost(
|
||||
usd_cost,
|
||||
input_usd,
|
||||
output_usd,
|
||||
input_tokens,
|
||||
cache_read_tokens,
|
||||
cache_creation_tokens,
|
||||
output_tokens,
|
||||
response_data,
|
||||
|
||||
return CostData(
|
||||
base_msats=-1,
|
||||
input_msats=-1, # Cost field doesn't break down by token type
|
||||
output_msats=-1,
|
||||
total_msats=cost_in_msats,
|
||||
total_usd=usd_cost,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
@@ -140,35 +114,69 @@ async def calculate_cost(
|
||||
"model": response_data.get("model", "unknown"),
|
||||
},
|
||||
)
|
||||
# Fall through to token-based calculation
|
||||
|
||||
# Fall back to token-based pricing
|
||||
try:
|
||||
pricing_rates = _get_pricing_rates(response_data)
|
||||
except ValueError as e:
|
||||
return CostDataError(message=str(e), code="pricing_error")
|
||||
MSATS_PER_1K_INPUT_TOKENS: float = (
|
||||
float(settings.fixed_per_1k_input_tokens) * 1000.0
|
||||
)
|
||||
MSATS_PER_1K_OUTPUT_TOKENS: float = (
|
||||
float(settings.fixed_per_1k_output_tokens) * 1000.0
|
||||
)
|
||||
|
||||
if pricing_rates is None:
|
||||
input_rate = float(settings.fixed_per_1k_input_tokens) * 1000.0
|
||||
output_rate = float(settings.fixed_per_1k_output_tokens) * 1000.0
|
||||
cache_read_rate = input_rate
|
||||
cache_creation_rate = input_rate
|
||||
else:
|
||||
input_rate, output_rate, cache_read_rate, cache_creation_rate = pricing_rates
|
||||
if not settings.fixed_pricing:
|
||||
response_model = response_data.get("model", "")
|
||||
logger.debug(
|
||||
"Using model-based pricing",
|
||||
extra={"model": response_model},
|
||||
)
|
||||
|
||||
if not (input_rate and output_rate):
|
||||
from ..proxy import get_model_instance
|
||||
|
||||
model_obj = get_model_instance(response_model)
|
||||
|
||||
if not model_obj:
|
||||
logger.error(
|
||||
"Invalid model in response",
|
||||
extra={"response_model": response_model},
|
||||
)
|
||||
return CostDataError(
|
||||
message=f"Invalid model in response: {response_model}",
|
||||
code="model_not_found",
|
||||
)
|
||||
|
||||
if not model_obj.sats_pricing:
|
||||
logger.error(
|
||||
"Model pricing not defined",
|
||||
extra={"model": response_model, "model_id": response_model},
|
||||
)
|
||||
return CostDataError(
|
||||
message="Model pricing not defined", code="pricing_not_found"
|
||||
)
|
||||
|
||||
try:
|
||||
mspp = float(model_obj.sats_pricing.prompt)
|
||||
mspc = float(model_obj.sats_pricing.completion)
|
||||
except Exception:
|
||||
return CostDataError(message="Invalid pricing data", code="pricing_invalid")
|
||||
|
||||
MSATS_PER_1K_INPUT_TOKENS = mspp * 1_000_000.0
|
||||
MSATS_PER_1K_OUTPUT_TOKENS = mspc * 1_000_000.0
|
||||
|
||||
logger.info(
|
||||
"Applied model-specific pricing",
|
||||
extra={
|
||||
"model": response_model,
|
||||
"input_price_msats_per_1k": MSATS_PER_1K_INPUT_TOKENS,
|
||||
"output_price_msats_per_1k": MSATS_PER_1K_OUTPUT_TOKENS,
|
||||
},
|
||||
)
|
||||
|
||||
if not (MSATS_PER_1K_OUTPUT_TOKENS and MSATS_PER_1K_INPUT_TOKENS):
|
||||
logger.warning(
|
||||
"No token pricing configured — billing at flat MaxCostData. "
|
||||
"Token counts %s in the upstream response but cannot be "
|
||||
"priced; the request will appear in dashboards with the "
|
||||
"raw counts and a fixed max-cost charge.",
|
||||
"are present"
|
||||
if (input_tokens > 0 or output_tokens > 0)
|
||||
else "are zero",
|
||||
"No token pricing configured, using base cost",
|
||||
extra={
|
||||
"base_cost_msats": max_cost,
|
||||
"model": response_data.get("model", "unknown"),
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
},
|
||||
)
|
||||
return MaxCostData(
|
||||
@@ -176,245 +184,35 @@ async def calculate_cost(
|
||||
input_msats=0,
|
||||
output_msats=0,
|
||||
total_msats=max_cost,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_input_tokens=cache_read_tokens,
|
||||
cache_creation_input_tokens=cache_creation_tokens,
|
||||
cache_read_msats=0,
|
||||
cache_creation_msats=0,
|
||||
)
|
||||
|
||||
return _calculate_from_tokens(
|
||||
input_tokens,
|
||||
output_tokens,
|
||||
cache_read_tokens,
|
||||
cache_creation_tokens,
|
||||
input_rate,
|
||||
output_rate,
|
||||
cache_read_rate,
|
||||
cache_creation_rate,
|
||||
response_data,
|
||||
input_tokens = usage_data.get("prompt_tokens", 0)
|
||||
output_tokens = usage_data.get("completion_tokens", 0)
|
||||
|
||||
# added for response api
|
||||
input_tokens = (
|
||||
input_tokens if input_tokens != 0 else usage_data.get("input_tokens", 0)
|
||||
)
|
||||
output_tokens = (
|
||||
output_tokens if output_tokens != 0 else usage_data.get("output_tokens", 0)
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Helper Functions (ordered by call sequence in calculate_cost)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
def parse_token_count(value: object) -> int:
|
||||
"""Parse a token count from various formats (int, float, str, bool)."""
|
||||
if isinstance(value, bool):
|
||||
return 0
|
||||
if isinstance(value, int):
|
||||
return max(0, value)
|
||||
if isinstance(value, float):
|
||||
return max(0, int(value))
|
||||
if isinstance(value, str):
|
||||
try:
|
||||
return max(0, int(float(value)))
|
||||
except ValueError:
|
||||
return 0
|
||||
return 0
|
||||
|
||||
|
||||
def _coerce_usd(value: object) -> float:
|
||||
"""Coerce a value to USD float, handling various formats safely."""
|
||||
if value is None or isinstance(value, bool):
|
||||
return 0.0
|
||||
if not isinstance(value, (int, float, str)):
|
||||
return 0.0
|
||||
try:
|
||||
return max(0.0, float(value))
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
|
||||
def _extract_token_pair(
|
||||
usage_data: dict, standard_field: str, alt_field: str
|
||||
) -> int:
|
||||
"""Extract token count trying two field names in order."""
|
||||
value = parse_token_count(usage_data.get(standard_field, 0))
|
||||
if value > 0:
|
||||
return value
|
||||
return parse_token_count(usage_data.get(alt_field, 0))
|
||||
|
||||
|
||||
def _extract_cache_tokens(usage_data: dict, input_tokens: int) -> tuple[int, int, int]:
|
||||
"""Extract cache tokens, handling OpenAI vs Anthropic formats.
|
||||
|
||||
Returns: (cache_read_tokens, cache_creation_tokens, adjusted_input_tokens)
|
||||
"""
|
||||
cache_read = parse_token_count(usage_data.get("cache_read_input_tokens", 0))
|
||||
cache_creation = parse_token_count(
|
||||
usage_data.get("cache_creation_input_tokens", 0)
|
||||
# added for response api
|
||||
input_tokens = (
|
||||
input_tokens
|
||||
if input_tokens != 0
|
||||
else response_data.get("usage", {}).get("input_tokens", 0)
|
||||
)
|
||||
output_tokens = (
|
||||
output_tokens
|
||||
if output_tokens != 0
|
||||
else response_data.get("usage", {}).get("output_tokens", 0)
|
||||
)
|
||||
|
||||
# OpenAI: cache is included in input_tokens, subtract it
|
||||
prompt_details = usage_data.get("prompt_tokens_details")
|
||||
if isinstance(prompt_details, dict) and not cache_read:
|
||||
openai_cached = parse_token_count(prompt_details.get("cached_tokens", 0))
|
||||
if openai_cached:
|
||||
cache_read = openai_cached
|
||||
input_tokens = max(0, input_tokens - cache_read)
|
||||
input_msats = round(input_tokens / 1000 * MSATS_PER_1K_INPUT_TOKENS, 3)
|
||||
|
||||
return cache_read, cache_creation, input_tokens
|
||||
|
||||
|
||||
def _resolve_usd_cost(usage_data: dict, response_data: dict) -> float:
|
||||
"""Resolve USD cost with clear priority order.
|
||||
|
||||
Priority: cost_details.total_cost → total_cost → cost (in both usage and response).
|
||||
"""
|
||||
cost_details = usage_data.get("cost_details")
|
||||
if isinstance(cost_details, dict):
|
||||
cost = _coerce_usd(cost_details.get("total_cost"))
|
||||
if cost > 0:
|
||||
return cost
|
||||
|
||||
for source in [usage_data, response_data]:
|
||||
if not isinstance(source, dict):
|
||||
continue
|
||||
for field in ("total_cost", "cost"):
|
||||
cost = _coerce_usd(source.get(field))
|
||||
if cost > 0:
|
||||
return cost
|
||||
|
||||
return 0.0
|
||||
|
||||
|
||||
def _get_pricing_rates(
|
||||
response_data: dict,
|
||||
) -> tuple[float, float, float, float] | None:
|
||||
"""Get model-based pricing rates or None if using fixed pricing.
|
||||
|
||||
Returns: (input_rate, output_rate, cache_read_rate, cache_write_rate)
|
||||
"""
|
||||
if settings.fixed_pricing:
|
||||
return None
|
||||
|
||||
from ..proxy import get_model_instance
|
||||
|
||||
response_model = response_data.get("model", "")
|
||||
model_obj = get_model_instance(response_model)
|
||||
|
||||
if not model_obj:
|
||||
logger.error("Invalid model in response", extra={"response_model": response_model})
|
||||
raise ValueError(f"Invalid model: {response_model}")
|
||||
|
||||
if not model_obj.sats_pricing:
|
||||
logger.error(
|
||||
"Model pricing not defined",
|
||||
extra={"model": response_model, "model_id": response_model},
|
||||
)
|
||||
raise ValueError("Model pricing not defined")
|
||||
|
||||
try:
|
||||
mspp = float(model_obj.sats_pricing.prompt)
|
||||
mspc = float(model_obj.sats_pricing.completion)
|
||||
mscr = float(model_obj.sats_pricing.input_cache_read or 0)
|
||||
mscw = float(model_obj.sats_pricing.input_cache_write or 0)
|
||||
|
||||
mspp_1k = mspp * 1_000_000.0
|
||||
mspc_1k = mspc * 1_000_000.0
|
||||
mscr_1k = mscr * 1_000_000.0 if mscr > 0 else mspp_1k
|
||||
mscw_1k = mscw * 1_000_000.0 if mscw > 0 else mspp_1k
|
||||
|
||||
logger.info(
|
||||
"Applied model-specific pricing",
|
||||
extra={
|
||||
"model": response_model,
|
||||
"input_price_msats_per_1k": mspp_1k,
|
||||
"output_price_msats_per_1k": mspc_1k,
|
||||
"cache_read_price_msats_per_1k": mscr_1k,
|
||||
"cache_write_price_msats_per_1k": mscw_1k,
|
||||
},
|
||||
)
|
||||
return mspp_1k, mspc_1k, mscr_1k, mscw_1k
|
||||
except Exception as e:
|
||||
logger.error("Invalid pricing data", extra={"error": str(e)})
|
||||
raise ValueError("Invalid pricing data") from e
|
||||
|
||||
|
||||
def _calculate_from_usd_cost(
|
||||
usd_cost: float,
|
||||
input_usd: float,
|
||||
output_usd: float,
|
||||
input_tokens: int,
|
||||
cache_read_tokens: int,
|
||||
cache_creation_tokens: int,
|
||||
output_tokens: int,
|
||||
response_data: dict,
|
||||
) -> CostData:
|
||||
"""Calculate cost from USD figures, deriving input/output split from tokens."""
|
||||
sats_per_usd = 1.0 / sats_usd_price()
|
||||
cost_in_sats = usd_cost * sats_per_usd
|
||||
cost_in_msats = math.ceil(cost_in_sats * 1000)
|
||||
|
||||
if input_usd > 0 or output_usd > 0:
|
||||
input_msats = int((input_usd * sats_per_usd) * 1000)
|
||||
output_msats = int((output_usd * sats_per_usd) * 1000)
|
||||
else:
|
||||
effective_input_tokens = (
|
||||
input_tokens + cache_read_tokens + cache_creation_tokens
|
||||
)
|
||||
total_tokens = effective_input_tokens + output_tokens
|
||||
input_msats = (
|
||||
int(cost_in_msats * effective_input_tokens / total_tokens)
|
||||
if total_tokens > 0
|
||||
else 0
|
||||
)
|
||||
output_msats = cost_in_msats - input_msats
|
||||
|
||||
logger.info(
|
||||
"Using cost from usage data/details",
|
||||
extra={
|
||||
"usd_cost": usd_cost,
|
||||
"cost_in_sats": cost_in_sats,
|
||||
"cost_in_msats": cost_in_msats,
|
||||
"model": response_data.get("model", "unknown"),
|
||||
},
|
||||
)
|
||||
|
||||
return CostData(
|
||||
base_msats=0,
|
||||
input_msats=input_msats,
|
||||
output_msats=output_msats,
|
||||
total_msats=cost_in_msats,
|
||||
total_usd=usd_cost,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_input_tokens=cache_read_tokens,
|
||||
cache_creation_input_tokens=cache_creation_tokens,
|
||||
cache_read_msats=0,
|
||||
cache_creation_msats=0,
|
||||
)
|
||||
|
||||
|
||||
def _calculate_from_tokens(
|
||||
input_tokens: int,
|
||||
output_tokens: int,
|
||||
cache_read_tokens: int,
|
||||
cache_creation_tokens: int,
|
||||
input_rate: float,
|
||||
output_rate: float,
|
||||
cache_read_rate: float,
|
||||
cache_creation_rate: float,
|
||||
response_data: dict,
|
||||
) -> CostData:
|
||||
"""Calculate cost from token counts using pricing rates."""
|
||||
calc_input_msats = round(input_tokens / 1000 * input_rate, 3)
|
||||
calc_output_msats = round(output_tokens / 1000 * output_rate, 3)
|
||||
calc_cache_read_msats = round(cache_read_tokens / 1000 * cache_read_rate, 3)
|
||||
calc_cache_write_msats = round(
|
||||
cache_creation_tokens / 1000 * cache_creation_rate, 3
|
||||
)
|
||||
token_based_cost = math.ceil(
|
||||
calc_input_msats
|
||||
+ calc_output_msats
|
||||
+ calc_cache_read_msats
|
||||
+ calc_cache_write_msats
|
||||
)
|
||||
output_msats = round(output_tokens / 1000 * MSATS_PER_1K_OUTPUT_TOKENS, 3)
|
||||
token_based_cost = math.ceil(input_msats + output_msats)
|
||||
total_usd = (token_based_cost / 1000.0) * sats_usd_price()
|
||||
|
||||
logger.info(
|
||||
@@ -422,12 +220,8 @@ def _calculate_from_tokens(
|
||||
extra={
|
||||
"input_tokens": input_tokens,
|
||||
"output_tokens": output_tokens,
|
||||
"cache_read_input_tokens": cache_read_tokens,
|
||||
"cache_creation_input_tokens": cache_creation_tokens,
|
||||
"input_cost_msats": calc_input_msats,
|
||||
"output_cost_msats": calc_output_msats,
|
||||
"cache_read_cost_msats": calc_cache_read_msats,
|
||||
"cache_creation_cost_msats": calc_cache_write_msats,
|
||||
"input_cost_msats": input_msats,
|
||||
"output_cost_msats": output_msats,
|
||||
"total_cost_msats": token_based_cost,
|
||||
"total_usd": total_usd,
|
||||
"model": response_data.get("model", "unknown"),
|
||||
@@ -436,14 +230,8 @@ def _calculate_from_tokens(
|
||||
|
||||
return CostData(
|
||||
base_msats=0,
|
||||
input_msats=int(calc_input_msats),
|
||||
output_msats=int(calc_output_msats),
|
||||
input_msats=int(input_msats),
|
||||
output_msats=int(output_msats),
|
||||
total_msats=token_based_cost,
|
||||
total_usd=total_usd,
|
||||
input_tokens=input_tokens,
|
||||
output_tokens=output_tokens,
|
||||
cache_read_input_tokens=cache_read_tokens,
|
||||
cache_creation_input_tokens=cache_creation_tokens,
|
||||
cache_read_msats=int(calc_cache_read_msats),
|
||||
cache_creation_msats=int(calc_cache_write_msats),
|
||||
)
|
||||
|
||||
@@ -29,13 +29,15 @@ def check_token_balance(headers: dict, body: dict, max_cost_for_model: int) -> N
|
||||
},
|
||||
)
|
||||
elif auth := headers.get("authorization", None):
|
||||
cashu_token = auth.split(" ")[1] if len(auth.split(" ")) > 1 else ""
|
||||
logger.debug(
|
||||
"Skipping preflight token balance check for Authorization header",
|
||||
"Using Authorization header token",
|
||||
extra={
|
||||
"auth_preview": auth[:20] + "..." if len(auth) > 20 else auth,
|
||||
"token_preview": cashu_token[:20] + "..."
|
||||
if len(cashu_token) > 20
|
||||
else cashu_token
|
||||
},
|
||||
)
|
||||
return
|
||||
else:
|
||||
logger.error("No authentication token provided")
|
||||
raise HTTPException(status_code=401, detail="Unauthorized")
|
||||
@@ -73,7 +75,7 @@ def check_token_balance(headers: dict, body: dict, max_cost_for_model: int) -> N
|
||||
|
||||
if max_cost_for_model > amount_msat:
|
||||
raise HTTPException(
|
||||
status_code=402,
|
||||
status_code=413,
|
||||
detail={
|
||||
"reason": "Insufficient balance",
|
||||
"amount_required_msat": max_cost_for_model,
|
||||
@@ -167,32 +169,9 @@ async def calculate_discounted_max_cost(
|
||||
|
||||
tol = settings.tolerance_percentage
|
||||
tol_factor = max(0.0, 1 - float(tol) / 100.0)
|
||||
|
||||
max_prompt_allowed_sats = model_pricing.max_prompt_cost * tol_factor
|
||||
max_completion_allowed_sats = model_pricing.max_completion_cost * tol_factor
|
||||
|
||||
if model_obj:
|
||||
prompt_token_limit: int | None = None
|
||||
if model_obj.top_provider and (
|
||||
model_obj.top_provider.context_length
|
||||
or model_obj.top_provider.max_completion_tokens
|
||||
):
|
||||
cl = model_obj.top_provider.context_length
|
||||
mct = model_obj.top_provider.max_completion_tokens
|
||||
if cl and mct:
|
||||
prompt_token_limit = max(0, cl - mct)
|
||||
elif cl:
|
||||
prompt_token_limit = cl
|
||||
elif mct:
|
||||
prompt_token_limit = 0
|
||||
elif model_obj.context_length:
|
||||
prompt_token_limit = model_obj.context_length
|
||||
|
||||
if prompt_token_limit is not None:
|
||||
max_prompt_allowed_sats = (
|
||||
prompt_token_limit * model_pricing.prompt * tol_factor
|
||||
)
|
||||
|
||||
adjusted = max_cost_for_model
|
||||
|
||||
if messages := body.get("messages"):
|
||||
|
||||
@@ -25,6 +25,82 @@ class LNURLError(Exception):
|
||||
"""LNURL related errors."""
|
||||
|
||||
|
||||
def parse_lightning_invoice_amount(invoice: str, currency: str = "sat") -> int:
|
||||
"""Parse Lightning invoice (BOLT-11) to extract amount in specified currency units.
|
||||
|
||||
Args:
|
||||
invoice: BOLT-11 Lightning invoice string
|
||||
currency: Target currency unit ("sat" or "msat")
|
||||
|
||||
Returns:
|
||||
Amount in the specified currency unit
|
||||
|
||||
Raises:
|
||||
LNURLError: If invoice format is invalid or amount cannot be parsed
|
||||
"""
|
||||
invoice = invoice.lower().strip()
|
||||
|
||||
if not invoice.startswith("ln"):
|
||||
raise LNURLError("Invalid Lightning invoice format")
|
||||
|
||||
# Find the network part (bc, tb, etc.)
|
||||
network_start = 2
|
||||
while network_start < len(invoice) and invoice[network_start] not in "0123456789":
|
||||
network_start += 1
|
||||
|
||||
if network_start >= len(invoice):
|
||||
raise LNURLError("Invalid Lightning invoice format")
|
||||
|
||||
# Parse amount and multiplier
|
||||
amount_str = ""
|
||||
multiplier = ""
|
||||
i = network_start
|
||||
|
||||
# Extract numeric part
|
||||
while i < len(invoice) and invoice[i].isdigit():
|
||||
amount_str += invoice[i]
|
||||
i += 1
|
||||
|
||||
# Extract multiplier if present
|
||||
if i < len(invoice) and invoice[i] in "munp":
|
||||
multiplier = invoice[i]
|
||||
i += 1
|
||||
|
||||
# Check if we have the required "1" separator
|
||||
if i >= len(invoice) or invoice[i] != "1":
|
||||
raise LNURLError("Invalid Lightning invoice format")
|
||||
|
||||
if not amount_str:
|
||||
raise LNURLError("Lightning invoice amount not specified")
|
||||
|
||||
# Convert to base units
|
||||
try:
|
||||
amount = int(amount_str)
|
||||
except ValueError:
|
||||
raise LNURLError("Invalid Lightning invoice amount")
|
||||
|
||||
# Apply multiplier to get millisatoshis
|
||||
if multiplier == "m": # milli = 10^-3
|
||||
amount_msat = amount * 100_000_000 # amount is in BTC * 10^-3
|
||||
elif multiplier == "u": # micro = 10^-6
|
||||
amount_msat = amount * 100_000 # amount is in BTC * 10^-6
|
||||
elif multiplier == "n": # nano = 10^-9
|
||||
amount_msat = amount * 100 # amount is in BTC * 10^-9
|
||||
elif multiplier == "p": # pico = 10^-12
|
||||
amount_msat = amount // 10 # amount is in BTC * 10^-12
|
||||
else:
|
||||
# No multiplier means the amount is in BTC
|
||||
amount_msat = amount * 100_000_000_000 # Convert BTC to msat
|
||||
|
||||
# Convert to target currency unit
|
||||
if currency == "msat":
|
||||
return amount_msat
|
||||
elif currency == "sat":
|
||||
return amount_msat // 1000
|
||||
else:
|
||||
raise LNURLError(f"Unsupported currency for Lightning: {currency}")
|
||||
|
||||
|
||||
async def decode_lnurl(lnurl: str) -> str:
|
||||
"""Decode LNURL to get the actual URL.
|
||||
|
||||
|
||||
+51
-93
@@ -4,11 +4,10 @@ import random
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends
|
||||
from pydantic import BaseModel as V2BaseModel
|
||||
from pydantic.v1 import BaseModel
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from ..core.db import ModelRow, UpstreamProviderRow, get_session
|
||||
from ..core.db import ModelRow, get_session
|
||||
from ..core.logging import get_logger
|
||||
from ..core.settings import settings
|
||||
from .price import sats_usd_price
|
||||
@@ -61,7 +60,6 @@ class Model(BaseModel):
|
||||
upstream_provider_id: int | str | None = None
|
||||
canonical_slug: str | None = None
|
||||
alias_ids: list[str] | None = None
|
||||
forwarded_model_id: str | None = None
|
||||
|
||||
def __hash__(self) -> int:
|
||||
return hash(self.id)
|
||||
@@ -145,6 +143,14 @@ async def async_fetch_openrouter_models(source_filter: str | None = None) -> lis
|
||||
return []
|
||||
|
||||
|
||||
def is_openrouter_upstream() -> bool:
|
||||
try:
|
||||
base = (settings.upstream_base_url or "").strip().rstrip("/")
|
||||
except Exception:
|
||||
return False
|
||||
return base.lower() == "https://openrouter.ai/api/v1"
|
||||
|
||||
|
||||
def _row_to_model(
|
||||
row: ModelRow, apply_provider_fee: bool = False, provider_fee: float = 1.01
|
||||
) -> Model:
|
||||
@@ -179,7 +185,6 @@ def _row_to_model(
|
||||
upstream_provider_id=row.upstream_provider_id,
|
||||
canonical_slug=getattr(row, "canonical_slug", None),
|
||||
alias_ids=json.loads(row.alias_ids) if row.alias_ids else None,
|
||||
forwarded_model_id=getattr(row, "forwarded_model_id", None) or row.id,
|
||||
)
|
||||
|
||||
if apply_provider_fee:
|
||||
@@ -198,6 +203,29 @@ def _row_to_model(
|
||||
return model
|
||||
|
||||
|
||||
def _model_to_row_payload(model: Model) -> dict[str, str | int | bool | None]:
|
||||
return {
|
||||
"id": model.id,
|
||||
"name": model.name,
|
||||
"created": model.created,
|
||||
"description": model.description,
|
||||
"context_length": model.context_length,
|
||||
"architecture": json.dumps(model.architecture.dict()),
|
||||
"pricing": json.dumps(model.pricing.dict()),
|
||||
"sats_pricing": json.dumps(model.sats_pricing.dict())
|
||||
if model.sats_pricing
|
||||
else None,
|
||||
"per_request_limits": json.dumps(model.per_request_limits)
|
||||
if model.per_request_limits is not None
|
||||
else None,
|
||||
"top_provider": json.dumps(model.top_provider.dict())
|
||||
if model.top_provider is not None
|
||||
else None,
|
||||
"enabled": model.enabled,
|
||||
"upstream_provider_id": model.upstream_provider_id,
|
||||
}
|
||||
|
||||
|
||||
async def list_models(
|
||||
session: AsyncSession,
|
||||
upstream_id: int,
|
||||
@@ -234,6 +262,21 @@ async def list_models(
|
||||
]
|
||||
|
||||
|
||||
async def get_model_by_id(
|
||||
model_id: str, provider_id: int, session: AsyncSession
|
||||
) -> Model | None:
|
||||
from ..core.db import UpstreamProviderRow
|
||||
|
||||
row = await session.get(ModelRow, (model_id, provider_id))
|
||||
if not row or not row.enabled:
|
||||
return None
|
||||
provider = await session.get(UpstreamProviderRow, provider_id)
|
||||
if not provider or not provider.enabled:
|
||||
return None
|
||||
provider_fee = provider.provider_fee if provider else 1.01
|
||||
return _row_to_model(row, apply_provider_fee=True, provider_fee=provider_fee)
|
||||
|
||||
|
||||
def _calculate_usd_max_costs(model: Model) -> tuple[float, float, float]:
|
||||
"""Calculate max costs in USD based on model context/token limits.
|
||||
|
||||
@@ -332,7 +375,6 @@ def _update_model_sats_pricing(model: Model, sats_to_usd: float) -> Model:
|
||||
upstream_provider_id=model.upstream_provider_id,
|
||||
canonical_slug=model.canonical_slug,
|
||||
alias_ids=model.alias_ids,
|
||||
forwarded_model_id=model.forwarded_model_id,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
@@ -351,9 +393,6 @@ async def _update_sats_pricing_once() -> None:
|
||||
from ..proxy import get_upstreams, refresh_model_maps
|
||||
|
||||
upstreams = get_upstreams()
|
||||
if not upstreams:
|
||||
return
|
||||
|
||||
sats_to_usd = sats_usd_price()
|
||||
|
||||
updated_count = 0
|
||||
@@ -363,14 +402,11 @@ async def _update_sats_pricing_once() -> None:
|
||||
for m in upstream.get_cached_models()
|
||||
]
|
||||
upstream._models_cache = updated_models
|
||||
upstream._models_by_id = {m.forwarded_model_id or m.id: m for m in updated_models}
|
||||
upstream._models_by_id = {m.id: m for m in updated_models}
|
||||
updated_count += len(updated_models)
|
||||
|
||||
if updated_count > 0:
|
||||
logger.info(
|
||||
f"Updated sats pricing for {updated_count} models",
|
||||
extra={"models_updated": updated_count},
|
||||
)
|
||||
logger.info("Updated sats pricing", extra={"models_updated": updated_count})
|
||||
await refresh_model_maps()
|
||||
|
||||
|
||||
@@ -412,89 +448,11 @@ async def update_sats_pricing() -> None:
|
||||
logger.error(f"Error updating sats pricing: {e}")
|
||||
|
||||
|
||||
class ModelTestRequest(V2BaseModel):
|
||||
model_id: str
|
||||
endpoint_type: str
|
||||
request_data: dict
|
||||
|
||||
|
||||
@models_router.post("/api/models/test")
|
||||
async def test_model(
|
||||
payload: ModelTestRequest,
|
||||
session: AsyncSession = Depends(get_session),
|
||||
) -> dict:
|
||||
"""Test a model by sending a request through its configured upstream provider."""
|
||||
from sqlmodel import select
|
||||
|
||||
result = await session.execute(
|
||||
select(ModelRow).where(ModelRow.id == payload.model_id)
|
||||
)
|
||||
model_row = result.scalars().first()
|
||||
|
||||
if not model_row:
|
||||
return {
|
||||
"success": False,
|
||||
"error": f"Model '{payload.model_id}' not found in database",
|
||||
"status_code": 404,
|
||||
}
|
||||
|
||||
provider = await session.get(UpstreamProviderRow, model_row.upstream_provider_id)
|
||||
if not provider:
|
||||
return {
|
||||
"success": False,
|
||||
"error": "Upstream provider not found",
|
||||
"status_code": 404,
|
||||
}
|
||||
|
||||
base_url = provider.base_url.rstrip("/")
|
||||
if payload.endpoint_type == "chat-completions":
|
||||
url = f"{base_url}/chat/completions"
|
||||
else:
|
||||
url = f"{base_url}/{payload.endpoint_type}"
|
||||
|
||||
actual_model_id = model_row.forwarded_model_id or model_row.id
|
||||
request_data = dict(payload.request_data)
|
||||
request_data["model"] = actual_model_id
|
||||
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {provider.api_key}",
|
||||
}
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
response = await client.post(url, json=request_data, headers=headers)
|
||||
try:
|
||||
response_data = response.json()
|
||||
except Exception:
|
||||
response_data = {"raw": response.text}
|
||||
|
||||
return {
|
||||
"success": response.status_code < 400,
|
||||
"data": response_data,
|
||||
"status_code": response.status_code,
|
||||
}
|
||||
except Exception as e:
|
||||
return {
|
||||
"success": False,
|
||||
"error": str(e),
|
||||
"status_code": 500,
|
||||
}
|
||||
|
||||
|
||||
@models_router.get("/v1/models")
|
||||
@models_router.get("/v1/models/", include_in_schema=False)
|
||||
@models_router.get("/models")
|
||||
@models_router.get("/models/", include_in_schema=False)
|
||||
@models_router.get("/models", include_in_schema=False)
|
||||
async def models(session: AsyncSession = Depends(get_session)) -> dict:
|
||||
"""Get all available models from all providers with database overrides applied."""
|
||||
from ..proxy import get_unique_models
|
||||
|
||||
items = get_unique_models()
|
||||
data = []
|
||||
for model in items:
|
||||
m = model.dict()
|
||||
if model.forwarded_model_id:
|
||||
m["id"] = model.forwarded_model_id
|
||||
data.append(m)
|
||||
return {"data": data}
|
||||
return {"data": items}
|
||||
|
||||
+15
-90
@@ -1,9 +1,8 @@
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request
|
||||
from fastapi.responses import HTMLResponse, JSONResponse, Response, StreamingResponse
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
from sqlmodel import select
|
||||
|
||||
from .algorithm import create_model_mappings
|
||||
@@ -18,7 +17,6 @@ from .core.db import (
|
||||
get_session,
|
||||
)
|
||||
from .core.exceptions import UpstreamError
|
||||
from .core.settings import settings
|
||||
from .payment.helpers import (
|
||||
calculate_discounted_max_cost,
|
||||
check_token_balance,
|
||||
@@ -70,25 +68,7 @@ def get_upstreams() -> list[BaseUpstreamProvider]:
|
||||
|
||||
def get_model_instance(model_id: str) -> Model | None:
|
||||
"""Get Model instance by ID from global cache."""
|
||||
if not model_id:
|
||||
return None
|
||||
|
||||
model_id_lower = model_id.lower()
|
||||
# Try exact match first
|
||||
if model := _model_instances.get(model_id_lower):
|
||||
return model
|
||||
|
||||
# Try stripping common version suffixes (e.g., -20251222)
|
||||
# This handles cases where upstream returns a specific version
|
||||
# but we only track the base model name.
|
||||
import re
|
||||
|
||||
base_model_id = re.sub(r"-\d{8}$", "", model_id_lower)
|
||||
if base_model_id != model_id_lower:
|
||||
if model := _model_instances.get(base_model_id):
|
||||
return model
|
||||
|
||||
return None
|
||||
return _model_instances.get(model_id.lower())
|
||||
|
||||
|
||||
def get_provider_for_model(model_id: str) -> list[BaseUpstreamProvider] | None:
|
||||
@@ -151,53 +131,17 @@ async def refresh_model_maps_periodically() -> None:
|
||||
)
|
||||
|
||||
|
||||
_API_PATH_PREFIXES = ("v1/", "responses")
|
||||
|
||||
_NOT_FOUND_HTML_FILE = Path(__file__).parent.parent / "ui_out" / "404.html"
|
||||
|
||||
|
||||
def _read_not_found_html() -> str | None:
|
||||
try:
|
||||
return _NOT_FOUND_HTML_FILE.read_text(encoding="utf-8")
|
||||
except OSError:
|
||||
return None
|
||||
|
||||
|
||||
_NOT_FOUND_HTML: str | None = _read_not_found_html()
|
||||
|
||||
|
||||
def _build_not_found_response(request: Request, path: str) -> Response:
|
||||
"""Return a 404 for unknown paths.
|
||||
"""
|
||||
accept = request.headers.get("accept", "").lower()
|
||||
prefers_json = "application/json" in accept and "text/html" not in accept
|
||||
request_id = getattr(request.state, "request_id", "unknown")
|
||||
|
||||
if not prefers_json and _NOT_FOUND_HTML is not None:
|
||||
return HTMLResponse(content=_NOT_FOUND_HTML, status_code=404)
|
||||
|
||||
return JSONResponse(
|
||||
status_code=404,
|
||||
content={
|
||||
"error": {
|
||||
"message": f"Path '/{path}' not found",
|
||||
"type": "not_found",
|
||||
"code": 404,
|
||||
},
|
||||
"request_id": request_id,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@proxy_router.api_route("/{path:path}", methods=["GET", "POST"], response_model=None)
|
||||
async def proxy(
|
||||
request: Request, path: str, session: AsyncSession = Depends(get_session)
|
||||
) -> Response | StreamingResponse:
|
||||
if not path.startswith(_API_PATH_PREFIXES):
|
||||
return _build_not_found_response(request, path)
|
||||
|
||||
headers = dict(request.headers)
|
||||
|
||||
if "x-cashu" not in headers and "authorization" not in headers.keys():
|
||||
return create_error_response(
|
||||
"unauthorized", "Unauthorized", 401, request=request
|
||||
)
|
||||
|
||||
is_responses_api = path.startswith("v1/responses") or path.startswith("responses")
|
||||
request_body = await request.body()
|
||||
request_body_dict = parse_request_body_json(request_body, path)
|
||||
@@ -208,7 +152,6 @@ async def proxy(
|
||||
model_id = request_body_dict.get("model", "unknown")
|
||||
|
||||
model_obj = get_model_instance(model_id)
|
||||
|
||||
if not model_obj:
|
||||
return create_error_response(
|
||||
"invalid_model", f"Model '{model_id}' not found", 400, request=request
|
||||
@@ -233,9 +176,6 @@ async def proxy(
|
||||
max_cost_for_model = await calculate_discounted_max_cost(
|
||||
_max_cost_for_model, request_body_dict, model_obj=model_obj
|
||||
)
|
||||
# Ensure max_cost_for_model is at least the minimum allowed request cost
|
||||
max_cost_for_model = max(max_cost_for_model, settings.min_request_msat)
|
||||
|
||||
check_token_balance(headers, request_body_dict, max_cost_for_model)
|
||||
|
||||
if x_cashu := headers.get("x-cashu", None):
|
||||
@@ -266,9 +206,7 @@ async def proxy(
|
||||
)
|
||||
|
||||
elif auth := headers.get("authorization", None):
|
||||
key = await get_bearer_token_key(
|
||||
headers, path, session, auth, max_cost_for_model, model_id
|
||||
)
|
||||
key = await get_bearer_token_key(headers, path, session, auth)
|
||||
|
||||
else:
|
||||
if request.method not in ["GET"]:
|
||||
@@ -356,13 +294,9 @@ async def proxy(
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
except UpstreamError:
|
||||
# Let the outer UpstreamError handler manage retry/revert
|
||||
raise
|
||||
except Exception as e:
|
||||
# Unexpected error (not an upstream failure) — revert and propagate
|
||||
logger.error(
|
||||
"Unexpected error in upstream request, reverting payment",
|
||||
"Upstream request failed, ensuring payment is reverted",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
@@ -447,16 +381,10 @@ async def proxy(
|
||||
|
||||
|
||||
async def get_bearer_token_key(
|
||||
headers: dict,
|
||||
path: str,
|
||||
session: AsyncSession,
|
||||
auth: str,
|
||||
min_cost: int = 0,
|
||||
model_id: str = "unknown",
|
||||
headers: dict, path: str, session: AsyncSession, auth: str
|
||||
) -> ApiKey:
|
||||
"""Handle bearer token authentication proxy requests."""
|
||||
parts = auth.split()
|
||||
bearer_key = parts[1] if len(parts) > 1 and parts[0].lower() == "bearer" else ""
|
||||
bearer_key = auth.replace("Bearer ", "") if auth.startswith("Bearer ") else ""
|
||||
refund_address = headers.get("Refund-LNURL", None)
|
||||
key_expiry_time = headers.get("Key-Expiry-Time", None)
|
||||
|
||||
@@ -469,7 +397,6 @@ async def get_bearer_token_key(
|
||||
"bearer_key_preview": bearer_key[:20] + "..."
|
||||
if len(bearer_key) > 20
|
||||
else bearer_key,
|
||||
"min_cost": min_cost,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -508,7 +435,6 @@ async def get_bearer_token_key(
|
||||
session,
|
||||
refund_address,
|
||||
key_expiry_time, # type: ignore
|
||||
min_cost=min_cost,
|
||||
)
|
||||
logger.info(
|
||||
"Bearer token validated successfully",
|
||||
@@ -520,16 +446,15 @@ async def get_bearer_token_key(
|
||||
)
|
||||
return key
|
||||
except Exception as e:
|
||||
key_preview = bearer_key[:20] + "..." if len(bearer_key) > 20 else bearer_key
|
||||
logger.error(
|
||||
f"Bearer token validation failed: {type(e).__name__}: {e} path={path} model={model_id!r} min_cost={min_cost} key={key_preview!r}",
|
||||
"Bearer token validation failed",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"path": path,
|
||||
"model_id": model_id,
|
||||
"min_cost_msat": min_cost,
|
||||
"bearer_key_preview": key_preview,
|
||||
"bearer_key_preview": bearer_key[:20] + "..."
|
||||
if len(bearer_key) > 20
|
||||
else bearer_key,
|
||||
},
|
||||
)
|
||||
raise
|
||||
|
||||
@@ -10,7 +10,6 @@ from .openai import OpenAIUpstreamProvider
|
||||
from .openrouter import OpenRouterUpstreamProvider
|
||||
from .perplexity import PerplexityUpstreamProvider
|
||||
from .ppqai import PPQAIUpstreamProvider
|
||||
from .routstr import RoutstrUpstreamProvider
|
||||
from .xai import XAIUpstreamProvider
|
||||
|
||||
upstream_provider_classes: list[type[BaseUpstreamProvider]] = [
|
||||
@@ -25,7 +24,6 @@ upstream_provider_classes: list[type[BaseUpstreamProvider]] = [
|
||||
OpenRouterUpstreamProvider,
|
||||
PerplexityUpstreamProvider,
|
||||
PPQAIUpstreamProvider,
|
||||
RoutstrUpstreamProvider,
|
||||
XAIUpstreamProvider,
|
||||
]
|
||||
"""List of all upstream classes"""
|
||||
|
||||
@@ -13,8 +13,6 @@ class AnthropicUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "anthropic"
|
||||
default_base_url = "https://api.anthropic.com/v1"
|
||||
platform_url = "https://console.anthropic.com/settings/keys"
|
||||
supports_anthropic_messages = True
|
||||
litellm_provider_prefix = "anthropic/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.01):
|
||||
super().__init__(
|
||||
|
||||
@@ -1,157 +0,0 @@
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
from sqlmodel import select
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.db import UpstreamProviderRow, create_session
|
||||
from ..wallet import send_token
|
||||
from .routstr import RoutstrUpstreamProvider
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Check every 60 seconds
|
||||
AUTO_TOPUP_INTERVAL_SECONDS = 60
|
||||
|
||||
|
||||
async def periodic_auto_topup() -> None:
|
||||
"""Background task that monitors Routstr provider balances and auto-tops up when below threshold.
|
||||
|
||||
For each Routstr provider with auto_topup enabled in provider_settings:
|
||||
1. Checks the upstream balance via get_balance()
|
||||
2. If balance < topup_threshold, creates a cashu token from the configured mint
|
||||
3. Sends the token to the upstream provider via topup()
|
||||
"""
|
||||
# Wait for initial startup to complete
|
||||
await asyncio.sleep(30)
|
||||
logger.info("Auto top-up worker started")
|
||||
|
||||
while True:
|
||||
try:
|
||||
await _run_auto_topup_cycle()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Auto top-up cycle failed",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
|
||||
await asyncio.sleep(AUTO_TOPUP_INTERVAL_SECONDS)
|
||||
|
||||
|
||||
async def _run_auto_topup_cycle() -> None:
|
||||
"""Single cycle: check all eligible providers and top up if needed."""
|
||||
async with create_session() as session:
|
||||
query = select(UpstreamProviderRow).where(
|
||||
UpstreamProviderRow.provider_type == "routstr",
|
||||
UpstreamProviderRow.enabled == True, # noqa: E712
|
||||
)
|
||||
result = await session.exec(query)
|
||||
providers = result.all()
|
||||
|
||||
for row in providers:
|
||||
try:
|
||||
await _check_and_topup(row)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Auto top-up failed for provider",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"base_url": row.base_url,
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _check_and_topup(row: UpstreamProviderRow) -> None:
|
||||
"""Check a single provider's balance and top up if below threshold."""
|
||||
# Parse provider settings
|
||||
settings: dict = {}
|
||||
if row.provider_settings:
|
||||
try:
|
||||
settings = json.loads(row.provider_settings)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return
|
||||
|
||||
if not settings.get("auto_topup"):
|
||||
return
|
||||
|
||||
threshold = settings.get("topup_threshold")
|
||||
amount = settings.get("topup_amount_limit")
|
||||
mint_url = settings.get("topup_mint_url")
|
||||
|
||||
if not threshold or not amount or not mint_url:
|
||||
logger.warning(
|
||||
"Auto top-up enabled but missing configuration",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"has_threshold": bool(threshold),
|
||||
"has_amount": bool(amount),
|
||||
"has_mint": bool(mint_url),
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
if not row.api_key:
|
||||
return
|
||||
|
||||
# Instantiate provider and check balance
|
||||
provider = RoutstrUpstreamProvider.from_db_row(row)
|
||||
balance = await provider.get_balance()
|
||||
|
||||
if balance is None:
|
||||
logger.warning(
|
||||
"Could not fetch balance for auto top-up",
|
||||
extra={"provider_id": row.id, "base_url": row.base_url},
|
||||
)
|
||||
return
|
||||
|
||||
if balance >= threshold * 1000:
|
||||
return
|
||||
|
||||
# Balance is below threshold - create token and top up
|
||||
logger.info(
|
||||
"Auto top-up triggered",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"balance": balance,
|
||||
"threshold": threshold,
|
||||
"topup_amount": amount,
|
||||
"mint_url": mint_url,
|
||||
},
|
||||
)
|
||||
|
||||
print(amount, mint_url)
|
||||
try:
|
||||
token = await send_token(amount, "sat", mint_url)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to create cashu token for auto top-up",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"amount": amount,
|
||||
"mint_url": mint_url,
|
||||
"error": str(e),
|
||||
},
|
||||
)
|
||||
return
|
||||
|
||||
result = await provider.topup(token)
|
||||
|
||||
if "error" in result:
|
||||
logger.error(
|
||||
"Auto top-up upstream call failed",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"error": result["error"],
|
||||
},
|
||||
)
|
||||
else:
|
||||
logger.info(
|
||||
"Auto top-up completed successfully",
|
||||
extra={
|
||||
"provider_id": row.id,
|
||||
"amount": amount,
|
||||
"new_balance_approx": balance + amount,
|
||||
},
|
||||
)
|
||||
+11
-46
@@ -4,7 +4,6 @@ from .base import BaseUpstreamProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.db import UpstreamProviderRow
|
||||
from ..payment.models import Model
|
||||
|
||||
|
||||
class AzureUpstreamProvider(BaseUpstreamProvider):
|
||||
@@ -13,7 +12,6 @@ class AzureUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "azure"
|
||||
default_base_url = None
|
||||
platform_url = "https://portal.azure.com/"
|
||||
litellm_provider_prefix = "azure/"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -60,52 +58,19 @@ class AzureUpstreamProvider(BaseUpstreamProvider):
|
||||
"platform_url": cls.platform_url,
|
||||
}
|
||||
|
||||
def prepare_headers(self, request_headers: dict) -> dict:
|
||||
"""Prepare headers for Azure OpenAI, adding api-key."""
|
||||
headers = super().prepare_headers(request_headers)
|
||||
if self.api_key:
|
||||
headers["api-key"] = self.api_key
|
||||
headers.pop("Authorization", None)
|
||||
headers.pop("authorization", None)
|
||||
return headers
|
||||
|
||||
def prepare_params(
|
||||
self, path: str, query_params: Mapping[str, str] | None
|
||||
) -> Mapping[str, str]:
|
||||
"""Prepare query parameters for Azure OpenAI, adding API version."""
|
||||
"""Prepare query parameters for Azure OpenAI, adding API version.
|
||||
|
||||
Args:
|
||||
path: Request path
|
||||
query_params: Original query parameters from the client
|
||||
|
||||
Returns:
|
||||
Query parameters dict with Azure API version added for chat completions
|
||||
"""
|
||||
params = dict(query_params or {})
|
||||
version = (self.api_version or "").replace("\ufeff", "").strip()
|
||||
if not version or version.lower() == "v1":
|
||||
version = "2024-02-15-preview"
|
||||
params["api-version"] = version
|
||||
if path.endswith("chat/completions"):
|
||||
params["api-version"] = self.api_version
|
||||
return params
|
||||
|
||||
def normalize_request_path(
|
||||
self, path: str, model_obj: "Model | None" = None
|
||||
) -> str:
|
||||
"""Build Azure deployment-specific request path."""
|
||||
clean_path = super().normalize_request_path(path, model_obj).lstrip("/")
|
||||
if model_obj is None:
|
||||
return clean_path
|
||||
|
||||
deployment_id = getattr(
|
||||
model_obj, "canonical_slug", None
|
||||
) or self.transform_model_name(model_obj.id)
|
||||
deployment_id = deployment_id.split("/")[-1]
|
||||
return f"openai/deployments/{deployment_id}/{clean_path}"
|
||||
|
||||
def get_request_base_url(
|
||||
self, path: str, model_obj: "Model | None" = None
|
||||
) -> str:
|
||||
"""Use endpoint root, stripping accidental /openai/v1 suffix if present."""
|
||||
base_url = self.base_url.rstrip("/")
|
||||
marker = "/openai/v1"
|
||||
if marker in base_url:
|
||||
base_url = base_url.split(marker, 1)[0].rstrip("/")
|
||||
return base_url
|
||||
|
||||
def transform_model_name(self, model_id: str) -> str:
|
||||
"""Extract deployment name from model ID."""
|
||||
if "/" in model_id:
|
||||
return model_id.split("/")[-1]
|
||||
return model_id
|
||||
|
||||
+238
-1778
File diff suppressed because it is too large
Load Diff
@@ -26,14 +26,14 @@ class GeminiClient(BaseAPIClient):
|
||||
max_tokens: int | None = None,
|
||||
**kwargs: Any,
|
||||
) -> dict[str, Any]:
|
||||
from openai import omit
|
||||
from openai import NOT_GIVEN
|
||||
|
||||
response = await self.client.chat.completions.create(
|
||||
model=model,
|
||||
messages=messages, # type: ignore
|
||||
temperature=temperature if temperature is not None else omit,
|
||||
max_tokens=max_tokens if max_tokens is not None else omit,
|
||||
top_p=kwargs.get("top_p", omit),
|
||||
temperature=temperature if temperature is not None else NOT_GIVEN,
|
||||
max_tokens=max_tokens if max_tokens is not None else NOT_GIVEN,
|
||||
top_p=kwargs.get("top_p", NOT_GIVEN),
|
||||
)
|
||||
return response.model_dump()
|
||||
|
||||
@@ -45,7 +45,7 @@ class GeminiClient(BaseAPIClient):
|
||||
max_tokens: int | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncGenerator[dict[str, Any], None]:
|
||||
from openai import omit
|
||||
from openai import NOT_GIVEN
|
||||
|
||||
usage_callback = kwargs.get("usage_callback")
|
||||
completion_callback = kwargs.get("completion_callback")
|
||||
@@ -55,9 +55,9 @@ class GeminiClient(BaseAPIClient):
|
||||
messages=messages, # type: ignore
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
temperature=temperature if temperature is not None else omit,
|
||||
max_tokens=max_tokens if max_tokens is not None else omit,
|
||||
top_p=kwargs.get("top_p", omit),
|
||||
temperature=temperature if temperature is not None else NOT_GIVEN,
|
||||
max_tokens=max_tokens if max_tokens is not None else NOT_GIVEN,
|
||||
top_p=kwargs.get("top_p", NOT_GIVEN),
|
||||
)
|
||||
|
||||
final_usage = None
|
||||
|
||||
@@ -12,7 +12,6 @@ class FireworksUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "fireworks"
|
||||
default_base_url = "https://api.fireworks.ai/inference/v1"
|
||||
platform_url = "https://app.fireworks.ai/settings/users/api-keys"
|
||||
litellm_provider_prefix = "fireworks_ai/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.01):
|
||||
super().__init__(
|
||||
|
||||
+235
-67
@@ -1,13 +1,17 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
from . import gemini_messages
|
||||
from fastapi import Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from .base import BaseUpstreamProvider
|
||||
from .clients.gemini import GeminiClient
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.db import UpstreamProviderRow
|
||||
from ..core.db import ApiKey, AsyncSession, UpstreamProviderRow
|
||||
from ..payment.models import Model
|
||||
|
||||
from ..core.logging import get_logger
|
||||
@@ -16,19 +20,9 @@ logger = get_logger(__name__)
|
||||
|
||||
|
||||
class GeminiUpstreamProvider(BaseUpstreamProvider):
|
||||
"""Gemini provider — proxies through Gemini's OpenAI-compat surface.
|
||||
|
||||
The chat-completions, embeddings, and models paths all flow through
|
||||
:meth:`BaseUpstreamProvider.forward_request`; we only override
|
||||
``get_request_base_url`` to redirect to ``{base}/openai/...`` and
|
||||
``_dispatch_anthropic_messages`` to inject thought-signatures on the
|
||||
/v1/messages path (see :mod:`gemini_messages` for that rationale).
|
||||
"""
|
||||
|
||||
provider_type = "gemini"
|
||||
default_base_url = "https://generativelanguage.googleapis.com/v1beta"
|
||||
platform_url = "https://aistudio.google.com/app/apikey"
|
||||
litellm_provider_prefix = "gemini/"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -45,7 +39,7 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
|
||||
|
||||
@property
|
||||
def client(self) -> GeminiClient:
|
||||
"""Get or create the Gemini API client (used for the models listing)."""
|
||||
"""Get or create the Gemini API client."""
|
||||
if self._client is None:
|
||||
self._client = GeminiClient(api_key=self.api_key)
|
||||
return self._client
|
||||
@@ -71,66 +65,240 @@ class GeminiUpstreamProvider(BaseUpstreamProvider):
|
||||
}
|
||||
|
||||
def transform_model_name(self, model_id: str) -> str:
|
||||
"""Reduce a routstr model id to the bare upstream Gemini name.
|
||||
return model_id.removeprefix("gemini/")
|
||||
|
||||
Gemini's OpenAI-compat surface expects the literal model id
|
||||
(e.g. ``gemini-3.1-flash-lite-preview``) — no ``gemini/`` provider
|
||||
prefix and no ``google/`` vendor sub-prefix. Take the last path
|
||||
segment so we tolerate any of:
|
||||
|
||||
``gemini-2.0-flash``
|
||||
``gemini/gemini-2.0-flash``
|
||||
``gemini/google/gemini-3.1-flash-lite-preview``
|
||||
"""
|
||||
return model_id.rsplit("/", 1)[-1]
|
||||
|
||||
@property
|
||||
def compat_base_url(self) -> str:
|
||||
"""Gemini's OpenAI-compat surface, regardless of what's stored.
|
||||
|
||||
Stored ``base_url`` may be ``.../v1beta`` (the native Gemini API
|
||||
root) or ``.../v1beta/openai`` (already pointed at the compat
|
||||
surface). Normalize to the latter.
|
||||
"""
|
||||
return self.base_url.rstrip("/").removesuffix("/openai") + "/openai"
|
||||
|
||||
def get_request_base_url(
|
||||
self, path: str, model_obj: "Model | None" = None
|
||||
) -> str:
|
||||
"""Route every proxied request to the OpenAI-compat surface.
|
||||
|
||||
Required because the stored ``base_url`` typically points at the
|
||||
native Gemini API (``/v1beta``), but :meth:`forward_request`
|
||||
forwards OpenAI-shaped paths (``/chat/completions``,
|
||||
``/embeddings``, ``/models``) which only exist under the
|
||||
``/openai`` subtree.
|
||||
"""
|
||||
return self.compat_base_url
|
||||
|
||||
async def _dispatch_anthropic_messages(
|
||||
async def forward_request(
|
||||
self,
|
||||
request: Request,
|
||||
path: str,
|
||||
headers: dict,
|
||||
request_body: bytes | None,
|
||||
model_obj: "Model",
|
||||
*,
|
||||
log_extra: dict[str, Any] | None = None,
|
||||
) -> tuple[bool, Any, str | None]:
|
||||
"""Dispatch /v1/messages via the gemini-specific httpx path.
|
||||
key: ApiKey,
|
||||
max_cost_for_model: int,
|
||||
session: AsyncSession,
|
||||
model_obj: Model,
|
||||
) -> Response | StreamingResponse:
|
||||
# Remove provider prefix from model ID for Gemini API
|
||||
if "/" in model_obj.id:
|
||||
model_obj.id = model_obj.id.split("/", 1)[1]
|
||||
|
||||
See :mod:`routstr.upstream.gemini_messages` for the full rationale
|
||||
(thought-signature injection, why litellm + the openai SDK can't
|
||||
carry the required ``extra_content`` field).
|
||||
"""
|
||||
return await gemini_messages.dispatch_gemini_messages(
|
||||
request_body=request_body,
|
||||
model_obj=model_obj,
|
||||
base_url=self.compat_base_url,
|
||||
api_key=self.api_key,
|
||||
transform_model_name=self.transform_model_name,
|
||||
log_extra=log_extra,
|
||||
)
|
||||
if not path.startswith("chat/completions"):
|
||||
return await super().forward_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
|
||||
if not request_body:
|
||||
return await super().forward_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
|
||||
try:
|
||||
openai_data = json.loads(request_body)
|
||||
messages = openai_data.get("messages", [])
|
||||
temperature = openai_data.get("temperature")
|
||||
max_tokens = openai_data.get("max_tokens")
|
||||
top_p = openai_data.get("top_p")
|
||||
is_streaming = openai_data.get("stream", False)
|
||||
|
||||
logger.info(
|
||||
"Processing Gemini request with client abstraction",
|
||||
extra={
|
||||
"model": model_obj.id,
|
||||
"is_streaming": is_streaming,
|
||||
"message_count": len(messages),
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
|
||||
if is_streaming:
|
||||
final_usage_data: dict | None = None
|
||||
|
||||
def usage_callback(usage_data: dict[str, Any]) -> None:
|
||||
"""Callback to capture usage data during streaming"""
|
||||
nonlocal final_usage_data
|
||||
final_usage_data = usage_data
|
||||
|
||||
async def completion_callback(
|
||||
model: str, usage_data: dict[str, Any] | None
|
||||
) -> None:
|
||||
"""Callback to handle payment when streaming completes"""
|
||||
nonlocal final_usage_data
|
||||
if usage_data:
|
||||
final_usage_data = usage_data
|
||||
|
||||
payment_data = {
|
||||
"model": model,
|
||||
"usage": final_usage_data,
|
||||
}
|
||||
|
||||
from ..auth import adjust_payment_for_tokens
|
||||
from ..core.db import create_session
|
||||
|
||||
async with create_session() as new_session:
|
||||
fresh_key = await new_session.get(key.__class__, key.hashed_key)
|
||||
if fresh_key:
|
||||
try:
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
fresh_key,
|
||||
payment_data,
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
|
||||
logger.info(
|
||||
"Gemini streaming payment finalized",
|
||||
extra={
|
||||
"cost_data": cost_data,
|
||||
"usage_data": final_usage_data,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
except Exception as cost_error:
|
||||
logger.error(
|
||||
"Error finalizing Gemini streaming payment",
|
||||
extra={
|
||||
"error": str(cost_error),
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
|
||||
response_generator = self.client.generate_content_stream(
|
||||
model=model_obj.id,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
top_p=top_p,
|
||||
usage_callback=usage_callback,
|
||||
completion_callback=completion_callback,
|
||||
)
|
||||
|
||||
async def stream_with_cost() -> AsyncGenerator[bytes, None]:
|
||||
payment_finalized = False
|
||||
|
||||
async def finalize_payment() -> None:
|
||||
nonlocal payment_finalized
|
||||
if payment_finalized:
|
||||
return
|
||||
from ..auth import adjust_payment_for_tokens
|
||||
from ..core.db import create_session
|
||||
|
||||
async with create_session() as new_session:
|
||||
fresh_key = await new_session.get(
|
||||
key.__class__, key.hashed_key
|
||||
)
|
||||
if fresh_key:
|
||||
try:
|
||||
await adjust_payment_for_tokens(
|
||||
fresh_key,
|
||||
{
|
||||
"model": model_obj.id,
|
||||
"usage": final_usage_data,
|
||||
},
|
||||
new_session,
|
||||
max_cost_for_model,
|
||||
)
|
||||
payment_finalized = True
|
||||
except Exception as cost_error:
|
||||
logger.error(
|
||||
"Error finalizing Gemini streaming payment in fallback",
|
||||
extra={
|
||||
"error": str(cost_error),
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
async for chunk in response_generator:
|
||||
sse_data = f"data: {json.dumps(chunk)}\n\n"
|
||||
yield sse_data.encode()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error in Gemini streaming response",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
raise
|
||||
finally:
|
||||
if not payment_finalized:
|
||||
await finalize_payment()
|
||||
|
||||
return StreamingResponse(
|
||||
stream_with_cost(),
|
||||
media_type="text/event-stream",
|
||||
headers={"Cache-Control": "no-cache", "Connection": "keep-alive"},
|
||||
)
|
||||
|
||||
else:
|
||||
openai_format_response = await self.client.generate_content(
|
||||
model=model_obj.id,
|
||||
messages=messages,
|
||||
temperature=temperature,
|
||||
max_tokens=max_tokens,
|
||||
top_p=top_p,
|
||||
)
|
||||
|
||||
from ..auth import adjust_payment_for_tokens
|
||||
|
||||
cost_data = await adjust_payment_for_tokens(
|
||||
key, openai_format_response, session, max_cost_for_model
|
||||
)
|
||||
openai_format_response["cost"] = cost_data
|
||||
|
||||
logger.info(
|
||||
"Gemini non-streaming payment completed",
|
||||
extra={
|
||||
"cost_data": cost_data,
|
||||
"model": model_obj.id,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
|
||||
return Response(
|
||||
content=json.dumps(openai_format_response),
|
||||
media_type="application/json",
|
||||
headers={"Cache-Control": "no-cache"},
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error in Gemini forward_request",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"path": path,
|
||||
"key_hash": key.hashed_key[:8] + "...",
|
||||
},
|
||||
)
|
||||
return await super().forward_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
|
||||
async def _fetch_provider_models(self) -> dict:
|
||||
"""Fetch models from Gemini API via the OpenAI-compat client."""
|
||||
"""Fetch models from Gemini API."""
|
||||
try:
|
||||
models_data = await self.client.list_models()
|
||||
|
||||
|
||||
@@ -1,480 +0,0 @@
|
||||
"""Custom /v1/messages dispatcher for Gemini's OpenAI-compat endpoint.
|
||||
|
||||
Why this exists
|
||||
---------------
|
||||
|
||||
Gemini 2.5 / 3 thinking models reject inbound ``functionCall`` parts that
|
||||
lack a ``thought_signature`` field once any prior turn in the conversation
|
||||
contains a function call. Anthropic-Messages clients (Claude Code etc.)
|
||||
have no concept of thought signatures, so multi-turn tool conversations
|
||||
fail with::
|
||||
|
||||
Function call is missing a thought_signature in functionCall parts.
|
||||
|
||||
Google's published escape hatch (https://ai.google.dev/gemini-api/docs/
|
||||
thought-signatures, FAQ #1) is the dummy signature
|
||||
``"skip_thought_signature_validator"`` placed at
|
||||
``tool_calls[i].extra_content.google.thought_signature`` for every tool
|
||||
call in the request. The hatch is documented specifically for
|
||||
"transferring a trace from a different model that does not include thought
|
||||
signatures" — exactly our case.
|
||||
|
||||
Why we can't reach the wire via litellm
|
||||
---------------------------------------
|
||||
|
||||
``litellm.anthropic.messages.acreate`` flows through the openai SDK, whose
|
||||
pydantic ``ChatCompletionMessageToolCall`` model silently drops unknown
|
||||
fields like ``extra_content``. Litellm has no openai-compat translator
|
||||
that emits ``extra_content.google.thought_signature``. So we bypass both
|
||||
litellm and the openai SDK at the transport layer.
|
||||
|
||||
Pipeline
|
||||
--------
|
||||
|
||||
1. Translate Anthropic body → OpenAI body via litellm's
|
||||
``AnthropicAdapter`` (the same translator
|
||||
``litellm.anthropic.messages.acreate`` uses internally).
|
||||
2. Inject ``extra_content.google.thought_signature`` on every
|
||||
``tool_calls[]`` entry.
|
||||
3. Set ``reasoning_effort="none"`` to disable Gemini's thinking pass.
|
||||
4. POST directly to ``{base_url}/chat/completions`` with ``stream=true``
|
||||
via ``httpx`` (preserves arbitrary fields verbatim).
|
||||
5. Translate OpenAI streaming chunks → Anthropic SSE events.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
from typing import Any, Callable
|
||||
|
||||
import httpx
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.exceptions import UpstreamError
|
||||
from ..payment.models import Model
|
||||
from .messages_dispatch import (
|
||||
ANTHROPIC_ONLY_FIELDS,
|
||||
aggregate_anthropic_events_to_message,
|
||||
)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
DUMMY_THOUGHT_SIGNATURE = "skip_thought_signature_validator"
|
||||
|
||||
# Mapping: OpenAI finish_reason → Anthropic stop_reason
|
||||
_FINISH_TO_STOP = {
|
||||
"stop": "end_turn",
|
||||
"length": "max_tokens",
|
||||
"tool_calls": "tool_use",
|
||||
"function_call": "tool_use",
|
||||
"content_filter": "refusal",
|
||||
}
|
||||
|
||||
|
||||
def inject_thought_signatures(messages: list[dict]) -> None:
|
||||
"""Add ``extra_content.google.thought_signature`` to every tool_call.
|
||||
|
||||
Mutates ``messages`` in place. Idempotent: existing signatures are not
|
||||
overwritten.
|
||||
"""
|
||||
for msg in messages:
|
||||
tool_calls = msg.get("tool_calls")
|
||||
if not isinstance(tool_calls, list):
|
||||
continue
|
||||
for tc in tool_calls:
|
||||
if not isinstance(tc, dict):
|
||||
continue
|
||||
extra = tc.get("extra_content")
|
||||
if not isinstance(extra, dict):
|
||||
extra = tc["extra_content"] = {}
|
||||
google_cfg = extra.get("google")
|
||||
if not isinstance(google_cfg, dict):
|
||||
google_cfg = extra["google"] = {}
|
||||
google_cfg.setdefault("thought_signature", DUMMY_THOUGHT_SIGNATURE)
|
||||
|
||||
|
||||
def _translate_anthropic_to_openai(body: dict, model: str) -> dict:
|
||||
"""Use litellm's translator to convert an Anthropic /messages body to
|
||||
OpenAI /chat/completions kwargs.
|
||||
|
||||
Imported lazily because the litellm internal path is heavy and not
|
||||
needed for any other code path in routstr.
|
||||
"""
|
||||
from litellm.llms.anthropic.experimental_pass_through.adapters.transformation import ( # noqa: E501
|
||||
AnthropicAdapter,
|
||||
)
|
||||
|
||||
kwargs = {"model": model, **body}
|
||||
translated = AnthropicAdapter().translate_completion_input_params(kwargs)
|
||||
if translated is None:
|
||||
raise UpstreamError(
|
||||
"Failed to translate Anthropic body to OpenAI format",
|
||||
status_code=500,
|
||||
)
|
||||
return dict(translated)
|
||||
|
||||
|
||||
def _sse_event(event_type: str, payload: dict) -> bytes:
|
||||
return f"event: {event_type}\ndata: {json.dumps(payload)}\n\n".encode()
|
||||
|
||||
|
||||
async def _openai_chunks_to_anthropic_events(
|
||||
line_iter: AsyncIterator[str], requested_model: str | None
|
||||
) -> AsyncGenerator[bytes, None]:
|
||||
"""Translate an OpenAI chat-completions SSE byte stream into the
|
||||
Anthropic-Messages SSE event sequence.
|
||||
|
||||
Maintains per-chunk state across:
|
||||
|
||||
* one optional text content block (lazy-opened on first text delta)
|
||||
* any number of tool_use blocks indexed by openai's ``delta.tool_calls[].index``
|
||||
* final ``stop_reason`` / ``usage`` carried out via ``message_delta`` /
|
||||
``message_stop``
|
||||
"""
|
||||
msg_id = f"msg_{uuid.uuid4().hex[:24]}"
|
||||
started = False
|
||||
text_block_idx: int | None = None
|
||||
tool_block_indices: dict[int, int] = {}
|
||||
next_block_idx = 0
|
||||
final_finish_reason: str | None = None
|
||||
final_usage: dict[str, int] = {"input_tokens": 0, "output_tokens": 0}
|
||||
|
||||
def open_text_block() -> bytes:
|
||||
nonlocal text_block_idx, next_block_idx
|
||||
text_block_idx = next_block_idx
|
||||
next_block_idx += 1
|
||||
return _sse_event(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": text_block_idx,
|
||||
"content_block": {"type": "text", "text": ""},
|
||||
},
|
||||
)
|
||||
|
||||
def open_tool_block(delta_idx: int, tc: dict) -> bytes:
|
||||
nonlocal next_block_idx
|
||||
block_idx = next_block_idx
|
||||
next_block_idx += 1
|
||||
tool_block_indices[delta_idx] = block_idx
|
||||
fn = tc.get("function") or {}
|
||||
return _sse_event(
|
||||
"content_block_start",
|
||||
{
|
||||
"type": "content_block_start",
|
||||
"index": block_idx,
|
||||
"content_block": {
|
||||
"type": "tool_use",
|
||||
"id": tc.get("id") or f"toolu_{uuid.uuid4().hex[:24]}",
|
||||
"name": fn.get("name") or "",
|
||||
"input": {},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
def close_block(idx: int) -> bytes:
|
||||
return _sse_event(
|
||||
"content_block_stop",
|
||||
{"type": "content_block_stop", "index": idx},
|
||||
)
|
||||
|
||||
async for raw_line in line_iter:
|
||||
line = raw_line.strip()
|
||||
if not line:
|
||||
continue
|
||||
if not line.startswith("data:"):
|
||||
continue
|
||||
payload = line[5:].lstrip()
|
||||
if not payload or payload == "[DONE]":
|
||||
continue
|
||||
try:
|
||||
chunk = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if not isinstance(chunk, dict):
|
||||
continue
|
||||
|
||||
if not started:
|
||||
started = True
|
||||
yield _sse_event(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": chunk.get("id") or msg_id,
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": requested_model or chunk.get("model") or "",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"stop_sequence": None,
|
||||
"usage": {
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0,
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
usage = chunk.get("usage")
|
||||
if isinstance(usage, dict):
|
||||
in_tok = usage.get("prompt_tokens") or usage.get("input_tokens") or 0
|
||||
out_tok = usage.get("completion_tokens") or usage.get("output_tokens") or 0
|
||||
if in_tok:
|
||||
final_usage["input_tokens"] = int(in_tok)
|
||||
if out_tok:
|
||||
final_usage["output_tokens"] = int(out_tok)
|
||||
|
||||
choices = chunk.get("choices") or []
|
||||
if not choices:
|
||||
continue
|
||||
choice = choices[0] if isinstance(choices[0], dict) else {}
|
||||
delta = choice.get("delta") or {}
|
||||
if not isinstance(delta, dict):
|
||||
delta = {}
|
||||
|
||||
# Text delta
|
||||
text = delta.get("content")
|
||||
if isinstance(text, str) and text:
|
||||
if text_block_idx is None:
|
||||
yield open_text_block()
|
||||
yield _sse_event(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": text_block_idx,
|
||||
"delta": {"type": "text_delta", "text": text},
|
||||
},
|
||||
)
|
||||
|
||||
# Tool call deltas
|
||||
tool_calls_delta = delta.get("tool_calls")
|
||||
if isinstance(tool_calls_delta, list):
|
||||
for tc in tool_calls_delta:
|
||||
if not isinstance(tc, dict):
|
||||
continue
|
||||
d_idx = int(tc.get("index") or 0)
|
||||
if d_idx not in tool_block_indices:
|
||||
yield open_tool_block(d_idx, tc)
|
||||
block_idx = tool_block_indices[d_idx]
|
||||
fn = tc.get("function") or {}
|
||||
args = fn.get("arguments")
|
||||
if isinstance(args, str) and args:
|
||||
yield _sse_event(
|
||||
"content_block_delta",
|
||||
{
|
||||
"type": "content_block_delta",
|
||||
"index": block_idx,
|
||||
"delta": {
|
||||
"type": "input_json_delta",
|
||||
"partial_json": args,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
finish = choice.get("finish_reason")
|
||||
if finish:
|
||||
final_finish_reason = finish
|
||||
|
||||
# Close any open content blocks
|
||||
if text_block_idx is not None:
|
||||
yield close_block(text_block_idx)
|
||||
for block_idx in tool_block_indices.values():
|
||||
yield close_block(block_idx)
|
||||
|
||||
# message_delta with stop_reason and usage
|
||||
stop_reason = _FINISH_TO_STOP.get(final_finish_reason or "", "end_turn")
|
||||
yield _sse_event(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
|
||||
"usage": final_usage,
|
||||
},
|
||||
)
|
||||
yield _sse_event("message_stop", {"type": "message_stop"})
|
||||
|
||||
|
||||
async def _post_and_stream(
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
payload: dict,
|
||||
log_extra: dict[str, Any] | None,
|
||||
) -> tuple[httpx.AsyncClient, httpx.Response]:
|
||||
"""POST to upstream chat-completions and return (client, response) for
|
||||
streaming. Caller is responsible for closing both."""
|
||||
url = f"{base_url.rstrip('/')}/chat/completions"
|
||||
client = httpx.AsyncClient(timeout=httpx.Timeout(120.0, read=120.0))
|
||||
try:
|
||||
request = client.build_request(
|
||||
"POST",
|
||||
url,
|
||||
json=payload,
|
||||
headers={
|
||||
"Authorization": f"Bearer {api_key}",
|
||||
"Content-Type": "application/json",
|
||||
"Accept": "text/event-stream",
|
||||
},
|
||||
)
|
||||
response = await client.send(request, stream=True)
|
||||
except Exception as exc:
|
||||
await client.aclose()
|
||||
logger.error(
|
||||
"Gemini messages dispatch HTTP error",
|
||||
extra={"error": str(exc), "url": url, **(log_extra or {})},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Failed to reach Gemini upstream: {exc}", status_code=502
|
||||
) from exc
|
||||
|
||||
if response.status_code >= 400:
|
||||
try:
|
||||
body_bytes = await response.aread()
|
||||
finally:
|
||||
await response.aclose()
|
||||
await client.aclose()
|
||||
body_text = body_bytes.decode("utf-8", errors="replace")
|
||||
logger.error(
|
||||
"Gemini messages dispatch upstream error",
|
||||
extra={
|
||||
"status_code": response.status_code,
|
||||
"body": body_text[:1000],
|
||||
"url": url,
|
||||
**(log_extra or {}),
|
||||
},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Upstream error via gemini compat: {body_text}",
|
||||
status_code=response.status_code,
|
||||
)
|
||||
|
||||
return client, response
|
||||
|
||||
|
||||
async def dispatch_gemini_messages(
|
||||
*,
|
||||
request_body: bytes | None,
|
||||
model_obj: Model,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
transform_model_name: Callable[[str], str],
|
||||
log_extra: dict[str, Any] | None = None,
|
||||
) -> tuple[bool, Any, str | None]:
|
||||
"""Dispatch a /v1/messages request to Gemini's OpenAI-compat endpoint
|
||||
with thought-signature injection.
|
||||
|
||||
Returns ``(client_stream, result, requested_model)`` where ``result``
|
||||
is either an ``AsyncIterator[bytes]`` of Anthropic-format SSE events
|
||||
(for streaming clients) or an Anthropic Message dict (after the caller
|
||||
aggregates).
|
||||
"""
|
||||
if not request_body:
|
||||
raise UpstreamError(
|
||||
"Missing request body for /v1/messages", status_code=400
|
||||
)
|
||||
|
||||
try:
|
||||
body: dict = json.loads(request_body)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise UpstreamError(
|
||||
f"Invalid JSON in /v1/messages body: {exc}", status_code=400
|
||||
) from exc
|
||||
|
||||
body.pop("model", None)
|
||||
client_stream = bool(body.pop("stream", False))
|
||||
|
||||
# Anthropic-Messages-only fields that don't translate to OpenAI
|
||||
# Chat Completions. litellm's translator passes through unknown
|
||||
# top-level fields verbatim and Gemini's compat surface 400s on
|
||||
# unknown names like ``context_management`` / ``output_config``.
|
||||
dropped: dict[str, Any] = {}
|
||||
for field in ANTHROPIC_ONLY_FIELDS:
|
||||
if field in body:
|
||||
dropped[field] = body.pop(field)
|
||||
if dropped:
|
||||
logger.debug(
|
||||
"Dropped anthropic-only fields before gemini compat dispatch",
|
||||
extra={"dropped_keys": sorted(dropped.keys())},
|
||||
)
|
||||
|
||||
requested_model = (
|
||||
(model_obj.forwarded_model_id or model_obj.id) if model_obj else None
|
||||
)
|
||||
upstream_model = transform_model_name(model_obj.id)
|
||||
|
||||
openai_kwargs = _translate_anthropic_to_openai(body, upstream_model)
|
||||
|
||||
messages = openai_kwargs.get("messages") or []
|
||||
if isinstance(messages, list):
|
||||
inject_thought_signatures(messages)
|
||||
|
||||
# Disable Gemini's thinking pass; the dummy signature already lifts
|
||||
# validation, but skipping thinking entirely avoids degraded model
|
||||
# output and keeps tool-calling deterministic.
|
||||
openai_kwargs.setdefault("reasoning_effort", "none")
|
||||
openai_kwargs["stream"] = True
|
||||
openai_kwargs["model"] = upstream_model
|
||||
# OpenAI-compat backends (including Gemini's) only emit a final
|
||||
# ``usage`` chunk when the request opts in via this flag. Without it
|
||||
# the cost-calculation pipeline can't read real token counts and
|
||||
# falls back to MaxCostData billing.
|
||||
existing_stream_options = openai_kwargs.get("stream_options")
|
||||
merged_stream_options = (
|
||||
dict(existing_stream_options)
|
||||
if isinstance(existing_stream_options, dict)
|
||||
else {}
|
||||
)
|
||||
merged_stream_options.setdefault("include_usage", True)
|
||||
openai_kwargs["stream_options"] = merged_stream_options
|
||||
|
||||
logger.info(
|
||||
"Dispatching /v1/messages via gemini compat (httpx)",
|
||||
extra={
|
||||
"model": upstream_model,
|
||||
"client_stream": client_stream,
|
||||
"messages_with_tool_calls": sum(
|
||||
1 for m in messages if isinstance(m, dict) and m.get("tool_calls")
|
||||
),
|
||||
**(log_extra or {}),
|
||||
},
|
||||
)
|
||||
|
||||
http_client, response = await _post_and_stream(
|
||||
base_url, api_key, openai_kwargs, log_extra
|
||||
)
|
||||
|
||||
async def line_iter() -> AsyncGenerator[str, None]:
|
||||
try:
|
||||
async for line in response.aiter_lines():
|
||||
yield line
|
||||
finally:
|
||||
await response.aclose()
|
||||
await http_client.aclose()
|
||||
|
||||
anthropic_event_iter = _openai_chunks_to_anthropic_events(
|
||||
line_iter(), requested_model
|
||||
)
|
||||
|
||||
if not client_stream:
|
||||
# Aggregate the Anthropic SSE byte stream into a single Message dict
|
||||
# so the rest of the pipeline (cost calc, metadata injection,
|
||||
# response building) can treat it identically to a non-streaming
|
||||
# litellm response.
|
||||
try:
|
||||
aggregated = await aggregate_anthropic_events_to_message(
|
||||
anthropic_event_iter
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to aggregate Gemini compat events into message",
|
||||
extra={"error": str(exc), **(log_extra or {})},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Failed to aggregate upstream stream: {exc}",
|
||||
status_code=502,
|
||||
) from exc
|
||||
return client_stream, aggregated, requested_model
|
||||
|
||||
return client_stream, anthropic_event_iter, requested_model
|
||||
@@ -12,7 +12,6 @@ class GroqUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "groq"
|
||||
default_base_url = "https://api.groq.com/openai/v1"
|
||||
platform_url = "https://console.groq.com/keys"
|
||||
litellm_provider_prefix = "groq/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.01):
|
||||
super().__init__(
|
||||
|
||||
@@ -3,7 +3,7 @@ from __future__ import annotations
|
||||
import asyncio
|
||||
import os
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Callable
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.settings import Settings
|
||||
@@ -122,16 +122,12 @@ async def get_all_models_with_overrides(
|
||||
|
||||
|
||||
async def refresh_upstreams_models_periodically(
|
||||
upstreams_provider: (
|
||||
Callable[[], list[BaseUpstreamProvider]] | list[BaseUpstreamProvider]
|
||||
),
|
||||
upstreams: list[BaseUpstreamProvider],
|
||||
) -> None:
|
||||
"""Background task to periodically refresh models cache for all providers.
|
||||
|
||||
Args:
|
||||
upstreams_provider: Either a callable returning the live upstream list
|
||||
(preferred — picks up providers added/changed via reinitialize_upstreams),
|
||||
or a static list (legacy, will go stale after reinitialize_upstreams).
|
||||
upstreams: List of upstream provider instances
|
||||
"""
|
||||
import asyncio
|
||||
import random
|
||||
@@ -143,14 +139,9 @@ async def refresh_upstreams_models_periodically(
|
||||
logger.info("Provider models refresh disabled (interval <= 0)")
|
||||
return
|
||||
|
||||
def _resolve_upstreams() -> list[BaseUpstreamProvider]:
|
||||
if callable(upstreams_provider):
|
||||
return upstreams_provider()
|
||||
return upstreams_provider
|
||||
|
||||
while True:
|
||||
try:
|
||||
for upstream in _resolve_upstreams():
|
||||
for upstream in upstreams:
|
||||
try:
|
||||
await upstream.refresh_models_cache()
|
||||
except Exception as e:
|
||||
@@ -197,14 +188,13 @@ async def init_upstreams() -> list[BaseUpstreamProvider]:
|
||||
existing_providers = result.all()
|
||||
|
||||
if not existing_providers:
|
||||
logger.info(
|
||||
"No upstream providers found in database, seeding from settings"
|
||||
)
|
||||
await _seed_providers_from_settings(session, settings)
|
||||
await session.commit()
|
||||
result = await session.exec(select(UpstreamProviderRow))
|
||||
existing_providers = result.all()
|
||||
if existing_providers:
|
||||
logger.info(
|
||||
f"Seeded {len(existing_providers)} upstream providers from settings"
|
||||
)
|
||||
|
||||
async def _init_single_provider(
|
||||
provider_row: UpstreamProviderRow,
|
||||
@@ -215,9 +205,6 @@ async def init_upstreams() -> list[BaseUpstreamProvider]:
|
||||
|
||||
provider = _instantiate_provider(provider_row)
|
||||
if provider:
|
||||
# Keep provider DB id on runtime instance so model mapping can
|
||||
# bind DB overrides to the correct upstream.
|
||||
setattr(provider, "db_id", provider_row.id)
|
||||
await provider.refresh_models_cache()
|
||||
logger.debug(
|
||||
f"Initialized {provider_row.provider_type} provider",
|
||||
|
||||
@@ -1,162 +0,0 @@
|
||||
"""Map an upstream `base_url` to the correct litellm provider prefix.
|
||||
|
||||
Used by `BaseUpstreamProvider.get_litellm_provider_prefix` so that custom /
|
||||
generic provider rows (which inherit the base class default) get routed to
|
||||
the right litellm backend instead of falling back to `openai/`.
|
||||
|
||||
The table is compiled from litellm 1.74's
|
||||
`litellm/litellm_core_utils/get_llm_provider_logic.py` (the
|
||||
`openai_compatible_endpoints` table) plus the providers documented at
|
||||
https://docs.litellm.ai/docs/providers. Substring match is used so that
|
||||
URLs with paths, ports, regional subdomains, etc. all resolve correctly.
|
||||
|
||||
Order matters: more specific needles must appear before more generic ones
|
||||
(e.g. ``openai.azure.com`` before ``api.openai.com``).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
import litellm
|
||||
|
||||
DEFAULT_PREFIX = "openai/"
|
||||
|
||||
LITELLM_HOST_PREFIX_MAP: tuple[tuple[str, str], ...] = (
|
||||
# Azure must win over api.openai.com because the host ends with
|
||||
# `openai.azure.com` and we don't want it picked up as plain OpenAI.
|
||||
("openai.azure.com", "azure/"),
|
||||
# Google
|
||||
("generativelanguage.googleapis.com", "gemini/"),
|
||||
("aiplatform.googleapis.com", "vertex_ai/"),
|
||||
# First-class providers with native litellm prefixes
|
||||
("api.openai.com", "openai/"),
|
||||
("api.anthropic.com", "anthropic/"),
|
||||
("api.groq.com", "groq/"),
|
||||
("api.fireworks.ai", "fireworks_ai/"),
|
||||
("api.x.ai", "xai/"),
|
||||
("api.perplexity.ai", "perplexity/"),
|
||||
("openrouter.ai", "openrouter/"),
|
||||
("api.deepseek.com", "deepseek/"),
|
||||
("api.together.xyz", "together_ai/"),
|
||||
("codestral.mistral.ai", "codestral/"),
|
||||
("api.mistral.ai", "mistral/"),
|
||||
("api.cohere.com", "cohere_chat/"),
|
||||
("api.cohere.ai", "cohere_chat/"),
|
||||
("api.deepinfra.com", "deepinfra/"),
|
||||
("api.endpoints.anyscale.com", "anyscale/"),
|
||||
("api.cerebras.ai", "cerebras/"),
|
||||
("inference.baseten.co", "baseten/"),
|
||||
("api.sambanova.ai", "sambanova/"),
|
||||
("api.ai21.com", "ai21_chat/"),
|
||||
("api.friendli.ai", "friendliai/"),
|
||||
("api.galadriel.com", "galadriel/"),
|
||||
("api.llama.com", "meta_llama/"),
|
||||
("api.featherless.ai", "featherless_ai/"),
|
||||
("inference.api.nscale.com", "nscale/"),
|
||||
("dashscope-intl.aliyuncs.com", "dashscope/"),
|
||||
("api.moonshot.ai", "moonshot/"),
|
||||
("api.moonshot.cn", "moonshot/"),
|
||||
("api.minimax.io", "minimax/"),
|
||||
("api.minimaxi.com", "minimax/"),
|
||||
("platform.publicai.co", "publicai/"),
|
||||
("api.synthetic.new", "synthetic/"),
|
||||
("api.stima.tech", "apertis/"),
|
||||
("nano-gpt.com", "nano-gpt/"),
|
||||
("api.poe.com", "poe/"),
|
||||
("llm.chutes.ai", "chutes/"),
|
||||
("api.v0.dev", "v0/"),
|
||||
("api.lambda.ai", "lambda_ai/"),
|
||||
("api.hyperbolic.xyz", "hyperbolic/"),
|
||||
("ai-gateway.vercel.sh", "vercel_ai_gateway/"),
|
||||
("api.inference.wandb.ai", "wandb/"),
|
||||
("integrate.api.nvidia.com", "nvidia_nim/"),
|
||||
("api.studio.nebius.com", "nebius/"),
|
||||
("api.novita.ai", "novita/"),
|
||||
("ark.cn-beijing.volces.com", "volcengine/"),
|
||||
("api.voyageai.com", "voyage/"),
|
||||
("api.jina.ai", "jina_ai/"),
|
||||
("api.aimlapi.com", "aiml/"),
|
||||
("api.snowflakecomputing.com", "snowflake/"),
|
||||
("databricks.com", "databricks/"),
|
||||
("huggingface.co", "huggingface/"),
|
||||
)
|
||||
|
||||
# Substrings that indicate an Ollama deployment regardless of port/scheme.
|
||||
OLLAMA_HOST_HINTS: tuple[str, ...] = (
|
||||
"localhost:11434",
|
||||
"127.0.0.1:11434",
|
||||
"ollama",
|
||||
)
|
||||
|
||||
|
||||
def detect_litellm_prefix(
|
||||
base_url: str | None, default: str = DEFAULT_PREFIX
|
||||
) -> str:
|
||||
"""Return the litellm provider prefix (`"<provider>/"`) for `base_url`.
|
||||
|
||||
Falls back to `default` when the host doesn't match any known provider.
|
||||
The default is `openai/` because every unmatched OpenAI-compatible
|
||||
server is, by definition, an OpenAI-compatible server.
|
||||
"""
|
||||
if not base_url:
|
||||
return default
|
||||
|
||||
parsed = urlsplit(base_url)
|
||||
host = parsed.netloc.lower() or base_url.lower()
|
||||
|
||||
for needle, prefix in LITELLM_HOST_PREFIX_MAP:
|
||||
if needle in host:
|
||||
return prefix
|
||||
|
||||
if any(hint in host for hint in OLLAMA_HOST_HINTS):
|
||||
return "ollama_chat/"
|
||||
|
||||
return default
|
||||
|
||||
|
||||
_configured = False
|
||||
|
||||
|
||||
def configure_litellm() -> None:
|
||||
"""Apply litellm global settings used by the messages-dispatch path.
|
||||
|
||||
Idempotent: safe to call from both app startup and module-level
|
||||
initializers without side effects on the second invocation.
|
||||
|
||||
Settings applied:
|
||||
|
||||
* ``LITELLM_DEBUG=1`` enables litellm's verbose debug logger.
|
||||
* Forces the Anthropic-messages adapter to call OpenAI Chat Completions
|
||||
(POST ``/chat/completions``) instead of the Responses API (POST
|
||||
``/responses``) for ``openai/``-prefixed providers. OpenAI-compatible
|
||||
upstreams like Google's generativelanguage compat endpoint expose
|
||||
``/chat/completions`` but not ``/responses``, which would 404. Set
|
||||
``LITELLM_USE_RESPONSES_API_FOR_ANTHROPIC_MESSAGES=1`` to opt out.
|
||||
* Silently drops Anthropic-Messages-only parameters (``thinking``,
|
||||
``cache_control``, ``context_management``, ...) when translating to
|
||||
providers that don't accept them, instead of raising
|
||||
``UnsupportedParamsError``. Set ``LITELLM_STRICT_PARAMS=1`` to opt
|
||||
out.
|
||||
"""
|
||||
global _configured
|
||||
if _configured:
|
||||
return
|
||||
|
||||
if os.getenv("LITELLM_DEBUG") == "1":
|
||||
try:
|
||||
litellm._turn_on_debug() # type: ignore[no-untyped-call]
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if os.getenv("LITELLM_USE_RESPONSES_API_FOR_ANTHROPIC_MESSAGES") != "1":
|
||||
try:
|
||||
litellm.use_chat_completions_url_for_anthropic_messages = True
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if os.getenv("LITELLM_STRICT_PARAMS") != "1":
|
||||
litellm.drop_params = True
|
||||
|
||||
_configured = True
|
||||
@@ -1,561 +0,0 @@
|
||||
"""Pure helpers for translating ``/v1/messages`` to upstream chat completions
|
||||
via litellm.
|
||||
|
||||
This module owns the litellm/Anthropic-Messages translation layer:
|
||||
|
||||
* SSE parsing (``parse_sse_blocks``, ``events_from_chunk``)
|
||||
* Payload coercion (``coerce_litellm_payload``)
|
||||
* Stream aggregation (``aggregate_anthropic_events_to_message``) — drains
|
||||
a streamed Anthropic event sequence into a single Message dict
|
||||
* Per-event annotation for streaming (``annotate_event``,
|
||||
``stream_annotated_events``) — handles the model-rewrite + token-tally
|
||||
bookkeeping shared by the bearer-key and x-cashu streaming paths
|
||||
* The dispatch entry point (``dispatch_anthropic_messages``)
|
||||
* Refund math (``compute_refund``)
|
||||
|
||||
Nothing in here touches ``BaseUpstreamProvider``; the thin instance methods
|
||||
on the provider class forward to these functions and only retain logic that
|
||||
genuinely needs ``self`` (cost adjustment, metadata injection, refund
|
||||
sending).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator, AsyncIterator
|
||||
from typing import Any, Callable, NamedTuple, cast
|
||||
|
||||
import litellm
|
||||
|
||||
from ..core import get_logger
|
||||
from ..core.exceptions import UpstreamError
|
||||
from ..payment.models import Model
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
# Anthropic-Messages-only fields that don't translate to OpenAI
|
||||
# Chat Completions. ``litellm.drop_params`` only filters *known*
|
||||
# unsupported params; these newer/extension fields get passed through
|
||||
# verbatim and the upstream rejects them with a 400. Pop them here so the
|
||||
# request reaches the upstream cleanly.
|
||||
ANTHROPIC_ONLY_FIELDS: tuple[str, ...] = (
|
||||
"thinking",
|
||||
"cache_control",
|
||||
"context_management",
|
||||
"output_config",
|
||||
"mcp_servers",
|
||||
"service_tier",
|
||||
"anthropic_version",
|
||||
"anthropic_beta",
|
||||
)
|
||||
|
||||
|
||||
def coerce_litellm_payload(payload: object) -> dict:
|
||||
"""Convert a litellm event into a plain dict.
|
||||
|
||||
Non-streaming responses come back as Anthropic-shaped pydantic models
|
||||
or dicts. Streaming may yield raw bytes/str (SSE-encoded); those go
|
||||
through ``events_from_chunk`` instead, not here.
|
||||
"""
|
||||
if isinstance(payload, dict):
|
||||
return dict(payload)
|
||||
if hasattr(payload, "model_dump"):
|
||||
return cast(dict, payload.model_dump())
|
||||
raise TypeError(f"Cannot coerce {type(payload).__name__} to dict")
|
||||
|
||||
|
||||
def parse_sse_blocks(buffer: bytes) -> tuple[list[dict], bytes]:
|
||||
"""Parse complete SSE event blocks out of a byte buffer.
|
||||
|
||||
Returns (events, remaining_buffer). Events are JSON objects parsed from
|
||||
one or more ``data:`` lines per block. Comments, blank lines, and
|
||||
``[DONE]`` sentinels are ignored. A trailing partial block is preserved
|
||||
in remaining_buffer.
|
||||
"""
|
||||
events: list[dict] = []
|
||||
while True:
|
||||
sep = buffer.find(b"\n\n")
|
||||
if sep < 0:
|
||||
sep_rn = buffer.find(b"\r\n\r\n")
|
||||
if sep_rn < 0:
|
||||
break
|
||||
block = buffer[:sep_rn]
|
||||
buffer = buffer[sep_rn + 4 :]
|
||||
else:
|
||||
block = buffer[:sep]
|
||||
buffer = buffer[sep + 2 :]
|
||||
|
||||
data_lines: list[str] = []
|
||||
for raw_line in block.replace(b"\r\n", b"\n").split(b"\n"):
|
||||
line = raw_line.decode("utf-8", errors="replace")
|
||||
if line.startswith(":"):
|
||||
continue
|
||||
if line.startswith("data:"):
|
||||
data_lines.append(line[5:].lstrip())
|
||||
if not data_lines:
|
||||
continue
|
||||
payload = "\n".join(data_lines).strip()
|
||||
if not payload or payload == "[DONE]":
|
||||
continue
|
||||
try:
|
||||
obj = json.loads(payload)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
if isinstance(obj, dict):
|
||||
events.append(obj)
|
||||
return events, buffer
|
||||
|
||||
|
||||
def events_from_chunk(
|
||||
chunk: object, sse_buffer: bytes
|
||||
) -> tuple[list[dict], bytes]:
|
||||
"""Normalize a stream chunk into one or more event dicts.
|
||||
|
||||
``litellm.anthropic.messages.acreate(stream=True)`` yields raw SSE
|
||||
bytes in practice; some adapters yield strings or typed events. Handle
|
||||
all three.
|
||||
"""
|
||||
if isinstance(chunk, (bytes, bytearray)):
|
||||
sse_buffer += bytes(chunk)
|
||||
events, sse_buffer = parse_sse_blocks(sse_buffer)
|
||||
return events, sse_buffer
|
||||
if isinstance(chunk, str):
|
||||
sse_buffer += chunk.encode("utf-8")
|
||||
events, sse_buffer = parse_sse_blocks(sse_buffer)
|
||||
return events, sse_buffer
|
||||
return [coerce_litellm_payload(chunk)], sse_buffer
|
||||
|
||||
|
||||
async def aggregate_anthropic_events_to_message(
|
||||
iterator: AsyncIterator[Any],
|
||||
) -> dict:
|
||||
"""Drain an Anthropic-Messages event iterator into a single Message dict.
|
||||
|
||||
Produces the shape ``litellm.anthropic.messages.acreate(stream=False)``
|
||||
would have returned. Used to transparently stream from upstream while
|
||||
still returning a non-streaming response to the client. Lets us
|
||||
sidestep upstream quirks (e.g. Fireworks rejects ``max_tokens > 4096``
|
||||
unless ``stream=true``) without leaking that into client-visible
|
||||
behavior.
|
||||
"""
|
||||
sse_buffer = b""
|
||||
message: dict = {}
|
||||
blocks: list[dict] = []
|
||||
partial_json: dict[int, str] = {}
|
||||
final_stop_reason: str | None = None
|
||||
final_stop_sequence: str | None = None
|
||||
final_usage: dict[str, Any] = {}
|
||||
final_model: str | None = None
|
||||
|
||||
async for chunk in iterator:
|
||||
events, sse_buffer = events_from_chunk(chunk, sse_buffer)
|
||||
for event in events:
|
||||
etype = event.get("type")
|
||||
if etype == "message_start":
|
||||
raw = event.get("message") or {}
|
||||
if isinstance(raw, dict):
|
||||
message = dict(raw)
|
||||
existing = message.get("content")
|
||||
blocks = list(existing) if isinstance(existing, list) else []
|
||||
usage = message.get("usage")
|
||||
if isinstance(usage, dict):
|
||||
final_usage = dict(usage)
|
||||
if isinstance(message.get("model"), str):
|
||||
final_model = message["model"]
|
||||
elif etype == "content_block_start":
|
||||
idx = int(event.get("index") or 0)
|
||||
cb = event.get("content_block") or {}
|
||||
cb_dict = dict(cb) if isinstance(cb, dict) else {}
|
||||
while len(blocks) <= idx:
|
||||
blocks.append({})
|
||||
blocks[idx] = cb_dict
|
||||
elif etype == "content_block_delta":
|
||||
idx = int(event.get("index") or 0)
|
||||
if idx >= len(blocks):
|
||||
continue
|
||||
delta = event.get("delta") or {}
|
||||
if not isinstance(delta, dict):
|
||||
continue
|
||||
dtype = delta.get("type")
|
||||
block = blocks[idx]
|
||||
if dtype == "text_delta":
|
||||
block["text"] = (block.get("text") or "") + (
|
||||
delta.get("text") or ""
|
||||
)
|
||||
elif dtype == "input_json_delta":
|
||||
partial_json[idx] = partial_json.get(idx, "") + (
|
||||
delta.get("partial_json") or ""
|
||||
)
|
||||
elif dtype == "thinking_delta":
|
||||
block["thinking"] = (block.get("thinking") or "") + (
|
||||
delta.get("thinking") or ""
|
||||
)
|
||||
elif dtype == "signature_delta":
|
||||
block["signature"] = (block.get("signature") or "") + (
|
||||
delta.get("signature") or ""
|
||||
)
|
||||
elif etype == "content_block_stop":
|
||||
idx = int(event.get("index") or 0)
|
||||
raw_json = partial_json.pop(idx, None)
|
||||
if raw_json is not None and idx < len(blocks):
|
||||
try:
|
||||
blocks[idx]["input"] = (
|
||||
json.loads(raw_json) if raw_json else {}
|
||||
)
|
||||
except json.JSONDecodeError:
|
||||
blocks[idx]["input"] = raw_json
|
||||
elif etype == "message_delta":
|
||||
delta = event.get("delta") or {}
|
||||
if isinstance(delta, dict):
|
||||
if "stop_reason" in delta:
|
||||
final_stop_reason = delta.get("stop_reason")
|
||||
if "stop_sequence" in delta:
|
||||
final_stop_sequence = delta.get("stop_sequence")
|
||||
usage = event.get("usage")
|
||||
if isinstance(usage, dict):
|
||||
final_usage.update(usage)
|
||||
# message_stop: nothing to merge
|
||||
|
||||
if not message:
|
||||
# Upstream returned no message_start; expose what we can so the
|
||||
# client at least sees the assembled content.
|
||||
message = {
|
||||
"id": "",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [],
|
||||
}
|
||||
|
||||
message["content"] = blocks
|
||||
if final_model and not message.get("model"):
|
||||
message["model"] = final_model
|
||||
if final_stop_reason is not None:
|
||||
message["stop_reason"] = final_stop_reason
|
||||
if final_stop_sequence is not None:
|
||||
message["stop_sequence"] = final_stop_sequence
|
||||
if final_usage:
|
||||
existing_usage = message.get("usage")
|
||||
merged = dict(existing_usage) if isinstance(existing_usage, dict) else {}
|
||||
merged.update(final_usage)
|
||||
message["usage"] = merged
|
||||
return message
|
||||
|
||||
|
||||
class AnnotatedEvent(NamedTuple):
|
||||
"""One Anthropic SSE event after model-rewrite + token-tally bookkeeping.
|
||||
|
||||
``sse_bytes`` is the wire-ready ``event:`` / ``data:`` block; the two
|
||||
streaming paths in ``BaseUpstreamProvider`` consume ``sse_bytes`` plus
|
||||
the tallies and only differ in whether they stream live or buffer
|
||||
first.
|
||||
|
||||
``cache_read_input_tokens`` and ``cache_creation_input_tokens`` are
|
||||
surfaced separately so the cost path can price them against the cache
|
||||
rate rather than fold them silently into the regular input bucket.
|
||||
|
||||
``total_cost`` / ``input_cost`` / ``output_cost`` carry any
|
||||
USD cost figures the upstream attached to this event (from
|
||||
``usage.cost``, ``usage.total_cost``, or ``usage.cost_details``) so the
|
||||
streaming paths can re-embed them in the rebuilt ``usage`` dict and let
|
||||
``calculate_cost`` convert directly USD→sats instead of falling back to
|
||||
token-based math.
|
||||
"""
|
||||
|
||||
event: dict
|
||||
sse_bytes: bytes
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
cache_read_input_tokens: int
|
||||
cache_creation_input_tokens: int
|
||||
total_cost: float
|
||||
input_cost: float
|
||||
output_cost: float
|
||||
model: str | None
|
||||
|
||||
|
||||
def annotate_event(event: dict, requested_model: str | None) -> AnnotatedEvent:
|
||||
"""Rewrite ``model`` fields and extract per-event token / model info.
|
||||
|
||||
Mutates ``event`` in place when ``requested_model`` is set so the
|
||||
upstream's true model name doesn't leak to the client.
|
||||
"""
|
||||
if requested_model:
|
||||
msg = event.get("message")
|
||||
if isinstance(msg, dict) and "model" in msg:
|
||||
msg["model"] = requested_model
|
||||
if "model" in event:
|
||||
event["model"] = requested_model
|
||||
|
||||
in_tokens = 0
|
||||
out_tokens = 0
|
||||
cache_read_tokens = 0
|
||||
cache_create_tokens = 0
|
||||
total_cost = 0.0
|
||||
input_cost = 0.0
|
||||
output_cost = 0.0
|
||||
model: str | None = None
|
||||
|
||||
def _coerce_float(value: object) -> float:
|
||||
if value is None or isinstance(value, bool):
|
||||
return 0.0
|
||||
if not isinstance(value, (int, float, str)):
|
||||
return 0.0
|
||||
try:
|
||||
return max(0.0, float(value))
|
||||
except (TypeError, ValueError):
|
||||
return 0.0
|
||||
|
||||
def _accumulate(usage: dict) -> None:
|
||||
nonlocal in_tokens, out_tokens, cache_read_tokens, cache_create_tokens
|
||||
nonlocal total_cost, input_cost, output_cost
|
||||
in_tokens += int(usage.get("input_tokens") or 0)
|
||||
out_tokens += int(usage.get("output_tokens") or 0)
|
||||
cache_read_tokens += int(usage.get("cache_read_input_tokens") or 0)
|
||||
cache_create_tokens += int(usage.get("cache_creation_input_tokens") or 0)
|
||||
total_cost += _coerce_float(usage.get("total_cost"))
|
||||
input_cost += _coerce_float(usage.get("input_cost"))
|
||||
output_cost += _coerce_float(usage.get("output_cost"))
|
||||
|
||||
msg_for_meta = event.get("message")
|
||||
if isinstance(msg_for_meta, dict):
|
||||
if msg_for_meta.get("model"):
|
||||
model = str(msg_for_meta["model"])
|
||||
usage = msg_for_meta.get("usage")
|
||||
if isinstance(usage, dict):
|
||||
_accumulate(usage)
|
||||
|
||||
if isinstance(event.get("usage"), dict):
|
||||
_accumulate(event["usage"])
|
||||
|
||||
# Some upstreams (notably OpenRouter-style proxies) attach cost fields
|
||||
# directly at the event root rather than inside ``usage``.
|
||||
for field in ("total_cost", "cost"):
|
||||
total_cost = max(total_cost, _coerce_float(event.get(field)))
|
||||
input_cost = max(input_cost, _coerce_float(event.get("input_cost")))
|
||||
output_cost = max(output_cost, _coerce_float(event.get("output_cost")))
|
||||
root_cost_details = event.get("cost_details")
|
||||
if isinstance(root_cost_details, dict):
|
||||
total_cost = max(
|
||||
total_cost,
|
||||
_coerce_float(root_cost_details.get("total_cost")),
|
||||
)
|
||||
input_cost = max(
|
||||
input_cost,
|
||||
_coerce_float(root_cost_details.get("input_cost")),
|
||||
)
|
||||
output_cost = max(
|
||||
output_cost,
|
||||
_coerce_float(root_cost_details.get("output_cost")),
|
||||
)
|
||||
|
||||
event_type = str(event.get("type") or "")
|
||||
payload = json.dumps(event)
|
||||
if event_type:
|
||||
sse_bytes = f"event: {event_type}\ndata: {payload}\n\n".encode()
|
||||
else:
|
||||
sse_bytes = f"data: {payload}\n\n".encode()
|
||||
|
||||
return AnnotatedEvent(
|
||||
event,
|
||||
sse_bytes,
|
||||
in_tokens,
|
||||
out_tokens,
|
||||
cache_read_tokens,
|
||||
cache_create_tokens,
|
||||
total_cost,
|
||||
input_cost,
|
||||
output_cost,
|
||||
model,
|
||||
)
|
||||
|
||||
|
||||
async def stream_annotated_events(
|
||||
iterator: AsyncIterator[Any],
|
||||
requested_model: str | None,
|
||||
) -> AsyncGenerator[AnnotatedEvent, None]:
|
||||
"""Yield annotated, SSE-serialized events from a litellm stream.
|
||||
|
||||
Both streaming paths in ``BaseUpstreamProvider`` consume this; the only
|
||||
divergence between them — yield-as-you-go vs buffer-then-replay — stays
|
||||
in the caller.
|
||||
"""
|
||||
sse_buffer = b""
|
||||
async for chunk in iterator:
|
||||
events, sse_buffer = events_from_chunk(chunk, sse_buffer)
|
||||
for event in events:
|
||||
yield annotate_event(event, requested_model)
|
||||
|
||||
|
||||
def embed_usd_costs(
|
||||
usage: dict,
|
||||
total_cost: float,
|
||||
input_cost: float,
|
||||
output_cost: float,
|
||||
) -> None:
|
||||
"""Mutate ``usage`` so ``calculate_cost`` will pick up the USD totals.
|
||||
|
||||
Mirrors the upstream shape: when any USD figure is present, attach
|
||||
``cost`` (used by the simple-fallback branch in ``calculate_cost``) and
|
||||
a ``cost_details`` block (used by the preferred branch — also gives the
|
||||
input/output USD split when we have one).
|
||||
"""
|
||||
if total_cost <= 0 and input_cost <= 0 and output_cost <= 0:
|
||||
return
|
||||
|
||||
cost_details: dict[str, float] = {}
|
||||
effective_total = total_cost
|
||||
if effective_total <= 0 and (input_cost > 0 or output_cost > 0):
|
||||
effective_total = input_cost + output_cost
|
||||
|
||||
if effective_total > 0:
|
||||
cost_details["total_cost"] = effective_total
|
||||
usage["cost"] = effective_total
|
||||
if input_cost > 0:
|
||||
cost_details["input_cost"] = input_cost
|
||||
if output_cost > 0:
|
||||
cost_details["output_cost"] = output_cost
|
||||
if cost_details:
|
||||
usage["cost_details"] = cost_details
|
||||
|
||||
|
||||
def compute_refund(amount: int, unit: str, cost_msats: int) -> int:
|
||||
if unit == "msat":
|
||||
return amount - cost_msats
|
||||
if unit == "sat":
|
||||
return amount - (cost_msats + 999) // 1000
|
||||
raise ValueError(f"Invalid unit: {unit}")
|
||||
|
||||
|
||||
async def dispatch_anthropic_messages(
|
||||
*,
|
||||
request_body: bytes | None,
|
||||
model_obj: Model,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
provider_prefix: str,
|
||||
transform_model_name: Callable[[str], str],
|
||||
log_extra: dict[str, Any] | None = None,
|
||||
) -> tuple[bool, Any, str | None]:
|
||||
"""Call ``litellm.anthropic.messages.acreate`` and return
|
||||
``(client_stream, result, requested_model)``.
|
||||
|
||||
Shared by the bearer-key and x-cashu paths. Raises :class:`UpstreamError`
|
||||
on bad input or upstream failure.
|
||||
"""
|
||||
if not request_body:
|
||||
raise UpstreamError(
|
||||
"Missing request body for /v1/messages", status_code=400
|
||||
)
|
||||
|
||||
try:
|
||||
body: dict = json.loads(request_body)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise UpstreamError(
|
||||
f"Invalid JSON in /v1/messages body: {exc}", status_code=400
|
||||
) from exc
|
||||
|
||||
body.pop("model", None)
|
||||
# `stream` here is what the **client** asked for. Upstream is always
|
||||
# streamed (see `upstream_stream` below); when the client asked for a
|
||||
# non-streaming response we drain and aggregate the events into a
|
||||
# single Anthropic Message dict before returning. This sidesteps
|
||||
# provider-specific non-streaming caps (e.g. Fireworks rejects
|
||||
# `max_tokens > 4096` unless `stream=true`).
|
||||
client_stream = bool(body.pop("stream", False))
|
||||
upstream_stream = True
|
||||
|
||||
dropped: dict[str, Any] = {}
|
||||
for field in ANTHROPIC_ONLY_FIELDS:
|
||||
if field in body:
|
||||
dropped[field] = body.pop(field)
|
||||
if dropped:
|
||||
logger.debug(
|
||||
"Dropped anthropic-only fields before litellm dispatch",
|
||||
extra={"dropped_keys": sorted(dropped.keys())},
|
||||
)
|
||||
|
||||
# Convention: `model.id` is the canonical upstream model name;
|
||||
# `forwarded_model_id` is the public alias the internal API exposes
|
||||
# and echoes back to the client.
|
||||
requested_model = (
|
||||
(model_obj.forwarded_model_id or model_obj.id) if model_obj else None
|
||||
)
|
||||
upstream_model = transform_model_name(model_obj.id)
|
||||
litellm_model = f"{provider_prefix}{upstream_model}"
|
||||
|
||||
kwargs: dict = {
|
||||
"model": litellm_model,
|
||||
"api_base": base_url,
|
||||
"api_key": api_key,
|
||||
"stream": upstream_stream,
|
||||
**body,
|
||||
}
|
||||
|
||||
logger.info(
|
||||
"Dispatching /v1/messages via litellm",
|
||||
extra={
|
||||
"model": litellm_model,
|
||||
"resolved_provider": provider_prefix.rstrip("/"),
|
||||
"client_stream": client_stream,
|
||||
"upstream_stream": upstream_stream,
|
||||
**(log_extra or {}),
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
result = await litellm.anthropic.messages.acreate(**kwargs)
|
||||
except Exception as exc:
|
||||
exc_message = getattr(exc, "message", None) or str(exc) or repr(exc)
|
||||
exc_status = getattr(exc, "status_code", None)
|
||||
exc_response = getattr(exc, "response", None)
|
||||
response_text = None
|
||||
if exc_response is not None:
|
||||
try:
|
||||
response_text = getattr(exc_response, "text", str(exc_response))
|
||||
except Exception:
|
||||
response_text = "<unreadable>"
|
||||
logger.error(
|
||||
"litellm dispatch failed",
|
||||
extra={
|
||||
"error": exc_message,
|
||||
"error_type": type(exc).__name__,
|
||||
"status_code": exc_status,
|
||||
"llm_provider": getattr(exc, "llm_provider", None),
|
||||
"body": getattr(exc, "body", None),
|
||||
"response_text": response_text,
|
||||
"model": litellm_model,
|
||||
"api_base": base_url,
|
||||
},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Upstream error via litellm: {exc_message}",
|
||||
status_code=exc_status if isinstance(exc_status, int) else 502,
|
||||
) from exc
|
||||
|
||||
if not client_stream and hasattr(result, "__aiter__"):
|
||||
# Client asked for a non-streaming response but we always stream
|
||||
# from upstream — drain the events into a single Anthropic Message
|
||||
# dict so the rest of the pipeline can treat it as if upstream had
|
||||
# returned non-streaming. Some litellm adapters return a
|
||||
# non-streaming dict even when ``stream=True``; in that case,
|
||||
# leave the result as-is.
|
||||
try:
|
||||
aggregated: Any = await aggregate_anthropic_events_to_message(
|
||||
cast(AsyncIterator[Any], result)
|
||||
)
|
||||
except Exception as exc:
|
||||
logger.error(
|
||||
"Failed to aggregate streamed events into message",
|
||||
extra={
|
||||
"error": str(exc),
|
||||
"error_type": type(exc).__name__,
|
||||
"model": litellm_model,
|
||||
},
|
||||
)
|
||||
raise UpstreamError(
|
||||
f"Failed to aggregate upstream stream: {exc}",
|
||||
status_code=502,
|
||||
) from exc
|
||||
return client_stream, aggregated, requested_model
|
||||
|
||||
return client_stream, result, requested_model
|
||||
@@ -3,11 +3,13 @@ from __future__ import annotations
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
from fastapi import Request
|
||||
from fastapi.responses import Response, StreamingResponse
|
||||
|
||||
from .base import BaseUpstreamProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.db import UpstreamProviderRow
|
||||
from ..core.db import ApiKey, AsyncSession, UpstreamProviderRow
|
||||
from ..payment.models import Model
|
||||
|
||||
from ..core.logging import get_logger
|
||||
@@ -21,7 +23,6 @@ class OllamaUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "ollama"
|
||||
default_base_url = "http://localhost:11434"
|
||||
platform_url = None
|
||||
litellm_provider_prefix = "ollama_chat/"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -66,11 +67,38 @@ class OllamaUpstreamProvider(BaseUpstreamProvider):
|
||||
"""Strip 'ollama/' prefix for Ollama API compatibility."""
|
||||
return model_id.removeprefix("ollama/")
|
||||
|
||||
def get_request_base_url(
|
||||
self, path: str, model_obj: Model | None = None
|
||||
) -> str:
|
||||
"""Route proxy traffic through Ollama's OpenAI-compatible /v1 endpoint."""
|
||||
return f"{self.base_url.rstrip('/')}/v1"
|
||||
async def forward_request(
|
||||
self,
|
||||
request: Request,
|
||||
path: str,
|
||||
headers: dict,
|
||||
request_body: bytes | None,
|
||||
key: ApiKey,
|
||||
max_cost_for_model: int,
|
||||
session: AsyncSession,
|
||||
model_obj: Model,
|
||||
) -> Response | StreamingResponse:
|
||||
"""Override to use OpenAI-compatible endpoint for proxy requests."""
|
||||
if path.startswith("v1/"):
|
||||
path = path.replace("v1/", "")
|
||||
|
||||
original_base_url = self.base_url
|
||||
self.base_url = f"{self.base_url}/v1"
|
||||
|
||||
try:
|
||||
result = await super().forward_request(
|
||||
request,
|
||||
path,
|
||||
headers,
|
||||
request_body,
|
||||
key,
|
||||
max_cost_for_model,
|
||||
session,
|
||||
model_obj,
|
||||
)
|
||||
return result
|
||||
finally:
|
||||
self.base_url = original_base_url
|
||||
|
||||
async def fetch_models(self) -> list[Model]:
|
||||
"""Fetch models from Ollama API using /api/tags endpoint."""
|
||||
@@ -185,7 +213,7 @@ class OllamaUpstreamProvider(BaseUpstreamProvider):
|
||||
except Exception:
|
||||
self._models_cache = models_with_fees
|
||||
|
||||
self._models_by_id = {m.forwarded_model_id or m.id: m for m in self._models_cache}
|
||||
self._models_by_id = {m.id: m for m in self._models_cache}
|
||||
logger.info(
|
||||
f"Refreshed models cache for {self.base_url}",
|
||||
extra={"model_count": len(models)},
|
||||
|
||||
@@ -15,8 +15,6 @@ class OpenRouterUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "openrouter"
|
||||
default_base_url = "https://openrouter.ai/api/v1"
|
||||
platform_url = "https://openrouter.ai/settings/keys"
|
||||
supports_anthropic_messages = True
|
||||
litellm_provider_prefix = "openrouter/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.06):
|
||||
"""Initialize OpenRouter provider with API key.
|
||||
|
||||
@@ -13,7 +13,6 @@ class PerplexityUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "perplexity"
|
||||
default_base_url = "https://api.perplexity.ai/"
|
||||
platform_url = "https://www.perplexity.ai/account/api/keys"
|
||||
litellm_provider_prefix = "perplexity/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.01):
|
||||
super().__init__(
|
||||
|
||||
+20
-44
@@ -1,9 +1,9 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import TYPE_CHECKING, Optional
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import httpx
|
||||
from pydantic.v1 import BaseModel, Field
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ..core.logging import get_logger
|
||||
from ..payment.models import Architecture, Model, Pricing, async_fetch_openrouter_models
|
||||
@@ -16,20 +16,18 @@ logger = get_logger(__name__)
|
||||
|
||||
|
||||
class PPQAIModelPricing(BaseModel):
|
||||
ui: Optional[dict[str, float]] = None
|
||||
api: Optional[dict[str, float]] = None
|
||||
input_per_1M_tokens: Optional[float] = Field(None, alias="input_per_1M_tokens")
|
||||
output_per_1M_tokens: Optional[float] = Field(None, alias="output_per_1M_tokens")
|
||||
ui: dict[str, float]
|
||||
api: dict[str, float]
|
||||
|
||||
|
||||
class PPQAIModel(BaseModel):
|
||||
id: str
|
||||
provider: Optional[str] = None
|
||||
provider: str
|
||||
name: str
|
||||
created_at: int
|
||||
context_length: int
|
||||
pricing: PPQAIModelPricing
|
||||
popular: bool = False
|
||||
popular: bool
|
||||
|
||||
|
||||
class PPQAIUpstreamProvider(BaseUpstreamProvider):
|
||||
@@ -136,54 +134,31 @@ class PPQAIUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
if or_model:
|
||||
input_price = None
|
||||
if ppqai_model.pricing.api:
|
||||
input_price = ppqai_model.pricing.api.get(
|
||||
"input_per_1M"
|
||||
)
|
||||
elif ppqai_model.pricing.input_per_1M_tokens:
|
||||
input_price = ppqai_model.pricing.input_per_1M_tokens
|
||||
|
||||
if input_price is not None:
|
||||
if input_price := ppqai_model.pricing.api.get(
|
||||
"input_per_1M"
|
||||
):
|
||||
or_model.pricing.prompt = input_price / 1_000_000
|
||||
|
||||
output_price = None
|
||||
if ppqai_model.pricing.api:
|
||||
output_price = ppqai_model.pricing.api.get(
|
||||
"output_per_1M"
|
||||
)
|
||||
elif ppqai_model.pricing.output_per_1M_tokens:
|
||||
output_price = ppqai_model.pricing.output_per_1M_tokens
|
||||
|
||||
if output_price is not None:
|
||||
if output_price := ppqai_model.pricing.api.get(
|
||||
"output_per_1M"
|
||||
):
|
||||
or_model.pricing.completion = output_price / 1_000_000
|
||||
|
||||
if cl := ppqai_model.context_length:
|
||||
or_model.context_length = cl
|
||||
models.append(or_model)
|
||||
else:
|
||||
input_price = 0.0
|
||||
if ppqai_model.pricing.api:
|
||||
input_price = ppqai_model.pricing.api.get(
|
||||
"input_per_1M", 0.0
|
||||
)
|
||||
elif ppqai_model.pricing.input_per_1M_tokens:
|
||||
input_price = ppqai_model.pricing.input_per_1M_tokens
|
||||
|
||||
output_price = 0.0
|
||||
if ppqai_model.pricing.api:
|
||||
output_price = ppqai_model.pricing.api.get(
|
||||
"output_per_1M", 0.0
|
||||
)
|
||||
elif ppqai_model.pricing.output_per_1M_tokens:
|
||||
output_price = ppqai_model.pricing.output_per_1M_tokens
|
||||
input_price = ppqai_model.pricing.api.get(
|
||||
"input_per_1M", 0.0
|
||||
)
|
||||
output_price = ppqai_model.pricing.api.get(
|
||||
"output_per_1M", 0.0
|
||||
)
|
||||
|
||||
models.append(
|
||||
Model(
|
||||
id=ppqai_model.id,
|
||||
name=ppqai_model.name,
|
||||
created=ppqai_model.created_at // 1000,
|
||||
description=f"{ppqai_model.provider or 'PPQ.AI'} model",
|
||||
description=f"{ppqai_model.provider} model",
|
||||
context_length=ppqai_model.context_length,
|
||||
architecture=Architecture(
|
||||
modality="text->text",
|
||||
@@ -308,6 +283,7 @@ class PPQAIUpstreamProvider(BaseUpstreamProvider):
|
||||
)
|
||||
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
print(f"Payload: {payload}", "sending to", url)
|
||||
response = await client.post(url, headers=headers, json=payload)
|
||||
response.raise_for_status()
|
||||
invoice_data = response.json()
|
||||
|
||||
@@ -1,180 +0,0 @@
|
||||
from typing import TYPE_CHECKING, Any
|
||||
|
||||
import httpx
|
||||
|
||||
from ..core import get_logger
|
||||
from ..payment.models import Model
|
||||
from .base import BaseUpstreamProvider
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..core.db import UpstreamProviderRow
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
class RoutstrUpstreamProvider(BaseUpstreamProvider):
|
||||
"""Upstream provider for communicating with another Routstr instance."""
|
||||
|
||||
provider_type = "routstr"
|
||||
default_base_url = None
|
||||
platform_url = None
|
||||
# Upstream Routstr nodes serve `/v1/messages` natively, so forward the
|
||||
# request as-is instead of round-tripping through litellm's
|
||||
# Anthropic→OpenAI translator.
|
||||
supports_anthropic_messages = True
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
base_url: str,
|
||||
api_key: str,
|
||||
provider_fee: float = 1.01,
|
||||
provider_settings: dict | None = None,
|
||||
):
|
||||
"""Initialize Routstr provider.
|
||||
|
||||
Args:
|
||||
base_url: Base URL of the upstream Routstr instance
|
||||
api_key: API key for the upstream Routstr instance
|
||||
provider_fee: Provider fee multiplier
|
||||
provider_settings: Provider-specific settings (auto-topup, etc.)
|
||||
"""
|
||||
# Ensure base_url doesn't end with /v1 as BaseUpstreamProvider appends it if needed
|
||||
# but Routstr paths are usually absolute from base.
|
||||
super().__init__(
|
||||
base_url=base_url.rstrip("/"),
|
||||
api_key=api_key,
|
||||
provider_fee=provider_fee,
|
||||
)
|
||||
self.settings = provider_settings or {}
|
||||
|
||||
def normalize_request_path(
|
||||
self, path: str, model_obj: "Model | None" = None
|
||||
) -> str:
|
||||
"""Preserve the ``v1/`` prefix when forwarding to an upstream Routstr.
|
||||
"""
|
||||
return path.lstrip("/")
|
||||
|
||||
@classmethod
|
||||
def from_db_row(
|
||||
cls, provider_row: "UpstreamProviderRow"
|
||||
) -> "RoutstrUpstreamProvider":
|
||||
import json
|
||||
|
||||
settings = {}
|
||||
if provider_row.provider_settings:
|
||||
try:
|
||||
settings = json.loads(provider_row.provider_settings)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return cls(
|
||||
base_url=provider_row.base_url,
|
||||
api_key=provider_row.api_key,
|
||||
provider_fee=provider_row.provider_fee,
|
||||
provider_settings=settings,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_provider_metadata(cls) -> dict[str, object]:
|
||||
return {
|
||||
"id": cls.provider_type,
|
||||
"name": "Routstr Node",
|
||||
"default_base_url": "",
|
||||
"fixed_base_url": False,
|
||||
"platform_url": cls.platform_url,
|
||||
"can_create_account": False,
|
||||
"can_topup": True,
|
||||
"can_show_balance": True,
|
||||
}
|
||||
|
||||
async def get_balance(self) -> float | None:
|
||||
"""Fetch balance from the upstream Routstr node.
|
||||
|
||||
Returns:
|
||||
Balance in satoshis, or None if failed
|
||||
"""
|
||||
url = f"{self.base_url}/v1/balance/info"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
response = await client.get(url, headers=headers, timeout=10.0)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
# Routstr balance info usually contains 'balance' in msats or sats
|
||||
# Check for msats and convert to sats
|
||||
if "balance_msats" in data:
|
||||
return float(data["balance_msats"]) / 1000.0
|
||||
return float(data.get("balance", 0))
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to fetch balance from upstream Routstr",
|
||||
extra={"url": url, "error": str(e)},
|
||||
)
|
||||
return None
|
||||
|
||||
async def topup(self, cashu_token: str) -> dict[str, Any]:
|
||||
"""Top up balance on the upstream Routstr node.
|
||||
|
||||
Args:
|
||||
cashu_token: Cashu token to deposit
|
||||
|
||||
Returns:
|
||||
Dict containing top-up result
|
||||
"""
|
||||
url = f"{self.base_url}/v1/balance/topup"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
payload = {"cashu_token": cashu_token}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
response = await client.post(
|
||||
url, headers=headers, json=payload, timeout=30.0
|
||||
)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to topup upstream Routstr",
|
||||
extra={"url": url, "error": str(e)},
|
||||
)
|
||||
return {"error": str(e)}
|
||||
|
||||
async def fetch_models(self) -> list[Model]:
|
||||
"""Fetch models from the upstream Routstr node."""
|
||||
url = f"{self.base_url}/v1/models"
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
response = await client.get(url, headers={}, timeout=15.0)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
models = data.get("data", [])
|
||||
return [Model(**m) for m in models]
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to fetch models from upstream Routstr",
|
||||
extra={"url": url, "error": str(e)},
|
||||
)
|
||||
return []
|
||||
|
||||
async def refund_balance(self) -> dict[str, Any]:
|
||||
"""Request a refund from the upstream Routstr node.
|
||||
|
||||
Returns:
|
||||
Dict containing refund result and token
|
||||
"""
|
||||
url = f"{self.base_url}/v1/balance/refund"
|
||||
headers = {"Authorization": f"Bearer {self.api_key}"}
|
||||
|
||||
async with httpx.AsyncClient() as client:
|
||||
try:
|
||||
response = await client.post(url, headers=headers, timeout=30.0)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Failed to request refund from upstream Routstr",
|
||||
extra={"url": url, "error": str(e)},
|
||||
)
|
||||
return {"error": str(e)}
|
||||
@@ -13,7 +13,6 @@ class XAIUpstreamProvider(BaseUpstreamProvider):
|
||||
provider_type = "x-ai"
|
||||
default_base_url = "https://api.x.ai/v1"
|
||||
platform_url = "https://console.x.ai/"
|
||||
litellm_provider_prefix = "xai/"
|
||||
|
||||
def __init__(self, api_key: str, provider_fee: float = 1.01):
|
||||
super().__init__(
|
||||
|
||||
+23
-320
@@ -1,32 +1,16 @@
|
||||
import asyncio
|
||||
import time
|
||||
import typing
|
||||
import math
|
||||
from typing import TypedDict
|
||||
|
||||
from cashu.core.base import Proof, Token
|
||||
from cashu.core.mint_info import MintInfo as _CashuMintInfo
|
||||
from cashu.wallet.helpers import deserialize_token_from_string
|
||||
from cashu.wallet.wallet import Wallet
|
||||
from pydantic_core import PydanticUndefined
|
||||
from sqlmodel import col, select, update
|
||||
from sqlmodel import col, update
|
||||
|
||||
from .core import db, get_logger
|
||||
from .core.db import store_cashu_transaction
|
||||
from .core.settings import settings
|
||||
from .payment.lnurl import raw_send_to_lnurl
|
||||
|
||||
# cashu still declares Optional[X] without explicit defaults on MintInfo.
|
||||
# Under pydantic v2 those are required, but real mints omit many of them.
|
||||
# Default Optional fields to None at import time so balance fetches don't 422.
|
||||
for _name, _field in _CashuMintInfo.model_fields.items():
|
||||
_annot = _field.annotation
|
||||
_is_optional = typing.get_origin(_annot) is typing.Union and type(
|
||||
None
|
||||
) in typing.get_args(_annot)
|
||||
if _is_optional and _field.default is PydanticUndefined:
|
||||
_field.default = None
|
||||
_CashuMintInfo.model_rebuild(force=True)
|
||||
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
@@ -50,7 +34,6 @@ async def recieve_token(
|
||||
|
||||
wallet.verify_proofs_dleq(token_obj.proofs)
|
||||
await wallet.split(proofs=token_obj.proofs, amount=0, include_fees=True)
|
||||
|
||||
return token_obj.amount, token_obj.unit, token_obj.mint
|
||||
|
||||
|
||||
@@ -75,86 +58,15 @@ async def send_token(amount: int, unit: str, mint_url: str | None = None) -> str
|
||||
return token
|
||||
|
||||
|
||||
async def _calculate_swap_amount(
|
||||
amount_msat: int,
|
||||
token_unit: str,
|
||||
token_mint_url: str,
|
||||
token_wallet: Wallet,
|
||||
primary_wallet: Wallet,
|
||||
) -> int:
|
||||
"""
|
||||
Calculate the amount to mint on the primary mint after accounting for
|
||||
potential swap fees (melt fees) on the foreign mint.
|
||||
"""
|
||||
if settings.primary_mint_unit == "sat":
|
||||
receive_amount = amount_msat // 1000
|
||||
else:
|
||||
receive_amount = amount_msat
|
||||
|
||||
if token_mint_url == settings.primary_mint:
|
||||
logger.info(
|
||||
"swap_to_primary_mint: skipping fee estimation (same mint)",
|
||||
extra={"minted_amount": receive_amount},
|
||||
)
|
||||
return int(receive_amount)
|
||||
|
||||
logger.info(
|
||||
"swap_to_primary_mint: estimating fees",
|
||||
extra={
|
||||
"dummy_amount": receive_amount,
|
||||
"unit": settings.primary_mint_unit,
|
||||
},
|
||||
)
|
||||
|
||||
try:
|
||||
dummy_mint_quote = await primary_wallet.request_mint(receive_amount)
|
||||
dummy_melt_quote = await token_wallet.melt_quote(dummy_mint_quote.request)
|
||||
|
||||
fee_reserve = dummy_melt_quote.fee_reserve
|
||||
if token_unit == "sat":
|
||||
fee_msat = fee_reserve * 1000
|
||||
else:
|
||||
fee_msat = fee_reserve
|
||||
|
||||
amount_msat_after_fee = amount_msat - fee_msat
|
||||
|
||||
if settings.primary_mint_unit == "sat":
|
||||
minted_amount = int(amount_msat_after_fee // 1000)
|
||||
else:
|
||||
minted_amount = int(amount_msat_after_fee)
|
||||
|
||||
if minted_amount <= 0:
|
||||
raise ValueError(f"Fees ({fee_reserve} {token_unit}) exceed token amount")
|
||||
|
||||
logger.info(
|
||||
"swap_to_primary_mint: fee estimation result",
|
||||
extra={
|
||||
"token_amount_sat": amount_msat // 1000,
|
||||
"estimated_fee_sat": fee_msat // 1000,
|
||||
"minted_amount": minted_amount,
|
||||
"minted_unit": settings.primary_mint_unit,
|
||||
},
|
||||
)
|
||||
return minted_amount
|
||||
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"swap_to_primary_mint: fee estimation failed",
|
||||
extra={"error": str(e)},
|
||||
)
|
||||
raise ValueError(f"Failed to estimate fees: {e}") from e
|
||||
|
||||
|
||||
async def swap_to_primary_mint(
|
||||
token_obj: Token, token_wallet: Wallet
|
||||
) -> tuple[int, str, str]:
|
||||
logger.info(
|
||||
"swap_to_primary_mint: starting",
|
||||
"swap_to_primary_mint",
|
||||
extra={
|
||||
"foreign_mint": token_obj.mint,
|
||||
"token_amount": token_obj.amount,
|
||||
"mint": token_obj.mint,
|
||||
"amount": token_obj.amount,
|
||||
"unit": token_obj.unit,
|
||||
"primary_mint": settings.primary_mint,
|
||||
},
|
||||
)
|
||||
# Ensure amount is an integer
|
||||
@@ -169,118 +81,24 @@ async def swap_to_primary_mint(
|
||||
amount_msat = token_amount
|
||||
else:
|
||||
raise ValueError("Invalid unit")
|
||||
estimated_fee_sat = math.ceil(max(amount_msat // 1000 * 0.01, 2)) + 1
|
||||
amount_msat_after_fee = amount_msat - estimated_fee_sat * 1000
|
||||
primary_wallet = await get_wallet(settings.primary_mint, settings.primary_mint_unit)
|
||||
|
||||
# If the token is already from the primary mint, we don't need to swap
|
||||
# and we definitely don't want to calculate or pay fees.
|
||||
if token_obj.mint == settings.primary_mint:
|
||||
logger.info(
|
||||
"swap_to_primary_mint: token already on primary mint, skipping swap",
|
||||
extra={
|
||||
"mint": token_obj.mint,
|
||||
"amount": token_amount,
|
||||
"unit": token_obj.unit,
|
||||
},
|
||||
)
|
||||
await token_wallet.split(proofs=token_obj.proofs, amount=0, include_fees=True)
|
||||
return token_amount, token_obj.unit, token_obj.mint
|
||||
|
||||
minted_amount = await _calculate_swap_amount(
|
||||
amount_msat,
|
||||
token_obj.unit,
|
||||
token_obj.mint,
|
||||
token_wallet,
|
||||
primary_wallet,
|
||||
)
|
||||
|
||||
if settings.primary_mint_unit == "sat":
|
||||
minted_amount = int(amount_msat_after_fee // 1000)
|
||||
else:
|
||||
minted_amount = int(amount_msat_after_fee)
|
||||
mint_quote = await primary_wallet.request_mint(minted_amount)
|
||||
logger.info(
|
||||
"swap_to_primary_mint: mint quote received",
|
||||
extra={"mint_quote_id": mint_quote.quote},
|
||||
)
|
||||
|
||||
melt_quote = await token_wallet.melt_quote(mint_quote.request)
|
||||
total_needed = melt_quote.amount + melt_quote.fee_reserve
|
||||
logger.info(
|
||||
"swap_to_primary_mint: melt quote received",
|
||||
extra={
|
||||
"melt_quote_id": melt_quote.quote,
|
||||
"melt_amount": melt_quote.amount,
|
||||
"melt_fee_reserve": melt_quote.fee_reserve,
|
||||
"total_needed": total_needed,
|
||||
"token_amount": token_amount,
|
||||
},
|
||||
)
|
||||
|
||||
if total_needed > token_amount:
|
||||
logger.warning(
|
||||
"swap_to_primary_mint: insufficient token amount for melt fees",
|
||||
extra={
|
||||
"token_amount": token_amount,
|
||||
"melt_amount": melt_quote.amount,
|
||||
"melt_fee_reserve": melt_quote.fee_reserve,
|
||||
"total_needed": total_needed,
|
||||
"shortfall": total_needed - token_amount,
|
||||
},
|
||||
)
|
||||
raise ValueError(
|
||||
f"Token amount ({token_amount} {token_obj.unit}) is insufficient to cover "
|
||||
f"melt fees. Needed: {total_needed} {token_obj.unit} "
|
||||
f"(amount: {melt_quote.amount} + fee: {melt_quote.fee_reserve})"
|
||||
)
|
||||
|
||||
try:
|
||||
_ = await token_wallet.melt(
|
||||
proofs=token_obj.proofs,
|
||||
invoice=mint_quote.request,
|
||||
fee_reserve_sat=melt_quote.fee_reserve,
|
||||
quote_id=melt_quote.quote,
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"swap_to_primary_mint: melt failed",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"foreign_mint": token_obj.mint,
|
||||
"token_amount": token_amount,
|
||||
"melt_quote_id": melt_quote.quote,
|
||||
"total_needed": total_needed,
|
||||
},
|
||||
)
|
||||
raise ValueError(
|
||||
f"Failed to melt token from foreign mint {token_obj.mint}: {e}"
|
||||
) from e
|
||||
|
||||
logger.info(
|
||||
"swap_to_primary_mint: melt succeeded, minting on primary",
|
||||
extra={"minted_amount": minted_amount, "mint_quote_id": mint_quote.quote},
|
||||
)
|
||||
|
||||
try:
|
||||
_ = await primary_wallet.mint(minted_amount, quote_id=mint_quote.quote)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"swap_to_primary_mint: mint on primary failed after successful melt",
|
||||
extra={
|
||||
"error": str(e),
|
||||
"error_type": type(e).__name__,
|
||||
"minted_amount": minted_amount,
|
||||
"mint_quote_id": mint_quote.quote,
|
||||
},
|
||||
)
|
||||
raise
|
||||
|
||||
logger.info(
|
||||
"swap_to_primary_mint: completed successfully",
|
||||
extra={
|
||||
"foreign_mint": token_obj.mint,
|
||||
"primary_mint": settings.primary_mint,
|
||||
"original_amount": token_amount,
|
||||
"minted_amount": minted_amount,
|
||||
"unit": settings.primary_mint_unit,
|
||||
},
|
||||
_ = await token_wallet.melt(
|
||||
proofs=token_obj.proofs,
|
||||
invoice=mint_quote.request,
|
||||
fee_reserve_sat=melt_quote.fee_reserve,
|
||||
quote_id=melt_quote.quote,
|
||||
)
|
||||
_ = await primary_wallet.mint(minted_amount, quote_id=mint_quote.quote)
|
||||
|
||||
return int(minted_amount), settings.primary_mint_unit, settings.primary_mint
|
||||
|
||||
@@ -295,8 +113,6 @@ async def credit_balance(
|
||||
|
||||
try:
|
||||
amount, unit, mint_url = await recieve_token(cashu_token)
|
||||
original_amount = amount
|
||||
original_unit = unit
|
||||
logger.info(
|
||||
"credit_balance: Token redeemed successfully",
|
||||
extra={"amount": amount, "unit": unit, "mint_url": mint_url},
|
||||
@@ -328,20 +144,7 @@ async def credit_balance(
|
||||
extra={"new_balance": key.balance},
|
||||
)
|
||||
|
||||
try:
|
||||
await store_cashu_transaction(
|
||||
token=cashu_token,
|
||||
amount=original_amount,
|
||||
unit=original_unit,
|
||||
mint_url=mint_url,
|
||||
typ="in",
|
||||
source="apikey",
|
||||
api_key_hashed_key=key.hashed_key,
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
logger.debug(
|
||||
logger.info(
|
||||
"Cashu token successfully redeemed and stored",
|
||||
extra={"amount": amount, "unit": unit, "mint_url": mint_url},
|
||||
)
|
||||
@@ -502,11 +305,11 @@ async def fetch_all_balances(
|
||||
|
||||
|
||||
async def periodic_payout() -> None:
|
||||
if not settings.receive_ln_address:
|
||||
logger.error("RECEIVE_LN_ADDRESS is not set, skipping payout")
|
||||
return
|
||||
while True:
|
||||
await asyncio.sleep(settings.payout_interval_seconds)
|
||||
print(settings.payout_interval_seconds)
|
||||
if not settings.receive_ln_address:
|
||||
continue
|
||||
await asyncio.sleep(60 * 15)
|
||||
try:
|
||||
async with db.create_session() as session:
|
||||
for mint_url in settings.cashu_mints:
|
||||
@@ -524,12 +327,7 @@ async def periodic_payout() -> None:
|
||||
user_balance = user_balance // 1000
|
||||
proofs_balance = sum(proof.amount for proof in proofs)
|
||||
available_balance = proofs_balance - user_balance
|
||||
# Threshold is configured in sats; convert for msat wallets.
|
||||
min_amount = (
|
||||
settings.min_payout_sat
|
||||
if unit == "sat"
|
||||
else settings.min_payout_sat * 1000
|
||||
)
|
||||
min_amount = 210 if unit == "sat" else 210000
|
||||
if available_balance > min_amount:
|
||||
amount_received = await raw_send_to_lnurl(
|
||||
wallet,
|
||||
@@ -554,101 +352,6 @@ async def periodic_payout() -> None:
|
||||
)
|
||||
|
||||
|
||||
async def periodic_refund_sweep() -> None:
|
||||
while True:
|
||||
await asyncio.sleep(60 * 60) # every hour
|
||||
try:
|
||||
cutoff = int(time.time()) - settings.refund_sweep_ttl_seconds
|
||||
async with db.create_session() as session:
|
||||
stmt = select(db.CashuTransaction).where(
|
||||
db.CashuTransaction.type == "out",
|
||||
db.CashuTransaction.collected == False, # noqa: E712
|
||||
db.CashuTransaction.swept == False, # noqa: E712
|
||||
db.CashuTransaction.created_at < cutoff,
|
||||
)
|
||||
results = await session.exec(stmt)
|
||||
refunds = results.all()
|
||||
|
||||
for refund in refunds:
|
||||
try:
|
||||
await recieve_token(refund.token)
|
||||
refund.swept = True
|
||||
session.add(refund)
|
||||
logger.info(
|
||||
"Swept uncollected refund",
|
||||
extra={
|
||||
"id": refund.id,
|
||||
"amount": refund.amount,
|
||||
"unit": refund.unit,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
error_msg = str(e).lower()
|
||||
if "already spent" in error_msg:
|
||||
refund.collected = True
|
||||
session.add(refund)
|
||||
logger.info(
|
||||
"Refund already spent (client collected), marking swept",
|
||||
extra={
|
||||
"id": refund.id,
|
||||
},
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"Failed to sweep refund",
|
||||
extra={
|
||||
"id": refund.id,
|
||||
"error": str(e),
|
||||
},
|
||||
)
|
||||
await session.commit()
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
"Error in periodic refund sweep",
|
||||
extra={"error": str(e), "error_type": type(e).__name__},
|
||||
)
|
||||
|
||||
|
||||
async def periodic_routstr_fee_payout() -> None:
|
||||
from .auth import (
|
||||
ROUTSTR_FEE_DEFAULT_PAYOUT,
|
||||
ROUTSTR_FEE_PAYOUT_INTERVAL_SECONDS,
|
||||
ROUTSTR_LN_ADDRESS,
|
||||
)
|
||||
|
||||
if not ROUTSTR_LN_ADDRESS:
|
||||
logger.info("ROUTSTR_LN_ADDRESS not set, skipping fee payout")
|
||||
return
|
||||
while True:
|
||||
await asyncio.sleep(ROUTSTR_FEE_PAYOUT_INTERVAL_SECONDS)
|
||||
try:
|
||||
async with db.create_session() as session:
|
||||
fee = await db.get_routstr_fee(session)
|
||||
accumulated_sats = fee.accumulated_msats // 1000
|
||||
if accumulated_sats >= ROUTSTR_FEE_DEFAULT_PAYOUT:
|
||||
wallet = await get_wallet(settings.primary_mint, "sat")
|
||||
proofs = get_proofs_per_mint_and_unit(
|
||||
wallet, settings.primary_mint, "sat", not_reserved=True
|
||||
)
|
||||
amount_received = await raw_send_to_lnurl(
|
||||
wallet, proofs, ROUTSTR_LN_ADDRESS, "sat", amount=accumulated_sats
|
||||
)
|
||||
paid_msats = accumulated_sats * 1000
|
||||
await db.reset_routstr_fee(session, paid_msats)
|
||||
logger.info(
|
||||
"Routstr fee payout sent",
|
||||
extra={
|
||||
"accumulated_sats": accumulated_sats,
|
||||
"amount_received": amount_received,
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
logger.error(
|
||||
f"Error in Routstr fee payout: {type(e).__name__}",
|
||||
extra={"error": str(e)},
|
||||
)
|
||||
|
||||
|
||||
async def send_to_lnurl(amount: int, unit: str, mint: str, address: str) -> int:
|
||||
wallet = await get_wallet(mint, unit)
|
||||
proofs = wallet._get_proofs_per_keyset(wallet.proofs)[wallet.keyset_id]
|
||||
|
||||
@@ -380,13 +380,6 @@ async def integration_session(
|
||||
yield session
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def patched_db_engine(integration_engine: Any) -> AsyncGenerator[None, None]:
|
||||
"""Patch the global db engine so create_session() uses the test engine."""
|
||||
with patch("routstr.core.db.engine", integration_engine):
|
||||
yield
|
||||
|
||||
|
||||
class DatabaseSnapshot:
|
||||
"""Utility to capture and compare database states"""
|
||||
|
||||
|
||||
@@ -1,80 +0,0 @@
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.admin import admin_sessions
|
||||
from routstr.core.db import UpstreamProviderRow
|
||||
|
||||
|
||||
async def _create_routstr_provider() -> UpstreamProviderRow:
|
||||
return UpstreamProviderRow(
|
||||
provider_type="routstr",
|
||||
base_url="https://upstream.example",
|
||||
api_key="",
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
|
||||
def _admin_headers() -> dict[str, str]:
|
||||
token = "test-admin-token"
|
||||
admin_sessions[token] = int(
|
||||
(datetime.now(timezone.utc) + timedelta(minutes=5)).timestamp()
|
||||
)
|
||||
return {"Authorization": f"Bearer {token}"}
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_routstr_provider_balance_timeout_returns_504(
|
||||
integration_client: httpx.AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
provider = await _create_routstr_provider()
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
|
||||
request = httpx.Request("GET", f"{provider.base_url}/v1/balance/info")
|
||||
timeout_error = httpx.ConnectTimeout("Connect timeout", request=request)
|
||||
|
||||
with patch(
|
||||
"httpx.AsyncHTTPTransport.handle_async_request",
|
||||
new=AsyncMock(side_effect=timeout_error),
|
||||
):
|
||||
response = await integration_client.get(
|
||||
f"/admin/api/upstream-providers/{provider.id}/balance",
|
||||
headers=_admin_headers(),
|
||||
)
|
||||
|
||||
assert response.status_code == 504
|
||||
assert response.json()["detail"] == "Timed out contacting upstream Routstr provider"
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_routstr_provider_balance_request_error_returns_502(
|
||||
integration_client: httpx.AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
provider = await _create_routstr_provider()
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
|
||||
request = httpx.Request("GET", f"{provider.base_url}/v1/balance/info")
|
||||
request_error = httpx.ConnectError("Connection failed", request=request)
|
||||
|
||||
with patch(
|
||||
"httpx.AsyncHTTPTransport.handle_async_request",
|
||||
new=AsyncMock(side_effect=request_error),
|
||||
):
|
||||
response = await integration_client.get(
|
||||
f"/admin/api/upstream-providers/{provider.id}/balance",
|
||||
headers=_admin_headers(),
|
||||
)
|
||||
|
||||
assert response.status_code == 502
|
||||
assert response.json()["detail"] == "Failed to contact upstream Routstr provider"
|
||||
@@ -1,369 +0,0 @@
|
||||
"""
|
||||
Integration tests for the balance-goes-negative bug in adjust_payment_for_tokens.
|
||||
|
||||
Root cause: when actual token cost exceeds the discounted reservation
|
||||
(cost_difference > 0, caused by tolerance_percentage discounting the reservation),
|
||||
the finalization UPDATE had no WHERE guard on balance, allowing balance to go negative.
|
||||
|
||||
Fix: added `.where(col(ApiKey.balance) >= total_cost_msats)` so the UPDATE is a no-op
|
||||
when balance is insufficient, then falls back to charging only deducted_max_cost.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.db import ApiKey
|
||||
from routstr.payment.cost_calculation import CostData
|
||||
|
||||
|
||||
def _make_key(balance: int, reserved: int) -> ApiKey:
|
||||
return ApiKey(
|
||||
hashed_key=f"test_{uuid.uuid4().hex}",
|
||||
balance=balance,
|
||||
reserved_balance=reserved,
|
||||
total_spent=0,
|
||||
total_requests=1,
|
||||
)
|
||||
|
||||
|
||||
async def _refresh(session: AsyncSession, key: ApiKey) -> ApiKey:
|
||||
await session.refresh(key)
|
||||
return key
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper: build a CostData where token cost > deducted_max_cost
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _cost_data(total_msats: int) -> CostData:
|
||||
return CostData(
|
||||
base_msats=0,
|
||||
input_msats=total_msats // 2,
|
||||
output_msats=total_msats - total_msats // 2,
|
||||
total_msats=total_msats,
|
||||
total_usd=0.0,
|
||||
input_tokens=100,
|
||||
output_tokens=100,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1 — exact reproduction of the bug
|
||||
#
|
||||
# Setup: balance == deducted_max_cost (user has just enough for the reservation,
|
||||
# nothing extra). Actual token cost is 1% higher (tolerance_percentage).
|
||||
#
|
||||
# Before fix: balance -= total_cost_msats → goes negative.
|
||||
# After fix: WHERE balance >= total_cost_msats fails → fallback charges
|
||||
# deducted_max_cost → balance reaches 0, never negative.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_balance_never_negative_when_cost_exceeds_reservation(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Balance must not go negative when actual token cost > discounted reservation."""
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
|
||||
deducted_max_cost = 990 # reserved (1% below true max of 1000)
|
||||
actual_token_cost = 1000 # actual cost at true max
|
||||
|
||||
# User has balance exactly equal to the reservation — tight budget
|
||||
key = _make_key(balance=deducted_max_cost, reserved=deducted_max_cost)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
response_data = {"model": "test-model", "usage": {"prompt_tokens": 100, "completion_tokens": 100}}
|
||||
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
|
||||
|
||||
await _refresh(integration_session, key)
|
||||
|
||||
assert key.balance >= 0, f"Balance went negative: {key.balance}"
|
||||
assert key.reserved_balance >= 0, f"Reserved balance went negative: {key.reserved_balance}"
|
||||
assert key.reserved_balance == 0, "Reservation must be fully released after finalization"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2 — balance is ZERO after the reservation is accounted for
|
||||
# (absolute floor case)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_balance_floor_at_zero_on_overrun(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""When balance exactly covers deducted_max_cost and cost overruns, balance reaches 0 not negative."""
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
|
||||
deducted_max_cost = 500
|
||||
actual_token_cost = 550 # 10% overrun
|
||||
|
||||
key = _make_key(balance=500, reserved=500)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 50}}
|
||||
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
|
||||
|
||||
await _refresh(integration_session, key)
|
||||
|
||||
assert key.balance == 0, (
|
||||
f"Expected balance=0 (charged deducted_max_cost fallback), got {key.balance}"
|
||||
)
|
||||
assert key.reserved_balance == 0, f"Reserved balance should be 0, got {key.reserved_balance}"
|
||||
# Fallback charges deducted_max_cost
|
||||
assert key.total_spent == deducted_max_cost, (
|
||||
f"Expected total_spent={deducted_max_cost}, got {key.total_spent}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 3 — balance has enough room: full token cost should be charged
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_full_cost_charged_when_balance_sufficient_for_overrun(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""When balance covers total_cost_msats, the full amount is charged (not just deducted_max_cost)."""
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
|
||||
deducted_max_cost = 990
|
||||
actual_token_cost = 1000
|
||||
|
||||
# User has extra balance beyond the reservation
|
||||
key = _make_key(balance=2000, reserved=990)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
response_data = {"model": "test-model", "usage": {"prompt_tokens": 100, "completion_tokens": 100}}
|
||||
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
|
||||
|
||||
await _refresh(integration_session, key)
|
||||
|
||||
assert key.balance >= 0, f"Balance went negative: {key.balance}"
|
||||
assert key.reserved_balance == 0, f"Reservation not released: {key.reserved_balance}"
|
||||
assert key.total_spent == actual_token_cost, (
|
||||
f"Expected full charge of {actual_token_cost}, got {key.total_spent}"
|
||||
)
|
||||
assert key.balance == 2000 - actual_token_cost, (
|
||||
f"Expected balance={2000 - actual_token_cost}, got {key.balance}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 4 — concurrent finalizations with cost overrun
|
||||
#
|
||||
# Multiple requests finish concurrently. Each has a small overrun.
|
||||
# None should drive balance negative.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_cost_overruns_never_negative(
|
||||
integration_session: AsyncSession,
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""Concurrent finalization with cost overruns must never produce negative balance."""
|
||||
import asyncio
|
||||
|
||||
from routstr.auth import adjust_payment_for_tokens, pay_for_request
|
||||
from routstr.core.db import create_session
|
||||
|
||||
deducted_max_cost = 990
|
||||
actual_token_cost = 1000
|
||||
n_requests = 5
|
||||
|
||||
# Fund the key with exactly enough for n_requests reservations + a tiny buffer
|
||||
starting_balance = deducted_max_cost * n_requests
|
||||
key_hash = f"test_concurrent_{uuid.uuid4().hex}"
|
||||
|
||||
async with create_session() as session:
|
||||
key = ApiKey(
|
||||
hashed_key=key_hash,
|
||||
balance=starting_balance,
|
||||
reserved_balance=0,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
)
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
# Reserve n_requests slots (sequentially, as pay_for_request is atomic)
|
||||
async with create_session() as session:
|
||||
key_to_reserve = await session.get(ApiKey, key_hash)
|
||||
assert key_to_reserve is not None
|
||||
for _ in range(n_requests):
|
||||
await pay_for_request(key_to_reserve, deducted_max_cost, session)
|
||||
await session.refresh(key_to_reserve)
|
||||
|
||||
# Now finalize all concurrently with cost overrun
|
||||
async def finalize() -> None:
|
||||
response_data = {
|
||||
"model": "test-model",
|
||||
"usage": {"prompt_tokens": 100, "completion_tokens": 100},
|
||||
}
|
||||
async with create_session() as session:
|
||||
fresh_key = await session.get(ApiKey, key_hash)
|
||||
assert fresh_key is not None
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(
|
||||
fresh_key, response_data, session, deducted_max_cost
|
||||
)
|
||||
|
||||
await asyncio.gather(*[finalize() for _ in range(n_requests)])
|
||||
|
||||
async with create_session() as session:
|
||||
final_key = await session.get(ApiKey, key_hash)
|
||||
assert final_key is not None
|
||||
|
||||
assert final_key.balance >= 0, (
|
||||
f"Balance went negative after concurrent overruns: {final_key.balance}"
|
||||
)
|
||||
assert final_key.reserved_balance == 0, (
|
||||
f"Reserved balance not fully released: {final_key.reserved_balance}"
|
||||
)
|
||||
assert final_key.total_spent <= starting_balance, (
|
||||
f"Total spent ({final_key.total_spent}) exceeds starting balance ({starting_balance})"
|
||||
)
|
||||
# Every request must have been charged at least deducted_max_cost — no free inference.
|
||||
assert final_key.total_spent == starting_balance, (
|
||||
f"Expected total_spent={starting_balance} (all {n_requests} reservations charged), "
|
||||
f"got {final_key.total_spent} — at least one request got free inference"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 5 — overrun with no balance at all (reserved_balance == balance)
|
||||
# simulates a user who topped up to exactly the reservation floor
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_free_balance_overrun_is_safe(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""User with zero free balance (all reserved) should never go negative on overrun."""
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
|
||||
deducted_max_cost = 1000
|
||||
actual_token_cost = 1050
|
||||
|
||||
# balance == reserved_balance: zero free balance
|
||||
key = _make_key(balance=1000, reserved=1000)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 100}}
|
||||
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(key, response_data, integration_session, deducted_max_cost)
|
||||
|
||||
await _refresh(integration_session, key)
|
||||
|
||||
assert key.balance >= 0, f"Balance went negative: {key.balance}"
|
||||
assert key.reserved_balance >= 0, f"Reserved balance went negative: {key.reserved_balance}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 6 — parallel requests: second finalization must not get free inference
|
||||
#
|
||||
# Root cause of the bug fixed in auth.py:
|
||||
# `.where(col(ApiKey.balance) >= total_cost_msats)` ignores other requests'
|
||||
# reservations, so after Request A charges total_cost_msats, balance can drop
|
||||
# below deducted_max_cost, causing Request B's fallback to release for free.
|
||||
#
|
||||
# Fix: use `balance - reserved_balance + deducted_max_cost >= total_cost_msats`
|
||||
# so the check accounts for concurrent reservations.
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_parallel_requests_no_free_inference(
|
||||
integration_session: AsyncSession,
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""Second parallel finalization must be charged even when first depleted free balance."""
|
||||
import asyncio
|
||||
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
from routstr.core.db import create_session
|
||||
|
||||
deducted_max_cost = 100
|
||||
actual_token_cost = 150 # overrun: 50 more than reserved
|
||||
|
||||
# Fund the key with exactly 2 * deducted_max_cost.
|
||||
# Both requests pre-reserved 100 each → balance=200, reserved=200, free=0.
|
||||
# Old check (balance >= total_cost_msats):
|
||||
# Request A: 200 >= 150 ✓ → charges 150 → balance=50, reserved=100
|
||||
# Request B: 50 >= 150 ✗ → fallback: 50 >= 100 ✗ → releases FREE
|
||||
# New check (balance - reserved + deducted >= total_cost_msats):
|
||||
# Both fall to fallback (0 free balance).
|
||||
# Both charge deducted_max_cost=100 → total_spent=200, balance=0.
|
||||
starting_balance = deducted_max_cost * 2
|
||||
key_hash = f"test_parallel_no_free_{uuid.uuid4().hex}"
|
||||
|
||||
async with create_session() as session:
|
||||
key = ApiKey(
|
||||
hashed_key=key_hash,
|
||||
balance=starting_balance,
|
||||
reserved_balance=deducted_max_cost * 2, # both slots pre-reserved
|
||||
total_spent=0,
|
||||
total_requests=2,
|
||||
)
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
async def finalize() -> None:
|
||||
response_data = {
|
||||
"model": "test-model",
|
||||
"usage": {"prompt_tokens": 50, "completion_tokens": 100},
|
||||
}
|
||||
async with create_session() as session:
|
||||
fresh_key = await session.get(ApiKey, key_hash)
|
||||
assert fresh_key is not None
|
||||
with patch(
|
||||
"routstr.auth.calculate_cost",
|
||||
return_value=_cost_data(actual_token_cost),
|
||||
):
|
||||
await adjust_payment_for_tokens(
|
||||
fresh_key, response_data, session, deducted_max_cost
|
||||
)
|
||||
|
||||
await asyncio.gather(finalize(), finalize())
|
||||
|
||||
async with create_session() as session:
|
||||
final_key = await session.get(ApiKey, key_hash)
|
||||
assert final_key is not None
|
||||
|
||||
assert final_key.balance >= 0, f"Balance went negative: {final_key.balance}"
|
||||
assert final_key.reserved_balance == 0, (
|
||||
f"Reserved balance not released: {final_key.reserved_balance}"
|
||||
)
|
||||
# Both requests must have been charged — no free inference.
|
||||
assert final_key.total_spent == starting_balance, (
|
||||
f"Expected total_spent={starting_balance} (both reservations charged), "
|
||||
f"got {final_key.total_spent} — one request got free inference"
|
||||
)
|
||||
@@ -1,97 +0,0 @@
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from httpx import AsyncClient
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_wallet_info_returns_child_keys(
|
||||
integration_client: AsyncClient,
|
||||
authenticated_client: AsyncClient,
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
"""Test that GET /v1/wallet/info returns child keys for a parent key"""
|
||||
|
||||
# 1. Get parent info to find its hashed_key
|
||||
response = await authenticated_client.get("/v1/wallet/info")
|
||||
assert response.status_code == 200
|
||||
parent_data = response.json()
|
||||
parent_data["api_key"]
|
||||
|
||||
# 2. Create child keys for this parent
|
||||
# We need to use the parent's authentication for this
|
||||
child_payload = {"count": 2, "balance_limit": 1000, "balance_limit_reset": "daily"}
|
||||
create_response = await authenticated_client.post(
|
||||
"/v1/wallet/child-key", json=child_payload
|
||||
)
|
||||
assert create_response.status_code == 200
|
||||
create_data = create_response.json()
|
||||
child_keys = create_data["api_keys"]
|
||||
assert len(child_keys) == 2
|
||||
|
||||
# 3. Call /info again and check for child_keys
|
||||
info_response = await authenticated_client.get("/v1/wallet/info")
|
||||
assert info_response.status_code == 200
|
||||
info_data = info_response.json()
|
||||
|
||||
assert "child_keys" in info_data
|
||||
assert len(info_data["child_keys"]) == 2
|
||||
|
||||
# Verify child key details
|
||||
for ck in info_data["child_keys"]:
|
||||
assert ck["api_key"] in child_keys
|
||||
assert ck["balance_limit"] == 1000
|
||||
assert ck["balance_limit_reset"] == "daily"
|
||||
assert "total_spent" in ck
|
||||
assert "total_requests" in ck
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_wallet_info_child_key_no_child_keys(
|
||||
integration_client: AsyncClient,
|
||||
authenticated_client: AsyncClient,
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
"""Test that GET /v1/wallet/info for a child key does NOT return child_keys"""
|
||||
|
||||
# 1. Create a child key
|
||||
child_payload = {"count": 1}
|
||||
create_response = await authenticated_client.post(
|
||||
"/v1/wallet/child-key", json=child_payload
|
||||
)
|
||||
assert create_response.status_code == 200
|
||||
child_key = create_response.json()["api_keys"][0]
|
||||
|
||||
# 2. Use the child key to get its info
|
||||
integration_client.headers["Authorization"] = f"Bearer {child_key}"
|
||||
info_response = await integration_client.get("/v1/wallet/info")
|
||||
assert info_response.status_code == 200
|
||||
info_data = info_response.json()
|
||||
|
||||
assert info_data["is_child"] is True
|
||||
assert "child_keys" not in info_data
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_account_info_root_returns_child_keys(
|
||||
authenticated_client: AsyncClient,
|
||||
) -> None:
|
||||
"""Test that GET / returns child keys for a parent key (root endpoint)"""
|
||||
|
||||
# 1. Create a child key
|
||||
child_payload = {"count": 1}
|
||||
await authenticated_client.post("/v1/wallet/child-key", json=child_payload)
|
||||
|
||||
# 2. Call root endpoint /v1/balance/
|
||||
# Note: routstr/balance.py defines router = APIRouter()
|
||||
# and it is included in balance_router with prefix /v1/balance
|
||||
# The endpoint is @router.get("/")
|
||||
response = await authenticated_client.get("/v1/balance/")
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
|
||||
assert "child_keys" in data
|
||||
assert len(data["child_keys"]) >= 1
|
||||
@@ -1,302 +0,0 @@
|
||||
"""Integration tests for CLI token management (/admin/api/cli-tokens).
|
||||
|
||||
Covers:
|
||||
- GET /admin/api/cli-tokens — list (preview only, no full token)
|
||||
- POST /admin/api/cli-tokens — create (returns full token once)
|
||||
- DELETE /admin/api/cli-tokens/{id} — revoke
|
||||
- Using a CLI token as Bearer auth against admin endpoints
|
||||
- Expiry enforcement (expired tokens are rejected by require_admin_api)
|
||||
- last_used_at bump on successful use
|
||||
- Auth failures: missing token, wrong token, revoked token
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import secrets
|
||||
import time
|
||||
from typing import AsyncGenerator
|
||||
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from httpx import AsyncClient
|
||||
from sqlmodel import select
|
||||
|
||||
from routstr.core.admin import admin_sessions
|
||||
from routstr.core.db import AsyncSession, CliToken
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Fixtures
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def admin_session_token() -> AsyncGenerator[str, None]:
|
||||
"""Inject a short-lived admin session token into admin_sessions."""
|
||||
token = secrets.token_urlsafe(24)
|
||||
admin_sessions[token] = int(time.time()) + 3600
|
||||
yield token
|
||||
admin_sessions.pop(token, None)
|
||||
|
||||
|
||||
@pytest_asyncio.fixture
|
||||
async def admin_client(
|
||||
integration_client: AsyncClient, admin_session_token: str
|
||||
) -> AsyncClient:
|
||||
"""An integration_client pre-authenticated with an admin session token."""
|
||||
integration_client.headers["Authorization"] = f"Bearer {admin_session_token}"
|
||||
return integration_client
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Creation
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_cli_token_returns_full_token_once(
|
||||
admin_client: AsyncClient,
|
||||
) -> None:
|
||||
"""POST /admin/api/cli-tokens returns the raw token only on creation."""
|
||||
resp = await admin_client.post(
|
||||
"/admin/api/cli-tokens",
|
||||
json={"name": "my-laptop"},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
|
||||
assert body["name"] == "my-laptop"
|
||||
assert isinstance(body["id"], str) and body["id"]
|
||||
assert isinstance(body["token"], str) and len(body["token"]) >= 32
|
||||
assert body["expires_at"] is None
|
||||
assert isinstance(body["created_at"], int)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_cli_token_with_expiry(admin_client: AsyncClient) -> None:
|
||||
"""expires_in_days sets expires_at ~= now + days * 86400."""
|
||||
before = int(time.time())
|
||||
resp = await admin_client.post(
|
||||
"/admin/api/cli-tokens",
|
||||
json={"name": "ci-runner", "expires_in_days": 7},
|
||||
)
|
||||
assert resp.status_code == 200
|
||||
body = resp.json()
|
||||
|
||||
assert body["expires_at"] is not None
|
||||
delta = body["expires_at"] - before
|
||||
# Allow 10s jitter around 7 * 86400
|
||||
assert 7 * 86400 - 10 <= delta <= 7 * 86400 + 10
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_cli_token_rejects_empty_name(
|
||||
admin_client: AsyncClient,
|
||||
) -> None:
|
||||
resp = await admin_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": " "}
|
||||
)
|
||||
assert resp.status_code == 400
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_cli_token_requires_admin(
|
||||
integration_client: AsyncClient,
|
||||
) -> None:
|
||||
"""No admin token / no bearer → 403."""
|
||||
resp = await integration_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": "no-auth"}
|
||||
)
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Listing
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_cli_tokens_returns_preview_not_full_token(
|
||||
admin_client: AsyncClient,
|
||||
) -> None:
|
||||
"""Listing never leaks the raw token."""
|
||||
create = await admin_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": "secret-keeper"}
|
||||
)
|
||||
assert create.status_code == 200
|
||||
full_token = create.json()["token"]
|
||||
|
||||
resp = await admin_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 200
|
||||
items = resp.json()
|
||||
assert any(t["name"] == "secret-keeper" for t in items)
|
||||
|
||||
for t in items:
|
||||
# No 'token' field, only 'token_preview'
|
||||
assert "token" not in t
|
||||
assert "token_preview" in t
|
||||
assert full_token not in t["token_preview"]
|
||||
assert "..." in t["token_preview"]
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_list_cli_tokens_requires_admin(
|
||||
integration_client: AsyncClient,
|
||||
) -> None:
|
||||
resp = await integration_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Using a CLI token as admin auth
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_cli_token_authorizes_admin_endpoints(
|
||||
admin_client: AsyncClient,
|
||||
integration_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""A freshly-created CLI token can be used as Bearer on admin endpoints."""
|
||||
create = await admin_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": "cli-auth"}
|
||||
)
|
||||
assert create.status_code == 200
|
||||
cli_token = create.json()["token"]
|
||||
token_id = create.json()["id"]
|
||||
|
||||
# Use a NEW client to isolate the header from admin_session_token
|
||||
integration_client.headers["Authorization"] = f"Bearer {cli_token}"
|
||||
resp = await integration_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 200
|
||||
|
||||
# last_used_at should be populated after use
|
||||
row = await integration_session.get(CliToken, token_id)
|
||||
assert row is not None
|
||||
assert row.last_used_at is not None
|
||||
assert row.last_used_at >= row.created_at
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_expired_cli_token_is_rejected(
|
||||
admin_client: AsyncClient,
|
||||
integration_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""A CLI token with expires_at in the past → 403."""
|
||||
create = await admin_client.post(
|
||||
"/admin/api/cli-tokens",
|
||||
json={"name": "will-expire", "expires_in_days": 1},
|
||||
)
|
||||
assert create.status_code == 200
|
||||
cli_token = create.json()["token"]
|
||||
token_id = create.json()["id"]
|
||||
|
||||
# Force-expire it in the DB
|
||||
row = await integration_session.get(CliToken, token_id)
|
||||
assert row is not None
|
||||
row.expires_at = int(time.time()) - 1
|
||||
integration_session.add(row)
|
||||
await integration_session.commit()
|
||||
|
||||
integration_client.headers["Authorization"] = f"Bearer {cli_token}"
|
||||
resp = await integration_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_invalid_bearer_token_is_rejected(
|
||||
integration_client: AsyncClient,
|
||||
) -> None:
|
||||
integration_client.headers["Authorization"] = "Bearer not-a-real-token"
|
||||
resp = await integration_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Revocation
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_revoke_cli_token_removes_auth(
|
||||
admin_client: AsyncClient,
|
||||
integration_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""After DELETE, the token no longer authorizes."""
|
||||
create = await admin_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": "to-revoke"}
|
||||
)
|
||||
token_id = create.json()["id"]
|
||||
cli_token = create.json()["token"]
|
||||
|
||||
revoke = await admin_client.delete(f"/admin/api/cli-tokens/{token_id}")
|
||||
assert revoke.status_code == 200
|
||||
assert revoke.json() == {"ok": True, "deleted_id": token_id}
|
||||
|
||||
# Row is gone
|
||||
row = await integration_session.get(CliToken, token_id)
|
||||
assert row is None
|
||||
|
||||
# Can no longer be used for auth
|
||||
integration_client.headers["Authorization"] = f"Bearer {cli_token}"
|
||||
resp = await integration_client.get("/admin/api/cli-tokens")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_revoke_unknown_cli_token_returns_404(
|
||||
admin_client: AsyncClient,
|
||||
) -> None:
|
||||
resp = await admin_client.delete("/admin/api/cli-tokens/does-not-exist")
|
||||
assert resp.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_revoke_cli_token_requires_admin(
|
||||
integration_client: AsyncClient,
|
||||
) -> None:
|
||||
resp = await integration_client.delete("/admin/api/cli-tokens/anything")
|
||||
assert resp.status_code == 403
|
||||
|
||||
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
# Lifecycle / uniqueness
|
||||
# ──────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_multiple_tokens_are_independent(
|
||||
admin_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Creating N tokens yields N unique tokens that all live in DB."""
|
||||
names = ["dev-a", "dev-b", "dev-c"]
|
||||
raw_tokens: list[str] = []
|
||||
ids: list[str] = []
|
||||
for name in names:
|
||||
r = await admin_client.post(
|
||||
"/admin/api/cli-tokens", json={"name": name}
|
||||
)
|
||||
assert r.status_code == 200
|
||||
raw_tokens.append(r.json()["token"])
|
||||
ids.append(r.json()["id"])
|
||||
|
||||
# All unique
|
||||
assert len(set(raw_tokens)) == len(raw_tokens)
|
||||
assert len(set(ids)) == len(ids)
|
||||
|
||||
# All in DB
|
||||
result = await integration_session.exec(
|
||||
select(CliToken).where(CliToken.name.in_(names)) # type: ignore[attr-defined]
|
||||
)
|
||||
rows = result.all()
|
||||
assert {r.name for r in rows} == set(names)
|
||||
@@ -267,9 +267,10 @@ async def test_admin_endpoint_unauthenticated(
|
||||
"""Test GET /admin/ endpoint redirects to /"""
|
||||
await db_snapshot.capture()
|
||||
|
||||
response = await integration_client.get("/admin/api/settings")
|
||||
response = await integration_client.get("/admin/")
|
||||
|
||||
assert response.status_code == 403
|
||||
assert response.status_code == 307
|
||||
assert response.headers.get("location") == "/"
|
||||
|
||||
diff = await db_snapshot.diff()
|
||||
assert len(diff["api_keys"]["added"]) == 0
|
||||
|
||||
@@ -1,240 +0,0 @@
|
||||
"""
|
||||
Tests showing how a user hits "Insufficient balance: X mSats required for this model"
|
||||
when their balance is too low for the model's cost.
|
||||
|
||||
The log line that triggered this:
|
||||
WARNING Insufficient billing balance during validation
|
||||
ERROR Bearer token validation failed: HTTPException: 402:
|
||||
{'error': {'message': 'Insufficient balance: 622888 mSats required
|
||||
for this model. 20320 available.', ...}}
|
||||
|
||||
This happens in validate_bearer_key (auth.py) when:
|
||||
billing_key.total_balance < min_cost (model's max cost)
|
||||
|
||||
and also in pay_for_request when the atomic UPDATE finds no available balance.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from httpx import AsyncClient
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.db import ApiKey
|
||||
|
||||
|
||||
def _key(balance: int, reserved: int = 0) -> ApiKey:
|
||||
return ApiKey(
|
||||
hashed_key=f"test_{uuid.uuid4().hex}",
|
||||
balance=balance,
|
||||
reserved_balance=reserved,
|
||||
total_spent=0,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1 — simplest case: balance < model cost → pay_for_request raises 402
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pay_for_request_raises_402_when_balance_too_low(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""
|
||||
User has 20_000 msats. Model costs 622_888 msats.
|
||||
pay_for_request must raise HTTP 402 with a clear message.
|
||||
"""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
model_cost = 622_888
|
||||
user_balance = 20_000
|
||||
|
||||
key = _key(balance=user_balance)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, model_cost, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
detail = exc_info.value.detail
|
||||
assert isinstance(detail, dict)
|
||||
error = detail["error"]
|
||||
assert error["code"] == "insufficient_balance"
|
||||
assert str(model_cost) in error["message"]
|
||||
assert str(user_balance) in error["message"]
|
||||
|
||||
# Balance must be untouched
|
||||
await integration_session.refresh(key)
|
||||
assert key.balance == user_balance
|
||||
assert key.reserved_balance == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2 — balance is zero
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pay_for_request_raises_402_on_zero_balance(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""User with zero balance cannot make any request."""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
key = _key(balance=0)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, 1_000, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
detail = exc_info.value.detail
|
||||
assert isinstance(detail, dict)
|
||||
assert detail["error"]["code"] == "insufficient_balance"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 3 — all balance is reserved (total_balance = balance - reserved = 0)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pay_for_request_raises_402_when_all_balance_reserved(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""
|
||||
User has 50_000 msats balance but 50_000 is already reserved for in-flight
|
||||
requests. Free balance (total_balance) = 0. Should get 402.
|
||||
"""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
key = _key(balance=50_000, reserved=50_000)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, 1_000, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
# Balance and reserved must be untouched
|
||||
await integration_session.refresh(key)
|
||||
assert key.balance == 50_000
|
||||
assert key.reserved_balance == 50_000
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 4 — balance just one msat below model cost
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pay_for_request_raises_402_one_msat_short(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Off-by-one: balance is exactly model_cost - 1."""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
model_cost = 10_000
|
||||
key = _key(balance=model_cost - 1)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, model_cost, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
await integration_session.refresh(key)
|
||||
assert key.balance == model_cost - 1 # untouched
|
||||
assert key.reserved_balance == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 5 — balance exactly equal to model cost → succeeds
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pay_for_request_succeeds_when_balance_equals_cost(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Balance == model cost: the request should be reserved successfully."""
|
||||
from routstr.auth import pay_for_request
|
||||
|
||||
model_cost = 10_000
|
||||
key = _key(balance=model_cost)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Should not raise
|
||||
await pay_for_request(key, model_cost, integration_session)
|
||||
|
||||
await integration_session.refresh(key)
|
||||
assert key.reserved_balance == model_cost
|
||||
assert key.balance == model_cost # balance unchanged, only reserved goes up
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 6 — HTTP layer returns 402 JSON with the right shape
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_http_402_response_shape_on_insufficient_balance(
|
||||
integration_client: AsyncClient,
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""
|
||||
End-to-end: POST /v1/chat/completions with a key whose balance is far below
|
||||
the mocked model cost returns HTTP 402 with the expected JSON error body.
|
||||
|
||||
Matches exactly the log snippet in the bug report:
|
||||
'Insufficient balance: X mSats required for this model. Y available.'
|
||||
"""
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
model_cost = 622_888
|
||||
user_balance = 20_320
|
||||
|
||||
key = _key(balance=user_balance)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Minimal model stub so proxy routing doesn't 400 before reaching balance check
|
||||
mock_model = MagicMock()
|
||||
mock_model.sats_pricing = None
|
||||
|
||||
# Upstream stub — never reached because balance check fires first
|
||||
mock_upstream = MagicMock()
|
||||
mock_upstream.prepare_headers = MagicMock(return_value={})
|
||||
|
||||
with (
|
||||
patch("routstr.proxy.get_model_instance", return_value=mock_model),
|
||||
patch("routstr.proxy.get_provider_for_model", return_value=[mock_upstream]),
|
||||
# Patch where it is used (proxy imports it at module level)
|
||||
patch(
|
||||
"routstr.proxy.get_max_cost_for_model",
|
||||
new=AsyncMock(return_value=model_cost),
|
||||
),
|
||||
):
|
||||
response = await integration_client.post(
|
||||
"/v1/chat/completions",
|
||||
headers={"Authorization": f"Bearer sk-{key.hashed_key}"},
|
||||
json={
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 402
|
||||
body = response.json()
|
||||
# FastAPI wraps HTTPException detail under "detail"
|
||||
error = body["detail"]["error"]
|
||||
assert error["code"] == "insufficient_balance"
|
||||
assert error["type"] == "insufficient_quota"
|
||||
assert str(model_cost) in error["message"]
|
||||
assert str(user_balance) in error["message"]
|
||||
|
||||
# Balance must be completely untouched
|
||||
await integration_session.refresh(key)
|
||||
assert key.balance == user_balance
|
||||
assert key.reserved_balance == 0
|
||||
assert key.total_spent == 0
|
||||
@@ -1,142 +0,0 @@
|
||||
import time
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import pytest
|
||||
from sqlmodel import select
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.auth import pay_for_request
|
||||
from routstr.core.db import ApiKey
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_validity_date(integration_session: AsyncSession) -> None:
|
||||
# 1. Create a key that is expired
|
||||
expired_time = int(time.time()) - 3600
|
||||
key = ApiKey(hashed_key="expired_key", balance=1000, validity_date=expired_time)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
# 2. Try to pay for a request - should fail
|
||||
with pytest.raises(Exception) as excinfo:
|
||||
await pay_for_request(key, 100, integration_session)
|
||||
assert "expired" in str(excinfo.value).lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_balance_limit(integration_session: AsyncSession) -> None:
|
||||
# 1. Create a key with a balance limit
|
||||
key = ApiKey(
|
||||
hashed_key="limited_key", balance=10000, balance_limit=500, total_spent=450
|
||||
)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
# 2. Try to pay for a request that exceeds the limit
|
||||
with pytest.raises(Exception) as excinfo:
|
||||
await pay_for_request(key, 100, integration_session)
|
||||
assert "limit exceeded" in str(excinfo.value).lower()
|
||||
|
||||
# 3. Try to pay for a request that fits
|
||||
await pay_for_request(key, 50, integration_session)
|
||||
await integration_session.refresh(key)
|
||||
# Note: total_spent is updated in adjust_payment_for_tokens,
|
||||
# but pay_for_request checks it.
|
||||
# In our current logic, pay_for_request checks (total_spent + cost) > balance_limit.
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_daily_reset_policy(integration_session: AsyncSession) -> None:
|
||||
# 1. Create a key with a daily reset policy and old reset date
|
||||
yesterday = int((datetime.now() - timedelta(days=1)).timestamp())
|
||||
key = ApiKey(
|
||||
hashed_key="daily_reset_key",
|
||||
balance=10000,
|
||||
balance_limit=1000,
|
||||
balance_limit_reset="daily",
|
||||
balance_limit_reset_date=yesterday,
|
||||
total_spent=900,
|
||||
)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
# 2. Pay for a request - should trigger reset first because it's a new day
|
||||
# Request is 200, total_spent is 900. 900+200 > 1000,
|
||||
# but reset should happen making total_spent 0, then 0+200 < 1000.
|
||||
await pay_for_request(key, 200, integration_session)
|
||||
|
||||
await integration_session.refresh(key)
|
||||
assert key.total_spent == 0 # Reset in pay_for_request happens before charging
|
||||
# Wait, the charging logic in pay_for_request increments parent/billing_key's total_requests,
|
||||
# but total_spent is updated in adjust_payment_for_tokens.
|
||||
# However, the reset logic sets total_spent to 0.
|
||||
assert key.balance_limit_reset_date is not None
|
||||
assert key.balance_limit_reset_date > yesterday
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_periodic_key_reset_job(integration_session: AsyncSession) -> None:
|
||||
# 1. Create multiple keys needing reset
|
||||
yesterday = int((datetime.now() - timedelta(days=1)).timestamp())
|
||||
key1 = ApiKey(
|
||||
hashed_key="job_reset_key_1",
|
||||
balance=1000,
|
||||
balance_limit=1000,
|
||||
balance_limit_reset="daily",
|
||||
balance_limit_reset_date=yesterday,
|
||||
total_spent=500,
|
||||
)
|
||||
key2 = ApiKey(
|
||||
hashed_key="job_reset_key_2",
|
||||
balance=1000,
|
||||
balance_limit=1000,
|
||||
balance_limit_reset="daily",
|
||||
balance_limit_reset_date=yesterday,
|
||||
total_spent=800,
|
||||
)
|
||||
integration_session.add(key1)
|
||||
integration_session.add(key2)
|
||||
await integration_session.commit()
|
||||
|
||||
# 2. Run the periodic reset logic manually (mocking the background task loop)
|
||||
# We can't easily run the actual loop because it has a sleep,
|
||||
# but we can test the logic inside.
|
||||
|
||||
# Implementation of periodic_key_reset logic for testing:
|
||||
stmt = select(ApiKey).where(ApiKey.balance_limit_reset != None) # noqa: E711
|
||||
keys = (await integration_session.exec(stmt)).all()
|
||||
now = int(time.time())
|
||||
for k in keys:
|
||||
if k.hashed_key in ["job_reset_key_1", "job_reset_key_2"]:
|
||||
k.total_spent = 0
|
||||
k.balance_limit_reset_date = now
|
||||
integration_session.add(k)
|
||||
await integration_session.commit()
|
||||
|
||||
# 3. Verify resets
|
||||
await integration_session.refresh(key1)
|
||||
await integration_session.refresh(key2)
|
||||
assert key1.total_spent == 0
|
||||
assert key2.total_spent == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_does_not_delete_key(integration_session: AsyncSession) -> None:
|
||||
# This requires mocking the router call or testing the logic in balance.py
|
||||
from routstr.balance import ApiKey
|
||||
|
||||
key = ApiKey(hashed_key="refund_test_key", balance=1000, reserved_balance=100)
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Logic from refund_wallet_endpoint:
|
||||
key.balance = 0
|
||||
key.reserved_balance = 0
|
||||
integration_session.add(key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Verify key still exists
|
||||
fetched_key = await integration_session.get(ApiKey, "refund_test_key")
|
||||
assert fetched_key is not None
|
||||
assert fetched_key.balance == 0
|
||||
assert fetched_key.reserved_balance == 0
|
||||
@@ -1,143 +0,0 @@
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Any, AsyncGenerator, cast
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from routstr.core.db import AsyncSession, ModelRow, UpstreamProviderRow
|
||||
from routstr.payment.models import Architecture, Model, Pricing
|
||||
from routstr.proxy import refresh_model_maps
|
||||
from routstr.upstream.base import BaseUpstreamProvider
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforce_lowest_provider_fee_for_same_url(
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
"""Test that the algorithm selects the provider with the lowest fee when URLs match."""
|
||||
|
||||
# 1. Create two providers with the same URL but different fees
|
||||
url = "https://api.example.com"
|
||||
p1 = UpstreamProviderRow(
|
||||
provider_type="custom",
|
||||
base_url=url,
|
||||
api_key="key1",
|
||||
enabled=True,
|
||||
provider_fee=1.01,
|
||||
)
|
||||
p2 = UpstreamProviderRow(
|
||||
provider_type="custom",
|
||||
base_url=url,
|
||||
api_key="key2",
|
||||
enabled=True,
|
||||
provider_fee=1.05,
|
||||
)
|
||||
|
||||
integration_session.add(p1)
|
||||
integration_session.add(p2)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(p1)
|
||||
await integration_session.refresh(p2)
|
||||
|
||||
assert p1.id is not None
|
||||
assert p2.id is not None
|
||||
|
||||
# 2. Add a model for each provider
|
||||
m1 = ModelRow(
|
||||
id="model-a",
|
||||
name="Model A",
|
||||
created=1,
|
||||
description="desc",
|
||||
context_length=100,
|
||||
architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}',
|
||||
pricing='{"prompt": 1.0, "completion": 1.0}',
|
||||
upstream_provider_id=p1.id,
|
||||
enabled=True,
|
||||
)
|
||||
m2 = ModelRow(
|
||||
id="model-a",
|
||||
name="Model A",
|
||||
created=1,
|
||||
description="desc",
|
||||
context_length=100,
|
||||
architecture='{"modality": "text", "input_modalities": ["text"], "output_modalities": ["text"], "tokenizer": "tiktoken", "instruct_type": "chat"}',
|
||||
pricing='{"prompt": 1.0, "completion": 1.0}',
|
||||
upstream_provider_id=p2.id,
|
||||
enabled=True,
|
||||
)
|
||||
|
||||
integration_session.add(m1)
|
||||
integration_session.add(m2)
|
||||
await integration_session.commit()
|
||||
|
||||
# 3. Create mock provider instances
|
||||
class MockProvider(BaseUpstreamProvider):
|
||||
db_id: int
|
||||
|
||||
def __init__(self, db_id: int, base_url: str, api_key: str, fee: float):
|
||||
super().__init__(base_url, api_key, fee)
|
||||
self.db_id = db_id
|
||||
self.provider_type = "custom"
|
||||
|
||||
def get_cached_models(self) -> list[Model]:
|
||||
return [
|
||||
Model(
|
||||
id="model-a",
|
||||
name="Model A",
|
||||
created=1,
|
||||
description="desc",
|
||||
context_length=100,
|
||||
architecture=Architecture(
|
||||
modality="text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="tiktoken",
|
||||
instruct_type="chat",
|
||||
),
|
||||
pricing=Pricing(prompt=1.0, completion=1.0),
|
||||
enabled=True,
|
||||
upstream_provider_id=self.db_id,
|
||||
)
|
||||
]
|
||||
|
||||
async def refresh_models_cache(self) -> None:
|
||||
pass
|
||||
|
||||
def prepare_headers(self, request_headers: dict[str, str]) -> dict[str, str]:
|
||||
return request_headers
|
||||
|
||||
# 4. Inject mock providers into the proxy
|
||||
from routstr import proxy
|
||||
|
||||
assert p1.id is not None
|
||||
assert p2.id is not None
|
||||
|
||||
# Need to patch proxy._upstreams and proxy.create_session
|
||||
mp1: MockProvider = MockProvider(p1.id, url, "key1", 1.01)
|
||||
mp2: MockProvider = MockProvider(p2.id, url, "key2", 1.05)
|
||||
|
||||
with (
|
||||
patch("routstr.proxy._upstreams", [mp1, mp2]),
|
||||
patch("routstr.proxy.create_session") as mock_session_factory,
|
||||
):
|
||||
# Configure mock_session_factory to return a session that uses the test engine
|
||||
@asynccontextmanager
|
||||
async def mock_create_session() -> AsyncGenerator[AsyncSession, None]:
|
||||
yield integration_session
|
||||
|
||||
mock_session_factory.return_value = mock_create_session()
|
||||
|
||||
await refresh_model_maps()
|
||||
|
||||
# 5. Check which provider is selected for 'model-a'
|
||||
provider_map = proxy.get_provider_for_model("model-a")
|
||||
|
||||
# Assertions
|
||||
assert provider_map is not None
|
||||
assert len(provider_map) >= 1
|
||||
|
||||
# Check the first one, cast to MockProvider to access db_id
|
||||
best_provider = cast(MockProvider, provider_map[0])
|
||||
assert best_provider.db_id == p1.id
|
||||
assert best_provider.provider_fee == 1.01
|
||||
@@ -3,17 +3,13 @@ Integration tests for provider management functionality.
|
||||
Tests GET /v1/providers/ endpoint for listing and managing providers.
|
||||
"""
|
||||
|
||||
import time
|
||||
from types import TracebackType
|
||||
from typing import Any, Generator
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
from httpx import AsyncClient
|
||||
|
||||
from routstr.core.admin import admin_sessions
|
||||
from routstr.core.db import UpstreamProviderRow
|
||||
from routstr.nostr.discovery import _PROVIDERS_CACHE
|
||||
from routstr.discovery import _PROVIDERS_CACHE
|
||||
|
||||
from .utils import ResponseValidator
|
||||
|
||||
@@ -75,10 +71,9 @@ async def test_providers_endpoint_default_response(
|
||||
}
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
with patch("routstr.nostr.discovery.fetch_provider_health") as mock_fetch:
|
||||
with patch("routstr.discovery.fetch_provider_health") as mock_fetch:
|
||||
# Configure mock to return appropriate responses
|
||||
mock_fetch.side_effect = lambda url: mock_fetch_responses.get(
|
||||
url, {"status_code": 500, "json": {"error": "Unknown provider"}}
|
||||
@@ -140,10 +135,9 @@ async def test_providers_endpoint_with_include_json(
|
||||
}
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
with patch("routstr.nostr.discovery.fetch_provider_health") as mock_fetch:
|
||||
with patch("routstr.discovery.fetch_provider_health") as mock_fetch:
|
||||
mock_fetch.return_value = {
|
||||
"status_code": 200,
|
||||
"json": mock_provider_response,
|
||||
@@ -215,10 +209,9 @@ async def test_providers_data_structure_validation(
|
||||
}
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
with patch("routstr.nostr.discovery.fetch_provider_health") as mock_fetch:
|
||||
with patch("routstr.discovery.fetch_provider_health") as mock_fetch:
|
||||
mock_fetch.return_value = mock_health_response
|
||||
|
||||
response = await integration_client.get("/v1/providers/?include_json=true")
|
||||
@@ -263,8 +256,7 @@ async def test_providers_endpoint_no_providers_found(
|
||||
]
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
response = await integration_client.get("/v1/providers/")
|
||||
|
||||
@@ -325,11 +317,10 @@ async def test_providers_endpoint_offline_providers(
|
||||
}
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
with patch(
|
||||
"routstr.nostr.discovery.fetch_provider_health",
|
||||
"routstr.discovery.fetch_provider_health",
|
||||
side_effect=mock_fetch_provider_health,
|
||||
):
|
||||
response = await integration_client.get("/v1/providers/?include_json=true")
|
||||
@@ -395,10 +386,9 @@ async def test_providers_endpoint_duplicate_urls(
|
||||
]
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
with patch("routstr.nostr.discovery.fetch_provider_health") as mock_fetch:
|
||||
with patch("routstr.discovery.fetch_provider_health") as mock_fetch:
|
||||
mock_fetch.return_value = {
|
||||
"status_code": 200,
|
||||
"endpoint": "root",
|
||||
@@ -435,8 +425,7 @@ async def test_providers_endpoint_nostr_relay_failures(
|
||||
raise Exception("Connection to relay failed")
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
side_effect=failing_query,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", side_effect=failing_query
|
||||
):
|
||||
response = await integration_client.get("/v1/providers/")
|
||||
|
||||
@@ -474,10 +463,9 @@ async def test_providers_endpoint_malformed_urls(
|
||||
]
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
with patch("routstr.nostr.discovery.fetch_provider_health") as mock_fetch:
|
||||
with patch("routstr.discovery.fetch_provider_health") as mock_fetch:
|
||||
mock_fetch.return_value = {"status_code": 200, "json": {"status": "online"}}
|
||||
|
||||
response = await integration_client.get("/v1/providers/")
|
||||
@@ -507,10 +495,9 @@ async def test_providers_endpoint_response_format(
|
||||
]
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
with patch("routstr.nostr.discovery.fetch_provider_health") as mock_fetch:
|
||||
with patch("routstr.discovery.fetch_provider_health") as mock_fetch:
|
||||
mock_fetch.return_value = {"status_code": 200, "json": {"status": "online"}}
|
||||
|
||||
# Test default format
|
||||
@@ -558,10 +545,9 @@ async def test_providers_endpoint_concurrent_requests(
|
||||
]
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
with patch("routstr.nostr.discovery.fetch_provider_health") as mock_fetch:
|
||||
with patch("routstr.discovery.fetch_provider_health") as mock_fetch:
|
||||
mock_fetch.return_value = {"status_code": 200, "json": {"status": "online"}}
|
||||
|
||||
# Create concurrent requests
|
||||
@@ -601,10 +587,9 @@ async def test_providers_endpoint_parameter_validation(
|
||||
]
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
with patch("routstr.nostr.discovery.fetch_provider_health") as mock_fetch:
|
||||
with patch("routstr.discovery.fetch_provider_health") as mock_fetch:
|
||||
mock_fetch.return_value = {"status_code": 200, "json": {"status": "online"}}
|
||||
|
||||
# Test various parameter values
|
||||
@@ -654,10 +639,9 @@ async def test_no_database_changes_during_provider_operations(
|
||||
]
|
||||
|
||||
with patch(
|
||||
"routstr.nostr.discovery.query_nostr_relay_for_providers",
|
||||
return_value=mock_events,
|
||||
"routstr.discovery.query_nostr_relay_for_providers", return_value=mock_events
|
||||
):
|
||||
with patch("routstr.nostr.discovery.fetch_provider_health") as mock_fetch:
|
||||
with patch("routstr.discovery.fetch_provider_health") as mock_fetch:
|
||||
mock_fetch.return_value = {"status_code": 200, "json": {"status": "online"}}
|
||||
|
||||
# Make multiple requests with different parameters
|
||||
@@ -682,88 +666,3 @@ async def test_no_database_changes_during_provider_operations(
|
||||
assert final_diff["api_keys"]["added"] == []
|
||||
assert final_diff["api_keys"]["modified"] == []
|
||||
assert final_diff["api_keys"]["removed"] == []
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_admin_routstr_topup_retries_transient_upstream_failure(
|
||||
integration_client: AsyncClient,
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
admin_token = "test-admin-token"
|
||||
admin_sessions[admin_token] = int(time.time()) + 3600
|
||||
integration_client.headers["Authorization"] = f"Bearer {admin_token}"
|
||||
|
||||
provider = UpstreamProviderRow(
|
||||
provider_type="routstr",
|
||||
base_url="https://node.example",
|
||||
api_key="sk-upstream-test",
|
||||
enabled=True,
|
||||
provider_fee=1.01,
|
||||
)
|
||||
integration_session.add(provider)
|
||||
await integration_session.commit()
|
||||
await integration_session.refresh(provider)
|
||||
|
||||
class MockResponse:
|
||||
def __init__(self, status_code: int, data: dict[str, Any] | None = None):
|
||||
self.status_code = status_code
|
||||
self._data = data or {}
|
||||
self.text = str(self._data)
|
||||
|
||||
def json(self) -> dict[str, Any]:
|
||||
return self._data
|
||||
|
||||
class MockAsyncClient:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
async def __aenter__(self) -> "MockAsyncClient":
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc: BaseException | None,
|
||||
tb: TracebackType | None,
|
||||
) -> None:
|
||||
return None
|
||||
|
||||
async def post(
|
||||
self, url: str, json: dict[str, Any], headers: dict[str, str]
|
||||
) -> MockResponse:
|
||||
self.calls += 1
|
||||
assert url == "https://node.example/v1/balance/lightning/invoice"
|
||||
assert json["amount_sats"] == 10
|
||||
assert json["purpose"] == "topup"
|
||||
assert json["api_key"] == "sk-upstream-test"
|
||||
assert headers["Authorization"] == "Bearer sk-upstream-test"
|
||||
|
||||
if self.calls == 1:
|
||||
return MockResponse(500, {"detail": "warmup failure"})
|
||||
|
||||
return MockResponse(
|
||||
200,
|
||||
{
|
||||
"bolt11": "lnbc1testinvoice",
|
||||
"invoice_id": "invoice-123",
|
||||
},
|
||||
)
|
||||
|
||||
mock_client = MockAsyncClient()
|
||||
|
||||
try:
|
||||
with patch("httpx.AsyncClient", return_value=mock_client):
|
||||
response = await integration_client.post(
|
||||
f"/admin/api/upstream-providers/{provider.id}/topup",
|
||||
json={"amount": 10},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
data = response.json()
|
||||
assert data["ok"] is True
|
||||
assert data["topup_data"]["payment_request"] == "lnbc1testinvoice"
|
||||
assert data["topup_data"]["invoice_id"] == "invoice-123"
|
||||
assert mock_client.calls == 2
|
||||
finally:
|
||||
admin_sessions.pop(admin_token, None)
|
||||
|
||||
@@ -171,11 +171,10 @@ async def test_proxy_get_unauthorized_access(integration_client: AsyncClient) ->
|
||||
assert response.status_code == 200 # GET requests are allowed
|
||||
|
||||
# Test 2: POST requests without auth should return 401
|
||||
# Note: Model validation happens before auth, so missing model returns 400
|
||||
response = await integration_client.post(
|
||||
"/v1/chat/completions", json={"test": "data"}
|
||||
)
|
||||
assert response.status_code in [400, 401] # Accept both for now
|
||||
assert response.status_code == 401
|
||||
|
||||
# Test 3: POST with invalid API key
|
||||
# Note: After refactor, model validation may happen before auth validation
|
||||
@@ -551,6 +550,9 @@ async def test_proxy_get_concurrent_requests(
|
||||
assert response.status_code == 200
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_proxy_get_response_format_preservation(
|
||||
|
||||
@@ -1,268 +0,0 @@
|
||||
"""
|
||||
Tests for the reservation lifecycle:
|
||||
|
||||
1. Reserve → reserved_balance increases, available (total_balance) decreases.
|
||||
2. Reserve → revert → reserved_balance restored, balance untouched.
|
||||
3. Reserve → finalise → reserved_balance released, balance charged.
|
||||
4. Two parallel reserves, only one fits → second blocked with 402.
|
||||
5. Three parallel reserves, two fit, third blocked with 402.
|
||||
6. Sequential reserves until balance exhausted → next request blocked.
|
||||
|
||||
Reservation invariant enforced by the atomic WHERE clause in pay_for_request:
|
||||
balance - reserved_balance >= cost_per_request
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.auth import pay_for_request, revert_pay_for_request
|
||||
from routstr.core.db import ApiKey, create_session
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def _make_key(balance: int, reserved: int = 0) -> ApiKey:
|
||||
return ApiKey(
|
||||
hashed_key=f"test_{uuid.uuid4().hex}",
|
||||
balance=balance,
|
||||
reserved_balance=reserved,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
)
|
||||
|
||||
|
||||
async def _persist(session: AsyncSession, key: ApiKey) -> ApiKey:
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
await session.refresh(key)
|
||||
return key
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 1 — Reserve: reserved_balance increases, available balance decreases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_reserve_increases_reserved_balance(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""pay_for_request must increment reserved_balance by cost_per_request."""
|
||||
cost = 100
|
||||
key = await _persist(integration_session, _make_key(balance=500))
|
||||
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
await integration_session.refresh(key)
|
||||
|
||||
assert key.reserved_balance == cost
|
||||
assert key.balance == 500 # balance column is NOT decremented on reserve
|
||||
assert key.total_balance == 500 - cost # available = balance - reserved
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 2 — Revert: reserved_balance restored, balance untouched
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_revert_releases_reservation(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""revert_pay_for_request must release the reservation without touching balance."""
|
||||
cost = 150
|
||||
key = await _persist(integration_session, _make_key(balance=300))
|
||||
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
await integration_session.refresh(key)
|
||||
assert key.reserved_balance == cost
|
||||
|
||||
await revert_pay_for_request(key, integration_session, cost)
|
||||
await integration_session.refresh(key)
|
||||
|
||||
assert key.reserved_balance == 0
|
||||
assert key.balance == 300 # balance unchanged after revert
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 3 — Finalise: reservation released + balance charged
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_finalise_releases_reservation_and_charges_balance(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""adjust_payment_for_tokens must zero reserved_balance and deduct actual cost."""
|
||||
from unittest.mock import patch
|
||||
|
||||
from routstr.auth import adjust_payment_for_tokens
|
||||
from routstr.payment.cost_calculation import CostData
|
||||
|
||||
cost = 100
|
||||
actual = 80 # actual < reserved → refund path
|
||||
key = await _persist(integration_session, _make_key(balance=500))
|
||||
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
await integration_session.refresh(key)
|
||||
assert key.reserved_balance == cost
|
||||
|
||||
cost_data = CostData(
|
||||
base_msats=0,
|
||||
input_msats=40,
|
||||
output_msats=40,
|
||||
total_msats=actual,
|
||||
total_usd=0.0,
|
||||
input_tokens=50,
|
||||
output_tokens=50,
|
||||
)
|
||||
response_data = {"model": "test-model", "usage": {"prompt_tokens": 50, "completion_tokens": 50}}
|
||||
|
||||
with patch("routstr.auth.calculate_cost", return_value=cost_data):
|
||||
await adjust_payment_for_tokens(key, response_data, integration_session, cost)
|
||||
|
||||
await integration_session.refresh(key)
|
||||
|
||||
assert key.reserved_balance == 0
|
||||
assert key.balance == 500 - actual
|
||||
assert key.total_spent == actual
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 4 — Concurrent: second parallel reserve blocked when balance exhausted
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_concurrent_second_reserve_blocked_when_balance_exhausted(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""When two requests race for the same balance, only one succeeds; the other gets 402."""
|
||||
cost = 300
|
||||
key_hash = f"test_concurrent_{uuid.uuid4().hex}"
|
||||
|
||||
async with create_session() as session:
|
||||
key = ApiKey(
|
||||
hashed_key=key_hash,
|
||||
balance=300, # exactly enough for ONE reservation
|
||||
reserved_balance=0,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
)
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
results: list[str] = []
|
||||
|
||||
async def attempt_reserve() -> None:
|
||||
async with create_session() as session:
|
||||
fresh_key = await session.get(ApiKey, key_hash)
|
||||
assert fresh_key is not None
|
||||
try:
|
||||
await pay_for_request(fresh_key, cost, session)
|
||||
results.append("success")
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 402
|
||||
results.append("blocked")
|
||||
|
||||
await asyncio.gather(attempt_reserve(), attempt_reserve())
|
||||
|
||||
assert sorted(results) == ["blocked", "success"], (
|
||||
f"Expected exactly one success and one 402, got: {results}"
|
||||
)
|
||||
|
||||
async with create_session() as session:
|
||||
final = await session.get(ApiKey, key_hash)
|
||||
assert final is not None
|
||||
|
||||
# reserved_balance must equal exactly one reservation (not two)
|
||||
assert final.reserved_balance == cost, (
|
||||
f"Expected reserved_balance={cost}, got {final.reserved_balance}"
|
||||
)
|
||||
assert final.balance == 300, "Balance column must not be modified by reservation"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 5 — Concurrent: three requests, two fit, third blocked
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_three_parallel_reserves_third_blocked(
|
||||
patched_db_engine: None,
|
||||
) -> None:
|
||||
"""Balance covers two reservations exactly; the third concurrent request must be blocked."""
|
||||
cost = 100
|
||||
key_hash = f"test_three_parallel_{uuid.uuid4().hex}"
|
||||
|
||||
async with create_session() as session:
|
||||
key = ApiKey(
|
||||
hashed_key=key_hash,
|
||||
balance=200, # fits exactly 2 reservations of 100
|
||||
reserved_balance=0,
|
||||
total_spent=0,
|
||||
total_requests=0,
|
||||
)
|
||||
session.add(key)
|
||||
await session.commit()
|
||||
|
||||
results: list[str] = []
|
||||
|
||||
async def attempt_reserve() -> None:
|
||||
async with create_session() as session:
|
||||
fresh_key = await session.get(ApiKey, key_hash)
|
||||
assert fresh_key is not None
|
||||
try:
|
||||
await pay_for_request(fresh_key, cost, session)
|
||||
results.append("success")
|
||||
except HTTPException as exc:
|
||||
assert exc.status_code == 402
|
||||
results.append("blocked")
|
||||
|
||||
await asyncio.gather(
|
||||
attempt_reserve(),
|
||||
attempt_reserve(),
|
||||
attempt_reserve(),
|
||||
)
|
||||
|
||||
successes = results.count("success")
|
||||
blocked = results.count("blocked")
|
||||
|
||||
assert successes == 2, f"Expected 2 successes, got {successes}: {results}"
|
||||
assert blocked == 1, f"Expected 1 blocked, got {blocked}: {results}"
|
||||
|
||||
async with create_session() as session:
|
||||
final = await session.get(ApiKey, key_hash)
|
||||
assert final is not None
|
||||
|
||||
assert final.reserved_balance == cost * 2, (
|
||||
f"Expected reserved_balance={cost * 2}, got {final.reserved_balance}"
|
||||
)
|
||||
assert final.balance == 200, "Balance column must not be modified by reservation"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test 6 — Sequential exhaustion: reserve until empty, next request blocked
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sequential_reserves_block_when_balance_exhausted(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Repeated reservations should block as soon as available balance drops below cost."""
|
||||
cost = 100
|
||||
key = await _persist(integration_session, _make_key(balance=250))
|
||||
|
||||
# First two succeed (100 + 100 = 200 ≤ 250)
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
await integration_session.refresh(key)
|
||||
assert key.reserved_balance == 200
|
||||
assert key.total_balance == 50 # 250 - 200
|
||||
|
||||
# Third: only 50 available, need 100 → blocked
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await pay_for_request(key, cost, integration_session)
|
||||
|
||||
assert exc_info.value.status_code == 402
|
||||
await integration_session.refresh(key)
|
||||
assert key.reserved_balance == 200 # unchanged after failed reserve
|
||||
@@ -133,16 +133,13 @@ async def test_reserved_balance_with_successful_requests(
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_revert_with_zero_reserved_balance_is_noop(
|
||||
async def test_insufficient_reserved_balance_for_revert(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that revert_pay_for_request is a no-op when reserved_balance is 0.
|
||||
|
||||
Previously this would drive reserved_balance negative. With the floor guard,
|
||||
it should return False and leave reserved_balance at 0.
|
||||
"""
|
||||
"""Test revert_pay_for_request behavior with insufficient reserved balance."""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
# Create key with zero reserved balance
|
||||
unique_key = f"test_revert_key_{uuid.uuid4().hex[:8]}"
|
||||
test_key = ApiKey(
|
||||
hashed_key=unique_key,
|
||||
@@ -152,211 +149,17 @@ async def test_revert_with_zero_reserved_balance_is_noop(
|
||||
integration_session.add(test_key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Try to revert more than available — should be a no-op
|
||||
result = await revert_pay_for_request(test_key, integration_session, 100)
|
||||
# Try to revert more than available
|
||||
# Note: Current implementation allows reserved_balance to go negative
|
||||
await revert_pay_for_request(test_key, integration_session, 100)
|
||||
|
||||
# Refresh to get updated values
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
assert result is False, "Revert should return False when reservation already released"
|
||||
assert test_key.reserved_balance == 0, (
|
||||
f"Reserved balance should remain 0, got: {test_key.reserved_balance}"
|
||||
# Current implementation allows negative reserved balance
|
||||
assert test_key.reserved_balance == -100, (
|
||||
f"Expected reserved_balance to be -100, got: {test_key.reserved_balance}"
|
||||
)
|
||||
assert test_key.total_requests == 0, (
|
||||
f"Total requests should remain 0, got: {test_key.total_requests}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_revert_with_sufficient_reserved_balance_succeeds(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that revert_pay_for_request works correctly when there is enough reserved balance."""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
unique_key = f"test_revert_ok_{uuid.uuid4().hex[:8]}"
|
||||
test_key = ApiKey(
|
||||
hashed_key=unique_key,
|
||||
balance=5000,
|
||||
reserved_balance=500,
|
||||
total_requests=3,
|
||||
)
|
||||
integration_session.add(test_key)
|
||||
await integration_session.commit()
|
||||
|
||||
result = await revert_pay_for_request(test_key, integration_session, 500)
|
||||
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
assert result is True, "Revert should return True on success"
|
||||
assert test_key.reserved_balance == 0, (
|
||||
f"Reserved balance should be 0, got: {test_key.reserved_balance}"
|
||||
)
|
||||
assert test_key.total_requests == 2, (
|
||||
f"Total requests should be 2, got: {test_key.total_requests}"
|
||||
)
|
||||
assert test_key.balance == 5000, "Balance should not change on revert"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_revert_partial_reserved_balance_is_noop(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that reverting more than the current reserved_balance is a no-op."""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
unique_key = f"test_revert_partial_{uuid.uuid4().hex[:8]}"
|
||||
test_key = ApiKey(
|
||||
hashed_key=unique_key,
|
||||
balance=5000,
|
||||
reserved_balance=50,
|
||||
total_requests=1,
|
||||
)
|
||||
integration_session.add(test_key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Try to revert 500 when only 50 is reserved — should be no-op
|
||||
result = await revert_pay_for_request(test_key, integration_session, 500)
|
||||
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
assert result is False, "Revert should fail when cost > reserved_balance"
|
||||
assert test_key.reserved_balance == 50, (
|
||||
f"Reserved balance should stay at 50, got: {test_key.reserved_balance}"
|
||||
)
|
||||
assert test_key.total_requests == 1, (
|
||||
f"Total requests should stay at 1, got: {test_key.total_requests}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_double_revert_prevented(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that calling revert twice doesn't drive reserved_balance negative.
|
||||
|
||||
This simulates the double-revert scenario where both upstream/base.py
|
||||
and proxy.py attempt to revert the same reservation.
|
||||
"""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
unique_key = f"test_double_revert_{uuid.uuid4().hex[:8]}"
|
||||
test_key = ApiKey(
|
||||
hashed_key=unique_key,
|
||||
balance=10000,
|
||||
reserved_balance=500,
|
||||
total_requests=5,
|
||||
)
|
||||
integration_session.add(test_key)
|
||||
await integration_session.commit()
|
||||
|
||||
# First revert — should succeed
|
||||
result1 = await revert_pay_for_request(test_key, integration_session, 500)
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
assert result1 is True
|
||||
assert test_key.reserved_balance == 0
|
||||
assert test_key.total_requests == 4
|
||||
|
||||
# Second revert of the same amount — should be no-op
|
||||
result2 = await revert_pay_for_request(test_key, integration_session, 500)
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
assert result2 is False, "Second revert should be a no-op"
|
||||
assert test_key.reserved_balance == 0, (
|
||||
f"Reserved balance should stay 0, got: {test_key.reserved_balance}"
|
||||
)
|
||||
assert test_key.total_requests == 4, (
|
||||
f"Total requests should stay 4, got: {test_key.total_requests}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_sequential_reverts_never_go_negative(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that multiple reverts don't cause negative reserved_balance.
|
||||
|
||||
Simulates the double-revert scenario where multiple code paths
|
||||
attempt to revert the same reservation.
|
||||
"""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
unique_key = f"test_multi_revert_{uuid.uuid4().hex[:8]}"
|
||||
test_key = ApiKey(
|
||||
hashed_key=unique_key,
|
||||
balance=10000,
|
||||
reserved_balance=500,
|
||||
total_requests=5,
|
||||
)
|
||||
integration_session.add(test_key)
|
||||
await integration_session.commit()
|
||||
|
||||
# Run 5 sequential reverts for the same 500 reservation
|
||||
results = []
|
||||
for _ in range(5):
|
||||
r = await revert_pay_for_request(test_key, integration_session, 500)
|
||||
results.append(r)
|
||||
|
||||
await integration_session.refresh(test_key)
|
||||
|
||||
# Exactly one should succeed, rest should be no-ops
|
||||
success_count = sum(1 for r in results if r is True)
|
||||
assert success_count == 1, (
|
||||
f"Exactly one revert should succeed, got {success_count} successes"
|
||||
)
|
||||
assert test_key.reserved_balance == 0, (
|
||||
f"Reserved balance should be 0, got: {test_key.reserved_balance}"
|
||||
)
|
||||
assert test_key.reserved_balance >= 0, (
|
||||
f"Reserved balance went negative: {test_key.reserved_balance}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_child_key_revert_floor_guard(
|
||||
integration_session: AsyncSession,
|
||||
) -> None:
|
||||
"""Test that child key reserved_balance also has floor guard on revert."""
|
||||
from routstr.auth import revert_pay_for_request
|
||||
|
||||
parent_key_hash = f"test_parent_{uuid.uuid4().hex[:8]}"
|
||||
child_key_hash = f"test_child_{uuid.uuid4().hex[:8]}"
|
||||
|
||||
parent_key = ApiKey(
|
||||
hashed_key=parent_key_hash,
|
||||
balance=10000,
|
||||
reserved_balance=500,
|
||||
total_requests=3,
|
||||
)
|
||||
child_key = ApiKey(
|
||||
hashed_key=child_key_hash,
|
||||
balance=0,
|
||||
reserved_balance=500,
|
||||
total_requests=3,
|
||||
parent_key_hash=parent_key_hash,
|
||||
)
|
||||
integration_session.add(parent_key)
|
||||
integration_session.add(child_key)
|
||||
await integration_session.commit()
|
||||
|
||||
# First revert succeeds
|
||||
result1 = await revert_pay_for_request(child_key, integration_session, 500)
|
||||
await integration_session.refresh(parent_key)
|
||||
await integration_session.refresh(child_key)
|
||||
|
||||
assert result1 is True
|
||||
assert parent_key.reserved_balance == 0
|
||||
assert child_key.reserved_balance == 0
|
||||
|
||||
# Second revert is a no-op for both parent and child
|
||||
result2 = await revert_pay_for_request(child_key, integration_session, 500)
|
||||
await integration_session.refresh(parent_key)
|
||||
await integration_session.refresh(child_key)
|
||||
|
||||
assert result2 is False
|
||||
assert parent_key.reserved_balance == 0, (
|
||||
f"Parent reserved_balance should stay 0, got: {parent_key.reserved_balance}"
|
||||
)
|
||||
assert child_key.reserved_balance == 0, (
|
||||
f"Child reserved_balance should stay 0, got: {child_key.reserved_balance}"
|
||||
assert test_key.total_requests == -1, (
|
||||
f"Expected total_requests to be -1, got: {test_key.total_requests}"
|
||||
)
|
||||
|
||||
@@ -13,7 +13,7 @@ import pytest
|
||||
from httpx import AsyncClient
|
||||
from sqlmodel import select
|
||||
|
||||
from routstr.core.db import ApiKey, CashuTransaction
|
||||
from routstr.core.db import ApiKey
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@@ -65,10 +65,9 @@ async def test_full_balance_refund_returns_cashu_token(
|
||||
except Exception as e:
|
||||
pytest.fail(f"Invalid Cashu token format: {e}")
|
||||
|
||||
# Try to use the API key - should still work but have 0 balance
|
||||
# Try to use the API key - should fail since it's been deleted
|
||||
response = await authenticated_client.get("/v1/wallet/")
|
||||
assert response.status_code == 200
|
||||
assert response.json()["balance"] == 0
|
||||
assert response.status_code == 401
|
||||
|
||||
# The refund token has been validated above by decoding it
|
||||
# The API key deletion has been verified by the 401 response
|
||||
@@ -262,16 +261,17 @@ async def test_database_state_after_refund(
|
||||
response = await authenticated_client.post("/v1/wallet/refund")
|
||||
assert response.status_code == 200
|
||||
|
||||
# Refresh the key to get the updated balance from the database
|
||||
await integration_session.refresh(key_before)
|
||||
# Verify key is deleted after refund
|
||||
result = await integration_session.execute(
|
||||
select(ApiKey).where(ApiKey.hashed_key == hashed_key) # type: ignore[arg-type]
|
||||
)
|
||||
assert result.scalar_one_or_none() is None
|
||||
|
||||
# Verify key balance is 0 after refund
|
||||
assert key_before.balance == 0
|
||||
|
||||
# Count total keys to ensure it wasn't deleted
|
||||
# Count total keys to ensure only the specific one was deleted
|
||||
result = await integration_session.execute(select(ApiKey))
|
||||
remaining_keys = result.scalars().all()
|
||||
assert len(remaining_keys) == 1
|
||||
# Should have no keys left (assuming clean test environment)
|
||||
assert len(remaining_keys) == 0
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@@ -356,70 +356,6 @@ async def test_concurrent_refund_requests(
|
||||
assert len(successful) + len(failed) == 5
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_rejects_concurrent_topup_on_same_key(
|
||||
authenticated_client: AsyncClient,
|
||||
testmint_wallet: Any,
|
||||
) -> None:
|
||||
"""Test refund returns 409 when a concurrent topup changes the balance first."""
|
||||
from routstr import balance as balance_module
|
||||
|
||||
wallet_response = await authenticated_client.get("/v1/wallet/")
|
||||
assert wallet_response.status_code == 200
|
||||
initial_balance = wallet_response.json()["balance"]
|
||||
|
||||
topup_amount_sat = 500
|
||||
topup_token = await testmint_wallet.mint_tokens(topup_amount_sat)
|
||||
|
||||
validate_called = asyncio.Event()
|
||||
allow_refund_to_continue = asyncio.Event()
|
||||
original_validate_bearer_key = balance_module.validate_bearer_key
|
||||
delayed_once = False
|
||||
|
||||
async def delayed_validate_bearer_key(*args: Any, **kwargs: Any) -> ApiKey:
|
||||
nonlocal delayed_once
|
||||
key = await original_validate_bearer_key(*args, **kwargs)
|
||||
if not delayed_once:
|
||||
delayed_once = True
|
||||
validate_called.set()
|
||||
await allow_refund_to_continue.wait()
|
||||
return key
|
||||
|
||||
async def issue_refund() -> Any:
|
||||
return await authenticated_client.post("/v1/wallet/refund")
|
||||
|
||||
async def issue_topup() -> Any:
|
||||
await validate_called.wait()
|
||||
try:
|
||||
return await authenticated_client.post(
|
||||
"/v1/wallet/topup", params={"cashu_token": topup_token}
|
||||
)
|
||||
finally:
|
||||
allow_refund_to_continue.set()
|
||||
|
||||
with patch(
|
||||
"routstr.balance.validate_bearer_key", new=delayed_validate_bearer_key
|
||||
):
|
||||
refund_response, topup_response = await asyncio.gather(
|
||||
issue_refund(), issue_topup()
|
||||
)
|
||||
|
||||
assert topup_response.status_code == 200
|
||||
assert topup_response.json()["msats"] == topup_amount_sat * 1000
|
||||
assert refund_response.status_code == 409
|
||||
assert (
|
||||
refund_response.json()["detail"]
|
||||
== "Balance changed concurrently. Please retry the refund."
|
||||
)
|
||||
|
||||
final_balance_response = await authenticated_client.get("/v1/wallet/")
|
||||
assert final_balance_response.status_code == 200
|
||||
assert final_balance_response.json()["balance"] == (
|
||||
initial_balance + topup_amount_sat * 1000
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_during_active_usage(
|
||||
@@ -452,46 +388,9 @@ async def test_refund_during_active_usage(
|
||||
# Refund should succeed
|
||||
assert refund_response.status_code == 200
|
||||
|
||||
# Further usage should return 200 but with 0 balance
|
||||
# Further usage should fail
|
||||
response = await authenticated_client.get("/v1/wallet/")
|
||||
assert response.status_code == 200
|
||||
assert response.json()["balance"] == 0
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
async def test_wallet_history_returns_apikey_transactions(
|
||||
authenticated_client: AsyncClient,
|
||||
testmint_wallet: Any,
|
||||
integration_session: Any,
|
||||
) -> None:
|
||||
wallet_response = await authenticated_client.get("/v1/wallet/")
|
||||
api_key = wallet_response.json()["api_key"]
|
||||
hashed_key = api_key[3:] if api_key.startswith("sk-") else api_key
|
||||
|
||||
topup_token = await testmint_wallet.mint_tokens(250)
|
||||
topup_response = await authenticated_client.post(
|
||||
"/v1/wallet/topup", params={"cashu_token": topup_token}
|
||||
)
|
||||
assert topup_response.status_code == 200
|
||||
|
||||
refund_response = await authenticated_client.post("/v1/wallet/refund")
|
||||
assert refund_response.status_code == 200
|
||||
|
||||
history_response = await authenticated_client.get("/v1/wallet/history")
|
||||
assert history_response.status_code == 200
|
||||
transactions = history_response.json()["transactions"]
|
||||
assert len(transactions) >= 2
|
||||
assert all("api_key_hashed_key" not in tx for tx in transactions)
|
||||
assert {tx["type"] for tx in transactions} >= {"in", "out"}
|
||||
|
||||
db_result = await integration_session.execute(
|
||||
select(CashuTransaction).where(
|
||||
CashuTransaction.api_key_hashed_key == hashed_key
|
||||
)
|
||||
)
|
||||
db_transactions = db_result.scalars().all()
|
||||
assert len(db_transactions) >= 2
|
||||
assert response.status_code == 401
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@@ -636,3 +535,6 @@ async def test_refund_with_expired_key(
|
||||
# Should still allow manual refund
|
||||
assert response.status_code == 200
|
||||
assert response.json()["recipient"] == "expired@ln.address"
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@ import pytest
|
||||
from httpx import AsyncClient
|
||||
from sqlmodel import select
|
||||
|
||||
from routstr.core.db import ApiKey, CashuTransaction
|
||||
from routstr.core.db import ApiKey
|
||||
|
||||
from .utils import (
|
||||
CashuTokenGenerator,
|
||||
@@ -71,16 +71,6 @@ async def test_topup_with_valid_token( # type: ignore[no-untyped-def]
|
||||
assert db_key.balance == new_balance
|
||||
assert db_key.balance == initial_balance + (topup_amount * 1000)
|
||||
|
||||
tx_result = await integration_session.execute(
|
||||
select(CashuTransaction).where(
|
||||
CashuTransaction.token == token,
|
||||
CashuTransaction.type == "in",
|
||||
)
|
||||
)
|
||||
tx = tx_result.scalar_one()
|
||||
assert tx.api_key_hashed_key == hashed_key
|
||||
assert tx.source == "apikey"
|
||||
|
||||
|
||||
@pytest.mark.integration
|
||||
@pytest.mark.asyncio
|
||||
|
||||
@@ -1,18 +1,14 @@
|
||||
"""Tests for the model prioritization algorithm."""
|
||||
|
||||
import os
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
# Set required env vars before importing
|
||||
os.environ["UPSTREAM_BASE_URL"] = "http://test"
|
||||
os.environ["UPSTREAM_API_KEY"] = "test"
|
||||
|
||||
from routstr.algorithm import ( # noqa: E402
|
||||
calculate_model_cost_score,
|
||||
create_model_mappings,
|
||||
get_provider_penalty,
|
||||
)
|
||||
from routstr.payment.models import Architecture, Model, Pricing # noqa: E402
|
||||
@@ -49,21 +45,11 @@ def create_test_model(
|
||||
)
|
||||
|
||||
|
||||
def create_test_provider(
|
||||
name: str,
|
||||
base_url: str = "http://test.com",
|
||||
*,
|
||||
db_id: int | None = None,
|
||||
models: list[Model] | None = None,
|
||||
upstream_name: str | None = None,
|
||||
) -> Mock:
|
||||
def create_test_provider(name: str, base_url: str = "http://test.com") -> Mock:
|
||||
"""Helper to create a test provider mock."""
|
||||
provider = Mock()
|
||||
provider.provider_type = name
|
||||
provider.base_url = base_url
|
||||
provider.db_id = db_id
|
||||
provider.upstream_name = upstream_name or name
|
||||
provider.get_cached_models.return_value = models or []
|
||||
return provider
|
||||
|
||||
|
||||
@@ -113,80 +99,3 @@ def test_get_provider_penalty_openrouter() -> None:
|
||||
provider = create_test_provider("openrouter", "https://openrouter.ai/api/v1")
|
||||
penalty = get_provider_penalty(provider)
|
||||
assert penalty == 1.001
|
||||
|
||||
|
||||
def test_create_model_mappings_includes_db_override_for_missing_cached_model(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Model overrides should still map when provider discovery misses the model."""
|
||||
provider = create_test_provider(
|
||||
"azure",
|
||||
"https://example.openai.azure.com/openai/v1",
|
||||
db_id=7,
|
||||
models=[],
|
||||
)
|
||||
override_model = create_test_model("azure/gpt-4o")
|
||||
override_model.canonical_slug = "azure-deployment"
|
||||
|
||||
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
|
||||
return override_model
|
||||
|
||||
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
|
||||
|
||||
override_row = SimpleNamespace(id="azure/gpt-4o", upstream_provider_id=7, enabled=True)
|
||||
|
||||
model_instances, provider_map, unique_models = create_model_mappings(
|
||||
upstreams=[provider],
|
||||
overrides_by_id={"azure/gpt-4o": (override_row, 1.01)},
|
||||
disabled_model_ids=set(),
|
||||
)
|
||||
|
||||
assert "azure/gpt-4o" in model_instances
|
||||
assert provider_map["azure/gpt-4o"] == [provider]
|
||||
assert "gpt-4o" in unique_models
|
||||
|
||||
|
||||
def test_create_model_mappings_dedupes_with_provider_identity_not_provider_type(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Different provider instances of same type should both survive dedupe."""
|
||||
provider_a_model = create_test_model(
|
||||
"azure/gpt-4o", prompt_price=0.01, completion_price=0.01
|
||||
)
|
||||
provider_a = create_test_provider(
|
||||
"azure",
|
||||
"https://a.openai.azure.com/openai/v1",
|
||||
db_id=1,
|
||||
models=[provider_a_model],
|
||||
upstream_name="azure-a",
|
||||
)
|
||||
provider_b = create_test_provider(
|
||||
"azure",
|
||||
"https://b.openai.azure.com/openai/v1",
|
||||
db_id=2,
|
||||
models=[],
|
||||
upstream_name="azure-b",
|
||||
)
|
||||
|
||||
override_model = create_test_model(
|
||||
"azure/gpt-4o", prompt_price=0.001, completion_price=0.001
|
||||
)
|
||||
override_model.canonical_slug = "azure-b-deployment"
|
||||
|
||||
def fake_row_to_model(*args, **kwargs) -> Model: # type: ignore[no-untyped-def]
|
||||
return override_model
|
||||
|
||||
monkeypatch.setattr("routstr.payment.models._row_to_model", fake_row_to_model)
|
||||
|
||||
override_row = SimpleNamespace(id="azure/gpt-4o", upstream_provider_id=2, enabled=True)
|
||||
|
||||
_, provider_map, _ = create_model_mappings(
|
||||
upstreams=[provider_a, provider_b],
|
||||
overrides_by_id={"azure/gpt-4o": (override_row, 1.01)},
|
||||
disabled_model_ids=set(),
|
||||
)
|
||||
|
||||
providers_for_alias = provider_map["azure/gpt-4o"]
|
||||
assert provider_a in providers_for_alias
|
||||
assert provider_b in providers_for_alias
|
||||
assert len(providers_for_alias) == 2
|
||||
|
||||
@@ -1,357 +0,0 @@
|
||||
import json
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi.responses import JSONResponse
|
||||
|
||||
from routstr.balance import refund_wallet_endpoint
|
||||
from routstr.core.db import ApiKey, CashuTransaction
|
||||
from routstr.wallet import credit_balance
|
||||
|
||||
|
||||
def _make_cashu_tx(
|
||||
token: str,
|
||||
amount: int,
|
||||
unit: str,
|
||||
type: str = "out",
|
||||
request_id: str | None = "req-abc",
|
||||
swept: bool = False,
|
||||
collected: bool = False,
|
||||
) -> CashuTransaction:
|
||||
tx = CashuTransaction(token=token, amount=amount, unit=unit, type=type, request_id=request_id)
|
||||
tx.swept = swept
|
||||
tx.collected = collected
|
||||
return tx
|
||||
|
||||
|
||||
def _exec_result(tx: CashuTransaction | None) -> MagicMock:
|
||||
result = MagicMock()
|
||||
result.first.return_value = tx
|
||||
return result
|
||||
|
||||
|
||||
def _update_result(rowcount: int) -> MagicMock:
|
||||
result = MagicMock()
|
||||
result.rowcount = rowcount
|
||||
return result
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_x_cashu_returns_token() -> None:
|
||||
x_cashu_token = "cashuAtest_token_value"
|
||||
in_tx = _make_cashu_tx(token=x_cashu_token, amount=0, unit="msat", type="in", request_id="req-abc")
|
||||
out_tx = _make_cashu_tx(token="cashuArefund_token", amount=1000, unit="msat", type="out", request_id="req-abc")
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(side_effect=[_exec_result(in_tx), _exec_result(out_tx)])
|
||||
session.add = MagicMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
result = await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-somekey",
|
||||
x_cashu=x_cashu_token,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert isinstance(result, JSONResponse)
|
||||
body = json.loads(result.body)
|
||||
assert body["token"] == "cashuArefund_token"
|
||||
assert body["msats"] == "1000"
|
||||
assert result.headers["X-Cashu"] == "cashuArefund_token"
|
||||
assert out_tx.collected is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_x_cashu_sat_unit() -> None:
|
||||
x_cashu_token = "cashuAsat_token"
|
||||
in_tx = _make_cashu_tx(token=x_cashu_token, amount=0, unit="sat", type="in", request_id="req-sat")
|
||||
out_tx = _make_cashu_tx(token="cashuArefund_sat", amount=500, unit="sat", type="out", request_id="req-sat")
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(side_effect=[_exec_result(in_tx), _exec_result(out_tx)])
|
||||
session.add = MagicMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
result = await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-somekey",
|
||||
x_cashu=x_cashu_token,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert isinstance(result, JSONResponse)
|
||||
body = json.loads(result.body)
|
||||
assert body["token"] == "cashuArefund_sat"
|
||||
assert body["sats"] == "500"
|
||||
assert "msats" not in body
|
||||
assert result.headers["X-Cashu"] == "cashuArefund_sat"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_x_cashu_not_found_raises_404() -> None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(return_value=_exec_result(None))
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-somekey",
|
||||
x_cashu="cashuAmissing_token",
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 404
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refund_x_cashu_swept_raises_410() -> None:
|
||||
from fastapi import HTTPException
|
||||
|
||||
in_tx = _make_cashu_tx(token="cashuAswept_token", amount=0, unit="msat", type="in", request_id="req-swept")
|
||||
out_tx = _make_cashu_tx(token="cashuAswept", amount=100, unit="msat", type="out", request_id="req-swept", swept=True)
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(side_effect=[_exec_result(in_tx), _exec_result(out_tx)])
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-somekey",
|
||||
x_cashu="cashuAswept_token",
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 410
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# source field defaults
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_cashu_transaction_source_defaults_to_x_cashu() -> None:
|
||||
tx = CashuTransaction(token="cashuAtest", amount=100, unit="msat")
|
||||
assert tx.source == "x-cashu"
|
||||
|
||||
|
||||
def test_cashu_transaction_source_can_be_apikey() -> None:
|
||||
tx = CashuTransaction(token="cashuAtest", amount=100, unit="msat", source="apikey")
|
||||
assert tx.source == "apikey"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# apikey-based refund: token logging and CashuTransaction storage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_api_key(
|
||||
balance: int = 5000,
|
||||
refund_currency: str | None = "sat",
|
||||
refund_mint_url: str | None = "https://mint.example.com",
|
||||
refund_address: str | None = None,
|
||||
parent_key_hash: str | None = None,
|
||||
) -> ApiKey:
|
||||
key = ApiKey(hashed_key="testhash")
|
||||
key.balance = balance
|
||||
key.reserved_balance = 0
|
||||
key.refund_currency = refund_currency
|
||||
key.refund_mint_url = refund_mint_url
|
||||
key.refund_address = refund_address
|
||||
key.parent_key_hash = parent_key_hash
|
||||
key.total_spent = 0
|
||||
key.total_requests = 0
|
||||
return key
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apikey_refund_stores_cashu_transaction_with_apikey_source() -> None:
|
||||
key = _make_api_key(balance=5000, refund_currency="sat")
|
||||
refund_token = "cashuArefund_apikey_token"
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(return_value=_update_result(1))
|
||||
session.add = MagicMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("routstr.balance.validate_bearer_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.get_billing_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.send_token", AsyncMock(return_value=refund_token)),
|
||||
patch("routstr.balance.store_cashu_transaction", AsyncMock()) as mock_store,
|
||||
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
|
||||
patch("routstr.balance._refund_cache_set", AsyncMock()),
|
||||
):
|
||||
result = await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-testhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert isinstance(result, dict)
|
||||
assert result["token"] == refund_token
|
||||
|
||||
mock_store.assert_awaited_once()
|
||||
call_kwargs = mock_store.call_args.kwargs
|
||||
assert call_kwargs["source"] == "apikey"
|
||||
assert call_kwargs["token"] == refund_token
|
||||
assert call_kwargs["typ"] == "out"
|
||||
assert call_kwargs["api_key_hashed_key"] == key.hashed_key
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apikey_refund_logs_token() -> None:
|
||||
key = _make_api_key(balance=5000, refund_currency="sat")
|
||||
refund_token = "cashuAlogged_token"
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(return_value=_update_result(1))
|
||||
session.add = MagicMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("routstr.balance.validate_bearer_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.get_billing_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.send_token", AsyncMock(return_value=refund_token)),
|
||||
patch("routstr.balance.store_cashu_transaction", AsyncMock()),
|
||||
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
|
||||
patch("routstr.balance._refund_cache_set", AsyncMock()),
|
||||
patch("routstr.balance.logger") as mock_logger,
|
||||
):
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-testhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
calls = [str(c) for c in mock_logger.info.call_args_list]
|
||||
assert any("cashu token issued" in c for c in calls)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apikey_refund_log_includes_path() -> None:
|
||||
key = _make_api_key(balance=5000, refund_currency="sat")
|
||||
refund_token = "cashuApath_token"
|
||||
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(return_value=_update_result(1))
|
||||
session.add = MagicMock()
|
||||
session.commit = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("routstr.balance.validate_bearer_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.get_billing_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.send_token", AsyncMock(return_value=refund_token)),
|
||||
patch("routstr.balance.store_cashu_transaction", AsyncMock()),
|
||||
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
|
||||
patch("routstr.balance._refund_cache_set", AsyncMock()),
|
||||
patch("routstr.balance.logger") as mock_logger,
|
||||
):
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-testhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
# Find the "cashu token issued" call and verify extra contains the path
|
||||
token_issued_calls = [
|
||||
c for c in mock_logger.info.call_args_list
|
||||
if c.args and "cashu token issued" in c.args[0]
|
||||
]
|
||||
assert len(token_issued_calls) == 1
|
||||
extra = token_issued_calls[0].kwargs.get("extra", {})
|
||||
assert extra.get("path") == "/v1/wallet/refund"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apikey_refund_rejects_on_concurrent_balance_change() -> None:
|
||||
"""When the debit CAS fails (rowcount=0), no token is minted and 409 is returned."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
key = _make_api_key(balance=5000, refund_currency="sat")
|
||||
|
||||
session = MagicMock()
|
||||
# Debit returns rowcount=0 → balance changed concurrently
|
||||
session.exec = AsyncMock(return_value=_update_result(0))
|
||||
session.commit = AsyncMock()
|
||||
|
||||
mock_send_token = AsyncMock(return_value="cashuAshould_not_be_minted")
|
||||
|
||||
with (
|
||||
patch("routstr.balance.validate_bearer_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.get_billing_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.send_token", mock_send_token),
|
||||
patch("routstr.balance.store_cashu_transaction", AsyncMock()),
|
||||
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
|
||||
patch("routstr.balance._refund_cache_set", AsyncMock()),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-testhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 409
|
||||
# Crucially: send_token must NOT have been called
|
||||
mock_send_token.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credit_balance_stores_apikey_transaction_history() -> None:
|
||||
key = _make_api_key(balance=1000)
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(return_value=_update_result(1))
|
||||
session.commit = AsyncMock()
|
||||
session.refresh = AsyncMock()
|
||||
|
||||
with (
|
||||
patch(
|
||||
"routstr.wallet.recieve_token",
|
||||
AsyncMock(return_value=(100, "sat", "https://mint.example")),
|
||||
),
|
||||
patch("routstr.wallet.store_cashu_transaction", AsyncMock()) as mock_store,
|
||||
):
|
||||
amount = await credit_balance("cashuAtopup_token", key, session)
|
||||
|
||||
assert amount == 100_000
|
||||
mock_store.assert_awaited_once()
|
||||
call_kwargs = mock_store.call_args.kwargs
|
||||
assert call_kwargs["typ"] == "in"
|
||||
assert call_kwargs["source"] == "apikey"
|
||||
assert call_kwargs["api_key_hashed_key"] == key.hashed_key
|
||||
assert call_kwargs["amount"] == 100
|
||||
assert call_kwargs["unit"] == "sat"
|
||||
assert call_kwargs["token"] == "cashuAtopup_token"
|
||||
assert call_kwargs["mint_url"] == "https://mint.example"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_apikey_refund_restores_balance_on_mint_failure() -> None:
|
||||
"""When debit succeeds but minting fails, balance must be restored."""
|
||||
from fastapi import HTTPException
|
||||
|
||||
key = _make_api_key(balance=5000, refund_currency="sat")
|
||||
|
||||
# First exec call = debit (succeeds), second = restore
|
||||
session = MagicMock()
|
||||
session.exec = AsyncMock(side_effect=[_update_result(1), _update_result(1)])
|
||||
session.commit = AsyncMock()
|
||||
|
||||
with (
|
||||
patch("routstr.balance.validate_bearer_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.get_billing_key", AsyncMock(return_value=key)),
|
||||
patch("routstr.balance.send_token", AsyncMock(side_effect=Exception("mint down"))),
|
||||
patch("routstr.balance.store_cashu_transaction", AsyncMock()),
|
||||
patch("routstr.balance._refund_cache_get", AsyncMock(return_value=None)),
|
||||
patch("routstr.balance._refund_cache_set", AsyncMock()),
|
||||
patch("routstr.balance.logger"),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await refund_wallet_endpoint(
|
||||
authorization="Bearer sk-testhash",
|
||||
x_cashu=None,
|
||||
session=session,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 503
|
||||
# Verify two exec calls: debit + restore
|
||||
assert session.exec.await_count == 2
|
||||
@@ -1,345 +0,0 @@
|
||||
"""Tests for cache token handling in cost calculation.
|
||||
|
||||
Covers OpenAI vs Anthropic caching formats, edge cases, and billing accuracy.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("UPSTREAM_BASE_URL", "http://test")
|
||||
os.environ.setdefault("UPSTREAM_API_KEY", "test")
|
||||
os.environ.setdefault("LIGHTNING_ADDRESS", "test@stm.to")
|
||||
|
||||
from routstr.core.settings import settings
|
||||
from routstr.payment.cost_calculation import CostData, MaxCostData, calculate_cost
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_session() -> AsyncMock:
|
||||
"""Mock AsyncSession for cost calculation tests."""
|
||||
return AsyncMock()
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def mock_fixed_pricing(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Mock settings and price to use fixed pricing."""
|
||||
monkeypatch.setattr(settings, "fixed_pricing", True)
|
||||
monkeypatch.setattr(settings, "fixed_per_1k_input_tokens", 0.001)
|
||||
monkeypatch.setattr(settings, "fixed_per_1k_output_tokens", 0.001)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def patch_sats_usd_price() -> None: # type: ignore[misc]
|
||||
"""Patch sats_usd_price to avoid initialization issues."""
|
||||
with patch("routstr.payment.cost_calculation.sats_usd_price", return_value=5.0e-5):
|
||||
yield
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 1: OpenAI Cache Format
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_cache_subtraction(mock_session: AsyncMock) -> None:
|
||||
"""OpenAI includes cached_tokens in prompt_tokens, subtract them."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 2000, # ← Includes 1000 cached
|
||||
"completion_tokens": 100,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 1000 # ← Extracted separately
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 1000 # 2000 - 1000
|
||||
assert result.cache_read_input_tokens == 1000
|
||||
assert result.output_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 2: Anthropic Cache Format
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_cache_additive(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Anthropic cache tokens are separate (additive) from input_tokens."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
"usage": {
|
||||
"input_tokens": 500, # ← Regular input only
|
||||
"output_tokens": 100,
|
||||
"cache_creation_input_tokens": 1500, # ← Additive, not included above
|
||||
"cache_read_input_tokens": 0,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 500
|
||||
assert result.cache_creation_input_tokens == 1500
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.output_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 3: Invalid Cache (Edge Case)
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_read_exceeds_prompt_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Handle buggy upstream reporting cached > prompt_tokens."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 150 # ← Invalid! Greater than prompt
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
# Should not go negative
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 0 # max(0, 100 - 150)
|
||||
assert result.cache_read_input_tokens == 150
|
||||
assert result.output_tokens == 50
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 4: Malformed Token Values
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_malformed_cache_tokens_coerce_to_zero(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Handle non-numeric cache token values."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"cache_read_input_tokens": "-50", # ← String, negative
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": "invalid" # ← Non-numeric string
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
# Both should coerce to 0
|
||||
assert isinstance(result, CostData)
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.input_tokens == 100 # No subtraction if cache_read = 0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 5: Anthropic Cache Not Subtracted
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_anthropic_cache_not_subtracted(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Anthropic cache fields should NOT be subtracted from input_tokens."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
"usage": {
|
||||
"input_tokens": 500,
|
||||
"completion_tokens": 100,
|
||||
"cache_read_input_tokens": 200, # ← Additive, don't subtract
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
# Anthropic: input_tokens stays as-is
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 500 # NOT 300
|
||||
assert result.cache_read_input_tokens == 200
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 6: Only Cache Read, No Regular Input
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_cache_read_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Handle response with only cache read tokens."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 0,
|
||||
"completion_tokens": 50,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 1000
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 0 # max(0, 0 - 1000)
|
||||
assert result.cache_read_input_tokens == 1000
|
||||
assert result.output_tokens == 50
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 7: Only Cache Creation
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_only_cache_creation_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Handle response with only cache creation tokens (Anthropic)."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
"usage": {
|
||||
"input_tokens": 500,
|
||||
"output_tokens": 100,
|
||||
"cache_creation_input_tokens": 2000,
|
||||
"cache_read_input_tokens": 0,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 500
|
||||
assert result.cache_creation_input_tokens == 2000
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.output_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 8: Both Cache Read and Creation
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_both_cache_read_and_creation(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Handle response with both cache read and creation."""
|
||||
response = {
|
||||
"model": "claude-3-5-sonnet",
|
||||
"usage": {
|
||||
"input_tokens": 300,
|
||||
"output_tokens": 100,
|
||||
"cache_creation_input_tokens": 2000,
|
||||
"cache_read_input_tokens": 500,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 300
|
||||
assert result.cache_creation_input_tokens == 2000
|
||||
assert result.cache_read_input_tokens == 500
|
||||
assert result.output_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 9: Token Field Fallback
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_token_field_fallback_order(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Verify fallback order for token extraction."""
|
||||
# When prompt_tokens is not present, fall back to input_tokens
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"input_tokens": 250,
|
||||
"completion_tokens": 50,
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 250
|
||||
assert result.output_tokens == 50
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 10: Float Token Values
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_float_token_values_coerced_to_int(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Handle float token values by converting to int."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100.7, # Float
|
||||
"completion_tokens": 50.3, # Float
|
||||
"cache_read_input_tokens": 25.9, # Float
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.input_tokens == 100 # Floored
|
||||
assert result.output_tokens == 50 # Floored
|
||||
assert result.cache_read_input_tokens == 25 # Floored
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 11: Boolean Cache Tokens
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_boolean_cache_tokens_coerced_to_zero(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Handle boolean cache token values by coercing to zero."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"cache_read_input_tokens": True, # Boolean
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.cache_read_input_tokens == 0 # Boolean coerced to 0
|
||||
assert result.input_tokens == 100 # No subtraction
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 12: Zero Cache Tokens
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_cache_tokens(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""Handle explicit zero cache tokens."""
|
||||
response = {
|
||||
"model": "gpt-4",
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 50,
|
||||
"prompt_tokens_details": {
|
||||
"cached_tokens": 0
|
||||
}
|
||||
}
|
||||
}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, CostData)
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.input_tokens == 100
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 13: Missing Usage Block
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_missing_usage_block(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""When usage is missing, return MaxCostData with zero tokens."""
|
||||
response = {"model": "gpt-4", "choices": [{"message": {"content": "test"}}]}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, MaxCostData)
|
||||
assert result.input_tokens == 0
|
||||
assert result.cache_read_input_tokens == 0
|
||||
assert result.output_tokens == 0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 14: Null Usage Block
|
||||
# ============================================================================
|
||||
@pytest.mark.asyncio
|
||||
async def test_null_usage_block(mock_session: AsyncMock, mock_fixed_pricing: None) -> None:
|
||||
"""When usage is null, return MaxCostData with zero tokens."""
|
||||
response = {"model": "gpt-4", "usage": None}
|
||||
result = await calculate_cost(response, max_cost=100000, session=mock_session)
|
||||
|
||||
assert isinstance(result, MaxCostData)
|
||||
assert result.input_tokens == 0
|
||||
assert result.cache_read_input_tokens == 0
|
||||
@@ -0,0 +1,155 @@
|
||||
"""Unit tests for model row payload conversion.
|
||||
|
||||
This module tests that _model_to_row_payload correctly serializes model data
|
||||
for database storage. Pricing is stored as-is without fee application.
|
||||
Fees are now applied per-provider when reading from the database.
|
||||
|
||||
Key behaviors tested:
|
||||
1. Pricing is stored as-is without fee application
|
||||
2. All model fields are correctly serialized to JSON
|
||||
3. Optional fields are handled correctly (None values)
|
||||
4. Pricing structure is preserved
|
||||
5. Original model objects are not mutated
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
# Set required env vars before importing
|
||||
os.environ["UPSTREAM_BASE_URL"] = "http://test"
|
||||
os.environ["UPSTREAM_API_KEY"] = "test"
|
||||
|
||||
from routstr.payment.models import ( # noqa: E402
|
||||
Architecture,
|
||||
Model,
|
||||
Pricing,
|
||||
_model_to_row_payload,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def base_architecture() -> Architecture:
|
||||
"""Provide standard architecture for test models."""
|
||||
return Architecture(
|
||||
modality="text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="gpt",
|
||||
instruct_type="chat",
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def standard_pricing() -> Pricing:
|
||||
"""Provide standard USD pricing with known values for testing."""
|
||||
return Pricing(
|
||||
prompt=0.001,
|
||||
completion=0.002,
|
||||
request=0.01,
|
||||
image=0.05,
|
||||
web_search=0.03,
|
||||
internal_reasoning=0.015,
|
||||
max_prompt_cost=10.0,
|
||||
max_completion_cost=20.0,
|
||||
max_cost=30.0,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def standard_model(base_architecture: Architecture, standard_pricing: Pricing) -> Model:
|
||||
"""Create a standard test model with known pricing."""
|
||||
return Model(
|
||||
id="test-model-standard",
|
||||
name="Test Model Standard",
|
||||
created=1234567890,
|
||||
description="A standard test model",
|
||||
context_length=8192,
|
||||
architecture=base_architecture,
|
||||
pricing=standard_pricing,
|
||||
)
|
||||
|
||||
|
||||
def test_pricing_stored_without_fees(standard_model: Model) -> None:
|
||||
"""Verify pricing is stored as-is without any fee application."""
|
||||
payload = _model_to_row_payload(standard_model)
|
||||
pricing_str = payload["pricing"]
|
||||
assert isinstance(pricing_str, str)
|
||||
pricing = json.loads(pricing_str)
|
||||
|
||||
assert pricing["prompt"] == pytest.approx(0.001, rel=1e-9)
|
||||
assert pricing["completion"] == pytest.approx(0.002, rel=1e-9)
|
||||
assert pricing["request"] == pytest.approx(0.01, rel=1e-9)
|
||||
assert pricing["image"] == pytest.approx(0.05, rel=1e-9)
|
||||
assert pricing["web_search"] == pytest.approx(0.03, rel=1e-9)
|
||||
assert pricing["internal_reasoning"] == pytest.approx(0.015, rel=1e-9)
|
||||
assert pricing["max_prompt_cost"] == pytest.approx(10.0, rel=1e-9)
|
||||
assert pricing["max_completion_cost"] == pytest.approx(20.0, rel=1e-9)
|
||||
assert pricing["max_cost"] == pytest.approx(30.0, rel=1e-9)
|
||||
|
||||
|
||||
def test_zero_value_pricing_fields(base_architecture: Architecture) -> None:
|
||||
"""Verify that zero-value pricing fields are stored correctly."""
|
||||
zero_pricing = Pricing(
|
||||
prompt=0.0,
|
||||
completion=0.0,
|
||||
request=0.0,
|
||||
image=0.0,
|
||||
web_search=0.0,
|
||||
internal_reasoning=0.0,
|
||||
max_prompt_cost=0.0,
|
||||
max_completion_cost=0.0,
|
||||
max_cost=0.0,
|
||||
)
|
||||
|
||||
model = Model(
|
||||
id="test-model-zero",
|
||||
name="Test Model Zero",
|
||||
created=1234567890,
|
||||
description="A model with zero pricing",
|
||||
context_length=8192,
|
||||
architecture=base_architecture,
|
||||
pricing=zero_pricing,
|
||||
)
|
||||
|
||||
payload = _model_to_row_payload(model)
|
||||
pricing_str = payload["pricing"]
|
||||
assert isinstance(pricing_str, str)
|
||||
pricing = json.loads(pricing_str)
|
||||
|
||||
assert pricing["prompt"] == pytest.approx(0.0, rel=1e-9)
|
||||
assert pricing["completion"] == pytest.approx(0.0, rel=1e-9)
|
||||
assert pricing["request"] == pytest.approx(0.0, rel=1e-9)
|
||||
|
||||
|
||||
def test_payload_structure_unchanged(standard_model: Model) -> None:
|
||||
"""Verify that payload structure matches expectations."""
|
||||
payload = _model_to_row_payload(standard_model)
|
||||
|
||||
assert "id" in payload
|
||||
assert "name" in payload
|
||||
assert "created" in payload
|
||||
assert "description" in payload
|
||||
assert "context_length" in payload
|
||||
assert "architecture" in payload
|
||||
assert "pricing" in payload
|
||||
assert "sats_pricing" in payload
|
||||
assert "per_request_limits" in payload
|
||||
assert "top_provider" in payload
|
||||
assert "enabled" in payload
|
||||
assert "upstream_provider_id" in payload
|
||||
|
||||
assert isinstance(payload["architecture"], str)
|
||||
assert isinstance(payload["pricing"], str)
|
||||
|
||||
|
||||
def test_original_model_not_mutated(standard_model: Model) -> None:
|
||||
"""Verify that the original model object is not mutated."""
|
||||
original_prompt = standard_model.pricing.prompt
|
||||
original_completion = standard_model.pricing.completion
|
||||
|
||||
_model_to_row_payload(standard_model)
|
||||
|
||||
assert standard_model.pricing.prompt == original_prompt
|
||||
assert standard_model.pricing.completion == original_completion
|
||||
@@ -1,93 +0,0 @@
|
||||
"""Tests for `routstr.upstream.litellm_routing.detect_litellm_prefix`."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from routstr.upstream.litellm_routing import detect_litellm_prefix
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"base_url,expected",
|
||||
[
|
||||
# The bug case: custom row pointing at Fireworks must NOT route to openai/.
|
||||
("https://api.fireworks.ai/inference/v1", "fireworks_ai/"),
|
||||
# Other OpenAI-compatible providers commonly plugged into the custom slot.
|
||||
("https://api.groq.com/openai/v1", "groq/"),
|
||||
("https://api.x.ai/v1", "xai/"),
|
||||
("https://api.deepseek.com/v1", "deepseek/"),
|
||||
("https://api.together.xyz/v1", "together_ai/"),
|
||||
("https://api.perplexity.ai", "perplexity/"),
|
||||
("https://openrouter.ai/api/v1", "openrouter/"),
|
||||
("https://api.mistral.ai/v1", "mistral/"),
|
||||
("https://codestral.mistral.ai/v1", "codestral/"),
|
||||
("https://api.cohere.com/v1", "cohere_chat/"),
|
||||
("https://api.cohere.ai/v1", "cohere_chat/"),
|
||||
("https://api.deepinfra.com/v1/openai", "deepinfra/"),
|
||||
("https://api.cerebras.ai/v1", "cerebras/"),
|
||||
("https://api.sambanova.ai/v1", "sambanova/"),
|
||||
("https://api.moonshot.cn/v1", "moonshot/"),
|
||||
("https://api.moonshot.ai/v1", "moonshot/"),
|
||||
("https://api.studio.nebius.com/v1", "nebius/"),
|
||||
("https://api.novita.ai/v3/openai", "novita/"),
|
||||
("https://api.lambda.ai/v1", "lambda_ai/"),
|
||||
("https://api.aimlapi.com/v1", "aiml/"),
|
||||
("https://api.featherless.ai/v1", "featherless_ai/"),
|
||||
("https://integrate.api.nvidia.com/v1", "nvidia_nim/"),
|
||||
("https://inference.baseten.co/v1", "baseten/"),
|
||||
("https://ai-gateway.vercel.sh/v1", "vercel_ai_gateway/"),
|
||||
("https://api.inference.wandb.ai/v1", "wandb/"),
|
||||
("https://api.poe.com/v1", "poe/"),
|
||||
("https://llm.chutes.ai/v1/", "chutes/"),
|
||||
("https://api.v0.dev/v1", "v0/"),
|
||||
("https://api.hyperbolic.xyz/v1", "hyperbolic/"),
|
||||
("https://api.synthetic.new/openai/v1", "synthetic/"),
|
||||
("https://api.stima.tech/v1", "apertis/"),
|
||||
("https://nano-gpt.com/api/v1", "nano-gpt/"),
|
||||
("https://api.friendli.ai/serverless/v1", "friendliai/"),
|
||||
("https://api.galadriel.com/v1", "galadriel/"),
|
||||
("https://api.llama.com/compat/v1", "meta_llama/"),
|
||||
("https://api.minimax.io/v1", "minimax/"),
|
||||
("https://api.minimaxi.com/v1", "minimax/"),
|
||||
("https://platform.publicai.co/v1", "publicai/"),
|
||||
("https://inference.api.nscale.com/v1", "nscale/"),
|
||||
("https://dashscope-intl.aliyuncs.com/compatible-mode/v1", "dashscope/"),
|
||||
("https://api.endpoints.anyscale.com/v1", "anyscale/"),
|
||||
("https://api.ai21.com/studio/v1", "ai21_chat/"),
|
||||
("https://ark.cn-beijing.volces.com/api/v3", "volcengine/"),
|
||||
("https://api.voyageai.com/v1", "voyage/"),
|
||||
("https://api.jina.ai/v1", "jina_ai/"),
|
||||
("https://api.snowflakecomputing.com", "snowflake/"),
|
||||
("https://my-workspace.databricks.com/serving-endpoints", "databricks/"),
|
||||
("https://huggingface.co/api/inference", "huggingface/"),
|
||||
# First-class providers — URL detection still produces the right prefix
|
||||
# so subclasses without an explicit override stay correct.
|
||||
("https://api.openai.com/v1", "openai/"),
|
||||
("https://api.anthropic.com/v1", "anthropic/"),
|
||||
("https://generativelanguage.googleapis.com/v1beta/openai", "gemini/"),
|
||||
("https://us-central1-aiplatform.googleapis.com/v1", "vertex_ai/"),
|
||||
# Azure ordering: must beat api.openai.com.
|
||||
("https://my-resource.openai.azure.com/openai/deployments/foo", "azure/"),
|
||||
# Ollama hints.
|
||||
("http://localhost:11434/v1", "ollama_chat/"),
|
||||
("http://127.0.0.1:11434/v1", "ollama_chat/"),
|
||||
("http://my-ollama-host:11434/v1", "ollama_chat/"),
|
||||
# Casing and trailing slash normalisation.
|
||||
("HTTPS://API.FIREWORKS.AI/INFERENCE/V1/", "fireworks_ai/"),
|
||||
# Unknown host falls back to openai/ (still OpenAI-compatible by convention).
|
||||
("https://example.com/v1", "openai/"),
|
||||
("", "openai/"),
|
||||
],
|
||||
)
|
||||
def test_detect_litellm_prefix(base_url: str, expected: str) -> None:
|
||||
assert detect_litellm_prefix(base_url) == expected
|
||||
|
||||
|
||||
def test_detect_litellm_prefix_none_uses_default() -> None:
|
||||
assert detect_litellm_prefix(None) == "openai/"
|
||||
|
||||
|
||||
def test_detect_litellm_prefix_custom_default() -> None:
|
||||
assert detect_litellm_prefix("https://example.com", default="anthropic/") == (
|
||||
"anthropic/"
|
||||
)
|
||||
@@ -1,333 +0,0 @@
|
||||
"""Tests for cost accumulation in streaming message dispatch.
|
||||
|
||||
Verifies that costs are correctly summed across multiple streaming events,
|
||||
not taking only the maximum.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("UPSTREAM_BASE_URL", "http://test")
|
||||
os.environ.setdefault("UPSTREAM_API_KEY", "test")
|
||||
|
||||
from routstr.upstream.messages_dispatch import annotate_event
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 1: Cost Accumulation via AnnotatedEvent
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_annotate_event_extracts_costs() -> None:
|
||||
"""Each event should report its own costs."""
|
||||
event = {
|
||||
"type": "content_block_start",
|
||||
"usage": {"total_cost": 0.010, "input_cost": 0.008, "output_cost": 0.002}
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
assert result.total_cost == 0.010
|
||||
assert result.input_cost == 0.008
|
||||
assert result.output_cost == 0.002
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 2: Multiple Events with Incremental Costs
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_multiple_events_sum_costs() -> None:
|
||||
"""When processing multiple events, costs should accumulate (not max)."""
|
||||
# Event 1: initial costs
|
||||
event1 = {
|
||||
"type": "message_start",
|
||||
"usage": {"input_tokens": 100, "total_cost": 0.010}
|
||||
}
|
||||
result1 = annotate_event(event1, None)
|
||||
|
||||
# Event 2: additional costs
|
||||
event2 = {
|
||||
"type": "content_block_start",
|
||||
"usage": {"output_tokens": 50, "total_cost": 0.005}
|
||||
}
|
||||
result2 = annotate_event(event2, None)
|
||||
|
||||
# Event 3: more costs
|
||||
event3 = {
|
||||
"type": "message_delta",
|
||||
"usage": {"output_tokens": 25, "total_cost": 0.008}
|
||||
}
|
||||
result3 = annotate_event(event3, None)
|
||||
|
||||
# Each event should report its own cost (before accumulation)
|
||||
assert result1.total_cost == 0.010
|
||||
assert result2.total_cost == 0.005
|
||||
assert result3.total_cost == 0.008
|
||||
|
||||
# When summed for billing: 0.010 + 0.005 + 0.008 = 0.023
|
||||
# This verifies the cost accumulation fix (using += instead of max())
|
||||
total_from_events = result1.total_cost + result2.total_cost + result3.total_cost
|
||||
assert total_from_events == 0.023
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 3: Token Accumulation Consistency
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_tokens_and_costs_extracted_independently() -> None:
|
||||
"""Tokens and costs should be extracted independently per event."""
|
||||
event = {
|
||||
"type": "content_block_delta",
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 50,
|
||||
"cache_read_input_tokens": 30,
|
||||
"total_cost": 0.007,
|
||||
"input_cost": 0.005,
|
||||
"output_cost": 0.002
|
||||
}
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
# All values should be extracted
|
||||
assert result.input_tokens == 100
|
||||
assert result.output_tokens == 50
|
||||
assert result.cache_read_input_tokens == 30
|
||||
assert result.total_cost == 0.007
|
||||
assert result.input_cost == 0.005
|
||||
assert result.output_cost == 0.002
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 4: Cost in Message vs Root
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_cost_extracted_from_message_usage() -> None:
|
||||
"""Costs in message.usage should be extracted correctly."""
|
||||
event = {
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"total_cost": 0.010,
|
||||
"input_cost": 0.008,
|
||||
"output_cost": 0.002
|
||||
}
|
||||
}
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
assert result.input_tokens == 100
|
||||
assert result.total_cost == 0.010
|
||||
assert result.input_cost == 0.008
|
||||
assert result.output_cost == 0.002
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 5: Cost at Event Root Level
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_cost_extracted_from_event_root() -> None:
|
||||
"""Costs at event root should be extracted (OpenRouter style)."""
|
||||
event = {
|
||||
"type": "content_block_delta",
|
||||
"usage": {"output_tokens": 25},
|
||||
"total_cost": 0.005, # ← At root level
|
||||
"cost": 0.005
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
assert result.output_tokens == 25
|
||||
assert result.total_cost == 0.005
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 6: Cost Details in Event Root
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_cost_details_extracted_from_event_root() -> None:
|
||||
"""cost_details at event root should be extracted correctly."""
|
||||
event = {
|
||||
"type": "message_delta",
|
||||
"cost_details": {
|
||||
"total_cost": 0.015,
|
||||
"input_cost": 0.010,
|
||||
"output_cost": 0.005
|
||||
}
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
assert result.total_cost == 0.015
|
||||
assert result.input_cost == 0.010
|
||||
assert result.output_cost == 0.005
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 7: No Duplicated Dict Lookups
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_annotate_event_no_duplicate_lookups() -> None:
|
||||
"""Verify that dict lookups are not duplicated (fix for copy-paste error)."""
|
||||
# This is tested implicitly through proper extraction
|
||||
event = {
|
||||
"type": "message_delta",
|
||||
"cost_details": {
|
||||
"total_cost": 0.020,
|
||||
"input_cost": 0.015,
|
||||
"output_cost": 0.005
|
||||
}
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
# Should extract each field exactly once, correctly
|
||||
assert result.total_cost == 0.020
|
||||
assert result.input_cost == 0.015
|
||||
assert result.output_cost == 0.005
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 8: Model Name Extraction
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_model_extracted_from_event() -> None:
|
||||
"""Model name should be extracted from event."""
|
||||
event = {
|
||||
"type": "message_start",
|
||||
"message": {"model": "claude-3-5-sonnet"}
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
assert result.model == "claude-3-5-sonnet"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 9: Model Name Override
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_model_name_override() -> None:
|
||||
"""Requested model should override actual model in event."""
|
||||
event = {
|
||||
"type": "message_start",
|
||||
"message": {"model": "actual-model"}
|
||||
}
|
||||
# Override with requested_model
|
||||
result = annotate_event(event, requested_model="alias-model")
|
||||
|
||||
# Event should be modified to use requested_model
|
||||
assert event["message"]["model"] == "alias-model" # type: ignore[index]
|
||||
assert result.model == "alias-model"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 10: SSE Encoding
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_sse_bytes_encoded() -> None:
|
||||
"""Event should be encoded as SSE bytes."""
|
||||
event = {
|
||||
"type": "content_block_start",
|
||||
"index": 0,
|
||||
"content_block": {"type": "text", "text": ""}
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
assert result.sse_bytes is not None
|
||||
assert isinstance(result.sse_bytes, bytes)
|
||||
assert b"event: content_block_start" in result.sse_bytes or b"data:" in result.sse_bytes
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 11: Missing Cost Fields Default to Zero
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_missing_cost_fields_default_to_zero() -> None:
|
||||
"""When cost fields are missing, should default to 0.0."""
|
||||
event = {
|
||||
"type": "content_block_delta",
|
||||
"usage": {"output_tokens": 25}
|
||||
# No cost fields
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
assert result.total_cost == 0.0
|
||||
assert result.input_cost == 0.0
|
||||
assert result.output_cost == 0.0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 12: Cache Tokens in Events
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_cache_tokens_extracted_from_event() -> None:
|
||||
"""Cache tokens should be extracted from event usage."""
|
||||
event = {
|
||||
"type": "message_delta",
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 50,
|
||||
"cache_read_input_tokens": 200,
|
||||
"cache_creation_input_tokens": 0
|
||||
}
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
assert result.input_tokens == 100
|
||||
assert result.output_tokens == 50
|
||||
assert result.cache_read_input_tokens == 200
|
||||
assert result.cache_creation_input_tokens == 0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 13: Malformed Cost Values
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_malformed_cost_values_coerced() -> None:
|
||||
"""Malformed cost values should be coerced to 0.0."""
|
||||
event = {
|
||||
"type": "message_delta",
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 50,
|
||||
"total_cost": "invalid", # ← Non-numeric
|
||||
}
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
# Invalid cost should default to 0.0
|
||||
assert result.total_cost == 0.0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 14: Negative Cost Values
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_negative_cost_values_clamped() -> None:
|
||||
"""Negative cost values should be clamped to 0.0."""
|
||||
event = {
|
||||
"type": "message_delta",
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 50,
|
||||
"total_cost": -0.05 # ← Negative
|
||||
}
|
||||
}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
# Negative cost should be clamped to 0.0
|
||||
assert result.total_cost == 0.0
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Test 15: Streaming Event Type Preserved
|
||||
# ============================================================================
|
||||
@pytest.mark.unit
|
||||
def test_event_type_preserved_in_sse() -> None:
|
||||
"""Event type should be preserved in SSE encoding."""
|
||||
event_types = ["message_start", "content_block_start", "content_block_delta", "message_delta"]
|
||||
for event_type in event_types:
|
||||
event = {"type": event_type}
|
||||
result = annotate_event(event, None)
|
||||
|
||||
assert result.event["type"] == event_type # type: ignore[index]
|
||||
# SSE should include the event type line if present
|
||||
if event_type:
|
||||
assert f"event: {event_type}".encode() in result.sse_bytes
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,117 +0,0 @@
|
||||
"""Regression tests for the periodic upstream models refresh loop."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
from typing import cast
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
os.environ.setdefault("UPSTREAM_BASE_URL", "http://test")
|
||||
os.environ.setdefault("UPSTREAM_API_KEY", "test")
|
||||
|
||||
from routstr.upstream.base import BaseUpstreamProvider # noqa: E402
|
||||
|
||||
|
||||
class _FakeUpstream:
|
||||
"""Minimal stand-in for BaseUpstreamProvider used by the refresh loop.
|
||||
|
||||
Only ``base_url`` (for error logging) and ``refresh_models_cache`` (the call
|
||||
under test) are exercised; everything else stays unused.
|
||||
"""
|
||||
|
||||
def __init__(self, name: str) -> None:
|
||||
self.base_url = f"http://{name}"
|
||||
self.refresh_models_cache = AsyncMock()
|
||||
|
||||
|
||||
def _make_fake_upstream(name: str) -> BaseUpstreamProvider:
|
||||
# The loop only uses duck-typed attributes — cast keeps the test type-clean
|
||||
# without dragging in BaseUpstreamProvider's full constructor.
|
||||
return cast(BaseUpstreamProvider, _FakeUpstream(name))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_loop_picks_up_providers_added_after_startup(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""If a provider is added after the loop starts (e.g. via reinitialize_upstreams),
|
||||
the next loop iteration must refresh it. Previously the loop captured the upstream
|
||||
list at startup and missed any later additions."""
|
||||
from routstr.core.settings import settings as global_settings
|
||||
from routstr.upstream.helpers import refresh_upstreams_models_periodically
|
||||
|
||||
# Tight interval so the test finishes quickly.
|
||||
monkeypatch.setattr(
|
||||
global_settings, "models_refresh_interval_seconds", 1, raising=False
|
||||
)
|
||||
|
||||
initial_upstream = _make_fake_upstream("initial")
|
||||
live_list: list[BaseUpstreamProvider] = [initial_upstream]
|
||||
|
||||
# Stub out the post-iteration sats-pricing refresh so the loop body has no DB deps.
|
||||
async def _noop_pricing_refresh() -> None: # pragma: no cover - trivial stub
|
||||
return None
|
||||
|
||||
monkeypatch.setattr(
|
||||
"routstr.payment.models._update_sats_pricing_once",
|
||||
_noop_pricing_refresh,
|
||||
)
|
||||
|
||||
task = asyncio.create_task(
|
||||
refresh_upstreams_models_periodically(lambda: live_list)
|
||||
)
|
||||
|
||||
try:
|
||||
# Wait for the first iteration to refresh the initial upstream.
|
||||
for _ in range(40):
|
||||
if initial_upstream.refresh_models_cache.await_count >= 1: # type: ignore[attr-defined]
|
||||
break
|
||||
await asyncio.sleep(0.05)
|
||||
assert initial_upstream.refresh_models_cache.await_count >= 1, ( # type: ignore[attr-defined]
|
||||
"loop did not refresh the initial upstream within the timeout"
|
||||
)
|
||||
|
||||
# Simulate reinitialize_upstreams: replace the live list contents with new
|
||||
# provider instances. The loop must observe the swap on its next tick.
|
||||
new_upstream = _make_fake_upstream("added-after-startup")
|
||||
live_list[:] = [new_upstream]
|
||||
|
||||
for _ in range(60):
|
||||
if new_upstream.refresh_models_cache.await_count >= 1: # type: ignore[attr-defined]
|
||||
break
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
assert new_upstream.refresh_models_cache.await_count >= 1, ( # type: ignore[attr-defined]
|
||||
"loop did not refresh the upstream added after startup — "
|
||||
"regression: list snapshot captured at startup"
|
||||
)
|
||||
finally:
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except asyncio.CancelledError:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_refresh_loop_disabled_when_interval_non_positive(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from routstr.core.settings import settings as global_settings
|
||||
from routstr.upstream.helpers import refresh_upstreams_models_periodically
|
||||
|
||||
monkeypatch.setattr(
|
||||
global_settings, "models_refresh_interval_seconds", 0, raising=False
|
||||
)
|
||||
|
||||
upstream = _make_fake_upstream("never-refreshed")
|
||||
|
||||
# Loop must return immediately without ever touching the upstream.
|
||||
await asyncio.wait_for(
|
||||
refresh_upstreams_models_periodically(lambda: [upstream]),
|
||||
timeout=1.0,
|
||||
)
|
||||
upstream.refresh_models_cache.assert_not_awaited() # type: ignore[attr-defined]
|
||||
@@ -1,267 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from routstr.nostr import analytics
|
||||
|
||||
|
||||
def test_aggregate_top_model_usage_sums_metrics() -> None:
|
||||
model_usage_mix = {
|
||||
"top_models": ["openai/gpt-4o", "anthropic/claude-3.5-sonnet"],
|
||||
"metrics": [
|
||||
{
|
||||
"model_counts": {
|
||||
"openai/gpt-4o": 4,
|
||||
"anthropic/claude-3.5-sonnet": 2,
|
||||
},
|
||||
"model_revenue_msats": {
|
||||
"openai/gpt-4o": 1500,
|
||||
"anthropic/claude-3.5-sonnet": 700,
|
||||
},
|
||||
"model_tokens": {
|
||||
"openai/gpt-4o": 1200,
|
||||
"anthropic/claude-3.5-sonnet": 600,
|
||||
},
|
||||
"others": 1,
|
||||
"others_revenue_msats": 300,
|
||||
"others_tokens": 200,
|
||||
},
|
||||
{
|
||||
"model_counts": {
|
||||
"openai/gpt-4o": 3,
|
||||
"anthropic/claude-3.5-sonnet": 1,
|
||||
},
|
||||
"model_revenue_msats": {
|
||||
"openai/gpt-4o": 1000,
|
||||
"anthropic/claude-3.5-sonnet": 500,
|
||||
},
|
||||
"model_tokens": {
|
||||
"openai/gpt-4o": 800,
|
||||
"anthropic/claude-3.5-sonnet": 300,
|
||||
},
|
||||
"others": 2,
|
||||
"others_revenue_msats": 450,
|
||||
"others_tokens": 350,
|
||||
},
|
||||
],
|
||||
}
|
||||
|
||||
rows, others = analytics._aggregate_top_model_usage(model_usage_mix)
|
||||
assert rows == [
|
||||
{
|
||||
"model": "openai/gpt-4o",
|
||||
"successful_requests": 7,
|
||||
"revenue_msats": 2500.0,
|
||||
"total_tokens": 2000,
|
||||
},
|
||||
{
|
||||
"model": "anthropic/claude-3.5-sonnet",
|
||||
"successful_requests": 3,
|
||||
"revenue_msats": 1200.0,
|
||||
"total_tokens": 900,
|
||||
},
|
||||
]
|
||||
assert others == {
|
||||
"successful_requests": 3,
|
||||
"revenue_msats": 750.0,
|
||||
"total_tokens": 550,
|
||||
}
|
||||
|
||||
|
||||
def test_build_stats_snapshot_payload_schema_and_shape(monkeypatch: Any) -> None:
|
||||
seen_windows: set[tuple[int, int]] = set()
|
||||
|
||||
def fake_usage_dashboard(
|
||||
*, interval: int, hours: int, error_limit: int, model_limit: int
|
||||
) -> dict[str, Any]:
|
||||
seen_windows.add((hours, interval))
|
||||
assert error_limit == 1
|
||||
assert model_limit == 20
|
||||
return {
|
||||
"summary": {
|
||||
"total_requests": hours,
|
||||
"successful_chat_completions": max(1, hours - 1),
|
||||
"failed_requests": 2,
|
||||
"success_rate": 90.0,
|
||||
"unique_models_count": 2,
|
||||
"input_tokens": 2000,
|
||||
"output_tokens": 1000,
|
||||
"total_tokens": 3000,
|
||||
"revenue_msats": 9000.0,
|
||||
"refunds_msats": 1000.0,
|
||||
"net_revenue_msats": 8000.0,
|
||||
"revenue_sats": 9.0,
|
||||
"refunds_sats": 1.0,
|
||||
"net_revenue_sats": 8.0,
|
||||
},
|
||||
"model_usage_mix": {
|
||||
"top_models": ["openai/gpt-4o"],
|
||||
"metrics": [
|
||||
{
|
||||
"timestamp": "2026-03-02 10:00:00",
|
||||
"model_counts": {"openai/gpt-4o": hours},
|
||||
"model_revenue_msats": {"openai/gpt-4o": float(hours * 100)},
|
||||
"model_tokens": {"openai/gpt-4o": hours * 10},
|
||||
"others": 4,
|
||||
"others_revenue_msats": 1800.0,
|
||||
"others_tokens": 400,
|
||||
}
|
||||
],
|
||||
},
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
analytics.log_manager, "get_usage_dashboard", fake_usage_dashboard
|
||||
)
|
||||
monkeypatch.setattr(analytics.settings, "npub", "npub1example")
|
||||
monkeypatch.setattr(analytics.settings, "http_url", "https://node.example.com")
|
||||
monkeypatch.setattr(analytics.settings, "onion_url", "")
|
||||
|
||||
payload = analytics.build_stats_snapshot_payload(
|
||||
"provider123",
|
||||
public_key_hex="ab" * 32,
|
||||
generated_at=1772451600,
|
||||
)
|
||||
|
||||
assert payload["schema"] == analytics.ANALYTICS_SCHEMA
|
||||
assert payload["provider_id"] == "provider123"
|
||||
assert payload["window_hours"] == 24
|
||||
assert payload["interval_minutes"] == 60
|
||||
assert payload["endpoint_urls"] == ["https://node.example.com"]
|
||||
assert seen_windows == {
|
||||
(24, 60),
|
||||
(7 * 24, 6 * 60),
|
||||
(30 * 24, 24 * 60),
|
||||
(90 * 24, 24 * 60),
|
||||
(365 * 24, 7 * 24 * 60),
|
||||
}
|
||||
assert set(payload["windows"].keys()) == {"24h", "7d", "30d", "3m", "1y"}
|
||||
assert payload["windows"]["1y"]["interval_minutes"] == 7 * 24 * 60
|
||||
assert payload["summary"]["total_requests"] == 24
|
||||
assert payload["top_model_usage"] == [
|
||||
{
|
||||
"model": "openai/gpt-4o",
|
||||
"successful_requests": 24,
|
||||
"revenue_msats": 2400.0,
|
||||
"total_tokens": 240,
|
||||
}
|
||||
]
|
||||
assert payload["others_usage"] == {
|
||||
"successful_requests": 4,
|
||||
"revenue_msats": 1800.0,
|
||||
"total_tokens": 400,
|
||||
}
|
||||
|
||||
|
||||
def test_create_stats_snapshot_event_tags() -> None:
|
||||
private_key_hex = "11" * 32
|
||||
event = analytics.create_stats_snapshot_event(
|
||||
private_key_hex,
|
||||
"provider123",
|
||||
payload_json='{"schema":"routstr.analytics.snapshot.v1"}',
|
||||
d_tag="provider123:stats",
|
||||
)
|
||||
|
||||
tags = event["tags"]
|
||||
assert ["d", "provider123:stats"] in tags
|
||||
assert ["provider", "provider123"] in tags
|
||||
assert ["schema", analytics.ANALYTICS_SCHEMA] in tags
|
||||
assert all(tag[0] != "period" for tag in tags)
|
||||
|
||||
|
||||
def test_fingerprint_payload_ignores_generated_at() -> None:
|
||||
a = {"schema": analytics.ANALYTICS_SCHEMA, "generated_at": 1000, "summary": {"x": 1}}
|
||||
b = {"schema": analytics.ANALYTICS_SCHEMA, "generated_at": 2000, "summary": {"x": 1}}
|
||||
|
||||
assert analytics._fingerprint_payload(a) == analytics._fingerprint_payload(b)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_usage_analytics_skips_when_disabled(monkeypatch: Any) -> None:
|
||||
delays: list[int] = []
|
||||
|
||||
async def fake_sleep(seconds: int) -> None:
|
||||
delays.append(seconds)
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
def fail_build(*args: Any, **kwargs: Any) -> dict[str, Any]:
|
||||
raise AssertionError("build_stats_snapshot_payload should not be called")
|
||||
|
||||
monkeypatch.setattr(analytics.settings, "enable_analytics_sharing", False)
|
||||
monkeypatch.setattr(analytics, "build_stats_snapshot_payload", fail_build)
|
||||
monkeypatch.setattr(analytics.asyncio, "sleep", fake_sleep)
|
||||
|
||||
await analytics.publish_usage_analytics()
|
||||
|
||||
assert delays == [analytics.DISABLED_POLL_SECONDS]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_usage_analytics_skips_without_nsec(monkeypatch: Any) -> None:
|
||||
delays: list[int] = []
|
||||
|
||||
async def fake_sleep(seconds: int) -> None:
|
||||
delays.append(seconds)
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
def fail_build(*args: Any, **kwargs: Any) -> dict[str, Any]:
|
||||
raise AssertionError("build_stats_snapshot_payload should not be called")
|
||||
|
||||
monkeypatch.setattr(analytics.settings, "enable_analytics_sharing", True)
|
||||
monkeypatch.setattr(analytics.settings, "nsec", "")
|
||||
monkeypatch.setattr(analytics, "build_stats_snapshot_payload", fail_build)
|
||||
monkeypatch.setattr(analytics.asyncio, "sleep", fake_sleep)
|
||||
|
||||
await analytics.publish_usage_analytics()
|
||||
|
||||
assert delays == [analytics.DISABLED_POLL_SECONDS]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_publish_usage_analytics_dedupes_unchanged_payload(monkeypatch: Any) -> None:
|
||||
published_events: list[dict[str, Any]] = []
|
||||
sleep_calls = 0
|
||||
|
||||
async def fake_sleep(seconds: int) -> None:
|
||||
nonlocal sleep_calls
|
||||
sleep_calls += 1
|
||||
if sleep_calls >= 2:
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
def fake_build_payload(
|
||||
provider_id: str,
|
||||
*,
|
||||
public_key_hex: str,
|
||||
generated_at: int,
|
||||
window_hours: int = 24,
|
||||
interval_minutes: int = 60,
|
||||
model_limit: int = 10,
|
||||
) -> dict[str, Any]:
|
||||
_ = (public_key_hex, generated_at, window_hours, interval_minutes, model_limit)
|
||||
return {
|
||||
"schema": analytics.ANALYTICS_SCHEMA,
|
||||
"generated_at": generated_at,
|
||||
"provider_id": provider_id,
|
||||
"summary": {"total_requests": 1},
|
||||
}
|
||||
|
||||
async def fake_publish(relay_url: str, event: dict[str, Any]) -> bool:
|
||||
_ = relay_url
|
||||
published_events.append(event)
|
||||
return True
|
||||
|
||||
monkeypatch.setattr(analytics.settings, "enable_analytics_sharing", True)
|
||||
monkeypatch.setattr(analytics.settings, "nsec", "11" * 32)
|
||||
monkeypatch.setattr(analytics.settings, "relays", ["wss://relay.example.com"])
|
||||
monkeypatch.setattr(analytics.settings, "provider_id", "")
|
||||
monkeypatch.setattr(analytics, "build_stats_snapshot_payload", fake_build_payload)
|
||||
monkeypatch.setattr(analytics, "publish_to_relay", fake_publish)
|
||||
monkeypatch.setattr(analytics.asyncio, "sleep", fake_sleep)
|
||||
|
||||
await analytics.publish_usage_analytics()
|
||||
|
||||
assert len(published_events) == 1
|
||||
assert ["schema", analytics.ANALYTICS_SCHEMA] in published_events[0].get("tags", [])
|
||||
@@ -1,64 +0,0 @@
|
||||
"""Tests for the built-in 404 handler in routstr.proxy."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
from routstr import proxy
|
||||
from routstr.proxy import proxy_router
|
||||
|
||||
|
||||
def _make_app() -> FastAPI:
|
||||
app = FastAPI()
|
||||
app.include_router(proxy_router)
|
||||
return app
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
proxy._NOT_FOUND_HTML is None,
|
||||
reason="UI bundle (ui_out/404.html) not present in this environment",
|
||||
)
|
||||
def test_unknown_path_returns_html_404_for_browser() -> None:
|
||||
client = TestClient(_make_app())
|
||||
response = client.get("/some/random/page", headers={"accept": "text/html"})
|
||||
assert response.status_code == 404
|
||||
assert response.headers["content-type"].startswith("text/html")
|
||||
assert "404" in response.text
|
||||
|
||||
|
||||
def test_unknown_path_returns_json_404_for_api_client() -> None:
|
||||
client = TestClient(_make_app())
|
||||
response = client.get(
|
||||
"/some/random/page", headers={"accept": "application/json"}
|
||||
)
|
||||
assert response.status_code == 404
|
||||
assert response.headers["content-type"].startswith("application/json")
|
||||
payload = response.json()
|
||||
assert payload["error"]["type"] == "not_found"
|
||||
assert payload["error"]["code"] == 404
|
||||
assert "/some/random/page" in payload["error"]["message"]
|
||||
|
||||
|
||||
def test_root_path_returns_404_for_proxy_router() -> None:
|
||||
client = TestClient(_make_app())
|
||||
response = client.get("/", headers={"accept": "application/json"})
|
||||
assert response.status_code == 404
|
||||
|
||||
|
||||
def test_v1_path_is_not_intercepted_by_404_handler() -> None:
|
||||
"""Paths starting with v1/ must reach the proxy logic, not the 404 handler."""
|
||||
client = TestClient(_make_app(), raise_server_exceptions=False)
|
||||
response = client.get("/v1/anything")
|
||||
if response.status_code == 404:
|
||||
# Any 404 here must come from inner proxy logic, not our HTML page.
|
||||
assert "<!DOCTYPE html>" not in response.text
|
||||
|
||||
|
||||
def test_json_returned_when_ui_html_missing(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(proxy, "_NOT_FOUND_HTML", None)
|
||||
client = TestClient(_make_app())
|
||||
response = client.get("/some/random/page", headers={"accept": "text/html"})
|
||||
assert response.status_code == 404
|
||||
assert response.headers["content-type"].startswith("application/json")
|
||||
@@ -1,19 +1,16 @@
|
||||
import os
|
||||
|
||||
import pytest
|
||||
from pydantic.v1 import ValidationError
|
||||
from sqlalchemy.ext.asyncio import create_async_engine
|
||||
from sqlmodel import text
|
||||
from sqlmodel.ext.asyncio.session import AsyncSession
|
||||
|
||||
from routstr.core.settings import Settings, SettingsService
|
||||
from routstr.core.settings import SettingsService
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_settings_seed_from_env_and_persist() -> None:
|
||||
os.environ["UPSTREAM_BASE_URL"] = "https://api.test/v1"
|
||||
os.environ.pop("ONION_URL", None)
|
||||
os.environ.pop("ENABLE_ANALYTICS_SHARING", None)
|
||||
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
@@ -22,101 +19,19 @@ async def test_settings_seed_from_env_and_persist() -> None:
|
||||
assert settings.upstream_base_url == "https://api.test/v1"
|
||||
# ONION_URL may be empty if not discoverable
|
||||
assert isinstance(settings.onion_url, str)
|
||||
assert settings.enable_analytics_sharing is True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_settings_db_precedence_over_env() -> None:
|
||||
os.environ["UPSTREAM_BASE_URL"] = "https://api.env/v1"
|
||||
os.environ["ENABLE_ANALYTICS_SHARING"] = "true"
|
||||
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
_ = await SettingsService.initialize(session)
|
||||
updated = await SettingsService.update(
|
||||
{"name": "DBName", "enable_analytics_sharing": False}, session
|
||||
)
|
||||
updated = await SettingsService.update({"name": "DBName"}, session)
|
||||
assert updated.name == "DBName"
|
||||
assert updated.enable_analytics_sharing is False
|
||||
|
||||
# Change env and re-initialize; DB should still win
|
||||
os.environ["NAME"] = "EnvName"
|
||||
os.environ["ENABLE_ANALYTICS_SHARING"] = "true"
|
||||
again = await SettingsService.initialize(session)
|
||||
assert again.name == "DBName"
|
||||
assert again.enable_analytics_sharing is False
|
||||
|
||||
|
||||
def test_payout_settings_have_sensible_defaults() -> None:
|
||||
s = Settings()
|
||||
assert s.min_payout_sat == 210
|
||||
assert s.payout_interval_seconds == 900
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"field,bad_value",
|
||||
[
|
||||
("min_payout_sat", 0),
|
||||
("min_payout_sat", -1),
|
||||
("payout_interval_seconds", 0),
|
||||
("payout_interval_seconds", -10),
|
||||
],
|
||||
)
|
||||
def test_payout_settings_reject_invalid_values(field: str, bad_value: int) -> None:
|
||||
kwargs: dict[str, object] = {field: bad_value}
|
||||
with pytest.raises(ValidationError):
|
||||
Settings(**kwargs) # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_payout_settings_accept_custom_positive_values() -> None:
|
||||
s = Settings(min_payout_sat=500, payout_interval_seconds=60)
|
||||
assert s.min_payout_sat == 500
|
||||
assert s.payout_interval_seconds == 60
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_payout_settings_persist_via_settings_service() -> None:
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
await SettingsService.initialize(session)
|
||||
updated = await SettingsService.update(
|
||||
{"min_payout_sat": 1000, "payout_interval_seconds": 300}, session
|
||||
)
|
||||
assert updated.min_payout_sat == 1000
|
||||
assert updated.payout_interval_seconds == 300
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_payout_settings_update_rejects_invalid() -> None:
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
await SettingsService.initialize(session)
|
||||
with pytest.raises(ValidationError):
|
||||
await SettingsService.update({"min_payout_sat": 0}, session)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_settings_initialize_discards_unknown_keys() -> None:
|
||||
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
||||
async with AsyncSession(engine, expire_on_commit=False) as session:
|
||||
_ = await SettingsService.initialize(session)
|
||||
|
||||
# Simulate older persisted key name and an unknown key.
|
||||
await session.exec( # type: ignore
|
||||
text(
|
||||
"UPDATE settings SET data = :data WHERE id = 1"
|
||||
).bindparams(
|
||||
data='{"name":"LegacyNode","nostr_analytics_enabled":false,"unknown_key":123}'
|
||||
)
|
||||
)
|
||||
await session.commit()
|
||||
|
||||
reloaded = await SettingsService.initialize(session)
|
||||
assert reloaded.name == "LegacyNode"
|
||||
assert reloaded.enable_analytics_sharing is True
|
||||
|
||||
row = await session.exec(text("SELECT data FROM settings WHERE id = 1")) # type: ignore
|
||||
stored_data = row.first()[0]
|
||||
assert '"enable_analytics_sharing": true' in stored_data
|
||||
assert "nostr_analytics_enabled" not in stored_data
|
||||
assert "unknown_key" not in stored_data
|
||||
|
||||
@@ -1,117 +0,0 @@
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from routstr.core.db import ApiKey
|
||||
from routstr.upstream.base import BaseUpstreamProvider
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_with_id_injection() -> None:
|
||||
"""Test that stream_with_cost correctly injects IDs into complete JSON chunks but skips partials."""
|
||||
provider = BaseUpstreamProvider(
|
||||
base_url="https://api.example.com", api_key="test_key"
|
||||
)
|
||||
|
||||
# Mock response with mixed chunks:
|
||||
# 1. Complete JSON without ID
|
||||
# 2. Partial JSON (should be passed through)
|
||||
# 3. Complete JSON with ID (should be preserved or updated if requested_model is set)
|
||||
# 4. [DONE] message
|
||||
chunks = [
|
||||
b'data: {"choices": [{"delta": {"content": "Hello"}}]}\n\n',
|
||||
b'data: {"choices": [{"delta": {"content": "', # Partial
|
||||
b'world"}}]}\n\n',
|
||||
b'data: {"id": "existing-id", "choices": [{"delta": {"content": "!"}}]}\n\n',
|
||||
b"data: [DONE]\n\n",
|
||||
]
|
||||
|
||||
async def aiter_bytes() -> AsyncGenerator[bytes, None]:
|
||||
for chunk in chunks:
|
||||
yield chunk
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 200
|
||||
mock_response.headers = {"content-type": "text/event-stream"}
|
||||
mock_response.aiter_bytes = aiter_bytes
|
||||
|
||||
key = MagicMock(spec=ApiKey)
|
||||
key.hashed_key = "test_hash"
|
||||
key.balance = 1000
|
||||
|
||||
background_tasks = MagicMock()
|
||||
|
||||
# We need to mock adjust_payment_for_tokens since it's called at the end
|
||||
with MagicMock():
|
||||
from routstr.upstream import base
|
||||
|
||||
# Mocking the module-level function used in the generator
|
||||
base.adjust_payment_for_tokens = AsyncMock(
|
||||
return_value={"total_usd": 0.1, "total_msats": 100}
|
||||
)
|
||||
# create_session() is used as an async context manager whose entered
|
||||
# value exposes an awaitable .get(). Build a mock that behaves that
|
||||
# way so the post-stream cost-chunk emission can run.
|
||||
mock_session = MagicMock()
|
||||
mock_session.get = AsyncMock(return_value=key)
|
||||
mock_ctx = MagicMock()
|
||||
mock_ctx.__aenter__ = AsyncMock(return_value=mock_session)
|
||||
mock_ctx.__aexit__ = AsyncMock(return_value=None)
|
||||
base.create_session = MagicMock(return_value=mock_ctx)
|
||||
|
||||
streaming_response = await provider.handle_streaming_chat_completion(
|
||||
response=mock_response,
|
||||
key=key,
|
||||
max_cost_for_model=100,
|
||||
background_tasks=background_tasks,
|
||||
requested_model="test-model",
|
||||
)
|
||||
|
||||
results = []
|
||||
async for chunk in streaming_response.body_iterator:
|
||||
results.append(chunk)
|
||||
|
||||
# Parse results
|
||||
parsed_results = []
|
||||
for r in results:
|
||||
if isinstance(r, bytes) and r.startswith(b"data: "):
|
||||
data = r[6:].decode().strip()
|
||||
if data == "[DONE]":
|
||||
parsed_results.append(data)
|
||||
else:
|
||||
try:
|
||||
parsed_results.append(json.loads(data))
|
||||
except (json.JSONDecodeError, UnicodeDecodeError):
|
||||
parsed_results.append(
|
||||
data
|
||||
) # Keep as string if it failed to parse
|
||||
|
||||
# Verifications
|
||||
# 1. First chunk should have an injected ID and the requested model
|
||||
assert isinstance(parsed_results[0], dict)
|
||||
assert "id" in parsed_results[0]
|
||||
assert parsed_results[0]["id"].startswith("chatcmpl-")
|
||||
assert parsed_results[0]["model"] == "test-model"
|
||||
|
||||
# 2. Second chunk was partial, should be passed as-is
|
||||
# In current implementation, re.split(b"data: ", b'data: {...') gives ['', '{...']
|
||||
# The first empty part is skipped. The second part is processed.
|
||||
|
||||
# Check that we have results
|
||||
assert len(parsed_results) >= 4
|
||||
|
||||
# Find the chunk that was "existing-id"
|
||||
id_chunk = next(
|
||||
r
|
||||
for r in parsed_results
|
||||
if isinstance(r, dict)
|
||||
and "choices" in r
|
||||
and r["choices"][0]["delta"].get("content") == "!"
|
||||
)
|
||||
assert id_chunk["id"] == parsed_results[0]["id"]
|
||||
assert id_chunk["model"] == "test-model"
|
||||
|
||||
# 4. [DONE] should be there
|
||||
assert "[DONE]" in parsed_results
|
||||
@@ -1,82 +0,0 @@
|
||||
"""Tests for Azure upstream provider request normalization."""
|
||||
|
||||
from routstr.payment.models import Architecture, Model, Pricing
|
||||
from routstr.upstream.azure import AzureUpstreamProvider
|
||||
|
||||
|
||||
def create_test_model(model_id: str, canonical_slug: str | None = None) -> Model:
|
||||
return Model(
|
||||
id=model_id,
|
||||
name=model_id,
|
||||
created=0,
|
||||
description="test",
|
||||
context_length=8192,
|
||||
architecture=Architecture(
|
||||
modality="text",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"],
|
||||
tokenizer="test",
|
||||
instruct_type=None,
|
||||
),
|
||||
pricing=Pricing(
|
||||
prompt=0.001,
|
||||
completion=0.001,
|
||||
request=0.0,
|
||||
image=0.0,
|
||||
web_search=0.0,
|
||||
internal_reasoning=0.0,
|
||||
),
|
||||
canonical_slug=canonical_slug,
|
||||
)
|
||||
|
||||
|
||||
def test_prepare_headers_uses_azure_api_key_header() -> None:
|
||||
provider = AzureUpstreamProvider(
|
||||
base_url="https://example.openai.azure.com",
|
||||
api_key="azure-key",
|
||||
api_version="2024-02-15-preview",
|
||||
)
|
||||
|
||||
headers = provider.prepare_headers({"Authorization": "Bearer user-token"})
|
||||
|
||||
assert headers["api-key"] == "azure-key"
|
||||
assert "Authorization" not in headers
|
||||
assert "authorization" not in headers
|
||||
|
||||
|
||||
def test_prepare_params_normalizes_azure_api_version() -> None:
|
||||
provider = AzureUpstreamProvider(
|
||||
base_url="https://example.openai.azure.com",
|
||||
api_key="azure-key",
|
||||
api_version="\ufeff v1 ",
|
||||
)
|
||||
|
||||
params = provider.prepare_params("chat/completions", {})
|
||||
|
||||
assert params["api-version"] == "2024-02-15-preview"
|
||||
|
||||
|
||||
def test_normalize_request_path_includes_deployment_id() -> None:
|
||||
provider = AzureUpstreamProvider(
|
||||
base_url="https://example.openai.azure.com/openai/v1",
|
||||
api_key="azure-key",
|
||||
api_version="2024-02-15-preview",
|
||||
)
|
||||
model = create_test_model("azure/gpt-4o", canonical_slug="deploy-gpt4o")
|
||||
|
||||
path = provider.normalize_request_path("v1/chat/completions", model)
|
||||
|
||||
assert path == "openai/deployments/deploy-gpt4o/chat/completions"
|
||||
|
||||
|
||||
def test_get_request_base_url_strips_openai_v1_suffix() -> None:
|
||||
provider = AzureUpstreamProvider(
|
||||
base_url="https://example.openai.azure.com/openai/v1",
|
||||
api_key="azure-key",
|
||||
api_version="2024-02-15-preview",
|
||||
)
|
||||
model = create_test_model("azure/gpt-4o")
|
||||
|
||||
base_url = provider.get_request_base_url("chat/completions", model)
|
||||
|
||||
assert base_url == "https://example.openai.azure.com"
|
||||
@@ -1,150 +0,0 @@
|
||||
"""Tests for ``BaseUpstreamProvider.forward_upstream_error_response``.
|
||||
|
||||
Upstream services (e.g. an Express server that doesn't expose ``/messages``)
|
||||
sometimes return a non-JSON error body. The proxy must surface those errors
|
||||
in a consistent JSON envelope so clients don't have to parse HTML.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from routstr.upstream.base import BaseUpstreamProvider, _is_json_content_type
|
||||
|
||||
|
||||
def _make_request(request_id: str = "req-123") -> Mock:
|
||||
request = Mock(spec=["method", "state"])
|
||||
request.method = "POST"
|
||||
request.state = Mock()
|
||||
request.state.request_id = request_id
|
||||
return request
|
||||
|
||||
|
||||
def _make_upstream_response(
|
||||
*,
|
||||
body: bytes,
|
||||
status_code: int = 404,
|
||||
content_type: str | None = "text/html",
|
||||
extra_headers: dict[str, str] | None = None,
|
||||
) -> httpx.Response:
|
||||
headers: dict[str, str] = {}
|
||||
if content_type is not None:
|
||||
headers["content-type"] = content_type
|
||||
if extra_headers:
|
||||
headers.update(extra_headers)
|
||||
return httpx.Response(status_code=status_code, headers=headers, content=body)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def provider() -> BaseUpstreamProvider:
|
||||
return BaseUpstreamProvider(
|
||||
base_url="https://privateprovider.xyz", api_key="k", provider_fee=1.0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"content_type,expected",
|
||||
[
|
||||
("application/json", True),
|
||||
("application/json; charset=utf-8", True),
|
||||
("text/json", True),
|
||||
("application/problem+json", True),
|
||||
("application/vnd.api+json", True),
|
||||
("text/html", False),
|
||||
("text/html; charset=utf-8", False),
|
||||
("text/plain", False),
|
||||
("", False),
|
||||
(None, False),
|
||||
],
|
||||
)
|
||||
def test_is_json_content_type(content_type: str | None, expected: bool) -> None:
|
||||
assert _is_json_content_type(content_type) is expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_html_error_is_normalized_to_json_envelope(
|
||||
provider: BaseUpstreamProvider,
|
||||
) -> None:
|
||||
html_body = (
|
||||
b"<!DOCTYPE html><html><head><title>Error</title></head>"
|
||||
b"<body><pre>Cannot POST /messages</pre></body></html>"
|
||||
)
|
||||
upstream = _make_upstream_response(body=html_body, status_code=404)
|
||||
|
||||
response = await provider.forward_upstream_error_response(
|
||||
_make_request(), "v1/messages", upstream
|
||||
)
|
||||
|
||||
assert response.status_code == 404
|
||||
assert response.media_type == "application/json"
|
||||
payload: dict[str, Any] = json.loads(bytes(response.body))
|
||||
assert payload["error"]["type"] == "upstream_error"
|
||||
assert payload["error"]["upstream_status"] == 404
|
||||
assert payload["error"]["upstream_content_type"] == "text/html"
|
||||
assert "Cannot POST /messages" in payload["error"]["upstream_body_preview"]
|
||||
assert payload["request_id"] == "req-123"
|
||||
# The upstream's text/html content-type must not survive — Response()
|
||||
# sets the JSON content-type for us via media_type.
|
||||
assert response.headers["content-type"].startswith("application/json")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_plain_text_error_is_normalized(
|
||||
provider: BaseUpstreamProvider,
|
||||
) -> None:
|
||||
upstream = _make_upstream_response(
|
||||
body=b"Service Unavailable", status_code=503, content_type="text/plain"
|
||||
)
|
||||
|
||||
response = await provider.forward_upstream_error_response(
|
||||
_make_request(), "v1/messages", upstream
|
||||
)
|
||||
|
||||
assert response.status_code == 503
|
||||
assert response.media_type == "application/json"
|
||||
payload = json.loads(bytes(response.body))
|
||||
assert payload["error"]["message"] == "Service Unavailable"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_empty_body_with_non_json_content_type_normalizes(
|
||||
provider: BaseUpstreamProvider,
|
||||
) -> None:
|
||||
upstream = _make_upstream_response(
|
||||
body=b"", status_code=502, content_type="text/html"
|
||||
)
|
||||
|
||||
response = await provider.forward_upstream_error_response(
|
||||
_make_request(), "v1/messages", upstream
|
||||
)
|
||||
|
||||
assert response.status_code == 502
|
||||
assert response.media_type == "application/json"
|
||||
payload = json.loads(bytes(response.body))
|
||||
assert payload["error"]["type"] == "upstream_error"
|
||||
assert payload["error"]["upstream_body_preview"] is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_json_error_body_is_passed_through_unchanged(
|
||||
provider: BaseUpstreamProvider,
|
||||
) -> None:
|
||||
json_body = json.dumps(
|
||||
{"error": {"message": "Invalid model", "type": "invalid_request_error"}}
|
||||
).encode()
|
||||
upstream = _make_upstream_response(
|
||||
body=json_body, status_code=400, content_type="application/json"
|
||||
)
|
||||
|
||||
response = await provider.forward_upstream_error_response(
|
||||
_make_request(), "v1/messages", upstream
|
||||
)
|
||||
|
||||
assert response.status_code == 400
|
||||
assert bytes(response.body) == json_body
|
||||
assert response.media_type == "application/json"
|
||||
@@ -1,313 +0,0 @@
|
||||
"""Unit tests for the Gemini /v1/messages dispatch path.
|
||||
|
||||
The Gemini upstream needs special handling because its OpenAI-compat
|
||||
surface rejects inbound ``functionCall`` parts that lack a
|
||||
``thought_signature``. We bypass litellm + openai SDK at the wire layer
|
||||
(see ``routstr/upstream/gemini_messages.py``) so we can inject Google's
|
||||
documented dummy signature (``"skip_thought_signature_validator"``).
|
||||
|
||||
These tests cover the two pure helpers that drive the dispatcher:
|
||||
|
||||
* ``inject_thought_signatures`` — request-side injection
|
||||
* ``_openai_chunks_to_anthropic_events`` — response-side translator
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from collections.abc import AsyncGenerator
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from routstr.upstream.gemini_messages import (
|
||||
DUMMY_THOUGHT_SIGNATURE,
|
||||
_openai_chunks_to_anthropic_events,
|
||||
inject_thought_signatures,
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# inject_thought_signatures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_inject_thought_signatures_adds_dummy_to_each_tool_call() -> None:
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "user", "content": "do thing"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "toolu_1",
|
||||
"type": "function",
|
||||
"function": {"name": "Bash", "arguments": "{}"},
|
||||
},
|
||||
{
|
||||
"id": "toolu_2",
|
||||
"type": "function",
|
||||
"function": {"name": "Read", "arguments": "{}"},
|
||||
},
|
||||
],
|
||||
},
|
||||
]
|
||||
|
||||
inject_thought_signatures(messages)
|
||||
|
||||
for tc in messages[1]["tool_calls"]:
|
||||
assert (
|
||||
tc["extra_content"]["google"]["thought_signature"]
|
||||
== DUMMY_THOUGHT_SIGNATURE
|
||||
)
|
||||
|
||||
|
||||
def test_inject_thought_signatures_preserves_existing_signature() -> None:
|
||||
"""Don't clobber a real signature that came back from a prior turn."""
|
||||
messages: list[dict[str, Any]] = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "tc1",
|
||||
"type": "function",
|
||||
"function": {"name": "fn", "arguments": "{}"},
|
||||
"extra_content": {
|
||||
"google": {"thought_signature": "real-signature"}
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
inject_thought_signatures(messages)
|
||||
|
||||
assert (
|
||||
messages[0]["tool_calls"][0]["extra_content"]["google"][
|
||||
"thought_signature"
|
||||
]
|
||||
== "real-signature"
|
||||
)
|
||||
|
||||
|
||||
def test_inject_thought_signatures_skips_messages_without_tool_calls() -> None:
|
||||
messages: list[dict[str, Any]] = [
|
||||
{"role": "user", "content": "hi"},
|
||||
{"role": "assistant", "content": "hello"},
|
||||
]
|
||||
|
||||
inject_thought_signatures(messages)
|
||||
|
||||
for m in messages:
|
||||
assert "extra_content" not in m
|
||||
|
||||
|
||||
def test_inject_thought_signatures_handles_malformed_extra_content() -> None:
|
||||
"""If a caller already set ``extra_content`` to a non-dict (defensive),
|
||||
we replace it instead of crashing."""
|
||||
messages: list[dict[str, Any]] = [
|
||||
{
|
||||
"role": "assistant",
|
||||
"tool_calls": [
|
||||
{
|
||||
"id": "tc",
|
||||
"type": "function",
|
||||
"function": {"name": "fn", "arguments": "{}"},
|
||||
"extra_content": "garbage",
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
inject_thought_signatures(messages)
|
||||
|
||||
extra = messages[0]["tool_calls"][0]["extra_content"]
|
||||
assert isinstance(extra, dict)
|
||||
assert extra["google"]["thought_signature"] == DUMMY_THOUGHT_SIGNATURE
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _openai_chunks_to_anthropic_events
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
async def _lines(*chunks: dict | str) -> AsyncGenerator[str, None]:
|
||||
"""Helper to wrap chunk dicts as SSE-style ``data:`` lines."""
|
||||
for c in chunks:
|
||||
if isinstance(c, dict):
|
||||
yield f"data: {json.dumps(c)}"
|
||||
else:
|
||||
yield c
|
||||
|
||||
|
||||
def _parse_anthropic_sse(blocks: list[bytes]) -> list[dict]:
|
||||
"""Flatten a list of Anthropic SSE byte chunks into event dicts."""
|
||||
events: list[dict] = []
|
||||
for blob in blocks:
|
||||
text = blob.decode()
|
||||
for entry in text.split("\n\n"):
|
||||
for line in entry.splitlines():
|
||||
if line.startswith("data:"):
|
||||
events.append(json.loads(line[5:].lstrip()))
|
||||
return events
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_translator_emits_text_only_response() -> None:
|
||||
"""Plain text response: message_start → content_block_* (text) →
|
||||
message_delta(end_turn) → message_stop."""
|
||||
chunks: list[dict] = [
|
||||
{
|
||||
"id": "chatcmpl-1",
|
||||
"model": "gemini-2.5-flash",
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant"}}],
|
||||
},
|
||||
{"choices": [{"delta": {"content": "Hello"}}]},
|
||||
{"choices": [{"delta": {"content": ", world"}}]},
|
||||
{
|
||||
"choices": [{"delta": {}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 5, "completion_tokens": 7},
|
||||
},
|
||||
]
|
||||
|
||||
out = []
|
||||
async for event_bytes in _openai_chunks_to_anthropic_events(
|
||||
_lines(*chunks), requested_model="gemini-2.5-flash"
|
||||
):
|
||||
out.append(event_bytes)
|
||||
|
||||
events = _parse_anthropic_sse(out)
|
||||
types = [e["type"] for e in events]
|
||||
assert types == [
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
"message_delta",
|
||||
"message_stop",
|
||||
]
|
||||
# Text deltas concatenate to "Hello, world".
|
||||
text_deltas = [
|
||||
e["delta"]["text"]
|
||||
for e in events
|
||||
if e["type"] == "content_block_delta"
|
||||
]
|
||||
assert "".join(text_deltas) == "Hello, world"
|
||||
# Stop reason was mapped from openai's "stop".
|
||||
msg_delta = next(e for e in events if e["type"] == "message_delta")
|
||||
assert msg_delta["delta"]["stop_reason"] == "end_turn"
|
||||
assert msg_delta["usage"]["input_tokens"] == 5
|
||||
assert msg_delta["usage"]["output_tokens"] == 7
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_translator_emits_tool_use_block() -> None:
|
||||
"""tool_calls split across deltas → tool_use content block with
|
||||
accumulated input_json_delta and stop_reason='tool_use'."""
|
||||
chunks: list[dict] = [
|
||||
{
|
||||
"id": "chatcmpl-2",
|
||||
"model": "gemini-2.5-flash",
|
||||
"choices": [{"delta": {"role": "assistant"}}],
|
||||
},
|
||||
{
|
||||
"choices": [
|
||||
{
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call-abc",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "Bash",
|
||||
"arguments": '{"cmd":',
|
||||
},
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"choices": [
|
||||
{
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"function": {"arguments": ' "ls"}'},
|
||||
}
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{"choices": [{"delta": {}, "finish_reason": "tool_calls"}]},
|
||||
]
|
||||
|
||||
out = []
|
||||
async for event_bytes in _openai_chunks_to_anthropic_events(
|
||||
_lines(*chunks), requested_model="gemini-2.5-flash"
|
||||
):
|
||||
out.append(event_bytes)
|
||||
|
||||
events = _parse_anthropic_sse(out)
|
||||
types = [e["type"] for e in events]
|
||||
assert types == [
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
"message_delta",
|
||||
"message_stop",
|
||||
]
|
||||
# Tool use block was opened with the right name.
|
||||
cb_start = next(e for e in events if e["type"] == "content_block_start")
|
||||
assert cb_start["content_block"]["type"] == "tool_use"
|
||||
assert cb_start["content_block"]["name"] == "Bash"
|
||||
assert cb_start["content_block"]["id"] == "call-abc"
|
||||
# Argument deltas were forwarded as input_json_delta partials.
|
||||
deltas = [e for e in events if e["type"] == "content_block_delta"]
|
||||
assert all(d["delta"]["type"] == "input_json_delta" for d in deltas)
|
||||
assert "".join(d["delta"]["partial_json"] for d in deltas) == (
|
||||
'{"cmd": "ls"}'
|
||||
)
|
||||
# tool_calls finish_reason → tool_use stop_reason.
|
||||
msg_delta = next(e for e in events if e["type"] == "message_delta")
|
||||
assert msg_delta["delta"]["stop_reason"] == "tool_use"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_translator_handles_done_sentinel_and_blank_lines() -> None:
|
||||
"""Spec edge cases from openai SSE: ``data: [DONE]``, blank lines,
|
||||
invalid JSON. Translator should skip them gracefully."""
|
||||
chunks: list[dict | str] = [
|
||||
{
|
||||
"id": "x",
|
||||
"model": "m",
|
||||
"choices": [{"delta": {"role": "assistant"}}],
|
||||
},
|
||||
{"choices": [{"delta": {"content": "ok"}}]},
|
||||
"",
|
||||
": comment",
|
||||
"data: not-json",
|
||||
"data: [DONE]",
|
||||
{"choices": [{"delta": {}, "finish_reason": "stop"}]},
|
||||
]
|
||||
|
||||
out = []
|
||||
async for event_bytes in _openai_chunks_to_anthropic_events(
|
||||
_lines(*chunks), requested_model=None
|
||||
):
|
||||
out.append(event_bytes)
|
||||
|
||||
events = _parse_anthropic_sse(out)
|
||||
assert events[0]["type"] == "message_start"
|
||||
assert events[-1]["type"] == "message_stop"
|
||||
text = "".join(
|
||||
e["delta"]["text"]
|
||||
for e in events
|
||||
if e["type"] == "content_block_delta"
|
||||
)
|
||||
assert text == "ok"
|
||||
@@ -1,112 +0,0 @@
|
||||
from types import TracebackType
|
||||
from unittest.mock import Mock
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from routstr.upstream.routstr import RoutstrUpstreamProvider
|
||||
|
||||
|
||||
class DummyAsyncClient:
|
||||
def __init__(self, response: Mock | None = None, error: Exception | None = None):
|
||||
self.response = response
|
||||
self.error = error
|
||||
self.calls: list[dict[str, object]] = []
|
||||
|
||||
async def __aenter__(self) -> "DummyAsyncClient":
|
||||
return self
|
||||
|
||||
async def __aexit__(
|
||||
self,
|
||||
exc_type: type[BaseException] | None,
|
||||
exc: BaseException | None,
|
||||
tb: TracebackType | None,
|
||||
) -> bool:
|
||||
return False
|
||||
|
||||
async def get(
|
||||
self, url: str, headers: dict[str, str], timeout: float
|
||||
) -> Mock:
|
||||
self.calls.append({"url": url, "headers": headers, "timeout": timeout})
|
||||
if self.error is not None:
|
||||
raise self.error
|
||||
assert self.response is not None
|
||||
return self.response
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_balance_omits_auth_header_when_api_key_missing(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
response = Mock()
|
||||
response.json.return_value = {"balance_msats": 42000}
|
||||
response.raise_for_status.return_value = None
|
||||
|
||||
client = DummyAsyncClient(response=response)
|
||||
monkeypatch.setattr("routstr.upstream.routstr.httpx.AsyncClient", lambda: client)
|
||||
|
||||
provider = RoutstrUpstreamProvider(base_url="https://node.example", api_key="")
|
||||
|
||||
balance = await provider.get_balance()
|
||||
|
||||
assert balance == 42.0
|
||||
assert client.calls == [
|
||||
{
|
||||
"url": "https://node.example/v1/balance/info",
|
||||
"headers": {},
|
||||
"timeout": 10.0,
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_balance_returns_none_on_connect_timeout(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
client = DummyAsyncClient(error=httpx.ConnectTimeout("timed out"))
|
||||
monkeypatch.setattr("routstr.upstream.routstr.httpx.AsyncClient", lambda: client)
|
||||
|
||||
provider = RoutstrUpstreamProvider(
|
||||
base_url="https://node.example",
|
||||
api_key="secret",
|
||||
)
|
||||
|
||||
balance = await provider.get_balance()
|
||||
|
||||
assert balance is None
|
||||
|
||||
|
||||
def test_normalize_request_path_keeps_v1_prefix() -> None:
|
||||
"""Routstr upstream stores ``base_url`` without ``/v1``; the prefix
|
||||
must stay on the path so ``build_request_url`` produces ``/v1/<endpoint>``
|
||||
instead of ``/<endpoint>`` (which the upstream Routstr 404s with HTML)."""
|
||||
provider = RoutstrUpstreamProvider(
|
||||
base_url="https://privateprovider.xyz", api_key="key"
|
||||
)
|
||||
|
||||
assert provider.normalize_request_path("v1/messages") == "v1/messages"
|
||||
assert provider.normalize_request_path("/v1/messages") == "v1/messages"
|
||||
assert (
|
||||
provider.normalize_request_path("v1/chat/completions")
|
||||
== "v1/chat/completions"
|
||||
)
|
||||
|
||||
|
||||
def test_build_request_url_for_v1_messages() -> None:
|
||||
"""Forwarding ``/v1/messages`` must hit the upstream's ``/v1/messages``."""
|
||||
provider = RoutstrUpstreamProvider(
|
||||
base_url="https://privateprovider.xyz", api_key="key"
|
||||
)
|
||||
|
||||
normalized = provider.normalize_request_path("v1/messages")
|
||||
|
||||
assert (
|
||||
provider.build_request_url(normalized)
|
||||
== "https://privateprovider.xyz/v1/messages"
|
||||
)
|
||||
|
||||
|
||||
def test_supports_anthropic_messages_natively() -> None:
|
||||
"""Routstr nodes serve ``/v1/messages`` directly, so the proxy must
|
||||
forward as-is instead of round-tripping through litellm."""
|
||||
assert RoutstrUpstreamProvider.supports_anthropic_messages is True
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user