Merge remote-tracking branch 'origin/main' into add-remi

CHANGELOG.md

@@ -1,11 +1,23 @@
 # Changelog
 
-## 4.3.8 (unreleased)
+## 4.3.10 (unreleased)
 
 
 - Feature: Add NUA REMi endpoints
 
 
+## 4.3.9 (2024-12-11)
+
+
+- Debug on ask and send to process
+
+
+## 4.3.8 (2024-12-09)
+
+
+- Support blanklineSplitter
+
+
 ## 4.3.7 (2024-12-02)
 
 

Makefile

@@ -11,6 +11,7 @@ lint:
 	mypy nuclia
 
 fmt:
+	ruff check --fix nuclia
 	ruff format nuclia
 
 test:

VERSION

@@ -1 +1 @@
-4.3.8.dev0
+4.3.10.dev0

nuclia/sdk/resource.py

@@ -1,18 +1,22 @@
 import json
 from typing import Any, Dict, List, Optional, Union
 from uuid import uuid4
+import os
+
+from nucliadb_sdk.v2 import exceptions
+import backoff
+import requests
+from nucliadb_models.search import AskRequest
+from nuclia.exceptions import RateLimitError
 
 from nucliadb_models.metadata import ResourceProcessingStatus
 from nucliadb_models.resource import Resource
-import os
 
-import requests
 from nuclia import get_list_parameter, get_regional_url
 from nuclia.decorators import kb, pretty
 from nuclia.lib.kb import NucliaDBClient
 from nuclia.sdk.logger import logger
 from nucliadb_models.search import (
-    AskRequest,
     Filter,
     RagStrategies,
     RagImagesStrategies,
@@ -34,6 +38,7 @@
     "summary",
     "metadata",
     "security",
+    "wait_for_commit",
 ]
 
 
@@ -52,6 +57,13 @@ class NucliaResource:
     All commands accept either `rid` or `slug` to identify the targeted resource.
     """
 
+    @backoff.on_exception(
+        backoff.expo,
+        RateLimitError,
+        jitter=backoff.random_jitter,
+        max_tries=5,
+        factor=10,
+    )
     @kb
     def create(*args, **kwargs) -> str:
         ndb = kwargs["ndb"]
@@ -63,7 +75,13 @@ def create(*args, **kwargs) -> str:
         for param in RESOURCE_ATTRIBUTES:
             if param in kwargs:
                 kw[param] = kwargs.get(param)
-        resource = ndb.ndb.create_resource(**kw)
+        try:
+            resource = ndb.ndb.create_resource(**kw)
+        except exceptions.RateLimitError as exc:
+            logger.debug(
+                "Rate limited while trying to create a resource. Waiting a bit before trying again..."
+            )
+            raise RateLimitError() from exc
         rid = resource.uuid
         return rid
 
@@ -233,6 +251,25 @@ def update(
         else:
             raise ValueError("Either rid or slug must be provided")
 
+    @kb
+    def send_to_process(
+        self, *, rid: Optional[str] = None, slug: Optional[str] = None, **kwargs
+    ):
+        ndb = kwargs["ndb"]
+        kw = {
+            "kbid": ndb.kbid,
+        }
+        if rid:
+            kw["rid"] = rid
+            if slug:
+                kw["slug"] = slug
+            ndb.ndb.reprocess_resource(**kw)
+        elif slug:
+            kw["rslug"] = slug
+            ndb.ndb.reprocess_resource_by_slug(**kw)
+        else:
+            raise ValueError("Either rid or slug must be provided")
+
     @kb
     def delete(
         self, *, rid: Optional[str] = None, slug: Optional[str] = None, **kwargs
@@ -355,6 +392,25 @@ async def download_file(
                 if chunk:
                     f.write(chunk)
 
+    @kb
+    async def send_to_process(
+        self, *, rid: Optional[str] = None, slug: Optional[str] = None, **kwargs
+    ):
+        ndb = kwargs["ndb"]
+        kw = {
+            "kbid": ndb.kbid,
+        }
+        if rid:
+            kw["rid"] = rid
+            if slug:
+                kw["slug"] = slug
+            await ndb.ndb.reprocess_resource(**kw)
+        elif slug:
+            kw["rslug"] = slug
+            await ndb.ndb.reprocess_resource_by_slug(**kw)
+        else:
+            raise ValueError("Either rid or slug must be provided")
+
     @kb
     async def update(
         self, *, rid: Optional[str] = None, slug: Optional[str] = None, **kwargs

nuclia/sdk/search.py

@@ -17,6 +17,7 @@
     SearchOptions,
     SearchRequest,
     SyncAskResponse,
+    ChatModel,
 )
 from pydantic import ValidationError
 
@@ -40,6 +41,11 @@ class AskAnswer:
     timings: Optional[Dict[str, float]]
     tokens: Optional[Dict[str, int]]
     retrieval_best_matches: Optional[List[str]]
+    status: Optional[str]
+    prompt_context: Optional[list[str]]
+    relations: Optional[Relations]
+    predict_request: Optional[ChatModel]
+    error_details: Optional[str]
 
     def __str__(self):
         if self.answer:
@@ -189,12 +195,23 @@ def ask(
             timings=None,
             tokens=None,
             object=ask_response.answer_json,
+            status=ask_response.status,
+            error_details=ask_response.error_details,
+            predict_request=ChatModel.model_validate(ask_response.predict_request)
+            if ask_response.predict_request is not None
+            else None,
+            relations=ask_response.relations,
+            prompt_context=ask_response.prompt_context,
         )
+
+        if ask_response.prompt_context:
+            result.prompt_context = ask_response.prompt_context
         if ask_response.metadata is not None:
             if ask_response.metadata.timings is not None:
                 result.timings = ask_response.metadata.timings.model_dump()
             if ask_response.metadata.tokens is not None:
                 result.tokens = ask_response.metadata.tokens.model_dump()
+
         return result
 
     @kb
@@ -264,6 +281,13 @@ def ask_json(
             timings=None,
             tokens=None,
             object=ask_response.answer_json,
+            status=ask_response.status,
+            error_details=ask_response.error_details,
+            predict_request=ChatModel.model_validate(ask_response.predict_request)
+            if ask_response.predict_request is not None
+            else None,
+            relations=ask_response.relations,
+            prompt_context=ask_response.prompt_context,
         )
         if ask_response.metadata is not None:
             if ask_response.metadata.timings is not None:
@@ -399,6 +423,11 @@ async def ask(
             timings=None,
             tokens=None,
             object=None,
+            status=None,
+            error_details=None,
+            predict_request=None,
+            relations=None,
+            prompt_context=None,
         )
         async for line in ask_stream_response.aiter_lines():
             try:
@@ -422,8 +451,16 @@ async def ask(
                 if ask_response_item.tokens:
                     result.tokens = ask_response_item.tokens.model_dump()
             elif ask_response_item.type == "status":
-                # Status is ignored
-                pass
+                result.status = ask_response_item.status
+            elif ask_response_item.type == "prequeries":
+                result.prequeries = ask_response_item.results
+            elif ask_response_item.type == "error":
+                result.error_details = ask_response_item.error
+            elif ask_response_item.type == "debug":
+                result.prompt_context = ask_response_item.metadata.get("prompt_context")
+                result.predict_request = ask_response_item.metadata.get(
+                    "predict_request"
+                )
             else:  # pragma: no cover
                 warnings.warn(f"Unknown ask stream item type: {ask_response_item.type}")
         return result
@@ -512,6 +549,11 @@ async def ask_json(
             timings=None,
             tokens=None,
             object=None,
+            status=None,
+            error_details=None,
+            predict_request=None,
+            relations=None,
+            prompt_context=None,
         )
         async for line in ask_stream_response.aiter_lines():
             try:
@@ -535,8 +577,16 @@ async def ask_json(
                 if ask_response_item.tokens:
                     result.tokens = ask_response_item.tokens.model_dump()
             elif ask_response_item.type == "status":
-                # Status is ignored
-                pass
+                result.status = ask_response_item.status
+            elif ask_response_item.type == "prequeries":
+                result.prequeries = ask_response_item.results
+            elif ask_response_item.type == "error":
+                result.error_details = ask_response_item.error
+            elif ask_response_item.type == "debug":
+                result.prompt_context = ask_response_item.metadata.get("prompt_context")
+                result.predict_request = ask_response_item.metadata.get(
+                    "predict_request"
+                )
             else:  # pragma: no cover
                 warnings.warn(f"Unknown ask stream item type: {ask_response_item.type}")
         return result

nuclia/sdk/upload.py

@@ -62,6 +62,7 @@ def file(
         rid: Optional[str] = None,
         field: Optional[str] = None,
         interpretTables: Optional[bool] = False,
+        blanklineSplitter: Optional[bool] = False,
         mimetype: Optional[str] = None,
         **kwargs,
     ) -> Optional[str]:
@@ -74,6 +75,8 @@ def file(
             mimetype = "application/octet-stream"
         if interpretTables:
             mimetype += "+aitable"
+        if blanklineSplitter:
+            mimetype += "+blankline"
         is_new_resource = False
         if rid is not None:
             rid, is_new_resource = self._get_or_create_resource(
@@ -244,6 +247,7 @@ def remote(
         rid: Optional[str] = None,
         field: Optional[str] = "file",
         interpretTables: Optional[bool] = False,
+        blanklineSplitter: Optional[bool] = False,
         **kwargs,
     ) -> str:
         """Upload a remote url to a Nuclia KnowledgeBox"""
@@ -263,6 +267,8 @@ def remote(
             mimetype = r.headers.get("Content-Type", "application/octet-stream")
             if interpretTables:
                 mimetype += "+aitable"
+            if blanklineSplitter:
+                mimetype += "+blankline"
             rid, is_new_resource = self._get_or_create_resource(
                 rid=rid, icon=mimetype, **kwargs
             )
@@ -368,6 +374,7 @@ async def file(
         field: Optional[str] = None,
         mimetype: Optional[str] = None,
         interpretTables: Optional[bool] = False,
+        blanklineSplitter: Optional[bool] = False,
         **kwargs,
     ) -> str:
         """Upload a file from filesystem to a Nuclia KnowledgeBox"""
@@ -379,6 +386,8 @@ async def file(
             mimetype = "application/octet-stream"
         if interpretTables:
             mimetype += "+aitable"
+        if blanklineSplitter:
+            mimetype += "+blankline"
         rid, is_new_resource = await self._get_or_create_resource(
             rid=rid, icon=mimetype, **kwargs
         )
@@ -542,6 +551,7 @@ async def remote(
         rid: Optional[str] = None,
         field: Optional[str] = "file",
         interpretTables: Optional[bool] = False,
+        blanklineSplitter: Optional[bool] = False,
         **kwargs,
     ) -> str:
         """Upload a remote url to a Nuclia KnowledgeBox"""
@@ -556,6 +566,8 @@ async def remote(
             mimetype = r.headers.get("Content-Type", "application/octet-stream")
             if interpretTables:
                 mimetype += "+aitable"
+            if blanklineSplitter:
+                mimetype += "+blankline"
             rid, is_new_resource = await self._get_or_create_resource(
                 rid=rid, icon=mimetype, **kwargs
             )

requirements.txt

@@ -5,8 +5,8 @@ requests
 httpx
 httpcore>=1.0.0
 prompt_toolkit
-nucliadb_sdk>=6.0.1,<7
-nucliadb_models>=6.0.1,<7
+nucliadb_sdk>=6.1.0,<7
+nucliadb_models>=6.1.0,<7
 nuclia-models>=0.20.1
 tqdm
 aiofiles