add /models endpoint

Signed-off-by: Peter Jausovec <peter.jausovec@solo.io>

Author: peterj

python/packages/autogen-studio/autogenstudio/web/app.py

@@ -15,7 +15,7 @@
 from .config import settings
 from .deps import cleanup_managers, init_auth_manager, init_managers, register_auth_dependencies
 from .initialization import AppInitializer
-from .routes import gallery, runs, sessions, settingsroute, teams, tool_servers, tools, validation, ws
+from .routes import gallery, models, runs, sessions, settingsroute, teams, tool_servers, tools, validation, ws
 
 # Initialize application
 app_file_path = os.path.dirname(os.path.abspath(__file__))
@@ -155,6 +155,13 @@ async def lifespan(app: FastAPI) -> AsyncGenerator[None, None]:
     responses={404: {"description": "Not found"}},
 )
 
+api.include_router(
+    models.router,
+    prefix="/models",
+    tags=["models"],
+    responses={404: {"description": "Not found"}},
+)
+
 # Version endpoint
 
 

python/packages/autogen-studio/autogenstudio/web/routes/models.py

@@ -0,0 +1,20 @@
+from typing import Dict, List
+
+from autogen_ext.models.anthropic._model_info import _MODEL_INFO as anthropic_models
+from autogen_ext.models.ollama._model_info import _MODEL_INFO as ollama_models
+from autogen_ext.models.openai._model_info import _MODEL_INFO as openai_models
+from fastapi import APIRouter
+
+router = APIRouter()
+
+
+@router.get("/")
+async def list_models() -> Dict[str, List[str]]:
+    response = {
+        "anthropic": list(anthropic_models.keys()),
+        "ollama": list(ollama_models.keys()),
+        "openAI": list(openai_models.keys()),
+        "azureOpenAI": list(openai_models.keys()),
+    }
+
+    return response

python/packages/autogen-studio/pyproject.toml

@@ -22,10 +22,10 @@ dependencies = [
     "pydantic",
     "pydantic-settings",
     "fastapi[standard]",
-    "typer", 
+    "typer",
     "aiofiles",
     "python-dotenv",
-    "websockets", 
+    "websockets",
     "sqlmodel",
     "psycopg",
     "alembic",
@@ -36,6 +36,7 @@ dependencies = [
     "autogen-agentchat>=0.4.9.2,<0.6",
     "autogen-ext[magentic-one, openai, azure]>=0.4.2,<0.6",
     "anthropic",
+    "ollama>=0.4.7",
 ]
 optional-dependencies = {web = ["fastapi", "uvicorn"], database = ["psycopg"]}