diff --git a/aieng-eval-agents/aieng/agent_evals/async_client_manager.py b/aieng-eval-agents/aieng/agent_evals/async_client_manager.py index a1f36f4..6491127 100644 --- a/aieng-eval-agents/aieng/agent_evals/async_client_manager.py +++ b/aieng-eval-agents/aieng/agent_evals/async_client_manager.py @@ -5,11 +5,15 @@ """ import logging +from typing import TYPE_CHECKING from aieng.agent_evals.configs import Configs -from aieng.agent_evals.tools import ReadOnlySqlDatabase from langfuse import Langfuse -from openai import AsyncOpenAI +from langfuse.openai import AsyncOpenAI + + +if TYPE_CHECKING: + from aieng.agent_evals.tools import ReadOnlySqlDatabase logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") @@ -94,11 +98,15 @@ def openai_client(self) -> AsyncOpenAI: if self._openai_client is None: api_key = self.configs.openai_api_key.get_secret_value() - self._openai_client = AsyncOpenAI(api_key=api_key, base_url=self.configs.openai_base_url) + self._openai_client = AsyncOpenAI( + api_key=api_key, + base_url=self.configs.openai_base_url, + max_retries=0, # Using custom retry logic (tenacity) elsewhere + ) self._initialized = True return self._openai_client - def report_generation_db(self, agent_name: str = "ReportGenerationAgent") -> ReadOnlySqlDatabase: + def report_generation_db(self, agent_name: str = "ReportGenerationAgent") -> "ReadOnlySqlDatabase": """Get or create Report Generation database connection. Returns @@ -107,6 +115,8 @@ def report_generation_db(self, agent_name: str = "ReportGenerationAgent") -> Rea The Report Generation database connection instance. """ if self._report_generation_db is None: + from aieng.agent_evals.tools import ReadOnlySqlDatabase + if self.configs.report_generation_db is None: raise ValueError("Report Generation database configuration is missing.") @@ -118,7 +128,7 @@ def report_generation_db(self, agent_name: str = "ReportGenerationAgent") -> Rea return self._report_generation_db - def aml_db(self, agent_name: str = "FraudInvestigationAnalyst") -> ReadOnlySqlDatabase: + def aml_db(self, agent_name: str = "FraudInvestigationAnalyst") -> "ReadOnlySqlDatabase": """Get or create AML database connection. Returns @@ -127,6 +137,8 @@ def aml_db(self, agent_name: str = "FraudInvestigationAnalyst") -> ReadOnlySqlDa The Report Generation database connection instance. """ if self._aml_db is None: + from aieng.agent_evals.tools import ReadOnlySqlDatabase + if self.configs.aml_db is None: raise ValueError("AML database configuration is missing.")