HF_MY_MODEL = {
    "catalog_id": "my_model",
    "category": CatalogCategory.HUGGINGFACE,
    "display_name": "My Custom Model",
    "description": "Description of what this model does",
    "provider_id": "huggingface",
    "model_type": "hf_local",
    "model_name": "organization/model-name",  # HuggingFace model ID
    "task_type": "text-classification",  # or "text-generation"
    "roles": ["triage"],  # triage, deep_scan, judge
    "parser_id": "hf_classification",  # Parser to use
    "parser_config": {
        "positive_labels": ["LABEL_1", "vulnerable"],
        "negative_labels": ["LABEL_0", "safe"],
        "threshold": 0.5,
    },
    "hf_kwargs": {
        "trust_remote_code": False,
    },
    "generation_kwargs": {
        "max_new_tokens": 512,
        "temperature": 0.1,
    },
    "size_mb": 500,
    "requires_gpu": False,
    "tags": ["classifier", "triage"],
}

# Add to MODEL_CATALOG list
MODEL_CATALOG.append(HF_MY_MODEL)

CLOUD_MY_MODEL = {
    "catalog_id": "my_cloud_model",
    "category": CatalogCategory.CLOUD,
    "display_name": "My Cloud Model",
    "description": "Cloud LLM for security analysis",
    "provider_id": "openai",  # openai, anthropic, google
    "model_type": "openai_cloud",  # openai_cloud, anthropic_cloud, google_cloud
    "model_name": "gpt-4o-mini",  # Provider's model name
    "task_type": "text-generation",
    "roles": ["deep_scan", "judge"],
    "parser_id": "json_schema",
    "generation_kwargs": {
        "max_tokens": 2048,
        "temperature": 0.1,
    },
    "requires_api_key": True,
    "tags": ["cloud", "llm", "deep_scan"],
}

from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.ensemble import RandomForestClassifier
import joblib

pipeline = Pipeline([
    ('tfidf', TfidfVectorizer(max_features=5000)),
    ('clf', RandomForestClassifier(n_estimators=100))
])
pipeline.fit(X_train, y_train)
joblib.dump(pipeline, 'my_model.pkl')

ML_MY_MODEL = {
    "catalog_id": "my_ml_model",
    "category": CatalogCategory.CLASSIC_ML,
    "display_name": "My ML Classifier",
    "description": "Custom sklearn model for vulnerability detection",
    "provider_id": "tool_ml",
    "model_type": "tool_ml",
    "model_name": "my_ml_model",
    "tool_id": "sklearn_classifier",
    "task_type": "classification",
    "roles": ["triage"],
    "parser_id": "tool_result",
    "settings": {
        "model_url": "https://github.com/user/repo/releases/download/v1/my_model.pkl",
        "threshold": 0.5,
    },
    "size_mb": 50,
    "requires_gpu": False,
    "tags": ["ml", "sklearn", "triage"],
}

# aegis/parsers/my_parser.py
from aegis.parsers.base import BaseParser
from aegis.models.schema import ParserResult, FindingCandidate

class MyParser(BaseParser):
    """Custom parser for specific model output format."""

    parser_id = "my_parser"

    def parse(self, raw_output: str, context: dict) -> ParserResult:
        """Parse model output into findings."""
        findings = []

        # Your parsing logic here
        # Extract vulnerabilities from raw_output

        for vuln in extracted_vulns:
            findings.append(FindingCandidate(
                file_path=context.get("file_path", "unknown"),
                line_start=vuln.get("line", 1),
                line_end=vuln.get("line", 1),
                snippet=vuln.get("code", ""),
                title=vuln.get("title", "Vulnerability"),
                category=vuln.get("category", "security"),
                cwe=vuln.get("cwe"),
                severity=vuln.get("severity", "medium"),
                description=vuln.get("description", ""),
                recommendation=vuln.get("fix", ""),
                confidence=vuln.get("confidence", 0.5),
            ))

        return ParserResult(
            findings=findings,
            raw_output=raw_output,
        )

from aegis.parsers.my_parser import MyParser

PARSER_REGISTRY = {
    # ... existing parsers
    "my_parser": MyParser,
}