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

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

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

@@ -1,12 +1,14 @@
-ygg-0.1.55.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+ygg-0.1.56.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
 yggdrasil/__init__.py,sha256=4-ghPak2S6zfMqmnlxW2GCgPb5s79znpKa2hGEGXcE4,24
-yggdrasil/version.py,sha256=RJcC5yver_ugtAU_HgwZ-2HWyszBXpsLZyyNma6G4Dc,22
+yggdrasil/exceptions.py,sha256=NEpbDFn-8ZRsLiEgJicCwrTHNMWAGtdrTJzosfAeVJo,82
+yggdrasil/version.py,sha256=c0ITmemMU7anWgDBUD3t_BAhA3li2gt3XgswCbHv1oU,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/ai/loki.py,sha256=HyVWxzJgfW03YO6TMOTJ1oNvrBJovnqYDn_MeNV6Ni0,11989
 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
+yggdrasil/databricks/compute/exceptions.py,sha256=Ug0ioxu5m2atdTX4OLH0s4R4dylHNxEdn7VhQI66b5M,209
+yggdrasil/databricks/compute/execution_context.py,sha256=mhcwSvKTxgcUHdb7huSEjCVU_feiXSGq0JLyLXldjQM,23952
 yggdrasil/databricks/compute/remote.py,sha256=yicEhyQypssRa2ByscO36s3cBkEgORFsRME9aaq91Pc,3045
 yggdrasil/databricks/jobs/__init__.py,sha256=snxGSJb0M5I39v0y3IR-uEeSlZR248cQ_4DJ1sYs-h8,154
 yggdrasil/databricks/jobs/config.py,sha256=9LGeHD04hbfy0xt8_6oobC4moKJh4_DTjZiK4Q2Tqjk,11557
@@ -19,7 +21,7 @@ yggdrasil/databricks/sql/warehouse.py,sha256=1J0dyQLJb-OS1_1xU1eAVZ4CoL2-FhFeowK
 yggdrasil/databricks/workspaces/__init__.py,sha256=dv2zotoFVhNFlTCdRq6gwf5bEzeZkOZszoNZMs0k59g,114
 yggdrasil/databricks/workspaces/filesytem.py,sha256=Z8JXU7_XUEbw9fpTQT1avRQKi-IAP2KemXBMPkUoY4w,9805
 yggdrasil/databricks/workspaces/io.py,sha256=PAoxIxYvTC162Dx2qL2hk8oAdt8BnYrQ3jJHcJm4VkA,33116
-yggdrasil/databricks/workspaces/path.py,sha256=KkvLFHrps3UFr4ogYdESbJHEMfQBcWfWfXjlrv_7rTU,55180
+yggdrasil/databricks/workspaces/path.py,sha256=h1j3bvjwKcDhJvlU_kAaLcLVz4jrdaWgjqPQMarZRHU,55233
 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=5DCPz5io_rmrpGNi5I6RChmyZ8kjlNUFGQl8mzQJThg,25511
@@ -59,8 +61,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.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.56.dist-info/METADATA,sha256=Rr4DBB8q39XEEzBTQtYB_pja448GR74cEh3gc8-BKvY,18528
+ygg-0.1.56.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+ygg-0.1.56.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
+ygg-0.1.56.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
+ygg-0.1.56.dist-info/RECORD,,

yggdrasil/databricks/ai/loki.py

@@ -1,44 +1,79 @@
-from typing import Optional, Dict, Any, List
-from dataclasses import field, dataclass
+"""
+loki.py
+
+Databricks Model Serving (OpenAI-compatible) wrapper with:
+- Loki.ask(): stateless call
+- TradingChatSession: stateful commodity trading analytics chat
+- SqlChatSession: stateful Databricks SQL generator chat
+
+Important constraint:
+- Gemini models only support ONE system prompt.
+  => We must NOT send multiple system messages.
+  => We fold summary + context blocks into a single system string.
+"""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Union
 
 from ..workspaces.workspace import WorkspaceService
 
 try:
+    from openai import OpenAI as _OpenAI  # noqa: F401
+except ImportError:
+    _OpenAI = None  # type: ignore
+
+
+def make_ai_client(api_key: str, base_url: str):
+    """Late import so module can load even if openai isn't installed."""
     from openai import OpenAI
+    return OpenAI(api_key=api_key, base_url=base_url)
+
+
+__all__ = ["Loki", "TradingChatSession", "SqlChatSession"]
 
-    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"
-]
+DEFAULT_TRADING_SYSTEM_PROMPT = """You are Loki: a conversational commodity trading analytics copilot.
+
+Scope:
+- Commodity trading analytics: curves, forwards, spreads, basis, hedging, risk, PnL explain, inventory, scheduling.
+- Databricks-first workflows: Spark/Delta/Unity Catalog, Databricks SQL, performant Python.
+
+Rules:
+- Do NOT invent prices, positions, PnL, risk, or market facts not provided.
+- State assumptions explicitly (units, time conventions, delivery months, calendars).
+- Prefer actionable output (SQL + efficient Python). Avoid slow patterns.
+- If data is missing, list exactly what you need and proceed with a reasonable template.
+
+Style:
+- Be concise, practical, performance-focused.
+"""
+
+DEFAULT_SQL_SYSTEM_PROMPT = """You are LokiSQL: a Databricks SQL generator for commodity trading analytics.
+
+Hard rules:
+- Output ONLY SQL unless the user explicitly asks for explanation.
+- Use Databricks SQL / Spark SQL dialect.
+- Prefer readable CTEs, explicit column lists, deterministic joins.
+- Do NOT invent table/column names. If missing, use placeholders like <table>, <col>.
+- If ambiguous, output best-effort SQL template with SQL comments (-- TODO ...) and placeholders.
+- Performance: push filters early, avoid exploding joins, avoid SELECT *.
+
+Default assumptions:
+- Dates UTC unless specified.
+"""
 
 
 @dataclass
 class Loki(WorkspaceService):
-    model: str = "databricks-gemini-2-5-flash"
+    """
+    Loki wraps an OpenAI-compatible client pointing at Databricks Model Serving endpoints.
+    """
 
-    _ai_client: Optional[OpenAI] = field(repr=False, hash=False, default=None)
+    model: str = "databricks-gemini-2-5-flash"
+    _ai_client: Optional[Any] = field(repr=False, hash=False, default=None)
 
     @property
     def ai_client(self):
@@ -47,9 +82,10 @@ class Loki(WorkspaceService):
         return self._ai_client
 
     def make_aiclient(self):
+        host = self.workspace.host.rstrip("/")
         return make_ai_client(
             api_key=self.workspace.current_token(),
-            base_url=self.workspace.host + "/serving-endpoints"
+            base_url=f"{host}/serving-endpoints",
         )
 
     def ask(
@@ -60,55 +96,25 @@ class Loki(WorkspaceService):
         max_tokens: int = 5000,
         temperature: Optional[float] = None,
         extra_messages: Optional[List[Dict[str, str]]] = None,
-        **kwargs,
+        **kwargs: Any,
     ) -> 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.).
+        Stateless single call to the model.
+
+        NOTE (Gemini constraint):
+        - Provide at most ONE system prompt (i.e., a single system message).
+        - Do not pass additional messages with role="system".
         """
         messages: List[Dict[str, str]] = []
-
         if system:
             messages.append({"role": "system", "content": system})
 
         if extra_messages:
+            # IMPORTANT: caller must not include any "system" roles here for Gemini models
             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,
@@ -120,3 +126,249 @@ class Loki(WorkspaceService):
 
         resp = self.ai_client.chat.completions.create(**params)
         return resp.choices[0].message.content or ""
+
+    def new_trading_chat(
+        self,
+        *,
+        system_prompt: str = DEFAULT_TRADING_SYSTEM_PROMPT,
+        max_context_turns: int = 20,
+        max_context_chars: int = 120_000,
+    ) -> "TradingChatSession":
+        return TradingChatSession(
+            loki=self,
+            system_prompt=system_prompt,
+            max_context_turns=max_context_turns,
+            max_context_chars=max_context_chars,
+        )
+
+    def new_sql_chat(
+        self,
+        *,
+        system_prompt: str = DEFAULT_SQL_SYSTEM_PROMPT,
+        max_context_turns: int = 20,
+        max_context_chars: int = 120_000,
+    ) -> "SqlChatSession":
+        return SqlChatSession(
+            loki=self,
+            system_prompt=system_prompt,
+            max_context_turns=max_context_turns,
+            max_context_chars=max_context_chars,
+        )
+
+
+@dataclass
+class _BaseChatSession:
+    """
+    Stateful session that maintains history + injected context blocks.
+
+    Gemini constraint:
+    - We must fold ALL system content into one system string.
+    - Therefore summary/context_blocks are concatenated into the system prompt.
+    """
+    loki: Loki
+    system_prompt: str
+
+    history: List[Dict[str, str]] = field(default_factory=list)
+    summary: Optional[str] = None
+    context_blocks: List[str] = field(default_factory=list)
+
+    max_context_turns: int = 20
+    max_context_chars: int = 120_000
+
+    def reset(self) -> None:
+        self.history.clear()
+        self.summary = None
+        self.context_blocks.clear()
+
+    def add_context(self, title: str, payload: Union[str, Dict[str, Any], List[Any]]) -> None:
+        if isinstance(payload, str):
+            payload_str = payload
+        else:
+            payload_str = json.dumps(payload, ensure_ascii=False, indent=2)
+
+        self.context_blocks.append(f"[Context: {title}]\n{payload_str}".strip())
+        self._trim()
+
+    def _estimate_chars(self, msgs: List[Dict[str, str]]) -> int:
+        return sum(len(m.get("content", "")) for m in msgs)
+
+    def _build_system(self, extra_system: Optional[str] = None) -> str:
+        parts: List[str] = [self.system_prompt.strip()]
+        if extra_system:
+            parts.append(extra_system.strip())
+        if self.summary:
+            parts.append(f"[ConversationSummary]\n{self.summary}".strip())
+        if self.context_blocks:
+            parts.append("\n\n".join(self.context_blocks).strip())
+        return "\n\n".join(p for p in parts if p)
+
+    def _trim(self) -> None:
+        # Turn trim (keep last N turns => N*2 messages)
+        if self.max_context_turns > 0:
+            max_msgs = self.max_context_turns * 2
+            if len(self.history) > max_msgs:
+                self.history = self.history[-max_msgs:]
+
+        # Char trim: shrink history first, then context blocks if needed
+        def total_chars() -> int:
+            sys_len = len(self._build_system())
+            return sys_len + self._estimate_chars(self.history)
+
+        while total_chars() > self.max_context_chars and self.history:
+            self.history = self.history[1:]
+
+        while total_chars() > self.max_context_chars and self.context_blocks:
+            self.context_blocks = self.context_blocks[1:]
+
+
+@dataclass
+class TradingChatSession(_BaseChatSession):
+    """
+    Commodity trading analytics chat session.
+    Optionally returns structured JSON for downstream automation.
+    """
+
+    def chat(
+        self,
+        user_text: str,
+        *,
+        structured: bool = True,
+        max_tokens: int = 12000,
+        temperature: Optional[float] = None,
+        **kwargs: Any,
+    ) -> Union[str, Dict[str, Any]]:
+        self._trim()
+
+        extra_system = None
+        if structured:
+            extra_system = (
+                "Respond ONLY as valid JSON with keys: "
+                "final_answer (string), assumptions (array of strings), data_needed (array of strings), "
+                "sql (string or null), python (string or null). "
+                "No markdown. No extra keys."
+            )
+
+        system = self._build_system(extra_system=extra_system)
+
+        assistant_text = self.loki.ask(
+            user_text,
+            system=system,
+            extra_messages=self.history,  # NOTE: history must contain no system roles
+            max_tokens=max_tokens,
+            temperature=temperature,
+            **kwargs,
+        )
+
+        self.history.append({"role": "user", "content": user_text})
+        self.history.append({"role": "assistant", "content": assistant_text})
+        self._trim()
+
+        if structured:
+            parsed = _try_parse_json_object(assistant_text)
+            if parsed is not None:
+                return parsed
+
+        return assistant_text
+
+
+@dataclass
+class SqlChatSession(_BaseChatSession):
+    """
+    SQL-only conversational session that generates Databricks SQL.
+
+    Uses a single system message with strict instructions to output SQL only.
+    """
+
+    def generate_sql(
+        self,
+        request: str,
+        *,
+        max_tokens: int = 12000,
+        temperature: Optional[float] = None,
+        sql_only: bool = True,
+        **kwargs: Any,
+    ) -> str:
+        self._trim()
+
+        extra_system = None
+        if sql_only:
+            extra_system = (
+                "Reminder: Output ONLY SQL. "
+                "If ambiguity exists, use SQL comments (-- TODO ...) and placeholders, but still output SQL only."
+            )
+
+        system = self._build_system(extra_system=extra_system)
+
+        sql = self.loki.ask(
+            request,
+            system=system,
+            extra_messages=self.history,  # history must contain no system roles
+            max_tokens=max_tokens,
+            temperature=temperature,
+            **kwargs,
+        ).strip()
+
+        self.history.append({"role": "user", "content": request})
+        self.history.append({"role": "assistant", "content": sql})
+        self._trim()
+
+        return _strip_sql_fences(sql)
+
+
+def _strip_sql_fences(text: str) -> str:
+    t = text.strip()
+    if t.startswith("```"):
+        lines = t.splitlines()
+        lines = lines[1:]  # drop ``` or ```sql
+        if lines and lines[-1].strip().startswith("```"):
+            lines = lines[:-1]
+        return "\n".join(lines).strip()
+    return t
+
+
+def _strip_markdown_fences(text: str) -> str:
+    """
+    Remove ```lang ... ``` fences if present.
+    Keeps inner content unchanged.
+    """
+    t = text.strip()
+    if not t.startswith("```"):
+        return t
+
+    lines = t.splitlines()
+    if not lines:
+        return t
+
+    # Drop first line: ``` or ```json
+    lines = lines[1:]
+
+    # Drop last line if it's ```
+    if lines and lines[-1].strip().startswith("```"):
+        lines = lines[:-1]
+
+    return "\n".join(lines).strip()
+
+
+def _try_parse_json_object(text: str) -> Optional[Dict[str, Any]]:
+    t = _strip_markdown_fences(text).strip()
+
+    # Best effort extraction if there's extra junk around JSON
+    if not t.startswith("{"):
+        start = t.find("{")
+        end = t.rfind("}")
+        if start != -1 and end != -1 and end > start:
+            t = t[start : end + 1]
+
+    try:
+        obj = json.loads(t)
+    except Exception:
+        return None
+
+    if not isinstance(obj, dict):
+        return None
+
+    required = {"final_answer", "assumptions", "data_needed", "sql", "python"}
+    if not required.issubset(set(obj.keys())):
+        return None
+
+    return obj

yggdrasil/databricks/compute/exceptions.py

@@ -0,0 +1,14 @@
+from ...exceptions import YGGException
+
+__all__ = [
+    "ComputeException",
+    "CommandAborted"
+]
+
+
+class ComputeException(YGGException):
+    pass
+
+
+class CommandAborted(YGGException):
+    pass

yggdrasil/databricks/compute/execution_context.py

@@ -16,6 +16,7 @@ from threading import Thread
 from types import ModuleType
 from typing import TYPE_CHECKING, Optional, Any, Callable, List, Dict, Union, Iterable, Tuple
 
+from .exceptions import CommandAborted
 from ...libs.databrickslib import databricks_sdk
 from ...pyutils.exceptions import raise_parsed_traceback
 from ...pyutils.expiring_dict import ExpiringDict
@@ -110,16 +111,12 @@ class ExecutionContext:
     def __exit__(self, exc_type, exc_val, exc_tb):
         """Exit the context manager and close the remote context if created."""
         if not self._was_connected:
-            self.close()
+            self.close(wait=False)
         self.cluster.__exit__(exc_type, exc_val=exc_val, exc_tb=exc_tb)
 
     def __del__(self):
         """Best-effort cleanup for the remote execution context."""
-        if self.context_id:
-            try:
-                Thread(target=self.close).start()
-            except BaseException:
-                pass
+        self.close(wait=False)
 
     @property
     def remote_metadata(self) -> RemoteMetadata:
@@ -180,7 +177,7 @@ print(json.dumps(meta))"""
         """
         return self.cluster.workspace.sdk()
 
-    def create_command(
+    def create(
         self,
         language: "Language",
     ) -> any:
@@ -197,15 +194,17 @@ print(json.dumps(meta))"""
             self.cluster
         )
 
+        client = self._workspace_client().command_execution
+
         try:
-            created = self._workspace_client().command_execution.create_and_wait(
+            created = client.create_and_wait(
                 cluster_id=self.cluster.cluster_id,
                 language=language,
             )
         except:
             self.cluster.ensure_running()
 
-            created = self._workspace_client().command_execution.create_and_wait(
+            created = client.create_and_wait(
                 cluster_id=self.cluster.cluster_id,
                 language=language,
             )
@@ -217,42 +216,38 @@ print(json.dumps(meta))"""
 
         created = getattr(created, "response", created)
 
-        return created
+        self.context_id = created.id
+
+        return self
 
     def connect(
         self,
-        language: Optional["Language"] = None
+        language: Optional["Language"] = None,
+        reset: bool = False
     ) -> "ExecutionContext":
         """Create a remote command execution context if not already open.
 
         Args:
             language: Optional language override for the context.
+            reset: Reset existing if connected
 
         Returns:
             The connected ExecutionContext instance.
         """
         if self.context_id is not None:
-            return self
+            if not reset:
+                return self
 
-        self.language = language or self.language
+            self.close(wait=False)
 
-        if self.language is None:
-            self.language = Language.PYTHON
+        language = language or self.language
 
-        ctx = self.create_command(language=self.language)
+        if language is None:
+            language = Language.PYTHON
 
-        context_id = ctx.id
-        if not context_id:
-            raise RuntimeError("Failed to create command execution context")
+        return self.create(language=language)
 
-        self.context_id = context_id
-        LOGGER.info(
-            "Opened execution context for %s",
-            self
-        )
-        return self
-
-    def close(self) -> None:
+    def close(self, wait: bool = True) -> None:
         """Destroy the remote command execution context if it exists.
 
         Returns:
@@ -261,12 +256,23 @@ print(json.dumps(meta))"""
         if not self.context_id:
             return
 
+        client = self._workspace_client()
+
         try:
-            self._workspace_client().command_execution.destroy(
-                cluster_id=self.cluster.cluster_id,
-                context_id=self.context_id,
-            )
-        except Exception:
+            if wait:
+                client.command_execution.destroy(
+                    cluster_id=self.cluster.cluster_id,
+                    context_id=self.context_id,
+                )
+            else:
+                Thread(
+                    target=client.command_execution.destroy,
+                    kwargs={
+                        "cluster_id": self.cluster.cluster_id,
+                        "context_id": self.context_id,
+                    }
+                ).start()
+        except BaseException:
             # non-fatal: context cleanup best-effort
             pass
         finally:
@@ -465,7 +471,18 @@ print(json.dumps(meta))"""
         )
 
         try:
-            return self._decode_result(result, result_tag=result_tag, print_stdout=print_stdout)
+            return self._decode_result(
+                result,
+                result_tag=result_tag,
+                print_stdout=print_stdout
+            )
+        except CommandAborted:
+            return self.connect(language=self.language, reset=True).execute_command(
+                command=command,
+                timeout=timeout,
+                result_tag=result_tag,
+                print_stdout=print_stdout
+            )
         except ModuleNotFoundError as remote_module_error:
             _MOD_NOT_FOUND_RE = re.compile(r"No module named ['\"]([^'\"]+)['\"]")
             module_name = _MOD_NOT_FOUND_RE.search(str(remote_module_error))
@@ -660,6 +677,9 @@ with zipfile.ZipFile(buf, "r") as zf:
         if res.result_type == ResultType.ERROR:
             message = res.cause or "Command execution failed"
 
+            if "client terminated the session" in message:
+                raise CommandAborted(message)
+
             if self.language == Language.PYTHON:
                 raise_parsed_traceback(message)
 
@@ -668,6 +688,7 @@ with zipfile.ZipFile(buf, "r") as zf:
                 or getattr(res, "stack_trace", None)
                 or getattr(res, "traceback", None)
             )
+
             if remote_tb:
                 message = f"{message}\n{remote_tb}"
 

yggdrasil/databricks/workspaces/path.py

@@ -30,7 +30,7 @@ from ...types.cast.registry import convert, register_converter
 from ...types.file_format import ExcelFileFormat
 
 if databricks is not None:
-    from databricks.sdk.service.catalog import VolumeType, PathOperation, VolumeInfo
+    from databricks.sdk.service.catalog import VolumeType, VolumeInfo
     from databricks.sdk.service.workspace import ObjectType
     from databricks.sdk.errors.platform import (
         NotFound,
@@ -1236,6 +1236,8 @@ class DatabricksPath:
         self,
         operation: Optional["PathOperation"] = None
     ):
+        from databricks.sdk.service.catalog import PathOperation
+
         if self.kind != DatabricksPathKind.VOLUME:
             raise ValueError(f"Cannot generate temporary credentials for {repr(self)}")
 

yggdrasil/exceptions.py

@@ -0,0 +1,7 @@
+__all__ = [
+    "YGGException"
+]
+
+
+class YGGException(Exception):
+    pass

yggdrasil/version.py

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