type annotations

Author: Kristian Nylund

text_2_sql/autogen/src/autogen_text_2_sql/state_store.py

@@ -15,8 +15,8 @@ class InMemoryStateStore(StateStore):
     def __init__(self):
         self.cache = TTLCache(maxsize=1000, ttl=4*60*60)  # 4 hours
 
-    def get_state(self, thread_id):
+    def get_state(self, thread_id: str) -> dict:
         return self.cache.get(thread_id)
 
-    def save_state(self, thread_id, state):
+    def save_state(self, thread_id: str, state: dict) -> None:
         self.cache[thread_id] = state