app.py

import gradio as gr
import os
from pathlib import Path
import autogen
import chromadb
import multiprocessing as mp
from autogen.retrieve_utils import TEXT_FORMATS, get_file_from_url, is_url
from autogen.agentchat.contrib.retrieve_assistant_agent import RetrieveAssistantAgent
from autogen.agentchat.contrib.retrieve_user_proxy_agent import (
 RetrieveUserProxyAgent,
 PROMPT_CODE,
)

TIMEOUT = 60


def initialize_agents(config_list, docs_path=None):
 if isinstance(config_list, gr.State):
 _config_list = config_list.value
 else:
 _config_list = config_list
 if docs_path is None:
 docs_path = "https://raw.githubusercontent.com/microsoft/autogen/main/README.md"

 assistant = RetrieveAssistantAgent(
 name="assistant",
 system_message="You are a helpful assistant.",
 )

 ragproxyagent = RetrieveUserProxyAgent(
 name="ragproxyagent",
 human_input_mode="NEVER",
 max_consecutive_auto_reply=5,
 retrieve_config={
 "task": "code",
 "docs_path": docs_path,
 "chunk_token_size": 2000,
 "model": _config_list[0]["model"],
 "client": chromadb.PersistentClient(path="/tmp/chromadb"),
 "embedding_model": "all-mpnet-base-v2",
 "customized_prompt": PROMPT_CODE,
 "get_or_create": True,
 "collection_name": "autogen_rag",
 },
 )

 return assistant, ragproxyagent


def initiate_chat(config_list, problem, queue, n_results=3):
 global assistant, ragproxyagent
 if isinstance(config_list, gr.State):
 _config_list = config_list.value
 else:
 _config_list = config_list
 if len(_config_list[0].get("api_key", "")) < 2:
 queue.put(
 ["Hi, nice to meet you! Please enter your API keys in below text boxs."]
 )
 return
 else:
 llm_config = (
 {
 "request_timeout": TIMEOUT,
 # "seed": 42,
 "config_list": _config_list,
 "use_cache": False,
 },
 )
 assistant.llm_config.update(llm_config[0])
 assistant.reset()
 try:
 ragproxyagent.initiate_chat(
 assistant, problem=problem, silent=False, n_results=n_results
 )
 messages = ragproxyagent.chat_messages
 messages = [messages[k] for k in messages.keys()][0]
 messages = [m["content"] for m in messages if m["role"] == "user"]
 print("messages: ", messages)
 except Exception as e:
 messages = [str(e)]
 queue.put(messages)


def chatbot_reply(input_text):
 """Chat with the agent through terminal."""
 queue = mp.Queue()
 process = mp.Process(
 target=initiate_chat,
 args=(config_list, input_text, queue),
 )
 process.start()
 try:
 # process.join(TIMEOUT+2)
 messages = queue.get(timeout=TIMEOUT)
 except Exception as e:
 messages = [
 str(e)
 if len(str(e)) > 0
 else "Invalid Request to OpenAI, please check your API keys."
 ]
 finally:
 try:
 process.terminate()
 except:
 pass
 return messages


def get_description_text():
 return """
 # Microsoft AutoGen: Retrieve Chat Demo

 This demo shows how to use the RetrieveUserProxyAgent and RetrieveAssistantAgent to build a chatbot.

 #### [AutoGen](https://github.com/microsoft/autogen) [Discord](https://discord.gg/pAbnFJrkgZ) [Blog](https://microsoft.github.io/autogen/blog/2023/10/18/RetrieveChat) [Paper](https://arxiv.org/abs/2308.08155) [SourceCode](https://github.com/thinkall/autogen-demos)
 """


global assistant, ragproxyagent

with gr.Blocks() as demo:
 config_list, assistant, ragproxyagent = (
 gr.State(
 [
 {
 "api_key": "",
 "api_base": "",
 "api_type": "azure",
 "api_version": "2023-07-01-preview",
 "model": "gpt-35-turbo",
 }
 ]
 ),
 None,
 None,
 )
 assistant, ragproxyagent = initialize_agents(config_list)

 gr.Markdown(get_description_text())
 chatbot = gr.Chatbot(
 [],
 elem_id="chatbot",
 bubble_full_width=False,
 avatar_images=(None, (os.path.join(os.path.dirname(__file__), "autogen.png"))),
 # height=600,
 )

 txt_input = gr.Textbox(
 scale=4,
 show_label=False,
 placeholder="Enter text and press enter",
 container=False,
 )

 with gr.Row():

 def update_config(config_list):
 global assistant, ragproxyagent
 config_list = autogen.config_list_from_models(
 model_list=[os.environ.get("MODEL", "gpt-35-turbo")],
 )
 if not config_list:
 config_list = [
 {
 "api_key": "",
 "api_base": "",
 "api_type": "azure",
 "api_version": "2023-07-01-preview",
 "model": "gpt-35-turbo",
 }
 ]
 llm_config = (
 {
 "request_timeout": TIMEOUT,
 # "seed": 42,
 "config_list": config_list,
 },
 )
 assistant.llm_config.update(llm_config[0])
 ragproxyagent._model = config_list[0]["model"]
 return config_list


 def set_params(model, oai_key, aoai_key, aoai_base):
 os.environ["MODEL"] = model
 os.environ["OPENAI_API_KEY"] = oai_key
 os.environ["AZURE_OPENAI_API_KEY"] = aoai_key
 os.environ["AZURE_OPENAI_API_BASE"] = aoai_base
 return model, oai_key, aoai_key, aoai_base


 txt_model = gr.Dropdown(
 label="Model",
 choices=[
 "gpt-4",
 "gpt-35-turbo",
 "gpt-3.5-turbo",
 ],
 allow_custom_value=True,
 value="gpt-35-turbo",
 container=True,
 )
 txt_oai_key = gr.Textbox(
 label="OpenAI API Key",
 placeholder="Enter key and press enter",
 max_lines=1,
 show_label=True,
 value=os.environ.get("OPENAI_API_KEY", ""),
 container=True,
 type="password",
 )
 txt_aoai_key = gr.Textbox(
 label="Azure OpenAI API Key",
 placeholder="Enter key and press enter",
 max_lines=1,
 show_label=True,
 value=os.environ.get("AZURE_OPENAI_API_KEY", ""),
 container=True,
 type="password",
 )
 txt_aoai_base_url = gr.Textbox(
 label="Azure OpenAI API Base",
 placeholder="Enter base url and press enter",
 max_lines=1,
 show_label=True,
 value=os.environ.get("AZURE_OPENAI_API_BASE", ""),
 container=True,
 type="password",
 )

 clear = gr.ClearButton([txt_input, chatbot])

 with gr.Row():

 def upload_file(file):
 return update_context_url(file.name)


 upload_button = gr.UploadButton(
 "Click to upload a context file or enter a url in the right textbox",
 file_types=[f".{i}" for i in TEXT_FORMATS],
 file_count="single",
 )

 txt_context_url = gr.Textbox(
 label="Enter the url to your context file and chat on the context",
 info=f"File must be in the format of [{', '.join(TEXT_FORMATS)}]",
 max_lines=1,
 show_label=True,
 value="https://raw.githubusercontent.com/microsoft/autogen/main/README.md",
 container=True,
 )

 txt_prompt = gr.Textbox(
 label="Enter your prompt for Retrieve Agent and press enter to replace the default prompt",
 max_lines=40,
 show_label=True,
 value=PROMPT_CODE,
 container=True,
 show_copy_button=True,
 )


 def respond(message, chat_history, model, oai_key, aoai_key, aoai_base):
 global config_list
 set_params(model, oai_key, aoai_key, aoai_base)
 config_list = update_config(config_list)
 messages = chatbot_reply(message)
 _msg = (
 messages[-1]
 if len(messages) > 0 and messages[-1] != "TERMINATE"
 else messages[-2]
 if len(messages) > 1
 else "Context is not enough for answering the question. Please press `enter` in the context url textbox to make sure the context is activated for the chat."
 )
 chat_history.append((message, _msg))
 return "", chat_history


 def update_prompt(prompt):
 ragproxyagent.customized_prompt = prompt
 return prompt


 def update_context_url(context_url):
 global assistant, ragproxyagent

 file_extension = Path(context_url).suffix
 print("file_extension: ", file_extension)
 if file_extension.lower() not in [f".{i}" for i in TEXT_FORMATS]:
 return f"File must be in the format of {TEXT_FORMATS}"

 if is_url(context_url):
 try:
 file_path = get_file_from_url(
 context_url,
 save_path=os.path.join("/tmp", os.path.basename(context_url)),
 )
 except Exception as e:
 return str(e)
 else:
 file_path = context_url
 context_url = os.path.basename(context_url)

 try:
 chromadb.PersistentClient(path="/tmp/chromadb").delete_collection(
 name="autogen_rag"
 )
 except:
 pass
 assistant, ragproxyagent = initialize_agents(config_list, docs_path=file_path)
 return context_url


 txt_input.submit(
 respond,
 [txt_input, chatbot, txt_model, txt_oai_key, txt_aoai_key, txt_aoai_base_url],
 [txt_input, chatbot],
 )
 txt_prompt.submit(update_prompt, [txt_prompt], [txt_prompt])
 txt_context_url.submit(update_context_url, [txt_context_url], [txt_context_url])
 upload_button.upload(upload_file, upload_button, [txt_context_url])

if __name__ == "__main__":
 demo.launch(share=True, server_name="0.0.0.0")