0
# RAG and Embeddings
1

2
Retrieval-augmented generation with text chunking, embeddings, and vector store integration. Supports document extraction, chunking strategies, embedding models, and vector database storage for semantic search.
3

4
## Capabilities
5

6
### RAG Configuration
7

8
Configuration for Retrieval-Augmented Generation pipelines.
9

10
```python { .api }
11
from kiln_ai.datamodel import RagConfig
12

13
class RagConfig:
14
    """
15
    Configuration for RAG (Retrieval-Augmented Generation).
16

17
    Properties:
18
    - vector_store_config (VectorStoreConfig): Vector database configuration
19
    - embedding_config (EmbeddingConfig): Embedding model configuration
20
    - chunker_config (ChunkerConfig): Text chunking configuration
21
    - top_k (int): Number of results to retrieve
22
    """
23
```
24

25
### Embedding Configuration
26

27
Configuration for embedding models and generation.
28

29
```python { .api }
30
from kiln_ai.datamodel import EmbeddingConfig, Embedding, ChunkEmbeddings
31

32
class EmbeddingConfig:
33
    """
34
    Configuration for embeddings.
35

36
    Properties:
37
    - model_id (str): Embedding model identifier
38
    - provider (str): Embedding provider name
39
    - dimensions (int): Embedding vector dimensions
40
    """
41

42
class Embedding:
43
    """
44
    Single embedding vector.
45

46
    Properties:
47
    - vector (list[float]): Embedding vector values
48
    - metadata (dict): Additional embedding metadata
49
    """
50

51
class ChunkEmbeddings:
52
    """
53
    Embeddings for document chunks.
54

55
    Properties:
56
    - embeddings (list[Embedding]): List of embedding vectors
57
    - chunk_ids (list[str]): Corresponding chunk identifiers
58
    """
59
```
60

61
### Embedding Adapters
62

63
Adapters for generating embeddings with various providers.
64

65
```python { .api }
66
from kiln_ai.adapters.embedding import (
67
    BaseEmbeddingAdapter,
68
    LitellmEmbeddingAdapter,
69
    EmbeddingOptions
70
)
71

72
class BaseEmbeddingAdapter:
73
    """
74
    Abstract embedding adapter interface.
75

76
    Methods:
77
    - embed(): Generate single embedding
78
    - embed_batch(): Generate batch embeddings
79
    """
80

81
    async def embed(self, text: str) -> 'Embedding':
82
        """
83
        Generate embedding for text.
84

85
        Parameters:
86
        - text (str): Input text
87

88
        Returns:
89
        Embedding: Embedding vector with metadata
90
        """
91

92
    async def embed_batch(self, texts: list) -> 'EmbeddingResult':
93
        """
94
        Generate embeddings for multiple texts.
95

96
        Parameters:
97
        - texts (list[str]): Input texts
98

99
        Returns:
100
        EmbeddingResult: Batch embeddings with usage info
101
        """
102

103
class LitellmEmbeddingAdapter(BaseEmbeddingAdapter):
104
    """
105
    LiteLLM embedding adapter supporting multiple providers.
106

107
    Supports:
108
    - OpenAI (text-embedding-3-small, text-embedding-3-large)
109
    - Cohere (embed-english-v3.0, embed-multilingual-v3.0)
110
    - Voyage AI (voyage-large-2, voyage-code-2)
111
    - Many more through LiteLLM
112
    """
113

114
    def __init__(self, model_name: str, provider: str, options: 'EmbeddingOptions' = None):
115
        """
116
        Initialize LiteLLM embedding adapter.
117

118
        Parameters:
119
        - model_name (str): Embedding model identifier
120
        - provider (str): Provider name
121
        - options (EmbeddingOptions | None): Embedding options
122
        """
123

124
class EmbeddingOptions:
125
    """
126
    Embedding configuration options.
127

128
    Properties:
129
    - dimensions (int | None): Vector dimensions (if configurable)
130
    - encoding_format (str | None): Encoding format (e.g., "float", "base64")
131
    """
132

133
class EmbeddingResult:
134
    """
135
    Batch embedding result.
136

137
    Properties:
138
    - embeddings (list[Embedding]): Generated embeddings
139
    - usage (dict): Token usage information
140
    """
141
```
142

143
### Embedding Registry
144

145
Get embedding adapters by model and provider.
146

147
```python { .api }
148
from kiln_ai.adapters.embedding.embedding_registry import embedding_adapter_from_type
149

150
def embedding_adapter_from_type(model_name: str, provider: str):
151
    """
152
    Get embedding adapter instance.
153

154
    Parameters:
155
    - model_name (str): Embedding model identifier
156
    - provider (str): Provider name
157

158
    Returns:
159
    BaseEmbeddingAdapter: Embedding adapter instance
160
    """
161
```
162

163
### Text Chunking
164

165
Configuration and adapters for text chunking strategies.
166

167
```python { .api }
168
from kiln_ai.datamodel import ChunkerConfig, ChunkerType, Chunk, ChunkedDocument
169

170
class ChunkerConfig:
171
    """
172
    Configuration for text chunking.
173

174
    Properties:
175
    - chunker_type (ChunkerType): Type of chunker to use
176
    - chunk_size (int): Size of each chunk in characters
177
    - chunk_overlap (int): Overlap between chunks in characters
178
    """
179

180
class ChunkerType:
181
    """
182
    Available chunker types.
183

184
    Values:
185
    - fixed_window: Fixed-size window chunking with overlap
186
    """
187
    fixed_window = "fixed_window"
188

189
class Chunk:
190
    """
191
    Single text chunk with metadata.
192

193
    Properties:
194
    - text (str): Chunk content
195
    - start_index (int): Start position in source document
196
    - end_index (int): End position in source document
197
    - metadata (dict): Additional chunk metadata (source, page, etc.)
198
    """
199

200
class ChunkedDocument:
201
    """
202
    Document split into chunks.
203

204
    Properties:
205
    - chunks (list[Chunk]): List of text chunks
206
    - source_document (str): Original document content
207
    """
208
```
209

210
### Chunking Adapters
211

212
Adapters for different chunking strategies.
213

214
```python { .api }
215
from kiln_ai.adapters.chunkers import (
216
    BaseChunker,
217
    FixedWindowChunker,
218
    TextChunk,
219
    ChunkingResult
220
)
221

222
class BaseChunker:
223
    """
224
    Abstract chunker interface.
225

226
    Methods:
227
    - chunk(): Chunk single document
228
    - chunk_documents(): Chunk multiple documents
229
    """
230

231
    def chunk(self, text: str) -> 'ChunkingResult':
232
        """
233
        Chunk single document.
234

235
        Parameters:
236
        - text (str): Document text
237

238
        Returns:
239
        ChunkingResult: Chunking result with metadata
240
        """
241

242
    def chunk_documents(self, documents: list) -> list:
243
        """
244
        Chunk multiple documents.
245

246
        Parameters:
247
        - documents (list[str]): Document texts
248

249
        Returns:
250
        list[ChunkingResult]: Chunking results
251
        """
252

253
class FixedWindowChunker(BaseChunker):
254
    """
255
    Fixed-size window chunking with overlap.
256

257
    Splits text into chunks of fixed size with configurable overlap
258
    to maintain context across chunk boundaries.
259
    """
260

261
    def __init__(self, chunk_size: int, chunk_overlap: int):
262
        """
263
        Initialize fixed window chunker.
264

265
        Parameters:
266
        - chunk_size (int): Size of each chunk in characters
267
        - chunk_overlap (int): Overlap between chunks in characters
268
        """
269

270
class TextChunk:
271
    """
272
    Text chunk with positional metadata.
273

274
    Properties:
275
    - text (str): Chunk text
276
    - start (int): Start position in document
277
    - end (int): End position in document
278
    - metadata (dict): Additional metadata
279
    """
280

281
class ChunkingResult:
282
    """
283
    Result of chunking operation.
284

285
    Properties:
286
    - chunks (list[TextChunk]): Text chunks
287
    - metadata (dict): Chunking metadata (strategy, params)
288
    """
289
```
290

291
### Chunker Registry
292

293
Get chunker adapters by type.
294

295
```python { .api }
296
from kiln_ai.adapters.chunkers.chunker_registry import chunker_adapter_from_type
297

298
def chunker_adapter_from_type(chunker_type: str, config: dict):
299
    """
300
    Get chunker adapter from type.
301

302
    Parameters:
303
    - chunker_type (str): Type of chunker
304
    - config (dict): Chunker configuration
305

306
    Returns:
307
    BaseChunker: Chunker adapter instance
308
    """
309
```
310

311
### Document Extraction
312

313
Extract and process documents for RAG pipelines.
314

315
```python { .api }
316
from kiln_ai.datamodel import (
317
    Document,
318
    Extraction,
319
    ExtractorConfig,
320
    FileInfo,
321
    Kind,
322
    OutputFormat,
323
    ExtractorType,
324
    ExtractionSource
325
)
326

327
class Document:
328
    """
329
    Document with extracted content.
330

331
    Properties:
332
    - id (str): Unique identifier
333
    - content (str): Extracted content
334
    - metadata (dict): Document metadata
335
    - kind (Kind): Type of document (text, pdf, image, html)
336
    """
337

338
    @staticmethod
339
    def load_from_file(path: str) -> 'Document':
340
        """Load document from .kiln file."""
341

342
    def save_to_file(self) -> None:
343
        """Save document to .kiln file."""
344

345
class Extraction:
346
    """
347
    Result of document extraction.
348

349
    Properties:
350
    - document (Document): Extracted document
351
    - extractor_config (ExtractorConfig): Configuration used
352
    """
353

354
class ExtractorConfig:
355
    """
356
    Configuration for document extraction.
357

358
    Properties:
359
    - extractor_type (ExtractorType): Type of extractor
360
    - options (dict): Extractor-specific options
361
    """
362

363
class FileInfo:
364
    """
365
    Metadata about source file.
366

367
    Properties:
368
    - filename (str): Name of file
369
    - path (str): File system path
370
    - size (int): File size in bytes
371
    - mime_type (str): MIME type
372
    """
373

374
class Kind:
375
    """
376
    Type of document.
377

378
    Values:
379
    - text: Plain text document
380
    - pdf: PDF document
381
    - image: Image file
382
    - html: HTML document
383
    """
384
    text = "text"
385
    pdf = "pdf"
386
    image = "image"
387
    html = "html"
388

389
class OutputFormat:
390
    """
391
    Format for extracted output.
392

393
    Values:
394
    - markdown: Markdown format
395
    - plain_text: Plain text format
396
    - structured: Structured data format
397
    """
398
    markdown = "markdown"
399
    plain_text = "plain_text"
400
    structured = "structured"
401

402
class ExtractorType:
403
    """
404
    Type of extractor to use.
405

406
    Values:
407
    - litellm: LiteLLM-based extraction
408
    - custom: Custom extractor
409
    """
410
    litellm = "litellm"
411
    custom = "custom"
412

413
class ExtractionSource:
414
    """
415
    Source type for extraction.
416

417
    Values:
418
    - file: Extract from file
419
    - url: Extract from URL
420
    - text: Extract from text
421
    """
422
    file = "file"
423
    url = "url"
424
    text = "text"
425
```
426

427
### Extraction Adapters
428

429
Adapters for extracting content from various sources.
430

431
```python { .api }
432
from kiln_ai.adapters.extractors import (
433
    BaseExtractor,
434
    LitellmExtractor,
435
    ExtractionInput,
436
    ExtractionOutput,
437
    encode_file_litellm_format
438
)
439

440
class BaseExtractor:
441
    """
442
    Abstract extractor interface.
443

444
    Methods:
445
    - extract(): Extract single document
446
    - extract_batch(): Extract multiple documents
447
    """
448

449
    async def extract(self, input_data: 'ExtractionInput') -> 'ExtractionOutput':
450
        """
451
        Extract content from source.
452

453
        Parameters:
454
        - input_data (ExtractionInput): Input specification
455

456
        Returns:
457
        ExtractionOutput: Extracted content
458
        """
459

460
    async def extract_batch(self, inputs: list) -> list:
461
        """
462
        Extract from multiple sources.
463

464
        Parameters:
465
        - inputs (list[ExtractionInput]): Input specifications
466

467
        Returns:
468
        list[ExtractionOutput]: Extracted contents
469
        """
470

471
class LitellmExtractor(BaseExtractor):
472
    """
473
    LiteLLM-based document extraction.
474

475
    Uses vision-capable models to extract text from:
476
    - PDFs
477
    - Images
478
    - Scanned documents
479
    """
480

481
    def __init__(self, model_name: str, provider: str):
482
        """
483
        Initialize LiteLLM extractor.
484

485
        Parameters:
486
        - model_name (str): Vision-capable model
487
        - provider (str): Provider name
488
        """
489

490
class ExtractionInput:
491
    """
492
    Input for extraction.
493

494
    Properties:
495
    - content (str): Source content (text, file path, URL)
496
    - content_type (str): Type of content
497
    - options (dict): Extraction options
498
    """
499

500
class ExtractionOutput:
501
    """
502
    Extracted content.
503

504
    Properties:
505
    - text (str): Extracted text
506
    - metadata (dict): Extraction metadata
507
    - format (OutputFormat): Output format
508
    """
509

510
def encode_file_litellm_format(file_path: str, mime_type: str) -> str:
511
    """
512
    Encode file for LiteLLM API.
513

514
    Parameters:
515
    - file_path (str): Path to file
516
    - mime_type (str): MIME type of file
517

518
    Returns:
519
    str: Base64 encoded file data
520
    """
521
```
522

523
### Extraction Runner
524

525
Execute extraction jobs.
526

527
```python { .api }
528
from kiln_ai.adapters.extractors import ExtractorRunner, ExtractorJob
529

530
class ExtractorRunner:
531
    """
532
    Execute document extractions.
533

534
    Methods:
535
    - run(): Execute single extraction
536
    - run_batch(): Execute batch extractions
537
    """
538

539
    def __init__(self, config: 'ExtractorConfig'):
540
        """
541
        Initialize extraction runner.
542

543
        Parameters:
544
        - config (ExtractorConfig): Extractor configuration
545
        """
546

547
    async def run(self, input_data: 'ExtractionInput') -> 'ExtractionOutput':
548
        """Execute single extraction."""
549

550
    async def run_batch(self, inputs: list) -> list:
551
        """Execute batch extractions."""
552

553
class ExtractorJob:
554
    """
555
    Extraction job configuration.
556

557
    Properties:
558
    - extractor_config (ExtractorConfig): Extractor settings
559
    - inputs (list[ExtractionInput]): Inputs to process
560
    """
561
```
562

563
### Extractor Registry
564

565
Get extractor adapters by type.
566

567
```python { .api }
568
from kiln_ai.adapters.extractors.extractor_registry import extractor_adapter_from_type
569

570
def extractor_adapter_from_type(extractor_type: str, config: dict):
571
    """
572
    Get extractor adapter from type.
573

574
    Parameters:
575
    - extractor_type (str): Type of extractor
576
    - config (dict): Extractor configuration
577

578
    Returns:
579
    BaseExtractor: Extractor adapter instance
580
    """
581
```
582

583
### Vector Store Configuration
584

585
Configuration for vector database integration.
586

587
```python { .api }
588
from kiln_ai.datamodel import VectorStoreConfig, VectorStoreType, LanceDBConfigBaseProperties
589

590
class VectorStoreConfig:
591
    """
592
    Configuration for vector database.
593

594
    Properties:
595
    - vector_store_type (VectorStoreType): Type of vector store
596
    - connection_params (dict): Connection parameters
597
    """
598

599
class VectorStoreType:
600
    """
601
    Type of vector store.
602

603
    Values:
604
    - lancedb: LanceDB vector database
605
    """
606
    lancedb = "lancedb"
607

608
class LanceDBConfigBaseProperties:
609
    """
610
    LanceDB-specific configuration.
611

612
    Properties:
613
    - uri (str): Database URI (file path or connection string)
614
    - table_name (str): Table name for storage
615
    """
616
```
617

618
## Usage Examples
619

620
### Basic Embedding Generation
621

622
```python
623
from kiln_ai.adapters.embedding import LitellmEmbeddingAdapter
624

625
# Create embedding adapter
626
adapter = LitellmEmbeddingAdapter(
627
    model_name="text-embedding-3-small",
628
    provider="openai"
629
)
630

631
# Generate single embedding
632
text = "This is a sample document for embedding."
633
embedding = await adapter.embed(text)
634
print(f"Embedding dimensions: {len(embedding.vector)}")
635
print(f"First few values: {embedding.vector[:5]}")
636

637
# Generate batch embeddings
638
texts = [
639
    "First document",
640
    "Second document",
641
    "Third document"
642
]
643
result = await adapter.embed_batch(texts)
644
print(f"Generated {len(result.embeddings)} embeddings")
645
print(f"Tokens used: {result.usage}")
646
```
647

648
### Text Chunking
649

650
```python
651
from kiln_ai.adapters.chunkers import FixedWindowChunker
652

653
# Create chunker with 500 character chunks and 50 character overlap
654
chunker = FixedWindowChunker(chunk_size=500, chunk_overlap=50)
655

656
# Load document
657
with open("long_document.txt", "r") as f:
658
    document_text = f.read()
659

660
# Chunk document
661
result = chunker.chunk(document_text)
662
print(f"Created {len(result.chunks)} chunks")
663

664
for i, chunk in enumerate(result.chunks):
665
    print(f"\nChunk {i+1}:")
666
    print(f"  Position: {chunk.start}-{chunk.end}")
667
    print(f"  Length: {len(chunk.text)} characters")
668
    print(f"  Preview: {chunk.text[:100]}...")
669
```
670

671
### Document Extraction
672

673
```python
674
from kiln_ai.datamodel import ExtractorConfig, ExtractorType, ExtractionSource
675
from kiln_ai.adapters.extractors import ExtractorRunner, ExtractionInput
676

677
# Configure extractor
678
config = ExtractorConfig(
679
    extractor_type=ExtractorType.litellm,
680
    options={
681
        "model": "gpt-4o",
682
        "provider": "openai"
683
    }
684
)
685

686
# Create extraction runner
687
runner = ExtractorRunner(config)
688

689
# Extract from PDF
690
pdf_input = ExtractionInput(
691
    content="/path/to/document.pdf",
692
    content_type="application/pdf",
693
    options={"output_format": "markdown"}
694
)
695

696
extraction = await runner.run(pdf_input)
697
print(f"Extracted text length: {len(extraction.text)}")
698
print(f"Output format: {extraction.format}")
699
```
700

701
### Complete RAG Pipeline
702

703
```python
704
from kiln_ai.datamodel import (
705
    RagConfig,
706
    EmbeddingConfig,
707
    ChunkerConfig,
708
    ChunkerType,
709
    VectorStoreConfig,
710
    VectorStoreType
711
)
712
from kiln_ai.adapters.chunkers import FixedWindowChunker
713
from kiln_ai.adapters.embedding import LitellmEmbeddingAdapter
714

715
# Configure RAG pipeline
716
rag_config = RagConfig(
717
    embedding_config=EmbeddingConfig(
718
        model_id="text-embedding-3-small",
719
        provider="openai",
720
        dimensions=1536
721
    ),
722
    chunker_config=ChunkerConfig(
723
        chunker_type=ChunkerType.fixed_window,
724
        chunk_size=500,
725
        chunk_overlap=50
726
    ),
727
    vector_store_config=VectorStoreConfig(
728
        vector_store_type=VectorStoreType.lancedb,
729
        connection_params={
730
            "uri": "/path/to/lancedb",
731
            "table_name": "documents"
732
        }
733
    ),
734
    top_k=5
735
)
736

737
# 1. Chunk documents
738
chunker = FixedWindowChunker(
739
    chunk_size=rag_config.chunker_config.chunk_size,
740
    chunk_overlap=rag_config.chunker_config.chunk_overlap
741
)
742

743
document_text = "Long document content..."
744
chunking_result = chunker.chunk(document_text)
745
chunks = [c.text for c in chunking_result.chunks]
746
print(f"Created {len(chunks)} chunks")
747

748
# 2. Generate embeddings
749
embedding_adapter = LitellmEmbeddingAdapter(
750
    model_name=rag_config.embedding_config.model_id,
751
    provider=rag_config.embedding_config.provider
752
)
753

754
embeddings_result = await embedding_adapter.embed_batch(chunks)
755
embeddings = embeddings_result.embeddings
756
print(f"Generated {len(embeddings)} embeddings")
757

758
# 3. Store in vector database (pseudocode - actual implementation varies)
759
# vector_store.add_documents(chunks, embeddings)
760

761
# 4. Query for relevant chunks
762
query = "What is the main topic?"
763
query_embedding = await embedding_adapter.embed(query)
764
# results = vector_store.search(query_embedding.vector, k=rag_config.top_k)
765
```
766

767
### Semantic Search
768

769
```python
770
from kiln_ai.adapters.embedding import LitellmEmbeddingAdapter
771
import numpy as np
772

773
# Setup
774
adapter = LitellmEmbeddingAdapter(
775
    model_name="text-embedding-3-small",
776
    provider="openai"
777
)
778

779
# Documents to search
780
documents = [
781
    "Python is a high-level programming language.",
782
    "Machine learning is a subset of artificial intelligence.",
783
    "Neural networks are inspired by biological neurons.",
784
    "Data science involves statistics and programming."
785
]
786

787
# Generate embeddings
788
doc_embeddings = await adapter.embed_batch(documents)
789

790
# Query
791
query = "What is AI?"
792
query_embedding = await adapter.embed(query)
793

794
# Calculate cosine similarity
795
def cosine_similarity(a, b):
796
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
797

798
# Find most similar documents
799
similarities = []
800
for i, doc_emb in enumerate(doc_embeddings.embeddings):
801
    sim = cosine_similarity(query_embedding.vector, doc_emb.vector)
802
    similarities.append((i, sim))
803

804
# Sort by similarity
805
similarities.sort(key=lambda x: -x[1])
806

807
print("Most relevant documents:")
808
for idx, sim in similarities[:3]:
809
    print(f"  [{sim:.3f}] {documents[idx]}")
810
```
811

812
### Multi-document Processing
813

814
```python
815
from kiln_ai.adapters.chunkers import FixedWindowChunker
816
from kiln_ai.adapters.embedding import LitellmEmbeddingAdapter
817
from kiln_ai.datamodel import ChunkedDocument, ChunkEmbeddings
818
import glob
819

820
# Initialize components
821
chunker = FixedWindowChunker(chunk_size=500, chunk_overlap=50)
822
embedder = LitellmEmbeddingAdapter(
823
    model_name="text-embedding-3-small",
824
    provider="openai"
825
)
826

827
# Process multiple documents
828
document_paths = glob.glob("/path/to/docs/*.txt")
829
all_chunks = []
830
chunk_metadata = []
831

832
for doc_path in document_paths:
833
    with open(doc_path, "r") as f:
834
        content = f.read()
835

836
    # Chunk document
837
    result = chunker.chunk(content)
838

839
    # Track chunks
840
    for chunk in result.chunks:
841
        all_chunks.append(chunk.text)
842
        chunk_metadata.append({
843
            "source": doc_path,
844
            "start": chunk.start,
845
            "end": chunk.end
846
        })
847

848
print(f"Total chunks: {len(all_chunks)}")
849

850
# Generate embeddings for all chunks
851
embeddings_result = await embedder.embed_batch(all_chunks)
852

853
# Store with metadata
854
for i, (chunk, embedding, metadata) in enumerate(zip(
855
    all_chunks,
856
    embeddings_result.embeddings,
857
    chunk_metadata
858
)):
859
    print(f"Chunk {i}: {metadata['source']}")
860
    # Store in vector database with metadata
861
```
862

863
### Image and PDF Extraction
864

865
```python
866
from kiln_ai.adapters.extractors import LitellmExtractor, ExtractionInput
867
from kiln_ai.datamodel import OutputFormat
868

869
# Create extractor with vision-capable model
870
extractor = LitellmExtractor(
871
    model_name="gpt-4o",
872
    provider="openai"
873
)
874

875
# Extract from image
876
image_input = ExtractionInput(
877
    content="/path/to/diagram.png",
878
    content_type="image/png",
879
    options={"output_format": OutputFormat.markdown}
880
)
881

882
image_extraction = await extractor.extract(image_input)
883
print("Extracted from image:")
884
print(image_extraction.text)
885

886
# Extract from PDF
887
pdf_input = ExtractionInput(
888
    content="/path/to/report.pdf",
889
    content_type="application/pdf",
890
    options={"output_format": OutputFormat.markdown}
891
)
892

893
pdf_extraction = await extractor.extract(pdf_input)
894
print("\nExtracted from PDF:")
895
print(pdf_extraction.text[:500])  # First 500 chars
896
```
897

898
### Comparing Embedding Models
899

900
```python
901
from kiln_ai.adapters.embedding import LitellmEmbeddingAdapter
902
from kiln_ai.adapters.ml_embedding_model_list import built_in_embedding_models_from_provider
903

904
# Get available models
905
openai_models = built_in_embedding_models_from_provider("openai")
906

907
test_text = "This is a test document for comparing embedding models."
908

909
print("Comparing embedding models:\n")
910
for model_info in openai_models:
911
    # Create adapter
912
    adapter = LitellmEmbeddingAdapter(
913
        model_name=model_info.name,
914
        provider="openai"
915
    )
916

917
    # Generate embedding
918
    embedding = await adapter.embed(test_text)
919

920
    print(f"{model_info.name}:")
921
    print(f"  Dimensions: {len(embedding.vector)}")
922
    print(f"  Max input tokens: {model_info.max_input_tokens}")
923
```
924

925
### Configurable Embedding Dimensions
926

927
```python
928
from kiln_ai.adapters.embedding import LitellmEmbeddingAdapter, EmbeddingOptions
929

930
# OpenAI text-embedding-3 models support configurable dimensions
931
options = EmbeddingOptions(
932
    dimensions=512,  # Reduce from default 1536
933
    encoding_format="float"
934
)
935

936
adapter = LitellmEmbeddingAdapter(
937
    model_name="text-embedding-3-small",
938
    provider="openai",
939
    options=options
940
)
941

942
embedding = await adapter.embed("Sample text")
943
print(f"Embedding dimensions: {len(embedding.vector)}")  # Should be 512
944
```
945