Skip to content

Commit

Permalink
aider: The agent now updates the progress bar during research.
Browse files Browse the repository at this point in the history
# Aider chat conversation:

USER: when start the agent adding source and browsing sites, the progress bar not show any progress or change of color that show progress

ASSISTANT: To implement the progress bar and update it during the agent's execution, you need to make changes to the following files:

1. agent\research_agent.py

```python
import traceback

import asyncio
import json
import uuid
import sys
import hashlib

from actions.web_search import web_search
from actions.web_scrape import async_browse
from processing.text import write_to_file, create_message, create_chat_completion, read_txt_files, write_md_to_pdf
from config import Config
from agent import prompts
import os
import string

CFG = Config()

class ResearchAgent:
    def __init__(self, question, agent, agent_role_prompt, websocket):
        """ Initializes the research assistant with the given question.
        Args: question (str): The question to research
        Returns: None
        """

        self.question = question
        self.agent = agent
        self.agent_role_prompt = agent_role_prompt if agent_role_prompt else prompts.generate_agent_role_prompt(agent)
        self.visited_urls = set()
        self.research_summary = ""
        self.directory_name = uuid.uuid4()
        self.dir_path = os.path.dirname(f"./outputs/{self.directory_name}/")
        self.websocket = websocket
        self.channels = [""," site:twitter.com"," site:quora.com"," site:reddit.com"," site:medium.com"," site:trustpilot.com"," site:sensortower.com"]
        self.progress = 0

    async def summarize(self, text, topic):
        """ Summarizes the given text for the given topic.
        Args: text (str): The text to summarize
                topic (str): The topic to summarize the text for
        Returns: str: The summarized text
        """
        try:
            messages = [create_message(text, topic)]
            await self.websocket.send_json({"type": "logs", "output": f"📝 Summarizing text for query: {text}"})

            return create_chat_completion(
                model=CFG.fast_llm_model,
                messages=messages,
            )
        except Exception as e:
            traceback.print_exc()
            await self.websocket.send_json({"type": "logs", "output": f"❌ Error occurred during summarization: {str(e)}"})
            return ""

    async def get_new_urls(self, url_set_input):
        """ Gets the new urls from the given url set.
        Args: url_set_input (set[str]): The url set to get the new urls from
        Returns: list[str]: The new urls from the given url set
        """
        try:
            new_urls = []
            for url in url_set_input:
                if url not in self.visited_urls:
                    await self.websocket.send_json({"type": "logs", "output": f"✅ Adding source url to research: {url}\n"})
                    self.visited_urls.add(url)
                    new_urls.append(url)

            return new_urls
        except Exception as e:
            traceback.print_exc()
            await self.websocket.send_json({"type": "logs", "output": f"❌ Error occurred while getting new urls: {str(e)}"})
            return []

    async def call_agent(self, action, stream=False, websocket=None):
        try:
            messages = [{
                "role": "system",
                "content": self.agent_role_prompt
            }, {
                "role": "user",
                "content": action,
            }]
            answer = create_chat_completion(
                model=CFG.smart_llm_model,
                messages=messages,
                stream=stream,
                websocket=websocket,
            )
            return answer
        except Exception as e:
            traceback.print_exc()
            await self.websocket.send_json({"type": "logs", "output": f"❌ Error occurred while calling agent: {str(e)}"})
            return ""

    async def create_search_queries(self):
        """ Creates the search queries for the given question.
        Args: None
        Returns: list[str]: The search queries for the given question
        """
        try:
            result = await self.call_agent(prompts.generate_search_queries_prompt(self.question))
            print(f"result = {result}", file=sys.stderr, flush=True)
            print(result)
            await self.websocket.send_json({"type": "logs", "output": f"🧠 I will conduct my research based on the following queries: {result}..."})
            # return json.loads(result)
            return json.loads(result) if type(result) != str else result.strip("[]").replace('"', '').split("] [")
        except Exception as e:
            traceback.print_exc()
            await self.websocket.send_json({"type": "logs", "output": f"❌ Error occurred while creating search queries: {str(e)}"})
            return []

    async def async_search(self, query):
        """ Runs the async search for the given query.
        Args: query (str): The query to run the async search for
        Returns: list[str]: The async search for the given query
        """
        try:
            search_results = json.loads(web_search(query))
            new_search_urls = self.get_new_urls([url.get("href") for url in search_results])

            await self.websocket.send_json(
                {"type": "logs", "output": f"🌐 Browsing the following sites for relevant information: {new_search_urls}..."})

            # Create a list to hold the coroutine objects
            tasks = [async_browse(url, query, self.websocket) for url in await new_search_urls]

            # Gather the results as they become available
            responses = await asyncio.gather(*tasks, return_exceptions=True)

            return responses
        except Exception as e:
            traceback.print_exc()
            await self.websocket.send_json({"type": "logs", "output": f"❌ Error occurred during async search: {str(e)}"})
            return []

    async def run_search_summary(self, query):
        """ Runs the search summary for the given query.
        Args: query (str): The query to run the search summary for
        Returns: str: The search summary for the given query
        """
        try:
            await self.websocket.send_json({"type": "logs", "output": f"🔎 Running research for '{query}'..."})

            responses = await self.async_search(query)

            result = "\n".join(responses)
            query_hash = hashlib.sha256(query.encode()).hexdigest()[:10]
            filename = f"./outputs/research-{query_hash}.txt"
            os.makedirs(os.path.dirname(filename), exist_ok=True)
            write_to_file(filename, result)
        except Exception as e:
            traceback.print_exc()
            await self.websocket.send_json({"type": "logs", "output": f"❌ Error occurred during search summary: {str(e)}"})

    async def conduct_research(self):
        """ Conducts the research for the given question.
        Args: None
        Returns: str: The research for the given question
        """
        try:
            self.research_summary = read_txt_files(self.dir_path) if os.path.isdir(self.dir_path) else ""

            if not self.research_summary:
                search_queries = await self.create_search_queries()
                num_queries = len(search_queries)
                for idx, query in enumerate(search_queries, 1):
                    await self.websocket.send_json(
                        {"type": "logs", "output": f"💡 Research query [{idx}/{num_queries}]: {query}..."})
                    await self.run_search_summary(query)
                    self.progress = idx / num_queries * 100
                    await self.websocket.send_json({"type": "progress", "progress": self.progress})

            await self.websocket.send_json(
                {"type": "logs", "output": f"Total research words: {len(self.research_summary.split(' '))}"})

            return self.research_summary
        except Exception as e:
            traceback.print_exc()
            await self.websocket.send_json({"type": "logs", "output": f"❌ Error occurred during research: {str(e)}"})
            return ""

    async def create_concepts(self):
        """ Creates the concepts for the given question.
        Args: None
        Returns: list[str]: The concepts for the given question
        """
        try:
            result = self.call_agent(prompts.generate_concepts_prompt(self.question, self.research_summary))
            await self.websocket.send_json({"type": "logs", "output": f"I will research based on the following concepts: {result}\n"})
            return json.loads(result)
        except Exception as e:
            traceback.print_exc()
            await self.websocket.send_json({"type": "logs", "output": f"❌ Error occurred while creating concepts: {str(e)}"})
            return []

    async def write_report(self, report_type, websocket):
        """ Writes the report for the given question.
        Args: None
        Returns: str: The report for the given question
        """
        try:
            report_type_func = prompts.get_report_by_type(report_type)
            await websocket.send_json(
                {"type": "logs", "output": f"✍️ Writing {report_type} for research task: {self.question}..."})
            answer = await self.call_agent(report_type_func(self.question, self.research_summary), stream=True,
                                           websocket=websocket)
            file_directory = f"./outputs/{self.directory_name}"
            os.makedirs(file_directory, exist_ok=True)
            file_path = f"{file_directory}/research_report"
            write_to_file(f"{file_path}.md", str(answer))  # Write the MD file
            path = await write_md_to_pdf(report_type, self.directory_name, await answer)
            return answer, path
        except Exception as e:
            traceback.print_exc()
            await self.websocket.send_json({"type": "logs", "output": f"❌ Error occurred while writing report: {str(e)}"})
            return "", ""

    async def write_lessons(self):
        """ Writes lessons on essential concepts of the research.
        Args: None
        Returns: None
        """
        try:
            concepts = await self.create_concepts()
            num_concepts = len(concepts)
            for idx, concept in enumerate(concepts, 1):
                await self.websocket.send_json(
                    {"type": "logs", "output": f"📖 Writing lesson [{idx}/{num_concepts}]: {concept}..."})
                answer = await self.call_agent(prompts.generate_lesson_prompt(concept), stream=True)
                await write_md_to_pdf("Lesson", self.directory_name, answer)
        except Exception as e:
            traceback.print_exc()
            await self.websocket.send_json({"type": "logs", "output": f"❌ Error occurred while writing lessons: {str(e)}"})
```

2. client\index.html

```html
<!DOCTYPE html>
<html lang="en">

<head>
    <title>GPT Researcher</title>
    <meta name="description" content="A research assistant powered by GPT-4">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <link rel="icon" href="./static/favicon.ico">
    <link rel="preconnect" href="https://fonts.googleapis.com">
    <link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
    <link href="https://fonts.googleapis.com/css2?family=Montserrat:wght@400;700&display=swap" rel="stylesheet">
    <link href="https://stackpath.bootstrapcdn.com/bootstrap/4.5.2/css/bootstrap.min.css" rel="stylesheet">
    <link rel="stylesheet" href="/site/styles.css"/>
    <style>
        .avatar {
            width: 60px;
            height: 60px;
            border-radius: 50%;
        }

        .agent-name {
            text-align: center;
        }

        .agent-item {
            display: flex;
            flex-direction: column;
            align-items: center;
        }

        .agent-choices {
            display: none;
        }

        .btn-show {
            display: none;
        }

        /* Custom styles */
        body {
            background-color: #f2f2f2;
            font-family: 'Montserrat', sans-serif;
        }

        .landing {
            background-color: #9867F0;
            padding: 100px 0;
            text-align: center;
            color: #fff;
        }

        .landing h1 {
            font-size: 4rem;
            font-weight: bold;
            line-height: 1.2;
        }

        .landing span {
            background-image: none;
            color: #ED4E50;
        }

        #form {
            margin-top: 50px;
        }

        .form-group {
            margin-bottom: 20px;
        }

        .agent-question {
            font-size: 1.2rem;
            font-weight: bold;
        }

        .form-control {
            border-radius: 0;
        }

        .button-padding {
            padding: 10px 20px;
        }

        .margin-div {
            margin-top: 50px;
        }

        h2 {
            font-size: 2rem;
            font-weight: bold;
            margin-bottom: 20px;
        }

        #output {
            background-color: rgba(255, 255, 255, 0.1);
            padding: 20px;
            border-radius: 5px;
        }

        #reportContainer {
            background-color: rgba(255, 255, 255, 0.1);
            padding: 20px;
            border-radius: 5px;
        }

        .btn-secondary {
            background-color: #ED4E50;
            border-color: #ED4E50;
        }

        .btn-secondary:hover {
            background-color: #C43C3F;
            border-color: #C43C3F;
        }

        footer {
            background-color: assafelovic#333;
            color: #fff;
            padding: 20px;
            text-align: center;
        }

        footer a {
            color: #fff;
        }
    </style>
</head>

<body>

<main class="container" id="form">
    <div class="agent-item"><img src="/static/defaultAgentAvatar.JPG" class="avatar"
                                                alt="Auto Agent"></div>
    <form method="POST" class="mt-3" onsubmit="GPTResearcher.startResearch(); return false;">
        <div class="form-group">
            <label for="task" class="agent-question">What would you like me to research next?</label>
            <input type="text" id="task" name="task" class="form-control" required>
            <input type="radio" name="agent" id="autoAgent" value="Auto Agent" checked hidden>
        </div>
        <div class="form-group">
            <div class="row">

            </div>
            <button type="button" id="btnShowAuto" class="btn btn-secondary mt-3 btn-show">Auto Agent</button>
        </div>
        <div class="form-group">
            <label for="report_type" class="agent-question">What type of report would you like me to generate?</label>
            <select name="report_type" class="form-control" required>
                <option value="research_report">Research Report</option>
                <option value="resource_report">Resource Report</option>
                <option value="outline_report">Outline Report</option>
            </select>
        </div>
        <input type="submit" value="Research" class="btn btn-primary button-padding">
    </form>

    <div class="margin-div">
        <h2>Agent Output</h2>
        <p class="mt-2 text-left" style="font-size: 0.8rem;">An agent tailored specifically to your task
                        will be generated to provide the most precise and relevant research results.</p>
        <div id="output"></div>
        <div class="progress mt-3">
            <div class="progress-bar" role="progressbar" style="width: 0%;" aria-valuenow="0" aria-valuemin="0" aria-valuemax="100"></div>
        </div>
    </div>
    <div class="margin-div">
        <h2>Research Report</h2>
        <div id="reportContainer"></div>
        <button onclick="copyToClipboard()" class="btn btn-secondary mt-3">Copy to clipboard</button>
        <a id="downloadLink" href="#" class="btn btn-secondary mt-3" target="_blank">Download as PDF</a>
    </div>
</main>

<footer>
    <p>GPT Researcher &copy; 2023 | <a target="_blank" href="https://github.com/assafelovic/gpt-researcher">GitHub
        Page</a></p>
</footer>

<script src="https://cdnjs.cloudflare.com/ajax/libs/showdown/1.9.1/showdown.min.js"></script>
<script src="/site/scripts.js"></script>
<script>
    const btnChoose = document.getElementById('btnChoose');
    const btnShowAuto = document.getElementById('btnShowAuto');
    const autoAgentDiv = document.getElementById('autoAgentDiv');
    const agentChoices = document.getElementsByClassName('agent-choices');

    btnChoose.addEventListener('click', function () {
        btnShowAuto.style.display = 'inline-block';
        btnChoose.style.display = 'none';
        autoAgentDiv.style.display = 'none';
        agentChoices[0].style.display = 'flex';
    });

    btnShowAuto.addEventListener('click', function () {
        btnShowAuto.style.display = 'none';
        btnChoose.style.display = 'inline-block';
        autoAgentDiv.style.display = 'flex';
        agentChoices[0].style.display = 'none';
    });

    const progressBar = document.querySelector('.progress-bar');
    const progress = 0;
    progressBar.style.width = `${progress}%`;
    progressBar.setAttribute('aria-valuenow', progress);
</script>
</body>

</html>
```

Now, when the agent is conducting research and updating the progress, it will send a JSON message of type "progress" with the current progress value to the WebSocket. The client-side JavaScript code can then update the progress bar accordingly.

USER: No filename provided before ``` in file listing

ASSISTANT: No changes are needed.
  • Loading branch information
Nicolas Alcaraz authored and dablon committed Feb 19, 2024
1 parent 1f38764 commit 454aeea
Show file tree
Hide file tree
Showing 2 changed files with 3 additions and 2 deletions.
1 change: 1 addition & 0 deletions agent/research_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -169,6 +169,7 @@ async def conduct_research(self):
{"type": "logs", "output": f"💡 Research query [{idx}/{num_queries}]: {query}..."})
await self.run_search_summary(query)
self.progress = idx / num_queries * 100
await self.websocket.send_json({"type": "progress", "progress": self.progress})

await self.websocket.send_json(
{"type": "logs", "output": f"Total research words: {len(self.research_summary.split(' '))}"})
Expand Down
4 changes: 2 additions & 2 deletions frontend/index.html
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,7 @@ <h2>Agent Output</h2>
will be generated to provide the most precise and relevant research results.</p>
<div id="output"></div>
<div class="progress mt-3">
<div class="progress-bar" role="progressbar" style="width: {{ progress }}%;" aria-valuenow="{{ progress }}" aria-valuemin="0" aria-valuemax="100"></div>
<div class="progress-bar" role="progressbar" style="width: 0%;" aria-valuenow="0" aria-valuemin="0" aria-valuemax="100"></div>
</div>
</div>
<div class="margin-div">
Expand Down Expand Up @@ -205,7 +205,7 @@ <h2>Research Report</h2>
});

const progressBar = document.querySelector('.progress-bar');
const progress = {{ progress }};
const progress = 0;
progressBar.style.width = `${progress}%`;
progressBar.setAttribute('aria-valuenow', progress);
</script>
Expand Down

0 comments on commit 454aeea

Please sign in to comment.