Merge pull request #1832 from h2oai/refactor_agents

Refactor agent code

openai_server/agent_prompting.py

@@ -8,15 +8,15 @@
     structure_to_messages
 
 
-def agent_system_prompt(agent_code_writer_system_message, autogen_system_site_packages):
+def agent_system_prompt(agent_code_writer_system_message, agent_system_site_packages):
     if agent_code_writer_system_message is None:
         cwd = os.path.abspath(os.getcwd())
         have_internet = get_have_internet()
         date_str = current_datetime()
 
         # The code writer agent's system message is to instruct the LLM on how to use
         # the code executor in the code executor agent.
-        if autogen_system_site_packages:
+        if agent_system_site_packages:
             # heavy packages only expect should use if system inherited
             extra_recommended_packages = """\n  * Image Processing: opencv-python
   * DataBase: pysqlite3
@@ -483,3 +483,30 @@ def get_mermaid_renderer_helper():
 * A png version of any svg is also created for use with image_query in order to analyze the svg (via the png).
 """
     return mmdc
+
+
+def get_full_system_prompt(agent_code_writer_system_message, agent_system_site_packages, system_prompt, base_url,
+                           api_key, model, text_context_list, image_file, temp_dir, query):
+    agent_code_writer_system_message = agent_system_prompt(agent_code_writer_system_message,
+                                                           agent_system_site_packages)
+
+    image_query_helper = get_image_query_helper(base_url, api_key, model)
+    mermaid_renderer_helper = get_mermaid_renderer_helper()
+
+    chat_doc_query, internal_file_names = get_chat_doc_context(text_context_list, image_file,
+                                                               temp_dir,
+                                                               # avoid text version of chat conversation, confuses LLM
+                                                               chat_conversation=None,
+                                                               system_prompt=system_prompt,
+                                                               prompt=query,
+                                                               model=model)
+
+    cwd = os.path.abspath(os.getcwd())
+    path_agent_tools = f'{cwd}/openai_server/agent_tools/'
+    list_dir = os.listdir('openai_server/agent_tools')
+    list_dir = [x for x in list_dir if not x.startswith('__')]
+
+    agent_tools_note = f"\nDo not hallucinate agent_tools tools. The only files in the {path_agent_tools} directory are as follows: {list_dir}\n"
+
+    system_message = agent_code_writer_system_message + image_query_helper + mermaid_renderer_helper + agent_tools_note + chat_doc_query
+    return system_message, internal_file_names, chat_doc_query, image_query_helper, mermaid_renderer_helper

openai_server/agent_utils.py

@@ -1,9 +1,15 @@
+import functools
 import inspect
 import os
+import shutil
 import sys
+import time
+
 import requests
 from PIL import Image
 
+from openai_server.backend_utils import get_user_dir, run_upload_api
+
 
 def get_have_internet():
     try:
@@ -46,9 +52,20 @@ def in_pycharm():
     return os.getenv("PYCHARM_HOSTED") is not None
 
 
+def get_inner_function_signature(func):
+    # Check if the function is a functools.partial object
+    if isinstance(func, functools.partial):
+        # Get the original function
+        assert func.keywords is not None and func.keywords, "The function must have keyword arguments."
+        func = func.keywords['run_agent_func']
+        return inspect.signature(func)
+    else:
+        return inspect.signature(func)
+
+
 def filter_kwargs(func, kwargs):
     # Get the parameter list of the function
-    sig = inspect.signature(func)
+    sig = get_inner_function_signature(func)
     valid_kwargs = {k: v for k, v in kwargs.items() if k in sig.parameters}
     return valid_kwargs
 
@@ -85,3 +102,133 @@ def current_datetime():
 
     # Print the formatted date, time, and time zone
     return "For current user query: Current Date, Time, and Local Time Zone: %s. Note some APIs may have data from different time zones, so may reflect a different date." % formatted_date_time
+
+
+def run_agent(run_agent_func=None,
+              **kwargs,
+              ) -> dict:
+    ret_dict = {}
+    try:
+        assert run_agent_func is not None, "run_agent_func must be provided."
+        ret_dict = run_agent_func(**kwargs)
+    finally:
+        if kwargs.get('agent_venv_dir') is None and 'agent_venv_dir' in ret_dict and ret_dict['agent_venv_dir']:
+            agent_venv_dir = ret_dict['agent_venv_dir']
+            if os.path.isdir(agent_venv_dir):
+                if kwargs.get('agent_verbose'):
+                    print("Clean-up: Removing agent_venv_dir: %s" % agent_venv_dir)
+                shutil.rmtree(agent_venv_dir)
+
+    return ret_dict
+
+
+def set_dummy_term():
+    # Disable color and advanced terminal features
+    os.environ['TERM'] = 'dumb'
+    os.environ['COLORTERM'] = ''
+    os.environ['CLICOLOR'] = '0'
+    os.environ['CLICOLOR_FORCE'] = '0'
+    os.environ['ANSI_COLORS_DISABLED'] = '1'
+
+    # force matplotlib to use terminal friendly backend
+    import matplotlib as mpl
+    mpl.use('Agg')
+
+
+def get_ret_dict_and_handle_files(chat_result, temp_dir, agent_verbose, internal_file_names, authorization,
+                                  autogen_run_code_in_docker, autogen_stop_docker_executor, executor,
+                                  agent_venv_dir, agent_code_writer_system_message, agent_system_site_packages,
+                                  chat_doc_query, image_query_helper, mermaid_renderer_helper,
+                                  autogen_code_restrictions_level, autogen_silent_exchange):
+    # DEBUG
+    if agent_verbose:
+        print("chat_result:", chat_result)
+        print("list_dir:", os.listdir(temp_dir))
+
+    # Get all files in the temp_dir and one level deep subdirectories
+    file_list = []
+    for root, dirs, files in os.walk(temp_dir):
+        # Exclude deeper directories by checking the depth
+        if root == temp_dir or os.path.dirname(root) == temp_dir:
+            file_list.extend([os.path.join(root, f) for f in files])
+
+    # Filter the list to include only files
+    file_list = [f for f in file_list if os.path.isfile(f)]
+    internal_file_names_norm_paths = [os.path.normpath(f) for f in internal_file_names]
+    # filter out internal files for RAG case
+    file_list = [f for f in file_list if os.path.normpath(f) not in internal_file_names_norm_paths]
+    if agent_verbose:
+        print("file_list:", file_list)
+
+    image_files, non_image_files = identify_image_files(file_list)
+    # keep no more than 10 image files:
+    image_files = image_files[:10]
+    file_list = image_files + non_image_files
+
+    # copy files so user can download
+    user_dir = get_user_dir(authorization)
+    if not os.path.isdir(user_dir):
+        os.makedirs(user_dir, exist_ok=True)
+    file_ids = []
+    for file in file_list:
+        new_path = os.path.join(user_dir, os.path.basename(file))
+        shutil.copy(file, new_path)
+        with open(new_path, "rb") as f:
+            content = f.read()
+        purpose = 'assistants'
+        response_dict = run_upload_api(content, new_path, purpose, authorization)
+        file_id = response_dict['id']
+        file_ids.append(file_id)
+
+    # temp_dir.cleanup()
+    if autogen_run_code_in_docker and autogen_stop_docker_executor:
+        t0 = time.time()
+        executor.stop()  # Stop the docker command line code executor (takes about 10 seconds, so slow)
+        if agent_verbose:
+            print(f"Executor Stop time taken: {time.time() - t0:.2f} seconds.")
+
+    ret_dict = {}
+    if file_list:
+        ret_dict.update(dict(files=file_list))
+    if file_ids:
+        ret_dict.update(dict(file_ids=file_ids))
+    if chat_result and hasattr(chat_result, 'chat_history'):
+        ret_dict.update(dict(chat_history=chat_result.chat_history))
+    if chat_result and hasattr(chat_result, 'cost'):
+        ret_dict.update(dict(cost=chat_result.cost))
+    if chat_result and hasattr(chat_result, 'summary') and chat_result.summary:
+        ret_dict.update(dict(summary=chat_result.summary))
+        print("Made summary: %s" % chat_result.summary, file=sys.stderr)
+    else:
+        if hasattr(chat_result, 'chat_history') and chat_result.chat_history:
+            summary = chat_result.chat_history[-1]['content']
+            if not summary and len(chat_result.chat_history) >= 2:
+                summary = chat_result.chat_history[-2]['content']
+            if summary:
+                print("Made summary from chat history: %s" % summary, file=sys.stderr)
+                ret_dict.update(dict(summary=summary))
+            else:
+                print("Did NOT make and could not make summary", file=sys.stderr)
+                ret_dict.update(dict(summary=''))
+        else:
+            print("Did NOT make any summary", file=sys.stderr)
+            ret_dict.update(dict(summary=''))
+    if agent_venv_dir is not None:
+        ret_dict.update(dict(agent_venv_dir=agent_venv_dir))
+    if agent_code_writer_system_message is not None:
+        ret_dict.update(dict(agent_code_writer_system_message=agent_code_writer_system_message))
+    if agent_system_site_packages is not None:
+        ret_dict.update(dict(agent_system_site_packages=agent_system_site_packages))
+    if chat_doc_query:
+        ret_dict.update(dict(chat_doc_query=chat_doc_query))
+    if image_query_helper:
+        ret_dict.update(dict(image_query_helper=image_query_helper))
+    if mermaid_renderer_helper:
+        ret_dict.update(dict(mermaid_renderer_helper=mermaid_renderer_helper))
+    ret_dict.update(dict(autogen_code_restrictions_level=autogen_code_restrictions_level))
+    ret_dict.update(dict(autogen_silent_exchange=autogen_silent_exchange))
+    # can re-use for chat continuation to avoid sending files over
+    # FIXME: Maybe just delete files and force send back to agent
+    ret_dict.update(dict(temp_dir=temp_dir))
+
+    return ret_dict

openai_server/autogen_2agent_backend.py

@@ -0,0 +1,148 @@
+import os
+import tempfile
+
+from openai_server.backend_utils import structure_to_messages
+from openai_server.agent_utils import get_ret_dict_and_handle_files
+from openai_server.agent_prompting import get_full_system_prompt
+
+from openai_server.autogen_utils import terminate_message_func, H2OConversableAgent
+
+
+def run_autogen_2agent(query=None,
+                       visible_models=None,
+                       stream_output=None,
+                       max_new_tokens=None,
+                       authorization=None,
+                       chat_conversation=None,
+                       text_context_list=None,
+                       system_prompt=None,
+                       image_file=None,
+                       # autogen/agent specific parameters
+                       agent_type=None,
+                       autogen_stop_docker_executor=None,
+                       autogen_run_code_in_docker=None,
+                       autogen_max_consecutive_auto_reply=None,
+                       autogen_max_turns=None,
+                       autogen_timeout=None,
+                       autogen_cache_seed=None,
+                       agent_venv_dir=None,
+                       agent_code_writer_system_message=None,
+                       agent_system_site_packages=None,
+                       autogen_code_restrictions_level=None,
+                       autogen_silent_exchange=None,
+                       agent_verbose=None) -> dict:
+    assert agent_type in ['autogen_2agent', 'auto'], "Invalid agent_type: %s" % agent_type
+    # raise openai.BadRequestError("Testing Error Handling")
+    # raise ValueError("Testing Error Handling")
+
+    # handle parameters from chatAPI and OpenAI -> h2oGPT transcription versions
+    assert visible_models is not None, "No visible_models specified"
+    model = visible_models  # transcribe early
+
+    if stream_output is None:
+        stream_output = False
+    assert max_new_tokens is not None, "No max_new_tokens specified"
+
+    # handle AutoGen specific parameters
+    if autogen_stop_docker_executor is None:
+        autogen_stop_docker_executor = False
+    if autogen_run_code_in_docker is None:
+        autogen_run_code_in_docker = False
+    if autogen_max_consecutive_auto_reply is None:
+        autogen_max_consecutive_auto_reply = 40
+    if autogen_max_turns is None:
+        autogen_max_turns = 40
+    if autogen_timeout is None:
+        autogen_timeout = 120
+    if agent_system_site_packages is None:
+        agent_system_site_packages = True
+    if autogen_code_restrictions_level is None:
+        autogen_code_restrictions_level = 2
+    if autogen_silent_exchange is None:
+        autogen_silent_exchange = True
+    if agent_verbose is None:
+        agent_verbose = False
+    if agent_verbose:
+        print("AutoGen using model=%s." % model, flush=True)
+
+    # Create a temporary directory to store the code files.
+    # temp_dir = tempfile.TemporaryDirectory().name
+    temp_dir = tempfile.mkdtemp()
+
+    # iostream = IOStream.get_default()
+    # iostream.print("\033[32m", end="")
+
+    from openai_server.autogen_agents import get_execution_agent
+    code_executor_agent, executor = \
+        get_execution_agent(autogen_run_code_in_docker, autogen_timeout, agent_system_site_packages,
+                            autogen_max_consecutive_auto_reply, autogen_code_restrictions_level,
+                            agent_venv_dir, temp_dir)
+
+    # FIXME:
+    # Auto-pip install
+    # Auto-return file list in each turn
+
+    base_url = os.environ['H2OGPT_OPENAI_BASE_URL']  # must exist
+    api_key = os.environ['H2OGPT_OPENAI_API_KEY']  # must exist
+    if agent_verbose:
+        print("base_url: %s" % base_url)
+        print("max_tokens: %s" % max_new_tokens)
+
+    system_message, internal_file_names, chat_doc_query, image_query_helper, mermaid_renderer_helper = \
+        get_full_system_prompt(agent_code_writer_system_message,
+                               agent_system_site_packages, system_prompt,
+                               base_url,
+                               api_key, model, text_context_list, image_file,
+                               temp_dir, query)
+
+    code_writer_agent = H2OConversableAgent(
+        "code_writer_agent",
+        system_message=system_message,
+        llm_config={"config_list": [{"model": model,
+                                     "api_key": api_key,
+                                     "base_url": base_url,
+                                     "stream": stream_output,
+                                     "cache_seed": autogen_cache_seed,
+                                     'max_tokens': max_new_tokens}]},
+        code_execution_config=False,  # Turn off code execution for this agent.
+        human_input_mode="NEVER",
+        is_termination_msg=terminate_message_func,
+        max_consecutive_auto_reply=autogen_max_consecutive_auto_reply,
+    )
+
+    # apply chat history
+    if chat_conversation:
+        chat_messages = structure_to_messages(None, None, chat_conversation, None)
+        for message in chat_messages:
+            if message['role'] == 'assistant':
+                code_writer_agent.send(message['content'], code_executor_agent, request_reply=False)
+            if message['role'] == 'user':
+                code_executor_agent.send(message['content'], code_writer_agent, request_reply=False)
+
+    chat_kwargs = dict(recipient=code_writer_agent,
+                       max_turns=autogen_max_turns,
+                       message=query,
+                       cache=None,
+                       silent=autogen_silent_exchange,
+                       clear_history=False,
+                       )
+    if autogen_cache_seed:
+        from autogen import Cache
+        # Use DiskCache as cache
+        cache_root_path = "./autogen_cache"
+        if not os.path.exists(cache_root_path):
+            os.makedirs(cache_root_path, exist_ok=True)
+        with Cache.disk(cache_seed=autogen_cache_seed, cache_path_root=cache_root_path) as cache:
+            chat_kwargs.update(dict(cache=cache))
+            chat_result = code_executor_agent.initiate_chat(**chat_kwargs)
+    else:
+        chat_result = code_executor_agent.initiate_chat(**chat_kwargs)
+
+    ret_dict = get_ret_dict_and_handle_files(chat_result, temp_dir, agent_verbose, internal_file_names, authorization,
+                                             autogen_run_code_in_docker, autogen_stop_docker_executor, executor,
+                                             agent_venv_dir, agent_code_writer_system_message,
+                                             agent_system_site_packages,
+                                             chat_doc_query, image_query_helper, mermaid_renderer_helper,
+                                             autogen_code_restrictions_level, autogen_silent_exchange)
+
+    return ret_dict