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},
+ }