nucliadb 6.5.0.post4426__py3-none-any.whl → 6.5.0.post4484__py3-none-any.whl

migrations/0037_backfill_catalog_facets.py

@@ -0,0 +1,74 @@
+# 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/>.
+#
+
+"""Migration #37
+
+Backfill catalog facets
+
+"""
+
+import logging
+from typing import cast
+
+from nucliadb.common.maindb.pg import PGDriver, PGTransaction
+from nucliadb.migrator.context import ExecutionContext
+
+logger = logging.getLogger(__name__)
+
+
+async def migrate(context: ExecutionContext) -> None:
+    driver = cast(PGDriver, context.kv_driver)
+
+    BATCH_SIZE = 1_000
+    async with driver.transaction() as txn:
+        txn = cast(PGTransaction, txn)
+        start_kbid = "00000000000000000000000000000000"
+        start_rid = "00000000000000000000000000000000"
+        while True:
+            async with txn.connection.cursor() as cur:
+                logger.info(f"Filling {BATCH_SIZE} catalog facets from {start_kbid}, {start_rid}")
+                # Get a batch of facets from the catalog table
+                await cur.execute(
+                    """
+                        WITH i AS (
+                            INSERT INTO catalog_facets (kbid, rid, facet)
+                            SELECT kbid, rid, unnest(extract_facets(labels)) FROM (
+                                SELECT * FROM catalog
+                                WHERE (kbid = %(kbid)s AND rid > %(rid)s) OR kbid > %(kbid)s
+                                ORDER BY kbid, rid
+                                LIMIT %(batch)s
+                            ) rs
+                            RETURNING kbid, rid
+                        )
+                        SELECT kbid, rid FROM i ORDER BY kbid DESC, rid DESC LIMIT 1;
+                    """,
+                    {"kbid": start_kbid, "rid": start_rid, "batch": BATCH_SIZE},
+                )
+
+                # Set the key for next iteration
+                results = await cur.fetchone()  # type: ignore
+                if results is None:
+                    break
+                (start_kbid, start_rid) = results
+
+                await txn.commit()
+
+
+async def migrate_kb(context: ExecutionContext, kbid: str) -> None: ...

migrations/pg/0001_bootstrap.py

@@ -26,7 +26,7 @@ async def migrate(txn: PGTransaction) -> None:
         # IF NOT EXISTS just for compatibility with older install predating the migration system
         await cur.execute("""
             CREATE TABLE IF NOT EXISTS resources (
-                key TEXT PRIMARY KEY,
+                key TEXT COLLATE ucs_basic PRIMARY KEY,
                 value BYTEA
             );
         """)

migrations/pg/0008_catalog_facets.py

@@ -0,0 +1,43 @@
+# 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.maindb.pg import PGTransaction
+
+
+async def migrate(txn: PGTransaction) -> None:
+    async with txn.connection.cursor() as cur:
+        await cur.execute(
+            """
+            CREATE TABLE catalog_facets (
+                id BIGINT PRIMARY KEY GENERATED ALWAYS AS IDENTITY,
+                kbid UUID,
+                rid UUID,
+                facet TEXT COLLATE ucs_basic,
+
+                FOREIGN KEY (kbid, rid) REFERENCES catalog (kbid, rid) ON DELETE CASCADE
+            );
+
+            -- For FK checks
+            CREATE INDEX ON catalog_facets(kbid, rid);
+
+            -- Best for per-facet aggregation, also used by search with facet filter
+            CREATE INDEX ON catalog_facets(kbid, facet);
+            """
+        )

migrations/pg/0009_extract_facets_safety.py

@@ -0,0 +1,26 @@
+# 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.maindb.pg import PGTransaction
+
+
+async def migrate(txn: PGTransaction) -> None:
+    async with txn.connection.cursor() as cur:
+        await cur.execute("ALTER FUNCTION extract_facets(text[]) PARALLEL SAFE;")

nucliadb/ingest/fields/base.py

@@ -19,8 +19,10 @@
 #
 from __future__ import annotations
 
+import asyncio
 import enum
 import logging
+from collections import defaultdict
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Generic, Optional, Type, TypeVar
 
@@ -113,6 +115,8 @@ class Field(Generic[PbType]):
                 raise InvalidPBClass(self.__class__, pb.__class__)
             self.value = pb
 
+        self.locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
+
     @property
     def kbid(self) -> str:
         return self.resource.kb.kbid
@@ -364,10 +368,13 @@ class Field(Generic[PbType]):
 
     async def get_extracted_text(self, force=False) -> Optional[ExtractedText]:
         if self.extracted_text is None or force:
-            sf = self.get_storage_field(FieldTypes.FIELD_TEXT)
-            payload = await self.storage.download_pb(sf, ExtractedText)
-            if payload is not None:
-                self.extracted_text = payload
+            async with self.locks["extracted_text"]:
+                # Value could have been fetched while waiting for the lock
+                if self.extracted_text is None or force:
+                    sf = self.get_storage_field(FieldTypes.FIELD_TEXT)
+                    payload = await self.storage.download_pb(sf, ExtractedText)
+                    if payload is not None:
+                        self.extracted_text = payload
         return self.extracted_text
 
     async def set_vectors(
@@ -499,10 +506,13 @@ class Field(Generic[PbType]):
 
     async def get_field_metadata(self, force: bool = False) -> Optional[FieldComputedMetadata]:
         if self.computed_metadata is None or force:
-            sf = self.get_storage_field(FieldTypes.FIELD_METADATA)
-            payload = await self.storage.download_pb(sf, FieldComputedMetadata)
-            if payload is not None:
-                self.computed_metadata = payload
+            async with self.locks["field_metadata"]:
+                # Value could have been fetched while waiting for the lock
+                if self.computed_metadata is None or force:
+                    sf = self.get_storage_field(FieldTypes.FIELD_METADATA)
+                    payload = await self.storage.download_pb(sf, FieldComputedMetadata)
+                    if payload is not None:
+                        self.computed_metadata = payload
         return self.computed_metadata
 
     async def set_large_field_metadata(self, payload: LargeComputedMetadataWrapper):

nucliadb/ingest/orm/processor/pgcatalog.py

@@ -40,6 +40,17 @@ def pgcatalog_enabled(kbid):
     return isinstance(get_driver(), PGDriver)
 
 
+def extract_facets(labels):
+    facets = set()
+    for label in labels:
+        parts = label.split("/")
+        facet = ""
+        for part in parts[1:]:
+            facet += f"/{part}"
+            facets.add(facet)
+    return facets
+
+
 @observer.wrap({"type": "update"})
 async def pgcatalog_update(txn: Transaction, kbid: str, resource: Resource, index_message: IndexMessage):
     if not pgcatalog_enabled(kbid):
@@ -76,6 +87,21 @@ async def pgcatalog_update(txn: Transaction, kbid: str, resource: Resource, inde
                 "slug": resource.basic.slug,
             },
         )
+        await cur.execute(
+            "DELETE FROM catalog_facets WHERE kbid = %(kbid)s AND rid = %(rid)s",
+            {
+                "kbid": resource.kb.kbid,
+                "rid": resource.uuid,
+            },
+        )
+        await cur.execute(
+            "INSERT INTO catalog_facets (kbid, rid, facet) SELECT %(kbid)s AS kbid, %(rid)s AS rid, unnest(%(facets)s::text[]) AS facet",
+            {
+                "kbid": resource.kb.kbid,
+                "rid": resource.uuid,
+                "facets": list(extract_facets(index_message.labels)),
+            },
+        )
 
 
 @observer.wrap({"type": "delete"})

nucliadb/ingest/orm/resource.py

@@ -19,6 +19,7 @@
 #
 from __future__ import annotations
 
+import asyncio
 import logging
 from collections import defaultdict
 from concurrent.futures import ThreadPoolExecutor
@@ -126,6 +127,7 @@ class Resource:
         self.disable_vectors = disable_vectors
         self._previous_status: Optional[Metadata.Status.ValueType] = None
         self.user_relations: Optional[PBRelations] = None
+        self.locks: dict[str, asyncio.Lock] = defaultdict(asyncio.Lock)
 
     async def set_slug(self):
         basic = await self.get_basic()
@@ -306,10 +308,13 @@ class Resource:
     async def get_field(self, key: str, type: FieldType.ValueType, load: bool = True):
         field = (type, key)
         if field not in self.fields:
-            field_obj: Field = KB_FIELDS[type](id=key, resource=self)
-            if load:
-                await field_obj.get_value()
-            self.fields[field] = field_obj
+            async with self.locks["field"]:
+                # Field could have been fetch while waiting for the lock
+                if field not in self.fields:
+                    field_obj: Field = KB_FIELDS[type](id=key, resource=self)
+                    if load:
+                        await field_obj.get_value()
+                    self.fields[field] = field_obj
         return self.fields[field]
 
     async def set_field(self, type: FieldType.ValueType, key: str, payload: Any):

nucliadb/search/api/v1/catalog.py

@@ -27,15 +27,13 @@ from pydantic import ValidationError
 
 from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
 from nucliadb.common.exceptions import InvalidQueryError
-from nucliadb.common.maindb.pg import PGDriver
-from nucliadb.common.maindb.utils import get_driver
 from nucliadb.models.responses import HTTPClientError
 from nucliadb.search import logger
 from nucliadb.search.api.v1.router import KB_PREFIX, api
 from nucliadb.search.api.v1.utils import fastapi_query
 from nucliadb.search.search import cache
 from nucliadb.search.search.merge import fetch_resources
-from nucliadb.search.search.pgcatalog import pgcatalog_search
+from nucliadb.search.search.pgcatalog import pgcatalog_facets, pgcatalog_search
 from nucliadb.search.search.query_parser.parsers import parse_catalog
 from nucliadb.search.search.utils import (
     maybe_log_request_payload,
@@ -45,6 +43,7 @@ from nucliadb_models.filters import CatalogFilterExpression
 from nucliadb_models.metadata import ResourceProcessingStatus
 from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_models.search import (
+    CatalogFacetsRequest,
     CatalogRequest,
     CatalogResponse,
     KnowledgeboxSearchResults,
@@ -157,9 +156,6 @@ async def catalog(
     returns bm25 results on titles and it does not support vector search.
     It is useful for listing resources in a knowledge box.
     """
-    if not pgcatalog_enabled():  # pragma: no cover
-        return HTTPClientError(status_code=501, detail="PG driver is needed for catalog search")
-
     maybe_log_request_payload(kbid, "/catalog", item)
     start_time = time()
     try:
@@ -196,5 +192,15 @@ async def catalog(
             )
 
 
-def pgcatalog_enabled():
-    return isinstance(get_driver(), PGDriver)
+@api.post(
+    f"/{KB_PREFIX}/{{kbid}}/catalog/facets",
+    status_code=200,
+    response_model=dict[str, int],
+    response_model_exclude_unset=True,
+    tags=["Search"],
+    include_in_schema=False,
+)
+@requires(NucliaDBRoles.READER)
+@version(1)
+async def catalog_facets(request: Request, kbid: str, item: CatalogFacetsRequest) -> dict[str, int]:
+    return await pgcatalog_facets(kbid, item)

nucliadb/search/search/chat/ask.py

@@ -78,6 +78,8 @@ from nucliadb_models.search import (
     AskRetrievalMatch,
     AskTimings,
     AskTokens,
+    AugmentedContext,
+    AugmentedContextResponseItem,
     ChatModel,
     ChatOptions,
     CitationsAskResponseItem,
@@ -143,6 +145,7 @@ class AskResult:
         metrics: AskMetrics,
         best_matches: list[RetrievalMatch],
         debug_chat_model: Optional[ChatModel],
+        augmented_context: AugmentedContext,
     ):
         # Initial attributes
         self.kbid = kbid
@@ -157,6 +160,7 @@ class AskResult:
         self.auditor: ChatAuditor = auditor
         self.metrics: AskMetrics = metrics
         self.best_matches: list[RetrievalMatch] = best_matches
+        self.augmented_context = augmented_context
 
         # Computed from the predict chat answer stream
         self._answer_text = ""
@@ -272,9 +276,13 @@ class AskResult:
             status_code=self.status_code,
         )
 
+        yield AugmentedContextResponseItem(augmented=self.augmented_context)
+
         # Stream out the citations
         if self._citations is not None:
-            yield CitationsAskResponseItem(citations=self._citations.citations)
+            yield CitationsAskResponseItem(
+                citations=self._citations.citations,
+            )
 
         # Stream out generic metadata about the answer
         if self._metadata is not None:
@@ -366,6 +374,7 @@ class AskResult:
             citations=citations,
             metadata=metadata,
             learning_id=self.nuclia_learning_id or "",
+            augmented_context=self.augmented_context,
         )
         if self.status_code == AnswerStatusCode.ERROR and self.status_error_details:
             response.error_details = self.status_error_details
@@ -569,6 +578,7 @@ async def ask(
             prompt_context,
             prompt_context_order,
             prompt_context_images,
+            augmented_context,
         ) = await prompt_context_builder.build()
 
     # Make the chat request to the predict API
@@ -631,6 +641,7 @@ async def ask(
         metrics=metrics,
         best_matches=retrieval_results.best_matches,
         debug_chat_model=chat_model,
+        augmented_context=augmented_context,
     )