Terminate aiohttp session if vector db exists

vocodedev · Aug 1, 2023 · ea1c504 · ea1c504
1 parent acc6069
commit ea1c504
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/vocode/streaming/streaming_conversation.py b/vocode/streaming/streaming_conversation.py
@@ -627,6 +627,9 @@ async def terminate(self):
         await self.synthesizer.tear_down()
         self.logger.debug("Terminating agent")
         self.agent.terminate()
+        if hasattr(self.agent.agent_config, "vector_db_config") and self.agent.agent_config.vector_db_config:
+            self.logger.debug("Terminating vector db")
+            self.agent.vector_db.tear_down()
         self.logger.debug("Terminating output device")
         self.output_device.terminate()
         self.logger.debug("Terminating speech transcriber")

diff --git a/vocode/streaming/vector_db/base_vector_db.py b/vocode/streaming/vector_db/base_vector_db.py
@@ -7,11 +7,17 @@ def __init__(self, aiohttp_session: Optional[aiohttp.ClientSession] = None,):
         if aiohttp_session:
             # the caller is responsible for closing the session
             self.aiohttp_session = aiohttp_session
+            self.should_close_session_on_tear_down = False
         else:
             self.aiohttp_session = aiohttp.ClientSession()
+            self.should_close_session_on_tear_down = True
 
     async def add_texts(self):
         raise NotImplementedError
 
     async def similarity_search_with_score(self):
         raise NotImplementedError
+
+    async def tear_down(self):
+        if self.should_close_session_on_tear_down:
+            await self.aiohttp_session.close()