Coverage for src/qdrant_loader/config/embedding.py: 100%

1"""Configuration for embedding generation.""" 

2 

3from pydantic import Field 

4 

5from qdrant_loader.config.base import BaseConfig 

6 

7 

8class EmbeddingConfig(BaseConfig): 

9 """Configuration for embedding generation.""" 

10 

11 model: str = Field( 

12 default="text-embedding-3-small", description="OpenAI embedding model to use" 

13 ) 

14 api_key: str | None = Field( 

15 default=None, description="API key for the embedding service" 

16 ) 

17 batch_size: int = Field( 

18 default=100, description="Number of texts to embed in a single batch" 

19 ) 

20 endpoint: str = Field( 

21 default="https://api.openai.com/v1", 

22 description="Base URL for the embedding API endpoint", 

23 ) 

24 tokenizer: str = Field( 

25 default="cl100k_base", # Default OpenAI tokenizer 

26 description="Tokenizer to use for token counting. Use 'cl100k_base' for OpenAI models or 'none' for other models", 

27 ) 

28 vector_size: int | None = Field( 

29 default=1536, 

30 description="Vector size for the embedding model (384 for BAAI/bge-small-en-v1.5, 1536 for OpenAI models)", 

31 ) 

32 max_tokens_per_request: int = Field( 

33 default=8000, 

34 description="Maximum total tokens allowed per embedding API request (leave buffer below model limit)", 

35 ) 

36 max_tokens_per_chunk: int = Field( 

37 default=8000, 

38 description="Maximum tokens allowed for a single chunk (should match or be below model's context limit)", 

39 )