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

nucliadb/reader/app.py

@@ -18,76 +18,46 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-import pkg_resources
+import importlib.metadata
+
 from fastapi import FastAPI
-from fastapi.responses import JSONResponse
 from starlette.middleware import Middleware
 from starlette.middleware.authentication import AuthenticationMiddleware
-from starlette.middleware.cors import CORSMiddleware
 from starlette.requests import ClientDisconnect, Request
 from starlette.responses import HTMLResponse
 
-from nucliadb.common.context.fastapi import set_app_context
 from nucliadb.reader import API_PREFIX
 from nucliadb.reader.api.v1.router import api as api_v1
-from nucliadb.reader.lifecycle import finalize, initialize
+from nucliadb.reader.lifecycle import lifespan
 from nucliadb_telemetry import errors
-from nucliadb_utils import const
+from nucliadb_telemetry.fastapi.utils import (
+    client_disconnect_handler,
+    global_exception_handler,
+)
+from nucliadb_utils.audit.stream import AuditMiddleware
 from nucliadb_utils.authentication import NucliaCloudAuthenticationBackend
 from nucliadb_utils.fastapi.openapi import extend_openapi
 from nucliadb_utils.fastapi.versioning import VersionedFastAPI
-from nucliadb_utils.settings import http_settings, running_settings
-from nucliadb_utils.utilities import has_feature
+from nucliadb_utils.settings import running_settings
+from nucliadb_utils.utilities import get_audit
 
 middleware = []
-
-if has_feature(const.Features.CORS_MIDDLEWARE, default=False):
-    middleware.append(
-        Middleware(
-            CORSMiddleware,
-            allow_origins=http_settings.cors_origins,
-            allow_methods=["*"],
-            # Authorization will be exluded from * in the future, (CORS non-wildcard request-header).
-            # Browsers already showing deprecation notices, so it needs to be specified explicitly
-            allow_headers=["*", "Authorization"],
-        )
-    )
-
 middleware.extend(
     [
         Middleware(
             AuthenticationMiddleware,
             backend=NucliaCloudAuthenticationBackend(),
-        )
+        ),
+        Middleware(AuditMiddleware, audit_utility_getter=get_audit),
     ]
 )
 
-errors.setup_error_handling(pkg_resources.get_distribution("nucliadb").version)
-
-on_startup = [initialize]
-on_shutdown = [finalize]
-
-
-async def global_exception_handler(request: Request, exc: Exception):
-    errors.capture_exception(exc)
-    return JSONResponse(
-        status_code=500,
-        content={"detail": "Something went wrong, please contact your administrator"},
-    )
-
-
-async def client_disconnect_handler(request: Request, exc: ClientDisconnect):
-    return JSONResponse(
-        status_code=200,
-        content={"detail": "Client disconnected while an operation was in course"},
-    )
-
+errors.setup_error_handling(importlib.metadata.distribution("nucliadb").version)
 
 fastapi_settings = dict(
     debug=running_settings.debug,
     middleware=middleware,
-    on_startup=on_startup,
-    on_shutdown=on_shutdown,
+    lifespan=lifespan,
     exception_handlers={
         Exception: global_exception_handler,
         ClientDisconnect: client_disconnect_handler,
@@ -117,7 +87,7 @@ def create_application() -> FastAPI:
     # Use raw starlette routes to avoid unnecessary overhead
     application.add_route("/", homepage)
 
-    # Inject application context into the fastapi app's state
-    set_app_context(application)
+    # # Inject application context into the fastapi app's state
+    # set_app_context(application)
 
     return application

nucliadb/reader/lifecycle.py

@@ -17,20 +17,34 @@
 # 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 contextlib import asynccontextmanager
+
+from fastapi import FastAPI
+
+from nucliadb.common.context.fastapi import inject_app_context
 from nucliadb.ingest.utils import start_ingest, stop_ingest
 from nucliadb.reader import SERVICE_NAME
 from nucliadb_telemetry.utils import clean_telemetry, setup_telemetry
-from nucliadb_utils.utilities import start_audit_utility, stop_audit_utility
+from nucliadb_utils.utilities import (
+    get_storage,
+    start_audit_utility,
+    stop_audit_utility,
+    teardown_storage,
+)
 
 
-async def initialize() -> None:
+@asynccontextmanager
+async def lifespan(app: FastAPI):
     await setup_telemetry(SERVICE_NAME)
-
+    await get_storage(service_name=SERVICE_NAME)
     await start_ingest(SERVICE_NAME)
     await start_audit_utility(SERVICE_NAME)
 
+    # Inject application context into the fastapi app's state
+    async with inject_app_context(app):
+        yield
 
-async def finalize() -> None:
     await stop_ingest()
     await stop_audit_utility()
+    await teardown_storage()
     await clean_telemetry(SERVICE_NAME)

nucliadb/reader/reader/notifications.py

@@ -59,9 +59,7 @@ RESOURCE_OP_PB_TO_MODEL = {
 }
 
 
-async def kb_notifications_stream(
-    context: ApplicationContext, kbid: str
-) -> AsyncGenerator[bytes, None]:
+async def kb_notifications_stream(context: ApplicationContext, kbid: str) -> AsyncGenerator[bytes, None]:
     """
     Returns an async generator that yields pubsub notifications for the given kbid.
     The generator will return after NOTIFICATIONS_TIMEOUT_S seconds.
@@ -70,13 +68,9 @@ async def kb_notifications_stream(
         resource_cache: dict[str, str] = {}
         async with async_timeout.timeout(NOTIFICATIONS_TIMEOUT_S):
             async for pb_notification in kb_notifications(kbid):
-                notification = await serialize_notification(
-                    context, pb_notification, resource_cache
-                )
+                notification = await serialize_notification(context, pb_notification, resource_cache)
                 line = encode_streamed_notification(notification) + b"\n"
-                logger.debug(
-                    f"Sending notification: {notification.type}", extra={"kbid": kbid}
-                )
+                logger.debug(f"Sending notification: {notification.type}", extra={"kbid": kbid})
                 yield line
     except asyncio.TimeoutError:
         return
@@ -101,17 +95,12 @@ async def kb_notifications(kbid: str) -> AsyncGenerator[writer_pb2.Notification,
         data = pubsub.parse(raw_data)
         notification = writer_pb2.Notification()
         notification.ParseFromString(data)
-        # We don't need the whole broker message, so we clear it to
-        # save space, as it can potentially be very big
-        notification.ClearField("message")
         try:
             queue.put_nowait(notification)
         except asyncio.QueueFull:  # pragma: no cover
             logger.warning("Queue is full, dropping notification", extra={"kbid": kbid})
 
-    async with managed_subscription(
-        pubsub, key=subscription_key, handler=subscription_handler
-    ):
+    async with managed_subscription(pubsub, key=subscription_key, handler=subscription_handler):
         try:
             while True:
                 notification: writer_pb2.Notification = await queue.get()
@@ -124,9 +113,7 @@ async def kb_notifications(kbid: str) -> AsyncGenerator[writer_pb2.Notification,
                 )
         except Exception as ex:
             capture_exception(ex)
-            logger.error(
-                "Error while streaming activity", exc_info=True, extra={"kbid": kbid}
-            )
+            logger.error("Error while streaming activity", exc_info=True, extra={"kbid": kbid})
             return
 
 
@@ -147,9 +134,7 @@ async def managed_subscription(pubsub: PubSubDriver, key: str, handler: Callback
         try:
             await pubsub.unsubscribe(key=key, subscription_id=subscription_id)
         except Exception:  # pragma: no cover
-            logger.warning(
-                "Error while unsubscribing from activity stream", exc_info=True
-            )
+            logger.warning("Error while unsubscribing from activity stream", exc_info=True)
 
 
 async def serialize_notification(
@@ -159,9 +144,7 @@ async def serialize_notification(
     resource_uuid = pb.uuid
     seqid = pb.seqid
 
-    resource_title = await get_resource_title_cached(
-        context.kv_driver, kbid, resource_uuid, cache
-    )
+    resource_title = await get_resource_title_cached(context.kv_driver, kbid, resource_uuid, cache)
     if pb.action == writer_pb2.Notification.Action.INDEXED:
         return ResourceIndexedNotification(
             data=ResourceIndexed(
@@ -216,18 +199,14 @@ async def get_resource_title_cached(
     return resource_title
 
 
-async def get_resource_title(
-    kv_driver: Driver, kbid: str, resource_uuid: str
-) -> Optional[str]:
+async def get_resource_title(kv_driver: Driver, kbid: str, resource_uuid: str) -> Optional[str]:
     async with kv_driver.transaction(read_only=True) as txn:
-        basic = await datamanagers.resources.get_resource_basic(
-            txn, kbid=kbid, rid=resource_uuid
-        )
+        basic = await datamanagers.resources.get_basic(txn, kbid=kbid, rid=resource_uuid)
         if basic is None:
             return None
         return basic.title
 
 
 def encode_streamed_notification(notification: Notification) -> bytes:
-    encoded_nofication = notification.json().encode("utf-8")
+    encoded_nofication = notification.model_dump_json().encode("utf-8")
     return encoded_nofication

nucliadb/search/__init__.py

@@ -34,9 +34,7 @@ API_PREFIX = "api"
 class EndpointFilter(logging.Filter):
     def filter(self, record: logging.LogRecord) -> bool:
         return (
-            record.args is not None
-            and len(record.args) >= 3
-            and record.args[2] not in ("/", "/metrics")  # type: ignore
+            record.args is not None and len(record.args) >= 3 and record.args[2] not in ("/", "/metrics")  # type: ignore
         )
 
 

nucliadb/search/api/v1/__init__.py

@@ -17,7 +17,8 @@
 # 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 . import chat  # noqa
+from . import ask  # noqa
+from . import catalog  # noqa
 from . import feedback  # noqa
 from . import find  # noqa
 from . import knowledgebox  # noqa
@@ -26,6 +27,5 @@ from . import search  # noqa
 from . import suggest  # noqa
 from . import summarize  # noqa
 from .resource import ask as ask_resource  # noqa
-from .resource import chat as chat_resource  # noqa
 from .resource import search as search_resource  # noqa
 from .router import api  # noqa

nucliadb/search/api/v1/ask.py

@@ -0,0 +1,112 @@
+# 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 typing import Optional, Union
+
+from fastapi import Header, Request, Response
+from fastapi_versioning import version
+from starlette.responses import StreamingResponse
+
+from nucliadb.models.responses import HTTPClientError
+from nucliadb.search.api.v1.router import KB_PREFIX, api
+from nucliadb.search.search import cache
+from nucliadb.search.search.chat.ask import AskResult, ask, handled_ask_exceptions
+from nucliadb.search.search.chat.exceptions import AnswerJsonSchemaTooLong
+from nucliadb.search.search.utils import maybe_log_request_payload
+from nucliadb_models.resource import NucliaDBRoles
+from nucliadb_models.search import (
+    AskRequest,
+    NucliaDBClientType,
+    SyncAskResponse,
+    parse_max_tokens,
+)
+from nucliadb_utils.authentication import requires
+
+
+@api.post(
+    f"/{KB_PREFIX}/{{kbid}}/ask",
+    status_code=200,
+    summary="Ask Knowledge Box",
+    description="Ask questions on a Knowledge Box",
+    tags=["Search"],
+    response_model=SyncAskResponse,
+)
+@requires(NucliaDBRoles.READER)
+@version(1)
+async def ask_knowledgebox_endpoint(
+    request: Request,
+    kbid: str,
+    item: AskRequest,
+    x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
+    x_nucliadb_user: str = Header(""),
+    x_forwarded_for: str = Header(""),
+    x_synchronous: bool = Header(
+        default=False,
+        description="When set to true, outputs response as JSON in a non-streaming way. "
+        "This is slower and requires waiting for entire answer to be ready.",
+    ),
+) -> Union[StreamingResponse, HTTPClientError, Response]:
+    return await create_ask_response(
+        kbid, item, x_nucliadb_user, x_ndb_client, x_forwarded_for, x_synchronous
+    )
+
+
+@handled_ask_exceptions
+async def create_ask_response(
+    kbid: str,
+    ask_request: AskRequest,
+    user_id: str,
+    client_type: NucliaDBClientType,
+    origin: str,
+    x_synchronous: bool,
+    resource: Optional[str] = None,
+) -> Response:
+    maybe_log_request_payload(kbid, "/ask", ask_request)
+    ask_request.max_tokens = parse_max_tokens(ask_request.max_tokens)
+    with cache.request_caches():
+        try:
+            ask_result: AskResult = await ask(
+                kbid=kbid,
+                ask_request=ask_request,
+                user_id=user_id,
+                client_type=client_type,
+                origin=origin,
+                resource=resource,
+            )
+        except AnswerJsonSchemaTooLong as err:
+            return HTTPClientError(status_code=400, detail=str(err))
+
+    headers = {
+        "NUCLIA-LEARNING-ID": ask_result.nuclia_learning_id or "unknown",
+        "Access-Control-Expose-Headers": "NUCLIA-LEARNING-ID",
+    }
+    if x_synchronous:
+        return Response(
+            content=await ask_result.json(),
+            status_code=200,
+            headers=headers,
+            media_type="application/json",
+        )
+    else:
+        return StreamingResponse(
+            content=ask_result.ndjson_stream(),
+            status_code=200,
+            headers=headers,
+            media_type="application/x-ndjson",
+        )

nucliadb/search/api/v1/catalog.py

@@ -0,0 +1,184 @@
+# 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 time import time
+from typing import Optional, Union
+
+from fastapi import Request, Response
+from fastapi_versioning import version
+
+from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
+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.exceptions import InvalidQueryError
+from nucliadb.search.search.merge import fetch_resources
+from nucliadb.search.search.pgcatalog import pgcatalog_search
+from nucliadb.search.search.query_parser.parser import parse_catalog
+from nucliadb.search.search.utils import (
+    maybe_log_request_payload,
+)
+from nucliadb_models.common import FieldTypeName
+from nucliadb_models.metadata import ResourceProcessingStatus
+from nucliadb_models.resource import NucliaDBRoles
+from nucliadb_models.search import (
+    CatalogRequest,
+    CatalogResponse,
+    KnowledgeboxSearchResults,
+    ResourceProperties,
+    SearchParamDefaults,
+    SortField,
+    SortOptions,
+    SortOrder,
+)
+from nucliadb_models.utils import DateTime
+from nucliadb_utils.authentication import requires
+from nucliadb_utils.exceptions import LimitsExceededError
+
+
+@api.get(
+    f"/{KB_PREFIX}/{{kbid}}/catalog",
+    status_code=200,
+    summary="List resources of a Knowledge Box",
+    description="List resources of a Knowledge Box",
+    response_model=KnowledgeboxSearchResults,
+    response_model_exclude_unset=True,
+    tags=["Search"],
+)
+@requires(NucliaDBRoles.READER)
+@version(1)
+async def catalog_get(
+    request: Request,
+    response: Response,
+    kbid: str,
+    query: str = fastapi_query(SearchParamDefaults.query),
+    filters: list[str] = fastapi_query(SearchParamDefaults.filters),
+    faceted: list[str] = fastapi_query(SearchParamDefaults.faceted),
+    sort_field: SortField = fastapi_query(SearchParamDefaults.sort_field),
+    sort_limit: Optional[int] = fastapi_query(SearchParamDefaults.sort_limit),
+    sort_order: SortOrder = fastapi_query(SearchParamDefaults.sort_order),
+    page_number: int = fastapi_query(SearchParamDefaults.catalog_page_number),
+    page_size: int = fastapi_query(SearchParamDefaults.catalog_page_size),
+    shards: list[str] = fastapi_query(SearchParamDefaults.shards, deprecated=True),
+    with_status: Optional[ResourceProcessingStatus] = fastapi_query(
+        SearchParamDefaults.with_status, deprecated="Use filters instead"
+    ),
+    debug: bool = fastapi_query(SearchParamDefaults.debug, include_in_schema=False),
+    range_creation_start: Optional[DateTime] = fastapi_query(SearchParamDefaults.range_creation_start),
+    range_creation_end: Optional[DateTime] = fastapi_query(SearchParamDefaults.range_creation_end),
+    range_modification_start: Optional[DateTime] = fastapi_query(
+        SearchParamDefaults.range_modification_start
+    ),
+    range_modification_end: Optional[DateTime] = fastapi_query(
+        SearchParamDefaults.range_modification_end
+    ),
+    hidden: Optional[bool] = fastapi_query(SearchParamDefaults.hidden),
+) -> Union[KnowledgeboxSearchResults, HTTPClientError]:
+    item = CatalogRequest(
+        query=query,
+        filters=filters,
+        faceted=faceted,
+        page_number=page_number,
+        page_size=page_size,
+        shards=shards,
+        debug=debug,
+        with_status=with_status,
+        range_creation_start=range_creation_start,
+        range_creation_end=range_creation_end,
+        range_modification_start=range_modification_start,
+        range_modification_end=range_modification_end,
+        hidden=hidden,
+    )
+    if sort_field:
+        item.sort = SortOptions(field=sort_field, limit=sort_limit, order=sort_order)
+    return await catalog(kbid, item)
+
+
+@api.post(
+    f"/{KB_PREFIX}/{{kbid}}/catalog",
+    status_code=200,
+    summary="List resources of a Knowledge Box",
+    description="List resources of a Knowledge Box",
+    response_model=KnowledgeboxSearchResults,
+    response_model_exclude_unset=True,
+    tags=["Search"],
+)
+@requires(NucliaDBRoles.READER)
+@version(1)
+async def catalog_post(
+    request: Request,
+    kbid: str,
+    item: CatalogRequest,
+) -> Union[CatalogResponse, HTTPClientError]:
+    return await catalog(kbid, item)
+
+
+async def catalog(
+    kbid: str,
+    item: CatalogRequest,
+):
+    """
+    Catalog endpoint is a simplified version of the search endpoint, it only
+    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:
+        with cache.request_caches():
+            query_parser = parse_catalog(kbid, item)
+
+            catalog_results = CatalogResponse()
+            catalog_results.fulltext = await pgcatalog_search(query_parser)
+            catalog_results.resources = await fetch_resources(
+                resources=[r.rid for r in catalog_results.fulltext.results],
+                kbid=kbid,
+                show=[ResourceProperties.BASIC, ResourceProperties.ERRORS],
+                field_type_filter=list(FieldTypeName),
+                extracted=[],
+            )
+            return catalog_results
+    except InvalidQueryError as exc:
+        return HTTPClientError(status_code=412, detail=str(exc))
+    except KnowledgeBoxNotFound:
+        return HTTPClientError(status_code=404, detail="Knowledge Box not found")
+    except LimitsExceededError as exc:
+        return HTTPClientError(status_code=exc.status_code, detail=exc.detail)
+    finally:
+        duration = time() - start_time
+        if duration > 2:  # pragma: no cover
+            logger.warning(
+                "Slow catalog request",
+                extra={
+                    "kbid": kbid,
+                    "duration": duration,
+                    "query": item.model_dump_json(),
+                },
+            )
+
+
+def pgcatalog_enabled():
+    return isinstance(get_driver(), PGDriver)

nucliadb/search/api/v1/feedback.py

@@ -18,24 +18,24 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 
-
 from fastapi import Header, Request, Response
 from fastapi_versioning import version
 
+from nucliadb.common.models_utils import to_proto
 from nucliadb.models.responses import HTTPClientError
-from nucliadb.search import logger, predict
+from nucliadb.search import logger
 from nucliadb.search.api.v1.router import KB_PREFIX, api
-from nucliadb.search.utilities import get_predict
 from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_models.search import FeedbackRequest, NucliaDBClientType
 from nucliadb_telemetry import errors
 from nucliadb_utils.authentication import requires
+from nucliadb_utils.utilities import get_audit
 
 
 @api.post(
     f"/{KB_PREFIX}/{{kbid}}/feedback",
     status_code=200,
-    name="Send Feedback",
+    summary="Send Feedback",
     description="Send feedback for a search operation in a Knowledge Box",
     tags=["Search"],
 )
@@ -51,28 +51,20 @@ async def send_feedback_endpoint(
     x_forwarded_for: str = Header(""),
 ):
     try:
-        return await send_feedback(
-            kbid, item, x_nucliadb_user, x_ndb_client, x_forwarded_for
-        )
-    except predict.ProxiedPredictAPIError as err:
-        return HTTPClientError(
-            status_code=503,
-            detail=f"Feedback service unavailable. {err.status}: {err.detail}",
-        )
+        audit = get_audit()
+        if audit is not None:
+            audit.feedback(
+                kbid=kbid,
+                user=x_nucliadb_user,
+                client_type=to_proto.client_type(x_ndb_client),
+                origin=x_forwarded_for,
+                learning_id=item.ident,
+                good=item.good,
+                task=to_proto.feedback_task(item.task),
+                feedback=item.feedback,
+                text_block_id=item.text_block_id,
+            )
     except Exception as ex:
         errors.capture_exception(ex)
         logger.exception("Unexpected error sending feedback", extra={"kbid": kbid})
         return HTTPClientError(status_code=500, detail=f"Internal server error")
-
-
-async def send_feedback(
-    kbid: str,
-    item: FeedbackRequest,
-    x_nucliadb_user: str,
-    x_ndb_client: NucliaDBClientType,
-    x_forwarded_for: str,
-):
-    predict = get_predict()
-    await predict.send_feedback(
-        kbid, item, x_nucliadb_user, x_ndb_client, x_forwarded_for
-    )