From 9544408601e27c792735c8b81a4589e389c113db Mon Sep 17 00:00:00 2001
From: Andrei Fajardo <andrei@nerdai.io>
Date: Wed, 9 Oct 2024 10:10:09 -0400
Subject: [PATCH] session endpoints

---
 llama_deploy/apiserver/routers/deployments.py | 24 +++++++++++++++++++
 tests/apiserver/routers/test_deployments.py   | 16 +++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/llama_deploy/apiserver/routers/deployments.py b/llama_deploy/apiserver/routers/deployments.py
index 8ff0ef34..fdd1c11c 100644
--- a/llama_deploy/apiserver/routers/deployments.py
+++ b/llama_deploy/apiserver/routers/deployments.py
@@ -137,3 +137,27 @@ async def create_deployment(config_file: UploadFile = File(...)) -> JSONResponse
             "name": config.name,
         }
     )
+
+
+@deployments_router.get("/{deployment_name}/sessions")
+async def get_sessions(
+    deployment_name: str,
+) -> JSONResponse:
+    """Get the active sessions in a deployment and service."""
+    deployment = manager.get_deployment(deployment_name)
+    if deployment is None:
+        raise HTTPException(status_code=404, detail="Deployment not found")
+
+    sessions = await deployment.client.list_sessions()
+    return JSONResponse(sessions)
+
+
+@deployments_router.get("/{deployment_name}/sessions/delete")
+async def delete_session(deployment_name: str, session_id: str) -> JSONResponse:
+    """Get the active sessions in a deployment and service."""
+    deployment = manager.get_deployment(deployment_name)
+    if deployment is None:
+        raise HTTPException(status_code=404, detail="Deployment not found")
+
+    await deployment.client.delete_session(session_id)
+    return JSONResponse({"session_id": session_id, "status": "Deleted"})
diff --git a/tests/apiserver/routers/test_deployments.py b/tests/apiserver/routers/test_deployments.py
index 5a6e68da..cbf09d24 100644
--- a/tests/apiserver/routers/test_deployments.py
+++ b/tests/apiserver/routers/test_deployments.py
@@ -177,3 +177,19 @@ def test_get_task_result(http_client: TestClient, data_path: Path) -> None:
         assert response.json() == "test_result"
         session.get_task_result.assert_called_with("test_task_id")
         deployment.client.get_session.assert_called_with("42")
+
+
+def test_get_sessions(http_client: TestClient, data_path: Path) -> None:
+    with mock.patch(
+        "llama_deploy.apiserver.routers.deployments.manager"
+    ) as mocked_manager:
+        deployment = mock.AsyncMock()
+        deployment.default_service = "TestService"
+        deployment.client.list_sessions.return_value = []
+        mocked_manager.get_deployment.return_value = deployment
+
+        response = http_client.get(
+            "/deployments/test-deployment/sessions/",
+        )
+        assert response.status_code == 200
+        assert response.json() == []