From b5d958631e52fc67e3bd22c8d825f3fb5e005153 Mon Sep 17 00:00:00 2001
From: Mandlin Sarah <mandlinsarah@gmail.com>
Date: Tue, 3 Sep 2024 03:20:07 -0700
Subject: [PATCH] Add timeout and retries for HTTP requests in
 AgenticSystemClient

---
 llama_toolchain/agentic_system/client.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/llama_toolchain/agentic_system/client.py b/llama_toolchain/agentic_system/client.py
index 56428c425..7cf867620 100644
--- a/llama_toolchain/agentic_system/client.py
+++ b/llama_toolchain/agentic_system/client.py
@@ -12,6 +12,7 @@ from typing import AsyncGenerator
 import fire
 
 import httpx
+from httpx import Timeout, Retry
 
 from llama_models.llama3.api.datatypes import (
     BuiltinTool,
@@ -47,7 +48,7 @@ class AgenticSystemClient(AgenticSystem):
     async def create_agentic_system(
         self, request: AgenticSystemCreateRequest
     ) -> AgenticSystemCreateResponse:
-        async with httpx.AsyncClient() as client:
+        async with httpx.AsyncClient(timeout=Timeout(10.0), retries=Retry(3)) as client:
             response = await client.post(
                 f"{self.base_url}/agentic_system/create",
                 data=request.json(),
@@ -60,7 +61,7 @@ class AgenticSystemClient(AgenticSystem):
         self,
         request: AgenticSystemSessionCreateRequest,
     ) -> AgenticSystemSessionCreateResponse:
-        async with httpx.AsyncClient() as client:
+        async with httpx.AsyncClient(timeout=Timeout(10.0), retries=Retry(3)) as client:
             response = await client.post(
                 f"{self.base_url}/agentic_system/session/create",
                 data=request.json(),
@@ -73,13 +74,12 @@ class AgenticSystemClient(AgenticSystem):
         self,
         request: AgenticSystemTurnCreateRequest,
     ) -> AsyncGenerator:
-        async with httpx.AsyncClient() as client:
+        async with httpx.AsyncClient(timeout=Timeout(20.0), retries=Retry(3)) as client:
             async with client.stream(
                 "POST",
                 f"{self.base_url}/agentic_system/turn/create",
                 data=request.json(),
                 headers={"Content-Type": "application/json"},
-                timeout=20,
             ) as response:
                 async for line in response.aiter_lines():
                     if line.startswith("data:"):
@@ -182,3 +182,4 @@ def main(host: str, port: int):
 
 if __name__ == "__main__":
     fire.Fire(main)
+