diff --git a/docsgpt/agents/headless_runner.py b/docsgpt/agents/headless_runner.py index cddf693e..4c1098c1 100644 --- a/docsgpt/agents/headless_runner.py +++ b/docsgpt/agents/headless_runner.py @@ -15,15 +15,16 @@ from docsgpt.api.answer.services.prompt_renderer import ( resolve_prompt_skeleton, ) from docsgpt.api.answer.services.stream_processor import ( + authorized_agent_sources, authorized_prompt_id, get_prompt, ) -from docsgpt.api.user.resource_access import ref_principal from docsgpt.core.settings import settings from docsgpt.quotas.service import QuotaExceededError, QuotaService +from docsgpt.retriever.dispatcher import build_dispatcher from docsgpt.retriever.retriever_creator import RetrieverCreator -from docsgpt.storage.db.repositories.sources import SourcesRepository from docsgpt.storage.db.session import db_readonly +from docsgpt.storage.db.source_config import SourceConfig logger = logging.getLogger(__name__) @@ -148,20 +149,22 @@ def _run_agent_headless( raise QuotaExceededError(exceeded) retriever_kind = agent_config.get("retriever", "classic") - source_id = agent_config.get("source_id") or agent_config.get("source") - source_active: Any = {} - if source_id: + # Every source a chat with this agent searches: the primary and the + # extras, each owned or team-shared to the owner, else attached by an + # editor who still qualifies. + sources_row = dict(agent_config) + sources_row["source_id"] = agent_config.get("source_id") or agent_config.get("source") + primary, source_docs = None, [] + if sources_row["source_id"] or sources_row.get("extra_source_ids"): with db_readonly() as conn: - # Owned or team-shared to the owner, else attached by an editor - # who still qualifies; read unscoped once authorized. - src_row = ( - SourcesRepository(conn).get_by_id(str(source_id)) - if ref_principal(conn, "agent", agent_config, "source", str(source_id)) - else None - ) - if src_row: - source_active = str(src_row["id"]) - retriever_kind = src_row.get("retriever", retriever_kind) + primary, source_docs = authorized_agent_sources(conn, sources_row) + if primary: + retriever_kind = primary.get("retriever") or retriever_kind + per_source = [ + {"id": str(doc["id"]), "retrieval": SourceConfig.parse(doc.get("config")).retrieval} + for doc in source_docs + ] + source_active: Any = [entry["id"] for entry in per_source] or {} source = {"active_docs": source_active} # ``chunks=0`` switches retrieval off; only a missing value takes the default. raw_chunks = agent_config.get("chunks") @@ -196,8 +199,7 @@ def _run_agent_headless( system_api_key = get_api_key_for_provider(provider or settings.LLM_PROVIDER) doc_token_limit = calculate_doc_token_budget(model_id=model_id, user_id=owner) - retriever = RetrieverCreator.create_retriever( - retriever_kind, + retriever_kwargs: Dict[str, Any] = dict( source=source, chat_history=chat_history or [], prompt=prompt, @@ -208,6 +210,13 @@ def _run_agent_headless( agent_id=agent_id, decoded_token=decoded_token, ) + # Routed per source like a chat's pre-fetch, so each source keeps its own + # retriever and retrieval settings. + retriever = build_dispatcher( + lambda: RetrieverCreator.create_retriever(retriever_kind, **retriever_kwargs), + sources=per_source, + **retriever_kwargs, + ) retrieved_docs: List[Dict[str, Any]] = [] try: docs = retriever.search(query) diff --git a/docsgpt/api/answer/services/stream_processor.py b/docsgpt/api/answer/services/stream_processor.py index 891bcca8..64395ca4 100644 --- a/docsgpt/api/answer/services/stream_processor.py +++ b/docsgpt/api/answer/services/stream_processor.py @@ -4,7 +4,7 @@ import json import logging import uuid from pathlib import Path -from typing import Any, Callable, Dict, Optional, Set, TypeVar +from typing import Any, Callable, Dict, List, Optional, Set, Tuple, TypeVar from flask import after_this_request @@ -183,6 +183,39 @@ def _agent_source_doc(conn: Any, sources_repo: Any, agent: dict, source_id: Any) return sources_repo.get_by_id(str(source_id)) +def authorized_agent_sources(conn: Any, agent: dict) -> Tuple[Optional[dict], List[dict]]: + """The source rows an agent run retrieves from: primary first, then extras. + + Each is authorized like :func:`_agent_source_doc` (the owner, else the + editor who attached it), and a source listed twice appears once. + + Args: + conn: An open database connection. + agent: The ``agents`` row. + + Returns: + The primary source row (None when unset or not usable) and every + usable row in run order. + """ + sources_repo = SourcesRepository(conn) + primary: Optional[dict] = None + rows: List[dict] = [] + seen: set = set() + refs = [(True, agent.get("source_id"))] + refs.extend((False, sid) for sid in agent.get("extra_source_ids") or []) + for is_primary, sid_raw in refs: + if not sid_raw: + continue + source_doc = _agent_source_doc(conn, sources_repo, agent, sid_raw) + if not source_doc or str(source_doc["id"]) in seen: + continue + if is_primary: + primary = source_doc + seen.add(str(source_doc["id"])) + rows.append(source_doc) + return primary, rows + + def _wiki_write_owner(conn: Any, source_id: str, caller: str) -> Optional[str]: """The owner id to write a wiki source as, when ``caller`` may edit it.""" from docsgpt.api.user.resource_access import resolve @@ -751,7 +784,6 @@ class StreamProcessor: agent = AgentsRepository(conn).find_by_key(api_key) if not agent: raise Exception("Invalid API Key, please generate a new key", 401) - sources_repo = SourcesRepository(conn) # The repo dict uses "user_id" — the streaming path expects # a "user" key (legacy Mongo shape) for identity propagation. data: Dict[str, Any] = dict(agent) @@ -761,68 +793,29 @@ class StreamProcessor: # ``_configure_source`` ignores an empty ``data["sources"]``, # so the primary must appear in the union too — not only in # the legacy ``data["source"]`` slot. - sources_list: list = [] - seen: set = set() - primary_id = agent.get("source_id") - # ``sources`` row may have NULL ``retriever``/``chunks`` — - # fall back to the agent's value (``dict.get`` returns None - # even when the key exists with value None). - if primary_id: - source_doc = _agent_source_doc(conn, sources_repo, agent, primary_id) - if source_doc: - sid = str(source_doc["id"]) - data["source"] = sid - src_retriever = source_doc.get("retriever") - if src_retriever: - data["retriever"] = src_retriever - src_chunks = source_doc.get("chunks") - if src_chunks is not None: - data["chunks"] = src_chunks - sources_list.append( - { - "id": sid, - "retriever": src_retriever or "classic", - "chunks": ( - src_chunks if src_chunks is not None - else data.get("chunks", "6") - ), - # Per-source behaviour contract (lenient read). - "retrieval": SourceConfig.parse( - source_doc.get("config") - ).retrieval, - } - ) - seen.add(sid) - else: - data["source"] = None - else: - data["source"] = None - - for sid_raw in agent.get("extra_source_ids") or []: - if not sid_raw: - continue - source_doc = _agent_source_doc(conn, sources_repo, agent, sid_raw) - if not source_doc: - continue - sid = str(source_doc["id"]) - if sid in seen: - continue - src_retriever = source_doc.get("retriever") - src_chunks = source_doc.get("chunks") - sources_list.append( - { - "id": sid, - "retriever": src_retriever or "classic", - "chunks": ( - src_chunks if src_chunks is not None - else data.get("chunks", "6") - ), - "retrieval": SourceConfig.parse( - source_doc.get("config") - ).retrieval, - } - ) - seen.add(sid) + primary, source_docs = authorized_agent_sources(conn, agent) + # ``sources`` row may have NULL ``retriever``/``chunks`` — fall back to + # the agent's value (``dict.get`` returns None even when the key + # exists with value None). The primary's own values win for the agent. + data["source"] = str(primary["id"]) if primary else None + if primary: + if primary.get("retriever"): + data["retriever"] = primary["retriever"] + if primary.get("chunks") is not None: + data["chunks"] = primary["chunks"] + sources_list: list = [ + { + "id": str(source_doc["id"]), + "retriever": source_doc.get("retriever") or "classic", + "chunks": ( + source_doc["chunks"] if source_doc.get("chunks") is not None + else data.get("chunks", "6") + ), + # Per-source behaviour contract (lenient read). + "retrieval": SourceConfig.parse(source_doc.get("config")).retrieval, + } + for source_doc in source_docs + ] data["sources"] = sources_list data["default_model_id"] = data.get("default_model_id", "") return data diff --git a/tests/api/user/test_resource_sponsors.py b/tests/api/user/test_resource_sponsors.py index 4762e852..6181c6c0 100644 --- a/tests/api/user/test_resource_sponsors.py +++ b/tests/api/user/test_resource_sponsors.py @@ -39,6 +39,7 @@ from docsgpt.storage.db.repositories.workflows import WorkflowsRepository OWNER, EDITOR, VIEWER, OTHER = "sp-owner", "sp-editor", "sp-viewer", "sp-other" _DB_MODULES = ( + "docsgpt.agents.headless_runner", "docsgpt.api.user.agents.routes", "docsgpt.api.user.workflows.routes", "docsgpt.api.user.base", @@ -376,6 +377,40 @@ class TestRunTime: data = processor._get_data_from_api_key(agent["key"]) assert [s["id"] for s in data["sources"]] == [source, shared] + def test_headless_run_retrieves_from_every_agent_source(self, app, pg_conn, monkeypatch): + """A scheduled or webhook run searches the same sources a chat does.""" + from unittest.mock import MagicMock + + from docsgpt.agents import headless_runner as hr + + agent, team_id, _, _, source = self._sponsored_agent(app, pg_conn) + shared = str(SourcesRepository(pg_conn).create("shared-src", user_id=OTHER)["id"]) + TeamMembersRepository(pg_conn).add_member(team_id, OWNER) + TeamResourceGrantsRepository(pg_conn).grant(team_id, "source", shared, OTHER, OTHER) + foreign = str(SourcesRepository(pg_conn).create("foreign-src", user_id="sp-stranger")["id"]) + AgentsRepository(pg_conn).update_by_id(str(agent["id"]), {"extra_source_ids": [shared, foreign]}) + + searched = [] + retriever = MagicMock(search=MagicMock(return_value=[])) + + def _create_retriever(cls, *_args, **kwargs): + searched.append(kwargs["source"].get("active_docs")) + return retriever + + run_agent = MagicMock(gen=MagicMock(return_value=iter([{"answer": "ok"}]))) + run_agent.llm.token_usage = {} + monkeypatch.setattr(hr.RetrieverCreator, "create_retriever", classmethod(_create_retriever)) + monkeypatch.setattr(hr, "ToolExecutor", lambda *a, **kw: MagicMock(headless_denials=[])) + monkeypatch.setattr(hr.AgentCreator, "create_agent", classmethod(lambda cls, *a, **kw: run_agent)) + monkeypatch.setattr(hr.QuotaService, "check", lambda *a, **kw: None) + with _patch_db(pg_conn), \ + patch("docsgpt.core.model_utils.validate_model_id", return_value=True), \ + patch("docsgpt.core.model_utils.get_provider_from_model_id", return_value="openai"), \ + patch("docsgpt.core.model_utils.get_api_key_for_provider", return_value="k"), \ + patch("docsgpt.utils.calculate_doc_token_budget", return_value=1000): + hr.run_agent_headless(_row(pg_conn, str(agent["id"])), "q") + assert searched == [[source, shared]] + def test_search_service_authorizes_sponsored_source(self, app, pg_conn): from docsgpt.services.search_service import _authorized_source_ids diff --git a/tests/tracing/test_entry_points.py b/tests/tracing/test_entry_points.py index 218432d2..ffeb8570 100644 --- a/tests/tracing/test_entry_points.py +++ b/tests/tracing/test_entry_points.py @@ -82,7 +82,8 @@ class TestHeadless: assert trace.request_id == "run-1" assert trace.user_id == "owner-1" assert trace.status == "ok" - assert [s.kind for s in trace.spans] == ["retrieval", "agent"] + # The per-source dispatcher's span wraps its retriever's own. + assert [s.kind for s in trace.spans] == ["retrieval", "retrieval", "agent"] def test_trace_can_belong_to_the_scheduling_user(self, monkeypatch, flushed): """A schedule on a shared agent is the scheduler's run, not the owner's.""" diff --git a/tests/worker/test_agent_workers.py b/tests/worker/test_agent_workers.py index 033f97d0..24c27129 100644 --- a/tests/worker/test_agent_workers.py +++ b/tests/worker/test_agent_workers.py @@ -375,4 +375,4 @@ class TestRunAgentHeadlessFromWebhook: outcome = headless_runner.run_agent_headless(agent_config, "hello") assert outcome["answer"] == "done" - assert captured_source.get("active_docs") == source_id + assert captured_source.get("active_docs") == [source_id]