ParthSadaria commited on
Commit
925b0de
·
verified ·
1 Parent(s): 32f3124

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +1365 -313
main.py CHANGED
@@ -1,395 +1,1447 @@
1
  import os
2
- import json
3
- import datetime
4
- import asyncio
5
  import re
6
- from functools import lru_cache
7
- from pathlib import Path
8
- from typing import List, Dict, Any, Tuple, Optional
9
-
10
- import httpx
11
- import uvicorn
12
  from dotenv import load_dotenv
13
- from fastapi import FastAPI, HTTPException, Request, Depends, Security, Query, APIRouter
14
- from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse, FileResponse, PlainTextResponse
15
  from fastapi.security import APIKeyHeader
16
- from pydantic_settings import BaseSettings
17
- from pydantic import BaseModel, Field
 
 
 
 
 
 
 
 
 
 
 
 
 
18
  from starlette.middleware.cors import CORSMiddleware
19
- from starlette.middleware.gzip import GZipMiddleware
20
- from starlette.status import HTTP_403_FORBIDDEN, HTTP_503_SERVICE_UNAVAILABLE
21
 
22
- # Use cloudscraper for specific endpoints that need it
23
- try:
24
- import cloudscraper
25
- except ImportError:
26
- cloudscraper = None
27
 
28
- from usage_tracker import UsageTracker
29
 
30
- # --- Initial Setup ---
31
  load_dotenv()
32
- # Use uvloop for better performance if available
33
- try:
34
- import uvloop
35
- asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
36
- except ImportError:
37
- pass
38
-
39
- # --- Configuration Management using Pydantic ---
40
- class Settings(BaseSettings):
41
- """Manages all application settings and environment variables in one place."""
42
- api_keys: List[str] = Field(..., env="API_KEYS")
43
-
44
- # Endpoints for various model providers
45
- secret_api_endpoint: str = Field(..., env="SECRET_API_ENDPOINT")
46
- secret_api_endpoint_2: str = Field(..., env="SECRET_API_ENDPOINT_2")
47
- secret_api_endpoint_3: str = Field(..., env="SECRET_API_ENDPOINT_3")
48
- secret_api_endpoint_4: str = "https://text.pollinations.ai/openai"
49
- secret_api_endpoint_5: str = Field(..., env="SECRET_API_ENDPOINT_5")
50
- secret_api_endpoint_6: str = Field(..., env="SECRET_API_ENDPOINT_6")
51
-
52
- # Specific provider keys and APIs
53
- mistral_api: str = "https://api.mistral.ai"
54
- mistral_key: str = Field(..., env="MISTRAL_KEY")
55
- gemini_key: str = Field(..., env="GEMINI_KEY")
56
- new_img_api: str = Field(..., env="NEW_IMG")
57
-
58
- endpoint_origin: Optional[str] = Field(None, env="ENDPOINT_ORIGIN")
59
- header_url: Optional[str] = Field(None, env="HEADER_URL")
60
 
61
- class Config:
62
- env_file = '.env'
63
- env_file_encoding = 'utf-8'
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
64
 
65
- @lru_cache()
66
- def get_settings():
67
- return Settings()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
68
 
69
- # --- Pydantic Models for Payloads ---
70
- class ChatPayload(BaseModel):
 
 
 
 
 
 
71
  model: str
72
- messages: List[Dict[str, Any]]
73
  stream: bool = False
74
 
75
  class ImageGenerationPayload(BaseModel):
 
76
  model: str
77
  prompt: str
78
- size: int
79
- number: int
80
-
81
- # --- Global Objects & State ---
82
- app = FastAPI(
83
- title="LokiAI API",
84
- version="2.5.0",
85
- description="A robust and scalable API proxy for various AI models, now fully rewritten.",
86
- )
87
- usage_tracker = UsageTracker()
88
- api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
89
- server_status = {"online": True}
90
-
91
- # --- Model & API Configuration ---
92
- MODEL_SETS = {
93
- "mistral": {"mistral-large-latest", "codestral-latest", "mistral-small-latest"},
94
- "pollinations": {"openai", "gemini", "phi", "llama"},
95
- "alternate": {"o1", "grok-3", "sonar-pro"},
96
- "claude": {"claude-3-7-sonnet", "claude 3.5 sonnet", "o3-mini-medium"},
97
- "gemini": {"gemini-1.5-pro", "gemini-1.5-flash", "gemini-2.0-flash"},
98
- "image": {"Flux Pro Ultra", "dall-e-3", "stable-diffusion-3-large-turbo"},
99
- }
100
 
101
- def get_api_details(model_name: str, settings: Settings) -> Tuple[str, Dict, str]:
102
- """Returns the endpoint, headers, and path for a given model."""
103
- if model_name in MODEL_SETS["mistral"]:
104
- return settings.mistral_api, {"Authorization": f"Bearer {settings.mistral_key}"}, "/v1/chat/completions"
105
- if model_name in MODEL_SETS["gemini"]:
106
- return settings.secret_api_endpoint_6, {"Authorization": f"Bearer {settings.gemini_key}"}, "/chat/completions"
107
- if model_name in MODEL_SETS["pollinations"]:
108
- return settings.secret_api_endpoint_4, {}, "/v1/chat/completions"
109
- if model_name in MODEL_SETS["claude"]:
110
- return settings.secret_api_endpoint_5, {}, "/v1/chat/completions"
111
- if model_name in MODEL_SETS["alternate"]:
112
- return settings.secret_api_endpoint_2, {}, "/v1/chat/completions"
113
- if model_name in MODEL_SETS["image"]:
114
- return settings.new_img_api, {}, ""
115
-
116
- # Default case
117
- headers = {
118
- "Origin": settings.header_url, "Referer": settings.header_url
119
- } if settings.header_url else {}
120
- return settings.secret_api_endpoint, headers, "/v1/chat/completions"
121
 
122
- # --- Dependencies & Security ---
123
- async def get_api_key(request: Request, api_key: str = Security(api_key_header)):
124
- """Validates the API key, allowing specific referers to bypass."""
125
- referer = request.headers.get("referer", "")
126
- if referer and "parthsadaria-lokiai.hf.space" in referer:
127
- return "hf_space_bypass"
 
128
 
129
- settings = get_settings()
130
- if not api_key or not api_key.startswith("Bearer "):
131
- raise HTTPException(status_code=HTTP_403_FORBIDDEN, detail="Invalid authorization format.")
132
-
133
- key = api_key.split(" ")[1]
134
- if key not in settings.api_keys:
135
- raise HTTPException(status_code=HTTP_403_FORBIDDEN, detail="Invalid API key.")
136
- return key
137
-
138
- @lru_cache()
139
- def get_http_client() -> httpx.AsyncClient:
140
- return httpx.AsyncClient(timeout=60.0, limits=httpx.Limits(max_connections=200))
141
-
142
- # --- API Routers ---
143
- chat_router = APIRouter(tags=["AI Models"])
144
- image_router = APIRouter(tags=["AI Models"])
145
- usage_router = APIRouter(tags=["Server Administration"])
146
- utility_router = APIRouter(tags=["Utilities & Pages"])
147
-
148
- # --- Chat Completions Router ---
149
- @chat_router.post("/chat/completions")
150
- async def chat_completions(
151
- payload: ChatPayload,
152
  request: Request,
153
- api_key: str = Depends(get_api_key),
154
- client: httpx.AsyncClient = Depends(get_http_client)
155
- ):
156
- if not server_status["online"]:
157
- raise HTTPException(status_code=HTTP_503_SERVICE_UNAVAILABLE, detail="Server under maintenance.")
 
 
 
 
 
158
 
159
- settings = get_settings()
160
- usage_tracker.record_request(request, payload.model, "/chat/completions")
161
- endpoint, headers, path = get_api_details(payload.model, settings)
 
 
162
 
163
- async def stream_generator():
164
- try:
165
- async with client.stream("POST", f"{endpoint}{path}", json=payload.dict(), headers=headers) as response:
166
- response.raise_for_status()
167
- async for chunk in response.aiter_bytes():
168
- yield chunk
169
- except httpx.HTTPStatusError as e:
170
- print(f"Upstream error: {e.response.status_code} - {e.response.text}")
171
- yield json.dumps({"error": {"code": 502, "message": "Bad Gateway: Upstream service error."}}).encode()
172
- except Exception as e:
173
- print(f"Streaming error: {e}")
174
- yield json.dumps({"error": {"code": 500, "message": "An internal error occurred."}}).encode()
175
 
176
- return StreamingResponse(stream_generator(), media_type="text/event-stream")
 
 
 
 
 
177
 
178
- # --- Image Generation Router ---
179
- @image_router.post("/images/generations")
180
- async def images_generations(
181
- payload: ImageGenerationPayload,
182
- request: Request,
183
- api_key: str = Depends(get_api_key),
184
- client: httpx.AsyncClient = Depends(get_http_client)
185
- ):
186
- if not server_status["online"]:
187
- raise HTTPException(status_code=HTTP_503_SERVICE_UNAVAILABLE, detail="Server under maintenance.")
188
-
189
- if payload.model not in MODEL_SETS["image"]:
190
- raise HTTPException(status_code=400, detail=f"Image model '{payload.model}' not supported.")
191
-
192
- settings = get_settings()
193
- usage_tracker.record_request(request, payload.model, "/images/generations")
194
- endpoint, headers, _ = get_api_details(payload.model, settings)
195
 
196
- try:
197
- response = await client.post(endpoint, json=payload.dict(), headers=headers)
198
- response.raise_for_status()
199
- return JSONResponse(content=response.json())
200
- except httpx.HTTPStatusError as e:
201
- raise HTTPException(status_code=e.response.status_code, detail=e.response.json().get("detail", "Upstream error"))
202
- except httpx.RequestError as e:
203
- raise HTTPException(status_code=502, detail=f"Failed to connect to image service: {e}")
204
-
205
- # --- Usage & Health Router ---
206
- @usage_router.get("/usage", response_class=HTMLResponse)
207
- async def get_usage_dashboard(days: int = Query(7, ge=1, le=30)):
208
- summary = usage_tracker.get_usage_summary(days=days)
209
- # The generate_usage_html function from the previous version can be used here directly
210
- # It has been moved to a separate file or helper for cleanliness in a real app
211
- # For this example, it's defined below for completeness.
212
- from usage_dashboard_generator import generate_usage_html
213
- return HTMLResponse(content=generate_usage_html(summary))
214
-
215
- @usage_router.get("/health")
216
- async def health_check():
217
- return {"status": "healthy" if server_status["online"] else "unhealthy", "version": app.version}
218
 
219
- @usage_router.get("/models")
220
- async def get_models():
 
221
  try:
222
- with open(Path(__file__).parent / 'models.json', 'r') as f:
 
223
  return json.load(f)
224
- except Exception:
225
- raise HTTPException(status_code=500, detail="models.json not found or is invalid.")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
226
 
227
- # --- Utility & Pages Router ---
228
  @lru_cache(maxsize=10)
229
- def read_static_file(file_path):
 
230
  try:
231
- with open(Path(__file__).parent / file_path, "r", encoding="utf-8") as file:
232
  return file.read()
233
  except FileNotFoundError:
234
  return None
235
 
236
- @utility_router.get("/", response_class=HTMLResponse)
237
- async def root_page():
238
- return HTMLResponse(content=read_static_file("index.html") or "<h1>Not Found</h1>")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
239
 
240
- @utility_router.get("/playground", response_class=HTMLResponse)
241
- async def playground_page():
242
- return HTMLResponse(content=read_static_file("playground.html") or "<h1>Not Found</h1>")
243
-
244
- @utility_router.get("/image-playground", response_class=HTMLResponse)
245
- async def image_playground_page():
246
- return HTMLResponse(content=read_static_file("image-playground.html") or "<h1>Not Found</h1>")
247
-
248
- @utility_router.get("/scraper", response_class=PlainTextResponse)
249
- async def scrape_url(url: str = Query(..., description="URL to scrape")):
250
- if not cloudscraper:
251
- raise HTTPException(status_code=501, detail="Scraper library not installed.")
252
  try:
253
- scraper = cloudscraper.create_scraper()
254
- response = scraper.get(url)
255
- response.raise_for_status()
256
- return PlainTextResponse(content=response.text)
257
- except Exception as e:
258
- raise HTTPException(status_code=500, detail=f"Failed to scrape URL: {e}")
 
 
 
 
259
 
 
 
 
 
 
 
 
 
 
 
260
 
261
- # --- Main Application Setup ---
262
- app.add_middleware(GZipMiddleware, minimum_size=1000)
263
- app.add_middleware(
264
- CORSMiddleware,
265
- allow_origins=["*"],
266
- allow_credentials=True,
267
- allow_methods=["*"],
268
- allow_headers=["*"],
269
- )
 
 
 
 
270
 
271
- # Include all the organized routers
272
- app.include_router(chat_router, prefix="/api/v1")
273
- app.include_router(chat_router) # For legacy /chat/completions
274
- app.include_router(image_router, prefix="/api/v1")
275
- app.include_router(image_router) # For legacy /images/generations
276
- app.include_router(usage_router)
277
- app.include_router(utility_router)
278
 
279
- @app.on_event("startup")
280
- async def startup_event():
281
- # Pre-load settings and client to catch config errors early
 
 
 
282
  try:
283
- get_settings()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
284
  except Exception as e:
285
- print(f"FATAL: Could not load settings from environment variables. Error: {e}")
286
- # In a real app, you might want to exit here
287
- get_http_client()
288
- print("--- LokiAI Server Started ---")
289
- print(f"Version: {app.version}")
290
- print("Usage tracking is active and will save data periodically.")
291
 
292
- @app.on_event("shutdown")
293
- async def shutdown_event():
294
- client = get_http_client()
295
- await client.aclose()
296
- usage_tracker.save_data()
297
- print("--- LokiAI Server Shutdown Complete ---")
 
298
 
 
 
 
 
 
 
 
299
 
300
- # Helper for usage dashboard - in a real project, this would be in its own file
301
- # I'm creating it here to make the example self-contained
302
- if not (Path(__file__).parent / "usage_dashboard_generator.py").exists():
303
- with open(Path(__file__).parent / "usage_dashboard_generator.py", "w") as f:
304
- f.write('''
305
- import json
306
- import datetime
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
307
 
308
- def generate_usage_html(usage_data: dict) -> str:
309
- model_labels = json.dumps(list(usage_data['model_usage'].keys()))
310
- model_values = json.dumps(list(usage_data['model_usage'].values()))
311
- daily_labels = json.dumps(list(usage_data['daily_usage'].keys()))
312
- daily_values = json.dumps([v['requests'] for v in usage_data['daily_usage'].values()])
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
313
 
314
- recent_requests_rows = "".join([
315
- f"""<tr>
316
- <td>{datetime.datetime.fromisoformat(req['timestamp']).strftime('%Y-%m-%d %H:%M:%S')}</td>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
317
  <td>{req['model']}</td>
318
  <td>{req['endpoint']}</td>
319
  <td>{req['ip_address']}</td>
320
- </tr>""" for req in usage_data['recent_requests']
 
 
321
  ])
322
 
323
- return f"""
324
  <!DOCTYPE html>
325
  <html lang="en">
326
  <head>
327
  <meta charset="UTF-8">
328
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
329
- <title>LokiAI - Usage Statistics</title>
330
- <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
331
- <link href="https://fonts.googleapis.com/css2?family=Inter:wght@400;600;700&display=swap" rel="stylesheet">
332
  <style>
333
- body {{ font-family: 'Inter', sans-serif; background-color: #0B0F19; color: #E0E0E0; margin: 0; padding: 20px; }}
334
- .container {{ max-width: 1400px; margin: auto; }}
335
- h1, h2 {{ color: #FFFFFF; }}
336
- .header {{ text-align: center; margin-bottom: 40px; }}
337
- .header h1 {{ font-size: 3em; font-weight: 700; }}
338
- .stats-grid {{ display: grid; grid-template-columns: repeat(auto-fit, minmax(280px, 1fr)); gap: 20px; margin-bottom: 40px; }}
339
- .chart-grid {{ display: grid; grid-template-columns: 1fr 1fr; gap: 20px; margin-bottom: 40px; }}
340
- .stat-card, .chart-container, .table-container {{ background: #1A2035; padding: 25px; border-radius: 12px; border: 1px solid #2A3045; }}
341
- .stat-card h3 {{ margin-top: 0; color: #8E95A9; font-size: 1em; font-weight: 600; text-transform: uppercase; }}
342
- .stat-card .value {{ font-size: 2.5em; font-weight: 700; color: #FFFFFF; }}
343
- table {{ width: 100%; border-collapse: collapse; }}
344
- th, td {{ padding: 14px; text-align: left; border-bottom: 1px solid #2A3045; }}
345
- th {{ background-color: #2A3045; font-weight: 600; }}
346
- @media (max-width: 768px) {{ .chart-grid {{ grid-template-columns: 1fr; }} }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
347
  </style>
348
  </head>
349
  <body>
350
  <div class="container">
351
- <div class="header"><h1>LokiAI Usage Dashboard</h1></div>
352
- <div class="stats-grid">
353
- <div class="stat-card"><h3>Total Requests</h3><p class="value">{usage_data['total_requests']}</p></div>
354
- <div class="stat-card"><h3>Unique IPs (All Time)</h3><p class="value">{usage_data['unique_ip_count']}</p></div>
355
- <div class="stat-card"><h3>Models Used (Last 7 Days)</h3><p class="value">{len(usage_data['model_usage'])}</p></div>
356
  </div>
357
- <div class="chart-grid">
358
- <div class="chart-container"><canvas id="dailyUsageChart"></canvas></div>
359
- <div class="chart-container"><canvas id="modelUsageChart"></canvas></div>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
360
  </div>
361
- <div class="table-container">
362
- <h2>Recent Requests</h2>
363
- <table>
364
- <thead><tr><th>Timestamp (UTC)</th><th>Model</th><th>Endpoint</th><th>IP Address</th></tr></thead>
365
- <tbody>{recent_requests_rows}</tbody>
366
- </table>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
367
  </div>
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
368
  </div>
 
369
  <script>
370
- const chartOptions = (ticksColor, gridColor) => ({{
371
- plugins: {{ legend: {{ labels: {{ color: ticksColor }} }} }},
372
- scales: {{
373
- y: {{ ticks: {{ color: ticksColor }}, grid: {{ color: gridColor }} }},
374
- x: {{ ticks: {{ color: ticksColor }}, grid: {{ color: 'transparent' }} }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
375
  }}
376
  }});
377
- new Chart(document.getElementById('dailyUsageChart'), {{
378
- type: 'line',
379
- data: {{ labels: {daily_labels}, datasets: [{{ label: 'Requests per Day', data: {daily_values}, borderColor: '#3a6ee0', tension: 0.1, backgroundColor: 'rgba(58, 110, 224, 0.2)', fill: true }}] }},
380
- options: chartOptions('#E0E0E0', '#2A3045')
381
- }});
382
- new Chart(document.getElementById('modelUsageChart'), {{
383
  type: 'doughnut',
384
- data: {{ labels: {model_labels}, datasets: [{{ label: 'Model Usage', data: {model_values}, backgroundColor: ['#3A6EE0', '#E94F37', '#44AF69', '#F4D35E', '#A06CD5'] }}] }},
385
- options: {{ plugins: {{ legend: {{ position: 'right', labels: {{ color: '#E0E0E0' }} }} }} }}
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
386
  }});
387
  </script>
388
  </body>
389
  </html>
390
  """
391
- ''')
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
392
 
393
  if __name__ == "__main__":
 
 
 
394
  uvicorn.run(app, host="0.0.0.0", port=7860)
395
 
 
1
  import os
 
 
 
2
  import re
 
 
 
 
 
 
3
  from dotenv import load_dotenv
4
+ from fastapi import FastAPI, HTTPException, Request, Depends, Security, Query
5
+ from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse, FileResponse, PlainTextResponse
6
  from fastapi.security import APIKeyHeader
7
+ from pydantic import BaseModel
8
+ import httpx
9
+ from functools import lru_cache
10
+ from pathlib import Path
11
+ import json
12
+ import datetime
13
+ import time
14
+ import threading
15
+ from typing import Optional, Dict, List, Any, Generator
16
+ import asyncio
17
+ from starlette.status import HTTP_403_FORBIDDEN
18
+ import cloudscraper
19
+ from concurrent.futures import ThreadPoolExecutor
20
+ import uvloop
21
+ from fastapi.middleware.gzip import GZipMiddleware
22
  from starlette.middleware.cors import CORSMiddleware
23
+ import contextlib
24
+ import requests
25
 
26
+ asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
 
 
 
27
 
28
+ executor = ThreadPoolExecutor(max_workers=16)
29
 
 
30
  load_dotenv()
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
31
 
32
+ api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
33
+
34
+ from usage_tracker import UsageTracker
35
+ usage_tracker = UsageTracker()
36
+
37
+ app = FastAPI()
38
+
39
+ app.add_middleware(GZipMiddleware, minimum_size=1000)
40
+ app.add_middleware(
41
+ CORSMiddleware,
42
+ allow_origins=["*"],
43
+ allow_credentials=True,
44
+ allow_methods=["*"],
45
+ allow_headers=["*"],
46
+ )
47
+
48
+ @lru_cache(maxsize=1)
49
+ def get_env_vars():
50
+ """
51
+ Loads and caches environment variables. This function is memoized
52
+ to avoid re-reading .env file on every call, improving performance.
53
+ """
54
+ return {
55
+ 'api_keys': os.getenv('API_KEYS', '').split(','),
56
+ 'secret_api_endpoint': os.getenv('SECRET_API_ENDPOINT'),
57
+ 'secret_api_endpoint_2': os.getenv('SECRET_API_ENDPOINT_2'),
58
+ 'secret_api_endpoint_3': os.getenv('SECRET_API_ENDPOINT_3'),
59
+ 'secret_api_endpoint_4': os.getenv('SECRET_API_ENDPOINT_4', "https://text.pollinations.ai/openai"),
60
+ 'secret_api_endpoint_5': os.getenv('SECRET_API_ENDPOINT_5'),
61
+ 'secret_api_endpoint_6': os.getenv('SECRET_API_ENDPOINT_6'), # New endpoint for Gemini
62
+ 'mistral_api': os.getenv('MISTRAL_API', "https://api.mistral.ai"),
63
+ 'mistral_key': os.getenv('MISTRAL_KEY'),
64
+ 'gemini_key': os.getenv('GEMINI_KEY'), # Gemini API Key
65
+ 'endpoint_origin': os.getenv('ENDPOINT_ORIGIN'),
66
+ 'new_img': os.getenv('NEW_IMG') # For image generation API
67
+ }
68
+
69
+ # Define sets of models for different API endpoints for easier routing
70
+ mistral_models = {
71
+ "mistral-large-latest", "pixtral-large-latest", "mistral-moderation-latest",
72
+ "ministral-3b-latest", "ministral-8b-latest", "open-mistral-nemo",
73
+ "mistral-small-latest", "mistral-saba-latest", "codestral-latest"
74
+ }
75
 
76
+ pollinations_models = {
77
+ "openai", "openai-large", "openai-fast", "openai-xlarge", "openai-reasoning",
78
+ "qwen-coder", "llama", "mistral", "searchgpt", "deepseek", "claude-hybridspace",
79
+ "deepseek-r1", "deepseek-reasoner", "llamalight", "gemini", "gemini-thinking",
80
+ "hormoz", "phi", "phi-mini", "openai-audio", "llama-scaleway"
81
+ }
82
+ alternate_models = {
83
+ "o1", "llama-4-scout", "o4-mini", "sonar", "sonar-pro", "sonar-reasoning",
84
+ "sonar-reasoning-pro", "grok-3", "grok-3-fast", "r1-1776", "o3"
85
+ }
86
+
87
+ claude_3_models = {
88
+ "claude-3-7-sonnet", "claude-3-7-sonnet-thinking", "claude 3.5 haiku",
89
+ "claude 3.5 sonnet", "claude 3.5 haiku", "o3-mini-medium", "o3-mini-high",
90
+ "grok-3", "grok-3-thinking", "grok 2"
91
+ }
92
+
93
+ gemini_models = {
94
+ "gemini-1.5-pro", "gemini-1.5-flash", "gemini-2.0-flash-lite-preview",
95
+ "gemini-2.0-flash", "gemini-2.0-flash-thinking", # aka Reasoning
96
+ "gemini-2.0-flash-preview-image-generation", "gemini-2.5-flash",
97
+ "gemini-2.5-pro-exp", "gemini-exp-1206"
98
+ }
99
 
100
+ supported_image_models = {
101
+ "Flux Pro Ultra", "grok-2-aurora", "Flux Pro", "Flux Pro Ultra Raw",
102
+ "Flux Dev", "Flux Schnell", "stable-diffusion-3-large-turbo",
103
+ "Flux Realism", "stable-diffusion-ultra", "dall-e-3", "sdxl-lightning-4step"
104
+ }
105
+
106
+ class Payload(BaseModel):
107
+ """Pydantic model for chat completion requests."""
108
  model: str
109
+ messages: list
110
  stream: bool = False
111
 
112
  class ImageGenerationPayload(BaseModel):
113
+ """Pydantic model for image generation requests."""
114
  model: str
115
  prompt: str
116
+ size: str = "1024x1024" # Default size, assuming models support it
117
+ number: int = 1
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
118
 
119
+ server_status = True # Global flag for server maintenance status
120
+ available_model_ids: List[str] = [] # List of all available model IDs
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
121
 
122
+ @lru_cache(maxsize=1)
123
+ def get_async_client():
124
+ """Returns a memoized httpx.AsyncClient instance for making async HTTP requests."""
125
+ return httpx.AsyncClient(
126
+ timeout=60.0,
127
+ limits=httpx.Limits(max_keepalive_connections=50, max_connections=200)
128
+ )
129
 
130
+ scraper_pool = []
131
+ MAX_SCRAPERS = 20
132
+
133
+ def get_scraper():
134
+ """Retrieves a cloudscraper instance from a pool for web scraping."""
135
+ if not scraper_pool:
136
+ # Initialize the pool if it's empty (should be done at startup)
137
+ for _ in range(MAX_SCRAPERS):
138
+ scraper_pool.append(cloudscraper.create_scraper())
139
+ # Simple round-robin selection from the pool
140
+ return scraper_pool[int(time.time() * 1000) % MAX_SCRAPERS]
141
+
142
+ async def verify_api_key(
 
 
 
 
 
 
 
 
 
 
143
  request: Request,
144
+ api_key: str = Security(api_key_header)
145
+ ) -> bool:
146
+ """
147
+ Verifies the API key provided in the Authorization header.
148
+ Allows access without API key if the request comes from specific Hugging Face spaces.
149
+ """
150
+ referer = request.headers.get("referer", "")
151
+ if referer.startswith(("https://parthsadaria-lokiai.hf.space/playground",
152
+ "https://parthsadaria-lokiai.hf.space/image-playground")):
153
+ return True
154
 
155
+ if not api_key:
156
+ raise HTTPException(
157
+ status_code=HTTP_403_FORBIDDEN,
158
+ detail="No API key provided"
159
+ )
160
 
161
+ if api_key.startswith('Bearer '):
162
+ api_key = api_key[7:]
 
 
 
 
 
 
 
 
 
 
163
 
164
+ valid_api_keys = get_env_vars().get('api_keys', [])
165
+ if not valid_api_keys or valid_api_keys == ['']:
166
+ raise HTTPException(
167
+ status_code=HTTP_403_FORBIDDEN,
168
+ detail="API keys not configured on server"
169
+ )
170
 
171
+ if api_key not in set(valid_api_keys):
172
+ raise HTTPException(
173
+ status_code=HTTP_403_FORBIDDEN,
174
+ detail="Invalid API key"
175
+ )
 
 
 
 
 
 
 
 
 
 
 
 
176
 
177
+ return True
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
178
 
179
+ @lru_cache(maxsize=1)
180
+ def load_models_data():
181
+ """Loads model data from 'models.json' and caches it."""
182
  try:
183
+ file_path = Path(__file__).parent / 'models.json'
184
+ with open(file_path, 'r') as f:
185
  return json.load(f)
186
+ except (FileNotFoundError, json.JSONDecodeError) as e:
187
+ print(f"Error loading models.json: {str(e)}")
188
+ return []
189
+
190
+ @app.get("/api/v1/models")
191
+ @app.get("/models")
192
+ async def get_models():
193
+ """Returns the list of available models."""
194
+ models_data = load_models_data()
195
+ if not models_data:
196
+ raise HTTPException(status_code=500, detail="Error loading available models")
197
+ return models_data
198
+
199
+ async def generate_search_async(query: str, systemprompt: Optional[str] = None, stream: bool = True):
200
+ """
201
+ Asynchronously generates a response using a search-based model.
202
+ Streams results if `stream` is True.
203
+ """
204
+ queue = asyncio.Queue()
205
+
206
+ async def _fetch_search_data():
207
+ """Internal helper to fetch data from the search API and put into queue."""
208
+ try:
209
+ headers = {"User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36"}
210
+ system_message = systemprompt or "Be Helpful and Friendly"
211
+ prompt = [{"role": "user", "content": query}]
212
+ prompt.insert(0, {"content": system_message, "role": "system"})
213
+ payload = {
214
+ "is_vscode_extension": True,
215
+ "message_history": prompt,
216
+ "requested_model": "searchgpt",
217
+ "user_input": prompt[-1]["content"],
218
+ }
219
+ secret_api_endpoint_3 = get_env_vars()['secret_api_endpoint_3']
220
+ if not secret_api_endpoint_3:
221
+ await queue.put({"error": "Search API endpoint not configured"})
222
+ return
223
+
224
+ async with httpx.AsyncClient(timeout=30.0) as client:
225
+ async with client.stream("POST", secret_api_endpoint_3, json=payload, headers=headers) as response:
226
+ if response.status_code != 200:
227
+ error_detail = await response.text()
228
+ await queue.put({"error": f"Search API returned status code {response.status_code}: {error_detail}"})
229
+ return
230
+
231
+ buffer = ""
232
+ async for line in response.aiter_lines():
233
+ if line.startswith("data: "):
234
+ try:
235
+ json_data = json.loads(line[6:])
236
+ content = json_data.get("choices", [{}])[0].get("delta", {}).get("content", "")
237
+ if content.strip():
238
+ cleaned_response = {
239
+ "created": json_data.get("created"),
240
+ "id": json_data.get("id"),
241
+ "model": "searchgpt",
242
+ "object": "chat.completion",
243
+ "choices": [
244
+ {
245
+ "message": {
246
+ "content": content
247
+ }
248
+ }
249
+ ]
250
+ }
251
+ await queue.put({"data": f"data: {json.dumps(cleaned_response)}\n\n", "text": content})
252
+ except json.JSONDecodeError:
253
+ # If line is not valid JSON, treat it as raw text and pass through if it's the end of stream
254
+ if line.strip() == "[DONE]":
255
+ continue # This is usually handled by the aiter_lines loop finishing
256
+ print(f"Warning: Could not decode JSON from search API stream: {line}")
257
+ await queue.put({"error": f"Invalid JSON from search API: {line}"})
258
+ break # Stop processing on bad JSON
259
+ await queue.put(None) # Signal end of stream
260
+ except Exception as e:
261
+ print(f"Error in _fetch_search_data: {e}")
262
+ await queue.put({"error": str(e)})
263
+ await queue.put(None)
264
+
265
+ asyncio.create_task(_fetch_search_data())
266
+ return queue
267
 
 
268
  @lru_cache(maxsize=10)
269
+ def read_html_file(file_path):
270
+ """Reads content of an HTML file and caches it."""
271
  try:
272
+ with open(file_path, "r") as file:
273
  return file.read()
274
  except FileNotFoundError:
275
  return None
276
 
277
+ # Static file routes for basic web assets
278
+ @app.get("/favicon.ico")
279
+ async def favicon():
280
+ favicon_path = Path(__file__).parent / "favicon.ico"
281
+ return FileResponse(favicon_path, media_type="image/x-icon")
282
+
283
+ @app.get("/banner.jpg")
284
+ async def banner():
285
+ banner_path = Path(__file__).parent / "banner.jpg"
286
+ return FileResponse(banner_path, media_type="image/jpeg")
287
+
288
+ @app.get("/ping")
289
+ async def ping():
290
+ """Simple health check endpoint."""
291
+ return {"message": "pong", "response_time": "0.000000 seconds"}
292
+
293
+ @app.get("/", response_class=HTMLResponse)
294
+ async def root():
295
+ """Serves the main index.html file."""
296
+ html_content = read_html_file("index.html")
297
+ if html_content is None:
298
+ raise HTTPException(status_code=404, detail="index.html not found")
299
+ return HTMLResponse(content=html_content)
300
+
301
+ @app.get("/script.js", response_class=HTMLResponse)
302
+ async def script():
303
+ """Serves script.js."""
304
+ html_content = read_html_file("script.js")
305
+ if html_content is None:
306
+ raise HTTPException(status_code=404, detail="script.js not found")
307
+ return HTMLResponse(content=html_content)
308
+
309
+ @app.get("/style.css", response_class=HTMLResponse)
310
+ async def style():
311
+ """Serves style.css."""
312
+ html_content = read_html_file("style.css")
313
+ if html_content is None:
314
+ raise HTTPException(status_code=404, detail="style.css not found")
315
+ return HTMLResponse(content=html_content)
316
+
317
+ @app.get("/dynamo", response_class=HTMLResponse)
318
+ async def dynamic_ai_page(request: Request):
319
+ """
320
+ Generates a dynamic HTML page using an AI model based on user-agent and IP.
321
+ Note: The hardcoded API endpoint and bearer token should ideally be managed
322
+ more securely, perhaps via environment variables and proper authentication.
323
+ """
324
+ user_agent = request.headers.get('user-agent', 'Unknown User')
325
+ client_ip = request.client.host if request.client else "Unknown IP"
326
+ location = f"IP: {client_ip}"
327
+
328
+ prompt = f"""
329
+ Generate a dynamic HTML page for a user with the following details: with name "LOKI.AI"
330
+ - User-Agent: {user_agent}
331
+ - Location: {location}
332
+ - Style: Cyberpunk, minimalist, or retro
333
+
334
+ Make sure the HTML is clean and includes a heading, also have cool animations a motivational message, and a cool background.
335
+ Wrap the generated HTML in triple backticks (```).
336
+ """
337
+
338
+ payload = {
339
+ "model": "mistral-small-latest",
340
+ "messages": [{"role": "user", "content": prompt}]
341
+ }
342
+
343
+ # Using the local /chat/completions endpoint for internal model call
344
+ # This assumes the current server can proxy to Mistral.
345
+ # For production, consider direct calls if not proxying is needed.
346
+ headers = {
347
+ "Authorization": "Bearer playground" # Use a dedicated internal token if available
348
+ }
349
 
 
 
 
 
 
 
 
 
 
 
 
 
350
  try:
351
+ # Use httpx.AsyncClient for making an async request
352
+ async with httpx.AsyncClient() as client:
353
+ response = await client.post(
354
+ f"http://localhost:7860/chat/completions", # Call self or internal API
355
+ json=payload,
356
+ headers=headers,
357
+ timeout=30.0
358
+ )
359
+ response.raise_for_status() # Raise an exception for bad status codes
360
+ data = response.json()
361
 
362
+ html_content = None
363
+ if data and 'choices' in data and len(data['choices']) > 0:
364
+ message_content = data['choices'][0].get('message', {}).get('content', '')
365
+ # Extract content within triple backticks
366
+ match = re.search(r"```(?:html)?(.*?)```", message_content, re.DOTALL)
367
+ if match:
368
+ html_content = match.group(1).strip()
369
+ else:
370
+ # Fallback: if no backticks, assume the whole content is HTML
371
+ html_content = message_content.strip()
372
 
373
+ if not html_content:
374
+ raise HTTPException(status_code=500, detail="Failed to generate HTML content from AI.")
375
+
376
+ return HTMLResponse(content=html_content)
377
+ except httpx.RequestError as e:
378
+ print(f"HTTPX Request Error in /dynamo: {e}")
379
+ raise HTTPException(status_code=500, detail=f"Failed to connect to internal AI service: {e}")
380
+ except httpx.HTTPStatusError as e:
381
+ print(f"HTTPX Status Error in /dynamo: {e.response.status_code} - {e.response.text}")
382
+ raise HTTPException(status_code=e.response.status_code, detail=f"Internal AI service responded with error: {e.response.text}")
383
+ except Exception as e:
384
+ print(f"An unexpected error occurred in /dynamo: {e}")
385
+ raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")
386
 
 
 
 
 
 
 
 
387
 
388
+ @app.get("/scraper", response_class=PlainTextResponse)
389
+ async def scrape_site(url: str = Query(..., description="URL to scrape")):
390
+ """
391
+ Scrapes the content of a given URL using cloudscraper.
392
+ Uses await in front of get_scraper().get() for async execution.
393
+ """
394
  try:
395
+ # get_scraper() returns a synchronous scraper object, but we are running
396
+ # it in an async endpoint. For CPU-bound tasks like this, it's better
397
+ # to offload to a thread pool to not block the event loop.
398
+ # However, cloudscraper's get method is typically synchronous.
399
+ # If cloudscraper were truly async, we'd use await.
400
+ # For now, running in executor to prevent blocking.
401
+ loop = asyncio.get_running_loop()
402
+ response_text = await loop.run_in_executor(
403
+ executor,
404
+ lambda: get_scraper().get(url).text
405
+ )
406
+
407
+ if response_text and len(response_text.strip()) > 0:
408
+ return PlainTextResponse(response_text)
409
+ else:
410
+ raise HTTPException(status_code=500, detail="Scraping returned empty content.")
411
  except Exception as e:
412
+ print(f"Cloudscraper failed: {e}")
413
+ raise HTTPException(status_code=500, detail=f"Cloudscraper failed: {e}")
 
 
 
 
414
 
415
+ @app.get("/playground", response_class=HTMLResponse)
416
+ async def playground():
417
+ """Serves the playground.html file."""
418
+ html_content = read_html_file("playground.html")
419
+ if html_content is None:
420
+ raise HTTPException(status_code=404, detail="playground.html not found")
421
+ return HTMLResponse(content=html_content)
422
 
423
+ @app.get("/image-playground", response_class=HTMLResponse)
424
+ async def image_playground():
425
+ """Serves the image-playground.html file."""
426
+ html_content = read_html_file("image-playground.html")
427
+ if html_content is None:
428
+ raise HTTPException(status_code=404, detail="image-playground.html not found")
429
+ return HTMLResponse(content=html_content)
430
 
431
+ GITHUB_BASE = "[https://raw.githubusercontent.com/Parthsadaria/Vetra/main](https://raw.githubusercontent.com/Parthsadaria/Vetra/main)"
432
+
433
+ FILES = {
434
+ "html": "index.html",
435
+ "css": "style.css",
436
+ "js": "script.js"
437
+ }
438
+
439
+ async def get_github_file(filename: str) -> Optional[str]:
440
+ """Fetches a file from a specified GitHub raw URL."""
441
+ url = f"{GITHUB_BASE}/{filename}"
442
+ async with httpx.AsyncClient() as client:
443
+ try:
444
+ res = await client.get(url, follow_redirects=True)
445
+ res.raise_for_status() # Raise an exception for HTTP errors (4xx or 5xx)
446
+ return res.text
447
+ except httpx.HTTPStatusError as e:
448
+ print(f"Error fetching {filename} from GitHub: {e.response.status_code} - {e.response.text}")
449
+ return None
450
+ except httpx.RequestError as e:
451
+ print(f"Request error fetching {filename} from GitHub: {e}")
452
+ return None
453
+
454
+ @app.get("/vetra", response_class=HTMLResponse)
455
+ async def serve_vetra():
456
+ """
457
+ Serves a dynamic HTML page by fetching HTML, CSS, and JS from GitHub
458
+ and embedding them into a single HTML response.
459
+ """
460
+ html = await get_github_file(FILES["html"])
461
+ css = await get_github_file(FILES["css"])
462
+ js = await get_github_file(FILES["js"])
463
+
464
+ if not html:
465
+ raise HTTPException(status_code=404, detail="index.html not found on GitHub")
466
+
467
+ final_html = html.replace(
468
+ "</head>",
469
+ f"<style>{css or '/* CSS not found */'}</style></head>"
470
+ ).replace(
471
+ "</body>",
472
+ f"<script>{js or '// JS not found'}</script></body>"
473
+ )
474
+
475
+ return HTMLResponse(content=final_html)
476
+
477
+ @app.get("/searchgpt")
478
+ async def search_gpt(q: str, request: Request, stream: Optional[bool] = False, systemprompt: Optional[str] = None):
479
+ """
480
+ Endpoint for search-based AI completion.
481
+ Records usage and streams results.
482
+ """
483
+ if not q:
484
+ raise HTTPException(status_code=400, detail="Query parameter 'q' is required")
485
+
486
+ # Record usage for searchgpt endpoint
487
+ usage_tracker.record_request(request=request, model="searchgpt", endpoint="/searchgpt")
488
 
489
+ queue = await generate_search_async(q, systemprompt=systemprompt, stream=True)
490
+
491
+ if stream:
492
+ async def stream_generator():
493
+ """Generator for streaming search results."""
494
+ collected_text = ""
495
+ while True:
496
+ item = await queue.get()
497
+ if item is None:
498
+ break
499
+
500
+ if "error" in item:
501
+ # Yield error as a data event so client can handle it gracefully
502
+ yield f"data: {json.dumps({'error': item['error']})}\n\n"
503
+ break
504
+
505
+ if "data" in item:
506
+ yield item["data"]
507
+ collected_text += item.get("text", "")
508
+
509
+ return StreamingResponse(
510
+ stream_generator(),
511
+ media_type="text/event-stream"
512
+ )
513
+ else:
514
+ # Non-streaming response: collect all chunks and return as JSON
515
+ collected_text = ""
516
+ while True:
517
+ item = await queue.get()
518
+ if item is None:
519
+ break
520
+
521
+ if "error" in item:
522
+ raise HTTPException(status_code=500, detail=item["error"])
523
+
524
+ collected_text += item.get("text", "")
525
+
526
+ return JSONResponse(content={"response": collected_text})
527
+
528
+ header_url = os.getenv('HEADER_URL') # This variable should be configured in .env
529
+
530
+ @app.post("/chat/completions")
531
+ @app.post("/api/v1/chat/completions")
532
+ async def get_completion(payload: Payload, request: Request, authenticated: bool = Depends(verify_api_key)):
533
+ """
534
+ Proxies chat completion requests to various AI model endpoints based on the model specified in the payload.
535
+ Records usage and handles streaming responses.
536
+ """
537
+ if not server_status:
538
+ raise HTTPException(
539
+ status_code=503,
540
+ detail="Server is under maintenance. Please try again later."
541
+ )
542
+
543
+ model_to_use = payload.model or "gpt-4o-mini" # Default model
544
+
545
+ # Validate if the requested model is available
546
+ if available_model_ids and model_to_use not in set(available_model_ids):
547
+ raise HTTPException(
548
+ status_code=400,
549
+ detail=f"Model '{model_to_use}' is not available. Check /models for the available model list."
550
+ )
551
+
552
+ # Record usage before making the external API call
553
+ usage_tracker.record_request(request=request, model=model_to_use, endpoint="/chat/completions")
554
+
555
+ payload_dict = payload.dict()
556
+ payload_dict["model"] = model_to_use # Ensure the payload has the resolved model name
557
+
558
+ stream_enabled = payload_dict.get("stream", True) # Default to streaming if not specified
559
+
560
+ env_vars = get_env_vars()
561
+
562
+ endpoint = None
563
+ custom_headers = {}
564
+ target_url_path = "/v1/chat/completions" # Default path for OpenAI-like APIs
565
+
566
+ # Determine the correct endpoint and headers based on the model
567
+ if model_to_use in mistral_models:
568
+ endpoint = env_vars['mistral_api']
569
+ custom_headers = {
570
+ "Authorization": f"Bearer {env_vars['mistral_key']}"
571
+ }
572
+ elif model_to_use in pollinations_models:
573
+ endpoint = env_vars['secret_api_endpoint_4']
574
+ custom_headers = {} # Pollinations.ai might not require auth
575
+ elif model_to_use in alternate_models:
576
+ endpoint = env_vars['secret_api_endpoint_2']
577
+ custom_headers = {}
578
+ elif model_to_use in claude_3_models:
579
+ endpoint = env_vars['secret_api_endpoint_5']
580
+ custom_headers = {} # Assuming no specific auth needed for this proxy
581
+ elif model_to_use in gemini_models:
582
+ endpoint = env_vars['secret_api_endpoint_6']
583
+ if not endpoint:
584
+ raise HTTPException(status_code=500, detail="Gemini API endpoint (SECRET_API_ENDPOINT_6) not configured.")
585
+ if not env_vars['gemini_key']:
586
+ raise HTTPException(status_code=500, detail="GEMINI_KEY not configured for Gemini models.")
587
+ custom_headers = {
588
+ "Authorization": f"Bearer {env_vars['gemini_key']}"
589
+ }
590
+ target_url_path = "/chat/completions" # Gemini's specific path
591
+ else:
592
+ # Default fallback for other models (e.g., OpenAI compatible APIs)
593
+ endpoint = env_vars['secret_api_endpoint']
594
+ custom_headers = {
595
+ "Origin": header_url,
596
+ "Priority": "u=1, i",
597
+ "Referer": header_url
598
+ }
599
+
600
+ if not endpoint:
601
+ raise HTTPException(status_code=500, detail=f"No API endpoint configured for model: {model_to_use}")
602
+
603
+ print(f"Proxying request for model '{model_to_use}' to endpoint: {endpoint}{target_url_path}")
604
+
605
+ async def real_time_stream_generator():
606
+ """Generator to stream responses from the upstream API."""
607
+ try:
608
+ async with httpx.AsyncClient(timeout=60.0) as client:
609
+ # Stream the request to the upstream API
610
+ async with client.stream("POST", f"{endpoint}{target_url_path}", json=payload_dict, headers=custom_headers) as response:
611
+ # Handle non-2xx responses from the upstream API
612
+ if response.status_code >= 400:
613
+ error_messages = {
614
+ 400: "Bad request. Verify input data.",
615
+ 401: "Unauthorized. Invalid API key for upstream service.",
616
+ 403: "Forbidden. You do not have access to this resource on upstream.",
617
+ 404: "The requested resource was not found on upstream.",
618
+ 422: "Unprocessable entity. Check your payload for upstream API.",
619
+ 500: "Internal server error from upstream API."
620
+ }
621
+ detail_message = error_messages.get(response.status_code, f"Upstream error code: {response.status_code}")
622
+
623
+ # Attempt to read upstream error response body for more detail
624
+ try:
625
+ error_body = await response.aread()
626
+ error_json = json.loads(error_body.decode('utf-8'))
627
+ if 'error' in error_json and 'message' in error_json['error']:
628
+ detail_message += f" - Upstream detail: {error_json['error']['message']}"
629
+ elif 'detail' in error_json:
630
+ detail_message += f" - Upstream detail: {error_json['detail']}"
631
+ else:
632
+ detail_message += f" - Upstream raw: {error_body.decode('utf-8')[:200]}..." # Limit for logging
633
+ except (json.JSONDecodeError, UnicodeDecodeError):
634
+ detail_message += f" - Upstream raw: {error_body.decode('utf-8', errors='ignore')[:200]}..."
635
+
636
+ raise HTTPException(status_code=response.status_code, detail=detail_message)
637
+
638
+ # Yield each line from the upstream stream
639
+ async for line in response.aiter_lines():
640
+ if line:
641
+ yield line + "\n"
642
+ except httpx.TimeoutException:
643
+ raise HTTPException(status_code=504, detail="Request to upstream AI service timed out.")
644
+ except httpx.RequestError as e:
645
+ raise HTTPException(status_code=502, detail=f"Failed to connect to upstream AI service: {str(e)}")
646
+ except Exception as e:
647
+ # Re-raise HTTPException if it's already one, otherwise wrap in a 500
648
+ if isinstance(e, HTTPException):
649
+ raise e
650
+ print(f"An unexpected error occurred during chat completion proxy: {e}")
651
+ raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")
652
+
653
+ if stream_enabled:
654
+ return StreamingResponse(
655
+ real_time_stream_generator(),
656
+ media_type="text/event-stream",
657
+ headers={
658
+ "Content-Type": "text/event-stream",
659
+ "Cache-Control": "no-cache",
660
+ "Connection": "keep-alive",
661
+ "X-Accel-Buffering": "no" # Disable buffering for SSE
662
+ }
663
+ )
664
+ else:
665
+ # For non-streaming requests, collect all parts and return a single JSON response
666
+ response_content_lines = []
667
+ async for line in real_time_stream_generator():
668
+ response_content_lines.append(line)
669
+
670
+ full_response_text = "".join(response_content_lines)
671
+
672
+ # Parse the concatenated stream data. This often involves stripping "data: " prefix
673
+ # and combining JSON objects from each line.
674
+ parsed_data = []
675
+ for line in full_response_text.splitlines():
676
+ if line.startswith("data: "):
677
+ try:
678
+ parsed_data.append(json.loads(line[6:]))
679
+ except json.JSONDecodeError:
680
+ print(f"Warning: Could not decode JSON line in non-streaming response: {line}")
681
+
682
+ # Attempt to reconstruct a single coherent JSON response
683
+ # This logic might need refinement based on actual API response format for non-streaming
684
+ final_json_response = {}
685
+ if parsed_data:
686
+ # Example: For OpenAI-like API, you might want the last 'choices' part
687
+ # This is a simplification and might need adjustment for other APIs
688
+ if 'choices' in parsed_data[-1]:
689
+ final_json_response = parsed_data[-1]
690
+ else:
691
+ # Fallback: just return the list of parsed objects
692
+ final_json_response = {"response_parts": parsed_data}
693
+
694
+ if not final_json_response:
695
+ # If nothing was parsed, indicate an issue
696
+ raise HTTPException(status_code=500, detail="No valid JSON response received from upstream API for non-streaming request.")
697
+
698
+ return JSONResponse(content=final_json_response)
699
+
700
+ @app.post("/images/generations")
701
+ async def create_image(payload: ImageGenerationPayload, request: Request, authenticated: bool = Depends(verify_api_key)):
702
+ """
703
+ Proxies image generation requests to a dedicated image generation API.
704
+ Records usage.
705
+ """
706
+ if not server_status:
707
+ raise HTTPException(
708
+ status_code=503,
709
+ detail="Server is under maintenance. Please try again later."
710
+ )
711
+
712
+ if payload.model not in supported_image_models:
713
+ raise HTTPException(
714
+ status_code=400,
715
+ detail=f"Model '{payload.model}' is not supported for image generation. Supported models are: {', '.join(supported_image_models)}"
716
+ )
717
+
718
+ # Record usage for image generation endpoint
719
+ usage_tracker.record_request(request=request, model=payload.model, endpoint="/images/generations")
720
+
721
+ api_payload = {
722
+ "model": payload.model,
723
+ "prompt": payload.prompt,
724
+ "size": payload.size,
725
+ "n": payload.number # Often 'n' for number of images in APIs
726
+ }
727
+
728
+ target_api_url = get_env_vars().get('new_img') # Get the image API URL from env vars
729
+ if not target_api_url:
730
+ raise HTTPException(status_code=500, detail="Image generation API endpoint (NEW_IMG) not configured.")
731
+
732
+ try:
733
+ async with httpx.AsyncClient(timeout=60.0) as client:
734
+ response = await client.post(target_api_url, json=api_payload)
735
+
736
+ response.raise_for_status() # Raise an exception for bad status codes (4xx or 5xx)
737
+
738
+ return JSONResponse(content=response.json())
739
+
740
+ except httpx.TimeoutException:
741
+ raise HTTPException(status_code=504, detail="Image generation request timed out.")
742
+ except httpx.RequestError as e:
743
+ raise HTTPException(status_code=502, detail=f"Error connecting to image generation service: {e}")
744
+ except httpx.HTTPStatusError as e:
745
+ error_detail = e.response.json().get("detail", f"Image generation failed with status code: {e.response.status_code}")
746
+ raise HTTPException(status_code=e.response.status_code, detail=error_detail)
747
+ except Exception as e:
748
+ print(f"An unexpected error occurred during image generation: {e}")
749
+ raise HTTPException(status_code=500, detail=f"An unexpected error occurred during image generation: {e}")
750
+
751
+ @app.get("/usage")
752
+ async def get_usage_json(days: int = 7):
753
+ """
754
+ Returns the raw usage data as JSON.
755
+ Can specify the number of days for the summary.
756
+ """
757
+ return usage_tracker.get_usage_summary(days)
758
+
759
+ def generate_usage_html(usage_data: Dict[str, Any]):
760
+ """
761
+ Generates an HTML page to display usage statistics.
762
+ Includes tables for model, API endpoint usage, daily usage, and recent requests.
763
+ Also includes placeholders for Chart.js to render graphs.
764
+ """
765
+ # Prepare data for Chart.js
766
+ # Model Usage Chart Data
767
+ model_labels = list(usage_data['model_usage_period'].keys())
768
+ model_counts = list(usage_data['model_usage_period'].values())
769
 
770
+ # Endpoint Usage Chart Data
771
+ endpoint_labels = list(usage_data['endpoint_usage_period'].keys())
772
+ endpoint_counts = list(usage_data['endpoint_usage_period'].values())
773
+
774
+ # Daily Usage Chart Data
775
+ daily_dates = list(usage_data['daily_usage_period'].keys())
776
+ daily_requests = [data['requests'] for data in usage_data['daily_usage_period'].values()]
777
+ daily_unique_ips = [data['unique_ips_count'] for data in usage_data['daily_usage_period'].values()]
778
+
779
+ # Format table rows for HTML
780
+ model_usage_all_time_rows = "\n".join([
781
+ f"""
782
+ <tr>
783
+ <td>{model}</td>
784
+ <td>{stats['total_requests']}</td>
785
+ <td>{datetime.datetime.fromisoformat(stats['first_used']).strftime("%Y-%m-%d %H:%M")}</td>
786
+ <td>{datetime.datetime.fromisoformat(stats['last_used']).strftime("%Y-%m-%d %H:%M")}</td>
787
+ </tr>
788
+ """ for model, stats in usage_data['all_time_model_usage'].items()
789
+ ])
790
+
791
+ api_usage_all_time_rows = "\n".join([
792
+ f"""
793
+ <tr>
794
+ <td>{endpoint}</td>
795
+ <td>{stats['total_requests']}</td>
796
+ <td>{datetime.datetime.fromisoformat(stats['first_used']).strftime("%Y-%m-%d %H:%M")}</td>
797
+ <td>{datetime.datetime.fromisoformat(stats['last_used']).strftime("%Y-%m-%d %H:%M")}</td>
798
+ </tr>
799
+ """ for endpoint, stats in usage_data['all_time_endpoint_usage'].items()
800
+ ])
801
+
802
+ daily_usage_table_rows = "\n".join([
803
+ f"""
804
+ <tr>
805
+ <td>{date}</td>
806
+ <td>{data['requests']}</td>
807
+ <td>{data['unique_ips_count']}</td>
808
+ </tr>
809
+ """ for date, data in usage_data['daily_usage_period'].items()
810
+ ])
811
+
812
+ recent_requests_rows = "\n".join([
813
+ f"""
814
+ <tr>
815
+ <td>{datetime.datetime.fromisoformat(req['timestamp']).strftime("%Y-%m-%d %H:%M:%S")}</td>
816
  <td>{req['model']}</td>
817
  <td>{req['endpoint']}</td>
818
  <td>{req['ip_address']}</td>
819
+ <td>{req['user_agent']}</td>
820
+ </tr>
821
+ """ for req in usage_data['recent_requests']
822
  ])
823
 
824
+ html_content = f"""
825
  <!DOCTYPE html>
826
  <html lang="en">
827
  <head>
828
  <meta charset="UTF-8">
829
  <meta name="viewport" content="width=device-width, initial-scale=1.0">
830
+ <title>Lokiai AI - Usage Statistics</title>
831
+ <link href="[https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&display=swap](https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&display=swap)" rel="stylesheet">
832
+ <script src="[https://cdn.jsdelivr.net/npm/chart.js](https://cdn.jsdelivr.net/npm/chart.js)"></script>
833
  <style>
834
+ :root {{
835
+ --bg-dark: #0f1011;
836
+ --bg-darker: #070708;
837
+ --text-primary: #e6e6e6;
838
+ --text-secondary: #8c8c8c;
839
+ --border-color: #2c2c2c;
840
+ --accent-color: #3a6ee0;
841
+ --accent-hover: #4a7ef0;
842
+ --chart-bg-light: rgba(58, 110, 224, 0.2);
843
+ --chart-border-light: #3a6ee0;
844
+ }}
845
+ body {{
846
+ font-family: 'Inter', sans-serif;
847
+ background-color: var(--bg-dark);
848
+ color: var(--text-primary);
849
+ max-width: 1200px;
850
+ margin: 0 auto;
851
+ padding: 40px 20px;
852
+ line-height: 1.6;
853
+ }}
854
+ .logo {{
855
+ display: flex;
856
+ align-items: center;
857
+ justify-content: center;
858
+ margin-bottom: 30px;
859
+ }}
860
+ .logo h1 {{
861
+ font-weight: 700;
862
+ font-size: 2.8em;
863
+ color: var(--text-primary);
864
+ margin-left: 15px;
865
+ }}
866
+ .logo img {{
867
+ width: 70px;
868
+ height: 70px;
869
+ border-radius: 12px;
870
+ box-shadow: 0 5px 15px rgba(0,0,0,0.2);
871
+ }}
872
+ .container {{
873
+ background-color: var(--bg-darker);
874
+ border-radius: 16px;
875
+ padding: 30px;
876
+ box-shadow: 0 20px 50px rgba(0,0,0,0.4);
877
+ border: 1px solid var(--border-color);
878
+ }}
879
+ h2, h3 {{
880
+ color: var(--text-primary);
881
+ border-bottom: 2px solid var(--border-color);
882
+ padding-bottom: 12px;
883
+ margin-top: 40px;
884
+ margin-bottom: 25px;
885
+ font-weight: 600;
886
+ font-size: 1.8em;
887
+ }}
888
+ .summary-grid {{
889
+ display: grid;
890
+ grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
891
+ gap: 20px;
892
+ margin-bottom: 30px;
893
+ }}
894
+ .summary-card {{
895
+ background-color: var(--bg-dark);
896
+ border-radius: 10px;
897
+ padding: 20px;
898
+ text-align: center;
899
+ border: 1px solid var(--border-color);
900
+ box-shadow: 0 8px 20px rgba(0,0,0,0.2);
901
+ transition: transform 0.2s ease-in-out;
902
+ }}
903
+ .summary-card:hover {{
904
+ transform: translateY(-5px);
905
+ }}
906
+ .summary-card h3 {{
907
+ margin-top: 0;
908
+ font-size: 1.1em;
909
+ color: var(--text-secondary);
910
+ border-bottom: none;
911
+ padding-bottom: 0;
912
+ margin-bottom: 10px;
913
+ }}
914
+ .summary-card p {{
915
+ font-size: 2.2em;
916
+ font-weight: 700;
917
+ color: var(--accent-color);
918
+ margin: 0;
919
+ }}
920
+ table {{
921
+ width: 100%;
922
+ border-collapse: separate;
923
+ border-spacing: 0;
924
+ margin-bottom: 40px;
925
+ background-color: var(--bg-dark);
926
+ border-radius: 10px;
927
+ overflow: hidden;
928
+ box-shadow: 0 8px 20px rgba(0,0,0,0.2);
929
+ }}
930
+ th, td {{
931
+ border: 1px solid var(--border-color);
932
+ padding: 15px;
933
+ text-align: left;
934
+ transition: background-color 0.3s ease;
935
+ }}
936
+ th {{
937
+ background-color: #1a1a1a;
938
+ color: var(--text-primary);
939
+ font-weight: 600;
940
+ text-transform: uppercase;
941
+ font-size: 0.95em;
942
+ }}
943
+ tr:nth-child(even) {{
944
+ background-color: rgba(255,255,255,0.03);
945
+ }}
946
+ tr:hover {{
947
+ background-color: rgba(62,100,255,0.1);
948
+ }}
949
+ .chart-container {{
950
+ background-color: var(--bg-dark);
951
+ border-radius: 10px;
952
+ padding: 20px;
953
+ margin-bottom: 40px;
954
+ border: 1px solid var(--border-color);
955
+ box-shadow: 0 8px 20px rgba(0,0,0,0.2);
956
+ max-height: 400px; /* Limit chart height */
957
+ position: relative; /* For responsive canvas */
958
+ }}
959
+ canvas {{
960
+ max-width: 100% !important;
961
+ height: auto !important;
962
+ }}
963
+ @media (max-width: 768px) {{
964
+ body {{
965
+ padding: 20px 10px;
966
+ }}
967
+ .container {{
968
+ padding: 20px;
969
+ }}
970
+ .logo h1 {{
971
+ font-size: 2em;
972
+ }}
973
+ .summary-card p {{
974
+ font-size: 1.8em;
975
+ }}
976
+ h2, h3 {{
977
+ font-size: 1.5em;
978
+ }}
979
+ table {{
980
+ font-size: 0.85em;
981
+ }}
982
+ th, td {{
983
+ padding: 10px;
984
+ }}
985
+ }}
986
  </style>
987
  </head>
988
  <body>
989
  <div class="container">
990
+ <div class="logo">
991
+ <img src="data:image/svg+xml;base64,PHN2ZyB3aWR0aD0iMjAwIiBoZWlnaHQ9IjIwMCIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMC9zdmciPjxwYXRoIGQ9Ik0xMDAgMzVMNTAgOTBoMTAwWiIgZmlsbD0iIzNhNmVlMCIvPjxjaXJjbGUgY3g9IjEwMCIgY3k9IjE0MCIgcj0iMzAiIGZpbGw9IiMzYTZlZTAiLz48L3N2Zz4=" alt="Lokiai AI Logo">
992
+ <h1>Lokiai AI Usage</h1>
 
 
993
  </div>
994
+
995
+ <div class="summary-grid">
996
+ <div class="summary-card">
997
+ <h3>Total Requests (All Time)</h3>
998
+ <p>{usage_data['total_requests']}</p>
999
+ </div>
1000
+ <div class="summary-card">
1001
+ <h3>Unique IPs (All Time)</h3>
1002
+ <p>{usage_data['unique_ips_total_count']}</p>
1003
+ </div>
1004
+ <div class="summary-card">
1005
+ <h3>Models Used (Last {days} Days)</h3>
1006
+ <p>{len(usage_data['model_usage_period'])}</p>
1007
+ </div>
1008
+ <div class="summary-card">
1009
+ <h3>Endpoints Used (Last {days} Days)</h3>
1010
+ <p>{len(usage_data['endpoint_usage_period'])}</p>
1011
+ </div>
1012
  </div>
1013
+
1014
+ <h2>Daily Usage (Last {days} Days)</h2>
1015
+ <div class="chart-container">
1016
+ <canvas id="dailyRequestsChart"></canvas>
1017
+ </div>
1018
+ <table>
1019
+ <thead>
1020
+ <tr>
1021
+ <th>Date</th>
1022
+ <th>Requests</th>
1023
+ <th>Unique IPs</th>
1024
+ </tr>
1025
+ </thead>
1026
+ <tbody>
1027
+ {daily_usage_table_rows}
1028
+ </tbody>
1029
+ </table>
1030
+
1031
+ <h2>Model Usage (Last {days} Days)</h2>
1032
+ <div class="chart-container">
1033
+ <canvas id="modelUsageChart"></canvas>
1034
+ </div>
1035
+ <h3>Model Usage (All Time Details)</h3>
1036
+ <table>
1037
+ <thead>
1038
+ <tr>
1039
+ <th>Model</th>
1040
+ <th>Total Requests</th>
1041
+ <th>First Used</th>
1042
+ <th>Last Used</th>
1043
+ </tr>
1044
+ </thead>
1045
+ <tbody>
1046
+ {model_usage_all_time_rows}
1047
+ </tbody>
1048
+ </table>
1049
+
1050
+ <h2>API Endpoint Usage (Last {days} Days)</h2>
1051
+ <div class="chart-container">
1052
+ <canvas id="endpointUsageChart"></canvas>
1053
  </div>
1054
+ <h3>API Endpoint Usage (All Time Details)</h3>
1055
+ <table>
1056
+ <thead>
1057
+ <tr>
1058
+ <th>Endpoint</th>
1059
+ <th>Total Requests</th>
1060
+ <th>First Used</th>
1061
+ <th>Last Used</th>
1062
+ </tr>
1063
+ </thead>
1064
+ <tbody>
1065
+ {api_usage_all_time_rows}
1066
+ </tbody>
1067
+ </table>
1068
+
1069
+ <h2>Recent Requests (Last 20)</h2>
1070
+ <table>
1071
+ <thead>
1072
+ <tr>
1073
+ <th>Timestamp</th>
1074
+ <th>Model</th>
1075
+ <th>Endpoint</th>
1076
+ <th>IP Address</th>
1077
+ <th>User Agent</th>
1078
+ </tr>
1079
+ </thead>
1080
+ <tbody>
1081
+ {recent_requests_rows}
1082
+ </tbody>
1083
+ </table>
1084
  </div>
1085
+
1086
  <script>
1087
+ // Chart.js data and rendering logic
1088
+ const modelLabels = {json.dumps(model_labels)};
1089
+ const modelCounts = {json.dumps(model_counts)};
1090
+
1091
+ const endpointLabels = {json.dumps(endpoint_labels)};
1092
+ const endpointCounts = {json.dumps(endpoint_counts)};
1093
+
1094
+ const dailyDates = {json.dumps(daily_dates)};
1095
+ const dailyRequests = {json.dumps(daily_requests)};
1096
+ const dailyUniqueIps = {json.dumps(daily_unique_ips)};
1097
+
1098
+ // Model Usage Chart (Bar Chart)
1099
+ new Chart(document.getElementById('modelUsageChart'), {{
1100
+ type: 'bar',
1101
+ data: {{
1102
+ labels: modelLabels,
1103
+ datasets: [{{
1104
+ label: 'Requests',
1105
+ data: modelCounts,
1106
+ backgroundColor: 'var(--chart-bg-light)',
1107
+ borderColor: 'var(--chart-border-light)',
1108
+ borderWidth: 1,
1109
+ borderRadius: 5,
1110
+ }}]
1111
+ }},
1112
+ options: {{
1113
+ responsive: true,
1114
+ maintainAspectRatio: false,
1115
+ plugins: {{
1116
+ legend: {{
1117
+ labels: {{
1118
+ color: 'var(--text-primary)'
1119
+ }}
1120
+ }},
1121
+ title: {{
1122
+ display: true,
1123
+ text: 'Model Usage',
1124
+ color: 'var(--text-primary)'
1125
+ }}
1126
+ }},
1127
+ scales: {{
1128
+ x: {{
1129
+ ticks: {{
1130
+ color: 'var(--text-secondary)'
1131
+ }},
1132
+ grid: {{
1133
+ color: 'var(--border-color)'
1134
+ }}
1135
+ }},
1136
+ y: {{
1137
+ beginAtZero: true,
1138
+ ticks: {{
1139
+ color: 'var(--text-secondary)'
1140
+ }},
1141
+ grid: {{
1142
+ color: 'var(--border-color)'
1143
+ }}
1144
+ }}
1145
+ }}
1146
  }}
1147
  }});
1148
+
1149
+ // Endpoint Usage Chart (Doughnut Chart)
1150
+ new Chart(document.getElementById('endpointUsageChart'), {{
 
 
 
1151
  type: 'doughnut',
1152
+ data: {{
1153
+ labels: endpointLabels,
1154
+ datasets: [{{
1155
+ label: 'Requests',
1156
+ data: endpointCounts,
1157
+ backgroundColor: [
1158
+ '#3a6ee0', '#5b8bff', '#8dc4ff', '#b3d8ff', '#d0e8ff',
1159
+ '#FF6384', '#36A2EB', '#FFCE56', '#4BC0C0', '#9966FF'
1160
+ ],
1161
+ hoverOffset: 4
1162
+ }}]
1163
+ }},
1164
+ options: {{
1165
+ responsive: true,
1166
+ maintainAspectRatio: false,
1167
+ plugins: {{
1168
+ legend: {{
1169
+ position: 'right',
1170
+ labels: {{
1171
+ color: 'var(--text-primary)'
1172
+ }}
1173
+ }},
1174
+ title: {{
1175
+ display: true,
1176
+ text: 'API Endpoint Usage',
1177
+ color: 'var(--text-primary)'
1178
+ }}
1179
+ }}
1180
+ }}
1181
+ }});
1182
+
1183
+ // Daily Requests Chart (Line Chart)
1184
+ new Chart(document.getElementById('dailyRequestsChart'), {{
1185
+ type: 'line',
1186
+ data: {{
1187
+ labels: dailyDates,
1188
+ datasets: [
1189
+ {{
1190
+ label: 'Total Requests',
1191
+ data: dailyRequests,
1192
+ borderColor: 'var(--accent-color)',
1193
+ backgroundColor: 'rgba(58, 110, 224, 0.1)',
1194
+ fill: true,
1195
+ tension: 0.3
1196
+ }},
1197
+ {{
1198
+ label: 'Unique IPs',
1199
+ data: dailyUniqueIps,
1200
+ borderColor: '#FFCE56', // A distinct color for unique IPs
1201
+ backgroundColor: 'rgba(255, 206, 86, 0.1)',
1202
+ fill: true,
1203
+ tension: 0.3
1204
+ }}
1205
+ ]
1206
+ }},
1207
+ options: {{
1208
+ responsive: true,
1209
+ maintainAspectRatio: false,
1210
+ plugins: {{
1211
+ legend: {{
1212
+ labels: {{
1213
+ color: 'var(--text-primary)'
1214
+ }}
1215
+ }},
1216
+ title: {{
1217
+ display: true,
1218
+ text: 'Daily Requests and Unique IPs',
1219
+ color: 'var(--text-primary)'
1220
+ }}
1221
+ }},
1222
+ scales: {{
1223
+ x: {{
1224
+ ticks: {{
1225
+ color: 'var(--text-secondary)'
1226
+ }},
1227
+ grid: {{
1228
+ color: 'var(--border-color)'
1229
+ }}
1230
+ }},
1231
+ y: {{
1232
+ beginAtZero: true,
1233
+ ticks: {{
1234
+ color: 'var(--text-secondary)'
1235
+ }},
1236
+ grid: {{
1237
+ color: 'var(--border-color)'
1238
+ }}
1239
+ }}
1240
+ }}
1241
+ }}
1242
  }});
1243
  </script>
1244
  </body>
1245
  </html>
1246
  """
1247
+ return html_content
1248
+
1249
+ @app.get("/usage/page", response_class=HTMLResponse)
1250
+ async def usage_page(days: int = 7):
1251
+ """
1252
+ Serves a detailed HTML page with usage statistics and charts.
1253
+ The 'days' query parameter can be used to specify the reporting period for charts.
1254
+ """
1255
+ usage_data = usage_tracker.get_usage_summary(days=days)
1256
+ html_content = generate_usage_html(usage_data)
1257
+ return HTMLResponse(content=html_content)
1258
+
1259
+ @app.get("/meme")
1260
+ async def get_meme():
1261
+ """
1262
+ Fetches a random meme from meme-api.com and streams the image content.
1263
+ Handles potential errors during fetching.
1264
+ """
1265
+ try:
1266
+ client = get_async_client()
1267
+ response = await client.get("[https://meme-api.com/gimme](https://meme-api.com/gimme)")
1268
+ response.raise_for_status() # Raise an exception for bad status codes
1269
+ response_data = response.json()
1270
+
1271
+ meme_url = response_data.get("url")
1272
+ if not meme_url:
1273
+ raise HTTPException(status_code=404, detail="No meme URL found in response.")
1274
+
1275
+ # Stream the image content back to the client
1276
+ image_response = await client.get(meme_url, follow_redirects=True)
1277
+ image_response.raise_for_status()
1278
+
1279
+ async def stream_with_larger_chunks():
1280
+ """Streams binary data in larger chunks for efficiency."""
1281
+ chunks = []
1282
+ size = 0
1283
+ # Define a larger chunk size for better streaming performance
1284
+ chunk_size = 65536 # 64 KB
1285
+ async for chunk in image_response.aiter_bytes(chunk_size=chunk_size):
1286
+ chunks.append(chunk)
1287
+ size += len(chunk)
1288
+ if size >= chunk_size * 2: # Send chunks when accumulated size is significant
1289
+ yield b''.join(chunks)
1290
+ chunks = []
1291
+ size = 0
1292
+ if chunks: # Yield any remaining chunks
1293
+ yield b''.join(chunks)
1294
+
1295
+ return StreamingResponse(
1296
+ stream_with_larger_chunks(),
1297
+ media_type=image_response.headers.get("content-type", "image/png"), # Fallback to png
1298
+ headers={'Cache-Control': 'max-age=3600'} # Cache memes for 1 hour
1299
+ )
1300
+ except httpx.HTTPStatusError as e:
1301
+ print(f"Error fetching meme from upstream: {e.response.status_code} - {e.response.text}")
1302
+ raise HTTPException(status_code=e.response.status_code, detail=f"Failed to fetch meme: {e.response.text}")
1303
+ except httpx.RequestError as e:
1304
+ print(f"Request error fetching meme: {e}")
1305
+ raise HTTPException(status_code=502, detail=f"Could not connect to meme service: {e}")
1306
+ except Exception as e:
1307
+ print(f"An unexpected error occurred while getting meme: {e}")
1308
+ raise HTTPException(status_code=500, detail="Failed to retrieve meme due to an unexpected error.")
1309
+
1310
+ def load_model_ids(json_file_path: str) -> List[str]:
1311
+ """
1312
+ Loads model IDs from a JSON file.
1313
+ This helps in dynamically determining available models.
1314
+ """
1315
+ try:
1316
+ with open(json_file_path, 'r') as f:
1317
+ models_data = json.load(f)
1318
+ return [model['id'] for model in models_data if 'id' in model]
1319
+ except Exception as e:
1320
+ print(f"Error loading model IDs from {json_file_path}: {str(e)}")
1321
+ return []
1322
+
1323
+ @app.on_event("startup")
1324
+ async def startup_event():
1325
+ """
1326
+ Actions to perform on application startup:
1327
+ - Load available model IDs.
1328
+ - Initialize scraper pool.
1329
+ - Check for missing environment variables and issue warnings.
1330
+ """
1331
+ global available_model_ids
1332
+ # Load models from a local models.json file first
1333
+ available_model_ids = load_model_ids("models.json")
1334
+ print(f"Loaded {len(available_model_ids)} model IDs from models.json")
1335
+
1336
+ # Extend with hardcoded model lists for various providers
1337
+ available_model_ids.extend(list(pollinations_models))
1338
+ available_model_ids.extend(list(alternate_models))
1339
+ available_model_ids.extend(list(mistral_models))
1340
+ available_model_ids.extend(list(claude_3_models))
1341
+ available_model_ids.extend(list(gemini_models)) # Add Gemini models explicitly
1342
+
1343
+ # Remove duplicates and store as a set for faster lookups
1344
+ available_model_ids = list(set(available_model_ids))
1345
+ print(f"Total unique available models after merging: {len(available_model_ids)}")
1346
+
1347
+ # Initialize scraper pool
1348
+ for _ in range(MAX_SCRAPERS):
1349
+ scraper_pool.append(cloudscraper.create_scraper())
1350
+ print(f"Initialized Cloudscraper pool with {MAX_SCRAPERS} instances.")
1351
+
1352
+ # Environment variable check for critical services
1353
+ env_vars = get_env_vars()
1354
+ missing_vars = []
1355
+
1356
+ if not env_vars['api_keys'] or env_vars['api_keys'] == ['']:
1357
+ missing_vars.append('API_KEYS')
1358
+ if not env_vars['secret_api_endpoint']:
1359
+ missing_vars.append('SECRET_API_ENDPOINT')
1360
+ if not env_vars['secret_api_endpoint_2']:
1361
+ missing_vars.append('SECRET_API_ENDPOINT_2')
1362
+ if not env_vars['secret_api_endpoint_3']:
1363
+ missing_vars.append('SECRET_API_ENDPOINT_3')
1364
+ if not env_vars['secret_api_endpoint_4'] and any(model in pollinations_models for model in available_model_ids):
1365
+ missing_vars.append('SECRET_API_ENDPOINT_4 (Pollinations.ai)')
1366
+ if not env_vars['secret_api_endpoint_5'] and any(model in claude_3_models for model in available_model_ids):
1367
+ missing_vars.append('SECRET_API_ENDPOINT_5 (Claude 3.x)')
1368
+ if not env_vars['secret_api_endpoint_6'] and any(model in gemini_models for model in available_model_ids):
1369
+ missing_vars.append('SECRET_API_ENDPOINT_6 (Gemini)')
1370
+ if not env_vars['mistral_api'] and any(model in mistral_models for model in available_model_ids):
1371
+ missing_vars.append('MISTRAL_API')
1372
+ if not env_vars['mistral_key'] and any(model in mistral_models for model in available_model_ids):
1373
+ missing_vars.append('MISTRAL_KEY')
1374
+ if not env_vars['gemini_key'] and any(model in gemini_models for model in available_model_ids):
1375
+ missing_vars.append('GEMINI_KEY')
1376
+ if not env_vars['new_img'] and len(supported_image_models) > 0:
1377
+ missing_vars.append('NEW_IMG (Image Generation)')
1378
+
1379
+ if missing_vars:
1380
+ print(f"WARNING: The following critical environment variables are missing or empty: {', '.join(missing_vars)}")
1381
+ print("Some server functionality (e.g., specific AI models, image generation) may be limited or unavailable.")
1382
+ else:
1383
+ print("All critical environment variables appear to be configured.")
1384
+
1385
+ print("Server started successfully!")
1386
+
1387
+ @app.on_event("shutdown")
1388
+ async def shutdown_event():
1389
+ """
1390
+ Actions to perform on application shutdown:
1391
+ - Close HTTPX client.
1392
+ - Clear scraper pool.
1393
+ - Save usage data to disk.
1394
+ """
1395
+ client = get_async_client()
1396
+ await client.aclose() # Ensure the httpx client connection pool is closed
1397
+ scraper_pool.clear() # Clear the scraper pool
1398
+ usage_tracker.save_data() # Persist usage data on shutdown
1399
+ print("Server shutdown complete!")
1400
+
1401
+ @app.get("/health")
1402
+ async def health_check():
1403
+ """
1404
+ Provides a health check endpoint, reporting server status and missing critical environment variables.
1405
+ """
1406
+ env_vars = get_env_vars()
1407
+ missing_critical_vars = []
1408
+
1409
+ # Re-check critical environment variables for health status
1410
+ if not env_vars['api_keys'] or env_vars['api_keys'] == ['']:
1411
+ missing_critical_vars.append('API_KEYS')
1412
+ if not env_vars['secret_api_endpoint']:
1413
+ missing_critical_vars.append('SECRET_API_ENDPOINT')
1414
+ if not env_vars['secret_api_endpoint_2']:
1415
+ missing_critical_vars.append('SECRET_API_ENDPOINT_2')
1416
+ if not env_vars['secret_api_endpoint_3']:
1417
+ missing_critical_vars.append('SECRET_API_ENDPOINT_3')
1418
+ # Check for specific service endpoints only if corresponding models are configured/supported
1419
+ if not env_vars['secret_api_endpoint_4'] and any(model in pollinations_models for model in available_model_ids):
1420
+ missing_critical_vars.append('SECRET_API_ENDPOINT_4 (Pollinations.ai)')
1421
+ if not env_vars['secret_api_endpoint_5'] and any(model in claude_3_models for model in available_model_ids):
1422
+ missing_critical_vars.append('SECRET_API_ENDPOINT_5 (Claude 3.x)')
1423
+ if not env_vars['secret_api_endpoint_6'] and any(model in gemini_models for model in available_model_ids):
1424
+ missing_critical_vars.append('SECRET_API_ENDPOINT_6 (Gemini)')
1425
+ if not env_vars['mistral_api'] and any(model in mistral_models for model in available_model_ids):
1426
+ missing_critical_vars.append('MISTRAL_API')
1427
+ if not env_vars['mistral_key'] and any(model in mistral_models for model in available_model_ids):
1428
+ missing_critical_vars.append('MISTRAL_KEY')
1429
+ if not env_vars['gemini_key'] and any(model in gemini_models for model in available_model_ids):
1430
+ missing_critical_vars.append('GEMINI_KEY')
1431
+ if not env_vars['new_img'] and len(supported_image_models) > 0:
1432
+ missing_critical_vars.append('NEW_IMG (Image Generation)')
1433
+
1434
+ health_status = {
1435
+ "status": "healthy" if not missing_critical_vars else "unhealthy",
1436
+ "missing_env_vars": missing_critical_vars,
1437
+ "server_status": server_status, # Reports global server status flag
1438
+ "message": "Everything's lit! 🚀" if not missing_critical_vars else "Uh oh, some env vars are missing. 😬"
1439
+ }
1440
+ return JSONResponse(content=health_status)
1441
 
1442
  if __name__ == "__main__":
1443
+ import uvicorn
1444
+ # When running directly, ensure startup_event is called to load models and check env vars
1445
+ # uvicorn handles startup/shutdown events automatically when run with `uvicorn.run()`
1446
  uvicorn.run(app, host="0.0.0.0", port=7860)
1447