Skip to content

Commit eb775c8

Browse files
MinuraPunchihewaMinuraPunchihewa
committed
updated existing unit tests
1 parent b165678 commit eb775c8

File tree

1 file changed

+59
-16
lines changed

1 file changed

+59
-16
lines changed

‎tests/test_mindsdb/test_knowledgebase.py‎

Lines changed: 59 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,21 @@ def test_create_knowledge_base(self):
2424
sql = """
2525
CREATE KNOWLEDGE_BASE my_knowledge_base
2626
USING
27-
MODEL=mindsdb.my_embedding_model,
27+
EMBEDDING_MODEL={
28+
"model": "text-embedding-3-small",
29+
"api_key": "sk-1234567890",
30+
},
2831
STORAGE = my_vector_database.some_table
2932
"""
3033
ast = parse_sql(sql)
3134
expected_ast = CreateKnowledgeBase(
3235
name=Identifier("my_knowledge_base"),
3336
if_not_exists=False,
34-
model=Identifier(parts=["mindsdb", "my_embedding_model"]),
37+
embedding_model={
38+
"model": "text-embedding-3-small",
39+
"api_key": "sk-1234567890",
40+
},
41+
reranking_model=None,
3542
storage=Identifier(parts=["my_vector_database", "some_table"]),
3643
from_select=None,
3744
params={},
@@ -42,18 +49,24 @@ def test_create_knowledge_base(self):
4249
sql = """
4350
CREATE KNOWLEDGE BASE my_knowledge_base
4451
USING
45-
MODEL=mindsdb.my_embedding_model,
52+
EMBEDDING_MODEL={
53+
"model": "text-embedding-3-small",
54+
"api_key": "sk-1234567890",
55+
},
4656
STORAGE = my_vector_database.some_table
4757
"""
4858
ast = parse_sql(sql)
4959
assert ast == expected_ast
5060

51-
# the order of MODEL and STORAGE should not matter
61+
# the order of EMBEDDING_MODEL and STORAGE should not matter
5262
sql = """
5363
CREATE KNOWLEDGE_BASE my_knowledge_base
5464
USING
5565
STORAGE = my_vector_database.some_table,
56-
MODEL = mindsdb.my_embedding_model
66+
EMBEDDING_MODEL={
67+
"model": "text-embedding-3-small",
68+
"api_key": "sk-1234567890",
69+
}
5770
"""
5871
ast = parse_sql(sql)
5972
assert ast == expected_ast
@@ -67,14 +80,21 @@ def test_create_knowledge_base(self):
6780
JOIN my_embedding_model
6881
)
6982
USING
70-
MODEL = mindsdb.my_embedding_model,
83+
EMBEDDING_MODEL={
84+
"model": "text-embedding-3-small",
85+
"api_key": "sk-1234567890",
86+
},
7187
STORAGE = my_vector_database.some_table
7288
"""
7389
ast = parse_sql(sql)
7490
expected_ast = CreateKnowledgeBase(
7591
name=Identifier("my_knowledge_base"),
7692
if_not_exists=False,
77-
model=Identifier(parts=["mindsdb", "my_embedding_model"]),
93+
embedding_model={
94+
"model": "text-embedding-3-small",
95+
"api_key": "sk-1234567890",
96+
},
97+
reranking_model=None,
7898
storage=Identifier(parts=["my_vector_database", "some_table"]),
7999
from_select=Select(
80100
targets=[
@@ -94,7 +114,7 @@ def test_create_knowledge_base(self):
94114

95115
assert ast == expected_ast
96116

97-
# create without MODEL
117+
# create without EMBEDDING_MODEL
98118
sql = """
99119
CREATE KNOWLEDGE_BASE my_knowledge_base
100120
USING
@@ -104,7 +124,8 @@ def test_create_knowledge_base(self):
104124
expected_ast = CreateKnowledgeBase(
105125
name=Identifier("my_knowledge_base"),
106126
if_not_exists=False,
107-
model=None,
127+
embedding_model=None,
128+
reranking_model=None,
108129
storage=Identifier(parts=["my_vector_database", "some_table"]),
109130
from_select=None,
110131
params={},
@@ -118,13 +139,20 @@ def test_create_knowledge_base(self):
118139
sql = """
119140
CREATE KNOWLEDGE_BASE my_knowledge_base
120141
USING
121-
MODEL = mindsdb.my_embedding_model
142+
EMBEDDING_MODEL={
143+
"model": "text-embedding-3-small",
144+
"api_key": "sk-1234567890",
145+
}
122146
"""
123147

124148
expected_ast = CreateKnowledgeBase(
125149
name=Identifier("my_knowledge_base"),
126150
if_not_exists=False,
127-
model=Identifier(parts=["mindsdb", "my_embedding_model"]),
151+
embedding_model={
152+
"model": "text-embedding-3-small",
153+
"api_key": "sk-1234567890",
154+
},
155+
reranking_model=None,
128156
from_select=None,
129157
params={},
130158
)
@@ -137,14 +165,21 @@ def test_create_knowledge_base(self):
137165
sql = """
138166
CREATE KNOWLEDGE_BASE IF NOT EXISTS my_knowledge_base
139167
USING
140-
MODEL = mindsdb.my_embedding_model,
168+
EMBEDDING_MODEL={
169+
"model": "text-embedding-3-small",
170+
"api_key": "sk-1234567890",
171+
},
141172
STORAGE = my_vector_database.some_table
142173
"""
143174
ast = parse_sql(sql)
144175
expected_ast = CreateKnowledgeBase(
145176
name=Identifier("my_knowledge_base"),
146177
if_not_exists=True,
147-
model=Identifier(parts=["mindsdb", "my_embedding_model"]),
178+
embedding_model={
179+
"model": "text-embedding-3-small",
180+
"api_key": "sk-1234567890",
181+
},
182+
reranking_model=None,
148183
storage=Identifier(parts=["my_vector_database", "some_table"]),
149184
from_select=None,
150185
params={},
@@ -160,7 +195,8 @@ def test_create_knowledge_base(self):
160195
expected_ast = CreateKnowledgeBase(
161196
name=Identifier("my_knowledge_base"),
162197
if_not_exists=False,
163-
model=None,
198+
embedding_model=None,
199+
reranking_model=None,
164200
storage=None,
165201
from_select=None,
166202
params={},
@@ -171,7 +207,10 @@ def test_create_knowledge_base(self):
171207
sql = """
172208
CREATE KNOWLEDGE_BASE my_knowledge_base
173209
USING
174-
MODEL = mindsdb.my_embedding_model,
210+
EMBEDDING_MODEL={
211+
"model": "text-embedding-3-small",
212+
"api_key": "sk-1234567890",
213+
},
175214
STORAGE = my_vector_database.some_table,
176215
some_param = 'some value',
177216
other_param = 'other value'
@@ -180,7 +219,11 @@ def test_create_knowledge_base(self):
180219
expected_ast = CreateKnowledgeBase(
181220
name=Identifier("my_knowledge_base"),
182221
if_not_exists=False,
183-
model=Identifier(parts=["mindsdb", "my_embedding_model"]),
222+
embedding_model={
223+
"model": "text-embedding-3-small",
224+
"api_key": "sk-1234567890",
225+
},
226+
reranking_model=None,
184227
storage=Identifier(parts=["my_vector_database", "some_table"]),
185228
from_select=None,
186229
params={"some_param": "some value", "other_param": "other value"},

0 commit comments

Comments
 (0)