ygg 0.1.53__tar.gz → 0.1.54__tar.gz

{ygg-0.1.53 → ygg-0.1.54}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.53
+Version: 0.1.54
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 License:                                  Apache License

{ygg-0.1.53 → ygg-0.1.54}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ygg"
-version = "0.1.53"
+version = "0.1.54"
 description = "Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks"
 readme = { file = "README.md", content-type = "text/markdown" }
 license = { file = "LICENSE" }

{ygg-0.1.53 → ygg-0.1.54}/src/ygg.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.53
+Version: 0.1.54
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 License:                                  Apache License

{ygg-0.1.53 → ygg-0.1.54}/src/ygg.egg-info/SOURCES.txt

@@ -10,6 +10,8 @@ src/ygg.egg-info/top_level.txt
 src/yggdrasil/__init__.py
 src/yggdrasil/version.py
 src/yggdrasil/databricks/__init__.py
+src/yggdrasil/databricks/ai/__init__.py
+src/yggdrasil/databricks/ai/loki.py
 src/yggdrasil/databricks/compute/__init__.py
 src/yggdrasil/databricks/compute/cluster.py
 src/yggdrasil/databricks/compute/execution_context.py

ygg-0.1.54/src/yggdrasil/databricks/ai/loki.py

@@ -0,0 +1,53 @@
+from typing import Optional
+from dataclasses import field, dataclass
+
+from ..workspaces.workspace import WorkspaceService
+
+try:
+    from openai import OpenAI
+
+    def make_openai_client(
+        api_key: str,
+        base_url: str
+    ):
+        return OpenAI(
+            api_key=api_key,
+            base_url=base_url
+        )
+except ImportError:
+    class OpenAI:
+        pass
+
+    def make_openai_client(
+        api_key: str,
+        base_url: str
+    ):
+        from openai import OpenAI
+
+        return OpenAI(
+            api_key=api_key,
+            base_url=base_url
+        )
+
+__all__ = [
+    "Loki"
+]
+
+
+@dataclass
+class Loki(WorkspaceService):
+    model: str = "databricks-gemini-2-5-flash"
+
+    _openai_client: Optional[OpenAI] = field(repr=False, hash=False, default=None)
+
+    @property
+    def openai_client(self):
+        if self._openai_client is None:
+            self._openai_client = self.make_openai_client()
+        return self._openai_client
+
+    def make_openai_client(self):
+        return make_openai_client(
+            api_key=self.workspace.current_token(),
+            base_url=self.workspace.host + "/serving-endpoints"
+        )

{ygg-0.1.53 → ygg-0.1.54}/src/yggdrasil/databricks/workspaces/io.py

@@ -12,7 +12,8 @@ import pyarrow.csv as pcsv
 import pyarrow.parquet as pq
 from pyarrow.dataset import (
     FileFormat,
-    ParquetFileFormat, CsvFileFormat,
+    ParquetFileFormat,
+    CsvFileFormat,
 )
 
 from .path_kind import DatabricksPathKind
@@ -71,10 +72,11 @@ class DatabricksIO(ABC, IO):
         self.close()
 
     def __del__(self):
-        try:
-            Thread(target=self.close).start()
-        except BaseException:
-            pass
+        if self._need_flush():
+            try:
+                Thread(target=self.close).start()
+            except BaseException:
+                pass
 
     def __next__(self):
         """Iterate over lines in the file."""
@@ -533,13 +535,16 @@ class DatabricksIO(ABC, IO):
 
         return size
 
+    def _need_flush(self):
+        return self._write_flag and self._buffer is not None
+
     def flush(self):
         """Flush buffered data to the remote path.
 
         Returns:
             None.
         """
-        if self._write_flag and self._buffer is not None:
+        if self._need_flush():
             self.write_all_bytes(data=self._buffer.getvalue())
             self._write_flag = False
 

{ygg-0.1.53 → ygg-0.1.54}/src/yggdrasil/databricks/workspaces/workspace.py

@@ -645,6 +645,7 @@ class Workspace:
 
     def clusters(
         self,
+        workspace: Optional["Workspace"] = None,
         cluster_id: Optional[str] = None,
         cluster_name: Optional[str] = None,
     ) -> "Cluster":
@@ -666,6 +667,16 @@ class Workspace:
             cluster_name=cluster_name,
         )
 
+    def loki(
+        self,
+        workspace: Optional["Workspace"] = None,
+    ):
+        from ..ai.loki import Loki
+
+        return Loki(
+            workspace=self,
+        )
+
 # ---------------------------------------------------------------------------
 # Workspace-bound base class
 # ---------------------------------------------------------------------------
@@ -717,13 +728,13 @@ class WorkspaceService(ABC):
         """
         return self.workspace.is_in_databricks_environment()
 
-    def connect(self):
+    def connect(self, clone: bool = False):
         """Connect the underlying workspace.
 
         Returns:
             The current WorkspaceService instance.
         """
-        self.workspace = self.workspace.connect()
+        self.workspace = self.workspace.connect(clone=clone)
         return self
 
     def dbfs_path(

ygg-0.1.54/src/yggdrasil/version.py

@@ -0,0 +1 @@
+__version__ = "0.1.54"

ygg-0.1.53/src/yggdrasil/version.py

@@ -1 +0,0 @@
-__version__ = "0.1.53"