✨ Feature: Add database, count the first character time.

.dockerignore

@@ -1,3 +1,7 @@
 api.yaml
 test
-json_str
+json_str
+*.jpg
+*.json
+*.png
+*.db

.gitignore

@@ -8,4 +8,5 @@ node_modules
 .pytest_cache
 *.jpg
 *.json
-*.png
+*.png
+*.db

README.md

@@ -133,7 +133,9 @@ Start the container
 
 ```bash
 docker run --user root -p 8001:8000 --name uni-api -dit \
--v ./api.yaml:/home/api.yaml \
+-e CONFIG_URL=http://file_url/api.yaml \ # If the local configuration file is already mounted, you do not need to set CONFIG_URL
+-v ./api.yaml:/home/api.yaml \ # If CONFIG_URL is already set, you do not need to mount the configuration file
+-v ./stats.db:/home/stats.db \ # If you do not want to save statistical data, you do not need to mount the stats.db file
 yym68686/uni-api:latest
 ```
 
@@ -145,14 +147,15 @@ services:
     container_name: uni-api
     image: yym68686/uni-api:latest
     environment:
-      - CONFIG_URL=http://file_url/api.yaml
+      - CONFIG_URL=http://file_url/api.yaml # If the local configuration file is already mounted, there is no need to set CONFIG_URL
     ports:
       - 8001:8000
     volumes:
-      - ./api.yaml:/home/api.yaml
+      - ./api.yaml:/home/api.yaml # If CONFIG_URL is already set, there is no need to mount the configuration file
+      - ./stats.db:/home/stats.db # If you do not want to save statistical data, there is no need to mount the stats.db file
 ```
 
-CONFIG_URL is a link that can automatically download a remote configuration file. For example, if you find it inconvenient to modify the configuration file on a certain platform, you can upload the configuration file to a hosting service that provides a direct link for uni-api to download. CONFIG_URL is this direct link.
+CONFIG_URL is used to automatically download remote configuration files. For example, if it is inconvenient to modify the configuration file on a certain platform, you can upload the configuration file to a hosting service and provide a direct link for uni-api to download. CONFIG_URL is this direct link. If you are using a locally mounted configuration file, you do not need to set CONFIG_URL. CONFIG_URL is used in situations where it is inconvenient to mount the configuration file.
 
 Run Docker Compose container in the background
 
@@ -178,6 +181,7 @@ docker rm -f uni-api
 docker run --user root -p 8001:8000 -dit --name uni-api \
 -e CONFIG_URL=http://file_url/api.yaml \
 -v ./api.yaml:/home/api.yaml \
+-v ./stats.db:/home/stats.db \
 yym68686/uni-api:latest
 docker logs -f uni-api
 ```

README_CN.md

@@ -133,7 +133,9 @@ Start the container
 
 ```bash
 docker run --user root -p 8001:8000 --name uni-api -dit \
--v ./api.yaml:/home/api.yaml \
+-e CONFIG_URL=http://file_url/api.yaml \ # 如果已经挂载了本地配置文件，不需要设置 CONFIG_URL
+-v ./api.yaml:/home/api.yaml \ # 如果已经设置 CONFIG_URL，不需要挂载配置文件
+-v ./stats.db:/home/stats.db \ # 如果不想保存统计数据，不需要挂载 stats.db 文件
 yym68686/uni-api:latest
 ```
 
@@ -145,14 +147,15 @@ services:
     container_name: uni-api
     image: yym68686/uni-api:latest
     environment:
-      - CONFIG_URL=http://file_url/api.yaml
+      - CONFIG_URL=http://file_url/api.yaml # 如果已经挂载了本地配置文件，不需要设置 CONFIG_URL
     ports:
       - 8001:8000
     volumes:
-      - ./api.yaml:/home/api.yaml
+      - ./api.yaml:/home/api.yaml # 如果已经设置 CONFIG_URL，不需要挂载配置文件
+      - ./stats.db:/home/stats.db # 如果不想保存统计数据，不需要挂载 stats.db 文件
 ```
 
-CONFIG_URL 就是可以自动下载远程的配置文件。比如你在某个平台不方便修改配置文件，可以把配置文件传到某个托管服务，可以提供直链给 uni-api 下载，CONFIG_URL 就是这个直链。
+CONFIG_URL 就是可以自动下载远程的配置文件。比如你在某个平台不方便修改配置文件，可以把配置文件传到某个托管服务，可以提供直链给 uni-api 下载，CONFIG_URL 就是这个直链。如果使用本地挂载的配置文件，不需要设置 CONFIG_URL。CONFIG_URL 是在不方便挂载配置文件的情况下使用。
 
 Run Docker Compose container in the background
 
@@ -178,6 +181,7 @@ docker rm -f uni-api
 docker run --user root -p 8001:8000 -dit --name uni-api \
 -e CONFIG_URL=http://file_url/api.yaml \
 -v ./api.yaml:/home/api.yaml \
+-v ./stats.db:/home/stats.db \
 yym68686/uni-api:latest
 docker logs -f uni-api
 ```

docker-compose.yml

@@ -7,4 +7,5 @@ services:
     ports:
       - 8001:8000
     volumes:
-      - ./api.yaml:/home/api.yaml
+      - ./api.yaml:/home/api.yaml
+      - ./stats.db:/home/stats.db

main.py

@@ -22,15 +22,16 @@ from typing import List, Dict, Union
 from urllib.parse import urlparse
 
 import os
-is_debug = os.getenv("DEBUG", False)
+is_debug = bool(os.getenv("DEBUG", False))
+
+async def create_tables():
+    async with engine.begin() as conn:
+        await conn.run_sync(Base.metadata.create_all)
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     # 启动时的代码
-
-    # # 启动事件
-    # routes = [{"path": route.path, "name": route.name} for route in app.routes]
-    # logger.info(f"Registered routes: {routes}")
+    await create_tables()
 
     TIMEOUT = float(os.getenv("TIMEOUT", 100))
     timeout = httpx.Timeout(connect=15.0, read=TIMEOUT, write=30.0, pool=30.0)
@@ -66,10 +67,7 @@ import asyncio
 from time import time
 from collections import defaultdict
 from starlette.middleware.base import BaseHTTPMiddleware
-from datetime import datetime
-from datetime import timedelta
 import json
-import aiofiles
 
 async def parse_request_body(request: Request):
     if request.method == "POST" and "application/json" in request.headers.get("content-type", ""):
@@ -79,30 +77,53 @@ async def parse_request_body(request: Request):
             return None
     return None
 
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
+from sqlalchemy.orm import declarative_base, sessionmaker
+from sqlalchemy import Column, Integer, String, Float, DateTime, select, Boolean
+from sqlalchemy.sql import func
+
+# 定义数据库模型
+Base = declarative_base()
+
+class RequestStat(Base):
+    __tablename__ = 'request_stats'
+    id = Column(Integer, primary_key=True)
+    endpoint = Column(String)
+    ip = Column(String)
+    token = Column(String)
+    total_time = Column(Float)
+    model = Column(String)
+    timestamp = Column(DateTime(timezone=True), server_default=func.now())
+
+class ChannelStat(Base):
+    __tablename__ = 'channel_stats'
+    id = Column(Integer, primary_key=True)
+    provider = Column(String)
+    model = Column(String)
+    api_key = Column(String)
+    success = Column(Boolean)
+    first_response_time = Column(Float)  # 新增: 记录首次响应时间
+    timestamp = Column(DateTime(timezone=True), server_default=func.now())
+
+# 创建异步引擎和会话
+engine = create_async_engine('sqlite+aiosqlite:///stats.db', echo=is_debug)
+async_session = sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)
+
 class StatsMiddleware(BaseHTTPMiddleware):
-    def __init__(self, app, exclude_paths=None, save_interval=3600, filename="stats.json"):
+    def __init__(self, app):
         super().__init__(app)
-        self.request_counts = defaultdict(int)
-        self.request_times = defaultdict(float)
-        self.ip_counts = defaultdict(lambda: defaultdict(int))
-        self.request_arrivals = defaultdict(list)
-        self.channel_success_counts = defaultdict(int)
-        self.model_counts = defaultdict(int)
-        self.channel_failure_counts = defaultdict(int)
-        self.lock = asyncio.Lock()
-        self.exclude_paths = set(exclude_paths or [])
-        self.save_interval = save_interval
-        self.filename = filename
-        self.last_save_time = time()
-
-        # 启动定期保存和清理任务
-        asyncio.create_task(self.periodic_save_and_cleanup())
+        self.db = async_session()
 
     async def dispatch(self, request: Request, call_next):
-        arrival_time = datetime.now()
+        if request.headers.get("x-api-key"):
+            token = request.headers.get("x-api-key")
+        elif request.headers.get("Authorization"):
+            token = request.headers.get("Authorization").split(" ")[1]
+        else:
+            token = None
+
         start_time = time()
 
-        # 使用依赖注入获取预解析的请求体
         request.state.parsed_body = await parse_request_body(request)
 
         model = "unknown"
@@ -121,86 +142,35 @@ class StatsMiddleware(BaseHTTPMiddleware):
         endpoint = f"{request.method} {request.url.path}"
         client_ip = request.client.host
 
-        if request.url.path not in self.exclude_paths:
-            async with self.lock:
-                self.request_counts[endpoint] += 1
-                self.request_times[endpoint] += process_time
-                self.ip_counts[endpoint][client_ip] += 1
-                self.request_arrivals[endpoint].append(arrival_time)
-                if model != "unknown":
-                    self.model_counts[model] += 1
+        # 异步更新数据库
+        await self.update_stats(endpoint, process_time, client_ip, model, token)
 
         return response
 
-    async def periodic_save_and_cleanup(self):
-        while True:
-            await asyncio.sleep(self.save_interval)
-            await self.save_stats()
-            await self.cleanup_old_data()
-
-    async def save_stats(self):
-        current_time = time()
-        if current_time - self.last_save_time < self.save_interval:
-            return
-
-        async with self.lock:
-            stats = {
-                "request_counts": dict(self.request_counts),
-                "request_times": dict(self.request_times),
-                "model_counts": dict(self.model_counts),
-                "ip_counts": {k: dict(v) for k, v in self.ip_counts.items()},
-                "request_arrivals": {k: [t.isoformat() for t in v] for k, v in self.request_arrivals.items()},
-                "channel_success_counts": dict(self.channel_success_counts),
-                "channel_failure_counts": dict(self.channel_failure_counts),
-                "channel_success_percentages": self.calculate_success_percentages(),
-                "channel_failure_percentages": self.calculate_failure_percentages()
-            }
-
-        filename = self.filename
-        async with aiofiles.open(filename, mode='w') as f:
-            await f.write(json.dumps(stats, indent=2))
-
-        self.last_save_time = current_time
-
-    def calculate_success_percentages(self):
-        percentages = {}
-        for channel, success_count in self.channel_success_counts.items():
-            total_count = success_count + self.channel_failure_counts[channel]
-            if total_count > 0:
-                percentages[channel] = success_count / total_count * 100
-            else:
-                percentages[channel] = 0
-
-        sorted_percentages = dict(sorted(percentages.items(), key=lambda item: item[1], reverse=True))
-        return sorted_percentages
-
-    def calculate_failure_percentages(self):
-        percentages = {}
-        for channel, failure_count in self.channel_failure_counts.items():
-            total_count = failure_count + self.channel_success_counts[channel]
-            if total_count > 0:
-                percentages[channel] = failure_count / total_count * 100
-            else:
-                percentages[channel] = 0
-
-        sorted_percentages = dict(sorted(percentages.items(), key=lambda item: item[1], reverse=True))
-        return sorted_percentages
-
-    async def cleanup_old_data(self):
-        cutoff_time = datetime.now() - timedelta(hours=24)
-        async with self.lock:
-            for endpoint in list(self.request_arrivals.keys()):
-                self.request_arrivals[endpoint] = [
-                    t for t in self.request_arrivals[endpoint] if t > cutoff_time
-                ]
-                if not self.request_arrivals[endpoint]:
-                    del self.request_arrivals[endpoint]
-                    self.request_counts.pop(endpoint, None)
-                    self.request_times.pop(endpoint, None)
-                    self.ip_counts.pop(endpoint, None)
-
-    async def cleanup(self):
-        await self.save_stats()
+    async def update_stats(self, endpoint, process_time, client_ip, model, token):
+        async with self.db as session:
+            # 为每个请求创建一条新的记录
+            new_request_stat = RequestStat(
+                endpoint=endpoint,
+                ip=client_ip,
+                token=token,
+                total_time=process_time,
+                model=model
+            )
+            session.add(new_request_stat)
+            await session.commit()
+
+    async def update_channel_stats(self, provider, model, api_key, success, first_response_time):
+        async with self.db as session:
+            channel_stat = ChannelStat(
+                provider=provider,
+                model=model,
+                api_key=api_key,
+                success=success,
+                first_response_time=first_response_time
+            )
+            session.add(channel_stat)
+            await session.commit()
 
 # 配置 CORS 中间件
 app.add_middleware(
@@ -211,10 +181,10 @@ app.add_middleware(
     allow_headers=["*"],  # 允许所有头部字段
 )
 
-app.add_middleware(StatsMiddleware, exclude_paths=["/stats", "/generate-api-key"])
+app.add_middleware(StatsMiddleware)
 
 # 在 process_request 函数中更新成功和失败计数
-async def process_request(request: Union[RequestModel, ImageGenerationRequest], provider: Dict, endpoint=None):
+async def process_request(request: Union[RequestModel, ImageGenerationRequest], provider: Dict, endpoint=None, token=None):
     url = provider['base_url']
     parsed_url = urlparse(url)
     # print("parsed_url", parsed_url)
@@ -269,25 +239,23 @@ async def process_request(request: Union[RequestModel, ImageGenerationRequest],
         if request.stream:
             model = provider['model'][request.model]
             generator = fetch_response_stream(app.state.client, url, headers, payload, engine, model)
-            wrapped_generator = await error_handling_wrapper(generator)
+            wrapped_generator, first_response_time = await error_handling_wrapper(generator)
             response = StreamingResponse(wrapped_generator, media_type="text/event-stream")
         else:
             generator = fetch_response(app.state.client, url, headers, payload)
-            wrapped_generator = await error_handling_wrapper(generator)
+            wrapped_generator, first_response_time = await error_handling_wrapper(generator)
             first_element = await anext(wrapped_generator)
             first_element = first_element.lstrip("data: ")
             first_element = json.loads(first_element)
             response = JSONResponse(first_element)
 
-        # 更新成功计数
-        async with app.middleware_stack.app.lock:
-            app.middleware_stack.app.channel_success_counts[provider['provider']] += 1
+        # 更新成功计数和首次响应时间
+        await app.middleware_stack.app.update_channel_stats(provider['provider'], request.model, token, success=True, first_response_time=first_response_time)
 
         return response
     except (Exception, HTTPException, asyncio.CancelledError, httpx.ReadError, httpx.RemoteProtocolError) as e:
-        # 更新失败计数
-        async with app.middleware_stack.app.lock:
-            app.middleware_stack.app.channel_failure_counts[provider['provider']] += 1
+        # 更新失败计数,首次响应时间为-1表示失败
+        await app.middleware_stack.app.update_channel_stats(provider['provider'], request.model, token, success=False, first_response_time=-1)
 
         raise e
 
@@ -421,10 +389,10 @@ class ModelRequestHandler:
         if safe_get(config, 'api_keys', api_index, "preferences", "AUTO_RETRY") == False:
             auto_retry = False
 
-        return await self.try_all_providers(request, matching_providers, use_round_robin, auto_retry, endpoint)
+        return await self.try_all_providers(request, matching_providers, use_round_robin, auto_retry, endpoint, token)
 
     # 在 try_all_providers 函数中处理失败的情况
-    async def try_all_providers(self, request: Union[RequestModel, ImageGenerationRequest], providers: List[Dict], use_round_robin: bool, auto_retry: bool, endpoint: str = None):
+    async def try_all_providers(self, request: Union[RequestModel, ImageGenerationRequest], providers: List[Dict], use_round_robin: bool, auto_retry: bool, endpoint: str = None, token: str = None):
         status_code = 500
         error_message = None
         num_providers = len(providers)
@@ -433,7 +401,7 @@ class ModelRequestHandler:
             self.last_provider_index = (start_index + i) % num_providers
             provider = providers[self.last_provider_index]
             try:
-                response = await process_request(request, provider, endpoint)
+                response = await process_request(request, provider, endpoint, token)
                 return response
             except HTTPException as e:
                 logger.error(f"Error with provider {provider['provider']}: {str(e)}")
@@ -510,6 +478,7 @@ async def get_user_rate_limit(api_index: str = None):
     return rate_limit
 
 security = HTTPBearer()
+
 async def rate_limit_dependency(request: Request, credentials: HTTPAuthorizationCredentials = Depends(security)):
     token = credentials.credentials if credentials else None
     api_list = app.state.api_list
@@ -576,24 +545,96 @@ def generate_api_key():
     return JSONResponse(content={"api_key": api_key})
 
 # 在 /stats 路由中返回成功和失败百分比
+from collections import defaultdict
+from sqlalchemy import func
+
+from collections import defaultdict
+from sqlalchemy import func, desc, case
+
 @app.get("/stats", dependencies=[Depends(rate_limit_dependency)])
 async def get_stats(request: Request, token: str = Depends(verify_admin_api_key)):
-    middleware = app.middleware_stack.app
-    if isinstance(middleware, StatsMiddleware):
-        async with middleware.lock:
-            stats = {
-                "channel_success_percentages": middleware.calculate_success_percentages(),
-                "channel_failure_percentages": middleware.calculate_failure_percentages(),
-                "model_counts": dict(middleware.model_counts),
-                "request_counts": dict(middleware.request_counts),
-                "request_times": dict(middleware.request_times),
-                "ip_counts": {k: dict(v) for k, v in middleware.ip_counts.items()},
-                "request_arrivals": {k: [t.isoformat() for t in v] for k, v in middleware.request_arrivals.items()},
-                "channel_success_counts": dict(middleware.channel_success_counts),
-                "channel_failure_counts": dict(middleware.channel_failure_counts),
-            }
-        return JSONResponse(content=stats)
-    return {"error": "StatsMiddleware not found"}
+    async with async_session() as session:
+        # 1. 每个渠道下面每个模型的成功率
+        channel_model_stats = await session.execute(
+            select(
+                ChannelStat.provider,
+                ChannelStat.model,
+                func.count().label('total'),
+                func.sum(case((ChannelStat.success == True, 1), else_=0)).label('success_count')
+            ).group_by(ChannelStat.provider, ChannelStat.model)
+        )
+        channel_model_stats = channel_model_stats.fetchall()
+
+        # 2. 每个渠道总的成功率
+        channel_stats = await session.execute(
+            select(
+                ChannelStat.provider,
+                func.count().label('total'),
+                func.sum(case((ChannelStat.success == True, 1), else_=0)).label('success_count')
+            ).group_by(ChannelStat.provider)
+        )
+        channel_stats = channel_stats.fetchall()
+
+        # 3. 每个模型在所有渠道总的请求次数
+        model_stats = await session.execute(
+            select(ChannelStat.model, func.count().label('count'))
+            .group_by(ChannelStat.model)
+            .order_by(desc('count'))
+        )
+        model_stats = model_stats.fetchall()
+
+        # 4. 每个端点的请求次数
+        endpoint_stats = await session.execute(
+            select(RequestStat.endpoint, func.count().label('count'))
+            .group_by(RequestStat.endpoint)
+            .order_by(desc('count'))
+        )
+        endpoint_stats = endpoint_stats.fetchall()
+
+        # 5. 每个ip请求的次数
+        ip_stats = await session.execute(
+            select(RequestStat.ip, func.count().label('count'))
+            .group_by(RequestStat.ip)
+            .order_by(desc('count'))
+        )
+        ip_stats = ip_stats.fetchall()
+
+    # 处理统计数据并返回
+    stats = {
+        "channel_model_success_rates": [
+            {
+                "provider": stat.provider,
+                "model": stat.model,
+                "success_rate": stat.success_count / stat.total if stat.total > 0 else 0
+            } for stat in sorted(channel_model_stats, key=lambda x: x.success_count / x.total if x.total > 0 else 0, reverse=True)
+        ],
+        "channel_success_rates": [
+            {
+                "provider": stat.provider,
+                "success_rate": stat.success_count / stat.total if stat.total > 0 else 0
+            } for stat in sorted(channel_stats, key=lambda x: x.success_count / x.total if x.total > 0 else 0, reverse=True)
+        ],
+        "model_request_counts": [
+            {
+                "model": stat.model,
+                "count": stat.count
+            } for stat in model_stats
+        ],
+        "endpoint_request_counts": [
+            {
+                "endpoint": stat.endpoint,
+                "count": stat.count
+            } for stat in endpoint_stats
+        ],
+        "ip_request_counts": [
+            {
+                "ip": stat.ip,
+                "count": stat.count
+            } for stat in ip_stats
+        ]
+    }
+
+    return JSONResponse(content=stats)
 
 # async def on_fetch(request, env):
 #     import asgi

requirements.txt

@@ -3,6 +3,9 @@ pytest
 uvicorn
 fastapi
 aiofiles
+greenlet
+aiosqlite
+sqlalchemy
 watchfiles
 httpx[http2]
 cryptography

test/provider_test.py

@@ -70,7 +70,8 @@ def test_request_model(test_client, api_key, get_model):
                     }
                 }
             }
-        ]
+        ],
+        "tool_choice": "auto"
     }
 
     headers = {

utils.py

@@ -116,9 +116,12 @@ def ensure_string(item):
         return str(item)
 
 import asyncio
+import time as time_module
 async def error_handling_wrapper(generator):
+    start_time = time_module.time()
     try:
         first_item = await generator.__anext__()
+        first_response_time = time_module.time() - start_time
         first_item_str = first_item
         # logger.info("first_item_str: %s", first_item_str)
         if isinstance(first_item_str, (bytes, bytearray)):
@@ -153,7 +156,7 @@ async def error_handling_wrapper(generator):
                 logger.error(f"Network error in new_generator: {e}")
                 raise
 
-        return new_generator()
+        return new_generator(), first_response_time
 
     except StopAsyncIteration:
         raise HTTPException(status_code=400, detail="data: {'error': 'No data returned'}")