Added litellm, models handler

Author: 3coins

jupyter_ai_litellm/chat_models_rest_api.py

@@ -0,0 +1,29 @@
+from jupyter_server.base.handlers import APIHandler as BaseAPIHandler
+from pydantic import BaseModel
+from tornado import web
+
+from .model_list import CHAT_MODELS
+
+
+class ChatModelsRestAPI(BaseAPIHandler):
+    """
+    A Tornado handler that defines the REST API served on the
+    `/api/ai/models/chat` endpoint.
+
+    - `GET /api/ai/models/chat`: returns list of all chat models.
+
+    - `GET /api/ai/models/chat?id=<model_id>`: returns info on that model (TODO)
+    """
+
+    @web.authenticated
+    def get(self):
+        response = ListChatModelsResponse(chat_models=CHAT_MODELS)
+        self.finish(response.model_dump_json())
+
+
+class ListChatModelsResponse(BaseModel):
+    chat_models: list[str]
+
+
+class ListEmbeddingModelsResponse(BaseModel):
+    embedding_models: list[str]

jupyter_ai_litellm/handlers.py

@@ -4,6 +4,8 @@
 from jupyter_server.utils import url_path_join
 import tornado
 
+from .chat_models_rest_api import ChatModelsRestAPI
+
 class RouteHandler(APIHandler):
     # The following decorator should be present on all verb methods (head, get, post,
     # patch, put, delete, options) to ensure only authorized user can request the
@@ -19,6 +21,10 @@ def setup_handlers(web_app):
     host_pattern = ".*$"
 
     base_url = web_app.settings["base_url"]
+    print(f"Base url is {base_url}")
     route_pattern = url_path_join(base_url, "jupyter-ai-litellm", "get-example")
-    handlers = [(route_pattern, RouteHandler)]
+    handlers = [
+        (route_pattern, RouteHandler),
+        (url_path_join(base_url, "api/models/chat") + r"(?:\?.*)?", ChatModelsRestAPI)
+    ]
     web_app.add_handlers(host_pattern, handlers)

jupyter_ai_litellm/model_list.py

@@ -0,0 +1,36 @@
+from litellm import all_embedding_models, models_by_provider
+
+chat_model_ids = []
+embedding_model_ids = []
+embedding_model_set = set(all_embedding_models)
+
+for provider_name in models_by_provider:
+    for model_name in models_by_provider[provider_name]:
+        model_name: str = model_name
+
+        if model_name.startswith(f"{provider_name}/"):
+            model_id = model_name
+        else:
+            model_id = f"{provider_name}/{model_name}"
+
+        is_embedding = (
+            model_name in embedding_model_set
+            or model_id in embedding_model_set
+            or "embed" in model_id
+        )
+
+        if is_embedding:
+            embedding_model_ids.append(model_id)
+        else:
+            chat_model_ids.append(model_id)
+
+
+CHAT_MODELS = sorted(chat_model_ids)
+"""
+List of chat model IDs, following the `litellm` syntax.
+"""
+
+EMBEDDING_MODELS = sorted(embedding_model_ids)
+"""
+List of embedding model IDs, following the `litellm` syntax.
+"""

jupyter_ai_litellm/tests/test_handlers.py

@@ -10,4 +10,16 @@ async def test_get_example(jp_fetch):
     payload = json.loads(response.body)
     assert payload == {
         "data": "This is /jupyter-ai-litellm/get-example endpoint!"
-    }
+    }
+
+async def test_get_chat_models(jp_fetch):
+    # When
+    response = await jp_fetch("api", "models", "chat")
+
+    # Then
+    assert response.code == 200
+    payload = json.loads(response.body)
+    chat_models = payload.get("chat_models")
+
+    assert chat_models
+    assert len(chat_models) > 0

pyproject.toml

@@ -27,7 +27,8 @@ authors = [
     { name = "Project Jupyter", email = "[email protected]" },
 ]
 dependencies = [
-    "jupyter_server>=2.4.0,<3"
+    "jupyter_server>=2.4.0,<3",
+    "litellm>=1.73,<2",
 ]
 dynamic = ["version"]