122 lines
3.9 KiB
Python
122 lines
3.9 KiB
Python
"""P6 种子：OpenAI provider + 2 个模型 + 默认路由规则。
|
|
|
|
API Key 从 `.env` 的 `OPENAI_API_KEY` 读取；未配置也能建 provider（走 mock）。
|
|
执行：python -m app.scripts.seed_p6
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
|
|
from sqlalchemy import delete, insert, select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.config import settings
|
|
from app.core.crypto import encrypt
|
|
from app.db.session import AsyncSessionLocal
|
|
from app.models import LLMModel, LLMProvider, LLMRoutingRule
|
|
|
|
|
|
DEFAULT_PROVIDER = {
|
|
"code": "openai",
|
|
"name": "OpenAI",
|
|
"vendor": "openai",
|
|
"api_base": settings.OPENAI_API_BASE or "https://api.openai.com/v1",
|
|
"description": "OpenAI 官方；联调阶段如 .env 未配 OPENAI_API_KEY，调用将走 mock 回显",
|
|
}
|
|
|
|
DEFAULT_MODELS = [
|
|
{
|
|
"model": "gpt-4o-mini",
|
|
"name": "GPT-4o mini",
|
|
"description": "通用场景推荐；响应快、成本低",
|
|
"max_context": 128_000,
|
|
"price_input_per_1k": 0.00015,
|
|
"price_output_per_1k": 0.00060,
|
|
},
|
|
{
|
|
"model": "gpt-4o",
|
|
"name": "GPT-4o",
|
|
"description": "高复杂度推理",
|
|
"max_context": 128_000,
|
|
"price_input_per_1k": 0.0025,
|
|
"price_output_per_1k": 0.01,
|
|
},
|
|
]
|
|
|
|
|
|
async def upsert(db: AsyncSession, model, key_field: str, value, **defaults):
|
|
stmt = select(model).where(getattr(model, key_field) == value)
|
|
obj = (await db.execute(stmt)).scalar_one_or_none()
|
|
if obj is None:
|
|
obj = model(**{key_field: value}, **defaults)
|
|
db.add(obj)
|
|
await db.flush()
|
|
else:
|
|
for k, v in defaults.items():
|
|
setattr(obj, k, v)
|
|
return obj
|
|
|
|
|
|
async def seed() -> None:
|
|
async with AsyncSessionLocal() as db:
|
|
prov = await upsert(
|
|
db, LLMProvider, "code", DEFAULT_PROVIDER["code"],
|
|
name=DEFAULT_PROVIDER["name"],
|
|
vendor=DEFAULT_PROVIDER["vendor"],
|
|
api_base=DEFAULT_PROVIDER["api_base"],
|
|
description=DEFAULT_PROVIDER["description"],
|
|
)
|
|
if settings.OPENAI_API_KEY:
|
|
prov.api_key_cipher = encrypt(settings.OPENAI_API_KEY)
|
|
await db.flush()
|
|
|
|
models: dict[str, LLMModel] = {}
|
|
for m in DEFAULT_MODELS:
|
|
stmt = select(LLMModel).where(
|
|
LLMModel.provider_id == prov.id, LLMModel.model == m["model"]
|
|
)
|
|
obj = (await db.execute(stmt)).scalar_one_or_none()
|
|
if obj is None:
|
|
obj = LLMModel(provider_id=prov.id, **m)
|
|
db.add(obj)
|
|
await db.flush()
|
|
else:
|
|
for k, v in m.items():
|
|
setattr(obj, k, v)
|
|
models[m["model"]] = obj
|
|
|
|
# 默认路由规则：无条件兜底 → gpt-4o-mini
|
|
await db.execute(
|
|
delete(LLMRoutingRule).where(LLMRoutingRule.name == "default")
|
|
)
|
|
db.add(LLMRoutingRule(
|
|
name="default",
|
|
priority=1000,
|
|
conditions=None,
|
|
model_id=models["gpt-4o-mini"].id,
|
|
enabled=True,
|
|
description="全局兜底路由",
|
|
))
|
|
# 高质量：category=ad 或 reasoning 命中 gpt-4o
|
|
await db.execute(
|
|
delete(LLMRoutingRule).where(LLMRoutingRule.name == "premium")
|
|
)
|
|
db.add(LLMRoutingRule(
|
|
name="premium",
|
|
priority=10,
|
|
conditions={"category": "reasoning"},
|
|
model_id=models["gpt-4o"].id,
|
|
enabled=True,
|
|
description="高复杂度推理使用 gpt-4o",
|
|
))
|
|
|
|
await db.commit()
|
|
print("[seed_p6] done.")
|
|
print(f" provider: {prov.code} (api_key_configured={bool(settings.OPENAI_API_KEY)})")
|
|
print(f" models: {list(models.keys())}")
|
|
print(f" rules: default(gpt-4o-mini) / premium(category=reasoning → gpt-4o)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(seed())
|