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

Code before:
    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:
            # In case of any errors, we'll proceed without the documentation
            log.warning("Failed to fetch available docsets due to an error.", exc_info=True)
            return []

        log.debug(f"Fetched {len(resp.json())} docsets.")
        return resp.json()

Review comment (suggestion):
Our telemetry is being collected at:

1. end of initial app
2. end of feature
3. end of Pythagora process

That means that every task will overwrite this telemetry value. I propose using `trace_code_event()` which we can manually trigger each time we collect external docs.

Code after:
    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]:
        url = urljoin(EXTERNAL_DOCUMENTATION_API, "docsets")
        client = httpx.Client(transport=httpx.HTTPTransport(retries=3))
        try:
            resp = client.get(url)
        except httpx.HTTPError:
            # In case of any errors, we'll proceed without the documentation