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

nucliadb/writer/api/v1/knowledgebox.py

@@ -18,127 +18,277 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 #
 import asyncio
-import json
+from functools import partial
 
 from fastapi import HTTPException, Response
 from fastapi_versioning import version
-from nucliadb_protos.knowledgebox_pb2 import (
-    DeleteKnowledgeBoxResponse,
-    KnowledgeBoxID,
-    KnowledgeBoxNew,
-    KnowledgeBoxResponseStatus,
-    KnowledgeBoxUpdate,
-    NewKnowledgeBoxResponse,
-    UpdateKnowledgeBoxResponse,
-)
 from starlette.requests import Request
 
+from nucliadb import learning_proxy
+from nucliadb.common import datamanagers
+from nucliadb.common.external_index_providers.exceptions import (
+    ExternalIndexCreationError,
+)
+from nucliadb.common.maindb.utils import get_driver
+from nucliadb.ingest.orm.exceptions import KnowledgeBoxConflict
+from nucliadb.ingest.orm.knowledgebox import KnowledgeBox
+from nucliadb.writer import logger, vectorsets
+from nucliadb.writer.api.utils import only_for_onprem
 from nucliadb.writer.api.v1.router import KB_PREFIX, KBS_PREFIX, api
 from nucliadb.writer.utilities import get_processing
+from nucliadb_models.external_index_providers import (
+    ExternalIndexProviderType,
+    PineconeServerlessCloud,
+)
 from nucliadb_models.resource import (
     KnowledgeBoxConfig,
     KnowledgeBoxObj,
     KnowledgeBoxObjID,
     NucliaDBRoles,
 )
+from nucliadb_protos import knowledgebox_pb2
 from nucliadb_utils.authentication import requires
-from nucliadb_utils.utilities import get_ingest
 
 
+@only_for_onprem
 @api.post(
     f"/{KBS_PREFIX}",
     status_code=201,
-    name="Create Knowledge Box",
-    response_model=KnowledgeBoxObj,
+    summary="Create Knowledge Box",
     tags=["Knowledge Boxes"],
     openapi_extra={"x-hidden-operation": True},
 )
 @requires(NucliaDBRoles.MANAGER)
 @version(1)
-async def create_kb(request: Request, item: KnowledgeBoxConfig):
-    ingest = get_ingest()
-    requestpb = KnowledgeBoxNew()
-    requestpb = parse_create_kb_request(item)
-    kbobj: NewKnowledgeBoxResponse = await ingest.NewKnowledgeBox(requestpb)  # type: ignore
-    if item.slug != "":
-        slug = item.slug
-    else:
-        slug = kbobj.uuid  # type: ignore
-    if kbobj.status == KnowledgeBoxResponseStatus.OK:
-        return KnowledgeBoxObj(uuid=kbobj.uuid, slug=slug)
-    elif kbobj.status == KnowledgeBoxResponseStatus.CONFLICT:
+async def create_kb_endpoint(request: Request, item: KnowledgeBoxConfig) -> KnowledgeBoxObj:
+    try:
+        kbid, slug = await create_kb(item)
+    except KnowledgeBoxConflict:
         raise HTTPException(status_code=419, detail="Knowledge box already exists")
-    elif kbobj.status == KnowledgeBoxResponseStatus.ERROR:
-        raise HTTPException(status_code=500, detail="Error on creating knowledge box")
-
-
-def parse_create_kb_request(item: KnowledgeBoxConfig) -> KnowledgeBoxNew:
-    requestpb = KnowledgeBoxNew()
-    if item.slug:
-        requestpb.slug = item.slug
-    if item.title:
-        requestpb.config.title = item.title
-    if item.description:
-        requestpb.config.description = item.description
-    if item.release_channel:
-        requestpb.release_channel = item.release_channel.to_pb()
+    except ExternalIndexCreationError as exc:
+        raise HTTPException(status_code=502, detail=str(exc))
+    except Exception:
+        logger.exception("Could not create KB")
+        raise HTTPException(status_code=500, detail="Error creating knowledge box")
+    else:
+        return KnowledgeBoxObj(uuid=kbid, slug=slug)
+
+
+async def create_kb(item: KnowledgeBoxConfig) -> tuple[str, str]:
+    driver = get_driver()
+    rollback_learning_config = None
+
+    kbid = KnowledgeBox.new_unique_kbid()
+
+    # Onprem KB creation doesn't have an existing learning configuration yet, so
+    # we need to call learning proxy to create it
     if item.learning_configuration:
-        requestpb.learning_config = json.dumps(item.learning_configuration)
-    return requestpb
+        user_learning_config = item.learning_configuration
+    else:
+        logger.warning(
+            "No learning configuration provided. Default will be used.",
+            extra={"kbid": kbid},
+        )
+        # learning will choose the default values
+        user_learning_config = {}
+
+    # We need to be backward compatible with the old "semantic_model" field where
+    # only one semantic model was allowed.
+    if "semantic_model" in user_learning_config:
+        user_learning_config["semantic_models"] = [user_learning_config.pop("semantic_model")]
 
+    # we rely on learning to return the updated configuration with defaults and
+    # any other needed values (e.g. matryoshka settings if available)
+    learning_config = await learning_proxy.set_configuration(kbid, config=user_learning_config)
 
+    # if KB creation fails, we'll have to delete its learning config
+    async def _rollback_learning_config(kbid: str):
+        try:
+            await learning_proxy.delete_configuration(kbid)
+        except Exception:
+            logger.warning(
+                "Could not rollback learning configuration",
+                exc_info=True,
+                extra={"kbid": kbid},
+            )
+
+    rollback_learning_config = partial(_rollback_learning_config, kbid)
+    semantic_models = learning_config.into_semantic_models_metadata()
+
+    external_index_provider = knowledgebox_pb2.CreateExternalIndexProviderMetadata(
+        type=knowledgebox_pb2.ExternalIndexProviderType.UNSET,
+    )
+    if (
+        item.external_index_provider
+        and item.external_index_provider.type == ExternalIndexProviderType.PINECONE
+    ):
+        pinecone_api_key = item.external_index_provider.api_key
+        serverless_pb = to_pinecone_serverless_cloud_pb(item.external_index_provider.serverless_cloud)
+        external_index_provider = knowledgebox_pb2.CreateExternalIndexProviderMetadata(
+            type=knowledgebox_pb2.ExternalIndexProviderType.PINECONE,
+            pinecone_config=knowledgebox_pb2.CreatePineconeConfig(
+                api_key=pinecone_api_key,
+                serverless_cloud=serverless_pb,
+            ),
+        )
+
+    try:
+        (kbid, slug) = await KnowledgeBox.create(
+            driver,
+            kbid=kbid,
+            slug=item.slug or kbid,
+            title=item.title or "",
+            description=item.description or "",
+            semantic_models=semantic_models,
+            external_index_provider=external_index_provider,
+            hidden_resources_enabled=item.hidden_resources_enabled,
+            hidden_resources_hide_on_creation=item.hidden_resources_hide_on_creation,
+        )
+
+    except Exception as exc:
+        logger.error("Unexpected error creating KB", exc_info=exc, extra={"slug": item.slug})
+        await rollback_learning_config()
+        raise
+
+    return (kbid, slug)
+
+
+@only_for_onprem
 @api.patch(
     f"/{KB_PREFIX}/{{kbid}}",
     status_code=200,
-    name="Update Knowledge Box",
+    summary="Update Knowledge Box",
     response_model=KnowledgeBoxObjID,
     tags=["Knowledge Boxes"],
     openapi_extra={"x-hidden-operation": True},
 )
 @requires(NucliaDBRoles.MANAGER)
 @version(1)
-async def update_kb(request: Request, kbid: str, item: KnowledgeBoxConfig):
-    ingest = get_ingest()
-    pbrequest = KnowledgeBoxUpdate(uuid=kbid)
-    if item.slug is not None:
-        pbrequest.slug = item.slug
-    if item.title:
-        pbrequest.config.title = item.title
-    if item.description:
-        pbrequest.config.description = item.description
-    kbobj: UpdateKnowledgeBoxResponse = await ingest.UpdateKnowledgeBox(pbrequest)  # type: ignore
-    if kbobj.status == KnowledgeBoxResponseStatus.OK:
-        return KnowledgeBoxObjID(uuid=kbobj.uuid)
-    elif kbobj.status == KnowledgeBoxResponseStatus.NOTFOUND:
+async def update_kb(request: Request, kbid: str, item: KnowledgeBoxConfig) -> KnowledgeBoxObjID:
+    driver = get_driver()
+    config = None
+    if (
+        item.slug
+        or item.title
+        or item.description
+        or item.hidden_resources_enabled
+        or item.hidden_resources_hide_on_creation
+    ):
+        config = knowledgebox_pb2.KnowledgeBoxConfig(
+            slug=item.slug or "",
+            title=item.title or "",
+            description=item.description or "",
+            hidden_resources_enabled=item.hidden_resources_enabled,
+            hidden_resources_hide_on_creation=item.hidden_resources_hide_on_creation,
+        )
+    try:
+        async with driver.transaction() as txn:
+            await KnowledgeBox.update(
+                txn,
+                uuid=kbid,
+                slug=item.slug,
+                config=config,
+            )
+            await txn.commit()
+    except datamanagers.exceptions.KnowledgeBoxNotFound:
         raise HTTPException(status_code=404, detail="Knowledge box does not exist")
-    elif kbobj.status == KnowledgeBoxResponseStatus.ERROR:
-        raise HTTPException(status_code=500, detail="Error on creating knowledge box")
+    except Exception as exc:
+        logger.exception("Could not update KB", exc_info=exc, extra={"kbid": kbid})
+        raise HTTPException(status_code=500, detail="Error updating knowledge box")
+    else:
+        return KnowledgeBoxObjID(uuid=kbid)
 
 
+@only_for_onprem
 @api.delete(
     f"/{KB_PREFIX}/{{kbid}}",
     status_code=200,
-    name="Delete Knowledge Box",
-    response_model=KnowledgeBoxObj,
+    summary="Delete Knowledge Box",
     tags=["Knowledge Boxes"],
     openapi_extra={"x-hidden-operation": True},
 )
 @requires(NucliaDBRoles.MANAGER)
 @version(1)
-async def delete_kb(request: Request, kbid: str):
-    ingest = get_ingest()
-
-    kbobj: DeleteKnowledgeBoxResponse = await ingest.DeleteKnowledgeBox(  # type: ignore
-        KnowledgeBoxID(uuid=kbid)
-    )
-    if kbobj.status == KnowledgeBoxResponseStatus.OK:
-        return KnowledgeBoxObj(uuid=kbid)
-    elif kbobj.status == KnowledgeBoxResponseStatus.NOTFOUND:
+async def delete_kb(request: Request, kbid: str) -> KnowledgeBoxObj:
+    driver = get_driver()
+    try:
+        await KnowledgeBox.delete(driver, kbid=kbid)
+    except datamanagers.exceptions.KnowledgeBoxNotFound:
         raise HTTPException(status_code=404, detail="Knowledge Box does not exists")
-    elif kbobj.status == KnowledgeBoxResponseStatus.ERROR:
-        raise HTTPException(status_code=500, detail="Error on deleting knowledge box")
+    except Exception as exc:
+        logger.exception("Could not delete KB", exc_info=exc, extra={"kbid": kbid})
+        raise HTTPException(status_code=500, detail="Error deleting knowledge box")
+
+    # onprem nucliadb must delete its learning configuration
+    try:
+        await learning_proxy.delete_configuration(kbid)
+        logger.info("Learning configuration deleted", extra={"kbid": kbid})
+    except Exception as exc:
+        logger.exception(
+            "Unexpected error deleting learning configuration",
+            exc_info=exc,
+            extra={"kbid": kbid},
+        )
 
+    # be nice and notify processing this KB is being deleted so we waste
+    # resources
     processing = get_processing()
     asyncio.create_task(processing.delete_from_processing(kbid=kbid))
 
-    return Response(status_code=204)
+    return KnowledgeBoxObj(uuid=kbid)
+
+
+def to_pinecone_serverless_cloud_pb(
+    serverless: PineconeServerlessCloud,
+) -> knowledgebox_pb2.PineconeServerlessCloud.ValueType:
+    return {
+        PineconeServerlessCloud.AWS_EU_WEST_1: knowledgebox_pb2.PineconeServerlessCloud.AWS_EU_WEST_1,
+        PineconeServerlessCloud.AWS_US_EAST_1: knowledgebox_pb2.PineconeServerlessCloud.AWS_US_EAST_1,
+        PineconeServerlessCloud.AWS_US_WEST_2: knowledgebox_pb2.PineconeServerlessCloud.AWS_US_WEST_2,
+        PineconeServerlessCloud.AZURE_EASTUS2: knowledgebox_pb2.PineconeServerlessCloud.AZURE_EASTUS2,
+        PineconeServerlessCloud.GCP_US_CENTRAL1: knowledgebox_pb2.PineconeServerlessCloud.GCP_US_CENTRAL1,
+    }[serverless]
+
+
+@api.post(
+    f"/{KB_PREFIX}/{{kbid}}/vectorsets/{{vectorset_id}}",
+    status_code=200,
+    summary="Add a vectorset to Knowledge Box",
+    tags=["Knowledge Boxes"],
+    # TODO: remove when the feature is mature
+    include_in_schema=False,
+)
+@requires(NucliaDBRoles.MANAGER)
+@version(1)
+async def add_vectorset(request: Request, kbid: str, vectorset_id: str) -> Response:
+    try:
+        await vectorsets.add(kbid, vectorset_id)
+    except learning_proxy.ProxiedLearningConfigError as err:
+        return Response(
+            status_code=err.status_code,
+            content=err.content,
+            media_type=err.content_type,
+        )
+    return Response(status_code=200)
+
+
+@api.delete(
+    f"/{KB_PREFIX}/{{kbid}}/vectorsets/{{vectorset_id}}",
+    status_code=200,
+    summary="Delete vectorset from Knowledge Box",
+    tags=["Knowledge Boxes"],
+    # TODO: remove when the feature is mature
+    include_in_schema=False,
+)
+@requires(NucliaDBRoles.MANAGER)
+@version(1)
+async def delete_vectorset(request: Request, kbid: str, vectorset_id: str) -> Response:
+    try:
+        await vectorsets.delete(kbid, vectorset_id)
+    except learning_proxy.ProxiedLearningConfigError as err:
+        return Response(
+            status_code=err.status_code,
+            content=err.content,
+            media_type=err.content_type,
+        )
+    return Response(status_code=200)

nucliadb/writer/api/v1/learning_config.py

@@ -29,7 +29,7 @@ from nucliadb_utils.authentication import requires
 @api.post(
     path=f"/{KB_PREFIX}/{{kbid}}/configuration",
     status_code=204,
-    name="Create Knowledge Box models configuration",
+    summary="Create Knowledge Box models configuration",
     description="Create configuration of models assigned to a Knowledge Box",
     response_model=None,
     tags=["Knowledge Boxes"],
@@ -46,7 +46,7 @@ async def set_configuration(
 @api.patch(
     path=f"/{KB_PREFIX}/{{kbid}}/configuration",
     status_code=204,
-    name="Update Knowledge Box models configuration",
+    summary="Update Knowledge Box models configuration",
     description="Update current configuration of models assigned to a Knowledge Box",
     response_model=None,
     tags=["Knowledge Boxes"],