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

hyperforge_nucliadb_agentic/ask/search/highlight.py

@@ -0,0 +1,93 @@
+import logging
+import re
+import string
+
+logger = logging.getLogger(__name__)
+PRE_WORD = string.punctuation + " "
+
+
+def highlight_paragraph(
+    text: str, words: list[str] | None = None, ematches: list[str] | None = None
+) -> str:
+    """
+    Highlight `text` with <mark></mark> tags around the words in `words` and `ematches`.
+
+    Parameters:
+    - text: The text to highlight.
+    - words: A list of words to highlight.
+    - ematches: A list of exact matches to highlight.
+
+    Returns:
+    - The highlighted text.
+    """
+    REGEX_TEMPLATE = r"(^|\s)({text})(\s|$)"
+    text_lower = text.lower()
+
+    marks = [0] * (len(text_lower) + 1)
+    ematches = ematches or []
+    for quote in ematches:
+        quote_regex = REGEX_TEMPLATE.format(text=re.escape(quote.lower()))
+        try:
+            for match in re.finditer(quote_regex, text_lower):
+                start, end = match.span(2)
+                marks[start] = 1
+                marks[end] = 2
+        except re.error:
+            logger.warning(
+                f"Regex errors while highlighting text. Regex: {quote_regex}"
+            )
+            continue
+
+    words = words or []
+    for word in words:
+        word_regex = REGEX_TEMPLATE.format(text=re.escape(word.lower()))
+        try:
+            for match in re.finditer(word_regex, text_lower):
+                start, end = match.span(2)
+                if marks[start] == 0 and marks[end] == 0:
+                    marks[start] = 1
+                    marks[end] = 2
+        except re.error:
+            logger.warning(f"Regex errors while highlighting text. Regex: {word_regex}")
+            continue
+
+    new_text = ""
+    actual = 0
+    mod = 0
+    skip = False
+
+    length = len(text)
+
+    for index, pos in enumerate(marks):
+        if skip:
+            skip = False
+            continue
+        if (index - mod) >= length:
+            char_pos = ""
+        else:
+            begining = True
+            if index > 0 and text[index - mod - 1] not in PRE_WORD:
+                begining = False
+            char_pos = text[index - mod]
+            if text[index - mod].lower() != text_lower[index]:
+                # May be incorrect positioning due to unicode lower
+                mod += 1
+                skip = True
+        if pos == 1 and actual == 0 and begining:
+            new_text += "<mark>"
+            new_text += char_pos
+            actual = 1
+        elif pos == 2 and actual == 1:
+            new_text += "</mark>"
+            new_text += char_pos
+            actual = 0
+        elif pos == 1 and actual > 0:
+            new_text += char_pos
+            actual += 1
+        elif pos == 2 and actual > 1:
+            new_text += char_pos
+            actual -= 1
+        else:
+            new_text += char_pos
+
+    return new_text

hyperforge_nucliadb_agentic/ask/search/hydrator.py

@@ -0,0 +1,29 @@
+from nucliadb_models.common import FieldTypeName
+from nucliadb_models.resource import ExtractedDataTypeName
+from nucliadb_models.search import ResourceProperties
+from pydantic import BaseModel
+
+
+class ResourceHydrationOptions(BaseModel):
+    """
+    Options for hydrating resources.
+    """
+
+    show: list[ResourceProperties] = []
+    extracted: list[ExtractedDataTypeName] = []
+    field_type_filter: list[FieldTypeName] = []
+
+
+class TextBlockHydrationOptions(BaseModel):
+    """
+    Options for hydrating text blocks (aka paragraphs).
+    """
+
+    # whether to highlight the text block with `<mark>...</mark>` tags or not
+    highlight: bool = False
+
+    # list of exact matches to highlight
+    ematches: list[str] | None = None
+
+    # If true, only hydrate the text block if its text is not already populated
+    only_hydrate_empty: bool = False

hyperforge_nucliadb_agentic/ask/search/metrics.py

@@ -0,0 +1,112 @@
+import contextlib
+import time
+from typing import Any
+
+from nucliadb_telemetry import metrics
+
+buckets = [
+    0.005,
+    0.01,
+    0.025,
+    0.05,
+    0.075,
+    0.1,
+    0.25,
+    0.5,
+    0.75,
+    1.0,
+    2.5,
+    5.0,
+    7.5,
+    10.0,
+    30.0,
+    60.0,
+    metrics.INF,
+]
+
+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,
+)
+rag_histogram = metrics.Histogram(
+    name="rag",
+    labels={"step": ""},
+    buckets=buckets,
+)
+
+MetricsData = dict[str, int | float]
+
+
+class Metrics:
+    def __init__(self: "Metrics", id: str):
+        self.id = id
+        self.child_spans: list[Metrics] = []
+        self._metrics: MetricsData = {}
+
+    @contextlib.contextmanager
+    def time(self, step: str):
+        start_time = time.monotonic()
+        try:
+            yield
+        finally:
+            elapsed = time.monotonic() - start_time
+            self._metrics[step] = elapsed
+            rag_histogram.observe(elapsed, labels={"step": step})
+
+    def child_span(self, id: str) -> "Metrics":
+        child_span = Metrics(id)
+        self.child_spans.append(child_span)
+        return child_span
+
+    def set(self, key: str, value: int | float):
+        self._metrics[key] = value
+
+    def get(self, key: str) -> int | float | None:
+        return self._metrics.get(key)
+
+    def to_dict(self) -> MetricsData:
+        return self._metrics
+
+    def dump(self) -> dict[str, Any]:
+        result = {}
+        for child in self.child_spans:
+            result.update(child.dump())
+        result[self.id] = self.to_dict()
+        return result
+
+    def __getitem__(self, key: str) -> int | float:
+        return self._metrics[key]
+
+
+class AskMetrics(Metrics):
+    def __init__(self: "AskMetrics"):
+        super().__init__(id="ask")
+        self.global_start = time.monotonic()
+        self.first_chunk_yielded_at: float | None = None
+        self.first_reasoning_chunk_yielded_at: float | None = 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) -> float | None:
+        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) -> float | None:
+        if self.first_reasoning_chunk_yielded_at is None:
+            return None
+        return self.first_reasoning_chunk_yielded_at - self.global_start

hyperforge_nucliadb_agentic/ask/search/parsers/ask.py

@@ -0,0 +1,70 @@
+from pydantic import BaseModel
+from typing_extensions import assert_never
+
+from hyperforge_nucliadb_agentic.ask.model import AskRequest, MaxTokens
+from hyperforge_nucliadb_agentic.ask.search.parsers.fetcher import (
+    Fetcher,
+)
+
+
+class Generation(BaseModel):
+    """Request field related with response generation"""
+
+    use_visual_llm: bool
+    max_context_tokens: int
+    max_answer_tokens: int | None
+
+
+class _AskParser:
+    def __init__(self, kbid: str, item: AskRequest, fetcher: Fetcher):
+        self.kbid = kbid
+        self.item = item
+        self.fetcher = fetcher
+
+    async def parse(self) -> Generation:
+        use_visual_llm = await self.fetcher.get_visual_llm_enabled()
+
+        if self.item.max_tokens is None:
+            max_tokens = None
+        elif isinstance(self.item.max_tokens, int):
+            max_tokens = MaxTokens(
+                context=None,
+                answer=self.item.max_tokens,
+            )
+        elif isinstance(self.item.max_tokens, MaxTokens):
+            max_tokens = self.item.max_tokens
+        else:  # pragma: no cover
+            assert_never(self.item.max_tokens)
+
+        max_context_tokens = await self.fetcher.get_max_context_tokens(max_tokens)
+        max_answer_tokens = self.fetcher.get_max_answer_tokens(max_tokens)
+
+        return Generation(
+            use_visual_llm=use_visual_llm,
+            max_context_tokens=max_context_tokens,
+            max_answer_tokens=max_answer_tokens,
+        )
+
+
+async def parse_ask(
+    kbid: str,
+    item: AskRequest,
+    *,
+    fetcher: Fetcher | None = None,
+) -> Generation:
+    fetcher = fetcher or fetcher_for_ask(kbid, item)
+    parser = _AskParser(kbid, item, fetcher)
+    return await parser.parse()
+
+
+def fetcher_for_ask(kbid: str, item: AskRequest) -> Fetcher:
+    return Fetcher(
+        kbid=kbid,
+        query=item.query,
+        user_vector=None,
+        vectorset=item.vectorset,
+        rephrase=item.rephrase,
+        rephrase_prompt=None,
+        generative_model=item.generative_model,
+        query_image=item.query_image,
+    )

hyperforge_nucliadb_agentic/ask/search/parsers/fetcher.py

@@ -0,0 +1,192 @@
+from google.protobuf.json_format import ParseDict
+from nucliadb_models.internal.predict import QueryInfo
+from nucliadb_protos import knowledgebox_pb2, utils_pb2
+from nucliadb_sdk import NucliaDBAsync
+
+from hyperforge_nucliadb_agentic.ask import logger
+from hyperforge_nucliadb_agentic.ask.exceptions import (
+    InvalidQueryError,
+)
+from hyperforge_nucliadb_agentic.ask.model import Image, MaxTokens
+from hyperforge_nucliadb_agentic.ask.predict import (
+    SendToPredictError,
+    convert_relations,
+    get_predict,
+)
+from hyperforge_nucliadb_agentic.ask.predict_models import QueryModel
+from hyperforge_nucliadb_agentic.ask.search import rpc
+
+
+class Fetcher:
+    """This class is an encapsulation of data gathering across different parts of
+    the system. Given the user query input, it aims to be as efficient as
+    possible removing redundant expensive calls to other parts of the system. An
+    instance of a fetcher caches it's results and it's thought to be used in the
+    context of a single request.
+
+    *DO NOT* use this as a global object!
+
+    """
+
+    def __init__(
+        self,
+        kbid: str,
+        *,
+        query: str,
+        user_vector: list[float] | None,
+        vectorset: str | None,
+        rephrase: bool,
+        rephrase_prompt: str | None,
+        generative_model: str | None,
+        query_image: Image | None,
+    ):
+        self.kbid = kbid
+        self.query = query
+        self.user_vector = user_vector
+        self.user_vectorset = vectorset
+        self.user_vectorset_validated = False
+        self.rephrase = rephrase
+        self.rephrase_prompt = rephrase_prompt
+        self.generative_model = generative_model
+        self.query_image = query_image
+
+        self._query_info: QueryInfo | None = None
+        self._vectorset: str | None = None
+
+    async def query_information(self) -> QueryInfo:
+        if self._query_info is None:
+            predict = get_predict()
+            item = QueryModel(
+                text=self.query,
+                semantic_models=[self.user_vectorset] if self.user_vectorset else None,
+                generative_model=self.generative_model,
+                rephrase=self.rephrase,
+                rephrase_prompt=self.rephrase_prompt,
+                query_image=self.query_image,
+            )
+            try:
+                self._query_info = await predict.query(self.kbid, item)
+            except TimeoutError as exc:
+                raise SendToPredictError(
+                    "timeout while requesting Predict API /query"
+                ) from exc
+
+        return self._query_info
+
+    # Retrieval
+
+    async def get_rephrased_query(self) -> str | None:
+        query_info = await self.query_information()
+        return query_info.rephrased_query
+
+    def get_cached_rephrased_query(self) -> str | None:
+        if self._query_info is None:
+            return None
+        return self._query_info.rephrased_query
+
+    async def get_detected_entities(self) -> list[utils_pb2.RelationNode]:
+        query_info = await self.query_information()
+        if query_info.entities is not None:
+            detected_entities = convert_relations(query_info.entities.model_dump())
+        else:
+            detected_entities = []
+        return detected_entities
+
+    async def get_semantic_min_score(self) -> float | None:
+        query_info = await self.query_information()
+        vectorset = await self.get_vectorset()
+        return query_info.semantic_thresholds.get(vectorset, None)
+
+    async def get_vectorset(self) -> str:
+        if self._vectorset is None:
+            if self.user_vectorset is not None:
+                self._vectorset = self.user_vectorset
+            else:
+                # when it's not provided, we get the default from Predict API
+                query_info = await self.query_information()
+                if query_info.sentence is None or len(query_info.sentence.vectors) == 0:
+                    logger.error(
+                        "Asking for a vectorset but /query didn't return one",
+                        extra={"kbid": self.kbid},
+                    )
+                    raise SendToPredictError(
+                        "Predict API didn't return a sentence vectorset"
+                    )
+                # vectors field is enforced by the data model to have at least one key
+                for vectorset in query_info.sentence.vectors.keys():
+                    self._vectorset = vectorset
+                    break
+        assert self._vectorset is not None
+        return self._vectorset
+
+    async def get_query_vector(self) -> list[float]:
+        if self.user_vector is not None:
+            return self.user_vector
+
+        query_info = await self.query_information()
+        if query_info.sentence is None:
+            logger.error(
+                "Asking for a semantic query vector but /query didn't return a sentence",
+                extra={"kbid": self.kbid},
+            )
+            raise SendToPredictError(
+                "Predict API didn't return a sentence for semantic search"
+            )
+
+        vectorset = await self.get_vectorset()
+        if vectorset not in query_info.sentence.vectors:
+            logger.error(
+                "Predict is not responding with a valid query nucliadb vectorset",
+                extra={
+                    "kbid": self.kbid,
+                    "vectorset": vectorset,
+                    "predict_vectorsets": ",".join(query_info.sentence.vectors.keys()),
+                },
+            )
+            raise SendToPredictError(
+                "Predict API didn't return the requested vectorset"
+            )
+
+        query_vector = query_info.sentence.vectors[vectorset]
+        return query_vector
+
+    async def get_classification_labels(
+        self, reader_sdk: NucliaDBAsync
+    ) -> knowledgebox_pb2.Labels:
+        labelsets = await rpc.labelsets(reader_sdk=reader_sdk, kbid=self.kbid)
+
+        # TODO(decoupled-ask): remove this conversion and refactor code to use API models instead of protobuf
+        kb_labels = knowledgebox_pb2.Labels()
+        for labelset, labels in labelsets.labelsets.items():
+            ParseDict(labels.model_dump(), kb_labels.labelset[labelset])
+
+        return kb_labels
+
+    # Generative
+
+    async def get_visual_llm_enabled(self) -> bool:
+        query_info = await self.query_information()
+        if query_info is None:
+            raise SendToPredictError("Error while using predict's query endpoint")
+
+        return query_info.visual_llm
+
+    async def get_max_context_tokens(self, max_tokens: MaxTokens | None) -> int:
+        query_info = await self.query_information()
+        if query_info is None:
+            raise SendToPredictError("Error while using predict's query endpoint")
+
+        model_max = query_info.max_context
+        if max_tokens is not None and max_tokens.context is not None:
+            if max_tokens.context > model_max:
+                raise InvalidQueryError(
+                    "max_tokens.context",
+                    f"Max context tokens is higher than the model's limit of {model_max}",
+                )
+            return max_tokens.context
+        return model_max
+
+    def get_max_answer_tokens(self, max_tokens: MaxTokens | None) -> int | None:
+        if max_tokens is not None and max_tokens.answer is not None:
+            return max_tokens.answer
+        return None