Files
LangBot/tests/integration/persistence/test_migrations.py
huanghuoguoguo 9ecb587ac0 refactor(provider): use LiteLLM as unified LLM requester backend (#2150)
* refactor(provider): use LiteLLM as unified LLM requester backend

  - Replace 23+ individual requester implementations with unified litellmchat.py
  - Add litellm_provider field to 27 YAML manifests for provider routing
  - Delete redundant requester subclasses
  - Add unit tests for LiteLLMRequester (29 tests)
  - Fix num_retries parameter name (was max_retries)
  - Fix exception handling order for subclass exceptions

  LiteLLM provides unified API for 100+ providers, eliminating need for
  provider-specific requesters.

* fix: ruff format provider.py

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* refactor(provider): simplify LiteLLM requester usage handling

  - Remove unused Anthropic-specific tool schema generation
  - Share completion argument construction between normal and streaming calls
  - Use LiteLLM/OpenAI native usage fields for monitoring
  - Collect stream token usage from LiteLLM stream_options
  - Update LiteLLM requester tests for unified usage fields

* restore: restore deleted provider requester files

Restore individual provider requester implementations that were
removed in de61b5d3. These files coexist with the unified
litellmchat.py backend.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

* feat: update requesters and improve provider selection UI

- Added `litellm_provider` field to various requesters' YAML configurations.
- Removed obsolete Python requester files for OpenRouter, PPIO, QHAIGC, ShengSuanYun, SiliconFlow, Space, TokenPony, VolcArk, and Xai.
- Introduced new requesters for Tencent and Together AI with corresponding YAML configurations and SVG icons.
- Enhanced the ProviderForm component to include a searchable dropdown for selecting providers, improving user experience.
- Updated localization files to include search provider text for both English and Chinese.

* fix(provider): align litellm rebase with master

* fix(provider): capture streaming token usage; add token observability

The LiteLLM streaming requester only captured usage when a chunk had an
empty `choices` list. Many OpenAI-compatible gateways (e.g. new-api) and
providers send the final usage payload in a chunk that still carries an
empty-delta choice, so streamed calls always recorded 0 tokens in the
monitoring logs/dashboard (non-streaming worked).

- Capture stream usage whenever a chunk carries it, regardless of choices
- Add robust _normalize_usage (dict/obj shapes, derive missing total_tokens)
- Register litellm in bootutils/deps.py (was in pyproject only)
- Add MonitoringService.get_token_statistics + /monitoring/token-statistics
  endpoint: summary, per-model breakdown, token timeseries, and a
  zero-token-success data-quality signal
- Add TokenMonitoring dashboard tab (summary tiles, stacked token chart,
  per-model table) + i18n (en/zh)
- Regression tests for stream usage capture and usage normalization

Verified end-to-end against a real OpenAI-compatible endpoint with
gpt-5.5 and claude-opus-4-8: tokens now recorded non-zero for both
streaming and non-streaming paths.

* refactor(provider): simplify litellm capabilities

* style: simplify wrapped expressions

* feat(models): persist context metadata

* fix(provider): handle dict embeddings and openai-compatible rerank in LiteLLMRequester

- invoke_embedding: support both object- and dict-shaped response.data
  entries (OpenAI-compatible gateways like new-api return dicts)
- invoke_rerank: litellm.arerank rejects the 'openai' provider, so for
  openai-compatible (or unspecified) providers call the standard
  Jina/Cohere-style POST /v1/rerank endpoint directly over HTTP
- accept both 'relevance_score' and 'score' fields in rerank results
- add unit tests for the openai-compatible HTTP rerank path

* feat(provider): enforce requester support_type when adding models

- frontend: AddModelPopover only shows model-type tabs (llm/embedding/
  rerank) that the provider's requester declares in its manifest
  support_type; ModelsDialog fetches requester manifests and maps
  requester -> support_type, passed down through ProviderCard
- backend: add _validate_provider_supports guard in create_llm_model /
  create_embedding_model / create_rerank_model so a model cannot be
  attached to a provider whose requester does not support that type,
  even if the frontend restriction is bypassed (manifests without
  support_type are allowed for backward compatibility)
- manifests: correct support_type for providers that do not offer all
  three model types:
  - llm only: anthropic, deepseek, groq, moonshot, openrouter, xai
  - llm + text-embedding: openai, gemini, mistral
  - add rerank to new-api (verified working via /v1/rerank)
  - set llm + text-embedding + rerank for aggregator/unknown gateways

* feat(provider): add searchable alias to requester manifests

- add a free-text 'alias' field to every requester manifest spec,
  containing the vendor's English/Chinese names, pinyin, common
  nicknames and flagship model-series names (e.g. moonshot -> kimi,
  月之暗面; zhipu -> glm, 智谱清言)
- frontend: ProviderForm requester search now also matches against
  alias (substring/contains), so searching 'kimi' surfaces Moonshot,
  '硅基' surfaces SiliconFlow, etc.
- also fix support_type: openrouter (relay) supports embedding+rerank;
  LangBot Space gains rerank (coming soon)

* fix(provider): make support_type guard defensive against incomplete model_mgr

- _validate_provider_supports now uses getattr to gracefully skip when
  model_mgr / provider_dict / manifest lookup is unavailable, instead of
  raising AttributeError (fixes unit tests that mock ap.model_mgr as a
  bare SimpleNamespace)
- add TestValidateProviderSupports covering: allow supported type,
  reject unsupported type, allow when support_type missing, allow when
  provider unknown, degrade safely when model_mgr is incomplete

* fix(persistence): guard 0004 migration against missing llm_models table

The 0004_add_llm_model_context_length migration called
inspector.get_columns('llm_models') unconditionally, raising
NoSuchTableError when the table does not exist (e.g. migrating a
fresh/empty DB, as exercised by the integration tests where
create_all() registers no tables because the ORM models are not
imported). Every other migration guards with a table-existence check
first; add the same guard here for both upgrade and downgrade.

Also restore the test head assertion to 0004 (it had been lowered to
0003 to mask this failure).

* Merge branch 'master' into feat/litellm

Resolve conflicts:
- uv.lock: regenerated via 'uv lock' to reconcile litellm/fastuuid
  (ours) with openai bump (master).
- Alembic migrations: master added 0004_add_mcp_readme while this
  branch added 0004_add_llm_model_context_length, both as children of
  0003 (would create multiple heads). Re-chain the litellm migration as
  0005_add_llm_model_context_length with down_revision=0004_add_mcp_readme
  for a single linear head. Update test head assertion accordingly.

* fix(persistence): shorten migration revision id to fit varchar(32)

PostgreSQL stores alembic_version.version_num as varchar(32).
'0005_add_llm_model_context_length' (33 chars) overflowed it, raising
StringDataRightTruncationError in the PG migration tests. Rename the
revision (and file) to '0005_add_llm_context_length' (27 chars) and
update the head assertions in both SQLite and PostgreSQL migration
tests.

---------

Co-authored-by: Claude Opus 4.7 <noreply@anthropic.com>
Co-authored-by: fdc310 <2213070223@qq.com>
Co-authored-by: RockChinQ <rockchinq@gmail.com>
2026-06-13 16:59:48 +08:00

251 lines
8.4 KiB
Python

"""
SQLite migration integration tests.
Tests real Alembic migration behavior using temporary SQLite databases.
Validates the migration workflow from .github/workflows/test-migrations.yml.
Run: uv run pytest tests/integration/persistence/test_migrations.py -q
"""
from __future__ import annotations
import pytest
from sqlalchemy.ext.asyncio import create_async_engine
from langbot.pkg.entity.persistence.base import Base
from langbot.pkg.persistence.alembic_runner import (
run_alembic_upgrade,
run_alembic_stamp,
get_alembic_current,
)
pytestmark = pytest.mark.integration
@pytest.fixture
def sqlite_db_url(tmp_path):
"""Create SQLite URL with temporary database file."""
db_file = tmp_path / "test_migrations.db"
return f"sqlite+aiosqlite:///{db_file}"
@pytest.fixture
async def sqlite_engine(sqlite_db_url):
"""Create async SQLite engine."""
engine = create_async_engine(sqlite_db_url)
yield engine
await engine.dispose()
class TestSQLiteMigrationBaseline:
"""Tests for baseline stamp workflow."""
@pytest.mark.asyncio
async def test_baseline_stamp_sets_revision(self, sqlite_engine):
"""
Stamp baseline on existing tables sets correct revision.
Workflow:
1. Create tables via Base.metadata.create_all
2. Stamp with '0001_baseline'
3. Verify current revision is '0001_baseline'
"""
# Create all tables (simulates existing DB created by ORM)
async with sqlite_engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
# Stamp baseline
await run_alembic_stamp(sqlite_engine, '0001_baseline')
# Verify revision
rev = await get_alembic_current(sqlite_engine)
assert rev == '0001_baseline', f"Expected '0001_baseline', got {rev}"
@pytest.mark.asyncio
async def test_baseline_stamp_on_empty_db(self, sqlite_engine):
"""
Stamp on empty database (no tables) still sets revision.
This is an edge case - stamping without tables.
"""
# Don't create tables - stamp directly
await run_alembic_stamp(sqlite_engine, '0001_baseline')
rev = await get_alembic_current(sqlite_engine)
assert rev == '0001_baseline'
class TestSQLiteMigrationUpgrade:
"""Tests for upgrade to head workflow."""
@pytest.mark.asyncio
async def test_upgrade_from_baseline_to_head(self, sqlite_engine):
"""
Upgrade from baseline to head applies all migrations.
Workflow:
1. Create tables
2. Stamp baseline
3. Upgrade to head
4. Verify current revision is head
"""
# Create tables
async with sqlite_engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
# Stamp baseline
await run_alembic_stamp(sqlite_engine, '0001_baseline')
# Upgrade to head
await run_alembic_upgrade(sqlite_engine, 'head')
# Verify revision
rev = await get_alembic_current(sqlite_engine)
assert rev is not None, "Expected a revision after upgrade"
# Head should be the latest migration
assert rev.startswith('0005'), f"Expected head to be 0005_*, got {rev}"
@pytest.mark.asyncio
async def test_upgrade_idempotent(self, sqlite_engine):
"""
Running upgrade to head multiple times is idempotent.
Workflow:
1. Upgrade to head
2. Get revision
3. Upgrade to head again
4. Verify same revision
"""
# Create tables
async with sqlite_engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
# Stamp and upgrade
await run_alembic_stamp(sqlite_engine, '0001_baseline')
await run_alembic_upgrade(sqlite_engine, 'head')
rev1 = await get_alembic_current(sqlite_engine)
# Upgrade again - should be idempotent
await run_alembic_upgrade(sqlite_engine, 'head')
rev2 = await get_alembic_current(sqlite_engine)
assert rev2 == rev1, f"Expected {rev1}, got {rev2}"
class TestSQLiteMigrationFreshDatabase:
"""Tests for fresh database workflow."""
@pytest.mark.asyncio
async def test_fresh_db_upgrade_from_scratch(self, tmp_path):
"""
Fresh database (no tables) can be upgraded directly to head.
Workflow:
1. Create fresh engine with new DB file
2. Create tables
3. Upgrade to head
4. Verify revision
"""
# Use different DB file for fresh test
fresh_db_file = tmp_path / "test_migrations_fresh.db"
fresh_url = f"sqlite+aiosqlite:///{fresh_db_file}"
fresh_engine = create_async_engine(fresh_url)
# Create tables on fresh DB
async with fresh_engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
# Upgrade to head directly (no baseline stamp)
await run_alembic_upgrade(fresh_engine, 'head')
# Verify revision
rev = await get_alembic_current(fresh_engine)
assert rev is not None, "Expected a revision on fresh DB"
await fresh_engine.dispose()
@pytest.mark.asyncio
async def test_fresh_db_without_create_all_behavior(self, tmp_path):
"""
Fresh database without create_all - test actual behavior.
This tests what happens when migrations run on truly empty DB.
The behavior is determined by Alembic and migration scripts.
EXPECTED: Either:
1. Migration succeeds (if scripts handle empty DB)
2. Migration fails with specific error (if scripts require tables)
IMPORTANT: This test verifies the ACTUAL behavior, not accepting
any arbitrary failure with try-except pass.
"""
fresh_db_file = tmp_path / "test_empty_migrations.db"
fresh_url = f"sqlite+aiosqlite:///{fresh_db_file}"
fresh_engine = create_async_engine(fresh_url)
# Capture the actual behavior
actual_result = None
actual_error = None
try:
await run_alembic_upgrade(fresh_engine, 'head')
rev = await get_alembic_current(fresh_engine)
actual_result = rev
except Exception as e:
actual_error = e
await fresh_engine.dispose()
# Verify specific behavior - one of two outcomes is expected
if actual_result is not None:
# Migration succeeded - verify revision exists
assert actual_result is not None, "Revision should exist after successful migration"
else:
# Migration failed - verify the error type is known
# Alembic typically raises specific errors for missing tables
assert actual_error is not None, "Error should be captured if migration failed"
# Log the error type for documentation (don't silently pass)
error_type = type(actual_error).__name__
# Acceptable error types for empty DB scenarios
acceptable_errors = [
'OperationalError', # SQLite table not found
'ProgrammingError', # SQLAlchemy errors
'CommandError', # Alembic command errors
]
assert error_type in acceptable_errors, (
f"Unexpected error type: {error_type}. "
f"This may indicate a regression in migration behavior. "
f"Error: {actual_error}"
)
class TestSQLiteMigrationGetCurrent:
"""Tests for get_alembic_current behavior."""
@pytest.mark.asyncio
async def test_get_current_on_unstamped_db_returns_none(self, sqlite_engine):
"""
get_alembic_current returns None for unstamped database.
"""
# Create tables but don't stamp
async with sqlite_engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
# No stamp - should return None
rev = await get_alembic_current(sqlite_engine)
assert rev is None, f"Expected None for unstamped DB, got {rev}"
@pytest.mark.asyncio
async def test_get_current_after_stamp_returns_revision(self, sqlite_engine):
"""
get_alembic_current returns correct revision after stamp.
"""
async with sqlite_engine.begin() as conn:
await conn.run_sync(Base.metadata.create_all)
await run_alembic_stamp(sqlite_engine, '0001_baseline')
rev = await get_alembic_current(sqlite_engine)
assert rev == '0001_baseline'