Added Cosmos state store

Author: Kristian Nylund

text_2_sql/autogen/pyproject.toml

@@ -18,6 +18,8 @@ dependencies = [
     "sqlparse>=0.4.4",
     "nltk>=3.8.1",
     "cachetools>=5.5.1",
+    "azure-cosmos>=4.9.0",
+    "azure-identity>=1.19.0",
 ]
 
 [dependency-groups]

text_2_sql/autogen/src/autogen_text_2_sql/state_store.py

@@ -1,14 +1,15 @@
 from abc import ABC, abstractmethod
 from cachetools import TTLCache
-
+from azure.cosmos import CosmosClient, exceptions
+from azure.identity import DefaultAzureCredential
 
 class StateStore(ABC):
     @abstractmethod
-    def get_state(self, thread_id):
+    def get_state(self, thread_id: str) -> dict:
         pass
 
     @abstractmethod
-    def save_state(self, thread_id, state):
+    def save_state(self, thread_id: str, state: dict) -> None:
         pass
 
 
@@ -21,3 +22,39 @@ def get_state(self, thread_id: str) -> dict:
 
     def save_state(self, thread_id: str, state: dict) -> None:
         self.cache[thread_id] = state
+
+
+class CosmosStateStore(StateStore):
+    def __init__(self, endpoint, database, container, partition_key = None):
+        client = CosmosClient(
+            url=endpoint,
+            credential=DefaultAzureCredential(),
+        )
+        database_client = client.get_database_client(database)
+        self._db = database_client.get_container_client(container)
+        self.partition_key = partition_key
+
+        # Set partition key field name
+        props = self._db.read()
+        pk_paths = props["partitionKey"]["paths"]
+        if (len(pk_paths) != 1):
+            raise ValueError("Only single partition key is supported")
+        self.partition_key_name = pk_paths[0].lstrip("/")
+        if ("/" in self.partition_key_name):
+            raise ValueError("Only top-level partition key is supported")
+
+    def get_state(self, thread_id: str) -> dict:
+        try:
+            item = self._db.read_item(item=thread_id, partition_key=self.partition_key)
+            return item["state"]
+        except exceptions.CosmosResourceNotFoundError:
+            return None
+
+    def save_state(self, thread_id: str, state: dict) -> None:
+        self._db.upsert_item(
+            body={
+                self.partition_key_name: self.partition_key,
+                "id": thread_id,
+                "state": state,
+            }
+        )