nucliadb 6.7.0.post4773__py3-none-any.whl → 6.7.0.post4786__py3-none-any.whl

nucliadb/search/predict.py

@@ -24,7 +24,7 @@ import os
 import random
 from dataclasses import dataclass
 from enum import Enum
-from typing import Any, AsyncGenerator, Optional
+from typing import AsyncGenerator, Optional
 from unittest.mock import AsyncMock, Mock
 
 import aiohttp
@@ -37,6 +37,7 @@ from nucliadb.search import logger
 from nucliadb.search.predict_models import (
     AppliedDataAugmentation,
     AugmentedField,
+    QueryModel,
     RunAgentsRequest,
     RunAgentsResponse,
 )
@@ -326,11 +327,7 @@ class PredictEngine:
     async def query(
         self,
         kbid: str,
-        sentence: str,
-        semantic_model: Optional[str] = None,
-        generative_model: Optional[str] = None,
-        rephrase: bool = False,
-        rephrase_prompt: Optional[str] = None,
+        item: QueryModel,
     ) -> QueryInfo:
         """
         Query endpoint: returns information to be used by NucliaDB at retrieval time, for instance:
@@ -354,21 +351,10 @@ class PredictEngine:
             logger.warning(error)
             raise SendToPredictError(error)
 
-        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
-
         resp = await self.make_request(
-            "GET",
+            "POST",
             url=self.get_predict_url(QUERY, kbid),
-            params=params,
+            json=item.model_dump(),
             headers=self.get_predict_headers(kbid),
         )
         await self.check_response(kbid, resp, expected_status=200)
@@ -501,16 +487,13 @@ class DummyPredictEngine(PredictEngine):
 
         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: bool = False,
-        rephrase_prompt: Optional[str] = None,
-    ) -> QueryInfo:
-        self.calls.append(("query", sentence))
+    async def query(self, kbid: str, item: QueryModel) -> QueryInfo:
+        self.calls.append(
+            (
+                "query",
+                item,
+            )
+        )
 
         if os.environ.get("TEST_SENTENCE_ENCODER") == "multilingual-2023-02-21":  # pragma: no cover
             base_vector = Qm2023
@@ -534,7 +517,7 @@ class DummyPredictEngine(PredictEngine):
                 timings[vectorset_id] = 0.010
 
         # and fake data with the passed one too
-        model = semantic_model or "<PREDICT-DEFAULT-SEMANTIC-MODEL>"
+        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
@@ -550,7 +533,7 @@ class DummyPredictEngine(PredictEngine):
                 vectors=vectors,
                 timings=timings,
             ),
-            query=sentence,
+            query=model,
         )
 
     async def detect_entities(self, kbid: str, sentence: str) -> list[RelationNode]:

nucliadb/search/predict_models.py

@@ -25,6 +25,7 @@ from google.protobuf.message import DecodeError, Message
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 from typing_extensions import TypedDict
 
+from nucliadb_models.search import Image
 from nucliadb_protos.resources_pb2 import FieldMetadata, FieldText, QuestionAnswers
 
 
@@ -163,3 +164,43 @@ class RunAgentsResponse(BaseModel):
         ...,
         title="Pairs of augmented FieldMetadata and Data Augmentation results by field id",
     )
+
+
+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(
+        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.",
+    )
+    rephrase: bool = Field(
+        default=False,
+        description="If true, the model will rephrase the input text before processing",
+    )
+    rephrase_prompt: Optional[str] = Field(
+        default=None,
+        description="Custom prompt for rephrasing the input text",
+        examples=[
+            """Rephrase this question so its better for retrieval, and keep the rephrased question in the same language as the original.
+QUESTION: {question}
+Please return ONLY the question without any explanation. Just the rephrased question.""",
+            """Rephrase this question so its better for retrieval, if in the image there are any machinery components with numeric identifiers, append them to the end of the question separated by a commas.
+QUESTION: {question}
+Please return ONLY the question without any explanation.""",
+        ],
+    )
+    generative_model: Optional[str] = Field(
+        default=None,
+        description="The generative model to use for rephrasing",
+    )
+    semantic_models: Optional[list[str]] = 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.",
+    )
+    agentic_entities: bool = Field(
+        default=False,
+        description="If true, the model will return the entities detected in the sentence guided by an already defined Graph Extraction Agent in the Knowledge Box.",
+    )

nucliadb/search/search/chat/ask.py

@@ -595,6 +595,7 @@ async def ask(
             image_strategies=ask_request.rag_images_strategies,
             max_context_characters=tokens_to_chars(generation.max_context_tokens),
             visual_llm=generation.use_visual_llm,
+            query_image=ask_request.query_image,
             metrics=metrics.child_span("context_building"),
         )
         (

nucliadb/search/search/chat/prompt.py

@@ -1042,6 +1042,7 @@ class PromptContextBuilder:
         image_strategies: Optional[Sequence[ImageRagStrategy]] = None,
         max_context_characters: Optional[int] = None,
         visual_llm: bool = False,
+        query_image: Optional[Image] = None,
         metrics: Metrics = Metrics("prompt_context_builder"),
     ):
         self.kbid = kbid
@@ -1054,6 +1055,7 @@ class PromptContextBuilder:
         self.max_context_characters = max_context_characters
         self.visual_llm = visual_llm
         self.metrics = metrics
+        self.query_image = query_image
         self.augmented_context = AugmentedContext(paragraphs={}, fields={})
 
     def prepend_user_context(self, context: CappedPromptContext):
@@ -1061,8 +1063,12 @@ class PromptContextBuilder:
         # it is added first, followed by the found text blocks in order of relevance
         for i, text_block in enumerate(self.user_context or []):
             context[f"USER_CONTEXT_{i}"] = text_block
-        for i, image in enumerate(self.user_image_context or []):
-            context.images[f"USER_IMAGE_CONTEXT_{i}"] = image
+        # Add the query image as part of the image context
+        if self.query_image is not None:
+            context.images["QUERY_IMAGE"] = self.query_image
+        else:
+            for i, image in enumerate(self.user_image_context or []):
+                context.images[f"USER_IMAGE_CONTEXT_{i}"] = image
 
     async def build(
         self,
@@ -1070,7 +1076,7 @@ class PromptContextBuilder:
         ccontext = CappedPromptContext(max_size=self.max_context_characters)
         self.prepend_user_context(ccontext)
         await self._build_context(ccontext)
-        if self.visual_llm:
+        if self.visual_llm and not self.query_image:
             await self._build_context_images(ccontext)
         context = ccontext.cap()
         context_images = ccontext.images

nucliadb/search/search/chat/query.py

@@ -206,6 +206,7 @@ def find_request_from_ask_request(item: AskRequest, query: str) -> FindRequest:
     find_request.debug = item.debug
     find_request.rephrase = item.rephrase
     find_request.rephrase_prompt = parse_rephrase_prompt(item)
+    find_request.query_image = item.query_image
     find_request.rank_fusion = item.rank_fusion
     find_request.reranker = item.reranker
     # We don't support pagination, we always get the top_k results.

nucliadb/search/search/query.py

@@ -85,6 +85,7 @@ async def paragraph_query_to_pb(
         rephrase=False,
         rephrase_prompt=None,
         generative_model=None,
+        query_image=None,
     )
     field_expr, paragraph_expr = await parse_old_filters(old, fetcher)
     if field_expr is not None:
@@ -225,6 +226,7 @@ async def suggest_query_to_pb(
         rephrase=False,
         rephrase_prompt=None,
         generative_model=None,
+        query_image=None,
     )
     field_expr, _ = await parse_old_filters(old, fetcher)
     if field_expr is not None and filter_expression is not None:

nucliadb/search/search/query_parser/fetcher.py

@@ -28,12 +28,14 @@ from nucliadb.common.exceptions import InvalidQueryError
 from nucliadb.common.maindb.utils import get_driver
 from nucliadb.search import logger
 from nucliadb.search.predict import SendToPredictError, convert_relations
+from nucliadb.search.predict_models import QueryModel
 from nucliadb.search.search.metrics import (
     query_parse_dependency_observer,
 )
 from nucliadb.search.utilities import get_predict
 from nucliadb_models.internal.predict import QueryInfo
 from nucliadb_models.search import (
+    Image,
     MaxTokens,
 )
 from nucliadb_protos import knowledgebox_pb2, utils_pb2
@@ -93,6 +95,7 @@ class Fetcher:
         rephrase: bool,
         rephrase_prompt: Optional[str],
         generative_model: Optional[str],
+        query_image: Optional[Image],
     ):
         self.kbid = kbid
         self.query = query
@@ -102,6 +105,7 @@ class Fetcher:
         self.rephrase = rephrase
         self.rephrase_prompt = rephrase_prompt
         self.generative_model = generative_model
+        self.query_image = query_image
 
         self.cache = FetcherCache()
         self.locks: dict[str, asyncio.Lock] = {}
@@ -327,6 +331,7 @@ class Fetcher:
                     self.generative_model,
                     self.rephrase,
                     self.rephrase_prompt,
+                    self.query_image,
                 )
             except (SendToPredictError, TimeoutError):
                 query_info = None
@@ -360,9 +365,18 @@ async def query_information(
     generative_model: Optional[str] = None,
     rephrase: bool = False,
     rephrase_prompt: Optional[str] = None,
+    query_image: Optional[Image] = None,
 ) -> QueryInfo:
     predict = get_predict()
-    return await predict.query(kbid, query, semantic_model, generative_model, rephrase, rephrase_prompt)
+    item = QueryModel(
+        text=query,
+        semantic_models=[semantic_model] if semantic_model else None,
+        generative_model=generative_model,
+        rephrase=rephrase,
+        rephrase_prompt=rephrase_prompt,
+        query_image=query_image,
+    )
+    return await predict.query(kbid, item)
 
 
 @query_parse_dependency_observer.wrap({"type": "detect_entities"})

nucliadb/search/search/query_parser/parsers/ask.py

@@ -41,6 +41,7 @@ def fetcher_for_ask(kbid: str, item: AskRequest) -> Fetcher:
         rephrase=item.rephrase,
         rephrase_prompt=None,
         generative_model=item.generative_model,
+        query_image=item.query_image,
     )
 
 

nucliadb/search/search/query_parser/parsers/find.py

@@ -83,6 +83,7 @@ def fetcher_for_find(kbid: str, item: FindRequest) -> Fetcher:
         rephrase=item.rephrase,
         rephrase_prompt=item.rephrase_prompt,
         generative_model=item.generative_model,
+        query_image=item.query_image,
     )
 
 

nucliadb/search/search/query_parser/parsers/search.py

@@ -79,6 +79,7 @@ def fetcher_for_search(kbid: str, item: SearchRequest) -> Fetcher:
         rephrase=item.rephrase,
         rephrase_prompt=item.rephrase_prompt,
         generative_model=None,
+        query_image=None,
     )
 
 

{nucliadb-6.7.0.post4773.dist-info → nucliadb-6.7.0.post4786.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nucliadb
-Version: 6.7.0.post4773
+Version: 6.7.0.post4786
 Summary: NucliaDB
 Author-email: Nuclia <nucliadb@nuclia.com>
 License-Expression: AGPL-3.0-or-later
@@ -19,11 +19,11 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Requires-Python: <4,>=3.9
 Description-Content-Type: text/markdown
-Requires-Dist: nucliadb-telemetry[all]>=6.7.0.post4773
-Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.7.0.post4773
-Requires-Dist: nucliadb-protos>=6.7.0.post4773
-Requires-Dist: nucliadb-models>=6.7.0.post4773
-Requires-Dist: nidx-protos>=6.7.0.post4773
+Requires-Dist: nucliadb-telemetry[all]>=6.7.0.post4786
+Requires-Dist: nucliadb-utils[cache,fastapi,storages]>=6.7.0.post4786
+Requires-Dist: nucliadb-protos>=6.7.0.post4786
+Requires-Dist: nucliadb-models>=6.7.0.post4786
+Requires-Dist: nidx-protos>=6.7.0.post4786
 Requires-Dist: nucliadb-admin-assets>=1.0.0.post1224
 Requires-Dist: nuclia-models>=0.46.0
 Requires-Dist: uvicorn[standard]

{nucliadb-6.7.0.post4773.dist-info → nucliadb-6.7.0.post4786.dist-info}/RECORD

@@ -213,8 +213,8 @@ nucliadb/search/__init__.py,sha256=tnypbqcH4nBHbGpkINudhKgdLKpwXQCvDtPchUlsyY4,1
 nucliadb/search/app.py,sha256=-WEX1AZRA8R_9aeOo9ovOTwjXW_7VfwWN7N2ccSoqXg,3387
 nucliadb/search/lifecycle.py,sha256=hiylV-lxsAWkqTCulXBg0EIfMQdejSr8Zar0L_GLFT8,2218
 nucliadb/search/openapi.py,sha256=t3Wo_4baTrfPftg2BHsyLWNZ1MYn7ZRdW7ht-wFOgRs,1016
-nucliadb/search/predict.py,sha256=xZtZaydg1pzXOSEDg0xyWNbbgA4zMQ59gbHi0wNuAxk,23770
-nucliadb/search/predict_models.py,sha256=pm4ykuWH9bTXxj5RlI2F6pmXSXOVt64WL_sRlc2u6Tk,6144
+nucliadb/search/predict.py,sha256=HK8E2N0Uru9FdVsNs8mbhKocN0m5kj32x4NtT0NZLRc,23111
+nucliadb/search/predict_models.py,sha256=ozuQZGWAK7v8W6UOk3xXQ_zW7YUNtRy_4l5LE5BuT_A,8172
 nucliadb/search/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nucliadb/search/run.py,sha256=aFb-CXRi_C8YMpP_ivNj8KW1BYhADj88y8K9Lr_nUPI,1402
 nucliadb/search/settings.py,sha256=vem3EcyYlTPSim0kEK-xe-erF4BZg0CT_LAb8ZRQAE8,1684
@@ -257,30 +257,30 @@ nucliadb/search/search/metrics.py,sha256=3I6IN0qDSmqIvUaWJmT3rt-Jyjs6LcvnKI8ZqCi
 nucliadb/search/search/paragraphs.py,sha256=pNAEiYqJGGUVcEf7xf-PFMVqz0PX4Qb-WNG-_zPGN2o,7799
 nucliadb/search/search/pgcatalog.py,sha256=0n_gDihZZhqrDLRHvHzS3IESvMRTcU6YShqizQMyE_Y,16807
 nucliadb/search/search/predict_proxy.py,sha256=Df8F5K-oS4TIXJc_y8UDViJTo7st5L0kMgxYPFZ39Vk,8806
-nucliadb/search/search/query.py,sha256=lIF1TXA35kTKC9iVJuMXDiiqgIj_BjxdG0oTOVn4kaI,11500
+nucliadb/search/search/query.py,sha256=lYCesbUv-B7IyVFQoCCurcxl_Azc5nq3jtVQJ9tk1Ao,11552
 nucliadb/search/search/rank_fusion.py,sha256=xZtXhbmKb_56gs73u6KkFm2efvTATOSMmpOV2wrAIqE,9613
 nucliadb/search/search/rerankers.py,sha256=E2J1QdKAojqbhHM3KAyaOXKf6tJyETUxKs4tf_BEyqk,7472
 nucliadb/search/search/shards.py,sha256=mc5DK-MoCv9AFhlXlOFHbPvetcyNDzTFOJ5rimK8PC8,2636
 nucliadb/search/search/summarize.py,sha256=S4-mUS8d-rvHFcsr8Pa8N5NTxU6ZTxLFZTMKTTOOpr4,5098
 nucliadb/search/search/utils.py,sha256=ajRIXfdTF67dBVahQCXW-rSv6gJpUMPt3QhJrWqArTQ,2175
 nucliadb/search/search/chat/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
-nucliadb/search/search/chat/ask.py,sha256=vJ3TSdr-cT_xh43UnoYugqxnHv_-LFSCYoU7o0NnI1M,39368
+nucliadb/search/search/chat/ask.py,sha256=c8gvGkspW5KUbDhVzcUfd7eF7SMfqzwKfCPidJniRLI,39417
 nucliadb/search/search/chat/exceptions.py,sha256=Siy4GXW2L7oPhIR86H3WHBhE9lkV4A4YaAszuGGUf54,1356
 nucliadb/search/search/chat/images.py,sha256=PA8VWxT5_HUGfW1ULhKTK46UBsVyINtWWqEM1ulzX1E,3095
-nucliadb/search/search/chat/prompt.py,sha256=gmYRC3aK03vrDoBElJP5H5Z7OEeu79k5yTxv3FEkN0I,53866
-nucliadb/search/search/chat/query.py,sha256=3jMPNbiFEOoS0ydMOPYkSx1qVlvAv51npzadWXDwkMs,16650
+nucliadb/search/search/chat/prompt.py,sha256=G8Z-8amdx7527-I17U4j80Tv8j9XGt2soYn3txsAWAE,54158
+nucliadb/search/search/chat/query.py,sha256=AhOPMf68p2BRjKz7CdkcUIDMANtxr00oGt42iKUUjAw,16698
 nucliadb/search/search/query_parser/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/search/search/query_parser/exceptions.py,sha256=sVl9gRNzhE-s480LBBVkiXzNRbKhYRQN5F3it5tNNp8,939
-nucliadb/search/search/query_parser/fetcher.py,sha256=nP4EySj2BvH10QgCvgzvp13Nf22wwfHsdLbDoPlH2cQ,16831
+nucliadb/search/search/query_parser/fetcher.py,sha256=dB_nQHbhdlAmztWqtuQ6Xx8UV1kozzGnpocrWwzW-X0,17253
 nucliadb/search/search/query_parser/models.py,sha256=kAslqX_-zaIdUpcpdNU2a5uQPQh7LC605qWLZ4aZ5T4,5064
 nucliadb/search/search/query_parser/old_filters.py,sha256=GsU3T3-WiSPvjucP7evHkshzAWZOli8qsuXChvWRCY0,9092
 nucliadb/search/search/query_parser/parsers/__init__.py,sha256=ySCNSdbesLXGZyR88919njulA6UE10_3PhqMG_Yj1o4,1034
-nucliadb/search/search/query_parser/parsers/ask.py,sha256=eTz8wS-EJHuAagR384h6TT64itymFZRpfZJGX8r6aZM,2771
+nucliadb/search/search/query_parser/parsers/ask.py,sha256=mc4Rx7jLM101tRZYOxEQC-txO5C-fDJid_oMAtZTRug,2809
 nucliadb/search/search/query_parser/parsers/catalog.py,sha256=JuDiBL2wdjAuEFEPo0e2nQ4VqWjF3FXakT0ziZk3Oes,7495
 nucliadb/search/search/query_parser/parsers/common.py,sha256=mJMPOKurBK7-A7s3oNlPLxHP_yIn4j5Uw8rh_OQtzS4,6339
-nucliadb/search/search/query_parser/parsers/find.py,sha256=lHVspg-i_eWXvu7BT9WfuFVGVKYhr380y4tDX5yfTD4,12735
+nucliadb/search/search/query_parser/parsers/find.py,sha256=4xQwa0BxNucenUrW_iZ2jCGd15Dm4AKS_B91BE8sDi4,12773
 nucliadb/search/search/query_parser/parsers/graph.py,sha256=zyqdUg5Afmhb2_-hvj9FUCaoLh026MUP1fgY2j-lD7c,9385
-nucliadb/search/search/query_parser/parsers/search.py,sha256=huhz3lk6y4n7fcHU2XB-90Q34sXIwP5-cwp2rqFtPh8,10477
+nucliadb/search/search/query_parser/parsers/search.py,sha256=78KSJ9t3I7nFVY2Qk2fMw2P1RHUdGRsWzBf59FdAeTA,10503
 nucliadb/search/search/query_parser/parsers/unit_retrieval.py,sha256=xiOQ7_X6MkcZs3W_0DjdVfyk-G1AY6RBx3oG5hsq7ig,11455
 nucliadb/standalone/__init__.py,sha256=cp15ZcFnHvpcu_5-aK2A4uUyvuZVV_MJn4bIXMa20ks,835
 nucliadb/standalone/api_router.py,sha256=zRSMlaRVHUDGTYA3zC03UV_aLLn-ch-kaeWn1tEjTXw,4338
@@ -376,8 +376,8 @@ nucliadb/writer/tus/local.py,sha256=7jYa_w9b-N90jWgN2sQKkNcomqn6JMVBOVeDOVYJHto,
 nucliadb/writer/tus/s3.py,sha256=vu1BGg4VqJ_x2P1u2BxqPKlSfw5orT_a3R-Ln5oPUpU,8483
 nucliadb/writer/tus/storage.py,sha256=ToqwjoYnjI4oIcwzkhha_MPxi-k4Jk3Lt55zRwaC1SM,2903
 nucliadb/writer/tus/utils.py,sha256=MSdVbRsRSZVdkaum69_0wku7X3p5wlZf4nr6E0GMKbw,2556
-nucliadb-6.7.0.post4773.dist-info/METADATA,sha256=QhixG_CdwX3ap13C7NraJfpH1mbqsYln_iin1WQz2oc,4158
-nucliadb-6.7.0.post4773.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nucliadb-6.7.0.post4773.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
-nucliadb-6.7.0.post4773.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
-nucliadb-6.7.0.post4773.dist-info/RECORD,,
+nucliadb-6.7.0.post4786.dist-info/METADATA,sha256=oU2Nm_poSfw_SxLfe0UStdqGzHQHe1Qul2NfXVtUfbs,4158
+nucliadb-6.7.0.post4786.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nucliadb-6.7.0.post4786.dist-info/entry_points.txt,sha256=XqGfgFDuY3zXQc8ewXM2TRVjTModIq851zOsgrmaXx4,1268
+nucliadb-6.7.0.post4786.dist-info/top_level.txt,sha256=hwYhTVnX7jkQ9gJCkVrbqEG1M4lT2F_iPQND1fCzF80,20
+nucliadb-6.7.0.post4786.dist-info/RECORD,,