nucliadb 6.7.2.post4874__py3-none-any.whl → 6.10.0.post5705__py3-none-any.whl

nucliadb/search/predict.py

@@ -22,9 +22,9 @@ import json
 import logging
 import os
 import random
+from collections.abc import AsyncGenerator
 from dataclasses import dataclass
 from enum import Enum
-from typing import AsyncGenerator, Optional
 from unittest.mock import AsyncMock, Mock
 
 import aiohttp
@@ -144,7 +144,7 @@ class AnswerStatusCode(str, Enum):
 @dataclass
 class RephraseResponse:
     rephrased_query: str
-    use_chat_history: Optional[bool]
+    use_chat_history: bool | None
 
 
 async def start_predict_engine():
@@ -176,18 +176,18 @@ def convert_relations(data: dict[str, list[dict[str, str]]]) -> list[RelationNod
 class PredictEngine:
     def __init__(
         self,
-        cluster_url: Optional[str] = None,
-        public_url: Optional[str] = None,
-        nuclia_service_account: Optional[str] = None,
-        zone: Optional[str] = None,
+        cluster_url: str | None = None,
+        public_url: str | None = None,
+        nuclia_service_account: str | None = None,
+        zone: str | None = None,
         onprem: bool = False,
         local_predict: bool = False,
-        local_predict_headers: Optional[dict[str, str]] = None,
+        local_predict_headers: dict[str, str] | None = None,
     ):
         self.nuclia_service_account = nuclia_service_account
         self.cluster_url = cluster_url
         if public_url is not None:
-            self.public_url: Optional[str] = public_url.format(zone=zone)
+            self.public_url: str | None = public_url.format(zone=zone)
         else:
             self.public_url = None
         self.zone = zone
@@ -294,7 +294,7 @@ class PredictEngine:
 
     @predict_observer.wrap({"type": "chat_ndjson"})
     async def chat_query_ndjson(
-        self, kbid: str, item: ChatModel, extra_headers: Optional[dict[str, str]] = None
+        self, kbid: str, item: ChatModel, extra_headers: dict[str, str] | None = None
     ) -> tuple[str, str, AsyncGenerator[GenerativeChunk, None]]:
         """
         Chat query using the new stream format
@@ -383,7 +383,7 @@ class PredictEngine:
 
     @predict_observer.wrap({"type": "summarize"})
     async def summarize(
-        self, kbid: str, item: SummarizeModel, extra_headers: Optional[dict[str, str]] = None
+        self, kbid: str, item: SummarizeModel, extra_headers: dict[str, str] | None = None
     ) -> SummarizedResponse:
         try:
             self.check_nua_key_is_configured_for_onprem()
@@ -447,6 +447,10 @@ class DummyPredictEngine(PredictEngine):
         self.cluster_url = "http://localhost:8000"
         self.public_url = "http://localhost:8000"
         self.calls = []
+        self.ndjson_reasoning = [
+            b'{"chunk": {"type": "reasoning", "text": "dummy "}}\n',
+            b'{"chunk": {"type": "reasoning", "text": "reasoning"}}\n',
+        ]
         self.ndjson_answer = [
             b'{"chunk": {"type": "text", "text": "valid "}}\n',
             b'{"chunk": {"type": "text", "text": "answer "}}\n',
@@ -477,13 +481,16 @@ class DummyPredictEngine(PredictEngine):
         return RephraseResponse(rephrased_query=DUMMY_REPHRASE_QUERY, use_chat_history=None)
 
     async def chat_query_ndjson(
-        self, kbid: str, item: ChatModel, extra_headers: Optional[dict[str, str]] = None
+        self, kbid: str, item: ChatModel, extra_headers: dict[str, str] | None = None
     ) -> tuple[str, str, AsyncGenerator[GenerativeChunk, None]]:
         self.calls.append(("chat_query_ndjson", item))
 
         async def generate():
-            for item in self.ndjson_answer:
-                yield GenerativeChunk.model_validate_json(item)
+            if item.reasoning is not False:
+                for chunk in self.ndjson_reasoning:
+                    yield GenerativeChunk.model_validate_json(chunk)
+            for chunk in self.ndjson_answer:
+                yield GenerativeChunk.model_validate_json(chunk)
 
         return (DUMMY_LEARNING_ID, DUMMY_LEARNING_MODEL, generate())
 
@@ -517,10 +524,17 @@ class DummyPredictEngine(PredictEngine):
                 timings[vectorset_id] = 0.010
 
         # and fake data with the passed one too
-        model = item.semantic_models[0] if item.semantic_models else "<PREDICT-DEFAULT-SEMANTIC-MODEL>"
-        semantic_thresholds[model] = self.default_semantic_threshold
-        vectors[model] = base_vector
-        timings[model] = 0.0
+        if item.semantic_models is not None:
+            for model in item.semantic_models:
+                semantic_thresholds[model] = self.default_semantic_threshold
+                vectors[model] = base_vector
+                timings[model] = 0.0
+
+        if len(vectors) == 0:
+            model = "<PREDICT-DEFAULT-SEMANTIC-MODEL>"
+            semantic_thresholds[model] = self.default_semantic_threshold
+            vectors[model] = base_vector
+            timings[model] = 0.0
 
         return QueryInfo(
             language="en",
@@ -533,7 +547,7 @@ class DummyPredictEngine(PredictEngine):
                 vectors=vectors,
                 timings=timings,
             ),
-            query=model,
+            query=item.text or "<PREDICT-QUERY>",
             rephrased_query="<REPHRASED-QUERY>" if item.rephrase or item.query_image else None,
         )
 
@@ -546,7 +560,7 @@ class DummyPredictEngine(PredictEngine):
             return DUMMY_RELATION_NODE
 
     async def summarize(
-        self, kbid: str, item: SummarizeModel, extra_headers: Optional[dict[str, str]] = None
+        self, kbid: str, item: SummarizeModel, extra_headers: dict[str, str] | None = None
     ) -> SummarizedResponse:
         self.calls.append(("summarize", (kbid, item)))
         response = SummarizedResponse(

nucliadb/search/predict_models.py

@@ -19,7 +19,6 @@
 
 from base64 import b64decode, b64encode
 from enum import Enum
-from typing import Optional
 
 from google.protobuf.message import DecodeError, Message
 from pydantic import BaseModel, ConfigDict, Field, field_validator
@@ -77,7 +76,7 @@ class RunAgentsRequest(BaseModel):
         default_factory=list,
         title="An optional list of Data Augmentation Agent IDs to run. If empty, all configured agents that match the filters are run.",
     )
-    filters: Optional[list[NameOperationFilter]] = Field(
+    filters: list[NameOperationFilter] | None = Field(
         default=None,
         title="Filters to select which Data Augmentation Agents are applied to the text. If empty, all configured agents for the Knowledge Box are applied.",
     )
@@ -93,7 +92,7 @@ class AppliedDataAugmentation(BaseModel):
         # Since we have protos as fields, we need to enable arbitrary_types_allowed
         arbitrary_types_allowed=True,
     )
-    qas: Optional[QuestionAnswers] = Field(
+    qas: QuestionAnswers | None = Field(
         default=None,
         description="Question and answers generated by the Question Answers agent",
     )
@@ -107,7 +106,7 @@ class AppliedDataAugmentation(BaseModel):
     )
 
     @field_validator("qas", mode="before")
-    def validate_qas(cls, qas: Optional[str]) -> Optional[QuestionAnswers]:
+    def validate_qas(cls, qas: str | None) -> QuestionAnswers | None:
         if qas is None:
             return None
         try:
@@ -171,8 +170,8 @@ class QueryModel(BaseModel):
     Model to represent a query request
     """
 
-    text: Optional[str] = Field(default=None, description="The query text to be processed")
-    query_image: Optional[Image] = Field(
+    text: str | None = Field(default=None, description="The query text to be processed")
+    query_image: Image | None = Field(
         default=None,
         description="Image to be considered as part of the query.  Even if the `rephrase` parameter is set to `false`, the rephrasing process will occur, combining the provided text with the image's visual features in the rephrased query.",
     )
@@ -180,7 +179,7 @@ class QueryModel(BaseModel):
         default=False,
         description="If true, the model will rephrase the input text before processing",
     )
-    rephrase_prompt: Optional[str] = Field(
+    rephrase_prompt: str | None = Field(
         default=None,
         description="Custom prompt for rephrasing the input text",
         examples=[
@@ -192,11 +191,11 @@ QUESTION: {question}
 Please return ONLY the question without any explanation.""",
         ],
     )
-    generative_model: Optional[str] = Field(
+    generative_model: str | None = Field(
         default=None,
         description="The generative model to use for rephrasing",
     )
-    semantic_models: Optional[list[str]] = Field(
+    semantic_models: list[str] | None = Field(
         default=None,
         description="Semantic models to compute the sentence vector for, if not provided, it will only compute the sentence vector for default semantic model in the Knowledge box's configuration.",
     )

nucliadb/search/requesters/utils.py

@@ -19,8 +19,9 @@
 
 import asyncio
 import json
+from collections.abc import Sequence
 from enum import Enum, auto
-from typing import Any, Optional, Sequence, TypeVar, Union, overload
+from typing import TypeVar, overload
 
 from fastapi import HTTPException
 from google.protobuf.json_format import MessageToDict
@@ -60,7 +61,7 @@ METHODS = {
     Method.GRAPH: graph_search_shard,
 }
 
-REQUEST_TYPE = Union[SuggestRequest, SearchRequest, GraphSearchRequest]
+REQUEST_TYPE = SuggestRequest | SearchRequest | GraphSearchRequest
 
 T = TypeVar(
     "T",
@@ -75,7 +76,7 @@ async def nidx_query(
     kbid: str,
     method: Method,
     pb_query: SuggestRequest,
-    timeout: Optional[float] = None,
+    timeout: float | None = None,
 ) -> tuple[list[SuggestResponse], list[str]]: ...
 
 
@@ -84,7 +85,7 @@ async def nidx_query(
     kbid: str,
     method: Method,
     pb_query: SearchRequest,
-    timeout: Optional[float] = None,
+    timeout: float | None = None,
 ) -> tuple[list[SearchResponse], list[str]]: ...
 
 
@@ -93,7 +94,7 @@ async def nidx_query(
     kbid: str,
     method: Method,
     pb_query: GraphSearchRequest,
-    timeout: Optional[float] = None,
+    timeout: float | None = None,
 ) -> tuple[list[GraphSearchResponse], list[str]]: ...
 
 
@@ -101,8 +102,8 @@ async def nidx_query(
     kbid: str,
     method: Method,
     pb_query: REQUEST_TYPE,
-    timeout: Optional[float] = None,
-) -> tuple[Sequence[Union[T, BaseException]], list[str]]:
+    timeout: float | None = None,
+) -> tuple[Sequence[T | BaseException], list[str]]:
     timeout = timeout or settings.search_timeout
     shard_manager = get_shard_manager()
     try:
@@ -133,7 +134,7 @@ async def nidx_query(
         )
 
     try:
-        results: list[Union[T, BaseException]] = await asyncio.wait_for(
+        results: list[T | BaseException] = await asyncio.wait_for(
             asyncio.gather(*ops, return_exceptions=True),
             timeout=timeout,
         )
@@ -159,13 +160,13 @@ async def nidx_query(
     return results, queried_shards
 
 
-def validate_nidx_query_results(results: list[Any]) -> Optional[HTTPException]:
+def validate_nidx_query_results(results: list[T | BaseException]) -> HTTPException | None:
     """
     Validate the results of a nidx query and return an exception if any error is found
 
     Handling of exception is responsibility of caller.
     """
-    if results is None or len(results) == 0:
+    if len(results) == 0:
         return HTTPException(status_code=500, detail=f"Error while executing shard queries. No results.")
 
     for result in results:

nucliadb/search/search/cache.py

@@ -19,9 +19,6 @@
 
 import contextlib
 import logging
-from typing import Optional
-
-import backoff
 
 from nucliadb.common.cache import (
     extracted_text_cache,
@@ -41,26 +38,35 @@ from nucliadb_utils.utilities import get_storage
 logger = logging.getLogger(__name__)
 
 
-async def get_resource(kbid: str, uuid: str) -> Optional[ResourceORM]:
+async def get_resource(kbid: str, uuid: str) -> ResourceORM | None:
     """
     Will try to get the resource from the cache, if it's not there it will fetch it from the ORM and cache it.
     """
     resource_cache = get_resource_cache()
     if resource_cache is None:
         logger.warning("Resource cache not set")
-        return await _orm_get_resource(kbid, uuid)
+        async with get_driver().ro_transaction() as txn:
+            storage = await get_storage(service_name=SERVICE_NAME)
+            kb = KnowledgeBoxORM(txn, storage, kbid)
+            return await kb.get(uuid)
 
     return await resource_cache.get(kbid, uuid)
 
 
-async def _orm_get_resource(kbid: str, uuid: str) -> Optional[ResourceORM]:
-    async with get_driver().ro_transaction() as txn:
-        storage = await get_storage(service_name=SERVICE_NAME)
-        kb = KnowledgeBoxORM(txn, storage, kbid)
-        return await kb.get(uuid)
+async def get_field(kbid: str, field_id: FieldId) -> Field | None:
+    rid = field_id.rid
+    orm_resource = await get_resource(kbid, rid)
+    if orm_resource is None:
+        return None
+    field_obj = await orm_resource.get_field(
+        key=field_id.key,
+        type=field_id.pb_type,
+        load=False,
+    )
+    return field_obj
 
 
-async def get_field_extracted_text(field: Field) -> Optional[ExtractedText]:
+async def get_field_extracted_text(field: Field) -> ExtractedText | None:
     if field.extracted_text is not None:
         return field.extracted_text
 
@@ -74,36 +80,6 @@ async def get_field_extracted_text(field: Field) -> Optional[ExtractedText]:
     return extracted_text
 
 
-@backoff.on_exception(backoff.expo, (Exception,), jitter=backoff.random_jitter, max_tries=3)
-async def field_get_extracted_text(field: Field) -> Optional[ExtractedText]:
-    try:
-        return await field.get_extracted_text()
-    except Exception:
-        logger.warning(
-            "Error getting extracted text for field. Retrying",
-            exc_info=True,
-            extra={
-                "kbid": field.kbid,
-                "resource_id": field.resource.uuid,
-                "field": f"{field.type}/{field.id}",
-            },
-        )
-        raise
-
-
-async def get_extracted_text_from_field_id(kbid: str, field: FieldId) -> Optional[ExtractedText]:
-    rid = field.rid
-    orm_resource = await get_resource(kbid, rid)
-    if orm_resource is None:
-        return None
-    field_obj = await orm_resource.get_field(
-        key=field.key,
-        type=field.pb_type,
-        load=False,
-    )
-    return await get_field_extracted_text(field_obj)
-
-
 @contextlib.contextmanager
 def request_caches():
     """
@@ -115,7 +91,8 @@ def request_caches():
     Makes sure to clean the caches at the end of the context manager.
     >>> with request_caches():
     ...     resource = await get_resource(kbid, uuid)
-    ...     extracted_text = await get_extracted_text_from_field_id(kbid, rid, field_id)
+    ...     field = await get_field(kbid, field_id)
+    ...     extracted_text = await get_field_extracted_text(field)
     """
 
     # This cache size is an arbitrary number, once we have a metric in place and