Skip to content

Commit 98a76cc

Browse files
hnwyllmmhnwyllmm
committed
reimplement openai and qwen embedding function based on openai package
1 parent f1b8038 commit 98a76cc

5 files changed

Lines changed: 853 additions & 198 deletions

File tree

‎src/pyseekdb/utils/embedding_functions/openai_base_embedding_function.py‎

Lines changed: 201 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,201 @@
1+
from pyseekdb.client.embedding_function import EmbeddingFunction, Embeddings, Documents
2+
from typing import Any, Optional
3+
import os
4+
5+
6+
class OpenAIBaseEmbeddingFunction(EmbeddingFunction[Documents]):
7+
"""
8+
Base embedding function for OpenAI-compatible embedding APIs.
9+
10+
This class provides a common implementation for embedding functions that use
11+
OpenAI-compatible APIs. It uses the `openai` package to make API calls.
12+
13+
Subclasses should override:
14+
- `_get_default_api_base()`: Return the default API base URL
15+
- `_get_default_api_key_env()`: Return the default API key environment variable name
16+
- `_get_model_dimensions()`: Return a dict mapping model names to their default dimensions
17+
- Optionally override `__init__` to set model-specific defaults
18+
19+
Example:
20+
.. code-block:: python
21+
class MyEmbeddingFunction(OpenAIBaseEmbeddingFunction):
22+
def _get_default_api_base(self):
23+
return "https://api.example.com/v1"
24+
25+
def _get_default_api_key_env(self):
26+
return "MY_API_KEY"
27+
28+
def _get_model_dimensions(self):
29+
return {"model-v1": 1536, "model-v2": 1024}
30+
"""
31+
32+
def __init__(
33+
self,
34+
model_name: str,
35+
api_key_env: Optional[str] = None,
36+
api_base: Optional[str] = None,
37+
dimensions: Optional[int] = None,
38+
**kwargs: Any,
39+
):
40+
"""Initialize OpenAIBaseEmbeddingFunction.
41+
42+
Args:
43+
model_name (str): Name of the embedding model.
44+
api_key_env (str, optional): Name of the environment variable containing the API key.
45+
Defaults to the value returned by `_get_default_api_key_env()`.
46+
api_base (str, optional): Base URL for the API endpoint.
47+
Defaults to the value returned by `_get_default_api_base()`.
48+
dimensions (int, optional): The number of dimensions the resulting embeddings should have.
49+
Can reduce dimensions from default for models that support it.
50+
**kwargs: Additional arguments to pass to the OpenAI client.
51+
Common options include:
52+
- timeout: Request timeout in seconds
53+
- max_retries: Maximum number of retries
54+
- See https://github.com/openai/openai-python for more options
55+
"""
56+
try:
57+
from openai import OpenAI
58+
except ImportError:
59+
raise ValueError(
60+
"The openai python package is not installed. Please install it with `pip install openai`"
61+
)
62+
63+
# Set defaults
64+
if api_key_env is None:
65+
api_key_env = self._get_default_api_key_env()
66+
if api_base is None:
67+
api_base = self._get_default_api_base()
68+
69+
# Validate that API key is available
70+
api_key = os.environ.get(api_key_env)
71+
if api_key is None:
72+
raise ValueError(
73+
f"API key environment variable '{api_key_env}' is not set. "
74+
f"Please set it before using {self.__class__.__name__}."
75+
)
76+
77+
# Store configuration
78+
self.model_name = model_name
79+
self.api_key_env = api_key_env
80+
self.api_base = api_base
81+
self._dimensions_param = dimensions
82+
self._client_kwargs = kwargs
83+
84+
# Initialize OpenAI client
85+
self._client = OpenAI(
86+
api_key=api_key,
87+
base_url=api_base,
88+
**kwargs
89+
)
90+
91+
# Store original model name for dimension lookup
92+
self._model_name = model_name
93+
94+
def _get_default_api_base(self) -> str:
95+
"""Get the default API base URL for this provider.
96+
97+
Subclasses should override this method.
98+
99+
Returns:
100+
str: Default API base URL
101+
"""
102+
raise NotImplementedError(
103+
"Subclasses must implement _get_default_api_base()"
104+
)
105+
106+
def _get_default_api_key_env(self) -> str:
107+
"""Get the default API key environment variable name for this provider.
108+
109+
Subclasses should override this method.
110+
111+
Returns:
112+
str: Default API key environment variable name
113+
"""
114+
raise NotImplementedError(
115+
"Subclasses must implement _get_default_api_key_env()"
116+
)
117+
118+
def _get_model_dimensions(self) -> dict[str, int]:
119+
"""Get a dictionary mapping model names to their default dimensions.
120+
121+
Subclasses should override this method.
122+
123+
Returns:
124+
dict[str, int]: Dictionary mapping model names to dimensions
125+
"""
126+
return {}
127+
128+
@property
129+
def dimension(self) -> int:
130+
"""Get the dimension of embeddings produced by this function.
131+
132+
Returns the known dimension for models without making an API call.
133+
If the dimensions parameter is specified, that value is returned.
134+
Otherwise, the default dimension for the model is returned.
135+
136+
If the model is not in the known dimensions list, falls back to calling
137+
the parent's dimension detection (which may make an API call).
138+
139+
Returns:
140+
int: The dimension of embeddings for this model.
141+
"""
142+
# If dimensions parameter is explicitly set, use it
143+
if self._dimensions_param is not None:
144+
return self._dimensions_param
145+
146+
# Check if we know the dimension for this model
147+
model_dimensions = self._get_model_dimensions()
148+
if self._model_name in model_dimensions:
149+
return model_dimensions[self._model_name]
150+
151+
# Fallback: make an API call to get the embedding and infer the dimension
152+
# This is done by actually generating an embedding for a dummy sentence
153+
test_input = "dimension probing"
154+
try:
155+
embeddings = self([test_input])
156+
except Exception as e:
157+
raise RuntimeError(f"Failed to determine embedding dimension via API call: {e}")
158+
if not embeddings or not isinstance(embeddings, list) or not isinstance(embeddings[0], list):
159+
raise RuntimeError("Could not get embedding dimension from API response")
160+
return len(embeddings[0])
161+
162+
def __call__(self, input: Documents) -> Embeddings:
163+
"""Generate embeddings for the given documents.
164+
165+
Args:
166+
input: Documents to generate embeddings for. Can be a single string or list of strings.
167+
168+
Returns:
169+
Embeddings for the documents as a list of lists of floats.
170+
"""
171+
# Handle single string input
172+
if isinstance(input, str):
173+
input = [input]
174+
175+
# Handle empty input
176+
if not input:
177+
return []
178+
179+
# Prepare request parameters
180+
request_params = {
181+
"model": self.model_name,
182+
"input": input,
183+
}
184+
185+
# Add dimensions parameter if specified
186+
if self._dimensions_param is not None:
187+
request_params["dimensions"] = self._dimensions_param
188+
189+
# Call OpenAI API
190+
response = self._client.embeddings.create(**request_params)
191+
192+
# Extract embeddings from response
193+
embeddings = [item.embedding for item in response.data]
194+
195+
# Validate that we got the expected number of embeddings
196+
if len(embeddings) != len(input):
197+
raise ValueError(
198+
f"Expected {len(input)} embeddings but got {len(embeddings)} from API"
199+
)
200+
201+
return embeddings

‎src/pyseekdb/utils/embedding_functions/openai_embedding_function.py‎

Lines changed: 56 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,25 @@
1-
from pyseekdb.utils.embedding_functions.litellm_embedding_function import LiteLLMEmbeddingFunction
1+
from pyseekdb.utils.embedding_functions.openai_base_embedding_function import OpenAIBaseEmbeddingFunction
22
from typing import Any, Optional
33

4+
# Known OpenAI embedding model dimensions
5+
# Source: https://platform.openai.com/docs/guides/embeddings
6+
_OPENAI_MODEL_DIMENSIONS = {
7+
"text-embedding-ada-002": 1536,
8+
"text-embedding-3-small": 1536,
9+
"text-embedding-3-large": 3072,
10+
}
411

5-
class OpenAIEmbeddingFunction(LiteLLMEmbeddingFunction):
12+
13+
class OpenAIEmbeddingFunction(OpenAIBaseEmbeddingFunction):
614
"""
715
A convenient embedding function for OpenAI embedding models.
816
9-
This class provides a simplified interface to OpenAI embedding models using LiteLLM.
10-
It sets default values for OpenAI-specific configurations.
17+
This class provides a simplified interface to OpenAI embedding models using the OpenAI API.
1118
1219
For more information about OpenAI models, see https://platform.openai.com/docs/guides/embeddings
1320
1421
Example:
15-
pip install pyseekdb litellm
22+
pip install pyseekdb openai
1623
1724
.. code-block:: python
1825
import pyseekdb
@@ -32,6 +39,12 @@ class OpenAIEmbeddingFunction(LiteLLMEmbeddingFunction):
3239
max_retries=3
3340
)
3441
42+
# Using text-embedding-3 with custom dimensions
43+
ef = OpenAIEmbeddingFunction(
44+
model_name="text-embedding-3-small",
45+
dimensions=512 # Reduce from default 1536 to 512
46+
)
47+
3548
db = pyseekdb.Client(path="./seekdb.db")
3649
collection = db.create_collection(name="my_collection", embedding_function=ef)
3750
# Add documents
@@ -46,6 +59,8 @@ def __init__(
4659
self,
4760
model_name: str = "text-embedding-ada-002",
4861
api_key_env: Optional[str] = None,
62+
api_base: Optional[str] = None,
63+
dimensions: Optional[int] = None,
4964
**kwargs: Any,
5065
):
5166
"""Initialize OpenAIEmbeddingFunction.
@@ -55,26 +70,50 @@ def __init__(
5570
Defaults to "text-embedding-ada-002".
5671
Other options include:
5772
- "text-embedding-ada-002" (1536 dimensions, default)
58-
- "text-embedding-3-small" (1536 dimensions)
59-
- "text-embedding-3-large" (3072 dimensions)
73+
- "text-embedding-3-small" (1536 dimensions by default, can be reduced via dimensions parameter)
74+
- "text-embedding-3-large" (3072 dimensions by default, can be reduced via dimensions parameter)
6075
api_key_env (str, optional): Name of the environment variable containing the OpenAI API key.
6176
Defaults to "OPENAI_API_KEY" if not provided.
62-
**kwargs: Additional arguments to pass to the LiteLLM embedding function.
77+
api_base (str, optional): Base URL for the API endpoint.
78+
Defaults to "https://api.openai.com/v1" if not provided.
79+
Useful for OpenAI-compatible proxies or custom endpoints.
80+
dimensions (int, optional): The number of dimensions the resulting embeddings should have.
81+
Only supported for text-embedding-3 models. Can reduce dimensions from
82+
default (1536 for text-embedding-3-small, 3072 for text-embedding-3-large).
83+
**kwargs: Additional arguments to pass to the OpenAI client.
6384
Common options include:
64-
- api_base: Base URL for the API endpoint (useful for OpenAI-compatible proxies)
6585
- timeout: Request timeout in seconds
6686
- max_retries: Maximum number of retries
67-
- api_version: API version
68-
- user: User identifier for usage tracking
69-
- See https://docs.litellm.ai/docs/embedding/supported_embedding for more options
87+
- See https://github.com/openai/openai-python for more options
7088
"""
71-
# Set default api_key_env if not provided
72-
if api_key_env is None:
73-
api_key_env = "OPENAI_API_KEY"
74-
75-
# Initialize parent class with OpenAI-specific defaults
7689
super().__init__(
7790
model_name=model_name,
7891
api_key_env=api_key_env,
92+
api_base=api_base,
93+
dimensions=dimensions,
7994
**kwargs
8095
)
96+
97+
def _get_default_api_base(self) -> str:
98+
"""Get the default API base URL for OpenAI.
99+
100+
Returns:
101+
str: Default OpenAI API base URL
102+
"""
103+
return "https://api.openai.com/v1"
104+
105+
def _get_default_api_key_env(self) -> str:
106+
"""Get the default API key environment variable name for OpenAI.
107+
108+
Returns:
109+
str: Default OpenAI API key environment variable name
110+
"""
111+
return "OPENAI_API_KEY"
112+
113+
def _get_model_dimensions(self) -> dict[str, int]:
114+
"""Get a dictionary mapping OpenAI model names to their default dimensions.
115+
116+
Returns:
117+
dict[str, int]: Dictionary mapping model names to dimensions
118+
"""
119+
return _OPENAI_MODEL_DIMENSIONS

0 commit comments

Comments
 (0)