nucliadb 6.7.2.post4874__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/search/search/chat/ask.py

@@ -20,18 +20,22 @@
 import dataclasses
 import functools
 import json
-from typing import AsyncGenerator, Optional, cast
+from collections.abc import AsyncGenerator
+from typing import cast
 
 from nuclia_models.common.consumption import Consumption
 from nuclia_models.predict.generative_responses import (
     CitationsGenerativeResponse,
+    FootnoteCitationsGenerativeResponse,
     GenerativeChunk,
     JSONGenerativeResponse,
     MetaGenerativeResponse,
+    ReasoningGenerativeResponse,
     StatusGenerativeResponse,
     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
@@ -47,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,
@@ -67,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,
@@ -90,6 +100,7 @@ from nucliadb_models.search import (
     FindOptions,
     FindParagraph,
     FindRequest,
+    FootnoteCitationsAskResponseItem,
     GraphStrategy,
     JSONAskResponseItem,
     KnowledgeboxFindResults,
@@ -102,6 +113,7 @@ from nucliadb_models.search import (
     PromptContext,
     PromptContextOrder,
     RagStrategyName,
+    ReasoningAskResponseItem,
     Relations,
     RelationsAskResponseItem,
     RetrievalAskResponseItem,
@@ -114,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
@@ -128,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)
 
 
@@ -139,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
@@ -167,12 +181,14 @@ class AskResult:
 
         # 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
-        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:
@@ -181,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
@@ -220,12 +236,21 @@ class AskResult:
     async def _stream(self) -> AsyncGenerator[AskResponseItemType, None]:
         # First, stream out the predict answer
         first_chunk_yielded = False
+        first_reasoning_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 isinstance(answer_chunk, TextGenerativeResponse):
+                    yield AnswerAskResponseItem(text=answer_chunk.text)
+                    if not first_chunk_yielded:
+                        self.metrics.record_first_chunk_yielded()
+                        first_chunk_yielded = True
+                elif isinstance(answer_chunk, ReasoningGenerativeResponse):
+                    yield ReasoningAskResponseItem(text=answer_chunk.text)
+                    if not first_reasoning_chunk_yielded:
+                        self.metrics.record_first_reasoning_chunk_yielded()
+                        first_reasoning_chunk_yielded = True
+                else:
+                    assert_never(answer_chunk)
 
         if self._object is not None:
             yield JSONAskResponseItem(object=self._object.object)
@@ -274,8 +299,10 @@ class AskResult:
             audit_answer = json.dumps(self._object.object).encode("utf-8")
         self.auditor.audit(
             text_answer=audit_answer,
+            text_reasoning=self._reasoning_text,
             generative_answer_time=self.metrics["stream_predict_answer"],
             generative_answer_first_chunk_time=self.metrics.get_first_chunk_time() or 0,
+            generative_reasoning_first_chunk_time=self.metrics.get_first_reasoning_chunk_time(),
             rephrase_time=self.metrics.get("rephrase"),
             status_code=self.status_code,
         )
@@ -287,6 +314,11 @@ class AskResult:
             yield CitationsAskResponseItem(
                 citations=self._citations.citations,
             )
+        # Stream out the footnote citations mapping
+        if self._footnote_citations is not None:
+            yield FootnoteCitationsAskResponseItem(
+                footnote_to_context=self._footnote_citations.footnote_to_context,
+            )
 
         # Stream out generic metadata about the answer
         if self._metadata is not None:
@@ -364,11 +396,15 @@ class AskResult:
         if self._citations is not None:
             citations = self._citations.citations
 
+        footnote_citations = {}
+        if self._footnote_citations is not None:
+            footnote_citations = self._footnote_citations.footnote_to_context
+
         answer_json = None
         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):
@@ -384,6 +420,7 @@ class AskResult:
 
         response = SyncAskResponse(
             answer=self._answer_text,
+            reasoning=self._reasoning_text,
             answer_json=answer_json,
             status=self.status_code.prettify(),
             relations=self._relations,
@@ -391,6 +428,7 @@ class AskResult:
             retrieval_best_matches=best_matches,
             prequeries=prequeries_results,
             citations=citations,
+            citation_footnote_to_context=footnote_citations,
             metadata=metadata,
             consumption=self._consumption,
             learning_id=self.nuclia_learning_id or "",
@@ -420,7 +458,9 @@ class AskResult:
                 )
         return self._relations
 
-    async def _stream_predict_answer_text(self) -> AsyncGenerator[str, None]:
+    async def _stream_predict_answer_text(
+        self,
+    ) -> 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.
@@ -435,13 +475,21 @@ class AskResult:
             item = generative_chunk.chunk
             if isinstance(item, TextGenerativeResponse):
                 self._answer_text += item.text
-                yield item.text
+                yield item
+            elif isinstance(item, ReasoningGenerativeResponse):
+                if self._reasoning_text is None:
+                    self._reasoning_text = item.text
+                else:
+                    self._reasoning_text += item.text
+                yield item
             elif isinstance(item, JSONGenerativeResponse):
                 self._object = item
             elif isinstance(item, StatusGenerativeResponse):
                 self._status = item
             elif isinstance(item, CitationsGenerativeResponse):
                 self._citations = item
+            elif isinstance(item, FootnoteCitationsGenerativeResponse):
+                self._footnote_citations = item
             elif isinstance(item, MetaGenerativeResponse):
                 self._metadata = item
             elif isinstance(item, Consumption):
@@ -456,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 []
@@ -507,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 []
@@ -559,11 +607,13 @@ async def ask(
             origin=origin,
             generative_answer_time=0,
             generative_answer_first_chunk_time=0,
+            generative_reasoning_first_chunk_time=None,
             rephrase_time=metrics.get("rephrase"),
             user_query=user_query,
             rephrased_query=rephrased_query,
             retrieval_rephrase_query=err.main_query.rephrased_query if err.main_query else None,
             text_answer=b"",
+            text_reasoning=None,
             status_code=AnswerStatusCode.NO_RETRIEVAL_DATA,
             chat_history=chat_history,
             query_context={},
@@ -585,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,
@@ -625,6 +692,7 @@ async def ask(
         json_schema=ask_request.answer_json_schema,
         rerank_context=False,
         top_k=ask_request.top_k,
+        reasoning=ask_request.reasoning,
     )
 
     nuclia_learning_id = None
@@ -632,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,
@@ -714,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:
@@ -733,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)
@@ -748,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.
@@ -787,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,
@@ -799,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,
@@ -835,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,
     )
@@ -875,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,
@@ -898,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,
     )
@@ -910,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]:
     """
@@ -925,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,
                         )
@@ -969,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.
@@ -1034,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