JXJBing commited on
Commit
07a2f32
·
verified ·
1 Parent(s): 010c7a1

Upload 102 files

Browse files
This view is limited to 50 files because it contains too many changes.   See raw diff
Files changed (50) hide show
  1. .gitattributes +8 -0
  2. Dockerfile +16 -0
  3. LICENSE +17 -0
  4. README.md +279 -10
  5. README_ZH.md +283 -0
  6. VERSION +1 -0
  7. app/config/config.py +513 -0
  8. app/core/application.py +156 -0
  9. app/core/constants.py +112 -0
  10. app/core/security.py +90 -0
  11. app/database/__init__.py +3 -0
  12. app/database/connection.py +71 -0
  13. app/database/initialization.py +77 -0
  14. app/database/models.py +129 -0
  15. app/database/services.py +775 -0
  16. app/domain/file_models.py +69 -0
  17. app/domain/gemini_models.py +115 -0
  18. app/domain/image_models.py +20 -0
  19. app/domain/openai_models.py +43 -0
  20. app/exception/exceptions.py +140 -0
  21. app/handler/error_handler.py +32 -0
  22. app/handler/message_converter.py +363 -0
  23. app/handler/response_handler.py +439 -0
  24. app/handler/retry_handler.py +51 -0
  25. app/handler/stream_optimizer.py +143 -0
  26. app/log/logger.py +349 -0
  27. app/main.py +15 -0
  28. app/middleware/middleware.py +81 -0
  29. app/middleware/request_logging_middleware.py +40 -0
  30. app/middleware/smart_routing_middleware.py +210 -0
  31. app/router/config_routes.py +225 -0
  32. app/router/error_log_routes.py +271 -0
  33. app/router/files_routes.py +296 -0
  34. app/router/gemini_routes.py +502 -0
  35. app/router/key_routes.py +83 -0
  36. app/router/openai_compatiable_routes.py +114 -0
  37. app/router/openai_routes.py +176 -0
  38. app/router/routes.py +287 -0
  39. app/router/scheduler_routes.py +57 -0
  40. app/router/stats_routes.py +56 -0
  41. app/router/version_routes.py +37 -0
  42. app/router/vertex_express_routes.py +147 -0
  43. app/scheduler/scheduled_tasks.py +195 -0
  44. app/service/chat/gemini_chat_service.py +555 -0
  45. app/service/chat/openai_chat_service.py +745 -0
  46. app/service/chat/vertex_express_chat_service.py +403 -0
  47. app/service/client/api_client.py +356 -0
  48. app/service/config/config_service.py +261 -0
  49. app/service/embedding/embedding_service.py +87 -0
  50. app/service/embedding/gemini_embedding_service.py +150 -0
.gitattributes CHANGED
@@ -33,3 +33,11 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
 
 
 
 
 
 
 
 
 
33
  *.zip filter=lfs diff=lfs merge=lfs -text
34
  *.zst filter=lfs diff=lfs merge=lfs -text
35
  *tfevents* filter=lfs diff=lfs merge=lfs -text
36
+ files/image.png filter=lfs diff=lfs merge=lfs -text
37
+ files/image1.png filter=lfs diff=lfs merge=lfs -text
38
+ files/image2.png filter=lfs diff=lfs merge=lfs -text
39
+ files/image3.png filter=lfs diff=lfs merge=lfs -text
40
+ files/image4.png filter=lfs diff=lfs merge=lfs -text
41
+ files/image5.png filter=lfs diff=lfs merge=lfs -text
42
+ files/image6.png filter=lfs diff=lfs merge=lfs -text
43
+ files/image7.png filter=lfs diff=lfs merge=lfs -text
Dockerfile ADDED
@@ -0,0 +1,16 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ FROM python:3.10-slim
2
+
3
+ WORKDIR /app
4
+
5
+ # 复制所需文件到容器中
6
+ COPY ./requirements.txt /app
7
+ COPY ./VERSION /app
8
+
9
+ RUN pip install --no-cache-dir -r requirements.txt
10
+ COPY ./app /app/app
11
+
12
+ # Expose port
13
+ EXPOSE 8000
14
+
15
+ # Run the application
16
+ CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000", "--no-access-log"]
LICENSE ADDED
@@ -0,0 +1,17 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ 知识共享署名-非商业性使用 4.0 国际 (CC BY-NC 4.0) 协议
2
+
3
+ 您可以自由地:
4
+ - 共享 — 在任何媒介以任何形式复制、发行本作品
5
+ - 演绎 — 修改、转换或以本作品为基础进行创作
6
+
7
+ 惟须遵守下列条件:
8
+ - 署名 — 您必须给出适当的署名,提供指向本协议的链接,并指明是否(对原作)作了修改。您可以以任何合理方式进行,但不得以任何方式暗示许可方认可您或您的使用。
9
+ - 非商业性使用 — 您不得将本作品用于商业目的,包括但不限于任何形式的商业倒卖、SaaS、API 付费接口、二次销售、打包出售、收费分发或其他直接或间接盈利行为。
10
+
11
+ 如需商业授权,请联系原作者获得书面许可。违者将承担相应法律责任。
12
+
13
+ Creative Commons Attribution-NonCommercial 4.0 International Public License
14
+
15
+ By exercising the Licensed Rights (defined below), You accept and agree to be bound by the terms and conditions of this Creative Commons Attribution-NonCommercial 4.0 International Public License ("Public License"). To the extent this Public License may be interpreted as a contract, You are granted the Licensed Rights in consideration of Your acceptance of these terms and conditions, and the Licensor grants You such rights in consideration of benefits the Licensor receives from making the Licensed Material available under these terms and conditions.
16
+
17
+ Full license text: https://creativecommons.org/licenses/by-nc/4.0/legalcode
README.md CHANGED
@@ -1,13 +1,282 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  ---
2
- title: Gemini
3
- emoji: 📚
4
- colorFrom: pink
5
- colorTo: red
6
- sdk: gradio
7
- sdk_version: 5.45.0
8
- app_file: app.py
9
- pinned: false
10
- license: mit
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
11
  ---
12
 
13
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ [Read this document in Chinese](README_ZH.md)
2
+
3
+ # Gemini Balance - Gemini API Proxy and Load Balancer
4
+
5
+ <p align="center">
6
+ <a href="https://trendshift.io/repositories/13692" target="_blank">
7
+ <img src="https://trendshift.io/api/badge/repositories/13692" alt="snailyp%2Fgemini-balance | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
8
+ </a>
9
+ </p>
10
+
11
+ <p align="center">
12
+ <a href="https://www.python.org/"><img src="https://img.shields.io/badge/Python-3.9%2B-blue.svg" alt="Python"></a>
13
+ <a href="https://fastapi.tiangolo.com/"><img src="https://img.shields.io/badge/FastAPI-0.100%2B-green.svg" alt="FastAPI"></a>
14
+ <a href="https://www.uvicorn.org/"><img src="https://img.shields.io/badge/Uvicorn-running-purple.svg" alt="Uvicorn"></a>
15
+ <a href="https://t.me/+soaHax5lyI0wZDVl"><img src="https://img.shields.io/badge/Telegram-Group-blue.svg?logo=telegram" alt="Telegram Group"></a>
16
+ </p>
17
+
18
+ > ⚠️ **Important**: This project is licensed under the [CC BY-NC 4.0](LICENSE) license. **Any form of commercial resale service is prohibited**.
19
+ > I have never sold this service on any platform. If you encounter someone selling this service, they are a reseller. Please do not be deceived.
20
+
21
+ ---
22
+
23
+ ## 📖 Project Introduction
24
+
25
+ **Gemini Balance** is an application built with Python FastAPI, designed to provide proxy and load balancing functions for the Google Gemini API. It allows you to manage multiple Gemini API Keys and implement key rotation, authentication, model filtering, and status monitoring through simple configuration. Additionally, the project integrates image generation and multiple image hosting upload functions, and supports proxying in the OpenAI API format.
26
+
27
+ <details>
28
+ <summary>📂 View Project Structure</summary>
29
+
30
+ ```plaintext
31
+ app/
32
+ ├── config/ # Configuration management
33
+ ├── core/ # Core application logic (FastAPI instance creation, middleware, etc.)
34
+ ├── database/ # Database models and connections
35
+ ├── domain/ # Business domain objects
36
+ ├── exception/ # Custom exceptions
37
+ ├── handler/ # Request handlers
38
+ ├── log/ # Logging configuration
39
+ ├── main.py # Application entry point
40
+ ├── middleware/ # FastAPI middleware
41
+ ├── router/ # API routes (Gemini, OpenAI, status page, etc.)
42
+ ├── scheduler/ # Scheduled tasks (e.g., Key status check)
43
+ ├── service/ # Business logic services (chat, Key management, statistics, etc.)
44
+ ├── static/ # Static files (CSS, JS)
45
+ ├── templates/ # HTML templates (e.g., Key status page)
46
+ └── utils/ # Utility functions
47
+ ```
48
+ </details>
49
+
50
+ ---
51
+
52
+ ## ✨ Feature Highlights
53
+
54
+ * **Multi-Key Load Balancing**: Supports configuring multiple Gemini API Keys (`API_KEYS`) for automatic sequential polling.
55
+ * **Visual Configuration**: Configurations modified through the admin backend take effect immediately without restarting.
56
+ ![Configuration Panel](files/image4.png)
57
+ * **Dual Protocol API Compatibility**: Supports both Gemini and OpenAI CHAT API formats.
58
+ * OpenAI Base URL: `http://localhost:8000(/hf)/v1`
59
+ * Gemini Base URL: `http://localhost:8000(/gemini)/v1beta`
60
+ * **Image-Text Chat & Modification**: Configure models with `IMAGE_MODELS` to support image-text chat and editing. Use the `configured_model-image` model name to invoke.
61
+ ![Chat with Image Generation](files/image6.png)
62
+ ![Modify Image](files/image7.png)
63
+ * **Web Search**: Configure models with `SEARCH_MODELS` to support web search. Use the `configured_model-search` model name to invoke.
64
+ ![Web Search](files/image8.png)
65
+ * **Key Status Monitoring**: Provides a `/keys_status` page (authentication required) for real-time monitoring.
66
+ ![Monitoring Panel](files/image.png)
67
+ * **Detailed Logging**: Provides detailed error logs for easy troubleshooting.
68
+ ![Call Details](files/image1.png)
69
+ ![Log List](files/image2.png)
70
+ ![Log Details](files/image3.png)
71
+ * **Flexible Key Addition**: Add keys in batches using the `gemini_key` regex, with automatic deduplication.
72
+ ![Add Key](files/image5.png)
73
+ * **Failure Retry & Auto-Disable**: Automatically retries failed API requests (`MAX_RETRIES`) and disables keys after excessive failures (`MAX_FAILURES`).
74
+ * **Comprehensive API Compatibility**:
75
+ * **Embeddings API**: Fully compatible with the OpenAI `embeddings` API format.
76
+ * **Image Generation API**: Adapts the `imagen-3.0-generate-002` model to the OpenAI image generation API format.
77
+ * **Automatic Model List Maintenance**: Automatically fetches and syncs the latest model lists from Gemini and OpenAI.
78
+ * **Proxy Support**: Supports HTTP/SOCKS5 proxies (`PROXIES`).
79
+ * **Docker Support**: Provides Docker images for both AMD and ARM architectures.
80
+ * Image Address: `ghcr.io/snailyp/gemini-balance:latest`
81
+
82
+ ---
83
+
84
+ ## 🚀 Quick Start
85
+
86
+ ### Option 1: Docker Compose (Recommended)
87
+
88
+ 1. **Get `docker-compose.yml`**:
89
+ Download the `docker-compose.yml` file from the project repository.
90
+ 2. **Prepare `.env` file**:
91
+ Copy `.env.example` to `.env` and configure it. Ensure `DATABASE_TYPE` is set to `mysql` and fill in the `MYSQL_*` details.
92
+ 3. **Start Services**:
93
+ In the directory containing `docker-compose.yml` and `.env`, run:
94
+ ```bash
95
+ docker-compose up -d
96
+ ```
97
+
98
+ ### Option 2: Docker Command
99
+
100
+ 1. **Pull Image**:
101
+ ```bash
102
+ docker pull ghcr.io/snailyp/gemini-balance:latest
103
+ ```
104
+ 2. **Prepare `.env` file**:
105
+ Copy `.env.example` to `.env` and configure it.
106
+ 3. **Run Container**:
107
+ ```bash
108
+ docker run -d -p 8000:8000 --name gemini-balance \
109
+ -v ./data:/app/data \
110
+ --env-file .env \
111
+ ghcr.io/snailyp/gemini-balance:latest
112
+ ```
113
+ * `-d`: Detached mode.
114
+ * `-p 8000:8000`: Map container port 8000 to host.
115
+ * `-v ./data:/app/data`: Mount volume for persistent data.
116
+ * `--env-file .env`: Load environment variables.
117
+
118
+ ### Option 3: Local Development
119
+
120
+ 1. **Clone and Install**:
121
+ ```bash
122
+ git clone https://github.com/snailyp/gemini-balance.git
123
+ cd gemini-balance
124
+ pip install -r requirements.txt
125
+ ```
126
+ 2. **Configure Environment**:
127
+ Copy `.env.example` to `.env` and configure it.
128
+ 3. **Start Application**:
129
+ ```bash
130
+ uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload
131
+ ```
132
+ Access the application at `http://localhost:8000`.
133
+
134
+ ---
135
+
136
+ ## ⚙️ API Endpoints
137
+
138
+ ### Gemini API Format (`/gemini/v1beta`)
139
+
140
+ This endpoint is directly forwarded to official Gemini API format endpoint, without advanced features.
141
+
142
+ * `GET /models`: List available Gemini models.
143
+ * `POST /models/{model_name}:generateContent`: Generate content.
144
+ * `POST /models/{model_name}:streamGenerateContent`: Stream content generation.
145
+
146
+ ### OpenAI API Format
147
+
148
+ #### Hugging Face (HF) Compatible
149
+
150
+ If you want to use advanced features, like fake streaming, please use this endpoint.
151
+
152
+ * `GET /hf/v1/models`: List models.
153
+ * `POST /hf/v1/chat/completions`: Chat completion.
154
+ * `POST /hf/v1/embeddings`: Create text embeddings.
155
+ * `POST /hf/v1/images/generations`: Generate images.
156
+
157
+ #### Standard OpenAI
158
+
159
+ This endpoint is directly forwarded to official OpenAI Compatible API format endpoint, without advanced features.
160
+
161
+ * `GET /openai/v1/models`: List models.
162
+ * `POST /openai/v1/chat/completions`: Chat completion (Recommended).
163
+ * `POST /openai/v1/embeddings`: Create text embeddings.
164
+ * `POST /openai/v1/images/generations`: Generate images.
165
+
166
  ---
167
+
168
+ <details>
169
+ <summary>📋 View Full Configuration List</summary>
170
+
171
+ | Configuration Item | Description | Default Value |
172
+ | :--- | :--- | :--- |
173
+ | **Database** | | |
174
+ | `DATABASE_TYPE` | `mysql` or `sqlite` | `mysql` |
175
+ | `SQLITE_DATABASE` | Path for SQLite database file | `default_db` |
176
+ | `MYSQL_HOST` | MySQL host address | `localhost` |
177
+ | `MYSQL_SOCKET` | MySQL socket address | `/var/run/mysqld/mysqld.sock` |
178
+ | `MYSQL_PORT` | MySQL port | `3306` |
179
+ | `MYSQL_USER` | MySQL username | `your_db_user` |
180
+ | `MYSQL_PASSWORD` | MySQL password | `your_db_password` |
181
+ | `MYSQL_DATABASE` | MySQL database name | `defaultdb` |
182
+ | **API** | | |
183
+ | `API_KEYS` | **Required**, list of Gemini API keys | `[]` |
184
+ | `ALLOWED_TOKENS` | **Required**, list of access tokens | `[]` |
185
+ | `AUTH_TOKEN` | Super admin token, defaults to the first of `ALLOWED_TOKENS` | `sk-123456` |
186
+ | `ADMIN_SESSION_EXPIRE` | Admin session expiration time in seconds (5 minutes to 24 hours) | `3600` |
187
+ | `TEST_MODEL` | Model for testing key validity | `gemini-1.5-flash` |
188
+ | `IMAGE_MODELS` | Models supporting image generation | `["gemini-2.0-flash-exp"]` |
189
+ | `SEARCH_MODELS` | Models supporting web search | `["gemini-2.0-flash-exp"]` |
190
+ | `FILTERED_MODELS` | Disabled models | `[]` |
191
+ | `TOOLS_CODE_EXECUTION_ENABLED` | Enable code execution tool | `false` |
192
+ | `SHOW_SEARCH_LINK` | Display search result links in response | `true` |
193
+ | `SHOW_THINKING_PROCESS` | Display model's thinking process | `true` |
194
+ | `THINKING_MODELS` | Models supporting thinking process | `[]` |
195
+ | `THINKING_BUDGET_MAP` | Budget map for thinking function (model:budget) | `{}` |
196
+ | `URL_NORMALIZATION_ENABLED` | Enable smart URL routing | `false` |
197
+ | `URL_CONTEXT_ENABLED` | Enable URL context understanding | `false` |
198
+ | `URL_CONTEXT_MODELS` | Models supporting URL context | `[]` |
199
+ | `BASE_URL` | Gemini API base URL | `https://generativelanguage.googleapis.com/v1beta` |
200
+ | `MAX_FAILURES` | Max failures allowed per key | `3` |
201
+ | `MAX_RETRIES` | Max retries for failed API requests | `3` |
202
+ | `CHECK_INTERVAL_HOURS` | Interval (hours) to re-check disabled keys | `1` |
203
+ | `TIMEZONE` | Application timezone | `Asia/Shanghai` |
204
+ | `TIME_OUT` | Request timeout (seconds) | `300` |
205
+ | `PROXIES` | List of proxy servers | `[]` |
206
+ | **Logging & Security** | | |
207
+ | `LOG_LEVEL` | Log level: `DEBUG`, `INFO`, `WARNING`, `ERROR` | `INFO` |
208
+ | `AUTO_DELETE_ERROR_LOGS_ENABLED` | Auto-delete error logs | `true` |
209
+ | `AUTO_DELETE_ERROR_LOGS_DAYS` | Error log retention period (days) | `7` |
210
+ | `AUTO_DELETE_REQUEST_LOGS_ENABLED`| Auto-delete request logs | `false` |
211
+ | `AUTO_DELETE_REQUEST_LOGS_DAYS` | Request log retention period (days) | `30` |
212
+ | `SAFETY_SETTINGS` | Content safety thresholds (JSON string) | `[{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"}, ...]` |
213
+ | **TTS** | | |
214
+ | `TTS_MODEL` | TTS model name | `gemini-2.5-flash-preview-tts` |
215
+ | `TTS_VOICE_NAME` | TTS voice name | `Zephyr` |
216
+ | `TTS_SPEED` | TTS speed | `normal` |
217
+ | **Image Generation** | | |
218
+ | `PAID_KEY` | Paid API Key for advanced features | `your-paid-api-key` |
219
+ | `CREATE_IMAGE_MODEL` | Image generation model | `imagen-3.0-generate-002` |
220
+ | `UPLOAD_PROVIDER` | Image upload provider: `smms`, `picgo`, `cloudflare_imgbed` | `smms` |
221
+ | `SMMS_SECRET_TOKEN` | SM.MS API Token | `your-smms-token` |
222
+ | `PICGO_API_KEY` | PicoGo API Key | `your-picogo-apikey` |
223
+ | `PICGO_API_URL` | PicoGo API Server URL | `https://www.picgo.net/api/1/upload` |
224
+ | `CLOUDFLARE_IMGBED_URL` | CloudFlare ImgBed upload URL | `https://xxxxxxx.pages.dev/upload` |
225
+ | `CLOUDFLARE_IMGBED_AUTH_CODE`| CloudFlare ImgBed auth key | `your-cloudflare-imgber-auth-code` |
226
+ | `CLOUDFLARE_IMGBED_UPLOAD_FOLDER`| CloudFlare ImgBed upload folder | `""` |
227
+ | **Stream Optimizer** | | |
228
+ | `STREAM_OPTIMIZER_ENABLED` | Enable stream output optimization | `false` |
229
+ | `STREAM_MIN_DELAY` | Minimum stream output delay | `0.016` |
230
+ | `STREAM_MAX_DELAY` | Maximum stream output delay | `0.024` |
231
+ | `STREAM_SHORT_TEXT_THRESHOLD`| Short text threshold | `10` |
232
+ | `STREAM_LONG_TEXT_THRESHOLD` | Long text threshold | `50` |
233
+ | `STREAM_CHUNK_SIZE` | Stream output chunk size | `5` |
234
+ | **Fake Stream** | | |
235
+ | `FAKE_STREAM_ENABLED` | Enable fake streaming | `false` |
236
+ | `FAKE_STREAM_EMPTY_DATA_INTERVAL_SECONDS` | Heartbeat interval for fake streaming (seconds) | `5` |
237
+
238
+ </details>
239
+
240
  ---
241
 
242
+ ## 🤝 Contributing
243
+
244
+ Pull Requests or Issues are welcome.
245
+
246
+ [![Contributors](https://contrib.rocks/image?repo=snailyp/gemini-balance)](https://github.com/snailyp/gemini-balance/graphs/contributors)
247
+
248
+ ## ⭐ Star History
249
+
250
+ [![Star History Chart](https://api.star-history.com/svg?repos=snailyp/gemini-balance&type=Date)](https://star-history.com/#snailyp/gemini-balance&Date)
251
+
252
+ ## 🎉 Special Thanks
253
+
254
+ * [PicGo](https://www.picgo.net/)
255
+ * [SM.MS](https://smms.app/)
256
+ * [CloudFlare-ImgBed](https://github.com/MarSeventh/CloudFlare-ImgBed)
257
+
258
+ ## 🙏 Our Supporters
259
+
260
+ A special shout-out to [DigitalOcean](https://m.do.co/c/b249dd7f3b4c) for providing the rock-solid and dependable cloud infrastructure that keeps this project humming!
261
+
262
+ <a href="https://m.do.co/c/b249dd7f3b4c">
263
+ <img src="files/dataocean.svg" alt="DigitalOcean Logo" width="200"/>
264
+ </a>
265
+
266
+ CDN acceleration and security protection for this project are sponsored by [Tencent EdgeOne](https://edgeone.ai/?from=github).
267
+
268
+ <a href="https://edgeone.ai/?from=github">
269
+ <img src="https://edgeone.ai/media/34fe3a45-492d-4ea4-ae5d-ea1087ca7b4b.png" alt="EdgeOne Logo" width="200"/>
270
+ </a>
271
+
272
+ ## 💖 Friendly Projects
273
+
274
+ * **[OneLine](https://github.com/chengtx809/OneLine)** by [chengtx809](https://github.com/chengtx809) - AI-driven hot event timeline generation tool.
275
+
276
+ ## 🎁 Project Support
277
+
278
+ If you find this project helpful, consider supporting me via [Afdian](https://afdian.com/a/snaily).
279
+
280
+ ## License
281
+
282
+ This project is licensed under the [CC BY-NC 4.0](LICENSE) (Attribution-NonCommercial) license.
README_ZH.md ADDED
@@ -0,0 +1,283 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Gemini Balance - Gemini API 代理和负载均衡器
2
+
3
+ <p align="center">
4
+ <a href="https://trendshift.io/repositories/13692" target="_blank">
5
+ <img src="https://trendshift.io/api/badge/repositories/13692" alt="snailyp%2Fgemini-balance | Trendshift" style="width: 250px; height: 55px;" width="250" height="55"/>
6
+ </a>
7
+ </p>
8
+
9
+ <p align="center">
10
+ <a href="https://www.python.org/"><img src="https://img.shields.io/badge/Python-3.9%2B-blue.svg" alt="Python"></a>
11
+ <a href="https://fastapi.tiangolo.com/"><img src="https://img.shields.io/badge/FastAPI-0.100%2B-green.svg" alt="FastAPI"></a>
12
+ <a href="https://www.uvicorn.org/"><img src="https://img.shields.io/badge/Uvicorn-running-purple.svg" alt="Uvicorn"></a>
13
+ <a href="https://t.me/+soaHax5lyI0wZDVl"><img src="https://img.shields.io/badge/Telegram-Group-blue.svg?logo=telegram" alt="Telegram Group"></a>
14
+ </p>
15
+
16
+ > ⚠️ **重要声明**: 本项目采用 [CC BY-NC 4.0](LICENSE) 协议,**禁止任何形式的商业倒卖服务**。
17
+ > 本人从未在任何平台售卖服务,如遇售卖,均为倒卖行为,请勿上当受骗。
18
+
19
+ ---
20
+
21
+ ## 📖 项目简介
22
+
23
+ **Gemini Balance** 是一个基于 Python FastAPI 构建的应用程序,旨在提供 Google Gemini API 的代理和负载均衡功能。它允许您管理多个 Gemini API Key,并通过简单的配置实现 Key 的轮询、认证、模型过滤和状态监控。此外,项目还集成了图像生成和多种图床上传功能,并支持 OpenAI API 格式的代理。
24
+
25
+ <details>
26
+ <summary>📂 查看项目结构</summary>
27
+
28
+ ```plaintext
29
+ app/
30
+ ├── config/ # 配置管理
31
+ ├── core/ # 核心应用逻辑 (FastAPI 实例创建, 中间件等)
32
+ ├── database/ # 数据库模型和连接
33
+ ├── domain/ # 业务领域对象
34
+ ├── exception/ # 自定义异常
35
+ ├── handler/ # 请求处理器
36
+ ├── log/ # 日志配置
37
+ ├── main.py # 应用入口
38
+ ├── middleware/ # FastAPI 中间件
39
+ ├── router/ # API 路由 (Gemini, OpenAI, 状态页等)
40
+ ├── scheduler/ # 定时任务 (如 Key 状态检查)
41
+ ├── service/ # 业务逻辑服务 (聊天, Key 管理, 统计等)
42
+ ├── static/ # 静态文件 (CSS, JS)
43
+ ├── templates/ # HTML 模板 (如 Key 状态页)
44
+ └── utils/ # 工具函数
45
+ ```
46
+ </details>
47
+
48
+ ---
49
+
50
+ ## ✨ 功能亮点
51
+
52
+ * **多 Key 负载均衡**: 支持配置多个 Gemini API Key (`API_KEYS`),自动按顺序轮询使用,提高可用性和并发能力。
53
+ * **可视化配置即时生效**: 通过管理后台修改配置后,无需重启服务即可生效。
54
+ ![配置面板](files/image4.png)
55
+ * **双协议 API 兼容**: 同时支持 Gemini 和 OpenAI 格式的 CHAT API 请求转发。
56
+ * OpenAI Base URL: `http://localhost:8000(/hf)/v1`
57
+ * Gemini Base URL: `http://localhost:8000(/gemini)/v1beta`
58
+ * **图文对话与修图**: 通过 `IMAGE_MODELS` 配置支持图文对话和修图功能的模型,调用时使用 `配置模型-image` 模型名。
59
+ ![对话生图](files/image6.png)
60
+ ![修改图片](files/image7.png)
61
+ * **联网搜索**: 通过 `SEARCH_MODELS` 配置支持联网搜索的模型,调用时使用 `配置模型-search` 模型名。
62
+ ![联网搜索](files/image8.png)
63
+ * **Key 状态监控**: 提供 `/keys_status` 页面(需要认证),实时查看各 Key 的状态和使用情况。
64
+ ![监控面板](files/image.png)
65
+ * **详细日志记录**: 提供详细的错误日志,方便排查问题。
66
+ ![调用详情](files/image1.png)
67
+ ![日志列表](files/image2.png)
68
+ ![日志详情](files/image3.png)
69
+ * **灵活的密钥添加**: 支持通过正则表达式 `gemini_key` 批量添加密钥,并自动去重。
70
+ ![添加密钥](files/image5.png)
71
+ * **失败重试与自动禁用**: 自动处理 API 请求失败,进行重试 (`MAX_RETRIES`),并在 Key 失效次数过多时自动禁用 (`MAX_FAILURES`),定时检查恢复 (`CHECK_INTERVAL_HOURS`)。
72
+ * **全面的 API 兼容**:
73
+ * **Embeddings 接口**: 完美适配 OpenAI 格式的 `embeddings` 接口。
74
+ * **画图接口**: 将 `imagen-3.0-generate-002` 模型接口改造为 OpenAI 画图接口格式。
75
+ * **模型列表自动维护**: 自动获取并同步 Gemini 和 OpenAI 的最新模型列表,兼容 New API。
76
+ * **代理支持**: 支持配置 HTTP/SOCKS5 代理 (`PROXIES`),方便在特殊网络环境下使用。
77
+ * **Docker 支持**: 提供 AMD 和 ARM 架构的 Docker 镜像,方便快速部署。
78
+ * 镜像地址: `ghcr.io/snailyp/gemini-balance:latest`
79
+
80
+ ---
81
+
82
+ ## 🚀 快速开始
83
+
84
+ ### 方式一:使用 Docker Compose (推荐)
85
+
86
+ 这是最推荐的部署方式,可以一键启动应用和数据库。
87
+
88
+ 1. **下载 `docker-compose.yml`**:
89
+ 从项目仓库获取 `docker-compose.yml` 文件。
90
+ 2. **准备 `.env` 文件**:
91
+ 从 `.env.example` 复制一份并重命名为 `.env`,然后根据需求修改配置。特别注意,`DATABASE_TYPE` 应设置为 `mysql`,并填写 `MYSQL_*` 相关配置。
92
+ 3. **启动服务**:
93
+ 在 `docker-compose.yml` 和 `.env` 文件所在的目录下,运行以下命令:
94
+ ```bash
95
+ docker-compose up -d
96
+ ```
97
+ 该命令会以后台模式启动 `gemini-balance` 应用和 `mysql` 数据库。
98
+
99
+ ### 方式二:使用 Docker 命令
100
+
101
+ 1. **拉取镜像**:
102
+ ```bash
103
+ docker pull ghcr.io/snailyp/gemini-balance:latest
104
+ ```
105
+ 2. **准备 `.env` 文件**:
106
+ 从 `.env.example` 复制一份并重命名为 `.env`,然后根据需求修改配置。
107
+ 3. **运行容器**:
108
+ ```bash
109
+ docker run -d -p 8000:8000 --name gemini-balance \
110
+ -v ./data:/app/data \
111
+ --env-file .env \
112
+ ghcr.io/snailyp/gemini-balance:latest
113
+ ```
114
+ * `-d`: 后台运行。
115
+ * `-p 8000:8000`: 将容器的 8000 端口映射到主机。
116
+ * `-v ./data:/app/data`: 挂载数据卷以持久化 SQLite 数据和日志。
117
+ * `--env-file .env`: 加载环境变量配置文件。
118
+
119
+ ### 方式三:本地运行 (适用于开发)
120
+
121
+ 1. **克隆仓库并安装依赖**:
122
+ ```bash
123
+ git clone https://github.com/snailyp/gemini-balance.git
124
+ cd gemini-balance
125
+ pip install -r requirements.txt
126
+ ```
127
+ 2. **配置环境变量**:
128
+ 从 `.env.example` 复制一份并重命名为 `.env`,然后根据需求修改配置。
129
+ 3. **启动应用**:
130
+ ```bash
131
+ uvicorn app.main:app --host 0.0.0.0 --port 8000 --reload
132
+ ```
133
+ 应用启动后,访问 `http://localhost:8000`。
134
+
135
+ ---
136
+
137
+ ## ⚙️ API 端点
138
+
139
+ ### Gemini API 格式 (`/gemini/v1beta`)
140
+
141
+ 此端点将请求直接转发到官方 Gemini API 格式的端点,不包含高级功能。
142
+
143
+ * `GET /models`: 列出可用的 Gemini 模型。
144
+ * `POST /models/{model_name}:generateContent`: 生成内容。
145
+ * `POST /models/{model_name}:streamGenerateContent`: 流式生成内容。
146
+
147
+ ### OpenAI API 格式
148
+
149
+ #### 兼容 huggingface (HF) 格式
150
+
151
+ 如果您需要使用高级功能(例如假流式输出),请使用此端点。
152
+
153
+ * `GET /hf/v1/models`: 列出模型。
154
+ * `POST /hf/v1/chat/completions`: 聊天补全。
155
+ * `POST /hf/v1/embeddings`: 创建文本嵌入。
156
+ * `POST /hf/v1/images/generations`: 生成图像。
157
+
158
+ #### 标准 OpenAI 格式
159
+
160
+ 此端点直接转发至官方的 OpenAI 兼容 API 格式端点,不包含高级功能。
161
+
162
+ * `GET /openai/v1/models`: 列出模型。
163
+ * `POST /openai/v1/chat/completions`: 聊天补全 (推荐,速度更快,防截断)。
164
+ * `POST /openai/v1/embeddings`: 创建文本嵌入。
165
+ * `POST /openai/v1/images/generations`: 生成图像。
166
+
167
+ ---
168
+
169
+ <details>
170
+ <summary>📋 查看完整配置项列表</summary>
171
+
172
+ | 配置项 | 说明 | 默认值 |
173
+ | :--- | :--- | :--- |
174
+ | **数据库配置** | | |
175
+ | `DATABASE_TYPE` | 数据库类型: `mysql` 或 `sqlite` | `mysql` |
176
+ | `SQLITE_DATABASE` | 当使用 `sqlite` 时必填,SQLite 数据库文件路径 | `default_db` |
177
+ | `MYSQL_HOST` | 当使用 `mysql` 时必填,MySQL 数据库主机地址 | `localhost` |
178
+ | `MYSQL_SOCKET` | 可选,MySQL 数据库 socket 地址 | `/var/run/mysqld/mysqld.sock` |
179
+ | `MYSQL_PORT` | 当使用 `mysql` 时必填,MySQL 数据库端口 | `3306` |
180
+ | `MYSQL_USER` | 当使用 `mysql` 时必填,MySQL 数据库用户名 | `your_db_user` |
181
+ | `MYSQL_PASSWORD` | 当使用 `mysql` 时必填,MySQL 数据库密码 | `your_db_password` |
182
+ | `MYSQL_DATABASE` | 当使用 `mysql` 时必填,MySQL 数据库名称 | `defaultdb` |
183
+ | **API 相关配置** | | |
184
+ | `API_KEYS` | **必填**, Gemini API 密钥列表,用于负载均衡 | `[]` |
185
+ | `ALLOWED_TOKENS` | **必填**, 允许访问的 Token 列表 | `[]` |
186
+ | `AUTH_TOKEN` | 超级管理员 Token,不填则使用 `ALLOWED_TOKENS` 的第一个 | `sk-123456` |
187
+ | `TEST_MODEL` | 用于测试密钥可用性的模型 | `gemini-1.5-flash` |
188
+ | `IMAGE_MODELS` | 支持绘图功能的模型列表 | `["gemini-2.0-flash-exp"]` |
189
+ | `SEARCH_MODELS` | 支持搜索功能的模型列表 | `["gemini-2.0-flash-exp"]` |
190
+ | `FILTERED_MODELS` | 被禁用的模型列表 | `[]` |
191
+ | `TOOLS_CODE_EXECUTION_ENABLED` | 是否启用代码执行工具 | `false` |
192
+ | `SHOW_SEARCH_LINK` | 是否在响应中显示搜索结果链接 | `true` |
193
+ | `SHOW_THINKING_PROCESS` | 是否显示模型思考过程 | `true` |
194
+ | `THINKING_MODELS` | 支持思考功能的模型列表 | `[]` |
195
+ | `THINKING_BUDGET_MAP` | 思考功能预算映射 (模型名:预算值) | `{}` |
196
+ | `URL_NORMALIZATION_ENABLED` | 是否启用智能路由映射功能 | `false` |
197
+ | `URL_CONTEXT_ENABLED` | 是否启用URL上下文理解功能 | `false` |
198
+ | `URL_CONTEXT_MODELS` | 支持URL上下文理解功能的模型列表 | `[]` |
199
+ | `BASE_URL` | Gemini API 基础 URL | `https://generativelanguage.googleapis.com/v1beta` |
200
+ | `MAX_FAILURES` | 单个 Key 允许的最大失败次数 | `3` |
201
+ | `MAX_RETRIES` | API 请求失败时的最大重试次数 | `3` |
202
+ | `CHECK_INTERVAL_HOURS` | 禁用 Key 恢复检查间隔 (小时) | `1` |
203
+ | `TIMEZONE` | 应用程序使用的时区 | `Asia/Shanghai` |
204
+ | `TIME_OUT` | 请求���时时间 (秒) | `300` |
205
+ | `PROXIES` | 代理服务器列表 (例如 `http://user:pass@host:port`) | `[]` |
206
+ | **日志与安全** | | |
207
+ | `LOG_LEVEL` | 日志级别: `DEBUG`, `INFO`, `WARNING`, `ERROR` | `INFO` |
208
+ | `AUTO_DELETE_ERROR_LOGS_ENABLED` | 是否自动删除错误日志 | `true` |
209
+ | `AUTO_DELETE_ERROR_LOGS_DAYS` | 错误日志保留天数 | `7` |
210
+ | `AUTO_DELETE_REQUEST_LOGS_ENABLED`| 是否自动删除请求日志 | `false` |
211
+ | `AUTO_DELETE_REQUEST_LOGS_DAYS` | 请求日志保留天数 | `30` |
212
+ | `SAFETY_SETTINGS` | 内容安全阈值 (JSON 字符串) | `[{"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"}, ...]` |
213
+ | **TTS 相关** | | |
214
+ | `TTS_MODEL` | TTS 模型名称 | `gemini-2.5-flash-preview-tts` |
215
+ | `TTS_VOICE_NAME` | TTS 语音名称 | `Zephyr` |
216
+ | `TTS_SPEED` | TTS 语速 | `normal` |
217
+ | **图像生成相关** | | |
218
+ | `PAID_KEY` | 付费版API Key,用于图片生成等高级功能 | `your-paid-api-key` |
219
+ | `CREATE_IMAGE_MODEL` | 图片生成模型 | `imagen-3.0-generate-002` |
220
+ | `UPLOAD_PROVIDER` | 图片上传提供商: `smms`, `picgo`, `cloudflare_imgbed` | `smms` |
221
+ | `SMMS_SECRET_TOKEN` | SM.MS图床的API Token | `your-smms-token` |
222
+ | `PICGO_API_KEY` | [PicoGo](https://www.picgo.net/)图床的API Key | `your-picogo-apikey` |
223
+ | `PICGO_API_URL` | [PicoGo](https://www.picgo.net/)图床的API服务器地址 | `https://www.picgo.net/api/1/upload` |
224
+ | `CLOUDFLARE_IMGBED_URL` | [CloudFlare](https://github.com/MarSeventh/CloudFlare-ImgBed) 图床上传地址 | `https://xxxxxxx.pages.dev/upload` |
225
+ | `CLOUDFLARE_IMGBED_AUTH_CODE`| CloudFlare图床的鉴权key | `your-cloudflare-imgber-auth-code` |
226
+ | `CLOUDFLARE_IMGBED_UPLOAD_FOLDER`| CloudFlare图床的上传文件夹路径 | `""` |
227
+ | **流式优化器相关** | | |
228
+ | `STREAM_OPTIMIZER_ENABLED` | 是否启用流式输出优化 | `false` |
229
+ | `STREAM_MIN_DELAY` | 流式输出最小延迟 | `0.016` |
230
+ | `STREAM_MAX_DELAY` | 流式输出最大延迟 | `0.024` |
231
+ | `STREAM_SHORT_TEXT_THRESHOLD`| 短文本阈值 | `10` |
232
+ | `STREAM_LONG_TEXT_THRESHOLD` | 长文本阈值 | `50` |
233
+ | `STREAM_CHUNK_SIZE` | 流式输出块大小 | `5` |
234
+ | **伪流式 (Fake Stream) 相关** | | |
235
+ | `FAKE_STREAM_ENABLED` | 是否启用伪流式传输 | `false` |
236
+ | `FAKE_STREAM_EMPTY_DATA_INTERVAL_SECONDS` | 伪流式传输时发送心跳空数据的间隔秒数 | `5` |
237
+
238
+ </details>
239
+
240
+ ---
241
+
242
+ ## 🤝 贡献
243
+
244
+ 欢迎通过提交 Pull Request 或 Issue 来为项目做出贡献。
245
+
246
+ [![Contributors](https://contrib.rocks/image?repo=snailyp/gemini-balance)](https://github.com/snailyp/gemini-balance/graphs/contributors)
247
+
248
+ ## ⭐ Star History
249
+
250
+ [![Star History Chart](https://api.star-history.com/svg?repos=snailyp/gemini-balance&type=Date)](https://star-history.com/#snailyp/gemini-balance&Date)
251
+
252
+ ## 🎉 特别鸣谢
253
+
254
+ * [PicGo](https://www.picgo.net/)
255
+ * [SM.MS](https://smms.app/)
256
+ * [CloudFlare-ImgBed](https://github.com/MarSeventh/CloudFlare-ImgBed)
257
+
258
+ ## 💖 友情项目
259
+
260
+ * **[OneLine](https://github.com/chengtx809/OneLine)** by [chengtx809](https://github.com/chengtx809) - AI 驱动的热点事件时间轴生成工具。
261
+
262
+ ## 🎁 项目支持
263
+
264
+ 如果你觉得这个项目对你有帮助,可以考虑通过 [爱发电](https://afdian.com/a/snaily) 支持我。
265
+
266
+ ## 许可证
267
+
268
+ 本项目采用 [CC BY-NC 4.0](LICENSE)(署名-非商业性使用)协议。
269
+
270
+
271
+ ## 赞助商
272
+
273
+ 特别感谢 [DigitalOcean](https://m.do.co/c/b249dd7f3b4c) 为本项目提供稳定可靠的云基础设施支持。
274
+
275
+ <a href="https://m.do.co/c/b249dd7f3b4c">
276
+ <img src="files/dataocean.svg" alt="DigitalOcean Logo" width="200"/>
277
+ </a>
278
+
279
+ 本项目的 CDN 加速和安全防护由 [Tencent EdgeOne](https://edgeone.ai/?from=github) 赞助。
280
+
281
+ <a href="https://edgeone.ai/?from=github">
282
+ <img src="https://edgeone.ai/media/34fe3a45-492d-4ea4-ae5d-ea1087ca7b4b.png" alt="EdgeOne Logo" width="200"/>
283
+ </a>
VERSION ADDED
@@ -0,0 +1 @@
 
 
1
+ 2.2.5
app/config/config.py ADDED
@@ -0,0 +1,513 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 应用程序配置模块
3
+ """
4
+
5
+ import datetime
6
+ import json
7
+ from typing import Any, Dict, List, Type, get_args, get_origin
8
+
9
+ from pydantic import ValidationError, ValidationInfo, field_validator, Field
10
+ from pydantic_settings import BaseSettings
11
+ from sqlalchemy import insert, select, update
12
+
13
+ from app.core.constants import (
14
+ API_VERSION,
15
+ DEFAULT_CREATE_IMAGE_MODEL,
16
+ DEFAULT_FILTER_MODELS,
17
+ DEFAULT_MODEL,
18
+ DEFAULT_SAFETY_SETTINGS,
19
+ DEFAULT_STREAM_CHUNK_SIZE,
20
+ DEFAULT_STREAM_LONG_TEXT_THRESHOLD,
21
+ DEFAULT_STREAM_MAX_DELAY,
22
+ DEFAULT_STREAM_MIN_DELAY,
23
+ DEFAULT_STREAM_SHORT_TEXT_THRESHOLD,
24
+ DEFAULT_TIMEOUT,
25
+ MAX_RETRIES,
26
+ )
27
+ from app.log.logger import Logger
28
+
29
+
30
+ class Settings(BaseSettings):
31
+ # 数据库配置
32
+ DATABASE_TYPE: str = "mysql" # sqlite 或 mysql
33
+ SQLITE_DATABASE: str = "default_db"
34
+ MYSQL_HOST: str = ""
35
+ MYSQL_PORT: int = 3306
36
+ MYSQL_USER: str = ""
37
+ MYSQL_PASSWORD: str = ""
38
+ MYSQL_DATABASE: str = ""
39
+ MYSQL_SOCKET: str = ""
40
+
41
+ # 验证 MySQL 配置
42
+ @field_validator(
43
+ "MYSQL_HOST", "MYSQL_PORT", "MYSQL_USER", "MYSQL_PASSWORD", "MYSQL_DATABASE"
44
+ )
45
+ def validate_mysql_config(cls, v: Any, info: ValidationInfo) -> Any:
46
+ if info.data.get("DATABASE_TYPE") == "mysql":
47
+ if v is None or v == "":
48
+ raise ValueError(
49
+ "MySQL configuration is required when DATABASE_TYPE is 'mysql'"
50
+ )
51
+ return v
52
+
53
+ # API相关配置
54
+ API_KEYS: List[str]=[]
55
+ ALLOWED_TOKENS: List[str]=[]
56
+ BASE_URL: str = f"https://generativelanguage.googleapis.com/{API_VERSION}"
57
+ AUTH_TOKEN: str = ""
58
+ MAX_FAILURES: int = 3
59
+ TEST_MODEL: str = DEFAULT_MODEL
60
+ TIME_OUT: int = DEFAULT_TIMEOUT
61
+ MAX_RETRIES: int = MAX_RETRIES
62
+ PROXIES: List[str] = []
63
+ PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: bool = True # 是否使用一致性哈希来选择代理
64
+ VERTEX_API_KEYS: List[str] = []
65
+ VERTEX_EXPRESS_BASE_URL: str = "https://aiplatform.googleapis.com/v1beta1/publishers/google"
66
+
67
+ # 智能路由配置
68
+ URL_NORMALIZATION_ENABLED: bool = False # 是否启用智能路由映射功能
69
+
70
+ # 自定义 Headers
71
+ CUSTOM_HEADERS: Dict[str, str] = {}
72
+
73
+ # 模型相关配置
74
+ SEARCH_MODELS: List[str] = ["gemini-2.0-flash-exp"]
75
+ IMAGE_MODELS: List[str] = ["gemini-2.0-flash-exp"]
76
+ FILTERED_MODELS: List[str] = DEFAULT_FILTER_MODELS
77
+ TOOLS_CODE_EXECUTION_ENABLED: bool = False
78
+ # 是否启用网址上下文
79
+ URL_CONTEXT_ENABLED: bool = False
80
+ URL_CONTEXT_MODELS: List[str] = ["gemini-2.5-pro","gemini-2.5-flash","gemini-2.5-flash-lite","gemini-2.0-flash","gemini-2.0-flash-live-001"]
81
+ SHOW_SEARCH_LINK: bool = True
82
+ SHOW_THINKING_PROCESS: bool = True
83
+ THINKING_MODELS: List[str] = []
84
+ THINKING_BUDGET_MAP: Dict[str, float] = {}
85
+
86
+ # TTS相关配置
87
+ TTS_MODEL: str = "gemini-2.5-flash-preview-tts"
88
+ TTS_VOICE_NAME: str = "Zephyr"
89
+ TTS_SPEED: str = "normal"
90
+
91
+ # 图像生成相关配置
92
+ PAID_KEY: str = ""
93
+ CREATE_IMAGE_MODEL: str = DEFAULT_CREATE_IMAGE_MODEL
94
+ UPLOAD_PROVIDER: str = "smms"
95
+ SMMS_SECRET_TOKEN: str = ""
96
+ PICGO_API_KEY: str = ""
97
+ PICGO_API_URL: str = "https://www.picgo.net/api/1/upload"
98
+ CLOUDFLARE_IMGBED_URL: str = ""
99
+ CLOUDFLARE_IMGBED_AUTH_CODE: str = ""
100
+ CLOUDFLARE_IMGBED_UPLOAD_FOLDER: str = ""
101
+
102
+ # 流式输出优化器配置
103
+ STREAM_OPTIMIZER_ENABLED: bool = False
104
+ STREAM_MIN_DELAY: float = DEFAULT_STREAM_MIN_DELAY
105
+ STREAM_MAX_DELAY: float = DEFAULT_STREAM_MAX_DELAY
106
+ STREAM_SHORT_TEXT_THRESHOLD: int = DEFAULT_STREAM_SHORT_TEXT_THRESHOLD
107
+ STREAM_LONG_TEXT_THRESHOLD: int = DEFAULT_STREAM_LONG_TEXT_THRESHOLD
108
+ STREAM_CHUNK_SIZE: int = DEFAULT_STREAM_CHUNK_SIZE
109
+
110
+ # 假流式配置 (Fake Streaming Configuration)
111
+ FAKE_STREAM_ENABLED: bool = False # 是否启用假流式输出
112
+ FAKE_STREAM_EMPTY_DATA_INTERVAL_SECONDS: int = 5 # 假流式发送空数据的间隔时间(秒)
113
+
114
+ # 调度器配置
115
+ CHECK_INTERVAL_HOURS: int = 1 # 默认检查间隔为1小时
116
+ TIMEZONE: str = "Asia/Shanghai" # 默认时区
117
+
118
+ # github
119
+ GITHUB_REPO_OWNER: str = "snailyp"
120
+ GITHUB_REPO_NAME: str = "gemini-balance"
121
+
122
+ # 日志配置
123
+ LOG_LEVEL: str = "INFO"
124
+ AUTO_DELETE_ERROR_LOGS_ENABLED: bool = True
125
+ AUTO_DELETE_ERROR_LOGS_DAYS: int = 7
126
+ AUTO_DELETE_REQUEST_LOGS_ENABLED: bool = False
127
+ AUTO_DELETE_REQUEST_LOGS_DAYS: int = 30
128
+ SAFETY_SETTINGS: List[Dict[str, str]] = DEFAULT_SAFETY_SETTINGS
129
+
130
+ # Files API
131
+ FILES_CLEANUP_ENABLED: bool = True
132
+ FILES_CLEANUP_INTERVAL_HOURS: int = 1
133
+ FILES_USER_ISOLATION_ENABLED: bool = True
134
+
135
+ # Admin Session Configuration
136
+ ADMIN_SESSION_EXPIRE: int = Field(
137
+ default=3600,
138
+ ge=300,
139
+ le=86400,
140
+ description="Admin session expiration time in seconds (5 minutes to 24 hours)"
141
+ )
142
+
143
+ def __init__(self, **kwargs):
144
+ super().__init__(**kwargs)
145
+ # 设置默认AUTH_TOKEN(如果未提供)
146
+ if not self.AUTH_TOKEN and self.ALLOWED_TOKENS:
147
+ self.AUTH_TOKEN = self.ALLOWED_TOKENS[0]
148
+
149
+
150
+ # 创建全局配置实例
151
+ settings = Settings()
152
+
153
+
154
+ def _parse_db_value(key: str, db_value: str, target_type: Type) -> Any:
155
+ """尝试将数据库字符串值解析为目标 Python 类型"""
156
+ from app.log.logger import get_config_logger
157
+
158
+ logger = get_config_logger()
159
+ try:
160
+ origin_type = get_origin(target_type)
161
+ args = get_args(target_type)
162
+
163
+ # 处理 List 类型
164
+ if origin_type is list:
165
+ # 处理 List[str]
166
+ if args and args[0] == str:
167
+ try:
168
+ parsed = json.loads(db_value)
169
+ if isinstance(parsed, list):
170
+ return [str(item) for item in parsed]
171
+ except json.JSONDecodeError:
172
+ return [item.strip() for item in db_value.split(",") if item.strip()]
173
+ logger.warning(
174
+ f"Could not parse '{db_value}' as List[str] for key '{key}', falling back to comma split or empty list."
175
+ )
176
+ return [item.strip() for item in db_value.split(",") if item.strip()]
177
+ # 处理 List[Dict[str, str]]
178
+ elif args and get_origin(args[0]) is dict:
179
+ try:
180
+ parsed = json.loads(db_value)
181
+ if isinstance(parsed, list):
182
+ valid = all(
183
+ isinstance(item, dict)
184
+ and all(isinstance(k, str) for k in item.keys())
185
+ and all(isinstance(v, str) for v in item.values())
186
+ for item in parsed
187
+ )
188
+ if valid:
189
+ return parsed
190
+ else:
191
+ logger.warning(
192
+ f"Invalid structure in List[Dict[str, str]] for key '{key}'. Value: {db_value}"
193
+ )
194
+ return []
195
+ else:
196
+ logger.warning(
197
+ f"Parsed DB value for key '{key}' is not a list type. Value: {db_value}"
198
+ )
199
+ return []
200
+ except json.JSONDecodeError:
201
+ logger.error(
202
+ f"Could not parse '{db_value}' as JSON for List[Dict[str, str]] for key '{key}'. Returning empty list."
203
+ )
204
+ return []
205
+ except Exception as e:
206
+ logger.error(
207
+ f"Error parsing List[Dict[str, str]] for key '{key}': {e}. Value: {db_value}. Returning empty list."
208
+ )
209
+ return []
210
+ # 处理 Dict 类型
211
+ elif origin_type is dict:
212
+ # 处理 Dict[str, str]
213
+ if args and args == (str, str):
214
+ parsed_dict = {}
215
+ try:
216
+ parsed = json.loads(db_value)
217
+ if isinstance(parsed, dict):
218
+ parsed_dict = {str(k): str(v) for k, v in parsed.items()}
219
+ else:
220
+ logger.warning(
221
+ f"Parsed DB value for key '{key}' is not a dictionary type. Value: {db_value}"
222
+ )
223
+ except json.JSONDecodeError:
224
+ logger.error(f"Could not parse '{db_value}' as Dict[str, str] for key '{key}'. Returning empty dict.")
225
+ return parsed_dict
226
+ # 处理 Dict[str, float]
227
+ elif args and args == (str, float):
228
+ parsed_dict = {}
229
+ try:
230
+ parsed = json.loads(db_value)
231
+ if isinstance(parsed, dict):
232
+ parsed_dict = {str(k): float(v) for k, v in parsed.items()}
233
+ else:
234
+ logger.warning(
235
+ f"Parsed DB value for key '{key}' is not a dictionary type. Value: {db_value}"
236
+ )
237
+ except (json.JSONDecodeError, ValueError, TypeError) as e1:
238
+ if isinstance(e1, json.JSONDecodeError) and "'" in db_value:
239
+ logger.warning(
240
+ f"Failed initial JSON parse for key '{key}'. Attempting to replace single quotes. Error: {e1}"
241
+ )
242
+ try:
243
+ corrected_db_value = db_value.replace("'", '"')
244
+ parsed = json.loads(corrected_db_value)
245
+ if isinstance(parsed, dict):
246
+ parsed_dict = {str(k): float(v) for k, v in parsed.items()}
247
+ else:
248
+ logger.warning(
249
+ f"Parsed DB value (after quote replacement) for key '{key}' is not a dictionary type. Value: {corrected_db_value}"
250
+ )
251
+ except (json.JSONDecodeError, ValueError, TypeError) as e2:
252
+ logger.error(
253
+ f"Could not parse '{db_value}' as Dict[str, float] for key '{key}' even after replacing quotes: {e2}. Returning empty dict."
254
+ )
255
+ else:
256
+ logger.error(
257
+ f"Could not parse '{db_value}' as Dict[str, float] for key '{key}': {e1}. Returning empty dict."
258
+ )
259
+ return parsed_dict
260
+ # 处理 bool
261
+ elif target_type == bool:
262
+ return db_value.lower() in ("true", "1", "yes", "on")
263
+ # 处理 int
264
+ elif target_type == int:
265
+ return int(db_value)
266
+ # 处理 float
267
+ elif target_type == float:
268
+ return float(db_value)
269
+ # 默认为 str 或其他 pydantic 能直接处理的类型
270
+ else:
271
+ return db_value
272
+ except (ValueError, TypeError, json.JSONDecodeError) as e:
273
+ logger.warning(
274
+ f"Failed to parse db_value '{db_value}' for key '{key}' as type {target_type}: {e}. Using original string value."
275
+ )
276
+ return db_value # 解析失败则返回原始字符串
277
+
278
+
279
+ async def sync_initial_settings():
280
+ """
281
+ 应用启动时同步配置:
282
+ 1. 从数据库加载设置。
283
+ 2. 将数据库设置合并到内存 settings (数据库优先)。
284
+ 3. 将最终的内存 settings 同步回数据库。
285
+ """
286
+ from app.log.logger import get_config_logger
287
+
288
+ logger = get_config_logger()
289
+ # 延迟导入以避免循环依赖和确保数据库连接已初始化
290
+ from app.database.connection import database
291
+ from app.database.models import Settings as SettingsModel
292
+
293
+ global settings
294
+ logger.info("Starting initial settings synchronization...")
295
+
296
+ if not database.is_connected:
297
+ try:
298
+ await database.connect()
299
+ logger.info("Database connection established for initial sync.")
300
+ except Exception as e:
301
+ logger.error(
302
+ f"Failed to connect to database for initial settings sync: {e}. Skipping sync."
303
+ )
304
+ return
305
+
306
+ try:
307
+ # 1. 从数据库加载设置
308
+ db_settings_raw: List[Dict[str, Any]] = []
309
+ try:
310
+ query = select(SettingsModel.key, SettingsModel.value)
311
+ results = await database.fetch_all(query)
312
+ db_settings_raw = [
313
+ {"key": row["key"], "value": row["value"]} for row in results
314
+ ]
315
+ logger.info(f"Fetched {len(db_settings_raw)} settings from database.")
316
+ except Exception as e:
317
+ logger.error(
318
+ f"Failed to fetch settings from database: {e}. Proceeding with environment/dotenv settings."
319
+ )
320
+ # 即使数据库读取失败,也要继续执行,确保基于 env/dotenv 的配置能同步到数据库
321
+
322
+ db_settings_map: Dict[str, str] = {
323
+ s["key"]: s["value"] for s in db_settings_raw
324
+ }
325
+
326
+ # 2. 将数据库设置合并到内存 settings (数据库优先)
327
+ updated_in_memory = False
328
+
329
+ for key, db_value in db_settings_map.items():
330
+ if key == "DATABASE_TYPE":
331
+ logger.debug(
332
+ f"Skipping update of '{key}' in memory from database. "
333
+ "This setting is controlled by environment/dotenv."
334
+ )
335
+ continue
336
+ if hasattr(settings, key):
337
+ target_type = Settings.__annotations__.get(key)
338
+ if target_type:
339
+ try:
340
+ parsed_db_value = _parse_db_value(key, db_value, target_type)
341
+ memory_value = getattr(settings, key)
342
+
343
+ # 比较解析后的值和内存中的值
344
+ # 注意:对于列表等复杂类型,直接比较可能不够健壮,但这里简化处理
345
+ if parsed_db_value != memory_value:
346
+ # 检查类型是否匹配,以防解析函数返回了不兼容的类型
347
+ type_match = False
348
+ origin_type = get_origin(target_type)
349
+ if origin_type: # It's a generic type
350
+ if isinstance(parsed_db_value, origin_type):
351
+ type_match = True
352
+ # It's a non-generic type, or a specific generic we want to handle
353
+ elif isinstance(parsed_db_value, target_type):
354
+ type_match = True
355
+
356
+ if type_match:
357
+ setattr(settings, key, parsed_db_value)
358
+ logger.debug(
359
+ f"Updated setting '{key}' in memory from database value ({target_type})."
360
+ )
361
+ updated_in_memory = True
362
+ else:
363
+ logger.warning(
364
+ f"Parsed DB value type mismatch for key '{key}'. Expected {target_type}, got {type(parsed_db_value)}. Skipping update."
365
+ )
366
+
367
+ except Exception as e:
368
+ logger.error(
369
+ f"Error processing database setting for key '{key}': {e}"
370
+ )
371
+ else:
372
+ logger.warning(
373
+ f"Database setting '{key}' not found in Settings model definition. Ignoring."
374
+ )
375
+
376
+ # 如果内存中有更新,重新验证 Pydantic 模型(可选但推荐)
377
+ if updated_in_memory:
378
+ try:
379
+ # 重新加载以确保类型转换和验证
380
+ settings = Settings(**settings.model_dump())
381
+ logger.info(
382
+ "Settings object re-validated after merging database values."
383
+ )
384
+ except ValidationError as e:
385
+ logger.error(
386
+ f"Validation error after merging database settings: {e}. Settings might be inconsistent."
387
+ )
388
+
389
+ # 3. 将最终的内存 settings 同步回数据库
390
+ final_memory_settings = settings.model_dump()
391
+ settings_to_update: List[Dict[str, Any]] = []
392
+ settings_to_insert: List[Dict[str, Any]] = []
393
+ now = datetime.datetime.now(datetime.timezone.utc)
394
+
395
+ existing_db_keys = set(db_settings_map.keys())
396
+
397
+ for key, value in final_memory_settings.items():
398
+ if key == "DATABASE_TYPE":
399
+ logger.debug(
400
+ f"Skipping synchronization of '{key}' to database. "
401
+ "This setting is controlled by environment/dotenv."
402
+ )
403
+ continue
404
+
405
+ # 序列化值为字符串或 JSON 字符串
406
+ if isinstance(value, (list, dict)):
407
+ db_value = json.dumps(
408
+ value, ensure_ascii=False
409
+ )
410
+ elif isinstance(value, bool):
411
+ db_value = str(value).lower()
412
+ elif value is None:
413
+ db_value = ""
414
+ else:
415
+ db_value = str(value)
416
+
417
+ data = {
418
+ "key": key,
419
+ "value": db_value,
420
+ "description": f"{key} configuration setting",
421
+ "updated_at": now,
422
+ }
423
+
424
+ if key in existing_db_keys:
425
+ # 仅当值与数据库中的不同时才更新
426
+ if db_settings_map[key] != db_value:
427
+ settings_to_update.append(data)
428
+ else:
429
+ # 如果键不在数据库中,则插入
430
+ data["created_at"] = now
431
+ settings_to_insert.append(data)
432
+
433
+ # 在事务中执行批量插入和更新
434
+ if settings_to_insert or settings_to_update:
435
+ try:
436
+ async with database.transaction():
437
+ if settings_to_insert:
438
+ # 获取现有描述以避免覆盖
439
+ query_existing = select(
440
+ SettingsModel.key, SettingsModel.description
441
+ ).where(
442
+ SettingsModel.key.in_(
443
+ [s["key"] for s in settings_to_insert]
444
+ )
445
+ )
446
+ existing_desc = {
447
+ row["key"]: row["description"]
448
+ for row in await database.fetch_all(query_existing)
449
+ }
450
+ for item in settings_to_insert:
451
+ item["description"] = existing_desc.get(
452
+ item["key"], item["description"]
453
+ )
454
+
455
+ query_insert = insert(SettingsModel).values(settings_to_insert)
456
+ await database.execute(query=query_insert)
457
+ logger.info(
458
+ f"Synced (inserted) {len(settings_to_insert)} settings to database."
459
+ )
460
+
461
+ if settings_to_update:
462
+ # 获取现有描述以避免覆盖
463
+ query_existing = select(
464
+ SettingsModel.key, SettingsModel.description
465
+ ).where(
466
+ SettingsModel.key.in_(
467
+ [s["key"] for s in settings_to_update]
468
+ )
469
+ )
470
+ existing_desc = {
471
+ row["key"]: row["description"]
472
+ for row in await database.fetch_all(query_existing)
473
+ }
474
+
475
+ for setting_data in settings_to_update:
476
+ setting_data["description"] = existing_desc.get(
477
+ setting_data["key"], setting_data["description"]
478
+ )
479
+ query_update = (
480
+ update(SettingsModel)
481
+ .where(SettingsModel.key == setting_data["key"])
482
+ .values(
483
+ value=setting_data["value"],
484
+ description=setting_data["description"],
485
+ updated_at=setting_data["updated_at"],
486
+ )
487
+ )
488
+ await database.execute(query=query_update)
489
+ logger.info(
490
+ f"Synced (updated) {len(settings_to_update)} settings to database."
491
+ )
492
+ except Exception as e:
493
+ logger.error(
494
+ f"Failed to sync settings to database during startup: {str(e)}"
495
+ )
496
+ else:
497
+ logger.info(
498
+ "No setting changes detected between memory and database during initial sync."
499
+ )
500
+
501
+ # 刷新日志等级
502
+ Logger.update_log_levels(final_memory_settings.get("LOG_LEVEL"))
503
+
504
+ except Exception as e:
505
+ logger.error(f"An unexpected error occurred during initial settings sync: {e}")
506
+ finally:
507
+ if database.is_connected:
508
+ try:
509
+ pass
510
+ except Exception as e:
511
+ logger.error(f"Error disconnecting database after initial sync: {e}")
512
+
513
+ logger.info("Initial settings synchronization finished.")
app/core/application.py ADDED
@@ -0,0 +1,156 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from contextlib import asynccontextmanager
2
+ from pathlib import Path
3
+
4
+ from fastapi import FastAPI
5
+ from fastapi.staticfiles import StaticFiles
6
+ from fastapi.templating import Jinja2Templates
7
+
8
+ from app.config.config import settings, sync_initial_settings
9
+ from app.database.connection import connect_to_db, disconnect_from_db
10
+ from app.database.initialization import initialize_database
11
+ from app.exception.exceptions import setup_exception_handlers
12
+ from app.log.logger import get_application_logger, setup_access_logging
13
+ from app.middleware.middleware import setup_middlewares
14
+ from app.router.routes import setup_routers
15
+ from app.scheduler.scheduled_tasks import start_scheduler, stop_scheduler
16
+ from app.service.key.key_manager import get_key_manager_instance
17
+ from app.service.update.update_service import check_for_updates
18
+ from app.utils.helpers import get_current_version
19
+
20
+ logger = get_application_logger()
21
+
22
+ PROJECT_ROOT = Path(__file__).resolve().parent.parent.parent
23
+ STATIC_DIR = PROJECT_ROOT / "app" / "static"
24
+ TEMPLATES_DIR = PROJECT_ROOT / "app" / "templates"
25
+
26
+ # 初始化模板引擎,并添加全局变量
27
+ templates = Jinja2Templates(directory="app/templates")
28
+
29
+
30
+ # 定义一个函数来更新模板全局变量
31
+ def update_template_globals(app: FastAPI, update_info: dict):
32
+ # Jinja2Templates 实例没有直接更新全局变量的方法
33
+ # 我们需要在请求上下文中传递这些变量,或者修改 Jinja 环境
34
+ # 更简单的方法是将其存储在 app.state 中,并在渲染时传递
35
+ app.state.update_info = update_info
36
+ logger.info(f"Update info stored in app.state: {update_info}")
37
+
38
+
39
+ # --- Helper functions for lifespan ---
40
+ async def _setup_database_and_config(app_settings):
41
+ """Initializes database, syncs settings, and initializes KeyManager."""
42
+ initialize_database()
43
+ logger.info("Database initialized successfully")
44
+ await connect_to_db()
45
+ await sync_initial_settings()
46
+ await get_key_manager_instance(app_settings.API_KEYS, app_settings.VERTEX_API_KEYS)
47
+ logger.info("Database, config sync, and KeyManager initialized successfully")
48
+
49
+
50
+ async def _shutdown_database():
51
+ """Disconnects from the database."""
52
+ await disconnect_from_db()
53
+
54
+
55
+ def _start_scheduler():
56
+ """Starts the background scheduler."""
57
+ try:
58
+ start_scheduler()
59
+ logger.info("Scheduler started successfully.")
60
+ except Exception as e:
61
+ logger.error(f"Failed to start scheduler: {e}")
62
+
63
+
64
+ def _stop_scheduler():
65
+ """Stops the background scheduler."""
66
+ stop_scheduler()
67
+
68
+
69
+ async def _perform_update_check(app: FastAPI):
70
+ """Checks for updates and stores the info in app.state."""
71
+ update_available, latest_version, error_message = await check_for_updates()
72
+ current_version = get_current_version()
73
+ update_info = {
74
+ "update_available": update_available,
75
+ "latest_version": latest_version,
76
+ "error_message": error_message,
77
+ "current_version": current_version,
78
+ }
79
+ if not hasattr(app, "state"):
80
+ from starlette.datastructures import State
81
+
82
+ app.state = State()
83
+ app.state.update_info = update_info
84
+ logger.info(f"Update check completed. Info: {update_info}")
85
+
86
+
87
+ @asynccontextmanager
88
+ async def lifespan(app: FastAPI):
89
+ """
90
+ Manages the application startup and shutdown events.
91
+
92
+ Args:
93
+ app: FastAPI应用实例
94
+ """
95
+ logger.info("Application starting up...")
96
+ try:
97
+ await _setup_database_and_config(settings)
98
+ await _perform_update_check(app)
99
+ _start_scheduler()
100
+
101
+ except Exception as e:
102
+ logger.critical(
103
+ f"Critical error during application startup: {str(e)}", exc_info=True
104
+ )
105
+
106
+ yield
107
+
108
+ logger.info("Application shutting down...")
109
+ _stop_scheduler()
110
+ await _shutdown_database()
111
+
112
+
113
+ def create_app() -> FastAPI:
114
+ """
115
+ 创建并配置FastAPI应用程序实例
116
+
117
+ Returns:
118
+ FastAPI: 配置好的FastAPI应用程序实例
119
+ """
120
+
121
+ # 创建FastAPI应用
122
+ current_version = get_current_version()
123
+ app = FastAPI(
124
+ title="Gemini Balance API",
125
+ description="Gemini API代理服务,支持负载均衡和密钥管理",
126
+ version=current_version,
127
+ lifespan=lifespan,
128
+ )
129
+
130
+ if not hasattr(app, "state"):
131
+ from starlette.datastructures import State
132
+
133
+ app.state = State()
134
+ app.state.update_info = {
135
+ "update_available": False,
136
+ "latest_version": None,
137
+ "error_message": "Initializing...",
138
+ "current_version": current_version,
139
+ }
140
+
141
+ # 配置静态文件
142
+ app.mount("/static", StaticFiles(directory=str(STATIC_DIR)), name="static")
143
+
144
+ # 配置中间件
145
+ setup_middlewares(app)
146
+
147
+ # 配置异常处理器
148
+ setup_exception_handlers(app)
149
+
150
+ # 配置路由
151
+ setup_routers(app)
152
+
153
+ # 配置访问日志API密钥隐藏
154
+ setup_access_logging()
155
+
156
+ return app
app/core/constants.py ADDED
@@ -0,0 +1,112 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 常量定义模块
3
+ """
4
+
5
+ # API相关常量
6
+ API_VERSION = "v1beta"
7
+ DEFAULT_TIMEOUT = 300 # 秒
8
+ MAX_RETRIES = 3 # 最大重试次数
9
+
10
+ # 模型相关常量
11
+ SUPPORTED_ROLES = ["user", "model", "system"]
12
+ DEFAULT_MODEL = "gemini-1.5-flash"
13
+ DEFAULT_TEMPERATURE = 0.7
14
+ DEFAULT_MAX_TOKENS = 8192
15
+ DEFAULT_TOP_P = 0.9
16
+ DEFAULT_TOP_K = 40
17
+ DEFAULT_FILTER_MODELS = [
18
+ "gemini-1.0-pro-vision-latest",
19
+ "gemini-pro-vision",
20
+ "chat-bison-001",
21
+ "text-bison-001",
22
+ "embedding-gecko-001",
23
+ ]
24
+ DEFAULT_CREATE_IMAGE_MODEL = "imagen-3.0-generate-002"
25
+
26
+ # 图像生成相关常量
27
+ VALID_IMAGE_RATIOS = ["1:1", "3:4", "4:3", "9:16", "16:9"]
28
+
29
+ # 上传提供商
30
+ UPLOAD_PROVIDERS = ["smms", "picgo", "cloudflare_imgbed"]
31
+ DEFAULT_UPLOAD_PROVIDER = "smms"
32
+
33
+ # 流式输出相关常量
34
+ DEFAULT_STREAM_MIN_DELAY = 0.016
35
+ DEFAULT_STREAM_MAX_DELAY = 0.024
36
+ DEFAULT_STREAM_SHORT_TEXT_THRESHOLD = 10
37
+ DEFAULT_STREAM_LONG_TEXT_THRESHOLD = 50
38
+ DEFAULT_STREAM_CHUNK_SIZE = 5
39
+
40
+ # 正则表达式模式
41
+ IMAGE_URL_PATTERN = r"!\[(.*?)\]\((.*?)\)"
42
+ DATA_URL_PATTERN = r"data:([^;]+);base64,(.+)"
43
+
44
+ # Audio/Video Settings
45
+ SUPPORTED_AUDIO_FORMATS = ["wav", "mp3", "flac", "ogg"]
46
+ SUPPORTED_VIDEO_FORMATS = ["mp4", "mov", "avi", "webm"]
47
+ MAX_AUDIO_SIZE_BYTES = 50 * 1024 * 1024 # Example: 50MB limit for Base64 payload
48
+ MAX_VIDEO_SIZE_BYTES = 200 * 1024 * 1024 # Example: 200MB limit
49
+
50
+ # Optional: Define MIME type mappings if needed, or handle directly in converter
51
+ AUDIO_FORMAT_TO_MIMETYPE = {
52
+ "wav": "audio/wav",
53
+ "mp3": "audio/mpeg",
54
+ "flac": "audio/flac",
55
+ "ogg": "audio/ogg",
56
+ }
57
+
58
+ VIDEO_FORMAT_TO_MIMETYPE = {
59
+ "mp4": "video/mp4",
60
+ "mov": "video/quicktime",
61
+ "avi": "video/x-msvideo",
62
+ "webm": "video/webm",
63
+ }
64
+
65
+ GEMINI_2_FLASH_EXP_SAFETY_SETTINGS = [
66
+ {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
67
+ {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},
68
+ {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"},
69
+ {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"},
70
+ {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "OFF"},
71
+ ]
72
+
73
+ DEFAULT_SAFETY_SETTINGS = [
74
+ {"category": "HARM_CATEGORY_HARASSMENT", "threshold": "OFF"},
75
+ {"category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "OFF"},
76
+ {"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "OFF"},
77
+ {"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "OFF"},
78
+ {"category": "HARM_CATEGORY_CIVIC_INTEGRITY", "threshold": "BLOCK_NONE"},
79
+ ]
80
+
81
+ TTS_VOICE_NAMES = [
82
+ "Zephyr",
83
+ "Puck",
84
+ "Charon",
85
+ "Kore",
86
+ "Fenrir",
87
+ "Leda",
88
+ "Orus",
89
+ "Aoede",
90
+ "Callirrhoe",
91
+ "Autonoe",
92
+ "Enceladus",
93
+ "Iapetus",
94
+ "Umbriel",
95
+ "Algieba",
96
+ "Despina",
97
+ "Erinome",
98
+ "Algenib",
99
+ "Rasalgethi",
100
+ "Laomedeia",
101
+ "Achernar",
102
+ "Alnilam",
103
+ "Schedar",
104
+ "Gacrux",
105
+ "Pulcherrima",
106
+ "Achird",
107
+ "Zubenelgenubi",
108
+ "Vindemiatrix",
109
+ "Sadachbia",
110
+ "Sadaltager",
111
+ "Sulafat",
112
+ ]
app/core/security.py ADDED
@@ -0,0 +1,90 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Optional
2
+
3
+ from fastapi import Header, HTTPException
4
+
5
+ from app.config.config import settings
6
+ from app.log.logger import get_security_logger
7
+
8
+ logger = get_security_logger()
9
+
10
+
11
+ def verify_auth_token(token: str) -> bool:
12
+ return token == settings.AUTH_TOKEN
13
+
14
+
15
+ class SecurityService:
16
+
17
+ async def verify_key(self, key: str):
18
+ if key not in settings.ALLOWED_TOKENS and key != settings.AUTH_TOKEN:
19
+ logger.error("Invalid key")
20
+ raise HTTPException(status_code=401, detail="Invalid key")
21
+ return key
22
+
23
+ async def verify_authorization(
24
+ self, authorization: Optional[str] = Header(None)
25
+ ) -> str:
26
+ if not authorization:
27
+ logger.error("Missing Authorization header")
28
+ raise HTTPException(status_code=401, detail="Missing Authorization header")
29
+
30
+ if not authorization.startswith("Bearer "):
31
+ logger.error("Invalid Authorization header format")
32
+ raise HTTPException(
33
+ status_code=401, detail="Invalid Authorization header format"
34
+ )
35
+
36
+ token = authorization.replace("Bearer ", "")
37
+ if token not in settings.ALLOWED_TOKENS and token != settings.AUTH_TOKEN:
38
+ logger.error("Invalid token")
39
+ raise HTTPException(status_code=401, detail="Invalid token")
40
+
41
+ return token
42
+
43
+ async def verify_goog_api_key(
44
+ self, x_goog_api_key: Optional[str] = Header(None)
45
+ ) -> str:
46
+ """验证Google API Key"""
47
+ if not x_goog_api_key:
48
+ logger.error("Missing x-goog-api-key header")
49
+ raise HTTPException(status_code=401, detail="Missing x-goog-api-key header")
50
+
51
+ if (
52
+ x_goog_api_key not in settings.ALLOWED_TOKENS
53
+ and x_goog_api_key != settings.AUTH_TOKEN
54
+ ):
55
+ logger.error("Invalid x-goog-api-key")
56
+ raise HTTPException(status_code=401, detail="Invalid x-goog-api-key")
57
+
58
+ return x_goog_api_key
59
+
60
+ async def verify_auth_token(
61
+ self, authorization: Optional[str] = Header(None)
62
+ ) -> str:
63
+ if not authorization:
64
+ logger.error("Missing auth_token header")
65
+ raise HTTPException(status_code=401, detail="Missing auth_token header")
66
+ token = authorization.replace("Bearer ", "")
67
+ if token != settings.AUTH_TOKEN:
68
+ logger.error("Invalid auth_token")
69
+ raise HTTPException(status_code=401, detail="Invalid auth_token")
70
+
71
+ return token
72
+
73
+ async def verify_key_or_goog_api_key(
74
+ self, key: Optional[str] = None , x_goog_api_key: Optional[str] = Header(None)
75
+ ) -> str:
76
+ """验证URL中的key或请求头中的x-goog-api-key"""
77
+ # 如果URL中的key有效,直接返回
78
+ if key in settings.ALLOWED_TOKENS or key == settings.AUTH_TOKEN:
79
+ return key
80
+
81
+ # 否则检查请求头中的x-goog-api-key
82
+ if not x_goog_api_key:
83
+ logger.error("Invalid key and missing x-goog-api-key header")
84
+ raise HTTPException(status_code=401, detail="Invalid key and missing x-goog-api-key header")
85
+
86
+ if x_goog_api_key not in settings.ALLOWED_TOKENS and x_goog_api_key != settings.AUTH_TOKEN:
87
+ logger.error("Invalid key and invalid x-goog-api-key")
88
+ raise HTTPException(status_code=401, detail="Invalid key and invalid x-goog-api-key")
89
+
90
+ return x_goog_api_key
app/database/__init__.py ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ """
2
+ 数据库模块
3
+ """
app/database/connection.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 数据库连接池模块
3
+ """
4
+ from pathlib import Path
5
+ from urllib.parse import quote_plus
6
+ from databases import Database
7
+ from sqlalchemy import create_engine, MetaData
8
+ from sqlalchemy.ext.declarative import declarative_base
9
+
10
+ from app.config.config import settings
11
+ from app.log.logger import get_database_logger
12
+
13
+ logger = get_database_logger()
14
+
15
+ # 数据库URL
16
+ if settings.DATABASE_TYPE == "sqlite":
17
+ # 确保 data 目录存在
18
+ data_dir = Path("data")
19
+ data_dir.mkdir(exist_ok=True)
20
+ db_path = data_dir / settings.SQLITE_DATABASE
21
+ DATABASE_URL = f"sqlite:///{db_path}"
22
+ elif settings.DATABASE_TYPE == "mysql":
23
+ if settings.MYSQL_SOCKET:
24
+ DATABASE_URL = f"mysql+pymysql://{settings.MYSQL_USER}:{quote_plus(settings.MYSQL_PASSWORD)}@/{settings.MYSQL_DATABASE}?unix_socket={settings.MYSQL_SOCKET}"
25
+ else:
26
+ DATABASE_URL = f"mysql+pymysql://{settings.MYSQL_USER}:{quote_plus(settings.MYSQL_PASSWORD)}@{settings.MYSQL_HOST}:{settings.MYSQL_PORT}/{settings.MYSQL_DATABASE}"
27
+ else:
28
+ raise ValueError("Unsupported database type. Please set DATABASE_TYPE to 'sqlite' or 'mysql'.")
29
+
30
+ # 创建数据库引擎
31
+ # pool_pre_ping=True: 在从连接池获取连接前执行简单的 "ping" 测试,确保连接有效
32
+ engine = create_engine(DATABASE_URL, pool_pre_ping=True)
33
+
34
+ # 创建元数据对象
35
+ metadata = MetaData()
36
+
37
+ # 创建基类
38
+ Base = declarative_base(metadata=metadata)
39
+
40
+ # 创建数据库连接池,并配置连接池参数,在sqlite中不使用连接池
41
+ # min_size/max_size: 连接池的最小/最大连接数
42
+ # pool_recycle=3600: 连接在池中允许存在的最大秒数(生命周期)。
43
+ # 设置为 3600 秒(1小时),确保在 MySQL 默认的 wait_timeout (通常8小时) 或其他网络超时之前回收连接。
44
+ # 如果遇到连接失效问题,可以尝试调低此值,使其小于实际的 wait_timeout 或网络超时时间。
45
+ # databases 库会自动处理连接失效后的重连尝试。
46
+ if settings.DATABASE_TYPE == "sqlite":
47
+ database = Database(DATABASE_URL)
48
+ else:
49
+ database = Database(DATABASE_URL, min_size=5, max_size=20, pool_recycle=1800)
50
+
51
+ async def connect_to_db():
52
+ """
53
+ 连接到数据库
54
+ """
55
+ try:
56
+ await database.connect()
57
+ logger.info(f"Connected to {settings.DATABASE_TYPE}")
58
+ except Exception as e:
59
+ logger.error(f"Failed to connect to database: {str(e)}")
60
+ raise
61
+
62
+
63
+ async def disconnect_from_db():
64
+ """
65
+ 断开数据库连接
66
+ """
67
+ try:
68
+ await database.disconnect()
69
+ logger.info(f"Disconnected from {settings.DATABASE_TYPE}")
70
+ except Exception as e:
71
+ logger.error(f"Failed to disconnect from database: {str(e)}")
app/database/initialization.py ADDED
@@ -0,0 +1,77 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 数据库初始化模块
3
+ """
4
+ from dotenv import dotenv_values
5
+
6
+ from sqlalchemy import inspect
7
+ from sqlalchemy.orm import Session
8
+
9
+ from app.database.connection import engine, Base
10
+ from app.database.models import Settings
11
+ from app.log.logger import get_database_logger
12
+
13
+ logger = get_database_logger()
14
+
15
+
16
+ def create_tables():
17
+ """
18
+ 创建数据库表
19
+ """
20
+ try:
21
+ # 创建所有表
22
+ Base.metadata.create_all(engine)
23
+ logger.info("Database tables created successfully")
24
+ except Exception as e:
25
+ logger.error(f"Failed to create database tables: {str(e)}")
26
+ raise
27
+
28
+
29
+ def import_env_to_settings():
30
+ """
31
+ 将.env文件中的配置项导入到t_settings表中
32
+ """
33
+ try:
34
+ # 获取.env文件中的所有配置项
35
+ env_values = dotenv_values(".env")
36
+
37
+ # 获取检查器
38
+ inspector = inspect(engine)
39
+
40
+ # 检查t_settings表是否存在
41
+ if "t_settings" in inspector.get_table_names():
42
+ # 使用Session进行数据库操作
43
+ with Session(engine) as session:
44
+ # 获取所有现有的配置项
45
+ current_settings = {setting.key: setting for setting in session.query(Settings).all()}
46
+
47
+ # 遍历所有配置项
48
+ for key, value in env_values.items():
49
+ # 检查配置项是否已存在
50
+ if key not in current_settings:
51
+ # 插入配置项
52
+ new_setting = Settings(key=key, value=value)
53
+ session.add(new_setting)
54
+ logger.info(f"Inserted setting: {key}")
55
+
56
+ # 提交事务
57
+ session.commit()
58
+
59
+ logger.info("Environment variables imported to settings table successfully")
60
+ except Exception as e:
61
+ logger.error(f"Failed to import environment variables to settings table: {str(e)}")
62
+ raise
63
+
64
+
65
+ def initialize_database():
66
+ """
67
+ 初始化数据库
68
+ """
69
+ try:
70
+ # 创建表
71
+ create_tables()
72
+
73
+ # 导入环境变量
74
+ import_env_to_settings()
75
+ except Exception as e:
76
+ logger.error(f"Failed to initialize database: {str(e)}")
77
+ raise
app/database/models.py ADDED
@@ -0,0 +1,129 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 数据库模型模块
3
+ """
4
+ import datetime
5
+ from sqlalchemy import Column, Integer, String, Text, DateTime, JSON, Boolean, BigInteger, Enum
6
+ import enum
7
+
8
+ from app.database.connection import Base
9
+
10
+
11
+ class Settings(Base):
12
+ """
13
+ 设置表,对应.env中的配置项
14
+ """
15
+ __tablename__ = "t_settings"
16
+
17
+ id = Column(Integer, primary_key=True, autoincrement=True)
18
+ key = Column(String(100), nullable=False, unique=True, comment="配置项键名")
19
+ value = Column(Text, nullable=True, comment="配置项值")
20
+ description = Column(String(255), nullable=True, comment="配置项描述")
21
+ created_at = Column(DateTime, default=datetime.datetime.now, comment="创建时间")
22
+ updated_at = Column(DateTime, default=datetime.datetime.now, onupdate=datetime.datetime.now, comment="更新时间")
23
+
24
+ def __repr__(self):
25
+ return f"<Settings(key='{self.key}', value='{self.value}')>"
26
+
27
+
28
+ class ErrorLog(Base):
29
+ """
30
+ 错误日志表
31
+ """
32
+ __tablename__ = "t_error_logs"
33
+
34
+ id = Column(Integer, primary_key=True, autoincrement=True)
35
+ gemini_key = Column(String(100), nullable=True, comment="Gemini API密钥")
36
+ model_name = Column(String(100), nullable=True, comment="模型名称")
37
+ error_type = Column(String(50), nullable=True, comment="错误类型")
38
+ error_log = Column(Text, nullable=True, comment="错误日志")
39
+ error_code = Column(Integer, nullable=True, comment="错误代码")
40
+ request_msg = Column(JSON, nullable=True, comment="请求消息")
41
+ request_time = Column(DateTime, default=datetime.datetime.now, comment="请求时间")
42
+
43
+ def __repr__(self):
44
+ return f"<ErrorLog(id='{self.id}', gemini_key='{self.gemini_key}')>"
45
+
46
+
47
+ class RequestLog(Base):
48
+ """
49
+ API 请求日志表
50
+ """
51
+
52
+ __tablename__ = "t_request_log"
53
+
54
+ id = Column(Integer, primary_key=True, autoincrement=True)
55
+ request_time = Column(DateTime, default=datetime.datetime.now, comment="请求时间")
56
+ model_name = Column(String(100), nullable=True, comment="模型名称")
57
+ api_key = Column(String(100), nullable=True, comment="使用的API密钥")
58
+ is_success = Column(Boolean, nullable=False, comment="请求是否成功")
59
+ status_code = Column(Integer, nullable=True, comment="API响应状态码")
60
+ latency_ms = Column(Integer, nullable=True, comment="请求耗时(毫秒)")
61
+
62
+ def __repr__(self):
63
+ return f"<RequestLog(id='{self.id}', key='{self.api_key[:4]}...', success='{self.is_success}')>"
64
+
65
+
66
+ class FileState(enum.Enum):
67
+ """文件状态枚举"""
68
+ PROCESSING = "PROCESSING"
69
+ ACTIVE = "ACTIVE"
70
+ FAILED = "FAILED"
71
+
72
+
73
+ class FileRecord(Base):
74
+ """
75
+ 文件记录表,用于存储上传到 Gemini 的文件信息
76
+ """
77
+ __tablename__ = "t_file_records"
78
+
79
+ id = Column(Integer, primary_key=True, autoincrement=True)
80
+
81
+ # 文件基本信息
82
+ name = Column(String(255), unique=True, nullable=False, comment="文件名称,格式: files/{file_id}")
83
+ display_name = Column(String(255), nullable=True, comment="用户上传时的原始文件名")
84
+ mime_type = Column(String(100), nullable=False, comment="MIME 类型")
85
+ size_bytes = Column(BigInteger, nullable=False, comment="文件大小(字节)")
86
+ sha256_hash = Column(String(255), nullable=True, comment="文件的 SHA256 哈希值")
87
+
88
+ # 状态信息
89
+ state = Column(Enum(FileState), nullable=False, default=FileState.PROCESSING, comment="文件状态")
90
+
91
+ # 时间戳
92
+ create_time = Column(DateTime, nullable=False, comment="创建时间")
93
+ update_time = Column(DateTime, nullable=False, comment="更新时间")
94
+ expiration_time = Column(DateTime, nullable=False, comment="过期时间")
95
+
96
+ # API 相关
97
+ uri = Column(String(500), nullable=False, comment="文件访问 URI")
98
+ api_key = Column(String(100), nullable=False, comment="上传时使用的 API Key")
99
+ upload_url = Column(Text, nullable=True, comment="临时上传 URL(用于分块上传)")
100
+
101
+ # 额外信息
102
+ user_token = Column(String(100), nullable=True, comment="上传用户的 token")
103
+ upload_completed = Column(DateTime, nullable=True, comment="上传完成时间")
104
+
105
+ def __repr__(self):
106
+ return f"<FileRecord(name='{self.name}', state='{self.state.value if self.state else 'None'}', api_key='{self.api_key[:8]}...')>"
107
+
108
+ def to_dict(self):
109
+ """转换为字典格式,用于 API 响应"""
110
+ return {
111
+ "name": self.name,
112
+ "displayName": self.display_name,
113
+ "mimeType": self.mime_type,
114
+ "sizeBytes": str(self.size_bytes),
115
+ "createTime": self.create_time.isoformat() + "Z",
116
+ "updateTime": self.update_time.isoformat() + "Z",
117
+ "expirationTime": self.expiration_time.isoformat() + "Z",
118
+ "sha256Hash": self.sha256_hash,
119
+ "uri": self.uri,
120
+ "state": self.state.value if self.state else "PROCESSING"
121
+ }
122
+
123
+ def is_expired(self):
124
+ """检查文件是否已过期"""
125
+ # 确保比较时都是 timezone-aware
126
+ expiration_time = self.expiration_time
127
+ if expiration_time.tzinfo is None:
128
+ expiration_time = expiration_time.replace(tzinfo=datetime.timezone.utc)
129
+ return datetime.datetime.now(datetime.timezone.utc) > expiration_time
app/database/services.py ADDED
@@ -0,0 +1,775 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 数据库服务模块
3
+ """
4
+
5
+ import json
6
+ from datetime import datetime, timedelta, timezone
7
+ from typing import Any, Dict, List, Optional, Union
8
+
9
+ from sqlalchemy import asc, delete, desc, func, insert, select, update
10
+
11
+ from app.database.connection import database
12
+ from app.database.models import ErrorLog, FileRecord, FileState, RequestLog, Settings
13
+ from app.log.logger import get_database_logger
14
+ from app.utils.helpers import redact_key_for_logging
15
+
16
+ logger = get_database_logger()
17
+
18
+
19
+ async def get_all_settings() -> List[Dict[str, Any]]:
20
+ """
21
+ 获取所有设置
22
+
23
+ Returns:
24
+ List[Dict[str, Any]]: 设置列表
25
+ """
26
+ try:
27
+ query = select(Settings)
28
+ result = await database.fetch_all(query)
29
+ return [dict(row) for row in result]
30
+ except Exception as e:
31
+ logger.error(f"Failed to get all settings: {str(e)}")
32
+ raise
33
+
34
+
35
+ async def get_setting(key: str) -> Optional[Dict[str, Any]]:
36
+ """
37
+ 获取指定键的设置
38
+
39
+ Args:
40
+ key: 设置键名
41
+
42
+ Returns:
43
+ Optional[Dict[str, Any]]: 设置信息,如果不存在则返回None
44
+ """
45
+ try:
46
+ query = select(Settings).where(Settings.key == key)
47
+ result = await database.fetch_one(query)
48
+ return dict(result) if result else None
49
+ except Exception as e:
50
+ logger.error(f"Failed to get setting {key}: {str(e)}")
51
+ raise
52
+
53
+
54
+ async def update_setting(
55
+ key: str, value: str, description: Optional[str] = None
56
+ ) -> bool:
57
+ """
58
+ 更新设置
59
+
60
+ Args:
61
+ key: 设置键名
62
+ value: 设置值
63
+ description: 设置描述
64
+
65
+ Returns:
66
+ bool: 是否更新成功
67
+ """
68
+ try:
69
+ # 检查设置是否存在
70
+ setting = await get_setting(key)
71
+
72
+ if setting:
73
+ # 更新设置
74
+ query = (
75
+ update(Settings)
76
+ .where(Settings.key == key)
77
+ .values(
78
+ value=value,
79
+ description=description if description else setting["description"],
80
+ updated_at=datetime.now(),
81
+ )
82
+ )
83
+ await database.execute(query)
84
+ logger.info(f"Updated setting: {key}")
85
+ return True
86
+ else:
87
+ # 插入设置
88
+ query = insert(Settings).values(
89
+ key=key,
90
+ value=value,
91
+ description=description,
92
+ created_at=datetime.now(),
93
+ updated_at=datetime.now(),
94
+ )
95
+ await database.execute(query)
96
+ logger.info(f"Inserted setting: {key}")
97
+ return True
98
+ except Exception as e:
99
+ logger.error(f"Failed to update setting {key}: {str(e)}")
100
+ return False
101
+
102
+
103
+ async def add_error_log(
104
+ gemini_key: Optional[str] = None,
105
+ model_name: Optional[str] = None,
106
+ error_type: Optional[str] = None,
107
+ error_log: Optional[str] = None,
108
+ error_code: Optional[int] = None,
109
+ request_msg: Optional[Union[Dict[str, Any], str]] = None,
110
+ request_datetime: Optional[datetime] = None,
111
+ ) -> bool:
112
+ """
113
+ 添加错误日志
114
+
115
+ Args:
116
+ gemini_key: Gemini API密钥
117
+ error_log: 错误日志
118
+ error_code: 错误代码 (例如 HTTP 状态码)
119
+ request_msg: 请求消息
120
+
121
+ Returns:
122
+ bool: 是否添加成功
123
+ """
124
+ try:
125
+ # 如果request_msg是字典,则转换为JSON字符串
126
+ if isinstance(request_msg, dict):
127
+ request_msg_json = request_msg
128
+ elif isinstance(request_msg, str):
129
+ try:
130
+ request_msg_json = json.loads(request_msg)
131
+ except json.JSONDecodeError:
132
+ request_msg_json = {"message": request_msg}
133
+ else:
134
+ request_msg_json = None
135
+
136
+ # 插入错误日志
137
+ query = insert(ErrorLog).values(
138
+ gemini_key=gemini_key,
139
+ error_type=error_type,
140
+ error_log=error_log,
141
+ model_name=model_name,
142
+ error_code=error_code,
143
+ request_msg=request_msg_json,
144
+ request_time=(request_datetime if request_datetime else datetime.now()),
145
+ )
146
+ await database.execute(query)
147
+ logger.info(f"Added error log for key: {redact_key_for_logging(gemini_key)}")
148
+ return True
149
+ except Exception as e:
150
+ logger.error(f"Failed to add error log: {str(e)}")
151
+ return False
152
+
153
+
154
+ async def get_error_logs(
155
+ limit: int = 20,
156
+ offset: int = 0,
157
+ key_search: Optional[str] = None,
158
+ error_search: Optional[str] = None,
159
+ error_code_search: Optional[str] = None,
160
+ start_date: Optional[datetime] = None,
161
+ end_date: Optional[datetime] = None,
162
+ sort_by: str = "id",
163
+ sort_order: str = "desc",
164
+ ) -> List[Dict[str, Any]]:
165
+ """
166
+ 获取错误日志,支持搜索、日期过滤和排序
167
+
168
+ Args:
169
+ limit (int): 限制数量
170
+ offset (int): 偏移量
171
+ key_search (Optional[str]): Gemini密钥搜索词 (模糊匹配)
172
+ error_search (Optional[str]): 错误类型或日志内容搜索词 (模糊匹配)
173
+ error_code_search (Optional[str]): 错误码搜索词 (精确匹配)
174
+ start_date (Optional[datetime]): 开始日期时间
175
+ end_date (Optional[datetime]): 结束日期时间
176
+ sort_by (str): 排序字段 (例如 'id', 'request_time')
177
+ sort_order (str): 排序顺序 ('asc' or 'desc')
178
+
179
+ Returns:
180
+ List[Dict[str, Any]]: 错误日志列表
181
+ """
182
+ try:
183
+ query = select(
184
+ ErrorLog.id,
185
+ ErrorLog.gemini_key,
186
+ ErrorLog.model_name,
187
+ ErrorLog.error_type,
188
+ ErrorLog.error_log,
189
+ ErrorLog.error_code,
190
+ ErrorLog.request_time,
191
+ )
192
+
193
+ if key_search:
194
+ query = query.where(ErrorLog.gemini_key.ilike(f"%{key_search}%"))
195
+ if error_search:
196
+ query = query.where(
197
+ (ErrorLog.error_type.ilike(f"%{error_search}%"))
198
+ | (ErrorLog.error_log.ilike(f"%{error_search}%"))
199
+ )
200
+ if start_date:
201
+ query = query.where(ErrorLog.request_time >= start_date)
202
+ if end_date:
203
+ query = query.where(ErrorLog.request_time < end_date)
204
+ if error_code_search:
205
+ try:
206
+ error_code_int = int(error_code_search)
207
+ query = query.where(ErrorLog.error_code == error_code_int)
208
+ except ValueError:
209
+ logger.warning(
210
+ f"Invalid format for error_code_search: '{error_code_search}'. Expected an integer. Skipping error code filter."
211
+ )
212
+
213
+ sort_column = getattr(ErrorLog, sort_by, ErrorLog.id)
214
+ if sort_order.lower() == "asc":
215
+ query = query.order_by(asc(sort_column))
216
+ else:
217
+ query = query.order_by(desc(sort_column))
218
+
219
+ query = query.limit(limit).offset(offset)
220
+
221
+ result = await database.fetch_all(query)
222
+ return [dict(row) for row in result]
223
+ except Exception as e:
224
+ logger.exception(f"Failed to get error logs with filters: {str(e)}")
225
+ raise
226
+
227
+
228
+ async def get_error_logs_count(
229
+ key_search: Optional[str] = None,
230
+ error_search: Optional[str] = None,
231
+ error_code_search: Optional[str] = None,
232
+ start_date: Optional[datetime] = None,
233
+ end_date: Optional[datetime] = None,
234
+ ) -> int:
235
+ """
236
+ 获取符合条件的错误日志总数
237
+
238
+ Args:
239
+ key_search (Optional[str]): Gemini密钥搜索词 (模糊匹配)
240
+ error_search (Optional[str]): 错误类型或日志内容搜索词 (模糊匹配)
241
+ error_code_search (Optional[str]): 错误码搜索词 (精确匹配)
242
+ start_date (Optional[datetime]): 开始日期时间
243
+ end_date (Optional[datetime]): 结束日期时间
244
+
245
+ Returns:
246
+ int: 日志总数
247
+ """
248
+ try:
249
+ query = select(func.count()).select_from(ErrorLog)
250
+
251
+ if key_search:
252
+ query = query.where(ErrorLog.gemini_key.ilike(f"%{key_search}%"))
253
+ if error_search:
254
+ query = query.where(
255
+ (ErrorLog.error_type.ilike(f"%{error_search}%"))
256
+ | (ErrorLog.error_log.ilike(f"%{error_search}%"))
257
+ )
258
+ if start_date:
259
+ query = query.where(ErrorLog.request_time >= start_date)
260
+ if end_date:
261
+ query = query.where(ErrorLog.request_time < end_date)
262
+ if error_code_search:
263
+ try:
264
+ error_code_int = int(error_code_search)
265
+ query = query.where(ErrorLog.error_code == error_code_int)
266
+ except ValueError:
267
+ logger.warning(
268
+ f"Invalid format for error_code_search in count: '{error_code_search}'. Expected an integer. Skipping error code filter."
269
+ )
270
+
271
+ count_result = await database.fetch_one(query)
272
+ return count_result[0] if count_result else 0
273
+ except Exception as e:
274
+ logger.exception(f"Failed to count error logs with filters: {str(e)}")
275
+ raise
276
+
277
+
278
+ # 新增函数:获取单条错误日志详情
279
+ async def get_error_log_details(log_id: int) -> Optional[Dict[str, Any]]:
280
+ """
281
+ 根据 ID 获取单个错误日志的详细信息
282
+
283
+ Args:
284
+ log_id (int): 错误日志的 ID
285
+
286
+ Returns:
287
+ Optional[Dict[str, Any]]: 包含日志详细信息的字典,如果未找到则返回 None
288
+ """
289
+ try:
290
+ query = select(ErrorLog).where(ErrorLog.id == log_id)
291
+ result = await database.fetch_one(query)
292
+ if result:
293
+ # 将 request_msg (JSONB) 转换为字符串以便在 API 中返回
294
+ log_dict = dict(result)
295
+ if "request_msg" in log_dict and log_dict["request_msg"] is not None:
296
+ # 确保即使是 None 或非 JSON 数据也能处理
297
+ try:
298
+ log_dict["request_msg"] = json.dumps(
299
+ log_dict["request_msg"], ensure_ascii=False, indent=2
300
+ )
301
+ except TypeError:
302
+ log_dict["request_msg"] = str(log_dict["request_msg"])
303
+ return log_dict
304
+ else:
305
+ return None
306
+ except Exception as e:
307
+ logger.exception(f"Failed to get error log details for ID {log_id}: {str(e)}")
308
+ raise
309
+
310
+
311
+ # 新增函数:通过 gemini_key / error_code / 时间窗口 查找最接近的错误日志
312
+ async def find_error_log_by_info(
313
+ gemini_key: str,
314
+ timestamp: datetime,
315
+ status_code: Optional[int] = None,
316
+ window_seconds: int = 1,
317
+ ) -> Optional[Dict[str, Any]]:
318
+ """
319
+ 在给定时间窗口内,根据 gemini_key(精确匹配)及可选的 status_code 查找最接近 timestamp 的错误日志。
320
+
321
+ 假设错误日志的 error_code 存储的是 HTTP 状态码或等价错误码。
322
+
323
+ Args:
324
+ gemini_key: 完整的 Gemini key 字符串。
325
+ timestamp: 目标时间(UTC 或本地,与存储一致)。
326
+ status_code: 可选的错误码,若提供则优先匹配该错误码。
327
+ window_seconds: 允许的时间偏差窗口,单位秒,默认为 1 秒。
328
+
329
+ Returns:
330
+ Optional[Dict[str, Any]]: 最匹配的一条错误日志的完整详情(字段与 get_error_log_details 一致),若未找到则返回 None。
331
+ """
332
+ try:
333
+ start_time = timestamp - timedelta(seconds=window_seconds)
334
+ end_time = timestamp + timedelta(seconds=window_seconds)
335
+
336
+ base_query = select(ErrorLog).where(
337
+ ErrorLog.gemini_key == gemini_key,
338
+ ErrorLog.request_time >= start_time,
339
+ ErrorLog.request_time <= end_time,
340
+ )
341
+
342
+ # 若提供了状态码,先尝试按状态码过滤
343
+ if status_code is not None:
344
+ query = base_query.where(ErrorLog.error_code == status_code).order_by(
345
+ ErrorLog.request_time.desc()
346
+ )
347
+ candidates = await database.fetch_all(query)
348
+ if not candidates:
349
+ # 回退:不按状态码,仅按时间窗口
350
+ query2 = base_query.order_by(ErrorLog.request_time.desc())
351
+ candidates = await database.fetch_all(query2)
352
+ else:
353
+ query = base_query.order_by(ErrorLog.request_time.desc())
354
+ candidates = await database.fetch_all(query)
355
+
356
+ if not candidates:
357
+ return None
358
+
359
+ # 在 Python 中选择与 timestamp 最接近的一条
360
+ def _to_dict(row: Any) -> Dict[str, Any]:
361
+ d = dict(row)
362
+ if "request_msg" in d and d["request_msg"] is not None:
363
+ try:
364
+ d["request_msg"] = json.dumps(
365
+ d["request_msg"], ensure_ascii=False, indent=2
366
+ )
367
+ except TypeError:
368
+ d["request_msg"] = str(d["request_msg"])
369
+ return d
370
+
371
+ best = min(
372
+ candidates,
373
+ key=lambda r: abs((r["request_time"] - timestamp).total_seconds()),
374
+ )
375
+ return _to_dict(best)
376
+ except Exception as e:
377
+ logger.exception(
378
+ f"Failed to find error log by info (key=***{gemini_key[-4:] if gemini_key else ''}, code={status_code}, ts={timestamp}, window={window_seconds}s): {str(e)}"
379
+ )
380
+ raise
381
+
382
+
383
+ async def delete_error_logs_by_ids(log_ids: List[int]) -> int:
384
+ """
385
+ 根据提供的 ID 列表批量删除错误日志 (异步)。
386
+
387
+ Args:
388
+ log_ids: 要删除的错误日志 ID 列表。
389
+
390
+ Returns:
391
+ int: 实际删除的日志数量。
392
+ """
393
+ if not log_ids:
394
+ return 0
395
+ try:
396
+ # 使用 databases 执行删除
397
+ query = delete(ErrorLog).where(ErrorLog.id.in_(log_ids))
398
+ # execute 返回受影响的行数,但 databases 库的 execute 不直接返回 rowcount
399
+ # 我们需要先查询是否存在,或者依赖数据库约束/触发器(如果适用)
400
+ # 或者,我们可以执行删除并假设成功,除非抛出异常
401
+ # 为了简单起见,我们执行删除并记录日志,不精确返回删除数量
402
+ # 如果需要精确数量,需要先执行 SELECT COUNT(*)
403
+ await database.execute(query)
404
+ # 注意:databases 的 execute 不返回 rowcount,所以我们不能直接返回删除的数量
405
+ # 返回 log_ids 的长度作为尝试删除的数量,或者返回 0/1 表示操作尝试
406
+ logger.info(f"Attempted bulk deletion for error logs with IDs: {log_ids}")
407
+ return len(log_ids) # 返回尝试删除的数量
408
+ except Exception as e:
409
+ # 数据库连接或执行错误
410
+ logger.error(
411
+ f"Error during bulk deletion of error logs {log_ids}: {e}", exc_info=True
412
+ )
413
+ raise
414
+
415
+
416
+ async def delete_error_log_by_id(log_id: int) -> bool:
417
+ """
418
+ 根据 ID 删除单个错误日志 (异步)。
419
+
420
+ Args:
421
+ log_id: 要删除的错误日志 ID。
422
+
423
+ Returns:
424
+ bool: 如果成功删除返回 True,否则返回 False。
425
+ """
426
+ try:
427
+ # 先检查是否存在 (可选,但更明确)
428
+ check_query = select(ErrorLog.id).where(ErrorLog.id == log_id)
429
+ exists = await database.fetch_one(check_query)
430
+
431
+ if not exists:
432
+ logger.warning(
433
+ f"Attempted to delete non-existent error log with ID: {log_id}"
434
+ )
435
+ return False
436
+
437
+ # 执行删除
438
+ delete_query = delete(ErrorLog).where(ErrorLog.id == log_id)
439
+ await database.execute(delete_query)
440
+ logger.info(f"Successfully deleted error log with ID: {log_id}")
441
+ return True
442
+ except Exception as e:
443
+ logger.error(f"Error deleting error log with ID {log_id}: {e}", exc_info=True)
444
+ raise
445
+
446
+
447
+ async def delete_all_error_logs() -> int:
448
+ """
449
+ 删除所有错误日志条目。
450
+
451
+ Returns:
452
+ int: 被删除的错误日志数量。如果使用的数据库驱动不支持返回受影响行数,则返回 -1 表示操作成功。
453
+ """
454
+ try:
455
+ # 直接执行删除操作,避免不必要的查询
456
+ delete_query = delete(ErrorLog)
457
+ await database.execute(delete_query)
458
+
459
+ logger.info("Successfully deleted all error logs.")
460
+
461
+ # 由于 databases 库的 execute 方法不返回受影响的行数,
462
+ # 返回 -1 表示删除操作成功执行,但具体删除数量未知
463
+ # 这比先查询再删除的方式更高效
464
+ return -1
465
+ except Exception as e:
466
+ logger.error(f"Failed to delete all error logs: {str(e)}", exc_info=True)
467
+ raise
468
+
469
+
470
+ # 新增函数:添加请求日志
471
+ async def add_request_log(
472
+ model_name: Optional[str],
473
+ api_key: Optional[str],
474
+ is_success: bool,
475
+ status_code: Optional[int] = None,
476
+ latency_ms: Optional[int] = None,
477
+ request_time: Optional[datetime] = None,
478
+ ) -> bool:
479
+ """
480
+ 添加 API 请求日志
481
+
482
+ Args:
483
+ model_name: 模型名称
484
+ api_key: 使用的 API 密钥
485
+ is_success: 请求是否成功
486
+ status_code: API 响应状态码
487
+ latency_ms: 请求耗时(毫秒)
488
+ request_time: 请求发生时间 (如果为 None, 则使用当前时间)
489
+
490
+ Returns:
491
+ bool: 是否添加成功
492
+ """
493
+ try:
494
+ log_time = request_time if request_time else datetime.now()
495
+
496
+ query = insert(RequestLog).values(
497
+ request_time=log_time,
498
+ model_name=model_name,
499
+ api_key=api_key,
500
+ is_success=is_success,
501
+ status_code=status_code,
502
+ latency_ms=latency_ms,
503
+ )
504
+ await database.execute(query)
505
+ return True
506
+ except Exception as e:
507
+ logger.error(f"Failed to add request log: {str(e)}")
508
+ return False
509
+
510
+
511
+ # ==================== 文件记录相关函数 ====================
512
+
513
+
514
+ async def create_file_record(
515
+ name: str,
516
+ mime_type: str,
517
+ size_bytes: int,
518
+ api_key: str,
519
+ uri: str,
520
+ create_time: datetime,
521
+ update_time: datetime,
522
+ expiration_time: datetime,
523
+ state: FileState = FileState.PROCESSING,
524
+ display_name: Optional[str] = None,
525
+ sha256_hash: Optional[str] = None,
526
+ upload_url: Optional[str] = None,
527
+ user_token: Optional[str] = None,
528
+ ) -> Dict[str, Any]:
529
+ """
530
+ 创建文件记录
531
+
532
+ Args:
533
+ name: 文件名称(格式: files/{file_id})
534
+ mime_type: MIME 类型
535
+ size_bytes: 文件大小(字节)
536
+ api_key: 上传时使用的 API Key
537
+ uri: 文件访问 URI
538
+ create_time: 创建时间
539
+ update_time: 更新时间
540
+ expiration_time: 过期时间
541
+ display_name: 显示名称
542
+ sha256_hash: SHA256 哈希值
543
+ upload_url: 临时上传 URL
544
+ user_token: 上传用户的 token
545
+
546
+ Returns:
547
+ Dict[str, Any]: 创建的文件记录
548
+ """
549
+ try:
550
+ query = insert(FileRecord).values(
551
+ name=name,
552
+ display_name=display_name,
553
+ mime_type=mime_type,
554
+ size_bytes=size_bytes,
555
+ sha256_hash=sha256_hash,
556
+ state=state,
557
+ create_time=create_time,
558
+ update_time=update_time,
559
+ expiration_time=expiration_time,
560
+ uri=uri,
561
+ api_key=api_key,
562
+ upload_url=upload_url,
563
+ user_token=user_token,
564
+ )
565
+ await database.execute(query)
566
+
567
+ # 返回创建的记录
568
+ return await get_file_record_by_name(name)
569
+ except Exception as e:
570
+ logger.error(f"Failed to create file record: {str(e)}")
571
+ raise
572
+
573
+
574
+ async def get_file_record_by_name(name: str) -> Optional[Dict[str, Any]]:
575
+ """
576
+ 根据文件名获取文件记录
577
+
578
+ Args:
579
+ name: 文件名称(格式: files/{file_id})
580
+
581
+ Returns:
582
+ Optional[Dict[str, Any]]: 文件记录,如果不存在则返回 None
583
+ """
584
+ try:
585
+ query = select(FileRecord).where(FileRecord.name == name)
586
+ result = await database.fetch_one(query)
587
+ return dict(result) if result else None
588
+ except Exception as e:
589
+ logger.error(f"Failed to get file record by name {name}: {str(e)}")
590
+ raise
591
+
592
+
593
+ async def update_file_record_state(
594
+ file_name: str,
595
+ state: FileState,
596
+ update_time: Optional[datetime] = None,
597
+ upload_completed: Optional[datetime] = None,
598
+ sha256_hash: Optional[str] = None,
599
+ ) -> bool:
600
+ """
601
+ 更新文件记录状态
602
+
603
+ Args:
604
+ file_name: 文件名
605
+ state: 新状态
606
+ update_time: 更新时间
607
+ upload_completed: 上传完成时间
608
+ sha256_hash: SHA256 哈希值
609
+
610
+ Returns:
611
+ bool: 是否更新成功
612
+ """
613
+ try:
614
+ values = {"state": state}
615
+ if update_time:
616
+ values["update_time"] = update_time
617
+ if upload_completed:
618
+ values["upload_completed"] = upload_completed
619
+ if sha256_hash:
620
+ values["sha256_hash"] = sha256_hash
621
+
622
+ query = update(FileRecord).where(FileRecord.name == file_name).values(**values)
623
+ result = await database.execute(query)
624
+
625
+ if result:
626
+ logger.info(f"Updated file record state for {file_name} to {state}")
627
+ return True
628
+
629
+ logger.warning(f"File record not found for update: {file_name}")
630
+ return False
631
+ except Exception as e:
632
+ logger.error(f"Failed to update file record state: {str(e)}")
633
+ return False
634
+
635
+
636
+ async def list_file_records(
637
+ user_token: Optional[str] = None,
638
+ api_key: Optional[str] = None,
639
+ page_size: int = 10,
640
+ page_token: Optional[str] = None,
641
+ ) -> tuple[List[Dict[str, Any]], Optional[str]]:
642
+ """
643
+ 列出文件记录
644
+
645
+ Args:
646
+ user_token: 用户 token(如果提供,只返回该用户的文件)
647
+ api_key: API Key(如果提供,只返回使用该 key 的文件)
648
+ page_size: 每页大小
649
+ page_token: 分页标记(偏移量)
650
+
651
+ Returns:
652
+ tuple[List[Dict[str, Any]], Optional[str]]: (文件列表, 下一页标记)
653
+ """
654
+ try:
655
+ logger.debug(
656
+ f"list_file_records called with page_size={page_size}, page_token={page_token}"
657
+ )
658
+ query = select(FileRecord).where(
659
+ FileRecord.expiration_time > datetime.now(timezone.utc)
660
+ )
661
+
662
+ if user_token:
663
+ query = query.where(FileRecord.user_token == user_token)
664
+ if api_key:
665
+ query = query.where(FileRecord.api_key == api_key)
666
+
667
+ # 使用偏移量进行分页
668
+ offset = 0
669
+ if page_token:
670
+ try:
671
+ offset = int(page_token)
672
+ except ValueError:
673
+ logger.warning(f"Invalid page token: {page_token}")
674
+ offset = 0
675
+
676
+ # 按ID升序排列,使用 OFFSET 和 LIMIT
677
+ query = query.order_by(FileRecord.id).offset(offset).limit(page_size + 1)
678
+
679
+ results = await database.fetch_all(query)
680
+
681
+ logger.debug(f"Query returned {len(results)} records")
682
+ if results:
683
+ logger.debug(
684
+ f"First record ID: {results[0]['id']}, Last record ID: {results[-1]['id']}"
685
+ )
686
+
687
+ # 处理分页
688
+ has_next = len(results) > page_size
689
+ if has_next:
690
+ results = results[:page_size]
691
+ # 下一页的偏移量是当前偏移量加上本页返回的记录数
692
+ next_offset = offset + page_size
693
+ next_page_token = str(next_offset)
694
+ logger.debug(
695
+ f"Has next page, offset={offset}, page_size={page_size}, next_page_token={next_page_token}"
696
+ )
697
+ else:
698
+ next_page_token = None
699
+ logger.debug(f"No next page, returning {len(results)} results")
700
+
701
+ return [dict(row) for row in results], next_page_token
702
+ except Exception as e:
703
+ logger.error(f"Failed to list file records: {str(e)}")
704
+ raise
705
+
706
+
707
+ async def delete_file_record(name: str) -> bool:
708
+ """
709
+ 删除文件记录
710
+
711
+ Args:
712
+ name: 文件名称
713
+
714
+ Returns:
715
+ bool: 是否删除成功
716
+ """
717
+ try:
718
+ query = delete(FileRecord).where(FileRecord.name == name)
719
+ await database.execute(query)
720
+ return True
721
+ except Exception as e:
722
+ logger.error(f"Failed to delete file record: {str(e)}")
723
+ return False
724
+
725
+
726
+ async def delete_expired_file_records() -> List[Dict[str, Any]]:
727
+ """
728
+ 删除已过期的文件记录
729
+
730
+ Returns:
731
+ List[Dict[str, Any]]: 删除的记录列表
732
+ """
733
+ try:
734
+ # 先获取要删除的记录
735
+ query = select(FileRecord).where(
736
+ FileRecord.expiration_time <= datetime.now(timezone.utc)
737
+ )
738
+ expired_records = await database.fetch_all(query)
739
+
740
+ if not expired_records:
741
+ return []
742
+
743
+ # 执行删除
744
+ delete_query = delete(FileRecord).where(
745
+ FileRecord.expiration_time <= datetime.now(timezone.utc)
746
+ )
747
+ await database.execute(delete_query)
748
+
749
+ logger.info(f"Deleted {len(expired_records)} expired file records")
750
+ return [dict(record) for record in expired_records]
751
+ except Exception as e:
752
+ logger.error(f"Failed to delete expired file records: {str(e)}")
753
+ raise
754
+
755
+
756
+ async def get_file_api_key(name: str) -> Optional[str]:
757
+ """
758
+ 获取文件对应的 API Key
759
+
760
+ Args:
761
+ name: 文件名称
762
+
763
+ Returns:
764
+ Optional[str]: API Key,如果文件不存在或已过期则返回 None
765
+ """
766
+ try:
767
+ query = select(FileRecord.api_key).where(
768
+ (FileRecord.name == name)
769
+ & (FileRecord.expiration_time > datetime.now(timezone.utc))
770
+ )
771
+ result = await database.fetch_one(query)
772
+ return result["api_key"] if result else None
773
+ except Exception as e:
774
+ logger.error(f"Failed to get file API key: {str(e)}")
775
+ raise
app/domain/file_models.py ADDED
@@ -0,0 +1,69 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Files API 相关的领域模型
3
+ """
4
+ from typing import Optional, Dict, Any, List
5
+ from datetime import datetime
6
+ from pydantic import BaseModel, Field
7
+
8
+
9
+ class FileUploadConfig(BaseModel):
10
+ """文件上传配置"""
11
+ mime_type: Optional[str] = Field(None, description="MIME 类型")
12
+ display_name: Optional[str] = Field(None, description="显示名称,最多40个字符")
13
+
14
+
15
+ class CreateFileRequest(BaseModel):
16
+ """创建文件请求(用于初始化上传)"""
17
+ file: Optional[Dict[str, Any]] = Field(None, description="文件元数据")
18
+
19
+
20
+ class FileMetadata(BaseModel):
21
+ """文件元数据响应"""
22
+ name: str = Field(..., description="文件名称,格式: files/{file_id}")
23
+ displayName: Optional[str] = Field(None, description="显示名称")
24
+ mimeType: str = Field(..., description="MIME 类型")
25
+ sizeBytes: str = Field(..., description="文件大小(字节)")
26
+ createTime: str = Field(..., description="创建时间 (RFC3339)")
27
+ updateTime: str = Field(..., description="更新时间 (RFC3339)")
28
+ expirationTime: str = Field(..., description="过期时间 (RFC3339)")
29
+ sha256Hash: Optional[str] = Field(None, description="SHA256 哈希值")
30
+ uri: str = Field(..., description="文件访问 URI")
31
+ state: str = Field(..., description="文件状态")
32
+
33
+ class Config:
34
+ json_encoders = {
35
+ datetime: lambda v: v.isoformat() + "Z"
36
+ }
37
+
38
+
39
+ class ListFilesRequest(BaseModel):
40
+ """列出文件请求参数"""
41
+ pageSize: Optional[int] = Field(10, ge=1, le=100, description="每页大小")
42
+ pageToken: Optional[str] = Field(None, description="分页标记")
43
+
44
+
45
+ class ListFilesResponse(BaseModel):
46
+ """列出文件响应"""
47
+ files: List[FileMetadata] = Field(default_factory=list, description="文件列表")
48
+ nextPageToken: Optional[str] = Field(None, description="下一页标记")
49
+
50
+
51
+ class UploadInitResponse(BaseModel):
52
+ """上传初始化响应(内部使用)"""
53
+ file_metadata: FileMetadata
54
+ upload_url: str
55
+
56
+
57
+ class FileKeyMapping(BaseModel):
58
+ """文件与 API Key 的映射关系(内部使用)"""
59
+ file_name: str
60
+ api_key: str
61
+ user_token: str
62
+ created_at: datetime
63
+ expires_at: datetime
64
+
65
+
66
+ class DeleteFileResponse(BaseModel):
67
+ """删除文件响应"""
68
+ success: bool = Field(..., description="是否删除成功")
69
+ message: Optional[str] = Field(None, description="消息")
app/domain/gemini_models.py ADDED
@@ -0,0 +1,115 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Any, Dict, List, Literal, Optional, Union
2
+
3
+ from pydantic import BaseModel, Field
4
+
5
+ from app.core.constants import DEFAULT_TEMPERATURE, DEFAULT_TOP_K, DEFAULT_TOP_P
6
+
7
+
8
+ class SafetySetting(BaseModel):
9
+ category: Optional[
10
+ Literal[
11
+ "HARM_CATEGORY_HATE_SPEECH",
12
+ "HARM_CATEGORY_DANGEROUS_CONTENT",
13
+ "HARM_CATEGORY_HARASSMENT",
14
+ "HARM_CATEGORY_SEXUALLY_EXPLICIT",
15
+ "HARM_CATEGORY_CIVIC_INTEGRITY",
16
+ ]
17
+ ] = None
18
+ threshold: Optional[
19
+ Literal[
20
+ "HARM_BLOCK_THRESHOLD_UNSPECIFIED",
21
+ "BLOCK_LOW_AND_ABOVE",
22
+ "BLOCK_MEDIUM_AND_ABOVE",
23
+ "BLOCK_ONLY_HIGH",
24
+ "BLOCK_NONE",
25
+ "OFF",
26
+ ]
27
+ ] = None
28
+
29
+
30
+ class GenerationConfig(BaseModel):
31
+ stopSequences: Optional[List[str]] = None
32
+ responseMimeType: Optional[str] = None
33
+ responseSchema: Optional[Dict[str, Any]] = None
34
+ candidateCount: Optional[int] = 1
35
+ maxOutputTokens: Optional[int] = None
36
+ temperature: Optional[float] = DEFAULT_TEMPERATURE
37
+ topP: Optional[float] = DEFAULT_TOP_P
38
+ topK: Optional[int] = DEFAULT_TOP_K
39
+ presencePenalty: Optional[float] = None
40
+ frequencyPenalty: Optional[float] = None
41
+ responseLogprobs: Optional[bool] = None
42
+ logprobs: Optional[int] = None
43
+ thinkingConfig: Optional[Dict[str, Any]] = None
44
+ # TTS相关字段
45
+ responseModalities: Optional[List[str]] = None
46
+ speechConfig: Optional[Dict[str, Any]] = None
47
+
48
+
49
+ class SystemInstruction(BaseModel):
50
+ role: Optional[str] = "system"
51
+ parts: Union[List[Dict[str, Any]], Dict[str, Any]]
52
+
53
+
54
+ class GeminiContent(BaseModel):
55
+ role: Optional[str] = None
56
+ parts: List[Dict[str, Any]]
57
+
58
+
59
+ class GeminiRequest(BaseModel):
60
+ contents: List[GeminiContent] = []
61
+ tools: Optional[Union[List[Dict[str, Any]], Dict[str, Any]]] = []
62
+ safetySettings: Optional[List[SafetySetting]] = Field(
63
+ default=None, alias="safety_settings"
64
+ )
65
+ generationConfig: Optional[GenerationConfig] = Field(
66
+ default=None, alias="generation_config"
67
+ )
68
+ systemInstruction: Optional[SystemInstruction] = Field(
69
+ default=None, alias="system_instruction"
70
+ )
71
+
72
+ class Config:
73
+ populate_by_name = True
74
+
75
+
76
+ class ResetSelectedKeysRequest(BaseModel):
77
+ keys: List[str]
78
+ key_type: str
79
+
80
+
81
+ class VerifySelectedKeysRequest(BaseModel):
82
+ keys: List[str]
83
+
84
+
85
+ class GeminiEmbedContent(BaseModel):
86
+ """嵌入内容模型"""
87
+
88
+ parts: List[Dict[str, str]]
89
+
90
+
91
+ class GeminiEmbedRequest(BaseModel):
92
+ """单一嵌入请求模型"""
93
+
94
+ content: GeminiEmbedContent
95
+ taskType: Optional[
96
+ Literal[
97
+ "TASK_TYPE_UNSPECIFIED",
98
+ "RETRIEVAL_QUERY",
99
+ "RETRIEVAL_DOCUMENT",
100
+ "SEMANTIC_SIMILARITY",
101
+ "CLASSIFICATION",
102
+ "CLUSTERING",
103
+ "QUESTION_ANSWERING",
104
+ "FACT_VERIFICATION",
105
+ "CODE_RETRIEVAL_QUERY",
106
+ ]
107
+ ] = None
108
+ title: Optional[str] = None
109
+ outputDimensionality: Optional[int] = None
110
+
111
+
112
+ class GeminiBatchEmbedRequest(BaseModel):
113
+ """批量嵌入请求模型"""
114
+
115
+ requests: List[GeminiEmbedRequest]
app/domain/image_models.py ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from typing import Union
2
+
3
+
4
+ class ImageMetadata:
5
+ def __init__(self, width: int, height: int, filename: str, size: int, url: str, delete_url: Union[str, None] = None):
6
+ self.width = width
7
+ self.height = height
8
+ self.filename = filename
9
+ self.size = size
10
+ self.url = url
11
+ self.delete_url = delete_url
12
+ class UploadResponse:
13
+ def __init__(self, success: bool, code: str, message: str, data: ImageMetadata):
14
+ self.success = success
15
+ self.code = code
16
+ self.message = message
17
+ self.data = data
18
+ class ImageUploader:
19
+ def upload(self, file: bytes, filename: str) -> UploadResponse:
20
+ raise NotImplementedError
app/domain/openai_models.py ADDED
@@ -0,0 +1,43 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from pydantic import BaseModel
2
+ from typing import Any, Dict, List, Optional, Union
3
+
4
+ from app.core.constants import DEFAULT_MODEL, DEFAULT_TEMPERATURE, DEFAULT_TOP_K, DEFAULT_TOP_P
5
+
6
+
7
+ class ChatRequest(BaseModel):
8
+ messages: List[dict]
9
+ model: str = DEFAULT_MODEL
10
+ temperature: Optional[float] = DEFAULT_TEMPERATURE
11
+ stream: Optional[bool] = False
12
+ max_tokens: Optional[int] = None
13
+ top_p: Optional[float] = DEFAULT_TOP_P
14
+ top_k: Optional[int] = DEFAULT_TOP_K
15
+ n: Optional[int] = 1
16
+ stop: Optional[Union[List[str],str]] = None
17
+ reasoning_effort: Optional[str] = None
18
+ tools: Optional[Union[List[Dict[str, Any]], Dict[str, Any]]] = []
19
+ tool_choice: Optional[str] = None
20
+ response_format: Optional[dict] = None
21
+
22
+
23
+ class EmbeddingRequest(BaseModel):
24
+ input: Union[str, List[str]]
25
+ model: str = "text-embedding-004"
26
+ encoding_format: Optional[str] = "float"
27
+
28
+
29
+ class ImageGenerationRequest(BaseModel):
30
+ model: str = "imagen-3.0-generate-002"
31
+ prompt: str = ""
32
+ n: int = 1
33
+ size: Optional[str] = "1024x1024"
34
+ quality: Optional[str] = None
35
+ style: Optional[str] = None
36
+ response_format: Optional[str] = "url"
37
+
38
+
39
+ class TTSRequest(BaseModel):
40
+ model: str = "gemini-2.5-flash-preview-tts"
41
+ input: str
42
+ voice: str = "Kore"
43
+ response_format: Optional[str] = "wav"
app/exception/exceptions.py ADDED
@@ -0,0 +1,140 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 异常处理模块,定义应用程序中使用的自定义异常和异常处理器
3
+ """
4
+
5
+ from fastapi import FastAPI, Request
6
+ from fastapi.exceptions import RequestValidationError
7
+ from fastapi.responses import JSONResponse
8
+ from starlette.exceptions import HTTPException as StarletteHTTPException
9
+
10
+ from app.log.logger import get_exceptions_logger
11
+
12
+ logger = get_exceptions_logger()
13
+
14
+
15
+ class APIError(Exception):
16
+ """API错误基类"""
17
+
18
+ def __init__(self, status_code: int, detail: str, error_code: str = None):
19
+ self.status_code = status_code
20
+ self.detail = detail
21
+ self.error_code = error_code or "api_error"
22
+ super().__init__(self.detail)
23
+
24
+
25
+ class AuthenticationError(APIError):
26
+ """认证错误"""
27
+
28
+ def __init__(self, detail: str = "Authentication failed"):
29
+ super().__init__(
30
+ status_code=401, detail=detail, error_code="authentication_error"
31
+ )
32
+
33
+
34
+ class AuthorizationError(APIError):
35
+ """授权错误"""
36
+
37
+ def __init__(self, detail: str = "Not authorized to access this resource"):
38
+ super().__init__(
39
+ status_code=403, detail=detail, error_code="authorization_error"
40
+ )
41
+
42
+
43
+ class ResourceNotFoundError(APIError):
44
+ """资源未找到错误"""
45
+
46
+ def __init__(self, detail: str = "Resource not found"):
47
+ super().__init__(
48
+ status_code=404, detail=detail, error_code="resource_not_found"
49
+ )
50
+
51
+
52
+ class ModelNotSupportedError(APIError):
53
+ """模型不支持错误"""
54
+
55
+ def __init__(self, model: str):
56
+ super().__init__(
57
+ status_code=400,
58
+ detail=f"Model {model} is not supported",
59
+ error_code="model_not_supported",
60
+ )
61
+
62
+
63
+ class APIKeyError(APIError):
64
+ """API密钥错误"""
65
+
66
+ def __init__(self, detail: str = "Invalid or expired API key"):
67
+ super().__init__(status_code=401, detail=detail, error_code="api_key_error")
68
+
69
+
70
+ class ServiceUnavailableError(APIError):
71
+ """服务不可用错误"""
72
+
73
+ def __init__(self, detail: str = "Service temporarily unavailable"):
74
+ super().__init__(
75
+ status_code=503, detail=detail, error_code="service_unavailable"
76
+ )
77
+
78
+
79
+ def setup_exception_handlers(app: FastAPI) -> None:
80
+ """
81
+ 设置应用程序的异常处理器
82
+
83
+ Args:
84
+ app: FastAPI应用程序实例
85
+ """
86
+
87
+ @app.exception_handler(APIError)
88
+ async def api_error_handler(request: Request, exc: APIError):
89
+ """处理API错误"""
90
+ logger.error(f"API Error: {exc.detail} (Code: {exc.error_code})")
91
+ return JSONResponse(
92
+ status_code=exc.status_code,
93
+ content={"error": {"code": exc.error_code, "message": exc.detail}},
94
+ )
95
+
96
+ @app.exception_handler(StarletteHTTPException)
97
+ async def http_exception_handler(request: Request, exc: StarletteHTTPException):
98
+ """处理HTTP异常"""
99
+ logger.error(f"HTTP Exception: {exc.detail} (Status: {exc.status_code})")
100
+ return JSONResponse(
101
+ status_code=exc.status_code,
102
+ content={"error": {"code": "http_error", "message": exc.detail}},
103
+ )
104
+
105
+ @app.exception_handler(RequestValidationError)
106
+ async def validation_exception_handler(
107
+ request: Request, exc: RequestValidationError
108
+ ):
109
+ """处理请求验证错误"""
110
+ error_details = []
111
+ for error in exc.errors():
112
+ error_details.append(
113
+ {"loc": error["loc"], "msg": error["msg"], "type": error["type"]}
114
+ )
115
+
116
+ logger.error(f"Validation Error: {error_details}")
117
+ return JSONResponse(
118
+ status_code=422,
119
+ content={
120
+ "error": {
121
+ "code": "validation_error",
122
+ "message": "Request validation failed",
123
+ "details": error_details,
124
+ }
125
+ },
126
+ )
127
+
128
+ @app.exception_handler(Exception)
129
+ async def general_exception_handler(request: Request, exc: Exception):
130
+ """处理通用异常"""
131
+ logger.exception(f"Unhandled Exception: {str(exc)}")
132
+ return JSONResponse(
133
+ status_code=500,
134
+ content={
135
+ "error": {
136
+ "code": "internal_server_error",
137
+ "message": "An unexpected error occurred",
138
+ }
139
+ },
140
+ )
app/handler/error_handler.py ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from contextlib import asynccontextmanager
2
+ from fastapi import HTTPException
3
+ import logging
4
+
5
+ @asynccontextmanager
6
+ async def handle_route_errors(logger: logging.Logger, operation_name: str, success_message: str = None, failure_message: str = None):
7
+ """
8
+ 一个异步上下文管理器,用于统一处理 FastAPI 路由中的常见错误和日志记录。
9
+
10
+ Args:
11
+ logger: 用于记录日志的 Logger 实例。
12
+ operation_name: 操作的名称,用于日志记录和错误详情。
13
+ success_message: 操作成功时记录的自定义消息 (可选)。
14
+ failure_message: 操作失败时记录的自定义消息 (可选)。
15
+ """
16
+ default_success_msg = f"{operation_name} request successful"
17
+ default_failure_msg = f"{operation_name} request failed"
18
+
19
+ logger.info("-" * 50 + operation_name + "-" * 50)
20
+ try:
21
+ yield
22
+ logger.info(success_message or default_success_msg)
23
+ except HTTPException as http_exc:
24
+ # 如果已经是 HTTPException,直接重新抛出,保留原始状态码和详情
25
+ logger.error(f"{failure_message or default_failure_msg}: {http_exc.detail} (Status: {http_exc.status_code})")
26
+ raise http_exc
27
+ except Exception as e:
28
+ # 对于其他所有异常,记录错误并抛出标准的 500 错误
29
+ logger.error(f"{failure_message or default_failure_msg}: {str(e)}")
30
+ raise HTTPException(
31
+ status_code=500, detail=f"Internal server error during {operation_name}"
32
+ ) from e
app/handler/message_converter.py ADDED
@@ -0,0 +1,363 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import base64
2
+ import json
3
+ import re
4
+ from abc import ABC, abstractmethod
5
+ from typing import Any, Dict, List, Optional
6
+
7
+ import requests
8
+
9
+ from app.core.constants import (
10
+ AUDIO_FORMAT_TO_MIMETYPE,
11
+ DATA_URL_PATTERN,
12
+ IMAGE_URL_PATTERN,
13
+ MAX_AUDIO_SIZE_BYTES,
14
+ MAX_VIDEO_SIZE_BYTES,
15
+ SUPPORTED_AUDIO_FORMATS,
16
+ SUPPORTED_ROLES,
17
+ SUPPORTED_VIDEO_FORMATS,
18
+ VIDEO_FORMAT_TO_MIMETYPE,
19
+ )
20
+ from app.log.logger import get_message_converter_logger
21
+
22
+ logger = get_message_converter_logger()
23
+
24
+
25
+ class MessageConverter(ABC):
26
+ """消息转换器基类"""
27
+
28
+ @abstractmethod
29
+ def convert(
30
+ self, messages: List[Dict[str, Any]], model: str
31
+ ) -> tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]:
32
+ pass
33
+
34
+
35
+ def _get_mime_type_and_data(base64_string):
36
+ """
37
+ 从 base64 字符串中提取 MIME 类型和数据。
38
+
39
+ 参数:
40
+ base64_string (str): 可能包含 MIME 类型信息的 base64 字符串
41
+
42
+ 返回:
43
+ tuple: (mime_type, encoded_data)
44
+ """
45
+ # 检查字符串是否以 "data:" 格式开始
46
+ if base64_string.startswith("data:"):
47
+ # 提取 MIME 类型和数据
48
+ pattern = DATA_URL_PATTERN
49
+ match = re.match(pattern, base64_string)
50
+ if match:
51
+ mime_type = (
52
+ "image/jpeg" if match.group(1) == "image/jpg" else match.group(1)
53
+ )
54
+ encoded_data = match.group(2)
55
+ return mime_type, encoded_data
56
+
57
+ # 如果不是预期格式,假定它只是数据部分
58
+ return None, base64_string
59
+
60
+
61
+ def _convert_image(image_url: str) -> Dict[str, Any]:
62
+ if image_url.startswith("data:image"):
63
+ mime_type, encoded_data = _get_mime_type_and_data(image_url)
64
+ return {"inline_data": {"mime_type": mime_type, "data": encoded_data}}
65
+ else:
66
+ encoded_data = _convert_image_to_base64(image_url)
67
+ return {"inline_data": {"mime_type": "image/png", "data": encoded_data}}
68
+
69
+
70
+ def _convert_image_to_base64(url: str) -> str:
71
+ """
72
+ 将图片URL转换为base64编码
73
+ Args:
74
+ url: 图片URL
75
+ Returns:
76
+ str: base64编码的图片数据
77
+ """
78
+ response = requests.get(url)
79
+ if response.status_code == 200:
80
+ # 将图片内容转换为base64
81
+ img_data = base64.b64encode(response.content).decode("utf-8")
82
+ return img_data
83
+ else:
84
+ raise Exception(f"Failed to fetch image: {response.status_code}")
85
+
86
+
87
+ def _process_text_with_image(text: str, model: str) -> List[Dict[str, Any]]:
88
+ """
89
+ 处理可能包含图片URL的文本,提取图片并转换为base64
90
+
91
+ Args:
92
+ text: 可能包含图片URL的文本
93
+
94
+ Returns:
95
+ List[Dict[str, Any]]: 包含文本和图片的部分列表
96
+ """
97
+ # 如果模型名中没有包含image,当作普通文本处理
98
+ if "image" not in model:
99
+ return [{"text": text}]
100
+ parts = []
101
+ img_url_match = re.search(IMAGE_URL_PATTERN, text)
102
+ if img_url_match:
103
+ # 提取URL
104
+ img_url = img_url_match.group(2)
105
+ # 先判断是否是base64url如果是,直接用,不过不是,再将URL对应的图片转换为base64
106
+ try:
107
+ base64_url_match = re.search(DATA_URL_PATTERN, img_url)
108
+ if base64_url_match:
109
+ parts.append(
110
+ {
111
+ "inline_data": {
112
+ "mimeType": base64_url_match.group(1),
113
+ "data": base64_url_match.group(2),
114
+ }
115
+ }
116
+ )
117
+ else:
118
+ base64_data = _convert_image_to_base64(img_url)
119
+ parts.append(
120
+ {"inline_data": {"mimeType": "image/png", "data": base64_data}}
121
+ )
122
+ except Exception:
123
+ # 如果转换失败,回退到文本模式
124
+ parts.append({"text": text})
125
+ else:
126
+ # 没有图片URL,作为纯文本处理
127
+ parts.append({"text": text})
128
+ return parts
129
+
130
+
131
+ class OpenAIMessageConverter(MessageConverter):
132
+ """OpenAI消息格式转换器"""
133
+
134
+ def _validate_media_data(
135
+ self, format: str, data: str, supported_formats: List[str], max_size: int
136
+ ) -> tuple[Optional[str], Optional[str]]:
137
+ """Validates format and size of Base64 media data."""
138
+ if format.lower() not in supported_formats:
139
+ logger.error(
140
+ f"Unsupported media format: {format}. Supported: {supported_formats}"
141
+ )
142
+ raise ValueError(f"Unsupported media format: {format}")
143
+
144
+ try:
145
+ decoded_data = base64.b64decode(data, validate=True)
146
+ if len(decoded_data) > max_size:
147
+ logger.error(
148
+ f"Media data size ({len(decoded_data)} bytes) exceeds limit ({max_size} bytes)."
149
+ )
150
+ raise ValueError(
151
+ f"Media data size exceeds limit of {max_size // 1024 // 1024}MB"
152
+ )
153
+ return data
154
+ except base64.binascii.Error as e:
155
+ logger.error(f"Invalid Base64 data provided: {e}")
156
+ raise ValueError("Invalid Base64 data")
157
+ except Exception as e:
158
+ logger.error(f"Error validating media data: {e}")
159
+ raise
160
+
161
+ def convert(
162
+ self, messages: List[Dict[str, Any]], model: str
163
+ ) -> tuple[List[Dict[str, Any]], Optional[Dict[str, Any]]]:
164
+ converted_messages = []
165
+ system_instruction_parts = []
166
+
167
+ for idx, msg in enumerate(messages):
168
+ role = msg.get("role", "")
169
+ parts = []
170
+
171
+ if "content" in msg and isinstance(msg["content"], list):
172
+ for content_item in msg["content"]:
173
+ if not isinstance(content_item, dict):
174
+ logger.warning(
175
+ f"Skipping unexpected content item format: {type(content_item)}"
176
+ )
177
+ continue
178
+
179
+ content_type = content_item.get("type")
180
+
181
+ if content_type == "text" and content_item.get("text"):
182
+ parts.append({"text": content_item["text"]})
183
+ elif content_type == "image_url" and content_item.get(
184
+ "image_url", {}
185
+ ).get("url"):
186
+ try:
187
+ parts.append(
188
+ _convert_image(content_item["image_url"]["url"])
189
+ )
190
+ except Exception as e:
191
+ logger.error(
192
+ f"Failed to convert image URL {content_item['image_url']['url']}: {e}"
193
+ )
194
+ parts.append(
195
+ {
196
+ "text": f"[Error processing image: {content_item['image_url']['url']}]"
197
+ }
198
+ )
199
+ elif content_type == "input_audio" and content_item.get(
200
+ "input_audio"
201
+ ):
202
+ audio_info = content_item["input_audio"]
203
+ audio_data = audio_info.get("data")
204
+ audio_format = audio_info.get("format", "").lower()
205
+
206
+ if not audio_data or not audio_format:
207
+ logger.warning(
208
+ "Skipping audio part due to missing data or format."
209
+ )
210
+ continue
211
+
212
+ try:
213
+ validated_data = self._validate_media_data(
214
+ audio_format,
215
+ audio_data,
216
+ SUPPORTED_AUDIO_FORMATS,
217
+ MAX_AUDIO_SIZE_BYTES,
218
+ )
219
+
220
+ # Get MIME type
221
+ mime_type = AUDIO_FORMAT_TO_MIMETYPE.get(audio_format)
222
+ if not mime_type:
223
+ # Should not happen if format validation passed, but double-check
224
+ logger.error(
225
+ f"Could not find MIME type for supported format: {audio_format}"
226
+ )
227
+ raise ValueError(
228
+ f"Internal error: MIME type mapping missing for {audio_format}"
229
+ )
230
+
231
+ parts.append(
232
+ {
233
+ "inline_data": {
234
+ "mimeType": mime_type,
235
+ "data": validated_data, # Use the validated Base64 data
236
+ }
237
+ }
238
+ )
239
+ logger.debug(
240
+ f"Successfully added audio part (format: {audio_format})"
241
+ )
242
+
243
+ except ValueError as e:
244
+ logger.error(
245
+ f"Skipping audio part due to validation error: {e}"
246
+ )
247
+ parts.append({"text": f"[Error processing audio: {e}]"})
248
+ except Exception:
249
+ logger.exception("Unexpected error processing audio part.")
250
+ parts.append(
251
+ {"text": "[Unexpected error processing audio]"}
252
+ )
253
+
254
+ elif content_type == "input_video" and content_item.get(
255
+ "input_video"
256
+ ):
257
+ video_info = content_item["input_video"]
258
+ video_data = video_info.get("data")
259
+ video_format = video_info.get("format", "").lower()
260
+
261
+ if not video_data or not video_format:
262
+ logger.warning(
263
+ "Skipping video part due to missing data or format."
264
+ )
265
+ continue
266
+
267
+ try:
268
+ validated_data = self._validate_media_data(
269
+ video_format,
270
+ video_data,
271
+ SUPPORTED_VIDEO_FORMATS,
272
+ MAX_VIDEO_SIZE_BYTES,
273
+ )
274
+ mime_type = VIDEO_FORMAT_TO_MIMETYPE.get(video_format)
275
+ if not mime_type:
276
+ raise ValueError(
277
+ f"Internal error: MIME type mapping missing for {video_format}"
278
+ )
279
+
280
+ parts.append(
281
+ {
282
+ "inline_data": {
283
+ "mimeType": mime_type,
284
+ "data": validated_data,
285
+ }
286
+ }
287
+ )
288
+ logger.debug(
289
+ f"Successfully added video part (format: {video_format})"
290
+ )
291
+
292
+ except ValueError as e:
293
+ logger.error(
294
+ f"Skipping video part due to validation error: {e}"
295
+ )
296
+ parts.append({"text": f"[Error processing video: {e}]"})
297
+ except Exception:
298
+ logger.exception("Unexpected error processing video part.")
299
+ parts.append(
300
+ {"text": "[Unexpected error processing video]"}
301
+ )
302
+
303
+ else:
304
+ # Log unrecognized but present types
305
+ if content_type:
306
+ logger.warning(
307
+ f"Unsupported content type or missing data in structured content: {content_type}"
308
+ )
309
+
310
+ elif (
311
+ "content" in msg and isinstance(msg["content"], str) and msg["content"]
312
+ ):
313
+ parts.extend(_process_text_with_image(msg["content"], model))
314
+ elif "tool_calls" in msg and isinstance(msg["tool_calls"], list):
315
+ # Keep existing tool call processing
316
+ for tool_call in msg["tool_calls"]:
317
+ function_call = tool_call.get("function", {})
318
+ # Sanitize arguments loading
319
+ arguments_str = function_call.get("arguments", "{}")
320
+ try:
321
+ function_call["args"] = json.loads(arguments_str)
322
+ except json.JSONDecodeError:
323
+ logger.warning(
324
+ f"Failed to decode tool call arguments: {arguments_str}"
325
+ )
326
+ function_call["args"] = {}
327
+ if "arguments" in function_call:
328
+ if "arguments" in function_call:
329
+ del function_call["arguments"]
330
+
331
+ parts.append({"functionCall": function_call})
332
+
333
+ if role not in SUPPORTED_ROLES:
334
+ if role == "tool":
335
+ role = "user"
336
+ else:
337
+ # 如果是最后一条消息,则认为是用户消息
338
+ if idx == len(messages) - 1:
339
+ role = "user"
340
+ else:
341
+ role = "model"
342
+ if parts:
343
+ if role == "system":
344
+ text_only_parts = [p for p in parts if "text" in p]
345
+ if len(text_only_parts) != len(parts):
346
+ logger.warning(
347
+ "Non-text parts found in system message; discarding them."
348
+ )
349
+ if text_only_parts:
350
+ system_instruction_parts.extend(text_only_parts)
351
+
352
+ else:
353
+ converted_messages.append({"role": role, "parts": parts})
354
+
355
+ system_instruction = (
356
+ None
357
+ if not system_instruction_parts
358
+ else {
359
+ "role": "system",
360
+ "parts": system_instruction_parts,
361
+ }
362
+ )
363
+ return converted_messages, system_instruction
app/handler/response_handler.py ADDED
@@ -0,0 +1,439 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import base64
2
+ import json
3
+ import random
4
+ import string
5
+ import time
6
+ import uuid
7
+ from abc import ABC, abstractmethod
8
+ from typing import Any, Dict, List, Optional
9
+
10
+ from app.config.config import settings
11
+ from app.log.logger import get_openai_logger
12
+ from app.utils.helpers import is_image_upload_configured
13
+ from app.utils.uploader import ImageUploaderFactory
14
+
15
+ logger = get_openai_logger()
16
+
17
+
18
+ class ResponseHandler(ABC):
19
+ """响应处理器基类"""
20
+
21
+ @abstractmethod
22
+ def handle_response(
23
+ self, response: Dict[str, Any], model: str, stream: bool = False
24
+ ) -> Dict[str, Any]:
25
+ pass
26
+
27
+
28
+ class GeminiResponseHandler(ResponseHandler):
29
+ """Gemini响应处理器"""
30
+
31
+ def __init__(self):
32
+ self.thinking_first = True
33
+ self.thinking_status = False
34
+
35
+ def handle_response(
36
+ self,
37
+ response: Dict[str, Any],
38
+ model: str,
39
+ stream: bool = False,
40
+ usage_metadata: Optional[Dict[str, Any]] = None,
41
+ ) -> Dict[str, Any]:
42
+ if stream:
43
+ return _handle_gemini_stream_response(response, model, stream)
44
+ return _handle_gemini_normal_response(response, model, stream)
45
+
46
+
47
+ def _handle_openai_stream_response(
48
+ response: Dict[str, Any],
49
+ model: str,
50
+ finish_reason: str,
51
+ usage_metadata: Optional[Dict[str, Any]],
52
+ ) -> Dict[str, Any]:
53
+ choices = []
54
+ candidates = response.get("candidates", [])
55
+
56
+ for candidate in candidates:
57
+ index = candidate.get("index", 0)
58
+ text, reasoning_content, tool_calls, _ = _extract_result(
59
+ {"candidates": [candidate]}, model, stream=True, gemini_format=False
60
+ )
61
+
62
+ if not text and not tool_calls and not reasoning_content:
63
+ delta = {}
64
+ else:
65
+ delta = {
66
+ "content": text,
67
+ "reasoning_content": reasoning_content,
68
+ "role": "assistant",
69
+ }
70
+ if tool_calls:
71
+ delta["tool_calls"] = tool_calls
72
+
73
+ choice = {"index": index, "delta": delta, "finish_reason": finish_reason}
74
+ choices.append(choice)
75
+
76
+ template_chunk = {
77
+ "id": f"chatcmpl-{uuid.uuid4()}",
78
+ "object": "chat.completion.chunk",
79
+ "created": int(time.time()),
80
+ "model": model,
81
+ "choices": choices,
82
+ }
83
+ if usage_metadata:
84
+ template_chunk["usage"] = {
85
+ "prompt_tokens": usage_metadata.get("promptTokenCount", 0),
86
+ "completion_tokens": usage_metadata.get("candidatesTokenCount", 0),
87
+ "total_tokens": usage_metadata.get("totalTokenCount", 0),
88
+ }
89
+ return template_chunk
90
+
91
+
92
+ def _handle_openai_normal_response(
93
+ response: Dict[str, Any],
94
+ model: str,
95
+ finish_reason: str,
96
+ usage_metadata: Optional[Dict[str, Any]],
97
+ ) -> Dict[str, Any]:
98
+ choices = []
99
+ candidates = response.get("candidates", [])
100
+
101
+ for i, candidate in enumerate(candidates):
102
+ text, reasoning_content, tool_calls, _ = _extract_result(
103
+ {"candidates": [candidate]}, model, stream=False, gemini_format=False
104
+ )
105
+ choice = {
106
+ "index": i,
107
+ "message": {
108
+ "role": "assistant",
109
+ "content": text,
110
+ "reasoning_content": reasoning_content,
111
+ "tool_calls": tool_calls,
112
+ },
113
+ "finish_reason": finish_reason,
114
+ }
115
+ choices.append(choice)
116
+
117
+ return {
118
+ "id": f"chatcmpl-{uuid.uuid4()}",
119
+ "object": "chat.completion",
120
+ "created": int(time.time()),
121
+ "model": model,
122
+ "choices": choices,
123
+ "usage": {
124
+ "prompt_tokens": usage_metadata.get("promptTokenCount", 0),
125
+ "completion_tokens": usage_metadata.get("candidatesTokenCount", 0),
126
+ "total_tokens": usage_metadata.get("totalTokenCount", 0),
127
+ },
128
+ }
129
+
130
+
131
+ class OpenAIResponseHandler(ResponseHandler):
132
+ """OpenAI响应处理器"""
133
+
134
+ def __init__(self, config):
135
+ self.config = config
136
+ self.thinking_first = True
137
+ self.thinking_status = False
138
+
139
+ def handle_response(
140
+ self,
141
+ response: Dict[str, Any],
142
+ model: str,
143
+ stream: bool = False,
144
+ finish_reason: str = None,
145
+ usage_metadata: Optional[Dict[str, Any]] = None,
146
+ ) -> Optional[Dict[str, Any]]:
147
+ if stream:
148
+ return _handle_openai_stream_response(
149
+ response, model, finish_reason, usage_metadata
150
+ )
151
+ return _handle_openai_normal_response(
152
+ response, model, finish_reason, usage_metadata
153
+ )
154
+
155
+ def handle_image_chat_response(
156
+ self, image_str: str, model: str, stream=False, finish_reason="stop"
157
+ ):
158
+ if stream:
159
+ return _handle_openai_stream_image_response(image_str, model, finish_reason)
160
+ return _handle_openai_normal_image_response(image_str, model, finish_reason)
161
+
162
+
163
+ def _handle_openai_stream_image_response(
164
+ image_str: str, model: str, finish_reason: str
165
+ ) -> Dict[str, Any]:
166
+ return {
167
+ "id": f"chatcmpl-{uuid.uuid4()}",
168
+ "object": "chat.completion.chunk",
169
+ "created": int(time.time()),
170
+ "model": model,
171
+ "choices": [
172
+ {
173
+ "index": 0,
174
+ "delta": {"content": image_str} if image_str else {},
175
+ "finish_reason": finish_reason,
176
+ }
177
+ ],
178
+ }
179
+
180
+
181
+ def _handle_openai_normal_image_response(
182
+ image_str: str, model: str, finish_reason: str
183
+ ) -> Dict[str, Any]:
184
+ return {
185
+ "id": f"chatcmpl-{uuid.uuid4()}",
186
+ "object": "chat.completion",
187
+ "created": int(time.time()),
188
+ "model": model,
189
+ "choices": [
190
+ {
191
+ "index": 0,
192
+ "message": {"role": "assistant", "content": image_str},
193
+ "finish_reason": finish_reason,
194
+ }
195
+ ],
196
+ "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
197
+ }
198
+
199
+
200
+ def _extract_result(
201
+ response: Dict[str, Any],
202
+ model: str,
203
+ stream: bool = False,
204
+ gemini_format: bool = False,
205
+ ) -> tuple[str, Optional[str], List[Dict[str, Any]], Optional[bool]]:
206
+ text, reasoning_content, tool_calls, thought = "", "", [], None
207
+
208
+ if stream:
209
+ if response.get("candidates"):
210
+ candidate = response["candidates"][0]
211
+ content = candidate.get("content", {})
212
+ parts = content.get("parts", [])
213
+ if not parts:
214
+ logger.warning("No parts found in stream response")
215
+ return "", None, [], None
216
+
217
+ if "text" in parts[0]:
218
+ text = parts[0].get("text")
219
+ if "thought" in parts[0]:
220
+ if not gemini_format and settings.SHOW_THINKING_PROCESS:
221
+ reasoning_content = text
222
+ text = ""
223
+ thought = parts[0].get("thought")
224
+ elif "executableCode" in parts[0]:
225
+ text = _format_code_block(parts[0]["executableCode"])
226
+ elif "codeExecution" in parts[0]:
227
+ text = _format_code_block(parts[0]["codeExecution"])
228
+ elif "executableCodeResult" in parts[0]:
229
+ text = _format_execution_result(parts[0]["executableCodeResult"])
230
+ elif "codeExecutionResult" in parts[0]:
231
+ text = _format_execution_result(parts[0]["codeExecutionResult"])
232
+ elif "inlineData" in parts[0]:
233
+ text = _extract_image_data(parts[0])
234
+ else:
235
+ text = ""
236
+ text = _add_search_link_text(model, candidate, text)
237
+ tool_calls = _extract_tool_calls(parts, gemini_format)
238
+ else:
239
+ if response.get("candidates"):
240
+ candidate = response["candidates"][0]
241
+ text, reasoning_content = "", ""
242
+
243
+ # 使用安全的访问方式
244
+ content = candidate.get("content", {})
245
+
246
+ if content and isinstance(content, dict):
247
+ parts = content.get("parts", [])
248
+
249
+ if parts:
250
+ for part in parts:
251
+ if "text" in part:
252
+ if "thought" in part and settings.SHOW_THINKING_PROCESS:
253
+ reasoning_content += part["text"]
254
+ else:
255
+ text += part["text"]
256
+ if "thought" in part and thought is None:
257
+ thought = part.get("thought")
258
+ elif "inlineData" in part:
259
+ text += _extract_image_data(part)
260
+ else:
261
+ logger.warning(f"No parts found in content for model: {model}")
262
+ else:
263
+ logger.error(f"Invalid content structure for model: {model}")
264
+
265
+ text = _add_search_link_text(model, candidate, text)
266
+
267
+ # 安全地获取 parts 用于工具调用提取
268
+ parts = candidate.get("content", {}).get("parts", [])
269
+ tool_calls = _extract_tool_calls(parts, gemini_format)
270
+ else:
271
+ logger.warning(f"No candidates found in response for model: {model}")
272
+ text = "暂无返回"
273
+
274
+ return text, reasoning_content, tool_calls, thought
275
+
276
+
277
+ def _has_inline_image_part(response: Dict[str, Any]) -> bool:
278
+ try:
279
+ for c in response.get("candidates", []):
280
+ for p in c.get("content", {}).get("parts", []):
281
+ if isinstance(p, dict) and ("inlineData" in p):
282
+ return True
283
+ except Exception:
284
+ return False
285
+ return False
286
+
287
+
288
+ def _extract_image_data(part: dict) -> str:
289
+ image_uploader = None
290
+ if settings.UPLOAD_PROVIDER == "smms":
291
+ image_uploader = ImageUploaderFactory.create(
292
+ provider=settings.UPLOAD_PROVIDER, api_key=settings.SMMS_SECRET_TOKEN
293
+ )
294
+ elif settings.UPLOAD_PROVIDER == "picgo":
295
+ image_uploader = ImageUploaderFactory.create(
296
+ provider=settings.UPLOAD_PROVIDER,
297
+ api_key=settings.PICGO_API_KEY,
298
+ api_url=settings.PICGO_API_URL
299
+ )
300
+ elif settings.UPLOAD_PROVIDER == "cloudflare_imgbed":
301
+ image_uploader = ImageUploaderFactory.create(
302
+ provider=settings.UPLOAD_PROVIDER,
303
+ base_url=settings.CLOUDFLARE_IMGBED_URL,
304
+ auth_code=settings.CLOUDFLARE_IMGBED_AUTH_CODE,
305
+ upload_folder=settings.CLOUDFLARE_IMGBED_UPLOAD_FOLDER,
306
+ )
307
+ current_date = time.strftime("%Y/%m/%d")
308
+ filename = f"{current_date}/{uuid.uuid4().hex[:8]}.png"
309
+ base64_data = part["inlineData"]["data"]
310
+ mime_type = part["inlineData"]["mimeType"]
311
+ # 将base64_data转成bytes数组
312
+ # Return empty string if no uploader is configured
313
+ if not is_image_upload_configured(settings):
314
+ return f"\n\n![image](data:{mime_type};base64,{base64_data})\n\n"
315
+ bytes_data = base64.b64decode(base64_data)
316
+ upload_response = image_uploader.upload(bytes_data, filename)
317
+ if upload_response.success:
318
+ text = f"\n\n![image]({upload_response.data.url})\n\n"
319
+ else:
320
+ text = f"\n\n![image](data:{mime_type};base64,{base64_data})\n\n"
321
+ return text
322
+
323
+
324
+ def _extract_tool_calls(
325
+ parts: List[Dict[str, Any]], gemini_format: bool
326
+ ) -> List[Dict[str, Any]]:
327
+ """提取工具调用信息"""
328
+ if not parts or not isinstance(parts, list):
329
+ return []
330
+
331
+ letters = string.ascii_lowercase + string.digits
332
+ tool_calls = list()
333
+
334
+ for i in range(len(parts)):
335
+ part = parts[i]
336
+ if not part or not isinstance(part, dict):
337
+ continue
338
+
339
+ item = part.get("functionCall", {})
340
+ if not item or not isinstance(item, dict):
341
+ continue
342
+
343
+ if gemini_format:
344
+ tool_calls.append(part)
345
+ else:
346
+ id = f"call_{''.join(random.sample(letters, 32))}"
347
+ name = item.get("name", "")
348
+ arguments = json.dumps(item.get("args", None) or {})
349
+
350
+ tool_calls.append(
351
+ {
352
+ "index": i,
353
+ "id": id,
354
+ "type": "function",
355
+ "function": {"name": name, "arguments": arguments},
356
+ }
357
+ )
358
+
359
+ return tool_calls
360
+
361
+
362
+ def _handle_gemini_stream_response(
363
+ response: Dict[str, Any], model: str, stream: bool
364
+ ) -> Dict[str, Any]:
365
+ # Early return raw Gemini response if no uploader configured and contains inline images
366
+ if not is_image_upload_configured(settings) and _has_inline_image_part(response):
367
+ return response
368
+
369
+ text, reasoning_content, tool_calls, thought = _extract_result(
370
+ response, model, stream=stream, gemini_format=True
371
+ )
372
+ if tool_calls:
373
+ content = {"parts": tool_calls, "role": "model"}
374
+ else:
375
+ part = {"text": text}
376
+ if thought is not None:
377
+ part["thought"] = thought
378
+ content = {"parts": [part], "role": "model"}
379
+ response["candidates"][0]["content"] = content
380
+ return response
381
+
382
+
383
+ def _handle_gemini_normal_response(
384
+ response: Dict[str, Any], model: str, stream: bool
385
+ ) -> Dict[str, Any]:
386
+ # Early return raw Gemini response if no uploader configured and contains inline images
387
+ if not is_image_upload_configured(settings) and _has_inline_image_part(response):
388
+ return response
389
+
390
+ text, reasoning_content, tool_calls, thought = _extract_result(
391
+ response, model, stream=stream, gemini_format=True
392
+ )
393
+ parts = []
394
+ if tool_calls:
395
+ parts = tool_calls
396
+ else:
397
+ if thought is not None:
398
+ parts.append({"text": reasoning_content, "thought": thought})
399
+ part = {"text": text}
400
+ parts.append(part)
401
+ content = {"parts": parts, "role": "model"}
402
+ response["candidates"][0]["content"] = content
403
+ return response
404
+
405
+
406
+ def _format_code_block(code_data: dict) -> str:
407
+ """格式化代码块输出"""
408
+ language = code_data.get("language", "").lower()
409
+ code = code_data.get("code", "").strip()
410
+ return f"""\n\n---\n\n【代码执行】\n```{language}\n{code}\n```\n"""
411
+
412
+
413
+ def _add_search_link_text(model: str, candidate: dict, text: str) -> str:
414
+ if (
415
+ settings.SHOW_SEARCH_LINK
416
+ and model.endswith("-search")
417
+ and "groundingMetadata" in candidate
418
+ and "groundingChunks" in candidate["groundingMetadata"]
419
+ ):
420
+ grounding_chunks = candidate["groundingMetadata"]["groundingChunks"]
421
+ text += "\n\n---\n\n"
422
+ text += "**【引用来源】**\n\n"
423
+ for _, grounding_chunk in enumerate(grounding_chunks, 1):
424
+ if "web" in grounding_chunk:
425
+ text += _create_search_link(grounding_chunk["web"])
426
+ return text
427
+ else:
428
+ return text
429
+
430
+
431
+ def _create_search_link(grounding_chunk: dict) -> str:
432
+ return f'\n- [{grounding_chunk["title"]}]({grounding_chunk["uri"]})'
433
+
434
+
435
+ def _format_execution_result(result_data: dict) -> str:
436
+ """格式化执行结果输出"""
437
+ outcome = result_data.get("outcome", "")
438
+ output = result_data.get("output", "").strip()
439
+ return f"""\n【执行结果】\n> outcome: {outcome}\n\n【输出结果】\n```plaintext\n{output}\n```\n\n---\n\n"""
app/handler/retry_handler.py ADDED
@@ -0,0 +1,51 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ from functools import wraps
3
+ from typing import Callable, TypeVar
4
+
5
+ from app.config.config import settings
6
+ from app.log.logger import get_retry_logger
7
+ from app.utils.helpers import redact_key_for_logging
8
+
9
+ T = TypeVar("T")
10
+ logger = get_retry_logger()
11
+
12
+
13
+ class RetryHandler:
14
+ """重试处理装饰器"""
15
+
16
+ def __init__(self, key_arg: str = "api_key"):
17
+ self.key_arg = key_arg
18
+
19
+ def __call__(self, func: Callable[..., T]) -> Callable[..., T]:
20
+ @wraps(func)
21
+ async def wrapper(*args, **kwargs) -> T:
22
+ last_exception = None
23
+
24
+ for attempt in range(settings.MAX_RETRIES):
25
+ retries = attempt + 1
26
+ try:
27
+ return await func(*args, **kwargs)
28
+ except Exception as e:
29
+ last_exception = e
30
+ logger.warning(
31
+ f"API call failed with error: {str(e)}. Attempt {retries} of {settings.MAX_RETRIES}"
32
+ )
33
+
34
+ # 从函数参数中获取 key_manager
35
+ key_manager = kwargs.get("key_manager")
36
+ if key_manager:
37
+ old_key = kwargs.get(self.key_arg)
38
+ new_key = await key_manager.handle_api_failure(old_key, retries)
39
+ if new_key:
40
+ kwargs[self.key_arg] = new_key
41
+ logger.info(f"Switched to new API key: {redact_key_for_logging(new_key)}")
42
+ else:
43
+ logger.error(f"No valid API key available after {retries} retries.")
44
+ break
45
+
46
+ logger.error(
47
+ f"All retry attempts failed, raising final exception: {str(last_exception)}"
48
+ )
49
+ raise last_exception
50
+
51
+ return wrapper
app/handler/stream_optimizer.py ADDED
@@ -0,0 +1,143 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ import asyncio
3
+ import math
4
+ from typing import Any, AsyncGenerator, Callable, List
5
+
6
+ from app.config.config import settings
7
+ from app.core.constants import (
8
+ DEFAULT_STREAM_CHUNK_SIZE,
9
+ DEFAULT_STREAM_LONG_TEXT_THRESHOLD,
10
+ DEFAULT_STREAM_MAX_DELAY,
11
+ DEFAULT_STREAM_MIN_DELAY,
12
+ DEFAULT_STREAM_SHORT_TEXT_THRESHOLD,
13
+ )
14
+ from app.log.logger import get_gemini_logger, get_openai_logger
15
+
16
+ logger_openai = get_openai_logger()
17
+ logger_gemini = get_gemini_logger()
18
+
19
+
20
+ class StreamOptimizer:
21
+ """流式输出优化器
22
+
23
+ 提供流式输出优化功能,包括智能延迟调整和长文本分块输出。
24
+ """
25
+
26
+ def __init__(
27
+ self,
28
+ logger=None,
29
+ min_delay: float = DEFAULT_STREAM_MIN_DELAY,
30
+ max_delay: float = DEFAULT_STREAM_MAX_DELAY,
31
+ short_text_threshold: int = DEFAULT_STREAM_SHORT_TEXT_THRESHOLD,
32
+ long_text_threshold: int = DEFAULT_STREAM_LONG_TEXT_THRESHOLD,
33
+ chunk_size: int = DEFAULT_STREAM_CHUNK_SIZE,
34
+ ):
35
+ """初始化流式输出优化器
36
+
37
+ 参数:
38
+ logger: 日志记录器
39
+ min_delay: 最小延迟时间(秒)
40
+ max_delay: 最大延迟时间(秒)
41
+ short_text_threshold: 短文本阈值(字符数)
42
+ long_text_threshold: 长文本阈值(字符数)
43
+ chunk_size: 长文本分块大小(字符数)
44
+ """
45
+ self.logger = logger
46
+ self.min_delay = min_delay
47
+ self.max_delay = max_delay
48
+ self.short_text_threshold = short_text_threshold
49
+ self.long_text_threshold = long_text_threshold
50
+ self.chunk_size = chunk_size
51
+
52
+ def calculate_delay(self, text_length: int) -> float:
53
+ """根据文本长度计算延迟时间
54
+
55
+ 参数:
56
+ text_length: 文本长度
57
+
58
+ 返回:
59
+ 延迟时间(秒)
60
+ """
61
+ if text_length <= self.short_text_threshold:
62
+ # 短文本使用较大延迟
63
+ return self.max_delay
64
+ elif text_length >= self.long_text_threshold:
65
+ # 长文本使用较小延迟
66
+ return self.min_delay
67
+ else:
68
+ # 中等长度文本使用线性插值计算延迟
69
+ # 使用对数函数使延迟变化更平滑
70
+ ratio = math.log(text_length / self.short_text_threshold) / math.log(
71
+ self.long_text_threshold / self.short_text_threshold
72
+ )
73
+ return self.max_delay - ratio * (self.max_delay - self.min_delay)
74
+
75
+ def split_text_into_chunks(self, text: str) -> List[str]:
76
+ """将文本分割成小块
77
+
78
+ 参数:
79
+ text: 要分割的文本
80
+
81
+ 返回:
82
+ 文本块列表
83
+ """
84
+ return [
85
+ text[i : i + self.chunk_size] for i in range(0, len(text), self.chunk_size)
86
+ ]
87
+
88
+ async def optimize_stream_output(
89
+ self,
90
+ text: str,
91
+ create_response_chunk: Callable[[str], Any],
92
+ format_chunk: Callable[[Any], str],
93
+ ) -> AsyncGenerator[str, None]:
94
+ """优化流式输出
95
+
96
+ 参数:
97
+ text: 要输出的文本
98
+ create_response_chunk: 创建响应块的函数,接收文本,返回响应块
99
+ format_chunk: 格式化响应块的函数,接收响应块,返回格式化后的字符串
100
+
101
+ 返回:
102
+ 异步生成器,生成格式化后的响应块
103
+ """
104
+ if not text:
105
+ return
106
+
107
+ # 计算智能延迟时间
108
+ delay = self.calculate_delay(len(text))
109
+
110
+ # 根据文本长度决定输出方式
111
+ if len(text) >= self.long_text_threshold:
112
+ # 长文本:分块输出
113
+ chunks = self.split_text_into_chunks(text)
114
+ for chunk_text in chunks:
115
+ chunk_response = create_response_chunk(chunk_text)
116
+ yield format_chunk(chunk_response)
117
+ await asyncio.sleep(delay)
118
+ else:
119
+ # 短文本:逐字符输出
120
+ for char in text:
121
+ char_chunk = create_response_chunk(char)
122
+ yield format_chunk(char_chunk)
123
+ await asyncio.sleep(delay)
124
+
125
+
126
+ # 创建默认的优化器实例,可以直接导入使用
127
+ openai_optimizer = StreamOptimizer(
128
+ logger=logger_openai,
129
+ min_delay=settings.STREAM_MIN_DELAY,
130
+ max_delay=settings.STREAM_MAX_DELAY,
131
+ short_text_threshold=settings.STREAM_SHORT_TEXT_THRESHOLD,
132
+ long_text_threshold=settings.STREAM_LONG_TEXT_THRESHOLD,
133
+ chunk_size=settings.STREAM_CHUNK_SIZE,
134
+ )
135
+
136
+ gemini_optimizer = StreamOptimizer(
137
+ logger=logger_gemini,
138
+ min_delay=settings.STREAM_MIN_DELAY,
139
+ max_delay=settings.STREAM_MAX_DELAY,
140
+ short_text_threshold=settings.STREAM_SHORT_TEXT_THRESHOLD,
141
+ long_text_threshold=settings.STREAM_LONG_TEXT_THRESHOLD,
142
+ chunk_size=settings.STREAM_CHUNK_SIZE,
143
+ )
app/log/logger.py ADDED
@@ -0,0 +1,349 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import logging
2
+ import platform
3
+ import re
4
+ import sys
5
+ from typing import Dict, Optional
6
+
7
+ # ANSI转义序列颜色代码
8
+ COLORS = {
9
+ "DEBUG": "\033[34m", # 蓝色
10
+ "INFO": "\033[32m", # 绿色
11
+ "WARNING": "\033[33m", # 黄色
12
+ "ERROR": "\033[31m", # 红色
13
+ "CRITICAL": "\033[1;31m", # 红色加粗
14
+ }
15
+
16
+
17
+ # Windows系统启用ANSI支持
18
+ if platform.system() == "Windows":
19
+ import ctypes
20
+
21
+ kernel32 = ctypes.windll.kernel32
22
+ kernel32.SetConsoleMode(kernel32.GetStdHandle(-11), 7)
23
+
24
+
25
+ class ColoredFormatter(logging.Formatter):
26
+ """
27
+ 自定义的日志格式化器,添加颜色支持
28
+ """
29
+
30
+ def format(self, record):
31
+ # 获取对应级别的颜色代码
32
+ color = COLORS.get(record.levelname, "")
33
+ # 添加颜色代码和重置代码
34
+ record.levelname = f"{color}{record.levelname}\033[0m"
35
+ # 创建包含文件名和行号的固定宽度字符串
36
+ record.fileloc = f"[{record.filename}:{record.lineno}]"
37
+ return super().format(record)
38
+
39
+
40
+ class AccessLogFormatter(logging.Formatter):
41
+ """
42
+ Custom access log formatter that redacts API keys in URLs
43
+ """
44
+
45
+ # API key patterns to match in URLs
46
+ API_KEY_PATTERNS = [
47
+ r"\bAIza[0-9A-Za-z_-]{35}", # Google API keys (like Gemini)
48
+ r"\bsk-[0-9A-Za-z_-]{20,}", # OpenAI and general sk- prefixed keys
49
+ ]
50
+
51
+ def __init__(self, *args, **kwargs):
52
+ super().__init__(*args, **kwargs)
53
+ # Compile regex patterns for better performance
54
+ self.compiled_patterns = [
55
+ re.compile(pattern) for pattern in self.API_KEY_PATTERNS
56
+ ]
57
+
58
+ def format(self, record):
59
+ # Format the record normally first
60
+ formatted_msg = super().format(record)
61
+
62
+ # Redact API keys in the formatted message
63
+ return self._redact_api_keys_in_message(formatted_msg)
64
+
65
+ def _redact_api_keys_in_message(self, message: str) -> str:
66
+ """
67
+ Replace API keys in log message with redacted versions
68
+ """
69
+ try:
70
+ for pattern in self.compiled_patterns:
71
+
72
+ def replace_key(match):
73
+ key = match.group(0)
74
+ return redact_key_for_logging(key)
75
+
76
+ message = pattern.sub(replace_key, message)
77
+
78
+ return message
79
+ except Exception as e:
80
+ # Log the error but don't expose the original message in case it contains keys
81
+ import logging
82
+
83
+ logger = logging.getLogger(__name__)
84
+ logger.error(f"Error redacting API keys in access log: {e}")
85
+ return "[LOG_REDACTION_ERROR]"
86
+
87
+
88
+ def redact_key_for_logging(key: str) -> str:
89
+ """
90
+ Redacts API key for secure logging by showing only first and last 6 characters.
91
+
92
+ Args:
93
+ key: API key to redact
94
+
95
+ Returns:
96
+ str: Redacted key in format "first6...last6" or descriptive placeholder for edge cases
97
+ """
98
+ if not key:
99
+ return key
100
+
101
+ if len(key) <= 12:
102
+ return f"{key[:3]}...{key[-3:]}"
103
+ else:
104
+ return f"{key[:6]}...{key[-6:]}"
105
+
106
+
107
+ # 日志格式 - 使用 fileloc 并设置固定宽度 (例如 30)
108
+ FORMATTER = ColoredFormatter(
109
+ "%(asctime)s | %(levelname)-17s | %(fileloc)-30s | %(message)s"
110
+ )
111
+
112
+ # 日志级别映射
113
+ LOG_LEVELS = {
114
+ "debug": logging.DEBUG,
115
+ "info": logging.INFO,
116
+ "warning": logging.WARNING,
117
+ "error": logging.ERROR,
118
+ "critical": logging.CRITICAL,
119
+ }
120
+
121
+
122
+ class Logger:
123
+ def __init__(self):
124
+ pass
125
+
126
+ _loggers: Dict[str, logging.Logger] = {}
127
+
128
+ @staticmethod
129
+ def setup_logger(name: str) -> logging.Logger:
130
+ """
131
+ 设置并获取logger
132
+ :param name: logger名称
133
+ :return: logger实例
134
+ """
135
+ # 导入 settings 对象
136
+ from app.config.config import settings
137
+
138
+ # 从全局配置获取日志级别
139
+ log_level_str = settings.LOG_LEVEL.lower()
140
+ level = LOG_LEVELS.get(log_level_str, logging.INFO)
141
+
142
+ if name in Logger._loggers:
143
+ # 如果 logger 已存在,检查并更新其级别(如果需要)
144
+ existing_logger = Logger._loggers[name]
145
+ if existing_logger.level != level:
146
+ existing_logger.setLevel(level)
147
+ return existing_logger
148
+
149
+ logger = logging.getLogger(name)
150
+ logger.setLevel(level)
151
+ logger.propagate = False
152
+
153
+ # 添加控制台输出
154
+ console_handler = logging.StreamHandler(sys.stdout)
155
+ console_handler.setFormatter(FORMATTER)
156
+ logger.addHandler(console_handler)
157
+
158
+ Logger._loggers[name] = logger
159
+ return logger
160
+
161
+ @staticmethod
162
+ def get_logger(name: str) -> Optional[logging.Logger]:
163
+ """
164
+ 获取已存在的logger
165
+ :param name: logger名称
166
+ :return: logger实例或None
167
+ """
168
+ return Logger._loggers.get(name)
169
+
170
+ @staticmethod
171
+ def update_log_levels(log_level: str):
172
+ """
173
+ 根据当前��全局配置更新所有已创建 logger 的日志级别。
174
+ """
175
+ log_level_str = log_level.lower()
176
+ new_level = LOG_LEVELS.get(log_level_str, logging.INFO)
177
+
178
+ updated_count = 0
179
+ for logger_name, logger_instance in Logger._loggers.items():
180
+ if logger_instance.level != new_level:
181
+ logger_instance.setLevel(new_level)
182
+ # 可选:记录级别变更日志,但注意避免在日志模块内部产生过多日志
183
+ # print(f"Updated log level for logger '{logger_name}' to {log_level_str.upper()}")
184
+ updated_count += 1
185
+
186
+
187
+ # 预定义的loggers
188
+ def get_openai_logger():
189
+ return Logger.setup_logger("openai")
190
+
191
+
192
+ def get_gemini_logger():
193
+ return Logger.setup_logger("gemini")
194
+
195
+
196
+ def get_chat_logger():
197
+ return Logger.setup_logger("chat")
198
+
199
+
200
+ def get_model_logger():
201
+ return Logger.setup_logger("model")
202
+
203
+
204
+ def get_security_logger():
205
+ return Logger.setup_logger("security")
206
+
207
+
208
+ def get_key_manager_logger():
209
+ return Logger.setup_logger("key_manager")
210
+
211
+
212
+ def get_main_logger():
213
+ return Logger.setup_logger("main")
214
+
215
+
216
+ def get_embeddings_logger():
217
+ return Logger.setup_logger("embeddings")
218
+
219
+
220
+ def get_request_logger():
221
+ return Logger.setup_logger("request")
222
+
223
+
224
+ def get_retry_logger():
225
+ return Logger.setup_logger("retry")
226
+
227
+
228
+ def get_image_create_logger():
229
+ return Logger.setup_logger("image_create")
230
+
231
+
232
+ def get_exceptions_logger():
233
+ return Logger.setup_logger("exceptions")
234
+
235
+
236
+ def get_application_logger():
237
+ return Logger.setup_logger("application")
238
+
239
+
240
+ def get_initialization_logger():
241
+ return Logger.setup_logger("initialization")
242
+
243
+
244
+ def get_middleware_logger():
245
+ return Logger.setup_logger("middleware")
246
+
247
+
248
+ def get_routes_logger():
249
+ return Logger.setup_logger("routes")
250
+
251
+
252
+ def get_config_routes_logger():
253
+ return Logger.setup_logger("config_routes")
254
+
255
+
256
+ def get_config_logger():
257
+ return Logger.setup_logger("config")
258
+
259
+
260
+ def get_database_logger():
261
+ return Logger.setup_logger("database")
262
+
263
+
264
+ def get_log_routes_logger():
265
+ return Logger.setup_logger("log_routes")
266
+
267
+
268
+ def get_stats_logger():
269
+ return Logger.setup_logger("stats")
270
+
271
+
272
+ def get_update_logger():
273
+ return Logger.setup_logger("update_service")
274
+
275
+
276
+ def get_scheduler_routes():
277
+ return Logger.setup_logger("scheduler_routes")
278
+
279
+
280
+ def get_message_converter_logger():
281
+ return Logger.setup_logger("message_converter")
282
+
283
+
284
+ def get_api_client_logger():
285
+ return Logger.setup_logger("api_client")
286
+
287
+
288
+ def get_openai_compatible_logger():
289
+ return Logger.setup_logger("openai_compatible")
290
+
291
+
292
+ def get_error_log_logger():
293
+ return Logger.setup_logger("error_log")
294
+
295
+
296
+ def get_request_log_logger():
297
+ return Logger.setup_logger("request_log")
298
+
299
+
300
+ def get_files_logger():
301
+ return Logger.setup_logger("files")
302
+
303
+
304
+ def get_vertex_express_logger():
305
+ return Logger.setup_logger("vertex_express")
306
+
307
+
308
+ def get_gemini_embedding_logger():
309
+ return Logger.setup_logger("gemini_embedding")
310
+
311
+
312
+ def setup_access_logging():
313
+ """
314
+ Configure uvicorn access logging with API key redaction
315
+
316
+ This function sets up a custom access log formatter that automatically
317
+ redacts API keys in HTTP access logs. It works by:
318
+
319
+ 1. Intercepting uvicorn's access log messages
320
+ 2. Using regex patterns to find API keys in URLs
321
+ 3. Replacing them with redacted versions (first6...last6)
322
+
323
+ Supported API key formats:
324
+ - Google/Gemini API keys: AIza[35 chars]
325
+ - OpenAI API keys: sk-[48 chars]
326
+ - General sk- prefixed keys: sk-[20+ chars]
327
+
328
+ Usage:
329
+ - Automatically called in main.py when running with uvicorn
330
+ - For production deployment with gunicorn, ensure this is called in startup
331
+ """
332
+ # Get the uvicorn access logger
333
+ access_logger = logging.getLogger("uvicorn.access")
334
+
335
+ # Remove existing handlers to avoid duplicate logs
336
+ for handler in access_logger.handlers[:]:
337
+ access_logger.removeHandler(handler)
338
+
339
+ # Create new handler with our custom formatter that includes timestamp and log level
340
+ handler = logging.StreamHandler(sys.stdout)
341
+ access_formatter = AccessLogFormatter("%(asctime)s | %(levelname)-8s | %(message)s")
342
+ handler.setFormatter(access_formatter)
343
+
344
+ # Add the handler to uvicorn access logger
345
+ access_logger.addHandler(handler)
346
+ access_logger.setLevel(logging.INFO)
347
+ access_logger.propagate = False
348
+
349
+ return access_logger
app/main.py ADDED
@@ -0,0 +1,15 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import uvicorn
2
+ from dotenv import load_dotenv
3
+
4
+ # 在导入应用程序配置之前加载 .env 文件到环境变量
5
+ load_dotenv()
6
+
7
+ from app.core.application import create_app
8
+ from app.log.logger import get_main_logger
9
+
10
+ app = create_app()
11
+
12
+ if __name__ == "__main__":
13
+ logger = get_main_logger()
14
+ logger.info("Starting application server...")
15
+ uvicorn.run(app, host="0.0.0.0", port=8001)
app/middleware/middleware.py ADDED
@@ -0,0 +1,81 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 中间件配置模块,负责设置和配置应用程序的中间件
3
+ """
4
+
5
+ from fastapi import FastAPI, Request
6
+ from fastapi.middleware.cors import CORSMiddleware
7
+ from fastapi.responses import RedirectResponse
8
+ from starlette.middleware.base import BaseHTTPMiddleware
9
+
10
+ # from app.middleware.request_logging_middleware import RequestLoggingMiddleware
11
+ from app.middleware.smart_routing_middleware import SmartRoutingMiddleware
12
+ from app.core.constants import API_VERSION
13
+ from app.core.security import verify_auth_token
14
+ from app.log.logger import get_middleware_logger
15
+
16
+ logger = get_middleware_logger()
17
+
18
+
19
+ class AuthMiddleware(BaseHTTPMiddleware):
20
+ """
21
+ 认证中间件,处理未经身份验证的请求
22
+ """
23
+
24
+ async def dispatch(self, request: Request, call_next):
25
+ # 允许特定路径绕过身份验证
26
+ if (
27
+ request.url.path not in ["/", "/auth"]
28
+ and not request.url.path.startswith("/static")
29
+ and not request.url.path.startswith("/gemini")
30
+ and not request.url.path.startswith("/v1")
31
+ and not request.url.path.startswith(f"/{API_VERSION}")
32
+ and not request.url.path.startswith("/health")
33
+ and not request.url.path.startswith("/hf")
34
+ and not request.url.path.startswith("/openai")
35
+ and not request.url.path.startswith("/api/version/check")
36
+ and not request.url.path.startswith("/vertex-express")
37
+ and not request.url.path.startswith("/upload")
38
+ ):
39
+
40
+ auth_token = request.cookies.get("auth_token")
41
+ if not auth_token or not verify_auth_token(auth_token):
42
+ logger.warning(f"Unauthorized access attempt to {request.url.path}")
43
+ return RedirectResponse(url="/")
44
+ logger.debug("Request authenticated successfully")
45
+
46
+ response = await call_next(request)
47
+ return response
48
+
49
+
50
+ def setup_middlewares(app: FastAPI) -> None:
51
+ """
52
+ 设置应用程序的中间件
53
+
54
+ Args:
55
+ app: FastAPI应用程序实例
56
+ """
57
+ # 添加智能路由中间件(必须在认证中间件之前)
58
+ app.add_middleware(SmartRoutingMiddleware)
59
+
60
+ # 添加认证中间件
61
+ app.add_middleware(AuthMiddleware)
62
+
63
+ # 添加请求日志中间件(可选,默认注释掉)
64
+ # app.add_middleware(RequestLoggingMiddleware)
65
+
66
+ # 配置CORS中间件
67
+ app.add_middleware(
68
+ CORSMiddleware,
69
+ allow_origins=["*"],
70
+ allow_credentials=True,
71
+ allow_methods=[
72
+ "GET",
73
+ "POST",
74
+ "PUT",
75
+ "DELETE",
76
+ "OPTIONS",
77
+ ],
78
+ allow_headers=["*"],
79
+ expose_headers=["*"],
80
+ max_age=600,
81
+ )
app/middleware/request_logging_middleware.py ADDED
@@ -0,0 +1,40 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import json
2
+
3
+ from fastapi import Request
4
+ from starlette.middleware.base import BaseHTTPMiddleware
5
+
6
+ from app.log.logger import get_request_logger
7
+
8
+ logger = get_request_logger()
9
+
10
+
11
+ # 添加中间件类
12
+ class RequestLoggingMiddleware(BaseHTTPMiddleware):
13
+ async def dispatch(self, request: Request, call_next):
14
+ # 记录请求路径
15
+ logger.info(f"Request path: {request.url.path}")
16
+
17
+ # 获取并记录请求体
18
+ try:
19
+ body = await request.body()
20
+ if body:
21
+ body_str = body.decode()
22
+ # 尝试格式化JSON
23
+ try:
24
+ formatted_body = json.loads(body_str)
25
+ logger.info(
26
+ f"Formatted request body:\n{json.dumps(formatted_body, indent=2, ensure_ascii=False)}"
27
+ )
28
+ except json.JSONDecodeError:
29
+ logger.error("Request body is not valid JSON.")
30
+ except Exception as e:
31
+ logger.error(f"Error reading request body: {str(e)}")
32
+
33
+ # 重置请求的接收器,以便后续处理器可以继续读取请求体
34
+ async def receive():
35
+ return {"type": "http.request", "body": body, "more_body": False}
36
+
37
+ request._receive = receive
38
+
39
+ response = await call_next(request)
40
+ return response
app/middleware/smart_routing_middleware.py ADDED
@@ -0,0 +1,210 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import Request
2
+ from starlette.middleware.base import BaseHTTPMiddleware
3
+ from app.config.config import settings
4
+ from app.log.logger import get_main_logger
5
+ import re
6
+
7
+ logger = get_main_logger()
8
+
9
+ class SmartRoutingMiddleware(BaseHTTPMiddleware):
10
+ def __init__(self, app):
11
+ super().__init__(app)
12
+ # 简化的路由规则 - 直接根据检测结果路由
13
+ pass
14
+
15
+ async def dispatch(self, request: Request, call_next):
16
+ if not settings.URL_NORMALIZATION_ENABLED:
17
+ return await call_next(request)
18
+ logger.debug(f"request: {request}")
19
+ original_path = str(request.url.path)
20
+ method = request.method
21
+
22
+ # 尝试修复URL
23
+ fixed_path, fix_info = self.fix_request_url(original_path, method, request)
24
+
25
+ if fixed_path != original_path:
26
+ logger.info(f"URL fixed: {method} {original_path} → {fixed_path}")
27
+ if fix_info:
28
+ logger.debug(f"Fix details: {fix_info}")
29
+
30
+ # 重写请求路径
31
+ request.scope["path"] = fixed_path
32
+ request.scope["raw_path"] = fixed_path.encode()
33
+
34
+ return await call_next(request)
35
+
36
+ def fix_request_url(self, path: str, method: str, request: Request) -> tuple:
37
+ """简化的URL修复逻辑"""
38
+
39
+ # 首先检查是否已经是正确的格式,如果是则不处理
40
+ if self.is_already_correct_format(path):
41
+ return path, None
42
+
43
+ # 1. 最高优先级:包含generateContent → Gemini格式
44
+ if "generatecontent" in path.lower() or "v1beta/models" in path.lower():
45
+ return self.fix_gemini_by_operation(path, method, request)
46
+
47
+ # 2. 第二优先级:包含/openai/ → OpenAI格式
48
+ if "/openai/" in path.lower():
49
+ return self.fix_openai_by_operation(path, method)
50
+
51
+ # 3. 第三优先级:包含/v1/ → v1格式
52
+ if "/v1/" in path.lower():
53
+ return self.fix_v1_by_operation(path, method)
54
+
55
+ # 4. 第四优先级:包含/chat/completions → chat功能
56
+ if "/chat/completions" in path.lower():
57
+ return "/v1/chat/completions", {"type": "v1_chat"}
58
+
59
+ # 5. 默认:原样传递
60
+ return path, None
61
+
62
+ def is_already_correct_format(self, path: str) -> bool:
63
+ """检查是否已经是正确的API格式"""
64
+ # 检查是否已经是正确的端点格式
65
+ correct_patterns = [
66
+ r"^/v1beta/models/[^/:]+:(generate|streamGenerate)Content$", # Gemini原生
67
+ r"^/gemini/v1beta/models/[^/:]+:(generate|streamGenerate)Content$", # Gemini带前缀
68
+ r"^/v1beta/models$", # Gemini模型列表
69
+ r"^/gemini/v1beta/models$", # Gemini带前缀的模型列表
70
+ r"^/v1/(chat/completions|models|embeddings|images/generations|audio/speech)$", # v1格式
71
+ r"^/openai/v1/(chat/completions|models|embeddings|images/generations|audio/speech)$", # OpenAI格式
72
+ r"^/hf/v1/(chat/completions|models|embeddings|images/generations|audio/speech)$", # HF格式
73
+ r"^/vertex-express/v1beta/models/[^/:]+:(generate|streamGenerate)Content$", # Vertex Express Gemini格式
74
+ r"^/vertex-express/v1beta/models$", # Vertex Express模型列表
75
+ r"^/vertex-express/v1/(chat/completions|models|embeddings|images/generations)$", # Vertex Express OpenAI格式
76
+ ]
77
+
78
+ for pattern in correct_patterns:
79
+ if re.match(pattern, path):
80
+ return True
81
+
82
+ return False
83
+
84
+ def fix_gemini_by_operation(
85
+ self, path: str, method: str, request: Request
86
+ ) -> tuple:
87
+ """根据Gemini操作修复,考虑端点偏好"""
88
+ if method == "GET":
89
+ return "/v1beta/models", {
90
+ "role": "gemini_models",
91
+ }
92
+
93
+ # 提取模型名称
94
+ try:
95
+ model_name = self.extract_model_name(path, request)
96
+ except ValueError:
97
+ # 无法提取模型名称,返回原路径不做处理
98
+ return path, None
99
+
100
+ # 检测是否为流式请求
101
+ is_stream = self.detect_stream_request(path, request)
102
+
103
+ # 检查是否有vertex-express偏好
104
+ if "/vertex-express/" in path.lower():
105
+ if is_stream:
106
+ target_url = (
107
+ f"/vertex-express/v1beta/models/{model_name}:streamGenerateContent"
108
+ )
109
+ else:
110
+ target_url = (
111
+ f"/vertex-express/v1beta/models/{model_name}:generateContent"
112
+ )
113
+
114
+ fix_info = {
115
+ "rule": (
116
+ "vertex_express_generate"
117
+ if not is_stream
118
+ else "vertex_express_stream"
119
+ ),
120
+ "preference": "vertex_express_format",
121
+ "is_stream": is_stream,
122
+ "model": model_name,
123
+ }
124
+ else:
125
+ # 标准Gemini端点
126
+ if is_stream:
127
+ target_url = f"/v1beta/models/{model_name}:streamGenerateContent"
128
+ else:
129
+ target_url = f"/v1beta/models/{model_name}:generateContent"
130
+
131
+ fix_info = {
132
+ "rule": "gemini_generate" if not is_stream else "gemini_stream",
133
+ "preference": "gemini_format",
134
+ "is_stream": is_stream,
135
+ "model": model_name,
136
+ }
137
+
138
+ return target_url, fix_info
139
+
140
+ def fix_openai_by_operation(self, path: str, method: str) -> tuple:
141
+ """根据操作类型修复OpenAI格式"""
142
+ if method == "POST":
143
+ if "chat" in path.lower() or "completion" in path.lower():
144
+ return "/openai/v1/chat/completions", {"type": "openai_chat"}
145
+ elif "embedding" in path.lower():
146
+ return "/openai/v1/embeddings", {"type": "openai_embeddings"}
147
+ elif "image" in path.lower():
148
+ return "/openai/v1/images/generations", {"type": "openai_images"}
149
+ elif "audio" in path.lower():
150
+ return "/openai/v1/audio/speech", {"type": "openai_audio"}
151
+ elif method == "GET":
152
+ if "model" in path.lower():
153
+ return "/openai/v1/models", {"type": "openai_models"}
154
+
155
+ return path, None
156
+
157
+ def fix_v1_by_operation(self, path: str, method: str) -> tuple:
158
+ """根据操作类型修复v1格式"""
159
+ if method == "POST":
160
+ if "chat" in path.lower() or "completion" in path.lower():
161
+ return "/v1/chat/completions", {"type": "v1_chat"}
162
+ elif "embedding" in path.lower():
163
+ return "/v1/embeddings", {"type": "v1_embeddings"}
164
+ elif "image" in path.lower():
165
+ return "/v1/images/generations", {"type": "v1_images"}
166
+ elif "audio" in path.lower():
167
+ return "/v1/audio/speech", {"type": "v1_audio"}
168
+ elif method == "GET":
169
+ if "model" in path.lower():
170
+ return "/v1/models", {"type": "v1_models"}
171
+
172
+ return path, None
173
+
174
+ def detect_stream_request(self, path: str, request: Request) -> bool:
175
+ """检测是否为流式请求"""
176
+ # 1. 路径中包含stream关键词
177
+ if "stream" in path.lower():
178
+ return True
179
+
180
+ # 2. 查询参数
181
+ if request.query_params.get("stream") == "true":
182
+ return True
183
+
184
+ return False
185
+
186
+ def extract_model_name(self, path: str, request: Request) -> str:
187
+ """从请求中提取模型名称,用于构建Gemini API URL"""
188
+ # 1. 从请求体中提取
189
+ try:
190
+ if hasattr(request, "_body") and request._body:
191
+ import json
192
+
193
+ body = json.loads(request._body.decode())
194
+ if "model" in body and body["model"]:
195
+ return body["model"]
196
+ except Exception:
197
+ pass
198
+
199
+ # 2. 从查询参数中提取
200
+ model_param = request.query_params.get("model")
201
+ if model_param:
202
+ return model_param
203
+
204
+ # 3. 从路径中提取(用于已包含模型名称的路径)
205
+ match = re.search(r"/models/([^/:]+)", path, re.IGNORECASE)
206
+ if match:
207
+ return match.group(1)
208
+
209
+ # 4. 如果无法提取模型名称,抛出异常
210
+ raise ValueError("Unable to extract model name from request")
app/router/config_routes.py ADDED
@@ -0,0 +1,225 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 配置路由模块
3
+ """
4
+
5
+ from typing import Any, Dict, List
6
+
7
+ from fastapi import APIRouter, HTTPException, Request
8
+ from fastapi.responses import RedirectResponse
9
+ from pydantic import BaseModel, Field
10
+
11
+ from app.core.security import verify_auth_token
12
+ from app.log.logger import Logger, get_config_routes_logger
13
+ from app.service.config.config_service import ConfigService
14
+ from app.service.proxy.proxy_check_service import get_proxy_check_service, ProxyCheckResult
15
+ from app.utils.helpers import redact_key_for_logging
16
+
17
+ router = APIRouter(prefix="/api/config", tags=["config"])
18
+
19
+ logger = get_config_routes_logger()
20
+
21
+
22
+ @router.get("", response_model=Dict[str, Any])
23
+ async def get_config(request: Request):
24
+ auth_token = request.cookies.get("auth_token")
25
+ if not auth_token or not verify_auth_token(auth_token):
26
+ logger.warning("Unauthorized access attempt to config page")
27
+ return RedirectResponse(url="/", status_code=302)
28
+ return await ConfigService.get_config()
29
+
30
+
31
+ @router.put("", response_model=Dict[str, Any])
32
+ async def update_config(config_data: Dict[str, Any], request: Request):
33
+ auth_token = request.cookies.get("auth_token")
34
+ if not auth_token or not verify_auth_token(auth_token):
35
+ logger.warning("Unauthorized access attempt to config page")
36
+ return RedirectResponse(url="/", status_code=302)
37
+ try:
38
+ result = await ConfigService.update_config(config_data)
39
+ # 配置更新成功后,立即更新所有 logger 的级别
40
+ Logger.update_log_levels(config_data["LOG_LEVEL"])
41
+ logger.info("Log levels updated after configuration change.")
42
+ return result
43
+ except Exception as e:
44
+ logger.error(f"Error updating config or log levels: {e}", exc_info=True)
45
+ raise HTTPException(status_code=400, detail=str(e))
46
+
47
+
48
+ @router.post("/reset", response_model=Dict[str, Any])
49
+ async def reset_config(request: Request):
50
+ auth_token = request.cookies.get("auth_token")
51
+ if not auth_token or not verify_auth_token(auth_token):
52
+ logger.warning("Unauthorized access attempt to config page")
53
+ return RedirectResponse(url="/", status_code=302)
54
+ try:
55
+ return await ConfigService.reset_config()
56
+ except Exception as e:
57
+ raise HTTPException(status_code=400, detail=str(e))
58
+
59
+
60
+ class DeleteKeysRequest(BaseModel):
61
+ keys: List[str] = Field(..., description="List of API keys to delete")
62
+
63
+
64
+ @router.delete("/keys/{key_to_delete}", response_model=Dict[str, Any])
65
+ async def delete_single_key(key_to_delete: str, request: Request):
66
+ auth_token = request.cookies.get("auth_token")
67
+ if not auth_token or not verify_auth_token(auth_token):
68
+ logger.warning(f"Unauthorized attempt to delete key: {redact_key_for_logging(key_to_delete)}")
69
+ return RedirectResponse(url="/", status_code=302)
70
+ try:
71
+ logger.info(f"Attempting to delete key: {redact_key_for_logging(key_to_delete)}")
72
+ result = await ConfigService.delete_key(key_to_delete)
73
+ if not result.get("success"):
74
+ raise HTTPException(
75
+ status_code=(
76
+ 404 if "not found" in result.get("message", "").lower() else 400
77
+ ),
78
+ detail=result.get("message"),
79
+ )
80
+ return result
81
+ except HTTPException as e:
82
+ raise e
83
+ except Exception as e:
84
+ logger.error(f"Error deleting key '{redact_key_for_logging(key_to_delete)}': {e}", exc_info=True)
85
+ raise HTTPException(status_code=500, detail=f"Error deleting key: {str(e)}")
86
+
87
+
88
+ @router.post("/keys/delete-selected", response_model=Dict[str, Any])
89
+ async def delete_selected_keys_route(
90
+ delete_request: DeleteKeysRequest, request: Request
91
+ ):
92
+ auth_token = request.cookies.get("auth_token")
93
+ if not auth_token or not verify_auth_token(auth_token):
94
+ logger.warning("Unauthorized attempt to bulk delete keys")
95
+ return RedirectResponse(url="/", status_code=302)
96
+
97
+ if not delete_request.keys:
98
+ logger.warning("Attempt to bulk delete keys with an empty list.")
99
+ raise HTTPException(status_code=400, detail="No keys provided for deletion.")
100
+
101
+ try:
102
+ logger.info(f"Attempting to bulk delete {len(delete_request.keys)} keys.")
103
+ result = await ConfigService.delete_selected_keys(delete_request.keys)
104
+ if not result.get("success") and result.get("deleted_count", 0) == 0:
105
+ raise HTTPException(
106
+ status_code=400, detail=result.get("message", "Failed to delete keys.")
107
+ )
108
+ return result
109
+ except HTTPException as e:
110
+ raise e
111
+ except Exception as e:
112
+ logger.error(f"Error bulk deleting keys: {e}", exc_info=True)
113
+ raise HTTPException(
114
+ status_code=500, detail=f"Error bulk deleting keys: {str(e)}"
115
+ )
116
+
117
+
118
+ @router.get("/ui/models")
119
+ async def get_ui_models(request: Request):
120
+ auth_token_cookie = request.cookies.get("auth_token")
121
+ if not auth_token_cookie or not verify_auth_token(auth_token_cookie):
122
+ logger.warning("Unauthorized access attempt to /api/config/ui/models")
123
+ raise HTTPException(status_code=403, detail="Not authenticated")
124
+
125
+ try:
126
+ models = await ConfigService.fetch_ui_models()
127
+ return models
128
+ except HTTPException as e:
129
+ raise e
130
+ except Exception as e:
131
+ logger.error(f"Unexpected error in /ui/models endpoint: {e}", exc_info=True)
132
+ raise HTTPException(
133
+ status_code=500,
134
+ detail=f"An unexpected error occurred while fetching UI models: {str(e)}",
135
+ )
136
+
137
+
138
+ class ProxyCheckRequest(BaseModel):
139
+ """Proxy check request"""
140
+ proxy: str = Field(..., description="Proxy address to check")
141
+ use_cache: bool = Field(True, description="Whether to use cached results")
142
+
143
+
144
+ class ProxyBatchCheckRequest(BaseModel):
145
+ """Batch proxy check request"""
146
+ proxies: List[str] = Field(..., description="List of proxy addresses to check")
147
+ use_cache: bool = Field(True, description="Whether to use cached results")
148
+ max_concurrent: int = Field(5, description="Maximum concurrent check count", ge=1, le=10)
149
+
150
+
151
+ @router.post("/proxy/check", response_model=ProxyCheckResult)
152
+ async def check_single_proxy(proxy_request: ProxyCheckRequest, request: Request):
153
+ """Check if a single proxy is available"""
154
+ auth_token = request.cookies.get("auth_token")
155
+ if not auth_token or not verify_auth_token(auth_token):
156
+ logger.warning("Unauthorized access attempt to proxy check")
157
+ return RedirectResponse(url="/", status_code=302)
158
+
159
+ try:
160
+ logger.info(f"Checking single proxy: {proxy_request.proxy}")
161
+ proxy_service = get_proxy_check_service()
162
+ result = await proxy_service.check_single_proxy(
163
+ proxy_request.proxy,
164
+ proxy_request.use_cache
165
+ )
166
+ return result
167
+ except Exception as e:
168
+ logger.error(f"Proxy check failed: {str(e)}", exc_info=True)
169
+ raise HTTPException(status_code=500, detail=f"Proxy check failed: {str(e)}")
170
+
171
+
172
+ @router.post("/proxy/check-all", response_model=List[ProxyCheckResult])
173
+ async def check_all_proxies(batch_request: ProxyBatchCheckRequest, request: Request):
174
+ """Check multiple proxies availability"""
175
+ auth_token = request.cookies.get("auth_token")
176
+ if not auth_token or not verify_auth_token(auth_token):
177
+ logger.warning("Unauthorized access attempt to batch proxy check")
178
+ return RedirectResponse(url="/", status_code=302)
179
+
180
+ try:
181
+ logger.info(f"Batch checking {len(batch_request.proxies)} proxies")
182
+ proxy_service = get_proxy_check_service()
183
+ results = await proxy_service.check_multiple_proxies(
184
+ batch_request.proxies,
185
+ batch_request.use_cache,
186
+ batch_request.max_concurrent
187
+ )
188
+ return results
189
+ except Exception as e:
190
+ logger.error(f"Batch proxy check failed: {str(e)}", exc_info=True)
191
+ raise HTTPException(status_code=500, detail=f"Batch proxy check failed: {str(e)}")
192
+
193
+
194
+ @router.get("/proxy/cache-stats")
195
+ async def get_proxy_cache_stats(request: Request):
196
+ """Get proxy check cache statistics"""
197
+ auth_token = request.cookies.get("auth_token")
198
+ if not auth_token or not verify_auth_token(auth_token):
199
+ logger.warning("Unauthorized access attempt to proxy cache stats")
200
+ return RedirectResponse(url="/", status_code=302)
201
+
202
+ try:
203
+ proxy_service = get_proxy_check_service()
204
+ stats = proxy_service.get_cache_stats()
205
+ return stats
206
+ except Exception as e:
207
+ logger.error(f"Get proxy cache stats failed: {str(e)}", exc_info=True)
208
+ raise HTTPException(status_code=500, detail=f"Get cache stats failed: {str(e)}")
209
+
210
+
211
+ @router.post("/proxy/clear-cache")
212
+ async def clear_proxy_cache(request: Request):
213
+ """Clear proxy check cache"""
214
+ auth_token = request.cookies.get("auth_token")
215
+ if not auth_token or not verify_auth_token(auth_token):
216
+ logger.warning("Unauthorized access attempt to clear proxy cache")
217
+ return RedirectResponse(url="/", status_code=302)
218
+
219
+ try:
220
+ proxy_service = get_proxy_check_service()
221
+ proxy_service.clear_cache()
222
+ return {"success": True, "message": "Proxy check cache cleared"}
223
+ except Exception as e:
224
+ logger.error(f"Clear proxy cache failed: {str(e)}", exc_info=True)
225
+ raise HTTPException(status_code=500, detail=f"Clear cache failed: {str(e)}")
app/router/error_log_routes.py ADDED
@@ -0,0 +1,271 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 日志路由模块
3
+ """
4
+
5
+ from datetime import datetime
6
+ from typing import Dict, List, Optional
7
+
8
+ from fastapi import (
9
+ APIRouter,
10
+ Body,
11
+ HTTPException,
12
+ Path,
13
+ Query,
14
+ Request,
15
+ Response,
16
+ status,
17
+ )
18
+ from pydantic import BaseModel
19
+
20
+ from app.core.security import verify_auth_token
21
+ from app.log.logger import get_log_routes_logger
22
+ from app.service.error_log import error_log_service
23
+
24
+ router = APIRouter(prefix="/api/logs", tags=["logs"])
25
+
26
+ logger = get_log_routes_logger()
27
+
28
+
29
+ class ErrorLogListItem(BaseModel):
30
+ id: int
31
+ gemini_key: Optional[str] = None
32
+ error_type: Optional[str] = None
33
+ error_code: Optional[int] = None
34
+ model_name: Optional[str] = None
35
+ request_time: Optional[datetime] = None
36
+
37
+
38
+ class ErrorLogListResponse(BaseModel):
39
+ logs: List[ErrorLogListItem]
40
+ total: int
41
+
42
+
43
+ @router.get("/errors", response_model=ErrorLogListResponse)
44
+ async def get_error_logs_api(
45
+ request: Request,
46
+ limit: int = Query(10, ge=1, le=1000),
47
+ offset: int = Query(0, ge=0),
48
+ key_search: Optional[str] = Query(
49
+ None, description="Search term for Gemini key (partial match)"
50
+ ),
51
+ error_search: Optional[str] = Query(
52
+ None, description="Search term for error type or log message"
53
+ ),
54
+ error_code_search: Optional[str] = Query(
55
+ None, description="Search term for error code"
56
+ ),
57
+ start_date: Optional[datetime] = Query(
58
+ None, description="Start datetime for filtering"
59
+ ),
60
+ end_date: Optional[datetime] = Query(
61
+ None, description="End datetime for filtering"
62
+ ),
63
+ sort_by: str = Query(
64
+ "id", description="Field to sort by (e.g., 'id', 'request_time')"
65
+ ),
66
+ sort_order: str = Query("desc", description="Sort order ('asc' or 'desc')"),
67
+ ):
68
+ """
69
+ 获取错误日志列表 (返回错误码),支持过滤和排序
70
+
71
+ Args:
72
+ request: 请求对象
73
+ limit: 限制数量
74
+ offset: 偏移量
75
+ key_search: 密钥搜索
76
+ error_search: 错误搜索 (可能搜索类型或日志内容,由DB层决定)
77
+ error_code_search: 错误码搜索
78
+ start_date: 开始日期
79
+ end_date: 结束日期
80
+ sort_by: 排序字段
81
+ sort_order: 排序顺序
82
+
83
+ Returns:
84
+ ErrorLogListResponse: An object containing the list of logs (with error_code) and the total count.
85
+ """
86
+ auth_token = request.cookies.get("auth_token")
87
+ if not auth_token or not verify_auth_token(auth_token):
88
+ logger.warning("Unauthorized access attempt to error logs list")
89
+ raise HTTPException(status_code=401, detail="Not authenticated")
90
+
91
+ try:
92
+ result = await error_log_service.process_get_error_logs(
93
+ limit=limit,
94
+ offset=offset,
95
+ key_search=key_search,
96
+ error_search=error_search,
97
+ error_code_search=error_code_search,
98
+ start_date=start_date,
99
+ end_date=end_date,
100
+ sort_by=sort_by,
101
+ sort_order=sort_order,
102
+ )
103
+ logs_data = result["logs"]
104
+ total_count = result["total"]
105
+
106
+ validated_logs = [ErrorLogListItem(**log) for log in logs_data]
107
+ return ErrorLogListResponse(logs=validated_logs, total=total_count)
108
+ except Exception as e:
109
+ logger.exception(f"Failed to get error logs list: {str(e)}")
110
+ raise HTTPException(
111
+ status_code=500, detail=f"Failed to get error logs list: {str(e)}"
112
+ )
113
+
114
+
115
+ class ErrorLogDetailResponse(BaseModel):
116
+ id: int
117
+ gemini_key: Optional[str] = None
118
+ error_type: Optional[str] = None
119
+ error_log: Optional[str] = None
120
+ request_msg: Optional[str] = None
121
+ model_name: Optional[str] = None
122
+ request_time: Optional[datetime] = None
123
+ error_code: Optional[int] = None
124
+
125
+
126
+ @router.get("/errors/{log_id}/details", response_model=ErrorLogDetailResponse)
127
+ async def get_error_log_detail_api(request: Request, log_id: int = Path(..., ge=1)):
128
+ """
129
+ 根据日志 ID 获取错误日志的详细信息 (包括 error_log 和 request_msg)
130
+ """
131
+ auth_token = request.cookies.get("auth_token")
132
+ if not auth_token or not verify_auth_token(auth_token):
133
+ logger.warning(
134
+ f"Unauthorized access attempt to error log details for ID: {log_id}"
135
+ )
136
+ raise HTTPException(status_code=401, detail="Not authenticated")
137
+
138
+ try:
139
+ log_details = await error_log_service.process_get_error_log_details(
140
+ log_id=log_id
141
+ )
142
+ if not log_details:
143
+ raise HTTPException(status_code=404, detail="Error log not found")
144
+
145
+ return ErrorLogDetailResponse(**log_details)
146
+ except HTTPException as http_exc:
147
+ raise http_exc
148
+ except Exception as e:
149
+ logger.exception(f"Failed to get error log details for ID {log_id}: {str(e)}")
150
+ raise HTTPException(
151
+ status_code=500, detail=f"Failed to get error log details: {str(e)}"
152
+ )
153
+
154
+
155
+ @router.get("/errors/lookup", response_model=ErrorLogDetailResponse)
156
+ async def lookup_error_log_by_info(
157
+ request: Request,
158
+ gemini_key: str = Query(..., description="完整的 Gemini key"),
159
+ timestamp: datetime = Query(..., description="请求时间 (ISO8601)"),
160
+ status_code: Optional[int] = Query(None, description="错误码 (可选)"),
161
+ window_seconds: int = Query(
162
+ 100, ge=1, le=300, description="时间窗口(秒), 默认100秒"
163
+ ),
164
+ ):
165
+ """
166
+ 通过 key / 错误码 / 时间窗口 查找最匹配的一条错误日志详情。
167
+ """
168
+ auth_token = request.cookies.get("auth_token")
169
+ if not auth_token or not verify_auth_token(auth_token):
170
+ logger.warning("Unauthorized access attempt to lookup error log by info")
171
+ raise HTTPException(status_code=401, detail="Not authenticated")
172
+
173
+ try:
174
+ detail = await error_log_service.process_find_error_log_by_info(
175
+ gemini_key=gemini_key,
176
+ timestamp=timestamp,
177
+ status_code=status_code,
178
+ window_seconds=window_seconds,
179
+ )
180
+ if not detail:
181
+ raise HTTPException(status_code=404, detail="No matching error log found")
182
+ return ErrorLogDetailResponse(**detail)
183
+ except HTTPException as http_exc:
184
+ raise http_exc
185
+ except Exception as e:
186
+ logger.exception(
187
+ f"Failed to lookup error log by info for key=***{gemini_key[-4:] if gemini_key else ''}: {str(e)}"
188
+ )
189
+ raise HTTPException(status_code=500, detail="Internal server error")
190
+
191
+
192
+ @router.delete("/errors", status_code=status.HTTP_204_NO_CONTENT)
193
+ async def delete_error_logs_bulk_api(
194
+ request: Request, payload: Dict[str, List[int]] = Body(...)
195
+ ):
196
+ """
197
+ 批量删除错误日志 (异步)
198
+ """
199
+ auth_token = request.cookies.get("auth_token")
200
+ if not auth_token or not verify_auth_token(auth_token):
201
+ logger.warning("Unauthorized access attempt to bulk delete error logs")
202
+ raise HTTPException(status_code=401, detail="Not authenticated")
203
+
204
+ log_ids = payload.get("ids")
205
+ if not log_ids:
206
+ raise HTTPException(status_code=400, detail="No log IDs provided for deletion.")
207
+
208
+ try:
209
+ deleted_count = await error_log_service.process_delete_error_logs_by_ids(
210
+ log_ids
211
+ )
212
+ # 注意:异步函数返回的是尝试删除的数量,可能不是精确值
213
+ logger.info(
214
+ f"Attempted bulk deletion for {deleted_count} error logs with IDs: {log_ids}"
215
+ )
216
+ return Response(status_code=status.HTTP_204_NO_CONTENT)
217
+ except Exception as e:
218
+ logger.exception(f"Error bulk deleting error logs with IDs {log_ids}: {str(e)}")
219
+ raise HTTPException(
220
+ status_code=500, detail="Internal server error during bulk deletion"
221
+ )
222
+
223
+
224
+ @router.delete("/errors/all", status_code=status.HTTP_204_NO_CONTENT)
225
+ async def delete_all_error_logs_api(request: Request):
226
+ """
227
+ 删除所有错误日志 (异步)
228
+ """
229
+ auth_token = request.cookies.get("auth_token")
230
+ if not auth_token or not verify_auth_token(auth_token):
231
+ logger.warning("Unauthorized access attempt to delete all error logs")
232
+ raise HTTPException(status_code=401, detail="Not authenticated")
233
+
234
+ try:
235
+ await error_log_service.process_delete_all_error_logs()
236
+ logger.info("Successfully deleted all error logs.")
237
+ # No body needed for 204 response
238
+ return Response(status_code=status.HTTP_204_NO_CONTENT)
239
+ except Exception as e:
240
+ logger.exception(f"Error deleting all error logs: {str(e)}")
241
+ raise HTTPException(
242
+ status_code=500, detail="Internal server error during deletion of all logs"
243
+ )
244
+
245
+
246
+ @router.delete("/errors/{log_id}", status_code=status.HTTP_204_NO_CONTENT)
247
+ async def delete_error_log_api(request: Request, log_id: int = Path(..., ge=1)):
248
+ """
249
+ 删除单个错误日志 (异步)
250
+ """
251
+ auth_token = request.cookies.get("auth_token")
252
+ if not auth_token or not verify_auth_token(auth_token):
253
+ logger.warning(f"Unauthorized access attempt to delete error log ID: {log_id}")
254
+ raise HTTPException(status_code=401, detail="Not authenticated")
255
+
256
+ try:
257
+ success = await error_log_service.process_delete_error_log_by_id(log_id)
258
+ if not success:
259
+ # 服务层现在在未找到时返回 False,我们在这里转换为 404
260
+ raise HTTPException(
261
+ status_code=404, detail=f"Error log with ID {log_id} not found"
262
+ )
263
+ logger.info(f"Successfully deleted error log with ID: {log_id}")
264
+ return Response(status_code=status.HTTP_204_NO_CONTENT)
265
+ except HTTPException as http_exc:
266
+ raise http_exc
267
+ except Exception as e:
268
+ logger.exception(f"Error deleting error log with ID {log_id}: {str(e)}")
269
+ raise HTTPException(
270
+ status_code=500, detail="Internal server error during deletion"
271
+ )
app/router/files_routes.py ADDED
@@ -0,0 +1,296 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Files API 路由
3
+ """
4
+ from typing import Optional
5
+ from fastapi import APIRouter, Request, Query, Depends, Header, HTTPException
6
+ from fastapi.responses import JSONResponse
7
+
8
+ from app.config.config import settings
9
+ from app.domain.file_models import (
10
+ FileMetadata,
11
+ ListFilesResponse,
12
+ DeleteFileResponse
13
+ )
14
+ from app.log.logger import get_files_logger
15
+ from app.core.security import SecurityService
16
+ from app.service.files.files_service import get_files_service
17
+ from app.service.files.file_upload_handler import get_upload_handler
18
+ from app.utils.helpers import redact_key_for_logging
19
+
20
+ logger = get_files_logger()
21
+
22
+ router = APIRouter()
23
+ security_service = SecurityService()
24
+
25
+
26
+ @router.post("/upload/v1beta/files")
27
+ async def upload_file_init(
28
+ request: Request,
29
+ auth_token: str = Depends(security_service.verify_key_or_goog_api_key),
30
+ x_goog_upload_protocol: Optional[str] = Header(None),
31
+ x_goog_upload_command: Optional[str] = Header(None),
32
+ x_goog_upload_header_content_length: Optional[str] = Header(None),
33
+ x_goog_upload_header_content_type: Optional[str] = Header(None),
34
+ ):
35
+ """初始化文件上传"""
36
+ logger.debug(f"Upload file request: {request.method=}, {request.url=}, {auth_token=}, {x_goog_upload_protocol=}, {x_goog_upload_command=}, {x_goog_upload_header_content_length=}, {x_goog_upload_header_content_type=}")
37
+
38
+ # 檢查是否是實際的上傳請求(有 upload_id)
39
+ if request.query_params.get("upload_id") and x_goog_upload_command in ["upload", "upload, finalize"]:
40
+ logger.debug("This is an upload request, not initialization. Redirecting to handle_upload.")
41
+ return await handle_upload(
42
+ upload_path="v1beta/files",
43
+ request=request,
44
+ key=request.query_params.get("key"),
45
+ auth_token=auth_token
46
+ )
47
+
48
+ try:
49
+ # 使用认证 token 作为 user_token
50
+ user_token = auth_token
51
+ # 获取请求体
52
+ body = await request.body()
53
+
54
+ # 构建请求主机 URL
55
+ request_host = f"{request.url.scheme}://{request.url.netloc}"
56
+ logger.info(f"Request host: {request_host}")
57
+
58
+ # 准备请求头
59
+ headers = {
60
+ "x-goog-upload-protocol": x_goog_upload_protocol or "resumable",
61
+ "x-goog-upload-command": x_goog_upload_command or "start",
62
+ }
63
+
64
+ if x_goog_upload_header_content_length:
65
+ headers["x-goog-upload-header-content-length"] = x_goog_upload_header_content_length
66
+ if x_goog_upload_header_content_type:
67
+ headers["x-goog-upload-header-content-type"] = x_goog_upload_header_content_type
68
+
69
+ # 调用服务
70
+ files_service = await get_files_service()
71
+ response_data, response_headers = await files_service.initialize_upload(
72
+ headers=headers,
73
+ body=body,
74
+ user_token=user_token,
75
+ request_host=request_host # 傳遞請求主機
76
+ )
77
+
78
+ logger.info(f"Upload initialization response: {response_data}")
79
+ logger.info(f"Upload initialization response headers: {response_headers}")
80
+
81
+ logger.info(f"Upload initialization response headers: {response_data}")
82
+ # 返回响应
83
+ return JSONResponse(
84
+ content=response_data,
85
+ headers=response_headers
86
+ )
87
+
88
+ except HTTPException as e:
89
+ logger.error(f"Upload initialization failed: {e.detail}")
90
+ return JSONResponse(
91
+ content={"error": {"message": e.detail}},
92
+ status_code=e.status_code
93
+ )
94
+ except Exception as e:
95
+ logger.error(f"Unexpected error in upload initialization: {str(e)}")
96
+ return JSONResponse(
97
+ content={"error": {"message": "Internal server error"}},
98
+ status_code=500
99
+ )
100
+
101
+
102
+ @router.get("/v1beta/files")
103
+ async def list_files(
104
+ page_size: int = Query(10, ge=1, le=100, description="每页大小", alias="pageSize"),
105
+ page_token: Optional[str] = Query(None, description="分页标记", alias="pageToken"),
106
+ auth_token: str = Depends(security_service.verify_key_or_goog_api_key)
107
+ ) -> ListFilesResponse:
108
+ """列出文件"""
109
+ logger.debug(f"List files: {page_size=}, {page_token=}, {auth_token=}")
110
+ try:
111
+ # 使用认证 token 作为 user_token(如果启用用户隔离)
112
+ user_token = auth_token if settings.FILES_USER_ISOLATION_ENABLED else None
113
+ # 调用服务
114
+ files_service = await get_files_service()
115
+ return await files_service.list_files(
116
+ page_size=page_size,
117
+ page_token=page_token,
118
+ user_token=user_token
119
+ )
120
+
121
+ except HTTPException as e:
122
+ logger.error(f"List files failed: {e.detail}")
123
+ return JSONResponse(
124
+ content={"error": {"message": e.detail}},
125
+ status_code=e.status_code
126
+ )
127
+ except Exception as e:
128
+ logger.error(f"Unexpected error in list files: {str(e)}")
129
+ return JSONResponse(
130
+ content={"error": {"message": "Internal server error"}},
131
+ status_code=500
132
+ )
133
+
134
+
135
+ @router.get("/v1beta/files/{file_id:path}")
136
+ async def get_file(
137
+ file_id: str,
138
+ auth_token: str = Depends(security_service.verify_key_or_goog_api_key)
139
+ ) -> FileMetadata:
140
+ """获取文件信息"""
141
+ logger.debug(f"Get file request: {file_id=}, {auth_token=}")
142
+ try:
143
+ # 使用认证 token 作为 user_token
144
+ user_token = auth_token
145
+ # 调用服务
146
+ files_service = await get_files_service()
147
+ return await files_service.get_file(f"files/{file_id}", user_token)
148
+
149
+ except HTTPException as e:
150
+ logger.error(f"Get file failed: {e.detail}")
151
+ return JSONResponse(
152
+ content={"error": {"message": e.detail}},
153
+ status_code=e.status_code
154
+ )
155
+ except Exception as e:
156
+ logger.error(f"Unexpected error in get file: {str(e)}")
157
+ return JSONResponse(
158
+ content={"error": {"message": "Internal server error"}},
159
+ status_code=500
160
+ )
161
+
162
+
163
+ @router.delete("/v1beta/files/{file_id:path}")
164
+ async def delete_file(
165
+ file_id: str,
166
+ auth_token: str = Depends(security_service.verify_key_or_goog_api_key)
167
+ ) -> DeleteFileResponse:
168
+ """删除文件"""
169
+ logger.info(f"Delete file: {file_id=}, {auth_token=}")
170
+ try:
171
+ # 使用认证 token 作为 user_token
172
+ user_token = auth_token
173
+ # 调用服务
174
+ files_service = await get_files_service()
175
+ success = await files_service.delete_file(f"files/{file_id}", user_token)
176
+
177
+ return DeleteFileResponse(
178
+ success=success,
179
+ message="File deleted successfully" if success else "Failed to delete file"
180
+ )
181
+
182
+ except HTTPException as e:
183
+ logger.error(f"Delete file failed: {e.detail}")
184
+ return JSONResponse(
185
+ content={"error": {"message": e.detail}},
186
+ status_code=e.status_code
187
+ )
188
+ except Exception as e:
189
+ logger.error(f"Unexpected error in delete file: {str(e)}")
190
+ return JSONResponse(
191
+ content={"error": {"message": "Internal server error"}},
192
+ status_code=500
193
+ )
194
+
195
+
196
+ # 处理上传请求的通配符路由
197
+ @router.api_route("/upload/{upload_path:path}", methods=["GET", "POST", "PUT"])
198
+ async def handle_upload(
199
+ upload_path: str,
200
+ request: Request,
201
+ key: Optional[str] = Query(None), # 從查詢參數獲取 key
202
+ auth_token: str = Depends(security_service.verify_key_or_goog_api_key)
203
+ ):
204
+ """处理文件上传请求"""
205
+ try:
206
+ logger.info(f"Handling upload request: {request.method} {upload_path}, key={redact_key_for_logging(key)}")
207
+
208
+ # 從查詢參數獲取 upload_id
209
+ upload_id = request.query_params.get("upload_id")
210
+ if not upload_id:
211
+ raise HTTPException(status_code=400, detail="Missing upload_id")
212
+
213
+ # 從 session 獲取真實的 API key
214
+ files_service = await get_files_service()
215
+ session_info = await files_service.get_upload_session(upload_id)
216
+ if not session_info:
217
+ logger.error(f"No session found for upload_id: {upload_id}")
218
+ raise HTTPException(status_code=404, detail="Upload session not found")
219
+
220
+ real_api_key = session_info["api_key"]
221
+ original_upload_url = session_info["upload_url"]
222
+
223
+ # 使用真實的 API key 構建完整的 Google 上傳 URL
224
+ # 保留原始 URL 的所有參數,但使用真實的 API key
225
+ upload_url = original_upload_url
226
+ logger.info(f"Using real API key for upload: {redact_key_for_logging(real_api_key)}")
227
+
228
+ # 代理上传请求
229
+ upload_handler = get_upload_handler()
230
+ return await upload_handler.proxy_upload_request(
231
+ request=request,
232
+ upload_url=upload_url,
233
+ files_service=files_service
234
+ )
235
+
236
+ except HTTPException as e:
237
+ logger.error(f"Upload handling failed: {e.detail}")
238
+ return JSONResponse(
239
+ content={"error": {"message": e.detail}},
240
+ status_code=e.status_code
241
+ )
242
+ except Exception as e:
243
+ logger.error(f"Unexpected error in upload handling: {str(e)}")
244
+ return JSONResponse(
245
+ content={"error": {"message": "Internal server error"}},
246
+ status_code=500
247
+ )
248
+
249
+
250
+ # 为兼容性添加 /gemini 前缀的路由
251
+ @router.post("/gemini/upload/v1beta/files")
252
+ async def gemini_upload_file_init(
253
+ request: Request,
254
+ auth_token: str = Depends(security_service.verify_key_or_goog_api_key),
255
+ x_goog_upload_protocol: Optional[str] = Header(None),
256
+ x_goog_upload_command: Optional[str] = Header(None),
257
+ x_goog_upload_header_content_length: Optional[str] = Header(None),
258
+ x_goog_upload_header_content_type: Optional[str] = Header(None),
259
+ ):
260
+ """初始化文件上传(Gemini 前缀)"""
261
+ return await upload_file_init(
262
+ request,
263
+ auth_token,
264
+ x_goog_upload_protocol,
265
+ x_goog_upload_command,
266
+ x_goog_upload_header_content_length,
267
+ x_goog_upload_header_content_type
268
+ )
269
+
270
+
271
+ @router.get("/gemini/v1beta/files")
272
+ async def gemini_list_files(
273
+ page_size: int = Query(10, ge=1, le=100, alias="pageSize"),
274
+ page_token: Optional[str] = Query(None, alias="pageToken"),
275
+ auth_token: str = Depends(security_service.verify_key_or_goog_api_key)
276
+ ) -> ListFilesResponse:
277
+ """列出文件(Gemini 前缀)"""
278
+ return await list_files(page_size, page_token, auth_token)
279
+
280
+
281
+ @router.get("/gemini/v1beta/files/{file_id:path}")
282
+ async def gemini_get_file(
283
+ file_id: str,
284
+ auth_token: str = Depends(security_service.verify_key_or_goog_api_key)
285
+ ) -> FileMetadata:
286
+ """获取文件信息(Gemini 前缀)"""
287
+ return await get_file(file_id, auth_token)
288
+
289
+
290
+ @router.delete("/gemini/v1beta/files/{file_id:path}")
291
+ async def gemini_delete_file(
292
+ file_id: str,
293
+ auth_token: str = Depends(security_service.verify_key_or_goog_api_key)
294
+ ) -> DeleteFileResponse:
295
+ """删除文件(Gemini 前缀)"""
296
+ return await delete_file(file_id, auth_token)
app/router/gemini_routes.py ADDED
@@ -0,0 +1,502 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import APIRouter, Depends, HTTPException
2
+ from fastapi.responses import StreamingResponse, JSONResponse
3
+ from copy import deepcopy
4
+ import asyncio
5
+ from app.config.config import settings
6
+ from app.log.logger import get_gemini_logger
7
+ from app.core.security import SecurityService
8
+ from app.domain.gemini_models import GeminiContent, GeminiRequest, ResetSelectedKeysRequest, VerifySelectedKeysRequest, GeminiEmbedRequest, GeminiBatchEmbedRequest
9
+ from app.service.chat.gemini_chat_service import GeminiChatService
10
+ from app.service.embedding.gemini_embedding_service import GeminiEmbeddingService
11
+ from app.service.key.key_manager import KeyManager, get_key_manager_instance
12
+ from app.service.tts.native.tts_routes import get_tts_chat_service
13
+ from app.service.model.model_service import ModelService
14
+ from app.handler.retry_handler import RetryHandler
15
+ from app.handler.error_handler import handle_route_errors
16
+ from app.core.constants import API_VERSION
17
+ from app.utils.helpers import redact_key_for_logging
18
+
19
+ router = APIRouter(prefix=f"/gemini/{API_VERSION}")
20
+ router_v1beta = APIRouter(prefix=f"/{API_VERSION}")
21
+ logger = get_gemini_logger()
22
+
23
+ security_service = SecurityService()
24
+ model_service = ModelService()
25
+
26
+
27
+ async def get_key_manager():
28
+ """获取密钥管理器实例"""
29
+ return await get_key_manager_instance()
30
+
31
+
32
+ async def get_next_working_key(key_manager: KeyManager = Depends(get_key_manager)):
33
+ """获取下一个可用的API密钥"""
34
+ return await key_manager.get_next_working_key()
35
+
36
+
37
+ async def get_chat_service(key_manager: KeyManager = Depends(get_key_manager)):
38
+ """获取Gemini聊天服务实例"""
39
+ return GeminiChatService(settings.BASE_URL, key_manager)
40
+
41
+
42
+ async def get_embedding_service(key_manager: KeyManager = Depends(get_key_manager)):
43
+ """获取Gemini嵌入服务实例"""
44
+ return GeminiEmbeddingService(settings.BASE_URL, key_manager)
45
+
46
+
47
+ @router.get("/models")
48
+ @router_v1beta.get("/models")
49
+ async def list_models(
50
+ _=Depends(security_service.verify_key_or_goog_api_key),
51
+ key_manager: KeyManager = Depends(get_key_manager)
52
+ ):
53
+ """获取可用的 Gemini 模型列表,并根据配置添加衍生模型(搜索、图像、非思考)。"""
54
+ operation_name = "list_gemini_models"
55
+ logger.info("-" * 50 + operation_name + "-" * 50)
56
+ logger.info("Handling Gemini models list request")
57
+
58
+ try:
59
+ api_key = await key_manager.get_random_valid_key()
60
+ if not api_key:
61
+ raise HTTPException(status_code=503, detail="No valid API keys available to fetch models.")
62
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
63
+
64
+ models_data = await model_service.get_gemini_models(api_key)
65
+ if not models_data or "models" not in models_data:
66
+ raise HTTPException(status_code=500, detail="Failed to fetch base models list.")
67
+
68
+ models_json = deepcopy(models_data)
69
+ model_mapping = {x.get("name", "").split("/", maxsplit=1)[-1]: x for x in models_json.get("models", [])}
70
+
71
+ def add_derived_model(base_name, suffix, display_suffix):
72
+ model = model_mapping.get(base_name)
73
+ if not model:
74
+ logger.warning(f"Base model '{base_name}' not found for derived model '{suffix}'.")
75
+ return
76
+ item = deepcopy(model)
77
+ item["name"] = f"models/{base_name}{suffix}"
78
+ display_name = f'{item.get("displayName", base_name)}{display_suffix}'
79
+ item["displayName"] = display_name
80
+ item["description"] = display_name
81
+ models_json["models"].append(item)
82
+
83
+ if settings.SEARCH_MODELS:
84
+ for name in settings.SEARCH_MODELS:
85
+ add_derived_model(name, "-search", " For Search")
86
+ if settings.IMAGE_MODELS:
87
+ for name in settings.IMAGE_MODELS:
88
+ add_derived_model(name, "-image", " For Image")
89
+ if settings.THINKING_MODELS:
90
+ for name in settings.THINKING_MODELS:
91
+ add_derived_model(name, "-non-thinking", " Non Thinking")
92
+
93
+ logger.info("Gemini models list request successful")
94
+ return models_json
95
+ except HTTPException as http_exc:
96
+ raise http_exc
97
+ except Exception as e:
98
+ logger.error(f"Error getting Gemini models list: {str(e)}")
99
+ raise HTTPException(
100
+ status_code=500, detail="Internal server error while fetching Gemini models list"
101
+ ) from e
102
+
103
+
104
+ @router.post("/models/{model_name}:generateContent")
105
+ @router_v1beta.post("/models/{model_name}:generateContent")
106
+ @RetryHandler(key_arg="api_key")
107
+ async def generate_content(
108
+ model_name: str,
109
+ request: GeminiRequest,
110
+ _=Depends(security_service.verify_key_or_goog_api_key),
111
+ api_key: str = Depends(get_next_working_key),
112
+ key_manager: KeyManager = Depends(get_key_manager),
113
+ chat_service: GeminiChatService = Depends(get_chat_service)
114
+ ):
115
+ """处理 Gemini 非流式内容生成请求。"""
116
+ operation_name = "gemini_generate_content"
117
+ async with handle_route_errors(logger, operation_name, failure_message="Content generation failed"):
118
+ logger.info(f"Handling Gemini content generation request for model: {model_name}")
119
+ logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
120
+
121
+ # 检测是否为原生Gemini TTS请求
122
+ is_native_tts = False
123
+ if "tts" in model_name.lower() and request.generationConfig:
124
+ # 直接从解析后的request对象获取TTS配置
125
+ response_modalities = request.generationConfig.responseModalities or []
126
+ speech_config = request.generationConfig.speechConfig or {}
127
+
128
+ # 如果包含AUDIO模态和语音配置,则认为是原生TTS请求
129
+ if "AUDIO" in response_modalities and speech_config:
130
+ is_native_tts = True
131
+ logger.info("Detected native Gemini TTS request")
132
+ logger.info(f"TTS responseModalities: {response_modalities}")
133
+ logger.info(f"TTS speechConfig: {speech_config}")
134
+
135
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
136
+
137
+ if not await model_service.check_model_support(model_name):
138
+ raise HTTPException(status_code=400, detail=f"Model {model_name} is not supported")
139
+
140
+ # 所有原生TTS请求都使用TTS增强服务
141
+ if is_native_tts:
142
+ try:
143
+ logger.info("Using native TTS enhanced service")
144
+ tts_service = await get_tts_chat_service(key_manager)
145
+ response = await tts_service.generate_content(
146
+ model=model_name,
147
+ request=request,
148
+ api_key=api_key
149
+ )
150
+ return response
151
+ except Exception as e:
152
+ logger.warning(f"Native TTS processing failed, falling back to standard service: {e}")
153
+
154
+ # 使用标准服务处理所有其他请求(非TTS)
155
+ response = await chat_service.generate_content(
156
+ model=model_name,
157
+ request=request,
158
+ api_key=api_key
159
+ )
160
+ return response
161
+
162
+
163
+ @router.post("/models/{model_name}:streamGenerateContent")
164
+ @router_v1beta.post("/models/{model_name}:streamGenerateContent")
165
+ @RetryHandler(key_arg="api_key")
166
+ async def stream_generate_content(
167
+ model_name: str,
168
+ request: GeminiRequest,
169
+ _=Depends(security_service.verify_key_or_goog_api_key),
170
+ api_key: str = Depends(get_next_working_key),
171
+ key_manager: KeyManager = Depends(get_key_manager),
172
+ chat_service: GeminiChatService = Depends(get_chat_service)
173
+ ):
174
+ """处理 Gemini 流式内容生成请求。"""
175
+ operation_name = "gemini_stream_generate_content"
176
+ async with handle_route_errors(logger, operation_name, failure_message="Streaming request initiation failed"):
177
+ logger.info(f"Handling Gemini streaming content generation for model: {model_name}")
178
+ logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
179
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
180
+
181
+ if not await model_service.check_model_support(model_name):
182
+ raise HTTPException(status_code=400, detail=f"Model {model_name} is not supported")
183
+
184
+ response_stream = chat_service.stream_generate_content(
185
+ model=model_name,
186
+ request=request,
187
+ api_key=api_key
188
+ )
189
+ return StreamingResponse(response_stream, media_type="text/event-stream")
190
+
191
+
192
+ @router.post("/models/{model_name}:countTokens")
193
+ @router_v1beta.post("/models/{model_name}:countTokens")
194
+ @RetryHandler(key_arg="api_key")
195
+ async def count_tokens(
196
+ model_name: str,
197
+ request: GeminiRequest,
198
+ _=Depends(security_service.verify_key_or_goog_api_key),
199
+ api_key: str = Depends(get_next_working_key),
200
+ key_manager: KeyManager = Depends(get_key_manager),
201
+ chat_service: GeminiChatService = Depends(get_chat_service)
202
+ ):
203
+ """处理 Gemini token 计数请求。"""
204
+ operation_name = "gemini_count_tokens"
205
+ async with handle_route_errors(logger, operation_name, failure_message="Token counting failed"):
206
+ logger.info(f"Handling Gemini token count request for model: {model_name}")
207
+ logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
208
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
209
+
210
+ if not await model_service.check_model_support(model_name):
211
+ raise HTTPException(status_code=400, detail=f"Model {model_name} is not supported")
212
+
213
+ response = await chat_service.count_tokens(
214
+ model=model_name,
215
+ request=request,
216
+ api_key=api_key
217
+ )
218
+ return response
219
+
220
+ @router.post("/models/{model_name}:embedContent")
221
+ @router_v1beta.post("/models/{model_name}:embedContent")
222
+ @RetryHandler(key_arg="api_key")
223
+ async def embed_content(
224
+ model_name: str,
225
+ request: GeminiEmbedRequest,
226
+ _=Depends(security_service.verify_key_or_goog_api_key),
227
+ api_key: str = Depends(get_next_working_key),
228
+ key_manager: KeyManager = Depends(get_key_manager),
229
+ embedding_service: GeminiEmbeddingService = Depends(get_embedding_service)
230
+ ):
231
+ """处理 Gemini 单一嵌入请求"""
232
+ operation_name = "gemini_embed_content"
233
+ async with handle_route_errors(logger, operation_name, failure_message="Embedding content generation failed"):
234
+ logger.info(f"Handling Gemini embedding request for model: {model_name}")
235
+ logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
236
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
237
+
238
+ if not await model_service.check_model_support(model_name):
239
+ raise HTTPException(status_code=400, detail=f"Model {model_name} is not supported")
240
+
241
+ response = await embedding_service.embed_content(
242
+ model=model_name,
243
+ request=request,
244
+ api_key=api_key
245
+ )
246
+ return response
247
+
248
+
249
+ @router.post("/models/{model_name}:batchEmbedContents")
250
+ @router_v1beta.post("/models/{model_name}:batchEmbedContents")
251
+ @RetryHandler(key_arg="api_key")
252
+ async def batch_embed_contents(
253
+ model_name: str,
254
+ request: GeminiBatchEmbedRequest,
255
+ _=Depends(security_service.verify_key_or_goog_api_key),
256
+ api_key: str = Depends(get_next_working_key),
257
+ key_manager: KeyManager = Depends(get_key_manager),
258
+ embedding_service: GeminiEmbeddingService = Depends(get_embedding_service)
259
+ ):
260
+ """处理 Gemini 批量嵌入请求"""
261
+ operation_name = "gemini_batch_embed_contents"
262
+ async with handle_route_errors(logger, operation_name, failure_message="Batch embedding content generation failed"):
263
+ logger.info(f"Handling Gemini batch embedding request for model: {model_name}")
264
+ logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
265
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
266
+
267
+ if not await model_service.check_model_support(model_name):
268
+ raise HTTPException(status_code=400, detail=f"Model {model_name} is not supported")
269
+
270
+ response = await embedding_service.batch_embed_contents(
271
+ model=model_name,
272
+ request=request,
273
+ api_key=api_key
274
+ )
275
+ return response
276
+
277
+
278
+ @router.post("/reset-all-fail-counts")
279
+ async def reset_all_key_fail_counts(key_type: str = None, key_manager: KeyManager = Depends(get_key_manager)):
280
+ """批量重置Gemini API密钥的失败计数,可选择性地仅重置有效或无效密钥"""
281
+ logger.info("-" * 50 + "reset_all_gemini_key_fail_counts" + "-" * 50)
282
+ logger.info(f"Received reset request with key_type: {key_type}")
283
+
284
+ try:
285
+ # 获取分类后的密钥
286
+ keys_by_status = await key_manager.get_keys_by_status()
287
+ valid_keys = keys_by_status.get("valid_keys", {})
288
+ invalid_keys = keys_by_status.get("invalid_keys", {})
289
+
290
+ # 根据类型选择要重置的密钥
291
+ keys_to_reset = []
292
+ if key_type == "valid":
293
+ keys_to_reset = list(valid_keys.keys())
294
+ logger.info(f"Resetting only valid keys, count: {len(keys_to_reset)}")
295
+ elif key_type == "invalid":
296
+ keys_to_reset = list(invalid_keys.keys())
297
+ logger.info(f"Resetting only invalid keys, count: {len(keys_to_reset)}")
298
+ else:
299
+ # 重置所有密钥
300
+ await key_manager.reset_failure_counts()
301
+ return JSONResponse({"success": True, "message": "所有密钥的失败计数已重置"})
302
+
303
+ # 批量重置指定类型的密钥
304
+ for key in keys_to_reset:
305
+ await key_manager.reset_key_failure_count(key)
306
+
307
+ return JSONResponse({
308
+ "success": True,
309
+ "message": f"{key_type}密钥的失败计数已重置",
310
+ "reset_count": len(keys_to_reset)
311
+ })
312
+ except Exception as e:
313
+ logger.error(f"Failed to reset key failure counts: {str(e)}")
314
+ return JSONResponse({"success": False, "message": f"批量重置失败: {str(e)}"}, status_code=500)
315
+
316
+
317
+ @router.post("/reset-selected-fail-counts")
318
+ async def reset_selected_key_fail_counts(
319
+ request: ResetSelectedKeysRequest,
320
+ key_manager: KeyManager = Depends(get_key_manager)
321
+ ):
322
+ """批量重置选定Gemini API密钥的失败计数"""
323
+ logger.info("-" * 50 + "reset_selected_gemini_key_fail_counts" + "-" * 50)
324
+ keys_to_reset = request.keys
325
+ key_type = request.key_type
326
+ logger.info(f"Received reset request for {len(keys_to_reset)} selected {key_type} keys.")
327
+
328
+ if not keys_to_reset:
329
+ return JSONResponse({"success": False, "message": "没有提供需要重置的密钥"}, status_code=400)
330
+
331
+ reset_count = 0
332
+ errors = []
333
+
334
+ try:
335
+ for key in keys_to_reset:
336
+ try:
337
+ result = await key_manager.reset_key_failure_count(key)
338
+ if result:
339
+ reset_count += 1
340
+ else:
341
+ logger.warning(f"Key not found during selective reset: {redact_key_for_logging(key)}")
342
+ except Exception as key_error:
343
+ logger.error(f"Error resetting key {redact_key_for_logging(key)}: {str(key_error)}")
344
+ errors.append(f"Key {key}: {str(key_error)}")
345
+
346
+ if errors:
347
+ error_message = f"批量重置完成,但出现错误: {'; '.join(errors)}"
348
+ final_success = reset_count > 0
349
+ status_code = 207 if final_success and errors else 500
350
+ return JSONResponse({
351
+ "success": final_success,
352
+ "message": error_message,
353
+ "reset_count": reset_count
354
+ }, status_code=status_code)
355
+
356
+ return JSONResponse({
357
+ "success": True,
358
+ "message": f"成功重置 {reset_count} 个选定 {key_type} 密钥的失败计数",
359
+ "reset_count": reset_count
360
+ })
361
+ except Exception as e:
362
+ logger.error(f"Failed to process reset selected key failure counts request: {str(e)}")
363
+ return JSONResponse({"success": False, "message": f"批量重置处理失败: {str(e)}"}, status_code=500)
364
+
365
+
366
+ @router.post("/reset-fail-count/{api_key}")
367
+ async def reset_key_fail_count(api_key: str, key_manager: KeyManager = Depends(get_key_manager)):
368
+ """重置指定Gemini API密钥的失败计数"""
369
+ logger.info("-" * 50 + "reset_gemini_key_fail_count" + "-" * 50)
370
+ logger.info(f"Resetting failure count for API key: {redact_key_for_logging(api_key)}")
371
+
372
+ try:
373
+ result = await key_manager.reset_key_failure_count(api_key)
374
+ if result:
375
+ return JSONResponse({"success": True, "message": "失败计数已重置"})
376
+ return JSONResponse({"success": False, "message": "未找到指定密钥"}, status_code=404)
377
+ except Exception as e:
378
+ logger.error(f"Failed to reset key failure count: {str(e)}")
379
+ return JSONResponse({"success": False, "message": f"重置失败: {str(e)}"}, status_code=500)
380
+
381
+
382
+ @router.post("/verify-key/{api_key}")
383
+ async def verify_key(api_key: str, chat_service: GeminiChatService = Depends(get_chat_service), key_manager: KeyManager = Depends(get_key_manager)):
384
+ """验证Gemini API密钥的有效性"""
385
+ logger.info("-" * 50 + "verify_gemini_key" + "-" * 50)
386
+ logger.info("Verifying API key validity")
387
+
388
+ try:
389
+ gemini_request = GeminiRequest(
390
+ contents=[
391
+ GeminiContent(
392
+ role="user",
393
+ parts=[{"text": "hi"}],
394
+ )
395
+ ],
396
+ generation_config={"temperature": 0.7, "topP": 1.0, "maxOutputTokens": 10}
397
+ )
398
+
399
+ response = await chat_service.generate_content(
400
+ settings.TEST_MODEL,
401
+ gemini_request,
402
+ api_key
403
+ )
404
+
405
+ if response:
406
+ # 如果密钥验证成功,则重置其失败计数
407
+ await key_manager.reset_key_failure_count(api_key)
408
+ return JSONResponse({"status": "valid"})
409
+ except Exception as e:
410
+ logger.error(f"Key verification failed: {str(e)}")
411
+
412
+ async with key_manager.failure_count_lock:
413
+ if api_key in key_manager.key_failure_counts:
414
+ key_manager.key_failure_counts[api_key] += 1
415
+ logger.warning(f"Verification exception for key: {redact_key_for_logging(api_key)}, incrementing failure count")
416
+
417
+ return JSONResponse({"status": "invalid", "error": str(e)})
418
+
419
+
420
+ @router.post("/verify-selected-keys")
421
+ async def verify_selected_keys(
422
+ request: VerifySelectedKeysRequest,
423
+ chat_service: GeminiChatService = Depends(get_chat_service),
424
+ key_manager: KeyManager = Depends(get_key_manager)
425
+ ):
426
+ """批量验证选定Gemini API密钥的有效性"""
427
+ logger.info("-" * 50 + "verify_selected_gemini_keys" + "-" * 50)
428
+ keys_to_verify = request.keys
429
+ logger.info(f"Received verification request for {len(keys_to_verify)} selected keys.")
430
+
431
+ if not keys_to_verify:
432
+ return JSONResponse({"success": False, "message": "没有提供需要验证的密钥"}, status_code=400)
433
+
434
+ successful_keys = []
435
+ failed_keys = {}
436
+
437
+ async def _verify_single_key(api_key: str):
438
+ """内部函数,用于验证单个密钥并处理异常"""
439
+ nonlocal successful_keys, failed_keys
440
+ try:
441
+ gemini_request = GeminiRequest(
442
+ contents=[GeminiContent(role="user", parts=[{"text": "hi"}])],
443
+ generation_config={"temperature": 0.7, "topP": 1.0, "maxOutputTokens": 10}
444
+ )
445
+ await chat_service.generate_content(
446
+ settings.TEST_MODEL,
447
+ gemini_request,
448
+ api_key
449
+ )
450
+ successful_keys.append(api_key)
451
+ # 如果密钥验证成功,则重置其失败计数
452
+ await key_manager.reset_key_failure_count(api_key)
453
+ return api_key, "valid", None
454
+ except Exception as e:
455
+ error_message = str(e)
456
+ logger.warning(f"Key verification failed for {redact_key_for_logging(api_key)}: {error_message}")
457
+ async with key_manager.failure_count_lock:
458
+ if api_key in key_manager.key_failure_counts:
459
+ key_manager.key_failure_counts[api_key] += 1
460
+ logger.warning(f"Bulk verification exception for key: {redact_key_for_logging(api_key)}, incrementing failure count")
461
+ else:
462
+ key_manager.key_failure_counts[api_key] = 1
463
+ logger.warning(f"Bulk verification exception for key: {redact_key_for_logging(api_key)}, initializing failure count to 1")
464
+ failed_keys[api_key] = error_message
465
+ return api_key, "invalid", error_message
466
+
467
+ tasks = [_verify_single_key(key) for key in keys_to_verify]
468
+ results = await asyncio.gather(*tasks, return_exceptions=True)
469
+
470
+ for result in results:
471
+ if isinstance(result, Exception):
472
+ logger.error(f"An unexpected error occurred during bulk verification task: {result}")
473
+ elif result:
474
+ if not isinstance(result, Exception) and result:
475
+ key, status, error = result
476
+ elif isinstance(result, Exception):
477
+ logger.error(f"Task execution error during bulk verification: {result}")
478
+
479
+ valid_count = len(successful_keys)
480
+ invalid_count = len(failed_keys)
481
+ logger.info(f"Bulk verification finished. Valid: {valid_count}, Invalid: {invalid_count}")
482
+
483
+ if failed_keys:
484
+ message = f"批量验证完成。成功: {valid_count}, 失败: {invalid_count}。"
485
+ return JSONResponse({
486
+ "success": True,
487
+ "message": message,
488
+ "successful_keys": successful_keys,
489
+ "failed_keys": failed_keys,
490
+ "valid_count": valid_count,
491
+ "invalid_count": invalid_count
492
+ })
493
+ else:
494
+ message = f"批量验证成功完成。所有 {valid_count} 个密钥均有效。"
495
+ return JSONResponse({
496
+ "success": True,
497
+ "message": message,
498
+ "successful_keys": successful_keys,
499
+ "failed_keys": {},
500
+ "valid_count": valid_count,
501
+ "invalid_count": 0
502
+ })
app/router/key_routes.py ADDED
@@ -0,0 +1,83 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import APIRouter, Depends, Request
2
+ from app.service.key.key_manager import KeyManager, get_key_manager_instance
3
+ from app.core.security import verify_auth_token
4
+ from fastapi.responses import JSONResponse
5
+
6
+ router = APIRouter()
7
+
8
+ @router.get("/api/keys")
9
+ async def get_keys_paginated(
10
+ request: Request,
11
+ page: int = 1,
12
+ limit: int = 10,
13
+ search: str = None,
14
+ fail_count_threshold: int = None,
15
+ status: str = "all", # 'valid', 'invalid', 'all'
16
+ key_manager: KeyManager = Depends(get_key_manager_instance),
17
+ ):
18
+ """
19
+ Get paginated, filtered, and searched keys.
20
+ """
21
+ auth_token = request.cookies.get("auth_token")
22
+ if not auth_token or not verify_auth_token(auth_token):
23
+ return JSONResponse(status_code=401, content={"detail": "Unauthorized"})
24
+
25
+ all_keys_with_status = await key_manager.get_all_keys_with_fail_count()
26
+
27
+ # Filter by status
28
+ if status == "valid":
29
+ keys_to_filter = all_keys_with_status["valid_keys"]
30
+ elif status == "invalid":
31
+ keys_to_filter = all_keys_with_status["invalid_keys"]
32
+ else:
33
+ # Combine both for 'all' status, which might be useful for a unified view if ever needed
34
+ keys_to_filter = {**all_keys_with_status["valid_keys"], **all_keys_with_status["invalid_keys"]}
35
+
36
+
37
+ # Further filtering (search and fail_count_threshold)
38
+ filtered_keys = {}
39
+ for key, fail_count in keys_to_filter.items():
40
+ search_match = True
41
+ if search:
42
+ search_match = search.lower() in key.lower()
43
+
44
+ fail_count_match = True
45
+ if fail_count_threshold is not None:
46
+ fail_count_match = fail_count >= fail_count_threshold
47
+
48
+ if search_match and fail_count_match:
49
+ filtered_keys[key] = fail_count
50
+
51
+ # Pagination
52
+ keys_list = list(filtered_keys.items())
53
+ total_items = len(keys_list)
54
+ start_index = (page - 1) * limit
55
+ end_index = start_index + limit
56
+ paginated_keys = dict(keys_list[start_index:end_index])
57
+
58
+ return {
59
+ "keys": paginated_keys,
60
+ "total_items": total_items,
61
+ "total_pages": (total_items + limit - 1) // limit,
62
+ "current_page": page,
63
+ }
64
+
65
+ @router.get("/api/keys/all")
66
+ async def get_all_keys(
67
+ request: Request,
68
+ key_manager: KeyManager = Depends(get_key_manager_instance),
69
+ ):
70
+ """
71
+ Get all keys (both valid and invalid) for bulk operations.
72
+ """
73
+ auth_token = request.cookies.get("auth_token")
74
+ if not auth_token or not verify_auth_token(auth_token):
75
+ return JSONResponse(status_code=401, content={"detail": "Unauthorized"})
76
+
77
+ all_keys_with_status = await key_manager.get_all_keys_with_fail_count()
78
+
79
+ return {
80
+ "valid_keys": list(all_keys_with_status["valid_keys"].keys()),
81
+ "invalid_keys": list(all_keys_with_status["invalid_keys"].keys()),
82
+ "total_count": len(all_keys_with_status["valid_keys"]) + len(all_keys_with_status["invalid_keys"])
83
+ }
app/router/openai_compatiable_routes.py ADDED
@@ -0,0 +1,114 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import APIRouter, Depends
2
+ from fastapi.responses import StreamingResponse
3
+
4
+ from app.config.config import settings
5
+ from app.core.security import SecurityService
6
+ from app.domain.openai_models import (
7
+ ChatRequest,
8
+ EmbeddingRequest,
9
+ ImageGenerationRequest,
10
+ )
11
+ from app.handler.retry_handler import RetryHandler
12
+ from app.handler.error_handler import handle_route_errors
13
+ from app.log.logger import get_openai_compatible_logger
14
+ from app.service.key.key_manager import KeyManager, get_key_manager_instance
15
+ from app.service.openai_compatiable.openai_compatiable_service import OpenAICompatiableService
16
+ from app.utils.helpers import redact_key_for_logging
17
+
18
+
19
+ router = APIRouter()
20
+ logger = get_openai_compatible_logger()
21
+
22
+ security_service = SecurityService()
23
+
24
+ async def get_key_manager():
25
+ return await get_key_manager_instance()
26
+
27
+
28
+ async def get_next_working_key_wrapper(
29
+ key_manager: KeyManager = Depends(get_key_manager),
30
+ ):
31
+ return await key_manager.get_next_working_key()
32
+
33
+
34
+ async def get_openai_service(key_manager: KeyManager = Depends(get_key_manager)):
35
+ """获取OpenAI聊天服务实例"""
36
+ return OpenAICompatiableService(settings.BASE_URL, key_manager)
37
+
38
+
39
+ @router.get("/openai/v1/models")
40
+ async def list_models(
41
+ _=Depends(security_service.verify_authorization),
42
+ key_manager: KeyManager = Depends(get_key_manager),
43
+ openai_service: OpenAICompatiableService = Depends(get_openai_service),
44
+ ):
45
+ """获取可用模型列表。"""
46
+ operation_name = "list_models"
47
+ async with handle_route_errors(logger, operation_name):
48
+ logger.info("Handling models list request")
49
+ api_key = await key_manager.get_random_valid_key()
50
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
51
+ return await openai_service.get_models(api_key)
52
+
53
+
54
+ @router.post("/openai/v1/chat/completions")
55
+ @RetryHandler(key_arg="api_key")
56
+ async def chat_completion(
57
+ request: ChatRequest,
58
+ _=Depends(security_service.verify_authorization),
59
+ api_key: str = Depends(get_next_working_key_wrapper),
60
+ key_manager: KeyManager = Depends(get_key_manager),
61
+ openai_service: OpenAICompatiableService = Depends(get_openai_service),
62
+ ):
63
+ """处理聊天补全请求,支持流式响应和特定模型切换。"""
64
+ operation_name = "chat_completion"
65
+ is_image_chat = request.model == f"{settings.CREATE_IMAGE_MODEL}-chat"
66
+ current_api_key = api_key
67
+ if is_image_chat:
68
+ current_api_key = await key_manager.get_paid_key()
69
+
70
+ async with handle_route_errors(logger, operation_name):
71
+ logger.info(f"Handling chat completion request for model: {request.model}")
72
+ logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
73
+ logger.info(f"Using API key: {redact_key_for_logging(current_api_key)}")
74
+
75
+ if is_image_chat:
76
+ response = await openai_service.create_image_chat_completion(request, current_api_key)
77
+ return response
78
+ else:
79
+ response = await openai_service.create_chat_completion(request, current_api_key)
80
+ if request.stream:
81
+ return StreamingResponse(response, media_type="text/event-stream")
82
+ return response
83
+
84
+
85
+ @router.post("/openai/v1/images/generations")
86
+ async def generate_image(
87
+ request: ImageGenerationRequest,
88
+ _=Depends(security_service.verify_authorization),
89
+ openai_service: OpenAICompatiableService = Depends(get_openai_service),
90
+ ):
91
+ """处理图像生成请求。"""
92
+ operation_name = "generate_image"
93
+ async with handle_route_errors(logger, operation_name):
94
+ logger.info(f"Handling image generation request for prompt: {request.prompt}")
95
+ request.model = settings.CREATE_IMAGE_MODEL
96
+ return await openai_service.generate_images(request)
97
+
98
+
99
+ @router.post("/openai/v1/embeddings")
100
+ async def embedding(
101
+ request: EmbeddingRequest,
102
+ _=Depends(security_service.verify_authorization),
103
+ key_manager: KeyManager = Depends(get_key_manager),
104
+ openai_service: OpenAICompatiableService = Depends(get_openai_service),
105
+ ):
106
+ """处理文本嵌入请求。"""
107
+ operation_name = "embedding"
108
+ async with handle_route_errors(logger, operation_name):
109
+ logger.info(f"Handling embedding request for model: {request.model}")
110
+ api_key = await key_manager.get_next_working_key()
111
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
112
+ return await openai_service.create_embeddings(
113
+ input_text=request.input, model=request.model, api_key=api_key
114
+ )
app/router/openai_routes.py ADDED
@@ -0,0 +1,176 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import APIRouter, Depends, HTTPException, Response
2
+ from fastapi.responses import StreamingResponse
3
+
4
+ from app.config.config import settings
5
+ from app.core.security import SecurityService
6
+ from app.domain.openai_models import (
7
+ ChatRequest,
8
+ EmbeddingRequest,
9
+ ImageGenerationRequest,
10
+ TTSRequest,
11
+ )
12
+ from app.handler.retry_handler import RetryHandler
13
+ from app.handler.error_handler import handle_route_errors
14
+ from app.log.logger import get_openai_logger
15
+ from app.service.chat.openai_chat_service import OpenAIChatService
16
+ from app.service.embedding.embedding_service import EmbeddingService
17
+ from app.service.image.image_create_service import ImageCreateService
18
+ from app.service.tts.tts_service import TTSService
19
+ from app.service.key.key_manager import KeyManager, get_key_manager_instance
20
+ from app.service.model.model_service import ModelService
21
+ from app.utils.helpers import redact_key_for_logging
22
+
23
+ router = APIRouter()
24
+ logger = get_openai_logger()
25
+
26
+ security_service = SecurityService()
27
+ model_service = ModelService()
28
+ embedding_service = EmbeddingService()
29
+ image_create_service = ImageCreateService()
30
+ tts_service = TTSService()
31
+
32
+
33
+ async def get_key_manager():
34
+ return await get_key_manager_instance()
35
+
36
+
37
+ async def get_next_working_key_wrapper(
38
+ key_manager: KeyManager = Depends(get_key_manager),
39
+ ):
40
+ return await key_manager.get_next_working_key()
41
+
42
+
43
+ async def get_openai_chat_service(key_manager: KeyManager = Depends(get_key_manager)):
44
+ """获取OpenAI聊天服务实例"""
45
+ return OpenAIChatService(settings.BASE_URL, key_manager)
46
+
47
+
48
+ async def get_tts_service():
49
+ """获取TTS服务实例"""
50
+ return tts_service
51
+
52
+
53
+ @router.get("/v1/models")
54
+ @router.get("/hf/v1/models")
55
+ async def list_models(
56
+ _=Depends(security_service.verify_authorization),
57
+ key_manager: KeyManager = Depends(get_key_manager),
58
+ ):
59
+ """获取可用的 OpenAI 模型列表 (兼容 Gemini 和 OpenAI)。"""
60
+ operation_name = "list_models"
61
+ async with handle_route_errors(logger, operation_name):
62
+ logger.info("Handling models list request")
63
+ api_key = await key_manager.get_random_valid_key()
64
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
65
+ return await model_service.get_gemini_openai_models(api_key)
66
+
67
+
68
+ @router.post("/v1/chat/completions")
69
+ @router.post("/hf/v1/chat/completions")
70
+ @RetryHandler(key_arg="api_key")
71
+ async def chat_completion(
72
+ request: ChatRequest,
73
+ _=Depends(security_service.verify_authorization),
74
+ api_key: str = Depends(get_next_working_key_wrapper),
75
+ key_manager: KeyManager = Depends(get_key_manager),
76
+ chat_service: OpenAIChatService = Depends(get_openai_chat_service),
77
+ ):
78
+ """处理 OpenAI 聊天补全请求,支持流式响应和特定模型切换。"""
79
+ operation_name = "chat_completion"
80
+ is_image_chat = request.model == f"{settings.CREATE_IMAGE_MODEL}-chat"
81
+ current_api_key = api_key
82
+ if is_image_chat:
83
+ current_api_key = await key_manager.get_paid_key()
84
+
85
+ async with handle_route_errors(logger, operation_name):
86
+ logger.info(f"Handling chat completion request for model: {request.model}")
87
+ logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
88
+ logger.info(f"Using API key: {redact_key_for_logging(current_api_key)}")
89
+
90
+ if not await model_service.check_model_support(request.model):
91
+ raise HTTPException(
92
+ status_code=400, detail=f"Model {request.model} is not supported"
93
+ )
94
+
95
+ if is_image_chat:
96
+ response = await chat_service.create_image_chat_completion(request, current_api_key)
97
+ if request.stream:
98
+ return StreamingResponse(response, media_type="text/event-stream")
99
+ return response
100
+ else:
101
+ response = await chat_service.create_chat_completion(request, current_api_key)
102
+ if request.stream:
103
+ return StreamingResponse(response, media_type="text/event-stream")
104
+ return response
105
+
106
+
107
+ @router.post("/v1/images/generations")
108
+ @router.post("/hf/v1/images/generations")
109
+ async def generate_image(
110
+ request: ImageGenerationRequest,
111
+ _=Depends(security_service.verify_authorization),
112
+ ):
113
+ """处理 OpenAI 图像生成请求。"""
114
+ operation_name = "generate_image"
115
+ async with handle_route_errors(logger, operation_name):
116
+ logger.info(f"Handling image generation request for prompt: {request.prompt}")
117
+ response = image_create_service.generate_images(request)
118
+ return response
119
+
120
+
121
+ @router.post("/v1/embeddings")
122
+ @router.post("/hf/v1/embeddings")
123
+ async def embedding(
124
+ request: EmbeddingRequest,
125
+ _=Depends(security_service.verify_authorization),
126
+ key_manager: KeyManager = Depends(get_key_manager),
127
+ ):
128
+ """处理 OpenAI 文本嵌入请求。"""
129
+ operation_name = "embedding"
130
+ async with handle_route_errors(logger, operation_name):
131
+ logger.info(f"Handling embedding request for model: {request.model}")
132
+ api_key = await key_manager.get_next_working_key()
133
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
134
+ response = await embedding_service.create_embedding(
135
+ input_text=request.input, model=request.model, api_key=api_key
136
+ )
137
+ return response
138
+
139
+
140
+ @router.get("/v1/keys/list")
141
+ @router.get("/hf/v1/keys/list")
142
+ async def get_keys_list(
143
+ _=Depends(security_service.verify_auth_token),
144
+ key_manager: KeyManager = Depends(get_key_manager),
145
+ ):
146
+ """获取有效和无效的API key列表 (需要管理 Token 认证)。"""
147
+ operation_name = "get_keys_list"
148
+ async with handle_route_errors(logger, operation_name):
149
+ logger.info("Handling keys list request")
150
+ keys_status = await key_manager.get_keys_by_status()
151
+ return {
152
+ "status": "success",
153
+ "data": {
154
+ "valid_keys": keys_status["valid_keys"],
155
+ "invalid_keys": keys_status["invalid_keys"],
156
+ },
157
+ "total": len(keys_status["valid_keys"]) + len(keys_status["invalid_keys"]),
158
+ }
159
+
160
+
161
+ @router.post("/v1/audio/speech")
162
+ @router.post("/hf/v1/audio/speech")
163
+ async def text_to_speech(
164
+ request: TTSRequest,
165
+ _=Depends(security_service.verify_authorization),
166
+ api_key: str = Depends(get_next_working_key_wrapper),
167
+ tts_service: TTSService = Depends(get_tts_service),
168
+ ):
169
+ """处理 OpenAI TTS 请求。"""
170
+ operation_name = "text_to_speech"
171
+ async with handle_route_errors(logger, operation_name):
172
+ logger.info(f"Handling TTS request for model: {request.model}")
173
+ logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
174
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
175
+ audio_data = await tts_service.create_tts(request, api_key)
176
+ return Response(content=audio_data, media_type="audio/wav")
app/router/routes.py ADDED
@@ -0,0 +1,287 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 路由配置模块,负责设置和配置应用程序的路由
3
+ """
4
+
5
+ from fastapi import FastAPI, Request
6
+ from fastapi.responses import HTMLResponse, RedirectResponse
7
+ from fastapi.templating import Jinja2Templates
8
+
9
+ from app.config.config import settings
10
+ from app.core.security import verify_auth_token
11
+ from app.log.logger import get_routes_logger
12
+ from app.router import (
13
+ config_routes,
14
+ error_log_routes,
15
+ files_routes,
16
+ gemini_routes,
17
+ key_routes,
18
+ openai_compatiable_routes,
19
+ openai_routes,
20
+ scheduler_routes,
21
+ stats_routes,
22
+ version_routes,
23
+ vertex_express_routes,
24
+ )
25
+ from app.service.key.key_manager import get_key_manager_instance
26
+ from app.service.stats.stats_service import StatsService
27
+
28
+ logger = get_routes_logger()
29
+
30
+ templates = Jinja2Templates(directory="app/templates")
31
+
32
+
33
+ def setup_routers(app: FastAPI) -> None:
34
+ """
35
+ 设置应用程序的路由
36
+
37
+ Args:
38
+ app: FastAPI应用程序实例
39
+ """
40
+ app.include_router(openai_routes.router)
41
+ app.include_router(gemini_routes.router)
42
+ app.include_router(gemini_routes.router_v1beta)
43
+ app.include_router(config_routes.router)
44
+ app.include_router(error_log_routes.router)
45
+ app.include_router(scheduler_routes.router)
46
+ app.include_router(stats_routes.router)
47
+ app.include_router(version_routes.router)
48
+ app.include_router(openai_compatiable_routes.router)
49
+ app.include_router(vertex_express_routes.router)
50
+ app.include_router(files_routes.router)
51
+ app.include_router(key_routes.router)
52
+
53
+ setup_page_routes(app)
54
+
55
+ setup_health_routes(app)
56
+ setup_api_stats_routes(app)
57
+
58
+
59
+ def setup_page_routes(app: FastAPI) -> None:
60
+ """
61
+ 设置页面相关的路由
62
+
63
+ Args:
64
+ app: FastAPI应用程序实例
65
+ """
66
+
67
+ @app.get("/", response_class=HTMLResponse)
68
+ async def auth_page(request: Request):
69
+ """认证页面"""
70
+ return templates.TemplateResponse("auth.html", {"request": request})
71
+
72
+ @app.post("/auth")
73
+ async def authenticate(request: Request):
74
+ """处理认证请求"""
75
+ try:
76
+ form = await request.form()
77
+ auth_token = form.get("auth_token")
78
+ if not auth_token:
79
+ logger.warning("Authentication attempt with empty token")
80
+ return RedirectResponse(url="/", status_code=302)
81
+
82
+ if verify_auth_token(auth_token):
83
+ logger.info("Successful authentication")
84
+ response = RedirectResponse(url="/keys", status_code=302)
85
+ response.set_cookie(
86
+ key="auth_token",
87
+ value=auth_token,
88
+ httponly=True,
89
+ max_age=settings.ADMIN_SESSION_EXPIRE,
90
+ )
91
+ return response
92
+ logger.warning("Failed authentication attempt with invalid token")
93
+ return RedirectResponse(url="/", status_code=302)
94
+ except Exception as e:
95
+ logger.error(f"Authentication error: {str(e)}")
96
+ return RedirectResponse(url="/", status_code=302)
97
+
98
+ @app.get("/keys", response_class=HTMLResponse)
99
+ async def keys_page(request: Request):
100
+ """密钥管理页面"""
101
+ try:
102
+ auth_token = request.cookies.get("auth_token")
103
+ if not auth_token or not verify_auth_token(auth_token):
104
+ logger.warning("Unauthorized access attempt to keys page")
105
+ return RedirectResponse(url="/", status_code=302)
106
+
107
+ key_manager = await get_key_manager_instance()
108
+ keys_status = await key_manager.get_keys_by_status()
109
+ total_keys = len(keys_status["valid_keys"]) + len(
110
+ keys_status["invalid_keys"]
111
+ )
112
+ valid_key_count = len(keys_status["valid_keys"])
113
+ invalid_key_count = len(keys_status["invalid_keys"])
114
+
115
+ stats_service = StatsService()
116
+ api_stats = await stats_service.get_api_usage_stats()
117
+ logger.info(f"API stats retrieved: {api_stats}")
118
+
119
+ logger.info(f"Keys status retrieved successfully. Total keys: {total_keys}")
120
+ return templates.TemplateResponse(
121
+ "keys_status.html",
122
+ {
123
+ "request": request,
124
+ "valid_keys": {},
125
+ "invalid_keys": {},
126
+ "total_keys": total_keys,
127
+ "valid_key_count": valid_key_count,
128
+ "invalid_key_count": invalid_key_count,
129
+ "api_stats": api_stats,
130
+ },
131
+ )
132
+ except Exception as e:
133
+ logger.error(f"Error retrieving keys status or API stats: {str(e)}")
134
+ # Even if there's an error, render the page with whatever data is available
135
+ # or with empty/default values, so the frontend can still load.
136
+ return templates.TemplateResponse(
137
+ "keys_status.html",
138
+ {
139
+ "request": request,
140
+ "valid_keys": {},
141
+ "invalid_keys": {},
142
+ "total_keys": 0,
143
+ "valid_key_count": 0,
144
+ "invalid_key_count": 0,
145
+ "api_stats": { # Provide a default structure for api_stats
146
+ "calls_1m": {"total": 0, "success": 0, "failure": 0},
147
+ "calls_1h": {"total": 0, "success": 0, "failure": 0},
148
+ "calls_24h": {"total": 0, "success": 0, "failure": 0},
149
+ "calls_month": {"total": 0, "success": 0, "failure": 0},
150
+ },
151
+ },
152
+ )
153
+
154
+ @app.get("/config", response_class=HTMLResponse)
155
+ async def config_page(request: Request):
156
+ """配置编辑页面"""
157
+ try:
158
+ auth_token = request.cookies.get("auth_token")
159
+ if not auth_token or not verify_auth_token(auth_token):
160
+ logger.warning("Unauthorized access attempt to config page")
161
+ return RedirectResponse(url="/", status_code=302)
162
+
163
+ logger.info("Config page accessed successfully")
164
+ return templates.TemplateResponse(
165
+ "config_editor.html", {"request": request}
166
+ )
167
+ except Exception as e:
168
+ logger.error(f"Error accessing config page: {str(e)}")
169
+ raise
170
+
171
+ @app.get("/logs", response_class=HTMLResponse)
172
+ async def logs_page(request: Request):
173
+ """错误日志页面"""
174
+ try:
175
+ auth_token = request.cookies.get("auth_token")
176
+ if not auth_token or not verify_auth_token(auth_token):
177
+ logger.warning("Unauthorized access attempt to logs page")
178
+ return RedirectResponse(url="/", status_code=302)
179
+
180
+ logger.info("Logs page accessed successfully")
181
+ return templates.TemplateResponse("error_logs.html", {"request": request})
182
+ except Exception as e:
183
+ logger.error(f"Error accessing logs page: {str(e)}")
184
+ raise
185
+
186
+
187
+ def setup_health_routes(app: FastAPI) -> None:
188
+ """
189
+ 设置健康检查相关的路由
190
+
191
+ Args:
192
+ app: FastAPI应用程序实例
193
+ """
194
+
195
+ @app.get("/health")
196
+ async def health_check(request: Request):
197
+ """健康检查端点"""
198
+ logger.info("Health check endpoint called")
199
+ return {"status": "healthy"}
200
+
201
+
202
+ def setup_api_stats_routes(app: FastAPI) -> None:
203
+ """
204
+ 设置 API 统计相关的路由
205
+
206
+ Args:
207
+ app: FastAPI应用程序实例
208
+ """
209
+
210
+ @app.get("/api/stats/details")
211
+ async def api_stats_details(request: Request, period: str):
212
+ """获取指定时间段内的 API 调用详情"""
213
+ try:
214
+ auth_token = request.cookies.get("auth_token")
215
+ if not auth_token or not verify_auth_token(auth_token):
216
+ logger.warning("Unauthorized access attempt to API stats details")
217
+ return {"error": "Unauthorized"}, 401
218
+
219
+ logger.info(f"Fetching API call details for period: {period}")
220
+ stats_service = StatsService()
221
+ details = await stats_service.get_api_call_details(period)
222
+ return details
223
+ except ValueError as e:
224
+ logger.warning(
225
+ f"Invalid period requested for API stats details: {period} - {str(e)}"
226
+ )
227
+ return {"error": str(e)}, 400
228
+ except Exception as e:
229
+ logger.error(
230
+ f"Error fetching API stats details for period {period}: {str(e)}"
231
+ )
232
+ return {"error": "Internal server error"}, 500
233
+
234
+ @app.get("/api/stats/attention-keys")
235
+ async def api_stats_attention_keys(
236
+ request: Request, limit: int = 20, status_code: int = 429
237
+ ):
238
+ """返回最近24小时指定错误码次数最多的Key(仅包含内存Key列表中的)。默认错误码429。"""
239
+ try:
240
+ auth_token = request.cookies.get("auth_token")
241
+ if not auth_token or not verify_auth_token(auth_token):
242
+ logger.warning("Unauthorized access attempt to attention-keys")
243
+ return {"error": "Unauthorized"}, 401
244
+
245
+ # 支持所有标准HTTP状态码范围
246
+ # if not isinstance(status_code, int) or status_code < 100 or status_code > 599:
247
+ # return {"error": f"Unsupported status_code: {status_code}"}, 400
248
+
249
+ key_manager = await get_key_manager_instance()
250
+ keys_status = await key_manager.get_keys_by_status()
251
+ in_memory_keys = set(keys_status.get("valid_keys", [])) | set(
252
+ keys_status.get("invalid_keys", [])
253
+ )
254
+ stats_service = StatsService()
255
+ data = await stats_service.get_attention_keys_last_24h(
256
+ in_memory_keys, limit, status_code
257
+ )
258
+ return data
259
+ except Exception as e:
260
+ logger.error(f"Error fetching attention keys: {e}")
261
+ return {"error": "Internal server error"}, 500
262
+
263
+ @app.get("/api/stats/key-details")
264
+ async def api_stats_key_details(request: Request, key: str, period: str):
265
+ """获取指定密钥在指定时间段内的调用详情"""
266
+ try:
267
+ auth_token = request.cookies.get("auth_token")
268
+ if not auth_token or not verify_auth_token(auth_token):
269
+ logger.warning("Unauthorized access attempt to API key stats details")
270
+ return {"error": "Unauthorized"}, 401
271
+
272
+ logger.info(
273
+ f"Fetching key call details for key=...{key[-4:] if key else ''}, period: {period}"
274
+ )
275
+ stats_service = StatsService()
276
+ details = await stats_service.get_key_call_details(key, period)
277
+ return details
278
+ except ValueError as e:
279
+ logger.warning(
280
+ f"Invalid period requested for key stats details: {period} - {str(e)}"
281
+ )
282
+ return {"error": str(e)}, 400
283
+ except Exception as e:
284
+ logger.error(
285
+ f"Error fetching key stats details for period {period}: {str(e)}"
286
+ )
287
+ return {"error": "Internal server error"}, 500
app/router/scheduler_routes.py ADDED
@@ -0,0 +1,57 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 定时任务控制路由模块
3
+ """
4
+
5
+ from fastapi import APIRouter, Request, HTTPException, status
6
+ from fastapi.responses import JSONResponse
7
+
8
+ from app.core.security import verify_auth_token
9
+ from app.scheduler.scheduled_tasks import start_scheduler, stop_scheduler
10
+ from app.log.logger import get_scheduler_routes
11
+
12
+ logger = get_scheduler_routes()
13
+
14
+ router = APIRouter(
15
+ prefix="/api/scheduler",
16
+ tags=["Scheduler"]
17
+ )
18
+
19
+ async def verify_token(request: Request):
20
+ auth_token = request.cookies.get("auth_token")
21
+ if not auth_token or not verify_auth_token(auth_token):
22
+ logger.warning("Unauthorized access attempt to scheduler API")
23
+ raise HTTPException(
24
+ status_code=status.HTTP_401_UNAUTHORIZED,
25
+ detail="Not authenticated",
26
+ headers={"WWW-Authenticate": "Bearer"},
27
+ )
28
+
29
+ @router.post("/start", summary="启动定时任务")
30
+ async def start_scheduler_endpoint(request: Request):
31
+ """Start the background scheduler task"""
32
+ await verify_token(request)
33
+ try:
34
+ logger.info("Received request to start scheduler.")
35
+ start_scheduler()
36
+ return JSONResponse(content={"message": "Scheduler started successfully."}, status_code=status.HTTP_200_OK)
37
+ except Exception as e:
38
+ logger.error(f"Error starting scheduler: {str(e)}", exc_info=True)
39
+ raise HTTPException(
40
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
41
+ detail=f"Failed to start scheduler: {str(e)}"
42
+ )
43
+
44
+ @router.post("/stop", summary="停止定时任务")
45
+ async def stop_scheduler_endpoint(request: Request):
46
+ """Stop the background scheduler task"""
47
+ await verify_token(request)
48
+ try:
49
+ logger.info("Received request to stop scheduler.")
50
+ stop_scheduler()
51
+ return JSONResponse(content={"message": "Scheduler stopped successfully."}, status_code=status.HTTP_200_OK)
52
+ except Exception as e:
53
+ logger.error(f"Error stopping scheduler: {str(e)}", exc_info=True)
54
+ raise HTTPException(
55
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
56
+ detail=f"Failed to stop scheduler: {str(e)}"
57
+ )
app/router/stats_routes.py ADDED
@@ -0,0 +1,56 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import APIRouter, Depends, HTTPException, Request
2
+ from starlette import status
3
+ from app.core.security import verify_auth_token
4
+ from app.service.stats.stats_service import StatsService
5
+ from app.log.logger import get_stats_logger
6
+ from app.utils.helpers import redact_key_for_logging
7
+
8
+ logger = get_stats_logger()
9
+
10
+
11
+ async def verify_token(request: Request):
12
+ auth_token = request.cookies.get("auth_token")
13
+ if not auth_token or not verify_auth_token(auth_token):
14
+ logger.warning("Unauthorized access attempt to scheduler API")
15
+ raise HTTPException(
16
+ status_code=status.HTTP_401_UNAUTHORIZED,
17
+ detail="Not authenticated",
18
+ headers={"WWW-Authenticate": "Bearer"},
19
+ )
20
+
21
+ router = APIRouter(
22
+ prefix="/api",
23
+ tags=["stats"],
24
+ dependencies=[Depends(verify_token)]
25
+ )
26
+
27
+ stats_service = StatsService()
28
+
29
+ @router.get("/key-usage-details/{key}",
30
+ summary="获取指定密钥最近24小时的模型调用次数",
31
+ description="根据提供的 API 密钥,返回过去24小时内每个模型被调用的次数统计。")
32
+ async def get_key_usage_details(key: str):
33
+ """
34
+ Retrieves the model usage count for a specific API key within the last 24 hours.
35
+
36
+ Args:
37
+ key: The API key to get usage details for.
38
+
39
+ Returns:
40
+ A dictionary with model names as keys and their call counts as values.
41
+ Example: {"gemini-pro": 10, "gemini-1.5-pro-latest": 5}
42
+
43
+ Raises:
44
+ HTTPException: If an error occurs during data retrieval.
45
+ """
46
+ try:
47
+ usage_details = await stats_service.get_key_usage_details_last_24h(key)
48
+ if usage_details is None:
49
+ return {}
50
+ return usage_details
51
+ except Exception as e:
52
+ logger.error(f"Error fetching key usage details for key {redact_key_for_logging(key)}: {e}")
53
+ raise HTTPException(
54
+ status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
55
+ detail=f"获取密钥使用详情时出错: {e}"
56
+ )
app/router/version_routes.py ADDED
@@ -0,0 +1,37 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import APIRouter, HTTPException
2
+ from pydantic import BaseModel, Field
3
+ from typing import Optional
4
+
5
+ from app.service.update.update_service import check_for_updates
6
+ from app.utils.helpers import get_current_version
7
+ from app.log.logger import get_update_logger
8
+
9
+ router = APIRouter(prefix="/api/version", tags=["Version"])
10
+ logger = get_update_logger()
11
+
12
+ class VersionInfo(BaseModel):
13
+ current_version: str = Field(..., description="当前应用程序版本")
14
+ latest_version: Optional[str] = Field(None, description="可用的最新版本")
15
+ update_available: bool = Field(False, description="是否有可用更新")
16
+ error_message: Optional[str] = Field(None, description="检查更新时发生的错误信息")
17
+
18
+ @router.get("/check", response_model=VersionInfo, summary="检查应用程序更新")
19
+ async def get_version_info():
20
+ """
21
+ 检查当前应用程序版本与最新的 GitHub release 版本。
22
+ """
23
+ try:
24
+ current_version = get_current_version()
25
+ update_available, latest_version, error_message = await check_for_updates()
26
+
27
+ logger.info(f"Version check API result: current={current_version}, latest={latest_version}, available={update_available}, error='{error_message}'")
28
+
29
+ return VersionInfo(
30
+ current_version=current_version,
31
+ latest_version=latest_version,
32
+ update_available=update_available,
33
+ error_message=error_message
34
+ )
35
+ except Exception as e:
36
+ logger.error(f"Error in /api/version/check endpoint: {e}", exc_info=True)
37
+ raise HTTPException(status_code=500, detail="检查版本信息时发生内部错误")
app/router/vertex_express_routes.py ADDED
@@ -0,0 +1,147 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from fastapi import APIRouter, Depends, HTTPException
2
+ from fastapi.responses import StreamingResponse
3
+ from copy import deepcopy
4
+ from app.config.config import settings
5
+ from app.log.logger import get_vertex_express_logger
6
+ from app.core.security import SecurityService
7
+ from app.domain.gemini_models import GeminiRequest
8
+ from app.service.chat.vertex_express_chat_service import GeminiChatService
9
+ from app.service.key.key_manager import KeyManager, get_key_manager_instance
10
+ from app.service.model.model_service import ModelService
11
+ from app.handler.retry_handler import RetryHandler
12
+ from app.handler.error_handler import handle_route_errors
13
+ from app.core.constants import API_VERSION
14
+ from app.utils.helpers import redact_key_for_logging
15
+
16
+ router = APIRouter(prefix=f"/vertex-express/{API_VERSION}")
17
+ logger = get_vertex_express_logger()
18
+
19
+ security_service = SecurityService()
20
+ model_service = ModelService()
21
+
22
+
23
+ async def get_key_manager():
24
+ """获取密钥管理器实例"""
25
+ return await get_key_manager_instance()
26
+
27
+
28
+ async def get_next_working_key(key_manager: KeyManager = Depends(get_key_manager)):
29
+ """获取下一个可用的API密钥"""
30
+ return await key_manager.get_next_working_vertex_key()
31
+
32
+
33
+ async def get_chat_service(key_manager: KeyManager = Depends(get_key_manager)):
34
+ """获取Gemini聊天服务实例"""
35
+ return GeminiChatService(settings.VERTEX_EXPRESS_BASE_URL, key_manager)
36
+
37
+
38
+ @router.get("/models")
39
+ async def list_models(
40
+ _=Depends(security_service.verify_key_or_goog_api_key),
41
+ key_manager: KeyManager = Depends(get_key_manager)
42
+ ):
43
+ """获取可用的 Gemini 模型列表,并根据配置添加衍生模型(搜索、图像、非思考)。"""
44
+ operation_name = "list_gemini_models"
45
+ logger.info("-" * 50 + operation_name + "-" * 50)
46
+ logger.info("Handling Gemini models list request")
47
+
48
+ try:
49
+ api_key = await key_manager.get_random_valid_key()
50
+ if not api_key:
51
+ raise HTTPException(status_code=503, detail="No valid API keys available to fetch models.")
52
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
53
+
54
+ models_data = await model_service.get_gemini_models(api_key)
55
+ if not models_data or "models" not in models_data:
56
+ raise HTTPException(status_code=500, detail="Failed to fetch base models list.")
57
+
58
+ models_json = deepcopy(models_data)
59
+ model_mapping = {x.get("name", "").split("/", maxsplit=1)[-1]: x for x in models_json.get("models", [])}
60
+
61
+ def add_derived_model(base_name, suffix, display_suffix):
62
+ model = model_mapping.get(base_name)
63
+ if not model:
64
+ logger.warning(f"Base model '{base_name}' not found for derived model '{suffix}'.")
65
+ return
66
+ item = deepcopy(model)
67
+ item["name"] = f"models/{base_name}{suffix}"
68
+ display_name = f'{item.get("displayName", base_name)}{display_suffix}'
69
+ item["displayName"] = display_name
70
+ item["description"] = display_name
71
+ models_json["models"].append(item)
72
+
73
+ if settings.SEARCH_MODELS:
74
+ for name in settings.SEARCH_MODELS:
75
+ add_derived_model(name, "-search", " For Search")
76
+ if settings.IMAGE_MODELS:
77
+ for name in settings.IMAGE_MODELS:
78
+ add_derived_model(name, "-image", " For Image")
79
+ if settings.THINKING_MODELS:
80
+ for name in settings.THINKING_MODELS:
81
+ add_derived_model(name, "-non-thinking", " Non Thinking")
82
+
83
+ logger.info("Gemini models list request successful")
84
+ return models_json
85
+ except HTTPException as http_exc:
86
+ raise http_exc
87
+ except Exception as e:
88
+ logger.error(f"Error getting Gemini models list: {str(e)}")
89
+ raise HTTPException(
90
+ status_code=500, detail="Internal server error while fetching Gemini models list"
91
+ ) from e
92
+
93
+
94
+ @router.post("/models/{model_name}:generateContent")
95
+ @RetryHandler(key_arg="api_key")
96
+ async def generate_content(
97
+ model_name: str,
98
+ request: GeminiRequest,
99
+ _=Depends(security_service.verify_key_or_goog_api_key),
100
+ api_key: str = Depends(get_next_working_key),
101
+ key_manager: KeyManager = Depends(get_key_manager),
102
+ chat_service: GeminiChatService = Depends(get_chat_service)
103
+ ):
104
+ """处理 Gemini 非流式内容生成请求。"""
105
+ operation_name = "gemini_generate_content"
106
+ async with handle_route_errors(logger, operation_name, failure_message="Content generation failed"):
107
+ logger.info(f"Handling Gemini content generation request for model: {model_name}")
108
+ logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
109
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
110
+
111
+ if not await model_service.check_model_support(model_name):
112
+ raise HTTPException(status_code=400, detail=f"Model {model_name} is not supported")
113
+
114
+ response = await chat_service.generate_content(
115
+ model=model_name,
116
+ request=request,
117
+ api_key=api_key
118
+ )
119
+ return response
120
+
121
+
122
+ @router.post("/models/{model_name}:streamGenerateContent")
123
+ @RetryHandler(key_arg="api_key")
124
+ async def stream_generate_content(
125
+ model_name: str,
126
+ request: GeminiRequest,
127
+ _=Depends(security_service.verify_key_or_goog_api_key),
128
+ api_key: str = Depends(get_next_working_key),
129
+ key_manager: KeyManager = Depends(get_key_manager),
130
+ chat_service: GeminiChatService = Depends(get_chat_service)
131
+ ):
132
+ """处理 Gemini 流式内容生成请求。"""
133
+ operation_name = "gemini_stream_generate_content"
134
+ async with handle_route_errors(logger, operation_name, failure_message="Streaming request initiation failed"):
135
+ logger.info(f"Handling Gemini streaming content generation for model: {model_name}")
136
+ logger.debug(f"Request: \n{request.model_dump_json(indent=2)}")
137
+ logger.info(f"Using API key: {redact_key_for_logging(api_key)}")
138
+
139
+ if not await model_service.check_model_support(model_name):
140
+ raise HTTPException(status_code=400, detail=f"Model {model_name} is not supported")
141
+
142
+ response_stream = chat_service.stream_generate_content(
143
+ model=model_name,
144
+ request=request,
145
+ api_key=api_key
146
+ )
147
+ return StreamingResponse(response_stream, media_type="text/event-stream")
app/scheduler/scheduled_tasks.py ADDED
@@ -0,0 +1,195 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ from apscheduler.schedulers.asyncio import AsyncIOScheduler
3
+
4
+ from app.config.config import settings
5
+ from app.domain.gemini_models import GeminiContent, GeminiRequest
6
+ from app.log.logger import Logger
7
+ from app.service.chat.gemini_chat_service import GeminiChatService
8
+ from app.service.error_log.error_log_service import delete_old_error_logs
9
+ from app.service.key.key_manager import get_key_manager_instance
10
+ from app.service.request_log.request_log_service import delete_old_request_logs_task
11
+ from app.service.files.files_service import get_files_service
12
+ from app.utils.helpers import redact_key_for_logging
13
+
14
+ logger = Logger.setup_logger("scheduler")
15
+
16
+
17
+ async def check_failed_keys():
18
+ """
19
+ 定时检查失败次数大于0的API密钥,并尝试验证它们。
20
+ 如果验证成功,重置失败计数;如果失败,增加失败计数。
21
+ """
22
+ logger.info("Starting scheduled check for failed API keys...")
23
+ try:
24
+ key_manager = await get_key_manager_instance()
25
+ # 确保 KeyManager 已经初始化
26
+ if not key_manager or not hasattr(key_manager, "key_failure_counts"):
27
+ logger.warning(
28
+ "KeyManager instance not available or not initialized. Skipping check."
29
+ )
30
+ return
31
+
32
+ # 创建 GeminiChatService 实例用于验证
33
+ # 注意:这里直接创建实例,而不是通过依赖注入,因为这是后台任务
34
+ chat_service = GeminiChatService(settings.BASE_URL, key_manager)
35
+
36
+ # 获取需要检查的 key 列表 (失败次数 > 0)
37
+ keys_to_check = []
38
+ async with key_manager.failure_count_lock: # 访问共享数据需要加锁
39
+ # 复制一份以避免在迭代时修改字典
40
+ failure_counts_copy = key_manager.key_failure_counts.copy()
41
+ keys_to_check = [
42
+ key for key, count in failure_counts_copy.items() if count > 0
43
+ ] # 检查所有失败次数大于0的key
44
+
45
+ if not keys_to_check:
46
+ logger.info("No keys with failure count > 0 found. Skipping verification.")
47
+ return
48
+
49
+ logger.info(
50
+ f"Found {len(keys_to_check)} keys with failure count > 0 to verify."
51
+ )
52
+
53
+ for key in keys_to_check:
54
+ # 隐藏部分 key 用于日志记录
55
+ log_key = redact_key_for_logging(key)
56
+ logger.info(f"Verifying key: {log_key}...")
57
+ try:
58
+ # 构造测试请求
59
+ gemini_request = GeminiRequest(
60
+ contents=[
61
+ GeminiContent(
62
+ role="user",
63
+ parts=[{"text": "hi"}],
64
+ )
65
+ ]
66
+ )
67
+ await chat_service.generate_content(
68
+ settings.TEST_MODEL, gemini_request, key
69
+ )
70
+ logger.info(
71
+ f"Key {log_key} verification successful. Resetting failure count."
72
+ )
73
+ await key_manager.reset_key_failure_count(key)
74
+ except Exception as e:
75
+ logger.warning(
76
+ f"Key {log_key} verification failed: {str(e)}. Incrementing failure count."
77
+ )
78
+ # 直接操作计数器,需要加锁
79
+ async with key_manager.failure_count_lock:
80
+ # 再次检查 key 是否存在且失败次数未达上限
81
+ if (
82
+ key in key_manager.key_failure_counts
83
+ and key_manager.key_failure_counts[key]
84
+ < key_manager.MAX_FAILURES
85
+ ):
86
+ key_manager.key_failure_counts[key] += 1
87
+ logger.info(
88
+ f"Failure count for key {log_key} incremented to {key_manager.key_failure_counts[key]}."
89
+ )
90
+ elif key in key_manager.key_failure_counts:
91
+ logger.warning(
92
+ f"Key {log_key} reached MAX_FAILURES ({key_manager.MAX_FAILURES}). Not incrementing further."
93
+ )
94
+
95
+ except Exception as e:
96
+ logger.error(
97
+ f"An error occurred during the scheduled key check: {str(e)}", exc_info=True
98
+ )
99
+
100
+
101
+ async def cleanup_expired_files():
102
+ """
103
+ 定时清理过期的文件记录
104
+ """
105
+ logger.info("Starting scheduled cleanup for expired files...")
106
+ try:
107
+ files_service = await get_files_service()
108
+ deleted_count = await files_service.cleanup_expired_files()
109
+
110
+ if deleted_count > 0:
111
+ logger.info(f"Successfully cleaned up {deleted_count} expired files.")
112
+ else:
113
+ logger.info("No expired files to clean up.")
114
+
115
+ except Exception as e:
116
+ logger.error(
117
+ f"An error occurred during the scheduled file cleanup: {str(e)}", exc_info=True
118
+ )
119
+
120
+
121
+ def setup_scheduler():
122
+ """设置并启动 APScheduler"""
123
+ scheduler = AsyncIOScheduler(timezone=str(settings.TIMEZONE)) # 从配置读取时区
124
+ # 添加检查失败密钥的定时任务
125
+ scheduler.add_job(
126
+ check_failed_keys,
127
+ "interval",
128
+ hours=settings.CHECK_INTERVAL_HOURS,
129
+ id="check_failed_keys_job",
130
+ name="Check Failed API Keys",
131
+ )
132
+ logger.info(
133
+ f"Key check job scheduled to run every {settings.CHECK_INTERVAL_HOURS} hour(s)."
134
+ )
135
+
136
+ # 新增:添加自动删除错误日志的定时任务,每天凌晨3点执行
137
+ scheduler.add_job(
138
+ delete_old_error_logs,
139
+ "cron",
140
+ hour=3,
141
+ minute=0,
142
+ id="delete_old_error_logs_job",
143
+ name="Delete Old Error Logs",
144
+ )
145
+ logger.info("Auto-delete error logs job scheduled to run daily at 3:00 AM.")
146
+
147
+ # 新增:添加自动删除请求日志的定时任务,每天凌晨3点05分执行
148
+ scheduler.add_job(
149
+ delete_old_request_logs_task,
150
+ "cron",
151
+ hour=3,
152
+ minute=5,
153
+ id="delete_old_request_logs_job",
154
+ name="Delete Old Request Logs",
155
+ )
156
+ logger.info(
157
+ f"Auto-delete request logs job scheduled to run daily at 3:05 AM, if enabled and AUTO_DELETE_REQUEST_LOGS_DAYS is set to {settings.AUTO_DELETE_REQUEST_LOGS_DAYS} days."
158
+ )
159
+
160
+ # 新增:添加文件过期清理的定时任务,每小时执行一次
161
+ if getattr(settings, 'FILES_CLEANUP_ENABLED', True):
162
+ cleanup_interval = getattr(settings, 'FILES_CLEANUP_INTERVAL_HOURS', 1)
163
+ scheduler.add_job(
164
+ cleanup_expired_files,
165
+ "interval",
166
+ hours=cleanup_interval,
167
+ id="cleanup_expired_files_job",
168
+ name="Cleanup Expired Files",
169
+ )
170
+ logger.info(
171
+ f"File cleanup job scheduled to run every {cleanup_interval} hour(s)."
172
+ )
173
+
174
+ scheduler.start()
175
+ logger.info("Scheduler started with all jobs.")
176
+ return scheduler
177
+
178
+
179
+ # 可以在这里添加一个全局的 scheduler 实例,以便在应用关闭时优雅地停止
180
+ scheduler_instance = None
181
+
182
+
183
+ def start_scheduler():
184
+ global scheduler_instance
185
+ if scheduler_instance is None or not scheduler_instance.running:
186
+ logger.info("Starting scheduler...")
187
+ scheduler_instance = setup_scheduler()
188
+ logger.info("Scheduler is already running.")
189
+
190
+
191
+ def stop_scheduler():
192
+ global scheduler_instance
193
+ if scheduler_instance and scheduler_instance.running:
194
+ scheduler_instance.shutdown()
195
+ logger.info("Scheduler stopped.")
app/service/chat/gemini_chat_service.py ADDED
@@ -0,0 +1,555 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app/services/chat_service.py
2
+
3
+ import datetime
4
+ import json
5
+ import re
6
+ import time
7
+ from typing import Any, AsyncGenerator, Dict, List
8
+
9
+ from app.config.config import settings
10
+ from app.core.constants import GEMINI_2_FLASH_EXP_SAFETY_SETTINGS
11
+ from app.database.services import add_error_log, add_request_log, get_file_api_key
12
+ from app.domain.gemini_models import GeminiRequest
13
+ from app.handler.response_handler import GeminiResponseHandler
14
+ from app.handler.stream_optimizer import gemini_optimizer
15
+ from app.log.logger import get_gemini_logger
16
+ from app.service.client.api_client import GeminiApiClient
17
+ from app.service.key.key_manager import KeyManager
18
+ from app.utils.helpers import redact_key_for_logging
19
+
20
+ logger = get_gemini_logger()
21
+
22
+
23
+ def _has_image_parts(contents: List[Dict[str, Any]]) -> bool:
24
+ """判断消息是否包含图片部分"""
25
+ for content in contents:
26
+ if "parts" in content:
27
+ for part in content["parts"]:
28
+ if "image_url" in part or "inline_data" in part:
29
+ return True
30
+ return False
31
+
32
+
33
+ def _extract_file_references(contents: List[Dict[str, Any]]) -> List[str]:
34
+ """從內容中提取文件引用"""
35
+ file_names = []
36
+ for content in contents:
37
+ if "parts" in content:
38
+ for part in content["parts"]:
39
+ if not isinstance(part, dict) or "fileData" not in part:
40
+ continue
41
+ file_data = part["fileData"]
42
+ if "fileUri" not in file_data:
43
+ continue
44
+ file_uri = file_data["fileUri"]
45
+ # 從 URI 中提取文件名
46
+ # 1. https://generativelanguage.googleapis.com/v1beta/files/{file_id}
47
+ match = re.match(
48
+ rf"{re.escape(settings.BASE_URL)}/(files/.*)", file_uri
49
+ )
50
+ if not match:
51
+ logger.warning(f"Invalid file URI: {file_uri}")
52
+ continue
53
+ file_id = match.group(1)
54
+ file_names.append(file_id)
55
+ logger.info(f"Found file reference: {file_id}")
56
+ return file_names
57
+
58
+
59
+ def _clean_json_schema_properties(obj: Any) -> Any:
60
+ """清理JSON Schema中Gemini API不支持的字段"""
61
+ if not isinstance(obj, dict):
62
+ return obj
63
+
64
+ # Gemini API不支持的JSON Schema字段
65
+ unsupported_fields = {
66
+ "exclusiveMaximum",
67
+ "exclusiveMinimum",
68
+ "const",
69
+ "examples",
70
+ "contentEncoding",
71
+ "contentMediaType",
72
+ "if",
73
+ "then",
74
+ "else",
75
+ "allOf",
76
+ "anyOf",
77
+ "oneOf",
78
+ "not",
79
+ "definitions",
80
+ "$schema",
81
+ "$id",
82
+ "$ref",
83
+ "$comment",
84
+ "readOnly",
85
+ "writeOnly",
86
+ }
87
+
88
+ cleaned = {}
89
+ for key, value in obj.items():
90
+ if key in unsupported_fields:
91
+ continue
92
+ if isinstance(value, dict):
93
+ cleaned[key] = _clean_json_schema_properties(value)
94
+ elif isinstance(value, list):
95
+ cleaned[key] = [_clean_json_schema_properties(item) for item in value]
96
+ else:
97
+ cleaned[key] = value
98
+
99
+ return cleaned
100
+
101
+
102
+ def _build_tools(model: str, payload: Dict[str, Any]) -> List[Dict[str, Any]]:
103
+ """构建工具"""
104
+
105
+ def _has_function_call(contents: List[Dict[str, Any]]) -> bool:
106
+ """检查内容中是否包含 functionCall"""
107
+ if not contents or not isinstance(contents, list):
108
+ return False
109
+ for content in contents:
110
+ if not content or not isinstance(content, dict) or "parts" not in content:
111
+ continue
112
+ parts = content.get("parts", [])
113
+ if not parts or not isinstance(parts, list):
114
+ continue
115
+ for part in parts:
116
+ if isinstance(part, dict) and "functionCall" in part:
117
+ return True
118
+ return False
119
+
120
+ def _merge_tools(tools: List[Dict[str, Any]]) -> Dict[str, Any]:
121
+ record = dict()
122
+ for item in tools:
123
+ if not item or not isinstance(item, dict):
124
+ continue
125
+
126
+ for k, v in item.items():
127
+ if k == "functionDeclarations" and v and isinstance(v, list):
128
+ functions = record.get("functionDeclarations", [])
129
+ # 清理每个函数声明中的不支持字段
130
+ cleaned_functions = []
131
+ for func in v:
132
+ if isinstance(func, dict):
133
+ cleaned_func = _clean_json_schema_properties(func)
134
+ cleaned_functions.append(cleaned_func)
135
+ else:
136
+ cleaned_functions.append(func)
137
+ functions.extend(cleaned_functions)
138
+ record["functionDeclarations"] = functions
139
+ else:
140
+ record[k] = v
141
+ return record
142
+
143
+ def _is_structured_output_request(payload: Dict[str, Any]) -> bool:
144
+ """检查请求是否要求结构化JSON输出"""
145
+ try:
146
+ generation_config = payload.get("generationConfig", {})
147
+ return generation_config.get("responseMimeType") == "application/json"
148
+ except (AttributeError, TypeError):
149
+ return False
150
+
151
+ tool = dict()
152
+ if payload and isinstance(payload, dict) and "tools" in payload:
153
+ if payload.get("tools") and isinstance(payload.get("tools"), dict):
154
+ payload["tools"] = [payload.get("tools")]
155
+ items = payload.get("tools", [])
156
+ if items and isinstance(items, list):
157
+ tool.update(_merge_tools(items))
158
+
159
+ # "Tool use with a response mime type: 'application/json' is unsupported"
160
+ # Gemini API限制:不支持同时使用tools和结构化输出(response_mime_type='application/json')
161
+ # 当请求指定了JSON响应格式时,跳过所有工具的添加以避免API错误
162
+ has_structured_output = _is_structured_output_request(payload)
163
+ if not has_structured_output:
164
+ if (
165
+ settings.TOOLS_CODE_EXECUTION_ENABLED
166
+ and not (model.endswith("-search") or "-thinking" in model)
167
+ and not _has_image_parts(payload.get("contents", []))
168
+ ):
169
+ tool["codeExecution"] = {}
170
+
171
+ if model.endswith("-search"):
172
+ tool["googleSearch"] = {}
173
+
174
+ real_model = _get_real_model(model)
175
+ if real_model in settings.URL_CONTEXT_MODELS and settings.URL_CONTEXT_ENABLED:
176
+ tool["urlContext"] = {}
177
+
178
+ # 解决 "Tool use with function calling is unsupported" 问题
179
+ if tool.get("functionDeclarations") or _has_function_call(
180
+ payload.get("contents", [])
181
+ ):
182
+ tool.pop("googleSearch", None)
183
+ tool.pop("codeExecution", None)
184
+ tool.pop("urlContext", None)
185
+
186
+ return [tool] if tool else []
187
+
188
+
189
+ def _get_real_model(model: str) -> str:
190
+ if model.endswith("-search"):
191
+ model = model[:-7]
192
+ if model.endswith("-image"):
193
+ model = model[:-6]
194
+ if model.endswith("-non-thinking"):
195
+ model = model[:-13]
196
+ if "-search" in model and "-non-thinking" in model:
197
+ model = model[:-20]
198
+ return model
199
+
200
+
201
+ def _get_safety_settings(model: str) -> List[Dict[str, str]]:
202
+ """获取安全设置"""
203
+ if model == "gemini-2.0-flash-exp":
204
+ return GEMINI_2_FLASH_EXP_SAFETY_SETTINGS
205
+ return settings.SAFETY_SETTINGS
206
+
207
+
208
+ def _filter_empty_parts(contents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
209
+ """Filters out contents with empty or invalid parts."""
210
+ if not contents:
211
+ return []
212
+
213
+ filtered_contents = []
214
+ for content in contents:
215
+ if (
216
+ not content
217
+ or "parts" not in content
218
+ or not isinstance(content.get("parts"), list)
219
+ ):
220
+ continue
221
+
222
+ valid_parts = [
223
+ part for part in content["parts"] if isinstance(part, dict) and part
224
+ ]
225
+
226
+ if valid_parts:
227
+ new_content = content.copy()
228
+ new_content["parts"] = valid_parts
229
+ filtered_contents.append(new_content)
230
+
231
+ return filtered_contents
232
+
233
+
234
+ def _build_payload(model: str, request: GeminiRequest) -> Dict[str, Any]:
235
+ """构建请求payload"""
236
+ request_dict = request.model_dump(exclude_none=False)
237
+ if request.generationConfig:
238
+ if request.generationConfig.maxOutputTokens is None:
239
+ # 如果未指定最大输出长度,则不传递该字段,解决截断的问题
240
+ if "maxOutputTokens" in request_dict["generationConfig"]:
241
+ request_dict["generationConfig"].pop("maxOutputTokens")
242
+
243
+ # 检查是否为TTS模型
244
+ is_tts_model = "tts" in model.lower()
245
+
246
+ if is_tts_model:
247
+ # TTS模型使用简化的payload,不包含tools和safetySettings
248
+ payload = {
249
+ "contents": _filter_empty_parts(request_dict.get("contents", [])),
250
+ "generationConfig": request_dict.get("generationConfig"),
251
+ }
252
+
253
+ # 只在有systemInstruction时才添加
254
+ if request_dict.get("systemInstruction"):
255
+ payload["systemInstruction"] = request_dict.get("systemInstruction")
256
+ else:
257
+ # 非TTS模型使用完整的payload
258
+ payload = {
259
+ "contents": _filter_empty_parts(request_dict.get("contents", [])),
260
+ "tools": _build_tools(model, request_dict),
261
+ "safetySettings": _get_safety_settings(model),
262
+ "generationConfig": request_dict.get("generationConfig"),
263
+ "systemInstruction": request_dict.get("systemInstruction"),
264
+ }
265
+
266
+ # 确保 generationConfig 不为 None
267
+ if payload["generationConfig"] is None:
268
+ payload["generationConfig"] = {}
269
+
270
+ if model.endswith("-image") or model.endswith("-image-generation"):
271
+ payload.pop("systemInstruction")
272
+ payload["generationConfig"]["responseModalities"] = ["Text", "Image"]
273
+
274
+ # 处理思考配置:优先使用客户端提供的配置,否则使���默认配置
275
+ client_thinking_config = None
276
+ if request.generationConfig and request.generationConfig.thinkingConfig:
277
+ client_thinking_config = request.generationConfig.thinkingConfig
278
+
279
+ if client_thinking_config is not None:
280
+ # 客户端提供了思考配置,直接使用
281
+ payload["generationConfig"]["thinkingConfig"] = client_thinking_config
282
+ else:
283
+ # 客户端没有提供思考配置,使用默认配置
284
+ if model.endswith("-non-thinking"):
285
+ if "gemini-2.5-pro" in model:
286
+ payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": 128}
287
+ else:
288
+ payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": 0}
289
+ elif _get_real_model(model) in settings.THINKING_BUDGET_MAP:
290
+ if settings.SHOW_THINKING_PROCESS:
291
+ payload["generationConfig"]["thinkingConfig"] = {
292
+ "thinkingBudget": settings.THINKING_BUDGET_MAP.get(model, 1000),
293
+ "includeThoughts": True,
294
+ }
295
+ else:
296
+ payload["generationConfig"]["thinkingConfig"] = {
297
+ "thinkingBudget": settings.THINKING_BUDGET_MAP.get(model, 1000)
298
+ }
299
+
300
+ return payload
301
+
302
+
303
+ class GeminiChatService:
304
+ """聊天服务"""
305
+
306
+ def __init__(self, base_url: str, key_manager: KeyManager):
307
+ self.api_client = GeminiApiClient(base_url, settings.TIME_OUT)
308
+ self.key_manager = key_manager
309
+ self.response_handler = GeminiResponseHandler()
310
+
311
+ def _extract_text_from_response(self, response: Dict[str, Any]) -> str:
312
+ """从响应中提取文本内容"""
313
+ if not response.get("candidates"):
314
+ return ""
315
+
316
+ candidate = response["candidates"][0]
317
+ content = candidate.get("content", {})
318
+ parts = content.get("parts", [])
319
+
320
+ if parts and "text" in parts[0]:
321
+ return parts[0].get("text", "")
322
+ return ""
323
+
324
+ def _create_char_response(
325
+ self, original_response: Dict[str, Any], text: str
326
+ ) -> Dict[str, Any]:
327
+ """创建包含指定文本的响应"""
328
+ response_copy = json.loads(json.dumps(original_response))
329
+ if response_copy.get("candidates") and response_copy["candidates"][0].get(
330
+ "content", {}
331
+ ).get("parts"):
332
+ response_copy["candidates"][0]["content"]["parts"][0]["text"] = text
333
+ return response_copy
334
+
335
+ async def generate_content(
336
+ self, model: str, request: GeminiRequest, api_key: str
337
+ ) -> Dict[str, Any]:
338
+ """生成内容"""
339
+ # 檢查並獲取文件專用的 API key(如果有文件)
340
+ file_names = _extract_file_references(request.model_dump().get("contents", []))
341
+ if file_names:
342
+ logger.info(f"Request contains file references: {file_names}")
343
+ file_api_key = await get_file_api_key(file_names[0])
344
+ if file_api_key:
345
+ logger.info(
346
+ f"Found API key for file {file_names[0]}: {redact_key_for_logging(file_api_key)}"
347
+ )
348
+ api_key = file_api_key # 使用文件的 API key
349
+ else:
350
+ logger.warning(
351
+ f"No API key found for file {file_names[0]}, using default key: {redact_key_for_logging(api_key)}"
352
+ )
353
+
354
+ payload = _build_payload(model, request)
355
+ start_time = time.perf_counter()
356
+ request_datetime = datetime.datetime.now()
357
+ is_success = False
358
+ status_code = None
359
+ response = None
360
+
361
+ try:
362
+ response = await self.api_client.generate_content(payload, model, api_key)
363
+ is_success = True
364
+ status_code = 200
365
+ return self.response_handler.handle_response(response, model, stream=False)
366
+ except Exception as e:
367
+ is_success = False
368
+ error_log_msg = str(e)
369
+ logger.error(f"Normal API call failed with error: {error_log_msg}")
370
+ match = re.search(r"status code (\d+)", error_log_msg)
371
+ if match:
372
+ status_code = int(match.group(1))
373
+ else:
374
+ status_code = 500
375
+
376
+ await add_error_log(
377
+ gemini_key=api_key,
378
+ model_name=model,
379
+ error_type="gemini-chat-non-stream",
380
+ error_log=error_log_msg,
381
+ error_code=status_code,
382
+ request_msg=payload,
383
+ request_datetime=request_datetime,
384
+ )
385
+ raise e
386
+ finally:
387
+ end_time = time.perf_counter()
388
+ latency_ms = int((end_time - start_time) * 1000)
389
+ await add_request_log(
390
+ model_name=model,
391
+ api_key=api_key,
392
+ is_success=is_success,
393
+ status_code=status_code,
394
+ latency_ms=latency_ms,
395
+ request_time=request_datetime,
396
+ )
397
+
398
+ async def count_tokens(
399
+ self, model: str, request: GeminiRequest, api_key: str
400
+ ) -> Dict[str, Any]:
401
+ """计算token数量"""
402
+ # countTokens API只需要contents
403
+ payload = {
404
+ "contents": _filter_empty_parts(request.model_dump().get("contents", []))
405
+ }
406
+ start_time = time.perf_counter()
407
+ request_datetime = datetime.datetime.now()
408
+ is_success = False
409
+ status_code = None
410
+ response = None
411
+
412
+ try:
413
+ response = await self.api_client.count_tokens(payload, model, api_key)
414
+ is_success = True
415
+ status_code = 200
416
+ return response
417
+ except Exception as e:
418
+ is_success = False
419
+ error_log_msg = str(e)
420
+ logger.error(f"Count tokens API call failed with error: {error_log_msg}")
421
+ match = re.search(r"status code (\d+)", error_log_msg)
422
+ if match:
423
+ status_code = int(match.group(1))
424
+ else:
425
+ status_code = 500
426
+
427
+ await add_error_log(
428
+ gemini_key=api_key,
429
+ model_name=model,
430
+ error_type="gemini-count-tokens",
431
+ error_log=error_log_msg,
432
+ error_code=status_code,
433
+ request_msg=payload,
434
+ )
435
+ raise e
436
+ finally:
437
+ end_time = time.perf_counter()
438
+ latency_ms = int((end_time - start_time) * 1000)
439
+ await add_request_log(
440
+ model_name=model,
441
+ api_key=api_key,
442
+ is_success=is_success,
443
+ status_code=status_code,
444
+ latency_ms=latency_ms,
445
+ request_time=request_datetime,
446
+ )
447
+
448
+ async def stream_generate_content(
449
+ self, model: str, request: GeminiRequest, api_key: str
450
+ ) -> AsyncGenerator[str, None]:
451
+ """流式生成内容"""
452
+ # 檢查並獲取文件專用的 API key(如果有文件)
453
+ file_names = _extract_file_references(request.model_dump().get("contents", []))
454
+ if file_names:
455
+ logger.info(f"Request contains file references: {file_names}")
456
+ file_api_key = await get_file_api_key(file_names[0])
457
+ if file_api_key:
458
+ logger.info(
459
+ f"Found API key for file {file_names[0]}: {redact_key_for_logging(file_api_key)}"
460
+ )
461
+ api_key = file_api_key # 使用文件的 API key
462
+ else:
463
+ logger.warning(
464
+ f"No API key found for file {file_names[0]}, using default key: {redact_key_for_logging(api_key)}"
465
+ )
466
+
467
+ retries = 0
468
+ max_retries = settings.MAX_RETRIES
469
+ payload = _build_payload(model, request)
470
+ is_success = False
471
+ status_code = None
472
+ final_api_key = api_key
473
+
474
+ while retries < max_retries:
475
+ request_datetime = datetime.datetime.now()
476
+ start_time = time.perf_counter()
477
+ current_attempt_key = api_key
478
+ final_api_key = current_attempt_key
479
+ try:
480
+ async for line in self.api_client.stream_generate_content(
481
+ payload, model, current_attempt_key
482
+ ):
483
+ # print(line)
484
+ if line.startswith("data:"):
485
+ line = line[6:]
486
+ response_data = self.response_handler.handle_response(
487
+ json.loads(line), model, stream=True
488
+ )
489
+ text = self._extract_text_from_response(response_data)
490
+ # 如果有文本内容,且开启了流式输出优化器,则使用流式输出优化器处理
491
+ if text and settings.STREAM_OPTIMIZER_ENABLED:
492
+ # 使用流式输出优化器处理文本输出
493
+ async for (
494
+ optimized_chunk
495
+ ) in gemini_optimizer.optimize_stream_output(
496
+ text,
497
+ lambda t: self._create_char_response(response_data, t),
498
+ lambda c: "data: " + json.dumps(c) + "\n\n",
499
+ ):
500
+ yield optimized_chunk
501
+ else:
502
+ # 如果没有文本内容(如工具调用等),整块输出
503
+ yield "data: " + json.dumps(response_data) + "\n\n"
504
+ logger.info("Streaming completed successfully")
505
+ is_success = True
506
+ status_code = 200
507
+ break
508
+ except Exception as e:
509
+ retries += 1
510
+ is_success = False
511
+ error_log_msg = str(e)
512
+ logger.warning(
513
+ f"Streaming API call failed with error: {error_log_msg}. Attempt {retries} of {max_retries}"
514
+ )
515
+ match = re.search(r"status code (\d+)", error_log_msg)
516
+ if match:
517
+ status_code = int(match.group(1))
518
+ else:
519
+ status_code = 500
520
+
521
+ await add_error_log(
522
+ gemini_key=current_attempt_key,
523
+ model_name=model,
524
+ error_type="gemini-chat-stream",
525
+ error_log=error_log_msg,
526
+ error_code=status_code,
527
+ request_msg=payload,
528
+ request_datetime=request_datetime,
529
+ )
530
+
531
+ api_key = await self.key_manager.handle_api_failure(
532
+ current_attempt_key, retries
533
+ )
534
+ if api_key:
535
+ logger.info(
536
+ f"Switched to new API key: {redact_key_for_logging(api_key)}"
537
+ )
538
+ else:
539
+ logger.error(f"No valid API key available after {retries} retries.")
540
+ break
541
+
542
+ if retries >= max_retries:
543
+ logger.error(f"Max retries ({max_retries}) reached for streaming.")
544
+ break
545
+ finally:
546
+ end_time = time.perf_counter()
547
+ latency_ms = int((end_time - start_time) * 1000)
548
+ await add_request_log(
549
+ model_name=model,
550
+ api_key=final_api_key,
551
+ is_success=is_success,
552
+ status_code=status_code,
553
+ latency_ms=latency_ms,
554
+ request_time=request_datetime,
555
+ )
app/service/chat/openai_chat_service.py ADDED
@@ -0,0 +1,745 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app/services/chat_service.py
2
+
3
+ import asyncio
4
+ import datetime
5
+ import json
6
+ import re
7
+ import time
8
+ from copy import deepcopy
9
+ from typing import Any, AsyncGenerator, Dict, List, Optional, Union
10
+
11
+ from app.config.config import settings
12
+ from app.core.constants import GEMINI_2_FLASH_EXP_SAFETY_SETTINGS
13
+ from app.database.services import (
14
+ add_error_log,
15
+ add_request_log,
16
+ )
17
+ from app.domain.openai_models import ChatRequest, ImageGenerationRequest
18
+ from app.handler.message_converter import OpenAIMessageConverter
19
+ from app.handler.response_handler import OpenAIResponseHandler
20
+ from app.handler.stream_optimizer import openai_optimizer
21
+ from app.log.logger import get_openai_logger
22
+ from app.service.client.api_client import GeminiApiClient
23
+ from app.service.image.image_create_service import ImageCreateService
24
+ from app.service.key.key_manager import KeyManager
25
+
26
+ logger = get_openai_logger()
27
+
28
+
29
+ def _has_media_parts(messages: List[Dict[str, Any]]) -> bool:
30
+ """判断消息是否包含多媒体部分"""
31
+ for message in messages:
32
+ if "parts" in message:
33
+ for part in message["parts"]:
34
+ if "image_url" in part or "inline_data" in part:
35
+ return True
36
+ return False
37
+
38
+
39
+ def _clean_json_schema_properties(obj: Any) -> Any:
40
+ """清理JSON Schema中Gemini API不支持的字段"""
41
+ if not isinstance(obj, dict):
42
+ return obj
43
+
44
+ # Gemini API不支持的JSON Schema字段
45
+ unsupported_fields = {
46
+ "exclusiveMaximum",
47
+ "exclusiveMinimum",
48
+ "const",
49
+ "examples",
50
+ "contentEncoding",
51
+ "contentMediaType",
52
+ "if",
53
+ "then",
54
+ "else",
55
+ "allOf",
56
+ "anyOf",
57
+ "oneOf",
58
+ "not",
59
+ "definitions",
60
+ "$schema",
61
+ "$id",
62
+ "$ref",
63
+ "$comment",
64
+ "readOnly",
65
+ "writeOnly",
66
+ }
67
+
68
+ cleaned = {}
69
+ for key, value in obj.items():
70
+ if key in unsupported_fields:
71
+ continue
72
+ if isinstance(value, dict):
73
+ cleaned[key] = _clean_json_schema_properties(value)
74
+ elif isinstance(value, list):
75
+ cleaned[key] = [_clean_json_schema_properties(item) for item in value]
76
+ else:
77
+ cleaned[key] = value
78
+
79
+ return cleaned
80
+
81
+
82
+ def _build_tools(
83
+ request: ChatRequest, messages: List[Dict[str, Any]]
84
+ ) -> List[Dict[str, Any]]:
85
+ """构建工具"""
86
+ tool = dict()
87
+ model = request.model
88
+
89
+ if (
90
+ settings.TOOLS_CODE_EXECUTION_ENABLED
91
+ and not (
92
+ model.endswith("-search")
93
+ or "-thinking" in model
94
+ or model.endswith("-image")
95
+ or model.endswith("-image-generation")
96
+ )
97
+ and not _has_media_parts(messages)
98
+ ):
99
+ tool["codeExecution"] = {}
100
+ logger.debug("Code execution tool enabled.")
101
+ elif _has_media_parts(messages):
102
+ logger.debug("Code execution tool disabled due to media parts presence.")
103
+
104
+ if model.endswith("-search"):
105
+ tool["googleSearch"] = {}
106
+
107
+ real_model = _get_real_model(model)
108
+ if real_model in settings.URL_CONTEXT_MODELS and settings.URL_CONTEXT_ENABLED:
109
+ tool["urlContext"] = {}
110
+
111
+ # 将 request 中的 tools 合并到 tools 中
112
+ if request.tools:
113
+ function_declarations = []
114
+ for item in request.tools:
115
+ if not item or not isinstance(item, dict):
116
+ continue
117
+
118
+ if item.get("type", "") == "function" and item.get("function"):
119
+ function = deepcopy(item.get("function"))
120
+ parameters = function.get("parameters", {})
121
+ if parameters.get("type") == "object" and not parameters.get(
122
+ "properties", {}
123
+ ):
124
+ function.pop("parameters", None)
125
+
126
+ # 清理函数中的不支持字段
127
+ function = _clean_json_schema_properties(function)
128
+ function_declarations.append(function)
129
+
130
+ if function_declarations:
131
+ # 按照 function 的 name 去重
132
+ names, functions = set(), []
133
+ for fc in function_declarations:
134
+ if fc.get("name") not in names:
135
+ if fc.get("name") == "googleSearch":
136
+ # cherry开启内置搜索时,添加googleSearch工具
137
+ tool["googleSearch"] = {}
138
+ else:
139
+ # 其他函数,添加到functionDeclarations中
140
+ names.add(fc.get("name"))
141
+ functions.append(fc)
142
+
143
+ tool["functionDeclarations"] = functions
144
+
145
+ # 解决 "Tool use with function calling is unsupported" 问题
146
+ if tool.get("functionDeclarations"):
147
+ tool.pop("googleSearch", None)
148
+ tool.pop("codeExecution", None)
149
+ tool.pop("urlContext", None)
150
+
151
+ return [tool] if tool else []
152
+
153
+
154
+ def _get_real_model(model: str) -> str:
155
+ if model.endswith("-search"):
156
+ model = model[:-7]
157
+ if model.endswith("-image"):
158
+ model = model[:-6]
159
+ if model.endswith("-non-thinking"):
160
+ model = model[:-13]
161
+ if "-search" in model and "-non-thinking" in model:
162
+ model = model[:-20]
163
+ return model
164
+
165
+
166
+ def _get_safety_settings(model: str) -> List[Dict[str, str]]:
167
+ """获取安全设置"""
168
+ # if (
169
+ # "2.0" in model
170
+ # and "gemini-2.0-flash-thinking-exp" not in model
171
+ # and "gemini-2.0-pro-exp" not in model
172
+ # ):
173
+ if model == "gemini-2.0-flash-exp":
174
+ return GEMINI_2_FLASH_EXP_SAFETY_SETTINGS
175
+ return settings.SAFETY_SETTINGS
176
+
177
+
178
+ def _validate_and_set_max_tokens(
179
+ payload: Dict[str, Any], max_tokens: Optional[int], logger_instance
180
+ ) -> None:
181
+ """验证并设置 max_tokens 参数"""
182
+ if max_tokens is None:
183
+ return
184
+
185
+ # 参数验证和处理
186
+ if max_tokens <= 0:
187
+ logger_instance.warning(
188
+ f"Invalid max_tokens value: {max_tokens}, will not set maxOutputTokens"
189
+ )
190
+ # 不设置 maxOutputTokens,让 Gemini API 使用默认值
191
+ else:
192
+ payload["generationConfig"]["maxOutputTokens"] = max_tokens
193
+
194
+
195
+ def _build_payload(
196
+ request: ChatRequest,
197
+ messages: List[Dict[str, Any]],
198
+ instruction: Optional[Dict[str, Any]] = None,
199
+ ) -> Dict[str, Any]:
200
+ """构建请求payload"""
201
+ payload = {
202
+ "contents": messages,
203
+ "generationConfig": {
204
+ "temperature": request.temperature,
205
+ "stopSequences": request.stop,
206
+ "topP": request.top_p,
207
+ "topK": request.top_k,
208
+ },
209
+ "tools": _build_tools(request, messages),
210
+ "safetySettings": _get_safety_settings(request.model),
211
+ }
212
+
213
+ # 处理 max_tokens 参数
214
+ _validate_and_set_max_tokens(payload, request.max_tokens, logger)
215
+
216
+ # 处理 n 参数
217
+ if request.n is not None and request.n > 0:
218
+ payload["generationConfig"]["candidateCount"] = request.n
219
+
220
+ if request.model.endswith("-image") or request.model.endswith("-image-generation"):
221
+ payload["generationConfig"]["responseModalities"] = ["Text", "Image"]
222
+
223
+ if request.model.endswith("-non-thinking"):
224
+ if "gemini-2.5-pro" in request.model:
225
+ payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": 128}
226
+ else:
227
+ payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": 0}
228
+
229
+ elif _get_real_model(request.model) in settings.THINKING_BUDGET_MAP:
230
+ if settings.SHOW_THINKING_PROCESS:
231
+ payload["generationConfig"]["thinkingConfig"] = {
232
+ "thinkingBudget": settings.THINKING_BUDGET_MAP.get(request.model, 1000),
233
+ "includeThoughts": True,
234
+ }
235
+ else:
236
+ payload["generationConfig"]["thinkingConfig"] = {
237
+ "thinkingBudget": settings.THINKING_BUDGET_MAP.get(request.model, 1000)
238
+ }
239
+
240
+ if (
241
+ instruction
242
+ and isinstance(instruction, dict)
243
+ and instruction.get("role") == "system"
244
+ and instruction.get("parts")
245
+ and not request.model.endswith("-image")
246
+ and not request.model.endswith("-image-generation")
247
+ ):
248
+ payload["systemInstruction"] = instruction
249
+
250
+ return payload
251
+
252
+
253
+ class OpenAIChatService:
254
+ """聊天服务"""
255
+
256
+ def __init__(self, base_url: str, key_manager: KeyManager = None):
257
+ self.message_converter = OpenAIMessageConverter()
258
+ self.response_handler = OpenAIResponseHandler(config=None)
259
+ self.api_client = GeminiApiClient(base_url, settings.TIME_OUT)
260
+ self.key_manager = key_manager
261
+ self.image_create_service = ImageCreateService()
262
+
263
+ def _extract_text_from_openai_chunk(self, chunk: Dict[str, Any]) -> str:
264
+ """从OpenAI响应块中提取文本内容"""
265
+ if not chunk.get("choices"):
266
+ return ""
267
+
268
+ choice = chunk["choices"][0]
269
+ if "delta" in choice and "content" in choice["delta"]:
270
+ return choice["delta"]["content"]
271
+ return ""
272
+
273
+ def _create_char_openai_chunk(
274
+ self, original_chunk: Dict[str, Any], text: str
275
+ ) -> Dict[str, Any]:
276
+ """创建包含指定文本的OpenAI响应块"""
277
+ chunk_copy = json.loads(json.dumps(original_chunk))
278
+ if chunk_copy.get("choices") and "delta" in chunk_copy["choices"][0]:
279
+ chunk_copy["choices"][0]["delta"]["content"] = text
280
+ return chunk_copy
281
+
282
+ async def create_chat_completion(
283
+ self,
284
+ request: ChatRequest,
285
+ api_key: str,
286
+ ) -> Union[Dict[str, Any], AsyncGenerator[str, None]]:
287
+ """创建聊天完成"""
288
+ messages, instruction = self.message_converter.convert(
289
+ request.messages, request.model
290
+ )
291
+
292
+ payload = _build_payload(request, messages, instruction)
293
+
294
+ if request.stream:
295
+ return self._handle_stream_completion(request.model, payload, api_key)
296
+ return await self._handle_normal_completion(request.model, payload, api_key)
297
+
298
+ async def _handle_normal_completion(
299
+ self, model: str, payload: Dict[str, Any], api_key: str
300
+ ) -> Dict[str, Any]:
301
+ """处理普通聊天完成"""
302
+ start_time = time.perf_counter()
303
+ request_datetime = datetime.datetime.now()
304
+ is_success = False
305
+ status_code = None
306
+ response = None
307
+
308
+ try:
309
+ response = await self.api_client.generate_content(payload, model, api_key)
310
+ usage_metadata = response.get("usageMetadata", {})
311
+ is_success = True
312
+ status_code = 200
313
+
314
+ # 尝试处理响应,捕获可能的响应处理异常
315
+ try:
316
+ result = self.response_handler.handle_response(
317
+ response,
318
+ model,
319
+ stream=False,
320
+ finish_reason="stop",
321
+ usage_metadata=usage_metadata,
322
+ )
323
+ return result
324
+ except Exception as response_error:
325
+ logger.error(
326
+ f"Response processing failed for model {model}: {str(response_error)}"
327
+ )
328
+
329
+ # 记录详细的错误信息
330
+ if "parts" in str(response_error):
331
+ logger.error("Response structure issue - missing or invalid parts")
332
+ if response.get("candidates"):
333
+ candidate = response["candidates"][0]
334
+ content = candidate.get("content", {})
335
+ logger.error(f"Content structure: {content}")
336
+
337
+ # 重新抛出异常
338
+ raise response_error
339
+
340
+ except Exception as e:
341
+ is_success = False
342
+ error_log_msg = str(e)
343
+ logger.error(f"API call failed for model {model}: {error_log_msg}")
344
+
345
+ # 特别记录 max_tokens 相关的错误
346
+ gen_config = payload.get("generationConfig", {})
347
+ if "maxOutputTokens" in gen_config:
348
+ logger.error(
349
+ f"Request had maxOutputTokens: {gen_config['maxOutputTokens']}"
350
+ )
351
+
352
+ # 如果是响应处理错误,记录更多信息
353
+ if "parts" in error_log_msg:
354
+ logger.error("This is likely a response processing error")
355
+
356
+ match = re.search(r"status code (\d+)", error_log_msg)
357
+ status_code = int(match.group(1)) if match else 500
358
+
359
+ await add_error_log(
360
+ gemini_key=api_key,
361
+ model_name=model,
362
+ error_type="openai-chat-non-stream",
363
+ error_log=error_log_msg,
364
+ error_code=status_code,
365
+ request_msg=payload,
366
+ request_datetime=request_datetime,
367
+ )
368
+ raise e
369
+ finally:
370
+ end_time = time.perf_counter()
371
+ latency_ms = int((end_time - start_time) * 1000)
372
+ logger.info(
373
+ f"Normal completion finished - Success: {is_success}, Latency: {latency_ms}ms"
374
+ )
375
+
376
+ await add_request_log(
377
+ model_name=model,
378
+ api_key=api_key,
379
+ is_success=is_success,
380
+ status_code=status_code,
381
+ latency_ms=latency_ms,
382
+ request_time=request_datetime,
383
+ )
384
+
385
+ async def _fake_stream_logic_impl(
386
+ self, model: str, payload: Dict[str, Any], api_key: str
387
+ ) -> AsyncGenerator[str, None]:
388
+ """处理伪流式 (fake stream) 的核心逻辑"""
389
+ logger.info(
390
+ f"Fake streaming enabled for model: {model}. Calling non-streaming endpoint."
391
+ )
392
+
393
+ api_response_task = asyncio.create_task(
394
+ self.api_client.generate_content(payload, model, api_key)
395
+ )
396
+
397
+ i = 0
398
+ try:
399
+ while not api_response_task.done():
400
+ i = i + 1
401
+ """定期发送空数据以保持连接"""
402
+ if i >= settings.FAKE_STREAM_EMPTY_DATA_INTERVAL_SECONDS:
403
+ i = 0
404
+ empty_chunk = self.response_handler.handle_response(
405
+ {},
406
+ model,
407
+ stream=True,
408
+ finish_reason="stop",
409
+ usage_metadata=None,
410
+ )
411
+ yield f"data: {json.dumps(empty_chunk)}\n\n"
412
+ logger.debug("Sent empty data chunk for fake stream heartbeat.")
413
+ await asyncio.sleep(1)
414
+ finally:
415
+ response = await api_response_task
416
+
417
+ if response and response.get("candidates"):
418
+ response = self.response_handler.handle_response(
419
+ response,
420
+ model,
421
+ stream=True,
422
+ finish_reason="stop",
423
+ usage_metadata=response.get("usageMetadata", {}),
424
+ )
425
+ yield f"data: {json.dumps(response)}\n\n"
426
+ logger.info(f"Sent full response content for fake stream: {model}")
427
+ else:
428
+ error_message = "Failed to get response from model"
429
+ if response and isinstance(response, dict) and response.get("error"):
430
+ error_details = response.get("error")
431
+ if isinstance(error_details, dict):
432
+ error_message = error_details.get("message", error_message)
433
+
434
+ logger.error(
435
+ f"No candidates or error in response for fake stream model {model}: {response}"
436
+ )
437
+ error_chunk = self.response_handler.handle_response(
438
+ {}, model, stream=True, finish_reason="stop", usage_metadata=None
439
+ )
440
+ yield f"data: {json.dumps(error_chunk)}\n\n"
441
+
442
+ async def _real_stream_logic_impl(
443
+ self, model: str, payload: Dict[str, Any], api_key: str
444
+ ) -> AsyncGenerator[str, None]:
445
+ """处理真实流式 (real stream) 的核心逻辑"""
446
+ tool_call_flag = False
447
+ usage_metadata = None
448
+ async for line in self.api_client.stream_generate_content(
449
+ payload, model, api_key
450
+ ):
451
+ if line.startswith("data:"):
452
+ chunk_str = line[6:]
453
+ if not chunk_str or chunk_str.isspace():
454
+ logger.debug(
455
+ f"Received empty data line for model {model}, skipping."
456
+ )
457
+ continue
458
+ try:
459
+ chunk = json.loads(chunk_str)
460
+ usage_metadata = chunk.get("usageMetadata", {})
461
+ except json.JSONDecodeError:
462
+ logger.error(
463
+ f"Failed to decode JSON from stream for model {model}: {chunk_str}"
464
+ )
465
+ continue
466
+ openai_chunk = self.response_handler.handle_response(
467
+ chunk,
468
+ model,
469
+ stream=True,
470
+ finish_reason=None,
471
+ usage_metadata=usage_metadata,
472
+ )
473
+ if openai_chunk:
474
+ text = self._extract_text_from_openai_chunk(openai_chunk)
475
+ if text and settings.STREAM_OPTIMIZER_ENABLED:
476
+ async for (
477
+ optimized_chunk_data
478
+ ) in openai_optimizer.optimize_stream_output(
479
+ text,
480
+ lambda t: self._create_char_openai_chunk(openai_chunk, t),
481
+ lambda c: f"data: {json.dumps(c)}\n\n",
482
+ ):
483
+ yield optimized_chunk_data
484
+ else:
485
+ if openai_chunk.get("choices") and openai_chunk["choices"][
486
+ 0
487
+ ].get("delta", {}).get("tool_calls"):
488
+ tool_call_flag = True
489
+
490
+ yield f"data: {json.dumps(openai_chunk)}\n\n"
491
+
492
+ if tool_call_flag:
493
+ yield f"data: {json.dumps(self.response_handler.handle_response({}, model, stream=True, finish_reason='tool_calls', usage_metadata=usage_metadata))}\n\n"
494
+ else:
495
+ yield f"data: {json.dumps(self.response_handler.handle_response({}, model, stream=True, finish_reason='stop', usage_metadata=usage_metadata))}\n\n"
496
+
497
+ async def _handle_stream_completion(
498
+ self, model: str, payload: Dict[str, Any], api_key: str
499
+ ) -> AsyncGenerator[str, None]:
500
+ """处理流式聊天完成,添加重试逻辑和假流式支持"""
501
+ retries = 0
502
+ max_retries = settings.MAX_RETRIES
503
+ is_success = False
504
+ status_code = None
505
+ final_api_key = api_key
506
+
507
+ while retries < max_retries:
508
+ start_time = time.perf_counter()
509
+ request_datetime = datetime.datetime.now()
510
+ current_attempt_key = final_api_key
511
+
512
+ try:
513
+ stream_generator = None
514
+ if settings.FAKE_STREAM_ENABLED:
515
+ logger.info(
516
+ f"Using fake stream logic for model: {model}, Attempt: {retries + 1}"
517
+ )
518
+ stream_generator = self._fake_stream_logic_impl(
519
+ model, payload, current_attempt_key
520
+ )
521
+ else:
522
+ logger.info(
523
+ f"Using real stream logic for model: {model}, Attempt: {retries + 1}"
524
+ )
525
+ stream_generator = self._real_stream_logic_impl(
526
+ model, payload, current_attempt_key
527
+ )
528
+
529
+ async for chunk_data in stream_generator:
530
+ yield chunk_data
531
+
532
+ yield "data: [DONE]\n\n"
533
+ logger.info(
534
+ f"Streaming completed successfully for model: {model}, FakeStream: {settings.FAKE_STREAM_ENABLED}, Attempt: {retries + 1}"
535
+ )
536
+ is_success = True
537
+ status_code = 200
538
+ break
539
+
540
+ except Exception as e:
541
+ retries += 1
542
+ is_success = False
543
+ error_log_msg = str(e)
544
+ logger.warning(
545
+ f"Streaming API call failed with error: {error_log_msg}. Attempt {retries} of {max_retries} with key {current_attempt_key}"
546
+ )
547
+
548
+ match = re.search(r"status code (\d+)", error_log_msg)
549
+ if match:
550
+ status_code = int(match.group(1))
551
+ else:
552
+ if isinstance(e, asyncio.TimeoutError):
553
+ status_code = 408
554
+ else:
555
+ status_code = 500
556
+
557
+ await add_error_log(
558
+ gemini_key=current_attempt_key,
559
+ model_name=model,
560
+ error_type="openai-chat-stream",
561
+ error_log=error_log_msg,
562
+ error_code=status_code,
563
+ request_msg=payload,
564
+ request_datetime=request_datetime,
565
+ )
566
+
567
+ if self.key_manager:
568
+ new_api_key = await self.key_manager.handle_api_failure(
569
+ current_attempt_key, retries
570
+ )
571
+ if new_api_key and new_api_key != current_attempt_key:
572
+ final_api_key = new_api_key
573
+ logger.info(
574
+ f"Switched to new API key for next attempt: {final_api_key}"
575
+ )
576
+ elif not new_api_key:
577
+ logger.error(
578
+ f"No valid API key available after {retries} retries, ceasing attempts for this request."
579
+ )
580
+ break
581
+ else:
582
+ logger.error(
583
+ "KeyManager not available, cannot switch API key. Ceasing attempts for this request."
584
+ )
585
+ break
586
+
587
+ if retries >= max_retries:
588
+ logger.error(
589
+ f"Max retries ({max_retries}) reached for streaming model {model}."
590
+ )
591
+ finally:
592
+ end_time = time.perf_counter()
593
+ latency_ms = int((end_time - start_time) * 1000)
594
+ await add_request_log(
595
+ model_name=model,
596
+ api_key=current_attempt_key,
597
+ is_success=is_success,
598
+ status_code=status_code,
599
+ latency_ms=latency_ms,
600
+ request_time=request_datetime,
601
+ )
602
+
603
+ if not is_success:
604
+ logger.error(
605
+ f"Streaming failed permanently for model {model} after {retries} attempts."
606
+ )
607
+ yield f"data: {json.dumps({'error': f'Streaming failed after {retries} retries.'})}\n\n"
608
+ yield "data: [DONE]\n\n"
609
+
610
+ async def create_image_chat_completion(
611
+ self, request: ChatRequest, api_key: str
612
+ ) -> Union[Dict[str, Any], AsyncGenerator[str, None]]:
613
+
614
+ image_generate_request = ImageGenerationRequest()
615
+ image_generate_request.prompt = request.messages[-1]["content"]
616
+ image_res = self.image_create_service.generate_images_chat(
617
+ image_generate_request
618
+ )
619
+
620
+ if request.stream:
621
+ return self._handle_stream_image_completion(
622
+ request.model, image_res, api_key
623
+ )
624
+ else:
625
+ return await self._handle_normal_image_completion(
626
+ request.model, image_res, api_key
627
+ )
628
+
629
+ async def _handle_stream_image_completion(
630
+ self, model: str, image_data: str, api_key: str
631
+ ) -> AsyncGenerator[str, None]:
632
+ logger.info(f"Starting stream image completion for model: {model}")
633
+ start_time = time.perf_counter()
634
+ request_datetime = datetime.datetime.now()
635
+ is_success = False
636
+ status_code = None
637
+
638
+ try:
639
+ if image_data:
640
+ openai_chunk = self.response_handler.handle_image_chat_response(
641
+ image_data, model, stream=True, finish_reason=None
642
+ )
643
+ if openai_chunk:
644
+ # 提取文本内容
645
+ text = self._extract_text_from_openai_chunk(openai_chunk)
646
+ if text:
647
+ # 使用流式输出优化器处理文本输出
648
+ async for (
649
+ optimized_chunk
650
+ ) in openai_optimizer.optimize_stream_output(
651
+ text,
652
+ lambda t: self._create_char_openai_chunk(openai_chunk, t),
653
+ lambda c: f"data: {json.dumps(c)}\n\n",
654
+ ):
655
+ yield optimized_chunk
656
+ else:
657
+ # 如果没有文本内容(如图片URL等),整块输出
658
+ yield f"data: {json.dumps(openai_chunk)}\n\n"
659
+ yield f"data: {json.dumps(self.response_handler.handle_response({}, model, stream=True, finish_reason='stop'))}\n\n"
660
+ logger.info(
661
+ f"Stream image completion finished successfully for model: {model}"
662
+ )
663
+ is_success = True
664
+ status_code = 200
665
+ yield "data: [DONE]\n\n"
666
+ except Exception as e:
667
+ is_success = False
668
+ error_log_msg = f"Stream image completion failed for model {model}: {e}"
669
+ logger.error(error_log_msg)
670
+ status_code = 500
671
+ await add_error_log(
672
+ gemini_key=api_key,
673
+ model_name=model,
674
+ error_type="openai-image-stream",
675
+ error_log=error_log_msg,
676
+ error_code=status_code,
677
+ request_msg={"image_data_truncated": image_data[:1000]},
678
+ request_datetime=request_datetime,
679
+ )
680
+ yield f"data: {json.dumps({'error': error_log_msg})}\n\n"
681
+ yield "data: [DONE]\n\n"
682
+ finally:
683
+ end_time = time.perf_counter()
684
+ latency_ms = int((end_time - start_time) * 1000)
685
+ logger.info(
686
+ f"Stream image completion for model {model} took {latency_ms} ms. Success: {is_success}"
687
+ )
688
+ await add_request_log(
689
+ model_name=model,
690
+ api_key=api_key,
691
+ is_success=is_success,
692
+ status_code=status_code,
693
+ latency_ms=latency_ms,
694
+ request_time=request_datetime,
695
+ )
696
+
697
+ async def _handle_normal_image_completion(
698
+ self, model: str, image_data: str, api_key: str
699
+ ) -> Dict[str, Any]:
700
+ logger.info(f"Starting normal image completion for model: {model}")
701
+ start_time = time.perf_counter()
702
+ request_datetime = datetime.datetime.now()
703
+ is_success = False
704
+ status_code = None
705
+ result = None
706
+
707
+ try:
708
+ result = self.response_handler.handle_image_chat_response(
709
+ image_data, model, stream=False, finish_reason="stop"
710
+ )
711
+ logger.info(
712
+ f"Normal image completion finished successfully for model: {model}"
713
+ )
714
+ is_success = True
715
+ status_code = 200
716
+ return result
717
+ except Exception as e:
718
+ is_success = False
719
+ error_log_msg = f"Normal image completion failed for model {model}: {e}"
720
+ logger.error(error_log_msg)
721
+ status_code = 500
722
+ await add_error_log(
723
+ gemini_key=api_key,
724
+ model_name=model,
725
+ error_type="openai-image-non-stream",
726
+ error_log=error_log_msg,
727
+ error_code=status_code,
728
+ request_msg={"image_data_truncated": image_data[:1000]},
729
+ request_datetime=request_datetime,
730
+ )
731
+ raise e
732
+ finally:
733
+ end_time = time.perf_counter()
734
+ latency_ms = int((end_time - start_time) * 1000)
735
+ logger.info(
736
+ f"Normal image completion for model {model} took {latency_ms} ms. Success: {is_success}"
737
+ )
738
+ await add_request_log(
739
+ model_name=model,
740
+ api_key=api_key,
741
+ is_success=is_success,
742
+ status_code=status_code,
743
+ latency_ms=latency_ms,
744
+ request_time=request_datetime,
745
+ )
app/service/chat/vertex_express_chat_service.py ADDED
@@ -0,0 +1,403 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app/services/chat_service.py
2
+
3
+ import datetime
4
+ import json
5
+ import re
6
+ import time
7
+ from typing import Any, AsyncGenerator, Dict, List
8
+
9
+ from app.config.config import settings
10
+ from app.core.constants import GEMINI_2_FLASH_EXP_SAFETY_SETTINGS
11
+ from app.database.services import add_error_log, add_request_log
12
+ from app.domain.gemini_models import GeminiRequest
13
+ from app.handler.response_handler import GeminiResponseHandler
14
+ from app.handler.stream_optimizer import gemini_optimizer
15
+ from app.log.logger import get_gemini_logger
16
+ from app.service.client.api_client import GeminiApiClient
17
+ from app.service.key.key_manager import KeyManager
18
+ from app.utils.helpers import redact_key_for_logging
19
+
20
+ logger = get_gemini_logger()
21
+
22
+
23
+ def _has_image_parts(contents: List[Dict[str, Any]]) -> bool:
24
+ """判断消息是否包含图片部分"""
25
+ for content in contents:
26
+ if "parts" in content:
27
+ for part in content["parts"]:
28
+ if "image_url" in part or "inline_data" in part:
29
+ return True
30
+ return False
31
+
32
+
33
+ def _clean_json_schema_properties(obj: Any) -> Any:
34
+ """清理JSON Schema中Gemini API不支持的字段"""
35
+ if not isinstance(obj, dict):
36
+ return obj
37
+
38
+ # Gemini API不支持的JSON Schema字段
39
+ unsupported_fields = {
40
+ "exclusiveMaximum",
41
+ "exclusiveMinimum",
42
+ "const",
43
+ "examples",
44
+ "contentEncoding",
45
+ "contentMediaType",
46
+ "if",
47
+ "then",
48
+ "else",
49
+ "allOf",
50
+ "anyOf",
51
+ "oneOf",
52
+ "not",
53
+ "definitions",
54
+ "$schema",
55
+ "$id",
56
+ "$ref",
57
+ "$comment",
58
+ "readOnly",
59
+ "writeOnly",
60
+ }
61
+
62
+ cleaned = {}
63
+ for key, value in obj.items():
64
+ if key in unsupported_fields:
65
+ continue
66
+ if isinstance(value, dict):
67
+ cleaned[key] = _clean_json_schema_properties(value)
68
+ elif isinstance(value, list):
69
+ cleaned[key] = [_clean_json_schema_properties(item) for item in value]
70
+ else:
71
+ cleaned[key] = value
72
+
73
+ return cleaned
74
+
75
+
76
+ def _build_tools(model: str, payload: Dict[str, Any]) -> List[Dict[str, Any]]:
77
+ """构建工具"""
78
+
79
+ def _has_function_call(contents: List[Dict[str, Any]]) -> bool:
80
+ """检查内容中是否包含 functionCall"""
81
+ if not contents or not isinstance(contents, list):
82
+ return False
83
+ for content in contents:
84
+ if not content or not isinstance(content, dict) or "parts" not in content:
85
+ continue
86
+ parts = content.get("parts", [])
87
+ if not parts or not isinstance(parts, list):
88
+ continue
89
+ for part in parts:
90
+ if isinstance(part, dict) and "functionCall" in part:
91
+ return True
92
+ return False
93
+
94
+ def _merge_tools(tools: List[Dict[str, Any]]) -> Dict[str, Any]:
95
+ record = dict()
96
+ for item in tools:
97
+ if not item or not isinstance(item, dict):
98
+ continue
99
+
100
+ for k, v in item.items():
101
+ if k == "functionDeclarations" and v and isinstance(v, list):
102
+ functions = record.get("functionDeclarations", [])
103
+ # 清理每个函数声明中的不支持字段
104
+ cleaned_functions = []
105
+ for func in v:
106
+ if isinstance(func, dict):
107
+ cleaned_func = _clean_json_schema_properties(func)
108
+ cleaned_functions.append(cleaned_func)
109
+ else:
110
+ cleaned_functions.append(func)
111
+ functions.extend(cleaned_functions)
112
+ record["functionDeclarations"] = functions
113
+ else:
114
+ record[k] = v
115
+ return record
116
+
117
+ def _is_structured_output_request(payload: Dict[str, Any]) -> bool:
118
+ """检查请求是否要求结构化JSON输出"""
119
+ try:
120
+ generation_config = payload.get("generationConfig", {})
121
+ return generation_config.get("responseMimeType") == "application/json"
122
+ except (AttributeError, TypeError):
123
+ return False
124
+
125
+ tool = dict()
126
+ if payload and isinstance(payload, dict) and "tools" in payload:
127
+ if payload.get("tools") and isinstance(payload.get("tools"), dict):
128
+ payload["tools"] = [payload.get("tools")]
129
+ items = payload.get("tools", [])
130
+ if items and isinstance(items, list):
131
+ tool.update(_merge_tools(items))
132
+
133
+ # "Tool use with a response mime type: 'application/json' is unsupported"
134
+ # Gemini API限制:不支持同时使用tools和结构化输出(response_mime_type='application/json')
135
+ # 当请求指定了JSON响应格式时,跳过所有工具的添加以避免API错误
136
+ has_structured_output = _is_structured_output_request(payload)
137
+ if not has_structured_output:
138
+ if (
139
+ settings.TOOLS_CODE_EXECUTION_ENABLED
140
+ and not (model.endswith("-search") or "-thinking" in model)
141
+ and not _has_image_parts(payload.get("contents", []))
142
+ ):
143
+ tool["codeExecution"] = {}
144
+
145
+ if model.endswith("-search"):
146
+ tool["googleSearch"] = {}
147
+
148
+ real_model = _get_real_model(model)
149
+ if real_model in settings.URL_CONTEXT_MODELS and settings.URL_CONTEXT_ENABLED:
150
+ tool["urlContext"] = {}
151
+
152
+ # 解决 "Tool use with function calling is unsupported" 问题
153
+ if tool.get("functionDeclarations") or _has_function_call(
154
+ payload.get("contents", [])
155
+ ):
156
+ tool.pop("googleSearch", None)
157
+ tool.pop("codeExecution", None)
158
+ tool.pop("urlContext", None)
159
+
160
+ return [tool] if tool else []
161
+
162
+
163
+ def _get_real_model(model: str) -> str:
164
+ if model.endswith("-search"):
165
+ model = model[:-7]
166
+ if model.endswith("-image"):
167
+ model = model[:-6]
168
+ if model.endswith("-non-thinking"):
169
+ model = model[:-13]
170
+ if "-search" in model and "-non-thinking" in model:
171
+ model = model[:-20]
172
+ return model
173
+
174
+
175
+ def _get_safety_settings(model: str) -> List[Dict[str, str]]:
176
+ """获取安全设置"""
177
+ if model == "gemini-2.0-flash-exp":
178
+ return GEMINI_2_FLASH_EXP_SAFETY_SETTINGS
179
+ return settings.SAFETY_SETTINGS
180
+
181
+
182
+ def _build_payload(model: str, request: GeminiRequest) -> Dict[str, Any]:
183
+ """构建请求payload"""
184
+ request_dict = request.model_dump(exclude_none=False)
185
+ if request.generationConfig:
186
+ if request.generationConfig.maxOutputTokens is None:
187
+ # 如果未指定最大输出长度,则不传递该字段,解决截断的问题
188
+ request_dict["generationConfig"].pop("maxOutputTokens")
189
+
190
+ payload = {
191
+ "contents": request_dict.get("contents", []),
192
+ "tools": _build_tools(model, request_dict),
193
+ "safetySettings": _get_safety_settings(model),
194
+ "generationConfig": request_dict.get("generationConfig"),
195
+ "systemInstruction": request_dict.get("systemInstruction"),
196
+ }
197
+
198
+ if model.endswith("-image") or model.endswith("-image-generation"):
199
+ payload.pop("systemInstruction")
200
+ payload["generationConfig"]["responseModalities"] = ["Text", "Image"]
201
+
202
+ # 处理思考配置:优先使用客户端提供的配置,否则使用默认配置
203
+ client_thinking_config = None
204
+ if request.generationConfig and request.generationConfig.thinkingConfig:
205
+ client_thinking_config = request.generationConfig.thinkingConfig
206
+
207
+ if client_thinking_config is not None:
208
+ # 客户端提供了思考配置,直接使用
209
+ payload["generationConfig"]["thinkingConfig"] = client_thinking_config
210
+ else:
211
+ # 客户端没有提供思考配置,使用默认配置
212
+ if model.endswith("-non-thinking"):
213
+ if "gemini-2.5-pro" in model:
214
+ payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": 128}
215
+ else:
216
+ payload["generationConfig"]["thinkingConfig"] = {"thinkingBudget": 0}
217
+ elif _get_real_model(model) in settings.THINKING_BUDGET_MAP:
218
+ if settings.SHOW_THINKING_PROCESS:
219
+ payload["generationConfig"]["thinkingConfig"] = {
220
+ "thinkingBudget": settings.THINKING_BUDGET_MAP.get(model, 1000),
221
+ "includeThoughts": True,
222
+ }
223
+ else:
224
+ payload["generationConfig"]["thinkingConfig"] = {
225
+ "thinkingBudget": settings.THINKING_BUDGET_MAP.get(model, 1000)
226
+ }
227
+
228
+ return payload
229
+
230
+
231
+ class GeminiChatService:
232
+ """聊天服务"""
233
+
234
+ def __init__(self, base_url: str, key_manager: KeyManager):
235
+ self.api_client = GeminiApiClient(base_url, settings.TIME_OUT)
236
+ self.key_manager = key_manager
237
+ self.response_handler = GeminiResponseHandler()
238
+
239
+ def _extract_text_from_response(self, response: Dict[str, Any]) -> str:
240
+ """从响应中提取文本内容"""
241
+ if not response.get("candidates"):
242
+ return ""
243
+
244
+ candidate = response["candidates"][0]
245
+ content = candidate.get("content", {})
246
+ parts = content.get("parts", [])
247
+
248
+ if parts and "text" in parts[0]:
249
+ return parts[0].get("text", "")
250
+ return ""
251
+
252
+ def _create_char_response(
253
+ self, original_response: Dict[str, Any], text: str
254
+ ) -> Dict[str, Any]:
255
+ """创建包含指定文本的响应"""
256
+ response_copy = json.loads(json.dumps(original_response)) # 深拷贝
257
+ if response_copy.get("candidates") and response_copy["candidates"][0].get(
258
+ "content", {}
259
+ ).get("parts"):
260
+ response_copy["candidates"][0]["content"]["parts"][0]["text"] = text
261
+ return response_copy
262
+
263
+ async def generate_content(
264
+ self, model: str, request: GeminiRequest, api_key: str
265
+ ) -> Dict[str, Any]:
266
+ """生成内容"""
267
+ payload = _build_payload(model, request)
268
+ start_time = time.perf_counter()
269
+ request_datetime = datetime.datetime.now()
270
+ is_success = False
271
+ status_code = None
272
+ response = None
273
+
274
+ try:
275
+ response = await self.api_client.generate_content(payload, model, api_key)
276
+ is_success = True
277
+ status_code = 200
278
+ return self.response_handler.handle_response(response, model, stream=False)
279
+ except Exception as e:
280
+ is_success = False
281
+ error_log_msg = str(e)
282
+ logger.error(f"Normal API call failed with error: {error_log_msg}")
283
+ match = re.search(r"status code (\d+)", error_log_msg)
284
+ if match:
285
+ status_code = int(match.group(1))
286
+ else:
287
+ status_code = 500
288
+
289
+ await add_error_log(
290
+ gemini_key=api_key,
291
+ model_name=model,
292
+ error_type="gemini-chat-non-stream",
293
+ error_log=error_log_msg,
294
+ error_code=status_code,
295
+ request_msg=payload,
296
+ request_datetime=request_datetime,
297
+ )
298
+ raise e
299
+ finally:
300
+ end_time = time.perf_counter()
301
+ latency_ms = int((end_time - start_time) * 1000)
302
+ await add_request_log(
303
+ model_name=model,
304
+ api_key=api_key,
305
+ is_success=is_success,
306
+ status_code=status_code,
307
+ latency_ms=latency_ms,
308
+ request_time=request_datetime,
309
+ )
310
+
311
+ async def stream_generate_content(
312
+ self, model: str, request: GeminiRequest, api_key: str
313
+ ) -> AsyncGenerator[str, None]:
314
+ """流式生成内容"""
315
+ retries = 0
316
+ max_retries = settings.MAX_RETRIES
317
+ payload = _build_payload(model, request)
318
+ is_success = False
319
+ status_code = None
320
+ final_api_key = api_key
321
+
322
+ while retries < max_retries:
323
+ request_datetime = datetime.datetime.now()
324
+ start_time = time.perf_counter()
325
+ current_attempt_key = api_key
326
+ final_api_key = current_attempt_key # Update final key used
327
+ try:
328
+ async for line in self.api_client.stream_generate_content(
329
+ payload, model, current_attempt_key
330
+ ):
331
+ # print(line)
332
+ if line.startswith("data:"):
333
+ line = line[6:]
334
+ response_data = self.response_handler.handle_response(
335
+ json.loads(line), model, stream=True
336
+ )
337
+ text = self._extract_text_from_response(response_data)
338
+ # 如果有文本内容,且开启了流式输出优化器,则使用流式输出优化器处理
339
+ if text and settings.STREAM_OPTIMIZER_ENABLED:
340
+ # 使用流式输出优化器处理文本输出
341
+ async for (
342
+ optimized_chunk
343
+ ) in gemini_optimizer.optimize_stream_output(
344
+ text,
345
+ lambda t: self._create_char_response(response_data, t),
346
+ lambda c: "data: " + json.dumps(c) + "\n\n",
347
+ ):
348
+ yield optimized_chunk
349
+ else:
350
+ # 如果没有文本内容(如工具调用等),整块输出
351
+ yield "data: " + json.dumps(response_data) + "\n\n"
352
+ logger.info("Streaming completed successfully")
353
+ is_success = True
354
+ status_code = 200
355
+ break
356
+ except Exception as e:
357
+ retries += 1
358
+ is_success = False
359
+ error_log_msg = str(e)
360
+ logger.warning(
361
+ f"Streaming API call failed with error: {error_log_msg}. Attempt {retries} of {max_retries}"
362
+ )
363
+ match = re.search(r"status code (\d+)", error_log_msg)
364
+ if match:
365
+ status_code = int(match.group(1))
366
+ else:
367
+ status_code = 500
368
+
369
+ await add_error_log(
370
+ gemini_key=current_attempt_key,
371
+ model_name=model,
372
+ error_type="gemini-chat-stream",
373
+ error_log=error_log_msg,
374
+ error_code=status_code,
375
+ request_msg=payload,
376
+ request_datetime=request_datetime,
377
+ )
378
+
379
+ api_key = await self.key_manager.handle_api_failure(
380
+ current_attempt_key, retries
381
+ )
382
+ if api_key:
383
+ logger.info(
384
+ f"Switched to new API key: {redact_key_for_logging(api_key)}"
385
+ )
386
+ else:
387
+ logger.error(f"No valid API key available after {retries} retries.")
388
+ break
389
+
390
+ if retries >= max_retries:
391
+ logger.error(f"Max retries ({max_retries}) reached for streaming.")
392
+ break
393
+ finally:
394
+ end_time = time.perf_counter()
395
+ latency_ms = int((end_time - start_time) * 1000)
396
+ await add_request_log(
397
+ model_name=model,
398
+ api_key=final_api_key,
399
+ is_success=is_success,
400
+ status_code=status_code,
401
+ latency_ms=latency_ms,
402
+ request_time=request_datetime,
403
+ )
app/service/client/api_client.py ADDED
@@ -0,0 +1,356 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app/services/chat/api_client.py
2
+
3
+ from typing import Dict, Any, AsyncGenerator, Optional
4
+ import httpx
5
+ import random
6
+ from abc import ABC, abstractmethod
7
+ from app.config.config import settings
8
+ from app.log.logger import get_api_client_logger
9
+ from app.core.constants import DEFAULT_TIMEOUT
10
+
11
+ logger = get_api_client_logger()
12
+
13
+ class ApiClient(ABC):
14
+ """API客户端基类"""
15
+
16
+ @abstractmethod
17
+ async def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
18
+ pass
19
+
20
+ @abstractmethod
21
+ async def stream_generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> AsyncGenerator[str, None]:
22
+ pass
23
+
24
+
25
+ class GeminiApiClient(ApiClient):
26
+ """Gemini API客户端"""
27
+
28
+ def __init__(self, base_url: str, timeout: int = DEFAULT_TIMEOUT):
29
+ self.base_url = base_url
30
+ self.timeout = timeout
31
+
32
+ def _get_real_model(self, model: str) -> str:
33
+ if model.endswith("-search"):
34
+ model = model[:-7]
35
+ if model.endswith("-image"):
36
+ model = model[:-6]
37
+ if model.endswith("-non-thinking"):
38
+ model = model[:-13]
39
+ if "-search" in model and "-non-thinking" in model:
40
+ model = model[:-20]
41
+ return model
42
+
43
+ def _prepare_headers(self) -> Dict[str, str]:
44
+ headers = {}
45
+ if settings.CUSTOM_HEADERS:
46
+ headers.update(settings.CUSTOM_HEADERS)
47
+ logger.info(f"Using custom headers: {settings.CUSTOM_HEADERS}")
48
+ return headers
49
+
50
+ async def get_models(self, api_key: str) -> Optional[Dict[str, Any]]:
51
+ """获取可用的 Gemini 模型列表"""
52
+ timeout = httpx.Timeout(timeout=5)
53
+
54
+ proxy_to_use = None
55
+ if settings.PROXIES:
56
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
57
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
58
+ else:
59
+ proxy_to_use = random.choice(settings.PROXIES)
60
+ logger.info(f"Using proxy for getting models: {proxy_to_use}")
61
+
62
+ headers = self._prepare_headers()
63
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
64
+ url = f"{self.base_url}/models?key={api_key}&pageSize=1000"
65
+ try:
66
+ response = await client.get(url, headers=headers)
67
+ response.raise_for_status()
68
+ return response.json()
69
+ except httpx.HTTPStatusError as e:
70
+ logger.error(f"获取模型列表失败: {e.response.status_code}")
71
+ logger.error(e.response.text)
72
+ return None
73
+ except httpx.RequestError as e:
74
+ logger.error(f"请求模型列表失败: {e}")
75
+ return None
76
+
77
+ async def generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
78
+ timeout = httpx.Timeout(self.timeout, read=self.timeout)
79
+ model = self._get_real_model(model)
80
+
81
+ proxy_to_use = None
82
+ if settings.PROXIES:
83
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
84
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
85
+ else:
86
+ proxy_to_use = random.choice(settings.PROXIES)
87
+ logger.info(f"Using proxy for getting models: {proxy_to_use}")
88
+
89
+ headers = self._prepare_headers()
90
+
91
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
92
+ url = f"{self.base_url}/models/{model}:generateContent?key={api_key}"
93
+
94
+ try:
95
+ response = await client.post(url, json=payload, headers=headers)
96
+
97
+ if response.status_code != 200:
98
+ error_content = response.text
99
+ logger.error(f"API call failed - Status: {response.status_code}, Content: {error_content}")
100
+ raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
101
+
102
+ response_data = response.json()
103
+
104
+ # 检查响应结构的基本信息
105
+ if not response_data.get("candidates"):
106
+ logger.warning("No candidates found in API response")
107
+
108
+ return response_data
109
+
110
+ except httpx.TimeoutException as e:
111
+ logger.error(f"Request timeout: {e}")
112
+ raise Exception(f"Request timeout: {e}")
113
+ except httpx.RequestError as e:
114
+ logger.error(f"Request error: {e}")
115
+ raise Exception(f"Request error: {e}")
116
+ except Exception as e:
117
+ logger.error(f"Unexpected error: {e}")
118
+ raise
119
+
120
+ async def stream_generate_content(self, payload: Dict[str, Any], model: str, api_key: str) -> AsyncGenerator[str, None]:
121
+ timeout = httpx.Timeout(self.timeout, read=self.timeout)
122
+ model = self._get_real_model(model)
123
+
124
+ proxy_to_use = None
125
+ if settings.PROXIES:
126
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
127
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
128
+ else:
129
+ proxy_to_use = random.choice(settings.PROXIES)
130
+ logger.info(f"Using proxy for getting models: {proxy_to_use}")
131
+
132
+ headers = self._prepare_headers()
133
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
134
+ url = f"{self.base_url}/models/{model}:streamGenerateContent?alt=sse&key={api_key}"
135
+ async with client.stream(method="POST", url=url, json=payload, headers=headers) as response:
136
+ if response.status_code != 200:
137
+ error_content = await response.aread()
138
+ error_msg = error_content.decode("utf-8")
139
+ raise Exception(f"API call failed with status code {response.status_code}, {error_msg}")
140
+ async for line in response.aiter_lines():
141
+ yield line
142
+
143
+ async def count_tokens(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
144
+ timeout = httpx.Timeout(self.timeout, read=self.timeout)
145
+ model = self._get_real_model(model)
146
+
147
+ proxy_to_use = None
148
+ if settings.PROXIES:
149
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
150
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
151
+ else:
152
+ proxy_to_use = random.choice(settings.PROXIES)
153
+ logger.info(f"Using proxy for counting tokens: {proxy_to_use}")
154
+
155
+ headers = self._prepare_headers()
156
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
157
+ url = f"{self.base_url}/models/{model}:countTokens?key={api_key}"
158
+ response = await client.post(url, json=payload, headers=headers)
159
+ if response.status_code != 200:
160
+ error_content = response.text
161
+ raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
162
+ return response.json()
163
+
164
+ async def embed_content(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
165
+ """单一嵌入内容生成"""
166
+ timeout = httpx.Timeout(self.timeout, read=self.timeout)
167
+ model = self._get_real_model(model)
168
+
169
+ proxy_to_use = None
170
+ if settings.PROXIES:
171
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
172
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
173
+ else:
174
+ proxy_to_use = random.choice(settings.PROXIES)
175
+ logger.info(f"Using proxy for embedding: {proxy_to_use}")
176
+
177
+ headers = self._prepare_headers()
178
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
179
+ url = f"{self.base_url}/models/{model}:embedContent?key={api_key}"
180
+
181
+ try:
182
+ response = await client.post(url, json=payload, headers=headers)
183
+
184
+ if response.status_code != 200:
185
+ error_content = response.text
186
+ logger.error(f"Embedding API call failed - Status: {response.status_code}, Content: {error_content}")
187
+ raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
188
+
189
+ return response.json()
190
+
191
+ except httpx.TimeoutException as e:
192
+ logger.error(f"Embedding request timeout: {e}")
193
+ raise Exception(f"Request timeout: {e}")
194
+ except httpx.RequestError as e:
195
+ logger.error(f"Embedding request error: {e}")
196
+ raise Exception(f"Request error: {e}")
197
+ except Exception as e:
198
+ logger.error(f"Unexpected embedding error: {e}")
199
+ raise
200
+
201
+ async def batch_embed_contents(self, payload: Dict[str, Any], model: str, api_key: str) -> Dict[str, Any]:
202
+ """批量嵌入内容生成"""
203
+ timeout = httpx.Timeout(self.timeout, read=self.timeout)
204
+ model = self._get_real_model(model)
205
+
206
+ proxy_to_use = None
207
+ if settings.PROXIES:
208
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
209
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
210
+ else:
211
+ proxy_to_use = random.choice(settings.PROXIES)
212
+ logger.info(f"Using proxy for batch embedding: {proxy_to_use}")
213
+
214
+ headers = self._prepare_headers()
215
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
216
+ url = f"{self.base_url}/models/{model}:batchEmbedContents?key={api_key}"
217
+
218
+ try:
219
+ response = await client.post(url, json=payload, headers=headers)
220
+
221
+ if response.status_code != 200:
222
+ error_content = response.text
223
+ logger.error(f"Batch embedding API call failed - Status: {response.status_code}, Content: {error_content}")
224
+ raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
225
+
226
+ return response.json()
227
+
228
+ except httpx.TimeoutException as e:
229
+ logger.error(f"Batch embedding request timeout: {e}")
230
+ raise Exception(f"Request timeout: {e}")
231
+ except httpx.RequestError as e:
232
+ logger.error(f"Batch embedding request error: {e}")
233
+ raise Exception(f"Request error: {e}")
234
+ except Exception as e:
235
+ logger.error(f"Unexpected batch embedding error: {e}")
236
+ raise
237
+
238
+
239
+ class OpenaiApiClient(ApiClient):
240
+ """OpenAI API客户端"""
241
+
242
+ def __init__(self, base_url: str, timeout: int = DEFAULT_TIMEOUT):
243
+ self.base_url = base_url
244
+ self.timeout = timeout
245
+
246
+ def _prepare_headers(self, api_key: str) -> Dict[str, str]:
247
+ headers = {"Authorization": f"Bearer {api_key}"}
248
+ if settings.CUSTOM_HEADERS:
249
+ headers.update(settings.CUSTOM_HEADERS)
250
+ logger.info(f"Using custom headers: {settings.CUSTOM_HEADERS}")
251
+ return headers
252
+
253
+ async def get_models(self, api_key: str) -> Dict[str, Any]:
254
+ timeout = httpx.Timeout(self.timeout, read=self.timeout)
255
+
256
+ proxy_to_use = None
257
+ if settings.PROXIES:
258
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
259
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
260
+ else:
261
+ proxy_to_use = random.choice(settings.PROXIES)
262
+ logger.info(f"Using proxy for getting models: {proxy_to_use}")
263
+
264
+ headers = self._prepare_headers(api_key)
265
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
266
+ url = f"{self.base_url}/openai/models"
267
+ response = await client.get(url, headers=headers)
268
+ if response.status_code != 200:
269
+ error_content = response.text
270
+ raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
271
+ return response.json()
272
+
273
+ async def generate_content(self, payload: Dict[str, Any], api_key: str) -> Dict[str, Any]:
274
+ timeout = httpx.Timeout(self.timeout, read=self.timeout)
275
+ logger.info(f"settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY: {settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY}")
276
+ proxy_to_use = None
277
+ if settings.PROXIES:
278
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
279
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
280
+ else:
281
+ proxy_to_use = random.choice(settings.PROXIES)
282
+ logger.info(f"Using proxy for getting models: {proxy_to_use}")
283
+
284
+ headers = self._prepare_headers(api_key)
285
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
286
+ url = f"{self.base_url}/openai/chat/completions"
287
+ response = await client.post(url, json=payload, headers=headers)
288
+ if response.status_code != 200:
289
+ error_content = response.text
290
+ raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
291
+ return response.json()
292
+
293
+ async def stream_generate_content(self, payload: Dict[str, Any], api_key: str) -> AsyncGenerator[str, None]:
294
+ timeout = httpx.Timeout(self.timeout, read=self.timeout)
295
+ proxy_to_use = None
296
+ if settings.PROXIES:
297
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
298
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
299
+ else:
300
+ proxy_to_use = random.choice(settings.PROXIES)
301
+ logger.info(f"Using proxy for getting models: {proxy_to_use}")
302
+
303
+ headers = self._prepare_headers(api_key)
304
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
305
+ url = f"{self.base_url}/openai/chat/completions"
306
+ async with client.stream(method="POST", url=url, json=payload, headers=headers) as response:
307
+ if response.status_code != 200:
308
+ error_content = await response.aread()
309
+ error_msg = error_content.decode("utf-8")
310
+ raise Exception(f"API call failed with status code {response.status_code}, {error_msg}")
311
+ async for line in response.aiter_lines():
312
+ yield line
313
+
314
+ async def create_embeddings(self, input: str, model: str, api_key: str) -> Dict[str, Any]:
315
+ timeout = httpx.Timeout(self.timeout, read=self.timeout)
316
+
317
+ proxy_to_use = None
318
+ if settings.PROXIES:
319
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
320
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
321
+ else:
322
+ proxy_to_use = random.choice(settings.PROXIES)
323
+ logger.info(f"Using proxy for getting models: {proxy_to_use}")
324
+
325
+ headers = self._prepare_headers(api_key)
326
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
327
+ url = f"{self.base_url}/openai/embeddings"
328
+ payload = {
329
+ "input": input,
330
+ "model": model,
331
+ }
332
+ response = await client.post(url, json=payload, headers=headers)
333
+ if response.status_code != 200:
334
+ error_content = response.text
335
+ raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
336
+ return response.json()
337
+
338
+ async def generate_images(self, payload: Dict[str, Any], api_key: str) -> Dict[str, Any]:
339
+ timeout = httpx.Timeout(self.timeout, read=self.timeout)
340
+
341
+ proxy_to_use = None
342
+ if settings.PROXIES:
343
+ if settings.PROXIES_USE_CONSISTENCY_HASH_BY_API_KEY:
344
+ proxy_to_use = settings.PROXIES[hash(api_key) % len(settings.PROXIES)]
345
+ else:
346
+ proxy_to_use = random.choice(settings.PROXIES)
347
+ logger.info(f"Using proxy for getting models: {proxy_to_use}")
348
+
349
+ headers = self._prepare_headers(api_key)
350
+ async with httpx.AsyncClient(timeout=timeout, proxy=proxy_to_use) as client:
351
+ url = f"{self.base_url}/openai/images/generations"
352
+ response = await client.post(url, json=payload, headers=headers)
353
+ if response.status_code != 200:
354
+ error_content = response.text
355
+ raise Exception(f"API call failed with status code {response.status_code}, {error_content}")
356
+ return response.json()
app/service/config/config_service.py ADDED
@@ -0,0 +1,261 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ 配置服务模块
3
+ """
4
+
5
+ import datetime
6
+ import json
7
+ from typing import Any, Dict, List
8
+
9
+ from dotenv import find_dotenv, load_dotenv
10
+ from fastapi import HTTPException
11
+ from sqlalchemy import insert, update
12
+
13
+ from app.config.config import Settings as ConfigSettings
14
+ from app.config.config import settings
15
+ from app.database.connection import database
16
+ from app.database.models import Settings
17
+ from app.database.services import get_all_settings
18
+ from app.log.logger import get_config_routes_logger
19
+ from app.service.key.key_manager import (
20
+ get_key_manager_instance,
21
+ reset_key_manager_instance,
22
+ )
23
+ from app.service.model.model_service import ModelService
24
+
25
+ logger = get_config_routes_logger()
26
+
27
+
28
+ class ConfigService:
29
+ """配置服务类,用于管理应用程序配置"""
30
+
31
+ @staticmethod
32
+ async def get_config() -> Dict[str, Any]:
33
+ return settings.model_dump()
34
+
35
+ @staticmethod
36
+ async def update_config(config_data: Dict[str, Any]) -> Dict[str, Any]:
37
+ for key, value in config_data.items():
38
+ if hasattr(settings, key):
39
+ setattr(settings, key, value)
40
+ logger.debug(f"Updated setting in memory: {key}")
41
+
42
+ # 获取现有设置
43
+ existing_settings_raw: List[Dict[str, Any]] = await get_all_settings()
44
+ existing_settings_map: Dict[str, Dict[str, Any]] = {
45
+ s["key"]: s for s in existing_settings_raw
46
+ }
47
+ existing_keys = set(existing_settings_map.keys())
48
+
49
+ settings_to_update: List[Dict[str, Any]] = []
50
+ settings_to_insert: List[Dict[str, Any]] = []
51
+ now = datetime.datetime.now(datetime.timezone(datetime.timedelta(hours=8)))
52
+
53
+ # 准备要更新或插入的数据
54
+ for key, value in config_data.items():
55
+ # 处理不同类型的值
56
+ if isinstance(value, list):
57
+ db_value = json.dumps(value)
58
+ elif isinstance(value, dict):
59
+ db_value = json.dumps(value)
60
+ elif isinstance(value, bool):
61
+ db_value = str(value).lower()
62
+ else:
63
+ db_value = str(value)
64
+
65
+ # 仅当值发生变化时才更新
66
+ if key in existing_keys and existing_settings_map[key]["value"] == db_value:
67
+ continue
68
+
69
+ description = f"{key}配置项"
70
+
71
+ data = {
72
+ "key": key,
73
+ "value": db_value,
74
+ "description": description,
75
+ "updated_at": now,
76
+ }
77
+
78
+ if key in existing_keys:
79
+ data["description"] = existing_settings_map[key].get(
80
+ "description", description
81
+ )
82
+ settings_to_update.append(data)
83
+ else:
84
+ data["created_at"] = now
85
+ settings_to_insert.append(data)
86
+
87
+ # 在事务中执行批量插入和更新
88
+ if settings_to_insert or settings_to_update:
89
+ try:
90
+ async with database.transaction():
91
+ if settings_to_insert:
92
+ query_insert = insert(Settings).values(settings_to_insert)
93
+ await database.execute(query=query_insert)
94
+ logger.info(
95
+ f"Bulk inserted {len(settings_to_insert)} settings."
96
+ )
97
+
98
+ if settings_to_update:
99
+ for setting_data in settings_to_update:
100
+ query_update = (
101
+ update(Settings)
102
+ .where(Settings.key == setting_data["key"])
103
+ .values(
104
+ value=setting_data["value"],
105
+ description=setting_data["description"],
106
+ updated_at=setting_data["updated_at"],
107
+ )
108
+ )
109
+ await database.execute(query=query_update)
110
+ logger.info(f"Updated {len(settings_to_update)} settings.")
111
+ except Exception as e:
112
+ logger.error(f"Failed to bulk update/insert settings: {str(e)}")
113
+ raise
114
+
115
+ # 重置并重新初始化 KeyManager
116
+ try:
117
+ await reset_key_manager_instance()
118
+ await get_key_manager_instance(settings.API_KEYS, settings.VERTEX_API_KEYS)
119
+ logger.info("KeyManager instance re-initialized with updated settings.")
120
+ except Exception as e:
121
+ logger.error(f"Failed to re-initialize KeyManager: {str(e)}")
122
+
123
+ return await ConfigService.get_config()
124
+
125
+ @staticmethod
126
+ async def delete_key(key_to_delete: str) -> Dict[str, Any]:
127
+ """删除单个API密钥"""
128
+ # 确保 settings.API_KEYS 是一个列表
129
+ if not isinstance(settings.API_KEYS, list):
130
+ settings.API_KEYS = []
131
+
132
+ original_keys_count = len(settings.API_KEYS)
133
+ # 创建一个不包含待删除密钥的新列表
134
+ updated_api_keys = [k for k in settings.API_KEYS if k != key_to_delete]
135
+
136
+ if len(updated_api_keys) < original_keys_count:
137
+ # 密钥已找到并从列表中移除
138
+ settings.API_KEYS = updated_api_keys # 首先更新内存中的 settings
139
+ # 使用 update_config 持久化更改,它同时处理数据库和 KeyManager
140
+ await ConfigService.update_config({"API_KEYS": settings.API_KEYS})
141
+ logger.info(f"密钥 '{key_to_delete}' 已成功删除。")
142
+ return {"success": True, "message": f"密钥 '{key_to_delete}' 已成功删除。"}
143
+ else:
144
+ # 未找到密钥
145
+ logger.warning(f"尝试删除密钥 '{key_to_delete}',但未找到该密钥。")
146
+ return {"success": False, "message": f"未找到密钥 '{key_to_delete}'。"}
147
+
148
+ @staticmethod
149
+ async def delete_selected_keys(keys_to_delete: List[str]) -> Dict[str, Any]:
150
+ """批量删除选定的API密钥"""
151
+ if not isinstance(settings.API_KEYS, list):
152
+ settings.API_KEYS = []
153
+
154
+ deleted_count = 0
155
+ not_found_keys: List[str] = []
156
+
157
+ current_api_keys = list(settings.API_KEYS)
158
+ keys_actually_removed: List[str] = []
159
+
160
+ for key_to_del in keys_to_delete:
161
+ if key_to_del in current_api_keys:
162
+ current_api_keys.remove(key_to_del)
163
+ keys_actually_removed.append(key_to_del)
164
+ deleted_count += 1
165
+ else:
166
+ not_found_keys.append(key_to_del)
167
+
168
+ if deleted_count > 0:
169
+ settings.API_KEYS = current_api_keys
170
+ await ConfigService.update_config({"API_KEYS": settings.API_KEYS})
171
+ logger.info(
172
+ f"成功删除 {deleted_count} 个密钥。密钥: {keys_actually_removed}"
173
+ )
174
+ message = f"成功删除 {deleted_count} 个密钥。"
175
+ if not_found_keys:
176
+ message += f" {len(not_found_keys)} 个密钥未找到: {not_found_keys}。"
177
+ return {
178
+ "success": True,
179
+ "message": message,
180
+ "deleted_count": deleted_count,
181
+ "not_found_keys": not_found_keys,
182
+ }
183
+ else:
184
+ message = "没有密钥被删除。"
185
+ if not_found_keys:
186
+ message = f"所有 {len(not_found_keys)} 个指定的密钥均未找到: {not_found_keys}。"
187
+ elif not keys_to_delete:
188
+ message = "未指定要删除的密钥。"
189
+ logger.warning(message)
190
+ return {
191
+ "success": False,
192
+ "message": message,
193
+ "deleted_count": 0,
194
+ "not_found_keys": not_found_keys,
195
+ }
196
+
197
+ @staticmethod
198
+ async def reset_config() -> Dict[str, Any]:
199
+ """
200
+ 重置配置:优先从系统环境变量加载,然后从 .env 文件加载,
201
+ 更新内存中的 settings 对象,并刷新 KeyManager。
202
+
203
+ Returns:
204
+ Dict[str, Any]: 重置后的配置字典
205
+ """
206
+ # 1. 重新加载配置对象,它应该处理环境变量和 .env 的优先级
207
+ _reload_settings()
208
+ logger.info(
209
+ "Settings object reloaded, prioritizing system environment variables then .env file."
210
+ )
211
+
212
+ # 2. 重置并重新初始化 KeyManager
213
+ try:
214
+ await reset_key_manager_instance()
215
+ # 确保使用更新后的 settings 中的 API_KEYS
216
+ await get_key_manager_instance(settings.API_KEYS)
217
+ logger.info("KeyManager instance re-initialized with reloaded settings.")
218
+ except Exception as e:
219
+ logger.error(f"Failed to re-initialize KeyManager during reset: {str(e)}")
220
+ # 根据需要决定是否抛出异常或继续
221
+ # 这里选择记录错误并继续
222
+
223
+ # 3. 返回更新后的配置
224
+ return await ConfigService.get_config()
225
+
226
+ @staticmethod
227
+ async def fetch_ui_models() -> List[Dict[str, Any]]:
228
+ """获取用于UI显示的模型列表"""
229
+ try:
230
+ key_manager = await get_key_manager_instance()
231
+ model_service = ModelService()
232
+
233
+ api_key = await key_manager.get_random_valid_key()
234
+ if not api_key:
235
+ logger.error("No valid API keys available to fetch model list for UI.")
236
+ raise HTTPException(
237
+ status_code=500,
238
+ detail="No valid API keys available to fetch model list.",
239
+ )
240
+
241
+ models = await model_service.get_gemini_openai_models(api_key)
242
+ return models
243
+ except HTTPException as e:
244
+ raise e
245
+ except Exception as e:
246
+ logger.error(
247
+ f"Failed to fetch models for UI in ConfigService: {e}", exc_info=True
248
+ )
249
+ raise HTTPException(
250
+ status_code=500, detail=f"Failed to fetch models for UI: {str(e)}"
251
+ )
252
+
253
+
254
+ # 重新加载配置的函数
255
+ def _reload_settings():
256
+ """重新加载环境变量并更新配置"""
257
+ # 显式加载 .env 文件,覆盖现有环境变量
258
+ load_dotenv(find_dotenv(), override=True)
259
+ # 更新现有 settings 对象的属性,而不是新建实例
260
+ for key, value in ConfigSettings().model_dump().items():
261
+ setattr(settings, key, value)
app/service/embedding/embedding_service.py ADDED
@@ -0,0 +1,87 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import datetime
2
+ import re
3
+ import time
4
+ from typing import List, Union
5
+
6
+ import openai
7
+ from openai import APIStatusError
8
+ from openai.types import CreateEmbeddingResponse
9
+
10
+ from app.config.config import settings
11
+ from app.database.services import add_error_log, add_request_log
12
+ from app.log.logger import get_embeddings_logger
13
+
14
+ logger = get_embeddings_logger()
15
+
16
+
17
+ class EmbeddingService:
18
+
19
+ async def create_embedding(
20
+ self, input_text: Union[str, List[str]], model: str, api_key: str
21
+ ) -> CreateEmbeddingResponse:
22
+ """Create embeddings using OpenAI API with database logging"""
23
+ start_time = time.perf_counter()
24
+ request_datetime = datetime.datetime.now()
25
+ is_success = False
26
+ status_code = None
27
+ response = None
28
+ error_log_msg = ""
29
+ if isinstance(input_text, list):
30
+ request_msg_log = {
31
+ "input_truncated": [
32
+ str(item)[:100] + "..." if len(str(item)) > 100 else str(item)
33
+ for item in input_text[:5]
34
+ ]
35
+ }
36
+ if len(input_text) > 5:
37
+ request_msg_log["input_truncated"].append("...")
38
+ else:
39
+ request_msg_log = {
40
+ "input_truncated": (
41
+ input_text[:1000] + "..." if len(input_text) > 1000 else input_text
42
+ )
43
+ }
44
+
45
+ try:
46
+ client = openai.OpenAI(api_key=api_key, base_url=settings.BASE_URL)
47
+ response = client.embeddings.create(input=input_text, model=model)
48
+ is_success = True
49
+ status_code = 200
50
+ return response
51
+ except APIStatusError as e:
52
+ is_success = False
53
+ status_code = e.status_code
54
+ error_log_msg = f"OpenAI API error: {e}"
55
+ logger.error(f"Error creating embedding (APIStatusError): {error_log_msg}")
56
+ raise e
57
+ except Exception as e:
58
+ is_success = False
59
+ error_log_msg = f"Generic error: {e}"
60
+ logger.error(f"Error creating embedding (Exception): {error_log_msg}")
61
+ match = re.search(r"status code (\d+)", str(e))
62
+ if match:
63
+ status_code = int(match.group(1))
64
+ else:
65
+ status_code = 500
66
+ raise e
67
+ finally:
68
+ end_time = time.perf_counter()
69
+ latency_ms = int((end_time - start_time) * 1000)
70
+ if not is_success:
71
+ await add_error_log(
72
+ gemini_key=api_key,
73
+ model_name=model,
74
+ error_type="openai-embedding",
75
+ error_log=error_log_msg,
76
+ error_code=status_code,
77
+ request_msg=request_msg_log,
78
+ request_datetime=request_datetime,
79
+ )
80
+ await add_request_log(
81
+ model_name=model,
82
+ api_key=api_key,
83
+ is_success=is_success,
84
+ status_code=status_code,
85
+ latency_ms=latency_ms,
86
+ request_time=request_datetime,
87
+ )
app/service/embedding/gemini_embedding_service.py ADDED
@@ -0,0 +1,150 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # app/service/embedding/gemini_embedding_service.py
2
+
3
+ import datetime
4
+ import re
5
+ import time
6
+ from typing import Any, Dict
7
+
8
+ from app.config.config import settings
9
+ from app.database.services import add_error_log, add_request_log
10
+ from app.domain.gemini_models import GeminiBatchEmbedRequest, GeminiEmbedRequest
11
+ from app.log.logger import get_gemini_embedding_logger
12
+ from app.service.client.api_client import GeminiApiClient
13
+ from app.service.key.key_manager import KeyManager
14
+
15
+ logger = get_gemini_embedding_logger()
16
+
17
+
18
+ def _build_embed_payload(request: GeminiEmbedRequest) -> Dict[str, Any]:
19
+ """构建嵌入请求payload"""
20
+ payload = {"content": request.content.model_dump()}
21
+
22
+ if request.taskType:
23
+ payload["taskType"] = request.taskType
24
+ if request.title:
25
+ payload["title"] = request.title
26
+ if request.outputDimensionality:
27
+ payload["outputDimensionality"] = request.outputDimensionality
28
+
29
+ return payload
30
+
31
+
32
+ def _build_batch_embed_payload(
33
+ request: GeminiBatchEmbedRequest, model: str
34
+ ) -> Dict[str, Any]:
35
+ """构建批量嵌入请求payload"""
36
+ requests = []
37
+ for embed_request in request.requests:
38
+ embed_payload = _build_embed_payload(embed_request)
39
+ embed_payload["model"] = (
40
+ f"models/{model}" # Gemini API要求每个请求包含model字段
41
+ )
42
+ requests.append(embed_payload)
43
+
44
+ return {"requests": requests}
45
+
46
+
47
+ class GeminiEmbeddingService:
48
+ """Gemini嵌入服务"""
49
+
50
+ def __init__(self, base_url: str, key_manager: KeyManager):
51
+ self.api_client = GeminiApiClient(base_url, settings.TIME_OUT)
52
+ self.key_manager = key_manager
53
+
54
+ async def embed_content(
55
+ self, model: str, request: GeminiEmbedRequest, api_key: str
56
+ ) -> Dict[str, Any]:
57
+ """生成单一嵌入内容"""
58
+ payload = _build_embed_payload(request)
59
+ start_time = time.perf_counter()
60
+ request_datetime = datetime.datetime.now()
61
+ is_success = False
62
+ status_code = None
63
+ response = None
64
+
65
+ try:
66
+ response = await self.api_client.embed_content(payload, model, api_key)
67
+ is_success = True
68
+ status_code = 200
69
+ return response
70
+ except Exception as e:
71
+ is_success = False
72
+ error_log_msg = str(e)
73
+ logger.error(f"Single embedding API call failed: {error_log_msg}")
74
+ match = re.search(r"status code (\d+)", error_log_msg)
75
+ if match:
76
+ status_code = int(match.group(1))
77
+ else:
78
+ status_code = 500
79
+
80
+ await add_error_log(
81
+ gemini_key=api_key,
82
+ model_name=model,
83
+ error_type="gemini-embed-single",
84
+ error_log=error_log_msg,
85
+ error_code=status_code,
86
+ request_msg=payload,
87
+ request_datetime=request_datetime,
88
+ )
89
+ raise e
90
+ finally:
91
+ end_time = time.perf_counter()
92
+ latency_ms = int((end_time - start_time) * 1000)
93
+ await add_request_log(
94
+ model_name=model,
95
+ api_key=api_key,
96
+ is_success=is_success,
97
+ status_code=status_code,
98
+ latency_ms=latency_ms,
99
+ request_time=request_datetime,
100
+ )
101
+
102
+ async def batch_embed_contents(
103
+ self, model: str, request: GeminiBatchEmbedRequest, api_key: str
104
+ ) -> Dict[str, Any]:
105
+ """生成批量嵌入内容"""
106
+ payload = _build_batch_embed_payload(request, model)
107
+ start_time = time.perf_counter()
108
+ request_datetime = datetime.datetime.now()
109
+ is_success = False
110
+ status_code = None
111
+ response = None
112
+
113
+ try:
114
+ response = await self.api_client.batch_embed_contents(
115
+ payload, model, api_key
116
+ )
117
+ is_success = True
118
+ status_code = 200
119
+ return response
120
+ except Exception as e:
121
+ is_success = False
122
+ error_log_msg = str(e)
123
+ logger.error(f"Batch embedding API call failed: {error_log_msg}")
124
+ match = re.search(r"status code (\d+)", error_log_msg)
125
+ if match:
126
+ status_code = int(match.group(1))
127
+ else:
128
+ status_code = 500
129
+
130
+ await add_error_log(
131
+ gemini_key=api_key,
132
+ model_name=model,
133
+ error_type="gemini-embed-batch",
134
+ error_log=error_log_msg,
135
+ error_code=status_code,
136
+ request_msg=payload,
137
+ request_datetime=request_datetime,
138
+ )
139
+ raise e
140
+ finally:
141
+ end_time = time.perf_counter()
142
+ latency_ms = int((end_time - start_time) * 1000)
143
+ await add_request_log(
144
+ model_name=model,
145
+ api_key=api_key,
146
+ is_success=is_success,
147
+ status_code=status_code,
148
+ latency_ms=latency_ms,
149
+ request_time=request_datetime,
150
+ )