nuclia 4.9.0__py3-none-any.whl → 4.9.2__py3-none-any.whl

nuclia/lib/kb.py

@@ -23,6 +23,7 @@ from nuclia_models.events.activity_logs import (  # type: ignore
 )
 from nuclia_models.events.remi import RemiQuery
 from nuclia_models.worker.tasks import TaskStartKB
+from nuclia_models.config.proto import ExtractConfig
 from nuclia.exceptions import RateLimitError
 from nuclia.lib.utils import handle_http_sync_errors, handle_http_async_errors
 from datetime import datetime
@@ -58,6 +59,8 @@ STOP_TASK = "/task/{task_id}/stop"
 DELETE_TASK = "/task/{task_id}"
 GET_TASK = "/task/{task_id}/inspect"
 RESTART_TASK = "/task/{task_id}/restart"
+EXTRACT_STRATEGIES = "/extract_strategies"
+DELETE_EXTRACT_STRATEGY = "/extract_strategies/strategy/{id}"
 
 DOWNLOAD_FORMAT_HEADERS = {
     DownloadFormat.CSV: "text/csv",
@@ -487,6 +490,37 @@ class NucliaDBClient(BaseNucliaDBClient):
         handle_http_sync_errors(response)
         return response
 
+    def list_extract_strategies(self) -> httpx.Response:
+        if self.reader_session is None:
+            raise Exception("KB not configured")
+
+        response: httpx.Response = self.reader_session.get(
+            f"{self.url}{EXTRACT_STRATEGIES}"
+        )
+        handle_http_sync_errors(response)
+        return response
+
+    def add_extract_strategy(self, config: ExtractConfig) -> httpx.Response:
+        if self.writer_session is None:
+            raise Exception("KB not configured")
+
+        response: httpx.Response = self.writer_session.post(
+            f"{self.url}{EXTRACT_STRATEGIES}",
+            json=config.model_dump(mode="json", exclude_unset=True),
+        )
+        handle_http_sync_errors(response)
+        return response
+
+    def delete_extract_strategy(self, strategy_id: str) -> httpx.Response:
+        if self.writer_session is None:
+            raise Exception("KB not configured")
+
+        response: httpx.Response = self.writer_session.delete(
+            f"{self.url}{DELETE_EXTRACT_STRATEGY.format(id=strategy_id)}",
+        )
+        handle_http_sync_errors(response)
+        return response
+
 
 class AsyncNucliaDBClient(BaseNucliaDBClient):
     reader_session: Optional[httpx.AsyncClient] = None
@@ -813,3 +847,34 @@ class AsyncNucliaDBClient(BaseNucliaDBClient):
         )
         await handle_http_async_errors(response)
         return response
+
+    async def list_extract_strategies(self) -> httpx.Response:
+        if self.reader_session is None:
+            raise Exception("KB not configured")
+
+        response: httpx.Response = await self.reader_session.get(
+            f"{self.url}{EXTRACT_STRATEGIES}"
+        )
+        await handle_http_async_errors(response)
+        return response
+
+    async def add_extract_strategy(self, config: ExtractConfig) -> httpx.Response:
+        if self.writer_session is None:
+            raise Exception("KB not configured")
+
+        response: httpx.Response = await self.writer_session.post(
+            f"{self.url}{EXTRACT_STRATEGIES}",
+            json=config.model_dump(mode="json", exclude_unset=True),
+        )
+        await handle_http_async_errors(response)
+        return response
+
+    async def delete_extract_strategy(self, strategy_id: str) -> httpx.Response:
+        if self.writer_session is None:
+            raise Exception("KB not configured")
+
+        response: httpx.Response = await self.writer_session.delete(
+            f"{self.url}{DELETE_EXTRACT_STRATEGY.format(id=strategy_id)}",
+        )
+        await handle_http_async_errors(response)
+        return response

nuclia/lib/nua_responses.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from enum import Enum
+from enum import Enum, IntEnum
 from typing import Any, Dict, List, Optional, Union, cast
 
 import pydantic
@@ -486,17 +486,37 @@ class ChatResponse(BaseModel):
         return ChatResponse(answer=obj.decode())
 
 
+class SimilarityFunction(IntEnum):
+    # Keep this in sync with SimilarityFunction in config.proto
+    # It's an IntEnum to match the protobuf definition
+    DOT = 0
+    COSINE = 1
+
+
+class SemanticConfig(BaseModel):
+    # Keep this in sync with SemanticConfig in config.proto
+    similarity: SimilarityFunction
+    size: int
+    threshold: float
+    max_tokens: Optional[int] = None
+    matryoshka_dims: list[int] = []
+    external: bool = False
+
+
 class StoredLearningConfiguration(BaseModel):
     resource_labelers_models: Optional[List[str]] = None
     paragraph_labelers_models: Optional[List[str]] = None
     intent_models: Optional[List[str]] = None
 
+    default_semantic_model: Optional[str] = None
+    semantic_models: Optional[list[str]] = Field(default_factory=list[str])
+    semantic_model_configs: dict[str, SemanticConfig] = {}
     semantic_model: str
     anonymization_model: str
     generative_model: str
     ner_model: str
     relation_model: str
-    visual_labeling: str
+    visual_labeling: Optional[str] = None
 
     user_keys: Optional[UserLearningKeys] = None
     user_prompts: Optional[UserPrompts] = None

nuclia/sdk/extract_strategy.py

@@ -0,0 +1,98 @@
+from nuclia.data import get_auth, get_async_auth
+from nuclia.decorators import kb, pretty
+from nuclia.lib.kb import NucliaDBClient, AsyncNucliaDBClient
+from nuclia.sdk.auth import NucliaAuth, AsyncNucliaAuth
+from nuclia_models.config.proto import ExtractConfig
+from typing import Dict
+
+
+class NucliaExtractStrategy:
+    @property
+    def _auth(self) -> NucliaAuth:
+        auth = get_auth()
+        return auth
+
+    @kb
+    @pretty
+    def list(self, *args, **kwargs) -> Dict[str, ExtractConfig]:
+        """
+        List extract strategies
+        """
+        ndb: NucliaDBClient = kwargs["ndb"]
+        response = ndb.list_extract_strategies()
+        return response.json()
+
+    @kb
+    def add(
+        self,
+        *args,
+        config: ExtractConfig,
+        **kwargs,
+    ) -> str:
+        """
+        Add extract strategy
+
+        :param config: strategy configuration
+        """
+        if isinstance(config, dict):
+            config = ExtractConfig.model_validate(config)
+
+        ndb: NucliaDBClient = kwargs["ndb"]
+        response = ndb.add_extract_strategy(config=config)
+        return response.json()
+
+    @kb
+    def delete(self, *args, id: str, **kwargs):
+        """
+        Delete extract strategy
+
+        :param id: ID of the strategy to delete
+        """
+        ndb: NucliaDBClient = kwargs["ndb"]
+        ndb.delete_extract_strategy(strategy_id=id)
+
+
+class AsyncNucliaExtractStrategy:
+    @property
+    def _auth(self) -> AsyncNucliaAuth:
+        auth = get_async_auth()
+        return auth
+
+    @kb
+    @pretty
+    async def list(self, *args, **kwargs) -> Dict[str, ExtractConfig]:
+        """
+        List extract strategies
+        """
+        ndb: AsyncNucliaDBClient = kwargs["ndb"]
+        response = await ndb.list_extract_strategies()
+        return response.json()
+
+    @kb
+    async def add(
+        self,
+        *args,
+        config: ExtractConfig,
+        **kwargs,
+    ) -> str:
+        """
+        Add extract strategy
+
+        :param config: strategy configuration
+        """
+        if isinstance(config, dict):
+            config = ExtractConfig.model_validate(config)
+
+        ndb: AsyncNucliaDBClient = kwargs["ndb"]
+        response = await ndb.add_extract_strategy(config=config)
+        return response.json()
+
+    @kb
+    async def delete(self, *args, id: str, **kwargs):
+        """
+        Delete extract strategy
+
+        :param id: ID of the strategy to delete
+        """
+        ndb: AsyncNucliaDBClient = kwargs["ndb"]
+        await ndb.delete_extract_strategy(strategy_id=id)

nuclia/sdk/kb.py

@@ -30,6 +30,10 @@ from nuclia.sdk.remi import NucliaRemi, AsyncNucliaRemi
 from nuclia.sdk.resource import AsyncNucliaResource, NucliaResource
 from nuclia.sdk.search import AsyncNucliaSearch, NucliaSearch
 from nuclia.sdk.task import NucliaTask, AsyncNucliaTask
+from nuclia.sdk.extract_strategy import (
+    NucliaExtractStrategy,
+    AsyncNucliaExtractStrategy,
+)
 from nuclia.sdk.upload import AsyncNucliaUpload, NucliaUpload
 
 
@@ -48,6 +52,7 @@ class NucliaKB:
         self.logs = NucliaLogs()
         self.task = NucliaTask()
         self.remi = NucliaRemi()
+        self.extract_strategies = NucliaExtractStrategy()
 
     @kb
     def list(
@@ -459,6 +464,7 @@ class AsyncNucliaKB:
         self.logs = AsyncNucliaLogs()
         self.task = AsyncNucliaTask()
         self.remi = AsyncNucliaRemi()
+        self.extract_strategies = AsyncNucliaExtractStrategy()
 
     @kb
     async def list(self, **kwargs) -> ResourceList:

nuclia/tests/test_kb/test_extract_strategies.py

@@ -0,0 +1,16 @@
+from nuclia.sdk.kb import NucliaKB
+
+
+def test_extract_strategies(testing_config):
+    nkb = NucliaKB()
+    # preventive clean up
+    for id in nkb.extract_strategies.list().keys():
+        nkb.extract_strategies.delete(id=id)
+
+    # tests
+    nkb.extract_strategies.add(config={"name": "strategy1", "vllm_config": {}})
+    all = nkb.extract_strategies.list()
+    assert len(all.keys()) == 1
+    nkb.extract_strategies.delete(id=list(all.keys())[0])
+    all = nkb.extract_strategies.list()
+    assert len(all.keys()) == 0

{nuclia-4.9.0.dist-info → nuclia-4.9.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nuclia
-Version: 4.9.0
+Version: 4.9.2
 Summary: Nuclia Python SDK
 Author-email: Nuclia <info@nuclia.com>
 License-Expression: MIT

{nuclia-4.9.0.dist-info → nuclia-4.9.2.dist-info}/RECORD

@@ -9,11 +9,11 @@ nuclia/cli/run.py,sha256=B1hP0upSbSCqqT89WAwsd93ZxkAoF6ajVyLOdYmo8fU,1560
 nuclia/cli/utils.py,sha256=iZ3P8juBdAGvaRUd2BGz7bpUXNDHdPrC5p876yyZ2Cs,1223
 nuclia/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nuclia/lib/conversations.py,sha256=M6qhL9NPEKroYF767S-Q2XWokRrjX02kpYTzRvZKwUE,149
-nuclia/lib/kb.py,sha256=jV-L8-a5oFVHIbFfeGO17hXf2IQk9UuW-9fq5D3LO6Y,28508
+nuclia/lib/kb.py,sha256=vSLfmV6HqPvWJwPVw4iIzkDnm0M96_ImTH3QMZdkd_I,30985
 nuclia/lib/models.py,sha256=ekEQrVIFU3aFvt60yQh-zpWkGNORBMSc7c5Hd_VzPzI,1564
 nuclia/lib/nua.py,sha256=sUVFdCjvLigTqUUhILywdHpiC0qKCtKPABn5kUXfuxQ,27064
 nuclia/lib/nua_chat.py,sha256=ApL1Y1FWvAVUt-Y9a_8TUSJIhg8-UmBSy8TlDPn6tD8,3874
-nuclia/lib/nua_responses.py,sha256=9WRGpvvhA2ZOyOv9kzX1zI7C3ypM8ySqciVYCSibUjo,12724
+nuclia/lib/nua_responses.py,sha256=RJB7bbgZEpYr29aOmaT2nb1X0g6wds4MOngq7N-ZcTg,13382
 nuclia/lib/utils.py,sha256=9l6DxBk-11WqUhXRn99cqeuVTUOJXj-1S6zckal7wOk,6312
 nuclia/sdk/__init__.py,sha256=-nAw8i53XBdmbfTa1FJZ0FNRMNakimDVpD6W4OdES-c,1374
 nuclia/sdk/accounts.py,sha256=7XQ3K9_jlSuk2Cez868FtazZ05xSGab6h3Mt1qMMwIE,647
@@ -21,7 +21,8 @@ nuclia/sdk/agent.py,sha256=ot_oA4yi7TkXahPnhcRIxztq9Dtskc85-A57WN1BNGg,1961
 nuclia/sdk/auth.py,sha256=o9CiWt3P_UHQW_L5Jq_9IOB-KpSCXFsVTouQcZp8BJM,25072
 nuclia/sdk/backup.py,sha256=adbPcNEbHGZW698o028toXKfDkDrmk5QRIDSiN6SPys,6529
 nuclia/sdk/export_import.py,sha256=y5cTOxhILwRPIvR2Ya12bk-ReGbeDzA3C9TPxgnOHD4,9756
-nuclia/sdk/kb.py,sha256=iDJimbzw_R9VeUu7a4F5wXoLdLxCim7W3UeEaAQ5dGI,26820
+nuclia/sdk/extract_strategy.py,sha256=NZBLLThdLyQYw8z1mT9iRhFjkE5sQP86-3QhsTiyV9o,2540
+nuclia/sdk/kb.py,sha256=2-H9FOvPgsG-ZYNUA8D4FYAJhX3K8m2VwbwSQy1JV7c,27044
 nuclia/sdk/kbs.py,sha256=nXEvg5ddZYdDS8Kie7TrN-s1meU9ecYLf9FlT5xr-ro,9131
 nuclia/sdk/logger.py,sha256=UHB81eS6IGmLrsofKxLh8cmF2AsaTj_HXP0tGqMr_HM,57
 nuclia/sdk/logs.py,sha256=3jfORpo8fzZiXFFSbGY0o3Bre1ZgJaKQCXgxP1keNHw,9614
@@ -42,6 +43,7 @@ nuclia/tests/assets/conversation.json,sha256=jLmVngHEW8QdAzSMa-MYgbJgYxOO8tHq6Ia
 nuclia/tests/test_kb/test_backup.py,sha256=67fAw1C_NB363G-wuI9LnSH6nG_YrYs-uBcFyBx6Peg,2932
 nuclia/tests/test_kb/test_conversation.py,sha256=plsWY_gEhK3P4CsIO9jkBmimZk3UDGeMYPu7YVFeuIQ,727
 nuclia/tests/test_kb/test_export_import.py,sha256=lQEww2jFNHZYcudFJqcHhoWAPrmtvvnPvcFqrijxLbo,1019
+nuclia/tests/test_kb/test_extract_strategies.py,sha256=aKNzE98H1sE4eFqh8XiwvzDT8-7SnMs0v3NTeubZ-XY,521
 nuclia/tests/test_kb/test_graph.py,sha256=ecAB-lWqm3_796HUuF-fSpGzlhsXV9NGZnAjvxXANps,2540
 nuclia/tests/test_kb/test_labels.py,sha256=IUdTq4mzv0OrOkwBWWy4UwKGKyJybtoHrgvXr676vyY,961
 nuclia/tests/test_kb/test_logs.py,sha256=Z9ELtiiU9NniITJzeWt92GCcERKYy9Nwc_fUVPboRU0,3121
@@ -61,9 +63,9 @@ nuclia/tests/test_nucliadb/test_crud.py,sha256=GuY76HRvt2DFaNgioKm5n0Aco1HnG7zzV
 nuclia/tests/unit/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nuclia/tests/unit/test_export_import.py,sha256=xo_wVbjUnNlVV65ZGH7LtZ38qy39EkJp2hjOuTHC1nU,980
 nuclia/tests/unit/test_nua_responses.py,sha256=t_hIdVztTi27RWvpfTJUYcCL0lpKdZFegZIwLdaPNh8,319
-nuclia-4.9.0.dist-info/licenses/LICENSE,sha256=Ops2LTti_HJtpmWcanuUTdTY3vKDR1myJ0gmGBKC0FA,1063
-nuclia-4.9.0.dist-info/METADATA,sha256=ZHT0wCR1AhZtfnWdcD-uSmfNkb8QznVDuWRt_2f6yP8,2337
-nuclia-4.9.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nuclia-4.9.0.dist-info/entry_points.txt,sha256=iZHOyXPNS54r3eQmdi5So20xO1gudI9K2oP4sQsCJRw,46
-nuclia-4.9.0.dist-info/top_level.txt,sha256=cqn_EitXOoXOSUvZnd4q6QGrhm04pg8tLAZtem-Zfdo,7
-nuclia-4.9.0.dist-info/RECORD,,
+nuclia-4.9.2.dist-info/licenses/LICENSE,sha256=Ops2LTti_HJtpmWcanuUTdTY3vKDR1myJ0gmGBKC0FA,1063
+nuclia-4.9.2.dist-info/METADATA,sha256=QqBc1WyVh98_k_OuLTJBoD5_Ig5VZuQ5ewQZIriu9p0,2337
+nuclia-4.9.2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nuclia-4.9.2.dist-info/entry_points.txt,sha256=iZHOyXPNS54r3eQmdi5So20xO1gudI9K2oP4sQsCJRw,46
+nuclia-4.9.2.dist-info/top_level.txt,sha256=cqn_EitXOoXOSUvZnd4q6QGrhm04pg8tLAZtem-Zfdo,7
+nuclia-4.9.2.dist-info/RECORD,,