AI_Watch_Platform/backend/app/api/endpoints/models_api.py

from typing import Any, List
from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.orm import Session
from backend.app.api import deps
from backend.app.core.database import get_db
from backend.app.models import sql_models
from backend.app.schemas import schemas
from backend.app.services.llm_agent import test_connection

router = APIRouter()

@router.get("", response_model=List[schemas.ModelConfig])
def read_models(
    db: Session = Depends(get_db),
    skip: int = 0,
    limit: int = 100,
    current_user: sql_models.User = Depends(deps.get_current_user)
) -> Any:
    return db.query(sql_models.ModelConfig).offset(skip).limit(limit).all()

@router.post("", response_model=schemas.ModelConfig)
def create_model(
    *,
    db: Session = Depends(get_db),
    model_in: schemas.ModelConfigCreate,
    current_user: sql_models.User = Depends(deps.get_current_user)
) -> Any:
    model = sql_models.ModelConfig(**model_in.model_dump())
    db.add(model)
    db.commit()
    db.refresh(model)
    return model

@router.put("/{id}", response_model=schemas.ModelConfig)
def update_model(
    *,
    db: Session = Depends(get_db),
    id: int,
    model_in: schemas.ModelConfigCreate,
    current_user: sql_models.User = Depends(deps.get_current_user)
) -> Any:
    model = db.query(sql_models.ModelConfig).filter(sql_models.ModelConfig.id == id).first()
    if not model:
        raise HTTPException(status_code=404, detail="Model config not found")
    
    for key, value in model_in.model_dump().items():
        setattr(model, key, value)
    
    db.commit()
    db.refresh(model)
    return model

@router.delete("/{id}", response_model=schemas.ModelConfig)
def delete_model(
    *,
    db: Session = Depends(get_db),
    id: int,
    current_user: sql_models.User = Depends(deps.get_current_user)
) -> Any:
    model = db.query(sql_models.ModelConfig).filter(sql_models.ModelConfig.id == id).first()
    if not model:
        raise HTTPException(status_code=404, detail="Model config not found")
    
    db.delete(model)
    db.commit()
    return model

@router.post("/test")
async def test_model_config(
    *,
    model_in: schemas.ModelConfigCreate,
    current_user: sql_models.User = Depends(deps.get_current_user)
) -> Any:
    try:
        result = await test_connection(
            base_url=model_in.base_url,
            api_key=model_in.api_key,
            model_name=model_in.model_name
        )
        return {"status": "success", "message": result}
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))