diff --git a/src/knowledge/capability.py b/src/knowledge/capability.py index 6edbcc7ce..36bbeafc9 100644 --- a/src/knowledge/capability.py +++ b/src/knowledge/capability.py @@ -16,14 +16,18 @@ from collections.abc import Sequence from dataclasses import dataclass -from typing import TYPE_CHECKING, Optional +from typing import Optional from pydantic_ai.capabilities import AbstractCapability -from pydantic_ai.tools import AgentDepsT +from pydantic_ai.messages import TextContent, UserContent +from pydantic_ai.models import ModelRequestContext +from pydantic_ai.tools import AgentDepsT, RunContext from sentence_transformers import CrossEncoder -if TYPE_CHECKING: - from knowledge.sources.base import KnowledgeSource +from knowledge.search import search_sources +from knowledge.sources.base import KnowledgeSource +from knowledge.types import KnowledgeMatch +from knowledge.utils import append_latest_message, extract_latest_message_text @dataclass @@ -38,3 +42,53 @@ class Knowledge(AbstractCapability[AgentDepsT]): `'cross-encoder/ms-marco-MiniLM-L-6-v2'`) used to rerank matches by relevance to the query before they're returned or injected. Loaded lazily on first use; requires the `sentence-transformers` package.""" + + async def before_model_request( + self, + ctx: RunContext[AgentDepsT], + request_context: ModelRequestContext, + ) -> ModelRequestContext: + """Search every auto-mode source concurrently, rank the merged matches, and inject them.""" + + def to_message(matches: list[KnowledgeMatch]) -> Optional[UserContent]: + if not matches: + return None + + match_blocks = [ + f'\n' + f"{match.content.strip()}\n" + f"" + for match in matches + ] + + body = "\n\n".join(match_blocks) + + return TextContent( + content=f"\n{body}\n", + metadata={ + "kind": "knowledge", + "knowledge_scores": {match.id: match.score for match in matches}, + "knowledge_metadata": { + match.id: match.metadata for match in matches + }, + "knowledge_sources": { + match.id: match.metadata.get("knowledge_source") + for match in matches + }, + }, + ) + + auto_sources = [source for source in self.sources if "auto" in source.mode] + if not auto_sources: + return request_context + + prompt = extract_latest_message_text(request_context.messages) + if not prompt: + return request_context + + matches = await search_sources(auto_sources, prompt, self.reranker) + if message := to_message(matches): + append_latest_message(request_context.messages, message) + + return request_context diff --git a/tests/unit/knowledge/test_capability.py b/tests/unit/knowledge/test_capability.py index 9727640fc..2e2769003 100644 --- a/tests/unit/knowledge/test_capability.py +++ b/tests/unit/knowledge/test_capability.py @@ -3,9 +3,20 @@ from __future__ import annotations from dataclasses import dataclass +from typing import Literal +import pytest from pydantic_ai.capabilities import AbstractCapability -from pytest_mock import MockerFixture, MockType +from pydantic_ai.messages import ( + ModelRequest, + ModelResponse, + SystemPromptPart, + TextContent, + TextPart, + UserPromptPart, +) +from pydantic_ai.models import ModelRequestContext +from pytest_mock import AsyncMockType, MockerFixture, MockType from knowledge.capability import Knowledge from knowledge.sources.base import KnowledgeSource @@ -49,3 +60,314 @@ def test_stores_optional_reranker(self, mocker: MockerFixture) -> None: capability = Knowledge(sources=[source], reranker=reranker) assert capability.reranker is reranker + + +def _source(name: str, *modes: Literal["tool", "auto"]) -> FakeKnowledgeSource: + """Build a fake knowledge source with the given exposure modes.""" + mode: set[Literal["tool", "auto"]] = set(modes) if modes else {"tool"} + return FakeKnowledgeSource(name=name, mode=mode) + + +def _match( + content: str, + score: float, + *, + match_id: str = "match-id", + source: str | None = "doc://source", + metadata: dict[str, object] | None = None, +) -> KnowledgeMatch: + """Build a knowledge match for auto-mode injection tests.""" + return KnowledgeMatch( + content=content, + score=score, + id=match_id, + source=source, + metadata=metadata or {}, + ) + + +def _user_request(text: str = "how do I restart?") -> ModelRequest: + """Build a model request whose latest part is a user prompt.""" + return ModelRequest(parts=[UserPromptPart(content=text)]) + + +def _request_context( + mocker: MockerFixture, + messages: list[ModelRequest | ModelResponse] | None = None, +) -> ModelRequestContext: + """Build a `ModelRequestContext` wrapping `messages`.""" + return ModelRequestContext( + model=mocker.Mock(), + messages=list(messages or [_user_request()]), + model_settings=None, + model_request_parameters=mocker.Mock(), + ) + + +def _injected_knowledge(request: ModelRequest) -> TextContent: + """Return the knowledge `TextContent` appended to `request`.""" + added = request.parts[-1] + assert isinstance(added, UserPromptPart) + content = added.content + assert isinstance(content, list) + text = content[0] + assert isinstance(text, TextContent) + return text + + +class TestKnowledgeBeforeModelRequestAutoMode: + """Tests for `Knowledge.before_model_request` auto-mode injection.""" + + @pytest.fixture(name="search_sources") + def search_sources_fixture(self, mocker: MockerFixture) -> AsyncMockType: + """Patch `search_sources` where the capability looks it up.""" + return mocker.patch( + "knowledge.capability.search_sources", + new_callable=mocker.AsyncMock, + return_value=[], + ) + + @pytest.mark.asyncio + async def test_returns_unchanged_when_no_auto_sources( + self, + mocker: MockerFixture, + search_sources: AsyncMockType, + ) -> None: + """Tool-only sources are not searched and the request is left untouched.""" + request = _user_request() + request_context = _request_context(mocker, [request]) + capability = Knowledge(sources=[_source("tools")]) + + result = await capability.before_model_request(mocker.Mock(), request_context) + + assert result is request_context + assert len(request.parts) == 1 + search_sources.assert_not_called() + + @pytest.mark.asyncio + async def test_returns_unchanged_when_sources_are_empty( + self, + mocker: MockerFixture, + search_sources: AsyncMockType, + ) -> None: + """An empty source list skips auto injection.""" + request = _user_request() + request_context = _request_context(mocker, [request]) + capability = Knowledge(sources=[]) + + result = await capability.before_model_request(mocker.Mock(), request_context) + + assert result is request_context + assert len(request.parts) == 1 + search_sources.assert_not_called() + + @pytest.mark.asyncio + async def test_returns_unchanged_when_prompt_is_missing( + self, + mocker: MockerFixture, + search_sources: AsyncMockType, + ) -> None: + """Auto sources are not searched when history has no user prompt.""" + request = ModelRequest(parts=[SystemPromptPart(content="system")]) + request_context = _request_context(mocker, [request]) + capability = Knowledge(sources=[_source("docs", "auto")]) + + result = await capability.before_model_request(mocker.Mock(), request_context) + + assert result is request_context + assert len(request.parts) == 1 + search_sources.assert_not_called() + + @pytest.mark.asyncio + async def test_returns_unchanged_when_prompt_is_empty( + self, + mocker: MockerFixture, + search_sources: AsyncMockType, + ) -> None: + """An empty user prompt is treated as missing and skips search.""" + request = _user_request("") + request_context = _request_context(mocker, [request]) + capability = Knowledge(sources=[_source("docs", "auto")]) + + result = await capability.before_model_request(mocker.Mock(), request_context) + + assert result is request_context + assert len(request.parts) == 1 + search_sources.assert_not_called() + + @pytest.mark.asyncio + async def test_returns_unchanged_when_search_finds_no_matches( + self, + mocker: MockerFixture, + search_sources: AsyncMockType, + ) -> None: + """An empty search result does not append a knowledge part.""" + request = _user_request() + request_context = _request_context(mocker, [request]) + capability = Knowledge(sources=[_source("docs", "auto")]) + + result = await capability.before_model_request(mocker.Mock(), request_context) + + assert result is request_context + assert len(request.parts) == 1 + search_sources.assert_awaited_once() + + @pytest.mark.asyncio + async def test_searches_only_auto_mode_sources_with_latest_prompt( + self, + mocker: MockerFixture, + search_sources: AsyncMockType, + ) -> None: + """Tool-only sources are skipped; auto sources keep declaration order.""" + tool = _source("tools") + auto_docs = _source("docs", "auto") + auto_wiki = _source("wiki", "auto") + earlier = _user_request("older question") + latest = _user_request("how do I restart?") + request_context = _request_context( + mocker, + [ + earlier, + ModelResponse(parts=[TextPart(content="reply")]), + latest, + ], + ) + capability = Knowledge(sources=[tool, auto_docs, auto_wiki]) + + await capability.before_model_request(mocker.Mock(), request_context) + + search_sources.assert_awaited_once_with( + [auto_docs, auto_wiki], + "how do I restart?", + None, + ) + assert len(earlier.parts) == 1 + + @pytest.mark.asyncio + async def test_includes_source_with_both_auto_and_tool_modes( + self, + mocker: MockerFixture, + search_sources: AsyncMockType, + ) -> None: + """A source exposing both modes is still searched automatically.""" + both = _source("docs", "auto", "tool") + request_context = _request_context(mocker) + capability = Knowledge(sources=[both]) + + await capability.before_model_request(mocker.Mock(), request_context) + + search_sources.assert_awaited_once_with([both], "how do I restart?", None) + + @pytest.mark.asyncio + async def test_forwards_reranker_to_search_sources( + self, + mocker: MockerFixture, + search_sources: AsyncMockType, + ) -> None: + """The capability's cross-encoder is passed through to search.""" + auto = _source("docs", "auto") + reranker: MockType = mocker.Mock() + request_context = _request_context(mocker) + capability = Knowledge(sources=[auto], reranker=reranker) + + await capability.before_model_request(mocker.Mock(), request_context) + + search_sources.assert_awaited_once_with([auto], "how do I restart?", reranker) + + @pytest.mark.asyncio + async def test_appends_formatted_knowledge_to_latest_request( + self, + mocker: MockerFixture, + search_sources: AsyncMockType, + ) -> None: + """Matches are injected as delimited XML on the latest user request.""" + original = UserPromptPart(content="how do I restart?") + request = ModelRequest(parts=[original]) + earlier = _user_request("older question") + request_context = _request_context( + mocker, + [earlier, ModelResponse(parts=[TextPart(content="reply")]), request], + ) + search_sources.return_value = [ + _match( + " Restart systemd ", + 0.9, + match_id="m1", + source="doc://restart", + metadata={"knowledge_source": "docs", "page": 1}, + ) + ] + capability = Knowledge(sources=[_source("docs", "auto")]) + + result = await capability.before_model_request(mocker.Mock(), request_context) + + assert result is request_context + assert request.parts[0] is original + assert len(earlier.parts) == 1 + text = _injected_knowledge(request) + assert text.content == ( + "\n" + '\n' + "Restart systemd\n" + "\n" + "" + ) + assert text.metadata == { + "kind": "knowledge", + "knowledge_scores": {"m1": 0.9}, + "knowledge_metadata": { + "m1": {"knowledge_source": "docs", "page": 1}, + }, + "knowledge_sources": {"m1": "docs"}, + } + + @pytest.mark.asyncio + async def test_formats_multiple_matches_and_missing_source_fields( + self, + mocker: MockerFixture, + search_sources: AsyncMockType, + ) -> None: + """Matches are joined with a blank line; missing fields get safe defaults.""" + request = _user_request() + request_context = _request_context(mocker, [request]) + search_sources.return_value = [ + _match( + "alpha", + 0.8, + match_id="m1", + source="doc://a", + metadata={"knowledge_source": "docs"}, + ), + _match( + "beta", + 0.2, + match_id="m2", + source=None, + metadata={}, + ), + ] + capability = Knowledge(sources=[_source("docs", "auto")]) + + await capability.before_model_request(mocker.Mock(), request_context) + + text = _injected_knowledge(request) + assert text.content == ( + "\n" + '\n' + "alpha\n" + "\n" + "\n" + '\n' + "beta\n" + "\n" + "" + ) + assert text.metadata == { + "kind": "knowledge", + "knowledge_scores": {"m1": 0.8, "m2": 0.2}, + "knowledge_metadata": { + "m1": {"knowledge_source": "docs"}, + "m2": {}, + }, + "knowledge_sources": {"m1": "docs", "m2": None}, + }