routers/cohere.py

from fastapi import APIRouter, Request, Depends
import requests

from global_state import get
from db_model.user import UserModel
from auth import get_current_user
from db.tbs_db import TbsDb

router = APIRouter()

db_module_filename = f"{get('project_root')}/db/cloudflare.py"

@router.post("/embed")
async def embed(request:Request, current_user: UserModel = Depends(get_current_user)):
    data = await request.json()
    model = data.get('model', '')  
    if (model=='')or(model is None):
        return {'error': 'model is empty'}
    api_key_info = await get_api_key(model)
    api_key = api_key_info.get('api_key', '')
    base_url = api_key_info.get('base_url', '')

    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {api_key}',
        'User-Agent': 'PostmanRuntime/7.43.0'
    }
    
    try:
        response = requests.post(url=f"{base_url}/embed", headers=headers, json=data)
        return response.json()
    except Exception as e:
        return {'error': e}


@router.post("/rerank")
async def rerank(request:Request, current_user: UserModel = Depends(get_current_user)):
    data = await request.json()
    model = data.get('model', '')  
    if (model=='')or(model is None):
        return {'error': 'model is empty'}
    api_key_info = await get_api_key(model)
    api_key = api_key_info.get('api_key', '')
    base_url = api_key_info.get('base_url', '')

    headers = {
        'Content-Type': 'application/json',
        'Authorization': f'Bearer {api_key}',
        'User-Agent': 'PostmanRuntime/7.43.0'
    }
    
    try:
        response = requests.post(url=f"{base_url}/rerank", headers=headers, json=data)
        return response.json()
    except Exception as e:
        return {'error': e}

async def get_api_key(model):
    query = f"""
SELECT an.api_name, ak.api_key, an.base_url, ag.group_name
FROM api_keys ak
JOIN api_groups ag ON ak.api_group_id = ag.id
JOIN api_names an ON an.api_group_id = ag.id
WHERE an.api_name='{model}' and disabled=0
ORDER BY ak.last_call_at
limit 1
"""
# WHERE ak.category='LLM' and an.api_name='{model}' and disabled=0

    response = TbsDb(db_module_filename, "Cloudflare").get_item(query)
    try:
        result = response['result'][0]['results'][0]
        api_key = result['api_key']
    except:
        api_key = ''
    
    query = f"update api_keys set last_call_at=datetime('now') where api_key='{api_key}'"
    TbsDb(db_module_filename, "Cloudflare").execute_query(query)
    return result