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

nucliadb/search/predict.py

@@ -19,31 +19,37 @@
 #
 import json
 import os
+import random
 from enum import Enum
-from typing import AsyncIterator, Optional
+from typing import Any, AsyncIterator, Optional
 from unittest.mock import AsyncMock, Mock
 
 import aiohttp
 import backoff
-from nucliadb_protos.utils_pb2 import RelationNode
+from nuclia_models.predict.generative_responses import GenerativeChunk
+from pydantic import ValidationError
 
-from nucliadb.ingest.tests.vectors import Q, Qm2023
+from nucliadb.common import datamanagers
 from nucliadb.search import logger
-from nucliadb_models.search import (
-    AskDocumentModel,
-    ChatModel,
-    FeedbackRequest,
+from nucliadb.tests.vectors import Q, Qm2023
+from nucliadb_models.internal.predict import (
     Ner,
     QueryInfo,
-    RephraseModel,
+    RerankModel,
+    RerankResponse,
     SentenceSearch,
+    TokenSearch,
+)
+from nucliadb_models.search import (
+    ChatModel,
+    RephraseModel,
     SummarizedResource,
     SummarizedResponse,
     SummarizeModel,
-    TokenSearch,
 )
-from nucliadb_telemetry import metrics
-from nucliadb_utils import const
+from nucliadb_protos.utils_pb2 import RelationNode
+from nucliadb_telemetry import errors, metrics
+from nucliadb_utils.const import Features
 from nucliadb_utils.exceptions import LimitsExceededError
 from nucliadb_utils.settings import nuclia_settings
 from nucliadb_utils.utilities import Utility, has_feature, set_utility
@@ -59,10 +65,6 @@ class ProxiedPredictAPIError(Exception):
         self.detail = detail
 
 
-class PredictVectorMissing(Exception):
-    pass
-
-
 class NUAKeyMissingError(Exception):
     pass
 
@@ -77,13 +79,12 @@ class RephraseMissingContextError(Exception):
 
 DUMMY_RELATION_NODE = [
     RelationNode(value="Ferran", ntype=RelationNode.NodeType.ENTITY, subtype="PERSON"),
-    RelationNode(
-        value="Joan Antoni", ntype=RelationNode.NodeType.ENTITY, subtype="PERSON"
-    ),
+    RelationNode(value="Joan Antoni", ntype=RelationNode.NodeType.ENTITY, subtype="PERSON"),
 ]
 
 DUMMY_REPHRASE_QUERY = "This is a rephrased query"
 DUMMY_LEARNING_ID = "00"
+DUMMY_LEARNING_MODEL = "chatgpt"
 
 
 PUBLIC_PREDICT = "/api/v1/predict"
@@ -94,11 +95,12 @@ TOKENS = "/tokens"
 QUERY = "/query"
 SUMMARIZE = "/summarize"
 CHAT = "/chat"
-ASK_DOCUMENT = "/ask_document"
 REPHRASE = "/rephrase"
 FEEDBACK = "/feedback"
+RERANK = "/rerank"
 
 NUCLIA_LEARNING_ID_HEADER = "NUCLIA-LEARNING-ID"
+NUCLIA_LEARNING_MODEL_HEADER = "NUCLIA-LEARNING-MODEL"
 
 
 predict_observer = metrics.Observer(
@@ -107,7 +109,6 @@ predict_observer = metrics.Observer(
     error_mappings={
         "over_limits": LimitsExceededError,
         "predict_api_error": SendToPredictError,
-        "empty_vectors": PredictVectorMissing,
     },
 )
 
@@ -121,6 +122,13 @@ class AnswerStatusCode(str, Enum):
     ERROR = "-1"
     NO_CONTEXT = "-2"
 
+    def prettify(self) -> str:
+        return {
+            AnswerStatusCode.SUCCESS: "success",
+            AnswerStatusCode.ERROR: "error",
+            AnswerStatusCode.NO_CONTEXT: "no_context",
+        }[self]
+
 
 async def start_predict_engine():
     if nuclia_settings.dummy_predict:
@@ -144,9 +152,7 @@ def convert_relations(data: dict[str, list[dict[str, str]]]) -> list[RelationNod
     for token in data["tokens"]:
         text = token["text"]
         klass = token["ner"]
-        result.append(
-            RelationNode(value=text, ntype=RelationNode.NodeType.ENTITY, subtype=klass)
-        )
+        result.append(RelationNode(value=text, ntype=RelationNode.NodeType.ENTITY, subtype=klass))
     return result
 
 
@@ -179,9 +185,7 @@ class PredictEngine:
         await self.session.close()
 
     def check_nua_key_is_configured_for_onprem(self):
-        if self.onprem and (
-            self.nuclia_service_account is None and self.local_predict is False
-        ):
+        if self.onprem and (self.nuclia_service_account is None and self.local_predict is False):
             raise NUAKeyMissingError()
 
     def get_predict_url(self, endpoint: str, kbid: str) -> str:
@@ -193,7 +197,7 @@ class PredictEngine:
             # /api/v1/predict/rephrase/{kbid}
             return f"{self.public_url}{PUBLIC_PREDICT}{endpoint}/{kbid}"
         else:
-            if has_feature(const.Features.VERSIONED_PRIVATE_PREDICT):
+            if has_feature(Features.VERSIONED_PRIVATE_PREDICT):
                 return f"{self.cluster_url}{VERSIONED_PRIVATE_PREDICT}{endpoint}"
             else:
                 return f"{self.cluster_url}{PRIVATE_PREDICT}{endpoint}"
@@ -207,16 +211,13 @@ class PredictEngine:
         else:
             return {"X-STF-KBID": kbid}
 
-    async def check_response(
-        self, resp: aiohttp.ClientResponse, expected_status: int = 200
-    ) -> None:
+    async def check_response(self, resp: aiohttp.ClientResponse, expected_status: int = 200) -> None:
         if resp.status == expected_status:
             return
 
         if resp.status == 402:
             data = await resp.json()
             raise LimitsExceededError(402, data["detail"])
-
         try:
             data = await resp.json()
             try:
@@ -228,7 +229,10 @@ class PredictEngine:
             aiohttp.client_exceptions.ContentTypeError,
         ):
             detail = await resp.text()
-        logger.error(f"Predict API error at {resp.url}: {detail}")
+        if str(resp.status).startswith("5"):
+            logger.error(f"Predict API error at {resp.url}: {detail}")
+        else:
+            logger.info(f"Predict API error at {resp.url}: {detail}")
         raise ProxiedPredictAPIError(status=resp.status, detail=detail)
 
     @backoff.on_exception(
@@ -241,36 +245,6 @@ class PredictEngine:
         func = getattr(self.session, method.lower())
         return await func(**request_args)
 
-    @predict_observer.wrap({"type": "feedback"})
-    async def send_feedback(
-        self,
-        kbid: str,
-        item: FeedbackRequest,
-        x_nucliadb_user: str,
-        x_ndb_client: str,
-        x_forwarded_for: str,
-    ):
-        try:
-            self.check_nua_key_is_configured_for_onprem()
-        except NUAKeyMissingError:
-            logger.warning(
-                "Nuclia Service account is not defined so could not send the feedback"
-            )
-            return
-
-        data = item.dict()
-        data["user_id"] = x_nucliadb_user
-        data["client"] = x_ndb_client
-        data["forwarded"] = x_forwarded_for
-
-        resp = await self.make_request(
-            "POST",
-            url=self.get_predict_url(FEEDBACK, kbid),
-            json=data,
-            headers=self.get_predict_headers(kbid),
-        )
-        await self.check_response(resp, expected_status=204)
-
     @predict_observer.wrap({"type": "rephrase"})
     async def rephrase_query(self, kbid: str, item: RephraseModel) -> str:
         try:
@@ -283,16 +257,20 @@ class PredictEngine:
         resp = await self.make_request(
             "POST",
             url=self.get_predict_url(REPHRASE, kbid),
-            json=item.dict(),
+            json=item.model_dump(),
             headers=self.get_predict_headers(kbid),
         )
         await self.check_response(resp, expected_status=200)
         return await _parse_rephrase_response(resp)
 
-    @predict_observer.wrap({"type": "chat"})
-    async def chat_query(
+    @predict_observer.wrap({"type": "chat_ndjson"})
+    async def chat_query_ndjson(
         self, kbid: str, item: ChatModel
-    ) -> tuple[str, AsyncIterator[bytes]]:
+    ) -> tuple[str, str, AsyncIterator[GenerativeChunk]]:
+        """
+        Chat query using the new stream format
+        Format specs: https://github.com/ndjson/ndjson-spec
+        """
         try:
             self.check_nua_key_is_configured_for_onprem()
         except NUAKeyMissingError:
@@ -300,60 +278,62 @@ class PredictEngine:
             logger.warning(error)
             raise SendToPredictError(error)
 
+        # The ndjson format is triggered by the Accept header
+        headers = self.get_predict_headers(kbid)
+        headers["Accept"] = "application/x-ndjson"
+
         resp = await self.make_request(
             "POST",
             url=self.get_predict_url(CHAT, kbid),
-            json=item.dict(),
-            headers=self.get_predict_headers(kbid),
+            json=item.model_dump(),
+            headers=headers,
             timeout=None,
         )
         await self.check_response(resp, expected_status=200)
         ident = resp.headers.get(NUCLIA_LEARNING_ID_HEADER)
-        return ident, get_answer_generator(resp)
-
-    @predict_observer.wrap({"type": "ask_document"})
-    async def ask_document(
-        self, kbid: str, question: str, blocks: list[list[str]], user_id: str
-    ) -> str:
-        try:
-            self.check_nua_key_is_configured_for_onprem()
-        except NUAKeyMissingError:
-            error = "Nuclia Service account is not defined so could not ask document"
-            logger.warning(error)
-            raise SendToPredictError(error)
-
-        item = AskDocumentModel(question=question, blocks=blocks, user_id=user_id)
-        resp = await self.make_request(
-            "POST",
-            url=self.get_predict_url(ASK_DOCUMENT, kbid),
-            json=item.dict(),
-            headers=self.get_predict_headers(kbid),
-            timeout=None,
-        )
-        await self.check_response(resp, expected_status=200)
-        return await resp.text()
+        model = resp.headers.get(NUCLIA_LEARNING_MODEL_HEADER)
+        return ident, model, get_chat_ndjson_generator(resp)
 
     @predict_observer.wrap({"type": "query"})
     async def query(
         self,
         kbid: str,
         sentence: str,
+        semantic_model: Optional[str] = None,
         generative_model: Optional[str] = None,
-        rephrase: Optional[bool] = False,
+        rephrase: bool = False,
+        rephrase_prompt: Optional[str] = None,
     ) -> QueryInfo:
+        """
+        Query endpoint: returns information to be used by NucliaDB at retrieval time, for instance:
+        - The embeddings
+        - The entities
+        - The stop words
+        - The semantic threshold
+        - etc.
+
+        :param kbid: KnowledgeBox ID
+        :param sentence: The query sentence
+        :param semantic_model: The semantic model to use to generate the embeddings
+        :param generative_model: The generative model that will be used to generate the answer
+        :param rephrase: If the query should be rephrased before calculating the embeddings for a better retrieval
+        :param rephrase_prompt: Custom prompt to use for rephrasing
+        """
         try:
             self.check_nua_key_is_configured_for_onprem()
         except NUAKeyMissingError:
-            error = (
-                "Nuclia Service account is not defined so could not ask query endpoint"
-            )
+            error = "Nuclia Service account is not defined so could not ask query endpoint"
             logger.warning(error)
             raise SendToPredictError(error)
 
-        params = {
+        params: dict[str, Any] = {
             "text": sentence,
             "rephrase": str(rephrase),
         }
+        if rephrase_prompt is not None:
+            params["rephrase_prompt"] = rephrase_prompt
+        if semantic_model is not None:
+            params["semantic_models"] = [semantic_model]
         if generative_model is not None:
             params["generative_model"] = generative_model
 
@@ -367,28 +347,6 @@ class PredictEngine:
         data = await resp.json()
         return QueryInfo(**data)
 
-    @predict_observer.wrap({"type": "sentence"})
-    async def convert_sentence_to_vector(self, kbid: str, sentence: str) -> list[float]:
-        try:
-            self.check_nua_key_is_configured_for_onprem()
-        except NUAKeyMissingError:
-            logger.warning(
-                "Nuclia Service account is not defined so could not retrieve vectors for the query"
-            )
-            return []
-
-        resp = await self.make_request(
-            "GET",
-            url=self.get_predict_url(SENTENCE, kbid),
-            params={"text": sentence},
-            headers=self.get_predict_headers(kbid),
-        )
-        await self.check_response(resp, expected_status=200)
-        data = await resp.json()
-        if len(data["data"]) == 0:
-            raise PredictVectorMissing()
-        return data["data"]
-
     @predict_observer.wrap({"type": "entities"})
     async def detect_entities(self, kbid: str, sentence: str) -> list[RelationNode]:
         try:
@@ -420,26 +378,46 @@ class PredictEngine:
         resp = await self.make_request(
             "POST",
             url=self.get_predict_url(SUMMARIZE, kbid),
-            json=item.dict(),
+            json=item.model_dump(),
             headers=self.get_predict_headers(kbid),
             timeout=None,
         )
         await self.check_response(resp, expected_status=200)
         data = await resp.json()
-        return SummarizedResponse.parse_obj(data)
+        return SummarizedResponse.model_validate(data)
+
+    @predict_observer.wrap({"type": "rerank"})
+    async def rerank(self, kbid: str, item: RerankModel) -> RerankResponse:
+        try:
+            self.check_nua_key_is_configured_for_onprem()
+        except NUAKeyMissingError:
+            error = "Nuclia Service account is not defined. Rerank operation could not be performed"
+            logger.warning(error)
+            raise SendToPredictError(error)
+        resp = await self.make_request(
+            "POST",
+            url=self.get_predict_url(RERANK, kbid),
+            json=item.model_dump(),
+            headers=self.get_predict_headers(kbid),
+        )
+        await self.check_response(resp, expected_status=200)
+        data = await resp.json()
+        return RerankResponse.model_validate(data)
 
 
 class DummyPredictEngine(PredictEngine):
+    default_semantic_threshold = 0.7
+
     def __init__(self):
         self.onprem = True
         self.cluster_url = "http://localhost:8000"
         self.public_url = "http://localhost:8000"
         self.calls = []
-        self.generated_answer = [
-            b"valid ",
-            b"answer ",
-            b" to",
-            AnswerStatusCode.SUCCESS.encode(),
+        self.ndjson_answer = [
+            b'{"chunk": {"type": "text", "text": "valid "}}\n',
+            b'{"chunk": {"type": "text", "text": "answer "}}\n',
+            b'{"chunk": {"type": "text", "text": "to"}}\n',
+            b'{"chunk": {"type": "status", "code": "0"}}\n',
         ]
         self.max_context = 1000
 
@@ -458,84 +436,72 @@ class DummyPredictEngine(PredictEngine):
         response.headers = {NUCLIA_LEARNING_ID_HEADER: DUMMY_LEARNING_ID}
         return response
 
-    async def send_feedback(
-        self,
-        kbid: str,
-        item: FeedbackRequest,
-        x_nucliadb_user: str,
-        x_ndb_client: str,
-        x_forwarded_for: str,
-    ):
-        self.calls.append(("send_feedback", item))
-        return
-
     async def rephrase_query(self, kbid: str, item: RephraseModel) -> str:
         self.calls.append(("rephrase_query", item))
         return DUMMY_REPHRASE_QUERY
 
-    async def chat_query(
+    async def chat_query_ndjson(
         self, kbid: str, item: ChatModel
-    ) -> tuple[str, AsyncIterator[bytes]]:
-        self.calls.append(("chat_query", item))
+    ) -> tuple[str, str, AsyncIterator[GenerativeChunk]]:
+        self.calls.append(("chat_query_ndjson", item))
 
         async def generate():
-            for i in self.generated_answer:
-                yield i
-
-        return (DUMMY_LEARNING_ID, generate())
+            for item in self.ndjson_answer:
+                yield GenerativeChunk.model_validate_json(item)
 
-    async def ask_document(
-        self, kbid: str, query: str, blocks: list[list[str]], user_id: str
-    ) -> str:
-        self.calls.append(("ask_document", (query, blocks, user_id)))
-        answer = os.environ.get("TEST_ASK_DOCUMENT") or "Answer to your question"
-        return answer
+        return (DUMMY_LEARNING_ID, DUMMY_LEARNING_MODEL, generate())
 
     async def query(
         self,
         kbid: str,
         sentence: str,
+        semantic_model: Optional[str] = None,
         generative_model: Optional[str] = None,
-        rephrase: Optional[bool] = False,
+        rephrase: bool = False,
+        rephrase_prompt: Optional[str] = None,
     ) -> QueryInfo:
         self.calls.append(("query", sentence))
-        if (
-            os.environ.get("TEST_SENTENCE_ENCODER") == "multilingual-2023-02-21"
-        ):  # pragma: no cover
-            return QueryInfo(
-                language="en",
-                stop_words=[],
-                semantic_threshold=0.7,
-                visual_llm=True,
-                max_context=self.max_context,
-                entities=TokenSearch(
-                    tokens=[Ner(text="text", ner="PERSON", start=0, end=2)], time=0.0
-                ),
-                sentence=SentenceSearch(data=Qm2023, time=0.0),
-                query=sentence,
-            )
-        else:
-            return QueryInfo(
-                language="en",
-                stop_words=[],
-                semantic_threshold=0.7,
-                visual_llm=True,
-                max_context=self.max_context,
-                entities=TokenSearch(
-                    tokens=[Ner(text="text", ner="PERSON", start=0, end=2)], time=0.0
-                ),
-                sentence=SentenceSearch(data=Q, time=0.0),
-                query=sentence,
-            )
 
-    async def convert_sentence_to_vector(self, kbid: str, sentence: str) -> list[float]:
-        self.calls.append(("convert_sentence_to_vector", sentence))
-        if (
-            os.environ.get("TEST_SENTENCE_ENCODER") == "multilingual-2023-02-21"
-        ):  # pragma: no cover
-            return Qm2023
+        if os.environ.get("TEST_SENTENCE_ENCODER") == "multilingual-2023-02-21":  # pragma: no cover
+            base_vector = Qm2023
         else:
-            return Q
+            base_vector = Q
+
+        # populate data with existing vectorsets
+        async with datamanagers.with_ro_transaction() as txn:
+            semantic_thresholds = {}
+            vectors = {}
+            timings = {}
+            async for vectorset_id, config in datamanagers.vectorsets.iter(txn, kbid=kbid):
+                semantic_thresholds[vectorset_id] = self.default_semantic_threshold
+                vectorset_dimension = config.vectorset_index_config.vector_dimension
+                if vectorset_dimension > len(base_vector):
+                    padding = vectorset_dimension - len(base_vector)
+                    vectors[vectorset_id] = base_vector + [random.random()] * padding
+                else:
+                    vectors[vectorset_id] = base_vector[:vectorset_dimension]
+
+                timings[vectorset_id] = 0.010
+
+        # and fake data with the passed one too
+        model = semantic_model or "<PREDICT-DEFAULT-SEMANTIC-MODEL>"
+        semantic_thresholds[model] = self.default_semantic_threshold
+        vectors[model] = base_vector
+        timings[model] = 0.0
+
+        return QueryInfo(
+            language="en",
+            stop_words=[],
+            semantic_thresholds=semantic_thresholds,
+            visual_llm=True,
+            max_context=self.max_context,
+            entities=TokenSearch(tokens=[Ner(text="text", ner="PERSON", start=0, end=2)], time=0.0),
+            sentence=SentenceSearch(
+                vectors=vectors,
+                timings=timings,
+            ),
+            query=sentence,
+        )
 
     async def detect_entities(self, kbid: str, sentence: str) -> list[RelationNode]:
         self.calls.append(("detect_entities", sentence))
@@ -554,9 +520,16 @@ class DummyPredictEngine(PredictEngine):
             rsummary = []
             for field_id, field_text in item.resources[rid].fields.items():
                 rsummary.append(f"{field_id}: {field_text}")
-            response.resources[rid] = SummarizedResource(
-                summary="\n\n".join(rsummary), tokens=10
-            )
+            response.resources[rid] = SummarizedResource(summary="\n\n".join(rsummary), tokens=10)
+        return response
+
+    async def rerank(self, kbid: str, item: RerankModel) -> RerankResponse:
+        self.calls.append(("rerank", (kbid, item)))
+        # as we don't have information about the retrieval scores, return a
+        # random score given by the dict iteration
+        response = RerankResponse(
+            context_scores={paragraph_id: i for i, paragraph_id in enumerate(item.context.keys())}
+        )
         return response
 
 
@@ -578,6 +551,21 @@ def get_answer_generator(response: aiohttp.ClientResponse):
     return _iter_answer_chunks(response.content.iter_chunks())
 
 
+def get_chat_ndjson_generator(
+    response: aiohttp.ClientResponse,
+) -> AsyncIterator[GenerativeChunk]:
+    async def _parse_generative_chunks(gen):
+        async for chunk in gen:
+            try:
+                yield GenerativeChunk.model_validate_json(chunk.strip())
+            except ValidationError as ex:
+                errors.capture_exception(ex)
+                logger.error(f"Invalid chunk received: {chunk}")
+                continue
+
+    return _parse_generative_chunks(response.content)
+
+
 async def _parse_rephrase_response(
     resp: aiohttp.ClientResponse,
 ) -> str: