Spaces:
Running
Running
File size: 5,225 Bytes
0231daa |
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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 |
"""
Common schemas shared across request and response models.
This module contains Pydantic models used by both requests and responses,
such as SparseEmbedding and ModelInfo.
"""
from typing import List, Optional, Literal
from pydantic import BaseModel, Field, ConfigDict
class SparseEmbedding(BaseModel):
"""
Sparse embedding representation.
Sparse embeddings are represented as two parallel arrays:
- indices: positions of non-zero values
- values: the actual values at those positions
Attributes:
indices: List of indices for non-zero elements
values: List of values corresponding to the indices
text: Optional original text that was embedded
"""
indices: List[int] = Field(
..., description="Indices of non-zero elements in the sparse vector"
)
values: List[float] = Field(..., description="Values corresponding to the indices")
text: Optional[str] = Field(None, description="Original text that was embedded")
class Config:
json_schema_extra = {
"example": {
"indices": [10, 25, 42, 100],
"values": [0.85, 0.62, 0.91, 0.73],
"text": "example query text",
}
}
class ModelInfo(BaseModel):
"""
Information about an available model.
Attributes:
id: Unique identifier for the model
name: Full model name (e.g., Hugging Face model path)
type: Model type ('embeddings' or 'sparse-embeddings')
loaded: Whether the model is currently loaded in memory
"""
id: str = Field(..., description="Unique model identifier")
name: str = Field(..., description="Full model name")
type: str = Field(..., description="Model type (embeddings or sparse-embeddings)")
loaded: bool = Field(..., description="Whether model is loaded in memory")
class Config:
json_schema_extra = {
"example": {
"id": "qwen3-0.6b",
"name": "Qwen/Qwen3-Embedding-0.6B",
"type": "embeddings",
"loaded": True,
}
}
class HealthStatus(BaseModel):
"""
Health check status information.
Attributes:
status: Overall status (ok or error)
total_models: Total number of configured models
loaded_models: Number of models currently loaded
startup_complete: Whether startup sequence is complete
"""
status: str = Field(..., description="Overall status")
total_models: int = Field(..., description="Total configured models")
loaded_models: int = Field(..., description="Currently loaded models")
startup_complete: bool = Field(..., description="Startup completion status")
class Config:
json_schema_extra = {
"example": {
"status": "ok",
"total_models": 2,
"loaded_models": 2,
"startup_complete": True,
}
}
class ErrorResponse(BaseModel):
"""
Standard error response format.
Attributes:
error: Error type/name
message: Detailed error message
detail: Additional error details (optional)
"""
error: str = Field(..., description="Error type")
message: str = Field(..., description="Error message")
detail: Optional[str] = Field(None, description="Additional details")
class Config:
json_schema_extra = {
"example": {
"error": "ModelNotFoundError",
"message": "Model 'unknown-model' not found in configuration",
"detail": "Available models: qwen3-0.6b, splade-pp-v2",
}
}
class EmbeddingOptions(BaseModel):
"""
Optional parameters for embedding generation.
These parameters are passed directly to the underlying sentence-transformers
model. Not all parameters work with all models - check model documentation.
Attributes:
normalize_embeddings: L2 normalize output embeddings
batch_size: Batch size for processing
convert_to_numpy: Return numpy arrays instead of lists
precision: Computation precision
"""
normalize_embeddings: Optional[bool] = Field(
None, description="L2 normalize the output embeddings"
)
batch_size: Optional[int] = Field(
None, ge=1, le=256, description="Batch size for processing texts"
)
convert_to_numpy: Optional[bool] = Field(
None, description="Return numpy arrays instead of Python lists"
)
precision: Optional[Literal["float32", "int8", "uint8", "binary", "ubinary"]] = (
Field(None, description="Computation precision for embeddings")
)
model_config = ConfigDict(
extra="forbid", # Prevent typos in field names
json_schema_extra={
"example": {
"normalize_embeddings": True,
"batch_size": 32,
}
},
)
def to_kwargs(self) -> dict:
"""
Convert options to kwargs dict, excluding None values.
Returns:
Dictionary of non-None parameters
"""
return {k: v for k, v in self.model_dump().items() if v is not None}
|