Update CosmosDBHandlers/cosmosChatHistoryHandler.py

CosmosDBHandlers/cosmosChatHistoryHandler.py

@@ -51,12 +51,12 @@ class ChatMemoryHandler():
             api_key=os.environ["AZURE_OPENAI_KEY"]
         )
 
-        self.database = self.cosmos_client.create_database_if_not_exists("TAL_ChatDB")
+        self.database = self.cosmos_client.create_database_if_not_exists("TAL_ChatData")
 
         # Container for chat history
         self.chat_container = self.database.create_container_if_not_exists(
             id="ChatHistory",
-            partition_key=PartitionKey(path="/sessionId"),
+            partition_key=PartitionKey(path="/functionUsed"),
             indexing_policy=self.indexing_policy,
             vector_embedding_policy=self.vector_embedding_policy
         )
@@ -64,7 +64,7 @@ class ChatMemoryHandler():
         # Container for SQL queries
         self.sql_container = self.database.create_container_if_not_exists(
             id="GeneratedQueries", 
-            partition_key=PartitionKey(path="/queryType")
+            partition_key=PartitionKey(path="/state")
         )
     
     async def _generate_embedding(self, query: str) -> List[float]:
@@ -91,14 +91,13 @@ class ChatMemoryHandler():
             self.logger.error(f"Failed to log chat interaction: {str(e)}")
 
 
-    async def log_sql_query(self, original_question: str, generated_sql: str):
+    async def log_sql_query(self, original_question: str, generated_sql: str, state: str="success"):
         try:
             sql_item = {
                 "id": str(uuid.uuid4()),
-                "queryType": "generated",
                 "originalQuestion": original_question,
                 "generatedSql": generated_sql,
-                "executionStatus": "generated",
+                "state": state,
                 "timestamp": datetime.now(timezone.utc).isoformat()
             }
             self.sql_container.create_item(body=sql_item)