run_rag.py

import gradio as gr
import os
from pathlib import Path
import asyncio
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,
)

assistant = RetrieveAssistantAgent(
    name="assistant",
    system_message="You are a helpful assistant.",
    llm_config={
        "timeout": 600,
        "config_list": [
            {
                "api_type": "openai",
                "model": "gpt-4o-mini",
                "api_key": "sk-xxxx",
            }
        ],
    },
)

ragproxyagent = RetrieveUserProxyAgent(
    name="ragproxyagent",
    human_input_mode="NEVER",
    max_consecutive_auto_reply=5,
    retrieve_config={
        "task": "code",
        "docs_path": [
            "https://raw.githubusercontent.com/microsoft/FLAML/main/website/docs/Examples/Integrate%20-%20Spark.md",
            "https://raw.githubusercontent.com/microsoft/FLAML/main/website/docs/Research.md",
            # os.path.join(os.path.abspath(""), "..", "website", "docs"),
        ],
        "chunk_token_size": 2000,
        "model": "gpt-4o-mini",
        "client": chromadb.PersistentClient(path="/tmp/chromadb"),
        "embedding_model": "all-mpnet-base-v2",
        "customized_prompt": PROMPT_CODE,
        "get_or_create": True,
        "collection_name": "autogen_rag",
    },
)


async def main() -> None:

    code_problem = "How can I use FLAML to perform a classification task and use spark to do parallel training. Train 30 seconds and force cancel jobs if time limit is reached."

    ragproxyagent.initiate_chat(
        assistant, problem=code_problem, silent=False, n_results=1
    )


asyncio.run(main())