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

nucliadb/search/search/hydrator/images.py

@@ -0,0 +1,130 @@
+# 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/>.
+#
+import base64
+from typing import Optional, cast
+
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_NAME, FieldId, ParagraphId
+from nucliadb.ingest.fields.base import Field
+from nucliadb.ingest.fields.file import File
+from nucliadb.search import SERVICE_NAME
+from nucliadb_models.common import FieldTypeName
+from nucliadb_models.search import Image
+from nucliadb_protos import resources_pb2
+from nucliadb_utils.utilities import get_storage
+
+
+async def paragraph_source_image(
+    kbid: str, paragraph_id: ParagraphId, paragraph: resources_pb2.Paragraph
+) -> Optional[Image]:
+    """Certain paragraphs are extracted from images using techniques like OCR or
+    inception. If that's the case, return the original image for this paragraph.
+
+    """
+    source_image = paragraph.representation.reference_file
+    if not source_image:
+        return None
+
+    if paragraph.kind not in (
+        resources_pb2.Paragraph.TypeParagraph.OCR,
+        resources_pb2.Paragraph.TypeParagraph.INCEPTION,
+    ):
+        return None
+
+    field_id = paragraph_id.field_id
+
+    # Paragraphs extracted from an image store its original image representation
+    # in the reference file. The path is incomplete though, as it's stored in
+    # the `generated` folder
+    image = await download_image(
+        kbid,
+        field_id,
+        f"generated/{source_image}",
+        # XXX: we assume all reference files are PNG images, but this actually
+        # depends on learning so it's a dangerous assumption. We should check it
+        # by ourselves
+        mime_type="image/png",
+    )
+    return image
+
+
+async def download_image(
+    kbid: str, field_id: FieldId, image_path: str, *, mime_type: str
+) -> Optional[Image]:
+    storage = await get_storage(service_name=SERVICE_NAME)
+    sf = storage.file_extracted(
+        kbid,
+        field_id.rid,
+        field_id.type,
+        field_id.key,
+        image_path,
+    )
+    raw_image = (await storage.downloadbytes(sf.bucket, sf.key)).getvalue()
+    if not raw_image:
+        return None
+    return Image(content_type=mime_type, b64encoded=base64.b64encode(raw_image).decode())
+
+
+async def download_page_preview(field: Field, page: int) -> Optional[Image]:
+    """Download a specific page preview for a field and return it as an Image.
+    As not all fields have previews, this function can return None.
+
+    Page previews are uploaded by learning and shared through a known path with.
+    nucliadb
+
+    """
+    field_type = FIELD_TYPE_STR_TO_NAME[field.type]
+
+    if field_type == FieldTypeName.FILE:
+        field = cast(File, field)
+        metadata = await field.get_file_extracted_data()
+
+        if metadata is None:
+            return None
+
+        assert page <= len(metadata.file_pages_previews.positions), (
+            f"paragraph page number {page} should be less or equal to the total file pages previews {len(metadata.file_pages_previews.positions)}"
+        )
+        image = await download_image(
+            field.kbid,
+            field.field_id,
+            f"generated/extracted_images_{page}.png",
+            mime_type="image/png",
+        )
+
+    elif field_type == FieldTypeName.LINK:
+        # TODO: in case of links, we want to return the link preview, that is a
+        # link converted to PDF and screenshotted
+        # REVIEW: link preview is an image or a PDF?
+        image = None
+
+    elif (
+        field_type == FieldTypeName.TEXT
+        or field_type == FieldTypeName.CONVERSATION
+        or field_type == FieldTypeName.GENERIC
+    ):
+        # these fields don't have previews
+        image = None
+
+    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"
+
+    return image

nucliadb/search/search/hydrator/paragraphs.py

@@ -0,0 +1,307 @@
+# 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/>.
+#
+import asyncio
+from dataclasses import dataclass
+from typing import Optional, Union
+
+from nucliadb.common.ids import FieldId, ParagraphId
+from nucliadb.ingest.fields.base import Field
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb.search.search import paragraphs
+from nucliadb.search.search.hydrator.fields import page_preview_id
+from nucliadb.search.search.hydrator.images import paragraph_source_image
+from nucliadb_models import hydration as hydration_models
+from nucliadb_protos import resources_pb2
+from nucliadb_protos.resources_pb2 import FieldComputedMetadata
+
+
+class ParagraphIndex:
+    """Small helper class to cache field paragraphs and its relations and be
+    used as an index.
+
+    """
+
+    NEXT = "next"
+    PREVIOUS = "previous"
+    PARENTS = "parents"
+    SIBLINGS = "siblings"
+    REPLACEMENTS = "replacements"
+
+    def __init__(self, field_id: FieldId) -> None:
+        self.field_id = field_id
+        self.paragraphs: dict[str, resources_pb2.Paragraph] = {}
+        self.neighbours: dict[tuple[str, str], str] = {}
+        self.related: dict[tuple[str, str], list[str]] = {}
+        self._lock = asyncio.Lock()
+        self._built = False
+
+    async def build(self, field: Field):
+        """Build the index if it hasn't been built yet.
+
+        This function is async-safe, multiple concurrent tasks can ask for a
+        built and it'll only be done once
+        """
+        if self._built:
+            return
+
+        async with self._lock:
+            # double check we haven't built the index meanwhile we waited for the
+            # lock
+            if self._built:
+                return
+
+            field_metadata = await field.get_field_metadata()
+
+            if field_metadata is None:
+                # field metadata may be still processing. As we want to provide a
+                # consistent view, even if it can appear meanwhile we hydrate, we
+                # consider we don't have it. We mark the index as built and any
+                # paragraph will be found for this field
+                self._built = True
+                return None
+
+            # REVIEW: this is a CPU-bound code, we may consider running this in an
+            # executor to not block the loop
+            self._build(field_metadata)
+            self._built = True
+
+    def _build(self, field_metadata: FieldComputedMetadata):
+        self.paragraphs.clear()
+        self.neighbours.clear()
+        self.related.clear()
+
+        if self.field_id.subfield_id is None:
+            field_paragraphs = field_metadata.metadata.paragraphs
+        else:
+            field_paragraphs = field_metadata.split_metadata[self.field_id.subfield_id].paragraphs
+
+        previous = None
+        for paragraph in field_paragraphs:
+            paragraph_id = self.field_id.paragraph_id(paragraph.start, paragraph.end).full()
+            self.paragraphs[paragraph_id] = paragraph
+
+            if previous is not None:
+                self.neighbours[(previous, ParagraphIndex.NEXT)] = paragraph_id
+                self.neighbours[(paragraph_id, ParagraphIndex.PREVIOUS)] = previous
+            previous = paragraph_id
+
+            self.related[(paragraph_id, ParagraphIndex.PARENTS)] = [
+                parent for parent in paragraph.relations.parents
+            ]
+            self.related[(paragraph_id, ParagraphIndex.SIBLINGS)] = [
+                sibling for sibling in paragraph.relations.siblings
+            ]
+            self.related[(paragraph_id, ParagraphIndex.REPLACEMENTS)] = [
+                replacement for replacement in paragraph.relations.replacements
+            ]
+
+    def get(self, paragraph_id: Union[str, ParagraphId]) -> Optional[resources_pb2.Paragraph]:
+        paragraph_id = str(paragraph_id)
+        return self.paragraphs.get(paragraph_id)
+
+    def previous(self, paragraph_id: Union[str, ParagraphId]) -> Optional[str]:
+        paragraph_id = str(paragraph_id)
+        return self.neighbours.get((paragraph_id, ParagraphIndex.PREVIOUS))
+
+    def next(self, paragraph_id: Union[str, ParagraphId]) -> Optional[str]:
+        paragraph_id = str(paragraph_id)
+        return self.neighbours.get((paragraph_id, ParagraphIndex.NEXT))
+
+    def n_previous(self, paragraph_id: Union[str, ParagraphId], count: int = 1) -> list[str]:
+        assert count >= 1, f"can't find negative previous {count}"
+        paragraph_id = str(paragraph_id)
+        previous: list[str] = []
+        current_id = paragraph_id
+        for _ in range(count):
+            previous_id = self.previous(current_id)
+            if previous_id is None:
+                # we've reached the first paragraph
+                break
+            previous.insert(0, previous_id)
+            current_id = previous_id
+        return previous
+
+    def n_next(self, paragraph_id: Union[str, ParagraphId], count: int = 1) -> list[str]:
+        assert count >= 1, f"can't find negative nexts {count}"
+        paragraph_id = str(paragraph_id)
+        nexts = []
+        current_id = paragraph_id
+        for _ in range(count):
+            next_id = self.next(current_id)
+            if next_id is None:
+                # we've reached the last paragraph
+                break
+            current_id = next_id
+            nexts.append(next_id)
+        return nexts
+
+    def parents(self, paragraph_id: Union[str, ParagraphId]) -> list[str]:
+        paragraph_id = str(paragraph_id)
+        return self.related.get((paragraph_id, ParagraphIndex.PARENTS), [])
+
+    def siblings(self, paragraph_id: Union[str, ParagraphId]) -> list[str]:
+        paragraph_id = str(paragraph_id)
+        return self.related.get((paragraph_id, ParagraphIndex.SIBLINGS), [])
+
+    def replacements(self, paragraph_id: Union[str, ParagraphId]) -> list[str]:
+        paragraph_id = str(paragraph_id)
+        return self.related.get((paragraph_id, ParagraphIndex.REPLACEMENTS), [])
+
+
+@dataclass
+class ExtraParagraphHydration:
+    field_page: Optional[int]
+    field_table_page: Optional[int]
+    related_paragraph_ids: list[ParagraphId]
+
+
+async def hydrate_paragraph(
+    resource: Resource,
+    field: Field,
+    paragraph_id: ParagraphId,
+    config: hydration_models.ParagraphHydration,
+    field_paragraphs_index: ParagraphIndex,
+) -> tuple[hydration_models.HydratedParagraph, ExtraParagraphHydration]:
+    """Hydrate a paragraph and return the extra hydration to built a coherent
+    hydration around this paragraph.
+
+    Although the resource and field exist, the paragraph doesn't necessarily
+    need to be a real one in the paragraph metadata, it can be made-up to
+    include more or less text than the originally extracted.
+
+    """
+    kbid = resource.kb.kbid
+
+    hydrated = hydration_models.HydratedParagraph(
+        id=paragraph_id.full(),
+        field=paragraph_id.field_id.full(),
+        resource=paragraph_id.rid,
+    )
+    extra_hydration = ExtraParagraphHydration(
+        field_page=None, field_table_page=None, related_paragraph_ids=[]
+    )
+
+    if config.text:
+        text = await paragraphs.get_paragraph_text(kbid=kbid, paragraph_id=paragraph_id)
+        hydrated.text = text
+
+    requires_paragraph_metadata = config.image or config.table or config.page or config.related
+    if requires_paragraph_metadata:
+        await field_paragraphs_index.build(field)
+        paragraph = field_paragraphs_index.get(paragraph_id)
+        if paragraph is not None:
+            # otherwise, this is a fake paragraph. We can't hydrate anything else here
+
+            if config.related:
+                hydrated.related, related_ids = await related_paragraphs_refs(
+                    paragraph_id, field_paragraphs_index, config.related
+                )
+                extra_hydration.related_paragraph_ids = related_ids
+
+            if config.image:
+                hydrated.image = hydration_models.HydratedParagraphImage()
+
+                if config.image.source_image:
+                    hydrated.image.source_image = await paragraph_source_image(
+                        kbid, paragraph_id, paragraph
+                    )
+
+            if config.page:
+                if hydrated.page is None:
+                    hydrated.page = hydration_models.HydratedParagraphPage()
+
+                if config.page.page_with_visual:
+                    if paragraph.page.page_with_visual:
+                        # Paragraphs can be found on pages with visual content. In this
+                        # case, we want to return the preview of the paragraph page as
+                        # an image
+                        page_number = paragraph.page.page
+                        # TODO: what should I do if I later find there's no page in the DB?
+                        hydrated.page.page_preview_ref = page_preview_id(page_number)
+                        extra_hydration.field_page = page_number
+
+            if config.table:
+                if hydrated.table is None:
+                    hydrated.table = hydration_models.HydratedParagraphTable()
+
+                if config.table.table_page_preview:
+                    if paragraph.representation.is_a_table:
+                        # When a paragraph comes with a table and table hydration is
+                        # enabled, we want to return the image representing that table.
+                        # Ideally we should hydrate the paragraph reference_file, but
+                        # table screenshots are not always perfect so we prefer to use
+                        # the page preview. If at some point the table images are good
+                        # enough, it'd be better to use those
+                        page_number = paragraph.page.page
+                        hydrated.table.page_preview_ref = page_preview_id(page_number)
+                        extra_hydration.field_table_page = page_number
+
+    return hydrated, extra_hydration
+
+
+async def related_paragraphs_refs(
+    paragraph_id: ParagraphId,
+    index: ParagraphIndex,
+    config: hydration_models.RelatedParagraphHydration,
+) -> tuple[hydration_models.RelatedParagraphRefs, list[ParagraphId]]:
+    """Compute the related paragraph references for a specific `paragraph_id`
+    and return them with the plain list of unique related paragraphs (to
+    facilitate work to the caller).
+
+    """
+    hydrated = hydration_models.RelatedParagraphRefs()
+    related = set()
+
+    if config.neighbours:
+        hydrated.neighbours = hydration_models.RelatedNeighbourParagraphRefs()
+
+        if config.neighbours.before is not None:
+            hydrated.neighbours.before = []
+            if config.neighbours.before > 0:
+                for previous_id in index.n_previous(paragraph_id, config.neighbours.before):
+                    hydrated.neighbours.before.insert(0, previous_id)
+                    related.add(ParagraphId.from_string(previous_id))
+
+        if config.neighbours.after is not None:
+            hydrated.neighbours.after = []
+            if config.neighbours.after > 0:
+                for next_id in index.n_next(paragraph_id, config.neighbours.after):
+                    hydrated.neighbours.after.append(next_id)
+                    related.add(ParagraphId.from_string(next_id))
+
+    if config.parents:
+        hydrated.parents = []
+        for parent_id in index.parents(paragraph_id):
+            hydrated.parents.append(parent_id)
+            related.add(ParagraphId.from_string(parent_id))
+
+    if config.siblings:
+        hydrated.siblings = []
+        for sibling_id in index.siblings(paragraph_id):
+            hydrated.siblings.append(sibling_id)
+            related.add(ParagraphId.from_string(sibling_id))
+
+    if config.replacements:
+        hydrated.replacements = []
+        for replacement_id in index.replacements(paragraph_id):
+            hydrated.replacements.append(replacement_id)
+            related.add(ParagraphId.from_string(replacement_id))
+
+    return hydrated, list(related)

nucliadb/search/search/hydrator/resources.py

@@ -0,0 +1,56 @@
+#!/usr/bin/env python3
+
+# 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.models_utils import from_proto
+from nucliadb.ingest.orm.resource import Resource
+from nucliadb_models import hydration as hydration_models
+from nucliadb_models.security import ResourceSecurity
+
+
+async def hydrate_resource(
+    resource: Resource, rid: str, config: hydration_models.ResourceHydration
+) -> hydration_models.HydratedResource:
+    basic = await resource.get_basic()
+
+    slug = basic.slug
+    hydrated = hydration_models.HydratedResource(id=rid, slug=slug)
+
+    if config.title:
+        hydrated.title = basic.title
+    if config.summary:
+        hydrated.summary = basic.summary
+
+    if config.security:
+        security = await resource.get_security()
+        hydrated.security = ResourceSecurity(access_groups=[])
+        if security is not None:
+            for group_id in security.access_groups:
+                hydrated.security.access_groups.append(group_id)
+
+    if config.origin:
+        origin = await resource.get_origin()
+        if origin is not None:
+            # TODO: we want a better hydration than proto to JSON
+            hydrated.origin = from_proto.origin(origin)
+
+    return hydrated

nucliadb/search/search/metrics.py

@@ -49,6 +49,10 @@ buckets = [
 ]
 
 generative_first_chunk_histogram = metrics.Histogram(
+    name="generative_reasoning_first_chunk",
+    buckets=buckets,
+)
+reasoning_first_chunk_histogram = metrics.Histogram(
     name="generative_first_chunk",
     buckets=buckets,
 )
@@ -107,12 +111,24 @@ class AskMetrics(Metrics):
         super().__init__(id="ask")
         self.global_start = time.monotonic()
         self.first_chunk_yielded_at: Optional[float] = None
+        self.first_reasoning_chunk_yielded_at: Optional[float] = None
 
     def record_first_chunk_yielded(self):
         self.first_chunk_yielded_at = time.monotonic()
         generative_first_chunk_histogram.observe(self.first_chunk_yielded_at - self.global_start)
 
+    def record_first_reasoning_chunk_yielded(self):
+        self.first_reasoning_chunk_yielded_at = time.monotonic()
+        reasoning_first_chunk_histogram.observe(
+            self.first_reasoning_chunk_yielded_at - self.global_start
+        )
+
     def get_first_chunk_time(self) -> Optional[float]:
         if self.first_chunk_yielded_at is None:
             return None
         return self.first_chunk_yielded_at - self.global_start
+
+    def get_first_reasoning_chunk_time(self) -> Optional[float]:
+        if self.first_reasoning_chunk_yielded_at is None:
+            return None
+        return self.first_reasoning_chunk_yielded_at - self.global_start

nucliadb/search/search/predict_proxy.py

@@ -28,6 +28,7 @@ from multidict import CIMultiDictProxy
 from nuclia_models.predict.generative_responses import (
     GenerativeChunk,
     JSONGenerativeResponse,
+    ReasoningGenerativeResponse,
     StatusGenerativeResponse,
     TextGenerativeResponse,
 )
@@ -87,6 +88,7 @@ async def predict_proxy(
     predict_headers = predict.get_predict_headers(kbid)
     user_headers = {k: v for k, v in headers.items() if k.capitalize() in ALLOWED_HEADERS}
 
+    metrics = AskMetrics()
     # Proxy the request to predict API
     predict_response = await predict.make_request(
         method=method,
@@ -109,7 +111,8 @@ async def predict_proxy(
                 client_type=client_type,
                 origin=origin,
                 user_query=user_query,
-                is_json="json" in (media_type or ""),
+                is_ndjson_stream="json" in (media_type or ""),
+                metrics=metrics,
             )
         else:
             streaming_generator = predict_response.content.iter_any()
@@ -120,7 +123,6 @@ async def predict_proxy(
             media_type=media_type,
         )
     else:
-        metrics = AskMetrics()
         with metrics.time(PREDICT_ANSWER_METRIC):
             content = await predict_response.read()
 
@@ -140,8 +142,10 @@ async def predict_proxy(
                 client_type=client_type,
                 origin=origin,
                 text_answer=content,
+                text_reasoning=None,
                 generative_answer_time=metrics[PREDICT_ANSWER_METRIC],
                 generative_answer_first_chunk_time=None,
+                generative_reasoning_first_chunk_time=None,
                 status_code=AnswerStatusCode(str(llm_status_code)),
             )
 
@@ -170,26 +174,35 @@ async def chat_streaming_generator(
     client_type: NucliaDBClientType,
     origin: str,
     user_query: str,
-    is_json: bool,
+    is_ndjson_stream: bool,
+    metrics: AskMetrics,
 ):
     first = True
+    first_reasoning = True
     status_code = AnswerStatusCode.ERROR.value
     text_answer = ""
+    text_reasoning = ""
     json_object = None
-    metrics = AskMetrics()
     with metrics.time(PREDICT_ANSWER_METRIC):
         async for chunk in predict_response.content:
-            if first:
-                metrics.record_first_chunk_yielded()
-                first = False
-
             yield chunk
-
-            if is_json:
+            if is_ndjson_stream:
                 try:
                     parsed_chunk = GenerativeChunk.model_validate_json(chunk).chunk
+                    if first and isinstance(
+                        parsed_chunk,
+                        (TextGenerativeResponse, JSONGenerativeResponse, StatusGenerativeResponse),
+                    ):
+                        metrics.record_first_chunk_yielded()
+                        first = False
+
                     if isinstance(parsed_chunk, TextGenerativeResponse):
                         text_answer += parsed_chunk.text
+                    elif isinstance(parsed_chunk, ReasoningGenerativeResponse):
+                        if first_reasoning:
+                            metrics.record_first_reasoning_chunk_yielded()
+                            first_reasoning = False
+                        text_reasoning += parsed_chunk.text
                     elif isinstance(parsed_chunk, JSONGenerativeResponse):
                         json_object = parsed_chunk.object
                     elif isinstance(parsed_chunk, StatusGenerativeResponse):
@@ -201,8 +214,11 @@ async def chat_streaming_generator(
                     )
             else:
                 text_answer += chunk.decode()
+                if first:
+                    metrics.record_first_chunk_yielded()
+                    first = False
 
-    if is_json is False and chunk:  # Ensure chunk is not empty before decoding
+    if is_ndjson_stream is False and chunk:  # Ensure chunk is not empty before decoding
         # If response is text the status_code comes at the last chunk of data
         last_chunk = chunk.decode()
         if last_chunk[-1] == "0":
@@ -218,8 +234,10 @@ async def chat_streaming_generator(
         client_type=client_type,
         origin=origin,
         text_answer=text_answer.encode() if json_object is None else json.dumps(json_object).encode(),
+        text_reasoning=text_reasoning if text_reasoning else None,
         generative_answer_time=metrics[PREDICT_ANSWER_METRIC],
         generative_answer_first_chunk_time=metrics.get_first_chunk_time(),
+        generative_reasoning_first_chunk_time=metrics.get_first_reasoning_chunk_time(),
         status_code=AnswerStatusCode(status_code),
     )
 
@@ -232,8 +250,10 @@ def audit_predict_proxy_endpoint(
     client_type: NucliaDBClientType,
     origin: str,
     text_answer: bytes,
+    text_reasoning: Optional[str],
     generative_answer_time: float,
     generative_answer_first_chunk_time: Optional[float],
+    generative_reasoning_first_chunk_time: Optional[float],
     status_code: AnswerStatusCode,
 ):
     maybe_audit_chat(
@@ -250,8 +270,10 @@ def audit_predict_proxy_endpoint(
         query_context_order={},
         model=headers.get(NUCLIA_LEARNING_MODEL_HEADER),
         text_answer=text_answer,
+        text_reasoning=text_reasoning,
         generative_answer_time=generative_answer_time,
         generative_answer_first_chunk_time=generative_answer_first_chunk_time or 0,
+        generative_reasoning_first_chunk_time=generative_reasoning_first_chunk_time,
         rephrase_time=None,
         status_code=status_code,
     )