From 1c9bc2d1af2aee83056196b8b0953ebcff82f82f Mon Sep 17 00:00:00 2001
From: AlpinDale <alpindale@gmail.com>
Date: Mon, 4 Nov 2024 12:35:08 +0000
Subject: [PATCH] feat: add serviceinfo URI

---
 endpoints/OAI/router.py | 34 +++++++++++++++++++++++++++++++++-
 endpoints/server.py     |  2 +-
 2 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/endpoints/OAI/router.py b/endpoints/OAI/router.py
index b6a44c9..fe10a9d 100644
--- a/endpoints/OAI/router.py
+++ b/endpoints/OAI/router.py
@@ -1,5 +1,6 @@
 import asyncio
 from fastapi import APIRouter, Depends, HTTPException, Request
+from fastapi.responses import JSONResponse
 from sse_starlette import EventSourceResponse
 from sys import maxsize
 
@@ -29,13 +30,19 @@ from endpoints.OAI.utils.embeddings import get_embeddings
 
 api_name = "OAI"
 router = APIRouter()
+host = None
+port = None
+
 urls = {
     "Completions": "http://{host}:{port}/v1/completions",
     "Chat completions": "http://{host}:{port}/v1/chat/completions",
 }
 
 
-def setup():
+def setup(server_host: str = None, server_port: int = None):
+    global host, port
+    host = server_host
+    port = server_port
     return router
 
 
@@ -166,3 +173,28 @@ async def embeddings(request: Request, data: EmbeddingsRequest) -> EmbeddingsRes
     )
 
     return response
+
+@router.get("/.well-known/serviceinfo")
+async def service_info():
+    return JSONResponse(content={
+        "version": 0.1,
+        "software": {
+            "name": "TabbyAPI",
+            "repository": "https://github.com/theroyallab/tabbyAPI",
+            "homepage": "https://github.com/theroyallab/tabbyAPI",
+        },
+        "api": {
+            "openai": {
+                "name": "OpenAI API",
+                "base_url": f"http://{host}:{port}/v1",
+                "documentation": "https://theroyallab.github.io/tabbyAPI",
+                "version": 1
+            },
+            "koboldai": {
+                "name": "KoboldAI API",
+                "base_url": f"http://{host}:{port}/api",
+                "documentation": "https://theroyallab.github.io/tabbyAPI",
+                "version": 1
+            }
+        }
+    })
diff --git a/endpoints/server.py b/endpoints/server.py
index 3555a5b..90bf6c4 100644
--- a/endpoints/server.py
+++ b/endpoints/server.py
@@ -52,7 +52,7 @@ def setup_app(host: Optional[str] = None, port: Optional[int] = None):
         selected_server = router_mapping.get(server.lower())
 
         if selected_server:
-            app.include_router(selected_server.setup())
+            app.include_router(selected_server.setup(host, port))
 
             logger.info(f"Starting {selected_server.api_name} API")
             for path, url in selected_server.urls.items():