feat(app):

- local llm support
- image file without ext support
- v0.3.0

.gitignore

@@ -25,4 +25,3 @@
 
 
 /temp
-/trails

requirements.txt

@@ -17,3 +17,4 @@ jieba>=0.42.1
 starlette>=0.37.2
 jinja2~=3.1.3
 apscheduler~=3.10.4
+cryptography~=41.0.3

setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = retk
-version = 0.2.9
+version = 0.3.0
 author = MorvanZhou
 author_email = [email protected]
 description = keep and reuse your thoughts
@@ -44,6 +44,7 @@ install_requires =
     starlette>=0.27.0
     jinja2~=3.1.3
     apscheduler~=3.10.4
+    cryptography~=41.0.3
 
 [options.packages.find]
 where = src
@@ -54,6 +55,8 @@ retk =
     .env.local
     models/search_engine/*.txt
     plugins/official_plugins/**/*
+    core/ai/llm/knowledge/*.md
+
 [options.extras_require]
 build =
     tox==3.24.3

src/retk/__init__.py

@@ -1,4 +1,4 @@
-from . import version_manager
+from . import local_manager
 from ._version import __version__
 from .core import scheduler
 from .models import tps

src/retk/config.py

@@ -43,6 +43,12 @@ class Settings(BaseSettings):
     COS_BUCKET_NAME: str = Field(env="COS_BUCKET_NAME", default="")
     COS_DOMAIN: Optional[str] = Field(env="COS_DOMAIN", default=None)
 
+    # llm knowledge service settings
+    LLM_KNOWLEDGE_SUMMARY_SERVICE: str = Field(env='LLM_KNOWLEDGE_SUMMARY_SERVICE', default="")
+    LLM_KNOWLEDGE_SUMMARY_MODEL: str = Field(env='LLM_KNOWLEDGE_SUMMARY_MODEL', default="")
+    LLM_KNOWLEDGE_EXTEND_SERVICE: str = Field(env='LLM_KNOWLEDGE_EXTEND_SERVICE', default="")
+    LLM_KNOWLEDGE_EXTEND_MODEL: str = Field(env='LLM_KNOWLEDGE_EXTEND_MODEL', default="")
+
     # hunyuan
     HUNYUAN_SECRET_ID: str = Field(env='HUNYUAN_SECRET_ID', default="")
     HUNYUAN_SECRET_KEY: str = Field(env='HUNYUAN_SECRET_KEY', default="")
@@ -58,7 +64,6 @@ class Settings(BaseSettings):
     OPENAI_API_KEY: str = Field(env='OPENAI_API_KEY', default="")
 
     # xfyun api
-    XFYUN_APP_ID: str = Field(env='XFYUN_APP_ID', default="")
     XFYUN_API_SECRET: str = Field(env='XFYUN_API_SECRET', default="")
     XFYUN_API_KEY: str = Field(env='XFYUN_API_KEY', default="")
 

src/retk/controllers/ai/llm_settings.py

@@ -0,0 +1,141 @@
+from retk import const, local_manager
+from retk.config import is_local_db
+from retk.controllers import schemas
+from retk.controllers.utils import json_exception, maybe_raise_json_exception
+from retk.core.ai.llm.api import LLM_DEFAULT_SERVICES
+from retk.models.tps import AuthedUser
+
+
+async def get_llm_api_settings(
+        au: AuthedUser,
+) -> schemas.ai.LLMApiSettingsResponse:
+    if not is_local_db():
+        raise json_exception(
+            request_id=au.request_id,
+            uid=au.u.id,
+            code=const.CodeEnum.NOT_PERMITTED,
+            language=au.language,
+        )
+    llm_settings = local_manager.llm.get_llm_api_settings()
+    if llm_settings is None:
+        raise json_exception(
+            request_id=au.request_id,
+            uid=au.u.id,
+            code=const.CodeEnum.INVALID_SETTING,
+            language=au.language,
+        )
+    return schemas.ai.LLMApiSettingsResponse(
+        requestId=au.request_id,
+        service=llm_settings.get("service", ""),
+        model=llm_settings.get("model", ""),
+        auth=llm_settings.get("auth", {}),
+    )
+
+
+async def change_llm_api_settings(
+        au: AuthedUser,
+        req: schemas.ai.LLMApiSettingsRequest,
+) -> schemas.ai.LLMApiSettingsResponse:
+    if not is_local_db():
+        raise json_exception(
+            request_id=au.request_id,
+            uid=au.u.id,
+            code=const.CodeEnum.NOT_PERMITTED,
+            language=au.language,
+        )
+    try:
+        local_manager.llm.change_llm_api(
+            llm_service=req.service,
+            llm_model=req.model,
+            llm_api_auth=req.auth,
+        )
+    except ValueError as e:
+        raise json_exception(
+            request_id=au.request_id,
+            uid=au.u.id,
+            code=const.CodeEnum.INVALID_SETTING,
+            language=au.language,
+            log_msg=str(e),
+        )
+    return schemas.ai.LLMApiSettingsResponse(
+        requestId=au.request_id,
+        service=req.service,
+        model=req.model,
+        auth=req.auth,
+    )
+
+
+async def delete_llm_api_settings(
+        au: AuthedUser,
+) -> schemas.RequestIdResponse:
+    if not is_local_db():
+        raise json_exception(
+            request_id=au.request_id,
+            uid=au.u.id,
+            code=const.CodeEnum.NOT_PERMITTED,
+            language=au.language,
+        )
+    local_manager.llm.delete_llm_api()
+    return schemas.RequestIdResponse(
+        requestId=au.request_id,
+    )
+
+
+async def llm_api_test(
+        au: AuthedUser,
+) -> schemas.RequestIdResponse:
+    if not is_local_db():
+        raise json_exception(
+            request_id=au.request_id,
+            uid=au.u.id,
+            code=const.CodeEnum.NOT_PERMITTED,
+            language=au.language,
+        )
+    llm_api = local_manager.llm.get_llm_api_settings()
+    if llm_api is None:
+        raise json_exception(
+            request_id=au.request_id,
+            uid=au.u.id,
+            code=const.CodeEnum.INVALID_SETTING,
+            language=au.language,
+            log_msg="llm_api is None",
+        )
+    service = llm_api.get("service")
+    if service not in LLM_DEFAULT_SERVICES:
+        raise json_exception(
+            request_id=au.request_id,
+            uid=au.u.id,
+            code=const.CodeEnum.INVALID_SETTING,
+            language=au.language,
+            log_msg=f"{service} not in LLM_DEFAULT_SERVICES: {LLM_DEFAULT_SERVICES.keys()}",
+        )
+    model = llm_api.get("model")
+    if model is None:
+        raise json_exception(
+            request_id=au.request_id,
+            uid=au.u.id,
+            code=const.CodeEnum.INVALID_SETTING,
+            language=au.language,
+            log_msg=f"model is None",
+        )
+    auth = llm_api.get("auth")
+    if auth is None or len(auth) == 0:
+        raise json_exception(
+            request_id=au.request_id,
+            uid=au.u.id,
+            code=const.CodeEnum.INVALID_SETTING,
+            language=au.language,
+            log_msg=f"auth is None or empty",
+        )
+    llm_service = LLM_DEFAULT_SERVICES[service]
+    resp, code = await llm_service.complete(
+        messages=[{"role": "user", "content": "hi"}],
+        model=model,
+        req_id=au.request_id,
+    )
+    print(resp)
+    maybe_raise_json_exception(au=au, code=code)
+
+    return schemas.RequestIdResponse(
+        requestId=au.request_id,
+    )

src/retk/controllers/schemas/ai.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Dict
 
 from pydantic import BaseModel
 
@@ -14,3 +14,16 @@ class Node(BaseModel):
 
     requestId: str
     nodes: List[Node]
+
+
+class LLMApiSettingsRequest(BaseModel):
+    service: str
+    model: str
+    auth: Dict[str, str]
+
+
+class LLMApiSettingsResponse(BaseModel):
+    requestId: str
+    service: str
+    model: str
+    auth: Dict[str, str]

src/retk/core/ai/llm/api/__init__.py

@@ -1,6 +1,51 @@
+from typing import Dict
+
 from .aliyun import AliyunService, AliyunModelEnum
 from .baidu import BaiduService, BaiduModelEnum
+from .base import BaseLLMService
 from .moonshot import MoonshotService, MoonshotModelEnum
 from .openai import OpenaiService, OpenaiModelEnum
 from .tencent import TencentService, TencentModelEnum
 from .xfyun import XfYunService, XfYunModelEnum
+
+LLM_SERVICES_CLASS = {
+    AliyunService.name: {
+        "service": AliyunService,
+        "models": AliyunModelEnum,
+    },
+    BaiduService.name: {
+        "service": BaiduService,
+        "models": BaiduModelEnum,
+    },
+    MoonshotService.name: {
+        "service": MoonshotService,
+        "models": MoonshotModelEnum,
+    },
+    OpenaiService.name: {
+        "service": OpenaiService,
+        "models": OpenaiModelEnum,
+    },
+    TencentService.name: {
+        "service": TencentService,
+        "models": TencentModelEnum,
+    },
+    XfYunService.name: {
+        "service": XfYunService,
+        "models": XfYunModelEnum,
+    },
+}
+
+TOP_P = 0.9
+TEMPERATURE = 0.6
+TIMEOUT = 60
+
+LLM_DEFAULT_SERVICES: Dict[str, BaseLLMService] = {
+    s.name: s(top_p=TOP_P, temperature=TEMPERATURE, timeout=TIMEOUT) for s in [
+        TencentService,
+        AliyunService,
+        OpenaiService,
+        MoonshotService,
+        XfYunService,
+        BaiduService,
+    ]
+}

src/retk/core/ai/llm/api/aliyun.py

@@ -10,6 +10,7 @@
 
 
 # https://help.aliyun.com/zh/dashscope/developer-reference/tongyi-thousand-questions-metering-and-billing
+# https://help.aliyun.com/zh/dashscope/developer-reference/model-introduction
 class AliyunModelEnum(Enum):
     QWEN1_5_05B = ModelConfig(
         key="qwen1.5-0.5b-chat",
@@ -72,6 +73,11 @@ def __init__(
         )
         self.concurrency = 5
 
+    @classmethod
+    def set_api_auth(cls, auth: Dict[str, str]):
+        settings = config.get_settings()
+        settings.ALIYUN_DASHSCOPE_API_KEY = auth.get("API-KEY", "")
+
     @staticmethod
     def get_headers(stream: bool) -> Dict[str, str]:
         k = config.get_settings().ALIYUN_DASHSCOPE_API_KEY

src/retk/core/ai/llm/api/baidu.py

@@ -13,13 +13,20 @@
 
 
 # https://cloud.baidu.com/doc/WENXINWORKSHOP/s/hlrk4akp7#tokens%E7%94%A8%E9%87%8F%E5%90%8E%E4%BB%98%E8%B4%B9
+# https://cloud.baidu.com/doc/WENXINWORKSHOP/s/Slkkydake#%E8%AF%A6%E6%83%85
 class BaiduModelEnum(Enum):
     ERNIE4_8K = ModelConfig(
         key="completions_pro",
         max_tokens=8000,
         RPM=120,
         TPM=120_000,
     )  # 0.12 / 0.12
+    ERNIE4_TURBO_8K = ModelConfig(
+        key="ernie-4.0-turbo-8k",
+        max_tokens=8000,
+        RPM=300,
+        TPM=300_000,
+    )  # 0.003 / 0.06
     ERNIE35_8K = ModelConfig(
         key="completions",
         max_tokens=8000,
@@ -85,6 +92,12 @@ def __init__(
         self.token_expires_at = datetime.now().timestamp()
         self.token = ""
 
+    @classmethod
+    def set_api_auth(cls, auth: Dict[str, str]):
+        settings = config.get_settings()
+        settings.BAIDU_QIANFAN_API_KEY = auth.get("API-KEY", "")
+        settings.BAIDU_QIANFAN_SECRET_KEY = auth.get("Secret Key", "")
+
     async def set_token(self, req_id: str = None):
         _s = config.get_settings()
         if _s.BAIDU_QIANFAN_API_KEY == "" or _s.BAIDU_QIANFAN_SECRET_KEY == "":
@@ -153,7 +166,7 @@ async def complete(
 
         if resp.get("error_code") is not None:
             logger.error(f"rid='{req_id}' | Baidu {model} | error: code={resp['error_code']} {resp['error_msg']}")
-            return resp["error_msg"], const.CodeEnum.LLM_SERVICE_ERROR
+            return resp["error_msg"], const.CodeEnum.INVALID_AUTH
         logger.info(f"rid='{req_id}' | Baidu {model} | usage: {resp['usage']}")
         return resp["result"], const.CodeEnum.OK
 

src/retk/core/ai/llm/api/base.py

@@ -48,6 +48,11 @@ def __init__(
         self.default_model: Optional[ModelConfig] = default_model
         self.endpoint = endpoint
 
+    @classmethod
+    @abstractmethod
+    def set_api_auth(cls, auth: Dict[str, str]):
+        ...
+
     async def _complete(
             self,
             url: str,
@@ -79,7 +84,11 @@ async def _complete(
         if resp.status_code != 200:
             txt = resp.text.replace('\n', '')
             logger.error(f"rid='{req_id}' Model error: {txt}")
-            return {}, const.CodeEnum.LLM_SERVICE_ERROR
+            if resp.status_code in [401, 403]:
+                code = const.CodeEnum.INVALID_AUTH
+            else:
+                code = const.CodeEnum.LLM_SERVICE_ERROR
+            return {}, code
 
         rj = resp.json()
         return rj, const.CodeEnum.OK

src/retk/core/ai/llm/api/moonshot.py

@@ -41,6 +41,10 @@ def __init__(
             timeout=timeout,
         )
 
+    @classmethod
+    def set_api_auth(cls, auth: Dict[str, str]):
+        config.get_settings().MOONSHOT_API_KEY = auth.get("API-KEY", "")
+
     @staticmethod
     def get_api_key():
         return config.get_settings().MOONSHOT_API_KEY

src/retk/core/ai/llm/api/openai.py

@@ -159,6 +159,10 @@ def __init__(
             timeout=timeout,
         )
 
+    @classmethod
+    def set_api_auth(cls, auth: Dict[str, str]):
+        config.get_settings().OPENAI_API_KEY = auth.get("API-KEY", "")
+
     @staticmethod
     def get_api_key():
         return config.get_settings().OPENAI_API_KEY