Refine resume parts and fix bugs in retrival using sql (#66)

api/apps/llm_app.py

@@ -21,11 +21,12 @@ from api.db.services.llm_service import LLMFactoriesService, TenantLLMService, L
 from api.db.services.user_service import TenantService, UserTenantService
 from api.utils.api_utils import server_error_response, get_data_error_result, validate_request
 from api.utils import get_uuid, get_format_time
-from api.db import StatusEnum, UserTenantRole
+from api.db import StatusEnum, UserTenantRole, LLMType
 from api.db.services.knowledgebase_service import KnowledgebaseService
 from api.db.db_models import Knowledgebase, TenantLLM
 from api.settings import stat_logger, RetCode
 from api.utils.api_utils import get_json_result
+from rag.llm import EmbeddingModel, CvModel, ChatModel
 
 
 @manager.route('/factories', methods=['GET'])
@@ -43,16 +44,37 @@ def factories():
 @validate_request("llm_factory", "api_key")
 def set_api_key():
     req = request.json
+    # test if api key works
+    msg = ""
+    for llm in LLMService.query(fid=req["llm_factory"]):
+        if llm.model_type == LLMType.EMBEDDING.value:
+            mdl = EmbeddingModel[req["llm_factory"]](
+                req["api_key"], llm.llm_name)
+            try:
+                arr, tc = mdl.encode(["Test if the api key is available"])
+                if len(arr[0]) == 0 or tc ==0: raise Exception("Fail")
+            except Exception as e:
+                msg += f"\nFail to access embedding model({llm.llm_name}) using this api key."
+        elif llm.model_type == LLMType.CHAT.value:
+            mdl = ChatModel[req["llm_factory"]](
+                req["api_key"], llm.llm_name)
+            try:
+                m, tc = mdl.chat(None, [{"role": "user", "content": "Hello! How are you doing!"}], {"temperature": 0.9})
+                if not tc: raise Exception(m)
+            except Exception as e:
+                msg += f"\nFail to access model({llm.llm_name}) using this api key." + str(e)
+
+    if msg: return get_data_error_result(retmsg=msg)
+
     llm = {
         "tenant_id": current_user.id,
         "llm_factory": req["llm_factory"],
         "api_key": req["api_key"]
     }
-    # TODO: Test api_key
     for n in ["model_type", "llm_name"]:
         if n in req: llm[n] = req[n]
 
-    TenantLLM.insert(**llm).on_conflict("replace").execute()
+    TenantLLMService.filter_update([TenantLLM.tenant_id==llm["tenant_id"], TenantLLM.llm_factory==llm["llm_factory"]], llm)
     return get_json_result(data=True)
 
 
@@ -69,6 +91,7 @@ def my_llms():
 @manager.route('/list', methods=['GET'])
 @login_required
 def list():
+    model_type = request.args.get("model_type")
     try:
         objs = TenantLLMService.query(tenant_id=current_user.id)
         mdlnms = set([o.to_dict()["llm_name"] for o in objs if o.api_key])
@@ -79,6 +102,7 @@ def list():
 
         res = {}
         for m in llms:
+            if model_type and m["model_type"] != model_type: continue
             if m["fid"] not in res: res[m["fid"]] = []
             res[m["fid"]].append(m)