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

nucliadb/search/api/v1/chat.py

@@ -1,263 +0,0 @@
-# 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 base64
-import json
-from typing import Any, Optional, Union
-
-import pydantic
-from fastapi import Body, Header, Request, Response
-from fastapi.openapi.models import Example
-from fastapi_versioning import version
-from starlette.responses import StreamingResponse
-
-from nucliadb.common.datamanagers.exceptions import KnowledgeBoxNotFound
-from nucliadb.models.responses import HTTPClientError
-from nucliadb.search import logger, predict
-from nucliadb.search.api.v1.router import KB_PREFIX, api
-from nucliadb.search.predict import AnswerStatusCode
-from nucliadb.search.search.chat.query import (
-    START_OF_CITATIONS,
-    chat,
-    get_relations_results,
-)
-from nucliadb.search.search.exceptions import (
-    IncompleteFindResultsError,
-    InvalidQueryError,
-)
-from nucliadb_models.resource import NucliaDBRoles
-from nucliadb_models.search import (
-    ChatOptions,
-    ChatRequest,
-    KnowledgeboxFindResults,
-    NucliaDBClientType,
-    PromptContext,
-    PromptContextOrder,
-    Relations,
-    parse_max_tokens,
-)
-from nucliadb_telemetry.errors import capture_exception
-from nucliadb_utils.authentication import requires
-from nucliadb_utils.exceptions import LimitsExceededError
-
-END_OF_STREAM = "_END_"
-
-
-class SyncChatResponse(pydantic.BaseModel):
-    answer: str
-    relations: Optional[Relations] = None
-    results: KnowledgeboxFindResults
-    status: AnswerStatusCode
-    citations: dict[str, Any] = {}
-    prompt_context: Optional[PromptContext] = None
-    prompt_context_order: Optional[PromptContextOrder] = None
-
-
-CHAT_EXAMPLES = {
-    "search_and_chat": Example(
-        summary="Ask who won the league final",
-        description="You can ask a question to your knowledge box",  # noqa
-        value={
-            "query": "Who won the league final?",
-        },
-    ),
-    "search_and_chat_with_custom_prompt": Example(
-        summary="Ask for the gold price evolution in 2023 in a very conscise way",
-        description="You can ask a question and specify a custom prompt to tweak the tone of the response",  # noqa
-        value={
-            "query": "How has the price of gold evolved during 2023?",
-            "prompt": "Given this context: {context}. Answer this {question} in a concise way using the provided context",  # noqa
-        },
-    ),
-}
-
-
-@api.post(
-    f"/{KB_PREFIX}/{{kbid}}/chat",
-    status_code=200,
-    summary="Chat on a Knowledge Box",
-    description="Chat on a Knowledge Box",
-    tags=["Search"],
-    response_model=None,
-)
-@requires(NucliaDBRoles.READER)
-@version(1)
-async def chat_knowledgebox_endpoint(
-    request: Request,
-    kbid: str,
-    item: ChatRequest = Body(openapi_examples=CHAT_EXAMPLES),
-    x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
-    x_nucliadb_user: str = Header(""),
-    x_forwarded_for: str = Header(""),
-    x_synchronous: bool = Header(
-        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]:
-    try:
-        return await create_chat_response(
-            kbid, item, x_nucliadb_user, x_ndb_client, x_forwarded_for, x_synchronous
-        )
-    except KnowledgeBoxNotFound:
-        return HTTPClientError(
-            status_code=404,
-            detail=f"Knowledge Box '{kbid}' not found.",
-        )
-    except LimitsExceededError as exc:
-        return HTTPClientError(status_code=exc.status_code, detail=exc.detail)
-    except predict.ProxiedPredictAPIError as err:
-        return HTTPClientError(
-            status_code=err.status,
-            detail=err.detail,
-        )
-    except IncompleteFindResultsError:
-        return HTTPClientError(
-            status_code=529,
-            detail="Temporary error on information retrieval. Please try again.",
-        )
-    except predict.RephraseMissingContextError:
-        return HTTPClientError(
-            status_code=412,
-            detail="Unable to rephrase the query with the provided context.",
-        )
-    except predict.RephraseError as err:
-        return HTTPClientError(
-            status_code=529,
-            detail=f"Temporary error while rephrasing the query. Please try again later. Error: {err}",
-        )
-    except InvalidQueryError as exc:
-        return HTTPClientError(status_code=412, detail=str(exc))
-
-
-async def create_chat_response(
-    kbid: str,
-    chat_request: ChatRequest,
-    user_id: str,
-    client_type: NucliaDBClientType,
-    origin: str,
-    x_synchronous: bool,
-    resource: Optional[str] = None,
-) -> Response:
-    chat_request.max_tokens = parse_max_tokens(chat_request.max_tokens)
-    chat_result = await chat(
-        kbid,
-        chat_request,
-        user_id,
-        client_type,
-        origin,
-        resource=resource,
-    )
-    if x_synchronous:
-        streamed_answer = b""
-        async for chunk in chat_result.answer_stream:
-            streamed_answer += chunk
-
-        answer, citations = parse_streamed_answer(
-            streamed_answer, chat_request.citations
-        )
-
-        relations_results = None
-        if ChatOptions.RELATIONS in chat_request.features:
-            # XXX should use query parser here
-            relations_results = await get_relations_results(
-                kbid=kbid, text_answer=answer, target_shard_replicas=chat_request.shards
-            )
-
-        sync_chat_resp = SyncChatResponse(
-            answer=answer,
-            relations=relations_results,
-            results=chat_result.find_results,
-            status=chat_result.status_code.value,
-            citations=citations,
-        )
-        if chat_request.debug:
-            sync_chat_resp.prompt_context = chat_result.prompt_context
-            sync_chat_resp.prompt_context_order = chat_result.prompt_context_order
-        return Response(
-            content=sync_chat_resp.json(exclude_unset=True),
-            headers={
-                "NUCLIA-LEARNING-ID": chat_result.nuclia_learning_id or "unknown",
-                "Access-Control-Expose-Headers": "NUCLIA-LEARNING-ID",
-                "Content-Type": "application/json",
-            },
-        )
-    else:
-
-        async def _streaming_response():
-            bytes_results = base64.b64encode(chat_result.find_results.json().encode())
-            yield len(bytes_results).to_bytes(length=4, byteorder="big", signed=False)
-            yield bytes_results
-
-            streamed_answer = b""
-            async for chunk in chat_result.answer_stream:
-                streamed_answer += chunk
-                yield chunk
-
-            answer, _ = parse_streamed_answer(streamed_answer, chat_request.citations)
-
-            yield END_OF_STREAM.encode()
-            if ChatOptions.RELATIONS in chat_request.features:
-                # XXX should use query parser here
-                relations_results = await get_relations_results(
-                    kbid=kbid,
-                    text_answer=answer,
-                    target_shard_replicas=chat_request.shards,
-                )
-                yield base64.b64encode(relations_results.json().encode())
-
-        return StreamingResponse(
-            _streaming_response(),
-            media_type="application/octet-stream",
-            headers={
-                "NUCLIA-LEARNING-ID": chat_result.nuclia_learning_id or "unknown",
-                "Access-Control-Expose-Headers": "NUCLIA-LEARNING-ID",
-            },
-        )
-
-
-def parse_streamed_answer(
-    streamed_bytes: bytes, requested_citations: bool
-) -> tuple[str, dict[str, Any]]:
-    try:
-        text_answer, tail = streamed_bytes.split(START_OF_CITATIONS, 1)
-    except ValueError:
-        if requested_citations:
-            logger.warning(
-                "Citations were requested but not found in the answer. "
-                "Returning the answer without citations."
-            )
-        return streamed_bytes.decode("utf-8"), {}
-    if not requested_citations:
-        logger.warning(
-            "Citations were not requested but found in the answer. "
-            "Returning the answer without citations."
-        )
-        return text_answer.decode("utf-8"), {}
-    try:
-        citations_length = int.from_bytes(tail[:4], byteorder="big", signed=False)
-        citations_bytes = tail[4 : 4 + citations_length]
-        citations = json.loads(base64.b64decode(citations_bytes).decode())
-        return text_answer.decode("utf-8"), citations
-    except Exception as exc:
-        capture_exception(exc)
-        logger.exception(
-            "Error parsing citations. Returning the answer without citations."
-        )
-        return text_answer.decode("utf-8"), {}

nucliadb/search/api/v1/resource/chat.py

@@ -1,174 +0,0 @@
-# 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.common import datamanagers
-from nucliadb.models.responses import HTTPClientError
-from nucliadb.search import predict
-from nucliadb.search.api.v1.router import KB_PREFIX, RESOURCE_SLUG_PREFIX, api
-from nucliadb.search.search.exceptions import (
-    IncompleteFindResultsError,
-    InvalidQueryError,
-)
-from nucliadb_models.resource import NucliaDBRoles
-from nucliadb_models.search import ChatRequest, NucliaDBClientType
-from nucliadb_utils.authentication import requires
-from nucliadb_utils.exceptions import LimitsExceededError
-
-from ..chat import create_chat_response
-
-
-@api.post(
-    f"/{KB_PREFIX}/{{kbid}}/resource/{{rid}}/chat",
-    status_code=200,
-    summary="Chat with a resource (by id)",
-    description="Chat with a resource",
-    tags=["Search"],
-    response_model=None,
-)
-@requires(NucliaDBRoles.READER)
-@version(1)
-async def resource_chat_endpoint_by_uuid(
-    request: Request,
-    kbid: str,
-    rid: str,
-    item: ChatRequest,
-    x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
-    x_nucliadb_user: str = Header(""),
-    x_forwarded_for: str = Header(""),
-    x_synchronous: bool = Header(
-        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 resource_chat_endpoint(
-        request,
-        kbid,
-        item,
-        x_ndb_client,
-        x_nucliadb_user,
-        x_forwarded_for,
-        x_synchronous,
-        resource_id=rid,
-    )
-
-
-@api.post(
-    f"/{KB_PREFIX}/{{kbid}}/{RESOURCE_SLUG_PREFIX}/{{slug}}/chat",
-    status_code=200,
-    summary="Chat with a resource (by slug)",
-    description="Chat with a resource",
-    tags=["Search"],
-    response_model=None,
-)
-@requires(NucliaDBRoles.READER)
-@version(1)
-async def resource_chat_endpoint_by_slug(
-    request: Request,
-    kbid: str,
-    slug: str,
-    item: ChatRequest,
-    x_ndb_client: NucliaDBClientType = Header(NucliaDBClientType.API),
-    x_nucliadb_user: str = Header(""),
-    x_forwarded_for: str = Header(""),
-    x_synchronous: bool = Header(
-        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 resource_chat_endpoint(
-        request,
-        kbid,
-        item,
-        x_ndb_client,
-        x_nucliadb_user,
-        x_forwarded_for,
-        x_synchronous,
-        resource_slug=slug,
-    )
-
-
-async def resource_chat_endpoint(
-    request: Request,
-    kbid: str,
-    item: ChatRequest,
-    x_ndb_client: NucliaDBClientType,
-    x_nucliadb_user: str,
-    x_forwarded_for: str,
-    x_synchronous: bool,
-    resource_id: Optional[str] = None,
-    resource_slug: Optional[str] = None,
-) -> Union[StreamingResponse, HTTPClientError, Response]:
-
-    if resource_id is None:
-        if resource_slug is None:
-            raise ValueError("Either resource_id or resource_slug must be provided")
-
-        resource_id = await get_resource_uuid_by_slug(kbid, resource_slug)
-        if resource_id is None:
-            return HTTPClientError(status_code=404, detail="Resource not found")
-
-    try:
-        return await create_chat_response(
-            kbid,
-            item,
-            x_nucliadb_user,
-            x_ndb_client,
-            x_forwarded_for,
-            x_synchronous,
-            resource=resource_id,
-        )
-    except LimitsExceededError as exc:
-        return HTTPClientError(status_code=exc.status_code, detail=exc.detail)
-    except predict.ProxiedPredictAPIError as err:
-        return HTTPClientError(
-            status_code=err.status,
-            detail=err.detail,
-        )
-    except IncompleteFindResultsError:
-        return HTTPClientError(
-            status_code=529,
-            detail="Temporary error on information retrieval. Please try again.",
-        )
-    except predict.RephraseMissingContextError:
-        return HTTPClientError(
-            status_code=412,
-            detail="Unable to rephrase the query with the provided context.",
-        )
-    except predict.RephraseError as err:
-        return HTTPClientError(
-            status_code=529,
-            detail=f"Temporary error while rephrasing the query. Please try again later. Error: {err}",
-        )
-    except InvalidQueryError as exc:
-        return HTTPClientError(status_code=412, detail=str(exc))
-
-
-async def get_resource_uuid_by_slug(kbid: str, slug: str) -> Optional[str]:
-    async with datamanagers.with_transaction() as txn:
-        return await datamanagers.resources.get_resource_uuid_from_slug(
-            txn, kbid=kbid, slug=slug
-        )

nucliadb/search/tests/__init__.py

@@ -1,19 +0,0 @@
-# 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/>.
-#

nucliadb/search/tests/conftest.py

@@ -1,33 +0,0 @@
-# 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/>.
-#
-pytest_plugins = [
-    "pytest_docker_fixtures",
-    "nucliadb.tests.fixtures",
-    "nucliadb.tests.tikv",
-    "nucliadb.ingest.tests.fixtures",  # should be refactored out
-    "nucliadb.search.tests.node",
-    "nucliadb.search.tests.fixtures",
-    "nucliadb_utils.tests.conftest",
-    "nucliadb_utils.tests.gcs",
-    "nucliadb_utils.tests.s3",
-    "nucliadb_utils.tests.nats",
-    "nucliadb_utils.tests.asyncbenchmark",
-    "nucliadb_utils.tests.indexing",
-]

nucliadb/search/tests/fixtures.py

@@ -1,199 +0,0 @@
-# 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 asyncio
-from enum import Enum
-from typing import Optional
-
-import pytest
-from httpx import AsyncClient
-from nucliadb_protos.nodereader_pb2 import GetShardRequest
-from nucliadb_protos.noderesources_pb2 import Shard
-from redis import asyncio as aioredis
-
-from nucliadb.common.cluster.manager import KBShardManager, get_index_node
-from nucliadb.common.maindb.utils import get_driver
-from nucliadb.ingest.cache import clear_ingest_cache
-from nucliadb.ingest.tests.fixtures import broker_resource
-from nucliadb.search import API_PREFIX
-from nucliadb_utils.tests import free_port
-from nucliadb_utils.utilities import clear_global_cache
-
-
-@pytest.fixture(scope="function")
-def test_settings_search(storage, natsd, node, maindb_driver):  # type: ignore
-    from nucliadb.ingest.settings import settings as ingest_settings
-    from nucliadb_utils.cache.settings import settings as cache_settings
-    from nucliadb_utils.settings import (
-        nuclia_settings,
-        nucliadb_settings,
-        running_settings,
-    )
-
-    cache_settings.cache_pubsub_nats_url = [natsd]
-
-    running_settings.debug = False
-
-    ingest_settings.disable_pull_worker = True
-
-    ingest_settings.nuclia_partitions = 1
-
-    nuclia_settings.dummy_processing = True
-    nuclia_settings.dummy_predict = True
-    nuclia_settings.dummy_learning_services = True
-
-    ingest_settings.grpc_port = free_port()
-
-    nucliadb_settings.nucliadb_ingest = f"localhost:{ingest_settings.grpc_port}"
-
-
-@pytest.mark.asyncio
-@pytest.fixture(scope="function")
-async def search_api(test_settings_search, transaction_utility, redis):  # type: ignore
-    from nucliadb.common.cluster import manager
-    from nucliadb.search.app import application
-
-    driver = aioredis.from_url(f"redis://{redis[0]}:{redis[1]}")
-    await driver.flushall()
-
-    await application.router.startup()
-
-    # Make sure is clean
-    await asyncio.sleep(1)
-    count = 0
-    while len(manager.INDEX_NODES) < 2:
-        print("awaiting cluster nodes - search fixtures.py")
-        await asyncio.sleep(1)
-        if count == 40:
-            raise Exception("No cluster")
-        count += 1
-
-    def make_client_fixture(
-        roles: Optional[list[Enum]] = None,
-        user: str = "",
-        version: str = "1",
-        root: bool = False,
-        extra_headers: Optional[dict[str, str]] = None,
-    ) -> AsyncClient:
-        roles = roles or []
-        client_base_url = "http://test"
-
-        if root is False:
-            client_base_url = f"{client_base_url}/{API_PREFIX}/v{version}"
-
-        client = AsyncClient(app=application, base_url=client_base_url)  # type: ignore
-        client.headers["X-NUCLIADB-ROLES"] = ";".join([role.value for role in roles])
-        client.headers["X-NUCLIADB-USER"] = user
-
-        extra_headers = extra_headers or {}
-        if len(extra_headers) == 0:
-            return client
-
-        for header, value in extra_headers.items():
-            client.headers[f"{header}"] = value
-
-        return client
-
-    yield make_client_fixture
-    await application.router.shutdown()
-    # Make sure nodes can sync
-    await asyncio.sleep(1)
-    await driver.flushall()
-    await driver.close(close_connection_pool=True)
-    clear_ingest_cache()
-    clear_global_cache()
-    manager.INDEX_NODES.clear()
-
-
-@pytest.fixture(scope="function")
-async def test_search_resource(
-    indexing_utility_registered,
-    processor,
-    knowledgebox_ingest,
-):
-    """
-    Create a resource that has every possible bit of information
-    """
-    message1 = broker_resource(knowledgebox_ingest, rid="foobar", slug="foobar-slug")
-    kbid = await inject_message(processor, knowledgebox_ingest, message1)
-    resource_field_count = 3
-    await wait_for_shard(knowledgebox_ingest, resource_field_count)
-    yield kbid
-
-
-@pytest.fixture(scope="function")
-async def multiple_search_resource(
-    indexing_utility_registered,
-    processor,
-    knowledgebox_ingest,
-):
-    """
-    Create 100 resources that have every possible bit of information
-    """
-    n_resources = 100
-    fields_per_resource = 3
-    for count in range(1, n_resources + 1):
-        message = broker_resource(knowledgebox_ingest)
-        await processor.process(message=message, seqid=count)
-
-    await wait_for_shard(knowledgebox_ingest, n_resources * fields_per_resource)
-    return knowledgebox_ingest
-
-
-async def inject_message(
-    processor, knowledgebox_ingest, message, count: int = 1
-) -> str:
-    await processor.process(message=message, seqid=count)
-    await wait_for_shard(knowledgebox_ingest, count)
-    return knowledgebox_ingest
-
-
-async def wait_for_shard(knowledgebox_ingest: str, count: int) -> str:
-    # Make sure is indexed
-    driver = get_driver()
-    txn = await driver.begin()
-    shard_manager = KBShardManager()
-    shard = await shard_manager.get_current_active_shard(txn, knowledgebox_ingest)
-    if shard is None:
-        raise Exception("Could not find shard")
-    await txn.abort()
-
-    checks: dict[str, bool] = {}
-    for replica in shard.replicas:
-        if replica.shard.id not in checks:
-            checks[replica.shard.id] = False
-
-    for i in range(30):
-        for replica in shard.replicas:
-            node_obj = get_index_node(replica.node)
-            if node_obj is not None:
-                req = GetShardRequest()
-                req.shard_id.id = replica.shard.id
-                count_shard: Shard = await node_obj.reader.GetShard(req)  # type: ignore
-                if count_shard.fields >= count:
-                    checks[replica.shard.id] = True
-                else:
-                    checks[replica.shard.id] = False
-
-        if all(checks.values()):
-            break
-        await asyncio.sleep(1)
-
-    assert all(checks.values())
-    return knowledgebox_ingest