From 884b6f5ecd726686e19508bc09b7345671a48052 Mon Sep 17 00:00:00 2001
From: kingbri <bdashore3@proton.me>
Date: Fri, 26 Jul 2024 21:32:05 -0400
Subject: [PATCH] API: Add log options for initialization

Make each API log their respective URLs to help inform users.

Signed-off-by: kingbri <bdashore3@proton.me>
---
 endpoints/Kobold/router.py |  5 +++++
 endpoints/OAI/router.py    |  5 +++++
 endpoints/server.py        | 30 ++++++++++++++++++++----------
 3 files changed, 30 insertions(+), 10 deletions(-)

diff --git a/endpoints/Kobold/router.py b/endpoints/Kobold/router.py
index c3265fc..6cfccf5 100644
--- a/endpoints/Kobold/router.py
+++ b/endpoints/Kobold/router.py
@@ -22,7 +22,12 @@ from endpoints.Kobold.utils.generation import (
 from endpoints.core.utils.model import get_current_model
 
 
+api_name = "KoboldAI"
 router = APIRouter(prefix="/api")
+urls = {
+    "Generation": "http://{host}:{port}/api/v1/generate",
+    "Streaming": "http://{host}:{port}/api/extra/generate/stream",
+}
 
 
 @router.post(
diff --git a/endpoints/OAI/router.py b/endpoints/OAI/router.py
index 771b7f3..55a2205 100644
--- a/endpoints/OAI/router.py
+++ b/endpoints/OAI/router.py
@@ -24,7 +24,12 @@ from endpoints.OAI.utils.completion import (
 )
 
 
+api_name = "OAI"
 router = APIRouter()
+urls = {
+    "Completions": "http://{host}:{port}/v1/completions",
+    "Chat completions": "http://{host}:{port}/v1/chat/completions",
+}
 
 
 # Completions endpoint
diff --git a/endpoints/server.py b/endpoints/server.py
index dfb1cdd..f8bfc8a 100644
--- a/endpoints/server.py
+++ b/endpoints/server.py
@@ -1,4 +1,5 @@
 import asyncio
+from typing import Optional
 import uvicorn
 from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
@@ -8,12 +9,12 @@ from common import config
 from common.logger import UVICORN_LOG_CONFIG
 from common.networking import get_global_depends
 from common.utils import unwrap
+from endpoints.Kobold import router as KoboldRouter
+from endpoints.OAI import router as OAIRouter
 from endpoints.core.router import router as CoreRouter
-from endpoints.Kobold.router import router as KoboldRouter
-from endpoints.OAI.router import router as OAIRouter
 
 
-def setup_app():
+def setup_app(host: Optional[str] = None, port: Optional[int] = None):
     """Includes the correct routers for startup"""
 
     app = FastAPI(
@@ -43,11 +44,20 @@ def setup_app():
     # Include the OAI api by default
     if api_servers:
         for server in api_servers:
-            server_name = server.lower()
-            if server_name in router_mapping:
-                app.include_router(router_mapping[server_name])
+            selected_server = router_mapping.get(server.lower())
+
+            if selected_server:
+                app.include_router(selected_server.router)
+
+                logger.info(f"Starting {selected_server.api_name} API")
+                for path, url in selected_server.urls.items():
+                    formatted_url = url.format(host=host, port=port)
+                    logger.info(f"{path}: {formatted_url}")
     else:
-        app.include_router(OAIRouter)
+        app.include_router(OAIRouter.router)
+        for path, url in OAIRouter.urls.items():
+            formatted_url = url.format(host=host, port=port)
+            logger.info(f"{path}: {formatted_url}")
 
     # Include core API request paths
     app.include_router(CoreRouter)
@@ -67,11 +77,11 @@ async def start_api(host: str, port: int):
 
     # TODO: Move OAI API to a separate folder
     logger.info(f"Developer documentation: http://{host}:{port}/redoc")
-    logger.info(f"Completions: http://{host}:{port}/v1/completions")
-    logger.info(f"Chat completions: http://{host}:{port}/v1/chat/completions")
+    # logger.info(f"Completions: http://{host}:{port}/v1/completions")
+    # logger.info(f"Chat completions: http://{host}:{port}/v1/chat/completions")
 
     # Setup app
-    app = setup_app()
+    app = setup_app(host, port)
 
     # Get the current event loop
     loop = asyncio.get_running_loop()