100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
from datetime import datetime, timezone
|
|
|
|
import httpx
|
|
from fastapi import APIRouter, Depends
|
|
from sqlmodel import Session, select
|
|
|
|
from ..database import get_session
|
|
from ..models import AIConfig
|
|
from ..schemas import AIConfigSave, AIConfigRead
|
|
|
|
router = APIRouter(prefix="/api/ai", tags=["ai-config"])
|
|
|
|
|
|
def _mask_key(key: str) -> str:
|
|
if len(key) <= 8:
|
|
return "****"
|
|
return key[:4] + "****" + key[-4:]
|
|
|
|
|
|
def _get_or_create_config(session: Session) -> AIConfig:
|
|
config = session.exec(select(AIConfig)).first()
|
|
if not config:
|
|
config = AIConfig()
|
|
session.add(config)
|
|
session.commit()
|
|
session.refresh(config)
|
|
return config
|
|
|
|
|
|
def _config_to_read(config: AIConfig) -> AIConfigRead:
|
|
return AIConfigRead(
|
|
id=config.id,
|
|
provider=config.provider,
|
|
base_url=config.base_url,
|
|
api_key_masked=_mask_key(config.api_key) if config.api_key else "",
|
|
model=config.model,
|
|
updated_at=config.updated_at,
|
|
)
|
|
|
|
|
|
@router.get("/config", response_model=AIConfigRead)
|
|
def get_config(session: Session = Depends(get_session)):
|
|
config = _get_or_create_config(session)
|
|
return _config_to_read(config)
|
|
|
|
|
|
@router.put("/config", response_model=AIConfigRead)
|
|
def save_config(data: AIConfigSave, session: Session = Depends(get_session)):
|
|
config = _get_or_create_config(session)
|
|
config.provider = data.provider
|
|
config.base_url = data.base_url
|
|
config.api_key = data.api_key
|
|
config.model = data.model
|
|
config.updated_at = datetime.now(timezone.utc)
|
|
session.add(config)
|
|
session.commit()
|
|
session.refresh(config)
|
|
return _config_to_read(config)
|
|
|
|
|
|
@router.post("/config/test")
|
|
async def test_config(data: AIConfigSave):
|
|
"""测试AI连接是否有效"""
|
|
try:
|
|
async with httpx.AsyncClient(timeout=15) as client:
|
|
if data.provider == "anthropic":
|
|
resp = await client.post(
|
|
f"{data.base_url.rstrip('/')}/messages",
|
|
headers={
|
|
"x-api-key": data.api_key,
|
|
"anthropic-version": "2023-06-01",
|
|
"content-type": "application/json",
|
|
},
|
|
json={
|
|
"model": data.model,
|
|
"max_tokens": 10,
|
|
"messages": [{"role": "user", "content": "Hi"}],
|
|
},
|
|
)
|
|
else:
|
|
resp = await client.post(
|
|
f"{data.base_url.rstrip('/')}/chat/completions",
|
|
headers={
|
|
"Authorization": f"Bearer {data.api_key}",
|
|
"Content-Type": "application/json",
|
|
},
|
|
json={
|
|
"model": data.model,
|
|
"max_tokens": 10,
|
|
"messages": [{"role": "user", "content": "Hi"}],
|
|
},
|
|
)
|
|
if resp.status_code < 400:
|
|
return {"ok": True, "message": "连接成功"}
|
|
return {"ok": False, "message": f"API返回错误: {resp.status_code}"}
|
|
except httpx.ConnectError:
|
|
return {"ok": False, "message": "无法连接到API地址"}
|
|
except Exception as e:
|
|
return {"ok": False, "message": str(e)}
|