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

databao_context_engine/introspection/property_extract.py

@@ -1,18 +1,33 @@
 import types
 from dataclasses import MISSING, fields, is_dataclass
-from typing import Annotated, Any, ForwardRef, Union, get_origin, get_type_hints
-
-from pydantic import BaseModel, _internal
+from typing import (
+    Annotated,
+    Any,
+    ForwardRef,
+    Iterable,
+    Mapping,
+    Union,
+    get_args,
+    get_origin,
+    get_type_hints,
+)
+
+from pydantic import BaseModel
 from pydantic_core import PydanticUndefinedType
 
-from databao_context_engine.pluginlib.config import ConfigPropertyAnnotation, ConfigPropertyDefinition
+from databao_context_engine.pluginlib.config import (
+    ConfigPropertyAnnotation,
+    ConfigPropertyDefinition,
+    ConfigSinglePropertyDefinition,
+    ConfigUnionPropertyDefinition,
+)
 
 
 def get_property_list_from_type(root_type: type) -> list[ConfigPropertyDefinition]:
-    return _get_property_list_from_type(parent_type=root_type, is_root_type=True)
+    return _get_property_list_from_type(parent_type=root_type)
 
 
-def _get_property_list_from_type(*, parent_type: type, is_root_type: bool) -> list[ConfigPropertyDefinition]:
+def _get_property_list_from_type(*, parent_type: type) -> list[ConfigPropertyDefinition]:
     if is_dataclass(parent_type):
         return _get_property_list_from_dataclass(parent_type=parent_type)
 
@@ -24,20 +39,15 @@ def _get_property_list_from_type(*, parent_type: type, is_root_type: bool) -> li
         # issubclass is raising a TypeError: issubclass() arg 1 must be a class
         pass
 
-    return _get_property_list_from_type_hints(parent_type=parent_type, is_root_type=is_root_type)
+    return _get_property_list_from_type_hints(parent_type=parent_type)
 
 
-def _get_property_list_from_type_hints(*, parent_type: type, is_root_type: bool) -> list[ConfigPropertyDefinition]:
+def _get_property_list_from_type_hints(*, parent_type: type) -> list[ConfigPropertyDefinition]:
     try:
         type_hints = get_type_hints(parent_type, include_extras=True)
-    except TypeError as e:
-        if is_root_type:
-            # Ignore root types that don't have type hints like dict or list
-            return []
-        else:
-            # If we're evaluating a nested property, we want to propagate the exception
-            # to let the parent property know that this type should be ignored
-            raise e
+    except TypeError:
+        # Return an empty list of properties for any type that is not an object (e.g: primitives like str or containers like dict, list, tuple, etc.
+        return []
 
     result = []
     for property_key, property_type in type_hints.items():
@@ -54,18 +64,14 @@ def _get_property_list_from_dataclass(parent_type: type) -> list[ConfigPropertyD
         raise ValueError(f"{parent_type} is not a dataclass")
 
     dataclass_fields = fields(parent_type)
+    type_hints = get_type_hints(parent_type, include_extras=True)
 
     result = []
     for field in dataclass_fields:
         has_field_default = field.default is not None and field.default != MISSING
 
-        if isinstance(field.type, str):
-            try:
-                property_type = _evaluate_type_string(field.type)
-            except Exception:
-                continue
-        else:
-            property_type = field.type
+        # Use the type hints if the field type wasn't resolved (aka. if it is a ForwardRef or a str)
+        property_type = type_hints[field.name] if isinstance(field.type, ForwardRef | str) else field.type
 
         property_for_field = _create_property(
             property_type=property_type,
@@ -84,6 +90,10 @@ def _get_property_list_from_pydantic_base_model(parent_type: type):
     if not issubclass(parent_type, BaseModel):
         raise ValueError(f"{parent_type} is not a Pydantic BaseModel")
 
+    if any(isinstance(field.annotation, ForwardRef) for field in parent_type.model_fields.values()):
+        # If any field's future type wasn't resolved yet, we rebuild the model to resolve them
+        parent_type.model_rebuild(force=True)
+
     pydantic_fields = parent_type.model_fields
     result = []
 
@@ -123,28 +133,58 @@ def _create_property(
     if annotation is not None and annotation.ignored_for_config_wizard:
         return None
 
-    actual_property_type = _read_actual_property_type(property_type)
+    actual_property_types = _read_actual_property_type(property_type)
 
-    try:
-        nested_properties = _get_property_list_from_type(parent_type=actual_property_type, is_root_type=False)
-    except TypeError:
+    required = annotation.required if annotation else is_property_required
+    secret = annotation.secret if annotation else False
+
+    if len(actual_property_types) > 1:
+        type_properties: dict[type, list[ConfigPropertyDefinition]] = {}
+
+        for union_type in actual_property_types:
+            nested_props = _get_property_list_from_type(parent_type=union_type)
+
+            type_properties[union_type] = nested_props
+
+        return ConfigUnionPropertyDefinition(
+            property_key=property_name,
+            types=actual_property_types,
+            type_properties=type_properties,
+        )
+
+    actual_property_type = actual_property_types[0]
+    nested_properties = _get_property_list_from_type(parent_type=actual_property_type)
+
+    if len(nested_properties) == 0 and _is_mapping_or_iterable(actual_property_type):
+        # Ignore Iterables and Mappings for which we didn't resolve nested properties
+        # (TypedDict is a Mapping but since we manage to resolve nested properties, it won't be ignored)
         return None
 
+    resolved_type = actual_property_type if not nested_properties else None
     default_value = compute_default_value(
         annotation=annotation,
         property_default=property_default,
         has_nested_properties=nested_properties is not None and len(nested_properties) > 0,
     )
 
-    return ConfigPropertyDefinition(
+    return ConfigSinglePropertyDefinition(
         property_key=property_name,
-        property_type=actual_property_type if not nested_properties else None,
-        required=annotation.required if annotation else is_property_required,
+        property_type=resolved_type,
+        required=required,
         default_value=default_value,
-        nested_properties=nested_properties if nested_properties else None,
+        nested_properties=nested_properties or None,
+        secret=secret,
     )
 
 
+def _is_mapping_or_iterable(property_type: type):
+    # For types like list[str], we need to get the origin (ie. list) to use in issubclass
+    origin = get_origin(property_type)
+
+    # We make sure to not return True for str, which is an Iterable
+    return property_type is not str and issubclass(origin if origin else property_type, (Mapping, Iterable))
+
+
 def _get_config_property_annotation(property_type) -> ConfigPropertyAnnotation | None:
     if get_origin(property_type) is Annotated:
         return next(
@@ -155,21 +195,15 @@ def _get_config_property_annotation(property_type) -> ConfigPropertyAnnotation |
     return None
 
 
-def _read_actual_property_type(property_type: type) -> type:
+def _read_actual_property_type(property_type: type) -> tuple[type, ...]:
     property_type_origin = get_origin(property_type)
 
     if property_type_origin is Annotated:
-        return property_type.__origin__  # type: ignore[attr-defined]
-    elif property_type_origin is Union or property_type_origin is types.UnionType:
-        type_args = property_type.__args__  # type: ignore[attr-defined]
-        if len(type_args) == 2 and type(None) in type_args:
-            # Uses the actual type T when the Union is "T | None" (or "None | T")
-            return next(arg for arg in type_args if arg is not None)
-        else:
-            # Ignoring Union types when it is not used as type | None as we wouldn't which type to pick
-            return type(None)
+        return _read_actual_property_type(property_type.__origin__)  # type: ignore[attr-defined]
+    if property_type_origin in (Union, types.UnionType):
+        return tuple(arg for arg in get_args(property_type) if arg is not type(None))
 
-    return property_type
+    return (property_type,)
 
 
 def compute_default_value(
@@ -185,18 +219,3 @@ def compute_default_value(
         return str(property_default)
 
     return None
-
-
-def _evaluate_type_string(property_type: str) -> type:
-    try:
-        # Using a pydantic internal function for this, to avoid having to implement type evaluation manually...
-        return _internal._typing_extra.eval_type(property_type)
-    except Exception as initial_error:
-        try:
-            # Try to convert it ourselves if Pydantic didn't work
-            return ForwardRef(property_type)._evaluate(  # type: ignore[return-value]
-                globalns=globals(), localns=locals(), recursive_guard=frozenset()
-            )
-        except Exception as e:
-            # Ignore if we didn't manage to convert the str to a type
-            raise e from initial_error

databao_context_engine/llm/__init__.py

@@ -0,0 +1,10 @@
+from databao_context_engine.llm.api import download_ollama_models_if_needed, install_ollama_if_needed
+from databao_context_engine.llm.errors import OllamaError, OllamaPermanentError, OllamaTransientError
+
+__all__ = [
+    "install_ollama_if_needed",
+    "download_ollama_models_if_needed",
+    "OllamaError",
+    "OllamaTransientError",
+    "OllamaPermanentError",
+]

databao_context_engine/llm/api.py

@@ -0,0 +1,57 @@
+from pathlib import Path
+
+from databao_context_engine.llm.errors import OllamaError
+from databao_context_engine.llm.factory import (
+    DEFAULT_DESCRIPTION_GENERATOR_MODEL,
+    DEFAULT_EMBED_MODEL_ID,
+    create_ollama_service,
+)
+from databao_context_engine.llm.install import resolve_ollama_bin
+
+
+def install_ollama_if_needed() -> Path:
+    """Install the Ollama CLI locally if needed.
+
+    This will look for any existing installation of Ollama on the system. If none is found, it will install it locally.
+
+    Here is the priority order of how it looks for an installed Ollama CLI binary:
+    1. Look at the path defined in the DCE_OLLAMA_BIN env var, if it is set
+    2. Look for `ollama` in the PATH
+    3. Look for a DCE-managed installation in the global DCE path
+
+    If Ollama is not found, it will get installed as a DCE-managed installation in the global DCE path.
+
+    Returns:
+        The path to the Ollama CLI executable.
+    """
+    return Path(resolve_ollama_bin())
+
+
+def download_ollama_models_if_needed(
+    *, download_embed_model: bool = True, download_description_generator_model: bool = False
+) -> None:
+    """Download the Ollama models required to run DCE if needed.
+
+    If the models were already downloaded, this method will do nothing.
+
+    If no Ollama CLI is found on the system, this method will install one as a DCE-managed installation in the global DCE path.
+
+    Args:
+        download_embed_model: Whether to download the embedding model.
+        download_description_generator_model: Whether to download the description generator model.
+
+    Raises:
+        OllamaError: If there is an error downloading one of the models.
+    """
+    ollama_service = create_ollama_service()
+
+    if download_embed_model:
+        try:
+            ollama_service.pull_model_if_needed(model=DEFAULT_EMBED_MODEL_ID)
+        except OllamaError as e:
+            raise e
+    if download_description_generator_model:
+        try:
+            ollama_service.pull_model_if_needed(model=DEFAULT_DESCRIPTION_GENERATOR_MODEL)
+        except OllamaError as e:
+            raise e

databao_context_engine/llm/descriptions/ollama.py

@@ -16,6 +16,4 @@ class OllamaDescriptionProvider(DescriptionProvider):
         return self._model_id
 
     def describe(self, text: str, context: str) -> str:
-        description = self._service.describe(model=self._model_id, text=text, context=context)
-
-        return description
+        return self._service.describe(model=self._model_id, text=text, context=context)

databao_context_engine/llm/errors.py

@@ -3,14 +3,8 @@ class OllamaError(Exception):
 
 
 class OllamaTransientError(OllamaError):
-    """
-    Errors that are likely temporary (network issues, timeouts, 5xx, etc.).
-    Typically worth retrying.
-    """
+    """Errors that are likely temporary (network issues, timeouts, 5xx, etc.), typically worth retrying."""
 
 
 class OllamaPermanentError(OllamaError):
-    """
-    Errors that are unlikely to succeed on retry without changing inputs
-    or configuration (4xx, bad response schema, etc.).
-    """
+    """Errors that are unlikely to succeed on retry without changing inputs or configuration (4xx, bad response schema, etc.)."""

databao_context_engine/llm/factory.py

@@ -5,6 +5,9 @@ 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
 
+DEFAULT_EMBED_MODEL_ID = "nomic-embed-text:v1.5"
+DEFAULT_DESCRIPTION_GENERATOR_MODEL = "llama3.2:1b"
+
 
 def _create_ollama_service_common(
     *,
@@ -39,7 +42,7 @@ def create_ollama_service(
 def create_ollama_embedding_provider(
     service: OllamaService,
     *,
-    model_id: str = "nomic-embed-text:v1.5",
+    model_id: str = DEFAULT_EMBED_MODEL_ID,
     dim: int = 768,
     pull_if_needed: bool = True,
 ) -> OllamaEmbeddingProvider:
@@ -52,7 +55,7 @@ def create_ollama_embedding_provider(
 def create_ollama_description_provider(
     service: OllamaService,
     *,
-    model_id: str = "llama3.2:1b",
+    model_id: str = DEFAULT_DESCRIPTION_GENERATOR_MODEL,
     pull_if_needed: bool = True,
 ):
     if pull_if_needed:

databao_context_engine/llm/install.py

@@ -49,14 +49,15 @@ ARTIFACTS: dict[str, ArtifactInfo] = {
 
 
 def resolve_ollama_bin() -> str:
-    """
-    Decide which `ollama` binary to use, in this order:
+    """Decide which `ollama` binary to use.
 
+    Here is the priority 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
+    Returns:
+        The full path to the binary
     """
     override = os.environ.get("DCE_OLLAMA_BIN")
     if override:
@@ -76,9 +77,7 @@ def resolve_ollama_bin() -> str:
 
 
 def _detect_platform() -> str:
-    """
-    Return one of: 'darwin', 'linux-amd64', 'linux-arm64', 'windows-amd64', 'windows-arm64'.
-    """
+    """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()
 
@@ -96,27 +95,25 @@ def _detect_platform() -> str:
     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.
-    """
+def _download_artifact_to_temp(artifact_version: str, artifact_name: str) -> Path:
+    """Download to a temporary file and return its path."""
     import urllib.request
 
+    artifact_url = f"https://github.com/ollama/ollama/releases/download/{artifact_version}/{artifact_name}"
+
     tmp_dir = Path(tempfile.mkdtemp(prefix="ollama-download-"))
-    file_name = url.rsplit("/", 1)[-1]
+    file_name = artifact_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:
+    logger.info("Downloading %s to %s", artifact_url, dest)
+    with urllib.request.urlopen(artifact_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
-    """
+    """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""):
@@ -127,36 +124,33 @@ def _verify_sha256(path: Path, expected_hex: str) -> None:
 
 
 def _extract_archive(archive: Path, target_dir: Path) -> None:
-    """
-    Extract archive into target_dir.
-    """
+    """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)
+            # There is no built-in protection against zip bombs in ZipFile.
+            # However, we previously checked the sha256 of the downloaded archive and we trust the origin (GitHub repo of Ollama)
+            zf.extractall(target_dir)  # noqa: S202
     elif name.endswith(".tgz") or name.endswith(".tar.gz"):
         with tarfile.open(archive, "r:gz") as tf:
-            tf.extractall(target_dir)
+            tf.extractall(target_dir, filter="data")
     else:
         raise RuntimeError(f"Unsupported archive format: {archive}")
 
 
 def _ensure_executable(path: Path) -> None:
-    """
-    Mark path as executable
-    """
+    """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
+        logger.debug("Failed to mark %s as executable", path, exc_info=True, stack_info=True)
 
 
 def install_ollama_to(target: Path) -> None:
-    """
-    Ensure an Ollama binary exists.
+    """Ensure an Ollama binary exist.
 
     If it doesn't exist, this will:
     - detect OS
@@ -164,6 +158,9 @@ def install_ollama_to(target: Path) -> None:
     - verify its SHA-256 checksum
     - extract into the installation directory
     - make the binary executable
+
+    Raises:
+        RuntimeError: If the user's platform is not supported
     """
     target = target.expanduser()
     if target.parent.name == "bin":
@@ -179,8 +176,7 @@ def install_ollama_to(target: Path) -> None:
     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)
+    archive_path = _download_artifact_to_temp(DEFAULT_VERSION, artifact.name)
 
     try:
         _verify_sha256(archive_path, artifact.sha256)
@@ -224,4 +220,4 @@ def install_ollama_to(target: Path) -> None:
         try:
             archive_path.unlink(missing_ok=True)
         except Exception:
-            pass
+            logger.debug("Failed to remove temporary archive %s", archive_path, exc_info=True, stack_info=True)

databao_context_engine/llm/runtime.py

@@ -26,7 +26,7 @@ class OllamaRuntime:
 
         stdout = subprocess.DEVNULL
 
-        proc = subprocess.Popen(
+        return subprocess.Popen(  # noqa: S603 We're always running Ollama
             cmd,
             cwd=str(self._config.work_dir) if self._config.work_dir else None,
             env=env,
@@ -36,8 +36,6 @@ class OllamaRuntime:
             close_fds=os.name != "nt",
         )
 
-        return proc
-
     def start_and_await(
         self,
         *,
@@ -62,11 +60,11 @@ class OllamaRuntime:
             try:
                 proc.terminate()
             except Exception:
-                pass
+                logger.debug("Failed to terminate Ollama server", exc_info=True, stack_info=True)
             try:
                 proc.kill()
             except Exception:
-                pass
+                logger.debug("Failed to kill Ollama server", exc_info=True, stack_info=True)
 
         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

@@ -36,9 +36,7 @@ class OllamaService:
         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`
-        """
+        """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}}

databao_context_engine/mcp/mcp_runner.py

@@ -2,15 +2,17 @@ import logging
 from pathlib import Path
 
 from databao_context_engine.mcp.mcp_server import McpServer, McpTransport
+from databao_context_engine.project.layout import ensure_project_dir
 
 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)
+    ensure_project_dir(project_dir=project_dir)
+
+    McpServer(project_dir, host, port).run(transport)

databao_context_engine/mcp/mcp_server.py

@@ -1,5 +1,6 @@
 import logging
 from contextlib import asynccontextmanager
+from datetime import date
 from pathlib import Path
 from typing import Literal
 
@@ -7,8 +8,6 @@ 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__)
 
@@ -26,12 +25,10 @@ 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)
 
@@ -43,19 +40,20 @@ class McpServer:
             annotations=ToolAnnotations(readOnlyHint=True, idempotentHint=True, openWorldHint=False),
         )
         def all_results_tool():
-            return run_all_results_tool(self._databao_context_engine, self._run_name)
+            return self._databao_context_engine.get_all_contexts_formatted()
 
         @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,
-            )
+            retrieve_results = self._databao_context_engine.search_context(retrieve_text=text, limit=limit)
+
+            display_results = [context_search_result.context_result for context_search_result in retrieve_results]
+
+            display_results.append(f"\nToday's date is {date.today()}")
+
+            return "\n".join(display_results)
 
         return mcp