Agent as tool

debug_gym/gym/tools/agent.py

@@ -0,0 +1,83 @@
+import logging
+import os
+from os.path import join as pjoin
+
+from debug_gym.agents.history_tracker import build_history_prompt
+from debug_gym.gym.entities import Observation
+from debug_gym.gym.tools.tool import EnvironmentTool
+from debug_gym.gym.tools.toolbox import Toolbox
+from debug_gym.gym.utils import is_subdirectory, show_line_number
+from debug_gym.llms.base import LLM
+
+
+@Toolbox.register()
+class AgentTool(EnvironmentTool):
+    name: str = "agent"
+    examples = [
+        """agent(query="") Executes query to send to this specialized agent.\n"""
+    ]
+    description = (
+        "This tool allows to call a large language model specialized agent, and therefore to query for information and exploit capabilities in the model.\n"
+        + "\n".join(examples)
+    )
+    arguments = {
+        "query": {
+            "type": ["string"],
+            "description": "The request to the agent.",
+        },
+    }
+
+    def __init__(
+        self, history, llm_name=None, llm_config=None, llm_config_file_path=None
+    ):
+        self._history = history
+        self.logger = logging.getLogger("agent_logger")
+        self._llm = LLM.instantiate(
+            llm_name=llm_name,
+            llm_config=llm_config,
+            llm_config_file_path=llm_config_file_path,
+            logger=self.logger,
+        )
+        self.llm_config = self._llm.config
+        super().__init__()
+
+    def use(
+        self,
+        environment,
+        query: str,
+    ) -> Observation:
+        messages = build_history_prompt(
+            self._history.filter_out(actions=[None]),
+            self._llm,
+            False,
+        )
+        messages.append(
+            {
+                "role": "user",
+                "content": query,
+            }
+        )
+        llm_response = self._llm.generate(messages, [], tool_choice="none")
+
+        return Observation(
+            self.name,
+            llm_response.response,
+        )
+
+    def __copy__(self):
+        cls = self.__class__
+        result = cls.__new__(cls)
+        result.__dict__.update(self.__dict__)
+        return result
+
+    def __deepcopy__(self, memo):
+        cls = self.__class__
+        result = cls.__new__(cls)
+
+        result.logger = logging.getLogger("agent_logger")
+        result._llm = LLM.instantiate(
+            llm_config=result.llm_config,
+            logger=result.logger,
+        )
+        result._history = self._history.copy()
+        return result

debug_gym/llms/base.py

@@ -207,29 +207,24 @@ def __init__(
         )
 
     @classmethod
-    def instantiate(
+    def instantiate_from_config(
         cls,
-        llm_name: str,
-        llm_config_file_path: str | None = None,
+        llm_config: LLMConfig,
         logger: DebugGymLogger | None = None,
     ) -> "LLM":
         """Creates an instance of the appropriate LLM class based on the configuration.
 
         Args:
-            llm_name: Name of the LLM model to instantiate.
-            llm_config_file_path: Optional path to the LLM configuration file.
+            llm_config: LLMConfig object containing the configuration.
             logger: Optional DebugGymLogger for logging.
 
         Returns:
             An instance of the appropriate LLM class.
         """
-        logger = logger or DebugGymLogger("debug-gym")
-        if llm_name == "human":
-            from debug_gym.llms import Human
-
-            return Human(llm_name, logger=logger)
-
-        llm_config = LLMConfigRegistry.from_file(llm_config_file_path)[llm_name]
+        if logger is not "none":
+            logger = logger or DebugGymLogger("debug-gym")
+        else:
+            logger = None
 
         tags = llm_config.tags
         if "azure openai" in tags:
@@ -247,6 +242,45 @@ def instantiate(
         llm = klass(llm_name, logger=logger, llm_config=llm_config)
         return llm
 
+    @classmethod
+    def instantiate(
+        cls,
+        llm_name: str | None = None,
+        llm_config: LLMConfig | None = None,
+        llm_config_file_path: str | None = None,
+        logger: DebugGymLogger | None = None,
+    ) -> "LLM":
+        """Creates an instance of the appropriate LLM class based on the configuration.
+
+        Args:
+            llm_name: Name of the LLM model to instantiate.
+            llm_config_file_path: Optional path to the LLM configuration file.
+            logger: Optional DebugGymLogger for logging.
+
+        Returns:
+            An instance of the appropriate LLM class.
+        """
+        if logger is not "none":
+            logger = logger or DebugGymLogger("debug-gym")
+        else:
+            logger = None
+
+        if llm_name == "human":
+            from debug_gym.llms import Human
+
+            return Human(llm_name, logger=logger)
+
+        if llm_config is None:
+            if llm_config_file_path is None:
+                raise ValueError(
+                    "Either llm_config or llm_config_file_path must be provided."
+                )
+
+            llm_config = LLMConfigRegistry.from_file(llm_config_file_path)[llm_name]
+
+        return cls.instantiate_from_config(llm_config, logger)
+
+
     @abstractmethod
     def generate(self, messages, tools, **kwargs) -> LLMResponse:
         """Generate a response given some messages and return it as an LLMResponse object.
@@ -363,4 +397,4 @@ def generate_with_drop_message_and_retry(messages, tools, **kwargs):
 
         self.logger.info(colored(llm_response.tool, "green"))
 
-        return llm_response
+        return llm_response

debug_gym/llms/openai.py

@@ -192,7 +192,7 @@ def generate(self, messages, tools, **kwargs) -> LLMResponse:
                 model=self.config.model,
                 messages=messages,
                 tools=self.define_tools(tools),
-                tool_choice="required",
+                tool_choice=kwargs.pop("tool_choice", "required"),
                 **kwargs,
             )
         except openai.BadRequestError as e: