dtx-models 0.18.2__py3-none-any.whl

dtx_models/providers/hf.py

@@ -0,0 +1,161 @@
+from enum import Enum
+from typing import Any, Dict, List, Literal, Optional
+
+from pydantic import BaseModel, Field, HttpUrl, field_serializer, model_validator
+
+from ..evaluator import  EvaluatorInScope
+from ..prompts import SupportedFormat
+
+
+class HuggingFaceTask(str, Enum):
+    TEXT_GENERATION = "text-generation"
+    TEXT2TEXT_GENERATION = "text2text-generation"
+    TEXT_CLASSIFICATION = "text-classification"
+    TOKEN_CLASSIFICATION = "token-classification"
+    FEATURE_EXTRACTION = "feature-extraction"
+    SENTENCE_SIMILARITY = "sentence-similarity"
+    FILL_MASK = "fill-mask"
+
+
+class DType(str, Enum):
+    FLOAT16 = "float16"
+    BFLOAT16 = "bfloat16"
+    FLOAT32 = "float32"
+
+
+class HFEndpoint(BaseModel):
+    api_endpoint: Optional[HttpUrl] = Field(
+        None, description="Custom API endpoint for the model."
+    )
+    api_key: Optional[str] = Field(None, description="Your HuggingFace API key.")
+
+
+class HuggingFaceProviderParams(BaseModel):
+    temperature: Optional[float] = Field(
+        None, ge=0, le=1, description="Controls randomness in generation."
+    )
+    top_k: Optional[int] = Field(
+        None, ge=1, description="Controls diversity via the top-k sampling strategy."
+    )
+    top_p: Optional[float] = Field(
+        None, ge=0, le=1, description="Controls diversity via nucleus sampling."
+    )
+    repetition_penalty: Optional[float] = Field(
+        None, ge=0, description="Penalty for repetition."
+    )
+    max_new_tokens: Optional[int] = Field(
+        None, ge=1, description="The maximum number of new tokens to generate."
+    )
+    max_time: Optional[float] = Field(
+        None, ge=0, description="The maximum time in seconds for the model to respond."
+    )
+    return_full_text: Optional[bool] = Field(
+        None, description="Whether to return the full text or just new text."
+    )
+    num_return_sequences: Optional[int] = Field(
+        None, ge=1, description="The number of sequences to return."
+    )
+    do_sample: Optional[bool] = Field(None, description="Whether to sample the output.")
+    use_cache: Optional[bool] = Field(None, description="Whether to use caching.")
+    wait_for_model: Optional[bool] = Field(
+        None, description="Whether to wait for the model to be ready."
+    )
+    extra_params: Optional[Dict[str, Any]] = Field(
+        None, description="Additional parameters to pass to HuggingFace API."
+    )
+    dtype: Optional[DType] = Field(
+        DType.BFLOAT16, description="Data type for model computation."
+    )
+
+    @field_serializer("dtype")
+    def serialize_dtype(self, dtype: DType) -> str:
+        """Serialize the enum to a string."""
+        if dtype:
+            return dtype.value
+        else:
+            return dtype
+
+
+class HuggingFaceProviderConfig(BaseModel):
+    model: str = Field(..., description="HuggingFace model name.")
+    task: HuggingFaceTask = Field(
+        ..., description="Task type for the HuggingFace model."
+    )
+    params: Optional[HuggingFaceProviderParams] = Field(
+        None, description="Configuration parameters for HuggingFace model."
+    )
+    endpoint: Optional[HFEndpoint] = Field(
+        None, description="HuggingFace API endpoint configuration."
+    )
+
+    support_multi_turn: bool = Field(
+        default=False,
+        description="Does it support Multi Turn",
+    )
+
+    supported_input_format: Optional[SupportedFormat] = Field(
+        default=SupportedFormat.TEXT,
+        description="Supported Input format",
+    )
+
+    id: Optional[str] = Field(
+        None,
+        description="HuggingFace model identifier computed as 'huggingface:<task>:<model>'",
+    )
+
+    preferred_evaluator: Optional[EvaluatorInScope] = Field(
+        None,
+        description="Preferred Evaluator for the provider",
+    )
+
+    @model_validator(mode="after")
+    def compute_id(cls, values):
+        if not values.id:
+            values.id = f"huggingface:{values.task.value}:{values.model}"
+        return values
+
+    @field_serializer("task")
+    def serialize_role(self, task: HuggingFaceTask) -> str:
+        """Serialize the enum to a string."""
+        return task.value
+
+    @field_serializer("supported_input_format")
+    def serialize_supported_format(
+        self, supported_input_format: SupportedFormat
+    ) -> str:
+        """Serialize the enum to a string."""
+        return str(supported_input_format)
+
+    def get_name(self) -> str:
+        """
+        Returns the model name as the provider's name.
+        """
+        return self.model
+
+
+class HuggingFaceGuardModelsProviderConfig(HuggingFaceProviderConfig):
+    safe_value: str = Field(
+        "",
+        description="JSONPath expression indicating the condition for input to be safe.",
+    )
+
+    """
+    Example Usage:
+    ----------------
+    provider = HuggingFaceGuardModelsProviderConfig(safe_value="$.score[?(@ < 0.6)]")
+    
+    This JSONPath expression checks if all scores in the input JSON are below 0.6.
+    """
+
+
+class HFProvider(BaseModel):
+    provider: Literal["huggingface"] = Field(
+        "huggingface", description="Provider ID, always set to 'http'."
+    )
+    config: HuggingFaceProviderConfig
+
+
+class HFModels(BaseModel):
+    huggingface: List[HuggingFaceProviderConfig | HuggingFaceGuardModelsProviderConfig] = Field(
+        default_factory=list, description="List of predefined models."
+    )

dtx_models/providers/http.py

@@ -0,0 +1,152 @@
+from enum import Enum
+from typing import Any, Callable, Dict, Literal, Optional, Union
+from dtx_models.utils.urls import url_2_name
+
+from pydantic import (
+    BaseModel,
+    Field,
+    HttpUrl,
+    field_serializer,
+    field_validator,
+)
+
+
+class HttpMethod(str, Enum):
+    GET = "GET"
+    POST = "POST"
+    PUT = "PUT"
+    DELETE = "DELETE"
+    PATCH = "PATCH"
+
+    def __str__(self):
+        return self.value  # Ensures correct YAML serialization
+
+    @classmethod
+    def values(cls):
+        return [member.value for member in cls]
+
+
+class BaseHttpProvider(BaseModel):
+    """Base class for common HTTP provider attributes."""
+
+    max_retries: int = Field(
+        4, ge=0, description="Maximum number of retries for failed requests."
+    )
+    validate_response: Optional[Union[str, Callable[[int], bool]]] = Field(
+        None, description="Validation function for HTTP response status."
+    )
+    transform_request: Optional[Union[str, Dict[str, Any]]] = Field(
+        None, description="Template or mapping to modify request before sending."
+    )
+    transform_response: Optional[Union[str, Dict[str, Any]]] = Field(
+        None, description="Transformation logic for processing API responses."
+    )
+
+    example_response: Optional[str] = Field(
+        None, description="Example response dumped as sample if available"
+    )
+
+
+class StructuredHttpProviderConfig(BaseHttpProvider):
+    """Defines an HTTP request provider with structured fields (URL, method, headers, body)."""
+
+    url: str = Field(..., description="The HTTP endpoint URL.")
+    method: HttpMethod = Field(..., description="HTTP method.")
+    headers: Optional[Dict[str, str]] = Field(
+        default_factory=dict, description="HTTP headers."
+    )
+    body: Optional[Union[str, Dict[str, Any]]] = Field(
+        None, description="HTTP request body as JSON or a form-urlencoded string."
+    )
+
+    @field_serializer("method")
+    def serialize_http_method(self, method: HttpMethod) -> Optional[str]:
+        """Serialize the HTTP method to a string."""
+        return str(method.value) if method else None
+
+    @field_validator("url", mode="before")
+    @classmethod
+    def validate_base_url(cls, value: str) -> str:
+        """Validate that the URL is well-formed."""
+        try:
+            HttpUrl(value)  # Validate the URL format
+        except ValueError as e:
+            raise ValueError(f"Invalid base_url: {value}") from e
+        return value
+
+    @field_validator("body", mode="before")
+    @classmethod
+    def validate_body(
+        cls, value: Optional[Union[str, Dict[str, Any]]]
+    ) -> Optional[Union[str, Dict[str, Any]]]:
+        """Ensure body is either a JSON dictionary or a form-urlencoded string."""
+        if value is not None and not isinstance(value, (str, dict)):
+            raise ValueError(
+                "Body must be a dictionary (JSON) or a string (form-urlencoded)."
+            )
+        return value
+
+
+    def get_name(self) -> str:
+        """Generate a name from the URL using scheme:host:port:/path"""
+        return url_2_name(self.url, level=3)
+
+
+class RawHttpProviderConfig(BaseHttpProvider):
+    """Defines an HTTP request provider using raw_request instead of structured fields."""
+
+    raw_request: str = Field(..., description="Full raw HTTP request in text format.")
+    use_https: bool = Field(default=False, description="Whether to use HTTPS.")
+
+    @field_validator("raw_request", mode="before")
+    @classmethod
+    def validate_raw_request(cls, value: str) -> str:
+        """Ensure raw_request is a valid string."""
+        if not isinstance(value, str):
+            raise ValueError("Raw HTTP request must be a string.")
+        return value
+
+    def get_name(self) -> str:
+        """
+        Extract the HTTP method and path from the first line of the raw request.
+        Returns something like 'GET /api/v1/resource'.
+        """
+        try:
+            first_line = self.raw_request.strip().splitlines()[0]
+            parts = first_line.split()
+            method = parts[0] if len(parts) > 0 else "UNKNOWN"
+            path = parts[1] if len(parts) > 1 else "/"
+            return f"{method} {path}"
+        except Exception:
+            return "INVALID RAW REQUEST"
+
+
+class HttpProvider(BaseModel):
+    provider: Literal["http"] = Field(
+        "http", description="Provider ID, always set to 'http'."
+    )
+    config: StructuredHttpProviderConfig | RawHttpProviderConfig
+
+
+class BaseHttpProviderResponse(BaseModel):
+    """
+    Represents a standardized response format for HTTP providers.
+
+    - `mutations`: Stores any changes or replacements applied to the request.
+    - `request`: Stores the raw request sent (either as a string or a structured dictionary).
+    - `response`: Stores the raw response received (either as a string or a structured dictionary).
+    """
+
+    mutations: Optional[Dict[str, Any]] = Field(
+        default_factory=dict, description="Replacements performed in the request"
+    )
+    # request: Optional[Union[str, Dict[str, Any]]] = Field(
+    #     default_factory=dict, description="Request Sent"
+    # )
+    response: Optional[Union[str, Dict[str, Any]]] = Field(
+        default_factory=dict, description="Response Received"
+    )
+
+
+# class Providers(BaseModel):
+#     providers: List[HttpProvider]

dtx_models/providers/litellm.py

@@ -0,0 +1,21 @@
+from typing import Literal, Optional
+
+from pydantic import BaseModel, Field
+
+from .openai import OpenaiProviderConfig
+
+
+class LitellmProviderConfig(OpenaiProviderConfig):
+    endpoint: Optional[str] = Field(
+        default=None,
+        description="Base URL of the OpenAI server or proxy endpoint.",
+    )
+
+
+class LitellmProvider(BaseModel):
+    """Wrapper for OpenAI provider configuration."""
+
+    provider: Literal["litellm"] = Field(
+        "litellm", description="Provider ID, always set to 'openai'."
+    )
+    config: LitellmProviderConfig

dtx_models/providers/models_spec.py

@@ -0,0 +1,229 @@
+from enum import Enum
+from typing import List, Optional
+
+from pydantic import BaseModel, Field
+
+from ..providers.base import ProviderType
+
+# --- Define Modalities ---
+
+
+class Modality(str, Enum):
+    TEXT = "text"
+    IMAGE = "image"
+    CODE = "code"
+
+
+# --- Define Task Types ---
+
+
+class ModelTaskType(str, Enum):
+    GENERATION = "generation"
+    CLASSIFICATION = "classification"
+    EMBEDDING = "embedding"
+
+
+# --- Define Model Schema ---
+
+
+class ModelSpec(BaseModel):
+    name: str = Field(..., description="Model name, e.g., gpt-4o, llama3.")
+    task: ModelTaskType = Field(..., description="Primary task type of the model.")
+    modalities: List[Modality] = Field(
+        ..., description="Supported modalities (text, image, code, etc.)."
+    )
+    description: Optional[str] = Field(
+        None, description="Optional description of the model."
+    )
+    context_window: Optional[int] = Field(None, description="Context window in tokens.")
+    max_completion_tokens: Optional[int] = Field(
+        None, description="Maximum completion tokens."
+    )
+
+
+# --- Define Provider Schema ---
+
+
+class ProviderModels(BaseModel):
+    provider: ProviderType = Field(
+        ..., description="Name of the provider, e.g., openai, groq, etc."
+    )
+    models: List[ModelSpec] = Field(
+        ..., description="List of models available under this provider."
+    )
+
+
+# --- Example Repository ---
+
+models_repository: List[ProviderModels] = [
+    ProviderModels(
+        provider=ProviderType.OPENAI,
+        models=[
+            ModelSpec(
+                name="gpt-4.5-preview",
+                task=ModelTaskType.GENERATION,
+                modalities=[Modality.TEXT, Modality.CODE],
+                description="Largest and most capable GPT model.",
+                context_window=128000,
+                max_completion_tokens=4096,
+            ),
+            ModelSpec(
+                name="gpt-4o",
+                task=ModelTaskType.GENERATION,
+                modalities=[Modality.TEXT, Modality.CODE],
+                description="Fast, intelligent, flexible GPT model.",
+                context_window=128000,
+                max_completion_tokens=4096,
+            ),
+            ModelSpec(
+                name="gpt-4o-mini",
+                task=ModelTaskType.GENERATION,
+                modalities=[Modality.TEXT, Modality.CODE],
+                description="Fast, affordable small model for focused tasks.",
+                context_window=128000,
+                max_completion_tokens=4096,
+            ),
+            ModelSpec(
+                name="gpt-4-turbo",
+                task=ModelTaskType.GENERATION,
+                modalities=[Modality.TEXT, Modality.CODE],
+                description="An older high-intelligence GPT model.",
+                context_window=128000,
+                max_completion_tokens=4096,
+            ),
+            ModelSpec(
+                name="gpt-3.5-turbo",
+                task=ModelTaskType.GENERATION,
+                modalities=[Modality.TEXT],
+                description="Legacy GPT model for cheaper chat and non-chat tasks.",
+                context_window=16384,
+                max_completion_tokens=4096,
+            ),
+            ModelSpec(
+                name="text-embedding-3-large",
+                task=ModelTaskType.EMBEDDING,
+                modalities=[Modality.TEXT],
+                description="Most capable embedding model.",
+                context_window=None,
+                max_completion_tokens=None,
+            ),
+            ModelSpec(
+                name="omni-moderation-latest",
+                task=ModelTaskType.CLASSIFICATION,
+                modalities=[Modality.TEXT, Modality.IMAGE],
+                description="Identify potentially harmful content in text and images.",
+                context_window=None,
+                max_completion_tokens=None,
+            ),
+        ],
+    ),
+    ProviderModels(
+        provider=ProviderType.GROQ,
+        models=[
+            ModelSpec(
+                name="llama-3.3-70b-versatile",
+                task=ModelTaskType.GENERATION,
+                modalities=[Modality.TEXT],
+                description="Meta's versatile 70B parameter model.",
+                context_window=128000,
+                max_completion_tokens=32768,
+            ),
+            ModelSpec(
+                name="llama-3.1-8b-instant",
+                task=ModelTaskType.GENERATION,
+                modalities=[Modality.TEXT],
+                description="Meta's 8B parameter instant model.",
+                context_window=128000,
+                max_completion_tokens=8192,
+            ),
+            ModelSpec(
+                name="llama-guard-3-8b",
+                task=ModelTaskType.CLASSIFICATION,
+                modalities=[Modality.TEXT],
+                description="Meta's guard model for moderation and classification.",
+                context_window=8192,
+                max_completion_tokens=None,
+            ),
+            ModelSpec(
+                name="llama3-70b-8192",
+                task=ModelTaskType.GENERATION,
+                modalities=[Modality.TEXT],
+                description="Meta's 70B model with 8k context window.",
+                context_window=8192,
+                max_completion_tokens=None,
+            ),
+            ModelSpec(
+                name="llama3-8b-8192",
+                task=ModelTaskType.GENERATION,
+                modalities=[Modality.TEXT],
+                description="Meta's 8B model with 8k context window.",
+                context_window=8192,
+                max_completion_tokens=None,
+            ),
+            ModelSpec(
+                name="whisper-large-v3",
+                task=ModelTaskType.CLASSIFICATION,
+                modalities=[Modality.TEXT],
+                description="OpenAI's general-purpose speech recognition model.",
+                context_window=None,
+                max_completion_tokens=None,
+            ),
+            ModelSpec(
+                name="whisper-large-v3-turbo",
+                task=ModelTaskType.CLASSIFICATION,
+                modalities=[Modality.TEXT],
+                description="Turbo version of OpenAI's Whisper speech recognition model.",
+                context_window=None,
+                max_completion_tokens=None,
+            ),
+        ],
+    ),
+]
+
+
+# --- ModelSpecRepo class ---
+
+
+class ModelSpecRepo:
+    def __init__(self, repository: Optional[List[ProviderModels]] = None):
+        self.repository = repository or models_repository
+
+    def get_model_by_name(self, name: str) -> Optional[ModelSpec]:
+        for provider in self.repository:
+            for model in provider.models:
+                if model.name == name:
+                    return model
+        return None
+
+    def get_models_by_provider(self, provider_name: ProviderType) -> List[ModelSpec]:
+        for provider in self.repository:
+            if provider.provider == provider_name:
+                return provider.models
+        return []
+
+    def list_all_models(self) -> List[ModelSpec]:
+        all_models = []
+        for provider in self.repository:
+            all_models.extend(provider.models)
+        return all_models
+
+    def get_models_by_task(self, task_type: ModelTaskType) -> List[ModelSpec]:
+        models = []
+        for provider in self.repository:
+            for model in provider.models:
+                if model.task == task_type:
+                    models.append(model)
+        return models
+
+
+# Example usage
+if __name__ == "__main__":
+    # Example: Use ModelSpecRepo
+    repo = ModelSpecRepo()
+    print("\nFetching model by name 'gpt-4o':")
+    model = repo.get_model_by_name("gpt-4o")
+    print(model)
+
+    print("\nListing all generation models:")
+    for model in repo.get_models_by_task(ModelTaskType.GENERATION):
+        print(f"- {model.name}")

dtx_models/providers/ollama.py

@@ -0,0 +1,107 @@
+from enum import Enum
+from typing import Any, Dict, Literal, Optional
+from dtx_models.utils.urls import url_2_name
+
+from pydantic import BaseModel, Field, field_serializer, model_validator
+
+## Supported Yaml
+
+# """
+# ollama:
+#   - model: llama3
+
+
+# ollama:
+#   - model: llama3
+#     task: text-generation
+#     endpoint: http://localhost:11434
+#     params:
+#       temperature: 1.0
+#       top_k: 50
+#       top_p: 1.0
+#       repeat_penalty: 1.1
+#       max_tokens: 512
+#       num_return_sequences: 1
+#       extra_params:
+#         stop: ["###", "User:"]
+# """
+
+
+class OllamaTask(str, Enum):
+    TEXT_GENERATION = "text-generation"
+    TEXT_CLASSIFICATION = "text-classification"
+    SAFETY_CLASSIFICATION = "safety-classification"
+
+    def __str__(self):
+        return self.value  # Ensures correct YAML serialization
+
+    @classmethod
+    def values(cls):
+        return [member.value for member in cls]
+
+
+class OllamaProviderParams(BaseModel):
+    temperature: Optional[float] = Field(
+        None, ge=0, le=1, description="Controls randomness in generation."
+    )
+    top_k: Optional[int] = Field(
+        None, ge=1, description="Top-k sampling strategy for generation."
+    )
+    top_p: Optional[float] = Field(
+        None, ge=0, le=1, description="Nucleus sampling (top-p)."
+    )
+    repeat_penalty: Optional[float] = Field(
+        None, ge=0, description="Penalty for repeating tokens."
+    )
+    max_tokens: Optional[int] = Field(
+        None, ge=1, description="Maximum number of tokens to generate."
+    )
+    num_return_sequences: Optional[int] = Field(
+        None, ge=1, description="Number of sequences to return."
+    )
+    extra_params: Optional[Dict[str, Any]] = Field(
+        default_factory=dict,
+        description="Additional parameters for Ollama model invocation.",
+    )
+
+
+class OllamaProviderConfig(BaseModel):
+    model: str = Field(..., description="Ollama model name (e.g. llama3, mistral).")
+    task: Optional[OllamaTask] = Field(
+        default=None,
+        description="Task type for the Ollama model. If not provided, will be inferred from the model name.",
+    )
+    params: Optional[OllamaProviderParams] = Field(
+        None, description="Optional parameters for customizing model behavior."
+    )
+    endpoint: Optional[str] = Field(
+        default="http://localhost:11434", description="Base URL of the Ollama server."
+    )
+
+    @model_validator(mode="after")
+    def compute_fields(cls, values):
+        if not values.task:
+            if "guard" in values.model:
+                values.task = OllamaTask.TEXT_CLASSIFICATION
+            else:
+                values.task = OllamaTask.TEXT_GENERATION
+
+        return values
+
+    @field_serializer("task")
+    def serialize_task(self, task: OllamaTask) -> str:
+        return task.value
+
+
+    def get_name(self) -> str:
+        """
+        Returns a name like 'llama3:http:localhost:11434:/' by combining model and formatted endpoint.
+        """
+        return f"{self.model}:{url_2_name(self.endpoint, level=3)}"
+
+
+class OllamaProvider(BaseModel):
+    provider: Literal["ollama"] = Field(
+        "ollama", description="Provider ID, always set to 'ollama'."
+    )
+    config: OllamaProviderConfig