nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__py3-none-any.whl

nucliadb/search/search/chat/ask.py

@@ -17,21 +17,32 @@
 # 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/>.
 #
+import dataclasses
 import functools
-from time import monotonic as time
-from typing import AsyncGenerator, Optional
+import json
+from typing import AsyncGenerator, Optional, cast
 
-from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
-from nucliadb.models.responses import HTTPClientError
-from nucliadb.search import logger, predict
-from nucliadb.search.predict import (
-    AnswerStatusCode,
+from nuclia_models.predict.generative_responses import (
     CitationsGenerativeResponse,
     GenerativeChunk,
+    JSONGenerativeResponse,
     MetaGenerativeResponse,
     StatusGenerativeResponse,
     TextGenerativeResponse,
 )
+from pydantic_core import ValidationError
+
+from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
+from nucliadb.models.responses import HTTPClientError
+from nucliadb.search import logger, predict
+from nucliadb.search.predict import (
+    AnswerStatusCode,
+    RephraseMissingContextError,
+)
+from nucliadb.search.search.chat.exceptions import (
+    AnswerJsonSchemaTooLong,
+    NoRetrievalResultsError,
+)
 from nucliadb.search.search.chat.prompt import PromptContextBuilder
 from nucliadb.search.search.chat.query import (
     NOT_ENOUGH_CONTEXT_ANSWER,
@@ -46,6 +57,7 @@ from nucliadb.search.search.exceptions import (
     IncompleteFindResultsError,
     InvalidQueryError,
 )
+from nucliadb.search.search.metrics import RAGMetrics
 from nucliadb.search.search.query import QueryParser
 from nucliadb.search.utilities import get_predict
 from nucliadb_models.search import (
@@ -53,6 +65,7 @@ from nucliadb_models.search import (
     AskRequest,
     AskResponseItem,
     AskResponseItemType,
+    AskRetrievalMatch,
     AskTimings,
     AskTokens,
     ChatModel,
@@ -60,48 +73,84 @@ from nucliadb_models.search import (
     CitationsAskResponseItem,
     DebugAskResponseItem,
     ErrorAskResponseItem,
+    FindParagraph,
+    FindRequest,
+    JSONAskResponseItem,
     KnowledgeboxFindResults,
     MetadataAskResponseItem,
     MinScore,
     NucliaDBClientType,
+    PrequeriesAskResponseItem,
+    PreQueriesStrategy,
+    PreQuery,
+    PreQueryResult,
     PromptContext,
     PromptContextOrder,
+    RagStrategyName,
     Relations,
     RelationsAskResponseItem,
     RetrievalAskResponseItem,
+    SearchOptions,
     StatusAskResponseItem,
     SyncAskMetadata,
     SyncAskResponse,
     UserPrompt,
+    parse_custom_prompt,
+    parse_rephrase_prompt,
 )
+from nucliadb_telemetry import errors
 from nucliadb_utils.exceptions import LimitsExceededError
 
 
+@dataclasses.dataclass
+class RetrievalMatch:
+    paragraph: FindParagraph
+    weighted_score: float
+
+
+@dataclasses.dataclass
+class RetrievalResults:
+    main_query: KnowledgeboxFindResults
+    query_parser: QueryParser
+    main_query_weight: float
+    prequeries: Optional[list[PreQueryResult]] = None
+    best_matches: list[RetrievalMatch] = dataclasses.field(default_factory=list)
+
+
 class AskResult:
     def __init__(
         self,
         *,
         kbid: str,
         ask_request: AskRequest,
-        find_results: KnowledgeboxFindResults,
+        main_results: KnowledgeboxFindResults,
+        prequeries_results: Optional[list[PreQueryResult]],
         nuclia_learning_id: Optional[str],
         predict_answer_stream: AsyncGenerator[GenerativeChunk, None],
         prompt_context: PromptContext,
         prompt_context_order: PromptContextOrder,
         auditor: ChatAuditor,
+        metrics: RAGMetrics,
+        best_matches: list[RetrievalMatch],
+        debug_chat_model: Optional[ChatModel],
     ):
         # Initial attributes
         self.kbid = kbid
         self.ask_request = ask_request
-        self.find_results = find_results
+        self.main_results = main_results
+        self.prequeries_results = prequeries_results or []
         self.nuclia_learning_id = nuclia_learning_id
         self.predict_answer_stream = predict_answer_stream
         self.prompt_context = prompt_context
+        self.debug_chat_model = debug_chat_model
         self.prompt_context_order = prompt_context_order
-        self.auditor = auditor
+        self.auditor: ChatAuditor = auditor
+        self.metrics: RAGMetrics = metrics
+        self.best_matches: list[RetrievalMatch] = best_matches
 
         # Computed from the predict chat answer stream
         self._answer_text = ""
+        self._object: Optional[JSONGenerativeResponse] = None
         self._status: Optional[StatusGenerativeResponse] = None
         self._citations: Optional[CitationsGenerativeResponse] = None
         self._metadata: Optional[MetaGenerativeResponse] = None
@@ -113,6 +162,12 @@ class AskResult:
             return AnswerStatusCode.SUCCESS
         return AnswerStatusCode(self._status.code)
 
+    @property
+    def status_error_details(self) -> Optional[str]:
+        if self._status is None:  # pragma: no cover
+            return None
+        return self._status.details
+
     @property
     def ask_request_with_relations(self) -> bool:
         return ChatOptions.RELATIONS in self.ask_request.features
@@ -128,34 +183,89 @@ class AskResult:
         except Exception as exc:
             # Handle any unexpected error that might happen
             # during the streaming and halt the stream
-            item = ErrorAskResponseItem(error=str(exc))
-            yield self._ndjson_encode(item)
-
-            staus = AnswerStatusCode.ERROR
-            item = StatusAskResponseItem(code=staus.value, status=staus.prettify())
+            errors.capture_exception(exc)
+            logger.error(
+                f"Unexpected error while generating the answer: {exc}",
+                extra={"kbid": self.kbid},
+            )
+            error_message = "Unexpected error while generating the answer. Please try again later."
+            if self.ask_request_with_debug_flag:
+                error_message += f" Error: {exc}"
+            item = ErrorAskResponseItem(error=error_message)
             yield self._ndjson_encode(item)
             return
 
     def _ndjson_encode(self, item: AskResponseItemType) -> str:
         result_item = AskResponseItem(item=item)
-        return result_item.json(exclude_unset=False, exclude_none=True) + "\n"
+        return result_item.model_dump_json(exclude_none=True, by_alias=True) + "\n"
 
     async def _stream(self) -> AsyncGenerator[AskResponseItemType, None]:
-        # First stream out the find results
-        yield RetrievalAskResponseItem(results=self.find_results)
+        # First, stream out the predict answer
+        first_chunk_yielded = False
+        with self.metrics.time("stream_predict_answer"):
+            async for answer_chunk in self._stream_predict_answer_text():
+                yield AnswerAskResponseItem(text=answer_chunk)
+                if not first_chunk_yielded:
+                    self.metrics.record_first_chunk_yielded()
+                    first_chunk_yielded = True
+
+        if self._object is not None:
+            yield JSONAskResponseItem(object=self._object.object)
+            if not first_chunk_yielded:
+                # When there is a JSON generative response, we consider the first chunk yielded
+                # to be the moment when the JSON object is yielded, not the text
+                self.metrics.record_first_chunk_yielded()
+                first_chunk_yielded = True
+
+        yield RetrievalAskResponseItem(
+            results=self.main_results,
+            best_matches=[
+                AskRetrievalMatch(
+                    id=match.paragraph.id,
+                )
+                for match in self.best_matches
+            ],
+        )
 
-        # Then stream out the predict answer
-        async for answer_chunk in self._stream_predict_answer_text():
-            yield AnswerAskResponseItem(text=answer_chunk)
+        if len(self.prequeries_results) > 0:
+            item = PrequeriesAskResponseItem()
+            for index, (prequery, result) in enumerate(self.prequeries_results):
+                prequery_id = prequery.id or f"prequery_{index}"
+                item.results[prequery_id] = result
+            yield item
+
+        # Then the status
+        if self.status_code == AnswerStatusCode.ERROR:
+            # If predict yielded an error status, we yield it too and halt the stream immediately
+            yield StatusAskResponseItem(
+                code=self.status_code.value,
+                status=self.status_code.prettify(),
+                details=self.status_error_details or "Unknown error",
+            )
+            return
 
-        # Then the status code
         yield StatusAskResponseItem(
-            code=self.status_code.value, status=self.status_code.prettify()
+            code=self.status_code.value,
+            status=self.status_code.prettify(),
         )
 
         # Audit the answer
-        await self.auditor.audit(
-            text_answer=self._answer_text.encode("utf-8"),
+        if self._object is None:
+            audit_answer = self._answer_text.encode("utf-8")
+        else:
+            audit_answer = json.dumps(self._object.object).encode("utf-8")
+
+        try:
+            rephrase_time = self.metrics.elapsed("rephrase")
+        except KeyError:
+            # Not all ask requests have a rephrase step
+            rephrase_time = None
+
+        self.auditor.audit(
+            text_answer=audit_answer,
+            generative_answer_time=self.metrics.elapsed("stream_predict_answer"),
+            generative_answer_first_chunk_time=self.metrics.get_first_chunk_time() or 0,
+            rephrase_time=rephrase_time,
             status_code=self.status_code,
         )
 
@@ -163,25 +273,24 @@ class AskResult:
         if self._citations is not None:
             yield CitationsAskResponseItem(citations=self._citations.citations)
 
-        # Stream out other metadata about the answer if available
+        # Stream out generic metadata about the answer
         if self._metadata is not None:
             yield MetadataAskResponseItem(
                 tokens=AskTokens(
                     input=self._metadata.input_tokens,
                     output=self._metadata.output_tokens,
+                    input_nuclia=self._metadata.input_nuclia_tokens,
+                    output_nuclia=self._metadata.output_nuclia_tokens,
                 ),
                 timings=AskTimings(
-                    generative_first_chunk=self._metadata.timings.get(
-                        "generative_first_chunk"
-                    ),
+                    generative_first_chunk=self._metadata.timings.get("generative_first_chunk"),
                     generative_total=self._metadata.timings.get("generative"),
                 ),
             )
 
         # Stream out the relations results
         should_query_relations = (
-            self.ask_request_with_relations
-            and self.status_code != AnswerStatusCode.NO_CONTEXT
+            self.ask_request_with_relations and self.status_code == AnswerStatusCode.SUCCESS
         )
         if should_query_relations:
             relations = await self.get_relations_results()
@@ -189,11 +298,15 @@ class AskResult:
 
         # Stream out debug information
         if self.ask_request_with_debug_flag:
+            predict_request = None
+            if self.debug_chat_model:
+                predict_request = self.debug_chat_model.model_dump(mode="json")
             yield DebugAskResponseItem(
                 metadata={
                     "prompt_context": sorted_prompt_context_list(
                         self.prompt_context, self.prompt_context_order
-                    )
+                    ),
+                    "predict_request": predict_request,
                 }
             )
 
@@ -208,40 +321,68 @@ class AskResult:
                 tokens=AskTokens(
                     input=self._metadata.input_tokens,
                     output=self._metadata.output_tokens,
+                    input_nuclia=self._metadata.input_nuclia_tokens,
+                    output_nuclia=self._metadata.output_nuclia_tokens,
                 ),
                 timings=AskTimings(
-                    generative_first_chunk=self._metadata.timings.get(
-                        "generative_first_chunk"
-                    ),
+                    generative_first_chunk=self._metadata.timings.get("generative_first_chunk"),
                     generative_total=self._metadata.timings.get("generative"),
                 ),
             )
         citations = {}
         if self._citations is not None:
             citations = self._citations.citations
+
+        answer_json = None
+        if self._object is not None:
+            answer_json = self._object.object
+
+        prequeries_results: Optional[dict[str, KnowledgeboxFindResults]] = None
+        if self.prequeries_results:
+            prequeries_results = {}
+            for index, (prequery, result) in enumerate(self.prequeries_results):
+                prequery_id = prequery.id or f"prequery_{index}"
+                prequeries_results[prequery_id] = result
+
+        best_matches = [
+            AskRetrievalMatch(
+                id=match.paragraph.id,
+            )
+            for match in self.best_matches
+        ]
+
         response = SyncAskResponse(
             answer=self._answer_text,
+            answer_json=answer_json,
             status=self.status_code.prettify(),
             relations=self._relations,
-            retrieval_results=self.find_results,
+            retrieval_results=self.main_results,
+            retrieval_best_matches=best_matches,
+            prequeries=prequeries_results,
             citations=citations,
             metadata=metadata,
             learning_id=self.nuclia_learning_id or "",
         )
+        if self.status_code == AnswerStatusCode.ERROR and self.status_error_details:
+            response.error_details = self.status_error_details
         if self.ask_request_with_debug_flag:
             sorted_prompt_context = sorted_prompt_context_list(
                 self.prompt_context, self.prompt_context_order
             )
             response.prompt_context = sorted_prompt_context
-        return response.json(exclude_unset=True)
+            if self.debug_chat_model:
+                response.predict_request = self.debug_chat_model.model_dump(mode="json")
+        return response.model_dump_json(exclude_none=True, by_alias=True)
 
     async def get_relations_results(self) -> Relations:
         if self._relations is None:
-            self._relations = await get_relations_results(
-                kbid=self.kbid,
-                text_answer=self._answer_text,
-                target_shard_replicas=self.ask_request.shards,
-            )
+            with self.metrics.time("relations"):
+                self._relations = await get_relations_results(
+                    kbid=self.kbid,
+                    text_answer=self._answer_text,
+                    target_shard_replicas=self.ask_request.shards,
+                    timeout=5.0,
+                )
         return self._relations
 
     async def _stream_predict_answer_text(self) -> AsyncGenerator[str, None]:
@@ -257,12 +398,12 @@ class AskResult:
             if isinstance(item, TextGenerativeResponse):
                 self._answer_text += item.text
                 yield item.text
+            elif isinstance(item, JSONGenerativeResponse):
+                self._object = item
             elif isinstance(item, StatusGenerativeResponse):
                 self._status = item
-                continue
             elif isinstance(item, CitationsGenerativeResponse):
                 self._citations = item
-                continue
             elif isinstance(item, MetaGenerativeResponse):
                 self._metadata = item
             else:
@@ -275,9 +416,11 @@ class AskResult:
 class NotEnoughContextAskResult(AskResult):
     def __init__(
         self,
-        find_results: KnowledgeboxFindResults,
+        main_results: Optional[KnowledgeboxFindResults] = None,
+        prequeries_results: Optional[list[PreQueryResult]] = None,
     ):
-        self.find_results = find_results
+        self.main_results = main_results or KnowledgeboxFindResults(resources={}, min_score=None)
+        self.prequeries_results = prequeries_results or []
         self.nuclia_learning_id = None
 
     async def ndjson_stream(self) -> AsyncGenerator[str, None]:
@@ -286,19 +429,17 @@ class NotEnoughContextAskResult(AskResult):
         return the find results and the messages indicating that there is not enough
         context in the corpus to answer.
         """
-        yield self._ndjson_encode(RetrievalAskResponseItem(results=self.find_results))
+        yield self._ndjson_encode(RetrievalAskResponseItem(results=self.main_results))
         yield self._ndjson_encode(AnswerAskResponseItem(text=NOT_ENOUGH_CONTEXT_ANSWER))
         status = AnswerStatusCode.NO_CONTEXT
-        yield self._ndjson_encode(
-            StatusAskResponseItem(code=status.value, status=status.prettify())
-        )
+        yield self._ndjson_encode(StatusAskResponseItem(code=status.value, status=status.prettify()))
 
     async def json(self) -> str:
         return SyncAskResponse(
             answer=NOT_ENOUGH_CONTEXT_ANSWER,
-            retrieval_results=self.find_results,
+            retrieval_results=self.main_results,
             status=AnswerStatusCode.NO_CONTEXT,
-        ).json(exclude_unset=True)
+        ).model_dump_json()
 
 
 async def ask(
@@ -310,7 +451,7 @@ async def ask(
     origin: str,
     resource: Optional[str] = None,
 ) -> AskResult:
-    start_time = time()
+    metrics = RAGMetrics()
     chat_history = ask_request.context or []
     user_context = ask_request.extra_context or []
     user_query = ask_request.query
@@ -318,117 +459,116 @@ async def ask(
     # Maybe rephrase the query
     rephrased_query = None
     if len(chat_history) > 0 or len(user_context) > 0:
-        rephrased_query = await rephrase_query(
-            kbid,
-            chat_history=chat_history,
-            query=user_query,
-            user_id=user_id,
-            user_context=user_context,
-            generative_model=ask_request.generative_model,
-        )
-
-    # Retrieval is not needed if we are chatting on a specific
-    # resource and the full_resource strategy is enabled
-    needs_retrieval = True
-    if resource is not None:
-        ask_request.resource_filters = [resource]
-        if any(
-            strategy.name == "full_resource" for strategy in ask_request.rag_strategies
-        ):
-            needs_retrieval = False
+        try:
+            with metrics.time("rephrase"):
+                rephrased_query = await rephrase_query(
+                    kbid,
+                    chat_history=chat_history,
+                    query=user_query,
+                    user_id=user_id,
+                    user_context=user_context,
+                    generative_model=ask_request.generative_model,
+                )
+        except RephraseMissingContextError:
+            logger.info("Failed to rephrase ask query, using original")
 
-    # Maybe do a retrieval query
-    if needs_retrieval:
-        find_results, query_parser = await get_find_results(
+    try:
+        retrieval_results = await retrieval_step(
             kbid=kbid,
-            # Prefer the rephrased query if available
-            query=rephrased_query or user_query,
-            chat_request=ask_request,
-            ndb_client=client_type,
-            user=user_id,
+            # Prefer the rephrased query for retrieval if available
+            main_query=rephrased_query or user_query,
+            ask_request=ask_request,
+            client_type=client_type,
+            user_id=user_id,
             origin=origin,
+            metrics=metrics,
+            resource=resource,
         )
-        if len(find_results.resources) == 0:
-            return NotEnoughContextAskResult(find_results=find_results)
-
-    else:
-        find_results = KnowledgeboxFindResults(resources={}, min_score=None)
-        query_parser = QueryParser(
-            kbid=kbid,
-            features=[],
-            query="",
-            filters=ask_request.filters,
-            page_number=0,
-            page_size=0,
-            min_score=MinScore(),
+    except NoRetrievalResultsError as err:
+        # If a retrieval was attempted but no results were found,
+        # early return the ask endpoint without querying the generative model
+        return NotEnoughContextAskResult(
+            main_results=err.main_query,
+            prequeries_results=err.prequeries,
         )
 
-    # Now we build the prompt context
-    query_parser.max_tokens = ask_request.max_tokens  # type: ignore
-    max_tokens_context = await query_parser.get_max_tokens_context()
-    prompt_context_builder = PromptContextBuilder(
-        kbid=kbid,
-        find_results=find_results,
-        resource=resource,
-        user_context=user_context,
-        strategies=ask_request.rag_strategies,
-        image_strategies=ask_request.rag_images_strategies,
-        max_context_characters=tokens_to_chars(max_tokens_context),
-        visual_llm=await query_parser.get_visual_llm_enabled(),
-    )
-    (
-        prompt_context,
-        prompt_context_order,
-        prompt_context_images,
-    ) = await prompt_context_builder.build()
+    query_parser = retrieval_results.query_parser
 
-    # Parse the user prompt (if any)
-    user_prompt = None
-    if ask_request.prompt is not None:
-        user_prompt = UserPrompt(prompt=ask_request.prompt)
+    # Now we build the prompt context
+    with metrics.time("context_building"):
+        query_parser.max_tokens = ask_request.max_tokens  # type: ignore
+        max_tokens_context = await query_parser.get_max_tokens_context()
+        prompt_context_builder = PromptContextBuilder(
+            kbid=kbid,
+            ordered_paragraphs=[match.paragraph for match in retrieval_results.best_matches],
+            resource=resource,
+            user_context=user_context,
+            strategies=ask_request.rag_strategies,
+            image_strategies=ask_request.rag_images_strategies,
+            max_context_characters=tokens_to_chars(max_tokens_context),
+            visual_llm=await query_parser.get_visual_llm_enabled(),
+        )
+        (
+            prompt_context,
+            prompt_context_order,
+            prompt_context_images,
+        ) = await prompt_context_builder.build()
 
     # Make the chat request to the predict API
+    custom_prompt = parse_custom_prompt(ask_request)
     chat_model = ChatModel(
         user_id=user_id,
+        system=custom_prompt.system,
+        user_prompt=UserPrompt(prompt=custom_prompt.user) if custom_prompt.user else None,
         query_context=prompt_context,
         query_context_order=prompt_context_order,
         chat_history=chat_history,
         question=user_query,
         truncate=True,
-        user_prompt=user_prompt,
         citations=ask_request.citations,
+        citation_threshold=ask_request.citation_threshold,
         generative_model=ask_request.generative_model,
         max_tokens=query_parser.get_max_tokens_answer(),
         query_context_images=prompt_context_images,
+        json_schema=ask_request.answer_json_schema,
+        rerank_context=False,
+        top_k=ask_request.top_k,
     )
-    predict = get_predict()
-    nuclia_learning_id, predict_answer_stream = await predict.chat_query_ndjson(
-        kbid, chat_model
-    )
+    with metrics.time("stream_start"):
+        predict = get_predict()
+        (
+            nuclia_learning_id,
+            nuclia_learning_model,
+            predict_answer_stream,
+        ) = await predict.chat_query_ndjson(kbid, chat_model)
+        debug_chat_model = chat_model
 
     auditor = ChatAuditor(
         kbid=kbid,
         user_id=user_id,
         client_type=client_type,
         origin=origin,
-        start_time=start_time,
         user_query=user_query,
         rephrased_query=rephrased_query,
         chat_history=chat_history,
         learning_id=nuclia_learning_id,
         query_context=prompt_context,
         query_context_order=prompt_context_order,
+        model=nuclia_learning_model,
     )
-
     return AskResult(
         kbid=kbid,
         ask_request=ask_request,
-        find_results=find_results,
+        main_results=retrieval_results.main_query,
+        prequeries_results=retrieval_results.prequeries,
         nuclia_learning_id=nuclia_learning_id,
-        predict_answer_stream=predict_answer_stream,
+        predict_answer_stream=predict_answer_stream,  # type: ignore
         prompt_context=prompt_context,
         prompt_context_order=prompt_context_order,
         auditor=auditor,
+        metrics=metrics,
+        best_matches=retrieval_results.best_matches,
+        debug_chat_model=debug_chat_model,
     )
 
 
@@ -468,3 +608,311 @@ def handled_ask_exceptions(func):
             return HTTPClientError(status_code=412, detail=str(exc))
 
     return wrapper
+
+
+def parse_prequeries(ask_request: AskRequest) -> Optional[PreQueriesStrategy]:
+    query_ids = []
+    for rag_strategy in ask_request.rag_strategies:
+        if rag_strategy.name == RagStrategyName.PREQUERIES:
+            prequeries = cast(PreQueriesStrategy, rag_strategy)
+            # Give each query a unique id if they don't have one
+            for index, query in enumerate(prequeries.queries):
+                if query.id is None:
+                    query.id = f"prequery_{index}"
+                if query.id in query_ids:
+                    raise InvalidQueryError(
+                        "rag_strategies",
+                        "Prequeries must have unique ids",
+                    )
+                query_ids.append(query.id)
+            return prequeries
+    return None
+
+
+async def retrieval_step(
+    kbid: str,
+    main_query: str,
+    ask_request: AskRequest,
+    client_type: NucliaDBClientType,
+    user_id: str,
+    origin: str,
+    metrics: RAGMetrics,
+    resource: Optional[str] = None,
+) -> RetrievalResults:
+    """
+    This function encapsulates all the logic related to retrieval in the ask endpoint.
+    """
+    if resource is None:
+        return await retrieval_in_kb(
+            kbid,
+            main_query,
+            ask_request,
+            client_type,
+            user_id,
+            origin,
+            metrics,
+        )
+    else:
+        return await retrieval_in_resource(
+            kbid,
+            resource,
+            main_query,
+            ask_request,
+            client_type,
+            user_id,
+            origin,
+            metrics,
+        )
+
+
+async def retrieval_in_kb(
+    kbid: str,
+    main_query: str,
+    ask_request: AskRequest,
+    client_type: NucliaDBClientType,
+    user_id: str,
+    origin: str,
+    metrics: RAGMetrics,
+) -> RetrievalResults:
+    prequeries = parse_prequeries(ask_request)
+    with metrics.time("retrieval"):
+        main_results, prequeries_results, query_parser = await get_find_results(
+            kbid=kbid,
+            query=main_query,
+            item=ask_request,
+            ndb_client=client_type,
+            user=user_id,
+            origin=origin,
+            metrics=metrics,
+            prequeries_strategy=prequeries,
+        )
+        if len(main_results.resources) == 0 and all(
+            len(prequery_result.resources) == 0 for (_, prequery_result) in prequeries_results or []
+        ):
+            raise NoRetrievalResultsError(main_results, prequeries_results)
+
+    main_query_weight = prequeries.main_query_weight if prequeries is not None else 1.0
+    best_matches = compute_best_matches(
+        main_results=main_results,
+        prequeries_results=prequeries_results,
+        main_query_weight=main_query_weight,
+    )
+    return RetrievalResults(
+        main_query=main_results,
+        prequeries=prequeries_results,
+        query_parser=query_parser,
+        main_query_weight=main_query_weight,
+        best_matches=best_matches,
+    )
+
+
+async def retrieval_in_resource(
+    kbid: str,
+    resource: str,
+    main_query: str,
+    ask_request: AskRequest,
+    client_type: NucliaDBClientType,
+    user_id: str,
+    origin: str,
+    metrics: RAGMetrics,
+) -> RetrievalResults:
+    if any(strategy.name == "full_resource" for strategy in ask_request.rag_strategies):
+        # Retrieval is not needed if we are chatting on a specific resource and the full_resource strategy is enabled
+        return RetrievalResults(
+            main_query=KnowledgeboxFindResults(resources={}, min_score=None),
+            prequeries=None,
+            query_parser=QueryParser(
+                kbid=kbid,
+                features=[],
+                query="",
+                label_filters=ask_request.filters,
+                keyword_filters=ask_request.keyword_filters,
+                top_k=0,
+                min_score=MinScore(),
+            ),
+            main_query_weight=1.0,
+        )
+
+    prequeries = parse_prequeries(ask_request)
+    if prequeries is None and ask_request.answer_json_schema is not None and main_query == "":
+        prequeries = calculate_prequeries_for_json_schema(ask_request)
+
+    # Make sure the retrieval is scoped to the resource if provided
+    ask_request.resource_filters = [resource]
+    if prequeries is not None:
+        for prequery in prequeries.queries:
+            if prequery.prefilter is True:
+                raise InvalidQueryError(
+                    "rag_strategies",
+                    "Prequeries with prefilter are not supported when asking on a resource",
+                )
+            prequery.request.resource_filters = [resource]
+
+    with metrics.time("retrieval"):
+        main_results, prequeries_results, query_parser = await get_find_results(
+            kbid=kbid,
+            query=main_query,
+            item=ask_request,
+            ndb_client=client_type,
+            user=user_id,
+            origin=origin,
+            metrics=metrics,
+            prequeries_strategy=prequeries,
+        )
+        if len(main_results.resources) == 0 and all(
+            len(prequery_result.resources) == 0 for (_, prequery_result) in prequeries_results or []
+        ):
+            raise NoRetrievalResultsError(main_results, prequeries_results)
+    main_query_weight = prequeries.main_query_weight if prequeries is not None else 1.0
+    best_matches = compute_best_matches(
+        main_results=main_results,
+        prequeries_results=prequeries_results,
+        main_query_weight=main_query_weight,
+    )
+    return RetrievalResults(
+        main_query=main_results,
+        prequeries=prequeries_results,
+        query_parser=query_parser,
+        main_query_weight=main_query_weight,
+        best_matches=best_matches,
+    )
+
+
+def compute_best_matches(
+    main_results: KnowledgeboxFindResults,
+    prequeries_results: Optional[list[PreQueryResult]] = None,
+    main_query_weight: float = 1.0,
+) -> list[RetrievalMatch]:
+    """
+    Returns the list of matches of the retrieval results, ordered by relevance (descending weighted score).
+
+    If prequeries_results is provided, the paragraphs of the prequeries are weighted according to the
+    normalized weight of the prequery. The paragraph score is not modified, but it is used to determine the order in which they
+    are presented in the LLM prompt context.
+
+    If a paragraph is matched in various prequeries, the final weighted score is the sum of the weighted scores for each prequery.
+
+    `main_query_weight` is the weight given to the paragraphs matching the main query when calculating the final score.
+    """
+
+    def iter_paragraphs(results: KnowledgeboxFindResults):
+        for resource in results.resources.values():
+            for field in resource.fields.values():
+                for paragraph in field.paragraphs.values():
+                    yield paragraph
+
+    total_weights = main_query_weight + sum(prequery.weight for prequery, _ in prequeries_results or [])
+    paragraph_id_to_match: dict[str, RetrievalMatch] = {}
+    for paragraph in iter_paragraphs(main_results):
+        normalized_weight = main_query_weight / total_weights
+        rmatch = RetrievalMatch(
+            paragraph=paragraph,
+            weighted_score=paragraph.score * normalized_weight,
+        )
+        paragraph_id_to_match[paragraph.id] = rmatch
+
+    for prequery, prequery_results in prequeries_results or []:
+        for paragraph in iter_paragraphs(prequery_results):
+            normalized_weight = prequery.weight / total_weights
+            weighted_score = paragraph.score * normalized_weight
+            if paragraph.id in paragraph_id_to_match:
+                rmatch = paragraph_id_to_match[paragraph.id]
+                # If a paragraph is matched in various prequeries, the final score is the
+                # sum of the weighted scores
+                rmatch.weighted_score += weighted_score
+            else:
+                paragraph_id_to_match[paragraph.id] = RetrievalMatch(
+                    paragraph=paragraph,
+                    weighted_score=weighted_score,
+                )
+
+    return sorted(
+        paragraph_id_to_match.values(),
+        key=lambda match: match.weighted_score,
+        reverse=True,
+    )
+
+
+def calculate_prequeries_for_json_schema(
+    ask_request: AskRequest,
+) -> Optional[PreQueriesStrategy]:
+    """
+    This function generates a PreQueriesStrategy with a query for each property in the JSON schema
+    found in ask_request.answer_json_schema.
+
+    This is useful for the use-case where the user is asking for a structured answer on a corpus
+    that is too big to send to the generative model.
+
+    For instance, a JSON schema like this:
+    {
+        "name": "book_ordering",
+        "description": "Structured answer for a book to order",
+        "parameters": {
+            "type": "object",
+            "properties": {
+                "title": {
+                    "type": "string",
+                    "description": "The title of the book"
+                },
+                "author": {
+                    "type": "string",
+                    "description": "The author of the book"
+                },
+            },
+            "required": ["title", "author"]
+        }
+    }
+    Will generate a PreQueriesStrategy with 2 queries, one for each property in the JSON schema, with equal weights
+    [
+        PreQuery(request=FindRequest(query="The title of the book", ...), weight=1.0),
+        PreQuery(request=FindRequest(query="The author of the book", ...), weight=1.0),
+    ]
+    """
+    prequeries: list[PreQuery] = []
+    json_schema = ask_request.answer_json_schema or {}
+    features = []
+    if ChatOptions.SEMANTIC in ask_request.features:
+        features.append(SearchOptions.SEMANTIC)
+    if ChatOptions.KEYWORD in ask_request.features:
+        features.append(SearchOptions.KEYWORD)
+
+    properties = json_schema.get("parameters", {}).get("properties", {})
+    if len(properties) == 0:  # pragma: no cover
+        return None
+    for prop_name, prop_def in properties.items():
+        query = prop_name
+        if prop_def.get("description"):
+            query += f": {prop_def['description']}"
+        req = FindRequest(
+            query=query,
+            features=features,
+            filters=[],
+            keyword_filters=[],
+            top_k=10,
+            min_score=ask_request.min_score,
+            vectorset=ask_request.vectorset,
+            highlight=False,
+            debug=False,
+            show=[],
+            with_duplicates=False,
+            with_synonyms=False,
+            resource_filters=[],  # to be filled with the resource filter
+            rephrase=ask_request.rephrase,
+            rephrase_prompt=parse_rephrase_prompt(ask_request),
+            security=ask_request.security,
+            autofilter=False,
+        )
+        prequery = PreQuery(
+            request=req,
+            weight=1.0,
+        )
+        prequeries.append(prequery)
+    try:
+        strategy = PreQueriesStrategy(queries=prequeries)
+    except ValidationError:
+        raise AnswerJsonSchemaTooLong(
+            "Answer JSON schema with too many properties generated too many prequeries"
+        )
+
+    ask_request.rag_strategies = [strategy]
+    return strategy