zwt13703
115 lines
3.5 KiB
Python
115 lines
3.5 KiB
Python
from fastapi import APIRouter, Depends, HTTPException
|
|
from sqlalchemy import select
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from database import get_db
|
|
from models.ai_model import AiModel
|
|
from schemas.schemas import AiModelCreate, AiModelUpdate, Response
|
|
from services.ai_service import call_ai
|
|
from services.security import decrypt_text, encrypt_text, mask_secret
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
def _serialize_model(model: AiModel) -> dict:
|
|
api_key = decrypt_text(model.api_key_encrypted)
|
|
return {
|
|
"id": model.id,
|
|
"name": model.name,
|
|
"provider": model.provider,
|
|
"api_format": model.api_format,
|
|
"api_endpoint": model.api_endpoint,
|
|
"api_key_preview": mask_secret(api_key),
|
|
"status": model.status,
|
|
"created_at": model.created_at,
|
|
}
|
|
|
|
|
|
@router.get("")
|
|
async def list_models(db: AsyncSession = Depends(get_db)):
|
|
result = await db.execute(select(AiModel).order_by(AiModel.id.desc()))
|
|
items = [_serialize_model(item) for item in result.scalars().all()]
|
|
return Response(data=items)
|
|
|
|
|
|
@router.post("")
|
|
async def create_model(body: AiModelCreate, db: AsyncSession = Depends(get_db)):
|
|
model = AiModel(
|
|
name=body.name,
|
|
provider=body.provider,
|
|
api_format=body.api_format,
|
|
api_endpoint=body.api_endpoint,
|
|
api_key_encrypted=encrypt_text(body.api_key),
|
|
status=body.status,
|
|
)
|
|
db.add(model)
|
|
await db.commit()
|
|
await db.refresh(model)
|
|
return Response(data=_serialize_model(model))
|
|
|
|
|
|
@router.put("/{model_id}")
|
|
async def update_model(model_id: int, body: AiModelUpdate, db: AsyncSession = Depends(get_db)):
|
|
model = await db.get(AiModel, model_id)
|
|
if model is None:
|
|
raise HTTPException(status_code=404, detail="模型不存在")
|
|
|
|
if body.name is not None:
|
|
model.name = body.name
|
|
if body.provider is not None:
|
|
model.provider = body.provider
|
|
if body.api_format is not None:
|
|
model.api_format = body.api_format
|
|
if body.api_endpoint is not None:
|
|
model.api_endpoint = body.api_endpoint
|
|
if body.status is not None:
|
|
model.status = body.status
|
|
if body.api_key:
|
|
model.api_key_encrypted = encrypt_text(body.api_key)
|
|
|
|
await db.commit()
|
|
await db.refresh(model)
|
|
return Response(data=_serialize_model(model))
|
|
|
|
|
|
@router.delete("/{model_id}")
|
|
async def delete_model(model_id: int, db: AsyncSession = Depends(get_db)):
|
|
model = await db.get(AiModel, model_id)
|
|
if model is None:
|
|
raise HTTPException(status_code=404, detail="模型不存在")
|
|
|
|
await db.delete(model)
|
|
await db.commit()
|
|
return Response(data={"id": model_id})
|
|
|
|
|
|
@router.post("/{model_id}/test")
|
|
async def test_model(model_id: int, db: AsyncSession = Depends(get_db)):
|
|
model = await db.get(AiModel, model_id)
|
|
if model is None:
|
|
raise HTTPException(status_code=404, detail="模型不存在")
|
|
|
|
class FakeParagraph:
|
|
title = "连接测试"
|
|
content = "请返回一段非常简短的测试文本。"
|
|
need_prompt = False
|
|
prompt_text = ""
|
|
output_format = "text"
|
|
|
|
try:
|
|
result = await call_ai(FakeParagraph(), model)
|
|
return Response(
|
|
data={
|
|
"id": model.id,
|
|
"success": True,
|
|
"message": f"模型 {model.name} 连接测试成功",
|
|
"preview": result.content,
|
|
}
|
|
)
|
|
except Exception as error:
|
|
return Response(
|
|
code=-1,
|
|
message=str(error),
|
|
data={"id": model.id, "success": False},
|
|
)
|