nucliadb 6.9.1.post5192__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/search/search/chat/ask.py

@@ -20,7 +20,8 @@
 import dataclasses
 import functools
 import json
-from typing import AsyncGenerator, Optional, Union, cast
+from collections.abc import AsyncGenerator
+from typing import cast
 
 from nuclia_models.common.consumption import Consumption
 from nuclia_models.predict.generative_responses import (
@@ -34,6 +35,7 @@ from nuclia_models.predict.generative_responses import (
     TextGenerativeResponse,
 )
 from pydantic_core import ValidationError
+from typing_extensions import assert_never
 
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
 from nucliadb.common.exceptions import InvalidQueryError
@@ -49,11 +51,13 @@ from nucliadb.search.search.chat.exceptions import (
     AnswerJsonSchemaTooLong,
     NoRetrievalResultsError,
 )
+from nucliadb.search.search.chat.old_prompt import PromptContextBuilder as OldPromptContextBuilder
 from nucliadb.search.search.chat.prompt import PromptContextBuilder
 from nucliadb.search.search.chat.query import (
     NOT_ENOUGH_CONTEXT_ANSWER,
     ChatAuditor,
     add_resource_filter,
+    get_answer_stream,
     get_find_results,
     get_relations_results,
     maybe_audit_chat,
@@ -69,11 +73,15 @@ from nucliadb.search.search.metrics import AskMetrics, Metrics
 from nucliadb.search.search.query_parser.fetcher import Fetcher
 from nucliadb.search.search.query_parser.parsers.ask import fetcher_for_ask, parse_ask
 from nucliadb.search.search.rank_fusion import WeightedCombSum
-from nucliadb.search.search.rerankers import (
-    get_reranker,
+from nucliadb_models.retrieval import (
+    GraphScore,
+    KeywordScore,
+    RerankerScore,
+    RrfScore,
+    SemanticScore,
 )
-from nucliadb.search.utilities import get_predict
 from nucliadb_models.search import (
+    SCORE_TYPE,
     AnswerAskResponseItem,
     AskRequest,
     AskResponseItem,
@@ -118,7 +126,9 @@ from nucliadb_models.search import (
     parse_rephrase_prompt,
 )
 from nucliadb_telemetry import errors
+from nucliadb_utils import const
 from nucliadb_utils.exceptions import LimitsExceededError
+from nucliadb_utils.utilities import has_feature
 
 
 @dataclasses.dataclass
@@ -132,7 +142,7 @@ class RetrievalResults:
     main_query: KnowledgeboxFindResults
     fetcher: Fetcher
     main_query_weight: float
-    prequeries: Optional[list[PreQueryResult]] = None
+    prequeries: list[PreQueryResult] | None = None
     best_matches: list[RetrievalMatch] = dataclasses.field(default_factory=list)
 
 
@@ -143,15 +153,15 @@ class AskResult:
         kbid: str,
         ask_request: AskRequest,
         main_results: KnowledgeboxFindResults,
-        prequeries_results: Optional[list[PreQueryResult]],
-        nuclia_learning_id: Optional[str],
-        predict_answer_stream: Optional[AsyncGenerator[GenerativeChunk, None]],
+        prequeries_results: list[PreQueryResult] | None,
+        nuclia_learning_id: str | None,
+        predict_answer_stream: AsyncGenerator[GenerativeChunk, None] | None,
         prompt_context: PromptContext,
         prompt_context_order: PromptContextOrder,
         auditor: ChatAuditor,
         metrics: AskMetrics,
         best_matches: list[RetrievalMatch],
-        debug_chat_model: Optional[ChatModel],
+        debug_chat_model: ChatModel | None,
         augmented_context: AugmentedContext,
     ):
         # Initial attributes
@@ -171,14 +181,14 @@ class AskResult:
 
         # Computed from the predict chat answer stream
         self._answer_text = ""
-        self._reasoning_text: Optional[str] = None
-        self._object: Optional[JSONGenerativeResponse] = None
-        self._status: Optional[StatusGenerativeResponse] = None
-        self._citations: Optional[CitationsGenerativeResponse] = None
-        self._footnote_citations: Optional[FootnoteCitationsGenerativeResponse] = None
-        self._metadata: Optional[MetaGenerativeResponse] = None
-        self._relations: Optional[Relations] = None
-        self._consumption: Optional[Consumption] = None
+        self._reasoning_text: str | None = None
+        self._object: JSONGenerativeResponse | None = None
+        self._status: StatusGenerativeResponse | None = None
+        self._citations: CitationsGenerativeResponse | None = None
+        self._footnote_citations: FootnoteCitationsGenerativeResponse | None = None
+        self._metadata: MetaGenerativeResponse | None = None
+        self._relations: Relations | None = None
+        self._consumption: Consumption | None = None
 
     @property
     def status_code(self) -> AnswerStatusCode:
@@ -187,7 +197,7 @@ class AskResult:
         return AnswerStatusCode(self._status.code)
 
     @property
-    def status_error_details(self) -> Optional[str]:
+    def status_error_details(self) -> str | None:
         if self._status is None:  # pragma: no cover
             return None
         return self._status.details
@@ -240,9 +250,7 @@ class AskResult:
                         self.metrics.record_first_reasoning_chunk_yielded()
                         first_reasoning_chunk_yielded = True
                 else:
-                    # This is a trick so mypy generates an error if this branch can be reached,
-                    # that is, if we are missing some ifs
-                    _a: int = "a"
+                    assert_never(answer_chunk)
 
         if self._object is not None:
             yield JSONAskResponseItem(object=self._object.object)
@@ -396,7 +404,7 @@ class AskResult:
         if self._object is not None:
             answer_json = self._object.object
 
-        prequeries_results: Optional[dict[str, KnowledgeboxFindResults]] = None
+        prequeries_results: dict[str, KnowledgeboxFindResults] | None = None
         if self.prequeries_results:
             prequeries_results = {}
             for index, (prequery, result) in enumerate(self.prequeries_results):
@@ -452,7 +460,7 @@ class AskResult:
 
     async def _stream_predict_answer_text(
         self,
-    ) -> AsyncGenerator[Union[TextGenerativeResponse, ReasoningGenerativeResponse], None]:
+    ) -> AsyncGenerator[TextGenerativeResponse | ReasoningGenerativeResponse, None]:
         """
         Reads the stream of the generative model, yielding the answer text but also parsing
         other items like status codes, citations and miscellaneous metadata.
@@ -496,8 +504,8 @@ class AskResult:
 class NotEnoughContextAskResult(AskResult):
     def __init__(
         self,
-        main_results: Optional[KnowledgeboxFindResults] = None,
-        prequeries_results: Optional[list[PreQueryResult]] = None,
+        main_results: KnowledgeboxFindResults | None = None,
+        prequeries_results: list[PreQueryResult] | None = None,
     ):
         self.main_results = main_results or KnowledgeboxFindResults(resources={}, min_score=None)
         self.prequeries_results = prequeries_results or []
@@ -547,8 +555,8 @@ async def ask(
     user_id: str,
     client_type: NucliaDBClientType,
     origin: str,
-    resource: Optional[str] = None,
-    extra_predict_headers: Optional[dict[str, str]] = None,
+    resource: str | None = None,
+    extra_predict_headers: dict[str, str] | None = None,
 ) -> AskResult:
     metrics = AskMetrics()
     chat_history = ask_request.chat_history or []
@@ -627,19 +635,36 @@ async def ask(
 
     # Now we build the prompt context
     with metrics.time("context_building"):
-        prompt_context_builder = PromptContextBuilder(
-            kbid=kbid,
-            ordered_paragraphs=[match.paragraph for match in retrieval_results.best_matches],
-            resource=resource,
-            user_context=user_context,
-            user_image_context=ask_request.extra_context_images,
-            strategies=ask_request.rag_strategies,
-            image_strategies=ask_request.rag_images_strategies,
-            max_context_characters=tokens_to_chars(generation.max_context_tokens),
-            visual_llm=generation.use_visual_llm,
-            query_image=ask_request.query_image,
-            metrics=metrics.child_span("context_building"),
-        )
+        prompt_context_builder: PromptContextBuilder | OldPromptContextBuilder
+        if has_feature(const.Features.ASK_DECOUPLED, context={"kbid": kbid}):
+            prompt_context_builder = PromptContextBuilder(
+                kbid=kbid,
+                ordered_paragraphs=[match.paragraph for match in retrieval_results.best_matches],
+                resource=resource,
+                user_context=user_context,
+                user_image_context=ask_request.extra_context_images,
+                strategies=ask_request.rag_strategies,
+                image_strategies=ask_request.rag_images_strategies,
+                max_context_characters=tokens_to_chars(generation.max_context_tokens),
+                visual_llm=generation.use_visual_llm,
+                query_image=ask_request.query_image,
+                metrics=metrics.child_span("context_building"),
+            )
+        else:
+            prompt_context_builder = OldPromptContextBuilder(
+                kbid=kbid,
+                ordered_paragraphs=[match.paragraph for match in retrieval_results.best_matches],
+                resource=resource,
+                user_context=user_context,
+                user_image_context=ask_request.extra_context_images,
+                strategies=ask_request.rag_strategies,
+                image_strategies=ask_request.rag_images_strategies,
+                max_context_characters=tokens_to_chars(generation.max_context_tokens),
+                visual_llm=generation.use_visual_llm,
+                query_image=ask_request.query_image,
+                metrics=metrics.child_span("context_building"),
+            )
+
         (
             prompt_context,
             prompt_context_order,
@@ -675,14 +700,11 @@ async def ask(
     predict_answer_stream = None
     if ask_request.generate_answer:
         with metrics.time("stream_start"):
-            predict = get_predict()
             (
                 nuclia_learning_id,
                 nuclia_learning_model,
                 predict_answer_stream,
-            ) = await predict.chat_query_ndjson(
-                kbid=kbid, item=chat_model, extra_headers=extra_predict_headers
-            )
+            ) = await get_answer_stream(kbid=kbid, item=chat_model, extra_headers=extra_predict_headers)
 
     auditor = ChatAuditor(
         kbid=kbid,
@@ -757,7 +779,7 @@ def handled_ask_exceptions(func):
     return wrapper
 
 
-def parse_prequeries(ask_request: AskRequest) -> Optional[PreQueriesStrategy]:
+def parse_prequeries(ask_request: AskRequest) -> PreQueriesStrategy | None:
     query_ids = []
     for rag_strategy in ask_request.rag_strategies:
         if rag_strategy.name == RagStrategyName.PREQUERIES:
@@ -776,7 +798,7 @@ def parse_prequeries(ask_request: AskRequest) -> Optional[PreQueriesStrategy]:
     return None
 
 
-def parse_graph_strategy(ask_request: AskRequest) -> Optional[GraphStrategy]:
+def parse_graph_strategy(ask_request: AskRequest) -> GraphStrategy | None:
     for rag_strategy in ask_request.rag_strategies:
         if rag_strategy.name == RagStrategyName.GRAPH:
             return cast(GraphStrategy, rag_strategy)
@@ -791,7 +813,7 @@ async def retrieval_step(
     user_id: str,
     origin: str,
     metrics: Metrics,
-    resource: Optional[str] = None,
+    resource: str | None = None,
 ) -> RetrievalResults:
     """
     This function encapsulates all the logic related to retrieval in the ask endpoint.
@@ -830,7 +852,7 @@ async def retrieval_in_kb(
 ) -> RetrievalResults:
     prequeries = parse_prequeries(ask_request)
     graph_strategy = parse_graph_strategy(ask_request)
-    main_results, prequeries_results, parsed_query = await get_find_results(
+    main_results, prequeries_results, fetcher, reranker = await get_find_results(
         kbid=kbid,
         query=main_query,
         item=ask_request,
@@ -842,10 +864,6 @@ async def retrieval_in_kb(
     )
 
     if graph_strategy is not None:
-        assert parsed_query.retrieval.reranker is not None, (
-            "find parser must provide a reranking algorithm"
-        )
-        reranker = get_reranker(parsed_query.retrieval.reranker)
         graph_results, graph_request = await get_graph_results(
             kbid=kbid,
             query=main_query,
@@ -878,7 +896,7 @@ async def retrieval_in_kb(
     return RetrievalResults(
         main_query=main_results,
         prequeries=prequeries_results,
-        fetcher=parsed_query.fetcher,
+        fetcher=fetcher,
         main_query_weight=main_query_weight,
         best_matches=best_matches,
     )
@@ -918,7 +936,7 @@ async def retrieval_in_resource(
                 )
             add_resource_filter(prequery.request, [resource])
 
-    main_results, prequeries_results, parsed_query = await get_find_results(
+    main_results, prequeries_results, fetcher, _ = await get_find_results(
         kbid=kbid,
         query=main_query,
         item=ask_request,
@@ -941,7 +959,7 @@ async def retrieval_in_resource(
     return RetrievalResults(
         main_query=main_results,
         prequeries=prequeries_results,
-        fetcher=parsed_query.fetcher,
+        fetcher=fetcher,
         main_query_weight=main_query_weight,
         best_matches=best_matches,
     )
@@ -953,7 +971,7 @@ class _FindParagraph(ScoredTextBlock):
 
 def compute_best_matches(
     main_results: KnowledgeboxFindResults,
-    prequeries_results: Optional[list[PreQueryResult]] = None,
+    prequeries_results: list[PreQueryResult] | None = None,
     main_query_weight: float = 1.0,
 ) -> list[RetrievalMatch]:
     """
@@ -968,15 +986,27 @@ def compute_best_matches(
     `main_query_weight` is the weight given to the paragraphs matching the main query when calculating the final score.
     """
 
+    score_type_map = {
+        SCORE_TYPE.VECTOR: SemanticScore,
+        SCORE_TYPE.BM25: KeywordScore,
+        SCORE_TYPE.BOTH: RrfScore,  # /find only exposes RRF as rank fusion algorithm
+        SCORE_TYPE.RERANKER: RerankerScore,
+        SCORE_TYPE.RELATION_RELEVANCE: GraphScore,
+    }
+
     def extract_paragraphs(results: KnowledgeboxFindResults) -> list[_FindParagraph]:
         paragraphs = []
         for resource in results.resources.values():
             for field in resource.fields.values():
                 for paragraph in field.paragraphs.values():
+                    # TODO(decoupled-ask): we don't know the score history, as
+                    # we are using find results. Once we move boolean queries
+                    # inside the new retrieval flow we'll move this and have the
+                    # proper information to do this rank fusion
                     paragraphs.append(
                         _FindParagraph(
                             paragraph_id=ParagraphId.from_string(paragraph.id),
-                            score=paragraph.score,
+                            scores=[score_type_map[paragraph.score_type](score=paragraph.score)],
                             score_type=paragraph.score_type,
                             original=paragraph,
                         )
@@ -1012,7 +1042,7 @@ def compute_best_matches(
 
 def calculate_prequeries_for_json_schema(
     ask_request: AskRequest,
-) -> Optional[PreQueriesStrategy]:
+) -> PreQueriesStrategy | None:
     """
     This function generates a PreQueriesStrategy with a query for each property in the JSON schema
     found in ask_request.answer_json_schema.
@@ -1077,7 +1107,6 @@ def calculate_prequeries_for_json_schema(
             rephrase=ask_request.rephrase,
             rephrase_prompt=parse_rephrase_prompt(ask_request),
             security=ask_request.security,
-            autofilter=False,
         )
         prequery = PreQuery(
             request=req,

nucliadb/search/search/chat/exceptions.py

@@ -19,17 +19,15 @@
 #
 
 
-from typing import Optional
-
 from nucliadb_models.search import KnowledgeboxFindResults, PreQueryResult
 
 
 class NoRetrievalResultsError(Exception):
     def __init__(
         self,
-        main: Optional[KnowledgeboxFindResults] = None,
-        prequeries: Optional[list[PreQueryResult]] = None,
-        prefilters: Optional[list[PreQueryResult]] = None,
+        main: KnowledgeboxFindResults | None = None,
+        prequeries: list[PreQueryResult] | None = None,
+        prefilters: list[PreQueryResult] | None = None,
     ):
         self.main_query = main
         self.prequeries = prequeries

nucliadb/search/search/chat/fetcher.py

@@ -0,0 +1,201 @@
+# Copyright (C) 2021 Bosutech XXI S.L.
+#
+# nucliadb is offered under the AGPL v3.0 and as commercial software.
+# For commercial licensing, contact us at info@nuclia.com.
+#
+# AGPL:
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU Affero General Public License as
+# published by the Free Software Foundation, either version 3 of the
+# License, or (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU Affero General Public License for more details.
+#
+# You should have received a copy of the GNU Affero General Public License
+# along with this program. If not, see <http://www.gnu.org/licenses/>.
+#
+
+
+from google.protobuf.json_format import ParseDict
+
+from nucliadb.common.exceptions import InvalidQueryError
+from nucliadb.search import logger
+from nucliadb.search.predict import SendToPredictError, convert_relations
+from nucliadb.search.predict_models import QueryModel
+from nucliadb.search.search.chat import rpc
+from nucliadb.search.search.query_parser.fetcher import Fetcher
+from nucliadb.search.utilities import get_predict
+from nucliadb_models.internal.predict import QueryInfo
+from nucliadb_models.search import Image, MaxTokens
+from nucliadb_protos import knowledgebox_pb2, utils_pb2
+
+
+class RAOFetcher(Fetcher):
+    def __init__(
+        self,
+        kbid: str,
+        *,
+        query: str,
+        user_vector: list[float] | None,
+        vectorset: str | None,
+        rephrase: bool,
+        rephrase_prompt: str | None,
+        generative_model: str | None,
+        query_image: Image | None,
+    ):
+        super().__init__(
+            kbid,
+            query=query,
+            user_vector=user_vector,
+            vectorset=vectorset,
+            rephrase=rephrase,
+            rephrase_prompt=rephrase_prompt,
+            generative_model=generative_model,
+            query_image=query_image,
+        )
+
+        self._query_info: QueryInfo | None = None
+        self._vectorset: str | None = None
+
+    async def query_information(self) -> QueryInfo:
+        if self._query_info is None:
+            self._query_info = await query_information(
+                kbid=self.kbid,
+                query=self.query,
+                semantic_model=self.user_vectorset,
+                generative_model=self.generative_model,
+                rephrase=self.rephrase,
+                rephrase_prompt=self.rephrase_prompt,
+                query_image=self.query_image,
+            )
+        return self._query_info
+
+    # Retrieval
+
+    async def get_rephrased_query(self) -> str | None:
+        query_info = await self.query_information()
+        return query_info.rephrased_query
+
+    async def get_detected_entities(self) -> list[utils_pb2.RelationNode]:
+        query_info = await self.query_information()
+        if query_info.entities is not None:
+            detected_entities = convert_relations(query_info.entities.model_dump())
+        else:
+            detected_entities = []
+        return detected_entities
+
+    async def get_semantic_min_score(self) -> float | None:
+        query_info = await self.query_information()
+        vectorset = await self.get_vectorset()
+        return query_info.semantic_thresholds.get(vectorset, None)
+
+    async def get_vectorset(self) -> str:
+        if self._vectorset is None:
+            if self.user_vectorset is not None:
+                self._vectorset = self.user_vectorset
+            else:
+                # when it's not provided, we get the default from Predict API
+                query_info = await self.query_information()
+                if query_info.sentence is None or len(query_info.sentence.vectors) == 0:
+                    logger.error(
+                        "Asking for a vectorset but /query didn't return one", extra={"kbid": self.kbid}
+                    )
+                    raise SendToPredictError("Predict API didn't return a sentence vectorset")
+                # vectors field is enforced by the data model to have at least one key
+                for vectorset in query_info.sentence.vectors.keys():
+                    self._vectorset = vectorset
+                    break
+        assert self._vectorset is not None
+        return self._vectorset
+
+    async def get_query_vector(self) -> list[float]:
+        if self.user_vector is not None:
+            return self.user_vector
+
+        query_info = await self.query_information()
+        if query_info.sentence is None:
+            logger.error(
+                "Asking for a semantic query vector but /query didn't return a sentence",
+                extra={"kbid": self.kbid},
+            )
+            raise SendToPredictError("Predict API didn't return a sentence for semantic search")
+
+        vectorset = await self.get_vectorset()
+        if vectorset not in query_info.sentence.vectors:
+            logger.error(
+                "Predict is not responding with a valid query nucliadb vectorset",
+                extra={
+                    "kbid": self.kbid,
+                    "vectorset": vectorset,
+                    "predict_vectorsets": ",".join(query_info.sentence.vectors.keys()),
+                },
+            )
+            raise SendToPredictError("Predict API didn't return the requested vectorset")
+
+        query_vector = query_info.sentence.vectors[vectorset]
+        return query_vector
+
+    async def get_classification_labels(self) -> knowledgebox_pb2.Labels:
+        labelsets = await rpc.labelsets(self.kbid)
+
+        # TODO(decoupled-ask): remove this conversion and refactor code to use API models instead of protobuf
+        kb_labels = knowledgebox_pb2.Labels()
+        for labelset, labels in labelsets.labelsets.items():
+            ParseDict(labels.model_dump(), kb_labels.labelset[labelset])
+
+        return kb_labels
+
+    # Generative
+
+    async def get_visual_llm_enabled(self) -> bool:
+        query_info = await self.query_information()
+        if query_info is None:
+            raise SendToPredictError("Error while using predict's query endpoint")
+
+        return query_info.visual_llm
+
+    async def get_max_context_tokens(self, max_tokens: MaxTokens | None) -> int:
+        query_info = await self.query_information()
+        if query_info is None:
+            raise SendToPredictError("Error while using predict's query endpoint")
+
+        model_max = query_info.max_context
+        if max_tokens is not None and max_tokens.context is not None:
+            if max_tokens.context > model_max:
+                raise InvalidQueryError(
+                    "max_tokens.context",
+                    f"Max context tokens is higher than the model's limit of {model_max}",
+                )
+            return max_tokens.context
+        return model_max
+
+    def get_max_answer_tokens(self, max_tokens: MaxTokens | None) -> int | None:
+        if max_tokens is not None and max_tokens.answer is not None:
+            return max_tokens.answer
+        return None
+
+
+async def query_information(
+    kbid: str,
+    query: str,
+    semantic_model: str | None,
+    generative_model: str | None = None,
+    rephrase: bool = False,
+    rephrase_prompt: str | None = None,
+    query_image: Image | None = None,
+) -> QueryInfo:
+    # NOTE: When moving /ask to RAO, this will need to change to whatever client/utility is used
+    # to call NUA predict (internally or externally in the case of onprem).
+    predict = get_predict()
+    item = QueryModel(
+        text=query,
+        semantic_models=[semantic_model] if semantic_model else None,
+        generative_model=generative_model,
+        rephrase=rephrase,
+        rephrase_prompt=rephrase_prompt,
+        query_image=query_image,
+    )
+    return await predict.query(kbid, item)

nucliadb/search/search/chat/images.py

@@ -19,7 +19,6 @@
 
 import base64
 from io import BytesIO
-from typing import Optional
 
 from nucliadb.common.ids import ParagraphId
 from nucliadb.ingest.fields.file import File
@@ -29,7 +28,8 @@ from nucliadb_utils.storages.storage import Storage
 from nucliadb_utils.utilities import get_storage
 
 
-async def get_page_image(kbid: str, paragraph_id: ParagraphId, page_number: int) -> Optional[Image]:
+# DEPRECATED(decoupled-ask): remove once old_prompt.py is removed
+async def get_page_image(kbid: str, paragraph_id: ParagraphId, page_number: int) -> Image | None:
     storage = await get_storage(service_name=SERVICE_NAME)
     sf = storage.file_extracted(
         kbid=kbid,
@@ -48,7 +48,8 @@ async def get_page_image(kbid: str, paragraph_id: ParagraphId, page_number: int)
     return image
 
 
-async def get_paragraph_image(kbid: str, paragraph_id: ParagraphId, reference: str) -> Optional[Image]:
+# DEPRECATED(decoupled-ask): remove once old_prompt.py is removed
+async def get_paragraph_image(kbid: str, paragraph_id: ParagraphId, reference: str) -> Image | None:
     storage = await get_storage(service_name=SERVICE_NAME)
     sf = storage.file_extracted(
         kbid=kbid,
@@ -67,7 +68,8 @@ async def get_paragraph_image(kbid: str, paragraph_id: ParagraphId, reference: s
     return image
 
 
-async def get_file_thumbnail_image(file: File) -> Optional[Image]:
+# DEPRECATED(decoupled-ask): remove once old_prompt.py is removed
+async def get_file_thumbnail_image(file: File) -> Image | None:
     fed = await file.get_file_extracted_data()
     if fed is None or not fed.HasField("file_thumbnail"):
         return None