diff --git a/docs/content/docs/documentation/API.mdx b/docs/content/docs/documentation/API.mdx index e800943bb..9d93a58f8 100644 --- a/docs/content/docs/documentation/API.mdx +++ b/docs/content/docs/documentation/API.mdx @@ -716,6 +716,7 @@ OpenAI-compatible text completion endpoint. | `spoken_style_answer` | `bool` | `false` | Generates a succinct spoken-style conversational answer based on the retrieved documents. | | `use_map_reduce` | `bool` | `false` | Uses a map-reduce strategy to aggregate information from multiple documents. See [map-reduce configuration](/openrag/documentation/env_vars/#map--reduce-configuration). | | `llm_override` | `object` | `null` | Overrides only the downstream LLM model name while still using OpenRAG's configured LLM endpoint and credentials. Accepts: `model` (string). Endpoint URL and API key are server-side configuration and cannot be changed by a client request. | +| `attachments` | `list[{"id": string}]` | `null` | Scopes RAG retrieval to a specific list of file IDs within the target partition, instead of searching the whole partition. Unknown or unindexed IDs are silently dropped; duplicates are deduplicated. The response's `extra.attachments` reports which IDs were actually used. | Examples: @@ -802,6 +803,30 @@ curl -X 'POST' 'http://localhost:8080/v1/chat/completions' \ }' ``` +```bash title="Scoping retrieval to specific attachments" +curl -X 'POST' 'http://localhost:8080/v1/chat/completions' \ + -H 'accept: application/json' \ + -H 'Authorization: Bearer YOUR_AUTH_TOKEN' \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "openrag-{partition_name}", + "messages": [ + { + "role": "user", + "content": "summarize these files" + } + ], + "stream": false, + "metadata": { + "attachments": [ + {"id": "file-001"}, + {"id": "file-002"}, + {"id": "file-003"} + ] + } +}' +``` + ### 🔧 Tools diff --git a/openrag/core/ports/workspace_repo.py b/openrag/core/ports/workspace_repo.py index 7d3361029..748cf9a1e 100644 --- a/openrag/core/ports/workspace_repo.py +++ b/openrag/core/ports/workspace_repo.py @@ -53,6 +53,14 @@ async def get_existing_file_ids(self, partition: str, file_ids: list[str]) -> se """Return the subset of ``file_ids`` that actually exist in ``partition``.""" ... + @abstractmethod + async def get_existing_file_ids_any_partition(self, file_ids: list[str]) -> set[str]: + """Return the subset of ``file_ids`` that exist in *any* partition. + + Unscoped by design — only for the ``SUPER_ADMIN_MODE`` ``"all"`` wildcard. + """ + ... + @abstractmethod async def remove_file_from_all_workspaces(self, file_id: str, partition: str) -> None: """Detach ``file_id`` from every workspace in ``partition``.""" diff --git a/openrag/core/utils/source_filtering.py b/openrag/core/utils/source_filtering.py index 7a5a693a5..9471ebc04 100644 --- a/openrag/core/utils/source_filtering.py +++ b/openrag/core/utils/source_filtering.py @@ -115,6 +115,7 @@ async def stream_with_source_filtering( citation_protocol_active: bool = True, all_sources: list | None = None, include_all_retrieved: bool = False, + extra_fields: dict | None = None, ): """Process an LLM SSE stream and, when active, strip source tags. @@ -292,6 +293,8 @@ async def stream_with_source_filtering( } if include_all_retrieved: extra_payload["all_retrieved_sources"] = all_sources if all_sources is not None else sources + if extra_fields: + extra_payload.update(extra_fields) if not saw_done: extra_payload["truncated"] = True logger.warning( diff --git a/openrag/services/orchestrators/query_service.py b/openrag/services/orchestrators/query_service.py index 882c8eeff..6292c4f4c 100644 --- a/openrag/services/orchestrators/query_service.py +++ b/openrag/services/orchestrators/query_service.py @@ -40,7 +40,7 @@ from collections.abc import AsyncIterator, Callable from datetime import datetime from enum import Enum -from typing import TYPE_CHECKING, Any +from typing import TYPE_CHECKING, Any, NamedTuple from core.models.preset import resolve_partition_chat_llm from core.models.query import Query, SearchQueries @@ -72,6 +72,21 @@ PrepareSources = Callable[[list, list], list] + +class _PrepareChatResult(NamedTuple): + """A named tuple stays positionally unpackable, so existing + ``a, b, c, ... = await self._prepare_chat(...)`` call sites still work. + """ + + payload: dict + docs: list + web_results: list + retrieved_docs: list + retrieved_web_results: list + citation_protocol_active: bool + indexed_attachment_ids: list[str] + + _MAP_SYSTEM_PROMPT = """You are an AI assistant specialized in extracting and synthesizing relevant information from text. Your task: @@ -445,10 +460,12 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L spoken_style = metadata.get("spoken_style_answer", False) use_websearch = metadata.get("websearch", False) workspace = metadata.get("workspace") + attachment_ids = _extract_attachment_ids(metadata) top_k = self._mr_max if use_map_reduce else None filter_params = None + indexed_attachment_ids: list[str] = [] if workspace and partition: scope = await self._workspace.resolve_scope(workspace, partition) if scope is None: @@ -459,8 +476,15 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L # partition the caller also has access to (#706). partition = [scope.partition] filter_params = {"file_id": scope.file_ids} - - force_retrieval = use_websearch or use_map_reduce + elif attachment_ids and partition: + # No ownership check needed: file_id is ANDed with the server-fixed + # partition (or, for the "all" wildcard, SUPER_ADMIN_MODE-only). + indexed_attachment_ids = await self._existing_file_ids(attachment_ids, partition) + filter_params = {"file_id": indexed_attachment_ids} + + # An attachment must never be dropped just because the classifier + # judged the turn conversational. + force_retrieval = use_websearch or use_map_reduce or bool(indexed_attachment_ids) if not queries.query_list: if not queries.requires_retrieval and not force_retrieval: # Resolved per request from the library (named -> default -> @@ -477,7 +501,7 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L context="", current_date=datetime.now().strftime("%A, %B %d, %Y, %H:%M:%S"), ) - return payload, [], [], [], [], True + return _PrepareChatResult(payload, [], [], [], [], True, indexed_attachment_ids) queries = SearchQueries(query_list=[Query(query=messages[-1]["content"])]) web_results: list = [] @@ -493,7 +517,7 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L chunks = [] if not chunks and not web_results and partition is None: - return payload, [], [], [], [], False + return _PrepareChatResult(payload, [], [], [], [], False, indexed_attachment_ids) docs = [c.to_langchain() for c in chunks] @@ -555,7 +579,24 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L }, ) payload["messages"] = new_messages - return payload, docs, web_results, retrieved_docs, retrieved_web_results, True + return _PrepareChatResult( + payload, docs, web_results, retrieved_docs, retrieved_web_results, True, indexed_attachment_ids + ) + + async def _existing_file_ids(self, file_ids: list[str], partitions: list[str]) -> list[str]: + """Order-preserving, deduplicated subset of ``file_ids`` indexed in ``partitions``. + + ``"all"`` (``SUPER_ADMIN_MODE`` wildcard) takes an unscoped lookup instead + of a per-partition one. + """ + if "all" in partitions: + if len(partitions) > 1: + raise ValueError("`partitions` cannot mix the wildcard with explicit values.") + found = set(await self._workspace.get_existing_file_ids_any_partition(file_ids)) + else: + results = await asyncio.gather(*(self._workspace.get_existing_file_ids(p, file_ids) for p in partitions)) + found = {fid for r in results for fid in r} + return [fid for fid in dict.fromkeys(file_ids) if fid in found] async def _gather_rag_and_web(self, queries, partition, top_k, filter_params): # Fuse the doc branch through retrieve_multi so a partition's rrf_k drives @@ -674,15 +715,16 @@ async def chat( citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): docs, web_results, retrieved_docs, retrieved_web_results = [], [], [], [] + attachments: list[str] = [] else: - ( - payload, - docs, - web_results, - retrieved_docs, - retrieved_web_results, - citation_protocol_active, - ) = await self._prepare_chat(partitions, payload, llm) + result = await self._prepare_chat(partitions, payload, llm) + payload = result.payload + docs = result.docs + web_results = result.web_results + retrieved_docs = result.retrieved_docs + retrieved_web_results = result.retrieved_web_results + citation_protocol_active = result.citation_protocol_active + attachments = result.indexed_attachment_ids sources = prepare_sources(docs, web_results) # `all_retrieved_sources` is debug/eval telemetry, not needed by most # callers — skip building it (and calling prepare_sources on the full, @@ -702,9 +744,11 @@ async def chat( else: clean, citations = content, None chunk["choices"][0]["message"]["content"] = clean - chunk["extra"] = json.dumps( - _build_extra_payload(sources, citations, all_sources, include_all_retrieved=include_all_retrieved) - ) + extra = _build_extra_payload(sources, citations, all_sources, include_all_retrieved=include_all_retrieved) + if metadata.get("attachments"): + # Indicate which attachments were actually searched to generate the answer. + extra["attachments"] = attachments + chunk["extra"] = json.dumps(extra) return chunk async def chat_stream( @@ -722,19 +766,22 @@ async def chat_stream( citation_protocol_active = False if partitions is None and not metadata.get("websearch", False): docs, web_results, retrieved_docs, retrieved_web_results = [], [], [], [] + attachments: list[str] = [] else: - ( - payload, - docs, - web_results, - retrieved_docs, - retrieved_web_results, - citation_protocol_active, - ) = await self._prepare_chat(partitions, payload, llm) + result = await self._prepare_chat(partitions, payload, llm) + payload = result.payload + docs = result.docs + web_results = result.web_results + retrieved_docs = result.retrieved_docs + retrieved_web_results = result.retrieved_web_results + citation_protocol_active = result.citation_protocol_active + attachments = result.indexed_attachment_ids sources = prepare_sources(docs, web_results) all_sources = prepare_sources(retrieved_docs, retrieved_web_results) if include_all_retrieved else None structured_output = _is_structured_output(payload) + extra_fields = {"attachments": attachments} if metadata.get("attachments") else None + payload["messages"] = self._sanitize_messages(payload["messages"]) llm_stream = llm.stream_chat(payload["messages"], **_sampling(payload)) async for sse_line in stream_with_source_filtering( @@ -744,6 +791,7 @@ async def chat_stream( all_sources=all_sources, include_all_retrieved=include_all_retrieved, citation_protocol_active=citation_protocol_active and not structured_output, + extra_fields=extra_fields, ): yield sse_line @@ -800,6 +848,14 @@ def _summary_doc(chunk, summary: str): return chunk.__class__(page_content=summary, metadata=chunk.metadata) +def _extract_attachment_ids(metadata: dict) -> list[str]: + """file_ids from ``metadata.attachments = [{"id": ...}, ...]``; malformed payloads dropped.""" + raw = metadata.get("attachments") + if not isinstance(raw, list): + return [] + return [a["id"] for a in raw if isinstance(a, dict) and isinstance(a.get("id"), str) and a["id"]] + + def _dedupe_web(web_lists: list[list]) -> list: seen: set[str] = set() out: list = [] diff --git a/openrag/services/orchestrators/workspace_service.py b/openrag/services/orchestrators/workspace_service.py index 78431d92f..8eda48165 100644 --- a/openrag/services/orchestrators/workspace_service.py +++ b/openrag/services/orchestrators/workspace_service.py @@ -82,6 +82,9 @@ async def create_workspace( async def get_existing_file_ids(self, partition: str, file_ids: list[str]) -> list[str]: return list(await self._workspace_repo.get_existing_file_ids(partition, file_ids)) + async def get_existing_file_ids_any_partition(self, file_ids: list[str]) -> list[str]: + return list(await self._workspace_repo.get_existing_file_ids_any_partition(file_ids)) + async def add_files(self, workspace_id: str, file_ids: list[str]) -> list[str]: """Associate files; returns any file_ids that were not found.""" return await self._workspace_repo.add_files_to_workspace(workspace_id, file_ids) diff --git a/openrag/services/persistence/workspace_repo.py b/openrag/services/persistence/workspace_repo.py index 28fefb742..ea8e01af2 100644 --- a/openrag/services/persistence/workspace_repo.py +++ b/openrag/services/persistence/workspace_repo.py @@ -250,6 +250,22 @@ async def get_existing_file_ids( ) return {r["file_id"] for r in rows} + async def get_existing_file_ids_any_partition(self, file_ids: list[str]) -> set[str]: + """Return the subset of ``file_ids`` that exist in *any* partition. + + Unscoped by design — only for the ``SUPER_ADMIN_MODE`` ``"all"`` wildcard. + """ + if not file_ids: + return set() + rows = await self.pool.fetch( + """ + SELECT DISTINCT file_id FROM files + WHERE file_id = ANY($1::text[]) + """, + file_ids, + ) + return {r["file_id"] for r in rows} + async def remove_file_from_all_workspaces( self, file_id: str, diff --git a/tests/integration/repos/test_workspace_repo.py b/tests/integration/repos/test_workspace_repo.py index d52507365..cf7296938 100644 --- a/tests/integration/repos/test_workspace_repo.py +++ b/tests/integration/repos/test_workspace_repo.py @@ -99,6 +99,46 @@ async def test_get_file_workspaces_is_partition_scoped( assert in_a == ["ws-a"] +class TestExistingFileIds: + """The two catalog-existence lookups behind attachment scoping. + + Exercised against real SQL rather than a fake, so the ``ANY($1::text[])`` + predicate and the presence/absence of the partition clause are actually + executed — the scoping assertions below both fail if either clause is + changed. + """ + + async def test_partition_scoped_ignores_other_partitions(self, postgres_store: PostgresStore): + await _seed_partition_and_files(postgres_store, partition="a", file_ids=("f1", "f2")) + await _seed_partition_and_files(postgres_store, partition="b", file_ids=("f3",)) + found = await postgres_store.workspace_repo.get_existing_file_ids("a", ["f1", "f3", "ghost"]) + # f3 exists, but in partition "b" — a partition-scoped lookup must not see it. + assert found == {"f1"} + + async def test_any_partition_spans_partitions(self, postgres_store: PostgresStore): + await _seed_partition_and_files(postgres_store, partition="a", file_ids=("f1",)) + await _seed_partition_and_files(postgres_store, partition="b", file_ids=("f3",)) + found = await postgres_store.workspace_repo.get_existing_file_ids_any_partition( + ["f1", "f3", "ghost"], + ) + assert found == {"f1", "f3"} + + async def test_any_partition_reports_a_shared_file_id_once(self, postgres_store: PostgresStore): + # The same file_id in two partitions is two ``files`` rows. Callers treat + # the result as a membership set, so it must come back as a single entry. + # (Both the query's DISTINCT and the set-building enforce this; the test + # pins the contract, not either mechanism.) + await _seed_partition_and_files(postgres_store, partition="a", file_ids=("shared",)) + await _seed_partition_and_files(postgres_store, partition="b", file_ids=("shared",)) + found = await postgres_store.workspace_repo.get_existing_file_ids_any_partition(["shared"]) + assert found == {"shared"} + + async def test_empty_input_short_circuits(self, postgres_store: PostgresStore): + repo = postgres_store.workspace_repo + assert await repo.get_existing_file_ids("a", []) == set() + assert await repo.get_existing_file_ids_any_partition([]) == set() + + class TestDeleteWorkspace: async def test_returns_orphan_file_ids(self, postgres_store: PostgresStore): await _seed_partition_and_files(postgres_store) diff --git a/tests/unit/services/orchestrators/test_query_service.py b/tests/unit/services/orchestrators/test_query_service.py index b2f0d3a77..25c74860b 100644 --- a/tests/unit/services/orchestrators/test_query_service.py +++ b/tests/unit/services/orchestrators/test_query_service.py @@ -110,8 +110,11 @@ async def search(self, query): class FakeWorkspace: - def __init__(self, scope=None): + def __init__(self, scope=None, existing=None): self._scope = scope + # None => every requested file_id is indexed in any partition (default); + # dict {partition: set(file_ids)} => partition-scoped existence. + self._existing = existing async def get_workspace(self, wid): return None @@ -119,6 +122,18 @@ async def get_workspace(self, wid): async def resolve_scope(self, workspace_id, allowed_partitions): return self._scope + async def get_existing_file_ids(self, partition, file_ids): + if self._existing is None: + return list(file_ids) + allowed = self._existing.get(partition, set()) + return [fid for fid in file_ids if fid in allowed] + + async def get_existing_file_ids_any_partition(self, file_ids): + if self._existing is None: + return list(file_ids) + allowed = {fid for partition_ids in self._existing.values() for fid in partition_ids} + return [fid for fid in file_ids if fid in allowed] + def _config(mode="SimpleRag"): return SimpleNamespace( @@ -863,9 +878,10 @@ async def test_websearch_with_partition_fuses_docs_via_retrieve_multi(): retrieval = FakeRetrieval() web_result = SimpleNamespace(url="https://ex.com", title="T", content="web body", snippet="") svc = _svc(retrieval=retrieval, web=FakeWeb(results=[web_result])) - _payload, _docs, web, _retrieved_docs, _retrieved_web, _citation_protocol_active = await svc._prepare_chat( + result = await svc._prepare_chat( ["p"], {"messages": [{"role": "user", "content": "q"}], "metadata": {"websearch": True}} ) + web = result.web_results assert len(retrieval.retrieve_multi_calls) == 1 # doc branch fused via the rrf_k-aware retrieve_multi assert retrieval.retrieve_per_query_calls == [] # legacy per-query + fuse()@60 path NOT used assert web and web[0].url == "https://ex.com" # websearch branch actually taken @@ -889,9 +905,8 @@ async def resolve_prompt(self, prompt_type, names=None): marker = MarkerPromptService() svc._prompt_service = marker - payload, _docs, _web, _retrieved_docs, _retrieved_web, _citation_protocol_active = await svc._prepare_chat( - ["p"], {"messages": [{"role": "user", "content": "q"}], "metadata": {}} - ) + result = await svc._prepare_chat(["p"], {"messages": [{"role": "user", "content": "q"}], "metadata": {}}) + payload = result.payload assert payload["messages"][0]["role"] == "system" assert payload["messages"][0]["content"].startswith("MARKER-SYS::") @@ -1089,6 +1104,240 @@ async def test_chat_without_workspace_unaffected(): assert call["filter_params"] is None +# --------------------------------------------------------------------------- # +# attachment scoping (metadata.attachments = [{"id": ...}]) +# --------------------------------------------------------------------------- # + + +@pytest.mark.asyncio +async def test_chat_with_valid_attachments_scopes_search_to_file_ids(): + retrieval = FakeRetrieval() + svc = _svc(retrieval=retrieval, llm=FakeLLM(chat_responses=["answer [Sources: none]"])) + await svc.chat( + partitions=["p1"], + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"attachments": [{"id": "fa"}, {"id": "fb"}]}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + call = retrieval.retrieve_multi_calls[0] + assert call["partitions"] == ["p1"] # attachments do NOT narrow the partition + assert call["filter_params"] == {"file_id": ["fa", "fb"]} + + +@pytest.mark.asyncio +async def test_chat_attachments_force_retrieval_even_when_classifier_skips(): + # Regression: an attached file must not be silently dropped just because + # the query-classifier judges the turn conversational. + query_json = json.dumps({"requires_retrieval": False, "query_list": []}) + llm = FakeLLM(chat_responses=[query_json, "answer [Sources: none]"]) + retrieval = FakeRetrieval() + svc = _svc(mode="ChatBotRag", llm=llm, retrieval=retrieval) + + await svc.chat( + partitions=["p1"], + payload={ + "messages": [{"role": "user", "content": "thanks!"}], + "metadata": {"attachments": [{"id": "fa"}]}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + + assert len(retrieval.retrieve_multi_calls) == 1 + assert retrieval.retrieve_multi_calls[0]["filter_params"] == {"file_id": ["fa"]} + + +@pytest.mark.asyncio +async def test_chat_attachments_malformed_ignored(): + # Items without a usable id (or not dicts) are skipped; an all-malformed + # attachments blob degrades to a normal unscoped chat, never raises. + retrieval = FakeRetrieval() + svc = _svc(retrieval=retrieval, llm=FakeLLM(chat_responses=["answer [Sources: none]"])) + await svc.chat( + partitions=["p1"], + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"attachments": [{"nope": "x"}, "raw-string", 123, {"id": ""}]}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + call = retrieval.retrieve_multi_calls[0] + assert call["filter_params"] is None + + +@pytest.mark.asyncio +@pytest.mark.parametrize("bad", [123, True, "abc", {"id": "x"}]) +async def test_chat_attachments_scalar_payload_is_unscoped(bad): + # A non-list attachments payload must degrade to a normal unscoped chat, + # never raise (a scalar like 123 is not iterable). + retrieval = FakeRetrieval() + svc = _svc(retrieval=retrieval, llm=FakeLLM(chat_responses=["answer [Sources: none]"])) + await svc.chat( + partitions=["p1"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {"attachments": bad}}, + prepare_sources=lambda d, w: [], + model_name="m", + ) + assert retrieval.retrieve_multi_calls[0]["filter_params"] is None + + +@pytest.mark.asyncio +async def test_chat_attachments_non_string_ids_ignored(): + retrieval = FakeRetrieval() + svc = _svc(retrieval=retrieval, llm=FakeLLM(chat_responses=["answer [Sources: none]"])) + await svc.chat( + partitions=["p1"], + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"attachments": [{"id": 123}, {"id": None}, {"id": ["x"]}]}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + assert retrieval.retrieve_multi_calls[0]["filter_params"] is None + + +@pytest.mark.asyncio +async def test_chat_empty_attachments_unaffected(): + retrieval = FakeRetrieval() + svc = _svc(retrieval=retrieval, llm=FakeLLM(chat_responses=["answer [Sources: none]"])) + await svc.chat( + partitions=["p1"], + payload={"messages": [{"role": "user", "content": "q"}], "metadata": {"attachments": []}}, + prepare_sources=lambda d, w: [], + model_name="m", + ) + call = retrieval.retrieve_multi_calls[0] + assert call["filter_params"] is None + + +@pytest.mark.asyncio +async def test_chat_workspace_and_attachments_both_present_workspace_wins(): + # workspace is checked first (elif) — when both are present the workspace + # scope wins and the attachments are ignored. + scope = WorkspaceScope(workspace_id="w1", partition="p1", file_ids=["wsa", "wsb"]) + retrieval = FakeRetrieval() + svc = _svc( + retrieval=retrieval, llm=FakeLLM(chat_responses=["answer [Sources: none]"]), workspace=FakeWorkspace(scope) + ) + await svc.chat( + partitions=["p1"], + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"workspace": "w1", "attachments": [{"id": "att"}]}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + call = retrieval.retrieve_multi_calls[0] + assert call["filter_params"] == {"file_id": ["wsa", "wsb"]} + + +@pytest.mark.asyncio +async def test_chat_attachments_drops_unindexed_and_reports_in_extra(): + retrieval = FakeRetrieval() + svc = _svc( + retrieval=retrieval, + llm=FakeLLM(chat_responses=["answer [Sources: none]"]), + workspace=FakeWorkspace(existing={"p1": {"fa"}}), # only fa is indexed; fb is not + ) + chunk = await svc.chat( + partitions=["p1"], + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"attachments": [{"id": "fa"}, {"id": "fb"}]}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + # Only the indexed id drives the filter... + assert retrieval.retrieve_multi_calls[0]["filter_params"] == {"file_id": ["fa"]} + # ...and the same validated list is reported back to the client in extra. + assert json.loads(chunk["extra"])["attachments"] == ["fa"] + + +@pytest.mark.asyncio +async def test_chat_attachments_duplicate_ids_deduped(): + retrieval = FakeRetrieval() + svc = _svc( + retrieval=retrieval, + llm=FakeLLM(chat_responses=["answer [Sources: none]"]), + workspace=FakeWorkspace(existing={"p1": {"fa"}}), + ) + chunk = await svc.chat( + partitions=["p1"], + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"attachments": [{"id": "fa"}, {"id": "fa"}]}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + assert retrieval.retrieve_multi_calls[0]["filter_params"] == {"file_id": ["fa"]} + assert json.loads(chunk["extra"])["attachments"] == ["fa"] + + +@pytest.mark.asyncio +async def test_chat_stream_reports_indexed_attachments_in_extra(): + svc = _svc(workspace=FakeWorkspace(existing={"p1": {"fa"}})) # only fa indexed in p1 + out = "".join( + [ + line + async for line in svc.chat_stream( + partitions=["p1"], + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"attachments": [{"id": "fa"}, {"id": "fb"}]}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + ] + ) + assert "attachments" in out and "fa" in out + assert "fb" not in out # unindexed id dropped, never reported + + +@pytest.mark.asyncio +async def test_chat_attachments_all_partition_looks_up_any_partition(): + retrieval = FakeRetrieval() + svc = _svc( + retrieval=retrieval, + llm=FakeLLM(chat_responses=["answer [Sources: none]"]), + workspace=FakeWorkspace(existing={"p1": {"fa"}}), + ) + chunk = await svc.chat( + partitions=["all"], + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"attachments": [{"id": "fa"}, {"id": "fb"}]}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + assert retrieval.retrieve_multi_calls[0]["filter_params"] == {"file_id": ["fa"]} + assert json.loads(chunk["extra"])["attachments"] == ["fa"] + + +@pytest.mark.asyncio +async def test_chat_attachments_all_mixed_with_explicit_partition_raises(): + svc = _svc(workspace=FakeWorkspace(existing={"p1": {"fa"}})) + with pytest.raises(ValueError): + await svc.chat( + partitions=["all", "p1"], + payload={ + "messages": [{"role": "user", "content": "q"}], + "metadata": {"attachments": [{"id": "fa"}]}, + }, + prepare_sources=lambda d, w: [], + model_name="m", + ) + + @pytest.mark.asyncio async def test_complete_direct_mode_preserves_literal_source_marker(): answer = "text body\n[Sources: none]" @@ -1417,9 +1666,8 @@ async def resolve_prompt(self, prompt_type, names=None): "p": SimpleNamespace(generation_prompt_names={"sys_prompt": "chatty"}, chat_history_depth=4) } - out, docs, web, _retrieved_docs, _retrieved_web, _ = await svc._prepare_chat( - ["p"], {"messages": [{"role": "user", "content": "hello!"}], "metadata": {}} - ) + result = await svc._prepare_chat(["p"], {"messages": [{"role": "user", "content": "hello!"}], "metadata": {}}) + out, docs, web = result.payload, result.docs, result.web_results # Resolved from the library, honouring the partition's selection, and no # retrieval happened.