Repository: Pythagora-io/gpt-pilot
File: core/agents/external_docs.py (Python)

Code before:
log = get_logger(__name__)


class ExternalDocumentation(BaseAgent):
    """Agent in charge of collecting and storing additional documentation.

    Docs are per task and are stores in the `tasks` variable in the project state.
    This agent ensures documentation is collected only once per task.

    Agent does 2 LLM interactions:
        1. Ask the LLM to select useful documentation from a predefined list.
        2. Ask the LLM to come up with a query to use to fetch the actual documentation snippets.

    Agent does 2 calls to our documentation API:
        1. Fetch all the available docsets. `docset` is a collection of documentation snippets
           for a single topic, eg. VueJS API Reference docs.
        2. Fetch the documentation snippets for given queries.

    """

    agent_type = "external-docs"
    display_name = "Documentation"

    async def run(self) -> AgentResponse:
        current_task = self.current_state.current_task
        if not current_task:
            # If we have no active task, there's no docs to collect
            return AgentResponse.done(self)

        available_docsets = await self._get_available_docsets()
        selected_docsets = await self._select_docsets(available_docsets)
        telemetry.set("docsets_used", selected_docsets)

        if not selected_docsets:
            log.info("No documentation selected for this task.")
            await self._store_docs([], available_docsets)
            return AgentResponse.done(self)

        queries = await self._create_queries(selected_docsets)
        doc_snippets = await self._fetch_snippets(queries)
        telemetry.set("doc_snippets_stored", len(doc_snippets))

        await self._store_docs(doc_snippets, available_docsets)
        return AgentResponse.done(self)

    async def _get_available_docsets(self) -> list[tuple]:
        url = urljoin(EXTERNAL_DOCUMENTATION_API, "docsets")
        client = httpx.Client(transport=httpx.HTTPTransport(retries=3))
        try:
            resp = client.get(url)
        except httpx.HTTPError:

Review comment (bug):
In case there is no current_task then this agent should not be called. Or is there a reason why we would want that?

In that case, it feels like we could call this agent without making any changes to the state and then the state will be committed and this agent will be called again... infinite loop?

Code after:
from core.telemetry import telemetry

log = get_logger(__name__)


class DocQueries(BaseModel):
    queries: list[str]


class SelectedDocsets(BaseModel):
    docsets: list[str]


class ExternalDocumentation(BaseAgent):
    """Agent in charge of collecting and storing additional documentation.

    Docs are per task and are stores in the `tasks` variable in the project state.
    This agent ensures documentation is collected only once per task.

    Agent does 2 LLM interactions:
        1. Ask the LLM to select useful documentation from a predefined list.
        2. Ask the LLM to come up with a query to use to fetch the actual documentation snippets.

    Agent does 2 calls to our documentation API:
        1. Fetch all the available docsets. `docset` is a collection of documentation snippets
           for a single topic, eg. VueJS API Reference docs.
        2. Fetch the documentation snippets for given queries.

    """

    agent_type = "external-docs"
    display_name = "Documentation"

    async def run(self) -> AgentResponse:
        available_docsets = await self._get_available_docsets()
        selected_docsets = await self._select_docsets(available_docsets)
        await telemetry.trace_code_event("docsets_used", selected_docsets)

        if not selected_docsets:
            log.info("No documentation selected for this task.")
            await self._store_docs([], available_docsets)
            return AgentResponse.done(self)

        queries = await self._create_queries(selected_docsets)
        doc_snippets = await self._fetch_snippets(queries)
        await telemetry.trace_code_event("doc_snippets", {"num_stored": len(doc_snippets)})

        await self._store_docs(doc_snippets, available_docsets)
        return AgentResponse.done(self)

    async def _get_available_docsets(self) -> list[tuple]: