allowing scope retrieval via a list of file_ids

openrag/core/utils/source_filtering.py

@@ -83,6 +83,7 @@ async def stream_with_source_filtering(
     sources: list,
     model_name: str,
     buffer_size: int | None = None,
+    extra_fields: dict | None = None,
 ):
     """Process an LLM SSE stream, stripping line-terminal source tags.
 
@@ -221,6 +222,8 @@ async def stream_with_source_filtering(
 
     filtered = filter_sources_by_citations(sources, citations)
     extra_payload = {"sources": filtered}
+    if extra_fields:
+        extra_payload.update(extra_fields)
     if not saw_done:
         extra_payload["truncated"] = True
         logger.warning(

openrag/services/orchestrators/query_service.py

@@ -397,10 +397,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:
@@ -411,6 +413,11 @@ 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}
+        elif attachment_ids and partition:
+            # No security check on the ids: file_id is ANDed with the
+            # server-fixed partition, so a foreign id can never match.
+            indexed_attachment_ids = await self._existing_file_ids(attachment_ids, partition)
+            filter_params = {"file_id": indexed_attachment_ids}
 
         web_results: list = []
         if partition is not None and use_websearch:
@@ -425,7 +432,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, [], []
+            return payload, [], [], []
 
         docs = [c.to_langchain() for c in chunks]
         if use_map_reduce and docs:
@@ -470,7 +477,22 @@ async def _prepare_chat(self, partition: list[str] | None, payload: dict, llm: L
             },
         )
         payload["messages"] = new_messages
-        return payload, docs, web_results
+        return payload, docs, web_results, indexed_attachment_ids
+
+    async def _existing_file_ids(self, file_ids: list[str], partitions: list[str]) -> list[str]:
+        """Order-preserving 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
@@ -569,9 +591,9 @@ async def chat(
         metadata = payload.get("metadata") or {}
         llm = self._resolve_llm(partitions)
         if partitions is None and not metadata.get("websearch", False):
-            docs, web_results = [], []
+            docs, web_results, attachments = [], [], []
         else:
-            payload, docs, web_results = await self._prepare_chat(partitions, payload, llm)
+            payload, docs, web_results, attachments = await self._prepare_chat(partitions, payload, llm)
         sources = prepare_sources(docs, web_results)
 
         payload["messages"] = self._sanitize_messages(payload["messages"])
@@ -580,7 +602,11 @@ async def chat(
         content = chunk.get("choices", [{}])[0].get("message", {}).get("content", "") or ""
         clean, citations = extract_and_strip_sources_block(content)
         chunk["choices"][0]["message"]["content"] = clean
-        chunk["extra"] = json.dumps({"sources": filter_sources_by_citations(sources, citations)})
+        extra = {"sources": filter_sources_by_citations(sources, citations)}
+        if metadata.get("attachments"):
+            # Indicate which attachments were actually leveraged to generate the answer.
+            extra["attachments"] = attachments
+        chunk["extra"] = json.dumps(extra)
         return chunk
 
     async def chat_stream(
@@ -595,14 +621,16 @@ async def chat_stream(
         metadata = payload.get("metadata") or {}
         llm = self._resolve_llm(partitions)
         if partitions is None and not metadata.get("websearch", False):
-            docs, web_results = [], []
+            docs, web_results, attachments = [], [], []
         else:
-            payload, docs, web_results = await self._prepare_chat(partitions, payload, llm)
+            payload, docs, web_results, attachments = await self._prepare_chat(partitions, payload, llm)
         sources = prepare_sources(docs, web_results)
 
+        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(llm_stream, sources, model_name):
+        async for sse_line in stream_with_source_filtering(llm_stream, sources, model_name, extra_fields=extra_fields):
             yield sse_line
 
     async def complete(
@@ -645,6 +673,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 = []

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)

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,