feat(llm):

- add moonshot
- add summary

src/retk/config.py

@@ -62,6 +62,9 @@ class Settings(BaseSettings):
     XFYUN_API_SECRET: str = Field(env='XFYUN_API_SECRET', default="")
     XFYUN_API_KEY: str = Field(env='XFYUN_API_KEY', default="")
 
+    # moonshot api
+    MOONSHOT_API_KEY: str = Field(env='MOONSHOT_API_KEY', default="")
+
     # Email client settings
     RETHINK_EMAIL: str = Field(env='RETHINK_EMAIL', default="")
     RETHINK_EMAIL_PASSWORD: str = Field(env='RETHINK_EMAIL_PASSWORD', default="")

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

@@ -1,5 +1 @@
-from .aliyun import Aliyun
-from .baidu import Baidu
-from .openai import OpenAI
-from .tencent import Tencent
-from .xfyun import XfYun
+from . import api, knowledge

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

@@ -0,0 +1,6 @@
+from .aliyun import AliyunService, AliyunModelEnum
+from .baidu import BaiduService, BaiduModelEnum
+from .moonshot import MoonshotService, MoonshotModelEnum
+from .openai import OpenaiService, OpenaiModelEnum
+from .tencent import TencentService, TencentModelEnum
+from .xfyun import XfYunService, XfYunModelEnum

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

@@ -4,18 +4,21 @@
 
 from retk import config, const
 from retk.logger import logger
-from .base import BaseLLM, MessagesType
+from .base import BaseLLMService, MessagesType, NoAPIKeyError
 
 
+# https://help.aliyun.com/zh/dashscope/developer-reference/tongyi-thousand-questions-metering-and-billing
 class AliyunModelEnum(str, Enum):
-    QWEN1_5_0_5B = "qwen1.5-0.5b-chat"
-    QWEN1_8B = "qwen-1.8b-chat"
-    BAICHUAN7BV1 = "baichuan-7b-v1"
-    LLAMA3_70B = "llama3-70b-instruct"
-    CHATGLM3_6B = "chatglm3-6b"
+    QWEN1_5_05B = "qwen1.5-0.5b-chat"  # free
+    QWEN_2B = "qwen-1.8b-chat"  # free
+    BAICHUAN7BV1 = "baichuan-7b-v1"  # free
+    QWEN_LONG = "qwen-long"  # in 0.0005/1000, out 0.002/1000
+    QWEN_TURBO = "qwen-turbo"  # in 0.002/1000, out 0.006/1000
+    QWEN_PLUS = "qwen-plus"  # in 0.004/1000, out 0.012/1000
+    QWEN_MAX = "qwen-max"  # in 0.04/1000, out 0.12/1000
 
 
-class Aliyun(BaseLLM):
+class AliyunService(BaseLLMService):
     def __init__(
             self,
             top_p: float = 0.9,
@@ -27,11 +30,11 @@ def __init__(
             top_p=top_p,
             temperature=temperature,
             timeout=timeout,
-            default_model=AliyunModelEnum.QWEN1_5_0_5B.value,
+            default_model=AliyunModelEnum.QWEN1_5_05B.value,
         )
         self.api_key = config.get_settings().ALIYUN_DASHSCOPE_API_KEY
         if self.api_key == "":
-            raise ValueError("Aliyun API key is empty")
+            raise NoAPIKeyError("Aliyun API key is empty")
 
     def get_headers(self, stream: bool) -> Dict[str, str]:
         h = {

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

@@ -5,7 +5,7 @@
 
 from retk import config, const, httpx_helper
 from retk.logger import logger
-from .base import BaseLLM, MessagesType
+from .base import BaseLLMService, MessagesType, NoAPIKeyError
 
 
 class BaiduModelEnum(str, Enum):
@@ -17,7 +17,7 @@ class BaiduModelEnum(str, Enum):
     YI_34B_CHAT = "yi_34b_chat"
 
 
-class Baidu(BaseLLM):
+class BaiduService(BaseLLMService):
     def __init__(
             self,
             top_p: float = 0.9,
@@ -34,7 +34,7 @@ def __init__(
         self.api_key = config.get_settings().BAIDU_QIANFAN_API_KEY
         self.secret_key = config.get_settings().BAIDU_QIANFAN_SECRET_KEY
         if self.api_key == "" or self.secret_key == "":
-            raise ValueError("Baidu api key or key is empty")
+            raise NoAPIKeyError("Baidu api key or key is empty")
 
         self.headers = {
             "Content-Type": "application/json",
@@ -69,6 +69,10 @@ async def set_token(self, req_id: str = None):
 
     @staticmethod
     def get_payload(messages: MessagesType, stream: bool) -> bytes:
+        if messages[0]["role"] == "system":
+            messages[0]["role"] = "user"
+            if messages[1]["role"] == "user":
+                messages.insert(1, {"role": "assistant", "content": "明白。"})
         return json.dumps(
             {
                 "messages": messages,

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

@@ -9,7 +9,11 @@
 MessagesType = List[Dict[Literal["role", "content"], str]]
 
 
-class BaseLLM(ABC):
+class NoAPIKeyError(Exception):
+    pass
+
+
+class BaseLLMService(ABC):
     default_timeout = 60.
 
     def __init__(

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

@@ -0,0 +1,27 @@
+from enum import Enum
+
+from retk import config
+from .openai import OpenaiLLMStyle
+
+
+class MoonshotModelEnum(str, Enum):
+    V1_8K = "moonshot-v1-8k"
+    V1_32K = "moonshot-v1-32k"
+    V1_128K = "moonshot-v1-128k"
+
+
+class MoonshotService(OpenaiLLMStyle):
+    def __init__(
+            self,
+            top_p: float = 0.9,
+            temperature: float = 0.7,
+            timeout: float = 60.,
+    ):
+        super().__init__(
+            api_key=config.get_settings().MOONSHOT_API_KEY,
+            endpoint="https://api.moonshot.cn/v1/chat/completions",
+            default_model=MoonshotModelEnum.V1_8K.value,
+            top_p=top_p,
+            temperature=temperature,
+            timeout=timeout,
+        )

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

@@ -4,34 +4,37 @@
 
 from retk import config, const
 from retk.logger import logger
-from .base import BaseLLM, MessagesType
+from .base import BaseLLMService, MessagesType, NoAPIKeyError
 
 
-class OpenAIModelEnum(str, Enum):
+class OpenaiModelEnum(str, Enum):
     GPT4 = "gpt-4"
     GPT4_TURBO_PREVIEW = "gpt-4-turbo-preview"
     GPT4_32K = "gpt-4-32k"
     GPT35_TURBO = "gpt-3.5-turbo"
     GPT35_TURBO_16K = "gpt-3.5-turbo-16k"
 
 
-class OpenAI(BaseLLM):
+class OpenaiLLMStyle(BaseLLMService):
     def __init__(
             self,
+            api_key: str,
+            endpoint: str,
+            default_model: str,
             top_p: float = 0.9,
             temperature: float = 0.7,
             timeout: float = 60.,
     ):
         super().__init__(
-            endpoint="https://api.openai.com/v1/chat/completions",
+            endpoint=endpoint,
             top_p=top_p,
             temperature=temperature,
             timeout=timeout,
-            default_model=OpenAIModelEnum.GPT35_TURBO.value,
+            default_model=default_model,
         )
-        self.api_key = config.get_settings().OPENAI_API_KEY
+        self.api_key = api_key
         if self.api_key == "":
-            raise ValueError("OpenAI API key is empty")
+            raise NoAPIKeyError(f"{self.__class__.__name__} API key is empty")
         self.headers = {
             "Content-Type": "application/json",
             "Authorization": f"Bearer {self.api_key}",
@@ -43,7 +46,7 @@ def get_payload(self, model: Optional[str], messages: MessagesType, stream: bool
         return json.dumps({
             "model": model,
             "messages": messages,
-            "max_tokens": 100,
+            # "max_tokens": 100,
             "temperature": self.temperature,
             "top_p": self.top_p,
             "stream": stream,
@@ -66,7 +69,7 @@ async def complete(
             return "", code
         if rj.get("error") is not None:
             return rj["error"]["message"], const.CodeEnum.LLM_SERVICE_ERROR
-        logger.info(f"ReqId={req_id} OpenAI model usage: {rj['usage']}")
+        logger.info(f"ReqId={req_id} {self.__class__.__name__} model usage: {rj['usage']}")
         return rj["choices"][0]["message"]["content"], code
 
     async def stream_complete(
@@ -86,19 +89,38 @@ async def stream_complete(
                 yield b, code
                 continue
             txt = ""
-            lines = b.splitlines()
+            lines = filter(lambda s: s != b"", b.split("\n\n".encode("utf-8")))
             for line in lines:
-                json_str = line.decode("utf-8").strip()
-                if json_str == "":
-                    continue
+                json_str = line.decode("utf-8")[5:].strip()
                 try:
                     json_data = json.loads(json_str)
                 except json.JSONDecodeError:
-                    logger.error(f"ReqId={req_id} OpenAI model stream error: json={json_str}")
+                    logger.error(f"ReqId={req_id} {self.__class__.__name__} model stream error: json={json_str}")
                     continue
                 choice = json_data["choices"][0]
                 if choice["finish_reason"] is not None:
-                    logger.info(f"ReqId={req_id} OpenAI model usage: {json_data['usage']}")
+                    try:
+                        usage = json_data["usage"]
+                    except KeyError:
+                        usage = choice["usage"]
+                    logger.info(f"ReqId={req_id} {self.__class__.__name__} model usage: {usage}")
                     break
                 txt += choice["delta"]["content"]
             yield txt.encode("utf-8"), code
+
+
+class OpenaiService(OpenaiLLMStyle):
+    def __init__(
+            self,
+            top_p: float = 0.9,
+            temperature: float = 0.7,
+            timeout: float = 60.,
+    ):
+        super().__init__(
+            api_key=config.get_settings().OPENAI_API_KEY,
+            endpoint="https://api.openai.com/v1/chat/completions",
+            default_model=OpenaiModelEnum.GPT35_TURBO.value,
+            top_p=top_p,
+            temperature=temperature,
+            timeout=timeout,
+        )

src/retk/core/ai/llm/tencent.py → src/retk/core/ai/llm/api/tencent.py

@@ -8,7 +8,7 @@
 
 from retk import config, const
 from retk.logger import logger
-from .base import BaseLLM, MessagesType
+from .base import BaseLLMService, MessagesType, NoAPIKeyError
 
 Headers = TypedDict("Headers", {
     "Authorization": str,
@@ -22,18 +22,18 @@
 
 
 class TencentModelEnum(str, Enum):
-    HUNYUAN_PRO = "hunyuan-pro"
-    HUNYUAN_STANDARD = "hunyuan-standard"
-    HUNYUAN_STANDARD_256K = "hunyuan-standard-256K"
-    HUNYUAN_LITE = "hunyuan-lite"
+    HUNYUAN_PRO = "hunyuan-pro"  # in 0.03/1000, out 0.10/1000
+    HUNYUAN_STANDARD = "hunyuan-standard"  # in 0.0045/1000, out 0.005/1000
+    HUNYUAN_STANDARD_256K = "hunyuan-standard-256K"  # in 0.015/1000, out 0.06/1000
+    HUNYUAN_LITE = "hunyuan-lite"  # free
 
 
 # 计算签名摘要函数
 def sign(key, msg):
     return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()
 
 
-class Tencent(BaseLLM):
+class TencentService(BaseLLMService):
     service = "hunyuan"
     host = "hunyuan.tencentcloudapi.com"
     version = "2023-09-01"
@@ -54,7 +54,7 @@ def __init__(
         self.secret_id = config.get_settings().HUNYUAN_SECRET_ID
         self.secret_key = config.get_settings().HUNYUAN_SECRET_KEY
         if self.secret_id == "" or self.secret_key == "":
-            raise ValueError("Tencent secret id or key is empty")
+            raise NoAPIKeyError("Tencent secret id or key is empty")
 
     def get_auth(self, action: str, payload: bytes, timestamp: int, content_type: str) -> str:
         algorithm = "TC3-HMAC-SHA256"

src/retk/core/ai/llm/xfyun.py → src/retk/core/ai/llm/api/xfyun.py

@@ -14,7 +14,7 @@
 
 from retk import config, const
 from retk.logger import logger
-from .base import BaseLLM, MessagesType
+from .base import BaseLLMService, MessagesType, NoAPIKeyError
 
 
 class XfYunModelEnum(str, Enum):
@@ -32,7 +32,7 @@ class XfYunModelEnum(str, Enum):
 }
 
 
-class XfYun(BaseLLM):
+class XfYunService(BaseLLMService):
     def __init__(
             self,
             top_p: float = 0.9,
@@ -51,7 +51,7 @@ def __init__(
         self.api_key = _s.XFYUN_API_KEY
         self.app_id = _s.XFYUN_APP_ID
         if self.api_secret == "" or self.api_key == "" or self.app_id == "":
-            raise ValueError("XfYun api secret or key is empty")
+            raise NoAPIKeyError("XfYun api secret or key is empty")
 
     def get_url(self, model: Optional[str], req_id: str = None) -> str:
         if model is None:

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

@@ -0,0 +1,52 @@
+from pathlib import Path
+from typing import Tuple
+
+from retk import const
+from ..api.base import BaseLLMService, MessagesType
+
+system_summary_prompt = (Path(__file__).parent / "system_summary.md").read_text(encoding="utf-8")
+system_extend_prompt = (Path(__file__).parent / "system_extend.md").read_text(encoding="utf-8")
+
+
+async def _send(
+        llm_service: BaseLLMService,
+        model: str,
+        system_prompt: str,
+        query: str,
+        req_id: str,
+) -> Tuple[str, const.CodeEnum]:
+    _msgs: MessagesType = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": query},
+    ]
+    return await llm_service.complete(messages=_msgs, model=model, req_id=req_id)
+
+
+async def summary(
+        llm_service: BaseLLMService,
+        model: str,
+        query: str,
+        req_id: str = None,
+) -> Tuple[str, const.CodeEnum]:
+    return await _send(
+        llm_service=llm_service,
+        model=model,
+        system_prompt=system_summary_prompt,
+        query=query,
+        req_id=req_id,
+    )
+
+
+async def extend(
+        llm_service: BaseLLMService,
+        model: str,
+        query: str,
+        req_id: str = None,
+) -> Tuple[str, const.CodeEnum]:
+    return await _send(
+        llm_service=llm_service,
+        model=model,
+        system_prompt=system_extend_prompt,
+        query=query,
+        req_id=req_id,
+    )