craft-ai-sdk 0.61.1__tar.gz → 0.62.0__tar.gz

{craft_ai_sdk-0.61.1 → craft_ai_sdk-0.62.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: craft-ai-sdk
-Version: 0.61.1
+Version: 0.62.0
 Summary: Craft AI MLOps platform SDK
 Home-page: https://www.craft.ai/
 License: Apache-2.0
@@ -16,6 +16,7 @@ Classifier: Programming Language :: Python :: 3.12
 Requires-Dist: PyJWT (>=2.3.0,<3.0.0)
 Requires-Dist: StrEnum (>=0.4.9,<0.5.0)
 Requires-Dist: requests (>=2.27.1,<3.0.0)
+Requires-Dist: weaviate-client (>=4.10.4,<5.0.0)
 Description-Content-Type: text/markdown
 
 # Craft AI Python SDK

{craft_ai_sdk-0.61.1 → craft_ai_sdk-0.62.0}/craft_ai_sdk/__init__.py

@@ -13,4 +13,4 @@ from .io import (  # noqa: F401
 )
 from .sdk import CraftAiSdk  # noqa: F401
 
-__version__ = "0.61.1"
+__version__ = "0.62.0"

craft_ai_sdk-0.62.0/craft_ai_sdk/core/vector_database.py

@@ -0,0 +1,67 @@
+from craft_ai_sdk.shared.environments import get_environment_id
+
+from ..sdk import BaseCraftAiSdk
+from ..shared.logger import log_action, log_func_result
+
+
+def get_vector_database_credentials(sdk: BaseCraftAiSdk):
+    """Get the credentials of the vector database.
+
+    Returns:
+        :obj:`dict`: The vector database credentials, with the following keys:
+            * ``"vector_database_url"`` (:obj:`str`): URL of the vector database.
+            * ``"vector_database_token"`` (:obj:`str`): Token to connect to the vector
+              database.
+    """
+    environment_id = get_environment_id(sdk)
+
+    vector_database_url = (
+        f"{sdk.base_control_api_url}/environments/{environment_id}/vector-database"
+    )
+
+    return sdk._get(vector_database_url)
+
+
+@log_func_result("Connecting to Weaviate")
+def get_weaviate_client(sdk: BaseCraftAiSdk):
+    """Get the Weaviate client.
+
+    Returns:
+        :obj:`weaviate.WeaviateClient`: The Weaviate client.
+    """
+    try:
+        import weaviate
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError(
+            "The 'weaviate' package is required to use the vector database. "
+            "You can install it with 'pip install weaviate-client'."
+        )
+    credentials = get_vector_database_credentials(sdk)
+
+    is_secure = credentials["vector_database_url"].startswith("https://")
+
+    vector_database_url = (
+        credentials["vector_database_url"]
+        .replace("http://", "")
+        .replace("https://", "")
+    ).rstrip("/")
+
+    log_action(sdk, "Connecting to Weaviate")
+
+    weaviate_client = weaviate.connect_to_custom(
+        http_host=vector_database_url,
+        http_port=8080,
+        grpc_host=vector_database_url,
+        grpc_port=8082,
+        http_secure=is_secure,
+        grpc_secure=is_secure,
+        headers={
+            "craft-vector-database-token": credentials["vector_database_token"],
+            "craft-vector-database-url": credentials["vector_database_url"],
+        },
+        auth_credentials=None,
+    )
+
+    log_action(sdk, "Connected to Weaviate, using version", weaviate.__version__)
+
+    return weaviate_client

{craft_ai_sdk-0.61.1 → craft_ai_sdk-0.62.0}/craft_ai_sdk/sdk.py

@@ -116,7 +116,11 @@ class CraftAiSdk(BaseCraftAiSdk):
         list_steps,
     )
     from .core.users import get_user
-    from .core.vector_database import get_vector_database_credentials
+
+    from .core.vector_database import (
+        get_vector_database_credentials,
+        get_weaviate_client,
+    )
 
     # Size (in bytes) from which datastore upload will switch to multipart
     # AWS: minimum part size is 5MiB
@@ -132,7 +136,7 @@ class CraftAiSdk(BaseCraftAiSdk):
         os.environ.get("CRAFT_AI__MULTIPART_PART_SIZE__B", str(38 * 256 * 1024))
     )
     _access_token_margin = timedelta(seconds=30)
-    _version = "0.61.1"  # Would be better to share it somewhere
+    _version = "0.62.0"  # Would be better to share it somewhere
 
     def __init__(
         self,

{craft_ai_sdk-0.61.1 → craft_ai_sdk-0.62.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "craft-ai-sdk"
-version = "0.61.1"
+version = "0.62.0"
 description = "Craft AI MLOps platform SDK"
 license = "Apache-2.0"
 authors = ["Craft AI <contact@craft.ai>"]
@@ -29,6 +29,7 @@ python = "^3.9"
 requests = "^2.27.1"
 PyJWT = "^2.3.0"
 StrEnum = "^0.4.9"
+weaviate-client = "^4.10.4"
 
 [tool.poetry.group.dev.dependencies]
 ipython = "^8.18.1"

craft_ai_sdk-0.61.1/craft_ai_sdk/core/vector_database.py

@@ -1,21 +0,0 @@
-from craft_ai_sdk.shared.environments import get_environment_id
-
-from ..sdk import BaseCraftAiSdk
-
-
-def get_vector_database_credentials(sdk: BaseCraftAiSdk):
-    """Get the credentials of the vector database.
-
-    Returns:
-        :obj:`dict`: The vector database credentials, with the following keys:
-            * ``"vector_database_url"`` (:obj:`str`): URL of the vector database.
-            * ``"vector_database_token"`` (:obj:`str`): Token to connect to the vector
-              database.
-    """
-    environment_id = get_environment_id(sdk)
-
-    vector_database_url = (
-        f"{sdk.base_control_api_url}/environments/{environment_id}/vector-database"
-    )
-
-    return sdk._get(vector_database_url)