Skip to content

Commit

Permalink
API: Add log options for initialization
Browse files Browse the repository at this point in the history
Make each API log their respective URLs to help inform users.

Signed-off-by: kingbri <[email protected]>
  • Loading branch information
kingbri1 committed Jul 27, 2024
1 parent e8fc13a commit 884b6f5
Show file tree
Hide file tree
Showing 3 changed files with 30 additions and 10 deletions.
5 changes: 5 additions & 0 deletions endpoints/Kobold/router.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,12 @@
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(
Expand Down
5 changes: 5 additions & 0 deletions endpoints/OAI/router.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,7 +24,12 @@
)


api_name = "OAI"
router = APIRouter()
urls = {
"Completions": "http://{host}:{port}/v1/completions",
"Chat completions": "http://{host}:{port}/v1/chat/completions",
}


# Completions endpoint
Expand Down
30 changes: 20 additions & 10 deletions endpoints/server.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import asyncio
from typing import Optional
import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
Expand All @@ -8,12 +9,12 @@
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(
Expand Down Expand Up @@ -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)
Expand All @@ -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()
Expand Down

0 comments on commit 884b6f5

Please sign in to comment.