Update main.py

main.py

@@ -38,8 +38,8 @@ api_key_header = APIKeyHeader(name="Authorization", auto_error=False)
 
 from usage_tracker import UsageTracker
 usage_tracker = UsageTracker()
-
-app = FastAPI()
+# Disable the basic docs because we are better than that
+app = FastAPI(docs_url=None, redoc_url=None)
 
 app.add_middleware(GZipMiddleware, minimum_size=1000)
 app.add_middleware(
@@ -311,80 +311,6 @@ async def root():
         raise HTTPException(status_code=404, detail="index.html not found")
     return HTMLResponse(content=html_content)
 
-@app.get("/script.js", response_class=HTMLResponse)
-async def script():
-    html_content = read_html_file("script.js")
-    if html_content is None:
-        raise HTTPException(status_code=404, detail="script.js not found")
-    return HTMLResponse(content=html_content)
-
-@app.get("/style.css", response_class=HTMLResponse)
-async def style():
-    html_content = read_html_file("style.css")
-    if html_content is None:
-        raise HTTPException(status_code=404, detail="style.css not found")
-    return HTMLResponse(content=html_content)
-
-@app.get("/dynamo", response_class=HTMLResponse)
-async def dynamic_ai_page(request: Request):
-    user_agent = request.headers.get('user-agent', 'Unknown User')
-    client_ip = request.client.host if request.client else "Unknown IP"
-    location = f"IP: {client_ip}"
-
-    prompt = f"""
-    Generate a dynamic HTML page for a user with the following details: with name "LOKI.AI"
-    - User-Agent: {user_agent}
-    - Location: {location}
-    - Style: Cyberpunk, minimalist, or retro
-
-    Make sure the HTML is clean and includes a heading, also have cool animations a motivational message, and a cool background.
-    Wrap the generated HTML in triple backticks (```).
-    """
-
-    payload_data = {
-        "model": "mistral-small-latest",
-        "messages": [{"role": "user", "content": prompt}],
-        "stream": False
-    }
-
-    headers = {
-        "Authorization": "Bearer playground"
-    }
-
-    try:
-        client = get_async_client()
-        response = await client.post(
-            f"http://localhost:7860/chat/completions",
-            json=payload_data,
-            headers=headers,
-            timeout=30.0
-        )
-        response.raise_for_status()
-        data = response.json()
-
-        html_content = None
-        if data and 'choices' in data and len(data['choices']) > 0:
-            message_content = data['choices'][0].get('message', {}).get('content', '')
-            match = re.search(r"```(?:html)?(.*?)```", message_content, re.DOTALL)
-            if match:
-                html_content = match.group(1).strip()
-            else:
-                html_content = message_content.strip()
-
-        if not html_content:
-            raise HTTPException(status_code=500, detail="Failed to generate HTML content from AI.")
-        
-        return HTMLResponse(content=html_content)
-    except httpx.RequestError as e:
-        print(f"HTTPX Request Error in /dynamo: {e}")
-        raise HTTPException(status_code=500, detail=f"Failed to connect to internal AI service: {e}")
-    except httpx.HTTPStatusError as e:
-        print(f"HTTPX Status Error in /dynamo: {e.response.status_code} - {e.response.text}")
-        raise HTTPException(status_code=e.response.status_code, detail=f"Internal AI service responded with error: {e.response.text}")
-    except Exception as e:
-        print(f"An unexpected error occurred in /dynamo: {e}")
-        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {e}")
-
 
 @app.get("/scraper", response_class=PlainTextResponse)
 async def scrape_site(url: str = Query(..., description="URL to scrape")):
@@ -425,70 +351,69 @@ FILES = {
     "js": "script.js"
 }
 
-async def get_github_file(filename: str) -> str | None:
-    url = f"{GITHUB_BASE}/{filename}"
-    client = get_async_client()
-    try:
-        res = await client.get(url, follow_redirects=True)
-        res.raise_for_status()
-        return res.text
-    except httpx.HTTPStatusError as e:
-        print(f"Error fetching {filename} from GitHub: {e.response.status_code} - {e.response.text}")
-        return None
-    except httpx.RequestError as e:
-        print(f"Request error fetching {filename} from GitHub: {e}")
-        return None
-
-@app.get("/vetra", response_class=HTMLResponse)
-async def serve_vetra():
-    html = await get_github_file(FILES["html"])
-    css = await get_github_file(FILES["css"])
-    js = await get_github_file(FILES["js"])
-
-    if not html:
-        raise HTTPException(status_code=404, detail="index.html not found on GitHub")
-
-    final_html = html.replace(
-        "</head>",
-        f"<style>{css or '/* CSS not found */'}</style></head>"
-    ).replace(
-        "</body>",
-        f"<script>{js or '// JS not found'}</script></body>"
-    )
-
-    return HTMLResponse(content=final_html)
-
-@app.get("/searchgpt")
-async def search_gpt(q: str, request: Request, stream: bool = False, systemprompt: str | None = None):
-    if not q:
-        raise HTTPException(status_code=400, detail="Query parameter 'q' is required")
-
-    usage_tracker.record_request(request=request, model="searchgpt", endpoint="/searchgpt")
-
-    if stream:
-        queue = await generate_search_async(q, systemprompt=systemprompt, stream=True)
-        async def stream_generator():
-            collected_text = ""
-            while True:
-                item = await queue.get()
-                if item is None:
-                    break
-
-                if "error" in item:
-                    yield f"data: {json.dumps({'error': item['error']})}\n\n"
-                    break
-
-                if "data" in item:
-                    yield item["data"]
-                    collected_text += item.get("text", "")
-
-        return StreamingResponse(
-            stream_generator(),
-            media_type="text/event-stream"
-        )
-    else:
-        response_data = await generate_search_async(q, systemprompt=systemprompt, stream=False)
-        return JSONResponse(content=response_data)
+@app.get("/docs", include_in_schema=False)
+async def custom_documentation():
+    return HTMLResponse("""
+<!DOCTYPE html>
+<html>
+  <head>
+    <title>Loki API - The Good Stuff</title>
+    <meta charset="utf-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1" />
+    <link href="https://fonts.googleapis.com/css2?family=Inter:wght@300;400;600;700&display=swap" rel="stylesheet">
+    <style>
+      body {
+        margin: 0;
+        background: #0f0f13; /* Deep dark background */
+      }
+      /* Custom Scrollbar because default is cringe */
+      ::-webkit-scrollbar { width: 8px; height: 8px; }
+      ::-webkit-scrollbar-track { background: #0f0f13; }
+      ::-webkit-scrollbar-thumb { background: #3b82f6; border-radius: 4px; }
+      ::-webkit-scrollbar-thumb:hover { background: #8b5cf6; }
+      
+      /* Overriding Scalar variables for that LOKI aesthetic */
+      :root {
+        --scalar-color-1: #3b82f6; /* Blue */
+        --scalar-color-2: #8b5cf6; /* Purple */
+        --scalar-color-3: #10b981; /* Green */
+        --scalar-color-accent: #8b5cf6;
+        
+        --scalar-background-1: #0b0b0e;
+        --scalar-background-2: #141419;
+        --scalar-background-3: #1f1f26;
+        
+        --scalar-border-color: #2d2d3b;
+        --scalar-font: 'Inter', sans-serif;
+      }
+
+      /* Smooth fade in for the whole page */
+      body {
+        animation: fadeIn 0.8s ease-in-out;
+      }
+
+      @keyframes fadeIn {
+        from { opacity: 0; transform: translateY(10px); }
+        to { opacity: 1; transform: translateY(0); }
+      }
+      
+      /* Let's hide the 'Made with Scalar' footer because we're mysterious */
+      .scalar-footer { display: none !important; }
+    </style>
+  </head>
+  <body>
+    <script
+      id="api-reference"
+      data-url="/openapi.json"
+      data-proxy-url="https://proxy.scalar.com"
+      data-layout="modern"
+      data-theme="deepSpace"
+      data-show-sidebar="true"
+    ></script>
+    <script src="https://cdn.jsdelivr.net/npm/@scalar/api-reference"></script>
+  </body>
+</html>
+    """)
 
 header_url = os.getenv('HEADER_URL')
 @app.post("/chat/completions")