Spaces:

ParthSadaria
/

lokiai

Running

App Files Files Community

ParthSadaria commited on Nov 27, 2024

Commit

ea75284

verified ·

1 Parent(s): 67a227e

Update main.py

Browse files

Files changed (1) hide show

main.py +91 -4

main.py CHANGED Viewed

@@ -11,7 +11,8 @@ import cloudscraper
 import json
 from typing import Optional
 import datetime
 load_dotenv() #idk why this shi
 app = FastAPI()
@@ -105,7 +106,7 @@ def generate_search(query: str, systemprompt: Optional[str] = None, stream: bool
 async def search_gpt(q: str, stream: Optional[bool] = False, systemprompt: Optional[str] = None):
     if not q:
         raise HTTPException(status_code=400, detail="Query parameter 'q' is required")
     if stream:
         return StreamingResponse(
             generate_search(q, systemprompt=systemprompt, stream=True),
@@ -161,7 +162,7 @@ async def get_completion(payload: Payload, request: Request):
             status_code=400,
             detail=f"Model '{model_to_use}' is not available. Check /models for the available model list."
         )
     # Proceed with the request handling
     payload_dict = payload.dict()
     payload_dict["model"] = model_to_use
@@ -224,7 +225,7 @@ async def generate_image(
     # Validate the image endpoint
     if not image_endpoint:
         raise HTTPException(status_code=500, detail="Image endpoint not configured in environment variables.")
     # Handle GET and POST prompts
     if request.method == "POST":
         try:
@@ -330,6 +331,92 @@ def load_model_ids(json_file_path):
     except json.JSONDecodeError:
         print("Error: Invalid JSON format in models.json.")
         return []
 @app.on_event("startup")
 async def startup_event():
     global available_model_ids

 import json
 from typing import Optional
 import datetime
+from usage_tracker import UsageTracker
+usage_tracker = UsageTracker()
 load_dotenv() #idk why this shi
 app = FastAPI()
 async def search_gpt(q: str, stream: Optional[bool] = False, systemprompt: Optional[str] = None):
     if not q:
         raise HTTPException(status_code=400, detail="Query parameter 'q' is required")
+    usage_tracker.record_request(endpoint="/searchgpt")
     if stream:
         return StreamingResponse(
             generate_search(q, systemprompt=systemprompt, stream=True),
             status_code=400,
             detail=f"Model '{model_to_use}' is not available. Check /models for the available model list."
         )
+    usage_tracker.record_request(model=model_to_use, endpoint="/chat/completions")
     # Proceed with the request handling
     payload_dict = payload.dict()
     payload_dict["model"] = model_to_use
     # Validate the image endpoint
     if not image_endpoint:
         raise HTTPException(status_code=500, detail="Image endpoint not configured in environment variables.")
+    usage_tracker.record_request(endpoint="/images/generations")
     # Handle GET and POST prompts
     if request.method == "POST":
         try:
     except json.JSONDecodeError:
         print("Error: Invalid JSON format in models.json.")
         return []
+@app.get("/usage")
+async def get_usage(days: int = 7):
+    """Retrieve usage statistics"""
+    return usage_tracker.get_usage_summary(days)
+@app.get("/usage/page", response_class=HTMLResponse)
+async def usage_page():
+    """Serve an HTML page showing usage statistics"""
+    # Retrieve usage data
+    usage_data = usage_tracker.get_usage_summary()
+    # Create an HTML template with usage information
+    html_content = f"""
+    <!DOCTYPE html>
+    <html lang="en">
+    <head>
+        <meta charset="UTF-8">
+        <title>API Usage Statistics</title>
+        <style>
+            body {{ font-family: Arial, sans-serif; max-width: 1000px; margin: 0 auto; padding: 20px; }}
+            h1 {{ color: #333; }}
+            table {{ width: 100%; border-collapse: collapse; margin-bottom: 20px; }}
+            th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
+            th {{ background-color: #f2f2f2; }}
+        </style>
+    </head>
+    <body>
+        <h1>API Usage Statistics</h1>
+        <h2>Total Requests: {usage_data['total_requests']}</h2>
+        <h3>Model Usage</h3>
+        <table>
+            <tr>
+                <th>Model</th>
+                <th>Total Requests</th>
+                <th>First Used</th>
+                <th>Last Used</th>
+            </tr>
+            {"".join([f"""
+            <tr>
+                <td>{model}</td>
+                <td>{model_data['total_requests']}</td>
+                <td>{model_data['first_used']}</td>
+                <td>{model_data['last_used']}</td>
+            </tr>""" for model, model_data in usage_data['models'].items()])}
+        </table>
+        <h3>API Endpoint Usage</h3>
+        <table>
+            <tr>
+                <th>Endpoint</th>
+                <th>Total Requests</th>
+                <th>First Used</th>
+                <th>Last Used</th>
+            </tr>
+            {"".join([f"""
+            <tr>
+                <td>{endpoint}</td>
+                <td>{endpoint_data['total_requests']}</td>
+                <td>{endpoint_data['first_used']}</td>
+                <td>{endpoint_data['last_used']}</td>
+            </tr>""" for endpoint, endpoint_data in usage_data['api_endpoints'].items()])}
+        </table>
+        <h3>Daily Usage (Last 7 Days)</h3>
+        <table>
+            <tr>
+                <th>Date</th>
+                <th>Entity</th>
+                <th>Requests</th>
+            </tr>
+            {"".join([f"""
+            {"".join([f"""
+            <tr>
+                <td>{date}</td>
+                <td>{entity}</td>
+                <td>{requests}</td>
+            </tr>""" for entity, requests in date_data.items()])}
+            """ for date, date_data in usage_data['recent_daily_usage'].items()])}
+        </table>
+    </body>
+    </html>
+    """
+    return HTMLResponse(content=html_content)
 @app.on_event("startup")
 async def startup_event():
     global available_model_ids