databao-context-engine 0.1.1__py3-none-any.whl

databao_context_engine/llm/factory.py

@@ -0,0 +1,61 @@
+from databao_context_engine.llm.config import OllamaConfig
+from databao_context_engine.llm.descriptions.ollama import OllamaDescriptionProvider
+from databao_context_engine.llm.embeddings.ollama import OllamaEmbeddingProvider
+from databao_context_engine.llm.install import resolve_ollama_bin
+from databao_context_engine.llm.runtime import OllamaRuntime
+from databao_context_engine.llm.service import OllamaService
+
+
+def _create_ollama_service_common(
+    *,
+    host: str,
+    port: int,
+    ensure_ready: bool,
+) -> OllamaService:
+    bin_path = resolve_ollama_bin()
+    config = OllamaConfig(host=host, port=port, bin_path=bin_path)
+    service = OllamaService(config)
+
+    if ensure_ready:
+        runtime = OllamaRuntime(config=config, service=service)
+        runtime.start_and_await(timeout=120)
+
+    return service
+
+
+def create_ollama_service(
+    *,
+    host: str = "127.0.0.1",
+    port: int = 11434,
+    ensure_ready: bool = True,
+) -> OllamaService:
+    return _create_ollama_service_common(
+        host=host,
+        port=port,
+        ensure_ready=ensure_ready,
+    )
+
+
+def create_ollama_embedding_provider(
+    service: OllamaService,
+    *,
+    model_id: str = "nomic-embed-text:v1.5",
+    dim: int = 768,
+    pull_if_needed: bool = True,
+) -> OllamaEmbeddingProvider:
+    if pull_if_needed:
+        service.pull_model_if_needed(model=model_id, timeout=900)
+
+    return OllamaEmbeddingProvider(service=service, model_id=model_id, dim=dim)
+
+
+def create_ollama_description_provider(
+    service: OllamaService,
+    *,
+    model_id: str = "llama3.2:1b",
+    pull_if_needed: bool = True,
+):
+    if pull_if_needed:
+        service.pull_model_if_needed(model=model_id, timeout=900)
+
+    return OllamaDescriptionProvider(service=service, model_id=model_id)

databao_context_engine/llm/install.py

@@ -0,0 +1,227 @@
+import hashlib
+import logging
+import os
+import shutil
+import stat
+import sys
+import tarfile
+import tempfile
+from pathlib import Path
+from typing import NamedTuple
+from zipfile import ZipFile
+
+from databao_context_engine.system.properties import get_dce_path
+
+MANAGED_OLLAMA_BIN = Path(get_dce_path() / "ollama/bin/ollama").expanduser()
+
+logger = logging.getLogger(__name__)
+
+
+class ArtifactInfo(NamedTuple):
+    name: str
+    sha256: str
+
+
+DEFAULT_VERSION = "v0.13.0"
+
+ARTIFACTS: dict[str, ArtifactInfo] = {
+    "darwin": ArtifactInfo(
+        "ollama-darwin.tgz",
+        "fa4ca04c48453c5ff81447d0630e996ee3e6b6af76a9eba52c69c0732f748161",
+    ),
+    "linux-amd64": ArtifactInfo(
+        "ollama-linux-amd64.tgz",
+        "c5e5b4840008d9c9bf955ec32c32b03afc57c986ac1c382d44c89c9f7dd2cc30",
+    ),
+    "linux-arm64": ArtifactInfo(
+        "ollama-linux-arm64.tgz",
+        "05eb97b87c690fa82626c6f4c7d656ae46ad5f2b7ee6aa324cc19dd88b89982b",
+    ),
+    "windows-amd64": ArtifactInfo(
+        "ollama-windows-amd64.zip",
+        "0fc913fc3763b8d2a490f2be90a51d474491ee22ea5a43ff31f1c58301a89656",
+    ),
+    "windows-arm64": ArtifactInfo(
+        "ollama-windows-arm64.zip",
+        "84c395e4187bd560cfc7c26b0142d970bcbdf0e0214b007bc527b7954430ea21",
+    ),
+}
+
+
+def resolve_ollama_bin() -> str:
+    """
+    Decide which `ollama` binary to use, in this order:
+
+    1. DCE_OLLAMA_BIN env var, if set and exists
+    2. `ollama` found on PATH
+    3. Managed installation under MANAGED_OLLAMA_BIN
+
+    Returns the full path to the binary
+    """
+    override = os.environ.get("DCE_OLLAMA_BIN")
+    if override:
+        p = Path(override).expanduser()
+        if p.is_file() and os.access(p, os.X_OK):
+            return str(p)
+
+    system_ollama = shutil.which("ollama")
+    if system_ollama:
+        return system_ollama
+
+    if not MANAGED_OLLAMA_BIN.exists():
+        logger.info("No existing Ollama installation detected. We will download and install Ollama.")
+        install_ollama_to(MANAGED_OLLAMA_BIN)
+
+    return str(MANAGED_OLLAMA_BIN)
+
+
+def _detect_platform() -> str:
+    """
+    Return one of: 'darwin', 'linux-amd64', 'linux-arm64', 'windows-amd64', 'windows-arm64'.
+    """
+    os_name = sys.platform.lower()
+    arch = (os.uname().machine if hasattr(os, "uname") else "").lower()
+
+    if os_name.startswith("darwin"):
+        return "darwin"
+    if os_name.startswith("win"):
+        if "arm" in arch or "aarch64" in arch:
+            return "windows-arm64"
+        return "windows-amd64"
+    if os_name.startswith("linux"):
+        if "arm" in arch or "aarch64" in arch:
+            return "linux-arm64"
+        return "linux-amd64"
+
+    raise RuntimeError(f"Unsupported OS/arch: os={os_name!r} arch={arch!r}")
+
+
+def _download_to_temp(url: str) -> Path:
+    """
+    Download to a temporary file and return its path.
+    """
+    import urllib.request
+
+    tmp_dir = Path(tempfile.mkdtemp(prefix="ollama-download-"))
+    file_name = url.rsplit("/", 1)[-1]
+    dest = tmp_dir / file_name
+
+    logger.info("Downloading %s to %s", url, dest)
+    with urllib.request.urlopen(url) as resp, dest.open("wb") as out:
+        shutil.copyfileobj(resp, out)
+
+    return dest
+
+
+def _verify_sha256(path: Path, expected_hex: str) -> None:
+    """
+    Verify SHA-256 of path matches expected_hex
+    """
+    h = hashlib.sha256()
+    with path.open("rb") as f:
+        for chunk in iter(lambda: f.read(8192), b""):
+            h.update(chunk)
+    actual = h.hexdigest()
+    if actual.lower() != expected_hex.lower():
+        raise RuntimeError(f"SHA256 mismatch for {path}: expected {expected_hex}, got {actual}")
+
+
+def _extract_archive(archive: Path, target_dir: Path) -> None:
+    """
+    Extract archive into target_dir.
+    """
+    name = archive.name.lower()
+    target_dir.mkdir(parents=True, exist_ok=True)
+
+    if name.endswith(".zip"):
+        with ZipFile(archive, "r") as zf:
+            zf.extractall(target_dir)
+    elif name.endswith(".tgz") or name.endswith(".tar.gz"):
+        with tarfile.open(archive, "r:gz") as tf:
+            tf.extractall(target_dir)
+    else:
+        raise RuntimeError(f"Unsupported archive format: {archive}")
+
+
+def _ensure_executable(path: Path) -> None:
+    """
+    Mark path as executable
+    """
+    try:
+        mode = path.stat().st_mode
+        path.chmod(mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH)
+    except Exception:
+        pass
+
+
+def install_ollama_to(target: Path) -> None:
+    """
+    Ensure an Ollama binary exists.
+
+    If it doesn't exist, this will:
+    - detect OS
+    - download the archive from GitHub
+    - verify its SHA-256 checksum
+    - extract into the installation directory
+    - make the binary executable
+    """
+    target = target.expanduser()
+    if target.parent.name == "bin":
+        install_root = target.parent.parent
+    else:
+        install_root = target.parent
+
+    install_root.mkdir(parents=True, exist_ok=True)
+
+    platform_key = _detect_platform()
+    try:
+        artifact = ARTIFACTS[platform_key]
+    except KeyError as e:
+        raise RuntimeError(f"Unsupported platform: {platform_key}") from e
+
+    url = f"https://github.com/ollama/ollama/releases/download/{DEFAULT_VERSION}/{artifact.name}"
+    archive_path = _download_to_temp(url)
+
+    try:
+        _verify_sha256(archive_path, artifact.sha256)
+        logger.info("Verified SHA256 for %s", archive_path.name)
+
+        _extract_archive(archive_path, install_root)
+
+        candidates: list[Path] = []
+        if sys.platform.startswith("win"):
+            candidates.extend(
+                [
+                    install_root / "ollama.exe",
+                    install_root / "bin" / "ollama.exe",
+                ]
+            )
+        else:
+            candidates.extend(
+                [
+                    install_root / "ollama",
+                    install_root / "bin" / "ollama",
+                ]
+            )
+
+        binary: Path | None = None
+        for c in candidates:
+            if c.exists():
+                binary = c
+                break
+
+        if binary is None:
+            raise RuntimeError(f"Installed Ollama archive but could not find binary under {install_root}")
+
+        if binary.resolve() != target.resolve():
+            target.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy2(binary, target)
+
+        _ensure_executable(target)
+        logger.info("Ollama installed at %s", target)
+
+    finally:
+        try:
+            archive_path.unlink(missing_ok=True)
+        except Exception:
+            pass

databao_context_engine/llm/runtime.py

@@ -0,0 +1,73 @@
+import logging
+import os
+import subprocess
+
+from databao_context_engine.llm.config import OllamaConfig
+from databao_context_engine.llm.service import OllamaService
+
+logger = logging.getLogger(__name__)
+
+
+class OllamaRuntime:
+    def __init__(self, service: OllamaService, config: OllamaConfig | None = None):
+        self._service = service
+        self._config = config or OllamaConfig()
+
+    def start_if_needed(self) -> subprocess.Popen | None:
+        if self._service.is_healthy():
+            return None
+
+        logger.info("Ollama server not running. Starting Ollama server...")
+        cmd = [self._config.bin_path, "serve"]
+        env = os.environ.copy()
+        env["OLLAMA_HOST"] = f"{self._config.host}:{self._config.port}"
+        if self._config.extra_env:
+            env.update(self._config.extra_env)
+
+        stdout = subprocess.DEVNULL
+
+        proc = subprocess.Popen(
+            cmd,
+            cwd=str(self._config.work_dir) if self._config.work_dir else None,
+            env=env,
+            stdout=stdout,
+            stderr=subprocess.STDOUT,
+            text=False,
+            close_fds=os.name != "nt",
+        )
+
+        return proc
+
+    def start_and_await(
+        self,
+        *,
+        timeout: float = 60.0,
+        poll_interval: float = 0.5,
+    ) -> subprocess.Popen | None:
+        already_healthy = self._service.is_healthy()
+        proc: subprocess.Popen | None = None
+
+        if not already_healthy:
+            proc = self.start_if_needed()
+
+        ok = self._service.wait_until_healthy(timeout=timeout, poll_interval=poll_interval)
+        if ok:
+            if proc is not None:
+                logger.info("Started Ollama server")
+            else:
+                logger.debug("Ollama server was already running")
+            return proc
+
+        if proc is not None:
+            try:
+                proc.terminate()
+            except Exception:
+                pass
+            try:
+                proc.kill()
+            except Exception:
+                pass
+
+        raise TimeoutError(
+            f"Timed out waiting for Ollama to become healthy at http://{self._config.host}:{self._config.port}"
+        )

databao_context_engine/llm/service.py

@@ -0,0 +1,159 @@
+import logging
+import textwrap
+import time
+from typing import Any
+
+import requests
+
+from databao_context_engine.llm.config import OllamaConfig
+from databao_context_engine.llm.errors import OllamaPermanentError, OllamaTransientError
+
+logger = logging.getLogger(__name__)
+
+
+class OllamaService:
+    def __init__(self, config: OllamaConfig, session: requests.Session | None = None):
+        self._base = config.base_url.rstrip("/")
+        self._timeout = config.timeout
+        self._headers = config.headers
+        self._session = session or requests.Session()
+
+    def embed(self, *, model: str, text: str) -> list[float]:
+        payload: dict[str, Any] = {
+            "model": model,
+            "prompt": text,
+        }
+        data = self._request_json(method="POST", path="/api/embeddings", json=payload)
+
+        vec = data.get("embedding")
+        if not isinstance(vec, list) or not all(isinstance(x, (int, float)) for x in vec):
+            alt = data.get("data")
+            if isinstance(alt, list) and alt and isinstance(alt[0], dict) and isinstance(alt[0].get("embedding"), list):
+                vec = alt[0]["embedding"]
+            else:
+                raise ValueError("Unexpected Ollama embedding response schema")
+
+        return [float(x) for x in vec]
+
+    def describe(self, *, model: str, text: str, context: str) -> str:
+        """
+        Ask Ollama to generate a short description for `text`
+        """
+        prompt = self._build_description_prompt(text=text, context=context)
+
+        payload: dict[str, Any] = {"model": model, "prompt": prompt, "stream": False, "options": {"temperature": 0.1}}
+        data = self._request_json(method="POST", path="/api/generate", json=payload)
+
+        response_text = data.get("response")
+        if not isinstance(response_text, str):
+            raise ValueError("Unexpected Ollama generate response schema (missing 'response' string)")
+        return response_text.strip()
+
+    def pull_model_if_needed(self, *, model: str, timeout: float = 900.0) -> None:
+        if self._is_model_available(model_name=model):
+            logger.debug(f"Ollama model {model} was already available, skipping pull")
+            return
+
+        logger.info("Ollama model %s not found locally. Pulling it (this may take several minutes)...", model)
+        self.pull_model(model=model, timeout=timeout)
+        logger.info("Ollama model %s pulled successfully", model)
+
+    def pull_model(self, *, model: str, timeout: float = 900.0) -> None:
+        payload: dict[str, Any] = {"name": model}
+        self._request(method="POST", path="/api/pull", json=payload, timeout=timeout)
+
+    def is_healthy(self, *, timeout: float = 3.0) -> bool:
+        url = f"{self._base}/api/tags"
+        try:
+            r = self._session.get(url, headers=self._headers, timeout=timeout)
+            return 200 <= r.status_code < 300
+        except requests.RequestException:
+            return False
+
+    def wait_until_healthy(self, *, timeout: float = 60.0, poll_interval: float = 0.5) -> bool:
+        deadline = time.monotonic() + float(timeout)
+        while time.monotonic() < deadline:
+            if self.is_healthy(timeout=min(poll_interval, timeout)):
+                return True
+            time.sleep(poll_interval)
+        return self.is_healthy(timeout=min(poll_interval, timeout))
+
+    def _is_model_available(self, *, model_name, timeout: float = 5.0) -> bool:
+        url = f"{self._base}/api/tags"
+        try:
+            r = self._session.get(url, headers=self._headers, timeout=timeout)
+
+            if 200 <= r.status_code < 300:
+                models = r.json().get("models")
+                if models and isinstance(models, list):
+                    local_model = next((model for model in models if model.get("name") == model_name), None)
+                    return local_model is not None
+
+            return False
+        except requests.RequestException:
+            return False
+
+    def _request(
+        self,
+        *,
+        method: str,
+        path: str,
+        timeout: float | None = None,
+        **kwargs,
+    ) -> requests.Response:
+        url = f"{self._base}{path}"
+        try:
+            resp = self._session.request(
+                method,
+                url,
+                headers=self._headers,
+                timeout=timeout or self._timeout,
+                **kwargs,
+            )
+        except requests.Timeout as e:
+            raise OllamaTransientError(f"Ollama request to {path} timed out after {timeout}s") from e
+        except requests.RequestException as e:
+            raise OllamaTransientError(f"Ollama request to {path} failed: {e}") from e
+
+        try:
+            resp.raise_for_status()
+        except requests.HTTPError as e:
+            raise OllamaPermanentError(f"Ollama error {resp.status_code} for {path}: {resp.text}") from e
+
+        return resp
+
+    def _request_json(
+        self,
+        *,
+        method: str,
+        path: str,
+        timeout: float | None = None,
+        **kwargs,
+    ) -> dict[str, Any]:
+        resp = self._request(method=method, path=path, timeout=timeout, **kwargs)
+        try:
+            return resp.json()
+        except ValueError as e:
+            raise OllamaPermanentError(f"Invalid JSON from Ollama for {path}") from e
+
+    @staticmethod
+    def _build_description_prompt(text: str, context: str) -> str:
+        base = """
+        You are a helpful assistant.
+        
+        I will give you some TEXT and CONTEXT.
+        Write a concise, human-readable description of the TEXT suitable for displaying in a UI.
+        - 1-2 sentences
+        - Be factual and avoid speculation
+        - No markdown
+        - No preambles or labels, just the description itself.
+        - Your entire reply MUST be only the description itself. No extra commentary.
+        
+        CONTEXT:
+        {context}
+        
+        TEXT:
+        {text}
+        """
+
+        return textwrap.dedent(base).format(context=context, text=text).strip()

databao_context_engine/main.py

@@ -0,0 +1,19 @@
+import logging
+
+from databao_context_engine.cli.commands import dce
+
+logger = logging.getLogger(__name__)
+
+
+def main() -> None:
+    try:
+        dce(obj={})
+    except Exception as e:
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.exception(e)
+        else:
+            logger.error(str(e))
+
+
+if __name__ == "__main__":
+    main()

databao_context_engine/mcp/all_results_tool.py

@@ -0,0 +1,5 @@
+from databao_context_engine import DatabaoContextEngine
+
+
+def run_all_results_tool(databao_context_engine: DatabaoContextEngine, run_name: str | None) -> str:
+    return databao_context_engine.get_all_contexts_formatted(run_name=run_name)

databao_context_engine/mcp/mcp_runner.py

@@ -0,0 +1,16 @@
+import logging
+from pathlib import Path
+
+from databao_context_engine.mcp.mcp_server import McpServer, McpTransport
+
+logger = logging.getLogger(__name__)
+
+
+def run_mcp_server(
+    project_dir: Path,
+    run_name: str | None,
+    transport: McpTransport,
+    host: str | None = None,
+    port: int | None = None,
+) -> None:
+    McpServer(project_dir, run_name, host, port).run(transport)

databao_context_engine/mcp/mcp_server.py

@@ -0,0 +1,63 @@
+import logging
+from contextlib import asynccontextmanager
+from pathlib import Path
+from typing import Literal
+
+from mcp.server import FastMCP
+from mcp.types import ToolAnnotations
+
+from databao_context_engine import DatabaoContextEngine
+from databao_context_engine.mcp.all_results_tool import run_all_results_tool
+from databao_context_engine.mcp.retrieve_tool import run_retrieve_tool
+
+logger = logging.getLogger(__name__)
+
+McpTransport = Literal["stdio", "streamable-http"]
+
+
+@asynccontextmanager
+async def mcp_server_lifespan(server: FastMCP):
+    logger.info(f"Starting MCP server on {server.settings.host}:{server.settings.port}...")
+    yield
+    logger.info("Stopping MCP server")
+
+
+class McpServer:
+    def __init__(
+        self,
+        project_dir: Path,
+        run_name: str | None,
+        host: str | None = None,
+        port: int | None = None,
+    ):
+        self._databao_context_engine = DatabaoContextEngine(project_dir)
+        self._run_name = run_name
+
+        self._mcp_server = self._create_mcp_server(host, port)
+
+    def _create_mcp_server(self, host: str | None = None, port: int | None = None) -> FastMCP:
+        mcp = FastMCP(host=host or "127.0.0.1", port=port or 8000, lifespan=mcp_server_lifespan)
+
+        @mcp.tool(
+            description="Retrieve the contents of the all_results file",
+            annotations=ToolAnnotations(readOnlyHint=True, idempotentHint=True, openWorldHint=False),
+        )
+        def all_results_tool():
+            return run_all_results_tool(self._databao_context_engine, self._run_name)
+
+        @mcp.tool(
+            description="Retrieve the context built from various resources, including databases, dbt tools, plain and structured files, to retrieve relevant information",
+            annotations=ToolAnnotations(readOnlyHint=True, idempotentHint=True, openWorldHint=False),
+        )
+        def retrieve_tool(text: str, limit: int | None):
+            return run_retrieve_tool(
+                databao_context_engine=self._databao_context_engine,
+                run_name=self._run_name,
+                text=text,
+                limit=limit or 50,
+            )
+
+        return mcp
+
+    def run(self, transport: McpTransport):
+        self._mcp_server.run(transport=transport)

databao_context_engine/mcp/retrieve_tool.py

@@ -0,0 +1,22 @@
+import datetime
+
+from databao_context_engine import DatabaoContextEngine
+
+
+def run_retrieve_tool(
+    *, databao_context_engine: DatabaoContextEngine, run_name: str | None, text: str, limit: int | None = None
+) -> str:
+    """
+    Execute the retrieve flow for MCP and return the matching display texts
+    Adds the current date to the end
+    """
+
+    retrieve_results = databao_context_engine.search_context(
+        retrieve_text=text, run_name=run_name, limit=limit, export_to_file=False
+    )
+
+    display_results = [context_search_result.context_result for context_search_result in retrieve_results]
+
+    display_results.append(f"\nToday's date is {datetime.date.today()}")
+
+    return "\n".join(display_results)