ygg 0.1.53__py3-none-any.whl → 0.1.55__py3-none-any.whl

{ygg-0.1.53.dist-info → ygg-0.1.55.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.53
+Version: 0.1.55
 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.dist-info → ygg-0.1.55.dist-info}/RECORD

@@ -1,7 +1,9 @@
-ygg-0.1.53.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+ygg-0.1.55.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
 yggdrasil/__init__.py,sha256=4-ghPak2S6zfMqmnlxW2GCgPb5s79znpKa2hGEGXcE4,24
-yggdrasil/version.py,sha256=OwqD5Bb-XPoHxyR3pJbfa8OkEcCnfpl_gLnvTYdwfkQ,22
+yggdrasil/version.py,sha256=RJcC5yver_ugtAU_HgwZ-2HWyszBXpsLZyyNma6G4Dc,22
 yggdrasil/databricks/__init__.py,sha256=skctY2c8W-hI81upx9F_PWRe5ishL3hrdiTuizgDjdw,152
+yggdrasil/databricks/ai/__init__.py,sha256=Mkp70UOVBzDQvdPNsqncHcyzxe5PnSGYE_bHnYxA1eA,21
+yggdrasil/databricks/ai/loki.py,sha256=iyekxctP6393LBN0PJOZgHBxQs9vDyQeRXPrru_krF0,3661
 yggdrasil/databricks/compute/__init__.py,sha256=NvdzmaJSNYY1uJthv1hHdBuNu3bD_-Z65DWnaJt9yXg,289
 yggdrasil/databricks/compute/cluster.py,sha256=YomLfvB0oxbgl6WDgBRxI1UXsxwlEbR6gq3FUbPHscY,44199
 yggdrasil/databricks/compute/execution_context.py,sha256=jIV6uru2NeX3O5lg-3KEqmXtLxxq45CFgkBQgQIIOHQ,23327
@@ -16,11 +18,11 @@ yggdrasil/databricks/sql/types.py,sha256=5G-BM9_eOsRKEMzeDTWUsWW5g4Idvs-czVCpOCr
 yggdrasil/databricks/sql/warehouse.py,sha256=1J0dyQLJb-OS1_1xU1eAVZ4CoL2-FhFeowKSvU3RzFc,9773
 yggdrasil/databricks/workspaces/__init__.py,sha256=dv2zotoFVhNFlTCdRq6gwf5bEzeZkOZszoNZMs0k59g,114
 yggdrasil/databricks/workspaces/filesytem.py,sha256=Z8JXU7_XUEbw9fpTQT1avRQKi-IAP2KemXBMPkUoY4w,9805
-yggdrasil/databricks/workspaces/io.py,sha256=4Tyrg2YPlyQMrdzE8Wxx5nTTWnFOmOfA-cFgorB57KA,32998
+yggdrasil/databricks/workspaces/io.py,sha256=PAoxIxYvTC162Dx2qL2hk8oAdt8BnYrQ3jJHcJm4VkA,33116
 yggdrasil/databricks/workspaces/path.py,sha256=KkvLFHrps3UFr4ogYdESbJHEMfQBcWfWfXjlrv_7rTU,55180
 yggdrasil/databricks/workspaces/path_kind.py,sha256=rhWe1ky7uPD0du0bZSv2S4fK4C5zWd7zAF3UeS2iiPU,283
 yggdrasil/databricks/workspaces/volumes_path.py,sha256=s8CA33cG3jpMVJy5MILLlkEBcFg_qInDCF2jozLj1Fg,2431
-yggdrasil/databricks/workspaces/workspace.py,sha256=zBlQdYNT_xKwUCYo3O4Q4g-8pfMvff3I26efyCfY_TY,24961
+yggdrasil/databricks/workspaces/workspace.py,sha256=5DCPz5io_rmrpGNi5I6RChmyZ8kjlNUFGQl8mzQJThg,25511
 yggdrasil/dataclasses/__init__.py,sha256=_RkhfF3KC1eSORby1dzvBXQ0-UGG3u6wyUQWX2jq1Pc,108
 yggdrasil/dataclasses/dataclass.py,sha256=LxrCjwvmBnb8yRI_N-c31RHHxB4XoJPixmKg9iBIuaI,1148
 yggdrasil/libs/__init__.py,sha256=zdC9OU0Xy36CLY9mg2drxN6S7isPR8aTLzJA6xVIeLE,91
@@ -31,14 +33,14 @@ yggdrasil/libs/sparklib.py,sha256=FQ3W1iz2EIpQreorOiQuFt15rdhq2QhGEAWp8Zrbl9A,10
 yggdrasil/libs/extensions/__init__.py,sha256=mcXW5Li3Cbprbs4Ci-b5A0Ju0wmLcfvEiFusTx6xNjU,117
 yggdrasil/libs/extensions/polars_extensions.py,sha256=RTkGi8llhPJjX7x9egix7-yXWo2X24zIAPSKXV37SSA,12397
 yggdrasil/libs/extensions/spark_extensions.py,sha256=E64n-3SFTDgMuXwWitX6vOYP9ln2lpGKb0htoBLEZgc,16745
-yggdrasil/pyutils/__init__.py,sha256=tl-LapAc71TV7RMgf2ftKwrzr8iiLOGHeJgA3RvO93w,293
+yggdrasil/pyutils/__init__.py,sha256=a7LKg9xWZGtv2o5jr_7wrSZJ9ZmwolKWc--8iRgKlL4,225
 yggdrasil/pyutils/callable_serde.py,sha256=1XckmFO-ThP0MedxgXwB71u9jWUuhM1btOzW9gJ8w9g,23117
 yggdrasil/pyutils/equality.py,sha256=Xyf8D1dLUCm3spDEir8Zyj7O4US_fBJwEylJCfJ9slI,3080
 yggdrasil/pyutils/exceptions.py,sha256=ssKNm-rjhavHUOZmGA7_1Gq9tSHDrb2EFI-cnBuWgng,3388
 yggdrasil/pyutils/expiring_dict.py,sha256=pr2u25LGwPVbLfsLptiHGovUtYRRo0AMjaJtCtJl7nQ,8477
 yggdrasil/pyutils/modules.py,sha256=B7IP99YqUMW6-DIESFzBx8-09V1d0a8qrIJUDFhhL2g,11424
 yggdrasil/pyutils/parallel.py,sha256=ubuq2m9dJzWYUyKCga4Y_9bpaeMYUrleYxdp49CHr44,6781
-yggdrasil/pyutils/python_env.py,sha256=Gh5geFK9ABpyWEfyegGUfIJUoPxKwcH0pqLBiMrW9Rw,51103
+yggdrasil/pyutils/python_env.py,sha256=nmXXcIUSSbpKWHBnYfCY1QIW-LQ8IILV9Ceijufu5Pg,51086
 yggdrasil/pyutils/retry.py,sha256=gXBtn1DdmIYIUmGKOUr8-SUT7MOu97LykN2YR4uocgc,11917
 yggdrasil/requests/__init__.py,sha256=dMesyzq97_DmI765x0TwaDPEfsxFtgGNgchk8LvEN-o,103
 yggdrasil/requests/msal.py,sha256=s2GCyzbgFdgdlJ1JqMrZ4qYVbmoG46-ZOTcaVQhZ-sQ,9220
@@ -57,8 +59,8 @@ yggdrasil/types/cast/registry.py,sha256=OOqIfbIjPH-a3figvu-zTvEtUDTEWhe2xIl3cCA4
 yggdrasil/types/cast/spark_cast.py,sha256=_KAsl1DqmKMSfWxqhVE7gosjYdgiL1C5bDQv6eP3HtA,24926
 yggdrasil/types/cast/spark_pandas_cast.py,sha256=BuTiWrdCANZCdD_p2MAytqm74eq-rdRXd-LGojBRrfU,5023
 yggdrasil/types/cast/spark_polars_cast.py,sha256=btmZNHXn2NSt3fUuB4xg7coaE0RezIBdZD92H8NK0Jw,9073
-ygg-0.1.53.dist-info/METADATA,sha256=U_RTDcmrHbLOrFiDxsO0yX4Lgi3qbjd009gh-uuohqY,18528
-ygg-0.1.53.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-ygg-0.1.53.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
-ygg-0.1.53.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
-ygg-0.1.53.dist-info/RECORD,,
+ygg-0.1.55.dist-info/METADATA,sha256=F84590C1dKd4ZllEbe1uzAsMuUw9we134qVVQmOcGNI,18528
+ygg-0.1.55.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+ygg-0.1.55.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
+ygg-0.1.55.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
+ygg-0.1.55.dist-info/RECORD,,

{ygg-0.1.53.dist-info → ygg-0.1.55.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.1)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

yggdrasil/databricks/ai/__init__.py

@@ -0,0 +1 @@
+from .loki import *

yggdrasil/databricks/ai/loki.py

@@ -0,0 +1,122 @@
+from typing import Optional, Dict, Any, List
+from dataclasses import field, dataclass
+
+from ..workspaces.workspace import WorkspaceService
+
+try:
+    from openai import OpenAI
+
+    def make_ai_client(
+        api_key: str,
+        base_url: str
+    ):
+        return OpenAI(
+            api_key=api_key,
+            base_url=base_url
+        )
+except ImportError:
+    class OpenAI:
+        pass
+
+    def make_ai_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"
+
+    _ai_client: Optional[OpenAI] = field(repr=False, hash=False, default=None)
+
+    @property
+    def ai_client(self):
+        if self._ai_client is None:
+            self._ai_client = self.make_aiclient()
+        return self._ai_client
+
+    def make_aiclient(self):
+        return make_ai_client(
+            api_key=self.workspace.current_token(),
+            base_url=self.workspace.host + "/serving-endpoints"
+        )
+
+    def ask(
+        self,
+        command: str,
+        *,
+        system: Optional[str] = None,
+        max_tokens: int = 5000,
+        temperature: Optional[float] = None,
+        extra_messages: Optional[List[Dict[str, str]]] = None,
+        **kwargs,
+    ) -> str:
+        """
+        Send a chat prompt to a Databricks Model Serving endpoint using the
+        OpenAI-compatible Chat Completions API and return the assistant's text.
+
+        This is a thin convenience wrapper around:
+            self.ai_client.chat.completions.create(...)
+
+        Parameters
+        ----------
+        command:
+            The user prompt text to send.
+        system:
+            Optional system instruction (prepended as the first message).
+        max_tokens:
+            Upper bound on the number of tokens generated in the response.
+        temperature:
+            Optional sampling temperature. If None, the client/model default is used.
+        extra_messages:
+            Optional list of additional chat messages to insert before the user prompt.
+            Each item should be a dict like {"role": "...", "content": "..."}.
+            Useful for few-shot examples or carrying prior context.
+        **kwargs:
+            Any additional parameters forwarded directly to
+            `chat.completions.create(...)` (e.g. top_p, presence_penalty, etc.).
+
+        Returns
+        -------
+        str
+            The assistant message content (empty string if the API returns no content).
+
+        Raises
+        ------
+        Exception
+            Propagates any exceptions raised by the OpenAI client (HTTP errors,
+            auth errors, invalid request errors, timeouts, etc.).
+        """
+        messages: List[Dict[str, str]] = []
+
+        if system:
+            messages.append({"role": "system", "content": system})
+
+        if extra_messages:
+            messages.extend(extra_messages)
+
+        messages.append({"role": "user", "content": command})
+
+        # Build params cleanly (only include temperature if provided)
+        params: Dict[str, Any] = dict(
+            model=self.model,
+            messages=messages,
+            max_tokens=max_tokens,
+            **kwargs,
+        )
+        if temperature is not None:
+            params["temperature"] = temperature
+
+        resp = self.ai_client.chat.completions.create(**params)
+        return resp.choices[0].message.content or ""

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
 

yggdrasil/databricks/workspaces/workspace.py

@@ -645,15 +645,16 @@ class Workspace:
 
     def clusters(
         self,
+        workspace: Optional["Workspace"] = None,
         cluster_id: Optional[str] = None,
         cluster_name: Optional[str] = None,
     ) -> "Cluster":
         """Return a Cluster helper bound to this workspace.
 
         Args:
+            workspace: Optional workspace override.
             cluster_id: Optional cluster id.
             cluster_name: Optional cluster name.
-            **kwargs: Additional Cluster parameters.
 
         Returns:
             A Cluster instance.
@@ -661,11 +662,30 @@ class Workspace:
         from ..compute.cluster import Cluster
 
         return Cluster(
-            workspace=self,
+            workspace=self if workspace is None else workspace,
             cluster_id=cluster_id,
             cluster_name=cluster_name,
         )
 
+    def loki(
+        self,
+        workspace: Optional["Workspace"] = None,
+    ):
+        """
+        Return a Cluster helper bound to this workspace.
+
+        Args:
+            workspace: Optional workspace override.
+
+        Returns:
+            A Loki AI instance.
+        """
+        from ..ai.loki import Loki
+
+        return Loki(
+            workspace=self if workspace is None else workspace,
+        )
+
 # ---------------------------------------------------------------------------
 # Workspace-bound base class
 # ---------------------------------------------------------------------------
@@ -717,13 +737,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(

yggdrasil/pyutils/__init__.py

@@ -4,5 +4,3 @@ from .retry import retry
 from .parallel import parallelize
 from .python_env import PythonEnv
 from .callable_serde import CallableSerde
-
-__all__ = ["retry", "parallelize", "PythonEnv", "CallableSerde"]

yggdrasil/pyutils/python_env.py

@@ -20,7 +20,7 @@ from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Any, Iterable, Iterator, Mapping, MutableMapping, Optional, Union, List, Tuple
 
-from yggdrasil.pyutils.modules import PipIndexSettings
+from .modules import PipIndexSettings
 
 log = logging.getLogger(__name__)
 

yggdrasil/version.py

@@ -1 +1 @@
-__version__ = "0.1.53"
+__version__ = "0.1.55"