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

yggdrasil/databricks/ai/loki.py

@@ -1,374 +0,0 @@
-"""
-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