Merge pull request #14 from RockChinQ/feat/heartbeat-checking

Feat: heartbeat and auto-disable

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

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

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

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

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

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__":

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)
+

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)

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)

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

main.py

@@ -1,16 +1,21 @@
 import os
 import sys
 import asyncio
+import logging
+
+logging.basicConfig(level=logging.INFO)
 
 if not os.path.exists('./data'):
     os.mkdir('./data')
 
 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()

requirements.txt

@@ -7,4 +7,5 @@ claude-api
 bardapi
 hugchat
 g4f
-revTongYi
+revTongYi
+colorlog