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

Code before:
        await self.send_message("Determining if external documentation is needed...")
        llm_response: str = await llm(convo)
        available_docsets = dict(available_docsets)
        if llm_response.strip().lower() == "done":
            return {}
        else:
            selected_keys = llm_response.splitlines()
            return {k: available_docsets[k] for k in selected_keys}

    async def _create_queries(self, docsets: dict[str, str]) -> dict[str, list[str]]:
        """Return queries we have to make to the docs API.

        Key is the docset_key and value is the list of queries for that docset.

        """
        queries = {}
        await self.send_message("Getting relevant documentation for the following topics:")
        for k, short_desc in docsets.items():
            llm = self.get_llm()
            convo = AgentConvo(self).template(
                "create_queries",
                short_description=short_desc,
                current_task=self.current_state.current_task,
            )
            llm_response: str = await llm(convo)
            if llm_response.strip().lower() == "done":
                continue
            else:
                queries[k] = llm_response.splitlines()

        return queries

    async def _fetch_snippets(self, queries: dict[str, list[str]]) -> list[tuple]:
        """Query the docs API and fetch the documentation snippets.

        Returns a list of tuples: (docset_key, snippets).

        """
        url = urljoin(EXTERNAL_DOCUMENTATION_API, "query")
        snippets: list[tuple] = []
        async with httpx.AsyncClient(transport=httpx.AsyncHTTPTransport(retries=3)) as client:
            reqs = []
            ordered_keys = []
            for docset_key, qs in queries.items():
                reqs.append(client.get(url, params={"q": qs, "doc_key": docset_key}))
                ordered_keys.append(docset_key)

            try:
                results = await asyncio.gather(*reqs)
            except httpx.HTTPError:
                log.warning("Failed to fetch documentation snippets", exc_info=True)

Review comment (suggestion):
We are expecting response `done` but nowhere in prompt I see instruction to respond with `done`. Also, see comment in `create_queries.prompt` regarding using schema and parser

Code after:
        convo = (
            AgentConvo(self)
            .template(
                "select_docset",
                current_task=self.current_state.current_task,
                available_docsets=available_docsets,
            )
            .require_schema(SelectedDocsets)
        )
        await self.send_message("Determining if external documentation is needed for the next task...")
        llm_response: SelectedDocsets = await llm(convo, parser=JSONParser(spec=SelectedDocsets))
        available_docsets = dict(available_docsets)
        return {k: available_docsets[k] for k in llm_response.docsets}

    async def _create_queries(self, docsets: dict[str, str]) -> dict[str, list[str]]:
        """Return queries we have to make to the docs API.

        Key is the docset_key and value is the list of queries for that docset.

        """
        queries = {}
        await self.send_message("Getting relevant documentation for the following topics:")
        for k, short_desc in docsets.items():
            llm = self.get_llm()
            convo = (
                AgentConvo(self)
                .template(
                    "create_docs_queries",
                    short_description=short_desc,
                    current_task=self.current_state.current_task,
                )
                .require_schema(DocQueries)
            )
            llm_response: DocQueries = await llm(convo, parser=JSONParser(spec=DocQueries))
            if llm_response.queries:
                queries[k] = llm_response.queries

        return queries

    async def _fetch_snippets(self, queries: dict[str, list[str]]) -> list[tuple]:
        """Query the docs API and fetch the documentation snippets.

        Returns a list of tuples: (docset_key, snippets).

        """
        url = urljoin(EXTERNAL_DOCUMENTATION_API, "query")
        snippets: list[tuple] = []
        async with httpx.AsyncClient(transport=httpx.AsyncHTTPTransport(retries=3)) as client:
            reqs = []
            ordered_keys = []
            for docset_key, qs in queries.items():