prometheusnn 0.1.0__py3-none-any.whl

prometheus/__init__.py

@@ -0,0 +1 @@
+__version__ = "1.0.0"

prometheus/core/config.py

@@ -0,0 +1,51 @@
+"""Configuration and directory management."""
+
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+
+@dataclass
+class PrometheusConfig:
+    """Global configuration for Prometheus."""
+
+    base_dir: Path = field(
+        default_factory=lambda: Path(os.environ.get("PROMETHEUS_HOME", Path.home() / ".prometheus"))
+    )
+    default_embedding_model: str = field(
+        default_factory=lambda: os.environ.get("PROMETHEUS_EMBEDDING_MODEL", "all-MiniLM-L6-v2")
+    )
+
+    def __post_init__(self):
+        self.base_dir = Path(self.base_dir)
+
+    @property
+    def models_dir(self) -> Path:
+        return self.base_dir / "models"
+
+    @property
+    def feedback_dir(self) -> Path:
+        return self.base_dir / "feedback"
+
+    @property
+    def logs_dir(self) -> Path:
+        return self.base_dir / "logs"
+
+    def ensure_dirs(self):
+        """Create all required directories."""
+        self.models_dir.mkdir(parents=True, exist_ok=True)
+        self.feedback_dir.mkdir(parents=True, exist_ok=True)
+        self.logs_dir.mkdir(parents=True, exist_ok=True)
+
+    def model_path(self, model_name: str) -> Path:
+        return self.models_dir / model_name
+
+    def list_models(self) -> list[str]:
+        """List all saved model names."""
+        if not self.models_dir.exists():
+            return []
+        return sorted(
+            d.name
+            for d in self.models_dir.iterdir()
+            if d.is_dir() and (d / "manifest.json").exists()
+        )

prometheus/core/embedding.py

@@ -0,0 +1,61 @@
+"""Thread-safe SentenceTransformer LRU cache and encoding utilities."""
+
+from __future__ import annotations
+
+import threading
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+from prometheus.core.lru import LRUCache
+
+if TYPE_CHECKING:
+    from sentence_transformers import SentenceTransformer
+
+_MODEL_CACHE_SIZE = 4
+
+_lock = threading.Lock()
+_models: LRUCache[str, SentenceTransformer] = LRUCache(maxsize=_MODEL_CACHE_SIZE)
+
+
+def get_embedding_model(model_name: str) -> SentenceTransformer:
+    """Get or load a SentenceTransformer model (LRU cached, thread-safe)."""
+    cached = _models.get(model_name)
+    if cached is not None:
+        return cached
+    with _lock:
+        cached = _models.get(model_name)
+        if cached is not None:
+            return cached
+        from sentence_transformers import SentenceTransformer
+
+        model = SentenceTransformer(model_name)
+        _models.put(model_name, model)
+        return model
+
+
+def encode_texts(
+    texts: list[str],
+    model_name: str,
+    batch_size: int = 64,
+    show_progress: bool = True,
+) -> np.ndarray:
+    """Encode texts to embeddings using a sentence-transformer model.
+
+    Returns:
+        (n_texts, embedding_dim) float32 array
+    """
+    model = get_embedding_model(model_name)
+    embeddings = model.encode(
+        texts,
+        batch_size=batch_size,
+        show_progress_bar=show_progress,
+        convert_to_numpy=True,
+    )
+    return embeddings.astype(np.float32)
+
+
+def get_embedding_dim(model_name: str) -> int:
+    """Get the embedding dimension for a model."""
+    model = get_embedding_model(model_name)
+    return model.get_sentence_embedding_dimension()

prometheus/core/logging.py

@@ -0,0 +1,24 @@
+"""Rotating file logger for PrometheusNN."""
+
+import logging
+from logging.handlers import RotatingFileHandler
+from pathlib import Path
+
+
+def setup_logging(logs_dir: Path, level: int = logging.INFO) -> None:
+    """Configure the prometheus logger with a rotating file handler."""
+    logs_dir.mkdir(parents=True, exist_ok=True)
+    handler = RotatingFileHandler(
+        logs_dir / "prometheus.log",
+        maxBytes=10 * 1024 * 1024,  # 10 MB
+        backupCount=3,
+    )
+    handler.setFormatter(
+        logging.Formatter(
+            "%(asctime)s %(name)s %(levelname)s %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+    )
+    root = logging.getLogger("prometheus")
+    root.setLevel(level)
+    root.addHandler(handler)

prometheus/core/lru.py

@@ -0,0 +1,47 @@
+"""Tiny LRU cache with explicit key eviction.
+
+Used for bounded in-memory caches where stdlib's ``functools.lru_cache``
+isn't a fit (no per-key invalidation, no in-place size tuning).
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from typing import Generic, TypeVar
+
+K = TypeVar("K")
+V = TypeVar("V")
+
+
+class LRUCache(Generic[K, V]):
+    """Fixed-capacity LRU cache. Not thread-safe; suitable for the MCP server's
+    single-threaded request loop.
+    """
+
+    def __init__(self, maxsize: int):
+        if maxsize < 1:
+            raise ValueError("maxsize must be >= 1")
+        self.maxsize = maxsize
+        self._data: OrderedDict[K, V] = OrderedDict()
+
+    def get(self, key: K) -> V | None:
+        if key in self._data:
+            self._data.move_to_end(key)
+            return self._data[key]
+        return None
+
+    def put(self, key: K, value: V) -> None:
+        if key in self._data:
+            self._data.move_to_end(key)
+        self._data[key] = value
+        while len(self._data) > self.maxsize:
+            self._data.popitem(last=False)
+
+    def pop(self, key: K, default: V | None = None) -> V | None:
+        return self._data.pop(key, default)
+
+    def __contains__(self, key: object) -> bool:
+        return key in self._data
+
+    def __len__(self) -> int:
+        return len(self._data)

prometheus/core/mcp_safety.py

@@ -0,0 +1,49 @@
+"""Decorator that wraps MCP tools with uniform error handling."""
+
+from __future__ import annotations
+
+import functools
+import logging
+from collections.abc import Callable
+from typing import Any
+
+from prometheus.core.responses import internal_error, validation_error
+
+logger = logging.getLogger(__name__)
+
+
+class ToolValidationError(Exception):
+    """Raise from inside a tool to return a validation_error response.
+
+    The mcp_tool_safe decorator catches this and converts it to the proper
+    response shape. Use this for bad-input failures (missing file, bad column
+    name, etc.) instead of returning validation_error() directly.
+    """
+
+    def __init__(self, message: str, field: str | None = None):
+        self.message = message
+        self.field = field
+        super().__init__(message)
+
+
+def mcp_tool_safe(func: Callable[..., Any]) -> Callable[..., Any]:
+    """Wrap an MCP tool with logging + uniform exception handling.
+
+    Catches:
+    - ToolValidationError → validation_error response (with field)
+    - Any other Exception → internal_error response, logged with traceback
+
+    Apply *inside* @mcp.tool() so the registered name and signature stay correct.
+    """
+
+    @functools.wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> dict:
+        try:
+            return func(*args, **kwargs)
+        except ToolValidationError as e:
+            return validation_error(e.message, field=e.field)
+        except Exception as e:
+            logger.error("%s failed: %s", func.__name__, e, exc_info=True)
+            return internal_error(str(e))
+
+    return wrapper

prometheus/core/responses.py

@@ -0,0 +1,61 @@
+"""Typed MCP response envelopes.
+
+Three discriminated variants — ``status`` field tells the client which one:
+
+- ``SuccessResponse`` (status="success") — wraps tool output payload
+- ``ErrorResponse`` (status="error") — internal/unexpected error
+- ``ValidationErrorResponse`` (status="validation_error") — bad input
+"""
+
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import BaseModel, ConfigDict
+
+
+class SuccessResponse(BaseModel):
+    """Successful tool response. Extra fields from the tool's payload are merged in."""
+
+    model_config = ConfigDict(extra="allow")
+
+    status: Literal["success"] = "success"
+    message: str | None = None
+
+
+class ErrorResponse(BaseModel):
+    """Internal error response — used when a tool raises unexpectedly."""
+
+    status: Literal["error"] = "error"
+    error: str
+
+
+class ValidationErrorResponse(BaseModel):
+    """Validation error response — used when tool inputs fail to validate."""
+
+    status: Literal["validation_error"] = "validation_error"
+    error: str
+    field: str | None = None
+
+
+def success_response(
+    data: dict[str, Any] | BaseModel | None = None,
+    message: str | None = None,
+) -> dict[str, Any]:
+    """Build a success response. Pass a dict or Pydantic model for payload, or None for message-only ack."""
+    payload: dict[str, Any] = {}
+    if isinstance(data, BaseModel):
+        payload = data.model_dump(exclude_none=True)
+    elif isinstance(data, dict):
+        payload = data
+    return SuccessResponse(message=message, **payload).model_dump(exclude_none=True)
+
+
+def internal_error(message: str) -> dict[str, Any]:
+    """Build an internal error response. Stack traces are logged server-side, not returned."""
+    return ErrorResponse(error=message).model_dump(exclude_none=True)
+
+
+def validation_error(message: str, field: str | None = None) -> dict[str, Any]:
+    """Build a validation error response."""
+    return ValidationErrorResponse(error=message, field=field).model_dump(exclude_none=True)

prometheus/data/loader.py

@@ -0,0 +1,267 @@
+"""CSV/JSON data loading, taxonomy extraction, embedding, and splitting."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from sklearn.model_selection import train_test_split
+
+from prometheus.core.embedding import encode_texts
+from prometheus.data.splits import DataSplit
+from prometheus.data.taxonomy import TaxonomyInfo
+
+
+def load_dataframe(file_path: str | Path) -> pd.DataFrame:
+    """Load a CSV or JSON file into a DataFrame."""
+    path = Path(file_path)
+    if path.suffix.lower() == ".json":
+        return pd.read_json(path)
+    return pd.read_csv(path)
+
+
+def validate_columns(
+    df: pd.DataFrame,
+    text_column: str,
+    level_columns: list[str],
+) -> list[str]:
+    """Validate that required columns exist. Returns list of warnings."""
+    warnings = []
+    missing = []
+
+    if text_column not in df.columns:
+        missing.append(text_column)
+    for col in level_columns:
+        if col not in df.columns:
+            missing.append(col)
+
+    if missing:
+        raise ValueError(f"Missing columns: {missing}. Available: {list(df.columns)}")
+
+    # Check for missing values
+    all_cols = [text_column] + level_columns
+    null_counts = df[all_cols].isnull().sum()
+    for col, count in null_counts.items():
+        if count > 0:
+            warnings.append(f"Column '{col}' has {count} missing values")
+
+    return warnings
+
+
+def check_hierarchy_consistency(
+    df: pd.DataFrame,
+    level_columns: list[str],
+) -> list[str]:
+    """Check that child labels map to consistent parents.
+
+    Returns list of warnings for inconsistencies.
+    """
+    warnings = []
+    for i in range(1, len(level_columns)):
+        child_col = level_columns[i]
+        parent_col = level_columns[i - 1]
+
+        # Group child values by parent
+        child_parents = df.groupby(child_col)[parent_col].nunique()
+        multi_parent = child_parents[child_parents > 1]
+
+        if len(multi_parent) > 0:
+            for child_val, n_parents in multi_parent.items():
+                parents = df[df[child_col] == child_val][parent_col].unique()
+                warnings.append(
+                    f"'{child_val}' at level '{child_col}' maps to {n_parents} "
+                    f"parents at '{parent_col}': {list(parents)}"
+                )
+    return warnings
+
+
+def extract_taxonomy(
+    df: pd.DataFrame,
+    model_name: str,
+    text_column: str,
+    level_columns: list[str],
+    embedding_model: str,
+) -> TaxonomyInfo:
+    """Extract taxonomy structure from a DataFrame."""
+    level_classes = []
+    for col in level_columns:
+        classes = sorted(df[col].dropna().unique().tolist())
+        level_classes.append(classes)
+
+    return TaxonomyInfo.build(
+        name=model_name,
+        text_column=text_column,
+        level_columns=level_columns,
+        level_classes=level_classes,
+        embedding_model=embedding_model,
+        total_samples=len(df),
+    )
+
+
+def prepare_data(
+    file_path: str | Path,
+    text_column: str,
+    level_columns: list[str],
+    model_name: str,
+    embedding_model: str,
+    train_ratio: float = 0.80,
+    val_ratio: float = 0.10,
+    test_ratio: float = 0.10,
+    random_state: int = 42,
+    show_progress: bool = True,
+) -> tuple[TaxonomyInfo, DataSplit, DataSplit, DataSplit, list[str]]:
+    """Full data preparation pipeline: load, validate, embed, split.
+
+    Returns:
+        (taxonomy, train_split, val_split, test_split, warnings)
+    """
+    df = load_dataframe(file_path)
+    warnings = validate_columns(df, text_column, level_columns)
+
+    # Drop rows with missing values in required columns
+    all_cols = [text_column] + level_columns
+    clean_df = df.dropna(subset=all_cols).reset_index(drop=True)
+    if len(clean_df) < len(df):
+        dropped = len(df) - len(clean_df)
+        warnings.append(f"Dropped {dropped} rows with missing values")
+
+    warnings.extend(check_hierarchy_consistency(clean_df, level_columns))
+
+    # Extract taxonomy
+    taxonomy = extract_taxonomy(clean_df, model_name, text_column, level_columns, embedding_model)
+
+    # Build parent→valid_children map for hierarchical masking
+    if len(level_columns) > 1:
+        pcm: dict[tuple[int, int], list[int]] = {}
+        for depth in range(1, len(level_columns)):
+            parent_col = level_columns[depth - 1]
+            child_col = level_columns[depth]
+            parent_level = taxonomy.get_level(depth - 1)
+            child_level = taxonomy.get_level(depth)
+            for parent_cls in parent_level.classes:
+                child_vals = clean_df.loc[clean_df[parent_col] == parent_cls, child_col].unique()
+                child_indices = sorted(child_level.encode(c) for c in child_vals)
+                parent_idx = parent_level.encode(parent_cls)
+                pcm[(depth, parent_idx)] = child_indices
+        taxonomy.parent_child_map = pcm
+
+    # Encode texts
+    texts = clean_df[text_column].tolist()
+    embeddings = encode_texts(texts, embedding_model, show_progress=show_progress)
+
+    # Encode labels by depth
+    labels_by_depth = {}
+    for depth, level in enumerate(taxonomy.levels):
+        labels_by_depth[depth] = np.array([level.encode(val) for val in clean_df[level.name]])
+
+    # Split — handle rare classes that can't be stratified
+    indices = np.arange(len(clean_df))
+    root_labels = labels_by_depth[0]
+
+    # Find classes too rare to stratify (need ≥2 samples in val+test,
+    # which means ≥ ceil(2 / val_test_ratio) total samples)
+    val_test_ratio = val_ratio + test_ratio
+    unique_labels, label_counts = np.unique(root_labels, return_counts=True)
+    min_samples_for_stratify = max(2, int(np.ceil(2 / val_test_ratio)))
+    rare_labels = set(unique_labels[label_counts < min_samples_for_stratify])
+
+    if rare_labels:
+        # Rare classes go straight into training; stratify the rest
+        rare_mask = np.isin(root_labels, list(rare_labels))
+        rare_idx = indices[rare_mask]
+        stratifiable_idx = indices[~rare_mask]
+        stratifiable_labels = root_labels[stratifiable_idx]
+
+        strat_train_idx, val_test_idx = train_test_split(
+            stratifiable_idx,
+            test_size=val_test_ratio,
+            random_state=random_state,
+            stratify=stratifiable_labels,
+        )
+        train_idx = np.concatenate([strat_train_idx, rare_idx])
+
+        test_proportion = test_ratio / val_test_ratio
+        val_idx, test_idx = train_test_split(
+            val_test_idx,
+            test_size=test_proportion,
+            random_state=random_state,
+            stratify=root_labels[val_test_idx],
+        )
+
+        rare_class_names = [taxonomy.levels[0].decode(int(lbl)) for lbl in rare_labels]
+        warnings.append(
+            f"Rare root classes forced into training set (too few samples to stratify): "
+            f"{rare_class_names}"
+        )
+    else:
+        # Standard stratified split
+        train_idx, val_test_idx = train_test_split(
+            indices,
+            test_size=val_test_ratio,
+            random_state=random_state,
+            stratify=root_labels,
+        )
+
+        test_proportion = test_ratio / val_test_ratio
+        val_idx, test_idx = train_test_split(
+            val_test_idx,
+            test_size=test_proportion,
+            random_state=random_state,
+            stratify=root_labels[val_test_idx],
+        )
+
+    def make_split(idx: np.ndarray) -> DataSplit:
+        return DataSplit(
+            embeddings=embeddings[idx],
+            labels_by_depth={d: lbls[idx] for d, lbls in labels_by_depth.items()},
+            texts=[texts[i] for i in idx],
+        )
+
+    return (
+        taxonomy,
+        make_split(train_idx),
+        make_split(val_idx),
+        make_split(test_idx),
+        warnings,
+    )
+
+
+def analyze_data(
+    file_path: str | Path,
+    text_column: str,
+    level_columns: list[str],
+) -> dict:
+    """Analyze training data without embedding. Returns summary stats."""
+    df = load_dataframe(file_path)
+    warnings = validate_columns(df, text_column, level_columns)
+
+    all_cols = [text_column] + level_columns
+    clean_df = df.dropna(subset=all_cols).reset_index(drop=True)
+    if len(clean_df) < len(df):
+        warnings.append(f"{len(df) - len(clean_df)} rows have missing values")
+
+    warnings.extend(check_hierarchy_consistency(clean_df, level_columns))
+
+    level_stats = []
+    for col in level_columns:
+        counts = clean_df[col].value_counts()
+        level_stats.append(
+            {
+                "column": col,
+                "num_classes": len(counts),
+                "class_counts": counts.to_dict(),
+                "min_count": int(counts.min()),
+                "max_count": int(counts.max()),
+            }
+        )
+
+    return {
+        "total_rows": len(df),
+        "clean_rows": len(clean_df),
+        "text_column": text_column,
+        "num_levels": len(level_columns),
+        "level_stats": level_stats,
+        "warnings": warnings,
+        "text_sample": clean_df[text_column].head(5).tolist(),
+    }

prometheus/data/splits.py

@@ -0,0 +1,62 @@
+"""Data splits and batch iteration for N-level training."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+
+@dataclass
+class DataSplit:
+    """A split of data for training/validation/calibration."""
+
+    embeddings: np.ndarray  # (N, embed_dim)
+    labels_by_depth: dict[int, np.ndarray]  # {0: (N,), 1: (N,), ...}
+    texts: list[str] | None = None
+
+    def __len__(self) -> int:
+        return len(self.embeddings)
+
+    @property
+    def num_levels(self) -> int:
+        return len(self.labels_by_depth)
+
+    def get_labels(self, depth: int) -> np.ndarray:
+        return self.labels_by_depth[depth]
+
+    @property
+    def leaf_labels(self) -> np.ndarray:
+        max_depth = max(self.labels_by_depth.keys())
+        return self.labels_by_depth[max_depth]
+
+
+class BatchIterator:
+    """Yields batches of (embeddings, label dicts) for training."""
+
+    def __init__(
+        self,
+        data: DataSplit,
+        batch_size: int = 64,
+        shuffle: bool = True,
+        random_state: int | None = None,
+    ):
+        self.data = data
+        self.batch_size = batch_size
+        self.shuffle = shuffle
+        self.rng = np.random.default_rng(random_state)
+
+    def __len__(self) -> int:
+        return (len(self.data) + self.batch_size - 1) // self.batch_size
+
+    def __iter__(self):
+        indices = np.arange(len(self.data))
+        if self.shuffle:
+            self.rng.shuffle(indices)
+
+        for start in range(0, len(indices), self.batch_size):
+            batch_idx = indices[start : start + self.batch_size]
+            batch = {"embeddings": self.data.embeddings[batch_idx]}
+            for depth, labels in self.data.labels_by_depth.items():
+                batch[depth] = labels[batch_idx]
+            yield batch