Adjust filename if was wrongly defined

openai_server/autogen_utils.py

@@ -5,6 +5,7 @@
 import logging
 import os
 import re
+import shutil
 import subprocess
 import sys
 import typing
@@ -69,7 +70,7 @@ def __init__(
             agent_tools_usage_hard_limits: Dict[str, int] = {},
             agent_tools_usage_soft_limits: Dict[str, int] = {},
             max_stream_length: int = 4096,
-            max_memory_usage: Optional[int] = 16*1024**3,  # 16GB
+            max_memory_usage: Optional[int] = 16 * 1024 ** 3,  # 16GB
     ):
         super().__init__(timeout, virtual_env_context, work_dir, functions, functions_module, execution_policies)
         self.autogen_code_restrictions_level = autogen_code_restrictions_level
@@ -294,8 +295,18 @@ def __execute_code_dont_check_setup(self, code_blocks: List[CodeBlock]) -> Comma
             try:
                 # Check if there is a filename comment
                 filename = self._get_file_name_from_content(code, self._work_dir)
+                # override filename and lang if tool use is detected
+                cwd = os.path.abspath(os.getcwd())
+                if code_block.execute and \
+                        f'python {cwd}/openai_server/agent_tools/' in code and \
+                        filename.endswith('.py'):
+                    # switch back to shell if was wrongly .py extension
+                    lang = 'shell'
+                    new_filename = filename.replace('.py', '.sh')
+                    shutil.move(filename, new_filename)
+                    filename = new_filename
                 # override lang if filename is detected, less error-prone than using code block lang
-                if filename.endswith('.sh'):
+                elif filename.endswith('.sh'):
                     lang = 'shell'
                 elif filename.endswith('.py'):
                     lang = 'python'
@@ -449,7 +460,7 @@ def _execute_code_dont_check_setup(self, code_blocks: List[CodeBlock]) -> Comman
                 ret = CommandLineCodeResult(exit_code=1, output=str(e))
             else:
                 raise
-        
+
         # Truncate output if it is too long
         ret = self.truncate_output(ret)
         # Add executed code note if needed
@@ -473,7 +484,8 @@ def update_agent_tool_usages(self, code_blocks: List[CodeBlock]) -> None:
             print(f"Step {self.turns} has agent tool usage: {self.agent_tools_usage}")
 
     @staticmethod
-    def executed_code_note(ret: CommandLineCodeResult, multiple_executable_code_detected: bool = False) -> CommandLineCodeResult:
+    def executed_code_note(ret: CommandLineCodeResult,
+                           multiple_executable_code_detected: bool = False) -> CommandLineCodeResult:
         if ret.exit_code == 0:
             if multiple_executable_code_detected:
                 executable_code_limitation_warning = """
@@ -582,7 +594,7 @@ def text_guardrail(text, any_fail=False, max_bad_lines=3, just_filter_out=True):
 
                         # Use the compiled regex to replace all api_key_values at once
                         line = regex.sub('', line)
-                        #for api_key_value in api_key_values:
+                        # for api_key_value in api_key_values:
                         #    line = line.replace(api_key_value, '')
                         lines.append(line)
                     else:
@@ -901,9 +913,9 @@ def __generate_code_execution_reply_using_executor(
             code_blocks = self._code_executor.code_extractor.extract_code_blocks(message["content"])
             stop_on_termination = False
             if (
-                len(code_blocks) == 0 or
+                    len(code_blocks) == 0 or
                     (stop_on_termination and "<FINISHED_ALL_TASKS>" in message["content"])
-                ):
+            ):
                 if self._confidence_level == 0:
                     self._confidence_level = 1
                     return True, self.confidence_level_guidelines()
@@ -1040,7 +1052,7 @@ def get_code_executor(
         agent_tools_usage_soft_limits={},
         max_stream_length=4096,
         # max memory per code execution process
-        max_memory_usage=16*1024**3,  # 16GB
+        max_memory_usage=16 * 1024 ** 3,  # 16GB
 ):
     if autogen_run_code_in_docker:
         from autogen.coding import DockerCommandLineCodeExecutor

src/version.py

@@ -1 +1 @@
-__version__ = "471ca9c2c8c218d0581582e48873fe22f42228a9"
+__version__ = "7a756e9b132fc32382803b9b0e7cc90467bec883"