hyperforge-nucliadb-agentic 1.0.0.post64__py3-none-any.whl

hyperforge_nucliadb_agentic/ask/search/prompt.py

@@ -0,0 +1,1298 @@
+import copy
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import cast
+
+import nucliadb_models
+import yaml
+from nucliadb_models.augment import (
+    AugmentedConversationField,
+    AugmentedField,
+    AugmentedFileField,
+    AugmentFields,
+    AugmentParagraph,
+    AugmentParagraphs,
+    AugmentRequest,
+    AugmentResourceFields,
+    AugmentResources,
+)
+from nucliadb_models.common import FieldTypeName
+from nucliadb_models.labels import translate_alias_to_system_label
+from nucliadb_models.search import SCORE_TYPE, FindParagraph
+from nucliadb_protos.resources_pb2 import FieldComputedMetadata
+from nucliadb_sdk import NucliaDBAsync
+from pydantic import BaseModel
+
+from hyperforge_nucliadb_agentic.ask import logger
+from hyperforge_nucliadb_agentic.ask.model import (
+    AugmentedContext,
+    AugmentedTextBlock,
+    ConversationalStrategy,
+    FieldExtensionStrategy,
+    FullResourceStrategy,
+    HierarchyResourceStrategy,
+    Image,
+    ImageRagStrategy,
+    ImageRagStrategyName,
+    MetadataExtensionStrategy,
+    MetadataExtensionType,
+    NeighbouringParagraphsStrategy,
+    PageImageStrategy,
+    ParagraphImageStrategy,
+    PromptContext,
+    PromptContextImages,
+    PromptContextOrder,
+    RagStrategy,
+    RagStrategyName,
+    TableImageStrategy,
+    TextBlockAugmentationType,
+    TextPosition,
+)
+from hyperforge_nucliadb_agentic.ask.search import rpc
+from hyperforge_nucliadb_agentic.ask.search.metrics import Metrics
+from hyperforge_nucliadb_agentic.ask.utils.ids import (
+    FIELD_TYPE_STR_TO_NAME,
+    FieldId,
+    ParagraphId,
+)
+
+TextBlockId = ParagraphId | FieldId
+
+
+class CappedPromptContext:
+    """
+    Class to keep track of the size (in number of characters) of the prompt context
+    and automatically trim data that exceeds the limit when it's being set on the dictionary.
+    """
+
+    def __init__(self, max_size: int | None):
+        self.output: PromptContext = {}
+        self.images: PromptContextImages = {}
+        self.max_size = max_size
+
+    def __setitem__(self, key: str, value: str) -> None:
+        self.output.__setitem__(key, value)
+
+    def __getitem__(self, key: str) -> str:
+        return self.output.__getitem__(key)
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.output
+
+    def __delitem__(self, key: str) -> None:
+        try:
+            self.output.__delitem__(key)
+        except KeyError:
+            pass
+
+    def text_block_ids(self) -> list[str]:
+        return list(self.output.keys())
+
+    @property
+    def size(self) -> int:
+        """
+        Returns the total size of the context in characters.
+        """
+        return sum(len(text) for text in self.output.values())
+
+    def cap(self) -> dict[str, str]:
+        """
+        This method will trim the context to the maximum size if it exceeds it.
+        It will remove text from the most recent entries first, until the size is below the limit.
+        """
+        if self.max_size is None:
+            return self.output
+
+        if self.size <= self.max_size:
+            return self.output
+
+        logger.info("Removing text from context to fit within the max size limit")
+        # Iterate the dictionary in reverse order of insertion
+        for key in reversed(list(self.output.keys())):
+            current_size = self.size
+            if current_size <= self.max_size:
+                break
+            # Remove text from the value
+            text = self.output[key]
+            # If removing the whole text still keeps the total size above the limit, remove it
+            if current_size - len(text) >= self.max_size:
+                del self.output[key]
+            else:
+                # Otherwise, trim the text to fit within the limit
+                excess_size = current_size - self.max_size
+                if excess_size > 0:
+                    trimmed_text = text[:-excess_size]
+                    self.output[key] = trimmed_text
+        return self.output
+
+
+async def default_prompt_context(
+    search_sdk: NucliaDBAsync,
+    context: CappedPromptContext,
+    kbid: str,
+    ordered_paragraphs: list[FindParagraph],
+) -> None:
+    """
+    - Updates context (which is an ordered dict of text_block_id -> context_text).
+    - text_block_id is typically the paragraph id, but has a special value for the
+      user context. (USER_CONTEXT_0, USER_CONTEXT_1, ...)
+    - Paragraphs are inserted in order of relevance, by increasing `order` field
+      of the find result paragraphs.
+    - User context is inserted first, in order of appearance.
+    - Using an dict prevents from duplicates pulled in through conversation expansion.
+    """
+
+    conversations = []
+
+    for paragraph in ordered_paragraphs:
+        context[paragraph.id] = _clean_paragraph_text(paragraph)
+
+        # If the paragraph is a conversation and it matches semantically, we
+        # assume we have matched with the question, therefore try to include the
+        # answer to the context by pulling the next few messages of the
+        # conversation field
+        rid, field_type, field_id, mident = paragraph.id.split("/")[:4]
+        # FIXME: a semantic paragraph can have reranker score. Once we
+        # refactor and have access to the score history, we can fix this
+        if field_type == "c" and paragraph.score_type in (
+            SCORE_TYPE.VECTOR,
+            SCORE_TYPE.BOTH,
+        ):
+            conversations.append(f"{rid}/{field_type}/{field_id}/{mident}")
+
+    augment = AugmentRequest(
+        fields=[
+            AugmentFields(
+                given=[id for id in conversations],
+                conversation_answer_or_messages_after=True,
+            ),
+        ]
+    )
+    augmented = await rpc.augment(search_sdk, kbid, augment)
+
+    for id in conversations:
+        conversation_id = FieldId.from_string(id)
+
+        augmented_field = augmented.fields.get(conversation_id.full_without_subfield())
+        if augmented_field is None or not isinstance(
+            augmented_field, AugmentedConversationField
+        ):
+            continue
+
+        for message in augmented_field.messages or []:
+            if message.text is None:
+                continue
+
+            message_id = copy.copy(conversation_id)
+            message_id.subfield_id = message.ident
+            pid = ParagraphId(
+                field_id=message_id, paragraph_start=0, paragraph_end=len(message.text)
+            ).full()
+            context[pid] = message.text
+
+
+async def full_resource_prompt_context(
+    search_sdk: NucliaDBAsync,
+    context: CappedPromptContext,
+    kbid: str,
+    ordered_paragraphs: list[FindParagraph],
+    rid: str | None,
+    strategy: FullResourceStrategy,
+    metrics: Metrics,
+    augmented_context: AugmentedContext,
+) -> None:
+    """
+    Algorithm steps:
+        - Collect the list of resources in the results (in order of relevance).
+        - For each resource, collect the extracted text from all its fields and craft the context.
+    Arguments:
+        context: The context to be updated.
+        kbid: The knowledge box id.
+        ordered_paragraphs: The results of the retrieval (find) operation.
+        resource: The resource to be included in the context. This is used only when chatting with a specific resource with no retrieval.
+        strategy: strategy instance containing, for example, the number of full resources to include in the context.
+    """
+    if rid is not None:
+        # The user has specified a resource to be included in the context.
+        ordered_resources = [rid]
+    else:
+        # Collect the list of resources in the results (in order of relevance).
+        ordered_resources = []
+        for paragraph in ordered_paragraphs:
+            rid = parse_text_block_id(paragraph.id).rid
+            if rid not in ordered_resources:
+                skip = False
+                if strategy.apply_to is not None:
+                    # decide whether the resource should be extended or not
+                    for label in strategy.apply_to.exclude:
+                        skip = skip or (
+                            translate_alias_to_system_label(label)
+                            in (paragraph.labels or [])
+                        )
+
+                if not skip:
+                    ordered_resources.append(rid)
+                    # skip when we have enough resource ids
+                    if (
+                        strategy.count is not None
+                        and len(ordered_resources) > strategy.count
+                    ):
+                        break
+
+    ordered_resources = ordered_resources[: strategy.count]
+
+    # For each resource, collect the extracted text from all its fields and
+    # include the title and summary as well
+    augmented = await rpc.augment(
+        search_sdk,
+        kbid,
+        AugmentRequest(
+            resources=[
+                AugmentResources(
+                    given=ordered_resources,
+                    title=True,
+                    summary=True,
+                    fields=AugmentResourceFields(
+                        text=True,
+                        filters=[],
+                    ),
+                )
+            ]
+        ),
+    )
+
+    extracted_texts = {}
+    for rid, resource in augmented.resources.items():
+        if resource.title is not None:
+            field_id = FieldId(rid=rid, type="a", key="title").full()
+            extracted_texts[field_id] = resource.title
+        if resource.summary is not None:
+            field_id = FieldId(rid=rid, type="a", key="summary").full()
+            extracted_texts[field_id] = resource.summary
+
+    for field_id, field in augmented.fields.items():
+        field = cast(AugmentedField, field)
+        if field.text is not None:
+            extracted_texts[field_id] = field.text
+
+    added_fields = set()
+    for field_id, extracted_text in extracted_texts.items():
+        # First off, remove the text block ids from paragraphs that belong to
+        # the same field, as otherwise the context will be duplicated.
+        for tb_id in context.text_block_ids():
+            if tb_id.startswith(field_id):
+                del context[tb_id]
+        # Add the extracted text of each field to the context.
+        context[field_id] = extracted_text
+        augmented_context.fields[field_id] = AugmentedTextBlock(
+            id=field_id,
+            text=extracted_text,
+            augmentation_type=TextBlockAugmentationType.FULL_RESOURCE,
+        )
+
+        added_fields.add(field_id)
+
+    metrics.set("full_resource_ops", len(added_fields))
+
+    if strategy.include_remaining_text_blocks:
+        for paragraph in ordered_paragraphs:
+            pid = cast(ParagraphId, parse_text_block_id(paragraph.id))
+            if pid.field_id.full() not in added_fields:
+                context[paragraph.id] = _clean_paragraph_text(paragraph)
+
+
+async def extend_prompt_context_with_metadata(
+    search_sdk: NucliaDBAsync,
+    context: CappedPromptContext,
+    kbid: str,
+    strategy: MetadataExtensionStrategy,
+    metrics: Metrics,
+    augmented_context: AugmentedContext,
+) -> None:
+    rids: list[str] = []
+    field_ids: list[str] = []
+    text_block_ids: list[TextBlockId] = []
+    for text_block_id in context.text_block_ids():
+        try:
+            tb_id = parse_text_block_id(text_block_id)
+        except ValueError:  # pragma: no cover
+            # Some text block ids are not paragraphs nor fields, so they are skipped
+            # (e.g. USER_CONTEXT_0, when the user provides extra context)
+            continue
+
+        field_id = tb_id if isinstance(tb_id, FieldId) else tb_id.field_id
+
+        text_block_ids.append(tb_id)
+        field_ids.append(field_id.full())
+        rids.append(tb_id.rid)
+
+    if len(text_block_ids) == 0:  # pragma: no cover
+        return
+
+    resource_origin = False
+    resource_extra = False
+    classification_labels = False
+    field_entities = False
+
+    ops = 0
+    if MetadataExtensionType.ORIGIN in strategy.types:
+        ops += 1
+        resource_origin = True
+
+    if MetadataExtensionType.CLASSIFICATION_LABELS in strategy.types:
+        ops += 1
+        classification_labels = True
+
+    if MetadataExtensionType.NERS in strategy.types:
+        ops += 1
+        field_entities = True
+
+    if MetadataExtensionType.EXTRA_METADATA in strategy.types:
+        ops += 1
+        resource_extra = True
+
+    metrics.set("metadata_extension_ops", ops * len(text_block_ids))
+
+    augment_req = AugmentRequest()
+    if resource_origin or resource_extra or classification_labels:
+        augment_req.resources = [
+            AugmentResources(
+                given=rids,
+                origin=resource_origin,
+                extra=resource_extra,
+                classification_labels=classification_labels,
+            )
+        ]
+    if classification_labels or field_entities:
+        augment_req.fields = [
+            AugmentFields(
+                given=field_ids,
+                classification_labels=classification_labels,
+                entities=field_entities,
+            )
+        ]
+
+    if augment_req.resources is None and augment_req.fields is None:
+        # nothing to augment
+        return
+
+    augmented = await rpc.augment(search_sdk, kbid, augment_req)
+
+    for tb_id in text_block_ids:
+        field_id = tb_id if isinstance(tb_id, FieldId) else tb_id.field_id
+
+        resource = augmented.resources.get(tb_id.rid)
+        field = augmented.fields.get(field_id.full())
+
+        if resource is not None:
+            if resource.origin is not None:
+                text = context.output.pop(tb_id.full())
+                extended_text = (
+                    text
+                    + f"\n\nDOCUMENT METADATA AT ORIGIN:\n{to_yaml(resource.origin)}"
+                )
+                context[tb_id.full()] = extended_text
+                augmented_context.paragraphs[tb_id.full()] = AugmentedTextBlock(
+                    id=tb_id.full(),
+                    text=extended_text,
+                    parent=tb_id.full(),
+                    augmentation_type=TextBlockAugmentationType.METADATA_EXTENSION,
+                )
+
+            if resource.extra is not None:
+                text = context.output.pop(tb_id.full())
+                extended_text = (
+                    text + f"\n\nDOCUMENT EXTRA METADATA:\n{to_yaml(resource.extra)}"
+                )
+                context[tb_id.full()] = extended_text
+                augmented_context.paragraphs[tb_id.full()] = AugmentedTextBlock(
+                    id=tb_id.full(),
+                    text=extended_text,
+                    parent=tb_id.full(),
+                    augmentation_type=TextBlockAugmentationType.METADATA_EXTENSION,
+                )
+
+        if tb_id.full() in context:
+            if (resource is not None and resource.classification_labels) or (
+                field is not None and field.classification_labels
+            ):
+                text = context.output.pop(tb_id.full())
+
+                labels_text = "DOCUMENT CLASSIFICATION LABELS:"
+                if resource is not None and resource.classification_labels:
+                    for labelset, labels in resource.classification_labels.items():
+                        for label in labels:
+                            labels_text += f"\n - {label} ({labelset})"
+
+                if field is not None and field.classification_labels:
+                    for labelset, labels in field.classification_labels.items():
+                        for label in labels:
+                            labels_text += f"\n - {label} ({labelset})"
+
+                extended_text = text + "\n\n" + labels_text
+
+                context[tb_id.full()] = extended_text
+                augmented_context.paragraphs[tb_id.full()] = AugmentedTextBlock(
+                    id=tb_id.full(),
+                    text=extended_text,
+                    parent=tb_id.full(),
+                    augmentation_type=TextBlockAugmentationType.METADATA_EXTENSION,
+                )
+
+            if field is not None and field.entities:
+                ners = field.entities
+
+                text = context.output.pop(tb_id.full())
+
+                ners_text = "DOCUMENT NAMED ENTITIES (NERs):"
+                for family, tokens in ners.items():
+                    ners_text += f"\n - {family}:"
+                    for token in sorted(list(tokens)):
+                        ners_text += f"\n   - {token}"
+
+                extended_text = text + "\n\n" + ners_text
+
+                context[tb_id.full()] = extended_text
+                augmented_context.paragraphs[tb_id.full()] = AugmentedTextBlock(
+                    id=tb_id.full(),
+                    text=extended_text,
+                    parent=tb_id.full(),
+                    augmentation_type=TextBlockAugmentationType.METADATA_EXTENSION,
+                )
+
+
+def parse_text_block_id(text_block_id: str) -> TextBlockId:
+    try:
+        # Typically, the text block id is a paragraph id
+        return ParagraphId.from_string(text_block_id)
+    except ValueError:
+        # When we're doing `full_resource` or `hierarchy` strategies,the text block id
+        # is a field id
+        return FieldId.from_string(text_block_id)
+
+
+def to_yaml(obj: BaseModel) -> str:
+    # FIXME: this dumps enums REALLY poorly, e.g.,
+    # `!!python/object/apply:nucliadb_models.metadata.Source\n- WEB` for
+    # Source.WEB instead of `WEB`
+    return yaml.dump(
+        obj.model_dump(exclude_none=True, exclude_defaults=True, exclude_unset=True),
+        default_flow_style=False,
+        indent=2,
+        sort_keys=True,
+    )
+
+
+async def field_extension_prompt_context(
+    search_sdk: NucliaDBAsync,
+    context: CappedPromptContext,
+    kbid: str,
+    ordered_paragraphs: list[FindParagraph],
+    strategy: FieldExtensionStrategy,
+    metrics: Metrics,
+    augmented_context: AugmentedContext,
+) -> None:
+    """
+    Algorithm steps:
+        - Collect the list of resources in the results (in order of relevance).
+        - For each resource, collect the extracted text from all its fields.
+        - Add the extracted text of each field to the beginning of the context.
+        - Add the extracted text of each paragraph to the end of the context.
+    """
+    ordered_resources = []
+    for paragraph in ordered_paragraphs:
+        resource_uuid = ParagraphId.from_string(paragraph.id).rid
+        if resource_uuid not in ordered_resources:
+            ordered_resources.append(resource_uuid)
+
+    resource_title = False
+    resource_summary = False
+    filters: list[
+        nucliadb_models.filters.Field | nucliadb_models.filters.Generated  # type: ignore
+    ] = []
+    # this strategy exposes a way to access resource title and summary using a
+    # field id. However, as they are resource properties, we must request it as
+    # that
+    for name in strategy.fields:
+        if name == "a/title":
+            resource_title = True
+        elif name == "a/summary":
+            resource_summary = True
+        else:
+            # model already enforces type/name format
+            field_type, field_name = name.split("/")
+            filters.append(
+                nucliadb_models.filters.Field(  # type: ignore
+                    type=FIELD_TYPE_STR_TO_NAME[field_type], name=field_name or None
+                )
+            )
+
+    for da_prefix in strategy.data_augmentation_field_prefixes:
+        filters.append(
+            nucliadb_models.filters.Generated(by="data-augmentation", da_task=da_prefix)  # type: ignore
+        )
+
+    augmented = await rpc.augment(
+        search_sdk,
+        kbid,
+        AugmentRequest(
+            resources=[
+                AugmentResources(
+                    given=ordered_resources,
+                    title=resource_title,
+                    summary=resource_summary,
+                    fields=AugmentResourceFields(
+                        text=True,
+                        filters=filters,
+                    )
+                    if len(filters) > 0
+                    else None,
+                )
+            ]
+        ),
+    )
+
+    # REVIEW(decoupled-ask): we don't have the field count anymore, is this good enough?
+    metrics.set("field_extension_ops", len(ordered_resources))
+
+    extracted_texts = {}
+    # now we need to expose title and summary as fields again, so it gets
+    # consistent with the view we are providing in the API
+    for rid, augmented_resource in augmented.resources.items():
+        if augmented_resource.title:
+            extracted_texts[f"{rid}/a/title"] = augmented_resource.title
+        if augmented_resource.summary:
+            extracted_texts[f"{rid}/a/summary"] = augmented_resource.summary
+
+    for fid, augmented_field in augmented.fields.items():
+        if augmented_field is None or augmented_field.text is None:  # pragma: no cover
+            continue
+        extracted_texts[fid] = augmented_field.text
+
+    for fid, extracted_text in extracted_texts.items():
+        # First off, remove the text block ids from paragraphs that belong to
+        # the same field, as otherwise the context will be duplicated.
+        for tb_id in context.text_block_ids():
+            if tb_id.startswith(fid):
+                del context[tb_id]
+        # Add the extracted text of each field to the beginning of the context.
+        if fid not in context:
+            context[fid] = extracted_text
+            augmented_context.fields[fid] = AugmentedTextBlock(
+                id=fid,
+                text=extracted_text,
+                augmentation_type=TextBlockAugmentationType.FIELD_EXTENSION,
+            )
+
+    # Add the extracted text of each paragraph to the end of the context.
+    for paragraph in ordered_paragraphs:
+        if paragraph.id not in context:
+            context[paragraph.id] = _clean_paragraph_text(paragraph)
+
+
+async def neighbouring_paragraphs_prompt_context(
+    search_sdk: NucliaDBAsync,
+    context: CappedPromptContext,
+    kbid: str,
+    ordered_text_blocks: list[FindParagraph],
+    strategy: NeighbouringParagraphsStrategy,
+    metrics: Metrics,
+    augmented_context: AugmentedContext,
+) -> None:
+    """
+    This function will get the paragraph texts and then craft a context with the neighbouring paragraphs of the
+    paragraphs in the ordered_paragraphs list.
+    """
+    retrieved_paragraphs_ids = [
+        ParagraphId.from_string(text_block.id) for text_block in ordered_text_blocks
+    ]
+
+    augmented = await rpc.augment(
+        search_sdk,
+        kbid,
+        AugmentRequest(
+            paragraphs=[
+                AugmentParagraphs(
+                    given=[
+                        AugmentParagraph(id=pid.full())
+                        for pid in retrieved_paragraphs_ids
+                    ],
+                    text=True,
+                    neighbours_before=strategy.before,
+                    neighbours_after=strategy.after,
+                )
+            ]
+        ),
+    )
+
+    for pid in retrieved_paragraphs_ids:
+        paragraph = augmented.paragraphs.get(pid.full())
+        if paragraph is None:
+            continue
+
+        ptext = paragraph.text or ""
+        if ptext and pid.full() not in context:
+            context[pid.full()] = ptext
+
+        # Now add the neighbouring paragraphs
+        neighbour_ids = [
+            *(paragraph.neighbours_before or []),
+            *(paragraph.neighbours_after or []),
+        ]
+        for npid in neighbour_ids:
+            neighbour = augmented.paragraphs.get(npid)
+            assert neighbour is not None, (
+                "augment should never return dangling paragraph references"
+            )
+
+            if (
+                ParagraphId.from_string(npid) in retrieved_paragraphs_ids
+                or npid in context
+            ):
+                # already added
+                continue
+
+            ntext = neighbour.text
+            if not ntext:
+                continue
+
+            context[npid] = ntext
+            augmented_context.paragraphs[npid] = AugmentedTextBlock(
+                id=npid,
+                text=ntext,
+                position=neighbour.position,
+                parent=pid.full(),
+                augmentation_type=TextBlockAugmentationType.NEIGHBOURING_PARAGRAPHS,
+            )
+
+    metrics.set("neighbouring_paragraphs_ops", len(augmented_context.paragraphs))
+
+
+def get_text_position(
+    paragraph_id: ParagraphId, index: int, field_metadata: FieldComputedMetadata
+) -> TextPosition | None:
+    if paragraph_id.field_id.subfield_id:
+        metadata = field_metadata.split_metadata[paragraph_id.field_id.subfield_id]
+    else:
+        metadata = field_metadata.metadata
+    try:
+        pmetadata = metadata.paragraphs[index]
+    except IndexError:
+        return None
+    page_number = None
+    if pmetadata.HasField("page"):
+        page_number = pmetadata.page.page
+    return TextPosition(
+        page_number=page_number,
+        index=index,
+        start=pmetadata.start,
+        end=pmetadata.end,
+        start_seconds=list(pmetadata.start_seconds),
+        end_seconds=list(pmetadata.end_seconds),
+    )
+
+
+def get_neighbouring_indices(
+    index: int, before: int, after: int, field_pids: list[ParagraphId]
+) -> list[int]:
+    lb_index = max(0, index - before)
+    ub_index = min(len(field_pids), index + after + 1)
+    return list(range(lb_index, index)) + list(range(index + 1, ub_index))
+
+
+async def conversation_prompt_context(
+    search_sdk: NucliaDBAsync,
+    reader_sdk: NucliaDBAsync,
+    context: CappedPromptContext,
+    kbid: str,
+    ordered_paragraphs: list[FindParagraph],
+    strategy: ConversationalStrategy,
+    visual_llm: bool,
+    metrics: Metrics,
+    augmented_context: AugmentedContext,
+):
+    analyzed_fields: set[str] = set()
+    ops = 0
+
+    conversation_paragraphs = []
+    for paragraph in ordered_paragraphs:
+        if paragraph.id not in context:
+            context[paragraph.id] = _clean_paragraph_text(paragraph)
+
+        parent_paragraph_id = ParagraphId.from_string(paragraph.id)
+
+        if (
+            parent_paragraph_id.field_id.type
+            != FieldTypeName.CONVERSATION.abbreviation()
+        ):
+            # conversational strategy only applies to conversation fields
+            continue
+
+        field_unique_id = parent_paragraph_id.field_id.full_without_subfield()
+        if field_unique_id in analyzed_fields:
+            continue
+
+        conversation_paragraphs.append((parent_paragraph_id, paragraph))
+        analyzed_fields.add(field_unique_id)
+
+    # augment conversation paragraphs
+
+    if strategy.full:
+        full_conversation = True
+        max_conversation_messages = None
+    else:
+        full_conversation = False
+        max_conversation_messages = strategy.max_messages
+
+    augment = AugmentRequest(
+        fields=[
+            AugmentFields(
+                given=[
+                    paragraph_id.field_id.full()
+                    for paragraph_id, _ in conversation_paragraphs
+                ],
+                full_conversation=full_conversation,
+                max_conversation_messages=max_conversation_messages,
+                conversation_text_attachments=strategy.attachments_text,
+                conversation_image_attachments=strategy.attachments_images,
+            )
+        ]
+    )
+    augmented = await rpc.augment(search_sdk, kbid, augment)
+
+    attachments: dict[ParagraphId, list[FieldId]] = {}
+    for parent_paragraph_id, paragraph in conversation_paragraphs:
+        fid = parent_paragraph_id.field_id
+        field = augmented.fields.get(fid.full_without_subfield())
+        if field is not None:
+            field = cast(AugmentedConversationField, field)
+            for _message in field.messages or []:
+                ops += 1
+                if not _message.text:
+                    continue
+
+                text = _message.text
+                pid = ParagraphId(
+                    field_id=FieldId(
+                        rid=fid.rid,
+                        type=fid.type,
+                        key=fid.key,
+                        subfield_id=_message.ident,
+                    ),
+                    paragraph_start=0,
+                    paragraph_end=len(text),
+                ).full()
+                if pid in context:
+                    continue
+                context[pid] = text
+
+                attachments.setdefault(parent_paragraph_id, []).extend(
+                    [
+                        FieldId.from_string(attachment_id)
+                        for attachment_id in field.attachments or []
+                    ]
+                )
+                augmented_context.paragraphs[pid] = AugmentedTextBlock(
+                    id=pid,
+                    text=text,
+                    parent=paragraph.id,
+                    augmentation_type=TextBlockAugmentationType.CONVERSATION,
+                )
+
+    # augment attachments
+
+    if strategy.attachments_text or (
+        (strategy.attachments_images and visual_llm) and len(attachments) > 0
+    ):
+        augment = AugmentRequest(
+            fields=[
+                AugmentFields(
+                    given=[
+                        id.full()
+                        for paragraph_attachments in attachments.values()
+                        for id in paragraph_attachments
+                    ],
+                    text=strategy.attachments_text,
+                    file_thumbnail=(strategy.attachments_images and visual_llm),
+                )
+            ]
+        )
+        augmented = await rpc.augment(search_sdk, kbid, augment)
+
+        for parent_paragraph_id, paragraph_attachments in attachments.items():
+            for attachment_id in paragraph_attachments:
+                attachment_field = augmented.fields.get(attachment_id.full())
+
+                if attachment_field is None:
+                    continue
+
+                if strategy.attachments_text and attachment_field.text:
+                    ops += 1
+
+                    pid = f"{attachment_id.full_without_subfield()}/0-{len(attachment_field.text)}"
+                    if pid not in context:
+                        text = f"Attachment {attachment_id.key}: {attachment_field.text}\n\n"
+                        context[pid] = text
+                        augmented_context.paragraphs[pid] = AugmentedTextBlock(
+                            id=pid,
+                            text=text,
+                            parent=parent_paragraph_id.full(),
+                            augmentation_type=TextBlockAugmentationType.CONVERSATION,
+                        )
+
+                if (
+                    (strategy.attachments_images and visual_llm)
+                    and isinstance(attachment_field, AugmentedFileField)
+                    and attachment_field.thumbnail_image
+                ):
+                    ops += 1
+
+                    image = await rpc.download_image(
+                        reader_sdk,
+                        kbid,
+                        attachment_id,
+                        attachment_field.thumbnail_image,
+                        # We assume the thumbnail is always generated as JPEG by Nuclia processing
+                        mime_type="image/jpeg",
+                    )
+                    if image is not None:
+                        pid = f"{attachment_id.rid}/f/{attachment_id.key}/0-0"
+                        context.images[pid] = image
+
+    metrics.set("conversation_ops", ops)
+
+
+async def hierarchy_prompt_context(
+    search_sdk: NucliaDBAsync,
+    context: CappedPromptContext,
+    kbid: str,
+    ordered_paragraphs: list[FindParagraph],
+    strategy: HierarchyResourceStrategy,
+    metrics: Metrics,
+    augmented_context: AugmentedContext,
+) -> None:
+    """
+    This function will get the paragraph texts (possibly with extra characters, if extra_characters > 0) and then
+    craft a context with all paragraphs of the same resource grouped together. Moreover, on each group of paragraphs,
+    it includes the resource title and summary so that the LLM can have a better understanding of the context.
+    """
+    paragraphs_extra_characters = max(strategy.count, 0)
+    # Make a copy of the ordered paragraphs to avoid modifying the original list, which is returned
+    # in the response to the user
+    ordered_paragraphs_copy = copy.deepcopy(ordered_paragraphs)
+    resources: dict[str, ExtraCharsParagraph] = {}
+
+    # Iterate paragraphs to get extended text
+    paragraphs_to_augment = []
+    for paragraph in ordered_paragraphs_copy:
+        paragraph_id = ParagraphId.from_string(paragraph.id)
+        rid = paragraph_id.rid
+
+        if paragraphs_extra_characters > 0:
+            paragraph_id.paragraph_end += paragraphs_extra_characters
+
+        paragraphs_to_augment.append(paragraph_id)
+
+        if rid not in resources:
+            # Get the title and the summary of the resource
+            title_paragraph_id = ParagraphId(
+                field_id=FieldId(
+                    rid=rid,
+                    type="a",
+                    key="title",
+                ),
+                paragraph_start=0,
+                paragraph_end=500,
+            )
+            summary_paragraph_id = ParagraphId(
+                field_id=FieldId(
+                    rid=rid,
+                    type="a",
+                    key="summary",
+                ),
+                paragraph_start=0,
+                paragraph_end=1000,
+            )
+            paragraphs_to_augment.append(title_paragraph_id)
+            paragraphs_to_augment.append(summary_paragraph_id)
+
+            resources[rid] = ExtraCharsParagraph(
+                title=title_paragraph_id,
+                summary=summary_paragraph_id,
+                paragraphs=[(paragraph, paragraph_id)],
+            )
+        else:
+            resources[rid].paragraphs.append((paragraph, paragraph_id))
+
+    metrics.set("hierarchy_ops", len(resources))
+
+    augmented = await rpc.augment(
+        search_sdk,
+        kbid,
+        AugmentRequest(
+            paragraphs=[
+                AugmentParagraphs(
+                    given=[
+                        AugmentParagraph(id=paragraph_id.full())
+                        for paragraph_id in paragraphs_to_augment
+                    ],
+                    text=True,
+                )
+            ]
+        ),
+    )
+
+    augmented_paragraphs = set()
+
+    # Modify the first paragraph of each resource to include the title and summary of the resource, as well as the
+    # extended paragraph text of all the paragraphs in the resource.
+    for values in resources.values():
+        augmented_title = augmented.paragraphs.get(values.title.full())
+        if augmented_title:
+            title_text = augmented_title.text or ""
+        else:
+            title_text = ""
+
+        augmented_summary = augmented.paragraphs.get(values.summary.full())
+        if augmented_summary:
+            summary_text = augmented_summary.text or ""
+        else:
+            summary_text = ""
+
+        first_paragraph = None
+        text_with_hierarchy = ""
+        for paragraph, paragraph_id in values.paragraphs:
+            augmented_paragraph = augmented.paragraphs.get(paragraph_id.full())
+            if augmented_paragraph:
+                extended_paragraph_text = augmented_paragraph.text or ""
+            else:
+                extended_paragraph_text = ""
+
+            if first_paragraph is None:
+                first_paragraph = paragraph
+            text_with_hierarchy += (
+                "\n EXTRACTED BLOCK: \n " + extended_paragraph_text + " \n\n "
+            )
+            # All paragraphs of the resource are cleared except the first one, which will be the
+            # one containing the whole hierarchy information
+            paragraph.text = ""
+
+        if first_paragraph is not None:
+            # The first paragraph is the only one holding the hierarchy information
+            first_paragraph.text = f"DOCUMENT: {title_text} \n SUMMARY: {summary_text} \n RESOURCE CONTENT: {text_with_hierarchy}"
+            augmented_paragraphs.add(first_paragraph.id)
+
+    # Now that the paragraphs have been modified, we can add them to the context
+    for paragraph in ordered_paragraphs_copy:
+        if paragraph.text == "":
+            # Skip paragraphs that were cleared in the hierarchy expansion
+            continue
+        paragraph_text = _clean_paragraph_text(paragraph)
+        context[paragraph.id] = paragraph_text
+        if paragraph.id in augmented_paragraphs:
+            pid = ParagraphId.from_string(paragraph.id)
+            augmented_context.paragraphs[pid.full()] = AugmentedTextBlock(
+                id=pid.full(),
+                text=paragraph_text,
+                augmentation_type=TextBlockAugmentationType.HIERARCHY,
+            )
+    return
+
+
+class PromptContextBuilder:
+    """
+    Builds the context for the LLM prompt.
+    """
+
+    def __init__(
+        self,
+        search_sdk: NucliaDBAsync,
+        reader_sdk: NucliaDBAsync,
+        kbid: str,
+        ordered_paragraphs: list[FindParagraph],
+        resource: str | None = None,
+        user_context: list[str] | None = None,
+        user_image_context: list[Image] | None = None,
+        strategies: Sequence[RagStrategy] | None = None,
+        image_strategies: Sequence[ImageRagStrategy] | None = None,
+        max_context_characters: int | None = None,
+        visual_llm: bool = False,
+        query_image: Image | None = None,
+        metrics: Metrics = Metrics("prompt_context_builder"),
+    ):
+        self.search_sdk = search_sdk
+        self.reader_sdk = reader_sdk
+        self.kbid = kbid
+        self.ordered_paragraphs = ordered_paragraphs
+        self.resource = resource
+        self.user_context = user_context
+        self.user_image_context = user_image_context
+        self.strategies = strategies
+        self.image_strategies = image_strategies
+        self.max_context_characters = max_context_characters
+        self.visual_llm = visual_llm
+        self.metrics = metrics
+        self.query_image = query_image
+        self.augmented_context = AugmentedContext(paragraphs={}, fields={})
+
+    def prepend_user_context(self, context: CappedPromptContext):
+        # Chat extra context passed by the user is the most important, therefore
+        # it is added first, followed by the found text blocks in order of relevance
+        for i, text_block in enumerate(self.user_context or []):
+            context[f"USER_CONTEXT_{i}"] = text_block
+        # Add the query image as part of the image context
+        if self.query_image is not None:
+            context.images["QUERY_IMAGE"] = self.query_image
+        else:
+            for i, image in enumerate(self.user_image_context or []):
+                context.images[f"USER_IMAGE_CONTEXT_{i}"] = image
+
+    async def build(
+        self,
+    ) -> tuple[
+        PromptContext, PromptContextOrder, PromptContextImages, AugmentedContext
+    ]:
+        ccontext = CappedPromptContext(max_size=self.max_context_characters)
+        self.prepend_user_context(ccontext)
+        await self._build_context(ccontext)
+        if self.visual_llm and not self.query_image:
+            await self._build_context_images(ccontext)
+        context = ccontext.cap()
+        context_images = ccontext.images
+        context_order = {
+            text_block_id: order for order, text_block_id in enumerate(context.keys())
+        }
+        return context, context_order, context_images, self.augmented_context
+
+    async def _build_context_images(self, context: CappedPromptContext) -> None:
+        ops = 0
+        if self.image_strategies is None or len(self.image_strategies) == 0:
+            # Nothing to do
+            return
+        page_image_strategy: PageImageStrategy | None = None
+        max_page_images = 5
+        table_image_strategy: TableImageStrategy | None = None
+        paragraph_image_strategy: ParagraphImageStrategy | None = None
+        for strategy in self.image_strategies:
+            if strategy.name == ImageRagStrategyName.PAGE_IMAGE:
+                if page_image_strategy is None:
+                    page_image_strategy = cast(PageImageStrategy, strategy)
+                    if page_image_strategy.count is not None:
+                        max_page_images = page_image_strategy.count
+            elif strategy.name == ImageRagStrategyName.TABLES:
+                if table_image_strategy is None:
+                    table_image_strategy = cast(TableImageStrategy, strategy)
+            elif strategy.name == ImageRagStrategyName.PARAGRAPH_IMAGE:
+                if paragraph_image_strategy is None:
+                    paragraph_image_strategy = cast(ParagraphImageStrategy, strategy)
+            else:  # pragma: no cover
+                logger.warning(
+                    "Unknown image strategy",
+                    extra={"strategy": strategy.name, "kbid": self.kbid},
+                )
+        page_images_added = 0
+        for paragraph in self.ordered_paragraphs:
+            pid = ParagraphId.from_string(paragraph.id)
+            paragraph_page_number = get_paragraph_page_number(paragraph)
+            if (
+                page_image_strategy is not None
+                and page_images_added < max_page_images
+                and paragraph_page_number is not None
+            ):
+                # page_image_id: rid/f/myfield/0
+                page_image_id = "/".join(
+                    [pid.field_id.full(), str(paragraph_page_number)]
+                )
+                if page_image_id not in context.images:
+                    image = await rpc.download_image(
+                        self.reader_sdk,
+                        self.kbid,
+                        pid.field_id,
+                        f"generated/extracted_images_{paragraph_page_number}.png",
+                        mime_type="image/png",
+                    )
+                    if image is not None:
+                        ops += 1
+                        context.images[page_image_id] = image
+                        page_images_added += 1
+                    else:
+                        logger.warning(
+                            "Could not retrieve image for paragraph from storage",
+                            extra={
+                                "kbid": self.kbid,
+                                "paragraph": pid.full(),
+                                "page_number": paragraph_page_number,
+                            },
+                        )
+
+            add_table = table_image_strategy is not None and paragraph.is_a_table
+            add_paragraph = (
+                paragraph_image_strategy is not None and not paragraph.is_a_table
+            )
+            if (add_table or add_paragraph) and (
+                paragraph.reference is not None and paragraph.reference != ""
+            ):
+                pimage = await rpc.download_image(
+                    self.reader_sdk,
+                    self.kbid,
+                    pid.field_id,
+                    f"generated/{paragraph.reference}",
+                    mime_type="image/png",
+                )
+                if pimage is not None:
+                    ops += 1
+                    context.images[paragraph.id] = pimage
+                else:
+                    logger.warning(
+                        "Could not retrieve image for paragraph from storage",
+                        extra={
+                            "kbid": self.kbid,
+                            "paragraph": pid.full(),
+                            "reference": paragraph.reference,
+                        },
+                    )
+        self.metrics.set("image_ops", ops)
+
+    async def _build_context(self, context: CappedPromptContext) -> None:
+        if self.strategies is None or len(self.strategies) == 0:
+            # When no strategy is specified, use the default one
+            await default_prompt_context(
+                self.search_sdk, context, self.kbid, self.ordered_paragraphs
+            )
+            return
+        else:
+            # Add the paragraphs to the context and then apply the strategies
+            for paragraph in self.ordered_paragraphs:
+                context[paragraph.id] = _clean_paragraph_text(paragraph)
+
+        strategies_not_handled_here = [
+            RagStrategyName.PREQUERIES,
+            RagStrategyName.GRAPH,
+        ]
+
+        full_resource: FullResourceStrategy | None = None
+        hierarchy: HierarchyResourceStrategy | None = None
+        neighbouring_paragraphs: NeighbouringParagraphsStrategy | None = None
+        field_extension: FieldExtensionStrategy | None = None
+        metadata_extension: MetadataExtensionStrategy | None = None
+        conversational_strategy: ConversationalStrategy | None = None
+        for strategy in self.strategies:
+            if strategy.name == RagStrategyName.FIELD_EXTENSION:
+                field_extension = cast(FieldExtensionStrategy, strategy)
+            elif strategy.name == RagStrategyName.CONVERSATION:
+                conversational_strategy = cast(ConversationalStrategy, strategy)
+            elif strategy.name == RagStrategyName.FULL_RESOURCE:
+                full_resource = cast(FullResourceStrategy, strategy)
+                if self.resource:  # pragma: no cover
+                    # When the retrieval is scoped to a specific resource
+                    # the full resource strategy only includes that resource
+                    full_resource.count = 1
+            elif strategy.name == RagStrategyName.HIERARCHY:
+                hierarchy = cast(HierarchyResourceStrategy, strategy)
+            elif strategy.name == RagStrategyName.NEIGHBOURING_PARAGRAPHS:
+                neighbouring_paragraphs = cast(NeighbouringParagraphsStrategy, strategy)
+            elif strategy.name == RagStrategyName.METADATA_EXTENSION:
+                metadata_extension = cast(MetadataExtensionStrategy, strategy)
+            elif strategy.name not in strategies_not_handled_here:  # pragma: no cover
+                # Prequeries and graph are not handled here
+                logger.warning(
+                    "Unknown rag strategy",
+                    extra={"strategy": strategy.name, "kbid": self.kbid},
+                )
+
+        if full_resource:
+            # When full resoure is enabled, only metadata extension is allowed.
+            await full_resource_prompt_context(
+                self.search_sdk,
+                context,
+                self.kbid,
+                self.ordered_paragraphs,
+                self.resource,
+                full_resource,
+                self.metrics,
+                self.augmented_context,
+            )
+            if metadata_extension:
+                await extend_prompt_context_with_metadata(
+                    self.search_sdk,
+                    context,
+                    self.kbid,
+                    metadata_extension,
+                    self.metrics,
+                    self.augmented_context,
+                )
+            return
+
+        if hierarchy:
+            await hierarchy_prompt_context(
+                self.search_sdk,
+                context,
+                self.kbid,
+                self.ordered_paragraphs,
+                hierarchy,
+                self.metrics,
+                self.augmented_context,
+            )
+        if neighbouring_paragraphs:
+            await neighbouring_paragraphs_prompt_context(
+                self.search_sdk,
+                context,
+                self.kbid,
+                self.ordered_paragraphs,
+                neighbouring_paragraphs,
+                self.metrics,
+                self.augmented_context,
+            )
+        if field_extension:
+            await field_extension_prompt_context(
+                self.search_sdk,
+                context,
+                self.kbid,
+                self.ordered_paragraphs,
+                field_extension,
+                self.metrics,
+                self.augmented_context,
+            )
+        if conversational_strategy:
+            await conversation_prompt_context(
+                self.search_sdk,
+                self.reader_sdk,
+                context,
+                self.kbid,
+                self.ordered_paragraphs,
+                conversational_strategy,
+                self.visual_llm,
+                self.metrics,
+                self.augmented_context,
+            )
+        if metadata_extension:
+            await extend_prompt_context_with_metadata(
+                self.search_sdk,
+                context,
+                self.kbid,
+                metadata_extension,
+                self.metrics,
+                self.augmented_context,
+            )
+
+
+def get_paragraph_page_number(paragraph: FindParagraph) -> int | None:
+    if not paragraph.page_with_visual:
+        return None
+    if paragraph.position is None:
+        return None
+    return paragraph.position.page_number
+
+
+@dataclass
+class ExtraCharsParagraph:
+    title: ParagraphId
+    summary: ParagraphId
+    paragraphs: list[tuple[FindParagraph, ParagraphId]]
+
+
+def _clean_paragraph_text(paragraph: FindParagraph) -> str:
+    text = paragraph.text.strip()
+    # Do not send highlight marks on prompt context
+    text = text.replace("<mark>", "").replace("</mark>", "")
+    return text