diff --git a/.github/workflows/docker-build.yml b/.github/workflows/docker-build.yml index 1601b86c2e0f..e931fab0bf4f 100644 --- a/.github/workflows/docker-build.yml +++ b/.github/workflows/docker-build.yml @@ -58,20 +58,25 @@ jobs: echo "${TAG}" > VERSION echo "Building tag: ${TAG} for ${{ matrix.arch }}" + - name: Normalize GHCR repository + run: echo "GHCR_REPOSITORY=${GITHUB_REPOSITORY,,}" >> $GITHUB_ENV + - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - - name: Log in to Docker Hub + - name: Log in to GHCR uses: docker/login-action@v3 with: - username: ${{ secrets.DOCKERHUB_USERNAME }} - password: ${{ secrets.DOCKERHUB_TOKEN }} + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} - name: Extract metadata (labels) id: meta uses: docker/metadata-action@v5 with: - images: calciumion/new-api + images: | + ghcr.io/${{ env.GHCR_REPOSITORY }} - name: Build & push id: build @@ -81,8 +86,8 @@ jobs: platforms: ${{ matrix.platform }} push: true tags: | - calciumion/new-api:${{ env.TAG }}-${{ matrix.arch }} - calciumion/new-api:latest-${{ matrix.arch }} + ghcr.io/${{ env.GHCR_REPOSITORY }}:${{ env.TAG }}-${{ matrix.arch }} + ghcr.io/${{ env.GHCR_REPOSITORY }}:latest-${{ matrix.arch }} labels: ${{ steps.meta.outputs.labels }} cache-from: type=gha cache-to: type=gha,mode=max @@ -93,13 +98,14 @@ jobs: uses: sigstore/cosign-installer@v3 - name: Sign image with cosign - run: cosign sign --yes calciumion/new-api@${{ steps.build.outputs.digest }} + run: | + cosign sign --yes ghcr.io/${{ env.GHCR_REPOSITORY }}@${{ steps.build.outputs.digest }} - name: Image summary run: | echo "### Docker Image Digest (${{ matrix.arch }})" >> $GITHUB_STEP_SUMMARY echo '```' >> $GITHUB_STEP_SUMMARY - echo "calciumion/new-api:${TAG}-${{ matrix.arch }}" >> $GITHUB_STEP_SUMMARY + echo "ghcr.io/${{ env.GHCR_REPOSITORY }}:${TAG}-${{ matrix.arch }}" >> $GITHUB_STEP_SUMMARY echo "${{ steps.build.outputs.digest }}" >> $GITHUB_STEP_SUMMARY echo '```' >> $GITHUB_STEP_SUMMARY @@ -108,34 +114,41 @@ jobs: needs: [build_single_arch] runs-on: ubuntu-latest if: startsWith(github.ref, 'refs/tags/') || github.event_name == 'workflow_dispatch' + permissions: + packages: write + contents: read steps: - name: Set version run: echo "TAG=${{ needs.build_single_arch.outputs.tag }}" >> $GITHUB_ENV - - name: Log in to Docker Hub + - name: Normalize GHCR repository + run: echo "GHCR_REPOSITORY=${GITHUB_REPOSITORY,,}" >> $GITHUB_ENV + + - name: Log in to GHCR uses: docker/login-action@v3 with: - username: ${{ secrets.DOCKERHUB_USERNAME }} - password: ${{ secrets.DOCKERHUB_TOKEN }} + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} - - name: Create & push manifest (version) + - name: Create & push manifest (GHCR version) run: | docker buildx imagetools create \ - -t calciumion/new-api:${TAG} \ - calciumion/new-api:${TAG}-amd64 \ - calciumion/new-api:${TAG}-arm64 + -t ghcr.io/${GHCR_REPOSITORY}:${TAG} \ + ghcr.io/${GHCR_REPOSITORY}:${TAG}-amd64 \ + ghcr.io/${GHCR_REPOSITORY}:${TAG}-arm64 - - name: Create & push manifest (latest) + - name: Create & push manifest (GHCR latest) run: | docker buildx imagetools create \ - -t calciumion/new-api:latest \ - calciumion/new-api:latest-amd64 \ - calciumion/new-api:latest-arm64 + -t ghcr.io/${GHCR_REPOSITORY}:latest \ + ghcr.io/${GHCR_REPOSITORY}:latest-amd64 \ + ghcr.io/${GHCR_REPOSITORY}:latest-arm64 - name: Manifest summary run: | echo "### Multi-arch Manifest" >> $GITHUB_STEP_SUMMARY echo '```' >> $GITHUB_STEP_SUMMARY - docker buildx imagetools inspect calciumion/new-api:${TAG} >> $GITHUB_STEP_SUMMARY + docker buildx imagetools inspect ghcr.io/${GHCR_REPOSITORY}:${TAG} >> $GITHUB_STEP_SUMMARY echo '```' >> $GITHUB_STEP_SUMMARY diff --git a/.github/workflows/docker-image-alpha.yml b/.github/workflows/docker-image-alpha.yml index 116dd1452152..612aa09b5efa 100644 --- a/.github/workflows/docker-image-alpha.yml +++ b/.github/workflows/docker-image-alpha.yml @@ -49,12 +49,6 @@ jobs: - name: Set up Docker Buildx uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3 - - name: Log in to Docker Hub - uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3 - with: - username: ${{ secrets.DOCKERHUB_USERNAME }} - password: ${{ secrets.DOCKERHUB_TOKEN }} - - name: Log in to GHCR uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3 with: @@ -67,7 +61,6 @@ jobs: uses: docker/metadata-action@c299e40c65443455700f0fdfc63efafe5b349051 # v5 with: images: | - calciumion/new-api ghcr.io/${{ env.GHCR_REPOSITORY }} - name: Build & push single-arch (to both registries) @@ -78,8 +71,6 @@ jobs: platforms: ${{ matrix.platform }} push: true tags: | - calciumion/new-api:alpha-${{ matrix.arch }} - calciumion/new-api:${{ steps.version.outputs.value }}-${{ matrix.arch }} ghcr.io/${{ env.GHCR_REPOSITORY }}:alpha-${{ matrix.arch }} ghcr.io/${{ env.GHCR_REPOSITORY }}:${{ steps.version.outputs.value }}-${{ matrix.arch }} labels: ${{ steps.meta.outputs.labels }} @@ -93,14 +84,12 @@ jobs: - name: Sign image with cosign run: | - cosign sign --yes calciumion/new-api@${{ steps.build.outputs.digest }} cosign sign --yes ghcr.io/${{ env.GHCR_REPOSITORY }}@${{ steps.build.outputs.digest }} - name: Output digest run: | echo "### Docker Image Digest (${{ matrix.arch }})" >> $GITHUB_STEP_SUMMARY echo '```' >> $GITHUB_STEP_SUMMARY - echo "calciumion/new-api:alpha-${{ matrix.arch }}" >> $GITHUB_STEP_SUMMARY echo "ghcr.io/${{ env.GHCR_REPOSITORY }}:alpha-${{ matrix.arch }}" >> $GITHUB_STEP_SUMMARY echo "${{ steps.build.outputs.digest }}" >> $GITHUB_STEP_SUMMARY echo '```' >> $GITHUB_STEP_SUMMARY @@ -128,26 +117,6 @@ jobs: echo "value=$VERSION" >> $GITHUB_OUTPUT echo "VERSION=$VERSION" >> $GITHUB_ENV - - name: Log in to Docker Hub - uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3 - with: - username: ${{ secrets.DOCKERHUB_USERNAME }} - password: ${{ secrets.DOCKERHUB_TOKEN }} - - - name: Create & push manifest (Docker Hub - alpha) - run: | - docker buildx imagetools create \ - -t calciumion/new-api:alpha \ - calciumion/new-api:alpha-amd64 \ - calciumion/new-api:alpha-arm64 - - - name: Create & push manifest (Docker Hub - versioned alpha) - run: | - docker buildx imagetools create \ - -t calciumion/new-api:${VERSION} \ - calciumion/new-api:${VERSION}-amd64 \ - calciumion/new-api:${VERSION}-arm64 - - name: Log in to GHCR uses: docker/login-action@c94ce9fb468520275223c153574b00df6fe4bcc9 # v3 with: @@ -173,7 +142,5 @@ jobs: run: | echo "### Multi-arch Manifest Digests" >> $GITHUB_STEP_SUMMARY echo '```' >> $GITHUB_STEP_SUMMARY - docker buildx imagetools inspect calciumion/new-api:alpha >> $GITHUB_STEP_SUMMARY - echo "---" >> $GITHUB_STEP_SUMMARY docker buildx imagetools inspect ghcr.io/${GHCR_REPOSITORY}:alpha >> $GITHUB_STEP_SUMMARY echo '```' >> $GITHUB_STEP_SUMMARY diff --git a/.github/workflows/docker-image-nightly.yml b/.github/workflows/docker-image-nightly.yml index 2125fa9dd925..2587ba684ce8 100644 --- a/.github/workflows/docker-image-nightly.yml +++ b/.github/workflows/docker-image-nightly.yml @@ -26,6 +26,7 @@ jobs: runs-on: ${{ matrix.runner }} permissions: + packages: write contents: read steps: @@ -43,21 +44,25 @@ jobs: echo "VERSION=$VERSION" >> $GITHUB_ENV echo "Publishing version: $VERSION for ${{ matrix.arch }}" + - name: Normalize GHCR repository + run: echo "GHCR_REPOSITORY=${GITHUB_REPOSITORY,,}" >> $GITHUB_ENV + - name: Set up Docker Buildx uses: docker/setup-buildx-action@v3 - - name: Log in to Docker Hub + - name: Log in to GHCR uses: docker/login-action@v3 with: - username: ${{ secrets.DOCKERHUB_USERNAME }} - password: ${{ secrets.DOCKERHUB_TOKEN }} + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} - name: Extract metadata (labels) id: meta uses: docker/metadata-action@v5 with: images: | - calciumion/new-api + ghcr.io/${{ env.GHCR_REPOSITORY }} - name: Build & push single-arch uses: docker/build-push-action@v6 @@ -66,8 +71,8 @@ jobs: platforms: ${{ matrix.platform }} push: true tags: | - calciumion/new-api:nightly-${{ matrix.arch }} - calciumion/new-api:${{ steps.version.outputs.value }}-${{ matrix.arch }} + ghcr.io/${{ env.GHCR_REPOSITORY }}:nightly-${{ matrix.arch }} + ghcr.io/${{ env.GHCR_REPOSITORY }}:${{ steps.version.outputs.value }}-${{ matrix.arch }} labels: ${{ steps.meta.outputs.labels }} cache-from: type=gha cache-to: type=gha,mode=max @@ -75,9 +80,12 @@ jobs: sbom: false create_manifests: - name: Create multi-arch manifests (Docker Hub) + name: Create multi-arch manifests (Docker Hub + GHCR) needs: [build_single_arch] runs-on: ubuntu-latest + permissions: + packages: write + contents: read steps: - name: Check out (shallow) @@ -92,22 +100,26 @@ jobs: echo "value=$VERSION" >> $GITHUB_OUTPUT echo "VERSION=$VERSION" >> $GITHUB_ENV - - name: Log in to Docker Hub + - name: Normalize GHCR repository + run: echo "GHCR_REPOSITORY=${GITHUB_REPOSITORY,,}" >> $GITHUB_ENV + + - name: Log in to GHCR uses: docker/login-action@v3 with: - username: ${{ secrets.DOCKERHUB_USERNAME }} - password: ${{ secrets.DOCKERHUB_TOKEN }} + registry: ghcr.io + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} - - name: Create & push manifest (Docker Hub - nightly) + - name: Create & push manifest (GHCR - nightly) run: | docker buildx imagetools create \ - -t calciumion/new-api:nightly \ - calciumion/new-api:nightly-amd64 \ - calciumion/new-api:nightly-arm64 + -t ghcr.io/${GHCR_REPOSITORY}:nightly \ + ghcr.io/${GHCR_REPOSITORY}:nightly-amd64 \ + ghcr.io/${GHCR_REPOSITORY}:nightly-arm64 - - name: Create & push manifest (Docker Hub - versioned nightly) + - name: Create & push manifest (GHCR - versioned nightly) run: | docker buildx imagetools create \ - -t calciumion/new-api:${VERSION} \ - calciumion/new-api:${VERSION}-amd64 \ - calciumion/new-api:${VERSION}-arm64 + -t ghcr.io/${GHCR_REPOSITORY}:${VERSION} \ + ghcr.io/${GHCR_REPOSITORY}:${VERSION}-amd64 \ + ghcr.io/${GHCR_REPOSITORY}:${VERSION}-arm64 diff --git a/.gitignore b/.gitignore index 75f5c4633874..d47b1b1d1d55 100644 --- a/.gitignore +++ b/.gitignore @@ -35,4 +35,5 @@ data/ .test token_estimator_test.go skills-lock.json +query-newapi/* .playwright-mcp diff --git a/Dockerfile b/Dockerfile index d01ab3f0f038..2c759433dc49 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,24 +1,28 @@ -FROM oven/bun:1@sha256:0733e50325078969732ebe3b15ce4c4be5082f18c4ac1a0f0ca4839c2e4e42a7 AS builder +FROM oven/bun:1@sha256:0733e50325078969732ebe3b15ce4c4be5082f18c4ac1a0f0ca4839c2e4e42a7 AS frontend-builder WORKDIR /build/web -COPY web/package.json web/bun.lock ./ -COPY web/default/package.json ./default/package.json -COPY web/classic/package.json ./classic/package.json -RUN bun install --frozen-lockfile -COPY ./web/default ./default -COPY ./VERSION /build/VERSION -RUN cd default && DISABLE_ESLINT_PLUGIN='true' VITE_REACT_APP_VERSION=$(cat /build/VERSION) bun run build - -FROM oven/bun:1@sha256:0733e50325078969732ebe3b15ce4c4be5082f18c4ac1a0f0ca4839c2e4e42a7 AS builder-classic +ARG FRONTEND_THEME=all +ARG FRONTEND_BUILD_GC_HEAP_SIZE=536870912 -WORKDIR /build/web COPY web/package.json web/bun.lock ./ COPY web/default/package.json ./default/package.json COPY web/classic/package.json ./classic/package.json RUN bun install --frozen-lockfile + +COPY ./web/default ./default COPY ./web/classic ./classic COPY ./VERSION /build/VERSION -RUN cd classic && VITE_REACT_APP_VERSION=$(cat /build/VERSION) bun run build +RUN set -eux; \ + mkdir -p /build/web/default/dist /build/web/classic/dist; \ + case "${FRONTEND_THEME}" in \ + default) cd /build/web/default && DISABLE_ESLINT_PLUGIN='true' BUN_JSC_gcMaxHeapSize="${FRONTEND_BUILD_GC_HEAP_SIZE}" VITE_REACT_APP_VERSION="$(cat /build/VERSION)" bun --smol run build ;; \ + classic) cd /build/web/classic && BUN_JSC_gcMaxHeapSize="${FRONTEND_BUILD_GC_HEAP_SIZE}" VITE_REACT_APP_VERSION="$(cat /build/VERSION)" bun --smol run build ;; \ + all) cd /build/web/default && DISABLE_ESLINT_PLUGIN='true' BUN_JSC_gcMaxHeapSize="${FRONTEND_BUILD_GC_HEAP_SIZE}" VITE_REACT_APP_VERSION="$(cat /build/VERSION)" bun --smol run build && cd /build/web/classic && BUN_JSC_gcMaxHeapSize="${FRONTEND_BUILD_GC_HEAP_SIZE}" VITE_REACT_APP_VERSION="$(cat /build/VERSION)" bun --smol run build ;; \ + "") ;; \ + *) echo "Invalid FRONTEND_THEME: ${FRONTEND_THEME} (use default|classic|all, or empty to skip frontend build)" >&2; exit 1 ;; \ + esac; \ + if [ ! -f /build/web/default/dist/index.html ]; then printf '%s\n' 'frontend not builtfrontend theme was not built in this image' > /build/web/default/dist/index.html; fi; \ + if [ ! -f /build/web/classic/dist/index.html ]; then printf '%s\n' 'frontend not builtfrontend theme was not built in this image' > /build/web/classic/dist/index.html; fi FROM golang:1.26.1-alpine@sha256:2389ebfa5b7f43eeafbd6be0c3700cc46690ef842ad962f6c5bd6be49ed82039 AS builder2 ENV GO111MODULE=on CGO_ENABLED=0 @@ -34,8 +38,8 @@ ADD go.mod go.sum ./ RUN go mod download COPY . . -COPY --from=builder /build/web/default/dist ./web/default/dist -COPY --from=builder-classic /build/web/classic/dist ./web/classic/dist +COPY --from=frontend-builder /build/web/default/dist ./web/default/dist +COPY --from=frontend-builder /build/web/classic/dist ./web/classic/dist RUN go build -ldflags "-s -w -X 'github.com/QuantumNous/new-api/common.Version=$(cat VERSION)'" -o new-api FROM debian:bookworm-slim@sha256:f06537653ac770703bc45b4b113475bd402f451e85223f0f2837acbf89ab020a diff --git a/bin/docker-local.sh b/bin/docker-local.sh new file mode 100755 index 000000000000..6d61be49fc09 --- /dev/null +++ b/bin/docker-local.sh @@ -0,0 +1,457 @@ +#!/usr/bin/env bash +set -Eeuo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ROOT_DIR="$(cd "${SCRIPT_DIR}/.." && pwd)" +cd "${ROOT_DIR}" + +PROJECT_NAME="${PROJECT_NAME:-new-api-local}" +IMAGE_NAME="${IMAGE_NAME:-new-api:local}" +CONTAINER_NAME="${CONTAINER_NAME:-${PROJECT_NAME}-app}" +POSTGRES_CONTAINER_NAME="${POSTGRES_CONTAINER_NAME:-${PROJECT_NAME}-postgres}" +REDIS_CONTAINER_NAME="${REDIS_CONTAINER_NAME:-${PROJECT_NAME}-redis}" +NETWORK_NAME="${NETWORK_NAME:-${PROJECT_NAME}-network}" +POSTGRES_VOLUME="${POSTGRES_VOLUME:-${PROJECT_NAME}-postgres-data}" +REDIS_VOLUME="${REDIS_VOLUME:-${PROJECT_NAME}-redis-data}" +APP_DATA_VOLUME="${APP_DATA_VOLUME:-${PROJECT_NAME}-app-data}" + +HOST_PORT="${HOST_PORT:-${PORT:-3000}}" +APP_PORT="${APP_PORT:-3000}" +POSTGRES_HOST_PORT="${POSTGRES_HOST_PORT:-}" +REDIS_HOST_PORT="${REDIS_HOST_PORT:-}" +LOCAL_TZ="${TZ:-Asia/Shanghai}" +ENV_FILE="${ENV_FILE:-}" +FOLLOW_LOGS="${FOLLOW_LOGS:-0}" +NO_CACHE="${NO_CACHE:-0}" +PLATFORM="${PLATFORM:-}" +ACTION="${1:-up}" + +STATE_DIR="${STATE_DIR:-${ROOT_DIR}/data/docker-local}" +LOG_DIR="${LOG_DIR:-${ROOT_DIR}/logs/docker-local}" +SECRETS_FILE="${SECRETS_FILE:-${STATE_DIR}/.env.generated}" + +POSTGRES_IMAGE="${POSTGRES_IMAGE:-postgres:15-alpine}" +REDIS_IMAGE="${REDIS_IMAGE:-redis:7-alpine}" +POSTGRES_DB="${POSTGRES_DB:-new-api}" +POSTGRES_USER="${POSTGRES_USER:-newapi}" +POSTGRES_PASSWORD="${POSTGRES_PASSWORD:-}" +REDIS_PASSWORD="${REDIS_PASSWORD:-}" +SESSION_SECRET="${SESSION_SECRET:-}" +CRYPTO_SECRET="${CRYPTO_SECRET:-}" +NODE_NAME="${NODE_NAME:-${PROJECT_NAME}-node-1}" +BUILD_ON_UP="${BUILD_ON_UP:-1}" +FRONTEND_BUILD_GC_HEAP_SIZE="${FRONTEND_BUILD_GC_HEAP_SIZE:-536870912}" +if [[ -z "${FRONTEND_THEME+x}" ]]; then + FRONTEND_THEME="classic" +fi + +usage() { + cat <<'USAGE' +Usage: + bash bin/docker-local.sh [up|build|run|stop|logs|status|clean|help] + +Default action: + up Build app image, then start PostgreSQL, Redis, and new-api. + +No manual config is required. The script persists generated secrets in: + ./data/docker-local/.env.generated + +Common environment overrides: + PROJECT_NAME=new-api-local Prefix for containers/network/volumes + IMAGE_NAME=new-api:local Docker image tag for the app + HOST_PORT=3000 Host port for new-api + POSTGRES_HOST_PORT=5432 Optional host port for PostgreSQL + REDIS_HOST_PORT=6379 Optional host port for Redis + BUILD_ON_UP=0 Skip docker build during up + NO_CACHE=1 Build without Docker cache + PLATFORM=linux/amd64 Optional docker build --platform value + FOLLOW_LOGS=1 Follow app logs after starting + ENV_FILE=.env.local Optional extra env file for the app + FRONTEND_THEME=classic Frontend theme for local deployment/build (default: classic; use default|classic; empty to skip build/theme) + FRONTEND_BUILD_GC_HEAP_SIZE=536870912 Bun/JSC GC heap limit for frontend build; lower is slower but uses less memory + +Advanced overrides: + POSTGRES_PASSWORD=... Override generated PostgreSQL password + REDIS_PASSWORD=... Override generated Redis password + SESSION_SECRET=... Override generated session secret + CRYPTO_SECRET=... Override generated crypto secret + +Examples: + bash bin/docker-local.sh + HOST_PORT=3001 bash bin/docker-local.sh up + BUILD_ON_UP=0 bash bin/docker-local.sh up + bash bin/docker-local.sh logs + bash bin/docker-local.sh status + bash bin/docker-local.sh stop + bash bin/docker-local.sh clean +USAGE +} + +log() { + printf '\033[1;34m==>\033[0m %s\n' "$*" +} + +warn() { + printf '\033[1;33mWARN\033[0m %s\n' "$*" >&2 +} + +require_cmd() { + if ! command -v "$1" >/dev/null 2>&1; then + echo "Missing required command: $1" >&2 + exit 1 + fi +} + +random_secret() { + if command -v openssl >/dev/null 2>&1; then + openssl rand -hex 32 + else + LC_ALL=C tr -dc 'A-Za-z0-9' "${SECRETS_FILE}" </dev/null 2>&1 +} + +container_running() { + [[ "$(docker inspect -f '{{.State.Running}}' "$1" 2>/dev/null || true)" == "true" ]] +} + +remove_container_if_exists() { + local name="$1" + if container_exists "${name}"; then + log "Removing existing container ${name}" + docker rm -f "${name}" >/dev/null + fi +} + +ensure_network() { + if ! docker network inspect "${NETWORK_NAME}" >/dev/null 2>&1; then + log "Creating Docker network ${NETWORK_NAME}" + docker network create "${NETWORK_NAME}" >/dev/null + fi +} + +build_image() { + require_cmd docker + + local build_args=() + if [[ "${NO_CACHE}" == "1" || "${NO_CACHE}" == "true" ]]; then + build_args+=(--no-cache) + fi + if [[ -n "${PLATFORM}" ]]; then + build_args+=(--platform "${PLATFORM}") + fi + build_args+=(--build-arg "FRONTEND_THEME=${FRONTEND_THEME}") + build_args+=(--build-arg "FRONTEND_BUILD_GC_HEAP_SIZE=${FRONTEND_BUILD_GC_HEAP_SIZE}") + + log "Building Docker image ${IMAGE_NAME} (frontend theme: ${FRONTEND_THEME:-none}, heap: ${FRONTEND_BUILD_GC_HEAP_SIZE})" + DOCKER_BUILDKIT="${DOCKER_BUILDKIT:-1}" docker build \ + "${build_args[@]}" \ + -f "${ROOT_DIR}/Dockerfile" \ + -t "${IMAGE_NAME}" \ + "${ROOT_DIR}" +} + +start_postgres() { + require_cmd docker + ensure_network + ensure_secrets + + if container_running "${POSTGRES_CONTAINER_NAME}"; then + log "PostgreSQL already running: ${POSTGRES_CONTAINER_NAME}" + return + fi + remove_container_if_exists "${POSTGRES_CONTAINER_NAME}" + + local port_args=() + if [[ -n "${POSTGRES_HOST_PORT}" ]]; then + port_args=(-p "${POSTGRES_HOST_PORT}:5432") + fi + + log "Starting PostgreSQL ${POSTGRES_CONTAINER_NAME}" + docker run -d \ + --name "${POSTGRES_CONTAINER_NAME}" \ + --restart unless-stopped \ + --network "${NETWORK_NAME}" \ + "${port_args[@]}" \ + -v "${POSTGRES_VOLUME}:/var/lib/postgresql/data" \ + -e "POSTGRES_DB=${POSTGRES_DB}" \ + -e "POSTGRES_USER=${POSTGRES_USER}" \ + -e "POSTGRES_PASSWORD=${POSTGRES_PASSWORD}" \ + -e "TZ=${LOCAL_TZ}" \ + "${POSTGRES_IMAGE}" >/dev/null +} + +start_redis() { + require_cmd docker + ensure_network + ensure_secrets + + if container_running "${REDIS_CONTAINER_NAME}"; then + log "Redis already running: ${REDIS_CONTAINER_NAME}" + return + fi + remove_container_if_exists "${REDIS_CONTAINER_NAME}" + + local port_args=() + if [[ -n "${REDIS_HOST_PORT}" ]]; then + port_args=(-p "${REDIS_HOST_PORT}:6379") + fi + + log "Starting Redis ${REDIS_CONTAINER_NAME}" + docker run -d \ + --name "${REDIS_CONTAINER_NAME}" \ + --restart unless-stopped \ + --network "${NETWORK_NAME}" \ + "${port_args[@]}" \ + -v "${REDIS_VOLUME}:/data" \ + -e "TZ=${LOCAL_TZ}" \ + "${REDIS_IMAGE}" \ + redis-server --appendonly yes --requirepass "${REDIS_PASSWORD}" >/dev/null +} + +wait_for_postgres() { + log "Waiting for PostgreSQL" + local i + for i in {1..60}; do + if docker exec "${POSTGRES_CONTAINER_NAME}" pg_isready -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" >/dev/null 2>&1; then + return + fi + sleep 1 + done + echo "PostgreSQL did not become ready in time" >&2 + docker logs "${POSTGRES_CONTAINER_NAME}" >&2 || true + exit 1 +} + +wait_for_redis() { + log "Waiting for Redis" + local i + for i in {1..60}; do + if docker exec "${REDIS_CONTAINER_NAME}" redis-cli -a "${REDIS_PASSWORD}" ping >/dev/null 2>&1; then + return + fi + sleep 1 + done + echo "Redis did not become ready in time" >&2 + docker logs "${REDIS_CONTAINER_NAME}" >&2 || true + exit 1 +} + +wait_for_options_table() { + local i + for i in {1..60}; do + if docker exec "${POSTGRES_CONTAINER_NAME}" psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -tAc "SELECT to_regclass('public.options') IS NOT NULL" 2>/dev/null | grep -q "t"; then + return + fi + sleep 1 + done + echo "options table did not become ready in time" >&2 + docker logs "${CONTAINER_NAME}" >&2 || true + exit 1 +} + +apply_frontend_theme() { + if [[ -z "${FRONTEND_THEME}" ]]; then + return + fi + if [[ "${FRONTEND_THEME}" != "default" && "${FRONTEND_THEME}" != "classic" ]]; then + echo "Invalid FRONTEND_THEME: ${FRONTEND_THEME} (use default|classic, or empty to skip)" >&2 + exit 1 + fi + + log "Setting frontend theme to ${FRONTEND_THEME}" + wait_for_options_table + docker exec "${POSTGRES_CONTAINER_NAME}" psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -v ON_ERROR_STOP=1 -c \ + "INSERT INTO options (\"key\", \"value\") VALUES ('theme.frontend','${FRONTEND_THEME}') ON CONFLICT (\"key\") DO UPDATE SET \"value\" = EXCLUDED.\"value\";" >/dev/null + + log "Restarting app to apply frontend theme" + docker restart "${CONTAINER_NAME}" >/dev/null +} + +run_container() { + require_cmd docker + ensure_network + ensure_secrets + start_postgres + start_redis + wait_for_postgres + wait_for_redis + mkdir -p "${LOG_DIR}" + + remove_container_if_exists "${CONTAINER_NAME}" + + local sql_dsn="postgresql://${POSTGRES_USER}:${POSTGRES_PASSWORD}@${POSTGRES_CONTAINER_NAME}:5432/${POSTGRES_DB}" + local redis_dsn="redis://:${REDIS_PASSWORD}@${REDIS_CONTAINER_NAME}:6379/0" + local env_args=( + -e "TZ=${LOCAL_TZ}" + -e "PORT=${APP_PORT}" + -e "SQL_DSN=${SQL_DSN:-${sql_dsn}}" + -e "REDIS_CONN_STRING=${REDIS_CONN_STRING:-${redis_dsn}}" + -e "SESSION_SECRET=${SESSION_SECRET}" + -e "CRYPTO_SECRET=${CRYPTO_SECRET}" + -e "ERROR_LOG_ENABLED=${ERROR_LOG_ENABLED:-true}" + -e "BATCH_UPDATE_ENABLED=${BATCH_UPDATE_ENABLED:-true}" + -e "MEMORY_CACHE_ENABLED=${MEMORY_CACHE_ENABLED:-true}" + -e "SYNC_FREQUENCY=${SYNC_FREQUENCY:-60}" + -e "NODE_NAME=${NODE_NAME}" + ) + + local pass_env_vars=( + LOG_SQL_DSN + RELAY_TIMEOUT + STREAMING_TIMEOUT + CHANNEL_UPDATE_FREQUENCY + GENERATE_DEFAULT_TOKEN + FRONTEND_BASE_URL + TRUSTED_REDIRECT_DOMAINS + ) + local name + for name in "${pass_env_vars[@]}"; do + if [[ -n "${!name:-}" ]]; then + env_args+=(-e "${name}=${!name}") + fi + done + + if [[ -n "${ENV_FILE}" ]]; then + if [[ ! -f "${ENV_FILE}" ]]; then + echo "ENV_FILE does not exist: ${ENV_FILE}" >&2 + exit 1 + fi + env_args+=(--env-file "${ENV_FILE}") + fi + + log "Starting app ${CONTAINER_NAME} on http://localhost:${HOST_PORT}" + docker run -d \ + --name "${CONTAINER_NAME}" \ + --restart unless-stopped \ + --network "${NETWORK_NAME}" \ + -p "${HOST_PORT}:${APP_PORT}" \ + -v "${APP_DATA_VOLUME}:/data" \ + -v "${LOG_DIR}:/app/logs" \ + "${env_args[@]}" \ + "${IMAGE_NAME}" \ + --log-dir /app/logs >/dev/null + + apply_frontend_theme + + log "Secrets file: ${SECRETS_FILE}" + log "PostgreSQL volume: ${POSTGRES_VOLUME}" + log "Redis volume: ${REDIS_VOLUME}" + log "App data volume: ${APP_DATA_VOLUME}" + log "Logs dir: ${LOG_DIR}" + log "Open: http://localhost:${HOST_PORT}" + + if [[ "${FOLLOW_LOGS}" == "1" || "${FOLLOW_LOGS}" == "true" ]]; then + docker logs -f "${CONTAINER_NAME}" + fi +} + +stop_container() { + require_cmd docker + remove_container_if_exists "${CONTAINER_NAME}" + remove_container_if_exists "${REDIS_CONTAINER_NAME}" + remove_container_if_exists "${POSTGRES_CONTAINER_NAME}" +} + +show_logs() { + require_cmd docker + local target="${2:-app}" + case "${target}" in + app) docker logs -f "${CONTAINER_NAME}" ;; + postgres|pg) docker logs -f "${POSTGRES_CONTAINER_NAME}" ;; + redis) docker logs -f "${REDIS_CONTAINER_NAME}" ;; + *) echo "Unknown logs target: ${target} (use app|postgres|redis)" >&2; exit 1 ;; + esac +} + +show_status() { + require_cmd docker + docker ps -a \ + --filter "name=^/${CONTAINER_NAME}$" \ + --filter "name=^/${POSTGRES_CONTAINER_NAME}$" \ + --filter "name=^/${REDIS_CONTAINER_NAME}$" +} + +clean_all() { + stop_container + log "Removing image ${IMAGE_NAME} if it exists" + docker image rm "${IMAGE_NAME}" >/dev/null 2>&1 || true + + if [[ "${KEEP_VOLUMES:-1}" == "0" || "${KEEP_VOLUMES:-1}" == "false" ]]; then + warn "Removing persistent volumes and generated secrets" + docker volume rm "${POSTGRES_VOLUME}" "${REDIS_VOLUME}" "${APP_DATA_VOLUME}" >/dev/null 2>&1 || true + rm -f "${SECRETS_FILE}" + else + log "Keeping volumes. Set KEEP_VOLUMES=0 bash bin/docker-local.sh clean to remove them." + fi +} + +case "${ACTION}" in + up) + if [[ "${BUILD_ON_UP}" == "1" || "${BUILD_ON_UP}" == "true" ]]; then + build_image + fi + run_container + ;; + build) + build_image + ;; + run) + run_container + ;; + stop) + stop_container + ;; + logs) + show_logs "$@" + ;; + status) + show_status + ;; + clean) + clean_all + ;; + help|-h|--help) + usage + ;; + *) + echo "Unknown action: ${ACTION}" >&2 + usage + exit 1 + ;; +esac diff --git a/controller/channel-billing.go b/controller/channel-billing.go index 751ee3600ac9..9691df266aeb 100644 --- a/controller/channel-billing.go +++ b/controller/channel-billing.go @@ -32,6 +32,10 @@ type OpenAISubscriptionResponse struct { AccessUntil int64 `json:"access_until"` } +type channelBalanceUpdateRequest struct { + Balance *float64 `json:"balance"` +} + type OpenAIUsageDailyCost struct { Timestamp float64 `json:"timestamp"` LineItems []struct { @@ -421,6 +425,28 @@ func updateChannelBalance(channel *model.Channel) (float64, error) { return balance, nil } +func ClearChannelUsedQuota(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + _, err = model.GetChannelById(id, false) + if err != nil { + common.ApiError(c, err) + return + } + if err = model.ResetChannelUsedQuota(id); err != nil { + common.ApiError(c, err) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "used_quota": 0, + }) +} + func UpdateChannelBalance(c *gin.Context) { id, err := strconv.Atoi(c.Param("id")) if err != nil { @@ -451,6 +477,35 @@ func UpdateChannelBalance(c *gin.Context) { }) } +func SetChannelBalance(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + var request channelBalanceUpdateRequest + if err := c.ShouldBindJSON(&request); err != nil { + common.ApiError(c, err) + return + } + if request.Balance == nil { + common.ApiErrorMsg(c, "剩余额度不能为空") + return + } + channel, err := model.GetChannelById(id, false) + if err != nil { + common.ApiError(c, err) + return + } + channel.UpdateBalance(*request.Balance) + model.InitChannelCache() + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "balance": *request.Balance, + }) +} + func updateAllChannelsBalance() error { channels, err := model.GetAllChannels(0, 0, true, false) if err != nil { diff --git a/controller/channel-test.go b/controller/channel-test.go index 37bf422b1ce3..5710ea6598ad 100644 --- a/controller/channel-test.go +++ b/controller/channel-test.go @@ -43,6 +43,84 @@ type testResult struct { newAPIError *types.NewAPIError } +var preferredAnthropicChannelTestModels = []string{ + "claude-sonnet-4-5-20250929", + "claude-sonnet-4-20250514", + "claude-3-7-sonnet-20250219", + "claude-3-5-sonnet-20241022", +} + +var deprecatedAnthropicChannelTestModels = map[string]bool{ + "claude-3-sonnet-20240229": true, +} + +func resolveChannelTestModel(channel *model.Channel, requestedModel string) string { + requestedModel = strings.TrimSpace(requestedModel) + if requestedModel != "" { + return requestedModel + } + if channel == nil { + return "gpt-4o-mini" + } + if channel.Type == constant.ChannelTypeAnthropic { + if modelName := defaultAnthropicChannelTestModel(channel); modelName != "" { + return modelName + } + } + if channel.TestModel != nil && strings.TrimSpace(*channel.TestModel) != "" { + return strings.TrimSpace(*channel.TestModel) + } + models := channel.GetModels() + if len(models) > 0 { + if modelName := strings.TrimSpace(models[0]); modelName != "" { + return modelName + } + } + return "gpt-4o-mini" +} + +func defaultAnthropicChannelTestModel(channel *model.Channel) string { + if channel == nil { + return "" + } + if channel.TestModel != nil { + if modelName := strings.TrimSpace(*channel.TestModel); isUsableAnthropicChannelTestModel(modelName) { + return modelName + } + } + + availableModels := make(map[string]bool) + orderedModels := make([]string, 0) + for _, modelName := range channel.GetModels() { + modelName = strings.TrimSpace(modelName) + if modelName == "" || availableModels[modelName] { + continue + } + availableModels[modelName] = true + orderedModels = append(orderedModels, modelName) + } + + for _, modelName := range preferredAnthropicChannelTestModels { + if availableModels[modelName] && isUsableAnthropicChannelTestModel(modelName) { + return modelName + } + } + for _, modelName := range orderedModels { + if isUsableAnthropicChannelTestModel(modelName) { + return modelName + } + } + return "" +} + +func isUsableAnthropicChannelTestModel(modelName string) bool { + modelName = strings.TrimSpace(modelName) + if modelName == "" || deprecatedAnthropicChannelTestModels[modelName] { + return false + } + return helper.HasModelBillingConfig(modelName) +} + func normalizeChannelTestEndpoint(channel *model.Channel, modelName, endpointType string) string { normalized := strings.TrimSpace(endpointType) if normalized != "" { @@ -94,20 +172,7 @@ func testChannel(channel *model.Channel, testUserID int, testModel string, endpo w := httptest.NewRecorder() c, _ := gin.CreateTestContext(w) - testModel = strings.TrimSpace(testModel) - if testModel == "" { - if channel.TestModel != nil && *channel.TestModel != "" { - testModel = strings.TrimSpace(*channel.TestModel) - } else { - models := channel.GetModels() - if len(models) > 0 { - testModel = strings.TrimSpace(models[0]) - } - if testModel == "" { - testModel = "gpt-4o-mini" - } - } - } + testModel = resolveChannelTestModel(channel, testModel) endpointType = normalizeChannelTestEndpoint(channel, testModel, endpointType) diff --git a/controller/channel.go b/controller/channel.go index c59e492a5a02..86fddfaaddf4 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -14,7 +14,6 @@ import ( "github.com/QuantumNous/new-api/dto" "github.com/QuantumNous/new-api/model" relaychannel "github.com/QuantumNous/new-api/relay/channel" - "github.com/QuantumNous/new-api/relay/channel/gemini" "github.com/QuantumNous/new-api/relay/channel/ollama" "github.com/QuantumNous/new-api/service" @@ -70,13 +69,7 @@ func clearChannelInfo(channel *model.Channel) { } func applyChannelStatusFilter(query *gorm.DB, statusFilter int) *gorm.DB { - if statusFilter == common.ChannelStatusEnabled { - return query.Where("status = ?", common.ChannelStatusEnabled) - } - if statusFilter == 0 { - return query.Where("status != ?", common.ChannelStatusEnabled) - } - return query + return model.ApplyChannelStatusFilter(query, statusFilter) } func buildChannelListQuery(group string, statusFilter int, typeFilter int) *gorm.DB { @@ -109,6 +102,12 @@ func GetAllChannels(c *gin.Context) { } var total int64 + stats, err := model.GetChannelListStats(buildChannelListQuery(groupFilter, statusFilter, typeFilter)) + if err != nil { + common.SysError("failed to calculate channel stats: " + err.Error()) + c.JSON(http.StatusOK, gin.H{"success": false, "message": "获取渠道统计失败，请稍后重试"}) + return + } if enableTagMode { tags, err := model.GetPaginatedChannelTags(buildChannelListQuery(groupFilter, statusFilter, typeFilter), pageInfo.GetStartIdx(), pageInfo.GetPageSize()) @@ -180,6 +179,7 @@ func GetAllChannels(c *gin.Context) { "total": total, "page": pageInfo.GetPage(), "page_size": pageInfo.GetPageSize(), + "stats": stats, "type_counts": typeCounts, }) return @@ -257,6 +257,149 @@ func FixChannelsAbilities(c *gin.Context) { }) } +func normalizeBatchChannelKeys(keys []string) []string { + normalized := make([]string, 0, len(keys)) + seen := make(map[string]struct{}, len(keys)) + for _, key := range keys { + key = strings.TrimSpace(key) + if key == "" { + continue + } + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + normalized = append(normalized, key) + } + return normalized +} + +type BatchChannelKeySearchRequest struct { + Keys []string `json:"keys"` + Keyword string `json:"keyword"` + Group string `json:"group"` + Model string `json:"model"` + Status string `json:"status"` + Type *json.RawMessage `json:"type"` + IDSort bool `json:"id_sort"` + SortBy string `json:"sort_by"` + SortOrder string `json:"sort_order"` + Page int `json:"p"` + PageSize int `json:"page_size"` + TagMode bool `json:"tag_mode"` +} + +func parseBatchChannelTypeFilter(raw *json.RawMessage) (int, bool) { + if raw == nil || len(*raw) == 0 || string(*raw) == "null" { + return -1, false + } + var numeric int + if err := common.Unmarshal(*raw, &numeric); err == nil { + return numeric, true + } + var text string + if err := common.Unmarshal(*raw, &text); err == nil { + text = strings.TrimSpace(text) + if text == "" { + return -1, false + } + if numeric, err := strconv.Atoi(text); err == nil { + return numeric, true + } + } + return -1, false +} + +func SearchChannelsByKeys(c *gin.Context) { + request := BatchChannelKeySearchRequest{} + if err := c.ShouldBindJSON(&request); err != nil { + common.ApiError(c, err) + return + } + if request.TagMode { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "批量密钥查询暂不支持标签模式", + }) + return + } + + keys := normalizeBatchChannelKeys(request.Keys) + if len(keys) == 0 { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "keys不能为空", + }) + return + } + + statusFilter := parseStatusFilter(request.Status) + sortOptions := model.NewChannelSortOptions(request.SortBy, request.SortOrder, request.IDSort) + channelData, err := model.SearchChannelsByExactKeys(keys, request.Keyword, request.Group, request.Model, statusFilter, request.IDSort, sortOptions) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + + typeCounts := make(map[int64]int64) + for _, channel := range channelData { + typeCounts[int64(channel.Type)]++ + } + + typeFilter, hasTypeFilter := parseBatchChannelTypeFilter(request.Type) + if hasTypeFilter && typeFilter >= 0 { + filtered := make([]*model.Channel, 0, len(channelData)) + for _, ch := range channelData { + if ch.Type == typeFilter { + filtered = append(filtered, ch) + } + } + channelData = filtered + } + + page := request.Page + if page < 1 { + page = 1 + } + pageSize := request.PageSize + if pageSize <= 0 { + pageSize = common.ItemsPerPage + } + if pageSize > 100 { + pageSize = 100 + } + + total := len(channelData) + startIdx := (page - 1) * pageSize + if startIdx > total { + startIdx = total + } + endIdx := startIdx + pageSize + if endIdx > total { + endIdx = total + } + pagedData := channelData[startIdx:endIdx] + + for _, datum := range pagedData { + clearChannelInfo(datum) + datum.Key = "" + } + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "data": gin.H{ + "items": pagedData, + "total": total, + "type_counts": typeCounts, + }, + }) + return +} + func SearchChannels(c *gin.Context) { keyword := c.Query("keyword") group := c.Query("group") @@ -342,6 +485,8 @@ func SearchChannels(c *gin.Context) { channelData = filtered } + stats := model.CalculateChannelListStats(channelData) + page, _ := strconv.Atoi(c.DefaultQuery("p", "1")) pageSize, _ := strconv.Atoi(c.DefaultQuery("page_size", "20")) if page < 1 { @@ -373,6 +518,7 @@ func SearchChannels(c *gin.Context) { "data": gin.H{ "items": pagedData, "total": total, + "stats": stats, "type_counts": typeCounts, }, }) @@ -455,14 +601,18 @@ func validateTwoFactorAuth(twoFA *model.TwoFA, code string) bool { // validateChannel 通用的渠道校验函数 func validateChannel(channel *model.Channel, isAdd bool) error { + if channel == nil { + return fmt.Errorf("channel cannot be empty") + } + // 校验 channel settings if err := channel.ValidateSettings(); err != nil { return fmt.Errorf("渠道额外设置[channel setting] 格式错误：%s", err.Error()) } - // 如果是添加操作，检查 channel 和 key 是否为空 + // 如果是添加操作，检查 key 是否为空 if isAdd { - if channel == nil || channel.Key == "" { + if channel.Key == "" { return fmt.Errorf("channel cannot be empty") } @@ -584,83 +734,64 @@ func getVertexArrayKeys(keys string) ([]string, error) { return cleanKeys, nil } -func AddChannel(c *gin.Context) { - addChannelRequest := AddChannelRequest{} - err := c.ShouldBindJSON(&addChannelRequest) - if err != nil { - common.ApiError(c, err) - return +func buildChannelsFromAddRequest(addChannelRequest *AddChannelRequest) ([]model.Channel, error) { + if addChannelRequest == nil { + return nil, fmt.Errorf("channel cannot be empty") } - - // 使用统一的校验函数 if err := validateChannel(addChannelRequest.Channel, true); err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - return + return nil, err } - addChannelRequest.Channel.CreatedTime = common.GetTimestamp() + baseChannel := *addChannelRequest.Channel + baseChannel.CreatedTime = common.GetTimestamp() keys := make([]string, 0) switch addChannelRequest.Mode { case "multi_to_single": - addChannelRequest.Channel.ChannelInfo.IsMultiKey = true - addChannelRequest.Channel.ChannelInfo.MultiKeyMode = addChannelRequest.MultiKeyMode - if addChannelRequest.Channel.Type == constant.ChannelTypeVertexAi && addChannelRequest.Channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey { - array, err := getVertexArrayKeys(addChannelRequest.Channel.Key) + baseChannel.ChannelInfo.IsMultiKey = true + baseChannel.ChannelInfo.MultiKeyMode = addChannelRequest.MultiKeyMode + if baseChannel.Type == constant.ChannelTypeVertexAi && baseChannel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey { + array, err := getVertexArrayKeys(baseChannel.Key) if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - return + return nil, err } - addChannelRequest.Channel.ChannelInfo.MultiKeySize = len(array) - addChannelRequest.Channel.Key = strings.Join(array, "\n") + baseChannel.ChannelInfo.MultiKeySize = len(array) + baseChannel.Key = strings.Join(array, "\n") } else { cleanKeys := make([]string, 0) - for _, key := range strings.Split(addChannelRequest.Channel.Key, "\n") { + for _, key := range strings.Split(baseChannel.Key, "\n") { + key = strings.TrimSpace(key) if key == "" { continue } - key = strings.TrimSpace(key) cleanKeys = append(cleanKeys, key) } - addChannelRequest.Channel.ChannelInfo.MultiKeySize = len(cleanKeys) - addChannelRequest.Channel.Key = strings.Join(cleanKeys, "\n") + baseChannel.ChannelInfo.MultiKeySize = len(cleanKeys) + baseChannel.Key = strings.Join(cleanKeys, "\n") } - keys = []string{addChannelRequest.Channel.Key} + keys = []string{baseChannel.Key} case "batch": - if addChannelRequest.Channel.Type == constant.ChannelTypeVertexAi && addChannelRequest.Channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey { - // multi json - keys, err = getVertexArrayKeys(addChannelRequest.Channel.Key) + if baseChannel.Type == constant.ChannelTypeVertexAi && baseChannel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey { + array, err := getVertexArrayKeys(baseChannel.Key) if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - return + return nil, err } + keys = array } else { - keys = strings.Split(addChannelRequest.Channel.Key, "\n") + keys = strings.Split(baseChannel.Key, "\n") } case "single": - keys = []string{addChannelRequest.Channel.Key} + keys = []string{baseChannel.Key} default: - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "不支持的添加模式", - }) - return + return nil, fmt.Errorf("不支持的添加模式") } channels := make([]model.Channel, 0, len(keys)) for _, key := range keys { + key = strings.TrimSpace(key) if key == "" { continue } - localChannel := addChannelRequest.Channel + localChannel := baseChannel localChannel.Key = key if addChannelRequest.BatchAddSetKeyPrefix2Name && len(keys) > 1 { keyPrefix := localChannel.Key @@ -669,13 +800,46 @@ func AddChannel(c *gin.Context) { } localChannel.Name = fmt.Sprintf("%s %s", localChannel.Name, keyPrefix) } - channels = append(channels, *localChannel) + channels = append(channels, localChannel) + } + if len(channels) == 0 { + return nil, fmt.Errorf("channel cannot be empty") + } + return channels, nil +} + +func createChannelsFromAddRequest(addChannelRequest *AddChannelRequest, tx *gorm.DB) ([]model.Channel, error) { + channels, err := buildChannelsFromAddRequest(addChannelRequest) + if err != nil { + return nil, err + } + if tx != nil { + err = model.CreateChannelsWithTx(tx, channels) + } else { + err = model.BatchInsertChannels(channels) } - err = model.BatchInsertChannels(channels) + if err != nil { + return nil, err + } + return channels, nil +} + +func AddChannel(c *gin.Context) { + addChannelRequest := AddChannelRequest{} + err := c.ShouldBindJSON(&addChannelRequest) if err != nil { common.ApiError(c, err) return } + + _, err = createChannelsFromAddRequest(&addChannelRequest, nil) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } service.ResetProxyClientCache() c.JSON(http.StatusOK, gin.H{ "success": true, @@ -993,9 +1157,10 @@ func UpdateChannel(c *gin.Context) { func FetchModels(c *gin.Context) { var req struct { - BaseURL string `json:"base_url"` - Type int `json:"type"` - Key string `json:"key"` + BaseURL string `json:"base_url"` + Type int `json:"type"` + Key string `json:"key"` + HeaderOverride string `json:"header_override"` } if err := c.ShouldBindJSON(&req); err != nil { @@ -1006,105 +1171,29 @@ func FetchModels(c *gin.Context) { return } - baseURL := req.BaseURL - if baseURL == "" { - baseURL = constant.ChannelBaseURLs[req.Type] + // 预览/新建渠道时还没有入库的 Channel，复用已保存渠道的上游拉取逻辑， + // 避免各渠道在鉴权头、特殊模型地址、代理等细节上出现分叉。 + channel := &model.Channel{ + Type: req.Type, + Key: strings.TrimSpace(strings.Split(strings.TrimSpace(req.Key), "\n")[0]), } - // remove line breaks and extra spaces. - key := strings.TrimSpace(req.Key) - key = strings.Split(key, "\n")[0] - - if req.Type == constant.ChannelTypeOllama { - models, err := ollama.FetchOllamaModels(baseURL, key) - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": fmt.Sprintf("获取Ollama模型失败: %s", err.Error()), - }) - return - } - - names := make([]string, 0, len(models)) - for _, modelInfo := range models { - names = append(names, modelInfo.Name) - } - - c.JSON(http.StatusOK, gin.H{ - "success": true, - "data": names, - }) - return + if baseURL := strings.TrimSpace(req.BaseURL); baseURL != "" { + channel.BaseURL = &baseURL } - - if req.Type == constant.ChannelTypeGemini { - models, err := gemini.FetchGeminiModels(baseURL, key, "") - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": fmt.Sprintf("获取Gemini模型失败: %s", err.Error()), - }) - return - } - - c.JSON(http.StatusOK, gin.H{ - "success": true, - "data": models, - }) - return + if headerOverride := strings.TrimSpace(req.HeaderOverride); headerOverride != "" { + channel.HeaderOverride = &headerOverride } - client := &http.Client{} - url := fmt.Sprintf("%s/v1/models", baseURL) - - request, err := http.NewRequest("GET", url, nil) + models, err := fetchChannelUpstreamModelIDs(channel) if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{ - "success": false, - "message": err.Error(), - }) - return - } - - request.Header.Set("Authorization", "Bearer "+key) - - response, err := client.Do(request) - if err != nil { - c.JSON(http.StatusInternalServerError, gin.H{ - "success": false, - "message": err.Error(), - }) - return - } - //check status code - if response.StatusCode != http.StatusOK { - c.JSON(http.StatusInternalServerError, gin.H{ - "success": false, - "message": "Failed to fetch models", - }) - return - } - defer response.Body.Close() - - var result struct { - Data []struct { - ID string `json:"id"` - } `json:"data"` - } - - if err := json.NewDecoder(response.Body).Decode(&result); err != nil { - c.JSON(http.StatusInternalServerError, gin.H{ + c.JSON(http.StatusOK, gin.H{ "success": false, - "message": err.Error(), + "message": fmt.Sprintf("获取模型列表失败: %s", err.Error()), }) return } - var models []string - for _, model := range result.Data { - models = append(models, model.ID) - } - c.JSON(http.StatusOK, gin.H{ "success": true, "data": models, diff --git a/controller/channel_batch_key_search_test.go b/controller/channel_batch_key_search_test.go new file mode 100644 index 000000000000..4e82d1f02d45 --- /dev/null +++ b/controller/channel_batch_key_search_test.go @@ -0,0 +1,144 @@ +package controller + +import ( + "bytes" + "encoding/json" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type channelKeySearchTestResponse struct { + Success bool `json:"success"` + Message string `json:"message"` + Data struct { + Items []model.Channel `json:"items"` + Total int `json:"total"` + TypeCounts map[int64]int64 `json:"type_counts"` + } `json:"data"` +} + +func postChannelKeySearch(t *testing.T, request BatchChannelKeySearchRequest) channelKeySearchTestResponse { + t.Helper() + + body, err := common.Marshal(request) + require.NoError(t, err) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/search/keys", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + SearchChannelsByKeys(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + var payload channelKeySearchTestResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &payload)) + return payload +} + +func seedChannelKeySearchChannels(t *testing.T) { + t.Helper() + + defaultGroup := "default,group-a" + otherGroup := "default,group-b" + channels := []model.Channel{ + {Id: 1, Type: 1, Key: "sk-match-a", Name: "needle alpha", Status: common.ChannelStatusEnabled, Group: defaultGroup, Models: "gpt-4o,gpt-4o-mini", Priority: common.GetPointer[int64](10)}, + {Id: 2, Type: 2, Key: "sk-match-b", Name: "needle beta", Status: common.ChannelStatusEnabled, Group: defaultGroup, Models: "gpt-4o,claude-3", Priority: common.GetPointer[int64](9)}, + {Id: 3, Type: 1, Key: "sk-disabled", Name: "needle disabled", Status: common.ChannelStatusManuallyDisabled, Group: defaultGroup, Models: "gpt-4o", Priority: common.GetPointer[int64](8)}, + {Id: 4, Type: 1, Key: "sk-other-group", Name: "needle other group", Status: common.ChannelStatusEnabled, Group: otherGroup, Models: "gpt-4o", Priority: common.GetPointer[int64](7)}, + {Id: 5, Type: 1, Key: "sk-not-requested", Name: "needle false positive", Status: common.ChannelStatusEnabled, Group: defaultGroup, Models: "gpt-4o", Priority: common.GetPointer[int64](6)}, + } + require.NoError(t, model.DB.Create(&channels).Error) +} + +func TestSearchChannelsByKeysExactFiltersCountsPaginationAndNoKeyLeak(t *testing.T) { + setupModelListControllerTestDB(t) + seedChannelKeySearchChannels(t) + + channelType := json.RawMessage(`"1"`) + payload := postChannelKeySearch(t, BatchChannelKeySearchRequest{ + Keys: []string{" sk-match-a ", "sk-match-b", "sk-other-group", "sk-match-a", ""}, + Keyword: "needle", + Group: "group-a", + Model: "gpt-4o", + Status: "enabled", + Type: &channelType, + SortBy: "id", + SortOrder: "asc", + Page: 1, + PageSize: 1, + }) + + require.True(t, payload.Success) + require.Equal(t, 1, payload.Data.Total) + require.Len(t, payload.Data.Items, 1) + require.Equal(t, 1, payload.Data.Items[0].Id) + require.Empty(t, payload.Data.Items[0].Key) + require.Equal(t, map[int64]int64{1: 1, 2: 1}, payload.Data.TypeCounts) +} + +func TestSearchChannelsByKeysComposesDisabledStatus(t *testing.T) { + setupModelListControllerTestDB(t) + seedChannelKeySearchChannels(t) + + payload := postChannelKeySearch(t, BatchChannelKeySearchRequest{ + Keys: []string{"sk-match-a", "sk-disabled"}, + Keyword: "needle", + Group: "group-a", + Model: "gpt-4o", + Status: "disabled", + Page: 1, + PageSize: 20, + }) + + require.True(t, payload.Success) + require.Equal(t, 1, payload.Data.Total) + require.Len(t, payload.Data.Items, 1) + require.Equal(t, 3, payload.Data.Items[0].Id) + require.Empty(t, payload.Data.Items[0].Key) + require.Equal(t, map[int64]int64{1: 1}, payload.Data.TypeCounts) +} + +func TestSearchChannelsByKeysHandlesMoreThanOneKeyChunk(t *testing.T) { + setupModelListControllerTestDB(t) + + channel := model.Channel{Id: 1001, Type: 1, Key: "sk-final-chunk", Name: "final chunk", Status: common.ChannelStatusEnabled, Group: "default", Models: "gpt-4o"} + require.NoError(t, model.DB.Create(&channel).Error) + + // The model chunks exact-key IN queries internally at 200 keys, so 250 keys crosses a chunk boundary. + keys := make([]string, 0, 250) + for i := 0; i < 249; i++ { + keys = append(keys, "sk-missing-"+common.GetRandomString(12)) + } + keys = append(keys, "sk-final-chunk") + + payload := postChannelKeySearch(t, BatchChannelKeySearchRequest{ + Keys: keys, + Page: 1, + PageSize: 20, + }) + + require.True(t, payload.Success) + require.Equal(t, 1, payload.Data.Total) + require.Len(t, payload.Data.Items, 1) + require.Equal(t, 1001, payload.Data.Items[0].Id) + require.Empty(t, payload.Data.Items[0].Key) +} + +func TestSearchChannelsByKeysRejectsTagModeV1(t *testing.T) { + setupModelListControllerTestDB(t) + + payload := postChannelKeySearch(t, BatchChannelKeySearchRequest{ + Keys: []string{"sk-any"}, + TagMode: true, + }) + + require.False(t, payload.Success) + require.Contains(t, payload.Message, "标签模式") +} diff --git a/controller/channel_preparation.go b/controller/channel_preparation.go new file mode 100644 index 000000000000..4a065f8f6bde --- /dev/null +++ b/controller/channel_preparation.go @@ -0,0 +1,513 @@ +package controller + +import ( + "fmt" + "net/http" + "strconv" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + + "github.com/gin-gonic/gin" +) + +type channelPreparationImportRequest struct { + Items []model.ChannelPreparation `json:"items"` +} + +type channelPreparationBatchRequest struct { + Ids []int `json:"ids"` +} + +type channelPreparationImportResult struct { + Index int `json:"index"` + Name string `json:"name"` + Data *model.ChannelPreparationResponse `json:"data,omitempty"` + Ok bool `json:"ok"` + Error string `json:"error,omitempty"` +} + +type channelPreparationPromoteResult struct { + Id int `json:"id"` + ChannelId int `json:"channel_id,omitempty"` + Ok bool `json:"ok"` + Error string `json:"error,omitempty"` +} + +func parseOptionalIntQuery(c *gin.Context, name string) (*int, error) { + value := strings.TrimSpace(c.Query(name)) + if value == "" { + return nil, nil + } + parsed, err := strconv.Atoi(value) + if err != nil { + return nil, err + } + return &parsed, nil +} + +func parseOptionalInt64Query(c *gin.Context, name string) (*int64, error) { + value := strings.TrimSpace(c.Query(name)) + if value == "" { + return nil, nil + } + parsed, err := strconv.ParseInt(value, 10, 64) + if err != nil { + return nil, err + } + return &parsed, nil +} + +func defaultChannelPreparationModels(channelType int) string { + if channelType == 0 { + channelType = constant.ChannelTypeAnthropic + } + models := channelId2Models[channelType] + if len(models) == 0 { + models = channelId2Models[constant.ChannelTypeAnthropic] + } + return strings.Join(models, ",") +} + +func applyChannelPreparationDefaults(preparation *model.ChannelPreparation) { + if preparation == nil { + return + } + if preparation.Type == 0 { + preparation.Type = constant.ChannelTypeAnthropic + } + if strings.TrimSpace(preparation.Models) == "" { + preparation.Models = defaultChannelPreparationModels(preparation.Type) + } +} + +func validateChannelPreparationInput(preparation *model.ChannelPreparation, isCreate bool) error { + if preparation == nil { + return fmt.Errorf("preparation cannot be empty") + } + preparation.Name = strings.TrimSpace(preparation.Name) + preparation.Key = strings.TrimSpace(preparation.Key) + if preparation.Name == "" { + return fmt.Errorf("name cannot be empty") + } + if isCreate && preparation.Key == "" { + return fmt.Errorf("key cannot be empty") + } + if strings.TrimSpace(preparation.Group) == "" { + preparation.Group = "default" + } + applyChannelPreparationDefaults(preparation) + if preparation.Remark != nil && len(*preparation.Remark) > 255 { + return fmt.Errorf("remark is too long") + } + if preparation.Setting != nil { + channel := preparation.ToChannel() + if err := channel.ValidateSettings(); err != nil { + return fmt.Errorf("渠道额外设置[channel setting] 格式错误：%s", err.Error()) + } + } + return nil +} + +func channelPreparationKeyConflictError(conflict model.ChannelPreparation) error { + statusText := "待晋升" + if conflict.Status == model.ChannelPreparationStatusPromoting { + statusText = "晋升中" + } + name := strings.TrimSpace(conflict.Name) + if name == "" { + name = "未命名" + } + return fmt.Errorf("Key 已存在于备货池%s候选渠道：%s（ID %d，%s）", statusText, conflict.KeyPreview(), conflict.Id, name) +} + +func checkChannelPreparationKeyConflict(key string, excludeID int) error { + conflicts, err := model.FindActiveChannelPreparationKeyConflicts([]string{key}, excludeID) + if err != nil { + return err + } + if conflict, ok := conflicts[strings.TrimSpace(key)]; ok { + return channelPreparationKeyConflictError(conflict) + } + return nil +} + +func GetChannelPreparations(c *gin.Context) { + page, _ := strconv.Atoi(c.Query("p")) + pageSize, _ := strconv.Atoi(c.Query("page_size")) + channelType, err := parseOptionalIntQuery(c, "type") + if err != nil { + common.ApiError(c, err) + return + } + status, err := parseOptionalIntQuery(c, "status") + if err != nil { + common.ApiError(c, err) + return + } + startTimestamp, err := parseOptionalInt64Query(c, "start_timestamp") + if err != nil { + common.ApiError(c, err) + return + } + endTimestamp, err := parseOptionalInt64Query(c, "end_timestamp") + if err != nil { + common.ApiError(c, err) + return + } + if status == nil { + pendingStatus := model.ChannelPreparationStatusPending + status = &pendingStatus + } + opts := model.ChannelPreparationListOptions{ + Page: page, + PageSize: pageSize, + Keyword: c.Query("keyword"), + Group: c.Query("group"), + Type: channelType, + Status: status, + StartTimestamp: startTimestamp, + EndTimestamp: endTimestamp, + IDSort: c.Query("id_sort") == "true" || c.Query("id_sort") == "1", + } + preparations, total, stats, statusCounts, typeCounts, err := model.GetChannelPreparations(opts) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{ + "items": model.ChannelPreparationResponses(preparations), + "total": total, + "page": opts.Page, + "page_size": opts.PageSize, + "stats": stats, + "status_counts": statusCounts, + "type_counts": typeCounts, + }) +} + +func GetChannelPreparation(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + var preparation model.ChannelPreparation + if err := model.DB.First(&preparation, "id = ?", id).Error; err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, preparation.ToResponse()) +} + +func TestChannelPreparation(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + var preparation model.ChannelPreparation + if err := model.DB.First(&preparation, "id = ?", id).Error; err != nil { + common.ApiError(c, err) + return + } + applyChannelPreparationDefaults(&preparation) + channel := preparation.ToChannel() + testModel := strings.TrimSpace(c.Query("model")) + endpointType := c.Query("endpoint_type") + isStream, _ := strconv.ParseBool(c.Query("stream")) + testUserID, err := resolveChannelTestUserID(c) + if err != nil { + common.ApiError(c, err) + return + } + tik := time.Now() + result := testChannel(channel, testUserID, testModel, endpointType, isStream) + milliseconds := time.Since(tik).Milliseconds() + consumedTime := float64(milliseconds) / 1000.0 + if result.localErr != nil { + message := result.localErr.Error() + go preparation.UpdateTestResult(milliseconds, model.ChannelPreparationTestStatusFailed, message) + resp := gin.H{ + "success": false, + "message": message, + "time": consumedTime, + } + if result.newAPIError != nil { + resp["error_code"] = result.newAPIError.GetErrorCode() + } + c.JSON(http.StatusOK, resp) + return + } + if result.newAPIError != nil { + message := result.newAPIError.Error() + go preparation.UpdateTestResult(milliseconds, model.ChannelPreparationTestStatusFailed, message) + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": message, + "time": consumedTime, + "error_code": result.newAPIError.GetErrorCode(), + }) + return + } + go preparation.UpdateResponseTime(milliseconds) + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "time": consumedTime, + }) +} + +func AddChannelPreparation(c *gin.Context) { + var preparation model.ChannelPreparation + if err := c.ShouldBindJSON(&preparation); err != nil { + common.ApiError(c, err) + return + } + if err := validateChannelPreparationInput(&preparation, true); err != nil { + common.ApiError(c, err) + return + } + if err := checkChannelPreparationKeyConflict(preparation.Key, 0); err != nil { + common.ApiError(c, err) + return + } + preparation.NormalizeForCreate() + if err := model.DB.Create(&preparation).Error; err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, preparation.ToResponse()) +} + +func UpdateChannelPreparation(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + var input model.ChannelPreparation + if err := c.ShouldBindJSON(&input); err != nil { + common.ApiError(c, err) + return + } + var existing model.ChannelPreparation + if err := model.DB.First(&existing, "id = ?", id).Error; err != nil { + common.ApiError(c, err) + return + } + if existing.Status != model.ChannelPreparationStatusPending { + common.ApiErrorMsg(c, "只有待晋升的候选渠道可以编辑") + return + } + input.NormalizeForUpdate(&existing) + if err := validateChannelPreparationInput(&input, false); err != nil { + common.ApiError(c, err) + return + } + if err := checkChannelPreparationKeyConflict(input.Key, existing.Id); err != nil { + common.ApiError(c, err) + return + } + if err := model.DB.Save(&input).Error; err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, input.ToResponse()) +} + +func DeleteChannelPreparation(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + result := model.DB.Delete(&model.ChannelPreparation{}, "id = ?", id) + if result.Error != nil { + common.ApiError(c, result.Error) + return + } + if result.RowsAffected == 0 { + common.ApiErrorMsg(c, "候选渠道不存在") + return + } + common.ApiSuccess(c, gin.H{"id": id}) +} + +func ImportChannelPreparations(c *gin.Context) { + var request channelPreparationImportRequest + if err := c.ShouldBindJSON(&request); err != nil { + common.ApiError(c, err) + return + } + + resultsByIndex := make([]channelPreparationImportResult, len(request.Items)) + resultSet := make([]bool, len(request.Items)) + normalizedItems := make([]model.ChannelPreparation, len(request.Items)) + validIndexes := make([]int, 0, len(request.Items)) + validKeys := make([]string, 0, len(request.Items)) + + for index, item := range request.Items { + if strings.TrimSpace(item.Source) == "" { + item.Source = "batch_import" + } + if err := validateChannelPreparationInput(&item, true); err != nil { + resultsByIndex[index] = channelPreparationImportResult{Index: index, Name: item.Name, Ok: false, Error: err.Error()} + resultSet[index] = true + continue + } + item.NormalizeForCreate() + normalizedItems[index] = item + validIndexes = append(validIndexes, index) + validKeys = append(validKeys, item.Key) + } + + dbConflicts, err := model.FindActiveChannelPreparationKeyConflicts(validKeys, 0) + if err != nil { + common.ApiError(c, err) + return + } + + seenImportKeys := make(map[string]int, len(validIndexes)) + for _, index := range validIndexes { + item := normalizedItems[index] + key := strings.TrimSpace(item.Key) + if firstIndex, ok := seenImportKeys[key]; ok { + resultsByIndex[index] = channelPreparationImportResult{ + Index: index, + Name: item.Name, + Ok: false, + Error: fmt.Sprintf("本次导入重复：第 %d 条已包含相同 Key", firstIndex+1), + } + resultSet[index] = true + continue + } + seenImportKeys[key] = index + + if conflict, ok := dbConflicts[key]; ok { + resultsByIndex[index] = channelPreparationImportResult{Index: index, Name: item.Name, Ok: false, Error: channelPreparationKeyConflictError(conflict).Error()} + resultSet[index] = true + continue + } + if err := model.DB.Create(&item).Error; err != nil { + resultsByIndex[index] = channelPreparationImportResult{Index: index, Name: item.Name, Ok: false, Error: err.Error()} + resultSet[index] = true + continue + } + response := item.ToResponse() + resultsByIndex[index] = channelPreparationImportResult{Index: index, Name: item.Name, Data: &response, Ok: true} + resultSet[index] = true + } + + results := make([]channelPreparationImportResult, 0, len(request.Items)) + for index := range request.Items { + if resultSet[index] { + results = append(results, resultsByIndex[index]) + } + } + common.ApiSuccess(c, gin.H{"results": results}) +} + +func promoteChannelPreparation(id int) (int, error) { + tx := model.DB.Begin() + if tx.Error != nil { + return 0, tx.Error + } + defer func() { + if r := recover(); r != nil { + tx.Rollback() + } + }() + + now := common.GetTimestamp() + lockResult := tx.Model(&model.ChannelPreparation{}). + Where("id = ? AND status = ?", id, model.ChannelPreparationStatusPending). + Updates(map[string]any{ + "status": model.ChannelPreparationStatusPromoting, + "updated_time": now, + }) + if lockResult.Error != nil { + tx.Rollback() + return 0, lockResult.Error + } + if lockResult.RowsAffected == 0 { + tx.Rollback() + return 0, fmt.Errorf("候选渠道不存在或不可晋升") + } + + var preparation model.ChannelPreparation + if err := tx.First(&preparation, "id = ?", id).Error; err != nil { + tx.Rollback() + return 0, err + } + applyChannelPreparationDefaults(&preparation) + channel := preparation.ToChannel() + channels, err := createChannelsFromAddRequest(&AddChannelRequest{Mode: "single", Channel: channel}, tx) + if err != nil { + tx.Rollback() + return 0, err + } + if len(channels) == 0 { + tx.Rollback() + return 0, fmt.Errorf("channel cannot be empty") + } + channelID := channels[0].Id + deleteResult := tx.Where("id = ? AND status = ?", id, model.ChannelPreparationStatusPromoting). + Delete(&model.ChannelPreparation{}) + if deleteResult.Error != nil { + tx.Rollback() + return 0, deleteResult.Error + } + if deleteResult.RowsAffected == 0 { + tx.Rollback() + return 0, fmt.Errorf("候选渠道晋升后删除失败") + } + if err := tx.Commit().Error; err != nil { + return 0, err + } + return channelID, nil +} + +func PromoteChannelPreparation(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil { + common.ApiError(c, err) + return + } + channelID, err := promoteChannelPreparation(id) + if err != nil { + common.ApiError(c, err) + return + } + model.InitChannelCache() + service.ResetProxyClientCache() + common.ApiSuccess(c, gin.H{"id": id, "channel_id": channelID}) +} + +func PromoteChannelPreparationsBatch(c *gin.Context) { + var request channelPreparationBatchRequest + if err := c.ShouldBindJSON(&request); err != nil { + common.ApiError(c, err) + return + } + results := make([]channelPreparationPromoteResult, 0, len(request.Ids)) + succeeded := false + for _, id := range request.Ids { + channelID, err := promoteChannelPreparation(id) + if err != nil { + results = append(results, channelPreparationPromoteResult{Id: id, Ok: false, Error: err.Error()}) + continue + } + succeeded = true + results = append(results, channelPreparationPromoteResult{Id: id, ChannelId: channelID, Ok: true}) + } + if succeeded { + model.InitChannelCache() + service.ResetProxyClientCache() + } + common.ApiSuccess(c, gin.H{"results": results}) +} diff --git a/controller/channel_preparation_auto_promotion.go b/controller/channel_preparation_auto_promotion.go new file mode 100644 index 000000000000..d51cada07bbf --- /dev/null +++ b/controller/channel_preparation_auto_promotion.go @@ -0,0 +1,705 @@ +package controller + +import ( + "fmt" + "io" + "math" + "math/rand" + "sort" + "strings" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/operation_setting" + + "github.com/gin-gonic/gin" +) + +const channelPreparationAutoPromotionTriggerManual = "manual" +const channelPreparationAutoPromotionTriggerScheduler = "scheduler" + +const ( + channelPreparationAutoPromotionShortageCount = "count" + channelPreparationAutoPromotionShortageCapacity = "capacity" +) + +type channelPreparationAutoPromotionRunRequest struct { + RuleId string `json:"rule_id"` +} + +type channelPreparationAutoPromotionCapacitySummary struct { + EligibleChannelCount int64 `json:"eligible_channel_count"` + UsableChannelCount int64 `json:"usable_channel_count"` + IgnoredNonPositiveBalanceChannelCount int64 `json:"ignored_non_positive_balance_channel_count"` + BalanceSumUSD float64 `json:"balance_sum_usd"` + UsedQuotaUSD float64 `json:"used_quota_usd"` + EffectiveCapacityUSD float64 `json:"effective_capacity_usd"` + RawEffectiveCapacityUSD float64 `json:"raw_effective_capacity_usd"` +} + +type channelPreparationAutoPromotionStep struct { + PreparationId int `json:"preparation_id"` + ChannelId int `json:"channel_id"` + CandidateBalanceUSD float64 `json:"candidate_balance_usd"` + CapacityBeforeUSD float64 `json:"capacity_before_usd"` + CapacityAfterUSD float64 `json:"capacity_after_usd"` + ShortageType string `json:"shortage_type"` + Strategy string `json:"strategy"` + UsableCountBefore int64 `json:"usable_count_before"` + UsableCountAfter int64 `json:"usable_count_after"` + CountDeficitBefore int64 `json:"count_deficit_before"` + CountDeficitAfter int64 `json:"count_deficit_after"` + CapacityDeficitBeforeUSD float64 `json:"capacity_deficit_before_usd"` + CapacityDeficitAfterUSD float64 `json:"capacity_deficit_after_usd"` +} + +type channelPreparationAutoPromotionRuleSummary struct { + Trigger string `json:"trigger"` + RuleId string `json:"rule_id"` + Group string `json:"group"` + Type int `json:"type"` + Strategy string `json:"strategy"` + ThresholdUSD float64 `json:"threshold_usd"` + MinimumUsableChannelCount int `json:"minimum_usable_channel_count"` + GuaranteePriority string `json:"guarantee_priority"` + CountShortageStrategy string `json:"count_shortage_strategy"` + CapacityShortageStrategy string `json:"capacity_shortage_strategy"` + InitialCapacity channelPreparationAutoPromotionCapacitySummary `json:"initial_capacity"` + FinalCapacity channelPreparationAutoPromotionCapacitySummary `json:"final_capacity"` + Promotions []channelPreparationAutoPromotionStep `json:"promotions"` + Failures []string `json:"failures"` + SkippedReason string `json:"skipped_reason,omitempty"` + RemainingDeficitUSD float64 `json:"remaining_deficit_usd"` + CountDeficit int64 `json:"count_deficit"` + CapacityDeficitUSD float64 `json:"capacity_deficit_usd"` + LimitReached bool `json:"limit_reached"` +} + +type channelPreparationAutoPromotionRunSummary struct { + Trigger string `json:"trigger"` + RuleId string `json:"rule_id,omitempty"` + StartedAt int64 `json:"started_at"` + FinishedAt int64 `json:"finished_at"` + MaxPromotions int `json:"max_promotions"` + TotalPromoted int `json:"total_promoted"` + LimitReached bool `json:"limit_reached"` + Rules []channelPreparationAutoPromotionRuleSummary `json:"rules"` + SkippedReason string `json:"skipped_reason,omitempty"` +} + +type channelPreparationAutoPromotionSchedulerStatus struct { + SchedulerEnabled bool `json:"scheduler_enabled"` + IntervalMinutes float64 `json:"interval_minutes"` + NextCheckAt int64 `json:"next_check_at"` + LastCheckAt int64 `json:"last_check_at"` + LastFinishedAt int64 `json:"last_finished_at"` + LastPromoted int `json:"last_promoted"` + Running bool `json:"running"` + IsMasterNode bool `json:"is_master_node"` + ServerTimestamp int64 `json:"server_timestamp"` +} + +type channelPreparationAutoPromotionCapacityAggregate struct { + EligibleChannelCount int64 `gorm:"column:eligible_channel_count"` + BalanceSumUSD float64 `gorm:"column:balance_sum_usd"` + UsedQuotaSum int64 `gorm:"column:used_quota_sum"` +} + +var ( + channelPreparationAutoPromotionRunMutex sync.Mutex + channelPreparationAutoPromotionTaskOnce sync.Once + channelPreparationAutoPromotionStatusMutex sync.RWMutex + channelPreparationAutoPromotionStatusSnapshot channelPreparationAutoPromotionSchedulerStatus +) + +func updateChannelPreparationAutoPromotionSchedulerStatus(update func(*channelPreparationAutoPromotionSchedulerStatus)) { + channelPreparationAutoPromotionStatusMutex.Lock() + defer channelPreparationAutoPromotionStatusMutex.Unlock() + update(&channelPreparationAutoPromotionStatusSnapshot) +} + +func getChannelPreparationAutoPromotionSchedulerStatus() channelPreparationAutoPromotionSchedulerStatus { + channelPreparationAutoPromotionStatusMutex.RLock() + status := channelPreparationAutoPromotionStatusSnapshot + channelPreparationAutoPromotionStatusMutex.RUnlock() + + setting := operation_setting.GetChannelPreparationAutoPromotionSetting() + status.SchedulerEnabled = setting.SchedulerEnabled + status.IntervalMinutes = setting.IntervalMinutes + status.IsMasterNode = common.IsMasterNode + status.ServerTimestamp = common.GetTimestamp() + if !setting.SchedulerEnabled && !status.Running { + status.NextCheckAt = 0 + } + if setting.SchedulerEnabled && common.IsMasterNode && !status.Running && status.NextCheckAt == 0 { + intervalMinutes := int(math.Round(setting.IntervalMinutes)) + if intervalMinutes <= 0 { + intervalMinutes = 10 + } + status.NextCheckAt = time.Now().Add(time.Duration(intervalMinutes) * time.Minute).Unix() + } + return status +} + +func GetChannelPreparationAutoPromotionSchedulerStatus(c *gin.Context) { + common.ApiSuccess(c, getChannelPreparationAutoPromotionSchedulerStatus()) +} + +func normalizeAutoPromotionDeficit(threshold float64, capacity float64) float64 { + deficit := threshold - capacity + if deficit < 0 { + return 0 + } + return deficit +} + +func channelPreparationAutoPromotionCountDeficit(minimum int, usable int64) int64 { + if minimum <= 0 { + return 0 + } + deficit := int64(minimum) - usable + if deficit < 0 { + return 0 + } + return deficit +} + +func chooseChannelPreparationAutoPromotionActiveShortage(rule operation_setting.ChannelPreparationAutoPromotionRule, countShort bool, capacityShort bool) string { + if countShort && capacityShort { + if rule.GuaranteePriority == operation_setting.ChannelPreparationAutoPromotionGuaranteePriorityCountFirst { + return channelPreparationAutoPromotionShortageCount + } + return channelPreparationAutoPromotionShortageCapacity + } + if countShort { + return channelPreparationAutoPromotionShortageCount + } + if capacityShort { + return channelPreparationAutoPromotionShortageCapacity + } + return "" +} + +func channelPreparationAutoPromotionStrategyForShortage(rule operation_setting.ChannelPreparationAutoPromotionRule, shortageType string) string { + if shortageType == channelPreparationAutoPromotionShortageCount { + return rule.CountShortageStrategy + } + return rule.CapacityShortageStrategy +} + +func safeQuotaToUSD(usedQuota int64) float64 { + if common.QuotaPerUnit <= 0 { + return 0 + } + return float64(usedQuota) / common.QuotaPerUnit +} + +func computeChannelPreparationAutoPromotionCapacity(group string, channelType int) (channelPreparationAutoPromotionCapacitySummary, error) { + query := model.DB.Model(&model.Channel{}) + query = model.ApplyChannelGroupFilter(query, group) + query = query.Where("status = ?", common.ChannelStatusEnabled). + Where("type = ?", channelType). + Where("balance > ?", 0) + + var aggregate channelPreparationAutoPromotionCapacityAggregate + if err := query.Select("COUNT(*) AS eligible_channel_count, COALESCE(SUM(balance), 0) AS balance_sum_usd, COALESCE(SUM(used_quota), 0) AS used_quota_sum").Scan(&aggregate).Error; err != nil { + return channelPreparationAutoPromotionCapacitySummary{}, err + } + + ignoredQuery := model.DB.Model(&model.Channel{}) + ignoredQuery = model.ApplyChannelGroupFilter(ignoredQuery, group) + ignoredQuery = ignoredQuery.Where("status = ?", common.ChannelStatusEnabled). + Where("type = ?", channelType). + Where("balance <= ?", 0) + var ignoredCount int64 + if err := ignoredQuery.Count(&ignoredCount).Error; err != nil { + return channelPreparationAutoPromotionCapacitySummary{}, err + } + + usedQuotaUSD := safeQuotaToUSD(aggregate.UsedQuotaSum) + rawCapacity := aggregate.BalanceSumUSD - usedQuotaUSD + capacity := rawCapacity + if capacity < 0 { + capacity = 0 + } + return channelPreparationAutoPromotionCapacitySummary{ + EligibleChannelCount: aggregate.EligibleChannelCount, + UsableChannelCount: aggregate.EligibleChannelCount, + IgnoredNonPositiveBalanceChannelCount: ignoredCount, + BalanceSumUSD: aggregate.BalanceSumUSD, + UsedQuotaUSD: usedQuotaUSD, + EffectiveCapacityUSD: capacity, + RawEffectiveCapacityUSD: rawCapacity, + }, nil +} + +func loadChannelPreparationAutoPromotionCandidates(group string, channelType int, excludedIds map[int]bool) ([]model.ChannelPreparation, error) { + query := model.DB.Model(&model.ChannelPreparation{}) + query = model.ApplyChannelGroupFilter(query, group) + query = query.Where("status = ?", model.ChannelPreparationStatusPending). + Where("type = ?", channelType). + Where("balance > ?", 0) + if len(excludedIds) > 0 { + ids := make([]int, 0, len(excludedIds)) + for id := range excludedIds { + ids = append(ids, id) + } + query = query.Where("id NOT IN ?", ids) + } + + var preparations []model.ChannelPreparation + if err := query.Order("priority DESC, id ASC").Find(&preparations).Error; err != nil { + return nil, err + } + return preparations, nil +} + +func preparationPriority(preparation model.ChannelPreparation) int64 { + if preparation.Priority == nil { + return 0 + } + return *preparation.Priority +} + +func preparationWeight(preparation model.ChannelPreparation) int64 { + weight := int64(0) + if preparation.Weight != nil { + weight = int64(*preparation.Weight) + } + return weight + 10 +} + +func channelPreparationAutoPromotionHighestPriorityTier(preparations []model.ChannelPreparation) []model.ChannelPreparation { + if len(preparations) == 0 { + return nil + } + sortedPreparations := append([]model.ChannelPreparation(nil), preparations...) + sort.SliceStable(sortedPreparations, func(i, j int) bool { + pi := preparationPriority(sortedPreparations[i]) + pj := preparationPriority(sortedPreparations[j]) + if pi == pj { + return sortedPreparations[i].Id < sortedPreparations[j].Id + } + return pi > pj + }) + topPriority := preparationPriority(sortedPreparations[0]) + tier := make([]model.ChannelPreparation, 0) + for _, preparation := range sortedPreparations { + if preparationPriority(preparation) != topPriority { + break + } + tier = append(tier, preparation) + } + return tier +} + +func chooseChannelPreparationAutoPromotionWeightedCandidate(tier []model.ChannelPreparation, rng *rand.Rand) (model.ChannelPreparation, bool) { + if len(tier) == 0 { + return model.ChannelPreparation{}, false + } + if len(tier) == 1 { + return tier[0], true + } + totalWeight := int64(0) + for _, preparation := range tier { + weight := preparationWeight(preparation) + if weight > 0 { + totalWeight += weight + } + } + if totalWeight <= 0 { + return tier[0], true + } + if rng == nil { + rng = rand.New(rand.NewSource(time.Now().UnixNano())) + } + pick := rng.Int63n(totalWeight) + for _, preparation := range tier { + weight := preparationWeight(preparation) + if weight <= 0 { + continue + } + if pick < weight { + return preparation, true + } + pick -= weight + } + return tier[len(tier)-1], true +} + +func chooseChannelPreparationAutoPromotionCandidate(preparations []model.ChannelPreparation, strategy string, rng *rand.Rand) (model.ChannelPreparation, bool) { + tier := channelPreparationAutoPromotionHighestPriorityTier(preparations) + if len(tier) == 0 { + return model.ChannelPreparation{}, false + } + + switch strategy { + case operation_setting.ChannelPreparationAutoPromotionStrategySmallBalanceFirst: + sort.SliceStable(tier, func(i, j int) bool { + if tier[i].Balance == tier[j].Balance { + return tier[i].Id < tier[j].Id + } + return tier[i].Balance < tier[j].Balance + }) + return tier[0], true + case operation_setting.ChannelPreparationAutoPromotionStrategyLargeBalanceFirst: + sort.SliceStable(tier, func(i, j int) bool { + if tier[i].Balance == tier[j].Balance { + return tier[i].Id < tier[j].Id + } + return tier[i].Balance > tier[j].Balance + }) + return tier[0], true + case operation_setting.ChannelPreparationAutoPromotionStrategyPriorityWeighted: + fallthrough + default: + return chooseChannelPreparationAutoPromotionWeightedCandidate(tier, rng) + } +} + +func normalizeChannelPreparationAutoPromotionRules(rules []operation_setting.ChannelPreparationAutoPromotionRule) []operation_setting.ChannelPreparationAutoPromotionRule { + normalized := make([]operation_setting.ChannelPreparationAutoPromotionRule, 0, len(rules)) + for _, rule := range rules { + operation_setting.NormalizeChannelPreparationAutoPromotionRule(&rule) + normalized = append(normalized, rule) + } + return normalized +} + +func recordChannelPreparationAutoPromotionManageLog(adminUserId *int, content string, channelId int, group string, adminInfo map[string]interface{}) { + logUserId := 0 + actor := "system" + if adminUserId != nil && *adminUserId > 0 { + logUserId = *adminUserId + actor = "admin" + } + + enrichedInfo := make(map[string]interface{}, len(adminInfo)+5) + for key, value := range adminInfo { + enrichedInfo[key] = value + } + enrichedInfo["event"] = "channel_preparation_auto_promotion" + enrichedInfo["actor"] = actor + enrichedInfo["node_name"] = common.NodeName + enrichedInfo["server_ip"] = common.GetIp() + enrichedInfo["version"] = common.Version + + model.RecordLogWithAdminInfoAndMetadata(logUserId, model.LogTypeManage, content, channelId, group, enrichedInfo) +} + +func runChannelPreparationAutoPromotionLocked(trigger string, optionalRuleId string, adminUserId *int) (channelPreparationAutoPromotionRunSummary, error) { + settingSnapshot := *operation_setting.GetChannelPreparationAutoPromotionSetting() + settingSnapshot.Rules = normalizeChannelPreparationAutoPromotionRules(settingSnapshot.Rules) + maxPromotions := settingSnapshot.MaxPromotionsPerRun + if maxPromotions <= 0 { + maxPromotions = 10 + } + + summary := channelPreparationAutoPromotionRunSummary{ + Trigger: trigger, + RuleId: strings.TrimSpace(optionalRuleId), + StartedAt: common.GetTimestamp(), + MaxPromotions: maxPromotions, + Rules: []channelPreparationAutoPromotionRuleSummary{}, + } + + if len(settingSnapshot.Rules) == 0 { + summary.SkippedReason = "没有配置自动晋升规则" + summary.FinishedAt = common.GetTimestamp() + return summary, nil + } + + rng := rand.New(rand.NewSource(time.Now().UnixNano())) + promotedAny := false + + for _, rule := range settingSnapshot.Rules { + if summary.TotalPromoted >= maxPromotions { + summary.LimitReached = true + break + } + if summary.RuleId != "" && rule.Id != summary.RuleId { + continue + } + + ruleSummary := channelPreparationAutoPromotionRuleSummary{ + Trigger: trigger, + RuleId: rule.Id, + Group: rule.Group, + Type: rule.Type, + Strategy: rule.Strategy, + ThresholdUSD: rule.ThresholdUSD, + MinimumUsableChannelCount: rule.MinimumUsableChannelCount, + GuaranteePriority: rule.GuaranteePriority, + CountShortageStrategy: rule.CountShortageStrategy, + CapacityShortageStrategy: rule.CapacityShortageStrategy, + Promotions: []channelPreparationAutoPromotionStep{}, + Failures: []string{}, + } + + if !rule.Enabled { + ruleSummary.SkippedReason = "规则未启用" + summary.Rules = append(summary.Rules, ruleSummary) + continue + } + if strings.TrimSpace(rule.Group) == "" { + ruleSummary.SkippedReason = "规则分组为空" + summary.Rules = append(summary.Rules, ruleSummary) + continue + } + if rule.Type <= 0 { + ruleSummary.SkippedReason = "渠道类型无效" + summary.Rules = append(summary.Rules, ruleSummary) + continue + } + if rule.ThresholdUSD <= 0 { + ruleSummary.SkippedReason = "阈值必须大于 0" + summary.Rules = append(summary.Rules, ruleSummary) + continue + } + if rule.MinimumUsableChannelCount < 0 { + ruleSummary.SkippedReason = "最低可用渠道数不能小于 0" + summary.Rules = append(summary.Rules, ruleSummary) + continue + } + if !operation_setting.IsSupportedChannelPreparationAutoPromotionGuaranteePriority(rule.GuaranteePriority) || + !operation_setting.IsSupportedChannelPreparationAutoPromotionStrategy(rule.CountShortageStrategy) || + !operation_setting.IsSupportedChannelPreparationAutoPromotionStrategy(rule.CapacityShortageStrategy) { + ruleSummary.SkippedReason = "自动晋升规则配置不支持" + summary.Rules = append(summary.Rules, ruleSummary) + continue + } + + capacity, err := computeChannelPreparationAutoPromotionCapacity(rule.Group, rule.Type) + if err != nil { + ruleSummary.Failures = append(ruleSummary.Failures, err.Error()) + summary.Rules = append(summary.Rules, ruleSummary) + continue + } + ruleSummary.InitialCapacity = capacity + ruleSummary.FinalCapacity = capacity + ruleSummary.CountDeficit = channelPreparationAutoPromotionCountDeficit(rule.MinimumUsableChannelCount, capacity.UsableChannelCount) + ruleSummary.CapacityDeficitUSD = normalizeAutoPromotionDeficit(rule.ThresholdUSD, capacity.EffectiveCapacityUSD) + if ruleSummary.CountDeficit == 0 && ruleSummary.CapacityDeficitUSD == 0 { + ruleSummary.SkippedReason = "容量和可用渠道数均已达标" + ruleSummary.RemainingDeficitUSD = 0 + summary.Rules = append(summary.Rules, ruleSummary) + continue + } + + failedCandidateIds := make(map[int]bool) + for summary.TotalPromoted < maxPromotions { + latestCapacity, err := computeChannelPreparationAutoPromotionCapacity(rule.Group, rule.Type) + if err != nil { + ruleSummary.Failures = append(ruleSummary.Failures, err.Error()) + break + } + ruleSummary.FinalCapacity = latestCapacity + countDeficitBefore := channelPreparationAutoPromotionCountDeficit(rule.MinimumUsableChannelCount, latestCapacity.UsableChannelCount) + capacityDeficitBefore := normalizeAutoPromotionDeficit(rule.ThresholdUSD, latestCapacity.EffectiveCapacityUSD) + shortageType := chooseChannelPreparationAutoPromotionActiveShortage(rule, countDeficitBefore > 0, capacityDeficitBefore > 0) + if shortageType == "" { + break + } + strategy := channelPreparationAutoPromotionStrategyForShortage(rule, shortageType) + if !operation_setting.IsSupportedChannelPreparationAutoPromotionStrategy(strategy) { + strategy = operation_setting.ChannelPreparationAutoPromotionStrategyPriorityWeighted + } + + candidates, err := loadChannelPreparationAutoPromotionCandidates(rule.Group, rule.Type, failedCandidateIds) + if err != nil { + ruleSummary.Failures = append(ruleSummary.Failures, err.Error()) + break + } + candidate, ok := chooseChannelPreparationAutoPromotionCandidate(candidates, strategy, rng) + if !ok { + ruleSummary.SkippedReason = "没有余额大于 0 的待晋升候选渠道" + break + } + + channelId, err := promoteChannelPreparation(candidate.Id) + if err != nil { + failedCandidateIds[candidate.Id] = true + ruleSummary.Failures = append(ruleSummary.Failures, fmt.Sprintf("候选渠道 %d 晋升失败：%s", candidate.Id, err.Error())) + continue + } + promotedAny = true + summary.TotalPromoted++ + afterCapacity, capacityErr := computeChannelPreparationAutoPromotionCapacity(rule.Group, rule.Type) + if capacityErr != nil { + ruleSummary.Failures = append(ruleSummary.Failures, fmt.Sprintf("候选渠道 %d 晋升后重新计算容量失败：%s", candidate.Id, capacityErr.Error())) + afterCapacity = latestCapacity + afterCapacity.EligibleChannelCount++ + afterCapacity.UsableChannelCount++ + afterCapacity.BalanceSumUSD += math.Max(candidate.Balance, 0) + afterCapacity.RawEffectiveCapacityUSD += math.Max(candidate.Balance, 0) + afterCapacity.EffectiveCapacityUSD = math.Max(afterCapacity.RawEffectiveCapacityUSD, 0) + } + ruleSummary.FinalCapacity = afterCapacity + countDeficitAfter := channelPreparationAutoPromotionCountDeficit(rule.MinimumUsableChannelCount, afterCapacity.UsableChannelCount) + capacityDeficitAfter := normalizeAutoPromotionDeficit(rule.ThresholdUSD, afterCapacity.EffectiveCapacityUSD) + ruleSummary.Promotions = append(ruleSummary.Promotions, channelPreparationAutoPromotionStep{ + PreparationId: candidate.Id, + ChannelId: channelId, + CandidateBalanceUSD: candidate.Balance, + CapacityBeforeUSD: latestCapacity.EffectiveCapacityUSD, + CapacityAfterUSD: afterCapacity.EffectiveCapacityUSD, + ShortageType: shortageType, + Strategy: strategy, + UsableCountBefore: latestCapacity.UsableChannelCount, + UsableCountAfter: afterCapacity.UsableChannelCount, + CountDeficitBefore: countDeficitBefore, + CountDeficitAfter: countDeficitAfter, + CapacityDeficitBeforeUSD: capacityDeficitBefore, + CapacityDeficitAfterUSD: capacityDeficitAfter, + }) + logContent := fmt.Sprintf("自动晋升候选渠道：规则=%s 分组=%s 类型=%d 不足=%s 策略=%s 候选ID=%d 渠道ID=%d 余额=%.4f 可用数 %d -> %d 容量 %.4f -> %.4f 缺口(count=%d->%d, capacity=%.4f->%.4f) 触发=%s", rule.Id, rule.Group, rule.Type, shortageType, strategy, candidate.Id, channelId, candidate.Balance, latestCapacity.UsableChannelCount, afterCapacity.UsableChannelCount, latestCapacity.EffectiveCapacityUSD, afterCapacity.EffectiveCapacityUSD, countDeficitBefore, countDeficitAfter, capacityDeficitBefore, capacityDeficitAfter, trigger) + common.SysLog(logContent) + recordChannelPreparationAutoPromotionManageLog(adminUserId, logContent, channelId, rule.Group, map[string]interface{}{ + "rule_id": rule.Id, + "group": rule.Group, + "type": rule.Type, + "preparation_id": candidate.Id, + "channel_id": channelId, + "candidate_balance": candidate.Balance, + "shortage_type": shortageType, + "strategy": strategy, + "usable_count_before": latestCapacity.UsableChannelCount, + "usable_count_after": afterCapacity.UsableChannelCount, + "capacity_before": latestCapacity.EffectiveCapacityUSD, + "capacity_after": afterCapacity.EffectiveCapacityUSD, + "count_deficit_before": countDeficitBefore, + "count_deficit_after": countDeficitAfter, + "capacity_deficit_before_usd": capacityDeficitBefore, + "capacity_deficit_after_usd": capacityDeficitAfter, + "trigger": trigger, + }) + } + + ruleSummary.CountDeficit = channelPreparationAutoPromotionCountDeficit(rule.MinimumUsableChannelCount, ruleSummary.FinalCapacity.UsableChannelCount) + ruleSummary.CapacityDeficitUSD = normalizeAutoPromotionDeficit(rule.ThresholdUSD, ruleSummary.FinalCapacity.EffectiveCapacityUSD) + if summary.TotalPromoted >= maxPromotions && (ruleSummary.CountDeficit > 0 || ruleSummary.CapacityDeficitUSD > 0) { + ruleSummary.LimitReached = true + summary.LimitReached = true + } + ruleSummary.RemainingDeficitUSD = ruleSummary.CapacityDeficitUSD + summary.Rules = append(summary.Rules, ruleSummary) + } + + if summary.RuleId != "" && len(summary.Rules) == 0 { + summary.SkippedReason = "未找到指定规则" + } + if promotedAny { + model.InitChannelCache() + service.ResetProxyClientCache() + } + summary.FinishedAt = common.GetTimestamp() + return summary, nil +} + +func RunChannelPreparationAutoPromotion(trigger string, optionalRuleId string, adminUserId *int) (channelPreparationAutoPromotionRunSummary, error) { + if trigger == "" { + trigger = channelPreparationAutoPromotionTriggerManual + } + if !channelPreparationAutoPromotionRunMutex.TryLock() { + return channelPreparationAutoPromotionRunSummary{}, fmt.Errorf("自动晋升正在执行中") + } + defer channelPreparationAutoPromotionRunMutex.Unlock() + return runChannelPreparationAutoPromotionLocked(trigger, optionalRuleId, adminUserId) +} + +func RunChannelPreparationAutoPromotionManually(c *gin.Context) { + var request channelPreparationAutoPromotionRunRequest + if err := c.ShouldBindJSON(&request); err != nil && err != io.EOF { + common.ApiError(c, err) + return + } + adminUserId := c.GetInt("id") + summary, err := RunChannelPreparationAutoPromotion(channelPreparationAutoPromotionTriggerManual, request.RuleId, &adminUserId) + if err != nil { + common.ApiErrorMsg(c, err.Error()) + return + } + recordChannelPreparationAutoPromotionManageLog(&adminUserId, fmt.Sprintf("手动执行渠道备货池自动晋升：晋升 %d 个渠道", summary.TotalPromoted), 0, "", map[string]interface{}{ + "rule_id": request.RuleId, + "total_promoted": summary.TotalPromoted, + "limit_reached": summary.LimitReached, + "trigger": channelPreparationAutoPromotionTriggerManual, + }) + common.ApiSuccess(c, summary) +} + +func StartChannelPreparationAutoPromotionTask() { + channelPreparationAutoPromotionTaskOnce.Do(func() { + if !common.IsMasterNode { + return + } + go func() { + common.SysLog("channel preparation auto promotion task started") + for { + setting := operation_setting.GetChannelPreparationAutoPromotionSetting() + if !setting.SchedulerEnabled { + updateChannelPreparationAutoPromotionSchedulerStatus(func(status *channelPreparationAutoPromotionSchedulerStatus) { + status.SchedulerEnabled = false + status.IntervalMinutes = setting.IntervalMinutes + status.NextCheckAt = 0 + status.Running = false + }) + time.Sleep(1 * time.Minute) + continue + } + intervalMinutes := int(math.Round(setting.IntervalMinutes)) + if intervalMinutes <= 0 { + intervalMinutes = 10 + } + intervalDuration := time.Duration(intervalMinutes) * time.Minute + nextCheckAt := time.Now().Add(intervalDuration).Unix() + updateChannelPreparationAutoPromotionSchedulerStatus(func(status *channelPreparationAutoPromotionSchedulerStatus) { + status.SchedulerEnabled = true + status.IntervalMinutes = float64(intervalMinutes) + status.NextCheckAt = nextCheckAt + status.Running = false + }) + time.Sleep(intervalDuration) + if !operation_setting.GetChannelPreparationAutoPromotionSetting().SchedulerEnabled { + updateChannelPreparationAutoPromotionSchedulerStatus(func(status *channelPreparationAutoPromotionSchedulerStatus) { + status.SchedulerEnabled = false + status.NextCheckAt = 0 + status.Running = false + }) + continue + } + common.SysLog(fmt.Sprintf("running channel preparation auto promotion with interval %d minutes", intervalMinutes)) + updateChannelPreparationAutoPromotionSchedulerStatus(func(status *channelPreparationAutoPromotionSchedulerStatus) { + status.Running = true + status.NextCheckAt = 0 + status.LastCheckAt = common.GetTimestamp() + }) + summary, err := RunChannelPreparationAutoPromotion(channelPreparationAutoPromotionTriggerScheduler, "", nil) + if err != nil { + updateChannelPreparationAutoPromotionSchedulerStatus(func(status *channelPreparationAutoPromotionSchedulerStatus) { + status.Running = false + status.LastFinishedAt = common.GetTimestamp() + }) + common.SysError("channel preparation auto promotion failed: " + err.Error()) + continue + } + updateChannelPreparationAutoPromotionSchedulerStatus(func(status *channelPreparationAutoPromotionSchedulerStatus) { + status.Running = false + status.LastFinishedAt = common.GetTimestamp() + status.LastPromoted = summary.TotalPromoted + }) + common.SysLog(fmt.Sprintf("channel preparation auto promotion finished: promoted=%d, limit_reached=%v", summary.TotalPromoted, summary.LimitReached)) + if summary.TotalPromoted > 0 || summary.LimitReached { + recordChannelPreparationAutoPromotionManageLog(nil, fmt.Sprintf("定时执行渠道备货池自动晋升：晋升 %d 个渠道", summary.TotalPromoted), 0, "", map[string]interface{}{ + "total_promoted": summary.TotalPromoted, + "limit_reached": summary.LimitReached, + "trigger": channelPreparationAutoPromotionTriggerScheduler, + }) + } + } + }() + }) +} diff --git a/controller/channel_preparation_auto_promotion_test.go b/controller/channel_preparation_auto_promotion_test.go new file mode 100644 index 000000000000..4a12e2d35aa9 --- /dev/null +++ b/controller/channel_preparation_auto_promotion_test.go @@ -0,0 +1,86 @@ +package controller + +import ( + "math/rand" + "testing" + + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/stretchr/testify/require" +) + +func int64Ptr(value int64) *int64 { return &value } +func uintPtr(value uint) *uint { return &value } + +func TestChooseChannelPreparationAutoPromotionCandidateRespectsHighestPriorityTier(t *testing.T) { + preparations := []model.ChannelPreparation{ + {Id: 1, Balance: 100, Priority: int64Ptr(1), Weight: uintPtr(100000)}, + {Id: 2, Balance: 10, Priority: int64Ptr(10), Weight: uintPtr(0)}, + } + + candidate, ok := chooseChannelPreparationAutoPromotionCandidate( + preparations, + operation_setting.ChannelPreparationAutoPromotionStrategyPriorityWeighted, + rand.New(rand.NewSource(1)), + ) + + require.True(t, ok) + require.Equal(t, 2, candidate.Id) +} + +func TestChooseChannelPreparationAutoPromotionCandidateSmallBalanceFirst(t *testing.T) { + preparations := []model.ChannelPreparation{ + {Id: 1, Balance: 1, Priority: int64Ptr(1)}, + {Id: 2, Balance: 5, Priority: int64Ptr(10)}, + {Id: 3, Balance: 2, Priority: int64Ptr(10)}, + {Id: 4, Balance: 2, Priority: int64Ptr(10)}, + } + + candidate, ok := chooseChannelPreparationAutoPromotionCandidate( + preparations, + operation_setting.ChannelPreparationAutoPromotionStrategySmallBalanceFirst, + nil, + ) + + require.True(t, ok) + require.Equal(t, 3, candidate.Id) +} + +func TestChooseChannelPreparationAutoPromotionCandidateLargeBalanceFirst(t *testing.T) { + preparations := []model.ChannelPreparation{ + {Id: 1, Balance: 100, Priority: int64Ptr(1)}, + {Id: 2, Balance: 5, Priority: int64Ptr(10)}, + {Id: 3, Balance: 10, Priority: int64Ptr(10)}, + {Id: 4, Balance: 10, Priority: int64Ptr(10)}, + } + + candidate, ok := chooseChannelPreparationAutoPromotionCandidate( + preparations, + operation_setting.ChannelPreparationAutoPromotionStrategyLargeBalanceFirst, + nil, + ) + + require.True(t, ok) + require.Equal(t, 3, candidate.Id) +} + +func TestChooseChannelPreparationAutoPromotionActiveShortage(t *testing.T) { + capacityFirst := operation_setting.ChannelPreparationAutoPromotionRule{ + GuaranteePriority: operation_setting.ChannelPreparationAutoPromotionGuaranteePriorityCapacityFirst, + } + countFirst := operation_setting.ChannelPreparationAutoPromotionRule{ + GuaranteePriority: operation_setting.ChannelPreparationAutoPromotionGuaranteePriorityCountFirst, + } + + require.Equal(t, channelPreparationAutoPromotionShortageCapacity, chooseChannelPreparationAutoPromotionActiveShortage(capacityFirst, true, true)) + require.Equal(t, channelPreparationAutoPromotionShortageCount, chooseChannelPreparationAutoPromotionActiveShortage(countFirst, true, true)) + require.Equal(t, channelPreparationAutoPromotionShortageCount, chooseChannelPreparationAutoPromotionActiveShortage(capacityFirst, true, false)) + require.Equal(t, channelPreparationAutoPromotionShortageCapacity, chooseChannelPreparationAutoPromotionActiveShortage(countFirst, false, true)) + require.Empty(t, chooseChannelPreparationAutoPromotionActiveShortage(countFirst, false, false)) +} + +func TestChannelPreparationAutoPromotionCountDeficit(t *testing.T) { + require.Equal(t, int64(0), channelPreparationAutoPromotionCountDeficit(0, 0)) + require.Equal(t, int64(3), channelPreparationAutoPromotionCountDeficit(5, 2)) + require.Equal(t, int64(0), channelPreparationAutoPromotionCountDeficit(2, 5)) +} diff --git a/controller/channel_query_key.go b/controller/channel_query_key.go new file mode 100644 index 000000000000..0eb973281845 --- /dev/null +++ b/controller/channel_query_key.go @@ -0,0 +1,139 @@ +package controller + +import ( + "errors" + "net/http" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" +) + +type QueryKeyReportRequest struct { + Keys []string `json:"keys"` +} + +type QueryKeyTestRequest struct { + Key string `json:"key"` + Source string `json:"source"` + TargetID int `json:"target_id"` + Model string `json:"model"` + EndpointType string `json:"endpoint_type"` + Stream bool `json:"stream"` +} + +func QueryChannelKeyReport(c *gin.Context) { + request := QueryKeyReportRequest{} + if err := common.DecodeJson(c.Request.Body, &request); err != nil { + common.ApiError(c, err) + return + } + + report, err := model.BuildChannelQueryKeyReport(request.Keys) + if err != nil { + common.ApiError(c, err) + return + } + + common.ApiSuccess(c, report) +} + +func QueryChannelKeyTest(c *gin.Context) { + request := QueryKeyTestRequest{} + if err := common.DecodeJson(c.Request.Body, &request); err != nil { + common.ApiError(c, err) + return + } + + key := strings.TrimSpace(request.Key) + source := strings.TrimSpace(request.Source) + if key == "" { + common.ApiErrorMsg(c, "key不能为空") + return + } + if request.TargetID <= 0 { + common.ApiErrorMsg(c, "target_id不能为空") + return + } + + channel, err := buildQueryKeyTestChannel(source, request.TargetID, key) + if err != nil { + common.ApiError(c, err) + return + } + + testUserID, err := resolveChannelTestUserID(c) + if err != nil { + common.ApiError(c, err) + return + } + + testModel := strings.TrimSpace(request.Model) + tik := time.Now() + result := testChannel(channel, testUserID, testModel, request.EndpointType, request.Stream) + milliseconds := time.Since(tik).Milliseconds() + consumedTime := float64(milliseconds) / 1000.0 + if result.localErr != nil { + resp := gin.H{ + "success": false, + "message": result.localErr.Error(), + "time": consumedTime, + } + if result.newAPIError != nil { + resp["error_code"] = result.newAPIError.GetErrorCode() + } + c.JSON(http.StatusOK, resp) + return + } + if result.newAPIError != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": result.newAPIError.Error(), + "time": consumedTime, + "error_code": result.newAPIError.GetErrorCode(), + }) + return + } + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "time": consumedTime, + }) +} + +func buildQueryKeyTestChannel(source string, targetID int, key string) (*model.Channel, error) { + switch source { + case model.QueryKeyReportSourceChannel: + channel, err := model.GetChannelById(targetID, true) + if err != nil { + return nil, err + } + if !model.QueryKeyReportStoredKeyContains(channel.Key, key) { + return nil, errors.New("key不属于该渠道") + } + testChannel := *channel + testChannel.Key = key + testChannel.Keys = nil + testChannel.ChannelInfo = model.ChannelInfo{} + return &testChannel, nil + case model.QueryKeyReportSourcePreparation: + var preparation model.ChannelPreparation + if err := model.DB.First(&preparation, "id = ?", targetID).Error; err != nil { + return nil, err + } + if !model.QueryKeyReportStoredKeyContains(preparation.Key, key) { + return nil, errors.New("key不属于该备货渠道") + } + applyChannelPreparationDefaults(&preparation) + testChannel := preparation.ToChannel() + testChannel.Id = preparation.Id + testChannel.Key = key + testChannel.Keys = nil + testChannel.ChannelInfo = model.ChannelInfo{} + return testChannel, nil + default: + return nil, errors.New("source不支持") + } +} diff --git a/controller/channel_query_key_report_test.go b/controller/channel_query_key_report_test.go new file mode 100644 index 000000000000..4bb103bb29af --- /dev/null +++ b/controller/channel_query_key_report_test.go @@ -0,0 +1,133 @@ +package controller + +import ( + "bytes" + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +type queryKeyReportTestResponse struct { + Success bool `json:"success"` + Message string `json:"message"` + Data *model.QueryKeyReport `json:"data"` +} + +func postQueryChannelKeyReport(t *testing.T, request QueryKeyReportRequest) queryKeyReportTestResponse { + t.Helper() + + body, err := common.Marshal(request) + require.NoError(t, err) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/query-key/report", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + QueryChannelKeyReport(ctx) + + require.Equal(t, http.StatusOK, recorder.Code) + var payload queryKeyReportTestResponse + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &payload)) + return payload +} + +func TestQueryChannelKeyReportReturnsApiSuccessReport(t *testing.T) { + setupModelListControllerTestDB(t) + + channel := model.Channel{Id: 2001, Type: 1, Key: "sk-report", Name: "report channel", Status: common.ChannelStatusEnabled, Group: "default", Models: "gpt-4o", UsedQuota: int64(common.QuotaPerUnit) * 2, Balance: 5} + require.NoError(t, model.DB.Create(&channel).Error) + + payload := postQueryChannelKeyReport(t, QueryKeyReportRequest{Keys: []string{" sk-report ", "sk-missing", "sk-report"}}) + + require.True(t, payload.Success) + require.Empty(t, payload.Message) + require.NotNil(t, payload.Data) + require.Equal(t, 3, payload.Data.TotalInput) + require.Equal(t, 2, payload.Data.UniqueKeys) + require.Equal(t, 1, payload.Data.DuplicateCount) + require.Equal(t, 1, payload.Data.FoundCount) + require.Equal(t, 1, payload.Data.NotFoundCount) + require.Len(t, payload.Data.Items, 2) + require.Equal(t, "sk-report", payload.Data.Items[0].Key) + require.True(t, payload.Data.Items[0].Found) + require.Len(t, payload.Data.Items[0].Channels, 1) +} + +func TestQueryChannelKeyReportRejectsEmptyInput(t *testing.T) { + setupModelListControllerTestDB(t) + + payload := postQueryChannelKeyReport(t, QueryKeyReportRequest{Keys: []string{"", " "}}) + + require.False(t, payload.Success) + require.Contains(t, payload.Message, "keys") + require.Nil(t, payload.Data) +} + +func TestQueryChannelKeyReportRejectsMoreThanTenThousandUniqueKeys(t *testing.T) { + setupModelListControllerTestDB(t) + + keys := make([]string, model.MaxQueryKeyReportKeys+1) + for i := range keys { + keys[i] = fmt.Sprintf("sk-%05d", i) + } + payload := postQueryChannelKeyReport(t, QueryKeyReportRequest{Keys: keys}) + + require.False(t, payload.Success) + require.Contains(t, payload.Message, "10000") + require.Nil(t, payload.Data) +} + +func TestBuildQueryKeyTestChannelUsesOnlyRequestedChannelKey(t *testing.T) { + setupModelListControllerTestDB(t) + + channel := model.Channel{ + Id: 2201, + Type: 1, + Key: "sk-a\nsk-b", + Name: "multi key channel", + Status: common.ChannelStatusEnabled, + Group: "default", + Models: "gpt-4o", + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + }, + } + require.NoError(t, model.DB.Create(&channel).Error) + + testChannel, err := buildQueryKeyTestChannel(model.QueryKeyReportSourceChannel, channel.Id, "sk-b") + require.NoError(t, err) + require.Equal(t, "sk-b", testChannel.Key) + require.False(t, testChannel.ChannelInfo.IsMultiKey) + + _, err = buildQueryKeyTestChannel(model.QueryKeyReportSourceChannel, channel.Id, "sk-missing") + require.Error(t, err) + require.Contains(t, err.Error(), "不属于") +} + +func TestBuildQueryKeyTestChannelSupportsPreparation(t *testing.T) { + setupModelListControllerTestDB(t) + + preparation := model.ChannelPreparation{ + Id: 2301, + Type: 2, + Key: "sk-prep-a\nsk-prep-b", + Name: "prep multi", + Status: model.ChannelPreparationStatusPending, + Group: "default", + Models: "claude-3", + } + require.NoError(t, model.DB.Create(&preparation).Error) + + testChannel, err := buildQueryKeyTestChannel(model.QueryKeyReportSourcePreparation, preparation.Id, "sk-prep-a") + require.NoError(t, err) + require.Equal(t, preparation.Id, testChannel.Id) + require.Equal(t, "sk-prep-a", testChannel.Key) + require.False(t, testChannel.ChannelInfo.IsMultiKey) +} diff --git a/controller/channel_test_internal_test.go b/controller/channel_test_internal_test.go index 025408010b1d..7291145bb962 100644 --- a/controller/channel_test_internal_test.go +++ b/controller/channel_test_internal_test.go @@ -5,9 +5,12 @@ import ( "testing" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/pkg/billingexpr" relaycommon "github.com/QuantumNous/new-api/relay/common" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/QuantumNous/new-api/types" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" @@ -80,3 +83,49 @@ func TestResolveChannelTestUserIDUsesRequestUser(t *testing.T) { require.NoError(t, err) require.Equal(t, 2, userID) } + +func TestResolveChannelTestModelSelectsConfiguredAnthropicModel(t *testing.T) { + withModelRatioConfig(t, map[string]float64{ + "claude-3-sonnet-20240229": 1.5, + "claude-3-7-sonnet-20250219": 1.5, + }) + + unconfiguredTestModel := "claude-unpriced-test-model" + channel := &model.Channel{ + Type: constant.ChannelTypeAnthropic, + TestModel: &unconfiguredTestModel, + Models: "claude-3-sonnet-20240229,claude-unpriced-test-model,claude-3-7-sonnet-20250219", + } + + require.Equal(t, "claude-3-7-sonnet-20250219", resolveChannelTestModel(channel, "")) + require.Equal(t, "claude-unpriced-test-model", resolveChannelTestModel(channel, " claude-unpriced-test-model ")) +} + +func TestResolveChannelTestModelUsesConfiguredAnthropicTestModel(t *testing.T) { + withModelRatioConfig(t, map[string]float64{ + "claude-3-5-sonnet-20241022": 1.5, + "claude-3-7-sonnet-20250219": 1.5, + }) + + testModel := "claude-3-5-sonnet-20241022" + channel := &model.Channel{ + Type: constant.ChannelTypeAnthropic, + TestModel: &testModel, + Models: "claude-3-7-sonnet-20250219,claude-3-5-sonnet-20241022", + } + + require.Equal(t, "claude-3-5-sonnet-20241022", resolveChannelTestModel(channel, "")) +} + +func withModelRatioConfig(t *testing.T, ratios map[string]float64) { + t.Helper() + + saved := ratio_setting.ModelRatio2JSONString() + t.Cleanup(func() { + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(saved)) + }) + + payload, err := common.Marshal(ratios) + require.NoError(t, err) + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(string(payload))) +} diff --git a/controller/cost_report.go b/controller/cost_report.go new file mode 100644 index 000000000000..22d36c94de50 --- /dev/null +++ b/controller/cost_report.go @@ -0,0 +1,291 @@ +package controller + +import ( + "fmt" + "net/http" + "net/url" + "strconv" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + costreport "github.com/QuantumNous/new-api/service/cost_report" + "github.com/gin-gonic/gin" +) + +func costReportService() *costreport.Service { + return costreport.NewService(model.DB, model.LOG_DB) +} + +func actorID(c *gin.Context) int { + if id, ok := c.Get("id"); ok { + if v, ok := id.(int); ok { + return v + } + } + return 0 +} + +func CostReportListTemplates(c *gin.Context) { + page, pageSize := parsePage(c) + details, total, err := costReportService().ListTemplates(c.Request.Context(), (page-1)*pageSize, pageSize) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": details, "total": total, "page": page, "page_size": pageSize}) +} + +func CostReportEnsureDefaultTemplate(c *gin.Context) { + detail, err := costReportService().EnsureDefaultTemplate(c.Request.Context(), actorID(c)) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, detail) +} + +func CostReportGetTemplate(c *gin.Context) { + id, err := parseIDParam(c, "id") + if err != nil { + common.ApiError(c, err) + return + } + detail, err := costReportService().GetTemplate(c.Request.Context(), id) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, detail) +} + +func CostReportSaveTemplate(c *gin.Context) { + var input costreport.TemplateSaveInput + if err := common.DecodeJson(c.Request.Body, &input); err != nil { + common.ApiError(c, err) + return + } + if idParam := c.Param("id"); idParam != "" { + id, err := strconv.Atoi(idParam) + if err != nil { + common.ApiError(c, err) + return + } + input.Id = id + } + input.ActorID = actorID(c) + detail, err := costReportService().SaveTemplate(c.Request.Context(), input) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, detail) +} + +func CostReportListTemplateVersions(c *gin.Context) { + id, err := parseIDParam(c, "id") + if err != nil { + common.ApiError(c, err) + return + } + versions, err := costReportService().ListTemplateVersions(c.Request.Context(), id) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": versions}) +} + +type costReportValidateRequest struct { + Config costreport.CostReportTemplateConfig `json:"config"` +} + +func CostReportValidateTemplate(c *gin.Context) { + var req costReportValidateRequest + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if err := costreport.ValidateTemplateConfig(req.Config); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"valid": true}) +} + +func CostReportPreview(c *gin.Context) { + var req costreport.PreviewRequest + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + resp, err := costReportService().Preview(c.Request.Context(), req) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, resp) +} + +func CostReportSaveRun(c *gin.Context) { + var req costreport.PreviewRequest + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if req.Config != nil { + common.ApiErrorMsg(c, "saving a run requires a persisted template version; preview unsaved config separately") + return + } + req.IncludeManual = true + preview, err := costReportService().Preview(c.Request.Context(), req) + if err != nil { + common.ApiError(c, err) + return + } + result, err := costReportService().SaveRunFromPreview(c.Request.Context(), preview, actorID(c)) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"run": result.Run, "row_count": len(result.Rows), "warnings": preview.Warnings}) +} + +func CostReportListRuns(c *gin.Context) { + page, pageSize := parsePage(c) + templateID, _ := strconv.Atoi(c.Query("template_id")) + periodKey := c.Query("period_key") + runs, total, err := costReportService().ListRuns(c.Request.Context(), templateID, periodKey, (page-1)*pageSize, pageSize) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{"items": runs, "total": total, "page": page, "page_size": pageSize}) +} + +func CostReportGetRun(c *gin.Context) { + id, err := parseIDParam(c, "id") + if err != nil { + common.ApiError(c, err) + return + } + detail, err := costReportService().GetRunDetail(c.Request.Context(), id) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, detail) +} + +func CostReportExportRun(c *gin.Context) { + id, err := parseIDParam(c, "id") + if err != nil { + common.ApiError(c, err) + return + } + data, filename, err := costReportService().ExportRunXLSX(c.Request.Context(), id) + if err != nil { + common.ApiError(c, err) + return + } + c.Header("Content-Disposition", fmt.Sprintf("attachment; filename*=UTF-8''%s", url.PathEscape(filename))) + c.Data(http.StatusOK, "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", data) +} + +func CostReportReadManualCells(c *gin.Context) { + templateID, _ := strconv.Atoi(c.Query("template_id")) + periodKey := c.Query("period_key") + rowKeys := splitQueryList(c.Query("row_keys")) + if rowKey := c.Query("row_key"); rowKey != "" { + rowKeys = append(rowKeys, rowKey) + } + manuals, err := costReportService().ReadManualCells(c.Request.Context(), templateID, periodKey, rowKeys) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, manuals) +} + +func CostReportUpsertManualCell(c *gin.Context) { + var input costreport.ManualCellInput + if err := common.DecodeJson(c.Request.Body, &input); err != nil { + common.ApiError(c, err) + return + } + input.UpdatedBy = actorID(c) + cell, err := costReportService().UpsertManualCell(c.Request.Context(), input) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, cell) +} + +type costReportClassificationPreviewRequest struct { + Config costreport.CostReportTemplateConfig `json:"config"` + Log model.Log `json:"log"` + Channel *model.Channel `json:"channel,omitempty"` + User *model.User `json:"user,omitempty"` + LogOther map[string]interface{} `json:"log_other,omitempty"` +} + +func CostReportClassificationPreview(c *gin.Context) { + var req costReportClassificationPreviewRequest + if err := common.DecodeJson(c.Request.Body, &req); err != nil { + common.ApiError(c, err) + return + } + if err := costreport.ValidateTemplateConfig(req.Config); err != nil { + common.ApiError(c, err) + return + } + result := costreport.Classify(req.Config, costreport.ClassificationInput{ + Log: &req.Log, + Channel: req.Channel, + User: req.User, + LogOther: req.LogOther, + }) + common.ApiSuccess(c, result) +} + +func parseIDParam(c *gin.Context, name string) (int, error) { + id, err := strconv.Atoi(c.Param(name)) + if err != nil || id <= 0 { + return 0, fmt.Errorf("invalid %s", name) + } + return id, nil +} + +func parsePage(c *gin.Context) (int, int) { + page, _ := strconv.Atoi(c.Query("p")) + if page <= 0 { + page, _ = strconv.Atoi(c.Query("page")) + } + if page <= 0 { + page = 1 + } + pageSize, _ := strconv.Atoi(c.Query("page_size")) + if pageSize <= 0 { + pageSize, _ = strconv.Atoi(c.Query("size")) + } + if pageSize <= 0 || pageSize > 100 { + pageSize = 20 + } + return page, pageSize +} + +func splitQueryList(value string) []string { + if strings.TrimSpace(value) == "" { + return nil + } + parts := strings.Split(value, ",") + out := make([]string, 0, len(parts)) + for _, part := range parts { + part = strings.TrimSpace(part) + if part != "" { + out = append(out, part) + } + } + return out +} diff --git a/controller/group.go b/controller/group.go index 6ba339a3f9bd..8d528d1e5a99 100644 --- a/controller/group.go +++ b/controller/group.go @@ -2,7 +2,10 @@ package controller import ( "net/http" + "sort" + "strings" + "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting" @@ -11,11 +14,43 @@ import ( "github.com/gin-gonic/gin" ) +func addGroupNames(groupSet map[string]bool, groupValues []string) { + for _, groupValue := range groupValues { + for _, groupName := range strings.Split(groupValue, ",") { + groupName = strings.TrimSpace(groupName) + if groupName != "" { + groupSet[groupName] = true + } + } + } +} + func GetGroups(c *gin.Context) { - groupNames := make([]string, 0) + groupSet := map[string]bool{"default": true} for groupName := range ratio_setting.GetGroupRatioCopy() { - groupNames = append(groupNames, groupName) + addGroupNames(groupSet, []string{groupName}) + } + channelGroups, err := model.GetDistinctChannelGroups() + if err != nil { + common.SysError("failed to get channel groups: " + err.Error()) + } else { + addGroupNames(groupSet, channelGroups) + } + preparationGroups, err := model.GetDistinctChannelPreparationGroups() + if err != nil { + common.SysError("failed to get channel preparation groups: " + err.Error()) + } else { + addGroupNames(groupSet, preparationGroups) + } + + groupNames := make([]string, 0, len(groupSet)) + for groupName := range groupSet { + if groupName != "default" { + groupNames = append(groupNames, groupName) + } } + sort.Strings(groupNames) + groupNames = append([]string{"default"}, groupNames...) c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", diff --git a/controller/log.go b/controller/log.go index a43f7b752278..10e05fb02c47 100644 --- a/controller/log.go +++ b/controller/log.go @@ -1,11 +1,15 @@ package controller import ( + "fmt" "net/http" + "net/url" "strconv" + "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" + usagelogexport "github.com/QuantumNous/new-api/service/usage_log_export" "github.com/gin-gonic/gin" ) @@ -55,6 +59,81 @@ func GetUserLogs(c *gin.Context) { return } +func GetLogExportFields(c *gin.Context) { + common.ApiSuccess(c, usagelogexport.FieldGroups(true)) +} + +func GetUserLogExportFields(c *gin.Context) { + common.ApiSuccess(c, usagelogexport.FieldGroups(false)) +} + +func ExportAllLogs(c *gin.Context) { + exportLogs(c, true) +} + +func ExportUserLogs(c *gin.Context) { + exportLogs(c, false) +} + +func exportLogs(c *gin.Context, isAdmin bool) { + filter := buildLogExportFilter(c, isAdmin) + file, filename, _, err := usagelogexport.BuildXLSX(c.Request.Context(), usagelogexport.ExportInput{ + Filter: filter, + Fields: splitLogExportFields(c.Query("fields")), + Timezone: c.Query("timezone"), + }) + if err != nil { + common.ApiError(c, err) + return + } + defer func() { _ = file.Close() }() + c.Header("Content-Disposition", fmt.Sprintf("attachment; filename*=UTF-8''%s", url.PathEscape(filename))) + c.Header("Content-Type", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet") + c.Status(http.StatusOK) + if err := file.Write(c.Writer); err != nil { + common.SysError("failed to write usage log export: " + err.Error()) + } +} + +func splitLogExportFields(value string) []string { + if value == "" { + return nil + } + parts := strings.Split(value, ",") + fields := make([]string, 0, len(parts)) + for _, part := range parts { + part = strings.TrimSpace(part) + if part != "" { + fields = append(fields, part) + } + } + return fields +} + +func buildLogExportFilter(c *gin.Context, isAdmin bool) model.LogExportFilter { + logType, _ := strconv.Atoi(c.Query("type")) + startTimestamp, _ := strconv.ParseInt(c.Query("start_timestamp"), 10, 64) + endTimestamp, _ := strconv.ParseInt(c.Query("end_timestamp"), 10, 64) + channel, _ := strconv.Atoi(c.Query("channel")) + filter := model.LogExportFilter{ + UserId: c.GetInt("id"), + IsAdmin: isAdmin, + LogType: logType, + StartTimestamp: startTimestamp, + EndTimestamp: endTimestamp, + ModelName: c.Query("model_name"), + TokenName: c.Query("token_name"), + Group: c.Query("group"), + RequestId: c.Query("request_id"), + UpstreamRequestId: c.Query("upstream_request_id"), + } + if isAdmin { + filter.Username = c.Query("username") + filter.Channel = channel + } + return filter +} + // Deprecated: SearchAllLogs 已废弃，前端未使用该接口。 func SearchAllLogs(c *gin.Context) { c.JSON(http.StatusOK, gin.H{ diff --git a/controller/model_list_test.go b/controller/model_list_test.go index 97d27cae5c6c..79e882efdcb3 100644 --- a/controller/model_list_test.go +++ b/controller/model_list_test.go @@ -43,7 +43,7 @@ func setupModelListControllerTestDB(t *testing.T) *gorm.DB { model.DB = db model.LOG_DB = db - require.NoError(t, db.AutoMigrate(&model.User{}, &model.Channel{}, &model.Ability{}, &model.Model{}, &model.Vendor{})) + require.NoError(t, db.AutoMigrate(&model.User{}, &model.Channel{}, &model.ChannelPreparation{}, &model.Ability{}, &model.Model{}, &model.Vendor{})) t.Cleanup(func() { sqlDB, err := db.DB() @@ -220,10 +220,12 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) { "zz-token-tiered-visible-model": `tier("base", p * 1 + c * 2)`, "zz-token-tiered-empty-expr-model": "", }) + setupModelListControllerTestDB(t) recorder := httptest.NewRecorder() ctx, _ := gin.CreateTestContext(recorder) ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") common.SetContextKey(ctx, constant.ContextKeyTokenModelLimitEnabled, true) common.SetContextKey(ctx, constant.ContextKeyTokenModelLimit, map[string]bool{ "zz-token-tiered-visible-model": true, diff --git a/controller/navigation.go b/controller/navigation.go new file mode 100644 index 000000000000..d7b3cc7e6d9f --- /dev/null +++ b/controller/navigation.go @@ -0,0 +1,400 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +package controller + +import ( + "net/http" + "strconv" + + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + "github.com/gin-gonic/gin" + "gorm.io/gorm" +) + +// GetNavigationTree 用户侧获取导航树 API +func GetNavigationTree(c *gin.Context) { + menuKey := c.DefaultQuery("menu_key", "default_web_top") + + // 从 Context 中提取语言，默认为 zh-CN + locale := c.GetString("lang") + if locale == "" { + locale = c.DefaultQuery("lang", "zh-CN") + } + + // 提取用户登录态与权限 + userID := c.GetInt("id") + userRole := c.GetInt("role") + userGroup := c.GetString("group") + + isAuthenticated := userID > 0 + + tree, err := service.NavService.GetVisibleNavigationTree(menuKey, locale, userRole, userGroup, isAuthenticated) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + + c.JSON(http.StatusOK, gin.H{ + "success": true, + "data": tree, + }) +} + +// ================= 管理侧菜单 CRUD 接口 ================= + +// AdminGetMenus 获取所有菜单容器列表 +func AdminGetMenus(c *gin.Context) { + var menus []model.NavigationMenu + if err := model.DB.Order("id asc").Find(&menus).Error; err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + c.JSON(http.StatusOK, gin.H{"success": true, "data": menus}) +} + +// AdminCreateMenu 创建新的菜单配置 +func AdminCreateMenu(c *gin.Context) { + var menu model.NavigationMenu + if err := c.ShouldBindJSON(&menu); err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + menu.IsSystem = false // 管理员手工创建的绝非系统菜单 + if err := model.DB.Create(&menu).Error; err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + service.NavService.InvalidateCache() + c.JSON(http.StatusOK, gin.H{"success": true, "data": menu}) +} + +// AdminUpdateMenu 更新菜单元数据 +func AdminUpdateMenu(c *gin.Context) { + idStr := c.Param("id") + id, err := strconv.Atoi(idStr) + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "invalid menu id"}) + return + } + + var menu model.NavigationMenu + if err := model.DB.First(&menu, id).Error; err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "menu not found"}) + return + } + + var input model.NavigationMenu + if err := c.ShouldBindJSON(&input); err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + // 仅允许修改名称、启用状态 + menu.Name = input.Name + menu.Enabled = input.Enabled + + if err := model.DB.Save(&menu).Error; err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + service.NavService.InvalidateCache() + c.JSON(http.StatusOK, gin.H{"success": true, "data": menu}) +} + +// AdminDeleteMenu 删除非系统级菜单 +func AdminDeleteMenu(c *gin.Context) { + idStr := c.Param("id") + id, err := strconv.Atoi(idStr) + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "invalid menu id"}) + return + } + + var menu model.NavigationMenu + if err := model.DB.First(&menu, id).Error; err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "menu not found"}) + return + } + + if menu.IsSystem { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "system menu cannot be deleted"}) + return + } + + // 在事务中级联清理 Menu 关联的所有 Item + err = model.DB.Transaction(func(tx *gorm.DB) error { + var items []model.NavigationItem + if err := tx.Where("menu_id = ?", menu.ID).Find(&items).Error; err != nil { + return err + } + + for _, item := range items { + // 触发级联物理删除 Translations & Rules + if err := tx.Delete(&item).Error; err != nil { + return err + } + } + + return tx.Delete(&menu).Error + }) + + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + service.NavService.InvalidateCache() + c.JSON(http.StatusOK, gin.H{"success": true, "message": "menu deleted successfully"}) +} + +// ================= 管理侧菜单节点 CRUD 接口 ================= + +// AdminGetItems 获取某个菜单下所有的平铺节点（含级联预加载翻译和规则，由前端还原树） +func AdminGetItems(c *gin.Context) { + menuIDStr := c.Query("menu_id") + if menuIDStr == "" { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "menu_id query parameter is required"}) + return + } + + menuID, err := strconv.Atoi(menuIDStr) + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "invalid menu_id"}) + return + } + + var items []model.NavigationItem + err = model.DB.Where("menu_id = ?", menuID). + Order("sort_order asc, id asc"). + Preload("Translations"). + Preload("Rules"). + Find(&items).Error + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + c.JSON(http.StatusOK, gin.H{"success": true, "data": items}) +} + +// AdminCreateItem 创建菜单节点（包含多语言和可见性规则的一体化保存） +func AdminCreateItem(c *gin.Context) { + var item model.NavigationItem + if err := c.ShouldBindJSON(&item); err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + // 1. URL 协议安全性拦截校验（防 XSS 注入） + if err := service.NavService.ValidateItemURL(item.Type, item.URL); err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + // 2. 事务级联创建节点及其子集合 + err := model.DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Omit("Translations", "Rules").Create(&item).Error; err != nil { + return err + } + + // 保存多语言 + for i := range item.Translations { + item.Translations[i].ItemID = item.ID + if err := tx.Create(&item.Translations[i]).Error; err != nil { + return err + } + } + + // 保存可见性规则 + for i := range item.Rules { + item.Rules[i].ItemID = item.ID + if err := tx.Create(&item.Rules[i]).Error; err != nil { + return err + } + } + + return nil + }) + + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + service.NavService.InvalidateCache() + c.JSON(http.StatusOK, gin.H{"success": true, "data": item}) +} + +// AdminUpdateItem 更新菜单节点及其子属性（采用 FullSaveAssociations 完整事务更新） +func AdminUpdateItem(c *gin.Context) { + idStr := c.Param("id") + id, err := strconv.Atoi(idStr) + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "invalid item id"}) + return + } + + var item model.NavigationItem + if err := model.DB.Preload("Translations").Preload("Rules").First(&item, id).Error; err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "item not found"}) + return + } + + var input model.NavigationItem + if err := c.ShouldBindJSON(&input); err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + // 1. 安全性 URL 校验 + if err := service.NavService.ValidateItemURL(input.Type, input.URL); err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + // 2. 级联事务全量覆盖更新 + err = model.DB.Transaction(func(tx *gorm.DB) error { + // 清理旧的多语言翻译和可见性规则，避免级联更新产生废弃记录 + if err := tx.Where("item_id = ?", item.ID).Delete(&model.NavigationItemTranslation{}).Error; err != nil { + return err + } + if err := tx.Where("item_id = ?", item.ID).Delete(&model.NavigationVisibilityRule{}).Error; err != nil { + return err + } + + // 更新字段 + item.ParentID = input.ParentID + item.Type = input.Type + item.ModuleKey = input.ModuleKey + item.Path = input.Path + item.URL = input.URL + item.IconKey = input.IconKey + item.SortOrder = input.SortOrder + item.Enabled = input.Enabled + item.OpenInNewTab = input.OpenInNewTab + item.ExactActive = input.ExactActive + + // 保存主体，忽略关联表的自动保存，避免与接下来的手动保存发生冲突 + if err := tx.Omit("Translations", "Rules").Save(&item).Error; err != nil { + return err + } + + // 创建新的 Translations + for i := range input.Translations { + input.Translations[i].ItemID = item.ID + input.Translations[i].ID = 0 // 重置 ID 确保插入 + if err := tx.Create(&input.Translations[i]).Error; err != nil { + return err + } + } + + // 创建新的 Rules + for i := range input.Rules { + input.Rules[i].ItemID = item.ID + input.Rules[i].ID = 0 // 重置 ID 确保插入 + if err := tx.Create(&input.Rules[i]).Error; err != nil { + return err + } + } + + return nil + }) + + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + service.NavService.InvalidateCache() + c.JSON(http.StatusOK, gin.H{"success": true, "data": item}) +} + +// AdminDeleteItem 删除节点（级联物理删除关联的 Translations 和 Rules） +func AdminDeleteItem(c *gin.Context) { + idStr := c.Param("id") + id, err := strconv.Atoi(idStr) + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "invalid item id"}) + return + } + + var item model.NavigationItem + if err := model.DB.First(&item, id).Error; err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": "item not found"}) + return + } + + err = model.DB.Transaction(func(tx *gorm.DB) error { + // 如果有子节点，将其父节点引用置为空，使子节点不致变成废弃不可达孤儿节点（或者可以选择级联删除子项） + // 在这里，按严谨级联规则，我们将子节点的 parent_id 设为 nil + if err := tx.Model(&model.NavigationItem{}).Where("parent_id = ?", item.ID).Update("parent_id", nil).Error; err != nil { + return err + } + + // 删除主体，触发外键约束自动级联删除 translations 和 rules + return tx.Delete(&item).Error + }) + + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + service.NavService.InvalidateCache() + c.JSON(http.StatusOK, gin.H{"success": true, "message": "item deleted successfully"}) +} + +type ReorderInput struct { + ItemID uint `json:"item_id"` + SortOrder int `json:"sort_order"` +} + +// AdminReorderItems 批量节点重新排序接口 +func AdminReorderItems(c *gin.Context) { + var inputs []ReorderInput + if err := c.ShouldBindJSON(&inputs); err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + err := model.DB.Transaction(func(tx *gorm.DB) error { + for _, input := range inputs { + if err := tx.Model(&model.NavigationItem{}).Where("id = ?", input.ItemID).Update("sort_order", input.SortOrder).Error; err != nil { + return err + } + } + return nil + }) + + if err != nil { + c.JSON(http.StatusOK, gin.H{"success": false, "message": err.Error()}) + return + } + + service.NavService.InvalidateCache() + c.JSON(http.StatusOK, gin.H{"success": true, "message": "items reordered successfully"}) +} diff --git a/controller/option.go b/controller/option.go index b5fdfdc1515b..84ddb818bc2c 100644 --- a/controller/option.go +++ b/controller/option.go @@ -295,6 +295,33 @@ func UpdateOption(c *gin.Context) { }) return } + case "channel_preparation_auto_promotion_setting.rules": + err = operation_setting.ValidateChannelPreparationAutoPromotionRulesJSONString(option.Value.(string)) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "自动晋升规则配置错误: " + err.Error(), + }) + return + } + case "channel_preparation_auto_promotion_setting.interval_minutes": + interval, parseErr := strconv.ParseFloat(option.Value.(string), 64) + if parseErr != nil || interval <= 0 { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "自动晋升检查间隔必须大于 0", + }) + return + } + case "channel_preparation_auto_promotion_setting.max_promotions_per_run": + limit, parseErr := strconv.Atoi(option.Value.(string)) + if parseErr != nil || limit <= 0 { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "自动晋升每次最大数量必须大于 0", + }) + return + } case "console_setting.api_info": err = console_setting.ValidateConsoleSettings(option.Value.(string), "ApiInfo") if err != nil { diff --git a/controller/twofa.go b/controller/twofa.go index ef86fd762ddc..7ce6383275fa 100644 --- a/controller/twofa.go +++ b/controller/twofa.go @@ -14,12 +14,35 @@ import ( // Setup2FARequest 设置2FA请求结构 type Setup2FARequest struct { - Code string `json:"code" binding:"required"` + Code string `json:"code" binding:"required"` + Password string `json:"password"` } // Verify2FARequest 验证2FA请求结构 type Verify2FARequest struct { - Code string `json:"code" binding:"required"` + Code string `json:"code" binding:"required"` + Password string `json:"password"` +} + +// verifyAccountPasswordForTwoFA 校验当前登录用户的账户密码，用于 2FA 管理类敏感操作的二次确认。 +// 防止仅持有被盗会话(cookie)、但不知道账户密码的攻击者自助开启/启用/禁用 2FA。 +// 说明：通过 OAuth 等方式注册、本身未设置密码(Password 为空)的账户跳过校验，避免被锁死。 +func verifyAccountPasswordForTwoFA(userId int, password string) error { + user, err := model.GetUserById(userId, true) + if err != nil { + return err + } + if user.Password == "" { + // 该账户未设置登录密码（如 OAuth 账户），无法用密码校验，放行。 + return nil + } + if password == "" { + return errors.New("请输入账户密码") + } + if !common.ValidatePasswordAndHash(password, user.Password) { + return errors.New("账户密码错误") + } + return nil } // Setup2FAResponse 设置2FA响应结构 @@ -33,6 +56,19 @@ type Setup2FAResponse struct { func Setup2FA(c *gin.Context) { userId := c.GetInt("id") + // 校验账户密码（防止被盗会话自助开启 2FA） + var pwReq struct { + Password string `json:"password"` + } + _ = c.ShouldBindJSON(&pwReq) + if err := verifyAccountPasswordForTwoFA(userId, pwReq.Password); err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + // 检查用户是否已经启用2FA existing, err := model.GetTwoFAByUserId(userId) if err != nil { @@ -146,6 +182,15 @@ func Enable2FA(c *gin.Context) { userId := c.GetInt("id") + // 校验账户密码（防止被盗会话自助启用 2FA） + if err := verifyAccountPasswordForTwoFA(userId, req.Password); err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + // 获取2FA记录 twoFA, err := model.GetTwoFAByUserId(userId) if err != nil { @@ -213,6 +258,15 @@ func Disable2FA(c *gin.Context) { userId := c.GetInt("id") + // 校验账户密码（防止被盗会话自助禁用 2FA） + if err := verifyAccountPasswordForTwoFA(userId, req.Password); err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + // 获取2FA记录 twoFA, err := model.GetTwoFAByUserId(userId) if err != nil { diff --git a/docker-compose.yml b/docker-compose.yml index 48c071f98feb..67b410c8a378 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -21,7 +21,7 @@ services: restart: always command: --log-dir /app/logs ports: - - "3000:3000" + - "3001:3000" volumes: - ./data:/data - ./logs:/app/logs diff --git a/docs/channel/channel-api.md b/docs/channel/channel-api.md new file mode 100644 index 000000000000..92cb94d4ba78 --- /dev/null +++ b/docs/channel/channel-api.md @@ -0,0 +1,659 @@ +# 渠道管理 API + +通过管理员鉴权调用 `/api/channel/*` 系列接口，可在不进入 Web 后台的前提下完成： + +- **状态统计** — 实时获取正常 / 手动禁用 / 自动禁用（被警用）渠道数（§3-§5） +- **新增渠道** — 单个、多 key 合并、批量拆分三种模式（§6） +- **修改 / 删除 / 测试 / 批量操作** — 完整的管理动作（§7） + +适用场景： + +- 外部监控/告警（正常渠道数低于阈值时通知值班） +- 巡检脚本（定期统计 / 导出健康状态） +- CI 部署后自检 +- 基础设施即代码（IaC）批量初始化或同步渠道配置 + +--- + +## 1. 状态枚举 + +定义位置：`common/constants.go:253-255` + +| 值 | 常量 | 含义 | +|---|---|---| +| `1` | `ChannelStatusEnabled` | 正常启用 | +| `2` | `ChannelStatusManuallyDisabled` | 管理员手动禁用 | +| `3` | `ChannelStatusAutoDisabled` | 系统自动禁用（被警用）— 通常因 key 失效、余额不足、连续报错触发 | + +--- + +## 2. 获取管理员 Access Token + +`/api/channel/*` 走 `AdminAuth` 中间件（`middleware/auth.go`），支持 **Session（Cookie）** 与 **Access Token（Bearer）** 两种鉴权。本文档使用 Token 方式，便于脚本调用。 + +### 方式 A：Web UI 一键生成（推荐） + +1. 使用 admin 账号登录 Web 控制台 +2. 进入「个人设置 / Personal Setting」 +3. 点击 **「生成系统访问令牌 / Generate System Access Token」** +4. 复制弹出的 token + +底层调用：`GET /api/user/token`（路由 `router/api-router.go:85`，需有效 admin session）。 + +### 方式 B：纯 CLI 流程 + +```bash +HOST="http://your-newapi-host:3000" +COOKIES=$(mktemp) + +# 1) 登录 admin 账号，保存 session cookie +curl -s -c "$COOKIES" \ + -H "Content-Type: application/json" \ + -X POST "$HOST/api/user/login" \ + -d '{"username":"admin","password":"YOUR_PASSWORD"}' + +# 2) 生成 / 刷新 access token +ADMIN_TOKEN=$(curl -s -b "$COOKIES" "$HOST/api/user/token" | jq -r .data) + +# 3) 获取自己的 user id（调用接口必须带） +ADMIN_UID=$(curl -s -b "$COOKIES" "$HOST/api/user/self" | jq -r .data.id) + +echo "Token: $ADMIN_TOKEN" +echo "UID: $ADMIN_UID" +``` + +> ⚠️ 重新调用 `/api/user/token` 会**覆盖**旧 token；用户被禁用后 token 同步失效。 + +--- + +## 3. 请求格式 + +### 端点 + +``` +GET /api/channel/ +``` + +### 必填请求头 + +| Header | 说明 | +|---|---| +| `Authorization: Bearer ` | 系统访问令牌（也兼容不带 `Bearer ` 前缀） | +| `New-Api-User: ` | 调用者 user id，未提供会返回 401 | + +### 查询参数 + +| 参数 | 类型 | 含义 | +|---|---|---| +| `status` | int | `1` 仅启用 / `0` 仅禁用（含 manual+auto）/ `-1` 全部 | +| `p` | int | 页码（从 1 开始） | +| `page_size` | int | 每页条数；最大会被后端限制为 `100`。仅统计时设为 `1` 即可（`total` 字段独立返回） | +| `group` | string | 按分组过滤 | +| `type` | int | 按渠道类型过滤（如 `1`=OpenAI，`14`=Anthropic） | +| `tag_mode` | bool | `true` 时返回按 tag 聚合 | +| `sort_by` | string | 排序字段：`id` / `name` / `priority` / `balance` / `response_time` / `test_time` | +| `sort_order` | string | `asc` 或 `desc`；无效值默认按 `desc` 处理 | +| `id_sort` | bool | 旧排序开关：未指定 `sort_by` 时可按 id 倒序 | + +参考实现：`controller/channel.go:92` (`GetAllChannels`)、`controller/channel.go:54` (`parseStatusFilter`)。 + +### 响应结构 + +```json +{ + "success": true, + "data": { + "items": [ { "id": 1, "name": "...", "type": 14, "status": 1, ... } ], + "total": 1, + "type_counts": { "14": 1 } + } +} +``` + +- `total` — 符合过滤条件的渠道总数 +- `type_counts` — 按 `type`（渠道类型）聚合的计数（map[type]count）。注意：当前实现会应用 `group/status` 过滤，但不会应用请求里的 `type` 过滤，因此可用于展示当前状态下所有类型分布。 +- `items` — 当前分页的渠道详情；接口不会返回 `key` 字段 + +--- + +## 4. 常用查询示例 + +### 4.1 仅取「正常启用」渠道数（最常用） + +```bash +curl -sH "Authorization: Bearer $ADMIN_TOKEN" \ + -H "New-Api-User: $ADMIN_UID" \ + "$HOST/api/channel/?status=1&p=1&page_size=1" \ + | jq '.data.total' +``` + +### 4.2 一次性获取三种状态分布 + +```bash +get_count() { + curl -sH "Authorization: Bearer $ADMIN_TOKEN" \ + -H "New-Api-User: $ADMIN_UID" \ + "$HOST/api/channel/?status=$1&p=1&page_size=1" \ + | jq -r '.data.total' +} + +ENABLED=$(get_count 1) +DISABLED=$(get_count 0) +TOTAL=$(get_count -1) + +echo "正常启用: $ENABLED" +echo "已禁用: $DISABLED" +echo "总计: $TOTAL" +``` + +### 4.3 精确区分「手动禁用」vs「自动禁用（被警用）」 + +接口的 `status=0` 把 manual+auto 合并返回，需在客户端按 `items[].status` 字段二次分组。注意后端会把 `page_size` 限制为最大 `100`，因此禁用渠道较多时必须分页拉完： + +```bash +page=1 +while :; do + resp=$(curl -sH "Authorization: Bearer $ADMIN_TOKEN" \ + -H "New-Api-User: $ADMIN_UID" \ + "$HOST/api/channel/?status=0&p=$page&page_size=100") + + echo "$resp" | jq -c '.data.items[]' + + total=$(echo "$resp" | jq -r '.data.total') + fetched=$((page * 100)) + [ "$fetched" -ge "$total" ] && break + page=$((page + 1)) +done | jq -s '[.[].status] | group_by(.) | map({status: .[0], count: length})' +``` + +输出示例： + +```json +[ + { "status": 2, "count": 3 }, // 手动禁用 3 个 + { "status": 3, "count": 7 } // 自动禁用（被警用）7 个 +] +``` + +### 4.4 按渠道类型分布 + +```bash +curl -sH "Authorization: Bearer $ADMIN_TOKEN" \ + -H "New-Api-User: $ADMIN_UID" \ + "$HOST/api/channel/?status=1&p=1&page_size=1" \ + | jq '.data.type_counts' +``` + +--- + +## 5. 完整健康检查脚本 + +将以下脚本保存为 `check_channels.sh`，配合 cron / systemd timer / Prometheus textfile collector 使用。 + +```bash +#!/usr/bin/env bash +# Usage: ADMIN_TOKEN=xxx ADMIN_UID=1 HOST=http://host:3000 ./check_channels.sh +# 退出码: 0 健康；1 正常渠道数 = 0；2 接口异常 +set -euo pipefail + +HOST="${HOST:?HOST is required}" +ADMIN_TOKEN="${ADMIN_TOKEN:?ADMIN_TOKEN is required}" +ADMIN_UID="${ADMIN_UID:?ADMIN_UID is required}" +MIN_HEALTHY="${MIN_HEALTHY:-1}" # 健康渠道最低阈值 + +call_page() { + local status="$1" + local page="$2" + curl -sf -m 10 \ + -H "Authorization: Bearer $ADMIN_TOKEN" \ + -H "New-Api-User: $ADMIN_UID" \ + "$HOST/api/channel/?status=$status&p=$page&page_size=100" \ + || { echo "API request failed for status=$status page=$page" >&2; exit 2; } +} + +fetch_all_items() { + local status="$1" + local page=1 + while :; do + resp=$(call_page "$status" "$page") + echo "$resp" | jq -c '.data.items[]' + total=$(echo "$resp" | jq -r '.data.total') + fetched=$((page * 100)) + [ "$fetched" -ge "$total" ] && break + page=$((page + 1)) + done +} + +ENABLED=$(call_page 1 1 | jq -r '.data.total') +DISABLED_ITEMS=$(fetch_all_items 0 | jq -s '.') +TOTAL=$(call_page -1 1 | jq -r '.data.total') + +MANUAL=$(echo "$DISABLED_ITEMS" | jq '[.[] | select(.status==2)] | length') +AUTO=$(echo "$DISABLED_ITEMS" | jq '[.[] | select(.status==3)] | length') + +cat <&2 + exit 1 +fi +exit 0 +``` + +--- + +## 6. 已验证响应（参考） + +针对实例 `http://74.50.127.249:3000/`（2026-05 测试）: + +``` +正常启用 (status=1) : 1 (type=14, Anthropic) +手动禁用 (status=2) : 0 +自动禁用 (status=3) : 1 ← 被警用 +渠道总数 : 2 +``` + +--- + +## 7. 注意事项 + +- `Authorization` 与 `New-Api-User` 两个 header **必须同时携带**，缺一返回 401；Header 名是标准 `Authorization`，不是 `authz` +- 普通用户的 access token 无法访问 `/api/channel/*`，必须为 admin / root 角色 +- access token 通过明文比对（`model.ValidateAccessToken`），**请按 secret 级别保管**，泄漏后立即重新生成覆盖 +- 接口返回的 `data.total` 是按过滤条件统计的**全量行数**，与分页参数 `page_size` 无关，因此用 `page_size=1` 取统计值最高效 +- `page_size` 最大会被限制为 `100`；需要遍历 `items` 做客户端统计时必须分页拉取 +- 大多数业务错误返回 HTTP 200 + `success:false`，但鉴权失败会返回 401，部分参数/上游错误可能返回 400/500；脚本应同时检查 HTTP 状态码和 `success` 字段 +- 自动禁用（status=3）的渠道，详细禁用原因记录在 Web 后台「渠道列表 → 已禁用」的 `tested_time` / `response_time` 与 channel 详情字段中；当前接口仅返回结构化字段，原始报错需另外查询 + +--- + +## 8. 批量按密钥精确查询渠道 + +### 端点 + +``` +POST /api/channel/search/keys +``` + +- 中间件：`AdminAuth` + `CriticalRateLimit` + `DisableCache` +- 用途：管理员粘贴多条渠道密钥后，按 **channel.key 精确匹配** 查询渠道行 +- 设计原因：密钥可能较多且属于敏感信息，因此使用 POST body，避免 URL 长度限制和 query string 泄漏 + +### 请求体 + +```jsonc +{ + "keys": ["sk-key-1", "sk-key-2"], // 必填；后端会 trim、去空行、去重 + "keyword": "OpenAI", // 可选；沿用 /api/channel/search 的 broad keyword 语义 + "group": "default", // 可选；与现有 group 过滤组合 + "model": "gpt-4o", // 可选；按 models LIKE 过滤 + "status": "enabled", // 可选：enabled / disabled / 空字符串(全部) + "type": 1, // 可选；用于类型 tab 过滤 + "id_sort": false, + "sort_by": "priority", // 可选：id/name/priority/balance/response_time/test_time + "sort_order": "desc", // 可选：asc / desc + "p": 1, + "page_size": 20, + "tag_mode": false // v1 不支持 true +} +``` + +### 响应 + +响应 shape 与 `GET /api/channel/search` 保持一致，仅返回：`items`、`total`、`type_counts`。 + +```json +{ + "success": true, + "message": "", + "data": { + "items": [ + { "id": 1, "name": "OpenAI-Primary", "key": "" } + ], + "total": 1, + "type_counts": { "1": 1, "14": 2 } + } +} +``` + +### 行为说明 + +- 精确匹配条件为 `channel.key IN keys`，不会把 `keyword` 命中的同名渠道当作密钥命中结果。 +- 过滤组合为：`精确密钥集合 AND keyword/group/model/status/type`。 +- `type_counts` 在应用 `type` 过滤前计算；也就是说类型 tab 统计反映当前密钥集合 + keyword/group/model/status 下的各类型数量。 +- 列表响应不会返回真实密钥；`items[].key` 会保持为空，真实密钥仍只能走受安全验证保护的单独密钥查看接口。 +- 请求没有硬性 key 数量限制；后端会把去重后的 keys 分块执行精确 `IN` 查询，避免 SQLite / MySQL / PostgreSQL 的 SQL 参数上限问题。 +- v1 不支持标签聚合模式：`tag_mode=true` 会返回 `success:false`，避免“某个 key 命中一个 tag 后返回同 tag 但 key 不匹配的渠道”的误导结果。 + +--- + +## 9. 添加渠道 + +### 端点 + +``` +POST /api/channel/ +``` + +- 中间件：`AdminAuth` +- 路由：`router/api-router.go:231` +- 处理函数：`controller.AddChannel`（`controller/channel.go:587`） +- 校验函数：`validateChannel`（`controller/channel.go:457`） + +### 请求体 + +```jsonc +{ + "mode": "single", // 必填，见下方"添加模式" + "multi_key_mode": "random", // 仅 multi_to_single 模式使用 + "batch_add_set_key_prefix_2_name": false, // batch 模式下是否给名字附加 key 前缀 + "channel": { // 必填，完整 Channel 对象 + "type": 1, + "name": "OpenAI-Primary", + "key": "sk-xxxxxxxx", + "models": "gpt-4o,gpt-4o-mini", + "group": "default", + "base_url": "https://api.openai.com", + "priority": 0, + "weight": 1, + "test_model": "gpt-4o-mini", + "model_mapping": "{}", // JSON 字符串 + "status_code_mapping": "{}", + "auto_ban": 1, // 1=失败时自动警用，0=保持启用 + "tag": "", + "remark": "", + "setting": "{}", + "param_override": "{}", + "header_override": "{}", + "other": "", // VertexAI 必填部署地区 JSON + "channel_info": { // 仅多 key 模式需要 + "is_multi_key": false, + "multi_key_mode": "", + "multi_key_size": 0 + } + } +} +``` + +完整 Channel 结构定义见 `model/channel.go:23`。 + +### 添加模式（`mode`） + +| mode | 含义 | +|---|---| +| `single` | 单 key 单渠道（最常用） | +| `multi_to_single` | 多个 key 合并到一个渠道，按 `multi_key_mode`（`random` / `polling`）轮询；keys 用 `\n` 分隔 | +| `batch` | 一次拆出多个独立渠道，每行 key 生成一条记录 | + +> ⚠️ 其他值会返回 `"不支持的添加模式"`。 + +### 关键校验规则 + +- `channel.key` 不能为空 +- 模型名长度必须 ≤ 255 +- `type=41`（VertexAI）：`other` 字段必填，且 JSON 中必须包含 `default` 区域；如使用 service account JSON，可用标准 JsonArray 批量导入 +- `type=57`（Codex）：`key` 必须是合法 JSON，包含 `access_token` 和 `account_id` + +### 常用 ChannelType 枚举 + +定义位置：`constant/channel.go` + +| Type | Provider | Type | Provider | +|---|---|---|---| +| 1 | OpenAI | 33 | AWS Bedrock | +| 3 | Azure | 34 | Cohere | +| 8 | Custom (OpenAI 兼容) | 37 | Dify | +| 14 | Anthropic Claude | 40 | SiliconFlow | +| 17 | 阿里通义 | 41 | VertexAI | +| 20 | OpenRouter | 42 | Mistral | +| 23 | 腾讯混元 | 43 | DeepSeek | +| 24 | Google Gemini | 45 | 火山豆包 | +| 25 | Moonshot | 48 | xAI | +| 27 | Perplexity | 57 | Codex | + +完整列表（含图像/视频类如 Suno、Kling、Jimeng、Vidu 等）见源码。 + +### 请求示例 + +#### 单渠道（最常用） + +```bash +curl -X POST "$HOST/api/channel/" \ + -H "Authorization: Bearer $ADMIN_TOKEN" \ + -H "New-Api-User: $ADMIN_UID" \ + -H "Content-Type: application/json" \ + -d '{ + "mode": "single", + "channel": { + "type": 1, + "name": "OpenAI-Primary", + "key": "sk-xxxxxxxx", + "models": "gpt-4o,gpt-4o-mini,gpt-3.5-turbo", + "group": "default", + "base_url": "https://api.openai.com", + "priority": 0, + "auto_ban": 1 + } + }' +``` + +#### 批量拆分（一次添加 3 个独立 Claude 渠道） + +```bash +curl -X POST "$HOST/api/channel/" \ + -H "Authorization: Bearer $ADMIN_TOKEN" \ + -H "New-Api-User: $ADMIN_UID" \ + -H "Content-Type: application/json" \ + -d '{ + "mode": "batch", + "batch_add_set_key_prefix_2_name": true, + "channel": { + "type": 14, + "name": "Claude", + "key": "sk-ant-key1\nsk-ant-key2\nsk-ant-key3", + "models": "claude-3-5-sonnet-20241022,claude-3-7-sonnet-latest", + "group": "default", + "priority": 0 + } + }' +``` + +启用 `batch_add_set_key_prefix_2_name` 后，每个新渠道名会变成 `Claude sk-ant-k`（取 key 前 8 字符），方便区分。 + +#### 多 Key 合并到一个渠道（轮询模式） + +```bash +curl -X POST "$HOST/api/channel/" \ + -H "Authorization: Bearer $ADMIN_TOKEN" \ + -H "New-Api-User: $ADMIN_UID" \ + -H "Content-Type: application/json" \ + -d '{ + "mode": "multi_to_single", + "multi_key_mode": "polling", + "channel": { + "type": 14, + "name": "Claude-Pool", + "key": "sk-ant-key1\nsk-ant-key2\nsk-ant-key3", + "models": "claude-3-5-sonnet-20241022", + "group": "default" + } + }' +``` + +### 响应 + +成功： + +```json +{ "success": true, "message": "" } +``` + +> 当前新增接口成功时不返回新渠道 ID。如需获取 ID，可按 `name` / `type` / `group` 等条件再查询列表匹配。 + +失败（如 key 为空、type 校验失败、JSON 格式错）： + +```json +{ "success": false, "message": "<中文错误描述>" } +``` + +> 多数业务校验失败会返回 HTTP 200 + `success:false`；鉴权失败或部分底层错误可能返回非 200。脚本应同时判断 HTTP 状态码和响应体的 `success`。 + +--- + +## 10. 其他管理接口 + +均在 `apiRouter.Group("/channel")` 下、`AdminAuth` 中间件保护（`router/api-router.go:218-257`）。其中 `POST /api/channel/:id/key` 额外需要 root 权限、安全验证和限流，`POST /api/channel/fetch_models` 额外需要 root 权限。 + +### 修改渠道 + +``` +PUT /api/channel/ +``` +处理：`controller.UpdateChannel`（`controller/channel.go:863`） + +请求体直接放 Channel 对象（不需要 `mode` 字段，必须带 `id`）： + +```bash +curl -X PUT "$HOST/api/channel/" \ + -H "Authorization: Bearer $ADMIN_TOKEN" \ + -H "New-Api-User: $ADMIN_UID" \ + -H "Content-Type: application/json" \ + -d '{ + "id": 7, + "type": 1, + "name": "OpenAI-Primary-Updated", + "key": "sk-newkey", + "models": "gpt-4o,gpt-4o-mini", + "priority": 10, + "status": 1 + }' +``` + +### 删除渠道 + +| 方法 | 路径 | 用途 | +|---|---|---| +| `DELETE` | `/api/channel/:id` | 删除单个渠道 | +| `POST` | `/api/channel/batch` | 批量删除，body: `{"ids":[1,2,3]}` | +| `DELETE` | `/api/channel/disabled` | 一键清理**所有已禁用**（手动+自动） | + +### 测试 / 余额 + +| 方法 | 路径 | 用途 | +|---|---|---| +| `GET` | `/api/channel/test/:id` | 测试指定渠道连通性 | +| `GET` | `/api/channel/test` | 测试全部启用中的渠道 | +| `GET` | `/api/channel/update_balance/:id` | 刷新单个渠道余额 | +| `GET` | `/api/channel/update_balance` | 刷新全部渠道余额 | + +### 按 tag 批量操作 + +| 方法 | 路径 | 用途 | +|---|---|---| +| `POST` | `/api/channel/tag/disabled` | 把某 tag 下全部渠道禁用 | +| `POST` | `/api/channel/tag/enabled` | 把某 tag 下全部渠道启用 | +| `PUT` | `/api/channel/tag` | 修改 tag 下渠道的批量配置 | + +### 模型与上游 + +| 方法 | 路径 | 用途 | +|---|---|---| +| `GET` | `/api/channel/models` | 全部内置模型列表 | +| `GET` | `/api/channel/models_enabled` | 启用中渠道支持的模型列表 | +| `GET` | `/api/channel/fetch_models/:id` | 拉取指定渠道上游真实支持的模型 | +| `POST` | `/api/channel/fetch_models` | 在创建前预拉取上游模型列表；额外需要 root 权限 | +| `POST` | `/api/channel/fix` | 修复 abilities 表（清理脏数据） | + +### 查询 / 复制 / 密钥查看 + +| 方法 | 路径 | 用途 | +|---|---|---| +| `GET` | `/api/channel/search` | 搜索渠道；支持 `keyword`、`group`、`model`、`status`、`type`、分页和排序参数 | +| `GET` | `/api/channel/:id` | 获取单个渠道详情；不返回 key | +| `POST` | `/api/channel/:id/key` | 查看渠道 key；额外需要 root 权限、安全验证、限流和禁用缓存 | +| `POST` | `/api/channel/copy/:id` | 复制渠道；支持 query：`suffix`、`reset_balance` | + +### tag / 多 key 扩展操作 + +| 方法 | 路径 | 用途 | +|---|---|---| +| `POST` | `/api/channel/batch/tag` | 批量设置渠道 tag，body: `{"ids":[1,2],"tag":"prod"}`；`tag:null` 可清空 | +| `GET` | `/api/channel/tag/models?tag=` | 获取某 tag 下模型列表最多的一条 `models` 字符串 | +| `POST` | `/api/channel/multi_key/manage` | 管理多 key 渠道，见下方动作列表 | + +`/api/channel/multi_key/manage` 请求体： + +```jsonc +{ + "channel_id": 1, + "action": "get_key_status", + "key_index": 0, + "page": 1, + "page_size": 50, + "status": 1 +} +``` + +支持的 `action`：`get_key_status`、`disable_key`、`enable_key`、`enable_all_keys`、`disable_all_keys`、`delete_key`、`delete_disabled_keys`。其中 `key_index` 仅单 key 操作需要；`delete_disabled_keys` 只删除自动禁用（status=3）的 key。 + +### Ollama 管理 + +| 方法 | 路径 | 用途 | +|---|---|---| +| `POST` | `/api/channel/ollama/pull` | 为 Ollama 渠道拉取模型，body: `{"channel_id":1,"model_name":"llama3"}` | +| `POST` | `/api/channel/ollama/pull/stream` | 流式拉取 Ollama 模型 | +| `DELETE` | `/api/channel/ollama/delete` | 删除 Ollama 模型 | +| `GET` | `/api/channel/ollama/version/:id` | 获取指定 Ollama 渠道版本 | + +### 上游模型更新 + +| 方法 | 路径 | 用途 | +|---|---|---| +| `POST` | `/api/channel/upstream_updates/apply` | 对单个渠道应用待处理的上游模型变更 | +| `POST` | `/api/channel/upstream_updates/apply_all` | 批量应用所有启用渠道的待处理上游模型变更 | + +### Codex OAuth + +| 方法 | 路径 | 用途 | +|---|---|---| +| `POST` | `/api/channel/codex/oauth/start` | 启动 Codex OAuth 流程（建渠道前） | +| `POST` | `/api/channel/codex/oauth/complete` | 完成回调 | +| `POST` | `/api/channel/:id/codex/oauth/start` | 已有渠道续期 OAuth | +| `POST` | `/api/channel/:id/codex/oauth/complete` | 已有渠道续期完成 | +| `POST` | `/api/channel/:id/codex/refresh` | 刷新已有 Codex 渠道凭证 | +| `GET` | `/api/channel/:id/codex/usage` | 获取 Codex 渠道用量 | + +--- + +## 10. 修改 / 删除接口的最小完整示例 + +```bash +# 0) 公共环境 +HOST="http://your-host:3000" +ADMIN_TOKEN="..." +ADMIN_UID=1 +H=(-H "Authorization: Bearer $ADMIN_TOKEN" -H "New-Api-User: $ADMIN_UID" -H "Content-Type: application/json") + +# 1) 新增（成功响应不返回 id，只返回 success/message） +ADD_OK=$(curl -sX POST "$HOST/api/channel/" "${H[@]}" -d '{ + "mode":"single", + "channel":{"type":1,"name":"tmp","key":"sk-test","models":"gpt-4o-mini","group":"default"} +}' | jq -r '.success') +echo "added: $ADD_OK" + +# 2) 再查询列表找到 id（page_size 最大 100） +ID=$(curl -s "${H[@]}" "$HOST/api/channel/search?keyword=tmp&status=1&page_size=100" \ + | jq -r '.data.items[] | select(.name=="tmp") | .id' | head -1) + +# 3) 测试连通性 +curl -s "${H[@]}" "$HOST/api/channel/test/$ID" | jq + +# 4) 删除 +curl -sX DELETE "${H[@]}" "$HOST/api/channel/$ID" | jq +``` diff --git a/docs/reviews/channel-cost-statistics-template-plan-critique-2026-06-03.md b/docs/reviews/channel-cost-statistics-template-plan-critique-2026-06-03.md new file mode 100644 index 000000000000..1838d6bf4bda --- /dev/null +++ b/docs/reviews/channel-cost-statistics-template-plan-critique-2026-06-03.md @@ -0,0 +1,35 @@ +# 成本统计模板计划 — 实施前批判（≤1 页） + +**范围**：审阅 `docs/plans/channel-cost-statistics-template-2026-06-03.md`，对照 context_builder 导出 `prompt-exports/oracle-plan-2026-06-03-132137-cost-report-plan-851-b6dc.md`。仅覆盖下列 5 点，不扩范围、不重写计划。 + +## 1. 三个最欠规格的接缝（实现者只能靠猜） + +1. **默认模板 → 样例 Excel 的字段/公式映射缺失。** 计划反复说"对齐样例 20 个表头"（plan:42, plan:87），Item 1 done-when 只写"包含核心字段"（plan:313），但从未把 20 列逐一标注为 dimension/metric/manual/formula，也未给出 成本/应收账款/中间利润/火力利润/利润比例 的具体公式。这是业务语义，实现者无从猜起，却是 Item 1（seed 默认模板）的交付前提。**这是最大的接缝。** +2. **`Log.Other` 的结构未定义。** 聚合第 3–4 步要"解析 Other"（plan:179），分类默认用 `log.other.claude == true`、field source 用 `log_other.*`（plan:131, plan:214, plan:234）——但 `Log.Other` 里并不存在 `claude` 布尔键。导出曾点名 `service/log_info_generate.go` / `types/price_data.go` 是 Other 的写入处（export:13, export:30），计划把这条线索丢了，实现者得逆向猜测 Other 的键名与含义。 +3. **`period_key` 与按时间分桶的语义未定。** `PeriodStart/PeriodEnd` 是区间，`PeriodKey` 是单值，row_key 里又含 `period:{period_key}`（plan:142, plan:192）。当 `period_mode=day` 而日期范围跨周时，一个 run 是产 1 个 period 还是 7 个 period 桶？聚合算法第 7–8 步生成 row_key 却从不按时间子窗口分桶（plan:184–195）。这决定 grouping 是否需要时间维度，影响 Item 1 的表结构。 + +## 2. 规格颗粒度失衡 + +- **过度规定（应交给实现代理）**：列出 10 个具体 React 组件文件名（plan:277–293）属于 UI 拆分的战术决策；REST 路由的精确嵌套（如 `.../versions/:version_id/validate`，plan:165）同样可由实现者定形。 +- **丢失了导出里有用的框架**：(a) 导出给了具体的公式环境变量与示例（`revenue_usd = quota / quota_per_unit` 等，export:624–642），计划 3.6 节抽象掉了，反而更难落地。(b) 导出明确建议**复用 usage-logs 既有组件**（`useUsageLogsData` 的列偏好/统计请求、`ColumnSelectorModal`、`UsageLogsColumnDefs` 成本渲染，export:17/33）；计划却凭空新建一整套组件树（plan:277–293），既过度规定又丢了复用锚点。 + +## 3. 矛盾与缺失依赖 + +- **重复造轮子 / 缺失依赖（已核实）**：计划称公式引擎用"现有 Go 依赖 `expr-lang/expr`"（plan:249），但**仓库已存在 `pkg/billingexpr`** —— 一套用同一引擎做好编译缓存、变量白名单、版本标签、settle 的计费表达式系统（CLAUDE.md Rule 7 强制先读 `pkg/billingexpr/expr.md`）。计划要新写 `service/cost_report/formula.go` 却从不提它。应先评估复用，而非平行实现。 +- **矛盾：手动值"跨模板升级存活" vs row_key 依赖 grouping。** 计划说 manual cell 不绑 `TemplateVersionId`，以便字段 key 不变时保留手动值（plan:152）；但 row_key 由"启用的 grouping 维度"派生并含 `template_id/period_key`（plan:192–195）。一旦模板改了 grouping（如新增 model 维度），row_key 格式即变，旧手动值静默孤立。存活性其实取决于 row_key 稳定，而非字段 key。 + +## 4. 过度规划（建议删减/简化） + +- **同一条顺序被编码三遍**：Phase A–D（plan:84–97）+ Work Items 1–5（plan:309–367）+ 导出的 15 步实施序（export:906–980）。保留 Work Items 即可，Phase 段可压成一句。 +- **Univer 被当作已交付特性来规划**：明确推迟（plan:81/301）的功能却铺了数据模型映射、迁移路径、References、Open Question（plan:95–97, 299–305, 371, 376–381）。收敛为"推迟 + OSS/Pro 边界待验证"一段即可。 +- Excel 的 Sheet 2/3/4 内部布局（plan:261–264）可简化为"附模板与规则元信息页"。 + +## 5. 会改变实施顺序的问题 + +1. **`pkg/billingexpr` 能否直接满足公式需求？** 若能，Item 2 的 `formula.go` 退化为"接入 billingexpr"，公式校验可大幅前移，改变 Item 1/2 的边界与排序。 +2. **样例 Excel 的成本/利润公式与业务语义到底是什么？** 若 V1 必须与样例一致，则"默认模板 seed"（实施早期步骤）在产品确认公式前无法落地，须把 seeding 排到公式语义敲定之后，阻塞 Item 1 的 done-when。 +3. **一个 run 是单时间桶还是多时间桶？** 答案决定 row_key 是否需含时间维度、聚合是否按 period 循环——必须在冻结 Item 1 表结构与 row_key 方案之前回答。 +4. **手动单元格身份是否必须跨 grouping 变更存活？** 若必须，则需在 Item 1 设计稳定的代理 row id（而非由 grouping 派生），把 schema 设计排到聚合之前。 + +--- +*结论：计划骨架合理，但落地前必须先钉死「默认模板字段+公式语义」「`Log.Other` 结构」「period 分桶」三处，并就「复用 `pkg/billingexpr`」与「row_key 稳定性」做出决定——这两项会直接重排工作项顺序。* diff --git a/docs/reviews/channel-preparation-auto-promotion-plan-critique-2026-06-02.md b/docs/reviews/channel-preparation-auto-promotion-plan-critique-2026-06-02.md new file mode 100644 index 000000000000..202ec6a20540 --- /dev/null +++ b/docs/reviews/channel-preparation-auto-promotion-plan-critique-2026-06-02.md @@ -0,0 +1,34 @@ +# Critique — Channel Preparation Auto-Promotion Plan (V1) + +**Scope:** Reviews `docs/plans/channel-preparation-auto-promotion-2026-06-02.md` against the context_builder export `prompt-exports/oracle-plan-2026-06-02-181619-auto-promotion-plan-e516.md`. Verdict: solid and executable; the issues below are gaps an implementer would otherwise guess at. The `promoteChannelPreparation` reuse + "reset cache once per run" claim was spot-checked against `controller/channel_preparation.go:249-346` and is **correct** (per-call tx, cache reset only in handlers) — not a contradiction. + +## 1. Top 3 under-specified seams + +1. **Cross-DB capacity aggregation by group token.** Item 2 says "reuse or mirror existing exact comma-delimited group filtering" (`controller/channel.go:53-109`), but that helper builds a *list* query; the capacity metric needs a `SUM(balance)`/`SUM(used_quota)` aggregate filtered to an exact token inside a comma-delimited `group` column. No SQL shape is given, and per CLAUDE.md Rule 2 this must work on SQLite/MySQL/PostgreSQL (`commonGroupCol`, no PG-only operators). Implementer will guess the aggregation query. +2. **Candidate selection over preparations.** The plan anchors to `model/ability.go:61-140` ("`weight + 10`"), but that path selects live *abilities*, not `ChannelPreparation` rows — it cannot be literally reused. Item 2 offers "or an equivalent deterministic testable helper" without specifying where the helper lives (model vs controller), the exact priority-tier-then-weighted algorithm, or the randomness-injection seam for tests. Pure guesswork as written. +3. **Concurrency contract.** Item 3 says a process-local mutex "returns an in-progress response rather than overlapping," but doesn't specify: where the mutex lives, the manual endpoint's busy response (HTTP 409 vs 200-with-status — Item 5 frontend needs this), or how a scheduler tick behaves when a manual run holds the lock (skip vs block). + +## 2. Specificity balance (plan vs. export) + +- **Over-specified, should be the impl agent's call:** UI layout ("above or near the action toolbar without covering table content"); the verbatim option-key strings and full Go struct field list; `MaxPromotionsPerRun *int` *per-rule* on top of the global cap — extra config surface with no V1 need (see §3). +- **Useful export framing dropped:** the export gave the capacity metric explicit *assumptions/failure cases* and a dedicated **Risks & migration** section; the plan compresses this to "Open Questions: None," which is over-confident given the metric reversal below. Worth restoring a 3-line risk note (stale balance, conservative metric, multi-node overshoot). +- **Reversal worth flagging:** the export recommended `balance_usd` and argued subtracting `UsedQuota` "risks double-counting after balance refresh." The plan defaults to `balance_minus_used_quota_usd` — and its own Background calls `balance` the provider *remaining* amount, making the subtraction internally tense. Keeping the metric an enum is good; the default choice deserves the export's caveat carried forward, not deleted. + +## 3. Contradictions / missing dependencies + +- **Global vs per-rule promotion limit precedence is undefined.** Item 1 has `MaxPromotionsPerRun` (global) and a per-rule `*int`; Item 3 says "run/rule limits" but never states whether the global cap is a per-run total across all rules or a per-rule fallback. Directly governs Item 3 control flow. +- **Permission split.** Manual trigger is `AdminAuth`; rule config persists via root-only `/api/option/`. The plan flags this but leaves the resulting UX (admins can *run* but not *see/configure* rules) unresolved — Item 5 needs a decision on gating the panel vs. graceful degradation. +- Item 6's per-candidate "capacity before/after" logging requires the run service (Item 3) to track capacity at candidate granularity; Item 3 only commits to per-*rule* initial/final. Make Item 3 expose per-step deltas or downgrade Item 6's logging granularity. + +## 4. Over-planning to cut/simplify + +- Per-rule `MaxPromotionsPerRun` (§3) — drop for V1; the global cap suffices. +- The `capacity_metric` future-migration narrative and "future strategy table" framing — keep the enum field, cut the prose; it's speculative for a V1 with one supported value each. +- Item 7's "concurrency tests cover double promotion attempts against the same candidate" — a full concurrency harness is heavy; the conditional `pending->promoting` update (`channel_preparation.go:261-273`) already guarantees this. Reduce to one focused test asserting `RowsAffected==0` on the second attempt. + +## 5. Questions that would change implementation order + +1. **Is `balance_minus_used_quota_usd` the final V1 metric, or `balance_usd`?** If the simpler metric wins, Item 2's quota conversion (`common.QuotaPerUnit`) and parts of Item 7 disappear and Item 2 shrinks. +2. **Is the hard-delete-on-promote invariant truly fixed?** The plan reverses the export's "promoted audit row" recommendation and therefore *deletes* the export's Item 6 (prep status visibility). If the user later wants audit rows, that frontend status-filter work returns and reorders Items 5–7. +3. **Do rules live on the prep-pool page (plan) or in monitoring settings (export left open)?** Monitoring settings would reuse existing save/load plumbing and let the settings work proceed independently of the prep-pool UI, changing Item 5's sequencing. +4. **Is multi-node overshoot acceptable?** If not, a Redis/distributed lock becomes a prerequisite for Items 3–4 rather than a deferred risk. diff --git a/docs/reviews/channel-preparation-pool-plan-critique-2026-06-02.md b/docs/reviews/channel-preparation-pool-plan-critique-2026-06-02.md new file mode 100644 index 000000000000..9b1a81b9eb2f --- /dev/null +++ b/docs/reviews/channel-preparation-pool-plan-critique-2026-06-02.md @@ -0,0 +1,39 @@ +# Plan Critique — Channel Preparation Pool (V1) + +**Scope:** Critique of `docs/plans/channel-preparation-pool-2026-06-02.md` against its source `context_builder` export (`prompt-exports/oracle-plan-2026-06-02-115813-prep-pool-plan-02d74-0ac7.md`). Five focus areas only; no scope expansion. + +## 1. Top 3 under-specified seams + +1. **Item 2 shared-helper fork is unresolved (plan line 103).** Done-when offers an OR — share `single/batch/multi_to_single` semantics in one helper *or* limit promotion to single-record while sharing validation. This is the highest-risk decision (refactor live `AddChannel` at `controller/channel.go:586-691` vs. a parallel helper) and the plan punts it to the implementer. The export was decisive ("only single prep → single live channel; batch loops over records"). Pin this. +2. **Multi-key staging/promotion is promised but never specified.** Background elevates "account/key-level staging and channel-level promotion" to a *user decision* (lines 11, 19) and hints at `ChannelInfo`/multi-key metadata (line 40, ref `model/channel.go:22-76`), yet the model field list, the promotion algorithm, and Item 4 done-when (line 141) only cover single-record → single channel+abilities. Implementer must guess the schema and promotion path. Either explicitly cut multi-key from V1 or specify it. +3. **Key sanitization shape + edit-preserve rule undefined.** Plan says "key preview" (line 43) and "do not expose full keys" (Item 3, line 124) but never defines the preview format nor the PUT behavior. Export specified "first 8–12 chars + ellipsis" and "blank/missing key preserves stored key." Without the preserve rule, the edit endpoint (line 52) risks silently wiping stored keys. Pin both. + +## 2. Specificity balance + +- **Over-specified (agent should own):** Item 6 done-when hardcodes the entire column list and filter set as acceptance criteria (lines 184-186); Item 4/API shape bakes exact route strings (`batch/promote`). Fine as guidance, wrong as gating criteria. +- **Dropped useful framing vs. export:** + - Status enum lost its numeric values + "promotable" matrix — keep at least the promotable mapping. + - List response dropped `status_counts`/`type_counts` even though Item 6 filters by status and would want count badges. + - Promotion concurrency guard compressed from export's explicit "row-lock / conditional status update, require `Status==pending`" down to a bare done-when assertion (line 143). Keep the mechanism hint. + +## 3. Contradictions / missing dependencies + +- **Item 2 dependency is mislabeled (line 112).** It lists "Item 1 for promotion consumers," but the live-channel creation helper only refactors `AddChannel` and never touches the prep model — it has no real dependency on Item 1 and can lead. This distorts the critical path. +- **Delete vs. archive inconsistency.** Item 3 route is "DELETE … or archive endpoint" (line 53) and done-when says "archive/delete" (line 121), but Item 6 row actions offer only "archive" (line 186). Unclear whether a hard delete exists; resolve to one. +- **Multi-key** (seam 2) is also a Background-vs-Work-Items contradiction. +- **Verified OK:** Item 4's `service/http_client.go` for cache reset is correct — `ResetProxyClientCache()` is defined there (`service/http_client.go:74`). + +## 4. Over-planning to cut / simplify + +- **Collapse UI Items 5/6/7** (shell S + table L + modals M). The shell-vs-table split forces artificial serialization for one classic page; merge to ≤2 items. +- **Trim Item 8:** drop the optional `docs/channel/channel-api.md` touch-up and the "if API documentation is updated…" clause (lines 239, 246) from V1 acceptance; keep only the routing-contract verification checklist. +- **Background provenance noise:** the "(reported by scout / seam probe)" tags (lines 17-25) and the References block (lines 258-267) duplicate the inline file:line refs — process artifacts that don't belong in a final plan. + +## 5. Questions that change implementation order + +1. **Is multi-key staging in V1 scope?** Yes → Item 1 schema gains `ChannelInfo`/multi-key metadata and Item 4 gains multi-key→channel logic; both grow and Item 1 must settle first. No → simplest path; pin it now. +2. **Refactor `AddChannel` vs. parallel helper (Item 2 fork)?** Refactor → Item 2 is a tested prerequisite on the live path, must land before Item 4, not parallelizable. Parallel helper → Item 2 is independent and can move late. +3. **Must promotion ship in the first deliverable?** Item 3 (CRUD/import) explicitly does *not* need Item 2 (line 132), so a storage-only MVP could ship first and defer Items 2+4 — reorders toward UI-early. +4. **Is default-frontend parity required in V1 (Open Q1)?** Yes → adds a parallel UI track and roughly doubles frontend scope; currently punted to classic-only. + +*(Open Q2 — archive terminal vs. restore — does not affect order, per the plan itself; excluded.)* diff --git a/docs/reviews/usage-log-excel-export-plan-critique-2026-06-03.md b/docs/reviews/usage-log-excel-export-plan-critique-2026-06-03.md new file mode 100644 index 000000000000..63e6ed3ced8d --- /dev/null +++ b/docs/reviews/usage-log-excel-export-plan-critique-2026-06-03.md @@ -0,0 +1,65 @@ +# Critique — Usage Log Excel Export Plan (2026-06-03) + +**Scope reviewed:** `docs/plans/usage-log-excel-export-2026-06-03.md` vs. its source export +`prompt-exports/oracle-plan-2026-06-03-161456-usage-log-export-0ab-8658.md`. Per request, this +covers only the five axes below; no scope expansion. + +## 1. Top 3 under-specified seams (implementer would guess) + +1. **Item 2 refactor boundary + total-count path.** "Extract reusable query builders" from + `model/log.go:328-480`, `model/midjourney.go`, `model/task.go` never says *what shape* to extract + (shared `applyFilters(db, params)` vs. a new `num=-1` flag vs. a parallel function). Worse, the + 100000-row cap needs a total count, but the only documented count path (`GetUserLogs`) wraps + `Count` in `Limit(logSearchCountLimit)` = 10000 (`model/log.go:415,446`). The plan never says how + export obtains an accurate total above 10000. An implementer must invent the counting strategy. +2. **Drawing/task filter-param contract.** Background pins the *common* param mapping precisely + (`buildApiParams`, `utils.ts:170-249`; `controller/log.go:13-55`) but gives **no** equivalent for + drawing/task — `buildBaseParams` and the midjourney/task controllers are named only in Item 6's + "Done when," with no file:line and no field list. Yet Items 2-4 require those exact params. The + export's `` *did* carry this (GetAllTasks/GetAllUserTask, tasksToDto, ms-vs-s + timestamps); the plan dropped it, leaving two of three categories under-specified. +3. **Default-selection vs. live column visibility.** Plan says defaults "mirror current table + columns," but `view-options.tsx:55-73` visibility is client state while `default_selected` is + backend-owned. Nobody is assigned to reconcile a user's toggled-off columns with the backend + schema, nor is the column-id ↔ backend `key` mapping defined. Implementer guesses the source of truth. + +## 2. Specificity balance (over-spec vs. dropped framing) + +- **Dropped useful framing:** the export listed concrete export-only field candidates per category + (e.g. common `request_id`/`upstream_request_id`/`other_json` admin-only; drawing `finish_time`/ + `video_url`/`raw_status`; task `result_url`/`data_json`). The plan deleted all of these and deferred + to "Item 1 评审." Re-attaching them as a non-binding starter list would save the implementer a + rediscovery pass. (Conversely, the plan *added* value the export lacked: file:line anchors, the + `logGroupCol` DB-compat note, the `logSearchCountLimit` note, the cost-stats critique reference.) +- **Over-specifies tactical choices the agent should own:** prescribing a 4-file service split + (`export.go`/`common.go`/`drawing.go`/`task.go`), the exact prop name `extraActions?: ReactNode`, + and batch size "1000" are implementation details, not plan-level decisions. One service file and a + free-hand batch size are fine to leave open. + +## 3. Contradictions / missing dependencies + +- **Count cap contradiction** (see seam 1): 10000 self-count limit vs. 100000 export cap is unresolved. +- **Timezone has no frontend owner.** Backend accepts `timezone` (Approach §API contract), but Item 6's + "Done when" never lists sending the browser tz, and no item sources it. Cross-item gap. +- **Spurious dependency:** Item 2 (query builders) "Dependencies: Item 1 (field schemas)" — these are + independent; query extraction does not need the schema. The false edge forces serial work. + +## 4. Over-planning to cut/simplify + +- **Six `export_fields` endpoints** (per category × admin/self). Validation only needs one schema + source; collapse to `/api/log/export_fields?category=&scope=` (or even ship schema client-side with + server-side key validation) — drops endpoint count by ~half. +- **Item 7 "记忆用户上次选择" (localStorage persistence)** is V1 scope creep; cut it. +- **All-three-categories acceptance in V1** while the plan itself permits "按 category 分层交付" — the + hard acceptance bar over-commits; see Q below. + +## 5. Questions that would change implementation order + +1. **Is V1 truly all three sections, or common-first?** Common-first shrinks Items 2-4 by ~2/3 and + makes drawing/task a follow-up — the single biggest ordering lever. (Open Question already half-asks this.) +2. **Does a global export/row-limit config already exist?** (Plan's Open Question #2.) If yes, resolve + before Item 1 — it removes the new-constant work and may reshape Item 3's cap. +3. **Must self export count/export beyond 10000?** Answer dictates whether a new count path must land in + Item 2 *before* Item 3 — reorders the backend chain. +4. **Should defaults follow currently-visible columns?** If yes, the frontend must pass visible column + ids, pulling the field-contract decision (Items 6-7) earlier than its current tail position. diff --git a/go.mod b/go.mod index eceb5e7d8889..1c6c1feb2574 100644 --- a/go.mod +++ b/go.mod @@ -27,6 +27,7 @@ require ( github.com/golang-jwt/jwt/v5 v5.3.0 github.com/google/uuid v1.6.0 github.com/gorilla/websocket v1.5.0 + github.com/xuri/excelize/v2 v2.10.1 github.com/grafana/pyroscope-go v1.2.7 github.com/jfreymuth/oggvorbis v1.0.5 github.com/jinzhu/copier v0.4.0 @@ -48,11 +49,11 @@ require ( github.com/tiktoken-go/tokenizer v0.6.2 github.com/waffo-com/waffo-go v1.3.1 github.com/yapingcat/gomedia v0.0.0-20240906162731-17feea57090c - golang.org/x/crypto v0.45.0 + golang.org/x/crypto v0.48.0 golang.org/x/image v0.38.0 - golang.org/x/net v0.47.0 + golang.org/x/net v0.50.0 golang.org/x/sync v0.20.0 - golang.org/x/sys v0.38.0 + golang.org/x/sys v0.41.0 golang.org/x/text v0.35.0 gopkg.in/yaml.v3 v3.0.1 gorm.io/driver/mysql v1.4.3 @@ -62,6 +63,14 @@ require ( require github.com/waffo-com/waffo-pancake-sdk-go v0.3.1 +require ( + github.com/richardlehane/mscfb v1.0.6 // indirect + github.com/richardlehane/msoleps v1.0.6 // indirect + github.com/tiendc/go-deepcopy v1.7.2 // indirect + github.com/xuri/efp v0.0.1 // indirect + github.com/xuri/nfp v0.0.2-0.20250530014748-2ddeb826f9a9 // indirect +) + require ( github.com/DmitriyVTitov/size v1.5.0 // indirect github.com/anknown/darts v0.0.0-20151216065714-83ff685239e6 // indirect diff --git a/go.sum b/go.sum index 1eb08878e607..c8e92abb7fab 100644 --- a/go.sum +++ b/go.sum @@ -246,6 +246,10 @@ github.com/prometheus/procfs v0.15.1 h1:YagwOFzUgYfKKHX6Dr+sHT7km/hxC76UB0leargg github.com/prometheus/procfs v0.15.1/go.mod h1:fB45yRUv8NstnjriLhBQLuOUt+WW4BsoGhij/e3PBqk= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec h1:W09IVJc94icq4NjY3clb7Lk8O1qJ8BdBEF8z0ibU0rE= github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec/go.mod h1:qqbHyh8v60DhA7CoWK5oRCqLrMHRGoxYCSS9EjAz6Eo= +github.com/richardlehane/mscfb v1.0.6 h1:eN3bvvZCp00bs7Zf52bxNwAx5lJDBK1tCuH19qq5aC8= +github.com/richardlehane/mscfb v1.0.6/go.mod h1:pe0+IUIc0AHh0+teNzBlJCtSyZdFOGgV4ZK9bsoV+Jo= +github.com/richardlehane/msoleps v1.0.6 h1:9BvkpjvD+iUBalUY4esMwv6uBkfOip/Lzvd93jvR9gg= +github.com/richardlehane/msoleps v1.0.6/go.mod h1:BWev5JBpU9Ko2WAgmZEuiz4/u3ZYTKbjLycmwiWUfWg= github.com/rogpeppe/go-internal v1.6.1/go.mod h1:xXDCJY+GAPziupqXw64V24skbSoqbTEfhy4qGm1nDQc= github.com/rogpeppe/go-internal v1.8.0/go.mod h1:WmiCO8CzOY8rg0OYDC4/i/2WRWAB6poM+XZ2dLUbcbE= github.com/rogpeppe/go-internal v1.10.0 h1:TMyTOH3F/DB16zRVcYyreMH6GnZZrwQVAoYjRBZyWFQ= @@ -292,6 +296,8 @@ github.com/tidwall/pretty v1.2.0 h1:RWIZEg2iJ8/g6fDDYzMpobmaoGh5OLl4AXtGUGPcqCs= github.com/tidwall/pretty v1.2.0/go.mod h1:ITEVvHYasfjBbM0u2Pg8T2nJnzm8xPwvNhhsoaGGjNU= github.com/tidwall/sjson v1.2.5 h1:kLy8mja+1c9jlljvWTlSazM7cKDRfJuR/bOJhcY5NcY= github.com/tidwall/sjson v1.2.5/go.mod h1:Fvgq9kS/6ociJEDnK0Fk1cpYF4FIW6ZF7LAe+6jwd28= +github.com/tiendc/go-deepcopy v1.7.2 h1:Ut2yYR7W9tWjTQitganoIue4UGxZwCcJy3orjrrIj44= +github.com/tiendc/go-deepcopy v1.7.2/go.mod h1:4bKjNC2r7boYOkD2IOuZpYjmlDdzjbpTRyCx+goBCJQ= github.com/tiktoken-go/tokenizer v0.6.2 h1:t0GN2DvcUZSFWT/62YOgoqb10y7gSXBGs0A+4VCQK+g= github.com/tiktoken-go/tokenizer v0.6.2/go.mod h1:6UCYI/DtOallbmL7sSy30p6YQv60qNyU/4aVigPOx6w= github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU= @@ -316,6 +322,12 @@ github.com/waffo-com/waffo-pancake-sdk-go v0.3.1 h1:ngQSN/oVB35xTwFPLfg++bxPC+Sp github.com/waffo-com/waffo-pancake-sdk-go v0.3.1/go.mod h1:OB2MyFIQaefoPO0FV3J+yu9sDP8RVFQ+sbFsXqGuObc= github.com/x448/float16 v0.8.4 h1:qLwI1I70+NjRFUR3zs1JPUCgaCXSh3SW62uAKT1mSBM= github.com/x448/float16 v0.8.4/go.mod h1:14CWIYCyZA/cWjXOioeEpHeN/83MdbZDRQHoFcYsOfg= +github.com/xuri/efp v0.0.1 h1:fws5Rv3myXyYni8uwj2qKjVaRP30PdjeYe2Y6FDsCL8= +github.com/xuri/efp v0.0.1/go.mod h1:ybY/Jr0T0GTCnYjKqmdwxyxn2BQf2RcQIIvex5QldPI= +github.com/xuri/excelize/v2 v2.10.1 h1:V62UlqopMqha3kOpnlHy2CcRVw1V8E63jFoWUmMzxN0= +github.com/xuri/excelize/v2 v2.10.1/go.mod h1:iG5tARpgaEeIhTqt3/fgXCGoBRt4hNXgCp3tfXKoOIc= +github.com/xuri/nfp v0.0.2-0.20250530014748-2ddeb826f9a9 h1:+C0TIdyyYmzadGaL/HBLbf3WdLgC29pgyhTjAT/0nuE= +github.com/xuri/nfp v0.0.2-0.20250530014748-2ddeb826f9a9/go.mod h1:WwHg+CVyzlv/TX9xqBFXEZAuxOPxn2k1GNHwG41IIUQ= github.com/xyproto/randomstring v1.0.5 h1:YtlWPoRdgMu3NZtP45drfy1GKoojuR7hmRcnhZqKjWU= github.com/xyproto/randomstring v1.0.5/go.mod h1:rgmS5DeNXLivK7YprL0pY+lTuhNQW3iGxZ18UQApw/E= github.com/yapingcat/gomedia v0.0.0-20240906162731-17feea57090c h1:xA2TJS9Hu/ivzaZIrDcwvpJ3Fnpsk5fDOJ4iSnL6J0w= @@ -331,6 +343,8 @@ golang.org/x/arch v0.21.0/go.mod h1:dNHoOeKiyja7GTvF9NJS1l3Z2yntpQNzgrjh1cU103A= golang.org/x/crypto v0.0.0-20210711020723-a769d52b0f97/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc= golang.org/x/crypto v0.45.0 h1:jMBrvKuj23MTlT0bQEOBcAE0mjg8mK9RXFhRH6nyF3Q= golang.org/x/crypto v0.45.0/go.mod h1:XTGrrkGJve7CYK7J8PEww4aY7gM3qMCElcJQ8n8JdX4= +golang.org/x/crypto v0.48.0 h1:/VRzVqiRSggnhY7gNRxPauEQ5Drw9haKdM0jqfcCFts= +golang.org/x/crypto v0.48.0/go.mod h1:r0kV5h3qnFPlQnBSrULhlsRfryS2pmewsg+XfMgkVos= golang.org/x/exp v0.0.0-20250620022241-b7579e27df2b h1:M2rDM6z3Fhozi9O7NWsxAkg/yqS/lQJ6PmkyIV3YP+o= golang.org/x/exp v0.0.0-20250620022241-b7579e27df2b/go.mod h1:3//PLf8L/X+8b4vuAfHzxeRUl04Adcb341+IGKfnqS8= golang.org/x/image v0.38.0 h1:5l+q+Y9JDC7mBOMjo4/aPhMDcxEptsX+Tt3GgRQRPuE= @@ -341,6 +355,8 @@ golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v golang.org/x/net v0.0.0-20210520170846-37e1c6afe023/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y= golang.org/x/net v0.47.0 h1:Mx+4dIFzqraBXUugkia1OOvlD6LemFo1ALMHjrXDOhY= golang.org/x/net v0.47.0/go.mod h1:/jNxtkgq5yWUGYkaZGqo27cfGZ1c5Nen03aYrrKpVRU= +golang.org/x/net v0.50.0 h1:ucWh9eiCGyDR3vtzso0WMQinm2Dnt8cFMuQa9K33J60= +golang.org/x/net v0.50.0/go.mod h1:UgoSli3F/pBgdJBHCTc+tp3gmrU4XswgGRgtnwWTfyM= golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= golang.org/x/sys v0.0.0-20190726091711-fc99dfbffb4e/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs= @@ -356,6 +372,8 @@ golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.38.0 h1:3yZWxaJjBmCWXqhN1qh02AkOnCQ1poK6oF+a7xWL6Gc= golang.org/x/sys v0.38.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= +golang.org/x/sys v0.41.0 h1:Ivj+2Cp/ylzLiEU89QhWblYnOE9zerudt9Ftecq2C6k= +golang.org/x/sys v0.41.0/go.mod h1:OgkHotnGiDImocRcuBABYBEXf8A9a87e/uXjp9XT3ks= golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo= golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8= golang.org/x/text v0.3.2/go.mod h1:bEr9sfX3Q8Zfm5fL9x+3itogRgK3+ptLWKqgva+5dAk= diff --git a/main.go b/main.go index 3361b8ce9338..2d4d0a4bc7d6 100644 --- a/main.go +++ b/main.go @@ -112,6 +112,7 @@ func main() { } go controller.AutomaticallyTestChannels() + controller.StartChannelPreparationAutoPromotionTask() // Codex credential auto-refresh check every 10 minutes, refresh when expires within 1 day service.StartCodexCredentialAutoRefreshTask() diff --git a/middleware/auth.go b/middleware/auth.go index 23d933fbe0c1..72eec75bd8b3 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -162,6 +162,16 @@ func TryUserAuth() func(c *gin.Context) { id := session.Get("id") if id != nil { c.Set("id", id) + if role := session.Get("role"); role != nil { + c.Set("role", role) + } + if username := session.Get("username"); username != nil { + c.Set("username", username) + } + if group := session.Get("group"); group != nil { + c.Set("group", group) + c.Set("user_group", group) + } } c.Next() } diff --git a/middleware/distributor.go b/middleware/distributor.go index 258aebb57037..fa302fdd1c88 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -101,33 +101,9 @@ func Distribute() func(c *gin.Context) { } } - if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { - affinityUsable := false - preferred, err := model.CacheGetChannel(preferredChannelID) - if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled { - if usingGroup == "auto" { - userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) - autoGroups := service.GetUserAutoGroup(userGroup) - for _, g := range autoGroups { - if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) { - selectGroup = g - common.SetContextKey(c, constant.ContextKeyAutoGroup, g) - channel = preferred - affinityUsable = true - service.MarkChannelAffinityUsed(c, g, preferred.Id) - break - } - } - } else if model.IsChannelEnabledForGroupModel(usingGroup, modelRequest.Model, preferred.Id) { - channel = preferred - selectGroup = usingGroup - affinityUsable = true - service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id) - } - } - if !affinityUsable && !service.ShouldKeepChannelAffinityOnChannelDisabled() { - service.ClearCurrentChannelAffinityCache(c) - } + if preferred, selectedGroup, found := service.GetUsablePreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found { + channel = preferred + selectGroup = selectedGroup } if channel == nil { diff --git a/middleware/distributor_channel_affinity_test.go b/middleware/distributor_channel_affinity_test.go new file mode 100644 index 000000000000..07f7f53cd40b --- /dev/null +++ b/middleware/distributor_channel_affinity_test.go @@ -0,0 +1,196 @@ +package middleware + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupDistributorAffinityTestDB(t *testing.T) *gorm.DB { + t.Helper() + + originalDB := model.DB + originalLogDB := model.LOG_DB + originalMemoryCacheEnabled := common.MemoryCacheEnabled + originalRedisEnabled := common.RedisEnabled + originalUsingSQLite := common.UsingSQLite + originalUsingMySQL := common.UsingMySQL + originalUsingPostgreSQL := common.UsingPostgreSQL + + gin.SetMode(gin.TestMode) + common.MemoryCacheEnabled = true + common.RedisEnabled = false + common.UsingSQLite = true + common.UsingMySQL = false + common.UsingPostgreSQL = false + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + + model.DB = db + model.LOG_DB = db + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{})) + service.ClearChannelAffinityCacheAll() + + t.Cleanup(func() { + service.ClearChannelAffinityCacheAll() + _ = db.Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(&model.Ability{}).Error + _ = db.Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(&model.Channel{}).Error + model.InitChannelCache() + + model.DB = originalDB + model.LOG_DB = originalLogDB + common.MemoryCacheEnabled = originalMemoryCacheEnabled + common.RedisEnabled = originalRedisEnabled + common.UsingSQLite = originalUsingSQLite + common.UsingMySQL = originalUsingMySQL + common.UsingPostgreSQL = originalUsingPostgreSQL + if originalMemoryCacheEnabled && originalDB != nil { + model.InitChannelCache() + } + if sqlDB, err := db.DB(); err == nil { + _ = sqlDB.Close() + } + }) + + return db +} + +func seedDistributorAffinityChannel(t *testing.T, db *gorm.DB, name string, status int, priority int64) *model.Channel { + t.Helper() + return seedDistributorAffinityChannelForModel(t, db, name, status, priority, "gpt-5") +} + +func seedDistributorAffinityChannelForModel(t *testing.T, db *gorm.DB, name string, status int, priority int64, modelName string) *model.Channel { + t.Helper() + + weight := uint(100) + autoBan := 1 + baseURL := "https://example.com" + channel := &model.Channel{ + Type: constant.ChannelTypeOpenAI, + Key: "sk-" + name, + Status: status, + Name: name, + Weight: &weight, + BaseURL: &baseURL, + Models: modelName, + Group: "default", + Priority: &priority, + AutoBan: &autoBan, + } + require.NoError(t, db.Create(channel).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: "default", + Model: modelName, + ChannelId: channel.Id, + Enabled: status == common.ChannelStatusEnabled, + Priority: &priority, + Weight: weight, + }).Error) + return channel +} + +func buildAffinityRequestContext(t *testing.T, body string) *gin.Context { + t.Helper() + + rec := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(rec) + ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + common.SetContextKey(ctx, constant.ContextKeyUsingGroup, "default") + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + return ctx +} + +func serveAffinityResponsesRequest(t *testing.T, body string) (int, int) { + t.Helper() + + var selectedChannelID int + router := gin.New() + router.Use(func(c *gin.Context) { + common.SetContextKey(c, constant.ContextKeyUsingGroup, "default") + common.SetContextKey(c, constant.ContextKeyUserGroup, "default") + common.SetContextKey(c, constant.ContextKeyTokenModelLimitEnabled, false) + c.Next() + }) + router.POST("/v1/responses", Distribute(), func(c *gin.Context) { + selectedChannelID = common.GetContextKeyInt(c, constant.ContextKeyChannelId) + c.Status(http.StatusOK) + }) + + recorder := httptest.NewRecorder() + request := httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(body)) + request.Header.Set("Content-Type", "application/json") + router.ServeHTTP(recorder, request) + return recorder.Code, selectedChannelID +} + +func TestDistributeInvalidatesDisabledAffinityChannelAndReselects(t *testing.T) { + db := setupDistributorAffinityTestDB(t) + + disabled := seedDistributorAffinityChannel(t, db, "affinity-disabled", common.ChannelStatusManuallyDisabled, 100) + available := seedDistributorAffinityChannel(t, db, "affinity-available", common.ChannelStatusEnabled, 90) + model.InitChannelCache() + + body := `{"model":"gpt-5","prompt_cache_key":"affinity-session-disabled"}` + bindCtx := buildAffinityRequestContext(t, body) + _, found := service.GetPreferredChannelByAffinity(bindCtx, "gpt-5", "default") + require.False(t, found) + service.RecordChannelAffinity(bindCtx, disabled.Id) + + checkCtx := buildAffinityRequestContext(t, body) + cachedChannelID, found := service.GetPreferredChannelByAffinity(checkCtx, "gpt-5", "default") + require.True(t, found) + require.Equal(t, disabled.Id, cachedChannelID) + + statusCode, selectedChannelID := serveAffinityResponsesRequest(t, body) + require.Equal(t, http.StatusOK, statusCode) + require.Equal(t, available.Id, selectedChannelID) + + refreshedCtx := buildAffinityRequestContext(t, body) + cachedChannelID, found = service.GetPreferredChannelByAffinity(refreshedCtx, "gpt-5", "default") + require.True(t, found) + require.Equal(t, available.Id, cachedChannelID) +} + +func TestDistributeInvalidatesModelMismatchedAffinityChannelAndReselects(t *testing.T) { + db := setupDistributorAffinityTestDB(t) + + mismatched := seedDistributorAffinityChannelForModel(t, db, "affinity-gpt5", common.ChannelStatusEnabled, 100, "gpt-5") + available := seedDistributorAffinityChannelForModel(t, db, "affinity-gpt4", common.ChannelStatusEnabled, 90, "gpt-4") + model.InitChannelCache() + + cacheBody := `{"model":"gpt-5","prompt_cache_key":"affinity-session-model-mismatch"}` + bindCtx := buildAffinityRequestContext(t, cacheBody) + _, found := service.GetPreferredChannelByAffinity(bindCtx, "gpt-5", "default") + require.False(t, found) + service.RecordChannelAffinity(bindCtx, mismatched.Id) + + requestBody := `{"model":"gpt-4","prompt_cache_key":"affinity-session-model-mismatch"}` + checkCtx := buildAffinityRequestContext(t, requestBody) + cachedChannelID, found := service.GetPreferredChannelByAffinity(checkCtx, "gpt-4", "default") + require.True(t, found) + require.Equal(t, mismatched.Id, cachedChannelID) + + statusCode, selectedChannelID := serveAffinityResponsesRequest(t, requestBody) + require.Equal(t, http.StatusOK, statusCode) + require.Equal(t, available.Id, selectedChannelID) + + refreshedCtx := buildAffinityRequestContext(t, requestBody) + cachedChannelID, found = service.GetPreferredChannelByAffinity(refreshedCtx, "gpt-4", "default") + require.True(t, found) + require.Equal(t, available.Id, cachedChannelID) +} diff --git a/model/channel.go b/model/channel.go index 78a1477c327e..b44a222ec44e 100644 --- a/model/channel.go +++ b/model/channel.go @@ -6,6 +6,7 @@ import ( "errors" "fmt" "math/rand" + "sort" "strings" "sync" @@ -20,6 +21,12 @@ import ( "gorm.io/gorm/clause" ) +type ChannelListStats struct { + UsedQuotaBalanceZero float64 `json:"used_quota_balance_zero" gorm:"column:used_quota_balance_zero"` + UsedQuotaBalanceNonzero float64 `json:"used_quota_balance_nonzero" gorm:"column:used_quota_balance_nonzero"` + BalanceTotal float64 `json:"balance_total" gorm:"column:balance_total"` +} + type Channel struct { Id int `json:"id"` Type int `json:"type" gorm:"default:0"` @@ -161,6 +168,16 @@ func ApplyChannelGroupFilter(query *gorm.DB, group string) *gorm.DB { return query.Where(channelGroupFilterCondition(), channelGroupFilterPattern(group)) } +func ApplyChannelStatusFilter(query *gorm.DB, statusFilter int) *gorm.DB { + if statusFilter == common.ChannelStatusEnabled { + return query.Where("status = ?", common.ChannelStatusEnabled) + } + if statusFilter == 0 { + return query.Where("status != ?", common.ChannelStatusEnabled) + } + return query +} + // Value implements driver.Valuer interface func (c ChannelInfo) Value() (driver.Value, error) { return common.Marshal(&c) @@ -365,6 +382,32 @@ func GetAllChannels(startIdx int, num int, selectAll bool, idSort bool, sortOpti return channels, err } +func CalculateChannelListStats(channels []*Channel) ChannelListStats { + stats := ChannelListStats{} + for _, channel := range channels { + if channel == nil { + continue + } + stats.BalanceTotal += channel.Balance + if channel.Balance == 0 { + stats.UsedQuotaBalanceZero += float64(channel.UsedQuota) + } else { + stats.UsedQuotaBalanceNonzero += float64(channel.UsedQuota) + } + } + return stats +} + +func GetChannelListStats(query *gorm.DB) (ChannelListStats, error) { + stats := ChannelListStats{} + err := query.Select(` + COALESCE(SUM(CASE WHEN balance = 0 THEN used_quota ELSE 0 END), 0) AS used_quota_balance_zero, + COALESCE(SUM(CASE WHEN balance <> 0 THEN used_quota ELSE 0 END), 0) AS used_quota_balance_nonzero, + COALESCE(SUM(balance), 0) AS balance_total + `).Scan(&stats).Error + return stats, err +} + func GetChannelsByTag(tag string, idSort bool, selectAll bool, sortOptions ...ChannelSortOptions) ([]*Channel, error) { var channels []*Channel order := resolveChannelSortOptions(idSort, sortOptions) @@ -376,39 +419,122 @@ func GetChannelsByTag(tag string, idSort bool, selectAll bool, sortOptions ...Ch return channels, err } -func SearchChannels(keyword string, group string, model string, idSort bool, sortOptions ...ChannelSortOptions) ([]*Channel, error) { - var channels []*Channel +func buildChannelSearchQuery(keyword string, group string, model string) *gorm.DB { modelsCol := "`models`" - - // 如果是 PostgreSQL，使用双引号 - if common.UsingPostgreSQL { - modelsCol = `"models"` - } - baseURLCol := "`base_url`" - // 如果是 PostgreSQL，使用双引号 if common.UsingPostgreSQL { + modelsCol = `"models"` baseURLCol = `"base_url"` } - order := resolveChannelSortOptions(idSort, sortOptions) - - // 构造基础查询 baseQuery := DB.Model(&Channel{}).Omit("key") - - // 构造WHERE子句 whereClause := "(id = ? OR name LIKE ? OR " + commonKeyCol + " = ? OR " + baseURLCol + " LIKE ?) AND " + modelsCol + " LIKE ?" args := []any{common.String2Int(keyword), "%" + keyword + "%", keyword, "%" + keyword + "%", "%" + model + "%"} - baseQuery = ApplyChannelGroupFilter(baseQuery.Where(whereClause, args...), group) + return ApplyChannelGroupFilter(baseQuery.Where(whereClause, args...), group) +} - // 执行查询 - err := order.Apply(baseQuery).Find(&channels).Error +func SearchChannels(keyword string, group string, model string, idSort bool, sortOptions ...ChannelSortOptions) ([]*Channel, error) { + var channels []*Channel + order := resolveChannelSortOptions(idSort, sortOptions) + err := order.Apply(buildChannelSearchQuery(keyword, group, model)).Find(&channels).Error if err != nil { return nil, err } return channels, nil } +const channelExactKeySearchChunkSize = 200 + +func SearchChannelsByExactKeys(keys []string, keyword string, group string, modelKeyword string, statusFilter int, idSort bool, sortOptions ...ChannelSortOptions) ([]*Channel, error) { + if len(keys) == 0 { + return []*Channel{}, nil + } + + order := resolveChannelSortOptions(idSort, sortOptions) + channelsByID := make(map[int]*Channel) + for start := 0; start < len(keys); start += channelExactKeySearchChunkSize { + end := start + channelExactKeySearchChunkSize + if end > len(keys) { + end = len(keys) + } + + var chunkChannels []*Channel + query := buildChannelSearchQuery(keyword, group, modelKeyword).Where(commonKeyCol+" IN ?", keys[start:end]) + query = ApplyChannelStatusFilter(query, statusFilter) + if err := query.Find(&chunkChannels).Error; err != nil { + return nil, err + } + for _, channel := range chunkChannels { + channelsByID[channel.Id] = channel + } + } + + channels := make([]*Channel, 0, len(channelsByID)) + for _, channel := range channelsByID { + channels = append(channels, channel) + } + sortChannels(channels, order) + return channels, nil +} + +func sortChannels(channels []*Channel, options ChannelSortOptions) { + less := func(a, b *Channel) bool { + sortBy := options.SortBy + if sortBy == "" { + sortBy = "priority" + } + desc := options.SortOrder != "asc" + switch sortBy { + case "id": + if a.Id != b.Id { + if desc { + return a.Id > b.Id + } + return a.Id < b.Id + } + case "name": + if a.Name != b.Name { + if desc { + return a.Name > b.Name + } + return a.Name < b.Name + } + case "balance": + if a.Balance != b.Balance { + if desc { + return a.Balance > b.Balance + } + return a.Balance < b.Balance + } + case "response_time": + if a.ResponseTime != b.ResponseTime { + if desc { + return a.ResponseTime > b.ResponseTime + } + return a.ResponseTime < b.ResponseTime + } + case "test_time": + if a.TestTime != b.TestTime { + if desc { + return a.TestTime > b.TestTime + } + return a.TestTime < b.TestTime + } + case "priority": + if a.GetPriority() != b.GetPriority() { + if desc { + return a.GetPriority() > b.GetPriority() + } + return a.GetPriority() < b.GetPriority() + } + } + return a.Id > b.Id + } + sort.SliceStable(channels, func(i, j int) bool { + return less(channels[i], channels[j]) + }) +} + func GetChannelById(id int, selectAll bool) (*Channel, error) { channel := &Channel{Id: id} var err error = nil @@ -423,6 +549,23 @@ func GetChannelById(id int, selectAll bool) (*Channel, error) { return channel, nil } +func CreateChannelsWithTx(tx *gorm.DB, channels []Channel) error { + if len(channels) == 0 { + return nil + } + for _, chunk := range lo.Chunk(channels, 50) { + if err := tx.Create(&chunk).Error; err != nil { + return err + } + for _, channel_ := range chunk { + if err := channel_.AddAbilities(tx); err != nil { + return err + } + } + } + return nil +} + func BatchInsertChannels(channels []Channel) error { if len(channels) == 0 { return nil @@ -437,17 +580,9 @@ func BatchInsertChannels(channels []Channel) error { } }() - for _, chunk := range lo.Chunk(channels, 50) { - if err := tx.Create(&chunk).Error; err != nil { - tx.Rollback() - return err - } - for _, channel_ := range chunk { - if err := channel_.AddAbilities(tx); err != nil { - tx.Rollback() - return err - } - } + if err := CreateChannelsWithTx(tx, channels); err != nil { + tx.Rollback() + return err } return tx.Commit().Error } @@ -867,6 +1002,10 @@ func updateChannelUsedQuota(id int, quota int) { } } +func ResetChannelUsedQuota(id int) error { + return DB.Model(&Channel{}).Where("id = ?", id).Update("used_quota", 0).Error +} + func DeleteChannelByStatus(status int64) (int64, error) { result := DB.Where("status = ?", status).Delete(&Channel{}) return result.RowsAffected, result.Error @@ -1071,6 +1210,15 @@ func CountChannelTags(query *gorm.DB) (int64, error) { return total, err } +func GetDistinctChannelGroups() ([]string, error) { + var groups []string + err := DB.Model(&Channel{}). + Where(commonGroupCol+" IS NOT NULL AND "+commonGroupCol+" != ''"). + Distinct(commonGroupCol). + Pluck(commonGroupCol, &groups).Error + return groups, err +} + // Get channels of specified type with pagination func GetChannelsByType(startIdx int, num int, idSort bool, channelType int) ([]*Channel, error) { var channels []*Channel diff --git a/model/channel_preparation.go b/model/channel_preparation.go new file mode 100644 index 000000000000..48d9c1ea3777 --- /dev/null +++ b/model/channel_preparation.go @@ -0,0 +1,402 @@ +package model + +import ( + "fmt" + "strings" + + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" +) + +const ( + ChannelPreparationStatusPending = 1 + ChannelPreparationStatusPromoted = 2 + ChannelPreparationStatusArchived = 3 + ChannelPreparationStatusPromoting = 4 +) + +const ( + ChannelPreparationTestStatusUntested = 0 + ChannelPreparationTestStatusSuccess = 1 + ChannelPreparationTestStatusFailed = 2 +) + +type ChannelPreparation struct { + Id int `json:"id"` + Type int `json:"type" gorm:"default:0"` + Key string `json:"key" gorm:"not null"` + OpenAIOrganization *string `json:"openai_organization"` + TestModel *string `json:"test_model"` + Name string `json:"name" gorm:"index"` + Weight *uint `json:"weight" gorm:"default:0"` + CreatedTime int64 `json:"created_time" gorm:"bigint"` + UpdatedTime int64 `json:"updated_time" gorm:"bigint"` + TestTime int64 `json:"test_time" gorm:"bigint;default:0"` + ResponseTime int `json:"response_time" gorm:"default:0"` + TestStatus int `json:"test_status" gorm:"default:0"` + TestMessage string `json:"test_message" gorm:"type:text"` + BaseURL *string `json:"base_url" gorm:"column:base_url;default:''"` + Other string `json:"other"` + Balance float64 `json:"balance"` + Models string `json:"models"` + Group string `json:"group" gorm:"type:varchar(64);default:'default'"` + ModelMapping *string `json:"model_mapping" gorm:"type:text"` + StatusCodeMapping *string `json:"status_code_mapping" gorm:"type:varchar(1024);default:''"` + Priority *int64 `json:"priority" gorm:"bigint;default:0"` + AutoBan *int `json:"auto_ban" gorm:"default:1"` + OtherInfo string `json:"other_info"` + Tag *string `json:"tag" gorm:"index"` + Setting *string `json:"setting" gorm:"type:text"` + ParamOverride *string `json:"param_override" gorm:"type:text"` + HeaderOverride *string `json:"header_override" gorm:"type:text"` + Remark *string `json:"remark" gorm:"type:varchar(255)" validate:"max=255"` + OtherSettings string `json:"settings" gorm:"column:settings"` + + Status int `json:"status" gorm:"default:1;index"` + Source string `json:"source" gorm:"type:varchar(64);index"` + Note string `json:"note" gorm:"type:text"` + PromotedTime *int64 `json:"promoted_time" gorm:"bigint"` + PromotedChannelId *int `json:"promoted_channel_id" gorm:"index"` +} + +type ChannelPreparationResponse struct { + Id int `json:"id"` + Type int `json:"type"` + KeyPreview string `json:"key_preview"` + OpenAIOrganization *string `json:"openai_organization"` + TestModel *string `json:"test_model"` + Name string `json:"name"` + Weight *uint `json:"weight"` + CreatedTime int64 `json:"created_time"` + UpdatedTime int64 `json:"updated_time"` + TestTime int64 `json:"test_time"` + ResponseTime int `json:"response_time"` + TestStatus int `json:"test_status"` + TestMessage string `json:"test_message"` + BaseURL *string `json:"base_url"` + Other string `json:"other"` + Balance float64 `json:"balance"` + Models string `json:"models"` + Group string `json:"group"` + ModelMapping *string `json:"model_mapping"` + StatusCodeMapping *string `json:"status_code_mapping"` + Priority *int64 `json:"priority"` + AutoBan *int `json:"auto_ban"` + OtherInfo string `json:"other_info"` + Tag *string `json:"tag"` + Setting *string `json:"setting"` + ParamOverride *string `json:"param_override"` + HeaderOverride *string `json:"header_override"` + Remark *string `json:"remark"` + OtherSettings string `json:"settings"` + Status int `json:"status"` + Source string `json:"source"` + Note string `json:"note"` + PromotedTime *int64 `json:"promoted_time"` + PromotedChannelId *int `json:"promoted_channel_id"` +} + +type ChannelPreparationListOptions struct { + Page int + PageSize int + Keyword string + Group string + Type *int + Status *int + StartTimestamp *int64 + EndTimestamp *int64 + IDSort bool +} + +type ChannelPreparationCountRow struct { + Value int `json:"value"` + Count int64 `json:"count"` +} + +type ChannelPreparationListStats struct { + BalanceTotal float64 `json:"balance_total" gorm:"column:balance_total"` +} + +func (p *ChannelPreparation) NormalizeForCreate() { + now := common.GetTimestamp() + p.Id = 0 + p.Status = ChannelPreparationStatusPending + p.CreatedTime = now + p.UpdatedTime = now + p.TestTime = 0 + p.ResponseTime = 0 + p.TestStatus = ChannelPreparationTestStatusUntested + p.TestMessage = "" + p.PromotedTime = nil + p.PromotedChannelId = nil + if strings.TrimSpace(p.Group) == "" { + p.Group = "default" + } + if p.AutoBan == nil { + defaultAutoBan := 1 + p.AutoBan = &defaultAutoBan + } +} + +func (p *ChannelPreparation) NormalizeForUpdate(existing *ChannelPreparation) { + p.Id = existing.Id + p.Status = existing.Status + p.CreatedTime = existing.CreatedTime + p.UpdatedTime = common.GetTimestamp() + p.TestTime = existing.TestTime + p.ResponseTime = existing.ResponseTime + p.TestStatus = existing.TestStatus + p.TestMessage = existing.TestMessage + p.PromotedTime = existing.PromotedTime + p.PromotedChannelId = existing.PromotedChannelId + if strings.TrimSpace(p.Key) == "" { + p.Key = existing.Key + } + if strings.TrimSpace(p.Group) == "" { + p.Group = "default" + } + if p.AutoBan == nil { + defaultAutoBan := 1 + p.AutoBan = &defaultAutoBan + } +} + +func (p *ChannelPreparation) KeyPreview() string { + key := strings.TrimSpace(p.Key) + if key == "" { + return "" + } + if len(key) <= 12 { + return key + } + return key[:8] + "..." + key[len(key)-4:] +} + +func (p *ChannelPreparation) ToResponse() ChannelPreparationResponse { + return ChannelPreparationResponse{ + Id: p.Id, + Type: p.Type, + KeyPreview: p.KeyPreview(), + OpenAIOrganization: p.OpenAIOrganization, + TestModel: p.TestModel, + Name: p.Name, + Weight: p.Weight, + CreatedTime: p.CreatedTime, + UpdatedTime: p.UpdatedTime, + TestTime: p.TestTime, + ResponseTime: p.ResponseTime, + TestStatus: p.TestStatus, + TestMessage: p.TestMessage, + BaseURL: p.BaseURL, + Other: p.Other, + Balance: p.Balance, + Models: p.Models, + Group: p.Group, + ModelMapping: p.ModelMapping, + StatusCodeMapping: p.StatusCodeMapping, + Priority: p.Priority, + AutoBan: p.AutoBan, + OtherInfo: p.OtherInfo, + Tag: p.Tag, + Setting: p.Setting, + ParamOverride: p.ParamOverride, + HeaderOverride: p.HeaderOverride, + Remark: p.Remark, + OtherSettings: p.OtherSettings, + Status: p.Status, + Source: p.Source, + Note: p.Note, + PromotedTime: p.PromotedTime, + PromotedChannelId: p.PromotedChannelId, + } +} + +func ChannelPreparationResponses(preparations []ChannelPreparation) []ChannelPreparationResponse { + responses := make([]ChannelPreparationResponse, 0, len(preparations)) + for _, preparation := range preparations { + responses = append(responses, preparation.ToResponse()) + } + return responses +} + +func FindActiveChannelPreparationKeyConflicts(keys []string, excludeID int) (map[string]ChannelPreparation, error) { + normalizedKeys := make([]string, 0, len(keys)) + seen := make(map[string]bool, len(keys)) + for _, key := range keys { + normalized := strings.TrimSpace(key) + if normalized == "" || seen[normalized] { + continue + } + seen[normalized] = true + normalizedKeys = append(normalizedKeys, normalized) + } + if len(normalizedKeys) == 0 { + return map[string]ChannelPreparation{}, nil + } + + activeStatuses := []int{ChannelPreparationStatusPending, ChannelPreparationStatusPromoting} + query := DB.Model(&ChannelPreparation{}). + Select("id, "+commonKeyCol+", name, status"). + Where("status IN ?", activeStatuses). + Where("TRIM("+commonKeyCol+") IN ?", normalizedKeys) + if excludeID > 0 { + query = query.Where("id <> ?", excludeID) + } + + var conflicts []ChannelPreparation + if err := query.Order("id asc").Find(&conflicts).Error; err != nil { + return nil, err + } + + result := make(map[string]ChannelPreparation, len(conflicts)) + for _, conflict := range conflicts { + normalized := strings.TrimSpace(conflict.Key) + if normalized == "" { + continue + } + if _, exists := result[normalized]; !exists { + result[normalized] = conflict + } + } + return result, nil +} + +func (p *ChannelPreparation) ToChannel() *Channel { + group := p.Group + if strings.TrimSpace(group) == "" { + group = "default" + } + autoBan := p.AutoBan + if autoBan == nil { + defaultAutoBan := 1 + autoBan = &defaultAutoBan + } + return &Channel{ + Type: p.Type, + Key: p.Key, + OpenAIOrganization: p.OpenAIOrganization, + TestModel: p.TestModel, + Status: common.ChannelStatusEnabled, + Name: p.Name, + TestTime: p.TestTime, + ResponseTime: p.ResponseTime, + Weight: p.Weight, + BaseURL: p.BaseURL, + Other: p.Other, + Balance: p.Balance, + Models: p.Models, + Group: group, + ModelMapping: p.ModelMapping, + StatusCodeMapping: p.StatusCodeMapping, + Priority: p.Priority, + AutoBan: autoBan, + OtherInfo: p.OtherInfo, + Tag: p.Tag, + Setting: p.Setting, + ParamOverride: p.ParamOverride, + HeaderOverride: p.HeaderOverride, + Remark: p.Remark, + OtherSettings: p.OtherSettings, + } +} + +func (p *ChannelPreparation) UpdateResponseTime(responseTime int64) { + p.UpdateTestResult(responseTime, ChannelPreparationTestStatusSuccess, "") +} + +func (p *ChannelPreparation) UpdateTestResult(responseTime int64, testStatus int, testMessage string) { + if len(testMessage) > 2048 { + testMessage = testMessage[:2048] + } + err := DB.Model(p).Select("response_time", "test_time", "test_status", "test_message").Updates(ChannelPreparation{ + TestTime: common.GetTimestamp(), + ResponseTime: int(responseTime), + TestStatus: testStatus, + TestMessage: testMessage, + }).Error + if err != nil { + common.SysLog(fmt.Sprintf("failed to update preparation test result: preparation_id=%d, error=%v", p.Id, err)) + } +} + +func applyChannelPreparationFilters(db *gorm.DB, opts ChannelPreparationListOptions, includeStatus bool, includeType bool) *gorm.DB { + keyword := strings.TrimSpace(opts.Keyword) + if keyword != "" { + like := "%" + keyword + "%" + db = db.Where("(id = ? OR name LIKE ? OR "+commonKeyCol+" = ? OR source LIKE ? OR note LIKE ?)", common.String2Int(keyword), like, keyword, like, like) + } + group := strings.TrimSpace(opts.Group) + if group != "" { + db = ApplyChannelGroupFilter(db, group) + } + if includeType && opts.Type != nil { + db = db.Where("type = ?", *opts.Type) + } + if includeStatus && opts.Status != nil { + db = db.Where("status = ?", *opts.Status) + } + if opts.StartTimestamp != nil { + db = db.Where("created_time >= ?", *opts.StartTimestamp) + } + if opts.EndTimestamp != nil { + db = db.Where("created_time <= ?", *opts.EndTimestamp) + } + return db +} + +func GetDistinctChannelPreparationGroups() ([]string, error) { + var groups []string + err := DB.Model(&ChannelPreparation{}). + Where(commonGroupCol+" IS NOT NULL AND "+commonGroupCol+" != ''"). + Distinct(commonGroupCol). + Pluck(commonGroupCol, &groups).Error + return groups, err +} + +func GetChannelPreparations(opts ChannelPreparationListOptions) ([]ChannelPreparation, int64, ChannelPreparationListStats, []ChannelPreparationCountRow, []ChannelPreparationCountRow, error) { + if opts.Page <= 0 { + opts.Page = 1 + } + if opts.PageSize <= 0 { + opts.PageSize = 20 + } + if opts.PageSize > 100 { + opts.PageSize = 100 + } + + base := applyChannelPreparationFilters(DB.Model(&ChannelPreparation{}), opts, true, true) + var total int64 + if err := base.Count(&total).Error; err != nil { + return nil, 0, ChannelPreparationListStats{}, nil, nil, err + } + + var stats ChannelPreparationListStats + statsQuery := applyChannelPreparationFilters(DB.Model(&ChannelPreparation{}), opts, true, true) + if err := statsQuery.Select("COALESCE(SUM(balance), 0) as balance_total").Scan(&stats).Error; err != nil { + return nil, 0, ChannelPreparationListStats{}, nil, nil, err + } + + var preparations []ChannelPreparation + order := "created_time desc, id desc" + if opts.IDSort { + order = "id desc" + } + err := base.Order(order).Limit(opts.PageSize).Offset((opts.Page - 1) * opts.PageSize).Find(&preparations).Error + if err != nil { + return nil, 0, ChannelPreparationListStats{}, nil, nil, err + } + + var statusCounts []ChannelPreparationCountRow + statusQuery := applyChannelPreparationFilters(DB.Model(&ChannelPreparation{}), opts, false, true) + if err := statusQuery.Select("status as value, count(*) as count").Group("status").Scan(&statusCounts).Error; err != nil { + return nil, 0, ChannelPreparationListStats{}, nil, nil, err + } + + var typeCounts []ChannelPreparationCountRow + typeQuery := applyChannelPreparationFilters(DB.Model(&ChannelPreparation{}), opts, true, false) + if err := typeQuery.Select("type as value, count(*) as count").Group("type").Scan(&typeCounts).Error; err != nil { + return nil, 0, ChannelPreparationListStats{}, nil, nil, err + } + + return preparations, total, stats, statusCounts, typeCounts, nil +} diff --git a/model/channel_preparation_test.go b/model/channel_preparation_test.go new file mode 100644 index 000000000000..a3b9fab285ad --- /dev/null +++ b/model/channel_preparation_test.go @@ -0,0 +1,224 @@ +package model + +import ( + "fmt" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupChannelPreparationModelTestDB(t *testing.T) *gorm.DB { + t.Helper() + + previousDB := DB + previousLogDB := LOG_DB + previousUsingSQLite := common.UsingSQLite + previousUsingMySQL := common.UsingMySQL + previousUsingPostgreSQL := common.UsingPostgreSQL + previousRedisEnabled := common.RedisEnabled + + common.UsingSQLite = true + common.UsingMySQL = false + common.UsingPostgreSQL = false + common.RedisEnabled = false + initCol() + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + DB = db + LOG_DB = db + require.NoError(t, db.AutoMigrate(&ChannelPreparation{})) + + t.Cleanup(func() { + DB = previousDB + LOG_DB = previousLogDB + common.UsingSQLite = previousUsingSQLite + common.UsingMySQL = previousUsingMySQL + common.UsingPostgreSQL = previousUsingPostgreSQL + common.RedisEnabled = previousRedisEnabled + initCol() + + sqlDB, err := db.DB() + if err == nil { + _ = sqlDB.Close() + } + }) + return db +} + +func TestChannelPreparationUpdateResponseTimeOnlyTouchesTestFields(t *testing.T) { + setupChannelPreparationModelTestDB(t) + + preparation := ChannelPreparation{ + Type: 1, + Key: "sk-test", + Name: "candidate", + Status: ChannelPreparationStatusPending, + Group: "default", + UpdatedTime: 12345, + TestTime: 111, + ResponseTime: 222, + } + require.NoError(t, DB.Create(&preparation).Error) + + before := common.GetTimestamp() + preparation.UpdateResponseTime(3456) + + var got ChannelPreparation + require.NoError(t, DB.First(&got, "id = ?", preparation.Id).Error) + require.Equal(t, 3456, got.ResponseTime) + require.Equal(t, ChannelPreparationTestStatusSuccess, got.TestStatus) + require.Empty(t, got.TestMessage) + require.GreaterOrEqual(t, got.TestTime, before) + require.LessOrEqual(t, got.TestTime, common.GetTimestamp()) + require.Equal(t, ChannelPreparationStatusPending, got.Status) + require.Equal(t, "sk-test", got.Key) + require.Equal(t, int64(12345), got.UpdatedTime) +} + +func TestChannelPreparationUpdateTestResultStoresFailure(t *testing.T) { + setupChannelPreparationModelTestDB(t) + + preparation := ChannelPreparation{ + Type: 1, + Key: "sk-test", + Name: "candidate", + Status: ChannelPreparationStatusPending, + Group: "default", + } + require.NoError(t, DB.Create(&preparation).Error) + + before := common.GetTimestamp() + preparation.UpdateTestResult(789, ChannelPreparationTestStatusFailed, "upstream timeout") + + var got ChannelPreparation + require.NoError(t, DB.First(&got, "id = ?", preparation.Id).Error) + require.Equal(t, 789, got.ResponseTime) + require.Equal(t, ChannelPreparationTestStatusFailed, got.TestStatus) + require.Equal(t, "upstream timeout", got.TestMessage) + require.GreaterOrEqual(t, got.TestTime, before) + require.LessOrEqual(t, got.TestTime, common.GetTimestamp()) +} + +func TestChannelPreparationNormalizePreservesAndResetsTestFields(t *testing.T) { + existing := &ChannelPreparation{ + Id: 7, + Status: ChannelPreparationStatusPending, + CreatedTime: 100, + UpdatedTime: 200, + Key: "existing-key", + Group: "vip", + TestTime: 300, + ResponseTime: 456, + TestStatus: ChannelPreparationTestStatusFailed, + TestMessage: "previous failure", + } + input := ChannelPreparation{Key: "", Group: ""} + input.NormalizeForUpdate(existing) + require.Equal(t, existing.Id, input.Id) + require.Equal(t, existing.Key, input.Key) + require.Equal(t, int64(300), input.TestTime) + require.Equal(t, 456, input.ResponseTime) + require.Equal(t, ChannelPreparationTestStatusFailed, input.TestStatus) + require.Equal(t, "previous failure", input.TestMessage) + + createInput := ChannelPreparation{ + TestTime: 300, + ResponseTime: 456, + TestStatus: ChannelPreparationTestStatusFailed, + TestMessage: "previous failure", + } + createInput.NormalizeForCreate() + require.Zero(t, createInput.TestTime) + require.Zero(t, createInput.ResponseTime) + require.Equal(t, ChannelPreparationTestStatusUntested, createInput.TestStatus) + require.Empty(t, createInput.TestMessage) +} + +func TestChannelPreparationResponseAndToChannelIncludeTestFields(t *testing.T) { + preparation := ChannelPreparation{ + Type: 1, + Key: "sk-test", + Name: "candidate", + Group: "default", + TestTime: 300, + ResponseTime: 456, + TestStatus: ChannelPreparationTestStatusFailed, + TestMessage: "failed", + } + + response := preparation.ToResponse() + require.Equal(t, int64(300), response.TestTime) + require.Equal(t, 456, response.ResponseTime) + require.Equal(t, ChannelPreparationTestStatusFailed, response.TestStatus) + require.Equal(t, "failed", response.TestMessage) + + channel := preparation.ToChannel() + require.Equal(t, int64(300), channel.TestTime) + require.Equal(t, 456, channel.ResponseTime) +} + +func TestFindActiveChannelPreparationKeyConflictsOnlyBlocksUnconsumedRecords(t *testing.T) { + setupChannelPreparationModelTestDB(t) + + preparations := []ChannelPreparation{ + {Id: 1, Type: 1, Key: "sk-pending", Name: "pending candidate", Status: ChannelPreparationStatusPending, Group: "default"}, + {Id: 2, Type: 1, Key: " sk-promoting ", Name: "promoting candidate", Status: ChannelPreparationStatusPromoting, Group: "default"}, + {Id: 3, Type: 1, Key: "sk-promoted", Name: "promoted candidate", Status: ChannelPreparationStatusPromoted, Group: "default"}, + {Id: 4, Type: 1, Key: "sk-archived", Name: "archived candidate", Status: ChannelPreparationStatusArchived, Group: "default"}, + } + require.NoError(t, DB.Create(&preparations).Error) + + conflicts, err := FindActiveChannelPreparationKeyConflicts([]string{"sk-pending", "sk-promoting", "sk-promoted", "sk-archived"}, 0) + require.NoError(t, err) + require.Contains(t, conflicts, "sk-pending") + require.Contains(t, conflicts, "sk-promoting") + require.NotContains(t, conflicts, "sk-promoted") + require.NotContains(t, conflicts, "sk-archived") + require.Equal(t, "pending candidate", conflicts["sk-pending"].Name) + require.Equal(t, "promoting candidate", conflicts["sk-promoting"].Name) + + conflicts, err = FindActiveChannelPreparationKeyConflicts([]string{"sk-pending"}, 1) + require.NoError(t, err) + require.Empty(t, conflicts) +} + +func TestGetChannelPreparationsFiltersGroupByExactToken(t *testing.T) { + setupChannelPreparationModelTestDB(t) + + preparations := []ChannelPreparation{ + {Id: 1, Type: 1, Key: "sk-vip", Name: "vip only", Status: ChannelPreparationStatusPending, Group: "vip", Balance: 10}, + {Id: 2, Type: 1, Key: "sk-svip", Name: "svip only", Status: ChannelPreparationStatusPending, Group: "svip", Balance: 20}, + {Id: 3, Type: 1, Key: "sk-default-vip", Name: "default vip", Status: ChannelPreparationStatusPending, Group: "default,vip", Balance: 30}, + {Id: 4, Type: 1, Key: "sk-vip2", Name: "vip2 only", Status: ChannelPreparationStatusPending, Group: "vip2", Balance: 40}, + } + require.NoError(t, DB.Create(&preparations).Error) + + items, total, stats, statusCounts, _, err := GetChannelPreparations(ChannelPreparationListOptions{Group: "vip", Page: 1, PageSize: 20}) + require.NoError(t, err) + require.Equal(t, int64(2), total) + require.InDelta(t, 40, stats.BalanceTotal, 0.000001) + require.Len(t, statusCounts, 1) + require.Equal(t, int64(2), statusCounts[0].Count) + + names := make(map[string]bool, len(items)) + for _, item := range items { + names[item.Name] = true + } + require.True(t, names["vip only"]) + require.True(t, names["default vip"]) + require.False(t, names["svip only"]) + require.False(t, names["vip2 only"]) + + items, total, stats, _, _, err = GetChannelPreparations(ChannelPreparationListOptions{Group: "svip", Page: 1, PageSize: 20}) + require.NoError(t, err) + require.Equal(t, int64(1), total) + require.InDelta(t, 20, stats.BalanceTotal, 0.000001) + require.Len(t, items, 1) + require.Equal(t, "svip only", items[0].Name) +} diff --git a/model/channel_query_key_report.go b/model/channel_query_key_report.go new file mode 100644 index 000000000000..47495f4409fc --- /dev/null +++ b/model/channel_query_key_report.go @@ -0,0 +1,446 @@ +package model + +import ( + "encoding/json" + "errors" + "math" + "strconv" + "strings" + + "github.com/QuantumNous/new-api/common" + "gorm.io/gorm" +) + +const MaxQueryKeyReportKeys = 10000 + +const ( + QueryKeyReportStatusNotFound = "not_found" + QueryKeyReportStatusFound = "found" + QueryKeyReportStatusOverBrushed = "over_brushed" +) + +const ( + QueryKeyReportSourceChannel = "channel" + QueryKeyReportSourcePreparation = "preparation" +) + +type QueryKeyReport struct { + TotalInput int `json:"total_input"` + UniqueKeys int `json:"unique_keys"` + DuplicateCount int `json:"duplicate_count"` + FoundCount int `json:"found_count"` + NotFoundCount int `json:"not_found_count"` + OverBrushedCount int `json:"over_brushed_count"` + TotalUsedQuota int64 `json:"total_used_quota"` + TotalUsedAmount float64 `json:"total_used_amount"` + TotalOriginalAmount float64 `json:"total_original_amount"` + TotalCurrentAmount float64 `json:"total_current_amount"` + TotalOverBrushAmount float64 `json:"total_over_brush_amount"` + Items []QueryKeyReportItem `json:"items"` +} + +type QueryKeyReportItem struct { + Key string `json:"key"` + Found bool `json:"found"` + Status string `json:"status"` + ChannelCount int `json:"channel_count"` + UsedQuota int64 `json:"used_quota"` + UsedAmount float64 `json:"used_amount"` + OriginalAmount float64 `json:"original_amount"` + CurrentAmount float64 `json:"current_amount"` + OverBrushAmount float64 `json:"over_brush_amount"` + OriginalAmountShared bool `json:"original_amount_shared"` + Channels []QueryKeyReportChannel `json:"channels"` +} + +type QueryKeyReportChannel struct { + Id int `json:"id"` + Source string `json:"source"` + Name string `json:"name"` + Type int `json:"type"` + Status int `json:"status"` + Group string `json:"group"` + Models string `json:"models"` + Tag *string `json:"tag"` + IsMultiKey bool `json:"is_multi_key"` + MatchedKeyCount int `json:"matched_key_count"` + UsedQuota int64 `json:"used_quota"` + UsedAmount float64 `json:"used_amount"` + MatchedUsedQuota int64 `json:"matched_used_quota"` + MatchedUsedAmount float64 `json:"matched_used_amount"` + OriginalAmount float64 `json:"original_amount"` + CurrentAmount float64 `json:"current_amount"` + OverBrushAmount float64 `json:"over_brush_amount"` + BalanceUpdatedTime int64 `json:"balance_updated_time"` +} + +type queryKeyReportInputRecord struct { + displayKey string + matchKey string +} + +type queryKeyReportItemAccumulator struct { + record queryKeyReportInputRecord + usedQuota int64 + usedAmount float64 + originalAmount float64 + overBrushAmount float64 + originalAmountShared bool + usesSharedMatchedUsage bool + sharedMatchedUsedAmount float64 + channels []QueryKeyReportChannel +} + +type queryKeyReportNonSharedTotal struct { + usedQuota int64 + usedAmount float64 + originalAmount float64 +} + +type queryKeyReportSharedTotal struct { + usedQuota int64 + usedAmount float64 + originalAmount float64 + overBrushAmount float64 +} + +type queryKeyReportChannelRecord struct { + source string + id int + key string + name string + channelType int + status int + group string + models string + tag *string + usedQuota int64 + balance float64 + balanceUpdatedTime int64 + isMultiKey bool +} + +func applyQueryKeyReportRecord(accumulators map[string]*queryKeyReportItemAccumulator, nonSharedTotals map[string]*queryKeyReportNonSharedTotal, sharedTotals map[string]queryKeyReportSharedTotal, record queryKeyReportChannelRecord) { + parsedKeys := parseQueryKeyReportChannelKeys(record.key) + if len(parsedKeys) == 0 { + return + } + + matchedKeys := make([]string, 0) + for _, parsedKey := range parsedKeys { + if _, ok := accumulators[parsedKey]; ok { + matchedKeys = append(matchedKeys, parsedKey) + } + } + if len(matchedKeys) == 0 { + return + } + + matchedKeyCount := len(matchedKeys) + usedAmount := quotaToAmount(record.usedQuota) + matchedUsedQuota := record.usedQuota * int64(matchedKeyCount) + matchedUsedAmount := usedAmount * float64(matchedKeyCount) + originalAmount := record.balance + isMultiKeyChannel := record.isMultiKey || len(parsedKeys) > 1 + sharedOriginal := matchedKeyCount > 1 + currentAmount := originalAmount - usedAmount + overBrushAmount := maxFloat(0, usedAmount-originalAmount) + if sharedOriginal { + currentAmount = originalAmount - matchedUsedAmount + overBrushAmount = maxFloat(0, matchedUsedAmount-originalAmount) + sharedTotals[record.source+":"+strconv.Itoa(record.id)] = queryKeyReportSharedTotal{ + usedQuota: matchedUsedQuota, + usedAmount: matchedUsedAmount, + originalAmount: originalAmount, + overBrushAmount: overBrushAmount, + } + } + + detail := QueryKeyReportChannel{ + Id: record.id, + Source: record.source, + Name: record.name, + Type: record.channelType, + Status: record.status, + Group: record.group, + Models: record.models, + Tag: record.tag, + IsMultiKey: isMultiKeyChannel, + MatchedKeyCount: matchedKeyCount, + UsedQuota: record.usedQuota, + UsedAmount: usedAmount, + MatchedUsedQuota: matchedUsedQuota, + MatchedUsedAmount: matchedUsedAmount, + OriginalAmount: originalAmount, + CurrentAmount: currentAmount, + OverBrushAmount: overBrushAmount, + BalanceUpdatedTime: record.balanceUpdatedTime, + } + + for _, matchedKey := range matchedKeys { + acc := accumulators[matchedKey] + acc.usedQuota += record.usedQuota + acc.usedAmount += usedAmount + acc.originalAmount = maxFloat(acc.originalAmount, originalAmount) + acc.channels = append(acc.channels, detail) + if isMultiKeyChannel { + acc.originalAmountShared = true + } + if sharedOriginal { + acc.usesSharedMatchedUsage = true + acc.sharedMatchedUsedAmount += matchedUsedAmount + acc.overBrushAmount += overBrushAmount + continue + } + + total := nonSharedTotals[matchedKey] + if total == nil { + total = &queryKeyReportNonSharedTotal{} + nonSharedTotals[matchedKey] = total + } + total.usedQuota += record.usedQuota + total.usedAmount += usedAmount + total.originalAmount = maxFloat(total.originalAmount, originalAmount) + } +} + +func BuildChannelQueryKeyReport(keys []string) (*QueryKeyReport, error) { + records, totalInput := normalizeQueryKeyReportInput(keys) + if len(records) == 0 { + return nil, errors.New("keys不能为空") + } + if len(records) > MaxQueryKeyReportKeys { + return nil, errors.New("最多支持10000个唯一密钥") + } + + accumulators := make(map[string]*queryKeyReportItemAccumulator, len(records)) + for _, record := range records { + accumulators[record.matchKey] = &queryKeyReportItemAccumulator{record: record} + } + + nonSharedTotals := make(map[string]*queryKeyReportNonSharedTotal) + sharedTotals := make(map[string]queryKeyReportSharedTotal) + + channelQuery := DB.Model(&Channel{}). + Select("id, " + commonKeyCol + ", name, type, status, " + commonGroupCol + ", models, tag, used_quota, balance, balance_updated_time, channel_info") + + channelResult := channelQuery.FindInBatches(&[]Channel{}, 500, func(tx *gorm.DB, batch int) error { + channels := tx.Statement.Dest.(*[]Channel) + for i := range *channels { + channel := &(*channels)[i] + applyQueryKeyReportRecord(accumulators, nonSharedTotals, sharedTotals, queryKeyReportChannelRecord{ + source: QueryKeyReportSourceChannel, + id: channel.Id, + key: channel.Key, + name: channel.Name, + channelType: channel.Type, + status: channel.Status, + group: channel.Group, + models: channel.Models, + tag: channel.Tag, + usedQuota: channel.UsedQuota, + balance: channel.Balance, + balanceUpdatedTime: channel.BalanceUpdatedTime, + isMultiKey: channel.ChannelInfo.IsMultiKey, + }) + } + return nil + }) + if channelResult.Error != nil { + return nil, channelResult.Error + } + + preparationQuery := DB.Model(&ChannelPreparation{}). + Select("id, " + commonKeyCol + ", name, type, status, " + commonGroupCol + ", models, tag, balance, updated_time") + + preparationResult := preparationQuery.FindInBatches(&[]ChannelPreparation{}, 500, func(tx *gorm.DB, batch int) error { + preparations := tx.Statement.Dest.(*[]ChannelPreparation) + for i := range *preparations { + preparation := &(*preparations)[i] + applyQueryKeyReportRecord(accumulators, nonSharedTotals, sharedTotals, queryKeyReportChannelRecord{ + source: QueryKeyReportSourcePreparation, + id: preparation.Id, + key: preparation.Key, + name: preparation.Name, + channelType: preparation.Type, + status: preparation.Status, + group: preparation.Group, + models: preparation.Models, + tag: preparation.Tag, + balance: preparation.Balance, + balanceUpdatedTime: preparation.UpdatedTime, + }) + } + return nil + }) + if preparationResult.Error != nil { + return nil, preparationResult.Error + } + + report := &QueryKeyReport{ + TotalInput: totalInput, + UniqueKeys: len(records), + DuplicateCount: totalInput - len(records), + Items: make([]QueryKeyReportItem, 0, len(records)), + } + + for _, total := range sharedTotals { + report.TotalUsedQuota += total.usedQuota + report.TotalUsedAmount += total.usedAmount + report.TotalOriginalAmount += total.originalAmount + report.TotalOverBrushAmount += total.overBrushAmount + } + for _, total := range nonSharedTotals { + report.TotalUsedQuota += total.usedQuota + report.TotalUsedAmount += total.usedAmount + report.TotalOriginalAmount += total.originalAmount + report.TotalOverBrushAmount += maxFloat(0, total.usedAmount-total.originalAmount) + } + report.TotalCurrentAmount = report.TotalOriginalAmount - report.TotalUsedAmount + + for _, record := range records { + acc := accumulators[record.matchKey] + item := QueryKeyReportItem{ + Key: record.displayKey, + Found: len(acc.channels) > 0, + Status: QueryKeyReportStatusNotFound, + ChannelCount: len(acc.channels), + UsedQuota: acc.usedQuota, + UsedAmount: acc.usedAmount, + OriginalAmount: acc.originalAmount, + CurrentAmount: acc.originalAmount - acc.usedAmount, + OverBrushAmount: acc.overBrushAmount, + OriginalAmountShared: acc.originalAmountShared, + Channels: acc.channels, + } + if item.Found { + report.FoundCount++ + if acc.usesSharedMatchedUsage { + item.CurrentAmount = item.OriginalAmount - acc.sharedMatchedUsedAmount + } else { + item.OverBrushAmount = maxFloat(0, item.UsedAmount-item.OriginalAmount) + } + if item.OverBrushAmount > 0 { + item.Status = QueryKeyReportStatusOverBrushed + report.OverBrushedCount++ + } else { + item.Status = QueryKeyReportStatusFound + } + } else { + report.NotFoundCount++ + } + report.Items = append(report.Items, item) + } + + return report, nil +} + +func normalizeQueryKeyReportInput(keys []string) ([]queryKeyReportInputRecord, int) { + records := make([]queryKeyReportInputRecord, 0, len(keys)) + seen := make(map[string]struct{}, len(keys)) + totalInput := 0 + for _, key := range keys { + displayKey := strings.TrimSpace(key) + if displayKey == "" { + continue + } + totalInput++ + matchKey := normalizeQueryKeyReportMatchKey(displayKey) + if matchKey == "" { + continue + } + if _, ok := seen[matchKey]; ok { + continue + } + seen[matchKey] = struct{}{} + records = append(records, queryKeyReportInputRecord{displayKey: displayKey, matchKey: matchKey}) + } + return records, totalInput +} + +func normalizeQueryKeyReportMatchKey(value string) string { + trimmed := strings.TrimSpace(value) + if trimmed == "" { + return "" + } + + var decoded any + if err := common.Unmarshal([]byte(trimmed), &decoded); err == nil { + switch typed := decoded.(type) { + case string: + return strings.TrimSpace(typed) + case nil: + return "" + default: + encoded, err := common.Marshal(typed) + if err == nil { + return string(encoded) + } + } + } + return trimmed +} + +func parseQueryKeyReportChannelKeys(key string) []string { + trimmed := strings.TrimSpace(key) + if trimmed == "" { + return []string{} + } + + keys := make([]string, 0) + if strings.HasPrefix(trimmed, "[") { + var values []json.RawMessage + if err := common.Unmarshal([]byte(trimmed), &values); err == nil { + for _, value := range values { + keys = append(keys, normalizeQueryKeyReportMatchKey(string(value))) + } + return uniqueNonBlankReportKeys(keys) + } + } + + for _, part := range strings.Split(strings.Trim(key, "\n"), "\n") { + keys = append(keys, normalizeQueryKeyReportMatchKey(part)) + } + return uniqueNonBlankReportKeys(keys) +} + +func uniqueNonBlankReportKeys(keys []string) []string { + seen := make(map[string]struct{}, len(keys)) + unique := make([]string, 0, len(keys)) + for _, key := range keys { + key = strings.TrimSpace(key) + if key == "" { + continue + } + if _, ok := seen[key]; ok { + continue + } + seen[key] = struct{}{} + unique = append(unique, key) + } + return unique +} + +func QueryKeyReportStoredKeyContains(storedKey string, inputKey string) bool { + matchKey := normalizeQueryKeyReportMatchKey(inputKey) + if matchKey == "" { + return false + } + for _, parsedKey := range parseQueryKeyReportChannelKeys(storedKey) { + if parsedKey == matchKey { + return true + } + } + return false +} + +func quotaToAmount(usedQuota int64) float64 { + if common.QuotaPerUnit == 0 { + return 0 + } + return float64(usedQuota) / common.QuotaPerUnit +} + +func maxFloat(a, b float64) float64 { + return math.Max(a, b) +} diff --git a/model/channel_query_key_report_test.go b/model/channel_query_key_report_test.go new file mode 100644 index 000000000000..1922ea7e670a --- /dev/null +++ b/model/channel_query_key_report_test.go @@ -0,0 +1,238 @@ +package model + +import ( + "fmt" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupQueryKeyReportModelTestDB(t *testing.T) *gorm.DB { + t.Helper() + + previousDB := DB + previousLogDB := LOG_DB + previousUsingSQLite := common.UsingSQLite + previousUsingMySQL := common.UsingMySQL + previousUsingPostgreSQL := common.UsingPostgreSQL + previousRedisEnabled := common.RedisEnabled + + common.UsingSQLite = true + common.UsingMySQL = false + common.UsingPostgreSQL = false + common.RedisEnabled = false + initCol() + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + DB = db + LOG_DB = db + require.NoError(t, db.AutoMigrate(&Channel{}, &ChannelPreparation{})) + + t.Cleanup(func() { + DB = previousDB + LOG_DB = previousLogDB + common.UsingSQLite = previousUsingSQLite + common.UsingMySQL = previousUsingMySQL + common.UsingPostgreSQL = previousUsingPostgreSQL + common.RedisEnabled = previousRedisEnabled + initCol() + + sqlDB, err := db.DB() + if err == nil { + _ = sqlDB.Close() + } + }) + return db +} + +func quotaUnits(units int64) int64 { + return int64(common.QuotaPerUnit) * units +} + +func reportItemByKey(t *testing.T, report *QueryKeyReport, key string) QueryKeyReportItem { + t.Helper() + for _, item := range report.Items { + if item.Key == key { + return item + } + } + t.Fatalf("missing report item for key %q", key) + return QueryKeyReportItem{} +} + +func TestBuildChannelQueryKeyReportAggregatesDuplicateRowsAndSharedMultiKeyBalance(t *testing.T) { + setupQueryKeyReportModelTestDB(t) + + channels := []Channel{ + {Id: 1, Type: 1, Key: "sk-repeat", Name: "repeat low balance", Status: common.ChannelStatusEnabled, Group: "default", Models: "gpt-4o", UsedQuota: quotaUnits(4), Balance: 5}, + {Id: 2, Type: 1, Key: "sk-repeat", Name: "repeat high balance", Status: common.ChannelStatusEnabled, Group: "default", Models: "gpt-4o", UsedQuota: quotaUnits(3), Balance: 6}, + {Id: 3, Type: 2, Key: "sk-shared-a\nsk-shared-b", Name: "shared multi", Status: common.ChannelStatusEnabled, Group: "default", Models: "gpt-4o-mini", UsedQuota: quotaUnits(6), Balance: 10}, + } + require.NoError(t, DB.Create(&channels).Error) + + report, err := BuildChannelQueryKeyReport([]string{" sk-repeat ", "sk-missing", "sk-shared-a", "sk-shared-b", "sk-repeat", ""}) + require.NoError(t, err) + + require.Equal(t, 5, report.TotalInput) + require.Equal(t, 4, report.UniqueKeys) + require.Equal(t, 1, report.DuplicateCount) + require.Equal(t, 3, report.FoundCount) + require.Equal(t, 1, report.NotFoundCount) + require.Equal(t, 3, report.OverBrushedCount) + require.Equal(t, quotaUnits(19), report.TotalUsedQuota) + require.InDelta(t, 19, report.TotalUsedAmount, 0.000001) + require.InDelta(t, 16, report.TotalOriginalAmount, 0.000001) + require.InDelta(t, -3, report.TotalCurrentAmount, 0.000001) + require.InDelta(t, 3, report.TotalOverBrushAmount, 0.000001) + + repeat := reportItemByKey(t, report, "sk-repeat") + require.True(t, repeat.Found) + require.Equal(t, QueryKeyReportStatusOverBrushed, repeat.Status) + require.Equal(t, 2, repeat.ChannelCount) + require.Equal(t, quotaUnits(7), repeat.UsedQuota) + require.InDelta(t, 7, repeat.UsedAmount, 0.000001) + require.InDelta(t, 6, repeat.OriginalAmount, 0.000001) + require.InDelta(t, -1, repeat.CurrentAmount, 0.000001) + require.InDelta(t, 1, repeat.OverBrushAmount, 0.000001) + require.False(t, repeat.OriginalAmountShared) + + missing := reportItemByKey(t, report, "sk-missing") + require.False(t, missing.Found) + require.Equal(t, QueryKeyReportStatusNotFound, missing.Status) + require.Equal(t, 0, missing.ChannelCount) + + sharedA := reportItemByKey(t, report, "sk-shared-a") + require.True(t, sharedA.Found) + require.Equal(t, QueryKeyReportStatusOverBrushed, sharedA.Status) + require.True(t, sharedA.OriginalAmountShared) + require.InDelta(t, 10, sharedA.OriginalAmount, 0.000001) + require.InDelta(t, -2, sharedA.CurrentAmount, 0.000001) + require.InDelta(t, 2, sharedA.OverBrushAmount, 0.000001) + require.Len(t, sharedA.Channels, 1) + require.Equal(t, 2, sharedA.Channels[0].MatchedKeyCount) + require.Equal(t, quotaUnits(12), sharedA.Channels[0].MatchedUsedQuota) + require.InDelta(t, 12, sharedA.Channels[0].MatchedUsedAmount, 0.000001) + require.InDelta(t, 10, sharedA.Channels[0].OriginalAmount, 0.000001) + require.InDelta(t, -2, sharedA.Channels[0].CurrentAmount, 0.000001) + require.InDelta(t, 2, sharedA.Channels[0].OverBrushAmount, 0.000001) + + sharedB := reportItemByKey(t, report, "sk-shared-b") + require.True(t, sharedB.OriginalAmountShared) + require.InDelta(t, 10, sharedB.OriginalAmount, 0.000001) +} + +func TestBuildChannelQueryKeyReportIncludesChannelPreparations(t *testing.T) { + setupQueryKeyReportModelTestDB(t) + + preparations := []ChannelPreparation{ + {Id: 20, Type: 2, Key: "sk-prep", Name: "prep single", Status: ChannelPreparationStatusPending, Group: "svip", Models: "claude-3", Balance: 20, UpdatedTime: 1717488000}, + {Id: 21, Type: 2, Key: "sk-prep-a\nsk-prep-b", Name: "prep multi", Status: ChannelPreparationStatusPending, Group: "default", Models: "claude-3", Balance: 9, UpdatedTime: 1717489000}, + } + require.NoError(t, DB.Create(&preparations).Error) + + report, err := BuildChannelQueryKeyReport([]string{"sk-prep", "sk-prep-a", "sk-prep-b", "sk-missing"}) + require.NoError(t, err) + + require.Equal(t, 3, report.FoundCount) + require.Equal(t, 1, report.NotFoundCount) + require.Equal(t, int64(0), report.TotalUsedQuota) + require.InDelta(t, 0, report.TotalUsedAmount, 0.000001) + require.InDelta(t, 29, report.TotalOriginalAmount, 0.000001) + require.InDelta(t, 29, report.TotalCurrentAmount, 0.000001) + require.InDelta(t, 0, report.TotalOverBrushAmount, 0.000001) + + prep := reportItemByKey(t, report, "sk-prep") + require.True(t, prep.Found) + require.Equal(t, QueryKeyReportStatusFound, prep.Status) + require.Equal(t, 1, prep.ChannelCount) + require.Equal(t, int64(0), prep.UsedQuota) + require.InDelta(t, 20, prep.OriginalAmount, 0.000001) + require.InDelta(t, 20, prep.CurrentAmount, 0.000001) + require.Equal(t, QueryKeyReportSourcePreparation, prep.Channels[0].Source) + require.Equal(t, ChannelPreparationStatusPending, prep.Channels[0].Status) + require.Equal(t, int64(1717488000), prep.Channels[0].BalanceUpdatedTime) + + multi := reportItemByKey(t, report, "sk-prep-a") + require.True(t, multi.Found) + require.True(t, multi.OriginalAmountShared) + require.Len(t, multi.Channels, 1) + require.Equal(t, QueryKeyReportSourcePreparation, multi.Channels[0].Source) + require.Equal(t, 2, multi.Channels[0].MatchedKeyCount) + require.InDelta(t, 9, multi.Channels[0].OriginalAmount, 0.000001) + require.InDelta(t, 9, multi.Channels[0].CurrentAmount, 0.000001) +} + +func TestBuildChannelQueryKeyReportMatchesMultiKeyFormatsAndSanitizesDetails(t *testing.T) { + setupQueryKeyReportModelTestDB(t) + + channels := []Channel{ + {Id: 10, Type: 1, Key: "sk-newline\nsk-dup\nsk-dup", Name: "newline multi", Status: common.ChannelStatusEnabled, Group: "default", Models: "gpt-4o", UsedQuota: quotaUnits(1), Balance: 3}, + {Id: 11, Type: 2, Key: `[ + "sk-json-string", + {"b":2,"a":"x"} + ]`, Name: "json multi", Status: common.ChannelStatusEnabled, Group: "default", Models: "claude-3", UsedQuota: quotaUnits(2), Balance: 10}, + } + require.NoError(t, DB.Create(&channels).Error) + + report, err := BuildChannelQueryKeyReport([]string{"sk-newline", "sk-dup", "sk-json-string", `{ "a": "x", "b": 2 }`}) + require.NoError(t, err) + + require.Equal(t, 4, report.FoundCount) + require.Equal(t, 0, report.NotFoundCount) + + newline := reportItemByKey(t, report, "sk-newline") + require.Len(t, newline.Channels, 1) + require.True(t, newline.Channels[0].IsMultiKey) + require.Equal(t, 2, newline.Channels[0].MatchedKeyCount) + require.True(t, newline.OriginalAmountShared) + + dup := reportItemByKey(t, report, "sk-dup") + require.Len(t, dup.Channels, 1) + require.Equal(t, 1, dup.ChannelCount) + + jsonString := reportItemByKey(t, report, "sk-json-string") + require.True(t, jsonString.OriginalAmountShared) + require.Len(t, jsonString.Channels, 1) + require.Equal(t, 2, jsonString.Channels[0].MatchedKeyCount) + + object := reportItemByKey(t, report, `{ "a": "x", "b": 2 }`) + require.True(t, object.Found) + require.True(t, object.OriginalAmountShared) + + detailBytes, err := common.Marshal(jsonString.Channels[0]) + require.NoError(t, err) + detailJSON := string(detailBytes) + require.NotContains(t, detailJSON, "sk-json-string") + require.NotContains(t, detailJSON, "sk-dup") + require.NotContains(t, detailJSON, "\"key\"") +} + +func TestQueryKeyReportStoredKeyContainsNormalizesInputFormats(t *testing.T) { + require.True(t, QueryKeyReportStoredKeyContains("sk-a\nsk-b", " sk-a ")) + require.True(t, QueryKeyReportStoredKeyContains(`["sk-json", {"b": 2, "a": "x"}]`, "sk-json")) + require.True(t, QueryKeyReportStoredKeyContains(`["sk-json", {"b": 2, "a": "x"}]`, `{ "a": "x", "b": 2 }`)) + require.False(t, QueryKeyReportStoredKeyContains("sk-a\nsk-b", "sk-c")) + require.False(t, QueryKeyReportStoredKeyContains("", "sk-a")) +} + +func TestBuildChannelQueryKeyReportRejectsEmptyAndTooManyUniqueKeys(t *testing.T) { + setupQueryKeyReportModelTestDB(t) + + _, err := BuildChannelQueryKeyReport([]string{"", " "}) + require.Error(t, err) + require.Contains(t, err.Error(), "keys") + + keys := make([]string, MaxQueryKeyReportKeys+1) + for i := range keys { + keys[i] = fmt.Sprintf("sk-%05d", i) + } + _, err = BuildChannelQueryKeyReport(keys) + require.Error(t, err) + require.Contains(t, err.Error(), "10000") +} diff --git a/model/cost_report.go b/model/cost_report.go new file mode 100644 index 000000000000..b23d8fa73af0 --- /dev/null +++ b/model/cost_report.go @@ -0,0 +1,86 @@ +package model + +const ( + CostReportTemplateStatusEnabled = 1 + CostReportTemplateStatusArchived = 2 + + CostReportTemplateVersionStatusActive = 1 + CostReportTemplateVersionStatusArchived = 2 + + CostReportRunStatusPending = 1 + CostReportRunStatusCompleted = 2 + CostReportRunStatusFailed = 3 +) + +// CostReportTemplate stores report template metadata and points at the current immutable version. +// Cost report tables live in the main DB only; do not add them to LOG_DB migrations. +type CostReportTemplate struct { + Id int `json:"id" gorm:"primaryKey"` + Key string `json:"key" gorm:"type:varchar(64);uniqueIndex;not null"` + Name string `json:"name" gorm:"type:varchar(128);not null"` + Description string `json:"description" gorm:"type:text"` + Status int `json:"status" gorm:"not null;default:1;index"` + CurrentVersionId *int `json:"current_version_id" gorm:"index"` + CreatedBy int `json:"created_by" gorm:"index"` + UpdatedBy int `json:"updated_by" gorm:"index"` + CreatedAt int64 `json:"created_at" gorm:"type:bigint;autoCreateTime"` + UpdatedAt int64 `json:"updated_at" gorm:"type:bigint;autoUpdateTime"` +} + +// CostReportTemplateVersion is an immutable snapshot of a template config. +type CostReportTemplateVersion struct { + Id int `json:"id" gorm:"primaryKey"` + TemplateId int `json:"template_id" gorm:"uniqueIndex:idx_cost_report_template_version;index;not null"` + Version int `json:"version" gorm:"uniqueIndex:idx_cost_report_template_version;not null"` + Status int `json:"status" gorm:"not null;default:1;index"` + ConfigJson string `json:"config_json" gorm:"type:text;not null"` + ConfigHash string `json:"config_hash" gorm:"type:varchar(64);index;not null"` + CreatedBy int `json:"created_by" gorm:"index"` + CreatedAt int64 `json:"created_at" gorm:"type:bigint;autoCreateTime"` +} + +// CostReportRun stores one saved report snapshot for one period. +type CostReportRun struct { + Id int `json:"id" gorm:"primaryKey"` + TemplateId int `json:"template_id" gorm:"index;not null"` + TemplateVersionId int `json:"template_version_id" gorm:"index;not null"` + PeriodStart int64 `json:"period_start" gorm:"type:bigint;index;not null"` + PeriodEnd int64 `json:"period_end" gorm:"type:bigint;index;not null"` + PeriodKey string `json:"period_key" gorm:"type:varchar(64);index;not null"` + Timezone string `json:"timezone" gorm:"type:varchar(64);not null"` + Status int `json:"status" gorm:"not null;default:1;index"` + ConfigSnapshotJson string `json:"config_snapshot_json" gorm:"type:text;not null"` + SourceLogMaxId int `json:"source_log_max_id" gorm:"index"` + SourceHash string `json:"source_hash" gorm:"type:varchar(64);index"` + RowCount int `json:"row_count" gorm:"not null;default:0"` + ErrorMessage string `json:"error_message" gorm:"type:text"` + CreatedBy int `json:"created_by" gorm:"index"` + CreatedAt int64 `json:"created_at" gorm:"type:bigint;autoCreateTime"` + UpdatedAt int64 `json:"updated_at" gorm:"type:bigint;autoUpdateTime"` +} + +// CostReportRowSnapshot freezes computed row data for a saved run. +type CostReportRowSnapshot struct { + Id int `json:"id" gorm:"primaryKey"` + RunId int `json:"run_id" gorm:"uniqueIndex:idx_cost_report_run_row;index;not null"` + RowKey string `json:"row_key" gorm:"type:varchar(64);uniqueIndex:idx_cost_report_run_row;not null"` + DimensionsJson string `json:"dimensions_json" gorm:"type:text;not null"` + MetricsJson string `json:"metrics_json" gorm:"type:text;not null"` + ManualValuesJson string `json:"manual_values_json" gorm:"type:text;not null"` + FormulaValuesJson string `json:"formula_values_json" gorm:"type:text;not null"` + CreatedAt int64 `json:"created_at" gorm:"type:bigint;autoCreateTime"` +} + +// CostReportManualCell stores reusable root-entered manual values by stable row identity. +type CostReportManualCell struct { + Id int `json:"id" gorm:"primaryKey"` + TemplateId int `json:"template_id" gorm:"uniqueIndex:idx_cost_report_manual_cell;index;not null"` + PeriodKey string `json:"period_key" gorm:"type:varchar(64);uniqueIndex:idx_cost_report_manual_cell;not null"` + RowKey string `json:"row_key" gorm:"type:varchar(64);uniqueIndex:idx_cost_report_manual_cell;not null"` + FieldKey string `json:"field_key" gorm:"type:varchar(64);uniqueIndex:idx_cost_report_manual_cell;not null"` + ValueType string `json:"value_type" gorm:"type:varchar(32);not null"` + ValueText string `json:"value_text" gorm:"type:text"` + UpdatedBy int `json:"updated_by" gorm:"index"` + CreatedAt int64 `json:"created_at" gorm:"type:bigint;autoCreateTime"` + UpdatedAt int64 `json:"updated_at" gorm:"type:bigint;autoUpdateTime"` +} diff --git a/model/log.go b/model/log.go index cbcc3983d924..17c8086cb072 100644 --- a/model/log.go +++ b/model/log.go @@ -68,20 +68,30 @@ const ( func formatUserLogs(logs []*Log, startIdx int) { for i := range logs { - logs[i].ChannelName = "" - var otherMap map[string]interface{} - otherMap, _ = common.StrToMap(logs[i].Other) - if otherMap != nil { - // Remove admin-only debug fields. - delete(otherMap, "admin_info") - // delete(otherMap, "reject_reason") - delete(otherMap, "stream_status") - } - logs[i].Other = common.MapToJsonStr(otherMap) + sanitizeUserLog(logs[i]) logs[i].Id = startIdx + i + 1 } } +func formatUserLogsForExport(logs []*Log) { + for i := range logs { + sanitizeUserLog(logs[i]) + } +} + +func sanitizeUserLog(log *Log) { + log.ChannelName = "" + var otherMap map[string]interface{} + otherMap, _ = common.StrToMap(log.Other) + if otherMap != nil { + // Remove admin-only debug fields. + delete(otherMap, "admin_info") + // delete(otherMap, "reject_reason") + delete(otherMap, "stream_status") + } + log.Other = common.MapToJsonStr(otherMap) +} + func GetLogByTokenId(tokenId int) (logs []*Log, err error) { err = LOG_DB.Model(&Log{}).Where("token_id = ?", tokenId).Order("id desc").Limit(common.MaxRecentItems).Find(&logs).Error formatUserLogs(logs, 0) @@ -108,16 +118,26 @@ func RecordLog(userId int, logType int, content string) { // RecordLogWithAdminInfo 记录操作日志，并将管理员相关信息存入 Other.admin_info， func RecordLogWithAdminInfo(userId int, logType int, content string, adminInfo map[string]interface{}) { + RecordLogWithAdminInfoAndMetadata(userId, logType, content, 0, "", adminInfo) +} + +// RecordLogWithAdminInfoAndMetadata 记录操作日志，并额外写入渠道、分组等可筛选字段。 +func RecordLogWithAdminInfoAndMetadata(userId int, logType int, content string, channelId int, group string, adminInfo map[string]interface{}) { if logType == LogTypeConsume && !common.LogConsumeEnabled { return } - username, _ := GetUsernameById(userId, false) + username := "system" + if userId > 0 { + username, _ = GetUsernameById(userId, false) + } log := &Log{ UserId: userId, Username: username, CreatedAt: common.GetTimestamp(), Type: logType, Content: content, + ChannelId: channelId, + Group: group, } if len(adminInfo) > 0 { other := map[string]interface{}{ @@ -315,7 +335,7 @@ func RecordTaskBillingLog(params RecordTaskBillingLogParams) { } } -func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) { +func buildAdminLogQuery(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, channel int, group string, requestId string, upstreamRequestId string) (*gorm.DB, error) { var tx *gorm.DB if logType == LogTypeUnknown { tx = LOG_DB @@ -323,11 +343,12 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName tx = LOG_DB.Where("logs.type = ?", logType) } + var err error if tx, err = applyExplicitLogTextFilter(tx, "logs.model_name", modelName); err != nil { - return nil, 0, err + return nil, err } if tx, err = applyExplicitLogTextFilter(tx, "logs.username", username); err != nil { - return nil, 0, err + return nil, err } if tokenName != "" { tx = tx.Where("logs.token_name = ?", tokenName) @@ -350,15 +371,10 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName if group != "" { tx = tx.Where("logs."+logGroupCol+" = ?", group) } - err = tx.Model(&Log{}).Count(&total).Error - if err != nil { - return nil, 0, err - } - err = tx.Order("logs.created_at desc, logs.id desc").Limit(num).Offset(startIdx).Find(&logs).Error - if err != nil { - return nil, 0, err - } + return tx, nil +} +func fillLogChannelNames(logs []*Log) error { channelIds := types.NewSet[int]() for _, log := range logs { if log.ChannelId != 0 { @@ -366,45 +382,66 @@ func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName } } - if channelIds.Len() > 0 { - var channels []struct { - Id int `gorm:"column:id"` - Name string `gorm:"column:name"` - } - if common.MemoryCacheEnabled { - // Cache get channel - for _, channelId := range channelIds.Items() { - if cacheChannel, err := CacheGetChannel(channelId); err == nil { - channels = append(channels, struct { - Id int `gorm:"column:id"` - Name string `gorm:"column:name"` - }{ - Id: channelId, - Name: cacheChannel.Name, - }) - } - } - } else { - // Bulk query channels from DB - if err = DB.Table("channels").Select("id, name").Where("id IN ?", channelIds.Items()).Find(&channels).Error; err != nil { - return logs, total, err + if channelIds.Len() == 0 { + return nil + } + + var channels []struct { + Id int `gorm:"column:id"` + Name string `gorm:"column:name"` + } + if common.MemoryCacheEnabled { + // Cache get channel + for _, channelId := range channelIds.Items() { + if cacheChannel, err := CacheGetChannel(channelId); err == nil { + channels = append(channels, struct { + Id int `gorm:"column:id"` + Name string `gorm:"column:name"` + }{ + Id: channelId, + Name: cacheChannel.Name, + }) } } - channelMap := make(map[int]string, len(channels)) - for _, channel := range channels { - channelMap[channel.Id] = channel.Name - } - for i := range logs { - logs[i].ChannelName = channelMap[logs[i].ChannelId] + } else { + // Bulk query channels from DB + if err := DB.Table("channels").Select("id, name").Where("id IN ?", channelIds.Items()).Find(&channels).Error; err != nil { + return err } } + channelMap := make(map[int]string, len(channels)) + for _, channel := range channels { + channelMap[channel.Id] = channel.Name + } + for i := range logs { + logs[i].ChannelName = channelMap[logs[i].ChannelId] + } + return nil +} + +func GetAllLogs(logType int, startTimestamp int64, endTimestamp int64, modelName string, username string, tokenName string, startIdx int, num int, channel int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) { + tx, err := buildAdminLogQuery(logType, startTimestamp, endTimestamp, modelName, username, tokenName, channel, group, requestId, upstreamRequestId) + if err != nil { + return nil, 0, err + } + err = tx.Model(&Log{}).Count(&total).Error + if err != nil { + return nil, 0, err + } + err = tx.Order("logs.created_at desc, logs.id desc").Limit(num).Offset(startIdx).Find(&logs).Error + if err != nil { + return nil, 0, err + } + if err = fillLogChannelNames(logs); err != nil { + return logs, total, err + } return logs, total, err } const logSearchCountLimit = 10000 -func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) { +func buildUserLogQuery(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, group string, requestId string, upstreamRequestId string) (*gorm.DB, error) { var tx *gorm.DB if logType == LogTypeUnknown { tx = LOG_DB.Where("logs.user_id = ?", userId) @@ -412,8 +449,9 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int tx = LOG_DB.Where("logs.user_id = ? and logs.type = ?", userId, logType) } + var err error if tx, err = applyExplicitLogTextFilter(tx, "logs.model_name", modelName); err != nil { - return nil, 0, err + return nil, err } if tokenName != "" { tx = tx.Where("logs.token_name = ?", tokenName) @@ -433,6 +471,14 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int if group != "" { tx = tx.Where("logs."+logGroupCol+" = ?", group) } + return tx, nil +} + +func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int64, modelName string, tokenName string, startIdx int, num int, group string, requestId string, upstreamRequestId string) (logs []*Log, total int64, err error) { + tx, err := buildUserLogQuery(userId, logType, startTimestamp, endTimestamp, modelName, tokenName, group, requestId, upstreamRequestId) + if err != nil { + return nil, 0, err + } err = tx.Model(&Log{}).Limit(logSearchCountLimit).Count(&total).Error if err != nil { common.SysError("failed to count user logs: " + err.Error()) @@ -448,6 +494,65 @@ func GetUserLogs(userId int, logType int, startTimestamp int64, endTimestamp int return logs, total, err } +type LogExportFilter struct { + UserId int + IsAdmin bool + LogType int + StartTimestamp int64 + EndTimestamp int64 + ModelName string + Username string + TokenName string + Channel int + Group string + RequestId string + UpstreamRequestId string +} + +func buildLogExportQuery(filter LogExportFilter) (*gorm.DB, error) { + if filter.IsAdmin { + return buildAdminLogQuery(filter.LogType, filter.StartTimestamp, filter.EndTimestamp, filter.ModelName, filter.Username, filter.TokenName, filter.Channel, filter.Group, filter.RequestId, filter.UpstreamRequestId) + } + return buildUserLogQuery(filter.UserId, filter.LogType, filter.StartTimestamp, filter.EndTimestamp, filter.ModelName, filter.TokenName, filter.Group, filter.RequestId, filter.UpstreamRequestId) +} + +func CountLogsForExport(ctx context.Context, filter LogExportFilter) (int64, error) { + tx, err := buildLogExportQuery(filter) + if err != nil { + return 0, err + } + var total int64 + if err = tx.WithContext(ctx).Model(&Log{}).Count(&total).Error; err != nil { + return 0, err + } + return total, nil +} + +func GetLogsForExportBatch(ctx context.Context, filter LogExportFilter, lastCreatedAt int64, lastID int, limit int, startIdx int) (logs []*Log, err error) { + if limit <= 0 { + return []*Log{}, nil + } + tx, err := buildLogExportQuery(filter) + if err != nil { + return nil, err + } + if lastCreatedAt > 0 || lastID > 0 { + tx = tx.Where("logs.created_at < ? OR (logs.created_at = ? AND logs.id < ?)", lastCreatedAt, lastCreatedAt, lastID) + } + err = tx.WithContext(ctx).Order("logs.created_at desc, logs.id desc").Limit(limit).Find(&logs).Error + if err != nil { + return nil, err + } + if filter.IsAdmin { + if err = fillLogChannelNames(logs); err != nil { + return logs, err + } + } else { + formatUserLogsForExport(logs) + } + return logs, nil +} + type Stat struct { Quota int `json:"quota"` Rpm int `json:"rpm"` diff --git a/model/main.go b/model/main.go index 6d9002462873..3adb785ca440 100644 --- a/model/main.go +++ b/model/main.go @@ -257,6 +257,7 @@ func migrateDB() error { err := DB.AutoMigrate( &Channel{}, + &ChannelPreparation{}, &Token{}, &User{}, &PasskeyCredential{}, @@ -281,10 +282,20 @@ func migrateDB() error { &CustomOAuthProvider{}, &UserOAuthBinding{}, &PerfMetric{}, + &NavigationMenu{}, + &NavigationItem{}, + &NavigationItemTranslation{}, + &NavigationVisibilityRule{}, + &CostReportTemplate{}, + &CostReportTemplateVersion{}, + &CostReportRun{}, + &CostReportRowSnapshot{}, + &CostReportManualCell{}, ) if err != nil { return err } + go seedDefaultNavigation() if common.UsingSQLite { if err := ensureSubscriptionPlanTableSQLite(); err != nil { return err @@ -306,6 +317,7 @@ func migrateDBFast() error { name string }{ {&Channel{}, "Channel"}, + {&ChannelPreparation{}, "ChannelPreparation"}, {&Token{}, "Token"}, {&User{}, "User"}, {&PasskeyCredential{}, "PasskeyCredential"}, @@ -330,6 +342,15 @@ func migrateDBFast() error { {&CustomOAuthProvider{}, "CustomOAuthProvider"}, {&UserOAuthBinding{}, "UserOAuthBinding"}, {&PerfMetric{}, "PerfMetric"}, + {&NavigationMenu{}, "NavigationMenu"}, + {&NavigationItem{}, "NavigationItem"}, + {&NavigationItemTranslation{}, "NavigationItemTranslation"}, + {&NavigationVisibilityRule{}, "NavigationVisibilityRule"}, + {&CostReportTemplate{}, "CostReportTemplate"}, + {&CostReportTemplateVersion{}, "CostReportTemplateVersion"}, + {&CostReportRun{}, "CostReportRun"}, + {&CostReportRowSnapshot{}, "CostReportRowSnapshot"}, + {&CostReportManualCell{}, "CostReportManualCell"}, } // 动态计算migration数量，确保errChan缓冲区足够大 errChan := make(chan error, len(migrations)) @@ -708,3 +729,137 @@ func PingDB() error { common.SysLog("Database pinged successfully") return nil } + +func seedDefaultNavigation() { + var count int64 + err := DB.Model(&NavigationMenu{}).Where("key = ?", "default_web_top").Count(&count).Error + if err != nil { + common.SysError("failed to query default_web_top menu: " + err.Error()) + return + } + if count > 0 { + return // 已经初始化过了，无需重复初始化 + } + + common.SysLog("Initializing default top navigation menu database records...") + + // 1. 创建默认顶部导航菜单 + menu := NavigationMenu{ + Key: "default_web_top", + Name: "默认顶部导航栏", + Client: "web_default", + Surface: "top", + Enabled: true, + IsSystem: true, + } + // 2. 初始内置模块定义 + type itemDef struct { + ModuleKey string + SortOrder int + IconKey string + Locales map[string]string + } + + defaultItems := []itemDef{ + { + ModuleKey: "home", + SortOrder: 1, + IconKey: "home", + Locales: map[string]string{ + "en": "Home", + "zh-CN": "首页", + "zh-TW": "首頁", + }, + }, + { + ModuleKey: "console", + SortOrder: 2, + IconKey: "layout-dashboard", + Locales: map[string]string{ + "en": "Console", + "zh-CN": "控制台", + "zh-TW": "控制台", + }, + }, + { + ModuleKey: "pricing", + SortOrder: 3, + IconKey: "credit-card", + Locales: map[string]string{ + "en": "Model Square", + "zh-CN": "模型广场", + "zh-TW": "模型廣場", + }, + }, + { + ModuleKey: "rankings", + SortOrder: 4, + IconKey: "trophy", + Locales: map[string]string{ + "en": "Rankings", + "zh-CN": "排行榜", + "zh-TW": "排行榜", + }, + }, + { + ModuleKey: "docs", + SortOrder: 5, + IconKey: "book-open", + Locales: map[string]string{ + "en": "Docs", + "zh-CN": "文档", + "zh-TW": "文檔", + }, + }, + { + ModuleKey: "about", + SortOrder: 6, + IconKey: "info", + Locales: map[string]string{ + "en": "About", + "zh-CN": "关于", + "zh-TW": "关于", + }, + }, + } + + // 开启事务进行菜单和菜单项的原子化创建 + err = DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Create(&menu).Error; err != nil { + return err + } + + for _, def := range defaultItems { + item := NavigationItem{ + MenuID: menu.ID, + Type: "builtin_module", + ModuleKey: def.ModuleKey, + IconKey: def.IconKey, + SortOrder: def.SortOrder, + Enabled: true, + } + if err := tx.Create(&item).Error; err != nil { + return err + } + + // 插入多语言翻译 + for locale, label := range def.Locales { + trans := NavigationItemTranslation{ + ItemID: item.ID, + Locale: locale, + Label: label, + } + if err := tx.Create(&trans).Error; err != nil { + return err + } + } + } + return nil + }) + + if err != nil { + common.SysError("failed to seed default navigation items: " + err.Error()) + } else { + common.SysLog("Default top navigation menu initialized successfully") + } +} diff --git a/model/navigation.go b/model/navigation.go new file mode 100644 index 000000000000..26eda28359c1 --- /dev/null +++ b/model/navigation.go @@ -0,0 +1,77 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +package model + +// NavigationMenu 定义导航菜单的集合类型（如顶部导航栏、侧边栏等） +type NavigationMenu struct { + ID uint `json:"id" gorm:"primaryKey"` + Key string `json:"key" gorm:"type:varchar(64);uniqueIndex;not null"` // 例如 "default_web_top" + Name string `json:"name" gorm:"type:varchar(128);not null"` // 菜单名称 + Client string `json:"client" gorm:"type:varchar(64);not null"` // "web_default", "mobile" 等 + Surface string `json:"surface" gorm:"type:varchar(64);not null"` // "top", "sidebar", "footer" 等 + Enabled bool `json:"enabled" gorm:"not null;default:true"` + IsSystem bool `json:"is_system" gorm:"not null;default:false"` // 系统置顶菜单，禁止删除 + CreatedAt int64 `json:"created_at" gorm:"type:bigint;autoCreateTime"` + UpdatedAt int64 `json:"updated_at" gorm:"type:bigint;autoUpdateTime"` +} + +// NavigationItem 树状嵌套的菜单节点 +type NavigationItem struct { + ID uint `json:"id" gorm:"primaryKey"` + MenuID uint `json:"menu_id" gorm:"index;not null"` + ParentID *uint `json:"parent_id" gorm:"index"` // 父节点ID，允许为 nil 表示顶级节点 + Type string `json:"type" gorm:"type:varchar(64);not null"` // builtin_module, internal_path, external_url, group, divider + ModuleKey string `json:"module_key" gorm:"type:varchar(128)"` // 内置模块对应的 Key（例如 "pricing"） + Path string `json:"path" gorm:"type:varchar(255)"` // 站内路径 + URL string `json:"url" gorm:"type:text"` // 外部链接 + IconKey string `json:"icon_key" gorm:"type:varchar(128)"` // Lucide/LobeHub 的图标对应键 + SortOrder int `json:"sort_order" gorm:"not null;default:0"` // 排序权重 + Enabled bool `json:"enabled" gorm:"not null;default:true"` + OpenInNewTab bool `json:"open_in_new_tab" gorm:"not null;default:false"` // 是否在新标签页打开 + ExactActive bool `json:"exact_active" gorm:"not null;default:false"` // 路由匹配时是否精确匹配 + CreatedAt int64 `json:"created_at" gorm:"type:bigint;autoCreateTime"` + UpdatedAt int64 `json:"updated_at" gorm:"type:bigint;autoUpdateTime"` + + // 关联字段，不写入数据库，由 GORM 自动处理级联操作 + Children []NavigationItem `json:"children,omitempty" gorm:"foreignKey:ParentID"` + Translations []NavigationItemTranslation `json:"translations,omitempty" gorm:"foreignKey:ItemID;constraint:OnDelete:CASCADE"` + Rules []NavigationVisibilityRule `json:"rules,omitempty" gorm:"foreignKey:ItemID;constraint:OnDelete:CASCADE"` +} + +// NavigationItemTranslation 支持导航节点多语言翻译的数据表 +type NavigationItemTranslation struct { + ID uint `json:"id" gorm:"primaryKey"` + ItemID uint `json:"item_id" gorm:"uniqueIndex:idx_item_locale;not null"` + Locale string `json:"locale" gorm:"type:varchar(32);uniqueIndex:idx_item_locale;not null"` // 区域标识，如 "zh-CN", "en-US", "zh-TW" + Label string `json:"label" gorm:"type:varchar(255);not null"` // 显示给用户的文字 + CreatedAt int64 `json:"created_at" gorm:"type:bigint;autoCreateTime"` + UpdatedAt int64 `json:"updated_at" gorm:"type:bigint;autoUpdateTime"` +} + +// NavigationVisibilityRule 控制导航节点精细可见性（如登录状态、角色等）的权限规则表 +type NavigationVisibilityRule struct { + ID uint `json:"id" gorm:"primaryKey"` + ItemID uint `json:"item_id" gorm:"index;not null"` + Effect string `json:"effect" gorm:"type:varchar(32);not null;default:'allow'"` // 作用效力："allow" 或 "deny" + SubjectType string `json:"subject_type" gorm:"type:varchar(64);not null"` // 主体类型：everyone, anonymous, authenticated, role, user_group + SubjectValue string `json:"subject_value" gorm:"type:varchar(255);not null"` // 主体对应的值（例如：role 时对应 "admin", "root"） + CreatedAt int64 `json:"created_at" gorm:"type:bigint;autoCreateTime"` + UpdatedAt int64 `json:"updated_at" gorm:"type:bigint;autoUpdateTime"` +} diff --git a/model/option.go b/model/option.go index ed1af72ebb12..183b9a57ebcd 100644 --- a/model/option.go +++ b/model/option.go @@ -1,6 +1,7 @@ package model import ( + "errors" "strconv" "strings" "time" @@ -182,6 +183,7 @@ func InitOptionMap() { common.OptionMapRWMutex.Unlock() loadOptionsFromDatabase() + backfillRequiredModelRatioOptions() } func loadOptionsFromDatabase() { @@ -194,6 +196,57 @@ func loadOptionsFromDatabase() { } } +var requiredModelRatioBackfills = map[string]float64{ + "claude-sonnet-4-6": 1.5, +} + +func backfillRequiredModelRatioOptions() { + var option Option + if err := DB.First(&option, "key = ?", "ModelRatio").Error; err != nil { + if !errors.Is(err, gorm.ErrRecordNotFound) { + common.SysLog("failed to load model ratio option for backfill: " + err.Error()) + } + return + } + + ratioMap := make(map[string]float64) + if err := common.Unmarshal([]byte(option.Value), &ratioMap); err != nil { + common.SysLog("failed to parse model ratio option for backfill: " + err.Error()) + return + } + + changed := false + for modelName, defaultRatio := range requiredModelRatioBackfills { + if _, ok := ratioMap[modelName]; ok { + continue + } + if ratio, ok := ratioMap["anthropic."+modelName]; ok { + defaultRatio = ratio + } else if ratio, ok := ratioMap["anthropic/"+modelName]; ok { + defaultRatio = ratio + } + ratioMap[modelName] = defaultRatio + changed = true + } + if !changed { + return + } + + payload, err := common.Marshal(ratioMap) + if err != nil { + common.SysLog("failed to marshal model ratio option for backfill: " + err.Error()) + return + } + option.Value = string(payload) + if err := DB.Save(&option).Error; err != nil { + common.SysLog("failed to save model ratio option backfill: " + err.Error()) + return + } + if err := updateOptionMap("ModelRatio", option.Value); err != nil { + common.SysLog("failed to reload model ratio option after backfill: " + err.Error()) + } +} + func SyncOptions(frequency int) { for { time.Sleep(time.Duration(frequency) * time.Second) diff --git a/model/option_test.go b/model/option_test.go new file mode 100644 index 000000000000..0a5128da492e --- /dev/null +++ b/model/option_test.go @@ -0,0 +1,88 @@ +package model + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/stretchr/testify/require" +) + +func setupOptionMapForTest(t *testing.T) { + t.Helper() + + common.OptionMapRWMutex.Lock() + previousOptionMap := common.OptionMap + common.OptionMap = make(map[string]string) + common.OptionMapRWMutex.Unlock() + + t.Cleanup(func() { + common.OptionMapRWMutex.Lock() + common.OptionMap = previousOptionMap + common.OptionMapRWMutex.Unlock() + }) +} + +func TestBackfillRequiredModelRatioOptionsAddsMissingClaudeSonnet46(t *testing.T) { + setupChannelPreparationModelTestDB(t) + setupOptionMapForTest(t) + require.NoError(t, DB.AutoMigrate(&Option{})) + + savedRatioConfig := ratio_setting.ModelRatio2JSONString() + t.Cleanup(func() { + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(savedRatioConfig)) + }) + + ratioConfig := map[string]float64{ + "anthropic.claude-sonnet-4-6": 1.7, + "claude-opus-4-7": 2.5, + } + payload, err := common.Marshal(ratioConfig) + require.NoError(t, err) + require.NoError(t, DB.Create(&Option{Key: "ModelRatio", Value: string(payload)}).Error) + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(string(payload))) + + backfillRequiredModelRatioOptions() + + var option Option + require.NoError(t, DB.First(&option, "key = ?", "ModelRatio").Error) + + updatedRatioConfig := map[string]float64{} + require.NoError(t, common.Unmarshal([]byte(option.Value), &updatedRatioConfig)) + require.Equal(t, 1.7, updatedRatioConfig["claude-sonnet-4-6"]) + require.Equal(t, 1.7, updatedRatioConfig["anthropic.claude-sonnet-4-6"]) + + ratio, ok, matchedModel := ratio_setting.GetModelRatio("claude-sonnet-4-6") + require.True(t, ok) + require.Equal(t, "claude-sonnet-4-6", matchedModel) + require.Equal(t, 1.7, ratio) +} + +func TestBackfillRequiredModelRatioOptionsPreservesExistingClaudeSonnet46(t *testing.T) { + setupChannelPreparationModelTestDB(t) + setupOptionMapForTest(t) + require.NoError(t, DB.AutoMigrate(&Option{})) + + savedRatioConfig := ratio_setting.ModelRatio2JSONString() + t.Cleanup(func() { + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(savedRatioConfig)) + }) + + ratioConfig := map[string]float64{ + "claude-sonnet-4-6": 1.2, + "anthropic.claude-sonnet-4-6": 1.7, + } + payload, err := common.Marshal(ratioConfig) + require.NoError(t, err) + require.NoError(t, DB.Create(&Option{Key: "ModelRatio", Value: string(payload)}).Error) + require.NoError(t, ratio_setting.UpdateModelRatioByJSONString(string(payload))) + + backfillRequiredModelRatioOptions() + + var option Option + require.NoError(t, DB.First(&option, "key = ?", "ModelRatio").Error) + + updatedRatioConfig := map[string]float64{} + require.NoError(t, common.Unmarshal([]byte(option.Value), &updatedRatioConfig)) + require.Equal(t, 1.2, updatedRatioConfig["claude-sonnet-4-6"]) +} diff --git a/relay/channel/claude/relay-claude.go b/relay/channel/claude/relay-claude.go index 18d7455e9f22..96b09c89d938 100644 --- a/relay/channel/claude/relay-claude.go +++ b/relay/channel/claude/relay-claude.go @@ -1,10 +1,13 @@ package claude import ( + "encoding/base64" "encoding/json" "fmt" "io" + "mime" "net/http" + "path/filepath" "strings" "github.com/QuantumNous/new-api/common" @@ -380,6 +383,48 @@ func RequestOpenAI2ClaudeMessage(c *gin.Context, textRequest dto.GeneralOpenAIRe Text: common.GetPointer[string](mediaMessage.Text), }) } + case dto.ContentTypeFile: + file := mediaMessage.GetFile() + if file == nil || file.FileData == "" { + continue + } + mimeTypeHint := mime.TypeByExtension(filepath.Ext(file.FileName)) + source := types.NewFileSourceFromData(file.FileData, mimeTypeHint) + base64Data, mimeType, err := service.GetBase64Data(c, source, "formatting file for Claude") + if err != nil { + return nil, fmt.Errorf("get file data failed: %s", err.Error()) + } + + switch { + case strings.HasPrefix(mimeType, "application/pdf"): + claudeMediaMessages = append(claudeMediaMessages, dto.ClaudeMediaMessage{ + Type: "document", + Source: &dto.ClaudeMessageSource{ + Type: "base64", + MediaType: mimeType, + Data: base64Data, + }, + }) + case strings.HasPrefix(mimeType, "text/"): + decoded, err := base64.StdEncoding.DecodeString(base64Data) + if err != nil { + return nil, fmt.Errorf("decode file data failed: %s", err.Error()) + } + claudeMediaMessages = append(claudeMediaMessages, dto.ClaudeMediaMessage{ + Type: "text", + Text: common.GetPointer[string](string(decoded)), + }) + case strings.HasPrefix(mimeType, "image/"): + claudeMediaMessages = append(claudeMediaMessages, dto.ClaudeMediaMessage{ + Type: "image", + Source: &dto.ClaudeMessageSource{ + Type: "base64", + MediaType: mimeType, + Data: base64Data, + }, + }) + } + continue default: source := mediaMessage.ToFileSource() if source == nil { @@ -389,20 +434,17 @@ func RequestOpenAI2ClaudeMessage(c *gin.Context, textRequest dto.GeneralOpenAIRe if err != nil { return nil, fmt.Errorf("get file data failed: %s", err.Error()) } - claudeMediaMessage := dto.ClaudeMediaMessage{ + if !strings.HasPrefix(mimeType, "image/") { + continue + } + claudeMediaMessages = append(claudeMediaMessages, dto.ClaudeMediaMessage{ + Type: "image", Source: &dto.ClaudeMessageSource{ - Type: "base64", + Type: "base64", + MediaType: mimeType, + Data: base64Data, }, - } - if strings.HasPrefix(mimeType, "application/pdf") { - claudeMediaMessage.Type = "document" - } else { - claudeMediaMessage.Type = "image" - } - - claudeMediaMessage.Source.MediaType = mimeType - claudeMediaMessage.Source.Data = base64Data - claudeMediaMessages = append(claudeMediaMessages, claudeMediaMessage) + }) continue } } diff --git a/router/api-router.go b/router/api-router.go index e98dc66ac048..1b18551cc5c2 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -54,6 +54,25 @@ func SetApiRouter(router *gin.Engine) { apiRouter.GET("/oauth/:provider", middleware.CriticalRateLimit(), controller.HandleOAuth) apiRouter.GET("/ratio_config", middleware.CriticalRateLimit(), controller.GetRatioConfig) + // 动态菜单树接口 + apiRouter.GET("/navigation/tree", middleware.TryUserAuth(), controller.GetNavigationTree) + + // 菜单管理后台路由（必须管理员及以上权限） + navigationAdminRoute := apiRouter.Group("/navigation/admin") + navigationAdminRoute.Use(middleware.AdminAuth()) + { + navigationAdminRoute.GET("/menus", controller.AdminGetMenus) + navigationAdminRoute.POST("/menus", controller.AdminCreateMenu) + navigationAdminRoute.PUT("/menus/:id", controller.AdminUpdateMenu) + navigationAdminRoute.DELETE("/menus/:id", controller.AdminDeleteMenu) + + navigationAdminRoute.GET("/items", controller.AdminGetItems) + navigationAdminRoute.POST("/items", controller.AdminCreateItem) + navigationAdminRoute.PUT("/items/:id", controller.AdminUpdateItem) + navigationAdminRoute.DELETE("/items/:id", controller.AdminDeleteItem) + navigationAdminRoute.POST("/items/reorder", controller.AdminReorderItems) + } + apiRouter.POST("/stripe/webhook", anonymousRequestBodyLimit, controller.StripeWebhook) apiRouter.POST("/creem/webhook", anonymousRequestBodyLimit, controller.CreemWebhook) apiRouter.POST("/waffo/webhook", anonymousRequestBodyLimit, controller.WaffoWebhook) @@ -225,19 +244,54 @@ func SetApiRouter(router *gin.Engine) { ratioSyncRoute.GET("/channels", controller.GetSyncableChannels) ratioSyncRoute.POST("/fetch", controller.FetchUpstreamRatios) } + costReportsRoute := apiRouter.Group("/cost_reports") + costReportsRoute.Use(middleware.RootAuth()) + { + costReportsRoute.GET("/templates", controller.CostReportListTemplates) + costReportsRoute.POST("/templates", controller.CostReportSaveTemplate) + costReportsRoute.POST("/templates/default", controller.CostReportEnsureDefaultTemplate) + costReportsRoute.POST("/templates/validate", controller.CostReportValidateTemplate) + costReportsRoute.GET("/templates/:id", controller.CostReportGetTemplate) + costReportsRoute.PUT("/templates/:id", controller.CostReportSaveTemplate) + costReportsRoute.GET("/templates/:id/versions", controller.CostReportListTemplateVersions) + costReportsRoute.POST("/preview", controller.CostReportPreview) + costReportsRoute.POST("/classification/preview", controller.CostReportClassificationPreview) + costReportsRoute.GET("/manual_cells", controller.CostReportReadManualCells) + costReportsRoute.POST("/manual_cells", controller.CostReportUpsertManualCell) + costReportsRoute.GET("/runs", controller.CostReportListRuns) + costReportsRoute.POST("/runs", controller.CostReportSaveRun) + costReportsRoute.GET("/runs/:id", controller.CostReportGetRun) + costReportsRoute.GET("/runs/:id/export", controller.CostReportExportRun) + } channelRoute := apiRouter.Group("/channel") channelRoute.Use(middleware.AdminAuth()) { channelRoute.GET("/", controller.GetAllChannels) channelRoute.GET("/search", controller.SearchChannels) + channelRoute.POST("/search/keys", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.SearchChannelsByKeys) + channelRoute.POST("/query-key/report", middleware.DisableCache(), controller.QueryChannelKeyReport) + channelRoute.POST("/query-key/test", middleware.DisableCache(), controller.QueryChannelKeyTest) channelRoute.GET("/models", controller.ChannelListModels) channelRoute.GET("/models_enabled", controller.EnabledListModels) + channelRoute.GET("/preparations", controller.GetChannelPreparations) + channelRoute.POST("/preparations", controller.AddChannelPreparation) + channelRoute.POST("/preparations/import", controller.ImportChannelPreparations) + channelRoute.POST("/preparations/batch/promote", controller.PromoteChannelPreparationsBatch) + channelRoute.GET("/preparations/auto-promotion/status", controller.GetChannelPreparationAutoPromotionSchedulerStatus) + channelRoute.POST("/preparations/auto-promotion/run", controller.RunChannelPreparationAutoPromotionManually) + channelRoute.GET("/preparations/:id/test", controller.TestChannelPreparation) + channelRoute.GET("/preparations/:id", controller.GetChannelPreparation) + channelRoute.PUT("/preparations/:id", controller.UpdateChannelPreparation) + channelRoute.DELETE("/preparations/:id", controller.DeleteChannelPreparation) + channelRoute.POST("/preparations/:id/promote", controller.PromoteChannelPreparation) channelRoute.GET("/:id", controller.GetChannel) channelRoute.POST("/:id/key", middleware.RootAuth(), middleware.CriticalRateLimit(), middleware.DisableCache(), middleware.SecureVerificationRequired(), controller.GetChannelKey) channelRoute.GET("/test", controller.TestAllChannels) channelRoute.GET("/test/:id", controller.TestChannel) channelRoute.GET("/update_balance", controller.UpdateAllChannelsBalance) channelRoute.GET("/update_balance/:id", controller.UpdateChannelBalance) + channelRoute.POST("/balance/:id", controller.SetChannelBalance) + channelRoute.POST("/used_quota/clear/:id", controller.ClearChannelUsedQuota) channelRoute.POST("/", controller.AddChannel) channelRoute.PUT("/", controller.UpdateChannel) channelRoute.DELETE("/disabled", controller.DeleteDisabledChannel) @@ -307,7 +361,11 @@ func SetApiRouter(router *gin.Engine) { logRoute.GET("/", middleware.AdminAuth(), controller.GetAllLogs) logRoute.DELETE("/", middleware.AdminAuth(), controller.DeleteHistoryLogs) logRoute.GET("/stat", middleware.AdminAuth(), controller.GetLogsStat) + logRoute.GET("/export_fields", middleware.AdminAuth(), controller.GetLogExportFields) + logRoute.GET("/export", middleware.AdminAuth(), controller.ExportAllLogs) logRoute.GET("/self/stat", middleware.UserAuth(), controller.GetLogsSelfStat) + logRoute.GET("/self/export_fields", middleware.UserAuth(), controller.GetUserLogExportFields) + logRoute.GET("/self/export", middleware.UserAuth(), controller.ExportUserLogs) logRoute.GET("/channel_affinity_usage_cache", middleware.AdminAuth(), controller.GetChannelAffinityUsageCacheStats) logRoute.GET("/search", middleware.AdminAuth(), controller.SearchAllLogs) logRoute.GET("/self", middleware.UserAuth(), controller.GetUserLogs) diff --git a/service/channel_affinity.go b/service/channel_affinity.go index 96ec13e248cc..8cde3f7f33d7 100644 --- a/service/channel_affinity.go +++ b/service/channel_affinity.go @@ -10,7 +10,9 @@ import ( "time" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/pkg/cachex" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/QuantumNous/new-api/types" @@ -623,6 +625,72 @@ func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup return 0, false } +func GetUsablePreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup string) (*model.Channel, string, bool) { + channelID, found := GetPreferredChannelByAffinity(c, modelName, usingGroup) + if !found { + return nil, "", false + } + + preferred, err := model.CacheGetChannel(channelID) + if err != nil || preferred == nil { + if !ShouldKeepChannelAffinityOnChannelDisabled() { + ClearCurrentChannelAffinityCache(c) + } + return nil, "", false + } + + selectedGroup, ok := validateChannelAffinityHit(c, preferred, modelName, usingGroup) + if !ok { + if !ShouldKeepChannelAffinityOnChannelDisabled() { + ClearCurrentChannelAffinityCache(c) + } + return nil, "", false + } + + MarkChannelAffinityUsed(c, selectedGroup, preferred.Id) + return preferred, selectedGroup, true +} + +func validateChannelAffinityHit(c *gin.Context, channel *model.Channel, modelName string, usingGroup string) (string, bool) { + if channel == nil || channel.Id <= 0 { + return "", false + } + if channel.Status != common.ChannelStatusEnabled { + return "", false + } + + if usingGroup == "auto" { + userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) + autoGroups := GetUserAutoGroup(userGroup) + for _, group := range autoGroups { + if model.IsChannelEnabledForGroupModel(group, modelName, channel.Id) { + common.SetContextKey(c, constant.ContextKeyAutoGroup, group) + return group, true + } + } + return "", false + } + + if model.IsChannelEnabledForGroupModel(usingGroup, modelName, channel.Id) { + return usingGroup, true + } + return "", false +} + +func DiscardChannelAffinityCacheForContext(c *gin.Context) bool { + cacheKey, _, ok := getChannelAffinityContext(c) + if !ok || cacheKey == "" { + return false + } + cache := getChannelAffinityCache() + deleted, err := cache.DeleteMany([]string{cacheKey}) + if err != nil { + common.SysError(fmt.Sprintf("channel affinity cache delete failed: key=%s, err=%v", cacheKey, err)) + return false + } + return deleted[cacheKey] +} + func ShouldSkipRetryAfterChannelAffinityFailure(c *gin.Context) bool { if c == nil { return false diff --git a/service/cost_report/aggregation.go b/service/cost_report/aggregation.go new file mode 100644 index 000000000000..992a7b0d9140 --- /dev/null +++ b/service/cost_report/aggregation.go @@ -0,0 +1,632 @@ +package cost_report + +import ( + "context" + "crypto/sha256" + "fmt" + "sort" + "strconv" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "gorm.io/gorm" +) + +type Service struct { + db *gorm.DB + logDB *gorm.DB +} + +func NewService(db *gorm.DB, logDB *gorm.DB) *Service { + if db == nil { + db = model.DB + } + if logDB == nil { + logDB = model.LOG_DB + } + return &Service{db: db, logDB: logDB} +} + +type PreviewRequest struct { + TemplateID int `json:"template_id"` + TemplateVersionID int `json:"template_version_id"` + Config *CostReportTemplateConfig `json:"config,omitempty"` + PeriodStart int64 `json:"period_start"` + PeriodEnd int64 `json:"period_end"` + PeriodKey string `json:"period_key"` + IncludeManual bool `json:"include_manual"` + MaxLogs int `json:"max_logs,omitempty"` +} + +type PreviewResponse struct { + TemplateID int `json:"template_id"` + TemplateVersionID int `json:"template_version_id"` + PeriodStart int64 `json:"period_start"` + PeriodEnd int64 `json:"period_end"` + PeriodKey string `json:"period_key"` + Timezone string `json:"timezone"` + SourceLogMaxID int `json:"source_log_max_id"` + Rows []PreviewRow `json:"rows"` + Warnings []string `json:"warnings,omitempty"` +} + +type PreviewRow struct { + RowKey string `json:"row_key"` + Dimensions map[string]interface{} `json:"dimensions"` + Metrics map[string]interface{} `json:"metrics"` + ManualValues map[string]interface{} `json:"manual_values"` + FormulaValues map[string]interface{} `json:"formula_values"` + Values map[string]interface{} `json:"values"` + ManualOverrides map[string]bool `json:"-"` +} + +type metricAccumulator struct { + aggregate string + count int + sum float64 + min float64 + max float64 + set bool +} + +type rowAccumulator struct { + row PreviewRow + metrics map[string]*metricAccumulator +} + +const consumeLogScanBatchSize = 1000 + +func (s *Service) Preview(ctx context.Context, request PreviewRequest) (*PreviewResponse, error) { + if s == nil || s.db == nil || s.logDB == nil { + return nil, fmt.Errorf("db and log_db are required") + } + if request.PeriodStart <= 0 || request.PeriodEnd <= request.PeriodStart { + return nil, fmt.Errorf("valid period_start and period_end are required") + } + + config, templateID, versionID, err := s.resolvePreviewConfig(ctx, request) + if err != nil { + return nil, err + } + if err := ValidateTemplateConfig(config); err != nil { + return nil, err + } + periodKey := request.PeriodKey + if periodKey == "" { + periodKey = defaultPeriodKey(config, request.PeriodStart) + } + loc, _ := time.LoadLocation(config.Timezone) + + sourceLogMaxID := 0 + channelIDs := map[int]bool{} + userIDs := map[int]bool{} + if err := s.scanConsumeLogs(ctx, request.PeriodStart, request.PeriodEnd, request.MaxLogs, func(logs []model.Log) error { + for i := range logs { + if logs[i].Id > sourceLogMaxID { + sourceLogMaxID = logs[i].Id + } + if logs[i].ChannelId > 0 { + channelIDs[logs[i].ChannelId] = true + } + if logs[i].UserId > 0 { + userIDs[logs[i].UserId] = true + } + } + return nil + }); err != nil { + return nil, err + } + channels, err := s.loadChannels(ctx, setKeys(channelIDs)) + if err != nil { + return nil, err + } + users, err := s.loadUsers(ctx, setKeys(userIDs)) + if err != nil { + return nil, err + } + + fieldsByKey := map[string]FieldConfig{} + for _, field := range config.Fields { + fieldsByKey[field.Key] = field + } + rowsByKey := map[string]*rowAccumulator{} + if err := s.scanConsumeLogs(ctx, request.PeriodStart, request.PeriodEnd, request.MaxLogs, func(logs []model.Log) error { + for i := range logs { + log := &logs[i] + channel := channels[log.ChannelId] + user := users[log.UserId] + other := parseLogOther(log.Other) + classification := Classify(config, ClassificationInput{Log: log, Channel: channel, User: user, LogOther: other}) + + baseValues := map[string]interface{}{} + for _, field := range config.Fields { + if field.Kind != FieldKindDimension { + continue + } + baseValues[field.Key] = dimensionValue(field.Source, log, channel, user, other, classification, loc) + } + rowKey := makeRowKey(config.Grouping, baseValues) + acc := rowsByKey[rowKey] + if acc == nil { + acc = &rowAccumulator{ + row: PreviewRow{ + RowKey: rowKey, + Dimensions: map[string]interface{}{}, + Metrics: map[string]interface{}{}, + ManualValues: map[string]interface{}{}, + FormulaValues: map[string]interface{}{}, + Values: map[string]interface{}{}, + ManualOverrides: map[string]bool{}, + }, + metrics: map[string]*metricAccumulator{}, + } + for key, value := range baseValues { + acc.row.Dimensions[key] = value + acc.row.Values[key] = value + } + rowsByKey[rowKey] = acc + } + for _, field := range config.Fields { + if field.Kind != FieldKindMetric { + continue + } + value, ok := metricSourceValue(field.Source, log, other) + if !ok { + continue + } + ma := acc.metrics[field.Key] + if ma == nil { + ma = &metricAccumulator{aggregate: field.Aggregate} + acc.metrics[field.Key] = ma + } + ma.add(value) + } + } + return nil + }); err != nil { + return nil, err + } + + rows := make([]PreviewRow, 0, len(rowsByKey)) + for _, acc := range rowsByKey { + for _, field := range config.Fields { + if field.Kind != FieldKindMetric { + continue + } + value := interface{}(float64(0)) + if ma := acc.metrics[field.Key]; ma != nil { + value = ma.value() + } + acc.row.Metrics[field.Key] = value + acc.row.Values[field.Key] = value + } + rows = append(rows, acc.row) + } + sortPreviewRows(rows, config.Sort) + for i := range rows { + if field, ok := fieldsByKey["row_index"]; ok && field.Source == "generated.row_index" { + rows[i].Dimensions["row_index"] = i + 1 + rows[i].Values["row_index"] = i + 1 + } + } + + applyManualDefaults(fieldsByKey, rows) + if request.IncludeManual && templateID > 0 { + if err := s.mergePersistedManualValues(ctx, templateID, periodKey, fieldsByKey, rows); err != nil { + return nil, err + } + } + + formulaEval, err := newFormulaEvaluator(config) + if err != nil { + return nil, err + } + warnings := formulaEval.evaluateRows(rows) + + return &PreviewResponse{ + TemplateID: templateID, + TemplateVersionID: versionID, + PeriodStart: request.PeriodStart, + PeriodEnd: request.PeriodEnd, + PeriodKey: periodKey, + Timezone: config.Timezone, + SourceLogMaxID: sourceLogMaxID, + Rows: rows, + Warnings: warnings, + }, nil +} + +func (s *Service) resolvePreviewConfig(ctx context.Context, request PreviewRequest) (CostReportTemplateConfig, int, int, error) { + if request.Config != nil { + return *request.Config, request.TemplateID, request.TemplateVersionID, nil + } + if request.TemplateVersionID > 0 { + var version model.CostReportTemplateVersion + if err := s.db.WithContext(ctx).First(&version, request.TemplateVersionID).Error; err != nil { + return CostReportTemplateConfig{}, 0, 0, err + } + var config CostReportTemplateConfig + if err := common.UnmarshalJsonStr(version.ConfigJson, &config); err != nil { + return CostReportTemplateConfig{}, 0, 0, err + } + return config, version.TemplateId, version.Id, nil + } + if request.TemplateID <= 0 { + return CostReportTemplateConfig{}, 0, 0, fmt.Errorf("template_id or config is required") + } + var template model.CostReportTemplate + if err := s.db.WithContext(ctx).First(&template, request.TemplateID).Error; err != nil { + return CostReportTemplateConfig{}, 0, 0, err + } + if template.CurrentVersionId == nil { + return CostReportTemplateConfig{}, 0, 0, fmt.Errorf("template has no current version") + } + return s.resolvePreviewConfig(ctx, PreviewRequest{TemplateVersionID: *template.CurrentVersionId}) +} + +func (s *Service) fetchConsumeLogs(ctx context.Context, start, end int64, maxLogs int) ([]model.Log, error) { + logs := []model.Log{} + err := s.scanConsumeLogs(ctx, start, end, maxLogs, func(batch []model.Log) error { + logs = append(logs, batch...) + return nil + }) + return logs, err +} + +func (s *Service) scanConsumeLogs(ctx context.Context, start, end int64, maxLogs int, handle func([]model.Log) error) error { + if handle == nil { + return fmt.Errorf("consume log handler is required") + } + remaining := maxLogs + lastCreatedAt := int64(-1) + lastID := 0 + for { + limit := consumeLogScanBatchSize + if remaining > 0 && remaining < limit { + limit = remaining + } + query := s.logDB.WithContext(ctx). + Where("type = ? AND created_at >= ? AND created_at < ?", model.LogTypeConsume, start, end) + if lastCreatedAt >= 0 { + query = query.Where("created_at > ? OR (created_at = ? AND id > ?)", lastCreatedAt, lastCreatedAt, lastID) + } + var batch []model.Log + if err := query.Order("created_at asc, id asc").Limit(limit).Find(&batch).Error; err != nil { + return err + } + if len(batch) == 0 { + return nil + } + if err := handle(batch); err != nil { + return err + } + last := batch[len(batch)-1] + lastCreatedAt = last.CreatedAt + lastID = last.Id + if remaining > 0 { + remaining -= len(batch) + if remaining <= 0 { + return nil + } + } + if len(batch) < limit { + return nil + } + } +} + +func (s *Service) loadChannels(ctx context.Context, ids []int) (map[int]*model.Channel, error) { + result := map[int]*model.Channel{} + if len(ids) == 0 { + return result, nil + } + var channels []model.Channel + if err := s.db.WithContext(ctx).Select("id", "type", "name", "models").Where("id IN ?", ids).Find(&channels).Error; err != nil { + return nil, err + } + for i := range channels { + channel := channels[i] + result[channel.Id] = &channel + } + return result, nil +} + +func (s *Service) loadUsers(ctx context.Context, ids []int) (map[int]*model.User, error) { + result := map[int]*model.User{} + if len(ids) == 0 { + return result, nil + } + var users []model.User + if err := s.db.WithContext(ctx).Select("id", "username", "display_name").Where("id IN ?", ids).Find(&users).Error; err != nil { + return nil, err + } + for i := range users { + user := users[i] + result[user.Id] = &user + } + return result, nil +} + +func applyManualDefaults(fields map[string]FieldConfig, rows []PreviewRow) { + for i := range rows { + for key, field := range fields { + if field.Kind != FieldKindManual { + continue + } + if _, exists := rows[i].Values[key]; exists { + continue + } + value := parseManualValue(field.ValueType, field.DefaultValue) + rows[i].ManualValues[key] = value + rows[i].Values[key] = value + } + } +} + +func (s *Service) mergePersistedManualValues(ctx context.Context, templateID int, periodKey string, fields map[string]FieldConfig, rows []PreviewRow) error { + if templateID <= 0 || periodKey == "" || len(rows) == 0 { + return nil + } + rowKeys := make([]string, 0, len(rows)) + for _, row := range rows { + rowKeys = append(rowKeys, row.RowKey) + } + manuals, err := s.ReadManualCells(ctx, templateID, periodKey, rowKeys) + if err != nil { + return err + } + for i := range rows { + for fieldKey, manual := range manuals[rows[i].RowKey] { + field, ok := fields[fieldKey] + if !ok || (field.Kind != FieldKindManual && !(field.Kind == FieldKindFormula && field.ManualOverride)) { + continue + } + rows[i].ManualValues[fieldKey] = manual.Value + rows[i].Values[fieldKey] = manual.Value + if field.Kind == FieldKindFormula { + rows[i].FormulaValues[fieldKey] = manual.Value + rows[i].ManualOverrides[fieldKey] = true + } + } + } + return nil +} + +func (ma *metricAccumulator) add(value float64) { + ma.count++ + ma.sum += value + if !ma.set || value < ma.min { + ma.min = value + } + if !ma.set || value > ma.max { + ma.max = value + } + ma.set = true +} + +func (ma *metricAccumulator) value() interface{} { + if ma == nil || !ma.set { + return float64(0) + } + switch ma.aggregate { + case "count": + return ma.count + case "avg": + if ma.count == 0 { + return float64(0) + } + return ma.sum / float64(ma.count) + case "min": + return ma.min + case "max": + return ma.max + default: + return ma.sum + } +} + +func dimensionValue(source string, log *model.Log, channel *model.Channel, user *model.User, other map[string]interface{}, classification ClassificationResult, loc *time.Location) interface{} { + switch source { + case "generated.row_index": + return 0 + case "period.date": + return time.Unix(log.CreatedAt, 0).In(loc).Format("2006-01-02") + case "log.username": + return log.Username + case "log.user_id": + return log.UserId + case "log.channel_id": + return log.ChannelId + case "log.model_name": + return log.ModelName + case "log.group": + return log.Group + case "classification.output": + return classification.Class + case "channel.name": + if channel == nil { + return "" + } + return channel.Name + case "channel.type": + if channel == nil { + return 0 + } + return channel.Type + case "user.display_name": + if user == nil { + return "" + } + return user.DisplayName + default: + if strings.HasPrefix(source, "log_other.") { + value, _ := nestedMapValue(other, strings.TrimPrefix(source, "log_other.")) + return value + } + return "" + } +} + +func metricSourceValue(source string, log *model.Log, other map[string]interface{}) (float64, bool) { + switch source { + case "log.created_at": + return float64(log.CreatedAt), true + case "log.quota": + return float64(log.Quota), true + case "log.quota_per_unit": + if common.QuotaPerUnit <= 0 { + return float64(log.Quota), true + } + return float64(log.Quota) / common.QuotaPerUnit, true + case "log.prompt_tokens": + return float64(log.PromptTokens), true + case "log.completion_tokens": + return float64(log.CompletionTokens), true + case "log.total_tokens": + return float64(log.PromptTokens + log.CompletionTokens), true + case "log.request_count": + return 1, true + default: + if strings.HasPrefix(source, "log_other.") { + value, ok := nestedMapValue(other, strings.TrimPrefix(source, "log_other.")) + if !ok { + return 0, false + } + return toFloat64(value), true + } + return 0, false + } +} + +func parseLogOther(text string) map[string]interface{} { + if strings.TrimSpace(text) == "" { + return map[string]interface{}{} + } + var out map[string]interface{} + if err := common.Unmarshal([]byte(text), &out); err != nil || out == nil { + return map[string]interface{}{} + } + return out +} + +func makeRowKey(grouping []string, values map[string]interface{}) string { + parts := normalizedRowKeyParts(grouping, values) + payload, err := common.Marshal(parts) + if err != nil { + payload = []byte(strings.Join(parts, "\x1f")) + } + hash := sha256.Sum256(payload) + return fmt.Sprintf("%x", hash) +} + +func normalizedRowKeyParts(grouping []string, values map[string]interface{}) []string { + parts := make([]string, 0, len(grouping)) + for _, key := range grouping { + parts = append(parts, key+"="+valueToString(values[key])) + } + return parts +} + +func sortPreviewRows(rows []PreviewRow, sortRules []SortConfig) { + sort.SliceStable(rows, func(i, j int) bool { + for _, rule := range sortRules { + cmp := compareValues(rows[i].Values[rule.Field], rows[j].Values[rule.Field]) + if cmp == 0 { + continue + } + if strings.ToLower(rule.Direction) == "desc" { + return cmp > 0 + } + return cmp < 0 + } + return rows[i].RowKey < rows[j].RowKey + }) +} + +func compareValues(a, b interface{}) int { + af, aok := numericValue(a) + bf, bok := numericValue(b) + if aok && bok { + if af < bf { + return -1 + } + if af > bf { + return 1 + } + return 0 + } + as := valueToString(a) + bs := valueToString(b) + if as < bs { + return -1 + } + if as > bs { + return 1 + } + return 0 +} + +func numericValue(value interface{}) (float64, bool) { + switch v := value.(type) { + case int: + return float64(v), true + case int64: + return float64(v), true + case float64: + return v, true + case float32: + return float64(v), true + default: + return 0, false + } +} + +func toFloat64(value interface{}) float64 { + if value == nil { + return 0 + } + if f, ok := numericValue(value); ok { + return f + } + switch v := value.(type) { + case string: + f, _ := strconv.ParseFloat(strings.TrimSpace(v), 64) + return f + default: + f, _ := strconv.ParseFloat(fmt.Sprint(v), 64) + return f + } +} + +func valueToString(value interface{}) string { + if value == nil { + return "" + } + switch v := value.(type) { + case string: + return v + case fmt.Stringer: + return v.String() + default: + return fmt.Sprint(v) + } +} + +func setKeys(m map[int]bool) []int { + keys := make([]int, 0, len(m)) + for key := range m { + keys = append(keys, key) + } + sort.Ints(keys) + return keys +} + +func defaultPeriodKey(config CostReportTemplateConfig, start int64) string { + loc, err := time.LoadLocation(config.Timezone) + if err != nil { + loc = time.UTC + } + return time.Unix(start, 0).In(loc).Format("2006-01-02") +} diff --git a/service/cost_report/classification.go b/service/cost_report/classification.go new file mode 100644 index 000000000000..648af15589b0 --- /dev/null +++ b/service/cost_report/classification.go @@ -0,0 +1,211 @@ +package cost_report + +import ( + "fmt" + "regexp" + "sort" + "strings" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" +) + +type ClassificationInput struct { + Log *model.Log + Channel *model.Channel + User *model.User + LogOther map[string]interface{} +} + +type ClassificationResult struct { + RuleKey string `json:"rule_key"` + Class string `json:"class"` +} + +func Classify(config CostReportTemplateConfig, input ClassificationInput) ClassificationResult { + rules := append([]ClassificationRuleConfig(nil), config.ClassificationRules...) + sort.SliceStable(rules, func(i, j int) bool { + return rules[i].Priority < rules[j].Priority + }) + + fallback := ClassificationResult{Class: "Other"} + for _, rule := range rules { + if !rule.Enabled { + continue + } + if rule.Fallback { + fallback = ClassificationResult{RuleKey: rule.Key, Class: rule.OutputClass} + continue + } + if classificationRuleMatches(rule, input) { + return ClassificationResult{RuleKey: rule.Key, Class: rule.OutputClass} + } + } + return fallback +} + +func classificationRuleMatches(rule ClassificationRuleConfig, input ClassificationInput) bool { + parts := make([]bool, 0, len(rule.Conditions)+len(rule.ConditionGroups)) + for _, condition := range rule.Conditions { + parts = append(parts, classificationConditionMatches(condition, input)) + } + for _, group := range rule.ConditionGroups { + parts = append(parts, classificationGroupMatches(group, input)) + } + return matchBools(rule.Match, parts) +} + +func classificationGroupMatches(group ClassificationConditionGroup, input ClassificationInput) bool { + parts := make([]bool, 0, len(group.Conditions)) + for _, condition := range group.Conditions { + parts = append(parts, classificationConditionMatches(condition, input)) + } + return matchBools(group.Match, parts) +} + +func matchBools(match string, values []bool) bool { + if len(values) == 0 { + return false + } + if match == "any" { + for _, value := range values { + if value { + return true + } + } + return false + } + for _, value := range values { + if !value { + return false + } + } + return true +} + +func classificationConditionMatches(condition ClassificationCondition, input ClassificationInput) bool { + actual, exists := classificationSourceValue(condition.Source, input) + if condition.Operator == "exists" { + return exists && strings.TrimSpace(actual) != "" + } + if !exists { + return false + } + + actualCmp := actual + want := condition.Value + values := condition.Values + if condition.CaseInsensitive { + actualCmp = strings.ToLower(actualCmp) + want = strings.ToLower(want) + values = make([]string, len(condition.Values)) + for i, value := range condition.Values { + values[i] = strings.ToLower(value) + } + } + + switch condition.Operator { + case "equals": + return actualCmp == want + case "contains": + return strings.Contains(actualCmp, want) + case "in": + for _, value := range values { + if actualCmp == value { + return true + } + } + return false + case "regex": + pattern := condition.Value + if condition.CaseInsensitive { + pattern = "(?i)" + pattern + } + re, err := regexp.Compile(pattern) + if err != nil { + return false + } + return re.MatchString(actual) + default: + return false + } +} + +func classificationSourceValue(source string, input ClassificationInput) (string, bool) { + switch source { + case "channel.type": + if input.Channel == nil { + return "", false + } + return fmt.Sprintf("%d", input.Channel.Type), true + case "channel.name": + if input.Channel == nil { + return "", false + } + return input.Channel.Name, true + case "channel.id": + if input.Log == nil || input.Log.ChannelId == 0 { + return "", false + } + return fmt.Sprintf("%d", input.Log.ChannelId), true + case "model_name": + if input.Log == nil { + return "", false + } + return input.Log.ModelName, true + case "group": + if input.Log == nil { + return "", false + } + return input.Log.Group, true + case "is_claude_related": + return fmt.Sprintf("%t", isClaudeRelated(input)), true + default: + if strings.HasPrefix(source, "log_other.") { + value, ok := nestedMapValue(input.LogOther, strings.TrimPrefix(source, "log_other.")) + if !ok || value == nil { + return "", false + } + return fmt.Sprint(value), true + } + return "", false + } +} + +func isClaudeRelated(input ClassificationInput) bool { + needles := []string{} + if input.Log != nil { + needles = append(needles, input.Log.ModelName, input.Log.Content) + } + if input.Channel != nil { + if input.Channel.Type == constant.ChannelTypeAnthropic || input.Channel.Type == constant.ChannelTypeAws { + return true + } + needles = append(needles, input.Channel.Name, input.Channel.Models) + } + for _, value := range needles { + if strings.Contains(strings.ToLower(value), "claude") { + return true + } + } + return false +} + +func nestedMapValue(root map[string]interface{}, dotted string) (interface{}, bool) { + if root == nil || dotted == "" { + return nil, false + } + parts := strings.Split(dotted, ".") + var current interface{} = root + for _, part := range parts { + m, ok := current.(map[string]interface{}) + if !ok { + return nil, false + } + current, ok = m[part] + if !ok { + return nil, false + } + } + return current, true +} diff --git a/service/cost_report/config.go b/service/cost_report/config.go new file mode 100644 index 000000000000..7e461380818d --- /dev/null +++ b/service/cost_report/config.go @@ -0,0 +1,671 @@ +package cost_report + +import ( + "crypto/sha256" + "fmt" + "math" + "regexp" + "sort" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/expr-lang/expr" + "github.com/expr-lang/expr/ast" + "gorm.io/gorm" +) + +const ( + DefaultTemplateKey = "claude_cost_default" + + PeriodModeDay = "day" + PeriodModeCustom = "custom" + + FieldKindDimension = "dimension" + FieldKindMetric = "metric" + FieldKindManual = "manual" + FieldKindFormula = "formula" + + FormulaModeStandard = "standard" + FormulaModeRunning = "running" +) + +var ( + identifierRE = regexp.MustCompile(`^[a-z][a-z0-9_]{0,63}$`) + logOtherSourceRE = regexp.MustCompile(`^log_other\.[A-Za-z_][A-Za-z0-9_]*(\.[A-Za-z_][A-Za-z0-9_]*)*$`) +) + +type CostReportTemplateConfig struct { + Timezone string `json:"timezone"` + PeriodMode string `json:"period_mode"` + Grouping []string `json:"grouping"` + Sort []SortConfig `json:"sort"` + Fields []FieldConfig `json:"fields"` + ClassificationRules []ClassificationRuleConfig `json:"classification_rules"` + ExportLayout ExportLayoutConfig `json:"export_layout"` +} + +type SortConfig struct { + Field string `json:"field"` + Direction string `json:"direction"` +} + +type FieldConfig struct { + Key string `json:"key"` + Label string `json:"label"` + Kind string `json:"kind"` + ValueType string `json:"value_type"` + Source string `json:"source,omitempty"` + Aggregate string `json:"aggregate,omitempty"` + Expression string `json:"expression,omitempty"` + InitialExpression string `json:"initial_expression,omitempty"` + FormulaMode string `json:"formula_mode,omitempty"` + DefaultValue string `json:"default_value,omitempty"` + Visible bool `json:"visible"` + Exportable bool `json:"exportable"` + Order int `json:"order"` + ManualOverride bool `json:"manual_override,omitempty"` + Generated bool `json:"generated,omitempty"` +} + +type ClassificationRuleConfig struct { + Key string `json:"key"` + Label string `json:"label"` + Priority int `json:"priority"` + Enabled bool `json:"enabled"` + Match string `json:"match,omitempty"` + Conditions []ClassificationCondition `json:"conditions,omitempty"` + ConditionGroups []ClassificationConditionGroup `json:"condition_groups,omitempty"` + OutputClass string `json:"output_class"` + Fallback bool `json:"fallback,omitempty"` +} + +type ClassificationConditionGroup struct { + Match string `json:"match"` + Conditions []ClassificationCondition `json:"conditions"` +} + +type ClassificationCondition struct { + Source string `json:"source"` + Operator string `json:"operator"` + Value string `json:"value,omitempty"` + Values []string `json:"values,omitempty"` + CaseInsensitive bool `json:"case_insensitive,omitempty"` +} + +type ExportLayoutConfig struct { + SheetName string `json:"sheet_name"` + FreezeHeader bool `json:"freeze_header"` + IncludeMeta bool `json:"include_meta"` + DateFormat string `json:"date_format"` + DecimalFormat string `json:"decimal_format"` +} + +func DefaultClaudeCostTemplateConfig() CostReportTemplateConfig { + return CostReportTemplateConfig{ + Timezone: "Asia/Shanghai", + PeriodMode: PeriodModeDay, + Grouping: []string{"report_date", "customer", "channel_class", "channel_id"}, + Sort: []SortConfig{ + {Field: "report_date", Direction: "asc"}, + {Field: "customer", Direction: "asc"}, + {Field: "channel_class", Direction: "asc"}, + {Field: "channel_id", Direction: "asc"}, + }, + Fields: []FieldConfig{ + {Key: "row_index", Label: "序号", Kind: FieldKindDimension, ValueType: "integer", Source: "generated.row_index", Visible: true, Exportable: true, Order: 10, Generated: true}, + {Key: "report_date", Label: "日期", Kind: FieldKindDimension, ValueType: "date", Source: "period.date", Visible: true, Exportable: true, Order: 20}, + {Key: "customer", Label: "客户", Kind: FieldKindDimension, ValueType: "string", Source: "log.username", Visible: true, Exportable: true, Order: 30}, + {Key: "channel_class", Label: "渠道类型", Kind: FieldKindDimension, ValueType: "string", Source: "classification.output", Visible: true, Exportable: true, Order: 40}, + {Key: "channel_id", Label: "渠道id", Kind: FieldKindDimension, ValueType: "integer", Source: "log.channel_id", Visible: true, Exportable: true, Order: 50}, + {Key: "start_time", Label: "开始使用时间", Kind: FieldKindMetric, ValueType: "date", Source: "log.created_at", Aggregate: "min", Visible: true, Exportable: true, Order: 60}, + {Key: "end_time", Label: "结束使用时间", Kind: FieldKindMetric, ValueType: "date", Source: "log.created_at", Aggregate: "max", Visible: true, Exportable: true, Order: 70}, + {Key: "payment", Label: "打款", Kind: FieldKindManual, ValueType: "currency", Visible: true, Exportable: true, Order: 80}, + {Key: "balance_status", Label: "余额状态", Kind: FieldKindFormula, ValueType: "currency", FormulaMode: FormulaModeRunning, InitialExpression: "payment - receivable", Expression: "previous_balance_status + payment - receivable", Visible: true, Exportable: true, Order: 90}, + {Key: "actual_consumption", Label: "实际消耗数", Kind: FieldKindMetric, ValueType: "decimal", Source: "log.quota_per_unit", Aggregate: "sum", Visible: true, Exportable: true, Order: 100}, + {Key: "unit_price", Label: "（单价）", Kind: FieldKindManual, ValueType: "currency", Visible: true, Exportable: true, Order: 110}, + {Key: "discount", Label: "折扣", Kind: FieldKindFormula, ValueType: "decimal", Expression: "unit_price / 6.8", Visible: true, Exportable: true, Order: 120}, + {Key: "cost", Label: "成本", Kind: FieldKindFormula, ValueType: "currency", Expression: "discount * actual_consumption", Visible: true, Exportable: true, Order: 130}, + {Key: "supply_discount", Label: "供货折扣", Kind: FieldKindManual, ValueType: "decimal", Visible: true, Exportable: true, Order: 140}, + {Key: "receivable", Label: "应收账款", Kind: FieldKindFormula, ValueType: "currency", Expression: "actual_consumption * supply_discount", Visible: true, Exportable: true, Order: 150}, + {Key: "unallocated_profit", Label: "利润（未分配中间方）", Kind: FieldKindFormula, ValueType: "currency", Expression: "receivable - cost", Visible: true, Exportable: true, Order: 160}, + {Key: "middle_profit_ratio", Label: "中间利润比例", Kind: FieldKindFormula, ValueType: "percent", Expression: "(0.73 - 0.67) * 0.55", Visible: true, Exportable: true, Order: 170, ManualOverride: true}, + {Key: "middle_profit", Label: "中间利润（居间）", Kind: FieldKindFormula, ValueType: "currency", Expression: "actual_consumption * middle_profit_ratio", Visible: true, Exportable: true, Order: 180}, + {Key: "xx_profit_ratio", Label: "xx利润比例", Kind: FieldKindFormula, ValueType: "percent", Expression: "(0.73 - 0.67) * 0.45", Visible: true, Exportable: true, Order: 190, ManualOverride: true}, + {Key: "xx_profit", Label: "xx利润", Kind: FieldKindFormula, ValueType: "currency", Expression: "actual_consumption * xx_profit_ratio", Visible: true, Exportable: true, Order: 200}, + }, + ClassificationRules: []ClassificationRuleConfig{ + { + Key: "aws_claude_by_type_or_name", + Label: "AWS Claude渠道", + Priority: 10, + Enabled: true, + Match: "any", + OutputClass: "AWS", + ConditionGroups: []ClassificationConditionGroup{ + {Match: "all", Conditions: []ClassificationCondition{{Source: "channel.type", Operator: "equals", Value: fmt.Sprintf("%d", constant.ChannelTypeAws)}}}, + {Match: "all", Conditions: []ClassificationCondition{{Source: "is_claude_related", Operator: "equals", Value: "true"}, {Source: "channel.name", Operator: "contains", Value: "aws", CaseInsensitive: true}}}, + }, + }, + { + Key: "claude_key", + Label: "Claude Key渠道", + Priority: 20, + Enabled: true, + Match: "all", + OutputClass: "Claude Key", + Conditions: []ClassificationCondition{{Source: "is_claude_related", Operator: "equals", Value: "true"}}, + }, + { + Key: "other", + Label: "其他渠道", + Priority: 1000, + Enabled: true, + OutputClass: "Other", + Fallback: true, + }, + }, + ExportLayout: ExportLayoutConfig{ + SheetName: "成本报表（总）", + FreezeHeader: true, + IncludeMeta: true, + DateFormat: "yyyy-mm-dd", + DecimalFormat: "0.00", + }, + } +} + +func ValidateTemplateConfig(config CostReportTemplateConfig) error { + if strings.TrimSpace(config.Timezone) == "" { + return fmt.Errorf("timezone is required") + } + if _, err := time.LoadLocation(config.Timezone); err != nil { + return fmt.Errorf("invalid timezone %q: %w", config.Timezone, err) + } + if config.PeriodMode != PeriodModeDay && config.PeriodMode != PeriodModeCustom { + return fmt.Errorf("invalid period_mode %q", config.PeriodMode) + } + if len(config.Fields) == 0 { + return fmt.Errorf("fields are required") + } + + fieldsByKey := make(map[string]FieldConfig, len(config.Fields)) + for i, field := range config.Fields { + if err := validateFieldConfig(field); err != nil { + return fmt.Errorf("fields[%d] %q: %w", i, field.Key, err) + } + if _, exists := fieldsByKey[field.Key]; exists { + return fmt.Errorf("duplicate field key %q", field.Key) + } + fieldsByKey[field.Key] = field + } + if err := validateGrouping(config.Grouping, fieldsByKey); err != nil { + return err + } + if err := validateSort(config.Sort, fieldsByKey); err != nil { + return err + } + if err := validateFormulas(config.Fields, fieldsByKey); err != nil { + return err + } + if err := validateClassificationRules(config.ClassificationRules); err != nil { + return err + } + if strings.TrimSpace(config.ExportLayout.SheetName) == "" { + return fmt.Errorf("export_layout.sheet_name is required") + } + return nil +} + +func validateFieldConfig(field FieldConfig) error { + if !identifierRE.MatchString(field.Key) { + return fmt.Errorf("invalid key") + } + if strings.TrimSpace(field.Label) == "" { + return fmt.Errorf("label is required") + } + if !validFieldKinds[field.Kind] { + return fmt.Errorf("invalid kind %q", field.Kind) + } + if !validValueTypes[field.ValueType] { + return fmt.Errorf("invalid value_type %q", field.ValueType) + } + if field.Kind == FieldKindFormula { + mode := field.FormulaMode + if mode == "" { + mode = FormulaModeStandard + } + if mode != FormulaModeStandard && mode != FormulaModeRunning { + return fmt.Errorf("invalid formula_mode %q", field.FormulaMode) + } + if strings.TrimSpace(field.Expression) == "" { + return fmt.Errorf("formula expression is required") + } + if mode == FormulaModeRunning && strings.TrimSpace(field.InitialExpression) == "" { + return fmt.Errorf("running formula initial_expression is required") + } + return nil + } + if field.Expression != "" || field.InitialExpression != "" { + return fmt.Errorf("only formula fields may define expressions") + } + if field.Kind == FieldKindMetric { + if !validMetricSources[field.Source] && !validLogOtherSource(field.Source) { + return fmt.Errorf("invalid metric source %q", field.Source) + } + if !validAggregates[field.Aggregate] { + return fmt.Errorf("invalid aggregate %q", field.Aggregate) + } + } + if field.Kind == FieldKindDimension { + if !validDimensionSources[field.Source] && !validLogOtherSource(field.Source) { + return fmt.Errorf("invalid dimension source %q", field.Source) + } + } + return nil +} + +func validateGrouping(grouping []string, fields map[string]FieldConfig) error { + if len(grouping) == 0 { + return fmt.Errorf("grouping is required") + } + seen := map[string]bool{} + for _, key := range grouping { + field, ok := fields[key] + if !ok { + return fmt.Errorf("grouping references unknown field %q", key) + } + if seen[key] { + return fmt.Errorf("grouping contains duplicate field %q", key) + } + seen[key] = true + if field.Kind != FieldKindDimension { + return fmt.Errorf("grouping field %q must be a dimension", key) + } + } + return nil +} + +func validateSort(sortRules []SortConfig, fields map[string]FieldConfig) error { + for i, rule := range sortRules { + if _, ok := fields[rule.Field]; !ok { + return fmt.Errorf("sort[%d] references unknown field %q", i, rule.Field) + } + dir := strings.ToLower(strings.TrimSpace(rule.Direction)) + if dir != "asc" && dir != "desc" { + return fmt.Errorf("sort[%d] has invalid direction %q", i, rule.Direction) + } + } + return nil +} + +func validateFormulas(fields []FieldConfig, fieldsByKey map[string]FieldConfig) error { + env := make(map[string]interface{}, len(fieldsByKey)+8) + for key := range fieldsByKey { + env[key] = float64(0) + env["previous_"+key] = float64(0) + } + env["max"] = math.Max + env["min"] = math.Min + env["abs"] = math.Abs + env["ceil"] = math.Ceil + env["floor"] = math.Floor + + deps := make(map[string][]string) + for _, field := range fields { + if field.Kind != FieldKindFormula { + continue + } + exprs := []string{field.Expression} + if field.InitialExpression != "" { + exprs = append(exprs, field.InitialExpression) + } + for _, expression := range exprs { + prog, err := expr.Compile(expression, expr.Env(env), expr.AsFloat64()) + if err != nil { + return fmt.Errorf("formula %q compile failed: %w", field.Key, err) + } + for ref, usage := range formulaRefs(prog.Node(), fieldsByKey) { + if ref == field.Key && usage.Direct { + return fmt.Errorf("formula %q references itself", field.Key) + } + if fieldsByKey[ref].Kind == FieldKindFormula && ref != field.Key { + deps[field.Key] = append(deps[field.Key], ref) + } + } + } + } + return validateFormulaAcyclic(deps) +} + +type formulaRefUsage struct { + Direct bool + Previous bool +} + +func formulaRefs(node ast.Node, fields map[string]FieldConfig) map[string]formulaRefUsage { + refs := map[string]formulaRefUsage{} + ast.Find(node, func(n ast.Node) bool { + id, ok := n.(*ast.IdentifierNode) + if !ok { + return false + } + name := id.Value + previous := false + if strings.HasPrefix(name, "previous_") { + name = strings.TrimPrefix(name, "previous_") + previous = true + } + if _, ok := fields[name]; ok { + usage := refs[name] + if previous { + usage.Previous = true + } else { + usage.Direct = true + } + refs[name] = usage + } + return false + }) + return refs +} + +func validateFormulaAcyclic(deps map[string][]string) error { + visiting := map[string]bool{} + visited := map[string]bool{} + var visit func(string) error + visit = func(key string) error { + if visiting[key] { + return fmt.Errorf("formula dependency cycle detected at %q", key) + } + if visited[key] { + return nil + } + visiting[key] = true + for _, dep := range deps[key] { + if err := visit(dep); err != nil { + return err + } + } + visiting[key] = false + visited[key] = true + return nil + } + keys := make([]string, 0, len(deps)) + for key := range deps { + keys = append(keys, key) + } + sort.Strings(keys) + for _, key := range keys { + if err := visit(key); err != nil { + return err + } + } + return nil +} + +func validateClassificationRules(rules []ClassificationRuleConfig) error { + if len(rules) == 0 { + return fmt.Errorf("classification_rules are required") + } + seen := map[string]bool{} + fallbackCount := 0 + for i, rule := range rules { + if !identifierRE.MatchString(rule.Key) { + return fmt.Errorf("classification_rules[%d] has invalid key %q", i, rule.Key) + } + if seen[rule.Key] { + return fmt.Errorf("duplicate classification rule key %q", rule.Key) + } + seen[rule.Key] = true + if strings.TrimSpace(rule.OutputClass) == "" { + return fmt.Errorf("classification rule %q output_class is required", rule.Key) + } + if rule.Fallback { + fallbackCount++ + continue + } + if !validMatch(rule.Match) { + return fmt.Errorf("classification rule %q has invalid match %q", rule.Key, rule.Match) + } + if len(rule.Conditions) == 0 && len(rule.ConditionGroups) == 0 { + return fmt.Errorf("classification rule %q requires conditions or condition_groups", rule.Key) + } + for j, condition := range rule.Conditions { + if err := validateClassificationCondition(condition); err != nil { + return fmt.Errorf("classification rule %q condition[%d]: %w", rule.Key, j, err) + } + } + for j, group := range rule.ConditionGroups { + if !validMatch(group.Match) { + return fmt.Errorf("classification rule %q group[%d] has invalid match %q", rule.Key, j, group.Match) + } + if len(group.Conditions) == 0 { + return fmt.Errorf("classification rule %q group[%d] requires conditions", rule.Key, j) + } + for k, condition := range group.Conditions { + if err := validateClassificationCondition(condition); err != nil { + return fmt.Errorf("classification rule %q group[%d] condition[%d]: %w", rule.Key, j, k, err) + } + } + } + } + if fallbackCount != 1 { + return fmt.Errorf("exactly one fallback classification rule is required") + } + return nil +} + +func validateClassificationCondition(condition ClassificationCondition) error { + if !validClassificationSources[condition.Source] && !validLogOtherSource(condition.Source) { + return fmt.Errorf("invalid source %q", condition.Source) + } + if !validClassificationOperators[condition.Operator] { + return fmt.Errorf("invalid operator %q", condition.Operator) + } + if condition.Operator == "in" { + if len(condition.Values) == 0 { + return fmt.Errorf("operator in requires values") + } + } else if condition.Operator != "exists" && condition.Value == "" { + return fmt.Errorf("operator %s requires value", condition.Operator) + } + if condition.Operator == "regex" { + if len(condition.Value) > 256 { + return fmt.Errorf("regex value is too long") + } + if _, err := regexp.Compile(condition.Value); err != nil { + return fmt.Errorf("invalid regex: %w", err) + } + } + return nil +} + +func validMatch(match string) bool { + return match == "" || match == "all" || match == "any" +} + +func validLogOtherSource(source string) bool { + return logOtherSourceRE.MatchString(source) +} + +func ConfigJSONAndHash(config CostReportTemplateConfig) (string, string, error) { + if err := ValidateTemplateConfig(config); err != nil { + return "", "", err + } + payload, err := common.Marshal(config) + if err != nil { + return "", "", err + } + hash := sha256.Sum256(payload) + return string(payload), fmt.Sprintf("%x", hash), nil +} + +func EnsureDefaultTemplates(db *gorm.DB, actorID int) error { + _, err := EnsureDefaultClaudeCostTemplate(db, actorID) + return err +} + +func EnsureDefaultClaudeCostTemplate(db *gorm.DB, actorID int) (*model.CostReportTemplate, error) { + if db == nil { + return nil, fmt.Errorf("db is nil") + } + configJSON, configHash, err := ConfigJSONAndHash(DefaultClaudeCostTemplateConfig()) + if err != nil { + return nil, err + } + + var template model.CostReportTemplate + err = db.Where("key = ?", DefaultTemplateKey).First(&template).Error + if err != nil && err != gorm.ErrRecordNotFound { + return nil, err + } + reusableVersionID := 0 + if err == nil { + if template.CurrentVersionId != nil { + var currentVersion model.CostReportTemplateVersion + versionErr := db.First(¤tVersion, *template.CurrentVersionId).Error + if versionErr == nil && currentVersion.ConfigHash == configHash { + return &template, nil + } + if versionErr != nil && versionErr != gorm.ErrRecordNotFound { + return nil, versionErr + } + } + var latestVersion model.CostReportTemplateVersion + latestErr := db.Where("template_id = ?", template.Id).Order("status asc, version desc, id desc").First(&latestVersion).Error + if latestErr == nil && latestVersion.ConfigHash == configHash { + reusableVersionID = latestVersion.Id + } + if latestErr != nil && latestErr != gorm.ErrRecordNotFound { + return nil, latestErr + } + } + + err = db.Transaction(func(tx *gorm.DB) error { + if template.Id == 0 { + template = model.CostReportTemplate{ + Key: DefaultTemplateKey, + Name: "Claude成本默认模板", + Description: "对齐样例Excel的Claude渠道成本统计模板", + Status: model.CostReportTemplateStatusEnabled, + CreatedBy: actorID, + UpdatedBy: actorID, + } + if err := tx.Create(&template).Error; err != nil { + return err + } + } + + if reusableVersionID > 0 { + if err := tx.Model(&model.CostReportTemplateVersion{}).Where("template_id = ? AND id <> ?", template.Id, reusableVersionID).Update("status", model.CostReportTemplateVersionStatusArchived).Error; err != nil { + return err + } + if err := tx.Model(&model.CostReportTemplateVersion{}).Where("id = ?", reusableVersionID).Update("status", model.CostReportTemplateVersionStatusActive).Error; err != nil { + return err + } + template.CurrentVersionId = &reusableVersionID + template.Name = "Claude成本默认模板" + template.Description = "对齐样例Excel的Claude渠道成本统计模板" + template.Status = model.CostReportTemplateStatusEnabled + template.UpdatedBy = actorID + return tx.Save(&template).Error + } + + var maxVersion int + if err := tx.Model(&model.CostReportTemplateVersion{}).Where("template_id = ?", template.Id).Select("COALESCE(MAX(version), 0)").Scan(&maxVersion).Error; err != nil { + return err + } + version := model.CostReportTemplateVersion{ + TemplateId: template.Id, + Version: maxVersion + 1, + Status: model.CostReportTemplateVersionStatusActive, + ConfigJson: configJSON, + ConfigHash: configHash, + CreatedBy: actorID, + } + if err := tx.Create(&version).Error; err != nil { + return err + } + if err := tx.Model(&model.CostReportTemplateVersion{}).Where("template_id = ? AND id <> ?", template.Id, version.Id).Update("status", model.CostReportTemplateVersionStatusArchived).Error; err != nil { + return err + } + template.CurrentVersionId = &version.Id + template.Name = "Claude成本默认模板" + template.Description = "对齐样例Excel的Claude渠道成本统计模板" + template.Status = model.CostReportTemplateStatusEnabled + template.UpdatedBy = actorID + return tx.Save(&template).Error + }) + if err != nil { + return nil, err + } + return &template, nil +} + +var validFieldKinds = map[string]bool{ + FieldKindDimension: true, + FieldKindMetric: true, + FieldKindManual: true, + FieldKindFormula: true, +} + +var validValueTypes = map[string]bool{ + "string": true, + "integer": true, + "decimal": true, + "currency": true, + "percent": true, + "date": true, +} + +var validAggregates = map[string]bool{ + "sum": true, + "count": true, + "avg": true, + "min": true, + "max": true, +} + +var validDimensionSources = map[string]bool{ + "generated.row_index": true, + "period.date": true, + "log.username": true, + "log.user_id": true, + "log.channel_id": true, + "log.model_name": true, + "log.group": true, + "classification.output": true, + "channel.name": true, + "channel.type": true, + "user.display_name": true, +} + +var validMetricSources = map[string]bool{ + "log.created_at": true, + "log.quota": true, + "log.quota_per_unit": true, + "log.prompt_tokens": true, + "log.completion_tokens": true, + "log.total_tokens": true, + "log.request_count": true, +} + +var validClassificationSources = map[string]bool{ + "channel.type": true, + "channel.name": true, + "channel.id": true, + "model_name": true, + "group": true, + "is_claude_related": true, +} + +var validClassificationOperators = map[string]bool{ + "equals": true, + "contains": true, + "regex": true, + "in": true, + "exists": true, +} diff --git a/service/cost_report/config_test.go b/service/cost_report/config_test.go new file mode 100644 index 000000000000..354bf334f1fe --- /dev/null +++ b/service/cost_report/config_test.go @@ -0,0 +1,149 @@ +package cost_report + +import ( + "strings" + "testing" + + "github.com/QuantumNous/new-api/model" + "github.com/glebarez/sqlite" + "gorm.io/gorm" +) + +func TestDefaultClaudeCostTemplateConfigValid(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + if err := ValidateTemplateConfig(cfg); err != nil { + t.Fatalf("default template should validate: %v", err) + } + if len(cfg.Fields) != 20 { + t.Fatalf("default template should contain 20 fields, got %d", len(cfg.Fields)) + } + if cfg.ClassificationRules[0].Key != "aws_claude_by_type_or_name" { + t.Fatalf("unexpected first classification rule: %s", cfg.ClassificationRules[0].Key) + } +} + +func TestValidateTemplateConfigRejectsDuplicateFieldKeys(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + cfg.Fields = append(cfg.Fields, cfg.Fields[0]) + assertValidationErrorContains(t, cfg, "duplicate field key") +} + +func TestValidateTemplateConfigRejectsInvalidFieldKind(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + cfg.Fields[0].Kind = "spreadsheet" + assertValidationErrorContains(t, cfg, "invalid kind") +} + +func TestValidateTemplateConfigRejectsInvalidFormula(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + for i := range cfg.Fields { + if cfg.Fields[i].Key == "cost" { + cfg.Fields[i].Expression = "missing_field + 1" + break + } + } + assertValidationErrorContains(t, cfg, "compile failed") +} + +func TestValidateTemplateConfigRejectsDirectSelfReferenceEvenWithPreviousReference(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + for i := range cfg.Fields { + if cfg.Fields[i].Key == "balance_status" { + cfg.Fields[i].Expression = "previous_balance_status - balance_status" + break + } + } + assertValidationErrorContains(t, cfg, "references itself") +} + +func TestValidateTemplateConfigRejectsUnsafeLogOtherSource(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + cfg.Fields[0].Source = "log_other." + assertValidationErrorContains(t, cfg, "invalid dimension source") +} + +func TestValidateTemplateConfigAcceptsDottedLogOtherSource(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + cfg.Fields[0].Source = "log_other.admin_info.request_path" + if err := ValidateTemplateConfig(cfg); err != nil { + t.Fatalf("expected dotted log_other source to validate: %v", err) + } +} + +func TestValidateTemplateConfigRejectsFormulaCycle(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + for i := range cfg.Fields { + switch cfg.Fields[i].Key { + case "cost": + cfg.Fields[i].Expression = "receivable + 1" + case "receivable": + cfg.Fields[i].Expression = "cost + 1" + } + } + assertValidationErrorContains(t, cfg, "cycle") +} + +func TestValidateTemplateConfigRejectsInvalidClassificationRule(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + cfg.ClassificationRules[0].ConditionGroups[0].Conditions[0].Operator = "starts_with" + assertValidationErrorContains(t, cfg, "invalid operator") +} + +func TestConfigJSONAndHashUsesCommonJSONWrapperSemantics(t *testing.T) { + jsonText, hash, err := ConfigJSONAndHash(DefaultClaudeCostTemplateConfig()) + if err != nil { + t.Fatalf("ConfigJSONAndHash failed: %v", err) + } + if jsonText == "" || hash == "" { + t.Fatalf("expected json and hash") + } + if !strings.Contains(jsonText, DefaultTemplateKey[:6]) && !strings.Contains(jsonText, "成本报表") { + t.Fatalf("json output does not look like default cost report config: %s", jsonText) + } +} + +func TestEnsureDefaultClaudeCostTemplateSeedsVersion(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate(&model.CostReportTemplate{}, &model.CostReportTemplateVersion{}); err != nil { + t.Fatalf("migrate: %v", err) + } + + template, err := EnsureDefaultClaudeCostTemplate(db, 100) + if err != nil { + t.Fatalf("seed default template: %v", err) + } + if template.Key != DefaultTemplateKey || template.CurrentVersionId == nil { + t.Fatalf("unexpected template after seed: %+v", template) + } + var versions int64 + if err := db.Model(&model.CostReportTemplateVersion{}).Where("template_id = ?", template.Id).Count(&versions).Error; err != nil { + t.Fatalf("count versions: %v", err) + } + if versions != 1 { + t.Fatalf("expected 1 version, got %d", versions) + } + + if _, err := EnsureDefaultClaudeCostTemplate(db, 100); err != nil { + t.Fatalf("second seed should be idempotent: %v", err) + } + if err := db.Model(&model.CostReportTemplateVersion{}).Where("template_id = ?", template.Id).Count(&versions).Error; err != nil { + t.Fatalf("count versions after second seed: %v", err) + } + if versions != 1 { + t.Fatalf("expected idempotent seed to keep 1 version, got %d", versions) + } +} + +func assertValidationErrorContains(t *testing.T, cfg CostReportTemplateConfig, want string) { + t.Helper() + err := ValidateTemplateConfig(cfg) + if err == nil { + t.Fatalf("expected validation error containing %q", want) + } + if !strings.Contains(err.Error(), want) { + t.Fatalf("expected error containing %q, got %v", want, err) + } +} diff --git a/service/cost_report/export_excel.go b/service/cost_report/export_excel.go new file mode 100644 index 000000000000..c76e5e6b712a --- /dev/null +++ b/service/cost_report/export_excel.go @@ -0,0 +1,205 @@ +package cost_report + +import ( + "bytes" + "context" + "fmt" + "regexp" + "sort" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/xuri/excelize/v2" +) + +var invalidSheetNameRE = regexp.MustCompile(`[\\/\?\*\[\]:]`) + +func (s *Service) ExportRunXLSX(ctx context.Context, runID int) ([]byte, string, error) { + detail, err := s.GetRunDetail(ctx, runID) + if err != nil { + return nil, "", err + } + file := excelize.NewFile() + dataSheet := sanitizeSheetName(detail.Config.ExportLayout.SheetName) + if dataSheet == "" { + dataSheet = "Cost Report" + } + defaultSheet := file.GetSheetName(0) + if defaultSheet == "" { + defaultSheet = "Sheet1" + } + if err := file.SetSheetName(defaultSheet, dataSheet); err != nil { + return nil, "", err + } + + fields := exportableFields(detail.Config) + for col, field := range fields { + cell, _ := excelize.CoordinatesToCellName(col+1, 1) + if err := file.SetCellValue(dataSheet, cell, field.Label); err != nil { + return nil, "", err + } + } + for rowIndex, row := range detail.Rows { + for col, field := range fields { + cell, _ := excelize.CoordinatesToCellName(col+1, rowIndex+2) + if err := file.SetCellValue(dataSheet, cell, exportCellValue(field, row.Values, detail.Run.Timezone)); err != nil { + return nil, "", err + } + } + } + if detail.Config.ExportLayout.FreezeHeader { + _ = file.SetPanes(dataSheet, &excelize.Panes{ + Freeze: true, + YSplit: 1, + TopLeftCell: "A2", + ActivePane: "bottomLeft", + Selection: []excelize.Selection{{ + Pane: "bottomLeft", + ActiveCell: "A2", + SQRef: "A2", + }}, + }) + } + if len(fields) > 0 { + lastCol, _ := excelize.ColumnNumberToName(len(fields)) + _ = file.SetColWidth(dataSheet, "A", lastCol, 16) + } + if detail.Config.ExportLayout.IncludeMeta { + if err := writeMetaSheet(file, detail); err != nil { + return nil, "", err + } + } + var buf bytes.Buffer + if err := file.Write(&buf); err != nil { + return nil, "", err + } + filename := fmt.Sprintf("cost-report-%s-run-%d.xlsx", safeFilenamePart(detail.Run.PeriodKey), detail.Run.Id) + return buf.Bytes(), filename, nil +} + +func exportableFields(config CostReportTemplateConfig) []FieldConfig { + fields := make([]FieldConfig, 0, len(config.Fields)) + for _, field := range config.Fields { + if field.Exportable { + fields = append(fields, field) + } + } + sort.SliceStable(fields, func(i, j int) bool { + if fields[i].Order == fields[j].Order { + return fields[i].Key < fields[j].Key + } + return fields[i].Order < fields[j].Order + }) + return fields +} + +func exportCellValue(field FieldConfig, values map[string]interface{}, timezone string) interface{} { + value := values[field.Key] + if field.ValueType != "date" { + return value + } + switch v := value.(type) { + case int: + return formatUnixTimestamp(int64(v), timezone) + case int64: + return formatUnixTimestamp(v, timezone) + case float64: + if v > 1000000000 { + return formatUnixTimestamp(int64(v), timezone) + } + return v + default: + return value + } +} + +func formatUnixTimestamp(ts int64, timezone string) string { + loc, err := time.LoadLocation(timezone) + if err != nil { + loc = time.UTC + } + return time.Unix(ts, 0).In(loc).Format("2006-01-02 15:04:05") +} + +func writeMetaSheet(file *excelize.File, detail *RunDetail) error { + metaSheet := uniqueSheetName(file, "Meta") + if _, err := file.NewSheet(metaSheet); err != nil { + return err + } + rulesJSON := "[]" + if payload, err := common.Marshal(detail.Config.ClassificationRules); err == nil { + rulesJSON = string(payload) + } + rows := [][]interface{}{ + {"template_id", detail.Run.TemplateId}, + {"template_version_id", detail.Run.TemplateVersionId}, + {"run_id", detail.Run.Id}, + {"period_key", detail.Run.PeriodKey}, + {"period_start", detail.Run.PeriodStart}, + {"period_end", detail.Run.PeriodEnd}, + {"timezone", detail.Run.Timezone}, + {"source_log_max_id", detail.Run.SourceLogMaxId}, + {"source_hash", detail.Run.SourceHash}, + {"row_count", detail.Run.RowCount}, + {"created_by", detail.Run.CreatedBy}, + {"created_at", detail.Run.CreatedAt}, + {"classification_rules_json", rulesJSON}, + } + for i, row := range rows { + for j, value := range row { + cell, _ := excelize.CoordinatesToCellName(j+1, i+1) + if err := file.SetCellValue(metaSheet, cell, value); err != nil { + return err + } + } + } + _ = file.SetColWidth(metaSheet, "A", "A", 28) + _ = file.SetColWidth(metaSheet, "B", "B", 80) + return nil +} + +func uniqueSheetName(file *excelize.File, base string) string { + used := map[string]bool{} + for _, name := range file.GetSheetList() { + used[name] = true + } + base = sanitizeSheetName(base) + if base == "" { + base = "Sheet" + } + if !used[base] { + return base + } + for i := 2; ; i++ { + suffix := fmt.Sprintf(" %d", i) + candidateBase := base + if len([]rune(candidateBase))+len([]rune(suffix)) > 31 { + runes := []rune(candidateBase) + candidateBase = string(runes[:31-len([]rune(suffix))]) + } + candidate := candidateBase + suffix + if !used[candidate] { + return candidate + } + } +} + +func sanitizeSheetName(name string) string { + name = strings.TrimSpace(invalidSheetNameRE.ReplaceAllString(name, "_")) + if len([]rune(name)) <= 31 { + return name + } + runes := []rune(name) + return string(runes[:31]) +} + +func safeFilenamePart(value string) string { + value = strings.TrimSpace(value) + if value == "" { + return "period" + } + value = invalidSheetNameRE.ReplaceAllString(value, "-") + value = strings.ReplaceAll(value, " ", "-") + return value +} diff --git a/service/cost_report/formula_adapter.go b/service/cost_report/formula_adapter.go new file mode 100644 index 000000000000..d8d2df6577a2 --- /dev/null +++ b/service/cost_report/formula_adapter.go @@ -0,0 +1,209 @@ +package cost_report + +import ( + "fmt" + "math" + "sort" + "strings" + + "github.com/expr-lang/expr" + "github.com/expr-lang/expr/vm" +) + +type formulaProgram struct { + field FieldConfig + program *vm.Program + initial *vm.Program +} + +type formulaEvaluator struct { + programs []formulaProgram + fieldsByKey map[string]FieldConfig + partitionFields []string +} + +func newFormulaEvaluator(config CostReportTemplateConfig) (*formulaEvaluator, error) { + fieldsByKey := make(map[string]FieldConfig, len(config.Fields)) + for _, field := range config.Fields { + fieldsByKey[field.Key] = field + } + env := formulaCompileEnv(fieldsByKey) + ordered, err := formulaEvaluationOrder(config.Fields, fieldsByKey) + if err != nil { + return nil, err + } + + programs := make([]formulaProgram, 0, len(ordered)) + for _, field := range ordered { + prog, err := expr.Compile(field.Expression, expr.Env(env), expr.AsFloat64()) + if err != nil { + return nil, fmt.Errorf("formula %q compile failed: %w", field.Key, err) + } + fp := formulaProgram{field: field, program: prog} + if field.InitialExpression != "" { + initial, err := expr.Compile(field.InitialExpression, expr.Env(env), expr.AsFloat64()) + if err != nil { + return nil, fmt.Errorf("formula %q initial compile failed: %w", field.Key, err) + } + fp.initial = initial + } + programs = append(programs, fp) + } + return &formulaEvaluator{programs: programs, fieldsByKey: fieldsByKey, partitionFields: defaultRunningPartitionFields(config.Grouping)}, nil +} + +func formulaCompileEnv(fields map[string]FieldConfig) map[string]interface{} { + env := make(map[string]interface{}, len(fields)*2+5) + for key := range fields { + env[key] = float64(0) + env["previous_"+key] = float64(0) + } + env["max"] = math.Max + env["min"] = math.Min + env["abs"] = math.Abs + env["ceil"] = math.Ceil + env["floor"] = math.Floor + return env +} + +func formulaEvaluationOrder(fields []FieldConfig, fieldsByKey map[string]FieldConfig) ([]FieldConfig, error) { + byKey := map[string]FieldConfig{} + deps := map[string][]string{} + for _, field := range fields { + if field.Kind != FieldKindFormula { + continue + } + byKey[field.Key] = field + exprs := []string{field.Expression} + if field.InitialExpression != "" { + exprs = append(exprs, field.InitialExpression) + } + seen := map[string]bool{} + for _, expression := range exprs { + prog, err := expr.Compile(expression, expr.Env(formulaCompileEnv(fieldsByKey)), expr.AsFloat64()) + if err != nil { + return nil, fmt.Errorf("formula %q compile failed: %w", field.Key, err) + } + for ref, usage := range formulaRefs(prog.Node(), fieldsByKey) { + if usage.Direct && fieldsByKey[ref].Kind == FieldKindFormula && ref != field.Key && !seen[ref] { + deps[field.Key] = append(deps[field.Key], ref) + seen[ref] = true + } + } + } + } + + visiting := map[string]bool{} + visited := map[string]bool{} + ordered := make([]FieldConfig, 0, len(byKey)) + var visit func(string) error + visit = func(key string) error { + if visiting[key] { + return fmt.Errorf("formula dependency cycle detected at %q", key) + } + if visited[key] { + return nil + } + visiting[key] = true + depKeys := append([]string(nil), deps[key]...) + sort.Strings(depKeys) + for _, dep := range depKeys { + if err := visit(dep); err != nil { + return err + } + } + visiting[key] = false + visited[key] = true + ordered = append(ordered, byKey[key]) + return nil + } + + fieldOrder := append([]FieldConfig(nil), fields...) + sort.SliceStable(fieldOrder, func(i, j int) bool { return fieldOrder[i].Order < fieldOrder[j].Order }) + for _, field := range fieldOrder { + if field.Kind != FieldKindFormula { + continue + } + if err := visit(field.Key); err != nil { + return nil, err + } + } + return ordered, nil +} + +func (e *formulaEvaluator) evaluateRows(rows []PreviewRow) []string { + if e == nil { + return nil + } + warnings := []string{} + previousByPartition := map[string]map[string]float64{} + for rowIndex := range rows { + partition := runningPartitionKey(rows[rowIndex], e.partitionFields) + previousByField := previousByPartition[partition] + if previousByField == nil { + previousByField = map[string]float64{} + previousByPartition[partition] = previousByField + } + for _, fp := range e.programs { + _, hasPreviousForField := previousByField[fp.field.Key] + if rows[rowIndex].ManualOverrides[fp.field.Key] { + previousByField[fp.field.Key] = toFloat64(rows[rowIndex].Values[fp.field.Key]) + continue + } + + env := formulaCompileEnv(e.fieldsByKey) + for key, value := range rows[rowIndex].Values { + env[key] = toFloat64(value) + } + for key, value := range previousByField { + env["previous_"+key] = value + } + + prog := fp.program + if fp.field.FormulaMode == FormulaModeRunning && !hasPreviousForField && fp.initial != nil { + prog = fp.initial + } + out, err := expr.Run(prog, env) + if err != nil { + warnings = append(warnings, fmt.Sprintf("row %q formula %q failed: %v", rows[rowIndex].RowKey, fp.field.Key, err)) + rows[rowIndex].FormulaValues[fp.field.Key] = float64(0) + rows[rowIndex].Values[fp.field.Key] = float64(0) + previousByField[fp.field.Key] = 0 + continue + } + value, ok := out.(float64) + if !ok || math.IsNaN(value) || math.IsInf(value, 0) { + warnings = append(warnings, fmt.Sprintf("row %q formula %q produced non-finite value", rows[rowIndex].RowKey, fp.field.Key)) + value = 0 + } + rows[rowIndex].FormulaValues[fp.field.Key] = value + rows[rowIndex].Values[fp.field.Key] = value + previousByField[fp.field.Key] = value + } + } + return warnings +} + +func defaultRunningPartitionFields(grouping []string) []string { + fields := make([]string, 0, len(grouping)) + for _, key := range grouping { + switch key { + case "row_index", "report_date": + continue + default: + fields = append(fields, key) + } + } + return fields +} + +func runningPartitionKey(row PreviewRow, fields []string) string { + if len(fields) == 0 { + return "" + } + parts := make([]string, 0, len(fields)) + for _, field := range fields { + parts = append(parts, field+"="+valueToString(row.Values[field])) + } + return strings.Join(parts, "|") +} diff --git a/service/cost_report/manual_cells.go b/service/cost_report/manual_cells.go new file mode 100644 index 000000000000..fb050ad7a060 --- /dev/null +++ b/service/cost_report/manual_cells.go @@ -0,0 +1,134 @@ +package cost_report + +import ( + "context" + "fmt" + "strconv" + "strings" + + "github.com/QuantumNous/new-api/model" + "gorm.io/gorm" + "gorm.io/gorm/clause" +) + +type ManualCellInput struct { + TemplateID int `json:"template_id"` + PeriodKey string `json:"period_key"` + RowKey string `json:"row_key"` + FieldKey string `json:"field_key"` + ValueType string `json:"value_type"` + ValueText string `json:"value_text"` + UpdatedBy int `json:"updated_by"` +} + +type ManualValue struct { + ValueType string `json:"value_type"` + ValueText string `json:"value_text"` + Value interface{} `json:"value"` + UpdatedBy int `json:"updated_by"` + UpdatedAt int64 `json:"updated_at"` +} + +func (s *Service) UpsertManualCell(ctx context.Context, input ManualCellInput) (*model.CostReportManualCell, error) { + if s == nil || s.db == nil { + return nil, fmt.Errorf("db is nil") + } + if input.TemplateID <= 0 { + return nil, fmt.Errorf("template_id is required") + } + if strings.TrimSpace(input.PeriodKey) == "" || strings.TrimSpace(input.RowKey) == "" || strings.TrimSpace(input.FieldKey) == "" { + return nil, fmt.Errorf("period_key, row_key and field_key are required") + } + if !validValueTypes[input.ValueType] { + return nil, fmt.Errorf("invalid value_type %q", input.ValueType) + } + cell := model.CostReportManualCell{ + TemplateId: input.TemplateID, + PeriodKey: input.PeriodKey, + RowKey: input.RowKey, + FieldKey: input.FieldKey, + ValueType: input.ValueType, + ValueText: input.ValueText, + UpdatedBy: input.UpdatedBy, + } + err := s.db.WithContext(ctx).Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "template_id"}, {Name: "period_key"}, {Name: "row_key"}, {Name: "field_key"}}, + DoUpdates: clause.AssignmentColumns([]string{ + "value_type", + "value_text", + "updated_by", + "updated_at", + }), + }).Create(&cell).Error + if err != nil { + return nil, err + } + if err := s.db.WithContext(ctx).Where("template_id = ? AND period_key = ? AND row_key = ? AND field_key = ?", input.TemplateID, input.PeriodKey, input.RowKey, input.FieldKey).First(&cell).Error; err != nil { + return nil, err + } + return &cell, nil +} + +func (s *Service) ReadManualCells(ctx context.Context, templateID int, periodKey string, rowKeys []string) (map[string]map[string]ManualValue, error) { + if s == nil || s.db == nil { + return nil, fmt.Errorf("db is nil") + } + if templateID <= 0 || strings.TrimSpace(periodKey) == "" { + return nil, fmt.Errorf("template_id and period_key are required") + } + query := s.db.WithContext(ctx).Where("template_id = ? AND period_key = ?", templateID, periodKey) + if len(rowKeys) > 0 { + query = query.Where("row_key IN ?", rowKeys) + } + var cells []model.CostReportManualCell + if err := query.Find(&cells).Error; err != nil { + if err == gorm.ErrRecordNotFound { + return map[string]map[string]ManualValue{}, nil + } + return nil, err + } + result := make(map[string]map[string]ManualValue, len(cells)) + for _, cell := range cells { + if result[cell.RowKey] == nil { + result[cell.RowKey] = map[string]ManualValue{} + } + result[cell.RowKey][cell.FieldKey] = ManualValue{ + ValueType: cell.ValueType, + ValueText: cell.ValueText, + Value: parseManualValue(cell.ValueType, cell.ValueText), + UpdatedBy: cell.UpdatedBy, + UpdatedAt: cell.UpdatedAt, + } + } + return result, nil +} + +func parseManualValue(valueType, text string) interface{} { + text = strings.TrimSpace(text) + if text == "" { + switch valueType { + case "integer": + return int64(0) + case "decimal", "currency", "percent": + return float64(0) + default: + return "" + } + } + switch valueType { + case "integer": + value, err := strconv.ParseInt(text, 10, 64) + if err != nil { + return int64(0) + } + return value + case "decimal", "currency", "percent": + value, err := strconv.ParseFloat(text, 64) + if err != nil { + return float64(0) + } + return value + default: + return text + } +} diff --git a/service/cost_report/runs.go b/service/cost_report/runs.go new file mode 100644 index 000000000000..2de507e0b69e --- /dev/null +++ b/service/cost_report/runs.go @@ -0,0 +1,275 @@ +package cost_report + +import ( + "context" + "crypto/sha256" + "fmt" + "sort" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "gorm.io/gorm" +) + +type SaveRunResult struct { + Run model.CostReportRun `json:"run"` + Rows []model.CostReportRowSnapshot `json:"rows,omitempty"` +} + +type RunDetail struct { + Run model.CostReportRun `json:"run"` + Config CostReportTemplateConfig `json:"config"` + Rows []RunSnapshotRow `json:"rows"` +} + +type RunSnapshotRow struct { + Id int `json:"id"` + RunId int `json:"run_id"` + RowKey string `json:"row_key"` + Dimensions map[string]interface{} `json:"dimensions"` + Metrics map[string]interface{} `json:"metrics"` + ManualValues map[string]interface{} `json:"manual_values"` + FormulaValues map[string]interface{} `json:"formula_values"` + Values map[string]interface{} `json:"values"` + CreatedAt int64 `json:"created_at"` +} + +func (s *Service) SaveRunFromPreview(ctx context.Context, preview *PreviewResponse, actorID int) (*SaveRunResult, error) { + if s == nil || s.db == nil { + return nil, fmt.Errorf("db is nil") + } + if preview == nil { + return nil, fmt.Errorf("preview is required") + } + if preview.TemplateID <= 0 || preview.TemplateVersionID <= 0 { + return nil, fmt.Errorf("template_id and template_version_id are required to save a run") + } + version, config, err := s.loadTemplateVersionConfig(ctx, preview.TemplateVersionID) + if err != nil { + return nil, err + } + if version.TemplateId != preview.TemplateID { + return nil, fmt.Errorf("template/version mismatch") + } + configJSON, _, err := ConfigJSONAndHash(*config) + if err != nil { + return nil, err + } + sourceHash, err := previewSourceHash(preview) + if err != nil { + return nil, err + } + + result := &SaveRunResult{} + err = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + run := model.CostReportRun{ + TemplateId: preview.TemplateID, + TemplateVersionId: preview.TemplateVersionID, + PeriodStart: preview.PeriodStart, + PeriodEnd: preview.PeriodEnd, + PeriodKey: preview.PeriodKey, + Timezone: preview.Timezone, + Status: model.CostReportRunStatusCompleted, + ConfigSnapshotJson: configJSON, + SourceLogMaxId: preview.SourceLogMaxID, + SourceHash: sourceHash, + RowCount: len(preview.Rows), + CreatedBy: actorID, + } + if err := tx.Create(&run).Error; err != nil { + return err + } + rows := make([]model.CostReportRowSnapshot, 0, len(preview.Rows)) + for _, row := range preview.Rows { + dimensionsJSON, err := marshalMap(row.Dimensions) + if err != nil { + return err + } + metricsJSON, err := marshalMap(row.Metrics) + if err != nil { + return err + } + manualJSON, err := marshalMap(row.ManualValues) + if err != nil { + return err + } + formulaJSON, err := marshalMap(row.FormulaValues) + if err != nil { + return err + } + rows = append(rows, model.CostReportRowSnapshot{ + RunId: run.Id, + RowKey: row.RowKey, + DimensionsJson: dimensionsJSON, + MetricsJson: metricsJSON, + ManualValuesJson: manualJSON, + FormulaValuesJson: formulaJSON, + }) + } + if len(rows) > 0 { + if err := tx.Create(&rows).Error; err != nil { + return err + } + } + result.Run = run + result.Rows = rows + return nil + }) + if err != nil { + return nil, err + } + return result, nil +} + +func (s *Service) ListRuns(ctx context.Context, templateID int, periodKey string, offset, limit int) ([]model.CostReportRun, int64, error) { + if s == nil || s.db == nil { + return nil, 0, fmt.Errorf("db is nil") + } + if limit <= 0 || limit > 100 { + limit = 20 + } + if offset < 0 { + offset = 0 + } + query := s.db.WithContext(ctx).Model(&model.CostReportRun{}) + if templateID > 0 { + query = query.Where("template_id = ?", templateID) + } + if periodKey != "" { + query = query.Where("period_key = ?", periodKey) + } + var total int64 + if err := query.Count(&total).Error; err != nil { + return nil, 0, err + } + var runs []model.CostReportRun + if err := query.Order("id desc").Offset(offset).Limit(limit).Find(&runs).Error; err != nil { + return nil, 0, err + } + return runs, total, nil +} + +func (s *Service) GetRunDetail(ctx context.Context, runID int) (*RunDetail, error) { + if s == nil || s.db == nil { + return nil, fmt.Errorf("db is nil") + } + if runID <= 0 { + return nil, fmt.Errorf("run id is required") + } + var run model.CostReportRun + if err := s.db.WithContext(ctx).First(&run, runID).Error; err != nil { + return nil, err + } + var config CostReportTemplateConfig + if err := common.UnmarshalJsonStr(run.ConfigSnapshotJson, &config); err != nil { + return nil, err + } + var snapshots []model.CostReportRowSnapshot + if err := s.db.WithContext(ctx).Where("run_id = ?", runID).Order("id asc").Find(&snapshots).Error; err != nil { + return nil, err + } + rows := make([]RunSnapshotRow, 0, len(snapshots)) + for _, snapshot := range snapshots { + row, err := decodeSnapshotRow(snapshot) + if err != nil { + return nil, err + } + rows = append(rows, row) + } + return &RunDetail{Run: run, Config: config, Rows: rows}, nil +} + +func decodeSnapshotRow(snapshot model.CostReportRowSnapshot) (RunSnapshotRow, error) { + row := RunSnapshotRow{ + Id: snapshot.Id, + RunId: snapshot.RunId, + RowKey: snapshot.RowKey, + Dimensions: map[string]interface{}{}, + Metrics: map[string]interface{}{}, + ManualValues: map[string]interface{}{}, + FormulaValues: map[string]interface{}{}, + Values: map[string]interface{}{}, + CreatedAt: snapshot.CreatedAt, + } + if err := unmarshalMap(snapshot.DimensionsJson, &row.Dimensions); err != nil { + return row, err + } + if err := unmarshalMap(snapshot.MetricsJson, &row.Metrics); err != nil { + return row, err + } + if err := unmarshalMap(snapshot.ManualValuesJson, &row.ManualValues); err != nil { + return row, err + } + if err := unmarshalMap(snapshot.FormulaValuesJson, &row.FormulaValues); err != nil { + return row, err + } + for key, value := range row.Dimensions { + row.Values[key] = value + } + for key, value := range row.Metrics { + row.Values[key] = value + } + for key, value := range row.ManualValues { + row.Values[key] = value + } + for key, value := range row.FormulaValues { + row.Values[key] = value + } + return row, nil +} + +func marshalMap(value map[string]interface{}) (string, error) { + if value == nil { + value = map[string]interface{}{} + } + payload, err := common.Marshal(value) + if err != nil { + return "", err + } + return string(payload), nil +} + +func unmarshalMap(text string, out *map[string]interface{}) error { + if text == "" { + *out = map[string]interface{}{} + return nil + } + return common.UnmarshalJsonStr(text, out) +} + +func previewSourceHash(preview *PreviewResponse) (string, error) { + type hashRow struct { + RowKey string `json:"row_key"` + Dimensions map[string]interface{} `json:"dimensions"` + Metrics map[string]interface{} `json:"metrics"` + ManualValues map[string]interface{} `json:"manual_values"` + FormulaValues map[string]interface{} `json:"formula_values"` + Values map[string]interface{} `json:"values"` + } + rows := make([]hashRow, 0, len(preview.Rows)) + for _, row := range preview.Rows { + rows = append(rows, hashRow{ + RowKey: row.RowKey, + Dimensions: row.Dimensions, + Metrics: row.Metrics, + ManualValues: row.ManualValues, + FormulaValues: row.FormulaValues, + Values: row.Values, + }) + } + sort.SliceStable(rows, func(i, j int) bool { return rows[i].RowKey < rows[j].RowKey }) + payload, err := common.Marshal(map[string]interface{}{ + "template_id": preview.TemplateID, + "template_version_id": preview.TemplateVersionID, + "period_start": preview.PeriodStart, + "period_end": preview.PeriodEnd, + "period_key": preview.PeriodKey, + "source_log_max_id": preview.SourceLogMaxID, + "rows": rows, + }) + if err != nil { + return "", err + } + hash := sha256.Sum256(payload) + return fmt.Sprintf("%x", hash), nil +} diff --git a/service/cost_report/service_test.go b/service/cost_report/service_test.go new file mode 100644 index 000000000000..ccf88786cfd2 --- /dev/null +++ b/service/cost_report/service_test.go @@ -0,0 +1,408 @@ +package cost_report + +import ( + "bytes" + "context" + "strings" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/glebarez/sqlite" + "github.com/xuri/excelize/v2" + "gorm.io/gorm" +) + +func TestClassifyDefaultRulesAWSClaudeOther(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + aws := Classify(cfg, ClassificationInput{ + Log: &model.Log{ModelName: "claude-3-5-sonnet", ChannelId: 1}, + Channel: &model.Channel{Id: 1, Type: constant.ChannelTypeAws, Name: "bedrock"}, + }) + if aws.Class != "AWS" { + t.Fatalf("expected AWS, got %+v", aws) + } + + claudeKey := Classify(cfg, ClassificationInput{ + Log: &model.Log{ModelName: "claude-3-haiku", ChannelId: 2}, + Channel: &model.Channel{Id: 2, Type: constant.ChannelTypeOpenAI, Name: "anthropic key"}, + }) + if claudeKey.Class != "Claude Key" { + t.Fatalf("expected Claude Key, got %+v", claudeKey) + } + + anthropicByType := Classify(cfg, ClassificationInput{ + Log: &model.Log{ModelName: "provider-default", ChannelId: 4}, + Channel: &model.Channel{Id: 4, Type: constant.ChannelTypeAnthropic, Name: "direct"}, + }) + if anthropicByType.Class != "Claude Key" { + t.Fatalf("expected Anthropic channel type to be Claude Key, got %+v", anthropicByType) + } + + other := Classify(cfg, ClassificationInput{ + Log: &model.Log{ModelName: "gpt-4o-mini", ChannelId: 3}, + Channel: &model.Channel{Id: 3, Type: constant.ChannelTypeOpenAI, Name: "openai"}, + }) + if other.Class != "Other" { + t.Fatalf("expected Other, got %+v", other) + } +} + +func TestFormulaEvaluatorDefaultRunningBalanceAndWarnings(t *testing.T) { + cfg := DefaultClaudeCostTemplateConfig() + evaluator, err := newFormulaEvaluator(cfg) + if err != nil { + t.Fatalf("new evaluator: %v", err) + } + rows := []PreviewRow{ + { + RowKey: "r1", + Values: map[string]interface{}{"customer": "alice", "channel_class": "AWS", "channel_id": 1, "payment": float64(100), "actual_consumption": float64(10), "unit_price": float64(6.8), "supply_discount": float64(1)}, + FormulaValues: map[string]interface{}{}, + ManualOverrides: map[string]bool{}, + }, + { + RowKey: "r2", + Values: map[string]interface{}{"customer": "alice", "channel_class": "AWS", "channel_id": 1, "payment": float64(5), "actual_consumption": float64(5), "unit_price": float64(6.8), "supply_discount": float64(2)}, + FormulaValues: map[string]interface{}{}, + ManualOverrides: map[string]bool{}, + }, + } + warnings := evaluator.evaluateRows(rows) + if len(warnings) != 0 { + t.Fatalf("unexpected warnings: %v", warnings) + } + if got := toFloat64(rows[0].Values["receivable"]); got != 10 { + t.Fatalf("row1 receivable = %v, want 10", got) + } + if got := toFloat64(rows[0].Values["balance_status"]); got != 90 { + t.Fatalf("row1 balance = %v, want 90", got) + } + if got := toFloat64(rows[1].Values["balance_status"]); got != 85 { + t.Fatalf("row2 running balance = %v, want 85", got) + } + + partitionRows := []PreviewRow{ + {RowKey: "p1", Values: map[string]interface{}{"customer": "alice", "channel_class": "AWS", "channel_id": 1, "payment": float64(100), "actual_consumption": float64(10), "unit_price": float64(6.8), "supply_discount": float64(1)}, FormulaValues: map[string]interface{}{}, ManualOverrides: map[string]bool{}}, + {RowKey: "p2", Values: map[string]interface{}{"customer": "bob", "channel_class": "AWS", "channel_id": 1, "payment": float64(50), "actual_consumption": float64(5), "unit_price": float64(6.8), "supply_discount": float64(2)}, FormulaValues: map[string]interface{}{}, ManualOverrides: map[string]bool{}}, + {RowKey: "p3", Values: map[string]interface{}{"customer": "alice", "channel_class": "AWS", "channel_id": 1, "payment": float64(2), "actual_consumption": float64(1), "unit_price": float64(6.8), "supply_discount": float64(10)}, FormulaValues: map[string]interface{}{}, ManualOverrides: map[string]bool{}}, + } + warnings = evaluator.evaluateRows(partitionRows) + if len(warnings) != 0 { + t.Fatalf("unexpected partition warnings: %v", warnings) + } + if got := toFloat64(partitionRows[1].Values["balance_status"]); got != 40 { + t.Fatalf("partitioned bob first balance = %v, want 40", got) + } + if got := toFloat64(partitionRows[2].Values["balance_status"]); got != 82 { + t.Fatalf("non-contiguous alice balance = %v, want 82", got) + } + + bad := cfg + bad.Fields = []FieldConfig{ + {Key: "actual_consumption", Label: "actual", Kind: FieldKindMetric, ValueType: "decimal", Source: "log.quota", Aggregate: "sum", Visible: true, Exportable: true, Order: 1}, + {Key: "unit_price", Label: "unit", Kind: FieldKindManual, ValueType: "decimal", Visible: true, Exportable: true, Order: 2}, + {Key: "bad_formula", Label: "bad", Kind: FieldKindFormula, ValueType: "decimal", Expression: "actual_consumption / unit_price", Visible: true, Exportable: true, Order: 3}, + } + bad.Grouping = []string{"actual_consumption"} + bad.Sort = nil + bad.ClassificationRules = cfg.ClassificationRules + badEval, err := newFormulaEvaluator(bad) + if err != nil { + t.Fatalf("new bad evaluator: %v", err) + } + badRows := []PreviewRow{{RowKey: "bad", Values: map[string]interface{}{"actual_consumption": float64(1), "unit_price": float64(0)}, FormulaValues: map[string]interface{}{}, ManualOverrides: map[string]bool{}}} + warnings = badEval.evaluateRows(badRows) + if len(warnings) == 0 || !strings.Contains(warnings[0], "non-finite") { + t.Fatalf("expected non-finite warning, got %v", warnings) + } +} + +func TestManualCellUpsertRead(t *testing.T) { + db := openCostReportTestDB(t) + svc := NewService(db, db) + ctx := context.Background() + + if _, err := svc.UpsertManualCell(ctx, ManualCellInput{TemplateID: 1, PeriodKey: "2026-06-03", RowKey: "row", FieldKey: "payment", ValueType: "currency", ValueText: "12.5", UpdatedBy: 7}); err != nil { + t.Fatalf("upsert manual: %v", err) + } + if _, err := svc.UpsertManualCell(ctx, ManualCellInput{TemplateID: 1, PeriodKey: "2026-06-03", RowKey: "row", FieldKey: "payment", ValueType: "currency", ValueText: "13.5", UpdatedBy: 8}); err != nil { + t.Fatalf("second upsert manual: %v", err) + } + manuals, err := svc.ReadManualCells(ctx, 1, "2026-06-03", []string{"row"}) + if err != nil { + t.Fatalf("read manual: %v", err) + } + if got := toFloat64(manuals["row"]["payment"].Value); got != 13.5 { + t.Fatalf("manual payment = %v, want 13.5", got) + } +} + +func TestPreviewUsesSeparateLogAndMainDBs(t *testing.T) { + oldQuotaPerUnit := common.QuotaPerUnit + common.QuotaPerUnit = 100 + t.Cleanup(func() { common.QuotaPerUnit = oldQuotaPerUnit }) + + mainDB := openCostReportTestDB(t) + logDB := openLogTestDB(t) + svc := NewService(mainDB, logDB) + ctx := context.Background() + + if err := mainDB.Create(&model.User{Id: 1, Username: "alice", DisplayName: "Alice", Password: "x"}).Error; err != nil { + t.Fatalf("create user: %v", err) + } + if err := mainDB.Create(&model.Channel{Id: 10, Type: constant.ChannelTypeAws, Name: "AWS Bedrock Claude", Key: "k", Models: "claude-3-5-sonnet"}).Error; err != nil { + t.Fatalf("create channel: %v", err) + } + start := time.Date(2026, 6, 3, 0, 0, 0, 0, time.UTC).Unix() + logs := []model.Log{ + {Id: 1, UserId: 1, Username: "alice", CreatedAt: start + 60, Type: model.LogTypeConsume, ModelName: "claude-3-5-sonnet", Quota: 250, ChannelId: 10, Group: "default"}, + {Id: 2, UserId: 1, Username: "alice", CreatedAt: start + 120, Type: model.LogTypeConsume, ModelName: "claude-3-5-sonnet", Quota: 150, ChannelId: 10, Group: "default"}, + } + if err := logDB.Create(&logs).Error; err != nil { + t.Fatalf("create logs: %v", err) + } + + cfg := DefaultClaudeCostTemplateConfig() + rowKey := makeRowKey(cfg.Grouping, map[string]interface{}{ + "report_date": "2026-06-03", + "customer": "alice", + "channel_class": "AWS", + "channel_id": 10, + }) + if _, err := svc.UpsertManualCell(ctx, ManualCellInput{TemplateID: 1, PeriodKey: "2026-06-03", RowKey: rowKey, FieldKey: "payment", ValueType: "currency", ValueText: "9", UpdatedBy: 1}); err != nil { + t.Fatalf("upsert manual before preview: %v", err) + } + + resp, err := svc.Preview(ctx, PreviewRequest{TemplateID: 1, Config: &cfg, PeriodStart: start, PeriodEnd: start + 24*3600, PeriodKey: "2026-06-03", IncludeManual: true}) + if err != nil { + t.Fatalf("preview: %v", err) + } + if len(resp.Rows) != 1 { + t.Fatalf("expected 1 row, got %d: %+v", len(resp.Rows), resp.Rows) + } + row := resp.Rows[0] + if row.RowKey != rowKey { + t.Fatalf("row key = %q, want %q", row.RowKey, rowKey) + } + if len(row.RowKey) != 64 { + t.Fatalf("row key length = %d, want 64", len(row.RowKey)) + } + if got := row.Dimensions["channel_class"]; got != "AWS" { + t.Fatalf("channel_class = %v, want AWS", got) + } + if got := toFloat64(row.Metrics["actual_consumption"]); got != 4 { + t.Fatalf("actual_consumption = %v, want 4", got) + } + if got := toFloat64(row.ManualValues["payment"]); got != 9 { + t.Fatalf("manual payment = %v, want 9", got) + } + if resp.SourceLogMaxID != 2 { + t.Fatalf("source log max id = %d, want 2", resp.SourceLogMaxID) + } +} + +func TestSaveRunAndExportXLSX(t *testing.T) { + db := openCostReportTestDB(t) + svc := NewService(db, db) + ctx := context.Background() + + detail, err := svc.EnsureDefaultTemplate(ctx, 1) + if err != nil { + t.Fatalf("ensure default template: %v", err) + } + if detail.CurrentVersion == nil { + t.Fatalf("default template has no version") + } + start := time.Date(2026, 6, 3, 0, 0, 0, 0, time.UTC).Unix() + preview := &PreviewResponse{ + TemplateID: detail.Template.Id, + TemplateVersionID: detail.CurrentVersion.Id, + PeriodStart: start, + PeriodEnd: start + 24*3600, + PeriodKey: "2026-06-03", + Timezone: "Asia/Shanghai", + SourceLogMaxID: 9, + Rows: []PreviewRow{{ + RowKey: makeRowKey(DefaultClaudeCostTemplateConfig().Grouping, map[string]interface{}{ + "report_date": "2026-06-03", + "customer": "alice", + "channel_class": "AWS", + "channel_id": 10, + }), + Dimensions: map[string]interface{}{"row_index": 1, "report_date": "2026-06-03", "customer": "alice", "channel_class": "AWS", "channel_id": 10}, + Metrics: map[string]interface{}{"start_time": float64(start + 60), "end_time": float64(start + 120), "actual_consumption": float64(4)}, + ManualValues: map[string]interface{}{"payment": float64(9), "unit_price": float64(6.8), "supply_discount": float64(1)}, + FormulaValues: map[string]interface{}{"discount": float64(1), "cost": float64(4), "receivable": float64(4)}, + Values: map[string]interface{}{"row_index": 1, "report_date": "2026-06-03", "customer": "alice", "channel_class": "AWS", "channel_id": 10, "start_time": float64(start + 60), "end_time": float64(start + 120), "payment": float64(9), "unit_price": float64(6.8), "actual_consumption": float64(4), "supply_discount": float64(1), "discount": float64(1), "cost": float64(4), "receivable": float64(4)}, + ManualOverrides: map[string]bool{}, + }}, + } + + saved, err := svc.SaveRunFromPreview(ctx, preview, 1) + if err != nil { + t.Fatalf("save run: %v", err) + } + if saved.Run.RowCount != 1 || len(saved.Rows) != 1 { + t.Fatalf("saved row count mismatch: run=%d rows=%d", saved.Run.RowCount, len(saved.Rows)) + } + runDetail, err := svc.GetRunDetail(ctx, saved.Run.Id) + if err != nil { + t.Fatalf("get run detail: %v", err) + } + if got := runDetail.Rows[0].Values["customer"]; got != "alice" { + t.Fatalf("snapshot customer = %v, want alice", got) + } + + data, filename, err := svc.ExportRunXLSX(ctx, saved.Run.Id) + if err != nil { + t.Fatalf("export xlsx: %v", err) + } + if len(data) == 0 || !strings.HasSuffix(filename, ".xlsx") { + t.Fatalf("bad export: len=%d filename=%q", len(data), filename) + } + book, err := excelize.OpenReader(bytes.NewReader(data)) + if err != nil { + t.Fatalf("open exported xlsx: %v", err) + } + defer func() { _ = book.Close() }() + sheet := sanitizeSheetName(DefaultClaudeCostTemplateConfig().ExportLayout.SheetName) + header, _ := book.GetCellValue(sheet, "A1") + customer, _ := book.GetCellValue(sheet, "C2") + metaKey, _ := book.GetCellValue("Meta", "A1") + if header != "序号" || customer != "alice" || metaKey != "template_id" { + t.Fatalf("unexpected xlsx cells: header=%q customer=%q meta=%q", header, customer, metaKey) + } +} + +func TestConsumeLogScanUsesBoundedBatchesAndMaxLogs(t *testing.T) { + logDB := openLogTestDB(t) + svc := NewService(openCostReportTestDB(t), logDB) + ctx := context.Background() + start := time.Date(2026, 6, 3, 0, 0, 0, 0, time.UTC).Unix() + logs := make([]model.Log, 0, consumeLogScanBatchSize+5) + for i := 1; i <= consumeLogScanBatchSize+5; i++ { + logs = append(logs, model.Log{Id: i, CreatedAt: start + int64(i%3), Type: model.LogTypeConsume, Quota: i}) + } + if err := logDB.Create(&logs).Error; err != nil { + t.Fatalf("create logs: %v", err) + } + + seen := 0 + batches := 0 + lastCreatedAt := int64(-1) + lastID := 0 + if err := svc.scanConsumeLogs(ctx, start, start+10, consumeLogScanBatchSize+3, func(batch []model.Log) error { + batches++ + if len(batch) > consumeLogScanBatchSize { + t.Fatalf("batch size = %d, want <= %d", len(batch), consumeLogScanBatchSize) + } + for _, log := range batch { + if lastCreatedAt > log.CreatedAt || (lastCreatedAt == log.CreatedAt && lastID >= log.Id) { + t.Fatalf("logs not ordered after created_at=%d id=%d: got created_at=%d id=%d", lastCreatedAt, lastID, log.CreatedAt, log.Id) + } + lastCreatedAt = log.CreatedAt + lastID = log.Id + seen++ + } + return nil + }); err != nil { + t.Fatalf("scan logs: %v", err) + } + if seen != consumeLogScanBatchSize+3 { + t.Fatalf("seen logs = %d, want %d", seen, consumeLogScanBatchSize+3) + } + if batches < 2 { + t.Fatalf("expected multiple batches, got %d", batches) + } +} + +func TestPreviewSourceHashIncludesRowContent(t *testing.T) { + base := &PreviewResponse{ + TemplateID: 1, + TemplateVersionID: 2, + PeriodStart: 10, + PeriodEnd: 20, + PeriodKey: "p", + SourceLogMaxID: 3, + Rows: []PreviewRow{{ + RowKey: "row", + Dimensions: map[string]interface{}{"customer": "alice"}, + Metrics: map[string]interface{}{"actual_consumption": float64(1)}, + ManualValues: map[string]interface{}{"payment": float64(1)}, + FormulaValues: map[string]interface{}{"receivable": float64(1)}, + Values: map[string]interface{}{"customer": "alice", "actual_consumption": float64(1), "payment": float64(1), "receivable": float64(1)}, + }}, + } + hash1, err := previewSourceHash(base) + if err != nil { + t.Fatalf("hash1: %v", err) + } + base.Rows[0].Values["payment"] = float64(2) + base.Rows[0].ManualValues["payment"] = float64(2) + hash2, err := previewSourceHash(base) + if err != nil { + t.Fatalf("hash2: %v", err) + } + if hash1 == hash2 { + t.Fatalf("source hash did not change when row content changed") + } +} + +func TestEnsureDefaultTemplateRepairsMissingCurrentVersion(t *testing.T) { + db := openCostReportTestDB(t) + ctx := context.Background() + configJSON, configHash, err := ConfigJSONAndHash(DefaultClaudeCostTemplateConfig()) + if err != nil { + t.Fatalf("default config hash: %v", err) + } + template := model.CostReportTemplate{Key: DefaultTemplateKey, Name: "existing", Status: model.CostReportTemplateStatusEnabled} + if err := db.Create(&template).Error; err != nil { + t.Fatalf("create template: %v", err) + } + version := model.CostReportTemplateVersion{TemplateId: template.Id, Version: 7, Status: model.CostReportTemplateVersionStatusActive, ConfigJson: configJSON, ConfigHash: configHash} + if err := db.Create(&version).Error; err != nil { + t.Fatalf("create version: %v", err) + } + + detail, err := NewService(db, db).EnsureDefaultTemplate(ctx, 42) + if err != nil { + t.Fatalf("ensure default template: %v", err) + } + if detail.Template.CurrentVersionId == nil || *detail.Template.CurrentVersionId != version.Id { + t.Fatalf("current version id = %v, want %d", detail.Template.CurrentVersionId, version.Id) + } + var count int64 + if err := db.Model(&model.CostReportTemplateVersion{}).Where("template_id = ?", template.Id).Count(&count).Error; err != nil { + t.Fatalf("count versions: %v", err) + } + if count != 1 { + t.Fatalf("version count = %d, want 1", count) + } +} + +func openCostReportTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate(&model.CostReportManualCell{}, &model.CostReportRun{}, &model.CostReportRowSnapshot{}, &model.Channel{}, &model.User{}, &model.CostReportTemplate{}, &model.CostReportTemplateVersion{}); err != nil { + t.Fatalf("migrate main test db: %v", err) + } + return db +} + +func openLogTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + if err != nil { + t.Fatalf("open sqlite: %v", err) + } + if err := db.AutoMigrate(&model.Log{}); err != nil { + t.Fatalf("migrate log test db: %v", err) + } + return db +} diff --git a/service/cost_report/templates.go b/service/cost_report/templates.go new file mode 100644 index 000000000000..6b682c12026e --- /dev/null +++ b/service/cost_report/templates.go @@ -0,0 +1,214 @@ +package cost_report + +import ( + "context" + "fmt" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "gorm.io/gorm" +) + +type TemplateSaveInput struct { + Id int `json:"id,omitempty"` + Key string `json:"key"` + Name string `json:"name"` + Description string `json:"description"` + Status int `json:"status"` + Config CostReportTemplateConfig `json:"config"` + ActorID int `json:"-"` +} + +type TemplateDetail struct { + Template model.CostReportTemplate `json:"template"` + CurrentVersion *model.CostReportTemplateVersion `json:"current_version,omitempty"` + Config *CostReportTemplateConfig `json:"config,omitempty"` +} + +func (s *Service) EnsureDefaultTemplate(ctx context.Context, actorID int) (*TemplateDetail, error) { + if s == nil || s.db == nil { + return nil, fmt.Errorf("db is nil") + } + template, err := EnsureDefaultClaudeCostTemplate(s.db.WithContext(ctx), actorID) + if err != nil { + return nil, err + } + return s.GetTemplate(ctx, template.Id) +} + +func (s *Service) ListTemplates(ctx context.Context, offset, limit int) ([]TemplateDetail, int64, error) { + if s == nil || s.db == nil { + return nil, 0, fmt.Errorf("db is nil") + } + if limit <= 0 || limit > 100 { + limit = 20 + } + if offset < 0 { + offset = 0 + } + var total int64 + if err := s.db.WithContext(ctx).Model(&model.CostReportTemplate{}).Count(&total).Error; err != nil { + return nil, 0, err + } + var templates []model.CostReportTemplate + if err := s.db.WithContext(ctx).Order("id desc").Offset(offset).Limit(limit).Find(&templates).Error; err != nil { + return nil, 0, err + } + details := make([]TemplateDetail, 0, len(templates)) + for i := range templates { + detail := TemplateDetail{Template: templates[i]} + if templates[i].CurrentVersionId != nil { + version, cfg, err := s.loadTemplateVersionConfig(ctx, *templates[i].CurrentVersionId) + if err != nil { + return nil, 0, err + } + detail.CurrentVersion = version + detail.Config = cfg + } + details = append(details, detail) + } + return details, total, nil +} + +func (s *Service) GetTemplate(ctx context.Context, id int) (*TemplateDetail, error) { + if s == nil || s.db == nil { + return nil, fmt.Errorf("db is nil") + } + if id <= 0 { + return nil, fmt.Errorf("template id is required") + } + var template model.CostReportTemplate + if err := s.db.WithContext(ctx).First(&template, id).Error; err != nil { + return nil, err + } + detail := &TemplateDetail{Template: template} + if template.CurrentVersionId != nil { + version, cfg, err := s.loadTemplateVersionConfig(ctx, *template.CurrentVersionId) + if err != nil { + return nil, err + } + detail.CurrentVersion = version + detail.Config = cfg + } + return detail, nil +} + +func (s *Service) ListTemplateVersions(ctx context.Context, templateID int) ([]model.CostReportTemplateVersion, error) { + if s == nil || s.db == nil { + return nil, fmt.Errorf("db is nil") + } + if templateID <= 0 { + return nil, fmt.Errorf("template_id is required") + } + var versions []model.CostReportTemplateVersion + if err := s.db.WithContext(ctx).Where("template_id = ?", templateID).Order("version desc").Find(&versions).Error; err != nil { + return nil, err + } + return versions, nil +} + +func (s *Service) SaveTemplate(ctx context.Context, input TemplateSaveInput) (*TemplateDetail, error) { + if s == nil || s.db == nil { + return nil, fmt.Errorf("db is nil") + } + input.Key = strings.TrimSpace(input.Key) + input.Name = strings.TrimSpace(input.Name) + if input.Key == "" || input.Name == "" { + return nil, fmt.Errorf("key and name are required") + } + if !identifierRE.MatchString(input.Key) { + return nil, fmt.Errorf("invalid template key") + } + status := input.Status + if status == 0 { + status = model.CostReportTemplateStatusEnabled + } + if status != model.CostReportTemplateStatusEnabled && status != model.CostReportTemplateStatusArchived { + return nil, fmt.Errorf("invalid template status") + } + configJSON, configHash, err := ConfigJSONAndHash(input.Config) + if err != nil { + return nil, err + } + + var savedID int + err = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error { + var template model.CostReportTemplate + if input.Id > 0 { + if err := tx.First(&template, input.Id).Error; err != nil { + return err + } + var dup model.CostReportTemplate + err := tx.Where("key = ? AND id <> ?", input.Key, input.Id).First(&dup).Error + if err == nil { + return fmt.Errorf("template key already exists") + } + if err != gorm.ErrRecordNotFound { + return err + } + template.Key = input.Key + template.Name = input.Name + template.Description = input.Description + template.Status = status + template.UpdatedBy = input.ActorID + if err := tx.Save(&template).Error; err != nil { + return err + } + } else { + template = model.CostReportTemplate{ + Key: input.Key, + Name: input.Name, + Description: input.Description, + Status: status, + CreatedBy: input.ActorID, + UpdatedBy: input.ActorID, + } + if err := tx.Create(&template).Error; err != nil { + return err + } + } + + var maxVersion int + if err := tx.Model(&model.CostReportTemplateVersion{}).Where("template_id = ?", template.Id).Select("COALESCE(MAX(version), 0)").Scan(&maxVersion).Error; err != nil { + return err + } + version := model.CostReportTemplateVersion{ + TemplateId: template.Id, + Version: maxVersion + 1, + Status: model.CostReportTemplateVersionStatusActive, + ConfigJson: configJSON, + ConfigHash: configHash, + CreatedBy: input.ActorID, + } + if err := tx.Create(&version).Error; err != nil { + return err + } + if err := tx.Model(&model.CostReportTemplateVersion{}).Where("template_id = ? AND id <> ?", template.Id, version.Id).Update("status", model.CostReportTemplateVersionStatusArchived).Error; err != nil { + return err + } + template.CurrentVersionId = &version.Id + template.UpdatedBy = input.ActorID + if err := tx.Save(&template).Error; err != nil { + return err + } + savedID = template.Id + return nil + }) + if err != nil { + return nil, err + } + return s.GetTemplate(ctx, savedID) +} + +func (s *Service) loadTemplateVersionConfig(ctx context.Context, versionID int) (*model.CostReportTemplateVersion, *CostReportTemplateConfig, error) { + var version model.CostReportTemplateVersion + if err := s.db.WithContext(ctx).First(&version, versionID).Error; err != nil { + return nil, nil, err + } + var config CostReportTemplateConfig + if err := common.UnmarshalJsonStr(version.ConfigJson, &config); err != nil { + return nil, nil, err + } + return &version, &config, nil +} diff --git a/service/navigation.go b/service/navigation.go new file mode 100644 index 000000000000..43090a94448a --- /dev/null +++ b/service/navigation.go @@ -0,0 +1,342 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +package service + +import ( + "errors" + "fmt" + "strings" + "sync" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "gorm.io/gorm" +) + +// NavigationItemDTO 下发给前端的统一菜单节点格式 +type NavigationItemDTO struct { + ID uint `json:"id"` + Type string `json:"type"` + ModuleKey string `json:"module_key,omitempty"` + Label string `json:"label"` + Path string `json:"path,omitempty"` + URL string `json:"url,omitempty"` + IconKey string `json:"icon_key,omitempty"` + OpenInNewTab bool `json:"open_in_new_tab"` + ExactActive bool `json:"exact_active"` + Children []NavigationItemDTO `json:"children,omitempty"` +} + +type NavigationService struct { + cache map[string][]NavigationItemDTO + cacheMu sync.RWMutex +} + +var NavService = &NavigationService{ + cache: make(map[string][]NavigationItemDTO), +} + +// GetVisibleNavigationTree 获取指定菜单可见的过滤树（线程安全，基于内存缓存） +func (s *NavigationService) GetVisibleNavigationTree(menuKey string, locale string, userRole int, userGroup string, isAuthenticated bool) ([]NavigationItemDTO, error) { + cacheKey := s.buildCacheKey(menuKey, locale, userRole, userGroup, isAuthenticated) + + // 1. 读缓存 + s.cacheMu.RLock() + if cachedData, ok := s.cache[cacheKey]; ok { + s.cacheMu.RUnlock() + return cachedData, nil + } + s.cacheMu.RUnlock() + + // 2. 查数据库并拼装 + var menu model.NavigationMenu + err := model.DB.Where("key = ? AND enabled = ?", menuKey, true).First(&menu).Error + if err != nil { + return nil, fmt.Errorf("menu not found: %w", err) + } + + var items []model.NavigationItem + err = model.DB.Where("menu_id = ? AND enabled = ?", menu.ID, true). + Order("sort_order asc, id asc"). + Preload("Translations"). + Preload("Rules"). + Find(&items).Error + if err != nil { + return nil, fmt.Errorf("failed to fetch menu items: %w", err) + } + + // 3. 过滤并翻译 + visibleItems := make([]model.NavigationItem, 0, len(items)) + for _, item := range items { + if s.checkVisibility(item.Rules, userRole, userGroup, isAuthenticated) { + visibleItems = append(visibleItems, item) + } + } + + // 4. 构建树形结构 + tree := s.buildTree(visibleItems, locale) + + // 5. 写入缓存 + s.cacheMu.Lock() + s.cache[cacheKey] = tree + s.cacheMu.Unlock() + + return tree, nil +} + +// InvalidateCache 清空所有缓存（在管理端 CRUD 修改导航后调用） +func (s *NavigationService) InvalidateCache() { + s.cacheMu.Lock() + s.cache = make(map[string][]NavigationItemDTO) + s.cacheMu.Unlock() + common.SysLog("Navigation memory cache invalidated") +} + +// buildCacheKey 构造唯一的缓存 Key +func (s *NavigationService) buildCacheKey(menuKey, locale string, userRole int, userGroup string, isAuthenticated bool) string { + return fmt.Sprintf("%s:%s:%d:%s:%t", menuKey, locale, userRole, userGroup, isAuthenticated) +} + +// checkVisibility 验证节点权限，实施 RBAC 可见性过滤规则 +func (s *NavigationService) checkVisibility(rules []model.NavigationVisibilityRule, userRole int, userGroup string, isAuthenticated bool) bool { + if len(rules) == 0 { + return true // 无规则限制，默认所有人可见 + } + + hasAllowRules := false + allowMatched := false + + for _, rule := range rules { + matched := s.evaluateRuleSubject(rule.SubjectType, rule.SubjectValue, userRole, userGroup, isAuthenticated) + + if rule.Effect == "deny" { + if matched { + return false // 只要命中任何一条 deny 规则，立即不可见 + } + } else if rule.Effect == "allow" { + hasAllowRules = true + if matched { + allowMatched = true + } + } + } + + // 如果配置了 allow 规则，必须命中至少一条 allow 规则才可见 + if hasAllowRules { + return allowMatched + } + + return true +} + +// evaluateRuleSubject 判断用户是否符合规则主体 +func (s *NavigationService) evaluateRuleSubject(subjectType, subjectValue string, userRole int, userGroup string, isAuthenticated bool) bool { + switch subjectType { + case "everyone": + return true + case "anonymous": + return !isAuthenticated + case "authenticated": + return isAuthenticated + case "role": + if !isAuthenticated { + return false + } + // 角色判断规范： + // "root" (100) -> 仅 root 匹配 + // "admin" (10) -> admin (10) 和 root (100) 匹配 + // "user" (1) -> 所有登录用户匹配 + switch strings.ToLower(subjectValue) { + case "root": + return userRole == common.RoleRootUser + case "admin": + return userRole >= common.RoleAdminUser + case "user": + return userRole >= common.RoleCommonUser + default: + return false + } + case "user_group": + if !isAuthenticated { + return false + } + return userGroup == subjectValue + default: + return false + } +} + +// buildTree 一次性遍历将扁平列表组装为树形结构，并应用翻译 fallback 规则 +func (s *NavigationService) buildTree(items []model.NavigationItem, locale string) []NavigationItemDTO { + // 初始化节点映射表 + dtoMap := make(map[uint]*NavigationItemDTO) + for _, item := range items { + dto := &NavigationItemDTO{ + ID: item.ID, + Type: item.Type, + ModuleKey: item.ModuleKey, + Path: item.Path, + URL: item.URL, + IconKey: item.IconKey, + OpenInNewTab: item.OpenInNewTab, + ExactActive: item.ExactActive, + Label: s.translateLabel(item.Translations, item.ModuleKey, locale), + Children: []NavigationItemDTO{}, + } + dtoMap[item.ID] = dto + } + + var rootDTOs []NavigationItemDTO + + // 二次遍历组装树状父子层级 + for _, item := range items { + dto := dtoMap[item.ID] + if dto == nil { + continue + } + + if item.ParentID == nil { + // 顶级菜单 + rootDTOs = append(rootDTOs, *dto) + } else { + // 子菜单，挂载到父节点下 + parentDTO := dtoMap[*item.ParentID] + if parentDTO != nil { + parentDTO.Children = append(parentDTO.Children, *dto) + } else { + // 父节点已在权限过滤中被裁剪或被禁用，降级作为顶级项（这里按严谨重构规范：无父节点的子项如果无有效父节点，不显示） + // 或者可以选择放入 rootDTOs。在此设计中，如果父节点被权限过滤掉，其子节点在软件工程规范中应该同步不可见 + } + } + } + + // 重新深拷贝或扁平复制以消除多级嵌套中由于引用的子对象在 map 树组装时的错乱 + var result []NavigationItemDTO + for _, rootItem := range rootDTOs { + result = append(result, s.deepCopyDTO(rootItem, dtoMap)) + } + + return result +} + +// deepCopyDTO 保证树的深拷贝以维持嵌套结构的正确格式 +func (s *NavigationService) deepCopyDTO(node NavigationItemDTO, dtoMap map[uint]*NavigationItemDTO) NavigationItemDTO { + actualNode := dtoMap[node.ID] + if actualNode == nil { + return node + } + + var copiedChildren []NavigationItemDTO + for _, child := range actualNode.Children { + copiedChildren = append(copiedChildren, s.deepCopyDTO(child, dtoMap)) + } + + node.Children = copiedChildren + return node +} + +// translateLabel 多语言 Fallback 精准解析 +func (s *NavigationService) translateLabel(translations []model.NavigationItemTranslation, moduleKey string, targetLocale string) string { + if len(translations) == 0 { + return moduleKey // 极端无翻译记录下的兜底，展示内置模块键名 + } + + transMap := make(map[string]string) + for _, t := range translations { + transMap[strings.ToLower(t.Locale)] = t.Label + } + + target := strings.ToLower(targetLocale) + + // 1. 精确匹配（如 zh-cn） + if label, ok := transMap[target]; ok { + return label + } + + // 2. 去除区域后缀的模糊匹配（如 zh-tw -> zh） + if parts := strings.Split(target, "-"); len(parts) > 1 { + if label, ok := transMap[parts[0]]; ok { + return label + } + } + + // 2.5 基础语言前缀模糊匹配（如 target="zh"，则匹配 "zh-cn" 或 "zh-tw"） + for k, v := range transMap { + if strings.HasPrefix(k, target+"-") { + return v + } + } + + // 3. Fallback 到英语 "en" + if label, ok := transMap["en"]; ok { + return label + } + if label, ok := transMap["en-us"]; ok { + return label + } + + // 4. Fallback 到中文 "zh-cn" + if label, ok := transMap["zh-cn"]; ok { + return label + } + if label, ok := transMap["zh"]; ok { + return label + } + + // 5. Fallback 到第一条已有翻译 + return translations[0].Label +} + +// SaveMenuWithTransaction 用于管理端安全保存（包含子节点和翻译等的事务性级联保存） +func (s *NavigationService) SaveMenuWithTransaction(menu *model.NavigationMenu) error { + // 可在此实现需要强事务绑定的复杂业务逻辑 + return model.DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Save(menu).Error; err != nil { + return err + } + s.InvalidateCache() + return nil + }) +} + +// ValidateItemURL 拦截恶意 URL 并防范 XSS 漏洞 +func (s *NavigationService) ValidateItemURL(itemType string, itemURL string) error { + if itemType != "external_url" { + return nil + } + + trimmedURL := strings.TrimSpace(itemURL) + if trimmedURL == "" { + return errors.New("external URL cannot be empty") + } + + lowerURL := strings.ToLower(trimmedURL) + // 拦截包含 javascript: 等具有运行脚本能力的恶意协议 + if strings.HasPrefix(lowerURL, "javascript:") || strings.HasPrefix(lowerURL, "data:") { + return errors.New("malicious URL protocol detected") + } + + // 必须以 http:// 或 https:// 开头 + if !strings.HasPrefix(lowerURL, "http://") && !strings.HasPrefix(lowerURL, "https://") { + return errors.New("external URL must start with http:// or https://") + } + + return nil +} diff --git a/service/navigation_test.go b/service/navigation_test.go new file mode 100644 index 000000000000..76c19f9a600b --- /dev/null +++ b/service/navigation_test.go @@ -0,0 +1,242 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +package service + +import ( + "fmt" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupNavigationTestDB(t *testing.T) *gorm.DB { + t.Helper() + + oldDB := model.DB + oldLogDB := model.LOG_DB + + common.UsingSQLite = true + common.UsingMySQL = false + common.UsingPostgreSQL = false + common.RedisEnabled = false + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + + model.DB = db + model.LOG_DB = db + + require.NoError(t, db.AutoMigrate( + &model.NavigationMenu{}, + &model.NavigationItem{}, + &model.NavigationItemTranslation{}, + &model.NavigationVisibilityRule{}, + )) + + t.Cleanup(func() { + sqlDB, err := db.DB() + if err == nil { + _ = sqlDB.Close() + } + model.DB = oldDB + model.LOG_DB = oldLogDB + }) + + return db +} + +func TestValidateItemURL(t *testing.T) { + tests := []struct { + name string + itemType string + url string + expectErr bool + }{ + {"Valid HTTPS URL", "external_url", "https://google.com/path?query=1", false}, + {"Valid HTTP URL", "external_url", "http://localhost:8080", false}, + {"Empty URL", "external_url", "", true}, + {"Malicious Javascript URL", "external_url", "javascript:alert(1)", true}, + {"Malicious Data URL", "external_url", "data:text/html;base64,PHNjcmlwdD5hbGVydCgxKTwvc2NyaXB0Pg==", true}, + {"No Protocol URL", "external_url", "www.google.com", true}, + {"Non-external type skipped", "builtin_module", "javascript:alert(1)", false}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + err := NavService.ValidateItemURL(tt.itemType, tt.url) + if tt.expectErr { + require.Error(t, err) + } else { + require.NoError(t, err) + } + }) + } +} + +func TestGetVisibleNavigationTree(t *testing.T) { + db := setupNavigationTestDB(t) + + // Create test menu + menu := model.NavigationMenu{ + Key: "test_menu", + Name: "Test Menu", + Client: "web_default", + Surface: "top", + Enabled: true, + IsSystem: false, + } + require.NoError(t, db.Create(&menu).Error) + + // Create test items + // Item 1: Builtin module - Everyone + item1 := model.NavigationItem{ + MenuID: menu.ID, + Type: "builtin_module", + ModuleKey: "home", + SortOrder: 1, + Enabled: true, + } + require.NoError(t, db.Create(&item1).Error) + require.NoError(t, db.Create(&model.NavigationItemTranslation{ + ItemID: item1.ID, + Locale: "zh-CN", + Label: "首页", + }).Error) + require.NoError(t, db.Create(&model.NavigationItemTranslation{ + ItemID: item1.ID, + Locale: "en", + Label: "Home", + }).Error) + + // Item 2: Admin only + item2 := model.NavigationItem{ + MenuID: menu.ID, + Type: "internal_path", + Path: "/admin/users", + SortOrder: 2, + Enabled: true, + } + require.NoError(t, db.Create(&item2).Error) + require.NoError(t, db.Create(&model.NavigationItemTranslation{ + ItemID: item2.ID, + Locale: "zh-CN", + Label: "用户管理", + }).Error) + require.NoError(t, db.Create(&model.NavigationVisibilityRule{ + ItemID: item2.ID, + Effect: "allow", + SubjectType: "role", + SubjectValue: "admin", + }).Error) + + // Item 3: VIP Group only + item3 := model.NavigationItem{ + MenuID: menu.ID, + Type: "external_url", + URL: "https://vip.example.com", + SortOrder: 3, + Enabled: true, + } + require.NoError(t, db.Create(&item3).Error) + require.NoError(t, db.Create(&model.NavigationItemTranslation{ + ItemID: item3.ID, + Locale: "zh-CN", + Label: "VIP专属", + }).Error) + require.NoError(t, db.Create(&model.NavigationVisibilityRule{ + ItemID: item3.ID, + Effect: "allow", + SubjectType: "user_group", + SubjectValue: "VIP", + }).Error) + + // Invalidate service cache to ensure fresh DB query + NavService.InvalidateCache() + + // Case 1: Anonymous user + t.Run("Anonymous visibility", func(t *testing.T) { + NavService.InvalidateCache() + tree, err := NavService.GetVisibleNavigationTree("test_menu", "zh-CN", 0, "", false) + require.NoError(t, err) + require.Len(t, tree, 1) + require.Equal(t, "首页", tree[0].Label) + }) + + // Case 2: Ordinary user (not admin, not VIP) + t.Run("Ordinary user visibility", func(t *testing.T) { + NavService.InvalidateCache() + tree, err := NavService.GetVisibleNavigationTree("test_menu", "zh-CN", common.RoleCommonUser, "default", true) + require.NoError(t, err) + require.Len(t, tree, 1) + require.Equal(t, "首页", tree[0].Label) + }) + + // Case 3: Admin user (role admin) + t.Run("Admin user visibility", func(t *testing.T) { + NavService.InvalidateCache() + tree, err := NavService.GetVisibleNavigationTree("test_menu", "zh-CN", common.RoleAdminUser, "default", true) + require.NoError(t, err) + require.Len(t, tree, 2) + require.Equal(t, "首页", tree[0].Label) + require.Equal(t, "用户管理", tree[1].Label) + }) + + // Case 4: VIP Group user (ordinary role) + t.Run("VIP group visibility", func(t *testing.T) { + NavService.InvalidateCache() + tree, err := NavService.GetVisibleNavigationTree("test_menu", "zh-CN", common.RoleCommonUser, "VIP", true) + require.NoError(t, err) + require.Len(t, tree, 2) + require.Equal(t, "首页", tree[0].Label) + require.Equal(t, "VIP专属", tree[1].Label) + }) +} + +func TestTranslateLabelFallback(t *testing.T) { + translations := []model.NavigationItemTranslation{ + {Locale: "zh-CN", Label: "中文简体"}, + {Locale: "zh-TW", Label: "中文繁體"}, + {Locale: "en", Label: "English"}, + } + + tests := []struct { + locale string + expected string + }{ + {"zh-CN", "中文简体"}, + {"zh-tw", "中文繁體"}, + {"zh-HK", "English"}, + {"en-US", "English"}, + {"fr-FR", "English"}, + } + + for _, tt := range tests { + t.Run(tt.locale, func(t *testing.T) { + label := NavService.translateLabel(translations, "fallback_module_key", tt.locale) + require.Equal(t, tt.expected, label) + }) + } +} diff --git a/service/usage_log_export/export_excel.go b/service/usage_log_export/export_excel.go new file mode 100644 index 000000000000..359da030c5ec --- /dev/null +++ b/service/usage_log_export/export_excel.go @@ -0,0 +1,586 @@ +package usage_log_export + +import ( + "bytes" + "context" + "fmt" + "io" + "strconv" + "strings" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/xuri/excelize/v2" +) + +const ( + exportBatchSize = 2000 + excelMaxRowsPerSheet = 1048576 + excelHeaderRows = 1 + excelMaxDataRows = excelMaxRowsPerSheet - excelHeaderRows + maxOtherJSONRunes = 8000 + usageLogSheetName = "Usage Logs" + createdAtTimeLayout = "2006-01-02 15:04:05" +) + +type FieldGroup struct { + Key string `json:"key"` + Label string `json:"label"` + Fields []FieldOption `json:"fields"` +} + +type FieldOption struct { + Key string `json:"key"` + Label string `json:"label"` + Group string `json:"group"` + Default bool `json:"default"` + AdminOnly bool `json:"admin_only,omitempty"` +} + +type ExportInput struct { + Filter model.LogExportFilter + Fields []string + Timezone string +} + +type fieldDefinition struct { + Key string + Label string + Group string + Default bool + AdminOnly bool + Value func(*model.Log, map[string]interface{}) interface{} +} + +var fieldDefinitions = []fieldDefinition{ + {Key: "created_at", Label: "Time", Group: "basic", Default: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + if log.CreatedAt == 0 { + return "" + } + return time.Unix(log.CreatedAt, 0).Local().Format(createdAtTimeLayout) + }}, + {Key: "type", Label: "Type", Group: "basic", Default: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return logTypeLabel(log.Type) + }}, + {Key: "channel", Label: "Channel", Group: "basic", Default: true, AdminOnly: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + if log.ChannelId == 0 { + return "" + } + if strings.TrimSpace(log.ChannelName) == "" { + return fmt.Sprintf("#%d", log.ChannelId) + } + return fmt.Sprintf("%s #%d", log.ChannelName, log.ChannelId) + }}, + {Key: "user", Label: "User", Group: "basic", Default: true, AdminOnly: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.Username + }}, + {Key: "token_name", Label: "Token", Group: "basic", Default: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.TokenName + }}, + {Key: "model_name", Label: "Model", Group: "basic", Default: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.ModelName + }}, + {Key: "group", Label: "Group", Group: "basic", Default: true, Value: func(log *model.Log, other map[string]interface{}) interface{} { + if log.Group != "" { + return log.Group + } + return stringValue(other["group"]) + }}, + {Key: "use_time", Label: "Timing", Group: "basic", Default: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.UseTime + }}, + {Key: "prompt_tokens", Label: "Input Tokens", Group: "basic", Default: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.PromptTokens + }}, + {Key: "completion_tokens", Label: "Output Tokens", Group: "basic", Default: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.CompletionTokens + }}, + {Key: "quota", Label: "Cost", Group: "basic", Default: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.Quota + }}, + {Key: "details_summary", Label: "Details", Group: "basic", Default: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.Content + }}, + {Key: "cache_read_tokens", Label: "Cache Read Tokens", Group: "cache", Default: true, Value: func(_ *model.Log, other map[string]interface{}) interface{} { + return numberValue(other["cache_tokens"]) + }}, + {Key: "cache_creation_tokens", Label: "Cache Creation Tokens", Group: "cache", Default: true, Value: func(_ *model.Log, other map[string]interface{}) interface{} { + return numberValue(other["cache_creation_tokens"]) + }}, + {Key: "cache_creation_tokens_5m", Label: "5m Cache Creation Tokens", Group: "cache", Default: true, Value: func(_ *model.Log, other map[string]interface{}) interface{} { + return numberValue(other["cache_creation_tokens_5m"]) + }}, + {Key: "cache_creation_tokens_1h", Label: "1h Cache Creation Tokens", Group: "cache", Default: true, Value: func(_ *model.Log, other map[string]interface{}) interface{} { + return numberValue(other["cache_creation_tokens_1h"]) + }}, + {Key: "record_id", Label: "Record ID", Group: "advanced", AdminOnly: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.Id + }}, + {Key: "request_id", Label: "Request ID", Group: "advanced", Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.RequestId + }}, + {Key: "upstream_request_id", Label: "Upstream Request ID", Group: "advanced", Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.UpstreamRequestId + }}, + {Key: "created_at_unix", Label: "Created At (Unix)", Group: "advanced", Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.CreatedAt + }}, + {Key: "ip", Label: "IP", Group: "advanced", AdminOnly: true, Value: func(log *model.Log, _ map[string]interface{}) interface{} { + return log.Ip + }}, + {Key: "other_json", Label: "Other JSON", Group: "advanced", AdminOnly: true, Value: func(log *model.Log, other map[string]interface{}) interface{} { + return sanitizedOtherJSON(other) + }}, +} + +var groupLabels = map[string]string{ + "basic": "Basic Fields", + "cache": "Cache Fields", + "advanced": "Advanced Fields", +} + +func FieldGroups(isAdmin bool) []FieldGroup { + groups := []FieldGroup{ + {Key: "basic", Label: groupLabels["basic"]}, + {Key: "cache", Label: groupLabels["cache"]}, + {Key: "advanced", Label: groupLabels["advanced"]}, + } + indexByKey := map[string]int{"basic": 0, "cache": 1, "advanced": 2} + for _, def := range fieldDefinitions { + if def.AdminOnly && !isAdmin { + continue + } + idx, ok := indexByKey[def.Group] + if !ok { + continue + } + groups[idx].Fields = append(groups[idx].Fields, FieldOption{ + Key: def.Key, + Label: def.Label, + Group: def.Group, + Default: def.Default, + AdminOnly: def.AdminOnly, + }) + } + out := make([]FieldGroup, 0, len(groups)) + for _, group := range groups { + if len(group.Fields) > 0 { + out = append(out, group) + } + } + return out +} + +func ExportXLSX(ctx context.Context, input ExportInput) ([]byte, string, int64, error) { + file, filename, total, err := BuildXLSX(ctx, input) + if err != nil { + return nil, "", total, err + } + defer func() { _ = file.Close() }() + var buf bytes.Buffer + if err := file.Write(&buf); err != nil { + return nil, "", total, err + } + return buf.Bytes(), filename, total, nil +} + +func WriteXLSX(ctx context.Context, input ExportInput, writer io.Writer) (string, int64, error) { + file, filename, total, err := BuildXLSX(ctx, input) + if err != nil { + return "", total, err + } + defer func() { _ = file.Close() }() + if err := file.Write(writer); err != nil { + return "", total, err + } + return filename, total, nil +} + +func BuildXLSX(ctx context.Context, input ExportInput) (*excelize.File, string, int64, error) { + total, err := model.CountLogsForExport(ctx, input.Filter) + if err != nil { + return nil, "", total, err + } + select { + case <-ctx.Done(): + return nil, "", total, ctx.Err() + default: + } + + fields, err := resolveFields(input.Fields, input.Filter.IsAdmin) + if err != nil { + return nil, "", total, err + } + if len(fields) == 0 { + return nil, "", total, fmt.Errorf("no export fields selected") + } + loc := exportLocation(input.Timezone) + + file := excelize.NewFile() + success := false + defer func() { + if !success { + _ = file.Close() + } + }() + defaultSheet := file.GetSheetName(0) + if defaultSheet == "" { + defaultSheet = "Sheet1" + } + + sheetIndex := 1 + sheetName := usageLogSheetNameForIndex(sheetIndex) + if err := file.SetSheetName(defaultSheet, sheetName); err != nil { + return nil, "", total, err + } + writer, err := newLogSheetWriter(file, sheetName, fields) + if err != nil { + return nil, "", total, err + } + + var lastCreatedAt int64 + var lastID int + written := 0 + dataRowsInSheet := 0 + + for { + select { + case <-ctx.Done(): + return nil, "", total, ctx.Err() + default: + } + + logs, err := model.GetLogsForExportBatch(ctx, input.Filter, lastCreatedAt, lastID, exportBatchSize, written) + if err != nil { + return nil, "", total, err + } + if len(logs) == 0 { + break + } + + for _, log := range logs { + if dataRowsInSheet >= excelMaxDataRows { + if err := writer.Flush(); err != nil { + return nil, "", total, err + } + sheetIndex++ + sheetName = usageLogSheetNameForIndex(sheetIndex) + if _, err := file.NewSheet(sheetName); err != nil { + return nil, "", total, err + } + writer, err = newLogSheetWriter(file, sheetName, fields) + if err != nil { + return nil, "", total, err + } + dataRowsInSheet = 0 + } + + row, err := exportRowValues(log, fields, loc) + if err != nil { + return nil, "", total, err + } + cell, _ := excelize.CoordinatesToCellName(1, dataRowsInSheet+excelHeaderRows+1) + if err := writer.SetRow(cell, row); err != nil { + return nil, "", total, err + } + dataRowsInSheet++ + written++ + } + + lastLog := logs[len(logs)-1] + lastCreatedAt = lastLog.CreatedAt + lastID = lastLog.Id + if len(logs) < exportBatchSize { + break + } + } + + if err := writer.Flush(); err != nil { + return nil, "", total, err + } + + filename := fmt.Sprintf("usage-logs-%s.xlsx", time.Now().Format("20060102-150405")) + success = true + return file, filename, total, nil +} + +func usageLogSheetNameForIndex(index int) string { + if index <= 1 { + return usageLogSheetName + } + return fmt.Sprintf("%s %d", usageLogSheetName, index) +} + +func newLogSheetWriter(file *excelize.File, sheetName string, fields []fieldDefinition) (*excelize.StreamWriter, error) { + writer, err := file.NewStreamWriter(sheetName) + if err != nil { + return nil, err + } + if len(fields) > 0 { + if err := writer.SetColWidth(1, len(fields), 16); err != nil { + return nil, err + } + } + if err := writer.SetPanes(&excelize.Panes{ + Freeze: true, + YSplit: 1, + TopLeftCell: "A2", + ActivePane: "bottomLeft", + Selection: []excelize.Selection{{ + Pane: "bottomLeft", + ActiveCell: "A2", + SQRef: "A2", + }}, + }); err != nil { + return nil, err + } + headers := make([]interface{}, 0, len(fields)) + for _, field := range fields { + headers = append(headers, field.Label) + } + if err := writer.SetRow("A1", headers); err != nil { + return nil, err + } + return writer, nil +} + +func exportRowValues(log *model.Log, fields []fieldDefinition, loc *time.Location) ([]interface{}, error) { + other := parseOther(log.Other) + row := make([]interface{}, 0, len(fields)) + for _, field := range fields { + value := field.Value(log, other) + if field.Key == "created_at" && log.CreatedAt != 0 { + value = time.Unix(log.CreatedAt, 0).In(loc).Format(createdAtTimeLayout) + } + row = append(row, value) + } + return row, nil +} + +func resolveFields(keys []string, isAdmin bool) ([]fieldDefinition, error) { + byKey := make(map[string]fieldDefinition, len(fieldDefinitions)) + allKeys := make(map[string]fieldDefinition, len(fieldDefinitions)) + for _, def := range fieldDefinitions { + allKeys[def.Key] = def + if def.AdminOnly && !isAdmin { + continue + } + byKey[def.Key] = def + } + if len(keys) == 0 { + fields := make([]fieldDefinition, 0, len(fieldDefinitions)) + for _, def := range fieldDefinitions { + if def.Default && (!def.AdminOnly || isAdmin) { + fields = append(fields, def) + } + } + return fields, nil + } + fields := make([]fieldDefinition, 0, len(keys)) + seen := make(map[string]bool, len(keys)) + for _, key := range keys { + key = strings.TrimSpace(key) + if key == "" || seen[key] { + continue + } + seen[key] = true + def, ok := byKey[key] + if !ok { + if blocked, exists := allKeys[key]; exists && blocked.AdminOnly && !isAdmin { + return nil, fmt.Errorf("field %s is not available for self export", key) + } + return nil, fmt.Errorf("unknown export field: %s", key) + } + fields = append(fields, def) + } + return fields, nil +} + +func parseOther(raw string) map[string]interface{} { + if strings.TrimSpace(raw) == "" { + return map[string]interface{}{} + } + var other map[string]interface{} + if err := common.UnmarshalJsonStr(raw, &other); err != nil || other == nil { + return map[string]interface{}{} + } + return other +} + +func numberValue(value interface{}) interface{} { + switch v := value.(type) { + case nil: + return 0 + case int: + return v + case int64: + return v + case float64: + if v == float64(int64(v)) { + return int64(v) + } + return v + case float32: + return float64(v) + case string: + if strings.TrimSpace(v) == "" { + return 0 + } + if i, err := strconv.ParseInt(v, 10, 64); err == nil { + return i + } + if f, err := strconv.ParseFloat(v, 64); err == nil { + return f + } + return v + default: + return v + } +} + +func stringValue(value interface{}) string { + if value == nil { + return "" + } + switch v := value.(type) { + case string: + return v + case fmt.Stringer: + return v.String() + default: + return fmt.Sprint(v) + } +} + +func sanitizedOtherJSON(other map[string]interface{}) string { + if len(other) == 0 { + return "" + } + clone := cloneMap(other) + redactOtherMap(clone) + payload, err := common.Marshal(clone) + if err != nil { + return "" + } + text := string(payload) + if len([]rune(text)) <= maxOtherJSONRunes { + return text + } + runes := []rune(text) + return string(runes[:maxOtherJSONRunes]) + "…" +} + +func cloneMap(in map[string]interface{}) map[string]interface{} { + out := make(map[string]interface{}, len(in)) + for k, v := range in { + if child, ok := v.(map[string]interface{}); ok { + out[k] = cloneMap(child) + continue + } + if child, ok := v.(map[interface{}]interface{}); ok { + mapped := make(map[string]interface{}, len(child)) + for ck, cv := range child { + mapped[fmt.Sprint(ck)] = cv + } + out[k] = cloneMap(mapped) + continue + } + if arr, ok := v.([]interface{}); ok { + out[k] = cloneSlice(arr) + continue + } + out[k] = v + } + return out +} + +func cloneSlice(in []interface{}) []interface{} { + out := make([]interface{}, len(in)) + for i, v := range in { + if child, ok := v.(map[string]interface{}); ok { + out[i] = cloneMap(child) + continue + } + if arr, ok := v.([]interface{}); ok { + out[i] = cloneSlice(arr) + continue + } + out[i] = v + } + return out +} + +func redactOtherMap(values map[string]interface{}) { + for key := range values { + if shouldRedactOtherKey(key) { + values[key] = "[redacted]" + } + } + for _, value := range values { + switch child := value.(type) { + case map[string]interface{}: + redactOtherMap(child) + case []interface{}: + for _, item := range child { + if itemMap, ok := item.(map[string]interface{}); ok { + redactOtherMap(itemMap) + } + } + } + } +} + +func shouldRedactOtherKey(key string) bool { + normalized := strings.ToLower(strings.ReplaceAll(strings.ReplaceAll(key, "-", "_"), " ", "_")) + if normalized == "token" || normalized == "key" || normalized == "secret" || normalized == "password" { + return true + } + redactParts := []string{ + "api_key", + "apikey", + "access_token", + "refresh_token", + "authorization", + "auth_header", + "bearer", + "secret", + "password", + "key_key", + "key_path", + } + for _, part := range redactParts { + if strings.Contains(normalized, part) { + return true + } + } + return false +} + +func exportLocation(timezone string) *time.Location { + if strings.TrimSpace(timezone) != "" { + if loc, err := time.LoadLocation(timezone); err == nil { + return loc + } + } + return time.Local +} + +func logTypeLabel(logType int) string { + switch logType { + case model.LogTypeTopup: + return "Top-up" + case model.LogTypeConsume: + return "Consume" + case model.LogTypeManage: + return "Manage" + case model.LogTypeSystem: + return "System" + case model.LogTypeError: + return "Error" + case model.LogTypeRefund: + return "Refund" + default: + return "Unknown" + } +} diff --git a/service/usage_log_export/export_excel_test.go b/service/usage_log_export/export_excel_test.go new file mode 100644 index 000000000000..666f152eb205 --- /dev/null +++ b/service/usage_log_export/export_excel_test.go @@ -0,0 +1,87 @@ +package usage_log_export + +import ( + "strings" + "testing" + + "github.com/QuantumNous/new-api/model" +) + +func TestResolveFieldsFiltersAdminOnly(t *testing.T) { + fields, err := resolveFields([]string{"created_at", "request_id"}, false) + if err != nil { + t.Fatalf("resolveFields returned unexpected error: %v", err) + } + keys := make([]string, 0, len(fields)) + for _, field := range fields { + keys = append(keys, field.Key) + } + joined := strings.Join(keys, ",") + if strings.Contains(joined, "record_id") || strings.Contains(joined, "ip") || strings.Contains(joined, "other_json") { + t.Fatalf("self export should filter admin-only fields, got %q", joined) + } + if joined != "created_at,request_id" { + t.Fatalf("unexpected self fields: %q", joined) + } +} + +func TestResolveFieldsRejectsUnauthorizedFields(t *testing.T) { + _, err := resolveFields([]string{"record_id"}, false) + if err == nil { + t.Fatal("expected self export to reject admin-only field") + } +} + +func TestCacheFieldsDoNotFallback(t *testing.T) { + other := map[string]interface{}{ + "cache_tokens": float64(11), + "cache_write_tokens": float64(99), + "cache_creation_tokens_5m": float64(5), + "cache_creation_tokens_1h": float64(7), + } + fields, err := resolveFields([]string{ + "cache_read_tokens", + "cache_creation_tokens", + "cache_creation_tokens_5m", + "cache_creation_tokens_1h", + }, true) + if err != nil { + t.Fatalf("resolveFields returned unexpected error: %v", err) + } + log := &model.Log{} + values := map[string]interface{}{} + for _, field := range fields { + values[field.Key] = field.Value(log, other) + } + if values["cache_read_tokens"] != int64(11) { + t.Fatalf("cache read should use cache_tokens directly, got %#v", values["cache_read_tokens"]) + } + if values["cache_creation_tokens"] != 0 { + t.Fatalf("cache_creation_tokens should not fallback to cache_write_tokens or split fields, got %#v", values["cache_creation_tokens"]) + } + if values["cache_creation_tokens_5m"] != int64(5) { + t.Fatalf("5m cache creation should use exact key, got %#v", values["cache_creation_tokens_5m"]) + } + if values["cache_creation_tokens_1h"] != int64(7) { + t.Fatalf("1h cache creation should use exact key, got %#v", values["cache_creation_tokens_1h"]) + } +} + +func TestSanitizedOtherJSONRedactsSecretsAndTruncates(t *testing.T) { + longSecret := strings.Repeat("x", maxOtherJSONRunes+100) + text := sanitizedOtherJSON(map[string]interface{}{ + "admin_info": map[string]interface{}{ + "key_key": longSecret, + }, + "headers": map[string]interface{}{ + "Authorization": longSecret, + "apiKey": longSecret, + }, + }) + if strings.Contains(text, longSecret) { + t.Fatal("sanitized other_json leaked secret value") + } + if !strings.Contains(text, "[redacted]") { + t.Fatalf("expected redaction marker, got %q", text) + } +} diff --git a/service/waffo_pancake_test.go b/service/waffo_pancake_test.go index 41c91a15ae23..fe285315727a 100644 --- a/service/waffo_pancake_test.go +++ b/service/waffo_pancake_test.go @@ -17,6 +17,9 @@ import ( func setupWaffoPancakeTestDB(t *testing.T) *gorm.DB { t.Helper() + oldDB := model.DB + oldLogDB := model.LOG_DB + common.UsingSQLite = true common.UsingMySQL = false common.UsingPostgreSQL = false @@ -36,6 +39,8 @@ func setupWaffoPancakeTestDB(t *testing.T) *gorm.DB { if err == nil { _ = sqlDB.Close() } + model.DB = oldDB + model.LOG_DB = oldLogDB }) return db diff --git a/setting/operation_setting/channel_preparation_auto_promotion_setting.go b/setting/operation_setting/channel_preparation_auto_promotion_setting.go new file mode 100644 index 000000000000..d20a61c12c9f --- /dev/null +++ b/setting/operation_setting/channel_preparation_auto_promotion_setting.go @@ -0,0 +1,188 @@ +package operation_setting + +import ( + "fmt" + "strings" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/setting/config" +) + +const ( + ChannelPreparationAutoPromotionGuaranteePriorityCapacityFirst = "capacity_first" + ChannelPreparationAutoPromotionGuaranteePriorityCountFirst = "count_first" +) + +const ( + ChannelPreparationAutoPromotionStrategyPriorityWeighted = "priority_weighted" + ChannelPreparationAutoPromotionStrategySmallBalanceFirst = "small_balance_first" + ChannelPreparationAutoPromotionStrategyLargeBalanceFirst = "large_balance_first" +) + +type ChannelPreparationAutoPromotionRule struct { + Id string `json:"id"` + Enabled bool `json:"enabled"` + Group string `json:"group"` + Type int `json:"type"` + ThresholdUSD float64 `json:"threshold_usd"` + MinimumUsableChannelCount int `json:"minimum_usable_channel_count"` + GuaranteePriority string `json:"guarantee_priority"` + CountShortageStrategy string `json:"count_shortage_strategy"` + CapacityShortageStrategy string `json:"capacity_shortage_strategy"` + // Strategy is kept for compatibility with older settings JSON and rollback. + Strategy string `json:"strategy,omitempty"` +} + +type ChannelPreparationAutoPromotionSetting struct { + SchedulerEnabled bool `json:"scheduler_enabled"` + IntervalMinutes float64 `json:"interval_minutes"` + MaxPromotionsPerRun int `json:"max_promotions_per_run"` + Rules []ChannelPreparationAutoPromotionRule `json:"rules"` +} + +var channelPreparationAutoPromotionSetting = ChannelPreparationAutoPromotionSetting{ + SchedulerEnabled: false, + IntervalMinutes: 10, + MaxPromotionsPerRun: 10, + Rules: []ChannelPreparationAutoPromotionRule{}, +} + +func init() { + config.GlobalConfig.Register("channel_preparation_auto_promotion_setting", &channelPreparationAutoPromotionSetting) +} + +func GetChannelPreparationAutoPromotionSetting() *ChannelPreparationAutoPromotionSetting { + NormalizeChannelPreparationAutoPromotionSetting(&channelPreparationAutoPromotionSetting) + return &channelPreparationAutoPromotionSetting +} + +func NormalizeChannelPreparationAutoPromotionSetting(setting *ChannelPreparationAutoPromotionSetting) { + if setting == nil { + return + } + if setting.IntervalMinutes <= 0 { + setting.IntervalMinutes = 10 + } + if setting.MaxPromotionsPerRun <= 0 { + setting.MaxPromotionsPerRun = 10 + } + for index := range setting.Rules { + NormalizeChannelPreparationAutoPromotionRule(&setting.Rules[index]) + } +} + +func NormalizeChannelPreparationAutoPromotionRule(rule *ChannelPreparationAutoPromotionRule) { + if rule == nil { + return + } + rule.Id = strings.TrimSpace(rule.Id) + rule.Group = strings.TrimSpace(rule.Group) + rule.Strategy = strings.TrimSpace(rule.Strategy) + rule.GuaranteePriority = strings.TrimSpace(rule.GuaranteePriority) + rule.CountShortageStrategy = strings.TrimSpace(rule.CountShortageStrategy) + rule.CapacityShortageStrategy = strings.TrimSpace(rule.CapacityShortageStrategy) + + legacyStrategy := rule.Strategy + if !IsSupportedChannelPreparationAutoPromotionStrategy(legacyStrategy) { + legacyStrategy = ChannelPreparationAutoPromotionStrategyPriorityWeighted + } + if !IsSupportedChannelPreparationAutoPromotionGuaranteePriority(rule.GuaranteePriority) { + rule.GuaranteePriority = ChannelPreparationAutoPromotionGuaranteePriorityCapacityFirst + } + if !IsSupportedChannelPreparationAutoPromotionStrategy(rule.CountShortageStrategy) { + rule.CountShortageStrategy = legacyStrategy + } + if !IsSupportedChannelPreparationAutoPromotionStrategy(rule.CapacityShortageStrategy) { + rule.CapacityShortageStrategy = legacyStrategy + } + // Keep legacy strategy readable by older code. Capacity-first is the V1-compatible path. + rule.Strategy = rule.CapacityShortageStrategy +} + +func IsSupportedChannelPreparationAutoPromotionGuaranteePriority(priority string) bool { + switch priority { + case ChannelPreparationAutoPromotionGuaranteePriorityCapacityFirst, + ChannelPreparationAutoPromotionGuaranteePriorityCountFirst: + return true + default: + return false + } +} + +func IsSupportedChannelPreparationAutoPromotionStrategy(strategy string) bool { + switch strategy { + case ChannelPreparationAutoPromotionStrategyPriorityWeighted, + ChannelPreparationAutoPromotionStrategySmallBalanceFirst, + ChannelPreparationAutoPromotionStrategyLargeBalanceFirst: + return true + default: + return false + } +} + +func ValidateChannelPreparationAutoPromotionRules(rules []ChannelPreparationAutoPromotionRule) error { + seenIds := make(map[string]bool) + for index := range rules { + rawRule := rules[index] + rawRule.GuaranteePriority = strings.TrimSpace(rawRule.GuaranteePriority) + rawRule.CountShortageStrategy = strings.TrimSpace(rawRule.CountShortageStrategy) + rawRule.CapacityShortageStrategy = strings.TrimSpace(rawRule.CapacityShortageStrategy) + rawRule.Strategy = strings.TrimSpace(rawRule.Strategy) + if rawRule.GuaranteePriority != "" && !IsSupportedChannelPreparationAutoPromotionGuaranteePriority(rawRule.GuaranteePriority) { + return fmt.Errorf("第 %d 条规则保障优先级无效", index+1) + } + if rawRule.CountShortageStrategy != "" && !IsSupportedChannelPreparationAutoPromotionStrategy(rawRule.CountShortageStrategy) { + return fmt.Errorf("第 %d 条规则数量不足策略无效", index+1) + } + if rawRule.CapacityShortageStrategy != "" && !IsSupportedChannelPreparationAutoPromotionStrategy(rawRule.CapacityShortageStrategy) { + return fmt.Errorf("第 %d 条规则容量不足策略无效", index+1) + } + if rawRule.Strategy != "" && !IsSupportedChannelPreparationAutoPromotionStrategy(rawRule.Strategy) && rawRule.CountShortageStrategy == "" && rawRule.CapacityShortageStrategy == "" { + return fmt.Errorf("第 %d 条规则策略无效", index+1) + } + + rule := rawRule + NormalizeChannelPreparationAutoPromotionRule(&rule) + if rule.Id == "" { + return fmt.Errorf("第 %d 条规则缺少 id", index+1) + } + if seenIds[rule.Id] { + return fmt.Errorf("规则 id 重复：%s", rule.Id) + } + seenIds[rule.Id] = true + if strings.TrimSpace(rule.Group) == "" { + return fmt.Errorf("第 %d 条规则缺少分组", index+1) + } + if rule.Type <= 0 { + return fmt.Errorf("第 %d 条规则渠道类型无效", index+1) + } + if rule.ThresholdUSD <= 0 { + return fmt.Errorf("第 %d 条规则阈值必须大于 0", index+1) + } + if rule.MinimumUsableChannelCount < 0 { + return fmt.Errorf("第 %d 条规则最低可用渠道数不能小于 0", index+1) + } + if !IsSupportedChannelPreparationAutoPromotionGuaranteePriority(rule.GuaranteePriority) { + return fmt.Errorf("第 %d 条规则保障优先级无效", index+1) + } + if !IsSupportedChannelPreparationAutoPromotionStrategy(rule.CountShortageStrategy) { + return fmt.Errorf("第 %d 条规则数量不足策略无效", index+1) + } + if !IsSupportedChannelPreparationAutoPromotionStrategy(rule.CapacityShortageStrategy) { + return fmt.Errorf("第 %d 条规则容量不足策略无效", index+1) + } + } + return nil +} + +func ValidateChannelPreparationAutoPromotionRulesJSONString(value string) error { + value = strings.TrimSpace(value) + if value == "" { + value = "[]" + } + var rules []ChannelPreparationAutoPromotionRule + if err := common.Unmarshal([]byte(value), &rules); err != nil { + return err + } + return ValidateChannelPreparationAutoPromotionRules(rules) +} diff --git a/setting/operation_setting/channel_preparation_auto_promotion_setting_test.go b/setting/operation_setting/channel_preparation_auto_promotion_setting_test.go new file mode 100644 index 000000000000..49e0ae96bf8f --- /dev/null +++ b/setting/operation_setting/channel_preparation_auto_promotion_setting_test.go @@ -0,0 +1,98 @@ +package operation_setting + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/require" +) + +func TestNormalizeChannelPreparationAutoPromotionRulePreservesLegacyJSON(t *testing.T) { + rule := ChannelPreparationAutoPromotionRule{ + Id: " legacy ", + Group: " default ", + Type: 14, + ThresholdUSD: 10, + Strategy: ChannelPreparationAutoPromotionStrategyLargeBalanceFirst, + } + + NormalizeChannelPreparationAutoPromotionRule(&rule) + + require.Equal(t, "legacy", rule.Id) + require.Equal(t, "default", rule.Group) + require.Equal(t, 0, rule.MinimumUsableChannelCount) + require.Equal(t, ChannelPreparationAutoPromotionGuaranteePriorityCapacityFirst, rule.GuaranteePriority) + require.Equal(t, ChannelPreparationAutoPromotionStrategyLargeBalanceFirst, rule.CountShortageStrategy) + require.Equal(t, ChannelPreparationAutoPromotionStrategyLargeBalanceFirst, rule.CapacityShortageStrategy) + require.Equal(t, rule.CapacityShortageStrategy, rule.Strategy) +} + +func TestValidateChannelPreparationAutoPromotionRulesAcceptsNewStrategies(t *testing.T) { + rules := []ChannelPreparationAutoPromotionRule{ + { + Id: "rule-1", + Enabled: true, + Group: "default", + Type: 14, + ThresholdUSD: 10, + MinimumUsableChannelCount: 2, + GuaranteePriority: ChannelPreparationAutoPromotionGuaranteePriorityCountFirst, + CountShortageStrategy: ChannelPreparationAutoPromotionStrategySmallBalanceFirst, + CapacityShortageStrategy: ChannelPreparationAutoPromotionStrategyLargeBalanceFirst, + }, + } + + require.NoError(t, ValidateChannelPreparationAutoPromotionRules(rules)) +} + +func TestValidateChannelPreparationAutoPromotionRulesRejectsInvalidRules(t *testing.T) { + baseRule := ChannelPreparationAutoPromotionRule{ + Id: "rule-1", + Enabled: true, + Group: "default", + Type: 14, + ThresholdUSD: 10, + } + + t.Run("duplicate id", func(t *testing.T) { + err := ValidateChannelPreparationAutoPromotionRules([]ChannelPreparationAutoPromotionRule{baseRule, baseRule}) + require.Error(t, err) + }) + + t.Run("negative minimum usable count", func(t *testing.T) { + rule := baseRule + rule.MinimumUsableChannelCount = -1 + err := ValidateChannelPreparationAutoPromotionRules([]ChannelPreparationAutoPromotionRule{rule}) + require.Error(t, err) + }) + + t.Run("unknown guarantee priority", func(t *testing.T) { + rule := baseRule + rule.GuaranteePriority = "speed_first" + err := ValidateChannelPreparationAutoPromotionRules([]ChannelPreparationAutoPromotionRule{rule}) + require.Error(t, err) + }) + + t.Run("unknown shortage strategy", func(t *testing.T) { + rule := baseRule + rule.CountShortageStrategy = "random" + err := ValidateChannelPreparationAutoPromotionRules([]ChannelPreparationAutoPromotionRule{rule}) + require.Error(t, err) + }) +} + +func TestValidateChannelPreparationAutoPromotionRulesJSONStringAcceptsOldRules(t *testing.T) { + payload, err := common.Marshal([]ChannelPreparationAutoPromotionRule{ + { + Id: "old-rule", + Enabled: true, + Group: "default", + Type: 14, + ThresholdUSD: 5, + Strategy: ChannelPreparationAutoPromotionStrategyPriorityWeighted, + }, + }) + require.NoError(t, err) + + require.NoError(t, ValidateChannelPreparationAutoPromotionRulesJSONString(string(payload))) +} diff --git a/setting/ratio_setting/model_ratio.go b/setting/ratio_setting/model_ratio.go index 23fd360e366f..11e057997744 100644 --- a/setting/ratio_setting/model_ratio.go +++ b/setting/ratio_setting/model_ratio.go @@ -140,6 +140,7 @@ var defaultModelRatio = map[string]float64{ "claude-3-7-sonnet-20250219-thinking": 1.5, "claude-sonnet-4-20250514": 1.5, "claude-sonnet-4-5-20250929": 1.5, + "claude-sonnet-4-6": 1.5, "claude-opus-4-5-20251101": 2.5, "claude-opus-4-6": 2.5, "claude-opus-4-6-max": 2.5, diff --git a/setting/ratio_setting/model_ratio_test.go b/setting/ratio_setting/model_ratio_test.go new file mode 100644 index 000000000000..8c4720a53ce9 --- /dev/null +++ b/setting/ratio_setting/model_ratio_test.go @@ -0,0 +1,13 @@ +package ratio_setting + +import ( + "testing" + + "github.com/stretchr/testify/require" +) + +func TestDefaultModelRatioIncludesClaudeSonnet46(t *testing.T) { + ratio, ok := defaultModelRatio["claude-sonnet-4-6"] + require.True(t, ok) + require.Equal(t, 1.5, ratio) +} diff --git a/web/bun.lock b/web/bun.lock index 18ec756e82d5..0bc56888622c 100644 --- a/web/bun.lock +++ b/web/bun.lock @@ -23,6 +23,8 @@ "history": "^5.3.0", "i18next": "^23.16.8", "i18next-browser-languagedetector": "^7.2.0", + "jspreadsheet-ce": "^5.0.4", + "jsuites": "^5.13.5", "katex": "^0.16.22", "lucide-react": "^0.511.0", "marked": "^4.1.1", @@ -477,6 +479,8 @@ "@jridgewell/trace-mapping": ["@jridgewell/trace-mapping@0.3.31", "", { "dependencies": { "@jridgewell/resolve-uri": "^3.1.0", "@jridgewell/sourcemap-codec": "^1.4.14" } }, "sha512-zzNR+SdQSDJzc8joaeP8QQoCQr8NuYx2dIIytl1QeBEZHJ9uW6hebsrYgbz8hJwUQao3TWCMtmfV8Nu1twOLAw=="], + "@jspreadsheet/formula": ["@jspreadsheet/formula@2.0.2", "", {}, "sha512-PDQYf9REQA53I7tVYkvkeyQxrd5jcjUeHgItYnRpjN2QiIQwawSqBDtGGEVQTSboTG+JwgGCuhvOpj7FxeKwew=="], + "@lit-labs/ssr-dom-shim": ["@lit-labs/ssr-dom-shim@1.6.0", "", {}, "sha512-VHb0ALPMTlgKjM6yIxxoQNnpKyUKLD04VzeQdsiXkMqkvYlAHxq9glGLmgbb889/1GsohSOAjvQYoiBppXFqrQ=="], "@lit/reactive-element": ["@lit/reactive-element@2.1.2", "", { "dependencies": { "@lit-labs/ssr-dom-shim": "^1.5.0" } }, "sha512-pbCDiVMnne1lYUIaYNN5wrwQXDtHaYtg7YEFPeW+hws6U47WeFvISGUWekPGKWOP1ygrs0ef0o1VJMk1exos5A=="], @@ -2049,6 +2053,10 @@ "jsonfile": ["jsonfile@6.2.1", "", { "dependencies": { "universalify": "^2.0.0" }, "optionalDependencies": { "graceful-fs": "^4.1.6" } }, "sha512-zwOTdL3rFQ/lRdBnntKVOX6k5cKJwEc1HdilT71BWEu7J41gXIB2MRp+vxduPSwZJPWBxEzv4yH1wYLJGUHX4Q=="], + "jspreadsheet-ce": ["jspreadsheet-ce@5.0.4", "", { "dependencies": { "@jspreadsheet/formula": "^2.0.2", "jsuites": "^5.12.0" } }, "sha512-ra1JI1n+tEGgRMzTzNkPZjG0HZz8W6bFGAiTiHl+eYarXdRmS5qDc/ua3l2ev7oZ6Og9kjfrXYHVLUWiVc308w=="], + + "jsuites": ["jsuites@5.13.5", "", {}, "sha512-cvkcpy/v5I3+IAcNPE4UP38PFCEfUQw9JI5NN61dlcXLwkD+2UTIOsRPvgMLeqI1eDWHL4AHfrbcE/+TFciUsw=="], + "katex": ["katex@0.16.47", "", { "dependencies": { "commander": "^8.3.0" }, "bin": { "katex": "cli.js" } }, "sha512-Eeo8Ys1doU1z+x8AZsPpQu+p/QcZBI5PeOo7QGQdy2x2m0MU/hYagBbGOmXwr5KVbEfVuWv9LpnQWeehogurjg=="], "keyv": ["keyv@4.5.4", "", { "dependencies": { "json-buffer": "3.0.1" } }, "sha512-oxVHkHR/EJf2CNXnWxRLW6mg7JyCCUcG0DtEGmL2ctUo1PNTin1PUil+r/+4r5MpVgC/fn1kjsx7mjSujKqIpw=="], diff --git a/web/classic/package.json b/web/classic/package.json index 8a840d67e7c3..7ee20a5779d1 100644 --- a/web/classic/package.json +++ b/web/classic/package.json @@ -18,6 +18,7 @@ "highlight.js": "^11.11.1", "i18next": "^23.16.8", "i18next-browser-languagedetector": "^7.2.0", + "jspreadsheet-ce": "^5.0.4", "katex": "^0.16.22", "lucide-react": "^0.511.0", "marked": "^4.1.1", @@ -41,7 +42,8 @@ "remark-math": "^6.0.0", "sse.js": "catalog:", "unist-util-visit": "^5.0.0", - "use-debounce": "^10.0.4" + "use-debounce": "^10.0.4", + "jsuites": "^5.13.5" }, "scripts": { "dev": "rsbuild dev", diff --git a/web/classic/src/App.jsx b/web/classic/src/App.jsx index a5d1ebc00b32..48c680a252a5 100644 --- a/web/classic/src/App.jsx +++ b/web/classic/src/App.jsx @@ -21,7 +21,7 @@ import React, { lazy, Suspense, useContext, useMemo } from 'react'; import { Route, Routes, useLocation, useParams } from 'react-router-dom'; import Loading from './components/common/ui/Loading'; import User from './pages/User'; -import { AuthRedirect, PrivateRoute, AdminRoute } from './helpers'; +import { AuthRedirect, PrivateRoute, AdminRoute, RootRoute } from './helpers'; import RegisterForm from './components/auth/RegisterForm'; import LoginForm from './components/auth/LoginForm'; import NotFound from './pages/NotFound'; @@ -32,6 +32,7 @@ import { StatusContext } from './context/Status'; import PasswordResetForm from './components/auth/PasswordResetForm'; import PasswordResetConfirm from './components/auth/PasswordResetConfirm'; import Channel from './pages/Channel'; +import ChannelPreparation from './pages/ChannelPreparation'; import Token from './pages/Token'; import Redemption from './pages/Redemption'; import TopUp from './pages/TopUp'; @@ -55,6 +56,8 @@ const Dashboard = lazy(() => import('./pages/Dashboard')); const About = lazy(() => import('./pages/About')); const UserAgreement = lazy(() => import('./pages/UserAgreement')); const PrivacyPolicy = lazy(() => import('./pages/PrivacyPolicy')); +const CostReport = lazy(() => import('./pages/CostReport')); +const QueryKey = lazy(() => import('./pages/QueryKey')); function DynamicOAuth2Callback() { const { provider } = useParams(); @@ -139,6 +142,34 @@ function App() { } /> + + + + } + /> + + } key={location.pathname}> + + + + } + /> + + } key={location.pathname}> + + + + } + /> } /> + + + + } + /> { '/console/midjourney', '/console/task', '/console/models', + '/console/query-key', '/pricing', ]; diff --git a/web/classic/src/components/layout/SiderBar.jsx b/web/classic/src/components/layout/SiderBar.jsx index a4375cf856cc..ba8391ca5fe6 100644 --- a/web/classic/src/components/layout/SiderBar.jsx +++ b/web/classic/src/components/layout/SiderBar.jsx @@ -33,6 +33,7 @@ import { Nav, Divider, Button } from '@douyinfe/semi-ui'; const routerMap = { home: '/', channel: '/console/channel', + channelPreparation: '/console/channel-preparations', token: '/console/token', redemption: '/console/redemption', topup: '/console/topup', @@ -153,6 +154,12 @@ const SiderBar = ({ onNavigate = () => {} }) => { to: '/channel', className: isAdmin() ? '' : 'tableHiddle', }, + { + text: t('渠道备货池'), + itemKey: 'channelPreparation', + to: '/console/channel-preparations', + className: isAdmin() ? '' : 'tableHiddle', + }, { text: t('订阅管理'), itemKey: 'subscription', @@ -205,7 +212,7 @@ const SiderBar = ({ onNavigate = () => {} }) => { { text: t('操练场'), itemKey: 'playground', - to: '/playground', + to: '/console/playground', }, { text: t('聊天'), diff --git a/web/classic/src/components/query-key/QueryKeyPage.jsx b/web/classic/src/components/query-key/QueryKeyPage.jsx new file mode 100644 index 000000000000..bc0c1d97643b --- /dev/null +++ b/web/classic/src/components/query-key/QueryKeyPage.jsx @@ -0,0 +1,1361 @@ +/* +Copyright (C) 2025 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ + +import React, { useMemo, useRef, useState } from 'react'; +import { useTranslation } from 'react-i18next'; +import { + Banner, + Button, + Card, + Collapse, + Dropdown, + Empty, + Space, + Spin, + Table, + Tag, + TextArea, + Tooltip, + Typography, +} from '@douyinfe/semi-ui'; +import { + IconAlertTriangle, + IconCopy, + IconRefresh, + IconSearch, +} from '@douyinfe/semi-icons'; +import { + API, + copy, + getChannelIcon, + renderGroup, + renderQuota, + renderQuotaWithAmount, + showError, + showInfo, + showSuccess, +} from '../../helpers'; +import { CHANNEL_OPTIONS } from '../../constants'; + +const { Text, Title } = Typography; + +const STATUS_CONFIG = { + found: { color: 'green', label: '已找到' }, + not_found: { color: 'grey', label: '未找到' }, + over_brushed: { color: 'red', label: '已超刷' }, +}; + +const SOURCE_CONFIG = { + channel: { color: 'green', label: '正式渠道' }, + preparation: { color: 'blue', label: '备货池' }, +}; + +const QUERY_KEY_TEST_STATUS = { + untested: { color: 'grey', label: '未测试' }, + testing: { color: 'blue', label: '测试中' }, + success: { color: 'green', label: '成功' }, + failed: { color: 'red', label: '失败' }, + partial: { color: 'orange', label: '部分成功' }, +}; + +const DEFAULT_BATCH_TEST_MODEL = ''; +const DEFAULT_BATCH_TEST_MODEL_LABEL = '使用渠道配置的测试模型'; + +const BUCKETS = [ + { key: 'all', label: '全部' }, + { key: 'found', label: '已找到' }, + { key: 'not_found', label: '未找到' }, + { key: 'over_brushed', label: '已超刷' }, +]; + +const stableStringify = (value) => { + if (Array.isArray(value)) return `[${value.map(stableStringify).join(',')}]`; + if (value && typeof value === 'object') { + return `{${Object.keys(value) + .sort() + .map((key) => `${JSON.stringify(key)}:${stableStringify(value[key])}`) + .join(',')}}`; + } + return JSON.stringify(value); +}; + +const normalizeMatchKey = (value) => { + const trimmed = String(value || '').trim(); + if (!trimmed) return ''; + try { + const parsed = JSON.parse(trimmed); + if (typeof parsed === 'string') return parsed.trim(); + if (parsed === null || parsed === undefined) return ''; + return stableStringify(parsed); + } catch (error) { + return trimmed; + } +}; + +const parseKeyInput = (text) => { + const lines = text + .split(/\r?\n/) + .map((line) => line.trim()) + .filter(Boolean); + const seen = new Set(); + const keys = []; + + lines.forEach((line) => { + const matchKey = normalizeMatchKey(line); + if (!matchKey || seen.has(matchKey)) return; + seen.add(matchKey); + keys.push(line); + }); + + return { + keys, + totalInput: lines.length, + duplicateCount: lines.length - keys.length, + }; +}; + +const formatDate = (timestamp) => { + if (!timestamp) return '-'; + return new Date(timestamp * 1000).toLocaleString(); +}; + +const channelTypeLabel = (type) => { + const option = CHANNEL_OPTIONS.find((item) => item.value === type); + return option?.label || type || '-'; +}; + +const getStatusConfig = (status) => + STATUS_CONFIG[status] || STATUS_CONFIG.not_found; + +const getSourceConfig = (source) => + SOURCE_CONFIG[source || 'channel'] || SOURCE_CONFIG.channel; + +const normalizeCopyCell = (value) => { + if (value === null || value === undefined) return ''; + return String(value).replace(/\t/g, ' ').replace(/\r?\n/g, ' '); +}; + +const buildTsv = (rows, columns, includeHeader) => { + const lines = []; + if (includeHeader) { + lines.push(columns.map((column) => normalizeCopyCell(column.label)).join('\t')); + } + rows.forEach((row) => { + lines.push( + columns.map((column) => normalizeCopyCell(column.getValue(row))).join('\t'), + ); + }); + return lines.join('\n'); +}; + +const buildQueryKeyTestId = (key, channel) => + `${normalizeMatchKey(key)}::${channel?.source || 'channel'}::${channel?.id}`; + +const MetricCard = ({ title, value, color }) => ( + +
{title}
+
+ {value} +
+ +); + +const QueryKeyPage = () => { + const { t } = useTranslation(); + const [inputText, setInputText] = useState(''); + const [loading, setLoading] = useState(false); + const [report, setReport] = useState(null); + const [activeBucket, setActiveBucket] = useState('all'); + const [queryKeyTestResults, setQueryKeyTestResults] = useState({}); + const [testingQueryKeyIds, setTestingQueryKeyIds] = useState(new Set()); + const [isQueryKeyBatchTesting, setIsQueryKeyBatchTesting] = useState(false); + const [queryKeyBatchProgress, setQueryKeyBatchProgress] = useState({ + finished: 0, + total: 0, + }); + const shouldStopQueryKeyBatchTestingRef = useRef(false); + + const parsed = useMemo(() => parseKeyInput(inputText), [inputText]); + const items = Array.isArray(report?.items) ? report.items : []; + + const filteredItems = useMemo(() => { + if (activeBucket === 'all') return items; + if (activeBucket === 'found') return items.filter((item) => item.found); + return items.filter((item) => item.status === activeBucket); + }, [activeBucket, items]); + + const bucketCounts = { + all: items.length, + found: report?.found_count || 0, + not_found: report?.not_found_count || 0, + over_brushed: report?.over_brushed_count || 0, + }; + + const submitReport = async () => { + if (loading || isQueryKeyBatchTesting) return; + if (parsed.keys.length === 0) { + showError(t('请输入密钥')); + return; + } + if (parsed.keys.length > 10000) { + showError(t('最多支持 10000 个唯一密钥')); + return; + } + + setLoading(true); + try { + const res = await API.post('/api/channel/query-key/report', { + keys: parsed.keys, + }); + const { success, message, data } = res.data || {}; + if (!success) { + showError(message || t('查询失败')); + return; + } + setReport(data); + setActiveBucket('all'); + setQueryKeyTestResults({}); + setTestingQueryKeyIds(new Set()); + setQueryKeyBatchProgress({ finished: 0, total: 0 }); + shouldStopQueryKeyBatchTestingRef.current = false; + showSuccess(t('查询完成')); + } catch (error) { + showError( + error?.response?.data?.message || error?.message || t('网络错误'), + ); + } finally { + setLoading(false); + } + }; + + const clearAll = () => { + if (loading || isQueryKeyBatchTesting) return; + setInputText(''); + setReport(null); + setActiveBucket('all'); + setQueryKeyTestResults({}); + setTestingQueryKeyIds(new Set()); + setQueryKeyBatchProgress({ finished: 0, total: 0 }); + shouldStopQueryKeyBatchTestingRef.current = false; + }; + + const copyKey = async (value) => { + const ok = await copy(value || ''); + if (ok) showSuccess(t('已复制')); + else showError(t('复制失败')); + }; + + const getStatusLabel = (item) => { + const config = getStatusConfig(item.status); + const labels = [t(config.label)]; + if (item.original_amount_shared) { + labels.push(t('原始额度为共享余额')); + } + return labels.join(' / '); + }; + + const getChannelStatusLabel = (channel) => { + return getChannelStatusMeta(channel).label; + }; + + const getChannelStatusMeta = (channel) => { + if (!channel) return { color: 'grey', label: t('未找到') }; + if (channel.source === 'preparation') { + if (channel.status === 2) return { color: 'green', label: t('已晋升') }; + if (channel.status === 3) return { color: 'grey', label: t('已归档') }; + if (channel.status === 4) return { color: 'orange', label: t('晋升中') }; + return { color: 'blue', label: t('待晋升') }; + } + return channel.status === 1 + ? { color: 'green', label: t('已启用') } + : { color: 'grey', label: t('已禁用') }; + }; + + const getItemChannels = (item) => + Array.isArray(item?.channels) ? item.channels : []; + + const buildQueryKeyBatchTasks = (rows) => + rows.flatMap((item) => + getItemChannels(item).map((channel) => ({ + item, + channel, + })), + ); + + const getItemChannelStatusText = (item) => { + const channels = getItemChannels(item); + if (channels.length === 0) return t('未找到'); + const counts = channels.reduce((acc, channel) => { + const label = getChannelStatusMeta(channel).label; + acc[label] = (acc[label] || 0) + 1; + return acc; + }, {}); + return Object.entries(counts) + .map(([label, count]) => (channels.length === 1 ? label : `${label} ${count}`)) + .join(' / '); + }; + + const renderItemChannelStatus = (item) => { + const channels = getItemChannels(item); + if (channels.length === 0) return {t('未找到')}; + const counts = channels.reduce((acc, channel) => { + const meta = getChannelStatusMeta(channel); + if (!acc[meta.label]) acc[meta.label] = { ...meta, count: 0 }; + acc[meta.label].count += 1; + return acc; + }, {}); + return ( + + {Object.values(counts).map((meta) => ( + + {meta.label} + {channels.length > 1 ? ` ${meta.count}` : ''} + + ))} + + ); + }; + + const getQueryKeyTestResult = (item, channel) => + queryKeyTestResults[buildQueryKeyTestId(item?.key, channel)]; + + const isQueryKeyTesting = (item, channel) => + testingQueryKeyIds.has(buildQueryKeyTestId(item?.key, channel)); + + const getItemTestSummary = (item) => { + const channels = getItemChannels(item); + const results = channels + .map((channel) => getQueryKeyTestResult(item, channel)) + .filter(Boolean); + const testingCount = channels.filter((channel) => + isQueryKeyTesting(item, channel), + ).length; + const successCount = results.filter((result) => result.success).length; + const failedCount = results.filter((result) => !result.success).length; + const responseTimes = results + .filter((result) => typeof result.time === 'number') + .map((result) => result.time); + return { + total: channels.length, + tested: results.length, + testingCount, + successCount, + failedCount, + fastestTime: + responseTimes.length > 0 ? Math.min(...responseTimes) : null, + firstFailure: results.find((result) => !result.success), + }; + }; + + const getQueryKeyTestStatusText = (item) => { + const summary = getItemTestSummary(item); + if (summary.total === 0) return '-'; + if (summary.testingCount > 0) return t('测试中'); + if (summary.tested === 0) return t('未测试'); + if (summary.total === 1) { + return summary.successCount === 1 ? t('成功') : t('失败'); + } + if (summary.failedCount === 0 && summary.tested === summary.total) { + return t('全部成功'); + } + if (summary.successCount > 0) { + return `${t('部分成功')} ${summary.successCount}/${summary.tested}`; + } + return `${t('全部失败')} ${summary.failedCount}/${summary.tested}`; + }; + + const getQueryKeyResponseTimeText = (item) => { + const summary = getItemTestSummary(item); + if (summary.testingCount > 0) return t('测试中'); + if (summary.fastestTime === null) return '-'; + return `${summary.fastestTime.toFixed(2)}s`; + }; + + const getSingleTestStatusText = (item, channel) => { + if (isQueryKeyTesting(item, channel)) return t('测试中'); + const result = getQueryKeyTestResult(item, channel); + if (!result) return t('未测试'); + return result.success ? t('成功') : t('失败'); + }; + + const getSingleResponseTimeText = (item, channel) => { + if (isQueryKeyTesting(item, channel)) return t('测试中'); + const result = getQueryKeyTestResult(item, channel); + if (!result || typeof result.time !== 'number') return '-'; + return `${result.time.toFixed(2)}s`; + }; + + const renderSingleTestStatus = (item, channel) => { + if (isQueryKeyTesting(item, channel)) { + return {t('测试中')}; + } + const result = getQueryKeyTestResult(item, channel); + if (!result) { + return {t('未测试')}; + } + const tag = ( + + {result.success ? t('成功') : t('失败')} + + ); + if (!result.success && result.message) { + return {tag}; + } + return tag; + }; + + const renderItemTestStatus = (item) => { + const summary = getItemTestSummary(item); + if (summary.total === 0) return -; + if (summary.testingCount > 0) { + return ( + + {t('测试中')} {summary.testingCount}/{summary.total} + + ); + } + if (summary.tested === 0) { + return {t('未测试')}; + } + if (summary.total === 1) { + const channel = getItemChannels(item)[0]; + return renderSingleTestStatus(item, channel); + } + const allSuccess = + summary.successCount === summary.total && summary.tested === summary.total; + const allFailed = summary.failedCount === summary.tested; + const color = allSuccess + ? QUERY_KEY_TEST_STATUS.success.color + : allFailed + ? QUERY_KEY_TEST_STATUS.failed.color + : QUERY_KEY_TEST_STATUS.partial.color; + const tag = ( + + {allSuccess + ? t('全部成功') + : allFailed + ? t('全部失败') + : t('部分成功')}{' '} + {summary.successCount}/{summary.tested} + + ); + if (summary.firstFailure?.message) { + return {tag}; + } + return tag; + }; + + const renderSingleResponseTime = (item, channel) => { + if (isQueryKeyTesting(item, channel)) return {t('测试中')}; + const result = getQueryKeyTestResult(item, channel); + if (!result || typeof result.time !== 'number') return -; + return {result.time.toFixed(2)}s; + }; + + const renderItemResponseTime = (item) => { + const summary = getItemTestSummary(item); + if (summary.testingCount > 0) return {t('测试中')}; + if (summary.fastestTime === null) return -; + return {summary.fastestTime.toFixed(2)}s; + }; + + const testQueryKeyChannel = async (item, channel, options = {}) => { + const testId = buildQueryKeyTestId(item?.key, channel); + if (!item?.key || !channel?.id || testingQueryKeyIds.has(testId)) return null; + + setTestingQueryKeyIds((previous) => { + const next = new Set(previous); + next.add(testId); + return next; + }); + + try { + const res = await API.post('/api/channel/query-key/test', { + key: item.key, + source: channel.source || 'channel', + target_id: channel.id, + model: options.model || DEFAULT_BATCH_TEST_MODEL, + endpoint_type: options.endpointType || '', + stream: Boolean(options.stream), + }); + const payload = res.data || {}; + const result = { + success: Boolean(payload.success), + message: payload.message || '', + time: typeof payload.time === 'number' ? payload.time : 0, + errorCode: payload.error_code || '', + }; + setQueryKeyTestResults((previous) => ({ + ...previous, + [testId]: result, + })); + if (!options.silent) { + if (result.success) { + showSuccess(t('测试成功')); + } else { + showError(result.message || t('测试失败')); + } + } + return result; + } catch (error) { + const result = { + success: false, + message: + error?.response?.data?.message || error?.message || t('网络错误'), + time: 0, + errorCode: '', + }; + setQueryKeyTestResults((previous) => ({ + ...previous, + [testId]: result, + })); + if (!options.silent) showError(result.message); + return result; + } finally { + setTestingQueryKeyIds((previous) => { + const next = new Set(previous); + next.delete(testId); + return next; + }); + } + }; + + const testQueryKeyItem = async (item) => { + const channels = getItemChannels(item); + if (channels.length === 0) { + showError(t('没有匹配的渠道')); + return; + } + if (channels.length === 1) { + await testQueryKeyChannel(item, channels[0], { + model: DEFAULT_BATCH_TEST_MODEL, + }); + return; + } + + const results = []; + for (const channel of channels) { + // Keep tests sequential to avoid creating an accidental upstream burst. + // eslint-disable-next-line no-await-in-loop + const result = await testQueryKeyChannel(item, channel, { + silent: true, + model: DEFAULT_BATCH_TEST_MODEL, + }); + if (result) results.push(result); + } + const successCount = results.filter((result) => result.success).length; + const failedCount = results.length - successCount; + if (failedCount === 0) { + showSuccess(t('测试完成：全部成功')); + } else { + showError( + t('测试完成：成功 {{success}} / 失败 {{failed}}') + .replace('{{success}}', successCount) + .replace('{{failed}}', failedCount), + ); + } + }; + + const batchTestQueryKeyItems = async (scope) => { + if (isQueryKeyBatchTesting) { + showInfo(t('批量测试正在进行中')); + return; + } + + const sourceRows = scope === 'filtered' ? filteredItems : items; + const tasks = buildQueryKeyBatchTasks(sourceRows); + if (tasks.length === 0) { + showError(t('没有可测试的渠道')); + return; + } + + const taskIds = new Set( + tasks.map(({ item, channel }) => buildQueryKeyTestId(item.key, channel)), + ); + setQueryKeyTestResults((previous) => { + const next = { ...previous }; + taskIds.forEach((testId) => { + delete next[testId]; + }); + return next; + }); + + setIsQueryKeyBatchTesting(true); + setQueryKeyBatchProgress({ finished: 0, total: tasks.length }); + shouldStopQueryKeyBatchTestingRef.current = false; + + let successCount = 0; + let failedCount = 0; + let finishedCount = 0; + const concurrencyLimit = 5; + + try { + for (let i = 0; i < tasks.length; i += concurrencyLimit) { + if (shouldStopQueryKeyBatchTestingRef.current) break; + const batch = tasks.slice(i, i + concurrencyLimit); + // eslint-disable-next-line no-await-in-loop + const results = await Promise.allSettled( + batch.map(({ item, channel }) => + testQueryKeyChannel(item, channel, { + silent: true, + model: DEFAULT_BATCH_TEST_MODEL, + }), + ), + ); + results.forEach((result) => { + finishedCount += 1; + if (result.status === 'fulfilled' && result.value?.success) { + successCount += 1; + } else { + failedCount += 1; + } + }); + setQueryKeyBatchProgress({ + finished: finishedCount, + total: tasks.length, + }); + } + + if (shouldStopQueryKeyBatchTestingRef.current) { + showInfo( + t('批量测试已停止：完成 {{finished}} / {{total}}') + .replace('{{finished}}', finishedCount) + .replace('{{total}}', tasks.length), + ); + } else if (failedCount === 0) { + showSuccess( + t('批量测试完成：全部成功，共 {{count}} 个') + .replace('{{count}}', successCount), + ); + } else { + showError( + t('批量测试完成：成功 {{success}} / 失败 {{failed}}') + .replace('{{success}}', successCount) + .replace('{{failed}}', failedCount), + ); + } + } finally { + setIsQueryKeyBatchTesting(false); + shouldStopQueryKeyBatchTestingRef.current = false; + } + }; + + const stopQueryKeyBatchTest = () => { + if (!isQueryKeyBatchTesting) return; + shouldStopQueryKeyBatchTestingRef.current = true; + showInfo(t('正在停止批量测试，已开始的请求会继续完成')); + }; + + const isItemTesting = (item) => + getItemChannels(item).some((channel) => isQueryKeyTesting(item, channel)); + + const mainCopyColumns = useMemo( + () => [ + { label: t('密钥'), getValue: (item) => item.key || '' }, + { label: t('结果'), getValue: getStatusLabel }, + { label: t('渠道状态'), getValue: getItemChannelStatusText }, + { label: t('测试状态'), getValue: getQueryKeyTestStatusText }, + { label: t('响应时间'), getValue: getQueryKeyResponseTimeText }, + { + label: t('渠道数'), + getValue: (item) => item.channel_count || 0, + }, + { + label: t('已用额度'), + getValue: (item) => renderQuota(item.used_quota || 0), + }, + { + label: t('已用金额'), + getValue: (item) => renderQuotaWithAmount(item.used_amount || 0), + }, + { + label: t('原始额度'), + getValue: (item) => { + const amount = renderQuotaWithAmount(item.original_amount || 0); + return item.original_amount_shared ? `${amount} (${t('共享')})` : amount; + }, + }, + { + label: t('理论当前额度'), + getValue: (item) => renderQuotaWithAmount(item.current_amount || 0), + }, + { + label: t('超刷金额'), + getValue: (item) => renderQuotaWithAmount(item.over_brush_amount || 0), + }, + ], + [t, queryKeyTestResults, testingQueryKeyIds], + ); + + const channelDetailCopyColumns = useMemo( + () => [ + { label: t('密钥'), getValue: ({ item }) => item.key || '' }, + { + label: t('来源'), + getValue: ({ channel }) => t(getSourceConfig(channel.source).label), + }, + { label: 'ID', getValue: ({ channel }) => channel.id || '' }, + { label: t('渠道'), getValue: ({ channel }) => channel.name || '' }, + { + label: t('类型'), + getValue: ({ channel }) => channelTypeLabel(channel.type), + }, + { + label: t('状态'), + getValue: ({ channel }) => getChannelStatusLabel(channel), + }, + { + label: t('测试状态'), + getValue: ({ item, channel }) => getSingleTestStatusText(item, channel), + }, + { + label: t('响应时间'), + getValue: ({ item, channel }) => getSingleResponseTimeText(item, channel), + }, + { label: t('分组'), getValue: ({ channel }) => channel.group || '' }, + { + label: t('匹配密钥数'), + getValue: ({ channel }) => channel.matched_key_count || 1, + }, + { + label: t('已用额度'), + getValue: ({ channel }) => renderQuota(channel.used_quota || 0), + }, + { + label: t('匹配已用金额'), + getValue: ({ channel }) => + renderQuotaWithAmount(channel.matched_used_amount || 0), + }, + { + label: t('原始额度'), + getValue: ({ channel }) => + renderQuotaWithAmount(channel.original_amount || 0), + }, + { + label: t('理论当前额度'), + getValue: ({ channel }) => + renderQuotaWithAmount(channel.current_amount || 0), + }, + { + label: t('超刷金额'), + getValue: ({ channel }) => + renderQuotaWithAmount(channel.over_brush_amount || 0), + }, + { + label: t('余额更新时间'), + getValue: ({ channel }) => formatDate(channel.balance_updated_time), + }, + ], + [t, queryKeyTestResults, testingQueryKeyIds], + ); + + const flattenChannelDetails = (rows) => + rows.flatMap((item) => + (Array.isArray(item.channels) ? item.channels : []).map((channel) => ({ + item, + channel, + })), + ); + + const copyRows = async (rows, copyColumns, includeHeader) => { + if (!rows.length) { + showError(t('暂无报告数据')); + return; + } + const ok = await copy(buildTsv(rows, copyColumns, includeHeader)); + if (ok) showSuccess(t('已复制')); + else showError(t('复制失败')); + }; + + const copyColumn = async (rows, copyColumnConfig, includeHeader) => { + await copyRows(rows, [copyColumnConfig], includeHeader); + }; + + const renderCopyMenu = (includeHeader) => ( + + copyRows(filteredItems, mainCopyColumns, includeHeader)} + > + {t('当前筛选结果')} + + copyRows(items, mainCopyColumns, includeHeader)}> + {t('全部结果')} + + + copyRows( + flattenChannelDetails(filteredItems), + channelDetailCopyColumns, + includeHeader, + ) + } + > + {t('当前筛选渠道明细')} + + + copyRows( + flattenChannelDetails(items), + channelDetailCopyColumns, + includeHeader, + ) + } + > + {t('全部渠道明细')} + + +
+ {t('单列（当前筛选）')} +
+ {mainCopyColumns.map((copyColumnConfig) => ( + + copyColumn(filteredItems, copyColumnConfig, includeHeader) + } + > + {copyColumnConfig.label} + + ))} + + ); + + const renderBatchTestMenu = () => ( + + batchTestQueryKeyItems('filtered')}> + {t('测试当前筛选')} + + batchTestQueryKeyItems('all')}> + {t('测试全部结果')} + + +
+ {t('默认模型：{{model}}').replace( + '{{model}}', + t(DEFAULT_BATCH_TEST_MODEL_LABEL), + )} +
+ + ); + + const channelColumns = [ + { + title: t('渠道'), + dataIndex: 'name', + width: 300, + render: (name, record) => { + const sourceConfig = getSourceConfig(record.source); + return ( +
+ {getChannelIcon(record.type)} + #{record.id} + {t(sourceConfig.label)} + {name || '-'} + {record.is_multi_key ? {t('多密钥')} : null} + {record.matched_key_count > 1 ? ( + {t('共享原始额度')} + ) : null} +
+ ); + }, + }, + { + title: t('类型'), + dataIndex: 'type', + width: 150, + render: (type) => channelTypeLabel(type), + }, + { + title: t('状态'), + dataIndex: 'status', + width: 110, + render: (_, record) => { + const meta = getChannelStatusMeta(record); + return {meta.label}; + }, + }, + { + title: t('测试状态'), + dataIndex: 'query_key_test_status', + width: 120, + render: (_, record) => renderSingleTestStatus(record.__item, record), + }, + { + title: t('响应时间'), + dataIndex: 'query_key_response_time', + width: 120, + render: (_, record) => renderSingleResponseTime(record.__item, record), + }, + { + title: t('操作'), + dataIndex: 'query_key_operate', + width: 120, + fixed: 'right', + render: (_, record) => { + const item = record.__item; + return ( + + ); + }, + }, + { + title: t('分组'), + dataIndex: 'group', + width: 140, + render: (group) => ( + + {String(group || '') + .split(',') + .map((item) => renderGroup(item))} + + ), + }, + { + title: t('匹配密钥数'), + dataIndex: 'matched_key_count', + width: 120, + render: (count) => count || 1, + }, + { + title: t('已用额度'), + dataIndex: 'used_quota', + width: 160, + render: (quota) => renderQuota(quota || 0), + }, + { + title: t('匹配已用金额'), + dataIndex: 'matched_used_amount', + width: 180, + render: (amount) => renderQuotaWithAmount(amount || 0), + }, + { + title: t('原始额度'), + dataIndex: 'original_amount', + width: 180, + render: (amount) => renderQuotaWithAmount(amount || 0), + }, + { + title: t('理论当前额度'), + dataIndex: 'current_amount', + width: 180, + render: (amount) => renderQuotaWithAmount(amount || 0), + }, + { + title: t('超刷金额'), + dataIndex: 'over_brush_amount', + width: 160, + render: (amount) => ( + 0 ? 'danger' : 'secondary'}> + {renderQuotaWithAmount(amount || 0)} + + ), + }, + { + title: t('余额更新时间'), + dataIndex: 'balance_updated_time', + width: 180, + render: formatDate, + }, + ]; + + const columns = [ + { + title: t('密钥'), + dataIndex: 'key', + width: 520, + render: (key) => ( +
+ + {key} + +
+ ), + }, + { + title: t('结果'), + dataIndex: 'status', + width: 180, + render: (status, record) => { + const config = getStatusConfig(status); + return ( + + {t(config.label)} + {record.original_amount_shared ? ( + {t('原始额度为共享余额')} + ) : null} + + ); + }, + }, + { + title: t('渠道状态'), + dataIndex: 'channel_status', + width: 180, + render: (_, record) => renderItemChannelStatus(record), + }, + { + title: t('测试状态'), + dataIndex: 'query_key_test_status', + width: 150, + render: (_, record) => renderItemTestStatus(record), + }, + { + title: t('响应时间'), + dataIndex: 'query_key_response_time', + width: 130, + render: (_, record) => renderItemResponseTime(record), + }, + { + title: t('渠道数'), + dataIndex: 'channel_count', + width: 100, + }, + { + title: t('已用额度'), + dataIndex: 'used_quota', + width: 160, + render: (quota) => renderQuota(quota || 0), + }, + { + title: t('已用金额'), + dataIndex: 'used_amount', + width: 160, + render: (amount) => renderQuotaWithAmount(amount || 0), + }, + { + title: t('原始额度'), + dataIndex: 'original_amount', + width: 190, + render: (amount, record) => ( + + {renderQuotaWithAmount(amount || 0)} + {record.original_amount_shared ? ( + {t('共享')} + ) : null} + + ), + }, + { + title: t('理论当前额度'), + dataIndex: 'current_amount', + width: 190, + render: (amount) => renderQuotaWithAmount(amount || 0), + }, + { + title: t('超刷金额'), + dataIndex: 'over_brush_amount', + width: 160, + render: (amount) => ( + 0 ? 'danger' : 'secondary'}> + {renderQuotaWithAmount(amount || 0)} + + ), + }, + { + title: t('操作'), + dataIndex: 'query_key_operate', + width: 130, + fixed: 'right', + render: (_, record) => { + const channels = getItemChannels(record); + return ( + + ); + }, + }, + ]; + + const expandedRowRender = (record) => { + const channels = Array.isArray(record.channels) ? record.channels : []; + if (channels.length === 0) { + return ; + } + const channelsWithItem = channels.map((channel) => ({ + ...channel, + __item: record, + })); + return ( +
+ + + `${record.key}-${channel.source || 'channel'}-${channel.id}` + } + pagination={false} + size='small' + scroll={{ x: 2200 }} + style={{ width: '100%' }} + /> + + ); + }; + + return ( +
+
+ + {t('批量密钥报告')} + + + {t('隐藏管理员页面，用于按密钥生成渠道用量与超刷报告。')} + +
+ + +
+ } + closeIcon={null} + description={t( + '每行一个渠道密钥，最多支持 10000 个唯一密钥。报告会匹配多密钥渠道，但不会展示任何渠道内的原始密钥。', + )} + /> +