Merge pull request #75 from MarketSquare/chore/update_dep_versions

Update dependency versions.
Use pydantic-ai-slim (instead of full package)
Use named arguments for UsageLimits
Adjust Tests

Author: manykarim

SelfhealingAgents/self_healing_system/agents/locator_agent/base_locator_agent.py

@@ -1,21 +1,27 @@
-from typing import Optional
 from abc import ABC, abstractmethod
+from typing import Optional
 
-from robot.api import logger as rf_logger
-from pydantic_ai.usage import UsageLimits
-from pydantic_ai.agent import AgentRunResult
 from pydantic_ai import Agent, ModelRetry, RunContext
+from pydantic_ai.agent import AgentRunResult
+from pydantic_ai.usage import UsageLimits
+from robot.api import logger as rf_logger
 
-from SelfhealingAgents.utils.cfg import Cfg
-from SelfhealingAgents.utils.logging import log
-from SelfhealingAgents.self_healing_system.llm.client_model import get_client_model
-from SelfhealingAgents.self_healing_system.schemas.api.locator_healing import LocatorHealingResponse
-from SelfhealingAgents.self_healing_system.schemas.internal_state.prompt_payload import PromptPayload
-from SelfhealingAgents.self_healing_system.context_retrieving.library_dom_utils.base_dom_utils import BaseDomUtils
 from SelfhealingAgents.self_healing_system.agents.prompts.locator.prompts_locator import (
     PromptsLocatorGenerationAgent,
-    PromptsLocatorSelectionAgent
+    PromptsLocatorSelectionAgent,
+)
+from SelfhealingAgents.self_healing_system.context_retrieving.library_dom_utils.base_dom_utils import (
+    BaseDomUtils,
+)
+from SelfhealingAgents.self_healing_system.llm.client_model import get_client_model
+from SelfhealingAgents.self_healing_system.schemas.api.locator_healing import (
+    LocatorHealingResponse,
+)
+from SelfhealingAgents.self_healing_system.schemas.internal_state.prompt_payload import (
+    PromptPayload,
 )
+from SelfhealingAgents.utils.cfg import Cfg
+from SelfhealingAgents.utils.logging import log
 
 
 class BaseLocatorAgent(ABC):
@@ -31,18 +37,17 @@ class BaseLocatorAgent(ABC):
                                                                                    generation.
         selection_agent (Optional[Agent[PromptPayload, str]]): Agent for DOM-based locator selection.
     """
-    def __init__(
-        self,
-        cfg: Cfg,
-        dom_utility: BaseDomUtils
-    ) -> None:
+
+    def __init__(self, cfg: Cfg, dom_utility: BaseDomUtils) -> None:
         """Initializes the BaseLocatorAgent.
 
         Args:
             cfg (Cfg): Instance of Cfg config class containing user-defined app configuration.
             dom_utility (BaseDomUtils): DOM utility instance for validation.
         """
-        self._usage_limits: UsageLimits = UsageLimits(cfg.request_limit, cfg.total_tokens_limit)
+        self._usage_limits: UsageLimits = UsageLimits(
+            request_limit=cfg.request_limit, total_tokens_limit=cfg.total_tokens_limit
+        )
         self._dom_utility: BaseDomUtils = dom_utility
         self._use_llm_for_locator_generation = cfg.use_llm_for_locator_generation
 
@@ -60,7 +65,9 @@ def __init__(
                     model=cfg.locator_agent_model,
                     cfg=cfg,
                 ),
-                system_prompt=PromptsLocatorGenerationAgent.get_system_msg(self._dom_utility),
+                system_prompt=PromptsLocatorGenerationAgent.get_system_msg(
+                    self._dom_utility
+                ),
                 deps_type=PromptPayload,
                 output_type=LocatorHealingResponse,
             )

SelfhealingAgents/self_healing_system/agents/orchestrator_agent/orchestrator_agent.py

@@ -1,18 +1,24 @@
-from robot.api import logger as rf_logger
-from pydantic_ai.usage import UsageLimits
-from pydantic_ai.agent import AgentRunResult
 from pydantic_ai import Agent, ModelRetry, RunContext
+from pydantic_ai.agent import AgentRunResult
+from pydantic_ai.usage import UsageLimits
+from robot.api import logger as rf_logger
 
-from SelfhealingAgents.utils.cfg import Cfg
-from SelfhealingAgents.utils.logging import log
+from SelfhealingAgents.self_healing_system.agents.locator_agent.base_locator_agent import (
+    BaseLocatorAgent,
+)
+from SelfhealingAgents.self_healing_system.agents.prompts.orchestrator.prompts_orchestrator import (
+    PromptsOrchestrator,
+)
 from SelfhealingAgents.self_healing_system.llm.client_model import get_client_model
-from SelfhealingAgents.self_healing_system.schemas.internal_state.prompt_payload import PromptPayload
-from SelfhealingAgents.self_healing_system.agents.locator_agent.base_locator_agent import BaseLocatorAgent
-from SelfhealingAgents.self_healing_system.agents.prompts.orchestrator.prompts_orchestrator import PromptsOrchestrator
 from SelfhealingAgents.self_healing_system.schemas.api.locator_healing import (
     LocatorHealingResponse,
     NoHealingNeededResponse,
 )
+from SelfhealingAgents.self_healing_system.schemas.internal_state.prompt_payload import (
+    PromptPayload,
+)
+from SelfhealingAgents.utils.cfg import Cfg
+from SelfhealingAgents.utils.logging import log
 
 
 class OrchestratorAgent:
@@ -24,6 +30,7 @@ class OrchestratorAgent:
         _usage_limits (UsageLimits): Usage limits for the orchestrator agent.
         _agent (Agent[PromptPayload, str]): The underlying agent for orchestrating healing.
     """
+
     def __init__(
         self,
         cfg: Cfg,
@@ -37,7 +44,9 @@ def __init__(
         """
         self._cfg = cfg
         self._locator_agent: BaseLocatorAgent = locator_agent
-        self._usage_limits: UsageLimits = UsageLimits(cfg.request_limit, cfg.total_tokens_limit)
+        self._usage_limits: UsageLimits = UsageLimits(
+            request_limit=cfg.request_limit, total_tokens_limit=cfg.total_tokens_limit
+        )
         self._agent: Agent[PromptPayload, str] = Agent[PromptPayload, str](
             model=get_client_model(
                 provider=cfg.orchestrator_agent_provider,
@@ -46,7 +55,7 @@ def __init__(
             ),
             system_prompt=PromptsOrchestrator.get_system_msg(),
             deps_type=PromptPayload,
-            output_type=[self._get_healed_locators, str],
+            output_type=[self._get_healed_locators],
         )
 
     @log
@@ -105,4 +114,4 @@ def _catch_token_limit_exceedance(response_output: str) -> None:
             response_output (str): The response from the LLM request.
         """
         if "error" in response_output:
-            rf_logger.info(response_output)
+            rf_logger.info(response_output)

pyproject.toml

@@ -10,15 +10,14 @@ license = "Apache-2.0"
 readme = "README.md"
 requires-python = ">=3.12,<4.0"
 dependencies = [
-    "pydantic-ai>=0.5.0",
+    "pydantic-ai-slim[logfire, openai, azure, litellm]",
     "robotframework >=7.1",
     "beautifulsoup4 >=4.13.4",
     "pydantic-settings >=2.9.1",
     "lxml>=6.0.0",
     "robotframework-browser",
     "robotframework-seleniumlibrary",
     "python-dotenv>=1.1.1",
-    "pydantic-logfire>=0.0.1",
 ]