nucliadb 2.46.1.post382__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/search/search/metrics.py

@@ -17,10 +17,81 @@
 # 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 contextlib
+import time
+from typing import Optional
+
 from nucliadb_telemetry import metrics
 
 merge_observer = metrics.Observer("merge_results", labels={"type": ""})
 node_features = metrics.Counter("nucliadb_node_features", labels={"type": ""})
-query_parse_dependency_observer = metrics.Observer(
-    "query_parse_dependency", labels={"type": ""}
+query_parse_dependency_observer = metrics.Observer("query_parse_dependency", labels={"type": ""})
+
+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_first_chunk",
+    buckets=buckets,
+)
+rag_histogram = metrics.Histogram(
+    name="rag",
+    labels={"step": ""},
+    buckets=buckets,
 )
+
+
+class RAGMetrics:
+    def __init__(self):
+        self.global_start = time.monotonic()
+        self._start_times: dict[str, float] = {}
+        self._end_times: dict[str, float] = {}
+        self.first_chunk_yielded_at: Optional[float] = None
+
+    @contextlib.contextmanager
+    def time(self, step: str):
+        self._start(step)
+        try:
+            yield
+        finally:
+            self._end(step)
+
+    def steps(self) -> dict[str, float]:
+        return {step: self.elapsed(step) for step in self._end_times.keys()}
+
+    def elapsed(self, step: str) -> float:
+        return self._end_times[step] - self._start_times[step]
+
+    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 get_first_chunk_time(self) -> Optional[float]:
+        if self.first_chunk_yielded_at is None:
+            return None
+        return self.first_chunk_yielded_at - self.global_start
+
+    def _start(self, step: str):
+        self._start_times[step] = time.monotonic()
+
+    def _end(self, step: str):
+        self._end_times[step] = time.monotonic()
+        elapsed = self.elapsed(step)
+        rag_histogram.observe(elapsed, labels={"step": step})

nucliadb/search/search/paragraphs.py

@@ -17,20 +17,16 @@
 # 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
 import logging
 import re
 import string
 from typing import Optional
 
-from nucliadb_protos.utils_pb2 import ExtractedText
-
+from nucliadb.common.ids import FIELD_TYPE_STR_TO_PB, ParagraphId
 from nucliadb.ingest.fields.base import Field
-from nucliadb.ingest.orm.resource import KB_REVERSE
 from nucliadb.ingest.orm.resource import Resource as ResourceORM
-from nucliadb_telemetry import metrics
-
-from .cache import get_resource_from_cache
+from nucliadb.search.search import cache
+from nucliadb_telemetry import errors, metrics
 
 logger = logging.getLogger(__name__)
 PRE_WORD = string.punctuation + " "
@@ -56,62 +52,6 @@ GET_PARAGRAPH_LATENCY = metrics.Observer(
 )
 
 
-EXTRACTED_CACHE_OPS = metrics.Counter(
-    "nucliadb_extracted_text_cache_ops", labels={"type": ""}
-)
-
-
-class ExtractedTextCache:
-    """
-    Used to cache extracted text from a resource in memory during
-    the process of search results serialization.
-    """
-
-    def __init__(self):
-        self.locks = {}
-        self.values = {}
-
-    def get_value(self, key: str) -> Optional[ExtractedText]:
-        return self.values.get(key)
-
-    def get_lock(self, key: str) -> asyncio.Lock:
-        return self.locks.setdefault(key, asyncio.Lock())
-
-    def set_value(self, key: str, value: ExtractedText) -> None:
-        self.values[key] = value
-
-    def clear(self):
-        self.values.clear()
-        self.locks.clear()
-
-
-async def get_field_extracted_text(
-    field: Field, cache: Optional[ExtractedTextCache] = None
-) -> Optional[ExtractedText]:
-    if cache is None:
-        return await field.get_extracted_text()
-
-    key = f"{field.kbid}/{field.uuid}/{field.id}"
-    extracted_text = cache.get_value(key)
-    if extracted_text is not None:
-        EXTRACTED_CACHE_OPS.inc({"type": "hit"})
-        return extracted_text
-
-    async with cache.get_lock(key):
-        # Check again in case another task already fetched it
-        extracted_text = cache.get_value(key)
-        if extracted_text is not None:
-            EXTRACTED_CACHE_OPS.inc({"type": "hit"})
-            return extracted_text
-
-        EXTRACTED_CACHE_OPS.inc({"type": "miss"})
-        extracted_text = await field.get_extracted_text()
-        if extracted_text is not None:
-            # Only cache if we actually have extracted text
-            cache.set_value(key, extracted_text)
-        return extracted_text
-
-
 @GET_PARAGRAPH_LATENCY.wrap({"type": "full"})
 async def get_paragraph_from_full_text(
     *,
@@ -119,16 +59,23 @@ async def get_paragraph_from_full_text(
     start: int,
     end: int,
     split: Optional[str] = None,
-    extracted_text_cache: Optional[ExtractedTextCache] = None,
+    log_on_missing_field: bool = True,
 ) -> str:
     """
     Pull paragraph from full text stored in database.
 
     This requires downloading the full text and then slicing it.
     """
-    extracted_text = await get_field_extracted_text(field, cache=extracted_text_cache)
+    extracted_text = await cache.get_field_extracted_text(field)
     if extracted_text is None:
-        logger.warning(f"{field} extracted_text does not exist on DB yet")
+        if log_on_missing_field:
+            logger.warning(
+                "Extracted_text for field does not exist on DB. This should not happen.",
+                extra={
+                    "field_id": field.resource_unique_id,
+                    "kbid": field.kbid,
+                },
+            )
         return ""
 
     if split not in (None, ""):
@@ -141,39 +88,46 @@ async def get_paragraph_from_full_text(
 async def get_paragraph_text(
     *,
     kbid: str,
-    rid: str,
-    field: str,
-    start: int,
-    end: int,
-    split: Optional[str] = None,
+    paragraph_id: ParagraphId,
     highlight: bool = False,
     ematches: Optional[list[str]] = None,
     matches: Optional[list[str]] = None,
     orm_resource: Optional[
         ResourceORM
     ] = None,  # allow passing in orm_resource to avoid extra DB calls or txn issues
-    extracted_text_cache: Optional[ExtractedTextCache] = None,
+    log_on_missing_field: bool = True,
 ) -> str:
+    rid = paragraph_id.rid
+    field_type = paragraph_id.field_id.type
+    field_key = paragraph_id.field_id.key
+
     if orm_resource is None:
-        orm_resource = await get_resource_from_cache(kbid, rid)
+        orm_resource = await cache.get_resource(kbid, rid)
         if orm_resource is None:
-            logger.error(f"{kbid}/{rid}:{field} does not exist on DB")
+            if log_on_missing_field:
+                logger.warning(
+                    "Resource does not exist on DB. This should not happen.",
+                    extra={"resource_id": rid, "kbid": kbid, "field": f"{field_type}/{field_key}"},
+                )
             return ""
 
-    _, field_type, field = field.split("/")
-    field_type_int = KB_REVERSE[field_type]
-    field_obj = await orm_resource.get_field(field, field_type_int, load=False)
+    field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
+    field_obj = await orm_resource.get_field(field_key, field_type_int, load=False)
 
     text = await get_paragraph_from_full_text(
         field=field_obj,
-        start=start,
-        end=end,
-        split=split,
-        extracted_text_cache=extracted_text_cache,
+        start=paragraph_id.paragraph_start,
+        end=paragraph_id.paragraph_end,
+        split=paragraph_id.field_id.subfield_id,
+        log_on_missing_field=log_on_missing_field,
     )
 
     if highlight:
-        text = highlight_paragraph(text, words=matches, ematches=ematches)
+        try:
+            text = highlight_paragraph(text, words=matches, ematches=ematches)
+        except Exception as ex:
+            errors.capture_exception(ex)
+            logger.exception("Error highlighting paragraph", extra={"kbid": kbid})
     return text
 
 
@@ -191,19 +145,17 @@ async def get_text_sentence(
     Leave separated from get paragraph for now until we understand the differences
     better.
     """
-    orm_resource = await get_resource_from_cache(kbid, rid)
+    orm_resource = await cache.get_resource(kbid, rid)
 
     if orm_resource is None:
         logger.warning(f"{rid} does not exist on DB")
         return ""
 
-    field_type_int = KB_REVERSE[field_type]
+    field_type_int = FIELD_TYPE_STR_TO_PB[field_type]
     field_obj = await orm_resource.get_field(field, field_type_int, load=False)
     extracted_text = await field_obj.get_extracted_text()
     if extracted_text is None:
-        logger.info(
-            f"{rid} {field} {field_type_int} extracted_text does not exist on DB"
-        )
+        logger.info(f"{rid} {field} {field_type_int} extracted_text does not exist on DB")
         return ""
     start = start - 1
     if start < 0:
@@ -216,36 +168,42 @@ async def get_text_sentence(
     return splitted_text
 
 
-def get_regex(some_string: str) -> str:
-    return r"\b" + some_string.lower() + r"\b"
-
-
 def highlight_paragraph(
     text: str, words: Optional[list[str]] = None, ematches: Optional[list[str]] = 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)
-    if ematches is not None:
-        for quote in ematches:
-            quote_regex = get_regex(quote.lower())
-            try:
-                for match in re.finditer(quote_regex, text_lower):
-                    start, end = match.span()
-                    marks[start] = 1
-                    marks[end] = 2
-            except re.error:
-                logger.warning(
-                    f"Regex errors while highlighting text. Regex: {quote_regex}"
-                )
-                continue
+    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 = get_regex(word.lower())
+        word_regex = REGEX_TEMPLATE.format(text=re.escape(word.lower()))
         try:
             for match in re.finditer(word_regex, text_lower):
-                start, end = match.span()
+                start, end = match.span(2)
                 if marks[start] == 0 and marks[end] == 0:
                     marks[start] = 1
                     marks[end] = 2

nucliadb/search/search/pgcatalog.py

@@ -0,0 +1,233 @@
+# 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 logging
+from collections import defaultdict
+from typing import Any, cast
+
+from psycopg.rows import dict_row
+
+from nucliadb.common.maindb.pg import PGDriver
+from nucliadb.common.maindb.utils import get_driver
+from nucliadb.search.search.query_parser.models import CatalogQuery
+from nucliadb_models.labels import translate_system_to_alias_label
+from nucliadb_models.metadata import ResourceProcessingStatus
+from nucliadb_models.search import (
+    ResourceResult,
+    Resources,
+    SortField,
+    SortOrder,
+)
+from nucliadb_telemetry import metrics
+
+from .filters import translate_label
+
+observer = metrics.Observer("pg_catalog_search", labels={"op": ""})
+logger = logging.getLogger(__name__)
+
+
+def _filter_operands(operands):
+    literals = []
+    nonliterals = []
+    for operand in operands:
+        op, params = next(iter(operand.items()))
+        if op == "literal":
+            literals.append(params)
+        else:
+            nonliterals.append(operand)
+
+    return literals, nonliterals
+
+
+def _convert_filter(filter, filter_params):
+    op, operands = next(iter(filter.items()))
+    if op == "literal":
+        param_name = f"param{len(filter_params)}"
+        filter_params[param_name] = [operands]
+        return f"labels @> %({param_name})s"
+    elif op in ("and", "or"):
+        array_op = "@>" if op == "and" else "&&"
+        sql = []
+        literals, nonliterals = _filter_operands(operands)
+        if literals:
+            param_name = f"param{len(filter_params)}"
+            filter_params[param_name] = literals
+            sql.append(f"labels {array_op} %({param_name})s")
+        for nonlit in nonliterals:
+            sql.append(_convert_filter(nonlit, filter_params))
+        return "(" + f" {op.upper()} ".join(sql) + ")"
+    elif op == "not":
+        return f"(NOT {_convert_filter(operands, filter_params)})"
+    else:
+        raise ValueError(f"Invalid operator {op}")
+
+
+def _prepare_query(catalog_query: CatalogQuery):
+    filter_sql = ["kbid = %(kbid)s"]
+    filter_params: dict[str, Any] = {"kbid": catalog_query.kbid}
+
+    if catalog_query.query:
+        # This is doing tokenization inside the SQL server (to keep the index updated). We could move it to
+        # the python code at update/query time if it ever becomes a problem but for now, a single regex
+        # executed per query is not a problem.
+        filter_sql.append(
+            "regexp_split_to_array(lower(title), '\\W') @> regexp_split_to_array(lower(%(query)s), '\\W')"
+        )
+        filter_params["query"] = catalog_query.query
+
+    if catalog_query.filters.creation.after:
+        filter_sql.append("created_at > %(created_at_start)s")
+        filter_params["created_at_start"] = catalog_query.filters.creation.after
+
+    if catalog_query.filters.creation.before:
+        filter_sql.append("created_at < %(created_at_end)s")
+        filter_params["created_at_end"] = catalog_query.filters.creation.before
+
+    if catalog_query.filters.modification.after:
+        filter_sql.append("modified_at > %(modified_at_start)s")
+        filter_params["modified_at_start"] = catalog_query.filters.modification.after
+
+    if catalog_query.filters.modification.before:
+        filter_sql.append("modified_at < %(modified_at_end)s")
+        filter_params["modified_at_end"] = catalog_query.filters.modification.before
+
+    if catalog_query.filters.labels:
+        filter_sql.append(_convert_filter(catalog_query.filters.labels, filter_params))
+
+    order_sql = ""
+    if catalog_query.sort:
+        if catalog_query.sort.field == SortField.CREATED:
+            order_field = "created_at"
+        elif catalog_query.sort.field == SortField.MODIFIED:
+            order_field = "modified_at"
+        elif catalog_query.sort.field == SortField.TITLE:
+            order_field = "title"
+        else:
+            # Deprecated order by score, use created_at instead
+            order_field = "created_at"
+
+        if catalog_query.sort.order == SortOrder.ASC:
+            order_dir = "ASC"
+        else:
+            order_dir = "DESC"
+
+        order_sql = f" ORDER BY {order_field} {order_dir}"
+
+    if catalog_query.filters.with_status:
+        filter_sql.append("labels && %(status)s")
+        if catalog_query.filters.with_status == ResourceProcessingStatus.PROCESSED:
+            filter_params["status"] = ["/n/s/PROCESSED", "/n/s/ERROR"]
+        else:
+            filter_params["status"] = ["/n/s/PENDING"]
+
+    return (
+        f"SELECT * FROM catalog WHERE {' AND '.join(filter_sql)}{order_sql}",
+        filter_params,
+    )
+
+
+def _pg_driver() -> PGDriver:
+    return cast(PGDriver, get_driver())
+
+
+@observer.wrap({"op": "search"})
+async def pgcatalog_search(catalog_query: CatalogQuery) -> Resources:
+    # Prepare SQL query
+    query, query_params = _prepare_query(catalog_query)
+
+    async with _pg_driver()._get_connection() as conn, conn.cursor(row_factory=dict_row) as cur:
+        facets = {}
+
+        # Faceted search
+        if catalog_query.faceted:
+            with observer({"op": "facets"}):
+                tmp_facets: dict[str, dict[str, int]] = {
+                    translate_label(f): defaultdict(int) for f in catalog_query.faceted
+                }
+                facet_filters = " OR ".join(f"label LIKE '{f}/%%'" for f in tmp_facets.keys())
+                for facet in tmp_facets.keys():
+                    if not (
+                        facet.startswith("/n/s") or facet.startswith("/n/i") or facet.startswith("/l")
+                    ):
+                        logger.warn(
+                            f"Unexpected facet used at catalog: {facet}, kbid={catalog_query.kbid}"
+                        )
+
+                await cur.execute(
+                    f"SELECT label, COUNT(*) FROM (SELECT unnest(labels) AS label FROM ({query}) fc) nl WHERE ({facet_filters}) GROUP BY 1 ORDER BY 1",
+                    query_params,
+                )
+
+                for row in await cur.fetchall():
+                    label = row["label"]
+                    label_parts = label.split("/")
+                    parent = "/".join(label_parts[:-1])
+                    count = row["count"]
+                    if parent in tmp_facets:
+                        tmp_facets[parent][translate_system_to_alias_label(label)] = count
+
+                    # No need to get recursive because our facets are at most 3 levels deep (e.g: /l/set/label)
+                    if len(label_parts) >= 3:
+                        grandparent = "/".join(label_parts[:-2])
+                        if grandparent in tmp_facets:
+                            tmp_facets[grandparent][translate_system_to_alias_label(parent)] += count
+
+                facets = {translate_system_to_alias_label(k): v for k, v in tmp_facets.items()}
+
+        # Totals
+        with observer({"op": "totals"}):
+            await cur.execute(
+                f"SELECT COUNT(*) FROM ({query}) fc",
+                query_params,
+            )
+            total = (await cur.fetchone())["count"]  # type: ignore
+
+        # Query
+        with observer({"op": "query"}):
+            offset = catalog_query.page_size * catalog_query.page_number
+            await cur.execute(
+                f"{query} LIMIT %(page_size)s OFFSET %(offset)s",
+                {
+                    **query_params,
+                    "page_size": catalog_query.page_size,
+                    "offset": offset,
+                },
+            )
+            data = await cur.fetchall()
+
+    return Resources(
+        facets=facets,
+        results=[
+            ResourceResult(
+                rid=str(r["rid"]).replace("-", ""),
+                field="title",
+                field_type="a",
+                labels=[label for label in r["labels"] if label.startswith("/l/")],
+                score=0,
+            )
+            for r in data
+        ],
+        query=catalog_query.query,
+        total=total,
+        page_number=catalog_query.page_number,
+        page_size=catalog_query.page_size,
+        next_page=(offset + len(data) < total),
+        min_score=0,
+    )

nucliadb/search/search/predict_proxy.py

@@ -84,5 +84,5 @@ async def predict_proxy(
 
 
 async def exists_kb(kbid: str) -> bool:
-    async with datamanagers.with_transaction(read_only=True) as txn:
+    async with datamanagers.with_ro_transaction() as txn:
         return await datamanagers.kb.exists_kb(txn, kbid=kbid)