diff --git a/core/agents/external_docs.py b/core/agents/external_docs.py index 7c7f1b0c9..6e818df94 100644 --- a/core/agents/external_docs.py +++ b/core/agents/external_docs.py @@ -92,7 +92,7 @@ async def _select_docsets(self, available_docsets: list[tuple]) -> dict[str, str 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} + return {k: available_docsets[k] for k in llm_response.docsets if k in available_docsets} async def _create_queries(self, docsets: dict[str, str]) -> dict[str, list[str]]: """Return queries we have to make to the docs API. diff --git a/core/prompts/external-docs/select_docset.prompt b/core/prompts/external-docs/select_docset.prompt index 8293c9080..e714c3288 100644 --- a/core/prompts/external-docs/select_docset.prompt +++ b/core/prompts/external-docs/select_docset.prompt @@ -12,4 +12,4 @@ Here is the list of available documentations: {{ docset[0], docset[1] }} {% endfor %} -Now, give me the list of the additional documentation that you would like to use to complete the task listed above. Return only the documentation that is absolutely required for the given task. If there is no additional documentation in the list that you would like to use, return an empty list. +Now, give me the list of the additional documentation that you would like to use to complete the task listed above. Return only the documentation that is absolutely required for the given task, and only from the list of available documentations provided above. If there is no additional documentation in the list that you would like to use, return an empty list. diff --git a/tests/agents/test_external_docs.py b/tests/agents/test_external_docs.py index d36ce0a8d..a3eed86e6 100644 --- a/tests/agents/test_external_docs.py +++ b/tests/agents/test_external_docs.py @@ -21,6 +21,21 @@ async def test_stores_documentation_snippets_for_task(agentcontext): assert ed.next_state.docs[0]["key"] == "vuejs-api-ref" +@pytest.mark.asyncio +async def test_continues_without_docs_for_invalid_docset(agentcontext): + sm, _, ui, mock_llm = agentcontext + + sm.current_state.tasks = [{"description": "Some VueJS task", "status": "todo"}] + await sm.commit() + + ed = ExternalDocumentation(sm, ui) + ed.get_llm = mock_llm( + side_effect=[SelectedDocsets(docsets=["doesnt-exist"]), DocQueries(queries=["VueJS component model"])] + ) + await ed.run() + assert ed.next_state.docs == [] + + @pytest.mark.asyncio async def test_continues_without_docs_if_api_is_down(agentcontext): sm, _, ui, _ = agentcontext