nucliadb 6.7.2.post4862__py3-none-any.whl → 6.9.2.post5282__py3-none-any.whl

nucliadb/search/search/chat/ask.py

@@ -20,14 +20,16 @@
 import dataclasses
 import functools
 import json
-from typing import AsyncGenerator, Optional, cast
+from typing import AsyncGenerator, Optional, Union, cast
 
 from nuclia_models.common.consumption import Consumption
 from nuclia_models.predict.generative_responses import (
     CitationsGenerativeResponse,
+    FootnoteCitationsGenerativeResponse,
     GenerativeChunk,
     JSONGenerativeResponse,
     MetaGenerativeResponse,
+    ReasoningGenerativeResponse,
     StatusGenerativeResponse,
     TextGenerativeResponse,
 )
@@ -90,6 +92,7 @@ from nucliadb_models.search import (
     FindOptions,
     FindParagraph,
     FindRequest,
+    FootnoteCitationsAskResponseItem,
     GraphStrategy,
     JSONAskResponseItem,
     KnowledgeboxFindResults,
@@ -102,6 +105,7 @@ from nucliadb_models.search import (
     PromptContext,
     PromptContextOrder,
     RagStrategyName,
+    ReasoningAskResponseItem,
     Relations,
     RelationsAskResponseItem,
     RetrievalAskResponseItem,
@@ -167,9 +171,11 @@ 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
@@ -220,12 +226,23 @@ 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:
+                    # 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"
 
         if self._object is not None:
             yield JSONAskResponseItem(object=self._object.object)
@@ -274,8 +291,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 +306,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,6 +388,10 @@ 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
@@ -384,6 +412,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 +420,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 +450,9 @@ class AskResult:
                 )
         return self._relations
 
-    async def _stream_predict_answer_text(self) -> AsyncGenerator[str, None]:
+    async def _stream_predict_answer_text(
+        self,
+    ) -> AsyncGenerator[Union[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 +467,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):
@@ -559,11 +599,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={},
@@ -625,6 +667,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
@@ -1034,7 +1077,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/prompt.py

@@ -26,6 +26,7 @@ from typing import Deque, Dict, List, Optional, Sequence, Tuple, Union, cast
 import yaml
 from pydantic import BaseModel
 
+from nucliadb.common import datamanagers
 from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB, FieldId, ParagraphId
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.common.models_utils import from_proto
@@ -246,7 +247,7 @@ async def default_prompt_context(
     - Using an dict prevents from duplicates pulled in through conversation expansion.
     """
     # Sort retrieved paragraphs by decreasing order (most relevant first)
-    async with get_driver().transaction(read_only=True) as txn:
+    async with get_driver().ro_transaction() as txn:
         storage = await get_storage()
         kb = KnowledgeBoxORM(txn, storage, kbid)
         for paragraph in ordered_paragraphs:
@@ -589,18 +590,7 @@ async def field_extension_prompt_context(
         if resource_uuid not in ordered_resources:
             ordered_resources.append(resource_uuid)
 
-    # Fetch the extracted texts of the specified fields for each resource
-    extend_fields = strategy.fields
-    extend_field_ids = []
-    for resource_uuid in ordered_resources:
-        for field_id in extend_fields:
-            try:
-                fid = FieldId.from_string(f"{resource_uuid}/{field_id.strip('/')}")
-                extend_field_ids.append(fid)
-            except ValueError:  # pragma: no cover
-                # Invalid field id, skiping
-                continue
-
+    extend_field_ids = await get_matching_field_ids(kbid, ordered_resources, strategy)
     tasks = [hydrate_field_text(kbid, fid) for fid in extend_field_ids]
     field_extracted_texts = await run_concurrently(tasks)
 
@@ -630,6 +620,43 @@ async def field_extension_prompt_context(
             context[paragraph.id] = _clean_paragraph_text(paragraph)
 
 
+async def get_matching_field_ids(
+    kbid: str, ordered_resources: list[str], strategy: FieldExtensionStrategy
+) -> list[FieldId]:
+    extend_field_ids: list[FieldId] = []
+    # Fetch the extracted texts of the specified fields for each resource
+    for resource_uuid in ordered_resources:
+        for field_id in strategy.fields:
+            try:
+                fid = FieldId.from_string(f"{resource_uuid}/{field_id.strip('/')}")
+                extend_field_ids.append(fid)
+            except ValueError:  # pragma: no cover
+                # Invalid field id, skiping
+                continue
+    if len(strategy.data_augmentation_field_prefixes) > 0:
+        for resource_uuid in ordered_resources:
+            all_field_ids = await datamanagers.atomic.resources.get_all_field_ids(
+                kbid=kbid, rid=resource_uuid, for_update=False
+            )
+            if all_field_ids is None:
+                continue
+            for fieldid in all_field_ids.fields:
+                # Generated fields are always text fields starting with "da-"
+                if any(
+                    (
+                        fieldid.field_type == resources_pb2.FieldType.TEXT
+                        and fieldid.field.startswith(f"da-{prefix}-")
+                    )
+                    for prefix in strategy.data_augmentation_field_prefixes
+                ):
+                    extend_field_ids.append(
+                        FieldId.from_pb(
+                            rid=resource_uuid, field_type=fieldid.field_type, key=fieldid.field
+                        )
+                    )
+    return extend_field_ids
+
+
 async def get_orm_field(kbid: str, field_id: FieldId) -> Optional[Field]:
     resource = await cache.get_resource(kbid, field_id.rid)
     if resource is None:  # pragma: no cover
@@ -779,7 +806,7 @@ async def conversation_prompt_context(
 ):
     analyzed_fields: List[str] = []
     ops = 0
-    async with get_driver().transaction(read_only=True) as txn:
+    async with get_driver().ro_transaction() as txn:
         storage = await get_storage()
         kb = KnowledgeBoxORM(txn, storage, kbid)
         for paragraph in ordered_paragraphs:
@@ -946,9 +973,14 @@ async def hierarchy_prompt_context(
         paragraph_id = ParagraphId.from_string(paragraph.id)
         extended_paragraph_text = paragraph.text
         if paragraphs_extra_characters > 0:
+            extended_paragraph_id = ParagraphId(
+                field_id=paragraph_id.field_id,
+                paragraph_start=paragraph_id.paragraph_start,
+                paragraph_end=paragraph_id.paragraph_end + paragraphs_extra_characters,
+            )
             extended_paragraph_text = await get_paragraph_text(
                 kbid=kbid,
-                paragraph_id=paragraph_id,
+                paragraph_id=extended_paragraph_id,
                 log_on_missing_field=True,
             )
         rid = paragraph_id.rid

nucliadb/search/search/chat/query.py

@@ -200,7 +200,6 @@ def find_request_from_ask_request(item: AskRequest, query: str) -> FindRequest:
     find_request.range_modification_end = item.range_modification_end
     find_request.show = item.show
     find_request.extracted = item.extracted
-    find_request.autofilter = item.autofilter
     find_request.highlight = item.highlight
     find_request.security = item.security
     find_request.debug = item.debug
@@ -308,11 +307,13 @@ def maybe_audit_chat(
     origin: str,
     generative_answer_time: float,
     generative_answer_first_chunk_time: float,
+    generative_reasoning_first_chunk_time: Optional[float],
     rephrase_time: Optional[float],
     user_query: str,
     rephrased_query: Optional[str],
     retrieval_rephrase_query: Optional[str],
     text_answer: bytes,
+    text_reasoning: Optional[str],
     status_code: AnswerStatusCode,
     chat_history: list[ChatContextMessage],
     query_context: PromptContext,
@@ -344,12 +345,14 @@ def maybe_audit_chat(
         question=user_query,
         generative_answer_time=generative_answer_time,
         generative_answer_first_chunk_time=generative_answer_first_chunk_time,
+        generative_reasoning_first_chunk_time=generative_reasoning_first_chunk_time,
         rephrase_time=rephrase_time,
         rephrased_question=rephrased_query,
         retrieval_rephrased_question=retrieval_rephrase_query,
         chat_context=chat_history_context,
         retrieved_context=chat_retrieved_context,
         answer=audit_answer,
+        reasoning=text_reasoning,
         learning_id=learning_id,
         status_code=int(status_code.value),
         model=model,
@@ -401,8 +404,10 @@ class ChatAuditor:
     def audit(
         self,
         text_answer: bytes,
+        text_reasoning: Optional[str],
         generative_answer_time: float,
         generative_answer_first_chunk_time: float,
+        generative_reasoning_first_chunk_time: Optional[float],
         rephrase_time: Optional[float],
         status_code: AnswerStatusCode,
     ):
@@ -416,8 +421,10 @@ class ChatAuditor:
             retrieval_rephrase_query=self.retrieval_rephrased_query,
             generative_answer_time=generative_answer_time,
             generative_answer_first_chunk_time=generative_answer_first_chunk_time,
+            generative_reasoning_first_chunk_time=generative_reasoning_first_chunk_time,
             rephrase_time=rephrase_time,
             text_answer=text_answer,
+            text_reasoning=text_reasoning,
             status_code=status_code,
             chat_history=self.chat_history,
             query_context=self.query_context,

nucliadb/search/search/fetch.py

@@ -55,7 +55,7 @@ async def fetch_resources(
         extracted = []
 
     result = {}
-    async with get_driver().transaction(read_only=True) as txn:
+    async with get_driver().ro_transaction() as txn:
         tasks = []
         for resource in resources:
             tasks.append(

nucliadb/search/search/find.py

@@ -100,7 +100,6 @@ async def _index_node_retrieval(
         (
             pb_query,
             incomplete_results,
-            autofilters,
             rephrased_query,
         ) = await legacy_convert_retrieval_to_proto(parsed)
 
@@ -137,7 +136,6 @@ async def _index_node_retrieval(
         )
 
     search_results.shards = queried_shards
-    search_results.autofilters = autofilters
 
     ndb_time = metrics["index_search"] + metrics["results_merge"]
     if metrics["index_search"] > settings.slow_node_query_log_threshold:
@@ -180,9 +178,7 @@ async def _external_index_retrieval(
     parsed = await parse_find(kbid, item)
     assert parsed.retrieval.reranker is not None, "find parser must provide a reranking algorithm"
     reranker = get_reranker(parsed.retrieval.reranker)
-    search_request, incomplete_results, _, rephrased_query = await legacy_convert_retrieval_to_proto(
-        parsed
-    )
+    search_request, incomplete_results, rephrased_query = await legacy_convert_retrieval_to_proto(parsed)
 
     # Query index
     query_results = await external_index_manager.query(search_request)  # noqa
@@ -220,7 +216,6 @@ async def _external_index_retrieval(
         page_number=0,
         page_size=item.top_k,
         relations=None,  # Not implemented for external indexes yet
-        autofilters=[],  # Not implemented for external indexes yet
         min_score=results_min_score,
         best_matches=best_matches,
         # These are not used for external indexes

nucliadb/search/search/{hydrator.py → hydrator/__init__.py}

@@ -28,7 +28,8 @@ from nucliadb.common.external_index_providers.base import TextBlockMatch
 from nucliadb.common.ids import FieldId
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.ingest.serialize import managed_serialize
-from nucliadb.search.search import cache, paragraphs
+from nucliadb.search.search import cache
+from nucliadb.search.search.paragraphs import get_paragraph_text
 from nucliadb_models.common import FieldTypeName
 from nucliadb_models.resource import ExtractedDataTypeName, Resource
 from nucliadb_models.search import (
@@ -79,7 +80,7 @@ async def hydrate_resource_text(
         return []
 
     # Schedule the extraction of the text of each field in the resource
-    async with get_driver().transaction(read_only=True) as txn:
+    async with get_driver().ro_transaction() as txn:
         resource.txn = txn
         runner = ConcurrentRunner(max_tasks=max_concurrent_tasks)
         for field_type, field_key in await resource.get_fields(force=True):
@@ -120,7 +121,7 @@ async def hydrate_resource_metadata(
         if concurrency_control is not None:
             await stack.enter_async_context(concurrency_control)
 
-        async with get_driver().transaction(read_only=True) as ro_txn:
+        async with get_driver().ro_transaction() as ro_txn:
             serialized_resource = await managed_serialize(
                 txn=ro_txn,
                 kbid=kbid,
@@ -170,7 +171,7 @@ async def hydrate_text_block(
         if concurrency_control is not None:
             await stack.enter_async_context(concurrency_control)
 
-        text_block.text = await paragraphs.get_paragraph_text(
+        text_block.text = await get_paragraph_text(
             kbid=kbid,
             paragraph_id=text_block.paragraph_id,
             highlight=options.highlight,

nucliadb/search/search/hydrator/fields.py

@@ -0,0 +1,175 @@
+# 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 nucliadb.common.ids import FIELD_TYPE_STR_TO_NAME, FieldId
+from nucliadb.common.models_utils import from_proto
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb.search.search.hydrator import hydrate_field_text
+from nucliadb_models import hydration as hydration_models
+from nucliadb_models.common import FieldTypeName
+
+
+def page_preview_id(page_number: int) -> str:
+    """Return the string page number for an specific page"""
+    return f"{page_number}"
+
+
+async def hydrate_field(resource: Resource, field_id: FieldId, config: hydration_models.FieldHydration):
+    field_type = FIELD_TYPE_STR_TO_NAME[field_id.type]
+
+    if field_type == FieldTypeName.TEXT:
+        if not config.text is not None:
+            return
+        return await hydrate_text_field(resource, field_id, config.text)
+
+    elif field_type == FieldTypeName.FILE is not None:
+        if not config.file:
+            return
+        return await hydrate_file_field(resource, field_id, config.file)
+
+    elif field_type == FieldTypeName.LINK is not None:
+        if not config.link:
+            return
+        return await hydrate_link_field(resource, field_id, config.link)
+
+    elif field_type == FieldTypeName.CONVERSATION is not None:
+        if not config.conversation:
+            return
+        return await hydrate_conversation_field(resource, field_id, config.conversation)
+
+    elif field_type == FieldTypeName.GENERIC is not None:
+        if not config.generic:
+            return
+        return await hydrate_generic_field(resource, field_id, config.generic)
+
+    else:  # pragma: no cover
+        # 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"
+
+
+async def hydrate_text_field(
+    resource: Resource,
+    field_id: FieldId,
+    config: hydration_models.TextFieldHydration,
+) -> hydration_models.HydratedTextField:
+    hydrated = hydration_models.HydratedTextField(
+        id=field_id.full(),
+        resource=field_id.rid,
+        field_type=FieldTypeName.TEXT,
+    )
+
+    if config.extracted_text:
+        field_text = await hydrate_field_text(resource.kb.kbid, field_id)
+        if field_text is not None:
+            (_, text) = field_text
+            hydrated.extracted = hydration_models.FieldExtractedData(text=text)
+
+    return hydrated
+
+
+async def hydrate_file_field(
+    resource: Resource,
+    field_id: FieldId,
+    config: hydration_models.FileFieldHydration,
+) -> hydration_models.HydratedFileField:
+    hydrated = hydration_models.HydratedFileField(
+        id=field_id.full(),
+        resource=field_id.rid,
+        field_type=FieldTypeName.FILE,
+    )
+
+    if config.value:
+        field = await resource.get_field(field_id.key, field_id.pb_type)
+        value = await field.get_value()
+        hydrated.value = from_proto.field_file(value)
+
+    if config.extracted_text:
+        field_text = await hydrate_field_text(resource.kb.kbid, field_id)
+        if field_text is not None:
+            (_, text) = field_text
+            hydrated.extracted = hydration_models.FieldExtractedData(text=text)
+
+    return hydrated
+
+
+async def hydrate_link_field(
+    resource: Resource,
+    field_id: FieldId,
+    config: hydration_models.LinkFieldHydration,
+) -> hydration_models.HydratedLinkField:
+    hydrated = hydration_models.HydratedLinkField(
+        id=field_id.full(),
+        resource=field_id.rid,
+        field_type=FieldTypeName.LINK,
+    )
+
+    if config.value:
+        field = await resource.get_field(field_id.key, field_id.pb_type)
+        value = await field.get_value()
+        hydrated.value = from_proto.field_link(value)
+
+    if config.extracted_text:
+        field_text = await hydrate_field_text(resource.kb.kbid, field_id)
+        if field_text is not None:
+            (_, text) = field_text
+            hydrated.extracted = hydration_models.FieldExtractedData(text=text)
+
+    return hydrated
+
+
+async def hydrate_conversation_field(
+    resource: Resource,
+    field_id: FieldId,
+    config: hydration_models.ConversationFieldHydration,
+) -> hydration_models.HydratedConversationField:
+    hydrated = hydration_models.HydratedConversationField(
+        id=field_id.full(),
+        resource=field_id.rid,
+        field_type=FieldTypeName.CONVERSATION,
+    )
+    # TODO: implement conversation fields
+    return hydrated
+
+
+async def hydrate_generic_field(
+    resource: Resource,
+    field_id: FieldId,
+    config: hydration_models.GenericFieldHydration,
+) -> hydration_models.HydratedGenericField:
+    hydrated = hydration_models.HydratedGenericField(
+        id=field_id.full(),
+        resource=field_id.rid,
+        field_type=FieldTypeName.GENERIC,
+    )
+
+    if config.value:
+        field = await resource.get_field(field_id.key, field_id.pb_type)
+        value = await field.get_value()
+        hydrated.value = value
+
+    if config.extracted_text:
+        field_text = await hydrate_field_text(resource.kb.kbid, field_id)
+        if field_text is not None:
+            (_, text) = field_text
+            hydrated.extracted = hydration_models.FieldExtractedData(text=text)
+
+    return hydrated