sf-vector-sdk 0.2.0__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- sf_vector_sdk-0.2.0.dist-info/METADATA +476 -0
- sf_vector_sdk-0.2.0.dist-info/RECORD +27 -0
- sf_vector_sdk-0.2.0.dist-info/WHEEL +4 -0
- vector_sdk/__init__.py +262 -0
- vector_sdk/client.py +538 -0
- vector_sdk/content_types.py +233 -0
- vector_sdk/generated/embedding_pipeline/content_types/v1/content_types_pb2.py +57 -0
- vector_sdk/generated/embedding_pipeline/content_types/v1/content_types_pb2.pyi +141 -0
- vector_sdk/generated/embedding_pipeline/db/vectors/v1/vectors_pb2.py +58 -0
- vector_sdk/generated/embedding_pipeline/db/vectors/v1/vectors_pb2.pyi +145 -0
- vector_sdk/generated/embedding_pipeline/query/v1/query_pb2.py +58 -0
- vector_sdk/generated/embedding_pipeline/query/v1/query_pb2.pyi +109 -0
- vector_sdk/generated/embedding_pipeline/tools/v1/tools_pb2.py +39 -0
- vector_sdk/generated/embedding_pipeline/tools/v1/tools_pb2.pyi +31 -0
- vector_sdk/hash/__init__.py +31 -0
- vector_sdk/hash/hasher.py +259 -0
- vector_sdk/hash/types.py +67 -0
- vector_sdk/namespaces/__init__.py +13 -0
- vector_sdk/namespaces/base.py +45 -0
- vector_sdk/namespaces/db.py +230 -0
- vector_sdk/namespaces/embeddings.py +268 -0
- vector_sdk/namespaces/search.py +258 -0
- vector_sdk/structured/__init__.py +60 -0
- vector_sdk/structured/router.py +190 -0
- vector_sdk/structured/structured_embeddings.py +431 -0
- vector_sdk/structured/tool_config.py +254 -0
- vector_sdk/types.py +864 -0
vector_sdk/types.py
ADDED
|
@@ -0,0 +1,864 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Type definitions for Vector SDK.
|
|
3
|
+
|
|
4
|
+
These types match the Go service's internal types and provide
|
|
5
|
+
type safety and documentation for Python users.
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
from dataclasses import dataclass, field
|
|
9
|
+
from datetime import datetime
|
|
10
|
+
from enum import Enum
|
|
11
|
+
from typing import Any, Optional
|
|
12
|
+
|
|
13
|
+
# ============================================================================
|
|
14
|
+
# Embedding Model Registry
|
|
15
|
+
# ============================================================================
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
class EmbeddingProvider(str, Enum):
|
|
19
|
+
"""Embedding service provider."""
|
|
20
|
+
GOOGLE = "google"
|
|
21
|
+
OPENAI = "openai"
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
@dataclass
|
|
25
|
+
class ModelConfig:
|
|
26
|
+
"""Configuration for a supported embedding model."""
|
|
27
|
+
name: str
|
|
28
|
+
provider: EmbeddingProvider
|
|
29
|
+
default_dimensions: int
|
|
30
|
+
max_dimensions: int
|
|
31
|
+
supports_custom_dimensions: bool
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
# Registry of all supported embedding models
|
|
35
|
+
SUPPORTED_MODELS: dict[str, ModelConfig] = {
|
|
36
|
+
# Google Vertex AI Models
|
|
37
|
+
"gemini-embedding-001": ModelConfig(
|
|
38
|
+
name="gemini-embedding-001",
|
|
39
|
+
provider=EmbeddingProvider.GOOGLE,
|
|
40
|
+
default_dimensions=3072,
|
|
41
|
+
max_dimensions=3072,
|
|
42
|
+
supports_custom_dimensions=False,
|
|
43
|
+
),
|
|
44
|
+
"text-embedding-004": ModelConfig(
|
|
45
|
+
name="text-embedding-004",
|
|
46
|
+
provider=EmbeddingProvider.GOOGLE,
|
|
47
|
+
default_dimensions=768,
|
|
48
|
+
max_dimensions=768,
|
|
49
|
+
supports_custom_dimensions=False,
|
|
50
|
+
),
|
|
51
|
+
"text-multilingual-embedding-002": ModelConfig(
|
|
52
|
+
name="text-multilingual-embedding-002",
|
|
53
|
+
provider=EmbeddingProvider.GOOGLE,
|
|
54
|
+
default_dimensions=768,
|
|
55
|
+
max_dimensions=768,
|
|
56
|
+
supports_custom_dimensions=False,
|
|
57
|
+
),
|
|
58
|
+
# OpenAI Models
|
|
59
|
+
"text-embedding-3-small": ModelConfig(
|
|
60
|
+
name="text-embedding-3-small",
|
|
61
|
+
provider=EmbeddingProvider.OPENAI,
|
|
62
|
+
default_dimensions=1536,
|
|
63
|
+
max_dimensions=1536,
|
|
64
|
+
supports_custom_dimensions=True,
|
|
65
|
+
),
|
|
66
|
+
"text-embedding-3-large": ModelConfig(
|
|
67
|
+
name="text-embedding-3-large",
|
|
68
|
+
provider=EmbeddingProvider.OPENAI,
|
|
69
|
+
default_dimensions=3072,
|
|
70
|
+
max_dimensions=3072,
|
|
71
|
+
supports_custom_dimensions=True,
|
|
72
|
+
),
|
|
73
|
+
}
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
class ModelValidationError(ValueError):
|
|
77
|
+
"""Raised when model validation fails."""
|
|
78
|
+
pass
|
|
79
|
+
|
|
80
|
+
|
|
81
|
+
def is_model_supported(model: str) -> bool:
|
|
82
|
+
"""Check if a model is supported."""
|
|
83
|
+
return model in SUPPORTED_MODELS
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def get_model_config(model: str) -> Optional[ModelConfig]:
|
|
87
|
+
"""Get the configuration for a model."""
|
|
88
|
+
return SUPPORTED_MODELS.get(model)
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def get_supported_models() -> list[str]:
|
|
92
|
+
"""Get all supported model names."""
|
|
93
|
+
return list(SUPPORTED_MODELS.keys())
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def get_google_models() -> list[str]:
|
|
97
|
+
"""Get all Google/Vertex AI model names."""
|
|
98
|
+
return [
|
|
99
|
+
name for name, cfg in SUPPORTED_MODELS.items()
|
|
100
|
+
if cfg.provider == EmbeddingProvider.GOOGLE
|
|
101
|
+
]
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
def get_openai_models() -> list[str]:
|
|
105
|
+
"""Get all OpenAI model names."""
|
|
106
|
+
return [
|
|
107
|
+
name for name, cfg in SUPPORTED_MODELS.items()
|
|
108
|
+
if cfg.provider == EmbeddingProvider.OPENAI
|
|
109
|
+
]
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
def validate_model(model: str, dimensions: Optional[int] = None) -> None:
|
|
113
|
+
"""
|
|
114
|
+
Validate a model name and optional dimensions.
|
|
115
|
+
|
|
116
|
+
Args:
|
|
117
|
+
model: The model name to validate
|
|
118
|
+
dimensions: Optional dimensions to validate
|
|
119
|
+
|
|
120
|
+
Raises:
|
|
121
|
+
ModelValidationError: If validation fails
|
|
122
|
+
"""
|
|
123
|
+
config = SUPPORTED_MODELS.get(model)
|
|
124
|
+
if config is None:
|
|
125
|
+
supported_list = ", ".join(get_supported_models())
|
|
126
|
+
raise ModelValidationError(
|
|
127
|
+
f'Unsupported embedding model: "{model}". '
|
|
128
|
+
f"Supported models are: {supported_list}"
|
|
129
|
+
)
|
|
130
|
+
|
|
131
|
+
if dimensions is not None and dimensions > 0:
|
|
132
|
+
if dimensions > config.max_dimensions:
|
|
133
|
+
raise ModelValidationError(
|
|
134
|
+
f"Dimensions {dimensions} exceeds maximum "
|
|
135
|
+
f"{config.max_dimensions} for model \"{model}\""
|
|
136
|
+
)
|
|
137
|
+
|
|
138
|
+
if not config.supports_custom_dimensions and dimensions != config.max_dimensions:
|
|
139
|
+
raise ModelValidationError(
|
|
140
|
+
f'Model "{model}" does not support custom dimensions '
|
|
141
|
+
f"(requested {dimensions}, must be {config.max_dimensions})"
|
|
142
|
+
)
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
# Priority constants
|
|
146
|
+
PRIORITY_CRITICAL = "critical"
|
|
147
|
+
PRIORITY_HIGH = "high"
|
|
148
|
+
PRIORITY_NORMAL = "normal"
|
|
149
|
+
PRIORITY_LOW = "low"
|
|
150
|
+
|
|
151
|
+
# Stream names
|
|
152
|
+
STREAM_CRITICAL = "embedding:critical"
|
|
153
|
+
STREAM_HIGH = "embedding:high"
|
|
154
|
+
STREAM_NORMAL = "embedding:normal"
|
|
155
|
+
STREAM_LOW = "embedding:low"
|
|
156
|
+
|
|
157
|
+
|
|
158
|
+
def get_stream_for_priority(priority: str) -> str:
|
|
159
|
+
"""Get the Redis Stream name for a given priority."""
|
|
160
|
+
mapping = {
|
|
161
|
+
PRIORITY_CRITICAL: STREAM_CRITICAL,
|
|
162
|
+
PRIORITY_HIGH: STREAM_HIGH,
|
|
163
|
+
PRIORITY_NORMAL: STREAM_NORMAL,
|
|
164
|
+
PRIORITY_LOW: STREAM_LOW,
|
|
165
|
+
}
|
|
166
|
+
return mapping.get(priority, STREAM_NORMAL)
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
@dataclass
|
|
170
|
+
class TextInput:
|
|
171
|
+
"""
|
|
172
|
+
A single text to embed.
|
|
173
|
+
|
|
174
|
+
Attributes:
|
|
175
|
+
id: External identifier for correlation (e.g., toolId, topicId)
|
|
176
|
+
text: The actual text content to embed
|
|
177
|
+
document: Optional full document to store alongside the embedding
|
|
178
|
+
"""
|
|
179
|
+
id: str
|
|
180
|
+
text: str
|
|
181
|
+
document: Optional[dict[str, Any]] = None
|
|
182
|
+
|
|
183
|
+
def to_dict(self) -> dict[str, Any]:
|
|
184
|
+
"""Convert to dictionary for JSON serialization."""
|
|
185
|
+
result = {"id": self.id, "text": self.text}
|
|
186
|
+
if self.document is not None:
|
|
187
|
+
result["document"] = self.document
|
|
188
|
+
return result
|
|
189
|
+
|
|
190
|
+
|
|
191
|
+
@dataclass
|
|
192
|
+
class MongoDBStorage:
|
|
193
|
+
"""
|
|
194
|
+
Configuration for storing embeddings in MongoDB.
|
|
195
|
+
|
|
196
|
+
Attributes:
|
|
197
|
+
database: MongoDB database name (e.g., "events_new")
|
|
198
|
+
collection: Collection name (e.g., "tool_vectors", "topic_vectors")
|
|
199
|
+
embedding_field: Field name to store the embedding vector
|
|
200
|
+
upsert_key: Field to use for upsert operations (e.g., "contentHash")
|
|
201
|
+
"""
|
|
202
|
+
database: str
|
|
203
|
+
collection: str
|
|
204
|
+
embedding_field: str
|
|
205
|
+
upsert_key: str
|
|
206
|
+
|
|
207
|
+
def to_dict(self) -> dict[str, Any]:
|
|
208
|
+
"""Convert to dictionary for JSON serialization."""
|
|
209
|
+
return {
|
|
210
|
+
"database": self.database,
|
|
211
|
+
"collection": self.collection,
|
|
212
|
+
"embeddingField": self.embedding_field,
|
|
213
|
+
"upsertKey": self.upsert_key,
|
|
214
|
+
}
|
|
215
|
+
|
|
216
|
+
|
|
217
|
+
@dataclass
|
|
218
|
+
class TurboPufferStorage:
|
|
219
|
+
"""
|
|
220
|
+
Configuration for storing embeddings in TurboPuffer.
|
|
221
|
+
|
|
222
|
+
Attributes:
|
|
223
|
+
namespace: TurboPuffer namespace (e.g., "tool_vectors")
|
|
224
|
+
id_field: Document field to use as the vector ID
|
|
225
|
+
metadata: List of document fields to include as TurboPuffer metadata
|
|
226
|
+
"""
|
|
227
|
+
namespace: str
|
|
228
|
+
id_field: str
|
|
229
|
+
metadata: Optional[list[str]] = None
|
|
230
|
+
|
|
231
|
+
def to_dict(self) -> dict[str, Any]:
|
|
232
|
+
"""Convert to dictionary for JSON serialization."""
|
|
233
|
+
result = {
|
|
234
|
+
"namespace": self.namespace,
|
|
235
|
+
"idField": self.id_field,
|
|
236
|
+
}
|
|
237
|
+
if self.metadata:
|
|
238
|
+
result["metadata"] = self.metadata
|
|
239
|
+
return result
|
|
240
|
+
|
|
241
|
+
|
|
242
|
+
@dataclass
|
|
243
|
+
class PineconeStorageConfig:
|
|
244
|
+
"""
|
|
245
|
+
Configuration for storing embeddings in Pinecone.
|
|
246
|
+
|
|
247
|
+
Attributes:
|
|
248
|
+
index_name: Pinecone index name
|
|
249
|
+
namespace: Namespace within the index (optional)
|
|
250
|
+
id_field: Document field to use as the vector ID
|
|
251
|
+
metadata: List of document fields to include as Pinecone metadata
|
|
252
|
+
"""
|
|
253
|
+
index_name: str
|
|
254
|
+
id_field: str
|
|
255
|
+
namespace: Optional[str] = None
|
|
256
|
+
metadata: Optional[list[str]] = None
|
|
257
|
+
|
|
258
|
+
def to_dict(self) -> dict[str, Any]:
|
|
259
|
+
"""Convert to dictionary for JSON serialization."""
|
|
260
|
+
result = {
|
|
261
|
+
"indexName": self.index_name,
|
|
262
|
+
"idField": self.id_field,
|
|
263
|
+
}
|
|
264
|
+
if self.namespace:
|
|
265
|
+
result["namespace"] = self.namespace
|
|
266
|
+
if self.metadata:
|
|
267
|
+
result["metadata"] = self.metadata
|
|
268
|
+
return result
|
|
269
|
+
|
|
270
|
+
|
|
271
|
+
@dataclass
|
|
272
|
+
class StorageConfig:
|
|
273
|
+
"""
|
|
274
|
+
Configuration for where to store generated embeddings.
|
|
275
|
+
|
|
276
|
+
If multiple storage backends are provided, embeddings are written to all.
|
|
277
|
+
If none is provided, embeddings are only returned via callback.
|
|
278
|
+
|
|
279
|
+
Attributes:
|
|
280
|
+
mongodb: MongoDB storage configuration
|
|
281
|
+
turbopuffer: TurboPuffer storage configuration
|
|
282
|
+
pinecone: Pinecone storage configuration
|
|
283
|
+
"""
|
|
284
|
+
mongodb: Optional[MongoDBStorage] = None
|
|
285
|
+
turbopuffer: Optional[TurboPufferStorage] = None
|
|
286
|
+
pinecone: Optional[PineconeStorageConfig] = None
|
|
287
|
+
|
|
288
|
+
def to_dict(self) -> dict[str, Any]:
|
|
289
|
+
"""Convert to dictionary for JSON serialization."""
|
|
290
|
+
result = {}
|
|
291
|
+
if self.mongodb:
|
|
292
|
+
result["mongodb"] = self.mongodb.to_dict()
|
|
293
|
+
if self.turbopuffer:
|
|
294
|
+
result["turbopuffer"] = self.turbopuffer.to_dict()
|
|
295
|
+
if self.pinecone:
|
|
296
|
+
result["pinecone"] = self.pinecone.to_dict()
|
|
297
|
+
return result
|
|
298
|
+
|
|
299
|
+
|
|
300
|
+
@dataclass
|
|
301
|
+
class CallbackConfig:
|
|
302
|
+
"""
|
|
303
|
+
Configuration for completion notifications.
|
|
304
|
+
|
|
305
|
+
Attributes:
|
|
306
|
+
type: Callback delivery method ("redis", "pubsub", "none")
|
|
307
|
+
topic: Pub/Sub topic for "pubsub" type callbacks
|
|
308
|
+
channel: Redis channel for "redis" type callbacks
|
|
309
|
+
"""
|
|
310
|
+
type: str = "redis"
|
|
311
|
+
topic: Optional[str] = None
|
|
312
|
+
channel: Optional[str] = None
|
|
313
|
+
|
|
314
|
+
def to_dict(self) -> dict[str, Any]:
|
|
315
|
+
"""Convert to dictionary for JSON serialization."""
|
|
316
|
+
result = {"type": self.type}
|
|
317
|
+
if self.topic:
|
|
318
|
+
result["topic"] = self.topic
|
|
319
|
+
if self.channel:
|
|
320
|
+
result["channel"] = self.channel
|
|
321
|
+
return result
|
|
322
|
+
|
|
323
|
+
|
|
324
|
+
@dataclass
|
|
325
|
+
class EmbeddingConfigOverride:
|
|
326
|
+
"""
|
|
327
|
+
Configuration for overriding embedding model and dimensions.
|
|
328
|
+
|
|
329
|
+
Attributes:
|
|
330
|
+
model: Embedding model name (e.g., "gemini-embedding-001", "text-embedding-3-small")
|
|
331
|
+
dimensions: Output embedding dimensions
|
|
332
|
+
"""
|
|
333
|
+
model: Optional[str] = None
|
|
334
|
+
dimensions: Optional[int] = None
|
|
335
|
+
|
|
336
|
+
def to_dict(self) -> dict[str, Any]:
|
|
337
|
+
"""Convert to dictionary for JSON serialization."""
|
|
338
|
+
result = {}
|
|
339
|
+
if self.model:
|
|
340
|
+
result["model"] = self.model
|
|
341
|
+
if self.dimensions:
|
|
342
|
+
result["dimensions"] = self.dimensions
|
|
343
|
+
return result
|
|
344
|
+
|
|
345
|
+
|
|
346
|
+
@dataclass
|
|
347
|
+
class EmbeddingRequest:
|
|
348
|
+
"""
|
|
349
|
+
A request to generate embeddings.
|
|
350
|
+
|
|
351
|
+
Attributes:
|
|
352
|
+
request_id: Unique identifier for tracking
|
|
353
|
+
content_type: Type of content (e.g., "topic", "flashcard")
|
|
354
|
+
priority: Queue priority ("critical", "high", "normal", "low")
|
|
355
|
+
texts: List of texts to embed
|
|
356
|
+
storage: Where to store the embeddings
|
|
357
|
+
callback: How to notify completion
|
|
358
|
+
embedding_config: Optional embedding model override
|
|
359
|
+
metadata: Arbitrary key-value pairs for tracking
|
|
360
|
+
created_at: When the request was created
|
|
361
|
+
"""
|
|
362
|
+
request_id: str
|
|
363
|
+
content_type: str
|
|
364
|
+
priority: str
|
|
365
|
+
texts: list[TextInput]
|
|
366
|
+
storage: Optional[StorageConfig] = None
|
|
367
|
+
callback: Optional[CallbackConfig] = None
|
|
368
|
+
embedding_config: Optional[EmbeddingConfigOverride] = None
|
|
369
|
+
metadata: dict[str, str] = field(default_factory=dict)
|
|
370
|
+
created_at: datetime = field(default_factory=datetime.utcnow)
|
|
371
|
+
|
|
372
|
+
def to_dict(self) -> dict[str, Any]:
|
|
373
|
+
"""Convert to dictionary for JSON serialization."""
|
|
374
|
+
result = {
|
|
375
|
+
"requestId": self.request_id,
|
|
376
|
+
"contentType": self.content_type,
|
|
377
|
+
"priority": self.priority,
|
|
378
|
+
"texts": [t.to_dict() for t in self.texts],
|
|
379
|
+
"metadata": self.metadata,
|
|
380
|
+
"createdAt": self.created_at.isoformat() + "Z",
|
|
381
|
+
}
|
|
382
|
+
if self.storage:
|
|
383
|
+
result["storage"] = self.storage.to_dict()
|
|
384
|
+
if self.callback:
|
|
385
|
+
result["callback"] = self.callback.to_dict()
|
|
386
|
+
if self.embedding_config:
|
|
387
|
+
result["embeddingConfig"] = self.embedding_config.to_dict()
|
|
388
|
+
return result
|
|
389
|
+
|
|
390
|
+
|
|
391
|
+
@dataclass
|
|
392
|
+
class EmbeddingError:
|
|
393
|
+
"""
|
|
394
|
+
Details about a failed embedding.
|
|
395
|
+
|
|
396
|
+
Attributes:
|
|
397
|
+
index: Position in the original texts array
|
|
398
|
+
id: The TextInput.id that failed
|
|
399
|
+
error: Error message
|
|
400
|
+
retryable: Whether this error can be retried
|
|
401
|
+
"""
|
|
402
|
+
index: int
|
|
403
|
+
id: str
|
|
404
|
+
error: str
|
|
405
|
+
retryable: bool
|
|
406
|
+
|
|
407
|
+
@classmethod
|
|
408
|
+
def from_dict(cls, data: dict[str, Any]) -> "EmbeddingError":
|
|
409
|
+
"""Create from dictionary."""
|
|
410
|
+
return cls(
|
|
411
|
+
index=data["index"],
|
|
412
|
+
id=data["id"],
|
|
413
|
+
error=data["error"],
|
|
414
|
+
retryable=data.get("retryable", False),
|
|
415
|
+
)
|
|
416
|
+
|
|
417
|
+
|
|
418
|
+
@dataclass
|
|
419
|
+
class TimingBreakdown:
|
|
420
|
+
"""
|
|
421
|
+
Processing duration breakdown.
|
|
422
|
+
|
|
423
|
+
Attributes:
|
|
424
|
+
queue_wait_ms: Time spent waiting in queue
|
|
425
|
+
vertex_ms: Time spent calling Vertex AI
|
|
426
|
+
mongodb_ms: Time spent writing to MongoDB
|
|
427
|
+
turbopuffer_ms: Time spent writing to TurboPuffer
|
|
428
|
+
total_ms: Total processing time
|
|
429
|
+
"""
|
|
430
|
+
queue_wait_ms: int
|
|
431
|
+
vertex_ms: int
|
|
432
|
+
mongodb_ms: int
|
|
433
|
+
turbopuffer_ms: int
|
|
434
|
+
total_ms: int
|
|
435
|
+
|
|
436
|
+
@classmethod
|
|
437
|
+
def from_dict(cls, data: dict[str, Any]) -> "TimingBreakdown":
|
|
438
|
+
"""Create from dictionary."""
|
|
439
|
+
return cls(
|
|
440
|
+
queue_wait_ms=data.get("queueWaitMs", 0),
|
|
441
|
+
vertex_ms=data.get("vertexMs", 0),
|
|
442
|
+
mongodb_ms=data.get("mongodbMs", 0),
|
|
443
|
+
turbopuffer_ms=data.get("turbopufferMs", 0),
|
|
444
|
+
total_ms=data.get("totalMs", 0),
|
|
445
|
+
)
|
|
446
|
+
|
|
447
|
+
|
|
448
|
+
@dataclass
|
|
449
|
+
class EmbeddingResult:
|
|
450
|
+
"""
|
|
451
|
+
Result of processing an embedding request.
|
|
452
|
+
|
|
453
|
+
Attributes:
|
|
454
|
+
request_id: The original request ID
|
|
455
|
+
status: "success", "partial", or "failed"
|
|
456
|
+
processed_count: Number of successfully processed embeddings
|
|
457
|
+
failed_count: Number of failed embeddings
|
|
458
|
+
errors: Details about failed items
|
|
459
|
+
timing: Processing duration breakdown
|
|
460
|
+
completed_at: When processing finished
|
|
461
|
+
"""
|
|
462
|
+
request_id: str
|
|
463
|
+
status: str
|
|
464
|
+
processed_count: int
|
|
465
|
+
failed_count: int
|
|
466
|
+
errors: list[EmbeddingError]
|
|
467
|
+
timing: Optional[TimingBreakdown]
|
|
468
|
+
completed_at: datetime
|
|
469
|
+
|
|
470
|
+
@classmethod
|
|
471
|
+
def from_dict(cls, data: dict[str, Any]) -> "EmbeddingResult":
|
|
472
|
+
"""Create from dictionary."""
|
|
473
|
+
errors = [EmbeddingError.from_dict(e) for e in data.get("errors", [])]
|
|
474
|
+
timing = None
|
|
475
|
+
if data.get("timing"):
|
|
476
|
+
timing = TimingBreakdown.from_dict(data["timing"])
|
|
477
|
+
|
|
478
|
+
return cls(
|
|
479
|
+
request_id=data["requestId"],
|
|
480
|
+
status=data["status"],
|
|
481
|
+
processed_count=data["processedCount"],
|
|
482
|
+
failed_count=data["failedCount"],
|
|
483
|
+
errors=errors,
|
|
484
|
+
timing=timing,
|
|
485
|
+
completed_at=datetime.fromisoformat(data["completedAt"].replace("Z", "+00:00")),
|
|
486
|
+
)
|
|
487
|
+
|
|
488
|
+
@property
|
|
489
|
+
def is_success(self) -> bool:
|
|
490
|
+
"""Check if the request was fully successful."""
|
|
491
|
+
return self.status == "success"
|
|
492
|
+
|
|
493
|
+
@property
|
|
494
|
+
def is_partial(self) -> bool:
|
|
495
|
+
"""Check if the request was partially successful."""
|
|
496
|
+
return self.status == "partial"
|
|
497
|
+
|
|
498
|
+
@property
|
|
499
|
+
def is_failed(self) -> bool:
|
|
500
|
+
"""Check if the request completely failed."""
|
|
501
|
+
return self.status == "failed"
|
|
502
|
+
|
|
503
|
+
|
|
504
|
+
# ============================================================================
|
|
505
|
+
# Query Types
|
|
506
|
+
# ============================================================================
|
|
507
|
+
|
|
508
|
+
# Query stream names
|
|
509
|
+
QUERY_STREAM_CRITICAL = "query:critical"
|
|
510
|
+
QUERY_STREAM_HIGH = "query:high"
|
|
511
|
+
QUERY_STREAM_NORMAL = "query:normal"
|
|
512
|
+
QUERY_STREAM_LOW = "query:low"
|
|
513
|
+
|
|
514
|
+
|
|
515
|
+
def get_query_stream_for_priority(priority: str) -> str:
|
|
516
|
+
"""Get the Redis Stream name for a query request."""
|
|
517
|
+
mapping = {
|
|
518
|
+
PRIORITY_CRITICAL: QUERY_STREAM_CRITICAL,
|
|
519
|
+
PRIORITY_HIGH: QUERY_STREAM_HIGH,
|
|
520
|
+
PRIORITY_NORMAL: QUERY_STREAM_NORMAL,
|
|
521
|
+
PRIORITY_LOW: QUERY_STREAM_LOW,
|
|
522
|
+
}
|
|
523
|
+
return mapping.get(priority, QUERY_STREAM_NORMAL)
|
|
524
|
+
|
|
525
|
+
|
|
526
|
+
# Vector database types
|
|
527
|
+
VECTOR_DATABASE_MONGODB = "mongodb"
|
|
528
|
+
VECTOR_DATABASE_TURBOPUFFER = "turbopuffer"
|
|
529
|
+
VECTOR_DATABASE_PINECONE = "pinecone"
|
|
530
|
+
|
|
531
|
+
|
|
532
|
+
@dataclass
|
|
533
|
+
class PineconeStorage:
|
|
534
|
+
"""
|
|
535
|
+
Configuration for storing embeddings in Pinecone.
|
|
536
|
+
|
|
537
|
+
Attributes:
|
|
538
|
+
index_name: Pinecone index name
|
|
539
|
+
namespace: Namespace within the index (optional)
|
|
540
|
+
id_field: Document field to use as the vector ID
|
|
541
|
+
metadata: List of document fields to include as Pinecone metadata
|
|
542
|
+
"""
|
|
543
|
+
index_name: str
|
|
544
|
+
id_field: str
|
|
545
|
+
namespace: Optional[str] = None
|
|
546
|
+
metadata: Optional[list[str]] = None
|
|
547
|
+
|
|
548
|
+
def to_dict(self) -> dict[str, Any]:
|
|
549
|
+
"""Convert to dictionary for JSON serialization."""
|
|
550
|
+
result = {
|
|
551
|
+
"indexName": self.index_name,
|
|
552
|
+
"idField": self.id_field,
|
|
553
|
+
}
|
|
554
|
+
if self.namespace:
|
|
555
|
+
result["namespace"] = self.namespace
|
|
556
|
+
if self.metadata:
|
|
557
|
+
result["metadata"] = self.metadata
|
|
558
|
+
return result
|
|
559
|
+
|
|
560
|
+
|
|
561
|
+
@dataclass
|
|
562
|
+
class QueryConfig:
|
|
563
|
+
"""
|
|
564
|
+
Configuration for query search parameters.
|
|
565
|
+
|
|
566
|
+
Attributes:
|
|
567
|
+
top_k: Number of results to return (default: 10)
|
|
568
|
+
min_score: Minimum similarity score threshold (0.0 to 1.0)
|
|
569
|
+
filters: Metadata filters for filtering results
|
|
570
|
+
namespace: Namespace for Pinecone/TurboPuffer
|
|
571
|
+
collection: Collection name for MongoDB
|
|
572
|
+
database: Database name for MongoDB
|
|
573
|
+
include_vectors: Whether to include vector values in response
|
|
574
|
+
include_metadata: Whether to include metadata in response
|
|
575
|
+
"""
|
|
576
|
+
top_k: int = 10
|
|
577
|
+
min_score: Optional[float] = None
|
|
578
|
+
filters: Optional[dict[str, str]] = None
|
|
579
|
+
namespace: Optional[str] = None
|
|
580
|
+
collection: Optional[str] = None
|
|
581
|
+
database: Optional[str] = None
|
|
582
|
+
include_vectors: bool = False
|
|
583
|
+
include_metadata: bool = True
|
|
584
|
+
|
|
585
|
+
def to_dict(self) -> dict[str, Any]:
|
|
586
|
+
"""Convert to dictionary for JSON serialization."""
|
|
587
|
+
result = {
|
|
588
|
+
"topK": self.top_k,
|
|
589
|
+
"includeVectors": self.include_vectors,
|
|
590
|
+
"includeMetadata": self.include_metadata,
|
|
591
|
+
}
|
|
592
|
+
if self.min_score is not None:
|
|
593
|
+
result["minScore"] = self.min_score
|
|
594
|
+
if self.filters:
|
|
595
|
+
result["filters"] = self.filters
|
|
596
|
+
if self.namespace:
|
|
597
|
+
result["namespace"] = self.namespace
|
|
598
|
+
if self.collection:
|
|
599
|
+
result["collection"] = self.collection
|
|
600
|
+
if self.database:
|
|
601
|
+
result["database"] = self.database
|
|
602
|
+
return result
|
|
603
|
+
|
|
604
|
+
|
|
605
|
+
@dataclass
|
|
606
|
+
class QueryRequest:
|
|
607
|
+
"""
|
|
608
|
+
A request to perform vector search.
|
|
609
|
+
|
|
610
|
+
Attributes:
|
|
611
|
+
request_id: Unique identifier for tracking
|
|
612
|
+
query_text: The text to search for (will be embedded)
|
|
613
|
+
database: Which vector database to search
|
|
614
|
+
priority: Queue priority
|
|
615
|
+
query_config: Query-specific configuration
|
|
616
|
+
embedding_config: Optional embedding configuration override
|
|
617
|
+
metadata: Arbitrary key-value pairs for tracking
|
|
618
|
+
created_at: When the request was created
|
|
619
|
+
"""
|
|
620
|
+
request_id: str
|
|
621
|
+
query_text: str
|
|
622
|
+
database: str
|
|
623
|
+
priority: str
|
|
624
|
+
query_config: QueryConfig
|
|
625
|
+
embedding_config: Optional[EmbeddingConfigOverride] = None
|
|
626
|
+
metadata: dict[str, str] = field(default_factory=dict)
|
|
627
|
+
created_at: datetime = field(default_factory=datetime.utcnow)
|
|
628
|
+
|
|
629
|
+
def to_dict(self) -> dict[str, Any]:
|
|
630
|
+
"""Convert to dictionary for JSON serialization."""
|
|
631
|
+
result = {
|
|
632
|
+
"requestId": self.request_id,
|
|
633
|
+
"queryText": self.query_text,
|
|
634
|
+
"database": self.database,
|
|
635
|
+
"priority": self.priority,
|
|
636
|
+
"queryConfig": self.query_config.to_dict(),
|
|
637
|
+
"metadata": self.metadata,
|
|
638
|
+
"createdAt": self.created_at.isoformat() + "Z",
|
|
639
|
+
}
|
|
640
|
+
if self.embedding_config:
|
|
641
|
+
result["embeddingConfig"] = self.embedding_config.to_dict()
|
|
642
|
+
return result
|
|
643
|
+
|
|
644
|
+
|
|
645
|
+
@dataclass
|
|
646
|
+
class VectorMatch:
|
|
647
|
+
"""
|
|
648
|
+
A single vector search result.
|
|
649
|
+
|
|
650
|
+
Attributes:
|
|
651
|
+
id: Vector ID in the database
|
|
652
|
+
score: Similarity score (higher is more similar)
|
|
653
|
+
metadata: Associated metadata
|
|
654
|
+
vector: The vector values (if requested)
|
|
655
|
+
"""
|
|
656
|
+
id: str
|
|
657
|
+
score: float
|
|
658
|
+
metadata: Optional[dict[str, Any]] = None
|
|
659
|
+
vector: Optional[list[float]] = None
|
|
660
|
+
|
|
661
|
+
@classmethod
|
|
662
|
+
def from_dict(cls, data: dict[str, Any]) -> "VectorMatch":
|
|
663
|
+
"""Create from dictionary."""
|
|
664
|
+
return cls(
|
|
665
|
+
id=data["id"],
|
|
666
|
+
score=data["score"],
|
|
667
|
+
metadata=data.get("metadata"),
|
|
668
|
+
vector=data.get("vector"),
|
|
669
|
+
)
|
|
670
|
+
|
|
671
|
+
|
|
672
|
+
@dataclass
|
|
673
|
+
class QueryTiming:
|
|
674
|
+
"""
|
|
675
|
+
Processing duration breakdown for queries.
|
|
676
|
+
|
|
677
|
+
Attributes:
|
|
678
|
+
queue_wait_ms: Time spent waiting in queue
|
|
679
|
+
embedding_ms: Time spent generating query embedding
|
|
680
|
+
search_ms: Time spent executing database search
|
|
681
|
+
total_ms: Total processing time
|
|
682
|
+
"""
|
|
683
|
+
queue_wait_ms: int
|
|
684
|
+
embedding_ms: int
|
|
685
|
+
search_ms: int
|
|
686
|
+
total_ms: int
|
|
687
|
+
|
|
688
|
+
@classmethod
|
|
689
|
+
def from_dict(cls, data: dict[str, Any]) -> "QueryTiming":
|
|
690
|
+
"""Create from dictionary."""
|
|
691
|
+
return cls(
|
|
692
|
+
queue_wait_ms=data.get("queueWaitMs", 0),
|
|
693
|
+
embedding_ms=data.get("embeddingMs", 0),
|
|
694
|
+
search_ms=data.get("searchMs", 0),
|
|
695
|
+
total_ms=data.get("totalMs", 0),
|
|
696
|
+
)
|
|
697
|
+
|
|
698
|
+
|
|
699
|
+
@dataclass
|
|
700
|
+
class QueryResult:
|
|
701
|
+
"""
|
|
702
|
+
Result of a vector search query.
|
|
703
|
+
|
|
704
|
+
Attributes:
|
|
705
|
+
request_id: The original request ID
|
|
706
|
+
status: "success" or "failed"
|
|
707
|
+
matches: Matching vectors with scores
|
|
708
|
+
error: Error message if status is "failed"
|
|
709
|
+
timing: Processing duration breakdown
|
|
710
|
+
completed_at: When processing finished
|
|
711
|
+
"""
|
|
712
|
+
request_id: str
|
|
713
|
+
status: str
|
|
714
|
+
matches: list[VectorMatch]
|
|
715
|
+
error: Optional[str]
|
|
716
|
+
timing: Optional[QueryTiming]
|
|
717
|
+
completed_at: datetime
|
|
718
|
+
|
|
719
|
+
@classmethod
|
|
720
|
+
def from_dict(cls, data: dict[str, Any]) -> "QueryResult":
|
|
721
|
+
"""Create from dictionary."""
|
|
722
|
+
matches = [VectorMatch.from_dict(m) for m in data.get("matches", [])]
|
|
723
|
+
timing = None
|
|
724
|
+
if data.get("timing"):
|
|
725
|
+
timing = QueryTiming.from_dict(data["timing"])
|
|
726
|
+
|
|
727
|
+
return cls(
|
|
728
|
+
request_id=data["requestId"],
|
|
729
|
+
status=data["status"],
|
|
730
|
+
matches=matches,
|
|
731
|
+
error=data.get("error"),
|
|
732
|
+
timing=timing,
|
|
733
|
+
completed_at=datetime.fromisoformat(data["completedAt"].replace("Z", "+00:00")),
|
|
734
|
+
)
|
|
735
|
+
|
|
736
|
+
@property
|
|
737
|
+
def is_success(self) -> bool:
|
|
738
|
+
"""Check if the query was successful."""
|
|
739
|
+
return self.status == "success"
|
|
740
|
+
|
|
741
|
+
@property
|
|
742
|
+
def is_failed(self) -> bool:
|
|
743
|
+
"""Check if the query failed."""
|
|
744
|
+
return self.status == "failed"
|
|
745
|
+
|
|
746
|
+
|
|
747
|
+
# ============================================================================
|
|
748
|
+
# Database Lookup Types (HTTP API)
|
|
749
|
+
# ============================================================================
|
|
750
|
+
|
|
751
|
+
|
|
752
|
+
@dataclass
|
|
753
|
+
class Document:
|
|
754
|
+
"""
|
|
755
|
+
A document retrieved from the database.
|
|
756
|
+
|
|
757
|
+
Attributes:
|
|
758
|
+
id: Document/vector ID
|
|
759
|
+
metadata: Document metadata
|
|
760
|
+
vector: Vector values (if requested)
|
|
761
|
+
"""
|
|
762
|
+
id: str
|
|
763
|
+
metadata: Optional[dict[str, Any]] = None
|
|
764
|
+
vector: Optional[list[float]] = None
|
|
765
|
+
|
|
766
|
+
@classmethod
|
|
767
|
+
def from_dict(cls, data: dict[str, Any]) -> "Document":
|
|
768
|
+
"""Create from dictionary."""
|
|
769
|
+
return cls(
|
|
770
|
+
id=data["id"],
|
|
771
|
+
metadata=data.get("metadata"),
|
|
772
|
+
vector=data.get("vector"),
|
|
773
|
+
)
|
|
774
|
+
|
|
775
|
+
|
|
776
|
+
@dataclass
|
|
777
|
+
class LookupTiming:
|
|
778
|
+
"""
|
|
779
|
+
Timing information for lookup operations.
|
|
780
|
+
|
|
781
|
+
Attributes:
|
|
782
|
+
total_ms: Total request duration in milliseconds
|
|
783
|
+
"""
|
|
784
|
+
total_ms: int
|
|
785
|
+
|
|
786
|
+
@classmethod
|
|
787
|
+
def from_dict(cls, data: dict[str, Any]) -> "LookupTiming":
|
|
788
|
+
"""Create from dictionary."""
|
|
789
|
+
return cls(total_ms=data.get("totalMs", 0))
|
|
790
|
+
|
|
791
|
+
|
|
792
|
+
@dataclass
|
|
793
|
+
class LookupResult:
|
|
794
|
+
"""
|
|
795
|
+
Result of a lookup or metadata search operation.
|
|
796
|
+
|
|
797
|
+
Attributes:
|
|
798
|
+
documents: Retrieved documents
|
|
799
|
+
timing: Timing information
|
|
800
|
+
"""
|
|
801
|
+
documents: list[Document]
|
|
802
|
+
timing: LookupTiming
|
|
803
|
+
|
|
804
|
+
@classmethod
|
|
805
|
+
def from_dict(cls, data: dict[str, Any]) -> "LookupResult":
|
|
806
|
+
"""Create from dictionary."""
|
|
807
|
+
documents = [Document.from_dict(d) for d in data.get("documents", [])]
|
|
808
|
+
timing = LookupTiming.from_dict(data.get("timing", {}))
|
|
809
|
+
return cls(documents=documents, timing=timing)
|
|
810
|
+
|
|
811
|
+
|
|
812
|
+
# ============================================================================
|
|
813
|
+
# Clone and Delete Types
|
|
814
|
+
# ============================================================================
|
|
815
|
+
|
|
816
|
+
|
|
817
|
+
@dataclass
|
|
818
|
+
class CloneResult:
|
|
819
|
+
"""
|
|
820
|
+
Result of a clone operation.
|
|
821
|
+
|
|
822
|
+
Attributes:
|
|
823
|
+
id: Document ID that was cloned
|
|
824
|
+
success: Whether the clone succeeded
|
|
825
|
+
timing: Timing information
|
|
826
|
+
"""
|
|
827
|
+
id: str
|
|
828
|
+
success: bool
|
|
829
|
+
timing: LookupTiming
|
|
830
|
+
|
|
831
|
+
@classmethod
|
|
832
|
+
def from_dict(cls, data: dict[str, Any]) -> "CloneResult":
|
|
833
|
+
"""Create from dictionary."""
|
|
834
|
+
timing = LookupTiming.from_dict(data.get("timing", {}))
|
|
835
|
+
return cls(
|
|
836
|
+
id=data["id"],
|
|
837
|
+
success=data["success"],
|
|
838
|
+
timing=timing,
|
|
839
|
+
)
|
|
840
|
+
|
|
841
|
+
|
|
842
|
+
@dataclass
|
|
843
|
+
class DeleteFromNamespaceResult:
|
|
844
|
+
"""
|
|
845
|
+
Result of a delete operation.
|
|
846
|
+
|
|
847
|
+
Attributes:
|
|
848
|
+
id: Document ID that was deleted
|
|
849
|
+
success: Whether the delete succeeded
|
|
850
|
+
timing: Timing information
|
|
851
|
+
"""
|
|
852
|
+
id: str
|
|
853
|
+
success: bool
|
|
854
|
+
timing: LookupTiming
|
|
855
|
+
|
|
856
|
+
@classmethod
|
|
857
|
+
def from_dict(cls, data: dict[str, Any]) -> "DeleteFromNamespaceResult":
|
|
858
|
+
"""Create from dictionary."""
|
|
859
|
+
timing = LookupTiming.from_dict(data.get("timing", {}))
|
|
860
|
+
return cls(
|
|
861
|
+
id=data["id"],
|
|
862
|
+
success=data["success"],
|
|
863
|
+
timing=timing,
|
|
864
|
+
)
|