Skip to content

Commit 95d5172

Browse files
committed
perf: Update breakpoint_threshold_type and breakpoint_threshold_amount para… …meters to ProcessingConfig for better code scalability and maintainability
1 parent 34b58b8 commit 95d5172

2 files changed

Lines changed: 11 additions & 4 deletions

File tree

embedder_service/models/chunker_with_embedder.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -28,5 +28,10 @@ def embed_query(self, text: str) -> List[float]:
2828
def chunker(config: ProcessingConfig):
2929
"""Helper function to get tools based on current request's config"""
3030
emb_model = SentenceTransformerEmbeddings(config.embedding_model)
31-
sem_chunker = SemanticChunker(emb_model, config.breakpoint_threshold_type, config.breakpoint_threshold_amount)
31+
sem_chunker = SemanticChunker(
32+
emb_model,
33+
breakpoint_threshold_type=config.breakpoint_threshold_type,
34+
breakpoint_threshold_amount=config.breakpoint_threshold_amount
35+
)
36+
3237
return sem_chunker

embedder_service/models/embed.py

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,11 +1,13 @@
11
from pydantic import BaseModel, Field
2-
from typing import List, Dict, Optional
2+
from typing import List, Dict, Optional, get_args
33
import os
44
from langchain_experimental.text_splitter import BreakpointThresholdType
55

66
_EMBEDDING_MODEL_DEFAULT = "all-MiniLM-L6-v2"
77
EMBEDDING_MODEL_NAME = os.getenv("EMBEDDING_MODEL", _EMBEDDING_MODEL_DEFAULT)
88

9+
breakpoints = get_args(BreakpointThresholdType)
10+
911

1012
class ProcessingConfig(BaseModel):
1113
"""Config model for embed API endpoint."""
@@ -18,9 +20,9 @@ class ProcessingConfig(BaseModel):
1820
embedding_model: str = Field(
1921
default=EMBEDDING_MODEL_NAME, description="Sentence Transformer model")
2022
breakpoint_threshold_type: BreakpointThresholdType = Field(
21-
default="percentile", description="Breakpoint threshold type")
23+
default=breakpoints[0], description="Breakpoint threshold type")
2224
breakpoint_threshold_amount: Optional[float] = Field(
23-
default=90, description="Breakpoint threshold amount")
25+
default=90.0, description="Breakpoint threshold amount")
2426
min_chunk_size: int = Field(default=100, description="Minimum chunk size")
2527
max_chunk_size: int = Field(default=1000, description="Maximum chunk size")
2628
store_in_chroma: bool = Field(

0 commit comments

Comments
 (0)