File size: 1,528 Bytes
2cb9dec
 
 
 
 
 
 
 
 
 
bc82930
 
2cb9dec
 
 
192ee60
 
 
 
6f4f307
 
 
 
 
 
 
 
 
 
2cb9dec
6f4f307
 
 
 
 
 
 
 
2cb9dec
 
 
 
6f4f307
 
 
 
 
 
5c144a1
6f4f307
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
from pydantic import BaseModel
from typing import List, Dict


# Pydantic models for request validation
class CreateEmbeddingRequest(BaseModel):
    query: str
    target_column: str = "product_type"
    output_column: str = "embedding"
    model: str = "text-embedding-3-small"
    batch_size: int = 10
    max_concurrent_requests: int = 10
    dataset_name: str = "re-mind/product_type_embedding"


class ReadEmbeddingRequest(BaseModel):
    dataset_name: str


# class UpdateEmbeddingRequest(BaseModel):
#     updates: Dict[str, List]  # Column name -> List of values
#     target_column: str = "product_type"
#     output_column: str = "embedding"
#     model: str = "text-embedding-3-small"
#     batch_size: int = 10
#     max_concurrent_requests: int = 10
#     dataset_name: str = "re-mind/product_type_embedding"


class UpdateEmbeddingRequest(BaseModel):
    dataset_name: str = "re-mind/product_type_embedding"
    updates: Dict[
        str, List
    ]  # Dictionary of column names and their corresponding values
    target_column: str = (
        "product_type"  # Column in the new data to generate embeddings for
    )
    output_column: str = "embedding"  # Column to store the generated embeddings


class DeleteEmbeddingRequest(BaseModel):
    dataset_name: str


# Request model for the /embed endpoint
class EmbedRequest(BaseModel):
    texts: List[str]  # List of strings to generate embeddings for
    output_column: str = (
        "embedding"  # Column to store embeddings (default: "embeddings")
    )