Added Cosmos state store (#152)

Author: krisnyl

image_processing/src/image_processing/requirements.txt

@@ -5,7 +5,7 @@ aiohttp==3.11.11
 aiosignal==1.3.2
 annotated-types==0.7.0
 anyio==4.8.0
-attrs==24.3.0
+attrs==25.1.0
 azure-ai-documentintelligence==1.0.0
 azure-ai-textanalytics==5.3.0
 azure-ai-vision-imageanalysis==1.0.0
@@ -15,7 +15,7 @@ azure-functions==1.21.3
 azure-identity==1.19.0
 azure-search==1.0.0b2
 azure-search-documents==11.6.0b8
-azure-storage-blob==12.24.0
+azure-storage-blob==12.24.1
 beautifulsoup4==4.12.3
 blis==0.7.11
 bs4==0.0.2
@@ -38,7 +38,7 @@ fsspec==2024.12.0
 h11==0.14.0
 httpcore==1.0.7
 httpx==0.28.1
-huggingface-hub==0.27.1
+huggingface-hub==0.28.0
 idna==3.10
 isodate==0.7.2
 jinja2==3.1.5
@@ -50,15 +50,15 @@ marisa-trie==1.2.1
 markdown-it-py==3.0.0
 markupsafe==3.0.2
 mdurl==0.1.2
-model2vec==0.3.7
+model2vec==0.3.8
 msal==1.31.1
 msal-extensions==1.2.0
 msrest==0.7.1
 multidict==6.1.0
 murmurhash==1.0.12
 numpy==1.26.4
 oauthlib==3.2.2
-openai==1.60.0
+openai==1.60.2
 openpyxl==3.1.5
 packaging==24.2
 pandas==2.2.3
@@ -67,7 +67,7 @@ portalocker==2.10.1
 preshed==3.0.9
 propcache==0.2.1
 pycparser==2.22 ; platform_python_implementation != 'PyPy'
-pydantic==2.10.5
+pydantic==2.10.6
 pydantic-core==2.27.2
 pygments==2.19.1
 pyjwt==2.10.1

text_2_sql/autogen/pyproject.toml

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

text_2_sql/autogen/src/autogen_text_2_sql/__init__.py

@@ -1,7 +1,7 @@
 # Copyright (c) Microsoft Corporation.
 # Licensed under the MIT License.
 from autogen_text_2_sql.autogen_text_2_sql import AutoGenText2Sql
-from autogen_text_2_sql.state_store import InMemoryStateStore
+from autogen_text_2_sql.state_store import InMemoryStateStore, CosmosStateStore
 
 from text_2_sql_core.payloads.interaction_payloads import (
     UserMessagePayload,
@@ -19,4 +19,5 @@
     "ProcessingUpdatePayload",
     "InteractionPayload",
     "InMemoryStateStore",
+    "CosmosStateStore",
 ]

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
 
 
 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,36 @@ 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, credential, partition_key=None):
+        client = CosmosClient(url=endpoint, credential=credential)
+        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,
+            }
+        )

text_2_sql/text_2_sql_core/src/text_2_sql_core/payloads/interaction_payloads.py

@@ -48,7 +48,7 @@ class PayloadBase(InteractionPayloadBase):
 class DismabiguationRequestsPayload(InteractionPayloadBase):
     class Body(InteractionPayloadBase):
         class DismabiguationRequest(InteractionPayloadBase):
-            assistant_question: str | None = Field(..., alias="AssistantQuestion")
+            assistant_question: str | None = Field(..., alias="assistantQuestion")
             user_choices: list[str] | None = Field(default=None, alias="userChoices")
 
         disambiguation_requests: list[DismabiguationRequest] | None = Field(