diff --git a/README.md b/README.md index e5767b0..c423224 100644 --- a/README.md +++ b/README.md @@ -83,7 +83,7 @@ python main.py 你可以在 `http://localhost:3000/` 打开管理页面。 -## Usage +## 用法 1. 创建一个 channel,按照说明填写配置,然后创建一个新的 key。 @@ -135,6 +135,33 @@ response = openai.ChatCompletion.create( print(response) ``` +### 配置文件 + +配置文件位于`data/config.yaml` + +```yaml +database: + # SQLite 数据库文件路径 + path: ./data/free_one_api.db + type: sqlite +router: + # 后端监听端口 + port: 3000 + # 管理页登录密码 + token: '12345678' +watchdog: + heartbeat: + # 自动停用渠道前的心跳失败次数 + fail_limit: 3 + # 心跳检测间隔(秒) + interval: 1800 + # 单个渠道心跳检测超时时间(秒) + timeout: 300 +web: + # 前端页面路径 + frontend_path: ./web/dist/ +``` + ## 快速体验 ### Demo diff --git a/README_en.md b/README_en.md index 2ce95eb..a0b52d0 100644 --- a/README_en.md +++ b/README_en.md @@ -135,6 +135,33 @@ response = openai.ChatCompletion.create( print(response) ``` +### Configurations + +Configuration file is saved at `data/config.yaml` + +```yaml +database: + # SQLite DB file path + path: ./data/free_one_api.db + type: sqlite +router: + # Backend listen port + port: 3000 + # Admin page login password + token: '12345678' +watchdog: + heartbeat: + # Max fail times + fail_limit: 3 + # Heartbeat check interval (seconds) + interval: 1800 + # Single channel heartbeat check timeout (seconds) + timeout: 300 +web: + # Frontend page path + frontend_path: ./web/dist/ +``` + ## Quick Test ### Demo diff --git a/free_one_api/entities/channel.py b/free_one_api/entities/channel.py index a10ea01..8bd281c 100644 --- a/free_one_api/entities/channel.py +++ b/free_one_api/entities/channel.py @@ -1,4 +1,6 @@ +import asyncio import json +import time import tiktoken @@ -21,6 +23,9 @@ class Channel: enabled: bool latency: int + + fail_count: int + """Amount of sequential failures. Only in memory.""" def __init__(self, id: int, name: str, adapter: llm.LLMLibAdapter, model_mapping: dict, enabled: bool, latency: int): self.id = id @@ -29,6 +34,7 @@ def __init__(self, id: int, name: str, adapter: llm.LLMLibAdapter, model_mapping self.model_mapping = model_mapping self.enabled = enabled self.latency = latency + self.fail_count = 0 @classmethod def dump_channel(cls, chan: 'Channel') -> dict: @@ -68,3 +74,21 @@ def count_tokens( num_tokens += len(encoding.encode(value)) num_tokens += 3 # every reply is primed with <|start|>assistant<|message|> return num_tokens + + async def heartbeat(self, timeout: int=300) -> int: + """Call adapter test, returns fail count.""" + + try: + start = time.time() + succ = await asyncio.wait_for(self.adapter.test(), timeout=timeout) + if succ: + latency = int((time.time() - start)*100)/100 + self.fail_count = 0 + self.latency = latency + return 0 + else: + self.fail_count += 1 + return self.fail_count + finally: + self.fail_count += 1 + return self.fail_count diff --git a/free_one_api/impls/adapter/hugchat.py b/free_one_api/impls/adapter/hugchat.py index 27a8baa..c6b6659 100644 --- a/free_one_api/impls/adapter/hugchat.py +++ b/free_one_api/impls/adapter/hugchat.py @@ -65,10 +65,10 @@ def chatbot(self) -> hugchat.ChatBot: sign = login.Login(self.config['email'], self.config['passwd']) cookie: requests.sessions.RequestsCookieJar = None try: - cookie = sign.loadCookiesFromDir("hugchatCookies") + cookie = sign.loadCookiesFromDir("data/hugchatCookies") except: cookie = sign.login() - sign.saveCookiesToDir("hugchatCookies") + sign.saveCookiesToDir("data/hugchatCookies") self._chatbot = hugchat.ChatBot(cookies=cookie.get_dict()) return self._chatbot diff --git a/free_one_api/impls/app.py b/free_one_api/impls/app.py index c54b3be..63da2e1 100644 --- a/free_one_api/impls/app.py +++ b/free_one_api/impls/app.py @@ -1,5 +1,6 @@ import os import sys +import asyncio import yaml @@ -11,6 +12,7 @@ from ..models.channel import mgr as chanmgr from ..models.key import mgr as keymgr from ..models.router import group as routergroup +from ..models.watchdog import wd as wdmgr from .adapter import revChatGPT from .adapter import claude @@ -35,21 +37,29 @@ class Application: key: keymgr.AbsAPIKeyManager """API Key manager.""" + watchdog: wdmgr.AbsWatchDog + def __init__( self, dbmgr: db.DatabaseInterface, router: routermgr.RouterManager, channel: chanmgr.AbsChannelManager, key: keymgr.AbsAPIKeyManager, + watchdog: wdmgr.AbsWatchDog, ): self.dbmgr = dbmgr self.router = router self.channel = channel self.key = key + self.watchdog = watchdog - def run(self): + async def run(self): """Run application.""" - return self.router.serve() + loop = asyncio.get_running_loop() + + loop.create_task(self.watchdog.run()) + + await self.router.serve(loop) default_config = { "database": { @@ -59,6 +69,7 @@ def run(self): "watchdog": { "heartbeat": { "interval": 1800, + "timeout": 300, "fail_limit": 3, }, }, @@ -137,11 +148,27 @@ async def make_application(config_path: str) -> Application: config=config['router'], ) + # watchdog and tasks + from .watchdog import wd as watchdog + + wdmgr = watchdog.WatchDog() + + # tasks + from .watchdog.tasks import heartbeat + + hbtask = heartbeat.HeartBeatTask( + channelmgr, + config['watchdog']['heartbeat'], + ) + + wdmgr.add_task(hbtask) + app = Application( dbmgr=dbmgr, router=routermgr, channel=channelmgr, key=apikeymgr, + watchdog=wdmgr, ) return app diff --git a/free_one_api/impls/router/mgr.py b/free_one_api/impls/router/mgr.py index 24688de..ba7a4f6 100644 --- a/free_one_api/impls/router/mgr.py +++ b/free_one_api/impls/router/mgr.py @@ -21,9 +21,9 @@ def __init__(self, routes: list[tuple[str, list[str], callable, dict]], config: for method in methods: self._app.route(route, methods=[method], **kwargs)(handler) - def serve(self): + async def serve(self, loop): """Serve API.""" - self._app.run(host="0.0.0.0", port=self.port) + return await self._app.run_task(host="0.0.0.0", port=self.port) if __name__ == "__main__": diff --git a/free_one_api/impls/watchdog/__init__.py b/free_one_api/impls/watchdog/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/free_one_api/impls/watchdog/tasks/__init__.py b/free_one_api/impls/watchdog/tasks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/free_one_api/impls/watchdog/tasks/heartbeat.py b/free_one_api/impls/watchdog/tasks/heartbeat.py new file mode 100644 index 0000000..1310df8 --- /dev/null +++ b/free_one_api/impls/watchdog/tasks/heartbeat.py @@ -0,0 +1,45 @@ +import asyncio +import traceback +import logging + +from ....models.watchdog import task +from ....models.channel import mgr as chanmgr +from ....entities import channel + + +class HeartBeatTask(task.AbsTask): + """HeartBeat task.""" + + channel: chanmgr.AbsChannelManager + + cfg: dict + + def __init__(self, chan: chanmgr.AbsChannelManager, cfg: dict): + self.channel = chan + self.cfg = cfg + + self.delay = 10 + self.interval = cfg['interval'] + + async def trigger(self): + """Trigger this task.""" + + process_task = [] + + for chan in self.channel.channels: + if chan.enabled: + async def process(ch: channel.Channel): + fail_count = await ch.heartbeat(timeout=self.cfg["timeout"]) + if fail_count > self.cfg["fail_limit"]: + try: + self.channel.disable_channel(ch.id) + logging.info(f"Disabled channel {ch.id} due to heartbeat failed {fail_count} times") + except Exception: + logging.warn(f"Failed to disable channel {ch.id}, traceback: {traceback.format_exc()}") + await self.channel.update_channel(ch) + + process_task.append(process(chan)) + + logging.info(f"Start heartbeat task, {len(process_task)} channels to process") + await asyncio.gather(*process_task) + \ No newline at end of file diff --git a/free_one_api/impls/watchdog/wd.py b/free_one_api/impls/watchdog/wd.py new file mode 100644 index 0000000..9a5b53d --- /dev/null +++ b/free_one_api/impls/watchdog/wd.py @@ -0,0 +1,23 @@ +import asyncio + +from ...models.watchdog import wd +from ...models.watchdog import task + + +class WatchDog(wd.AbsWatchDog): + """WatchDog implementation.""" + + def __init__(self): + self.tasks = [] + + async def run(self): + cor = [] + + for task in self.tasks: + cor.append(task.loop()) + + await asyncio.gather(*cor) + + def add_task(self, task: task.AbsTask): + """Add a task.""" + self.tasks.append(task) diff --git a/free_one_api/models/watchdog/__init__.py b/free_one_api/models/watchdog/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/free_one_api/models/watchdog/task.py b/free_one_api/models/watchdog/task.py new file mode 100644 index 0000000..97179c1 --- /dev/null +++ b/free_one_api/models/watchdog/task.py @@ -0,0 +1,33 @@ +import abc +import asyncio +import traceback +import logging + + +class AbsTask(metaclass=abc.ABCMeta): + """Task of WatchDog. + + A task may need instances of other modules to work. These dependencies should be set by outer constructor. + Task's delay and interval may be set by outer constructor also, and be scheduled by WatchDog implementation. + """ + + delay: int = 0 + """Delay before first trigger.""" + + interval: int = 60 + """Interval between two triggers.""" + + @abc.abstractmethod + async def trigger(self): + """Trigger this task.""" + raise NotImplementedError + + async def loop(self): + """Loop this task.""" + await asyncio.sleep(self.delay) + while True: + try: + await self.trigger() + except Exception: + logging.warn(f"Failed to trigger task {self.__class__.__name__}, traceback: {traceback.format_exc()}") + await asyncio.sleep(self.interval) diff --git a/free_one_api/models/watchdog/wd.py b/free_one_api/models/watchdog/wd.py new file mode 100644 index 0000000..4bfbdda --- /dev/null +++ b/free_one_api/models/watchdog/wd.py @@ -0,0 +1,14 @@ +import abc + +from . import task + +class AbsWatchDog(metaclass=abc.ABCMeta): + """Model of WatchDog.""" + + tasks: list[task.AbsTask] + """Added tasks.""" + + @abc.abstractmethod + async def run(self): + """Run WatchDog system.""" + raise NotImplementedError diff --git a/main.py b/main.py index add8f01..c1e5dae 100644 --- a/main.py +++ b/main.py @@ -1,6 +1,9 @@ import os import sys import asyncio +import logging + +logging.basicConfig(level=logging.INFO) if not os.path.exists('./data'): os.mkdir('./data') @@ -8,9 +11,11 @@ from free_one_api.impls import app def main(): - application = asyncio.run(app.make_application("./data/config.yaml")) + loop = asyncio.get_event_loop() + + application = loop.run_until_complete(app.make_application("./data/config.yaml")) - application.run() + loop.run_until_complete(application.run()) if __name__ == "__main__": main() diff --git a/requirements.txt b/requirements.txt index bf31d84..94340d5 100644 --- a/requirements.txt +++ b/requirements.txt @@ -7,4 +7,5 @@ claude-api bardapi hugchat g4f -revTongYi \ No newline at end of file +revTongYi +colorlog \ No newline at end of file