add a default session that gets reused for direct calls

Libertai · Jan 9, 2025 · b86fd33 · b86fd33
1 parent eebad33
commit b86fd33
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/libertai_agents/libertai_agents/agents.py b/libertai_agents/libertai_agents/agents.py
@@ -62,6 +62,7 @@ def __init__(
         self.system_prompt = system_prompt
         self.tools = tools
         self.llamacpp_params = llamacpp_params
+        self.call_session = None
 
         if expose_api:
             # Define API routes
@@ -159,6 +160,9 @@ async def __api_generate_answer(
         Generate an answer based on an existing conversation.
         The response messages can be streamed or sent in a single block.
         """
+        if self.call_session is None:
+            self.call_session = ClientSession()
+
         if stream:
             return StreamingResponse(
                 self.__dump_api_generate_streamed_answer(
@@ -169,7 +173,7 @@ async def __api_generate_answer(
 
         response_messages: list[Message] = []
         async for message in self.generate_answer(
-            messages, only_final_answer=only_final_answer
+            messages, only_final_answer=only_final_answer, session=self.call_session
         ):
             response_messages.append(message)
         return response_messages
@@ -184,8 +188,11 @@ async def __dump_api_generate_streamed_answer(
         :param only_final_answer: Param to pass to generate_answer
         :return: Iterable of each messages from generate_answer dumped to JSON
         """
+        if self.call_session is None:
+            self.call_session = ClientSession()
+
         async for message in self.generate_answer(
-            messages, only_final_answer=only_final_answer
+            messages, only_final_answer=only_final_answer, session=self.call_session
         ):
             yield json.dumps(message.model_dump(), indent=4)