ygg 0.1.54__tar.gz → 0.1.56__tar.gz

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

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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.54
+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.54 → ygg-0.1.56}/src/ygg.egg-info/SOURCES.txt

@@ -8,12 +8,14 @@ src/ygg.egg-info/entry_points.txt
 src/ygg.egg-info/requires.txt
 src/ygg.egg-info/top_level.txt
 src/yggdrasil/__init__.py
+src/yggdrasil/exceptions.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/exceptions.py
 src/yggdrasil/databricks/compute/execution_context.py
 src/yggdrasil/databricks/compute/remote.py
 src/yggdrasil/databricks/jobs/__init__.py

ygg-0.1.56/src/yggdrasil/databricks/ai/__init__.py

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

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

@@ -0,0 +1,374 @@
+"""
+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"]
+
+
+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):
+    """
+    Loki wraps an OpenAI-compatible client pointing at Databricks Model Serving endpoints.
+    """
+
+    model: str = "databricks-gemini-2-5-flash"
+    _ai_client: Optional[Any] = 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):
+        host = self.workspace.host.rstrip("/")
+        return make_ai_client(
+            api_key=self.workspace.current_token(),
+            base_url=f"{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: Any,
+    ) -> str:
+        """
+        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})
+
+        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 ""
+
+    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

ygg-0.1.56/src/yggdrasil/databricks/compute/exceptions.py

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

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

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

{ygg-0.1.54 → ygg-0.1.56}/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.56/src/yggdrasil/exceptions.py

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

{ygg-0.1.54 → ygg-0.1.56}/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.56}/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.56/src/yggdrasil/version.py

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

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"