Skip to content

yggdrasil.databricks.loki

loki

DatabricksLoki — the specialized Databricks agent.

A :class:~yggdrasil.loki.Loki that detects its workspace only from the ygg databricks configure session, reasons through a Databricks serving endpoint, and can deploy itself to run on Databricks compute.

The specialized service skills live next to the code they drive — one loki.py per service module (databricks/sql/loki.py, databricks/table/loki.py, …) for isolation. Importing this package imports them all, so they register into the global Loki catalog (ygg loki skills).

DatabricksLoki

DatabricksLoki(*, serving_endpoint: Optional[str] = None)

Bases: Loki

Loki specialized to a single Databricks workspace.

agent_id property

agent_id: int

Stable int64 id derived from user@host (xxhash, not crypto); cached.

aws property

aws: Any

The configured :class:~yggdrasil.aws.AWSClient when AWS is reachable.

Loki as an AWS token provider — the AWS skill fleet rides this client (its resolved credentials / region / role). None when no AWS session is detected.

databricks property

databricks: 'Optional[DatabricksClient]'

Client built from the configure profile (cached). None if unset.

current classmethod

current() -> 'Loki'

The process-global Loki (created on first use).

load_specialists

load_specialists() -> list[str]

Import the specialized skill fleets for every reachable backend.

Databricks problems get the databricks-* skills, AWS problems the aws-* skills — registered only when their backend is detected, so ygg loki skills shows the fleet that actually applies here. Returns the backends whose fleet loaded.

token_info

token_info() -> dict[str, Any]

Non-secret summary of the Databricks credentials Loki provides.

whoami

whoami(*, probe: bool = False) -> 'Optional[str]'

The Databricks user, if reachable. probe allows one network call.

engines

engines() -> 'list[TokenEngine]'

Every known reasoning engine (call .available() to filter).

available_engines

available_engines(*, refresh: bool = False) -> 'dict[str, TokenEngine]'

Reachable engines (name → instance), availability probed in parallel.

Several engines gate on a network round-trip — the Ollama liveness probe, the Databricks backend check — so probing them one after another stacks up their latencies on the startup path. Fanning the available() checks across a small thread pool collapses that to the slowest single probe. Each engine memoizes its own result, so this also warms the caches that later serial available() calls (the status line, :meth:engine, :meth:select) reuse for free.

engine

engine(name: 'Optional[str]' = None) -> 'Optional[TokenEngine]'

Resolve a reasoning engine by name, or the best available one.

select

select(
    text: "Optional[str]" = None,
    *,
    tier: "Optional[str]" = None,
    base: "Optional[str]" = None,
    confirm: "Optional[Callable[[TokenEngine, Optional[str]], bool]]" = None
) -> "Optional[TokenEngine]"

Resource-aware engine choice: keep light work cheap, escalate heavy work to a capable remote — asking first.

Two axes drive the pick:

  • Complexity — an explicit tier (deep → heavy), else the prompt itself (long or reasoning-heavy text → heavy).
  • Resources — a CUDA GPU, or enough CPU + RAM (≥ 4 cores, ≥ 8 GB), decides whether this box can comfortably run a local model.

A session pins a base provider, but complexity moves the choice both ways:

  • remote → local (demote): light work on a capable box drops to a free local model even when the base is a remote API — saving money, silently (no downside to ask about).
  • local → remote (escalate): heavy work climbs to the most capable remote (the base remote if it is one). When this means switching from a free local model up to a paid remote one, confirm(engine, model) is asked first; a falsy answer keeps the work on the cheap/local path.

With no local engine the base simply stands; with no remote the local engine carries even heavy work. Returns an available engine, or None when nothing is reachable.

can_run_local

can_run_local() -> bool

Whether this box can comfortably host a local model — cached.

Wraps :func:yggdrasil.loki.resources.can_run_local but memoizes the result for the process: the probe imports torch to check for a CUDA GPU (slow on the first call, on a box that has it), and engine selection asks this on every turn. Hardware doesn't change within a session.

bootstrap_local

bootstrap_local(
    *,
    model: "Optional[str]" = None,
    pull: bool = True,
    on_progress: "Optional[Callable[[dict[str, Any]], None]]" = None
) -> dict[str, Any]

Ready a free local reasoning engine, lazily installing on demand.

Prefers a reachable Ollama server — ensures the model sized to this workstation (the more RAM/GPU, the larger the default) is pulled, only if missing. Falls back to the HF transformers engine (weights lazy-download on first use). When neither is present, returns what to install. This is the "free local brain" entry point — sized to the box, smart enough for basic setup/config, and able to hand heavier work up to a remote model.

reason

reason(
    prompt: str,
    *,
    system: "Optional[str]" = None,
    engine: "Optional[str]" = None,
    tier: "Optional[str]" = None,
    base: "Optional[str]" = None,
    confirm: "Optional[Callable[[TokenEngine, Optional[str]], bool]]" = None,
    **options: Any
) -> str

Reason about prompt with the best (or named) engine → reply text.

tier ("fast" / "deep") forces the model tier; the default (None) lets the engine pick adaptively from the prompt. A pinned engine is used as-is; otherwise the choice is resource-aware (:meth:select), sticking to the session base and asking confirm before escalating to a paid remote model.

reason_stream

reason_stream(
    prompt: str,
    *,
    system: "Optional[str]" = None,
    engine: "Optional[str]" = None,
    tier: "Optional[str]" = None,
    base: "Optional[str]" = None,
    confirm: "Optional[Callable[[TokenEngine, Optional[str]], bool]]" = None,
    **options: Any
) -> "Iterator[str]"

Stream a reply to prompt — yields text chunks as they arrive.

Same engine/tier resolution as :meth:reason, but live: the chosen engine streams token deltas so the terminal prints them as they come.

act

act(
    task: str,
    *,
    root: str = ".",
    engine: "Optional[str]" = None,
    tier: "Optional[str]" = None,
    max_steps: int = 12,
    read_only: bool = False,
    allow_shell: bool = False,
    allow_web: bool = False,
    confirm: "Optional[Callable[[str], bool]]" = None,
    toolbox: "Optional[Toolbox]" = None,
    on_think: "Optional[Callable[[int], None]]" = None,
    on_step: "Optional[Callable[[dict[str, Any]], None]]" = None
) -> dict[str, Any]

Pursue task autonomously: discover, decide, and modify files.

This is Loki acting on its own — the reason→act→observe loop. The agent's engine plans against a tool catalog (filesystem discovery + edits, optionally a shell), emits one JSON tool call per turn, and Loki runs it and feeds the observation back, until the engine declares it's done or max_steps is hit. The tools are confined to root (the working tree the agent was pointed at).

tier pins the model tier for every turn; left None the engine adapts per turn — cheap early scouting turns, the capable model once the transcript (and the reasoning) grows.

Returns a transcript: the resolved engine, every step (its thought/tool/args/observation), the final answer, whether it completed, and the files_changed list. Pass on_step to stream each completed turn, and on_think(n) to learn when turn n is about to call the (slow) model — the CLI uses the pair to keep a live spinner + step-budget bar running through the otherwise silent reasoning between tool calls.

delegate

delegate(
    tasks: list[str],
    *,
    root: str = ".",
    engine: "Optional[str]" = None,
    tier: "Optional[str]" = None,
    max_steps: int = 8,
    allow_web: bool = True,
    allow_shell: bool = False,
    read_only: bool = False,
    timeout: "Optional[float]" = None,
    on_update: "Optional[Callable[[list[Any]], None]]" = None
) -> list[dict[str, Any]]

Fan tasks out to background process agents and wait for them.

Each task runs as its own ygg loki do subprocess (an isolated act loop), so independent work proceeds in parallel while Loki monitors them — its autonomy multiplier. Returns one summary row per agent (status, elapsed, answer, files_changed). on_update(agents) streams live progress (the CLI renders the dashboard from it).

decompose

decompose(
    goal: str,
    *,
    engine: "Optional[str]" = None,
    tier: "Optional[str]" = None,
    max_tasks: int = 6
) -> list[str]

Break goal into independent subtasks an agent fleet can run in parallel.

Asks the reasoning engine for a JSON array of self-contained tasks (parallel-safe — no ordering between them). Returns the parsed list, capped at max_tasks; falls back to [goal] if nothing parses.

plan

plan(text: str) -> 'AgentPlan'

Classify a request into a structured :class:AgentPlan.

Loki reasoning, optimized: rather than throwing every prompt at one engine, it classifies the problem — a category and solution action (answer / act on files / fetch tabular / ask Genie), the persona to embody (data engineer, analyst, software engineer, trader, confessor, companion, …), the skills likely required, and whether to isolate the work to a specialist (the "databricks on databricks" scheme). Returns an :class:AgentPlan (mapping-compatible).

classify_data

classify_data(text: str) -> dict[str, Any]

Global context: is this request data- or time-series-shaped?

Drives the data path — a positive classification routes a sourced request to tabular fetching + caching (:class:TabularSkill) instead of a plain page fetch. Returns {"data", "timeseries", "why"}.

specialist

specialist(name: str) -> 'Optional[Loki]'

A specialized agent to isolate a category of work, or None.

"databricks" resolves the workspace-bound :class:~yggdrasil.databricks.loki.DatabricksLoki when the SDK and a session are present; otherwise falls back to self.

Cached per name: the REPL asks for the same specialist on every databricks turn, and the resolution (import + singleton lookup + backend check) is stable for the process.

run

run(skill_name: str, **kwargs: Any) -> Any

Dispatch skill skill_name with kwargs, using self as the provider.

The first parameter is skill_name (not name) so a skill that itself takes a name= kwarg — e.g. scaffold(name=…) — can be dispatched as loki.run("scaffold", name="acme") without a clash.

card

card(*, refresh: bool = False) -> dict[str, Any]

Everything Loki knows about itself — identity, reach, skills.

deploy

deploy(
    *,
    name: Optional[str] = None,
    behavior: str = "reason",
    job: bool = True,
    **params: Any
) -> "Job"

Deploy this agent to run on Databricks compute.

Creates (or updates) a Databricks Job that runs the agent on serverless against the seeded ygg image through the single ygg wheel entry point — ygg loki reason ... for the default reason behavior, ygg loki run <behavior> --kwarg k=v ... otherwise (on the runtime ygg loki resolves to this DatabricksLoki). The job is upserted by name and returned; trigger it with job.run().

Requires a workspace seeded with ygg databricks deploy (for the serverless ygg environment) and a ygg databricks configure session.