nucliadb 4.0.0.post542__py3-none-any.whl → 6.2.1.post2798__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 Any, AsyncIterator, Literal, Optional, Union
+from typing import Any, AsyncIterator, Optional
 from unittest.mock import AsyncMock, Mock
 
 import aiohttp
 import backoff
-from nucliadb_protos.utils_pb2 import RelationNode
-from pydantic import BaseModel, Field, ValidationError
+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 (
-    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_protos.utils_pb2 import RelationNode
 from nucliadb_telemetry import errors, metrics
-from nucliadb_utils import const
+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
@@ -73,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"
@@ -92,8 +97,10 @@ SUMMARIZE = "/summarize"
 CHAT = "/chat"
 REPHRASE = "/rephrase"
 FEEDBACK = "/feedback"
+RERANK = "/rerank"
 
 NUCLIA_LEARNING_ID_HEADER = "NUCLIA-LEARNING-ID"
+NUCLIA_LEARNING_MODEL_HEADER = "NUCLIA-LEARNING-MODEL"
 
 
 predict_observer = metrics.Observer(
@@ -123,41 +130,6 @@ class AnswerStatusCode(str, Enum):
         }[self]
 
 
-class TextGenerativeResponse(BaseModel):
-    type: Literal["text"] = "text"
-    text: str
-
-
-class MetaGenerativeResponse(BaseModel):
-    type: Literal["meta"] = "meta"
-    input_tokens: int
-    output_tokens: int
-    timings: dict[str, float]
-
-
-class CitationsGenerativeResponse(BaseModel):
-    type: Literal["citations"] = "citations"
-    citations: dict[str, Any]
-
-
-class StatusGenerativeResponse(BaseModel):
-    type: Literal["status"] = "status"
-    code: str
-    details: Optional[str] = None
-
-
-GenerativeResponse = Union[
-    TextGenerativeResponse,
-    MetaGenerativeResponse,
-    CitationsGenerativeResponse,
-    StatusGenerativeResponse,
-]
-
-
-class GenerativeChunk(BaseModel):
-    chunk: GenerativeResponse = Field(..., discriminator="type")
-
-
 async def start_predict_engine():
     if nuclia_settings.dummy_predict:
         predict_util = DummyPredictEngine()
@@ -180,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
 
 
@@ -215,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:
@@ -229,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}"
@@ -243,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:
@@ -264,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(
@@ -277,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:
@@ -319,38 +257,16 @@ 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(
-        self, kbid: str, item: ChatModel
-    ) -> tuple[str, AsyncIterator[bytes]]:
-        try:
-            self.check_nua_key_is_configured_for_onprem()
-        except NUAKeyMissingError:
-            error = "Nuclia Service account is not defined so the chat operation could not be performed"
-            logger.warning(error)
-            raise SendToPredictError(error)
-
-        resp = await self.make_request(
-            "POST",
-            url=self.get_predict_url(CHAT, kbid),
-            json=item.dict(),
-            headers=self.get_predict_headers(kbid),
-            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": "chat_ndjson"})
     async def chat_query_ndjson(
         self, kbid: str, item: ChatModel
-    ) -> tuple[str, AsyncIterator[GenerativeChunk]]:
+    ) -> tuple[str, str, AsyncIterator[GenerativeChunk]]:
         """
         Chat query using the new stream format
         Format specs: https://github.com/ndjson/ndjson-spec
@@ -369,35 +285,55 @@ class PredictEngine:
         resp = await self.make_request(
             "POST",
             url=self.get_predict_url(CHAT, kbid),
-            json=item.dict(),
+            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_chat_ndjson_generator(resp)
+        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
 
@@ -442,27 +378,41 @@ 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',
@@ -486,79 +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(
-        self, kbid: str, item: ChatModel
-    ) -> tuple[str, AsyncIterator[bytes]]:
-        self.calls.append(("chat_query", item))
-
-        async def generate():
-            for i in self.generated_answer:
-                yield i
-
-        return (DUMMY_LEARNING_ID, generate())
-
     async def chat_query_ndjson(
         self, kbid: str, item: ChatModel
-    ) -> tuple[str, AsyncIterator[bytes]]:
+    ) -> tuple[str, str, AsyncIterator[GenerativeChunk]]:
         self.calls.append(("chat_query_ndjson", item))
 
         async def generate():
             for item in self.ndjson_answer:
-                yield GenerativeChunk.parse_raw(item)
+                yield GenerativeChunk.model_validate_json(item)
 
-        return (DUMMY_LEARNING_ID, generate())
+        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,
-            )
+
+        if os.environ.get("TEST_SENTENCE_ENCODER") == "multilingual-2023-02-21":  # pragma: no cover
+            base_vector = Qm2023
         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,
-            )
+            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))
@@ -577,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
 
 
@@ -604,11 +554,10 @@ def get_answer_generator(response: aiohttp.ClientResponse):
 def get_chat_ndjson_generator(
     response: aiohttp.ClientResponse,
 ) -> AsyncIterator[GenerativeChunk]:
-
     async def _parse_generative_chunks(gen):
         async for chunk in gen:
             try:
-                yield GenerativeChunk.parse_raw(chunk.strip())
+                yield GenerativeChunk.model_validate_json(chunk.strip())
             except ValidationError as ex:
                 errors.capture_exception(ex)
                 logger.error(f"Invalid chunk received: {chunk}")