ygg 0.1.54__tar.gz → 0.1.55__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.54
+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.54 → ygg-0.1.55}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ygg"
-version = "0.1.54"
+version = "0.1.55"
 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.54 → ygg-0.1.55}/src/ygg.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.54
+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.55/src/yggdrasil/databricks/ai/__init__.py

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

ygg-0.1.55/src/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 ""

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

@@ -652,9 +652,9 @@ class Workspace:
         """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.
@@ -662,7 +662,7 @@ 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,
         )
@@ -671,10 +671,19 @@ class Workspace:
         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,
+            workspace=self if workspace is None else workspace,
         )
 
 # ---------------------------------------------------------------------------

{ygg-0.1.54 → ygg-0.1.55}/src/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"]

{ygg-0.1.54 → ygg-0.1.55}/src/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__)
 

ygg-0.1.55/src/yggdrasil/version.py

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

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

@@ -1,53 +0,0 @@
-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.54/src/yggdrasil/version.py

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