classifyre-cli 0.4.9__tar.gz → 0.4.10__tar.gz

{classifyre_cli-0.4.9 → classifyre_cli-0.4.10}/.turbo/turbo-build.log

@@ -1,3 +1,3 @@
 $ uv sync
-Resolved 256 packages in 201ms
+Resolved 265 packages in 156ms
 Checked 50 packages in 1ms

{classifyre_cli-0.4.9 → classifyre_cli-0.4.10}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: classifyre-cli
-Version: 0.4.9
+Version: 0.4.10
 Summary: Classifyre CLI — scan and classify unstructured data sources
 License: MIT
 Keywords: data,ingestion,metadata,pii,secrets,unstructured

{classifyre_cli-0.4.9 → classifyre_cli-0.4.10}/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@classifyre/cli",
-  "version": "0.4.9",
+  "version": "0.4.10",
   "private": true,
   "scripts": {
     "build": "uv sync",

{classifyre_cli-0.4.9 → classifyre_cli-0.4.10}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "classifyre-cli"
-version = "0.4.9"
+version = "0.4.10"
 description = "Classifyre CLI — scan and classify unstructured data sources"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -91,6 +91,9 @@ custom = [
 regex = [
     "google-re2>=1.1",
 ]
+llm = [
+    "litellm>=1.86.2",
+]
 detectors = [
     { include-group = "file-processing" },
     { include-group = "privacy" },
@@ -101,6 +104,7 @@ detectors = [
     { include-group = "classification" },
     { include-group = "custom" },
     { include-group = "regex" },
+    { include-group = "llm" },
 ]
 file-processing = [
     "filetype>=1.2.0",
@@ -264,6 +268,8 @@ module = [
     "datasets",
     "setfit.*",
     "setfit",
+    "litellm.*",
+    "litellm",
     "sklearn.*",
     "sklearn",
     "numpy",

classifyre_cli-0.4.10/src/detectors/custom/runners/_llm.py

@@ -0,0 +1,230 @@
+"""AI/LLM pipeline runner — prompt-driven classification and field extraction."""
+
+from __future__ import annotations
+
+import json
+import logging
+import os
+from datetime import UTC, datetime
+from typing import Any
+
+# Quiet litellm's import-time provider preload warnings (bedrock/sagemaker need
+# botocore, which we don't install) before the library is ever imported.
+os.environ.setdefault("LITELLM_LOG", "ERROR")
+
+from ....models.generated_detectors import LLMPipelineSchema, Severity
+from ....models.generated_single_asset_scan_results import (
+    DetectionResult,
+    DetectorType,
+)
+from ...dependencies import require_module
+from ._base import _TEXT_CONTENT_TYPES, BaseRunner, _resolve_pipeline_severity
+
+logger = logging.getLogger(__name__)
+
+# Map the stored AI provider type onto the litellm model-string convention.
+_PROVIDER_PREFIX: dict[str, str] = {
+    "CLAUDE": "anthropic",
+    "GEMINI": "gemini",
+    "OPENAI_COMPATIBLE": "openai",
+}
+
+
+class LLMRunner(BaseRunner):
+    """AI detector — sends content to a configured LLM provider for classification + extraction."""
+
+    def __init__(
+        self, schema: LLMPipelineSchema, detector_key: str = "", detector_name: str = ""
+    ) -> None:
+        self._schema = schema
+        self._detector_key = detector_key
+        self._detector_name = detector_name
+
+        runtime = schema.provider_runtime
+        if runtime is None:
+            raise ValueError(
+                f"AI detector '{detector_key}' is missing provider_runtime — the API must "
+                "inject resolved provider credentials before dispatch."
+            )
+        self._runtime = runtime
+        self._litellm = require_module("litellm", "llm", ["llm"])
+        # Let litellm silently drop params an endpoint doesn't support (e.g.
+        # response_format / temperature on some OpenAI-compatible gateways)
+        # instead of raising. Keep its own logging quiet.
+        self._litellm.drop_params = True
+        self._litellm.suppress_debug_info = True
+        logging.getLogger("LiteLLM").setLevel(logging.ERROR)
+
+    def run(self, text: str) -> None:  # type: ignore[override]  # pragma: no cover
+        raise NotImplementedError("LLMRunner uses detect() directly")
+
+    def detect(self, content: str | bytes, content_type: str) -> list[DetectionResult]:
+        if isinstance(content, bytes):
+            return []
+        if content_type not in _TEXT_CONTENT_TYPES:
+            return []
+        text = content.strip()
+        if not text:
+            return []
+
+        schema = self._schema
+        content_limit = schema.content_limit or 8000
+        snippet = text[:content_limit]
+
+        messages = [
+            {"role": "system", "content": self._build_system_prompt()},
+            {"role": "user", "content": snippet},
+        ]
+
+        try:
+            response = self._litellm.completion(
+                model=self._model_string(),
+                api_key=self._runtime.api_key,
+                api_base=self._runtime.base_url or None,
+                temperature=schema.temperature if schema.temperature is not None else 0.0,
+                max_tokens=self._max_tokens(),
+                messages=messages,
+                response_format={"type": "json_object"},
+            )
+            raw = response.choices[0].message.content or "{}"
+            parsed = self._parse_json(raw)
+        except Exception as exc:
+            logger.error(
+                "llm detector error (detector=%s, model=%s): %s",
+                self._detector_key,
+                self._runtime.model,
+                exc,
+                exc_info=True,
+            )
+            return []
+
+        return self._results_from_payload(snippet, parsed)
+
+    def get_supported_content_types(self) -> list[str]:
+        return list(_TEXT_CONTENT_TYPES)
+
+    # ── Internals ────────────────────────────────────────────────────────────
+
+    def _max_tokens(self) -> int | None:
+        # `max_tokens` is generated as a RootModel[int] wrapper, so unwrap `.root`
+        # before handing it to litellm — passing the model object serialises to an
+        # invalid request body and fails the whole completion.
+        raw = self._schema.max_tokens
+        if raw is None:
+            return None
+        return getattr(raw, "root", raw)
+
+    def _model_string(self) -> str:
+        prefix = _PROVIDER_PREFIX.get(self._runtime.provider.value, "openai")
+        return f"{prefix}/{self._runtime.model}"
+
+    def _build_system_prompt(self) -> str:
+        schema = self._schema
+        parts: list[str] = [schema.system_prompt.strip()]
+
+        labels = schema.labels or []
+        if labels:
+            label_lines = "\n".join(
+                f"- {lbl.name}: {lbl.description}" if lbl.description else f"- {lbl.name}"
+                for lbl in labels
+            )
+            parts.append(
+                "Classify the content using these labels:\n"
+                + label_lines
+                + (
+                    "\nMultiple labels may apply."
+                    if schema.multi_label
+                    else "\nChoose the single best label."
+                )
+            )
+
+        fields = schema.output_fields or []
+        if fields:
+            field_lines = "\n".join(
+                f"- {f.name} ({f.type.value if f.type else 'string'}): {f.description}"
+                if f.description
+                else f"- {f.name} ({f.type.value if f.type else 'string'})"
+                for f in fields
+            )
+            parts.append("Also extract these fields:\n" + field_lines)
+
+        parts.append(
+            "Respond with a JSON object of the form: "
+            '{"labels": [{"name": "<label>", "confidence": <0-1>, '
+            '"matched_content": "<relevant snippet>"}], "fields": {<field name>: <value>}}. '
+            "Use only the labels listed above. Return an empty labels array when none apply."
+        )
+
+        if schema.response_example:
+            parts.append("Example response:\n" + schema.response_example.strip())
+
+        return "\n\n".join(parts)
+
+    @staticmethod
+    def _parse_json(raw: str) -> dict[str, Any]:
+        try:
+            parsed = json.loads(raw)
+        except json.JSONDecodeError:
+            start = raw.find("{")
+            end = raw.rfind("}")
+            if start == -1 or end == -1 or end <= start:
+                return {}
+            try:
+                parsed = json.loads(raw[start : end + 1])
+            except json.JSONDecodeError:
+                return {}
+        return parsed if isinstance(parsed, dict) else {}
+
+    def _results_from_payload(self, snippet: str, payload: dict[str, Any]) -> list[DetectionResult]:
+        schema = self._schema
+        threshold = schema.confidence_threshold if schema.confidence_threshold is not None else 0.5
+        default_severity = schema.severity or Severity.info
+        extracted = self._coerce_fields(payload.get("fields"))
+
+        raw_labels = payload.get("labels")
+        label_entries: list[dict[str, Any]] = (
+            [lbl for lbl in raw_labels if isinstance(lbl, dict)]
+            if isinstance(raw_labels, list)
+            else []
+        )
+
+        results: list[DetectionResult] = []
+        for entry in label_entries:
+            label = str(entry.get("name", "")).strip()
+            if not label:
+                continue
+            confidence = float(entry.get("confidence", 1.0) or 0.0)
+            if confidence < threshold:
+                continue
+            severity = _resolve_pipeline_severity(label, schema.severity_map, default_severity)
+            matched = str(entry.get("matched_content") or "").strip() or snippet[:320]
+            results.append(
+                DetectionResult(
+                    detector_type=DetectorType.CUSTOM,
+                    finding_type=f"llm:{label}",
+                    category="CLASSIFICATION",
+                    severity=severity,
+                    confidence=min(0.99, confidence),
+                    matched_content=matched,
+                    location=None,
+                    custom_detector_key=self._detector_key,
+                    custom_detector_name=self._detector_name,
+                    detected_at=datetime.now(UTC),
+                    metadata={
+                        "runner": "LLM",
+                        "provider": self._runtime.provider.value,
+                        "model": self._runtime.model,
+                        "label": label,
+                        "fields": extracted,
+                    },
+                    extracted_data=extracted or None,
+                    extraction_method="LLM",
+                )
+            )
+
+        results.sort(key=lambda r: r.confidence, reverse=True)
+        return results
+
+    @staticmethod
+    def _coerce_fields(raw: Any) -> dict[str, Any]:
+        return {str(k): v for k, v in raw.items()} if isinstance(raw, dict) else {}

{classifyre_cli-0.4.9 → classifyre_cli-0.4.10}/src/models/generated_detectors.py

@@ -189,7 +189,7 @@ class DetectorCatalog(RootModel[list[DetectorCatalogEntry]]):
                 'categories': ['CLASSIFICATION', 'COMPLIANCE'],
                 'supported_asset_types': ['TXT', 'TABLE', 'URL', 'IMAGE'],
                 'recommended_model': 'mDeBERTa-v3 + SetFit + GLiNER + HuggingFace transformers',
-                'notes': 'User-defined rules and pipelines tailored to specific business needs. Supports regex, GLiNER2, LLM, text classification, image classification, feature extraction, and object detection pipelines.',
+                'notes': 'User-defined rules and pipelines tailored to specific business needs. Supports regex, GLiNER2, AI/LLM (prompt-driven classification + extraction via a configured provider), text classification, image classification, feature extraction, and object detection pipelines.',
             },
         ],
         description='Detector capability catalog used for planning and runtime routing',
@@ -954,18 +954,156 @@ class RegexPipelineSchema(BaseModel):
     validation: PipelineValidationConfig | None = None
 
 
+class LLMLabelDefinition(BaseModel):
+    """
+    One classification label the AI detector may assign to content.
+    """
+
+    model_config = ConfigDict(
+        extra='forbid',
+    )
+    name: str = Field(
+        ...,
+        description="Label name returned by the model (e.g. 'good', 'bad', 'violent').",
+    )
+    description: str | None = Field(
+        '', description='Guidance describing when this label applies.'
+    )
+
+
 class Type3(StrEnum):
+    string = 'string'
+    number = 'number'
+    boolean = 'boolean'
+    list_string_ = 'list[string]'
+    list_number_ = 'list[number]'
+
+
+class LLMOutputField(BaseModel):
+    """
+    One structured property the AI detector extracts and stores in finding metadata and extracted_data.
+    """
+
+    model_config = ConfigDict(
+        extra='forbid',
+    )
+    name: str = Field(
+        ..., description='Output field name — becomes a key in extracted_data JSON.'
+    )
+    description: str | None = Field(
+        '', description='Hint for what this field captures.'
+    )
+    type: Type3 | None = 'string'
+
+
+class Provider(StrEnum):
+    """
+    Resolved AI provider type.
+    """
+
+    OPENAI_COMPATIBLE = 'OPENAI_COMPATIBLE'
+    CLAUDE = 'CLAUDE'
+    GEMINI = 'GEMINI'
+
+
+class LLMProviderRuntime(BaseModel):
+    """
+    Runtime-only provider credentials injected by the API at dispatch time. Never persisted with the detector config and rejected on create/update.
+    """
+
+    model_config = ConfigDict(
+        extra='forbid',
+    )
+    provider: Provider = Field(..., description='Resolved AI provider type.')
+    model: str = Field(
+        ...,
+        description='Resolved model identifier (e.g. gpt-4o, claude-sonnet-4-5, gemini-2.0-flash).',
+    )
+    api_key: str = Field(..., description='Decrypted provider API key.')
+    base_url: str | None = Field(
+        None,
+        description='Base URL for OpenAI-compatible endpoints. Null for managed providers.',
+    )
+    context_size: int | None = Field(
+        None, description='Optional context window size configured for the provider.'
+    )
+
+
+class Type4(StrEnum):
     LLM = 'LLM'
 
 
+class MaxTokens(RootModel[int]):
+    root: int = Field(
+        None,
+        description='Maximum tokens to generate. Provider default when null.',
+        ge=1,
+    )
+
+
 class LLMPipelineSchema(BaseModel):
+    """
+    AI detector pipeline. Sends content to a configured LLM provider with a system prompt, classifies it against a label set, and extracts structured fields. Predicted labels become findings (severity via severity_map); extracted fields are stored in finding metadata and extracted_data.
+    """
+
     model_config = ConfigDict(
         extra='forbid',
     )
     type: Literal['LLM'] = 'LLM'
+    system_prompt: str = Field(
+        ...,
+        description='Instruction describing what the model should detect, classify, and extract.',
+    )
+    response_example: str | None = Field(
+        None,
+        description='Optional few-shot example of the JSON the model should return.',
+    )
+    temperature: float | None = Field(
+        0.0,
+        description='Sampling temperature. Lower is more deterministic.',
+        ge=0.0,
+        le=2.0,
+    )
+    max_tokens: MaxTokens | None = Field(
+        None, description='Maximum tokens to generate. Provider default when null.'
+    )
+    labels: list[LLMLabelDefinition] | None = Field(
+        [],
+        description='Classification taxonomy the model assigns to content.',
+        validate_default=True,
+    )
+    multi_label: bool | None = Field(
+        False, description='Allow more than one label per asset.'
+    )
+    severity: Severity | None = Field(
+        'info',
+        description='Default severity when no severity_map rule matches a predicted label.',
+    )
+    severity_map: list[PipelineSeverityRule] | None = Field(
+        None,
+        description='Ordered rules mapping predicted labels to severity levels. First matching rule wins.',
+    )
+    confidence_threshold: float | None = Field(
+        0.5,
+        description='Minimum model confidence to report a label as a finding (0-1).',
+        ge=0.0,
+        le=1.0,
+    )
+    output_fields: list[LLMOutputField] | None = Field(
+        [],
+        description='Structured properties the model extracts. Stored in finding metadata and extracted_data.',
+        validate_default=True,
+    )
+    content_limit: int | None = Field(
+        8000, description='Maximum characters of content sent to the model.', ge=1
+    )
+    provider_runtime: LLMProviderRuntime | None = Field(
+        None,
+        description='Runtime-only credentials injected by the API at dispatch. Never persisted; rejected on create/update.',
+    )
 
 
-class Type4(StrEnum):
+class Type5(StrEnum):
     TEXT_CLASSIFICATION = 'TEXT_CLASSIFICATION'
 
 
@@ -1055,7 +1193,7 @@ class TextClassificationPipelineSchema(BaseModel):
     )
 
 
-class Type5(StrEnum):
+class Type6(StrEnum):
     IMAGE_CLASSIFICATION = 'IMAGE_CLASSIFICATION'
 
 
@@ -1108,7 +1246,7 @@ class ImageClassificationPipelineSchema(BaseModel):
     )
 
 
-class Type6(StrEnum):
+class Type7(StrEnum):
     FEATURE_EXTRACTION = 'FEATURE_EXTRACTION'
 
 
@@ -1180,7 +1318,7 @@ class FeatureExtractionPipelineSchema(BaseModel):
     )
 
 
-class Type7(StrEnum):
+class Type8(StrEnum):
     OBJECT_DETECTION = 'OBJECT_DETECTION'
 
 

classifyre_cli-0.4.10/tests/detectors/custom/test_llm_runner.py

@@ -0,0 +1,157 @@
+"""Unit tests for the AI/LLM detector runner.
+
+litellm.completion is mocked so the tests never hit a real provider.
+"""
+
+from __future__ import annotations
+
+import json
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+
+from src.detectors.custom.runners._llm import LLMRunner
+from src.models.generated_detectors import (
+    LLMLabelDefinition,
+    LLMOutputField,
+    LLMPipelineSchema,
+    LLMProviderRuntime,
+    PipelineSeverityRule,
+    Severity,
+)
+
+TEXT = "Dieses Produkt ist schrecklich und gewalttätig."
+
+
+def _runtime(provider: str = "CLAUDE", model: str = "claude-sonnet-4-5") -> LLMProviderRuntime:
+    return LLMProviderRuntime(provider=provider, model=model, api_key="sk-test")
+
+
+def _schema(**overrides) -> LLMPipelineSchema:
+    base: dict = {
+        "system_prompt": "Classify sentiment.",
+        "labels": [
+            LLMLabelDefinition(name="good"),
+            LLMLabelDefinition(name="bad", description="negative sentiment"),
+            LLMLabelDefinition(name="violent"),
+        ],
+        "severity_map": [
+            PipelineSeverityRule(pattern="violent", severity=Severity.critical),
+            PipelineSeverityRule(pattern="bad", severity=Severity.medium),
+        ],
+        "output_fields": [LLMOutputField(name="language", type="string")],
+        "confidence_threshold": 0.5,
+        "provider_runtime": _runtime(),
+    }
+    base.update(overrides)
+    return LLMPipelineSchema(**base)
+
+
+def _mock_completion(payload: dict) -> MagicMock:
+    message = SimpleNamespace(content=json.dumps(payload))
+    choice = SimpleNamespace(message=message)
+    return MagicMock(return_value=SimpleNamespace(choices=[choice]))
+
+
+def _runner(schema: LLMPipelineSchema, completion: MagicMock) -> LLMRunner:
+    runner = LLMRunner(schema, detector_key="sentiment", detector_name="Sentiment")
+    runner._litellm = SimpleNamespace(completion=completion)
+    return runner
+
+
+def test_missing_provider_runtime_raises() -> None:
+    schema = LLMPipelineSchema(system_prompt="x")
+    with pytest.raises(ValueError, match="provider_runtime"):
+        LLMRunner(schema, detector_key="d")
+
+
+def test_label_to_severity_and_extraction() -> None:
+    completion = _mock_completion(
+        {
+            "labels": [
+                {"name": "violent", "confidence": 0.95, "matched_content": "gewalttätig"},
+                {"name": "bad", "confidence": 0.8},
+            ],
+            "fields": {"language": "de"},
+        }
+    )
+    runner = _runner(_schema(), completion)
+    results = runner.detect(TEXT, "text/plain")
+
+    assert {r.finding_type for r in results} == {"llm:violent", "llm:bad"}
+    violent = next(r for r in results if r.finding_type == "llm:violent")
+    assert violent.severity == Severity.critical
+    assert violent.matched_content == "gewalttätig"
+    assert violent.extracted_data == {"language": "de"}
+    assert violent.extraction_method == "LLM"
+    assert violent.metadata["provider"] == "CLAUDE"
+
+    bad = next(r for r in results if r.finding_type == "llm:bad")
+    assert bad.severity == Severity.medium
+
+
+def test_confidence_threshold_filters() -> None:
+    completion = _mock_completion({"labels": [{"name": "bad", "confidence": 0.3}], "fields": {}})
+    runner = _runner(_schema(confidence_threshold=0.5), completion)
+    assert runner.detect(TEXT, "text/plain") == []
+
+
+def test_default_severity_when_no_rule_matches() -> None:
+    completion = _mock_completion({"labels": [{"name": "good", "confidence": 0.9}]})
+    runner = _runner(_schema(), completion)
+    results = runner.detect(TEXT, "text/plain")
+    assert results[0].severity == Severity.info
+
+
+def test_model_string_per_provider() -> None:
+    completion = _mock_completion({"labels": []})
+    runner = _runner(_schema(provider_runtime=_runtime("GEMINI", "gemini-2.0-flash")), completion)
+    runner.detect(TEXT, "text/plain")
+    assert completion.call_args.kwargs["model"] == "gemini/gemini-2.0-flash"
+
+
+def test_openai_compatible_uses_api_base_and_prefix() -> None:
+    runtime = LLMProviderRuntime(
+        provider="OPENAI_COMPATIBLE",
+        model="moonshotai/kimi-k2.6",
+        api_key="nv-key",
+        base_url="https://integrate.api.nvidia.com/v1",
+    )
+    completion = _mock_completion({"labels": []})
+    runner = _runner(_schema(provider_runtime=runtime), completion)
+    runner.detect(TEXT, "text/plain")
+    kwargs = completion.call_args.kwargs
+    assert kwargs["model"] == "openai/moonshotai/kimi-k2.6"
+    assert kwargs["api_base"] == "https://integrate.api.nvidia.com/v1"
+
+
+def test_max_tokens_passed_as_int() -> None:
+    # max_tokens is a RootModel[int] wrapper; it must be unwrapped to a plain int
+    # before being handed to litellm, otherwise the completion request fails.
+    completion = _mock_completion({"labels": []})
+    runner = _runner(_schema(max_tokens=256), completion)
+    runner.detect(TEXT, "text/plain")
+    assert completion.call_args.kwargs["max_tokens"] == 256
+
+
+def test_max_tokens_defaults_to_none() -> None:
+    completion = _mock_completion({"labels": []})
+    runner = _runner(_schema(), completion)
+    runner.detect(TEXT, "text/plain")
+    assert completion.call_args.kwargs["max_tokens"] is None
+
+
+def test_non_text_content_skipped() -> None:
+    completion = _mock_completion({"labels": [{"name": "bad", "confidence": 0.9}]})
+    runner = _runner(_schema(), completion)
+    assert runner.detect(b"bytes", "image/png") == []
+    assert runner.detect("text", "image/png") == []
+    completion.assert_not_called()
+
+
+def test_malformed_json_returns_no_findings() -> None:
+    message = SimpleNamespace(content="not json at all")
+    completion = MagicMock(return_value=SimpleNamespace(choices=[SimpleNamespace(message=message)]))
+    runner = _runner(_schema(), completion)
+    assert runner.detect(TEXT, "text/plain") == []

{classifyre_cli-0.4.9 → classifyre_cli-0.4.10}/tests/test_outputs.py

@@ -158,7 +158,7 @@ class _FakeSession:
         self.responses = responses
         self.calls: list[dict[str, Any]] = []
 
-    def mount(self, prefix: str, adapter: Any) -> None:  # noqa: ARG002
+    def mount(self, prefix: str, adapter: Any) -> None:
         """Absorb HTTPAdapter.mount() calls from RestOutputSink.__init__."""
 
     def request(