From b31f0f6e4e245f32a35949a4c9ca0a5e249eb5f9 Mon Sep 17 00:00:00 2001
From: Ashwin Bharambe <ashwin@meta.com>
Date: Mon, 16 Sep 2024 21:41:24 -0700
Subject: [PATCH] add control plane API

---
 llama_stack/apis/control_plane/__init__.py    |  7 ++++
 .../apis/control_plane/control_plane.py       | 35 +++++++++++++++++++
 llama_stack/stack.py                          |  2 ++
 3 files changed, 44 insertions(+)
 create mode 100644 llama_stack/apis/control_plane/__init__.py
 create mode 100644 llama_stack/apis/control_plane/control_plane.py

diff --git a/llama_stack/apis/control_plane/__init__.py b/llama_stack/apis/control_plane/__init__.py
new file mode 100644
index 000000000..5abb4e730
--- /dev/null
+++ b/llama_stack/apis/control_plane/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from .control_plane import *  # noqa: F401 F403
diff --git a/llama_stack/apis/control_plane/control_plane.py b/llama_stack/apis/control_plane/control_plane.py
new file mode 100644
index 000000000..db79e91cd
--- /dev/null
+++ b/llama_stack/apis/control_plane/control_plane.py
@@ -0,0 +1,35 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the terms described in the LICENSE file in
+# the root directory of this source tree.
+
+from datetime import datetime
+from typing import Any, List, Optional, Protocol
+
+from llama_models.schema_utils import json_schema_type, webmethod
+from pydantic import BaseModel
+
+
+@json_schema_type
+class ControlPlaneValue(BaseModel):
+    key: str
+    value: Any
+    expiration: Optional[datetime] = None
+
+
+@json_schema_type
+class ControlPlane(Protocol):
+    @webmethod(route="/control_plane/set")
+    async def set(
+        self, key: str, value: Any, expiration: Optional[datetime] = None
+    ) -> None: ...
+
+    @webmethod(route="/control_plane/get", method="GET")
+    async def get(self, key: str) -> Optional[ControlPlaneValue]: ...
+
+    @webmethod(route="/control_plane/delete")
+    async def delete(self, key: str) -> None: ...
+
+    @webmethod(route="/control_plane/range", method="GET")
+    async def range(self, start_key: str, end_key: str) -> List[ControlPlaneValue]: ...
diff --git a/llama_stack/stack.py b/llama_stack/stack.py
index f6c66d23b..226050679 100644
--- a/llama_stack/stack.py
+++ b/llama_stack/stack.py
@@ -6,6 +6,7 @@
 
 from llama_models.llama3.api.datatypes import *  # noqa: F403
 from llama_stack.apis.agents import *  # noqa: F403
+from llama_stack.apis.control_plane import *  # noqa: F403
 from llama_stack.apis.dataset import *  # noqa: F403
 from llama_stack.apis.evals import *  # noqa: F403
 from llama_stack.apis.inference import *  # noqa: F403
@@ -21,6 +22,7 @@ from llama_stack.apis.safety import *  # noqa: F403
 class LlamaStack(
     Inference,
     BatchInference,
+    ControlPlane,
     Agents,
     RewardScoring,
     Safety,