Skip to content

Commit bb967d7

Browse files
authored
Add SingleStoreEmbeddings (#29)
* Add SingleStoreEmbeddings
1 parent fcd6092 commit bb967d7

File tree

2 files changed

+25
-0
lines changed

2 files changed

+25
-0
lines changed

singlestoredb/ai/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
from .embeddings import SingleStoreEmbeddings # noqa: F401

singlestoredb/ai/embeddings.py

Lines changed: 24 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
import os as _os
2+
from typing import Any
3+
4+
try:
5+
from langchain_community.embeddings.ollama import OllamaEmbeddings
6+
except ImportError:
7+
raise ImportError(
8+
'Could not import langchain_community python package. '
9+
'Please install it with `pip install langchain_community`.',
10+
)
11+
12+
13+
class SingleStoreEmbeddings(OllamaEmbeddings):
14+
15+
def __init__(self, **kwargs: Any):
16+
url = _os.getenv('SINGLESTORE_AI_EXPERIMENTAL_URL')
17+
if not url:
18+
raise ValueError(
19+
"Environment variable 'SINGLESTORE_AI_EXPERIMENTAL_URL' must be set",
20+
)
21+
22+
base_url = url.strip('/v1')
23+
kwargs = {'model': 'nomic-embed-text', **kwargs}
24+
super().__init__(base_url=base_url, **kwargs)

0 commit comments

Comments
 (0)