Add citations to run result (#65)

Author: vhaldemar

examples/async/assistants/assistant_with_search_index.py

@@ -44,6 +44,19 @@ async def main() -> None:
     print('Question:', search_query)
     print('Answer:', result.text)
 
+    # You could access .citations attribute for debug purposes
+    for citation in result.citations:
+        for source in citation.sources:
+            # In future there will be more source types
+            if source.type != 'filechunk':
+                continue
+            print('Example source:', source)
+            # One source will be enough for example, it takes too much screen space to print
+            break
+        else:
+            continue
+        break
+
     search_query = "Cколько пошлина в Анталье"
     await thread.write(search_query)
 

examples/sync/assistants/assistant_with_search_index.py

@@ -42,6 +42,19 @@ def main() -> None:
     print('Question:', search_query)
     print('Answer:', result.text)
 
+    # You could access .citations attribute for debug purposes
+    for citation in result.citations:
+        for source in citation.sources:
+            # In future there will be more source types
+            if source.type != 'filechunk':
+                continue
+            print('Example source:', source)
+            # One source will be enough for example, it takes too much screen space to print
+            break
+        else:
+            continue
+        break
+
     search_query = "Cколько пошлина в Анталье"
     thread.write(search_query)
 

src/yandex_cloud_ml_sdk/_messages/base.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+import dataclasses
+from typing import Any
+
+from yandex_cloud_ml_sdk._types.result import BaseResult
+
+
+@dataclasses.dataclass(frozen=True)
+class BaseMessage(BaseResult):
+    parts: tuple[Any, ...]
+
+    @property
+    def text(self):
+        return '\n'.join(
+            part for part in self.parts
+            if isinstance(part, str)
+        )

src/yandex_cloud_ml_sdk/_messages/citations.py

@@ -0,0 +1,98 @@
+# pylint: disable=no-name-in-module
+
+from __future__ import annotations
+
+import abc
+import dataclasses
+from typing import TYPE_CHECKING
+
+from yandex.cloud.ai.assistants.v1.threads.message_pb2 import Citation as ProtoCitation
+from yandex.cloud.ai.assistants.v1.threads.message_pb2 import Source as ProtoSource
+
+from yandex_cloud_ml_sdk._files.file import BaseFile
+from yandex_cloud_ml_sdk._search_indexes.search_index import BaseSearchIndex
+from yandex_cloud_ml_sdk._types.result import BaseResult
+
+from .base import BaseMessage
+
+if TYPE_CHECKING:
+    from yandex_cloud_ml_sdk._sdk import BaseSDK
+
+
+@dataclasses.dataclass(frozen=True)
+class Citation(BaseResult):
+    sources: tuple[Source, ...]
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoCitation, sdk: BaseSDK) -> Citation:  # type: ignore[override]
+        return cls(
+            sources=tuple(
+                Source._from_proto(proto=source, sdk=sdk)
+                for source in proto.sources
+            )
+        )
+
+class Source(BaseResult):
+    @property
+    @abc.abstractmethod
+    def type(self) -> str:
+        pass
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoSource, sdk: BaseSDK) -> Source:  # type: ignore[override]
+        if proto.HasField('chunk'):
+            return FileChunk._from_proto(proto=proto, sdk=sdk)
+
+        return UnknownSource._from_proto(proto=proto, sdk=sdk)
+
+
+@dataclasses.dataclass(frozen=True)
+class FileChunk(Source, BaseMessage):
+    search_index: BaseSearchIndex
+    file: BaseFile | None
+
+    @property
+    def type(self) -> str:
+        return 'filechunk'
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoSource, sdk: BaseSDK) -> FileChunk | UnknownSource:  # type: ignore[override]
+        # pylint: disable=protected-access
+        chunk = proto.chunk
+        assert chunk
+
+        raw_parts = (part.text.content for part in chunk.content.content)
+        parts = tuple(part for part in raw_parts if part)
+
+        search_index = sdk.search_indexes._impl._from_proto(proto=chunk.search_index, sdk=sdk)
+        file: BaseFile | None = None
+
+        # NB: at the moment backend always returns non-empty source_file field
+        # but in case it deleted, source_file will content empty File structure
+        if (
+            chunk.HasField('source_file') and
+            chunk.source_file and
+            chunk.source_file.id
+        ):
+            file = sdk.files._file_impl._from_proto(proto=chunk.source_file, sdk=sdk)
+
+        return cls(
+            search_index=search_index,
+            file=file,
+            parts=parts,
+        )
+
+
+@dataclasses.dataclass(frozen=True)
+class UnknownSource(Source):
+    text: str
+
+    @property
+    def type(self) -> str:
+        return 'unknown'
+
+    @classmethod
+    def _from_proto(cls, proto: ProtoSource, sdk: BaseSDK) -> UnknownSource:  # type: ignore[override]
+        return cls(
+            text="Source's protobuf have unknown fields; try to update yandex-cloud-ml-sdk"
+        )

src/yandex_cloud_ml_sdk/_messages/message.py

@@ -10,6 +10,9 @@
 
 from yandex_cloud_ml_sdk._types.resource import BaseResource
 
+from .base import BaseMessage
+from .citations import Citation
+
 if TYPE_CHECKING:
     from yandex_cloud_ml_sdk._sdk import BaseSDK
 
@@ -21,24 +24,13 @@ class Author:
 
 
 @dataclasses.dataclass(frozen=True)
-class BaseMessage(BaseResource):
-    parts: tuple[Any, ...]
-
-    @property
-    def text(self):
-        return '\n'.join(
-            part for part in self.parts
-            if isinstance(part, str)
-        )
-
-
-@dataclasses.dataclass(frozen=True)
-class Message(BaseMessage):
+class Message(BaseMessage, BaseResource):
     thread_id: str
     created_by: str
     created_at: datetime
     labels: dict[str, str] | None
     author: Author
+    citations: tuple[Citation, ...]
 
     @classmethod
     def _kwargs_from_message(cls, proto: ProtoMessage, sdk: BaseSDK) -> dict[str, Any]:  # type: ignore[override]
@@ -58,12 +50,17 @@ def _kwargs_from_message(cls, proto: ProtoMessage, sdk: BaseSDK) -> dict[str, An
             role=proto.author.role,
             id=proto.author.id
         )
+        raw_citations = proto.citations or ()
+        kwargs['citations'] = tuple(
+            Citation._from_proto(proto=citation, sdk=sdk)
+            for citation in raw_citations
+        )
 
         return kwargs
 
 
 @dataclasses.dataclass(frozen=True)
-class PartialMessage(BaseMessage):
+class PartialMessage(BaseMessage, BaseResource):
     @classmethod
     def _kwargs_from_message(cls, proto: MessageContent, sdk: BaseSDK) -> dict[str, Any]:  # type: ignore[override]
         kwargs = super()._kwargs_from_message(proto, sdk=sdk)

src/yandex_cloud_ml_sdk/_runs/result.py

@@ -8,6 +8,7 @@
 from yandex.cloud.ai.assistants.v1.runs.run_pb2 import Run as ProtoRun
 from yandex.cloud.ai.assistants.v1.runs.run_service_pb2 import StreamEvent as ProtoStreamEvent
 
+from yandex_cloud_ml_sdk._messages.citations import Citation
 from yandex_cloud_ml_sdk._messages.message import BaseMessage, Message, PartialMessage
 from yandex_cloud_ml_sdk._models.completions.result import Usage
 from yandex_cloud_ml_sdk._types.result import BaseResult, ProtoMessage
@@ -60,7 +61,7 @@ def text(self) -> str:
         return self.message.text
 
     @property
-    def parts(self) -> tuple[Any]:
+    def parts(self) -> tuple[Any, ...]:
         return self.message.parts
 
 
@@ -101,6 +102,10 @@ def _from_proto(cls, *, proto: ProtoMessage, sdk: BaseSDK) -> RunResult:
             usage=usage,
         )
 
+    @property
+    def citations(self) -> tuple[Citation, ...]:
+        return self.message.citations
+
 
 @dataclasses.dataclass(frozen=True)
 class RunStreamEvent(BaseRunResult[StreamEvent, BaseMessage]):