mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-13 10:15:21 +00:00
arc53-machine
335241e2e6
Sponsoring a tool, source or prompt the agent's owner can't use now takes owning it or having edit access to it; use access alone no longer extends it to the agent's audience. A save that would make the caller a new sponsor is refused with 409 sponsor_confirmation_required (the resources and the agent's audience) until it is retried with confirm_sponsor listing them. When a sponsor loses access, the resource stops instead of passing to whoever saves next; another editor takes it over only by confirming. Workflows follow the same rules. sponsor_details now reports each sponsorship's state, the reason it stopped, and whether the reader can take it over.
289 lines
9.7 KiB
Python
289 lines
9.7 KiB
Python
"""Shared retrieval service used by the HTTP search route and the MCP tool.
|
|
|
|
Flask-free. Raises domain exceptions (``InvalidAPIKey``, ``SearchFailed``)
|
|
that callers translate into their own wire protocol (HTTP status codes,
|
|
MCP error responses, etc.).
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import logging
|
|
from typing import Any, Dict, List, Optional
|
|
|
|
from docsgpt import tracing
|
|
from docsgpt.core.settings import settings
|
|
from docsgpt.retriever.fanout import fetch_per_source
|
|
from docsgpt.storage.db.repositories.agents import AgentsRepository
|
|
from docsgpt.storage.db.session import db_readonly
|
|
from docsgpt.tracing.retrieval import describe_documents, start_retrieval_span
|
|
from docsgpt.vectorstore.vector_creator import VectorCreator
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class InvalidAPIKey(Exception):
|
|
"""The supplied ``api_key`` does not resolve to an agent."""
|
|
|
|
|
|
class SearchFailed(Exception):
|
|
"""Unexpected error during retrieval (e.g. DB outage). Caller maps to 5xx."""
|
|
|
|
|
|
def _collect_source_ids(agent: Dict[str, Any]) -> List[str]:
|
|
"""Extract the ordered list of source UUIDs to search.
|
|
|
|
Prefers ``extra_source_ids`` (PG ARRAY(UUID) of multi-source agents);
|
|
falls back to the legacy single ``source_id`` field.
|
|
"""
|
|
source_ids: List[str] = []
|
|
extra = agent.get("extra_source_ids") or []
|
|
for src in extra:
|
|
if src:
|
|
source_ids.append(str(src))
|
|
if not source_ids:
|
|
single = agent.get("source_id")
|
|
if single:
|
|
source_ids.append(str(single))
|
|
return source_ids
|
|
|
|
|
|
def _authorized_source_ids(conn, agent: Dict[str, Any], source_ids: List[str]) -> List[str]:
|
|
"""Drop source ids the agent may not read.
|
|
|
|
``_collect_source_ids`` trusts whatever the agent row carries, and this
|
|
service searches those ids directly. That made it the second half of a
|
|
real disclosure: ``/api/share`` resolved a client-supplied source with no
|
|
ownership predicate and baked it into the agent it created, after which
|
|
this path returned another tenant's documents. The share route is fixed,
|
|
but a row written before that — or by any future write path with the same
|
|
gap — is still live here, so re-resolve against the owner rather than
|
|
trusting the stored value.
|
|
|
|
Args:
|
|
conn: Open read connection.
|
|
agent: The agent row resolved from the API key.
|
|
source_ids: Ids extracted from that row.
|
|
|
|
A source the owner can't read still searches while the editor who
|
|
attached it (its sponsor) can edit the agent and still owns or edits
|
|
the source.
|
|
|
|
Returns:
|
|
list: The subset the agent's owner (or a live sponsor) may read.
|
|
"""
|
|
owner = agent.get("user_id")
|
|
if not owner:
|
|
logger.warning("Agent %s has no owner; refusing to search its sources.", agent.get("id"))
|
|
return []
|
|
|
|
from docsgpt.api.user.resource_access import ref_principal
|
|
|
|
allowed = []
|
|
for sid in source_ids:
|
|
try:
|
|
permitted = ref_principal(conn, "agent", agent, "source", str(sid)) is not None
|
|
except Exception:
|
|
# Fail closed, matching the answer path.
|
|
logger.warning("Access check failed for source %s; dropping it.", sid)
|
|
continue
|
|
if permitted:
|
|
allowed.append(sid)
|
|
else:
|
|
logger.warning(
|
|
"Agent %s references source %s its owner %s cannot read; dropping it.",
|
|
agent.get("id"), sid, owner,
|
|
)
|
|
return allowed
|
|
|
|
|
|
def _search_one(
|
|
source_id: str,
|
|
docsearch: Any,
|
|
query: str,
|
|
k: int,
|
|
query_vector: Optional[List[float]],
|
|
) -> Optional[List[Any]]:
|
|
"""Search one source, returning its hits or ``None`` when it fails.
|
|
|
|
Builds the vector store when not supplied, so each worker thread owns its
|
|
own store instance (and therefore its own DB connection). Errors are logged
|
|
and reported as ``None`` so one broken index cannot take the rest down.
|
|
"""
|
|
try:
|
|
if docsearch is None:
|
|
docsearch = VectorCreator.create_vectorstore(
|
|
settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
|
|
)
|
|
search_kwargs: Dict[str, Any] = {"k": k}
|
|
if query_vector is not None:
|
|
search_kwargs["query_vector"] = query_vector
|
|
return docsearch.search(query, **search_kwargs)
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error searching vectorstore {source_id}: {e}",
|
|
exc_info=True,
|
|
)
|
|
return None
|
|
|
|
|
|
def _fetch_sources(
|
|
query: str, source_ids: List[str], k: int
|
|
) -> List[Optional[List[Any]]]:
|
|
"""Fetch every source's hits: one query embedding, one bounded fan-out.
|
|
|
|
Shares :func:`~docsgpt.retriever.fanout.fetch_per_source` with
|
|
``ClassicRAG`` so the search route and the answer path order, embed and
|
|
degrade identically.
|
|
"""
|
|
return fetch_per_source(
|
|
source_ids,
|
|
lambda source_id: VectorCreator.create_vectorstore(
|
|
settings.VECTOR_STORE, source_id, settings.EMBEDDINGS_KEY
|
|
),
|
|
lambda source_id, docsearch, query_vector: _search_one(
|
|
source_id, docsearch, query, k, query_vector
|
|
),
|
|
lambda _source_id: query,
|
|
)
|
|
|
|
|
|
def _search_sources(
|
|
query: str, source_ids: List[str], chunks: int
|
|
) -> List[Dict[str, Any]]:
|
|
"""Search across each source's vectorstore and return up to ``chunks`` hits.
|
|
|
|
Per-source errors are logged and skipped so one broken index doesn't
|
|
take down the whole search. Results are de-duplicated by content hash.
|
|
"""
|
|
if chunks <= 0 or not source_ids:
|
|
return []
|
|
|
|
results: List[Dict[str, Any]] = []
|
|
# Blank ids build no store but still count here, so the per-source budget
|
|
# matches what the serial implementation handed each real source.
|
|
chunks_per_source = max(1, chunks // len(source_ids))
|
|
seen_texts: set[int] = set()
|
|
|
|
active_ids = [sid for sid in source_ids if sid and sid.strip()]
|
|
if not active_ids:
|
|
return []
|
|
|
|
# Fetch every source up front, then merge serially in source order so the
|
|
# dedupe / cap semantics are exactly what they were.
|
|
fetched = _fetch_sources(query, active_ids, chunks_per_source * 2)
|
|
|
|
for source_id, docs in zip(active_ids, fetched):
|
|
if docs is None:
|
|
continue
|
|
|
|
try:
|
|
for doc in docs:
|
|
if len(results) >= chunks:
|
|
break
|
|
|
|
if hasattr(doc, "page_content") and hasattr(doc, "metadata"):
|
|
page_content = doc.page_content
|
|
metadata = doc.metadata
|
|
else:
|
|
page_content = doc.get("text", doc.get("page_content", ""))
|
|
metadata = doc.get("metadata", {})
|
|
|
|
text_hash = hash(page_content[:200])
|
|
if text_hash in seen_texts:
|
|
continue
|
|
seen_texts.add(text_hash)
|
|
|
|
title = metadata.get("title", metadata.get("post_title", ""))
|
|
if not isinstance(title, str):
|
|
title = str(title) if title else ""
|
|
|
|
if title:
|
|
title = title.split("/")[-1]
|
|
else:
|
|
title = metadata.get("filename", page_content[:50] + "...")
|
|
|
|
source = metadata.get("source", source_id)
|
|
|
|
results.append(
|
|
{
|
|
"text": page_content,
|
|
"title": title,
|
|
"source": source,
|
|
}
|
|
)
|
|
|
|
if len(results) >= chunks:
|
|
break
|
|
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error searching vectorstore {source_id}: {e}",
|
|
exc_info=True,
|
|
)
|
|
continue
|
|
|
|
return results[:chunks]
|
|
|
|
|
|
def search(
|
|
api_key: str, query: str, chunks: int = 5, *, source: str = "search"
|
|
) -> List[Dict[str, Any]]:
|
|
"""Resolve an agent by API key and search its sources.
|
|
|
|
Every search that reaches the sources is recorded as an execution trace
|
|
owned by the agent's owner, under ``source``.
|
|
|
|
Args:
|
|
api_key: Agent API key (the opaque string stored on
|
|
``agents.key`` in Postgres).
|
|
query: Free-text search query.
|
|
chunks: Max number of hits to return.
|
|
source: Trace source name: ``search`` for ``/api/search``, ``mcp``
|
|
for the MCP ``search_docs`` tool.
|
|
|
|
Returns:
|
|
List of hit dicts with ``text``, ``title``, ``source`` keys.
|
|
Empty list if the agent has no sources configured.
|
|
|
|
Raises:
|
|
InvalidAPIKey: if ``api_key`` does not resolve to an agent.
|
|
SearchFailed: on unexpected DB / infrastructure errors.
|
|
"""
|
|
if chunks <= 0:
|
|
return []
|
|
|
|
try:
|
|
with db_readonly() as conn:
|
|
agent = AgentsRepository(conn).find_by_key(api_key)
|
|
if not agent:
|
|
raise InvalidAPIKey()
|
|
# Authorize inside the same connection the agent was read on.
|
|
source_ids = _authorized_source_ids(
|
|
conn, agent, _collect_source_ids(agent)
|
|
)
|
|
except InvalidAPIKey:
|
|
raise
|
|
except Exception as e:
|
|
raise SearchFailed("agent lookup failed") from e
|
|
|
|
if not source_ids:
|
|
return []
|
|
|
|
trace = tracing.start_trace(
|
|
source=source,
|
|
user_id=agent.get("user_id"),
|
|
agent_id=str(agent.get("id")) if agent.get("id") else None,
|
|
)
|
|
with tracing.activate(trace):
|
|
try:
|
|
with start_retrieval_span(
|
|
f"retrieval {source}",
|
|
sources=source_ids,
|
|
**{"docsgpt.top_k": chunks},
|
|
) as span:
|
|
results = _search_sources(query, source_ids, chunks)
|
|
describe_documents(span, results, query=query)
|
|
return results
|
|
finally:
|
|
tracing.flush(trace)
|