nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2777__py3-none-any.whl

nucliadb/reader/api/v1/services.py

@@ -22,26 +22,17 @@ from typing import Optional, Union
 
 from fastapi import HTTPException
 from fastapi.responses import StreamingResponse
-from fastapi_versioning import version  # type: ignore
+from fastapi_versioning import version
 from google.protobuf.json_format import MessageToDict
-from nucliadb_protos.knowledgebox_pb2 import Synonyms
-from nucliadb_protos.writer_pb2 import (
-    GetEntitiesGroupRequest,
-    GetEntitiesGroupResponse,
-    GetLabelSetRequest,
-    GetLabelSetResponse,
-    GetLabelsRequest,
-    GetLabelsResponse,
-    ListEntitiesGroupsRequest,
-    ListEntitiesGroupsResponse,
-)
 from starlette.requests import Request
 
 from nucliadb.common import datamanagers
 from nucliadb.common.cluster.settings import in_standalone_mode
 from nucliadb.common.context.fastapi import get_app_context
+from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
 from nucliadb.common.http_clients import processing
 from nucliadb.common.maindb.utils import get_driver
+from nucliadb.common.models_utils import from_proto
 from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
 from nucliadb.models.responses import HTTPClientError
 from nucliadb.reader import SERVICE_NAME
@@ -49,12 +40,19 @@ from nucliadb.reader.api.v1.router import KB_PREFIX, api
 from nucliadb.reader.reader.notifications import kb_notifications_stream
 from nucliadb_models.entities import (
     EntitiesGroup,
-    EntitiesGroupSummary,
     KnowledgeBoxEntities,
 )
 from nucliadb_models.labels import KnowledgeBoxLabels, LabelSet
 from nucliadb_models.resource import NucliaDBRoles
 from nucliadb_models.synonyms import KnowledgeBoxSynonyms
+from nucliadb_protos import writer_pb2
+from nucliadb_protos.knowledgebox_pb2 import Synonyms
+from nucliadb_protos.writer_pb2 import (
+    GetEntitiesGroupRequest,
+    GetEntitiesGroupResponse,
+    ListEntitiesGroupsRequest,
+    ListEntitiesGroupsResponse,
+)
 from nucliadb_utils.authentication import requires
 from nucliadb_utils.utilities import get_ingest, get_storage
 
@@ -88,19 +86,15 @@ async def list_entities_groups(kbid: str):
     if entities_groups.status == ListEntitiesGroupsResponse.Status.OK:
         response = KnowledgeBoxEntities(uuid=kbid)
         for key, eg_summary in entities_groups.groups.items():
-            entities_group = EntitiesGroupSummary.from_message(eg_summary)
+            entities_group = from_proto.entities_group_summary(eg_summary)
             response.groups[key] = entities_group
         return response
     elif entities_groups.status == ListEntitiesGroupsResponse.Status.NOTFOUND:
         raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
     elif entities_groups.status == ListEntitiesGroupsResponse.Status.ERROR:
-        raise HTTPException(
-            status_code=500, detail="Error while listing entities groups"
-        )
+        raise HTTPException(status_code=500, detail="Error while listing entities groups")
     else:
-        raise HTTPException(
-            status_code=500, detail="Error on listing Knowledge box entities"
-        )
+        raise HTTPException(status_code=500, detail="Error on listing Knowledge box entities")
 
 
 @api.get(
@@ -120,20 +114,14 @@ async def get_entity(request: Request, kbid: str, group: str) -> EntitiesGroup:
 
     kbobj: GetEntitiesGroupResponse = await ingest.GetEntitiesGroup(l_request)  # type: ignore
     if kbobj.status == GetEntitiesGroupResponse.Status.OK:
-        response = EntitiesGroup.from_message(kbobj.group)
+        response = from_proto.entities_group(kbobj.group)
         return response
     elif kbobj.status == GetEntitiesGroupResponse.Status.KB_NOT_FOUND:
-        raise HTTPException(
-            status_code=404, detail=f"Knowledge Box '{kbid}' does not exist"
-        )
+        raise HTTPException(status_code=404, detail=f"Knowledge Box '{kbid}' does not exist")
     elif kbobj.status == GetEntitiesGroupResponse.Status.ENTITIES_GROUP_NOT_FOUND:
-        raise HTTPException(
-            status_code=404, detail=f"Entities group '{group}' does not exist"
-        )
+        raise HTTPException(status_code=404, detail=f"Entities group '{group}' does not exist")
     else:
-        raise HTTPException(
-            status_code=500, detail="Error on getting entities group on a Knowledge box"
-        )
+        raise HTTPException(status_code=500, detail="Error on getting entities group on a Knowledge box")
 
 
 @api.get(
@@ -145,30 +133,29 @@ async def get_entity(request: Request, kbid: str, group: str) -> EntitiesGroup:
 )
 @requires(NucliaDBRoles.READER)
 @version(1)
-async def get_labelsets(request: Request, kbid: str) -> KnowledgeBoxLabels:
-    ingest = get_ingest()
-    l_request: GetLabelsRequest = GetLabelsRequest()
-    l_request.kb.uuid = kbid
+async def get_labelsets_endoint(request: Request, kbid: str) -> KnowledgeBoxLabels:
+    try:
+        return await get_labelsets(kbid)
+    except KnowledgeBoxNotFound:
+        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
+
 
-    kbobj: GetLabelsResponse = await ingest.GetLabels(l_request)  # type: ignore
-    if kbobj.status == GetLabelsResponse.Status.OK:
-        response = KnowledgeBoxLabels(uuid=kbid)
-        for labelset, labelset_data in kbobj.labels.labelset.items():
-            labelset_response = LabelSet(
-                **MessageToDict(
-                    labelset_data,
-                    preserving_proto_field_name=True,
-                    including_default_value_fields=True,
-                )
+async def get_labelsets(kbid: str) -> KnowledgeBoxLabels:
+    kb_exists = await datamanagers.atomic.kb.exists_kb(kbid=kbid)
+    if not kb_exists:
+        raise KnowledgeBoxNotFound()
+    labelsets: writer_pb2.Labels = await datamanagers.atomic.labelset.get_all(kbid=kbid)
+    response = KnowledgeBoxLabels(uuid=kbid)
+    for labelset, labelset_data in labelsets.labelset.items():
+        labelset_response = LabelSet(
+            **MessageToDict(
+                labelset_data,
+                preserving_proto_field_name=True,
+                including_default_value_fields=True,
             )
-            response.labelsets[labelset] = labelset_response
-        return response
-    elif kbobj.status == GetLabelsResponse.Status.NOTFOUND:
-        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-    else:
-        raise HTTPException(
-            status_code=500, detail="Error on getting Knowledge box labels"
         )
+        response.labelsets[labelset] = labelset_response
+    return response
 
 
 @api.get(
@@ -180,28 +167,31 @@ async def get_labelsets(request: Request, kbid: str) -> KnowledgeBoxLabels:
 )
 @requires(NucliaDBRoles.READER)
 @version(1)
-async def get_labelset(request: Request, kbid: str, labelset: str) -> LabelSet:
-    ingest = get_ingest()
-    l_request: GetLabelSetRequest = GetLabelSetRequest()
-    l_request.kb.uuid = kbid
-    l_request.labelset = labelset
+async def get_labelset_endpoint(request: Request, kbid: str, labelset: str) -> LabelSet:
+    try:
+        return await get_labelset(kbid, labelset)
+    except KnowledgeBoxNotFound:
+        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
+
 
-    kbobj: GetLabelSetResponse = await ingest.GetLabelSet(l_request)  # type: ignore
-    if kbobj.status == GetLabelSetResponse.Status.OK:
+async def get_labelset(kbid: str, labelset_id: str) -> LabelSet:
+    kb_exists = await datamanagers.atomic.kb.exists_kb(kbid=kbid)
+    if not kb_exists:
+        raise KnowledgeBoxNotFound()
+    labelset: Optional[writer_pb2.LabelSet] = await datamanagers.atomic.labelset.get(
+        kbid=kbid, labelset_id=labelset_id
+    )
+    if labelset is None:
+        response = LabelSet()
+    else:
         response = LabelSet(
             **MessageToDict(
-                kbobj.labelset,
+                labelset,
                 preserving_proto_field_name=True,
                 including_default_value_fields=True,
             )
         )
-        return response
-    elif kbobj.status == GetLabelSetResponse.Status.NOTFOUND:
-        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-    else:
-        raise HTTPException(
-            status_code=500, detail="Error on getting labelset on a Knowledge box"
-        )
+    return response
 
 
 @api.get(
@@ -215,15 +205,10 @@ async def get_labelset(request: Request, kbid: str, labelset: str) -> LabelSet:
 @requires(NucliaDBRoles.READER)
 @version(1)
 async def get_custom_synonyms(request: Request, kbid: str):
-    async with datamanagers.with_transaction(read_only=True) as txn:
-        if not datamanagers.kb.exists_kb(txn, kbid=kbid):
-            raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
-        synonyms = await datamanagers.synonyms.get(txn, kbid=kbid)
-
-    if synonyms is None:
-        synonyms = Synonyms()
-
-    return KnowledgeBoxSynonyms.from_message(synonyms)
+    if not await datamanagers.atomic.kb.exists_kb(kbid=kbid):
+        raise HTTPException(status_code=404, detail="Knowledge Box does not exist")
+    synonyms = await datamanagers.atomic.synonyms.get(kbid=kbid) or Synonyms()
+    return from_proto.kb_synonyms(synonyms)
 
 
 @api.get(
@@ -262,7 +247,7 @@ async def notifications_endpoint(
 
 
 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)
 
 
@@ -290,14 +275,12 @@ async def processing_status(
         return HTTPClientError(status_code=404, detail="Knowledge Box not found")
 
     async with processing.ProcessingHTTPClient() as client:
-        results = await client.requests(
-            cursor=cursor, scheduled=scheduled, kbid=kbid, limit=limit
-        )
+        results = await client.requests(cursor=cursor, scheduled=scheduled, kbid=kbid, limit=limit)
 
     storage = await get_storage(service_name=SERVICE_NAME)
     driver = get_driver()
 
-    async with driver.transaction(wait_for_abort=False, read_only=True) as txn:
+    async with driver.transaction(read_only=True) as txn:
         kb = KnowledgeBox(txn, storage, kbid)
 
         max_simultaneous = asyncio.Semaphore(10)
@@ -319,9 +302,7 @@ async def processing_status(
 
         result_items = [
             item
-            for item in await asyncio.gather(
-                *[_composition(result) for result in results.results]
-            )
+            for item in await asyncio.gather(*[_composition(result) for result in results.results])
             if item is not None
         ]
 

nucliadb/reader/app.py

@@ -18,63 +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 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_telemetry.fastapi.utils import (
     client_disconnect_handler,
     global_exception_handler,
 )
-from nucliadb_utils import const
+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]
+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,
@@ -104,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
@@ -106,9 +100,7 @@ async def kb_notifications(kbid: str) -> AsyncGenerator[writer_pb2.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()
@@ -121,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
 
 
@@ -144,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(
@@ -156,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(
@@ -213,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_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

@@ -18,7 +18,7 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 from . import ask  # noqa
-from . import chat  # noqa
+from . import catalog  # noqa
 from . import feedback  # noqa
 from . import find  # noqa
 from . import knowledgebox  # noqa
@@ -27,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

@@ -25,7 +25,10 @@ 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,
@@ -43,8 +46,6 @@ from nucliadb_utils.authentication import requires
     description="Ask questions on a Knowledge Box",
     tags=["Search"],
     response_model=SyncAskResponse,
-    # Add this to OpenAPI schema when endpoint is not in beta anymore
-    include_in_schema=False,
 )
 @requires(NucliaDBRoles.READER)
 @version(1)
@@ -76,15 +77,21 @@ async def create_ask_response(
     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)
-    ask_result: AskResult = await ask(
-        kbid=kbid,
-        ask_request=ask_request,
-        user_id=user_id,
-        client_type=client_type,
-        origin=origin,
-        resource=resource,
-    )
+    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",