nuclia 4.8.10__py3-none-any.whl → 4.9.1__py3-none-any.whl

nuclia/lib/models.py

@@ -48,7 +48,7 @@ class GraphRelation(BaseModel):
     destination: GraphEntity
 
     def to_relation(self) -> Relation:
-        return Relation.parse_obj(
+        return Relation.model_validate(
             {
                 "from": self.source.to_dict(),
                 "to": self.destination.to_dict(),
@@ -60,6 +60,6 @@ class GraphRelation(BaseModel):
 
 def get_relation(relation: Union[GraphRelation, dict]) -> GraphRelation:
     if isinstance(relation, dict):
-        return GraphRelation.parse_obj(relation)
+        return GraphRelation.model_validate(relation)
     else:
         return relation

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/search.py

@@ -20,6 +20,8 @@ from nucliadb_models.search import (
     SyncAskResponse,
     ChatModel,
 )
+from nucliadb_models.graph.requests import GraphSearchRequest
+from nucliadb_models.graph.responses import GraphSearchResponse
 from pydantic import ValidationError
 
 from nuclia.data import get_async_auth, get_auth
@@ -331,6 +333,33 @@ class NucliaSearch:
                 result.tokens = ask_response.metadata.tokens.model_dump()
         return result
 
+    @kb
+    def graph(
+        self,
+        *,
+        query: Union[dict, GraphSearchRequest],
+        **kwargs,
+    ) -> GraphSearchResponse:
+        """
+        Perform a graph path query.
+
+        See https://docs.nuclia.dev/docs/api#tag/Search/operation/graph_search_knowledgebox_kb__kbid__graph_post
+        """
+        ndb: NucliaDBClient = kwargs["ndb"]
+
+        if isinstance(query, GraphSearchRequest):
+            req = query
+        elif isinstance(query, dict):
+            try:
+                req = GraphSearchRequest.model_validate(query)
+            except ValidationError:
+                logger.exception("Error validating query")
+                raise
+        else:
+            raise Exception(f"Invalid query: '{query}'")
+
+        return ndb.ndb.graph_search(req, kbid=ndb.kbid)
+
 
 class AsyncNucliaSearch:
     """
@@ -659,3 +688,30 @@ class AsyncNucliaSearch:
             else:  # pragma: no cover
                 warnings.warn(f"Unknown ask stream item type: {ask_response_item.type}")
         return result
+
+    @kb
+    async def graph(
+        self,
+        *,
+        query: Union[dict, GraphSearchRequest],
+        **kwargs,
+    ) -> GraphSearchResponse:
+        """
+        Perform a graph path query.
+
+        See https://docs.nuclia.dev/docs/api#tag/Search/operation/graph_search_knowledgebox_kb__kbid__graph_post
+        """
+        ndb: AsyncNucliaDBClient = kwargs["ndb"]
+
+        if isinstance(query, GraphSearchRequest):
+            req = query
+        elif isinstance(query, dict):
+            try:
+                req = GraphSearchRequest.model_validate(query)
+            except ValidationError:
+                logger.exception("Error validating query")
+                raise
+        else:
+            raise Exception(f"Invalid query: '{query}'")
+
+        return await ndb.ndb.graph_search(req, kbid=ndb.kbid)

nuclia/tests/fixtures.py

@@ -49,8 +49,8 @@ def testing_config(testing_kb, testing_nua, testing_user):
         nuclia_auth._config.set_default_kb(TESTING_KBID)
         nuclia_auth._config.set_default_nua(client_id)
 
-    yield
-    reset_config_file()
+        yield
+        reset_config_file()
 
 
 @pytest.fixture(autouse=True)

nuclia/tests/test_kb/test_conversation.py

@@ -2,7 +2,6 @@ import os
 
 from nucliadb_sdk.v2.exceptions import NotFoundError
 
-from nuclia.sdk.kb import NucliaKB
 from nuclia.sdk.resource import NucliaResource
 from nuclia.sdk.upload import NucliaUpload
 
@@ -10,7 +9,6 @@ path = f"{os.path.dirname(__file__)}/../assets/conversation.json"
 
 
 def test_conversation(testing_config):
-    nkb = NucliaKB()
     nresource = NucliaResource()
     try:
         res = nresource.get(slug="conversation1")
@@ -18,9 +16,6 @@ def test_conversation(testing_config):
     except NotFoundError:
         pass
 
-    res = nkb.list()
-    assert len(res.resources) == 2
-
     nu = NucliaUpload()
     nu.conversation(path=path, slug="conversation1", field="c1")
 

nuclia/tests/test_kb/test_graph.py

@@ -6,6 +6,13 @@ from time import sleep
 
 def test_graph(testing_config):
     nkb = NucliaKB()
+
+    # ensure KB is clean
+    try:
+        nkb.delete_graph(slug="graph1")
+    except NotFoundError:
+        pass
+
     nkb.add_graph(
         slug="graph1",
         graph=[
@@ -28,6 +35,25 @@ def test_graph(testing_config):
     )
     relations = nkb.get_graph(slug="graph1")
     assert len(relations) == 3
+
+    # XXX: sometimes, nucliadb needs some more time to index the graph. We retry
+    # some times and fail if that keeps hapenning
+    RETRIES = 5
+    DELAY = 0.5
+    for _ in range(RETRIES):
+        paths = nkb.search.graph(
+            query={"query": {"prop": "path", "source": {"value": "Alice"}}}
+        )
+        try:
+            assert len(paths.paths) == 2
+        except AssertionError:
+            # wait for a bit before retrying
+            print("Graph was not indexed yet, waiting a bit...")
+            sleep(DELAY)
+        else:
+            break
+    assert len(paths.paths) == 2
+
     nkb.update_graph(
         slug="graph1",
         graph=[
@@ -40,6 +66,20 @@ def test_graph(testing_config):
     )
     relations = nkb.get_graph(slug="graph1")
     assert len(relations) == 4
+
+    for _ in range(RETRIES):
+        paths = nkb.search.graph(
+            query={"query": {"prop": "path", "source": {"value": "Victor"}}}
+        )
+        try:
+            assert len(paths.paths) == 1
+        except AssertionError:
+            print("Graph was not indexed yet, waiting a bit...")
+            sleep(DELAY)
+        else:
+            break
+    assert len(paths.paths) == 1
+
     nkb.delete_graph(slug="graph1")
     try:
         sleep(0.5)

{nuclia-4.8.10.dist-info → nuclia-4.9.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: nuclia
-Version: 4.8.10
+Version: 4.9.1
 Summary: Nuclia Python SDK
 Author-email: Nuclia <info@nuclia.com>
 License-Expression: MIT
@@ -24,8 +24,8 @@ Requires-Dist: requests
 Requires-Dist: httpx
 Requires-Dist: httpcore>=1.0.0
 Requires-Dist: prompt_toolkit
-Requires-Dist: nucliadb_sdk<7,>=6.4
-Requires-Dist: nucliadb_models<7,>=6.4
+Requires-Dist: nucliadb_sdk<7,>=6.5
+Requires-Dist: nucliadb_models<7,>=6.5
 Requires-Dist: nuclia-models>=0.41.1
 Requires-Dist: tqdm
 Requires-Dist: aiofiles

{nuclia-4.8.10.dist-info → nuclia-4.9.1.dist-info}/RECORD

@@ -10,10 +10,10 @@ 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/models.py,sha256=lO3TZa4jTQ6C8ptGO1oX-5pxDgv8ocO13z7j5jXOC0Y,1554
+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
@@ -31,18 +31,18 @@ nuclia/sdk/predict.py,sha256=KF7iT2aasaB9DIEAwqktXbOl2H_Y_ne-6-SEErN7YOk,9095
 nuclia/sdk/process.py,sha256=WuNnqaWprp-EABWDC_z7O2woesGIlYWnDUKozh7Ibr4,2241
 nuclia/sdk/remi.py,sha256=BEb3O9R2jOFlOda4vjFucKKGO1c2eTkqYZdFlIy3Zmo,4357
 nuclia/sdk/resource.py,sha256=0lSvD4e1FpN5iM9W295dOKLJ8hsXfIe8HKdo0HsVg20,13976
-nuclia/sdk/search.py,sha256=OpW2Tl0RQwG6e2Ytw6Iz7L0SMzaX06qTpb-f4mzw6_g,23435
+nuclia/sdk/search.py,sha256=1mLJzDO-W3ObReDL1xK8zgjIxkdpuIhyf87-L4bISKQ,25180
 nuclia/sdk/task.py,sha256=UawH-7IneRIGVOiLdDQ2vDBwe5eI51UXRbLRRVeu3C4,6095
 nuclia/sdk/upload.py,sha256=ZBzYROF3yP-77HcaR06OBsFjJAbTOCvF-nlxaqQZsT4,22720
 nuclia/sdk/zones.py,sha256=1ARWrTsTuzj8zguanpX3OaIw-3Qq_ULS_g4GG2mHxOA,342
 nuclia/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 nuclia/tests/conftest.py,sha256=-Gv9Q20tBVQNNWqWHGNmzkXJpijmgTdcvNnaRAUluGk,115
-nuclia/tests/fixtures.py,sha256=ppjT1VgEdkrqvRPoqt1LU7T6rUf9qZ1UkplQz4hT0IQ,1787
+nuclia/tests/fixtures.py,sha256=rcGGbtxQbGK9vR9efL7BsE1ugJl15QQDIxUyFZLfv-E,1795
 nuclia/tests/assets/conversation.json,sha256=jLmVngHEW8QdAzSMa-MYgbJgYxOO8tHq6Ia029-CD0E,348
 nuclia/tests/test_kb/test_backup.py,sha256=67fAw1C_NB363G-wuI9LnSH6nG_YrYs-uBcFyBx6Peg,2932
-nuclia/tests/test_kb/test_conversation.py,sha256=L__KMBucqChPGeYR_XwXl7Elt4LDlp1mcmFdD71U8cQ,840
+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_graph.py,sha256=PQSdqvPvOHr71pfdbUpsAK-_DxBQh7spxtJaaleF87A,1434
+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
 nuclia/tests/test_kb/test_remi.py,sha256=OX5N-MHbgcwpLg6fBjrAK_KhqkMspJo_VKQHCBCayZ8,2080
@@ -61,9 +61,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.8.10.dist-info/licenses/LICENSE,sha256=Ops2LTti_HJtpmWcanuUTdTY3vKDR1myJ0gmGBKC0FA,1063
-nuclia-4.8.10.dist-info/METADATA,sha256=Vcg3SQFhjDTSZ9akWQevdKqRVZyt8ModTnX5BcwgmsI,2338
-nuclia-4.8.10.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nuclia-4.8.10.dist-info/entry_points.txt,sha256=iZHOyXPNS54r3eQmdi5So20xO1gudI9K2oP4sQsCJRw,46
-nuclia-4.8.10.dist-info/top_level.txt,sha256=cqn_EitXOoXOSUvZnd4q6QGrhm04pg8tLAZtem-Zfdo,7
-nuclia-4.8.10.dist-info/RECORD,,
+nuclia-4.9.1.dist-info/licenses/LICENSE,sha256=Ops2LTti_HJtpmWcanuUTdTY3vKDR1myJ0gmGBKC0FA,1063
+nuclia-4.9.1.dist-info/METADATA,sha256=bScM2QXDcmlYLGp0pDG1ij02jclIV8u0Rk_CTFwkjvA,2337
+nuclia-4.9.1.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nuclia-4.9.1.dist-info/entry_points.txt,sha256=iZHOyXPNS54r3eQmdi5So20xO1gudI9K2oP4sQsCJRw,46
+nuclia-4.9.1.dist-info/top_level.txt,sha256=cqn_EitXOoXOSUvZnd4q6QGrhm04pg8tLAZtem-Zfdo,7
+nuclia-4.9.1.dist-info/RECORD,,