diff --git a/src/retk/config.py b/src/retk/config.py index 9e578ef..8aa7423 100644 --- a/src/retk/config.py +++ b/src/retk/config.py @@ -79,6 +79,10 @@ class Settings(BaseSettings): VOLCENGINE_API_KEY: str = Field(env='VOLCENGINE_API_KEY', default="") VOLCENGINE_ENDPOINT_ID: str = Field(env='VOLCENGINE_ENDPOINT_ID', default="") + # zhipu bigmodel + BIGMODEL_API_KEY: str = Field(env='BIGMODEL_API_KEY', default="") + BIGMODEL_CONCURRENCY: int = Field(env='BIGMODEL_CONCURRENCY', default=1) + # Email client settings RETHINK_EMAIL: str = Field(env='RETHINK_EMAIL', default="") RETHINK_EMAIL_PASSWORD: str = Field(env='RETHINK_EMAIL_PASSWORD', default="") diff --git a/src/retk/const/user_behavior_types.py b/src/retk/const/user_behavior_types.py index a2e7db3..5141f18 100644 --- a/src/retk/const/user_behavior_types.py +++ b/src/retk/const/user_behavior_types.py @@ -39,6 +39,7 @@ class UserBehaviorTypeEnum(IntEnum): NODE_FILE_UPLOAD = 23 # backend LLM_KNOWLEDGE_RESPONSE = 24 # backend + NODE_PAGE_VIEW = 25 # backend USER_BEHAVIOR_TYPE_MAP = { diff --git a/src/retk/controllers/node/node_ops.py b/src/retk/controllers/node/node_ops.py index 0cf9f20..ea74eff 100644 --- a/src/retk/controllers/node/node_ops.py +++ b/src/retk/controllers/node/node_ops.py @@ -102,6 +102,11 @@ async def get_node( n, code = await core.node.get(au=au, nid=nid) maybe_raise_json_exception(au=au, code=code) + await core.statistic.add_user_behavior( + uid=au.u.id, + type_=const.UserBehaviorTypeEnum.NODE_PAGE_VIEW, + remark=nid, + ) return schemas.node.NodeResponse( requestId=au.request_id, node=get_node_data(n), diff --git a/src/retk/core/ai/llm/api/__init__.py b/src/retk/core/ai/llm/api/__init__.py index 80fd4ba..03ceec3 100644 --- a/src/retk/core/ai/llm/api/__init__.py +++ b/src/retk/core/ai/llm/api/__init__.py @@ -3,6 +3,7 @@ from .aliyun import AliyunService, AliyunModelEnum from .baidu import BaiduService, BaiduModelEnum from .base import BaseLLMService +from .bigmodel import GLMService, GLMModelEnum from .moonshot import MoonshotService, MoonshotModelEnum from .openai import OpenaiService, OpenaiModelEnum from .tencent import TencentService, TencentModelEnum @@ -38,10 +39,14 @@ "service": VolcEngineService, "models": VolcEngineModelEnum, }, + GLMService.name: { + "service": GLMService, + "models": GLMModelEnum, + }, } TOP_P = 0.9 -TEMPERATURE = 0.6 +TEMPERATURE = 0.9 TIMEOUT = 60 LLM_DEFAULT_SERVICES: Dict[str, BaseLLMService] = { diff --git a/src/retk/core/ai/llm/api/aliyun.py b/src/retk/core/ai/llm/api/aliyun.py index 7b6376c..64c39e0 100644 --- a/src/retk/core/ai/llm/api/aliyun.py +++ b/src/retk/core/ai/llm/api/aliyun.py @@ -58,7 +58,7 @@ class AliyunService(BaseLLMService): def __init__( self, top_p: float = 0.9, - temperature: float = 0.4, + temperature: float = 0.9, timeout: float = 60., ): super().__init__( diff --git a/src/retk/core/ai/llm/api/baidu.py b/src/retk/core/ai/llm/api/baidu.py index 7adc487..60d287a 100644 --- a/src/retk/core/ai/llm/api/baidu.py +++ b/src/retk/core/ai/llm/api/baidu.py @@ -71,7 +71,7 @@ class BaiduService(BaseLLMService): def __init__( self, top_p: float = 0.9, - temperature: float = 0.4, + temperature: float = 0.9, timeout: float = 60., ): super().__init__( diff --git a/src/retk/core/ai/llm/api/base.py b/src/retk/core/ai/llm/api/base.py index d0afd8a..6a73899 100644 --- a/src/retk/core/ai/llm/api/base.py +++ b/src/retk/core/ai/llm/api/base.py @@ -36,7 +36,7 @@ def __init__( model_enum: Any, endpoint: str, top_p: float = 1., - temperature: float = 0.4, + temperature: float = 0.9, timeout: float = None, default_model: Optional[ModelConfig] = None, ): diff --git a/src/retk/core/ai/llm/api/bigmodel.py b/src/retk/core/ai/llm/api/bigmodel.py new file mode 100644 index 0000000..5ce11c5 --- /dev/null +++ b/src/retk/core/ai/llm/api/bigmodel.py @@ -0,0 +1,97 @@ +import asyncio +from enum import Enum +from typing import List, Tuple, Callable, Union, Dict + +from retk import config, const +from retk.core.utils import ratelimiter +from .base import ModelConfig, MessagesType +from .openai import OpenaiLLMStyle + + +# https://open.bigmodel.cn/dev/howuse/rate-limits/tiers?tab=0 +class GLMModelEnum(Enum): + GLM4_PLUS = ModelConfig( + key="GLM-4-Plus", + max_tokens=128_000, + ) + GLM4_LONG = ModelConfig( + key="GLM-4-Long", + max_tokens=1_000_000, + ) + GLM4_FLASH = ModelConfig( + key="GLM-4-Flash", + max_tokens=128_000, + ) + + +class GLMService(OpenaiLLMStyle): + name = "glm" + + def __init__( + self, + top_p: float = 0.9, + temperature: float = 0.9, + timeout: float = 60., + ): + super().__init__( + model_enum=GLMModelEnum, + endpoint="https://open.bigmodel.cn/api/paas/v4/chat/completions", + default_model=GLMModelEnum.GLM4_FLASH.value, + top_p=top_p, + temperature=temperature, + timeout=timeout, + ) + + @classmethod + def set_api_auth(cls, auth: Dict[str, str]): + config.get_settings().BIGMODEL_API_KEY = auth.get("API-KEY", "") + + @staticmethod + def get_api_key(): + return config.get_settings().BIGMODEL_API_KEY + + @staticmethod + async def _batch_complete_union( + messages: List[MessagesType], + func: Callable, + model: str = None, + req_id: str = None, + ) -> List[Tuple[Union[str, Dict[str, str]], const.CodeEnum]]: + settings = config.get_settings() + concurrent_limiter = ratelimiter.ConcurrentLimiter(n=settings.BIGMODEL_CONCURRENCY) + + tasks = [ + func( + limiters=[concurrent_limiter], + messages=m, + model=model, + req_id=req_id, + ) for m in messages + ] + return await asyncio.gather(*tasks) + + async def batch_complete( + self, + messages: List[MessagesType], + model: str = None, + req_id: str = None, + ) -> List[Tuple[str, const.CodeEnum]]: + return await self._batch_complete_union( + messages=messages, + func=self._batch_complete, + model=model, + req_id=req_id, + ) + + async def batch_complete_json_detect( + self, + messages: List[MessagesType], + model: str = None, + req_id: str = None, + ) -> List[Tuple[Dict[str, str], const.CodeEnum]]: + return await self._batch_complete_union( + messages=messages, + func=self._batch_stream_complete_json_detect, + model=model, + req_id=req_id, + ) diff --git a/src/retk/core/ai/llm/api/moonshot.py b/src/retk/core/ai/llm/api/moonshot.py index f06ee34..4e73ddf 100644 --- a/src/retk/core/ai/llm/api/moonshot.py +++ b/src/retk/core/ai/llm/api/moonshot.py @@ -30,7 +30,7 @@ class MoonshotService(OpenaiLLMStyle): def __init__( self, top_p: float = 0.9, - temperature: float = 0.3, + temperature: float = 0.9, timeout: float = 60., ): super().__init__( diff --git a/src/retk/core/ai/llm/api/openai.py b/src/retk/core/ai/llm/api/openai.py index 8104888..9d06431 100644 --- a/src/retk/core/ai/llm/api/openai.py +++ b/src/retk/core/ai/llm/api/openai.py @@ -43,7 +43,7 @@ def __init__( endpoint: str, default_model: ModelConfig, top_p: float = 0.9, - temperature: float = 0.4, + temperature: float = 0.9, timeout: float = 60., ): super().__init__( @@ -148,7 +148,7 @@ class OpenaiService(OpenaiLLMStyle): def __init__( self, top_p: float = 0.9, - temperature: float = 0.7, + temperature: float = 0.9, timeout: float = 60., ): super().__init__( diff --git a/src/retk/core/ai/llm/api/tencent.py b/src/retk/core/ai/llm/api/tencent.py index 0d1640b..e5bc989 100644 --- a/src/retk/core/ai/llm/api/tencent.py +++ b/src/retk/core/ai/llm/api/tencent.py @@ -58,7 +58,7 @@ class TencentService(BaseLLMService): def __init__( self, top_p: float = 0.9, - temperature: float = 0.4, + temperature: float = 0.9, timeout: float = 60., ): super().__init__( diff --git a/src/retk/core/ai/llm/api/volcengine.py b/src/retk/core/ai/llm/api/volcengine.py index e853c82..b99772b 100644 --- a/src/retk/core/ai/llm/api/volcengine.py +++ b/src/retk/core/ai/llm/api/volcengine.py @@ -59,7 +59,7 @@ class VolcEngineService(OpenaiLLMStyle): def __init__( self, top_p: float = 0.9, - temperature: float = 0.3, + temperature: float = 0.9, timeout: float = 60., ): super().__init__( diff --git a/src/retk/core/ai/llm/api/xfyun.py b/src/retk/core/ai/llm/api/xfyun.py index d78652b..7790a0e 100644 --- a/src/retk/core/ai/llm/api/xfyun.py +++ b/src/retk/core/ai/llm/api/xfyun.py @@ -44,7 +44,7 @@ class XfYunService(BaseLLMService): def __init__( self, top_p: float = 0.9, - temperature: float = 0.4, + temperature: float = 0.9, timeout: float = 60., ): super().__init__( diff --git a/src/retk/core/scheduler/schedule.py b/src/retk/core/scheduler/schedule.py index 9e13561..653a21c 100644 --- a/src/retk/core/scheduler/schedule.py +++ b/src/retk/core/scheduler/schedule.py @@ -107,6 +107,13 @@ def init_tasks(): kwargs={"delta_days": 30}, hour=1, ) + + if not config.is_local_db(): + run_every_at( + job_id="auto_daily_report", + func=tasks.auto_daily_report.auto_daily_report, + hour=0, + ) return diff --git a/src/retk/core/scheduler/tasks/__init__.py b/src/retk/core/scheduler/tasks/__init__.py index c7bb046..e4b5296 100644 --- a/src/retk/core/scheduler/tasks/__init__.py +++ b/src/retk/core/scheduler/tasks/__init__.py @@ -3,4 +3,5 @@ notice, extend_node, auto_clean_trash, + auto_daily_report, ) diff --git a/src/retk/core/scheduler/tasks/auto_daily_report.py b/src/retk/core/scheduler/tasks/auto_daily_report.py new file mode 100644 index 0000000..ce81467 --- /dev/null +++ b/src/retk/core/scheduler/tasks/auto_daily_report.py @@ -0,0 +1,93 @@ +import asyncio +import json +import os +import time +from datetime import datetime, timedelta + +from bson.objectid import ObjectId + +from retk import config, const +from retk.core.statistic import __write_line_date, __manage_files +from retk.models.client import init_mongo +from retk.models.coll import CollNameEnum + +try: + import aiofiles +except ImportError: + aiofiles = None + + +def auto_daily_report(): + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + res = loop.run_until_complete(_auto_daily_report()) + loop.close() + return res + + +async def _auto_daily_report(): + if config.is_local_db() or aiofiles is None: + return + file = const.settings.ANALYTICS_DIR / "daily_report" / "report.log" + lock = asyncio.Lock() + now = datetime.now() + # get last line efficiently + if file.exists(): + async with aiofiles.open(file, "r", encoding="utf-8") as f: + async with lock: + try: + await f.seek(-2, os.SEEK_END) + while f.read(1) != "\n": + await f.seek(-2, os.SEEK_CUR) + last_line = await f.readline() + except OSError: + last_line = "" + try: + last_record = json.loads(last_line) + except json.JSONDecodeError: + last_record = {} + else: + last_record = {} + yesterday = (now - timedelta(days=1)).date() + last_record_date = last_record.get("date", None) + + if last_record_date is not None: + last_record_date = datetime.strptime(last_record_date, '%Y-%m-%d').date() + else: + last_record_date = yesterday - timedelta(days=1) + + if last_record_date >= yesterday: + return + + await __manage_files(now, file, lock) + + _, db = init_mongo(connection_timeout=5) + total_email_users = await db[CollNameEnum.users.value].count_documents( + {"source": const.UserSourceEnum.EMAIL.value} + ) + total_google_users = await db[CollNameEnum.users.value].count_documents( + {"source": const.UserSourceEnum.GOOGLE.value} + ) + total_github_users = await db[CollNameEnum.users.value].count_documents( + {"source": const.UserSourceEnum.GITHUB.value} + ) + total_users = total_email_users + total_google_users + total_github_users + # date to int timestamp + timestamp = time.mktime(last_record_date.timetuple()) + time_filter = ObjectId.from_datetime(datetime.utcfromtimestamp(timestamp)) + + await __write_line_date( + data={ + "date": now.strftime('%Y-%m-%d'), + "totalUsers": total_users, + "totalEmailUsers": total_email_users, + "totalGoogleUsers": total_google_users, + "totalGithubUsers": total_github_users, + "newUsers": await db[CollNameEnum.users.value].count_documents({"_id": {"$gt": time_filter}}), + "totalNodes": await db[CollNameEnum.nodes.value].count_documents({}), + "newNodes": await db[CollNameEnum.nodes.value].count_documents({"_id": {"$gt": time_filter}}), + "totalFiles": await db[CollNameEnum.user_file.value].count_documents({}), + }, + path=file, + lock=lock + ) diff --git a/src/retk/core/statistic.py b/src/retk/core/statistic.py index 8b6df48..a8c62d6 100644 --- a/src/retk/core/statistic.py +++ b/src/retk/core/statistic.py @@ -18,6 +18,28 @@ async def __write_new(path: Path, lock: asyncio.Lock): await f.write("") +async def __write_line_date(data: dict, path: Path, lock: asyncio.Lock): + async with aiofiles.open(path, "a", encoding="utf-8") as f: + record_str = json.dumps(data, ensure_ascii=False) + async with lock: + await f.write(f"{record_str}\n") + + +async def __manage_files(now: datetime, path: Path, lock: asyncio.Lock): + par = path.parent + if not par.exists(): + par.mkdir(parents=True, exist_ok=True) + await __write_new(path, lock) + + # backup if too large + + # if the file is too large, rename it to current time and create a new one + if path.stat().st_size > const.settings.MAX_USER_BEHAVIOR_LOG_SIZE: + backup_file = path.parent / f"{now.strftime('%Y%m%d-%H%M%S')}.log" + path.rename(backup_file) + await __write_new(path, lock) + + async def add_user_behavior( uid: str, type_: const.UserBehaviorTypeEnum, @@ -25,27 +47,17 @@ async def add_user_behavior( ): if is_local_db() or aiofiles is None: return - data_dir = const.settings.ANALYTICS_DIR / "user_behavior" - current_log_file = data_dir / f"behavior.log" + file = const.settings.ANALYTICS_DIR / "user_behavior" / f"behavior.log" lock = asyncio.Lock() - if not current_log_file.exists(): - data_dir.mkdir(parents=True, exist_ok=True) - await __write_new(current_log_file, lock) - - time_now = datetime.now() - # if the file is too large, rename it to current time and create a new one - if current_log_file.stat().st_size > const.settings.MAX_USER_BEHAVIOR_LOG_SIZE: - backup_file = data_dir / f"{time_now.strftime('%Y%m%d-%H%M%S')}.log" - current_log_file.rename(backup_file) - await __write_new(current_log_file, lock) - - async with aiofiles.open(current_log_file, "a", encoding="utf-8") as f: - record = { - "time": time_now.strftime('%Y-%m-%d %H:%M:%S'), + now = datetime.now() + await __manage_files(now, file, lock) + await __write_line_date( + data={ + "time": now.strftime('%Y-%m-%d %H:%M:%S'), "uid": uid, "type": type_.value, "remark": remark, - } - record_str = json.dumps(record) - async with lock: - await f.write(f"{record_str}\n") + }, + path=file, + lock=lock + ) diff --git a/tests/test_account.py b/tests/test_account.py index be17688..c5a1a1a 100644 --- a/tests/test_account.py +++ b/tests/test_account.py @@ -8,7 +8,7 @@ from retk.models.client import client from retk.models.tps import convert_user_dict_to_authed_user from retk.utils import jwt_decode -from . import utils +from tests import utils class AccountTest(unittest.IsolatedAsyncioTestCase): diff --git a/tests/test_ai_llm_knowledge.py b/tests/test_ai_llm_knowledge.py index 2aaba0b..80df47e 100644 --- a/tests/test_ai_llm_knowledge.py +++ b/tests/test_ai_llm_knowledge.py @@ -5,8 +5,8 @@ from retk import const from retk.core.ai import llm from retk.core.ai.llm.knowledge.ops import ExtendCase -from . import utils -from .test_ai_llm_api import skip_no_api_key, clear_all_api_key +from tests import utils +from tests.test_ai_llm_api import skip_no_api_key, clear_all_api_key md_source = ["""\ 广东猪脚饭特点 @@ -81,6 +81,7 @@ async def test_summary(self): (llm.api.XfYunService.name, llm.api.XfYunModelEnum.SPARK_LITE), (llm.api.MoonshotService.name, llm.api.MoonshotModelEnum.V1_8K), # 这个总结比较好 (llm.api.VolcEngineService.name, llm.api.VolcEngineModelEnum.DOUBAO_PRO_32K), + (llm.api.GLMService.name, llm.api.GLMModelEnum.GLM4_FLASH), ]: cases = [ ExtendCase( @@ -112,6 +113,7 @@ async def test_extend(self): (llm.api.XfYunService.name, llm.api.XfYunModelEnum.SPARK_PRO), (llm.api.MoonshotService.name, llm.api.MoonshotModelEnum.V1_8K), # 这个延伸比较好 (llm.api.VolcEngineService.name, llm.api.VolcEngineModelEnum.DOUBAO_PRO_32K), # 这个延伸比较好 + (llm.api.GLMService.name, llm.api.GLMModelEnum.GLM4_PLUS), ]: cases = [ ExtendCase( diff --git a/tests/test_api.py b/tests/test_api.py index 76de20b..3c0f0db 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -23,7 +23,7 @@ from retk.models.tps.llm import ExtendedNode from retk.plugins.register import register_official_plugins, unregister_official_plugins from retk.utils import jwt_decode -from . import utils +from tests import utils class PublicApiTest(unittest.IsolatedAsyncioTestCase): diff --git a/tests/test_core_files_obsidian.py b/tests/test_core_files_obsidian.py index 0461aa0..f2df552 100644 --- a/tests/test_core_files_obsidian.py +++ b/tests/test_core_files_obsidian.py @@ -10,7 +10,7 @@ from retk.core.files import saver from retk.core.files.importing.async_tasks.obsidian import ops from retk.models.client import client -from . import utils +from tests import utils class ObsidianTest(unittest.IsolatedAsyncioTestCase): diff --git a/tests/test_core_local.py b/tests/test_core_local.py index 2c29f3c..0f474e6 100644 --- a/tests/test_core_local.py +++ b/tests/test_core_local.py @@ -24,7 +24,7 @@ from retk.models.client import client from retk.models.tps import ImportData, AuthedUser, convert_user_dict_to_authed_user from retk.utils import short_uuid -from . import utils +from tests import utils @patch("retk.core.ai.llm.knowledge.ops._batch_send", new_callable=AsyncMock, return_value=["", const.CodeEnum.OK]) diff --git a/tests/test_core_remote.py b/tests/test_core_remote.py index b1a791e..53062ed 100644 --- a/tests/test_core_remote.py +++ b/tests/test_core_remote.py @@ -18,7 +18,7 @@ from retk.models.client import client from retk.models.tps import AuthedUser, convert_user_dict_to_authed_user from retk.utils import get_token -from . import utils +from tests import utils @patch("retk.core.ai.llm.knowledge.ops._batch_send", new_callable=AsyncMock, return_value=["", const.CodeEnum.OK]) diff --git a/tests/test_data_restore.py b/tests/test_data_restore.py index 4ea69ed..13e1252 100644 --- a/tests/test_data_restore.py +++ b/tests/test_data_restore.py @@ -3,7 +3,7 @@ from retk import const, core from retk.models.client import client from retk.models.tps import AuthedUser, convert_user_dict_to_authed_user -from . import utils +from tests import utils class DataRestoreTest(unittest.IsolatedAsyncioTestCase): diff --git a/tests/test_local_manager.py b/tests/test_local_manager.py index 7b062d1..6e28011 100644 --- a/tests/test_local_manager.py +++ b/tests/test_local_manager.py @@ -5,7 +5,7 @@ from retk import const, config, local_manager, __version__ from retk.local_manager import recover -from . import utils +from tests import utils class RecoverTest(unittest.TestCase): diff --git a/tests/test_plugins.py b/tests/test_plugins.py index 05d55b7..264f60a 100644 --- a/tests/test_plugins.py +++ b/tests/test_plugins.py @@ -8,7 +8,7 @@ from retk.models.client import client from retk.models.tps import convert_user_dict_to_authed_user, AuthedUser from retk.plugins.base import event_plugin_map -from . import utils +from tests import utils class TestPlugin(retk.Plugin): diff --git a/tests/test_search_es.py b/tests/test_search_es.py index e812b56..90bd2c7 100644 --- a/tests/test_search_es.py +++ b/tests/test_search_es.py @@ -8,7 +8,7 @@ from retk import const, config from retk.models.search_engine.engine_es import ESSearcher, SearchDoc from retk.models.tps import AuthedUser -from . import utils +from tests import utils class ESTest(unittest.IsolatedAsyncioTestCase): diff --git a/tests/test_search_local.py b/tests/test_search_local.py index 9b332e2..0b91950 100644 --- a/tests/test_search_local.py +++ b/tests/test_search_local.py @@ -7,7 +7,7 @@ from retk import const from retk.models.search_engine.engine_local import LocalSearcher, SearchDoc from retk.models.tps import AuthedUser -from . import utils +from tests import utils class LocalSearchTest(unittest.IsolatedAsyncioTestCase):