File size: 2,247 Bytes
6165ba9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 | from enum import Enum
from typing import Any, Dict, List, Optional
from datetime import datetime
from pydantic import BaseModel, Field
# --- Enums (from enhanced_extractor.py) ---
class DataSource(str, Enum):
"""Enumeration of data sources for provenance tracking"""
HF_API = "huggingface_api"
MODEL_CARD = "model_card_yaml"
README_TEXT = "readme_text"
CONFIG_FILE = "config_file"
REPOSITORY_FILES = "repository_files"
EXTERNAL_REFERENCE = "external_reference"
INTELLIGENT_DEFAULT = "intelligent_default"
PLACEHOLDER = "placeholder"
REGISTRY_DRIVEN = "registry_driven"
class ConfidenceLevel(str, Enum):
"""Confidence levels for extracted data"""
HIGH = "high" # Direct API data, official sources
MEDIUM = "medium" # Inferred from reliable patterns
LOW = "low" # Weak inference or pattern matching
NONE = "none" # Placeholder values
# --- internal Models ---
class ExtractionResult(BaseModel):
"""Container for extraction results with full provenance"""
value: Any
source: DataSource
confidence: ConfidenceLevel
extraction_method: str
timestamp: str = Field(default_factory=lambda: datetime.utcnow().isoformat())
fallback_chain: List[str] = Field(default_factory=list)
def __str__(self):
return f"{self.value} (source: {self.source.value}, confidence: {self.confidence.value})"
# --- API Request Models ---
class GenerateRequest(BaseModel):
model_id: str
include_inference: bool = True
use_best_practices: bool = True
hf_token: Optional[str] = None
class BatchRequest(BaseModel):
model_ids: List[str]
include_inference: bool = True
use_best_practices: bool = True
hf_token: Optional[str] = None
# --- API Response Models ---
class AIBOMResponse(BaseModel):
aibom: Dict[str, Any]
model_id: str
generated_at: str
request_id: str
download_url: str
completeness_score: Optional[Dict[str, Any]] = None
class EnhancementReport(BaseModel):
ai_enhanced: bool = False
ai_model: Optional[str] = None
original_score: Dict[str, Any]
final_score: Dict[str, Any]
improvement: float = 0
|