From 4dfc42b65a413503bc6009a38d9d26c7fc209c76 Mon Sep 17 00:00:00 2001
From: Will Beebe <651833+WillBeebe@users.noreply.github.com>
Date: Sun, 28 Jul 2024 16:08:56 -0700
Subject: [PATCH 1/2] adding awaits

---
 src/agents/agent.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/agents/agent.py b/src/agents/agent.py
index 56435c8..d5afa70 100644
--- a/src/agents/agent.py
+++ b/src/agents/agent.py
@@ -34,7 +34,7 @@ def get_history(self):
         logger.debug("No storage manager found, returning empty history")
         return []
 
-    def generate_text(self, prompt: str) -> PromptResponse:
+    async def generate_text(self, prompt: str) -> PromptResponse:
         # generate_text_counter.add(1)
         logger.debug("Generating text for prompt: '%s'", prompt)
         past_messages = []
@@ -43,14 +43,14 @@ def generate_text(self, prompt: str) -> PromptResponse:
             # if self.storage_manager.get_past_messages_callback is not None:
             #     past_messages = self.storage_manager.get_past_messages_callback()
             # else:
-            past_messages = self.storage_manager.get_past_messages()
+            past_messages = await self.storage_manager.get_past_messages()
             logger.debug("Fetched %d past messages", len(past_messages))
         # todo: push down to core llm class, leave for now while scripting
 
         try:
             logger.debug("passing %d past messages", len(past_messages))
             if self.storage_manager is not None:
-                self.storage_manager.store_message("user", prompt)
+                await self.storage_manager.store_message("user", prompt)
             response = self.client.generate_text(prompt, past_messages, self.tools)
         except Exception as e:
             logger.error("Error generating text: %s", e, exc_info=True)
@@ -61,7 +61,7 @@ def generate_text(self, prompt: str) -> PromptResponse:
         if self.storage_manager is not None:
             try:
                 # translated = self._translate_response(response)
-                self.storage_manager.store_message("assistant", response.content)
+                await self.storage_manager.store_message("assistant", response.content)
             except Exception as e:
                 logger.error("Error storing messages: %s", e, exc_info=True)
                 raise e

From eb6e0e40f509417c3bcf4d5c7096a39e2570f210 Mon Sep 17 00:00:00 2001
From: Will <651833+WillBeebe@users.noreply.github.com>
Date: Sun, 28 Jul 2024 16:10:02 -0700
Subject: [PATCH 2/2] misc

---
 src/agents/agent.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/src/agents/agent.py b/src/agents/agent.py
index d5afa70..8cbbf66 100644
--- a/src/agents/agent.py
+++ b/src/agents/agent.py
@@ -55,7 +55,7 @@ async def generate_text(self, prompt: str) -> PromptResponse:
         except Exception as e:
             logger.error("Error generating text: %s", e, exc_info=True)
             if self.storage_manager is not None:
-                self.storage_manager.remove_last()
+                await self.storage_manager.remove_last()
             raise e
 
         if self.storage_manager is not None:
@@ -99,15 +99,15 @@ async def generate_text_stream(self,
         """Generates streaming text based on the given prompt and additional arguments."""
         past_messages = []
         if self.storage_manager is not None:
-            past_messages = self.storage_manager.get_past_messages()
+            past_messages = await self.storage_manager.get_past_messages()
             logger.debug("Fetched %d past messages", len(past_messages))
         if self.storage_manager is not None:
-            self.storage_manager.store_message("user", prompt)
+            await self.storage_manager.store_message("user", prompt)
         try:
             response = await self.client.generate_text_stream(prompt, past_messages, self.tools)
         except Exception as err:
             if self.storage_manager is not None:
-                self.storage_manager.remove_last()
+                await self.storage_manager.remove_last()
             raise err
 
         # TODO: can't do this with streaming. have to handle this in the API