feat: add citation-aware RAG

chatlas/__init__.py

@@ -10,6 +10,8 @@
 from ._content import (
     ContentToolRequest,
     ContentToolResult,
+    SearchResult,
+    ToolSearchResults,
 )
 from ._content_document import content_document_file, content_document_url
 from ._content_image import content_image_file, content_image_plot, content_image_url
@@ -93,11 +95,13 @@
     "interpolate",
     "interpolate_file",
     "Provider",
+    "SearchResult",
     "StreamController",
     "token_usage",
     "Tool",
     "ToolBuiltIn",
     "ToolRejectError",
+    "ToolSearchResults",
     "tool_web_fetch",
     "tool_web_search",
     "Turn",

chatlas/_chat.py

@@ -70,6 +70,7 @@
     start_tool_span,
 )
 from ._provider import ModelInfo, Provider, StandardModelParams, SubmitInputArgsT
+from ._rag import SegmentedAnswer, SegmentsDecoder
 from ._stream_controller import StreamController
 from ._tokens import tokens_log
 from ._tools import Tool, ToolBuiltIn, ToolRejectError
@@ -98,6 +99,7 @@
 
     from ._content import ToolAnnotations
     from ._files import FileManager
+    from ._rag import RagManager
 
 
 class TokensDict(TypedDict):
@@ -194,6 +196,7 @@ def __init__(
         self.kwargs_chat: SubmitInputArgsT = kwargs_chat or {}
 
         self._tools: dict[str, Tool | ToolBuiltIn] = {}
+        self._rag: Optional["RagManager"] = None
         self._on_tool_request_callbacks = CallbackManager()
         self._on_tool_result_callbacks = CallbackManager()
         self._current_display: Optional[MarkdownDisplay] = None
@@ -463,6 +466,21 @@ def files(self) -> "FileManager":
 
         return FileManager(self.provider)
 
+    @property
+    def rag(self) -> "RagManager":
+        """
+        Configure retrieval-augmented answers with citations.
+
+        Register a retrieval store (e.g. a raghilda store) once; afterwards
+        `.chat()` and `.stream()` answers cite the store's documents. See the
+        RAG article for details, including per-provider citation fidelity.
+        """
+        from ._rag import RagManager
+
+        if self._rag is None:
+            self._rag = RagManager(self)
+        return self._rag
+
     @property
     def model(self) -> str:
         """
@@ -2866,6 +2884,15 @@ def _submit_turns(
         if any(isinstance(x, Tool) and x._is_async for x in self._tools.values()):
             raise ValueError("Cannot use async tools in a synchronous chat")
 
+        # A user-supplied `data_model` always wins: only fall back to the
+        # hand-rolled RAG tier's segments schema when the caller hasn't asked
+        # for structured output of their own.
+        rag = self._rag
+        rag_decoder: Optional[SegmentsDecoder] = None
+        if data_model is None and rag is not None and rag.uses_segments_schema():
+            data_model = SegmentedAnswer
+            rag_decoder = SegmentsDecoder(rag._chunks)
+
         chat_span = start_chat_span(
             self.provider, [*self._turns, user_turn], _otel_parent
         )
@@ -2924,17 +2951,32 @@ def emit(x: str | Content):
                             break
                         result = self.provider.stream_merge_chunks(result, chunk)
                         for content in self.provider.stream_content(chunk, result):
-                            yield from acc.process_content(
-                                content, display_text(content), content_mode, emit
-                            )
+                            if rag_decoder is not None and isinstance(
+                                content, ContentText
+                            ):
+                                for c in rag_decoder.feed(content.text):
+                                    yield from acc.process_content(
+                                        c, display_text(c), content_mode, emit
+                                    )
+                            else:
+                                yield from acc.process_content(
+                                    content, display_text(content), content_mode, emit
+                                )
 
                     yield from acc.flush_thinking(content_mode, emit)
+                    if rag_decoder is not None:
+                        for c in rag_decoder.finish():
+                            yield from acc.process_content(
+                                c, display_text(c), content_mode, emit
+                            )
 
                     if not controller.cancelled:
                         turn = self.provider.stream_turn(
                             result,
                             has_data_model=data_model is not None,
                         )
+                        if rag is not None and rag_decoder is not None:
+                            turn = rag.transform_turn(turn)
                         if echo == "all":
                             emit_other_contents(turn, emit)
                         turn = finalize_assistant_turn(self.provider, turn)
@@ -2957,6 +2999,8 @@ def emit(x: str | Content):
                 turn = self.provider.value_turn(
                     response, has_data_model=data_model is not None
                 )
+                if rag is not None and rag_decoder is not None:
+                    turn = rag.transform_turn(turn)
                 emit_thinking_contents(turn, emit)
                 emit_web_contents(turn, emit)
 
@@ -3016,6 +3060,15 @@ async def _submit_turns_async(
         *,
         controller: StreamController,
     ) -> AsyncGenerator[str | Content, None]:
+        # A user-supplied `data_model` always wins: only fall back to the
+        # hand-rolled RAG tier's segments schema when the caller hasn't asked
+        # for structured output of their own.
+        rag = self._rag
+        rag_decoder: Optional[SegmentsDecoder] = None
+        if data_model is None and rag is not None and rag.uses_segments_schema():
+            data_model = SegmentedAnswer
+            rag_decoder = SegmentsDecoder(rag._chunks)
+
         chat_span = start_chat_span(
             self.provider, [*self._turns, user_turn], _otel_parent
         )
@@ -3074,19 +3127,36 @@ def emit(x: str | Content):
                             break
                         result = self.provider.stream_merge_chunks(result, chunk)
                         for content in self.provider.stream_content(chunk, result):
-                            for item in acc.process_content(
-                                content, display_text(content), content_mode, emit
+                            if rag_decoder is not None and isinstance(
+                                content, ContentText
                             ):
-                                yield item
+                                for c in rag_decoder.feed(content.text):
+                                    for item in acc.process_content(
+                                        c, display_text(c), content_mode, emit
+                                    ):
+                                        yield item
+                            else:
+                                for item in acc.process_content(
+                                    content, display_text(content), content_mode, emit
+                                ):
+                                    yield item
 
                     for item in acc.flush_thinking(content_mode, emit):
                         yield item
+                    if rag_decoder is not None:
+                        for c in rag_decoder.finish():
+                            for item in acc.process_content(
+                                c, display_text(c), content_mode, emit
+                            ):
+                                yield item
 
                     if not controller.cancelled:
                         turn = self.provider.stream_turn(
                             result,
                             has_data_model=data_model is not None,
                         )
+                        if rag is not None and rag_decoder is not None:
+                            turn = rag.transform_turn(turn)
                         if echo == "all":
                             emit_other_contents(turn, emit)
                         turn = finalize_assistant_turn(self.provider, turn)
@@ -3109,6 +3179,8 @@ def emit(x: str | Content):
                 turn = self.provider.value_turn(
                     response, has_data_model=data_model is not None
                 )
+                if rag is not None and rag_decoder is not None:
+                    turn = rag.transform_turn(turn)
                 emit_thinking_contents(turn, emit)
                 emit_web_contents(turn, emit)
 

chatlas/_content.py

@@ -246,15 +246,13 @@ def _repr_markdown_(self):
         return self.__str__()
 
 
-SourceTypeEnum = Literal["web"]
+SourceTypeEnum = Literal["web", "document"]
 
 
 class Source(BaseModel):
     """Identity of a piece of evidence a citation or search result points to.
 
-    Subclasses set a distinct ``type`` and add their identity fields. Today the
-    only concrete source is :class:`WebSource`; file/document/RAG variants are
-    added when that support lands.
+    Subclasses set a distinct ``type`` and add their identity fields.
     """
 
     type: SourceTypeEnum
@@ -274,6 +272,17 @@ def __str__(self) -> str:
         return self.url or self.title or "[web source]"
 
 
+class DocumentSource(Source):
+    """A document (file, store chunk, or upload) a citation points to."""
+
+    type: SourceTypeEnum = "document"
+    id: Optional[str] = None
+    title: Optional[str] = None
+
+    def __str__(self) -> str:
+        return self.id or self.title or "[document source]"
+
+
 class ContentText(Content):
     """
     Text content for a [](`~chatlas.Turn`)
@@ -724,6 +733,45 @@ def _arguments_str(self) -> str:
         return str(self.arguments)
 
 
+class SearchResult(BaseModel):
+    """One retrieved chunk, normalized for citation plumbing.
+
+    `id` must be unique across the whole conversation (RagManager assigns
+    them); it is the handle citations use to refer back to the chunk.
+    """
+
+    id: str
+    text: str
+    source: Optional[str] = None
+    title: Optional[str] = None
+    extra: dict[str, Any] = Field(default_factory=dict)
+
+
+class ToolSearchResults(BaseModel):
+    """Search results returned from a tool, opted into citability.
+
+    Return this from any tool (`ContentToolResult(value=ToolSearchResults(...))`
+    or directly) to let providers with native search-result citations
+    (Anthropic) cite individual results. Other providers receive the tagged
+    JSON from `to_dict()`.
+    """
+
+    results: list[SearchResult]
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "results": [
+                {
+                    "chunk_id": r.id,
+                    "source": r.source,
+                    "title": r.title,
+                    "text": r.text,
+                }
+                for r in self.results
+            ]
+        }
+
+
 class ContentJson(Content):
     """
     JSON content
@@ -1209,6 +1257,8 @@ def create_source(data: dict[str, Any]) -> Source:
     t = data.get("type")
     if t == "web":
         return WebSource.model_validate(data)
+    if t == "document":
+        return DocumentSource.model_validate(data)
     raise ValueError(f"Unknown source type: {t}")
 
 

chatlas/_provider.py

@@ -506,6 +506,14 @@ def _no_file_support(self) -> NotImplementedError:
             "Supported providers: ChatOpenAI, ChatAnthropic, ChatGoogle."
         )
 
+    def supports_native_search_results(self) -> bool:
+        """Whether tool results can be sent as natively-citable search results."""
+        return False
+
+    def supports_tools_with_data_model(self) -> bool:
+        """Whether tools and a `data_model` can coexist in one request."""
+        return True
+
 
 ProviderClassT = TypeVar("ProviderClassT", bound=type[Provider[Any, Any, Any, Any]])