nucliadb 6.8.1.post4981__py3-none-any.whl → 6.8.1.post4988__py3-none-any.whl

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

nucliadb/search/search/hydrator/images.py

@@ -0,0 +1,126 @@
+# 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: 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 paragraph.kind not in (
+        resources_pb2.Paragraph.TypeParagraph.OCR,
+        resources_pb2.Paragraph.TypeParagraph.INCEPTION,
+    ):
+        return None
+
+    field_id = ParagraphId.from_string(paragraph.key).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,305 @@
+# 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)
+
+            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/writer/api/v1/learning_config.py

@@ -17,7 +17,7 @@
 # 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 fastapi import Request
+from fastapi import Header, Request
 from fastapi_versioning import version
 from nuclia_models.config.proto import ExtractConfig, SplitConfiguration
 
@@ -55,10 +55,11 @@ async def set_configuration(
 @requires_one([NucliaDBRoles.MANAGER, NucliaDBRoles.WRITER])
 @version(1)
 async def patch_configuration(
-    request: Request,
-    kbid: str,
+    request: Request, kbid: str, x_nucliadb_account: str = Header(default="", include_in_schema=False)
 ):
-    return await learning_config_proxy(request, "PATCH", f"/config/{kbid}")
+    return await learning_config_proxy(
+        request, "PATCH", f"/config/{kbid}", headers={"account-id": x_nucliadb_account}
+    )
 
 
 @api.post(