add cohere api

app.py

@@ -17,12 +17,14 @@ set('project_root', parent_dir)
 # from routers.webtools_v1 import router as webtools_router
 from routers.users_v1 import router as users_router
 from routers.openai_v1_1 import router as openai_router
+from routers.cohere import router as cohere_router
 
 app = FastAPI()
 
 # app.include_router(webtools_router, prefix="/airs/v1", tags=["webtools"])
 app.include_router(users_router, prefix="/airs/v1", tags=["users"])
 app.include_router(openai_router, prefix="/airs/v1", tags=["openai"])
+app.include_router(cohere_router, prefix="/airs/cohere/v2", tags=["cohere"])
 
 @app.get("/")
 def greet_json():

routers/cohere.py

@@ -0,0 +1,105 @@
+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"
+
+print('db_module_filename',db_module_filename)
+
+@router.post("/embed")
+async def embed(request:Request, current_user: UserModel = Depends(get_current_user)):
+    data = await request.json()
+    print(data)
+    model = data.get('model', '')  
+    print(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', '')
+    print(api_key_info)
+
+    headers = {
+        'Content-Type': 'application/json',
+        'Authorization': f'Bearer {api_key}',
+        'User-Agent': 'PostmanRuntime/7.43.0'
+    }
+
+    print(f'base_url: {base_url}')
+    
+    try:
+        response = requests.post(url=f"{base_url}/embed", headers=headers, json=data)
+        print(response.json())
+        return response.json()
+    except Exception as e:
+        print(e)
+        return {'error': e}
+
+
+@router.post("/rerank")
+async def rerank(request:Request, current_user: UserModel = Depends(get_current_user)):
+    data = await request.json()
+    print(data)
+    model = data.get('model', '')  
+    print(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', '')
+    print(api_key_info)
+
+    headers = {
+        'Content-Type': 'application/json',
+        'Authorization': f'Bearer {api_key}',
+        'User-Agent': 'PostmanRuntime/7.43.0'
+    }
+
+    print(f'base_url: {base_url}')
+    
+    try:
+        response = requests.post(url=f"{base_url}/rerank", headers=headers, json=data)
+        print(response.json())
+        return response.json()
+    except Exception as e:
+        print(e)
+        return {'error': e}
+
+# # 从数据库获取默认模型
+# async def get_default_model():
+#     query = f"SELECT * FROM api_names order by default_order limit 1"
+#     response = TbsDb(db_module_filename, "Cloudflare").get_item(query)
+#     try:
+#         result = response['result'][0]['results'][0]['api_name']
+#     except:
+#         result = ''
+#     return result
+
+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

routers/openai_v1_1.py

@@ -26,6 +26,8 @@ async def chat_completions(request:Request, current_user: UserModel = Depends(ge
         'Authorization': f'Bearer {api_key}',
         'User-Agent': 'PostmanRuntime/7.43.0'
     }
+
+    print(f'base_url: {base_url}')
     
     try:
         response = requests.post(url=f"{base_url}/chat/completions", headers=headers, json=data)