Fix open ai embedding large compatibility with pgvector

This commit is contained in:
Adityavardhan Agrawal 2025-04-08 17:30:43 -07:00
parent 9533924bf8
commit 7bdc09828a
2 changed files with 6 additions and 2 deletions

View File

@ -7,6 +7,7 @@ from core.models.chunk import Chunk
from core.config import get_settings
logger = logging.getLogger(__name__)
PGVECTOR_MAX_DIMENSIONS = 2000 # Maximum dimensions for pgvector
class LiteLLMEmbeddingModel(BaseEmbeddingModel):
@ -33,7 +34,7 @@ class LiteLLMEmbeddingModel(BaseEmbeddingModel):
raise ValueError(f"Model '{model_key}' not found in registered_models configuration")
self.model_config = settings.REGISTERED_MODELS[model_key]
self.dimensions = settings.VECTOR_DIMENSIONS
self.dimensions = min(settings.VECTOR_DIMENSIONS, 2000)
logger.info(
f"Initialized LiteLLM embedding model with model_key={model_key}, config={self.model_config}"
)
@ -53,6 +54,8 @@ class LiteLLMEmbeddingModel(BaseEmbeddingModel):
try:
model_params = {"model": self.model_config["model_name"]}
if self.model_config["model_name"] == "text-embedding-3-large":
model_params["dimensions"] = PGVECTOR_MAX_DIMENSIONS
# Add all model-specific parameters from the config
for key, value in self.model_config.items():

View File

@ -15,6 +15,7 @@ from core.models.chunk import DocumentChunk
logger = logging.getLogger(__name__)
Base = declarative_base()
PGVECTOR_MAX_DIMENSIONS = 2000 # Maximum dimensions for pgvector
class Vector(UserDefinedType):
@ -132,7 +133,7 @@ class PGVectorStore(BaseVectorStore):
# Import config to get vector dimensions
from core.config import get_settings
settings = get_settings()
dimensions = settings.VECTOR_DIMENSIONS
dimensions = min(settings.VECTOR_DIMENSIONS, PGVECTOR_MAX_DIMENSIONS)
logger.info(f"Initializing PGVector store with {dimensions} dimensions")