ygg 0.1.57__py3-none-any.whl → 0.1.64__py3-none-any.whl

{ygg-0.1.57.dist-info → ygg-0.1.64.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ygg
-Version: 0.1.57
+Version: 0.1.64
 Summary: Type-friendly utilities for moving data between Python objects, Arrow, Polars, Pandas, Spark, and Databricks
 Author: Yggdrasil contributors
 License:                                  Apache License
@@ -207,7 +207,7 @@ License:                                  Apache License
         
 Project-URL: Homepage, https://github.com/Platob/Yggdrasil
 Project-URL: Repository, https://github.com/Platob/Yggdrasil
-Project-URL: Documentation, https://github.com/Platob/Yggdrasil/tree/main/python/docs
+Project-URL: Documentation, https://github.com/Platob/Yggdrasil
 Keywords: arrow,polars,pandas,spark,databricks,typing,dataclass,serialization
 Classifier: Development Status :: 3 - Alpha
 Classifier: Programming Language :: Python

ygg-0.1.64.dist-info/RECORD

@@ -0,0 +1,74 @@
+ygg-0.1.64.dist-info/licenses/LICENSE,sha256=HrhfyXIkWY2tGFK11kg7vPCqhgh5DcxleloqdhrpyMY,11558
+yggdrasil/__init__.py,sha256=4-ghPak2S6zfMqmnlxW2GCgPb5s79znpKa2hGEGXcE4,24
+yggdrasil/exceptions.py,sha256=NEpbDFn-8ZRsLiEgJicCwrTHNMWAGtdrTJzosfAeVJo,82
+yggdrasil/version.py,sha256=4MIEvUQSlMm4HeF033QtQTWKTLCJXVFMfXEyriOrkl4,22
+yggdrasil/ai/__init__.py,sha256=YEOVsyuvEOvPaZT8XN9xNysS_WOpHTbKgXgnA8up7x0,52
+yggdrasil/ai/session.py,sha256=10ATAnw8FOCpfIg9sNR4meki_MRckUzKZ9Uft4IXwLA,2515
+yggdrasil/ai/sql_session.py,sha256=n92tQjHUBIey6c3EJProiEEwfAtQm07Dtmei4WXzeG0,10812
+yggdrasil/databricks/__init__.py,sha256=0GRBP930ManOvyo-Y5E7bz7F2msnvU677OH6rxzPwd8,87
+yggdrasil/databricks/exceptions.py,sha256=-ZULt0wD5_Rxww11nk4Z46DvS5j18RdKR5ISmbQfUQA,142
+yggdrasil/databricks/compute/__init__.py,sha256=NvdzmaJSNYY1uJthv1hHdBuNu3bD_-Z65DWnaJt9yXg,289
+yggdrasil/databricks/compute/cluster.py,sha256=5frvxXyeJdAdba7vsGqk2dgfxZEveWsswgrCXZIOOSc,42145
+yggdrasil/databricks/compute/command_execution.py,sha256=osLykTE8S_2gle3XCUiEz8nB4dLKywkiJU4DdEQPT5g,20888
+yggdrasil/databricks/compute/exceptions.py,sha256=OgiAzicmSaAyMzPoeQASdHlvOHnfaNwOt-081_XxCHQ,340
+yggdrasil/databricks/compute/execution_context.py,sha256=MoZwuIHIJJaqo--xJcYIGVpA-eOPFaAjOTbLeVnUzrg,29414
+yggdrasil/databricks/compute/remote.py,sha256=sF99i7GXZcC0GiNgO9VO0I26rFbrtnDhK9vrC2gajuw,2623
+yggdrasil/databricks/jobs/__init__.py,sha256=snxGSJb0M5I39v0y3IR-uEeSlZR248cQ_4DJ1sYs-h8,154
+yggdrasil/databricks/jobs/config.py,sha256=9LGeHD04hbfy0xt8_6oobC4moKJh4_DTjZiK4Q2Tqjk,11557
+yggdrasil/databricks/sql/__init__.py,sha256=PetgRp1jEj5K3TgN09FwNUVjVN8YYuGq0cDIOTqsbns,144
+yggdrasil/databricks/sql/engine.py,sha256=nFWeegs91CtjCLzxgZsJwOlAXNVI1v_lfecuFVfKFFY,49979
+yggdrasil/databricks/sql/exceptions.py,sha256=srMR3Y9LQm45rkyxfyCgpgcoGtRRvGKWBEoUHf4kxsg,1762
+yggdrasil/databricks/sql/statement_result.py,sha256=01DzFX1bGDIGHj0OW2ngfVVJ1w1KHlZEfAI934E35CU,15549
+yggdrasil/databricks/sql/types.py,sha256=5G-BM9_eOsRKEMzeDTWUsWW5g4Idvs-czVCpOCrMhdA,6412
+yggdrasil/databricks/sql/warehouse.py,sha256=6rZAQ98avExuhllT-pKrjbIS3RFBMCC2RBVzXouNFo4,18976
+yggdrasil/databricks/workspaces/__init__.py,sha256=dv2zotoFVhNFlTCdRq6gwf5bEzeZkOZszoNZMs0k59g,114
+yggdrasil/databricks/workspaces/filesytem.py,sha256=Z8JXU7_XUEbw9fpTQT1avRQKi-IAP2KemXBMPkUoY4w,9805
+yggdrasil/databricks/workspaces/io.py,sha256=RdgN5lmEYNF5phPRkRMCVHbUl-t3ZUGkKbzgYSTKpII,37420
+yggdrasil/databricks/workspaces/path.py,sha256=k3UB0LhF4hQI-Iza50D5dVjhqNsAdP4KabKiWK7bTWM,55775
+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=f8Ihv3fqo6YYno1yvkXnMpZhQQWwMkklhg-C9MSVyrE,30103
+yggdrasil/dataclasses/__init__.py,sha256=_RkhfF3KC1eSORby1dzvBXQ0-UGG3u6wyUQWX2jq1Pc,108
+yggdrasil/dataclasses/dataclass.py,sha256=LxrCjwvmBnb8yRI_N-c31RHHxB4XoJPixmKg9iBIuaI,1148
+yggdrasil/libs/__init__.py,sha256=zdC9OU0Xy36CLY9mg2drxN6S7isPR8aTLzJA6xVIeLE,91
+yggdrasil/libs/databrickslib.py,sha256=Y99ARtrVKVBTH0qZ0njYr1Oa_757wtsVY4ywH07IdQ4,1109
+yggdrasil/libs/pandaslib.py,sha256=_U4sdFvLAFD16_65RG-RFmcx4c3fvVnALESFaAlT71M,887
+yggdrasil/libs/polarslib.py,sha256=WnnERtMTl__ZPidcZkoV7mb8-c680zcAnJgzAoD3ZE8,1437
+yggdrasil/libs/sparklib.py,sha256=FQ3W1iz2EIpQreorOiQuFt15rdhq2QhGEAWp8Zrbl9A,10177
+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=ESap-WP4A03bv3kvQK2S9OdplQC79YSpTYRJYJkH8EA,16749
+yggdrasil/pyutils/__init__.py,sha256=AOCLAn9ogFYjXxA6AiUudgNQlzM4A4qQIqJ8MJsm5i0,251
+yggdrasil/pyutils/callable_serde.py,sha256=rGCM0gSejwDDS2xbQimH1q_PggAdSpUsa65XKXMj9DI,22766
+yggdrasil/pyutils/equality.py,sha256=Xyf8D1dLUCm3spDEir8Zyj7O4US_fBJwEylJCfJ9slI,3080
+yggdrasil/pyutils/exceptions.py,sha256=1c0xxFvGML5gkDPGzD_Tgw1ff9bGMVygH8ASgeoII2E,3889
+yggdrasil/pyutils/expiring_dict.py,sha256=pr2u25LGwPVbLfsLptiHGovUtYRRo0AMjaJtCtJl7nQ,8477
+yggdrasil/pyutils/mimetypes.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+yggdrasil/pyutils/modules.py,sha256=a0YWunsuA-D-Ho41LrwkN6o7e88NMk7oOzrSzlG0kPQ,11488
+yggdrasil/pyutils/parallel.py,sha256=ubuq2m9dJzWYUyKCga4Y_9bpaeMYUrleYxdp49CHr44,6781
+yggdrasil/pyutils/python_env.py,sha256=d9s6i_fXz6j4f5BnigKxR0dEUJ-4BmqjBD2u2ybX29U,51025
+yggdrasil/pyutils/retry.py,sha256=gXBtn1DdmIYIUmGKOUr8-SUT7MOu97LykN2YR4uocgc,11917
+yggdrasil/pyutils/waiting_config.py,sha256=WiMOiKyGR5iKr83YK4dljn7OCaDpxXMUx8cz-bUNGMg,6255
+yggdrasil/requests/__init__.py,sha256=dMesyzq97_DmI765x0TwaDPEfsxFtgGNgchk8LvEN-o,103
+yggdrasil/requests/msal.py,sha256=XSuKsxEIApfygiWOBBOok_trQk3eeNb5P0f3RAUrtss,6666
+yggdrasil/requests/session.py,sha256=SLnrgHY0Lby7ZxclRFUjHdfM8euN_8bSQEWl7TkJY2U,1461
+yggdrasil/types/__init__.py,sha256=CrLiDeYNM9fO975sE5ufeVKcy7Ca702IsaG2Pk8T3YU,139
+yggdrasil/types/dummy_class.py,sha256=XXM3_ljL4XfY5LeF-WTj-myqHaKAUmWZ23cPDrXAnBM,2327
+yggdrasil/types/file_format.py,sha256=P-3JTa9FzhHj-ndWMGgsF0zxlR_V2q3a_p2R2CwSoRs,273
+yggdrasil/types/python_arrow.py,sha256=mOhyecAxa5u8JWsyTO26OMOWimHHgwLKWlkNSAyIVas,25636
+yggdrasil/types/python_defaults.py,sha256=kT7vuNDxzP_5tsy0aOkzVh1sZN7rKR7mky9nrYiFkl0,11063
+yggdrasil/types/cast/__init__.py,sha256=Oft3pTs2bRM5hT7YqJAuOKTYYk-SACLaMOXUVdafy_I,311
+yggdrasil/types/cast/arrow_cast.py,sha256=IZstOcHjLKPy62TFGgjMSW3ttPGt3hMi6RmDw-92T0E,41623
+yggdrasil/types/cast/cast_options.py,sha256=nDaEvCCs7TBamhTWyDrYf3LVaBWzioIP2Q5_LXrChF4,15532
+yggdrasil/types/cast/pandas_cast.py,sha256=6PaHgQyq06XM3_lebmB4PUlSHATL_0l8GFXMFvZGwvc,8890
+yggdrasil/types/cast/polars_cast.py,sha256=7qs8QC0Kn3GgaKfrzqr-O-2_9uy9ZUrcJV4HIaGwoUM,27584
+yggdrasil/types/cast/polars_pandas_cast.py,sha256=CS0P7teVv15IdX5g7v40RfkH1VMg6b-HM0V_gOfacm8,5071
+yggdrasil/types/cast/registry.py,sha256=OOqIfbIjPH-a3figvu-zTvEtUDTEWhe2xIl3cCA4PRM,20941
+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.64.dist-info/METADATA,sha256=GXhjqjZQyvPanFP8JLjl3BvZU-pLZB_NM6pR2It7Q74,18506
+ygg-0.1.64.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+ygg-0.1.64.dist-info/entry_points.txt,sha256=6q-vpWG3kvw2dhctQ0LALdatoeefkN855Ev02I1dKGY,70
+ygg-0.1.64.dist-info/top_level.txt,sha256=iBe9Kk4VIVbLpgv_p8OZUIfxgj4dgJ5wBg6vO3rigso,10
+ygg-0.1.64.dist-info/RECORD,,

yggdrasil/ai/__init__.py

@@ -0,0 +1,2 @@
+from .session import *
+from .sql_session import *

yggdrasil/ai/session.py

@@ -0,0 +1,87 @@
+# ai_session.py
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional
+
+from ..types.dummy_class import DummyModuleClass
+
+try:
+    from openai import OpenAI
+except:
+    OpenAI = DummyModuleClass
+
+__all__ = ["AISession"]
+
+
+@dataclass
+class AISession(ABC):
+    api_key: str
+    base_url: str
+    model: str = "gemini-2.5-flash"
+
+    client: OpenAI = field(init=False)
+
+    def __post_init__(self) -> None:
+        self.client = OpenAI(api_key=self.api_key, base_url=self.base_url)
+
+    @abstractmethod
+    def system_prompt(self) -> str:
+        raise NotImplementedError
+
+    def build_messages(
+        self,
+        user_prompt: str,
+        *,
+        context_system: Optional[str] = None,
+        extra_instructions: Optional[str] = None,
+    ) -> List[Dict[str, str]]:
+        # Minimal message stack: system + optional context + user
+        sys = self.system_prompt().strip()
+        msgs: List[Dict[str, str]] = [{"role": "system", "content": sys}]
+
+        if context_system:
+            msgs.append({"role": "system", "name": "context", "content": context_system.strip()})
+
+        msg = user_prompt.strip()
+        if extra_instructions:
+            msg = f"{msg}\n\nConstraints:\n{extra_instructions.strip()}"
+
+        msgs.append({"role": "user", "content": msg})
+        return msgs
+
+    def chat(
+        self,
+        user_prompt: str,
+        *,
+        context_system: Optional[str] = None,
+        extra_instructions: Optional[str] = None,
+        temperature: float = 0.0,
+        max_output_tokens: int = 320,
+        strip_code_fences: bool = True,
+    ) -> str:
+        messages = self.build_messages(
+            user_prompt,
+            context_system=context_system,
+            extra_instructions=extra_instructions,
+        )
+
+        resp = self.client.chat.completions.create(
+            model=self.model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_output_tokens,
+        )
+
+        out = (resp.choices[0].message.content or "").strip()
+
+        if strip_code_fences and out.startswith("```"):
+            out = out.split("```", 2)[1].strip()
+            low = out.lower()
+            if low.startswith("sql"):
+                out = out[3:].strip()
+            elif low.startswith("json"):
+                out = out[4:].strip()
+
+        return out

yggdrasil/ai/sql_session.py

@@ -0,0 +1,310 @@
+# sql_session.py
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Dict, Iterable, List, Optional, Tuple
+
+import pyarrow as pa
+
+from .session import AISession
+
+__all__ = ["SQLFlavor", "SQLAISession"]
+
+
+class SQLFlavor(str, Enum):
+    DATABRICKS = "databricks"
+    POSTGRESQL = "postgresql"
+    MONGODB = "mongodb"  # aggregation pipeline JSON array, not SQL
+
+
+def _b2s(x: object) -> str:
+    if x is None:
+        return ""
+    if isinstance(x, bytes):
+        return x.decode("utf-8", errors="ignore").strip()
+    return str(x).strip()
+
+
+def _try_parse_json_blob(s: str) -> Optional[dict]:
+    s = (s or "").strip()
+    if not s or not (s.startswith("{") and s.endswith("}")):
+        return None
+    try:
+        obj = json.loads(s)
+        return obj if isinstance(obj, dict) else None
+    except Exception:
+        return None
+
+
+def parse_qualified_name(name: str) -> Tuple[Optional[str], Optional[str], Optional[str]]:
+    n = (name or "").strip().strip('"').strip()
+    if not n:
+        return None, None, None
+    parts = [p.strip().strip('"') for p in n.split(".") if p.strip()]
+    if len(parts) == 1:
+        return None, None, parts[0]
+    if len(parts) == 2:
+        return None, parts[0], parts[1]
+    return parts[-3], parts[-2], parts[-1]
+
+
+def extract_object_name_from_metadata(metadata: Optional[Dict[object, object]], flavor: SQLFlavor) -> Optional[str]:
+    md_raw = metadata or {}
+    md: Dict[str, str] = {}
+    for k, v in md_raw.items():
+        ks = _b2s(k).lower()
+        vs = _b2s(v)
+        if ks:
+            md[ks] = vs
+
+    for k in ("table_ref", "full_table_name", "qualified_name", "table", "table_name", "name", "object", "object_name"):
+        if k in md:
+            obj = _try_parse_json_blob(md[k])
+            if obj:
+                for kk, vv in obj.items():
+                    if kk is None:
+                        continue
+                    md[str(kk).lower()] = _b2s(vv)
+
+    if flavor == SQLFlavor.MONGODB:
+        for ck in ("collection", "collection_name", "mongo.collection", "mongodb.collection"):
+            if md.get(ck):
+                return md[ck]
+        for tk in ("table", "table_name", "name", "full_table_name", "table_ref"):
+            if md.get(tk):
+                _, _, coll = parse_qualified_name(md[tk])
+                return coll or md[tk]
+        return None
+
+    catalog = md.get("catalog") or md.get("unity_catalog") or md.get("uc_catalog") or md.get("db_catalog") or ""
+    database = md.get("database") or md.get("db") or ""
+    if not catalog and database:
+        catalog = database
+
+    schema_name = md.get("schema") or md.get("namespace") or md.get("database_schema") or ""
+    table = md.get("table") or md.get("table_name") or md.get("relation") or md.get("object") or md.get("object_name") or ""
+
+    if not table:
+        for fk in ("table_ref", "full_table_name", "qualified_name", "name"):
+            if md.get(fk):
+                c, s, t = parse_qualified_name(md[fk])
+                catalog = catalog or (c or "")
+                schema_name = schema_name or (s or "")
+                table = table or (t or "")
+                if table:
+                    break
+
+    if not table:
+        return None
+
+    parts = [p for p in (catalog, schema_name, table) if p]
+    return ".".join(parts) if parts else None
+
+
+def extract_column_comment(field: pa.Field) -> Optional[str]:
+    md = field.metadata or {}
+    if not md:
+        return None
+
+    norm: Dict[str, str] = {}
+    for k, v in md.items():
+        ks = _b2s(k).lower()
+        vs = _b2s(v)
+        if ks:
+            norm[ks] = vs
+
+    for key in ("comment", "description", "doc", "column_comment", "spark.comment", "delta.comment"):
+        val = norm.get(key, "").strip()
+        if val:
+            return val
+
+    for k in ("meta", "metadata", "attrs", "properties"):
+        obj = _try_parse_json_blob(norm.get(k, ""))
+        if isinstance(obj, dict):
+            for kk in ("comment", "description", "doc"):
+                vv = _b2s(obj.get(kk, "")).strip()
+                if vv:
+                    return vv
+
+    return None
+
+
+# Token-light types for context (not for DDL)
+def _arrow_type_compact(dt: pa.DataType) -> str:
+    if pa.types.is_int64(dt):
+        return "i64"
+    if pa.types.is_int32(dt):
+        return "i32"
+    if pa.types.is_int16(dt):
+        return "i16"
+    if pa.types.is_int8(dt):
+        return "i8"
+    if pa.types.is_float64(dt):
+        return "f64"
+    if pa.types.is_float32(dt) or pa.types.is_float16(dt):
+        return "f32"
+    if pa.types.is_boolean(dt):
+        return "bool"
+    if pa.types.is_string(dt) or pa.types.is_large_string(dt):
+        return "str"
+    if pa.types.is_timestamp(dt):
+        return "ts"
+    if pa.types.is_date32(dt) or pa.types.is_date64(dt):
+        return "date"
+    if pa.types.is_decimal(dt):
+        return f"dec({dt.precision},{dt.scale})"
+    if pa.types.is_list(dt) or pa.types.is_large_list(dt) or pa.types.is_fixed_size_list(dt):
+        return f"arr<{_arrow_type_compact(dt.value_type)}>"
+    if pa.types.is_struct(dt):
+        return "struct"
+    if pa.types.is_map(dt):
+        return "map"
+    return "any"
+
+
+def _fields_signature_compact(fields: Iterable[pa.Field], *, include_comments: bool) -> str:
+    # Super compact: col:type[, ...] + optional short comment
+    parts: List[str] = []
+    for f in fields:
+        t = _arrow_type_compact(f.type)
+        if include_comments:
+            c = extract_column_comment(f)
+            if c:
+                c = c.replace("\n", " ").strip()
+                if len(c) > 60:
+                    c = c[:57] + "..."
+                parts.append(f"{f.name}:{t}#{c}")
+                continue
+        parts.append(f"{f.name}:{t}")
+    return ",".join(parts)
+
+
+@dataclass
+class SQLAISession(AISession):
+    flavor: SQLFlavor = SQLFlavor.DATABRICKS
+    max_context_objects: int = 25
+
+    # Registry
+    _fields: Dict[str, List[pa.Field]] = field(default_factory=dict, init=False)
+    _objects: Dict[str, str] = field(default_factory=dict, init=False)  # alias -> table/collection name
+    _meta: Dict[str, Dict[object, object]] = field(default_factory=dict, init=False)  # alias -> schema-level metadata
+
+    # Token controls
+    include_comments_in_context: bool = True
+    max_tables_in_context: int = 16  # big token saver
+
+    def system_prompt(self) -> str:
+        if self.flavor == SQLFlavor.MONGODB:
+            return (
+                "Return ONLY a MongoDB aggregation pipeline as a strict JSON array. No prose.\n"
+                "Use only fields in context. Prefer $match then $group/$project then $sort then $limit."
+            )
+
+        # SQL
+        dialect = "Databricks/Spark SQL" if self.flavor == SQLFlavor.DATABRICKS else "PostgreSQL"
+        return (
+            f"Return ONLY {dialect} query text. No prose, no markdown.\n"
+            "Use only columns in context. Prefer fully-qualified table names when provided.\n"
+            "Keep it short: avoid CTE unless necessary; return a single statement."
+        )
+
+    def set_flavor(self, flavor: SQLFlavor) -> None:
+        self.flavor = flavor
+
+    def register_fields(
+        self,
+        alias: str,
+        fields: List[pa.Field],
+        *,
+        schema_metadata: Optional[Dict[object, object]] = None,
+        object_name: Optional[str] = None,
+        prefer_metadata_object_name: bool = True,
+    ) -> None:
+        if len(self._fields) >= self.max_context_objects and alias not in self._fields:
+            oldest = next(iter(self._fields.keys()))
+            self._fields.pop(oldest, None)
+            self._objects.pop(oldest, None)
+            self._meta.pop(oldest, None)
+
+        self._fields[alias] = fields
+        self._meta[alias] = dict(schema_metadata or {})
+
+        resolved = object_name
+        if resolved is None and prefer_metadata_object_name:
+            resolved = extract_object_name_from_metadata(self._meta[alias], self.flavor)
+        if resolved:
+            self._objects[alias] = resolved
+
+    def _pick_relevant_aliases(self, prompt: str) -> List[str]:
+        """
+        Cheap heuristic to keep context tiny:
+        score aliases by presence of alias/object/column names in prompt.
+        """
+        p = prompt.lower()
+        scores: List[Tuple[int, str]] = []
+
+        for alias, cols in self._fields.items():
+            s = 0
+            if alias.lower() in p:
+                s += 5
+            obj = self._objects.get(alias, "")
+            if obj and obj.lower() in p:
+                s += 5
+            # column hits (cap influence)
+            hit = 0
+            for f in cols:
+                n = f.name.lower()
+                if n in p:
+                    hit += 1
+                    if hit >= 6:
+                        break
+            s += min(hit, 6)
+            scores.append((s, alias))
+
+        scores.sort(reverse=True)
+        picked = [a for s, a in scores if s > 0][: self.max_tables_in_context]
+
+        # fallback: if nothing matched, include first table only (still minimal)
+        if not picked and self._fields:
+            picked = [next(iter(self._fields.keys()))]
+
+        return picked
+
+    def _build_schema_context(self, aliases: List[str]) -> str:
+        lines = ["ctx:"]
+        for a in aliases:
+            obj = self._objects.get(a, "")
+            sig = _fields_signature_compact(
+                self._fields[a],
+                include_comments=self.include_comments_in_context and self.flavor != SQLFlavor.MONGODB,
+            )
+            # tiny format: alias=>obj|cols
+            if obj:
+                lines.append(f"{a}=>{obj}|{sig}")
+            else:
+                lines.append(f"{a}|{sig}")
+        return "\n".join(lines)
+
+    def generate_query(
+        self,
+        user_prompt: str,
+        *,
+        temperature: float = 0.0,
+        max_output_tokens: int = 4200,
+        extra_instructions: Optional[str] = None,
+        tables: Optional[List[str]] = None,
+    ) -> str:
+        # decide which tables to include (token saver)
+        aliases = tables if tables else self._pick_relevant_aliases(user_prompt)
+        context_system = self._build_schema_context(aliases)
+
+        return self.chat(
+            user_prompt,
+            context_system=context_system,
+            extra_instructions=extra_instructions,
+            temperature=temperature,
+            max_output_tokens=max_output_tokens,
+            strip_code_fences=True,
+        )

yggdrasil/databricks/__init__.py

@@ -1,6 +1,3 @@
 """Databricks integrations and helpers for Yggdrasil."""
 
 from .workspaces import *
-from .compute import *
-from .sql import *
-from .jobs import *