Repository: run-llama/llama_deploy

Issue:
Refactor Python SDK
- [x] #290 
- [x] #335
- [x] #336
- [x] #337
Consolidate typing for responses in apiserver/control plane
There's a bit of confusion around types used internally, types used as payloads in the various HTTP api we have (apiserver, control plane...) and the more complex types used in the Python SDK implementation. On top of this, some portion of the code don't used types at all, relying on plain dictionaries (see [this comment](https://github.com/run-llama/llama_deploy/pull/327#pullrequestreview-2389063299)). 
Let's consolidate.


Discussion:

Patch:
diff --git a/llama_deploy/apiserver/deployment.py b/llama_deploy/apiserver/deployment.py
index 830c3365..251d1c53 100644
--- a/llama_deploy/apiserver/deployment.py
+++ b/llama_deploy/apiserver/deployment.py
@@ -7,7 +7,7 @@
 from typing import Any
 
 from llama_deploy import (
-    AsyncLlamaDeployClient,
+    Client,
     ControlPlaneServer,
     SimpleMessageQueue,
     SimpleOrchestrator,
@@ -61,7 +61,7 @@ def __init__(self, *, config: Config, root_path: Path) -> None:
             config=config.control_plane,
         )
         self._workflow_services: list[WorkflowService] = self._load_services(config)
-        self._client = AsyncLlamaDeployClient(config.control_plane)
+        self._client = Client(control_plane_url=config.control_plane.url)
         self._default_service = config.default_service
 
     @property
@@ -69,7 +69,7 @@ def default_service(self) -> str | None:
         return self._default_service
 
     @property
-    def client(self) -> AsyncLlamaDeployClient:
+    def client(self) -> Client:
         """Returns an async client to interact with this deployment."""
         return self._client
 
diff --git a/llama_deploy/apiserver/routers/deployments.py b/llama_deploy/apiserver/routers/deployments.py
index eec538f2..d65281b0 100644
--- a/llama_deploy/apiserver/routers/deployments.py
+++ b/llama_deploy/apiserver/routers/deployments.py
@@ -6,7 +6,8 @@
 
 from llama_deploy.apiserver.config_parser import Config
 from llama_deploy.apiserver.server import manager
-from llama_deploy.types import TaskDefinition
+from llama_deploy.types import DeploymentDefinition, SessionDefinition, TaskDefinition
+from llama_deploy.types.core import TaskResult
 
 deployments_router = APIRouter(
     prefix="/deployments",
@@ -14,40 +15,29 @@
 
 
 @deployments_router.get("/")
-async def read_deployments() -> JSONResponse:
+async def read_deployments() -> list[DeploymentDefinition]:
     """Returns a list of active deployments."""
-    return JSONResponse(
-        {
-            "deployments": list(manager._deployments.keys()),
-        }
-    )
+    return [DeploymentDefinition(name=k) for k in manager._deployments.keys()]
 
 
 @deployments_router.get("/{deployment_name}")
-async def read_deployment(deployment_name: str) -> JSONResponse:
+async def read_deployment(deployment_name: str) -> DeploymentDefinition:
     """Returns the details of a specific deployment."""
     if deployment_name not in manager.deployment_names:
         raise HTTPException(status_code=404, detail="Deployment not found")
 
-    return JSONResponse(
-        {
-            f"{deployment_name}": "Up!",
-        }
-    )
+    return DeploymentDefinition(name=deployment_name)
 
 
 @deployments_router.post("/create")
-async def create_deployment(config_file: UploadFile = File(...)) -> JSONResponse:
+async def create_deployment(
+    config_file: UploadFile = File(...),
+) -> DeploymentDefinition:
     """Creates a new deployment by uploading a configuration file."""
     config = Config.from_yaml_bytes(await config_file.read())
     manager.deploy(config)
 
-    # Return some details about the file
-    return JSONResponse(
-        {
-            "name": config.name,
-        }
-    )
+    return DeploymentDefinition(name=config.name)
 
 
 @deployments_router.post("/{deployment_name}/tasks/run")
@@ -67,7 +57,10 @@ async def create_deployment_task(
             )
         task_definition.agent_id = deployment.default_service
 
-    session = await deployment.client.get_or_create_session(session_id or "none")
+    if session_id:
+        session = await deployment.client.core.sessions.get(session_id)
+    else:
+        session = await deployment.client.core.sessions.create()
 
     result = await session.run(
         task_definition.agent_id or "", **json.loads(task_definition.input)
@@ -75,7 +68,7 @@ async def create_deployment_task(
 
     # Assume the request does not care about the session if no session_id is provided
     if session_id is None:
-        await deployment.client.delete_session(session.session_id)
+        await deployment.client.core.sessions.delete(session.id)
 
     return JSONResponse(result)
 
@@ -83,7 +76,7 @@ async def create_deployment_task(
 @deployments_router.post("/{deployment_name}/tasks/create")
 async def create_deployment_task_nowait(
     deployment_name: str, task_definition: TaskDefinition, session_id: str | None = None
-) -> JSONResponse:
+) -> TaskDefinition:
     """Create a task for the deployment but don't wait for result."""
     deployment = manager.get_deployment(deployment_name)
     if deployment is None:
@@ -97,13 +90,17 @@ async def create_deployment_task_nowait(
             )
         task_definition.agent_id = deployment.default_service
 
-    session = await deployment.client.get_or_create_session(session_id or "none")
+    if session_id:
+        session = await deployment.client.core.sessions.get(session_id)
+    else:
+        session = await deployment.client.core.sessions.create()
 
-    task_id = await session.run_nowait(
+    task_definition.session_id = session_id
+    task_definition.task_id = await session.run_nowait(
         task_definition.agent_id or "", **json.loads(task_definition.input)
     )
 
-    return JSONResponse({"session_id": session.session_id, "task_id": task_id})
+    return task_definition
 
 
 @deployments_router.get("/{deployment_name}/tasks/{task_id}/events")
@@ -115,7 +112,7 @@ async def get_events(
     if deployment is None:
         raise HTTPException(status_code=404, detail="Deployment not found")
 
-    session = await deployment.client.get_session(session_id)
+    session = await deployment.client.core.sessions.get(session_id)
 
     async def event_stream() -> AsyncGenerator[str, None]:
         # need to convert back to str to use SSE
@@ -131,76 +128,72 @@ async def event_stream() -> AsyncGenerator[str, None]:
 @deployments_router.get("/{deployment_name}/tasks/{task_id}/results")
 async def get_task_result(
     deployment_name: str, session_id: str, task_id: str
-) -> JSONResponse:
+) -> TaskResult | None:
     """Get the task result associated with a task and session."""
     deployment = manager.get_deployment(deployment_name)
     if deployment is None:
         raise HTTPException(status_code=404, detail="Deployment not found")
 
-    session = await deployment.client.get_session(session_id)
-    result = await session.get_task_result(task_id)
-
-    return JSONResponse(result.result if result else "")
+    session = await deployment.client.core.sessions.get(session_id)
+    return await session.get_task_result(task_id)
 
 
 @deployments_router.get("/{deployment_name}/tasks")
 async def get_tasks(
     deployment_name: str,
-) -> JSONResponse:
+) -> list[TaskDefinition]:
     """Get all the tasks from all the sessions in a given deployment."""
     deployment = manager.get_deployment(deployment_name)
     if deployment is None:
         raise HTTPException(status_code=404, detail="Deployment not found")
 
     tasks: list[TaskDefinition] = []
-    for session_def in await deployment.client.list_sessions():
-        session = await deployment.client.get_session(session_id=session_def.session_id)
+    for session in await deployment.client.core.sessions.list():
         for task_def in await session.get_tasks():
             tasks.append(task_def)
-    return JSONResponse(tasks)
+    return tasks
 
 
 @deployments_router.get("/{deployment_name}/sessions")
 async def get_sessions(
     deployment_name: str,
-) -> JSONResponse:
+) -> list[SessionDefinition]:
     """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)
+    sessions = await deployment.client.core.sessions.list()
+    return [SessionDefinition(session_id=s.id) for s in sessions]
 
 
 @deployments_router.get("/{deployment_name}/sessions/{session_id}")
-async def get_session(deployment_name: str, session_id: str) -> JSONResponse:
+async def get_session(deployment_name: str, session_id: str) -> SessionDefinition:
     """Get the definition of a session by ID."""
     deployment = manager.get_deployment(deployment_name)
     if deployment is None:
         raise HTTPException(status_code=404, detail="Deployment not found")
 
-    session = await deployment.client.get_session_definition(session_id)
-    return JSONResponse(session.model_dump())
+    session = await deployment.client.core.sessions.get(session_id)
+    return SessionDefinition(session_id=session.id)
 
 
 @deployments_router.post("/{deployment_name}/sessions/create")
-async def create_session(deployment_name: str) -> JSONResponse:
+async def create_session(deployment_name: str) -> SessionDefinition:
     """Create a new session for a deployment."""
     deployment = manager.get_deployment(deployment_name)
     if deployment is None:
         raise HTTPException(status_code=404, detail="Deployment not found")
 
-    session = await deployment.client.create_session()
-    return JSONResponse({"session_id": session.session_id})
+    session = await deployment.client.core.sessions.create()
+    return SessionDefinition(session_id=session.id)
 
 
 @deployments_router.post("/{deployment_name}/sessions/delete")
-async def delete_session(deployment_name: str, session_id: str) -> JSONResponse:
+async def delete_session(deployment_name: str, session_id: str) -> None:
     """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"})
+    await deployment.client.core.sessions.delete(session_id)
diff --git a/llama_deploy/apiserver/routers/status.py b/llama_deploy/apiserver/routers/status.py
index 54727e58..0bb6894f 100644
--- a/llama_deploy/apiserver/routers/status.py
+++ b/llama_deploy/apiserver/routers/status.py
@@ -1,8 +1,7 @@
 from fastapi import APIRouter
-from fastapi.responses import JSONResponse
 
 from llama_deploy.apiserver.server import manager
-
+from llama_deploy.types.apiserver import Status, StatusEnum
 
 status_router = APIRouter(
     prefix="/status",
@@ -10,11 +9,10 @@
 
 
 @status_router.get("/")
-async def status() -> JSONResponse:
-    return JSONResponse(
-        {
-            "status": "Up!",
-            "max_deployments": manager._max_deployments,
-            "deployments": list(manager._deployments.keys()),
-        }
+async def status() -> Status:
+    return Status(
+        status=StatusEnum.HEALTHY,
+        max_deployments=manager._max_deployments,
+        deployments=list(manager._deployments.keys()),
+        status_message="",
     )
diff --git a/llama_deploy/client/models/apiserver.py b/llama_deploy/client/models/apiserver.py
index e69a7950..2197d7f6 100644
--- a/llama_deploy/client/models/apiserver.py
+++ b/llama_deploy/client/models/apiserver.py
@@ -10,12 +10,6 @@
 from .model import Collection, Model
 
 
-class Session(Model):
-    """A model representing a session."""
-
-    pass
-
-
 class SessionCollection(Collection):
     """A model representing a collection of session for a given deployment."""
 
@@ -40,7 +34,7 @@ async def delete(self, session_id: str) -> None:
             timeout=self.client.timeout,
         )
 
-    async def create(self) -> Session:
+    async def create(self) -> SessionDefinition:
         """"""
         create_url = f"{self.client.api_server_url}/deployments/{self.deployment_id}/sessions/create"
 
@@ -51,12 +45,9 @@ async def create(self) -> Session:
             timeout=self.client.timeout,
         )
 
-        session_def = SessionDefinition(**r.json())
+        return SessionDefinition(**r.json())
 
-        model_class = self._prepare(Session)
-        return model_class(client=self.client, id=session_def.session_id)
-
-    async def list(self) -> list[Session]:  # type: ignore
+    async def list(self) -> list[SessionDefinition]:
         """Returns a collection of all the sessions in the given deployment."""
         sessions_url = (
             f"{self.client.api_server_url}/deployments/{self.deployment_id}/sessions"
@@ -67,12 +58,8 @@ async def list(self) -> list[Session]:  # type: ignore
             verify=not self.client.disable_ssl,
             timeout=self.client.timeout,
         )
-        model_class = self._prepare(Session)
-        items = [
-            model_class(client=self.client, id=session_def.session_id)
-            for session_def in r.json()
-        ]
-        return items
+
+        return r.json()
 
 
 class Task(Model):
@@ -163,13 +150,11 @@ async def create(self, task: TaskDefinition) -> Task:
             session_id=response_fields["session_id"],
         )
 
-
-class Deployment(Model):
-    """A model representing a deployment."""
-
-    async def tasks(self) -> TaskCollection:
-        """Returns a collection of tasks from all the sessions in the given deployment."""
-        tasks_url = f"{self.client.api_server_url}/deployments/{self.id}/tasks"
+    async def list(self) -> list[Task]:
+        """Returns the list of tasks from this collection."""
+        tasks_url = (
+            f"{self.client.api_server_url}/deployments/{self.deployment_id}/tasks"
+        )
         r = await self.client.request(
             "GET",
             tasks_url,
@@ -182,29 +167,32 @@ async def tasks(self) -> TaskCollection:
                 client=self.client,
                 id=task_def.task_id,
                 session_id=task_def.session_id,
-                deployment_id=self.id,
+                deployment_id=self.deployment_id,
             )
             for task_def in r.json()
         }
         model_class = self._prepare(TaskCollection)
-        return model_class(client=self.client, deployment_id=self.id, items=items)
+        return model_class(
+            client=self.client, deployment_id=self.deployment_id, items=items
+        )
+
 
-    async def sessions(self) -> SessionCollection:
+class Deployment(Model):
+    """A model representing a deployment."""
+
+    @property
+    def tasks(self) -> TaskCollection:
+        """Returns a collection of tasks from all the sessions in the given deployment."""
+
+        model_class = self._prepare(TaskCollection)
+        return model_class(client=self.client, deployment_id=self.id, items={})
+
+    @property
+    def sessions(self) -> SessionCollection:
         """Returns a collection of all the sessions in the given deployment."""
-        sessions_url = f"{self.client.api_server_url}/deployments/{self.id}/sessions"
-        r = await self.client.request(
-            "GET",
-            sessions_url,
-            verify=not self.client.disable_ssl,
-            timeout=self.client.timeout,
-        )
-        model_class = self._prepare(Session)
-        items = {
-            "id": model_class(client=self.client, id=session_def.session_id)
-            for session_def in r.json()
-        }
+
         coll_model_class = self._prepare(SessionCollection)
-        return coll_model_class(client=self.client, deployment_id=self.id, items=items)
+        return coll_model_class(client=self.client, deployment_id=self.id, items={})
 
 
 class DeploymentCollection(Collection):
diff --git a/llama_deploy/types/__init__.py b/llama_deploy/types/__init__.py
index e25a5ea7..d959bbf4 100644
--- a/llama_deploy/types/__init__.py
+++ b/llama_deploy/types/__init__.py
@@ -1,3 +1,4 @@
+from .apiserver import DeploymentDefinition, Status, StatusEnum
 from .core import (
     CONTROL_PLANE_NAME,
     ActionTypes,
@@ -34,4 +35,7 @@
     "ToolCallBundle",
     "ToolCallResult",
     "generate_id",
+    "DeploymentDefinition",
+    "Status",
+    "StatusEnum",
 ]
diff --git a/llama_deploy/types/apiserver.py b/llama_deploy/types/apiserver.py
index ea977d69..6996888e 100644
--- a/llama_deploy/types/apiserver.py
+++ b/llama_deploy/types/apiserver.py
@@ -14,3 +14,7 @@ class Status(BaseModel):
     status_message: str
     max_deployments: int | None = None
     deployments: list[str] | None = None
+
+
+class DeploymentDefinition(BaseModel):
+    name: str


Test patch:
diff --git a/tests/apiserver/routers/test_deployments.py b/tests/apiserver/routers/test_deployments.py
index 03dc555f..f3ed7478 100644
--- a/tests/apiserver/routers/test_deployments.py
+++ b/tests/apiserver/routers/test_deployments.py
@@ -1,20 +1,20 @@
 import json
-import pytest
 from pathlib import Path
 from unittest import mock
 
+import pytest
 from fastapi.testclient import TestClient
+from llama_index.core.workflow import Event
 
 from llama_deploy.apiserver import Config
 from llama_deploy.types import TaskResult
-
-from llama_index.core.workflow import Event
+from llama_deploy.types.core import TaskDefinition
 
 
 def test_read_deployments(http_client: TestClient) -> None:
     response = http_client.get("/deployments")
     assert response.status_code == 200
-    assert response.json() == {"deployments": []}
+    assert response.json() == []
 
 
 def test_read_deployment(http_client: TestClient) -> None:
@@ -25,7 +25,7 @@ def test_read_deployment(http_client: TestClient) -> None:
 
         response = http_client.get("/deployments/test-deployment")
         assert response.status_code == 200
-        assert response.json() == {"test-deployment": "Up!"}
+        assert response.json() == {"name": "test-deployment"}
 
         response = http_client.get("/deployments/does-not-exist")
         assert response.status_code == 404
@@ -90,16 +90,16 @@ def test_run_deployment_task(http_client: TestClient, data_path: Path) -> None:
         deployment = mock.AsyncMock()
         deployment.default_service = "TestService"
         session = mock.AsyncMock()
-        deployment.client.get_or_create_session.return_value = session
+        deployment.client.core.sessions.create.return_value = session
         session.run.return_value = {"result": "test_result"}
-        session.session_id = "42"
+        session.id = "42"
         mocked_manager.get_deployment.return_value = deployment
         response = http_client.post(
             "/deployments/test-deployment/tasks/run/",
             json={"input": "{}"},
         )
         assert response.status_code == 200
-        deployment.client.delete_session.assert_called_with("42")
+        deployment.client.core.sessions.delete.assert_called_with("42")
 
 
 def test_create_deployment_task(http_client: TestClient, data_path: Path) -> None:
@@ -109,8 +109,8 @@ def test_create_deployment_task(http_client: TestClient, data_path: Path) -> Non
         deployment = mock.AsyncMock()
         deployment.default_service = "TestService"
         session = mock.AsyncMock()
-        deployment.client.get_or_create_session.return_value = session
-        session.session_id = "42"
+        deployment.client.core.sessions.create.return_value = session
+        session.id = "42"
         session.run_nowait.return_value = "test_task_id"
         mocked_manager.get_deployment.return_value = deployment
         response = http_client.post(
@@ -118,7 +118,8 @@ def test_create_deployment_task(http_client: TestClient, data_path: Path) -> Non
             json={"input": "{}"},
         )
         assert response.status_code == 200
-        assert response.json() == {"session_id": "42", "task_id": "test_task_id"}
+        td = TaskDefinition(**response.json())
+        assert td.task_id == "test_task_id"
 
 
 @pytest.mark.asyncio
@@ -135,7 +136,7 @@ async def test_get_event_stream(http_client: TestClient, data_path: Path) -> Non
         deployment = mock.AsyncMock()
         deployment.default_service = "TestService"
         session = mock.MagicMock()
-        deployment.client.get_session.return_value = session
+        deployment.client.core.sessions.get.return_value = session
         mocked_manager.get_deployment.return_value = deployment
         mocked_get_task_result_stream = mock.MagicMock()
         mocked_get_task_result_stream.__aiter__.return_value = (
@@ -152,7 +153,7 @@ async def test_get_event_stream(http_client: TestClient, data_path: Path) -> Non
             data = json.loads(line)
             assert data == mock_events[ix].dict()
             ix += 1
-        deployment.client.get_session.assert_called_with("42")
+        deployment.client.core.sessions.get.assert_called_with("42")
         session.get_task_result_stream.assert_called_with("test_task_id")
 
 
@@ -163,7 +164,7 @@ def test_get_task_result(http_client: TestClient, data_path: Path) -> None:
         deployment = mock.AsyncMock()
         deployment.default_service = "TestService"
         session = mock.AsyncMock()
-        deployment.client.get_session.return_value = session
+        deployment.client.core.sessions.get.return_value = session
         session.get_task_result.return_value = TaskResult(
             result="test_result", history=[], task_id="test_task_id"
         )
@@ -173,9 +174,9 @@ def test_get_task_result(http_client: TestClient, data_path: Path) -> None:
             "/deployments/test-deployment/tasks/test_task_id/results/?session_id=42",
         )
         assert response.status_code == 200
-        assert response.json() == "test_result"
+        assert TaskResult(**response.json()).result == "test_result"
         session.get_task_result.assert_called_with("test_task_id")
-        deployment.client.get_session.assert_called_with("42")
+        deployment.client.core.sessions.get.assert_called_with("42")
 
 
 def test_get_sessions(http_client: TestClient, data_path: Path) -> None:
@@ -206,5 +207,4 @@ def test_delete_session(http_client: TestClient, data_path: Path) -> None:
             "/deployments/test-deployment/sessions/delete/?session_id=42",
         )
         assert response.status_code == 200
-        assert response.json() == {"session_id": "42", "status": "Deleted"}
-        deployment.client.delete_session.assert_called_with("42")
+        deployment.client.core.sessions.delete.assert_called_with("42")
diff --git a/tests/apiserver/routers/test_status.py b/tests/apiserver/routers/test_status.py
index 9b8a412b..78c86e54 100644
--- a/tests/apiserver/routers/test_status.py
+++ b/tests/apiserver/routers/test_status.py
@@ -7,5 +7,6 @@ def test_read_main(http_client: TestClient) -> None:
     assert response.json() == {
         "max_deployments": 10,
         "deployments": [],
-        "status": "Up!",
+        "status": "Healthy",
+        "status_message": "",
     }
diff --git a/tests/client/models/test_apiserver.py b/tests/client/models/test_apiserver.py
index b49a3e73..4e658a78 100644
--- a/tests/client/models/test_apiserver.py
+++ b/tests/client/models/test_apiserver.py
@@ -9,7 +9,6 @@
     ApiServer,
     Deployment,
     DeploymentCollection,
-    Session,
     SessionCollection,
     Task,
     TaskCollection,
@@ -21,7 +20,7 @@
 async def test_session_collection_delete(client: Any) -> None:
     coll = SessionCollection(
         client=client,
-        items={"a_session": Session(id="a_session", client=client)},
+        items={},
         deployment_id="a_deployment",
     )
     await coll.delete("a_session")
@@ -83,9 +82,9 @@ async def test_session_collection_list(client: Any) -> None:
 
     # Verify returned sessions
     assert len(sessions) == 2
-    assert all(isinstance(session, Session) for session in sessions)
-    assert sessions[0].id == "session1"
-    assert sessions[1].id == "session2"
+    assert all(isinstance(session, SessionDefinition) for session in sessions)
+    assert sessions[0].session_id == "session1"
+    assert sessions[1].session_id == "session2"
 
 
 @pytest.mark.asyncio
@@ -182,7 +181,7 @@ async def test_task_deployment_tasks(client: Any) -> None:
     ]
     client.request.return_value = mock.MagicMock(json=lambda: res)
 
-    await d.tasks()
+    await d.tasks.list()
 
     client.request.assert_awaited_with(
         "GET",
@@ -198,7 +197,7 @@ async def test_task_deployment_sessions(client: Any) -> None:
     res: list[SessionDefinition] = [SessionDefinition(session_id="a_session")]
     client.request.return_value = mock.MagicMock(json=lambda: res)
 
-    await d.sessions()
+    await d.sessions.list()
 
     client.request.assert_awaited_with(
         "GET",