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

nucliadb/search/{tests/unit/test_run.py → search/chat/exceptions.py}

@@ -17,23 +17,24 @@
 # 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 unittest.mock import patch
 
-import pytest
 
-from nucliadb.search.app import application
-from nucliadb.search.run import run
+from typing import Optional
 
+from nucliadb_models.search import KnowledgeboxFindResults, PreQueryResult
 
-@pytest.fixture(scope="function")
-def run_fastapi_with_metrics():
-    with patch("nucliadb.search.run.run_fastapi_with_metrics") as mocked, patch(
-        "nucliadb.search.run.instrument_app"
-    ):
-        yield mocked
 
+class NoRetrievalResultsError(Exception):
+    def __init__(
+        self,
+        main: Optional[KnowledgeboxFindResults] = None,
+        prequeries: Optional[list[PreQueryResult]] = None,
+        prefilters: Optional[list[PreQueryResult]] = None,
+    ):
+        self.main_query = main
+        self.prequeries = prequeries
+        self.prefilters = prefilters
 
-def test_run_with_metrics(run_fastapi_with_metrics):
-    run()
 
-    run_fastapi_with_metrics.assert_called_once_with(application)
+class AnswerJsonSchemaTooLong(Exception):
+    pass

nucliadb/search/search/chat/images.py

@@ -18,43 +18,67 @@
 # along with this program. If not, see <http://www.gnu.org/licenses/>.
 
 import base64
+from io import BytesIO
+from typing import Optional
 
+from nucliadb.common.ids import ParagraphId
+from nucliadb.ingest.fields.file import File
 from nucliadb.search import SERVICE_NAME
 from nucliadb_models.search import Image
+from nucliadb_utils.storages.storage import Storage
 from nucliadb_utils.utilities import get_storage
 
 
-async def get_page_image(kbid: str, paragraph_id: str, page: int) -> Image:
+async def get_page_image(kbid: str, paragraph_id: ParagraphId, page_number: int) -> Optional[Image]:
     storage = await get_storage(service_name=SERVICE_NAME)
-
-    rid, field_type_letter, field_id, _ = paragraph_id.split("/")[:4]
-
     sf = storage.file_extracted(
-        kbid, rid, field_type_letter, field_id, f"generated/extracted_images_{page}.png"
+        kbid=kbid,
+        uuid=paragraph_id.rid,
+        field_type=paragraph_id.field_id.type,
+        field=paragraph_id.field_id.key,
+        key=f"generated/extracted_images_{page_number}.png",
     )
+    image_bytes = (await sf.storage.downloadbytes(sf.bucket, sf.key)).read()
+    if not image_bytes:
+        return None
     image = Image(
-        b64encoded=base64.b64encode(
-            (await sf.storage.downloadbytes(sf.bucket, sf.key)).read()
-        ).decode(),
+        b64encoded=base64.b64encode(image_bytes).decode(),
         content_type="image/png",
     )
-
     return image
 
 
-async def get_paragraph_image(kbid: str, paragraph_id: str, reference: str) -> Image:
+async def get_paragraph_image(kbid: str, paragraph_id: ParagraphId, reference: str) -> Optional[Image]:
     storage = await get_storage(service_name=SERVICE_NAME)
-
-    rid, field_type_letter, field_id, _ = paragraph_id.split("/")[:4]
-
     sf = storage.file_extracted(
-        kbid, rid, field_type_letter, field_id, f"generated/{reference}"
+        kbid=kbid,
+        uuid=paragraph_id.rid,
+        field_type=paragraph_id.field_id.type,
+        field=paragraph_id.field_id.key,
+        key=f"generated/{reference}",
     )
+    image_bytes = (await sf.storage.downloadbytes(sf.bucket, sf.key)).read()
+    if not image_bytes:
+        return None
     image = Image(
-        b64encoded=base64.b64encode(
-            (await sf.storage.downloadbytes(sf.bucket, sf.key)).read()
-        ).decode(),
+        b64encoded=base64.b64encode(image_bytes).decode(),
         content_type="image/png",
     )
+    return image
+
 
+async def get_file_thumbnail_image(file: File) -> Optional[Image]:
+    fed = await file.get_file_extracted_data()
+    if fed is None or not fed.HasField("file_thumbnail"):
+        return None
+    storage: Storage = await get_storage(service_name=SERVICE_NAME)
+    image_bytes: BytesIO = await storage.downloadbytescf(fed.file_thumbnail)
+    value = image_bytes.getvalue()
+    if len(value) == 0:
+        return None
+    image = Image(
+        b64encoded=base64.b64encode(value).decode(),
+        # We assume the thumbnail is always generated as jpeg by Nuclia processing
+        content_type="image/jpeg",
+    )
     return image