Skip to content

Commit 0c45176

Browse files
authored
Changed embedding model (#1850)
1 parent 0360856 commit 0c45176

File tree

5 files changed

+33
-33
lines changed

5 files changed

+33
-33
lines changed

kairon/shared/llm/processor.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@
3333
class LLMProcessor(LLMBase):
3434
_sparse_embedding = None
3535
_rerank_embedding = None
36-
__embedding__ = 1536
36+
__embedding__ = 3072
3737

3838
def __init__(self, bot: Text, llm_type: str):
3939
super().__init__(bot)
@@ -155,7 +155,7 @@ async def get_embedding(self, texts: Union[Text, List[Text]], user, **kwargs):
155155
truncated_texts = self.truncate_text(texts)
156156

157157
result = await litellm.aembedding(
158-
model="text-embedding-3-small",
158+
model="text-embedding-3-large",
159159
input=truncated_texts,
160160
metadata={'user': user, 'bot': self.bot, 'invocation': kwargs.get("invocation")},
161161
api_key=self.llm_secret_embedding.get('api_key'),

tests/integration_test/action_service_test.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -4170,7 +4170,7 @@ def test_vectordb_action_execution_embedding_search_from_value(mock_embedding):
41704170
log = ActionServerLogs.objects(action=action_name, bot='5f50fd0a56b698ca10d75d2e').get().to_mongo().to_dict()
41714171
log.pop('_id')
41724172
log.pop('timestamp')
4173-
expected = {"model": "text-embedding-3-small",
4173+
expected = {"model": "text-embedding-3-large",
41744174
"input": [payload_body],
41754175
'metadata': {'user': 'default', 'bot': '5f50fd0a56b698ca10d75d2e', 'invocation': 'db_action_qdrant'},
41764176
"api_key": 'key_value',
@@ -4422,7 +4422,7 @@ def test_vectordb_action_execution_embedding_search_from_slot(mock_embedding):
44224422
log = ActionServerLogs.objects(action=action_name, bot='5f50fx0a56b698ca10d35d2e').get().to_mongo().to_dict()
44234423
log.pop('_id')
44244424
log.pop('timestamp')
4425-
expected = {"model": "text-embedding-3-small",
4425+
expected = {"model": "text-embedding-3-large",
44264426
"input": [payload],
44274427
'metadata': {'user': 'default', 'bot': '5f50fx0a56b698ca10d35d2e', 'invocation': 'db_action_qdrant'},
44284428
"api_key": 'key_value',
@@ -4527,7 +4527,7 @@ def test_vectordb_action_execution_embedding_search_no_response_dispatch(mock_em
45274527
log = ActionServerLogs.objects(action=action_name, bot='5f50fd0a56v098ca10d75d2e').get().to_mongo().to_dict()
45284528
log.pop('_id')
45294529
log.pop('timestamp')
4530-
expected = {"model": "text-embedding-3-small",
4530+
expected = {"model": "text-embedding-3-large",
45314531
"input": [payload_body],
45324532
'metadata': {'user': 'default', 'bot': '5f50fd0a56v098ca10d75d2e', 'invocation': 'db_action_qdrant'},
45334533
"api_key": 'key_value',
@@ -13831,7 +13831,7 @@ def test_vectordb_action_execution_embedding_payload_search(mock_embedding):
1383113831
log = ActionServerLogs.objects(action=action_name, bot=bot).get().to_mongo().to_dict()
1383213832
log.pop('_id')
1383313833
log.pop('timestamp')
13834-
expected = {"model": "text-embedding-3-small",
13834+
expected = {"model": "text-embedding-3-large",
1383513835
"input": ["Hi How are you"],
1383613836
'metadata': {'user': 'default', 'bot': bot, 'invocation': 'db_action_qdrant'},
1383713837
"api_key": 'key_value',

tests/integration_test/services_test.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -2012,14 +2012,14 @@ def test_knowledge_vault_sync_push_menu(mock_embedding, mock_collection_exists,
20122012

20132013
expected_calls = [
20142014
{
2015-
"model": "text-embedding-3-small",
2015+
"model": "text-embedding-3-large",
20162016
"input": ['{"id":1,"item":"Juice","price":2.5,"quantity":10}'], # First input
20172017
"metadata": {'user': '[email protected]', 'bot': pytest.bot, 'invocation': 'knowledge_vault_sync'},
20182018
"api_key": "common_openai_key",
20192019
"num_retries": 3
20202020
},
20212021
{
2022-
"model": "text-embedding-3-small",
2022+
"model": "text-embedding-3-large",
20232023
"input": ['{"id":2,"item":"Apples","price":1.2,"quantity":20}'], # Second input
20242024
"metadata": {'user': '[email protected]', 'bot': pytest.bot, 'invocation': 'knowledge_vault_sync'},
20252025
"api_key": "common_openai_key",
@@ -2150,14 +2150,14 @@ def test_knowledge_vault_sync_field_update(mock_embedding, mock_collection_exist
21502150

21512151
expected_calls = [
21522152
{
2153-
"model": "text-embedding-3-small",
2153+
"model": "text-embedding-3-large",
21542154
"input": ['{"id":1,"item":"Juice","price":80.5,"quantity":56}'],
21552155
"metadata": {'user': '[email protected]', 'bot': pytest.bot, 'invocation': 'knowledge_vault_sync'},
21562156
"api_key": "common_openai_key",
21572157
"num_retries": 3
21582158
},
21592159
{
2160-
"model": "text-embedding-3-small",
2160+
"model": "text-embedding-3-large",
21612161
"input": ['{"id":2,"item":"Milk","price":27.0,"quantity":12}'], # Second input
21622162
"metadata": {'user': '[email protected]', 'bot': pytest.bot, 'invocation': 'knowledge_vault_sync'},
21632163
"api_key": "common_openai_key",

tests/unit_test/data_processor/data_processor_test.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -2131,7 +2131,7 @@ async def test_sync_with_qdrant_success(self, mock_collection_upsert, mock_embed
21312131
)
21322132

21332133
mock_embedding.assert_called_once_with(
2134-
model="text-embedding-3-small",
2134+
model="text-embedding-3-large",
21352135
input=['{"id":2,"item":"Milk","price":2.8,"quantity":5}'],
21362136
metadata={'user': user, 'bot': bot, 'invocation': 'knowledge_vault_sync'},
21372137
api_key="openai_key",
@@ -2229,7 +2229,7 @@ async def test_sync_with_qdrant_upsert_failure(self, mock_request, mock_embeddin
22292229
)
22302230

22312231
mock_embedding.assert_called_once_with(
2232-
model="text-embedding-3-small",
2232+
model="text-embedding-3-large",
22332233
input=['{"id":2,"item":"Milk","price":2.8,"quantity":5}'],
22342234
metadata={'user': user, 'bot': bot, 'invocation': 'knowledge_vault_sync'},
22352235
api_key="openai_key",

tests/unit_test/llm_test.py

+21-21
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ async def test_gpt3_faq_embedding_train(self, mock_embedding, aioresponses):
9797
'payload': {'content': test_content.data}
9898
}]}
9999

100-
expected = {"model": "text-embedding-3-small",
100+
expected = {"model": "text-embedding-3-large",
101101
"input": [test_content.data], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
102102
"api_key": value,
103103
"num_retries": 3}
@@ -223,7 +223,7 @@ async def test_gpt3_faq_embedding_train_payload_text(self, mock_embedding, aiore
223223
'payload': {'name': 'Nupur'}}]}
224224
assert response['faq'] == 3
225225

226-
expected = {"model": "text-embedding-3-small",
226+
expected = {"model": "text-embedding-3-large",
227227
"input": [json.dumps(test_content.data)],
228228
'metadata': {'user': user, 'bot': bot, 'invocation': None},
229229
"api_key": value,
@@ -294,7 +294,7 @@ async def test_gpt3_faq_embedding_train_payload_with_int(self, mock_embedding, a
294294
'payload': {'name': 'Ram', 'age': 23, 'color': 'red'}
295295
}]}
296296

297-
expected = {"model": "text-embedding-3-small",
297+
expected = {"model": "text-embedding-3-large",
298298
"input": [json.dumps(input)], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
299299
"api_key": value,
300300
"num_retries": 3}
@@ -375,7 +375,7 @@ async def test_gpt3_faq_embedding_train_int(self, mock_embedding, aioresponses):
375375
'payload': expected_payload
376376
}]}
377377

378-
expected = {"model": "text-embedding-3-small",
378+
expected = {"model": "text-embedding-3-large",
379379
"input": [json.dumps(input)], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
380380
"api_key": value,
381381
"num_retries": 3}
@@ -445,7 +445,7 @@ async def test_gpt3_faq_embedding_train_upsert_error(self, mock_embedding, aiore
445445
'points': [{'id': test_content.vector_id,
446446
'vector': embedding, 'payload': {'content': test_content.data}}]}
447447

448-
expected = {"model": "text-embedding-3-small",
448+
expected = {"model": "text-embedding-3-large",
449449
"input": [test_content.data], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
450450
"api_key": value,
451451
"num_retries": 3}
@@ -525,7 +525,7 @@ async def test_gpt3_faq_embedding_train_payload_upsert_error_json(self, mock_emb
525525
'payload': expected_payload
526526
}]}
527527

528-
expected = {"model": "text-embedding-3-small",
528+
expected = {"model": "text-embedding-3-large",
529529
"input": [json.dumps(test_content.data)], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
530530
"api_key": value,
531531
"num_retries": 3}
@@ -597,7 +597,7 @@ async def test_gpt3_faq_embedding_predict(self, mock_embedding, aioresponses):
597597
'score_threshold': 0.70}
598598
assert isinstance(time_elapsed, float) and time_elapsed > 0.0
599599

600-
expected = {"model": "text-embedding-3-small",
600+
expected = {"model": "text-embedding-3-large",
601601
"input": [query], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
602602
"api_key": value,
603603
"num_retries": 3}
@@ -672,7 +672,7 @@ async def test_gpt3_faq_embedding_predict_with_default_collection(self, mock_emb
672672

673673
assert isinstance(time_elapsed, float) and time_elapsed > 0.0
674674

675-
expected = {"model": "text-embedding-3-small",
675+
expected = {"model": "text-embedding-3-large",
676676
"input": [query], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
677677
"api_key": value,
678678
"num_retries": 3}
@@ -749,7 +749,7 @@ async def test_gpt3_faq_embedding_predict_with_values(self, mock_embedding, aior
749749

750750
assert isinstance(time_elapsed, float) and time_elapsed > 0.0
751751

752-
expected = {"model": "text-embedding-3-small",
752+
expected = {"model": "text-embedding-3-large",
753753
"input": [query], 'metadata': {'user': user, 'bot': gpt3.bot, 'invocation': None},
754754
"api_key": key,
755755
"num_retries": 3}
@@ -825,7 +825,7 @@ async def test_gpt3_faq_embedding_predict_with_values_and_stream(self, mock_embe
825825

826826
assert isinstance(time_elapsed, float) and time_elapsed > 0.0
827827

828-
expected = {"model": "text-embedding-3-small",
828+
expected = {"model": "text-embedding-3-large",
829829
"input": [query], 'metadata': {'user': user, 'bot': gpt3.bot, 'invocation': None},
830830
"api_key": key,
831831
"num_retries": 3}
@@ -938,7 +938,7 @@ async def test_gpt3_faq_embedding_predict_with_values_with_instructions(self,
938938

939939
assert isinstance(time_elapsed, float) and time_elapsed > 0.0
940940

941-
expected = {"model": "text-embedding-3-small",
941+
expected = {"model": "text-embedding-3-large",
942942
"input": [query], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
943943
"api_key": key,
944944
"num_retries": 3}
@@ -1013,7 +1013,7 @@ async def test_gpt3_faq_embedding_predict_completion_connection_error(self, mock
10131013
'score_threshold': 0.70}
10141014
assert isinstance(time_elapsed, float) and time_elapsed > 0.0
10151015

1016-
expected = {"model": "text-embedding-3-small",
1016+
expected = {"model": "text-embedding-3-large",
10171017
"input": [query], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
10181018
"api_key": key,
10191019
"num_retries": 3}
@@ -1067,7 +1067,7 @@ async def test_gpt3_faq_embedding_predict_exact_match(self, mock_embedding, mock
10671067
{'error': 'Retrieving chat completion for the provided query. Failed to connect to service: localhost'}]
10681068
assert isinstance(time_elapsed, float) and time_elapsed > 0.0
10691069

1070-
expected = {"model": "text-embedding-3-small",
1070+
expected = {"model": "text-embedding-3-large",
10711071
"input": [query], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
10721072
"api_key": key,
10731073
"num_retries": 3}
@@ -1113,7 +1113,7 @@ async def test_gpt3_faq_embedding_predict_embedding_connection_error(self, mock_
11131113
assert gpt3.logs == [{'error': 'Creating a new embedding for the provided query. Connection reset by peer!'}]
11141114
assert isinstance(time_elapsed, float) and time_elapsed > 0.0
11151115

1116-
expected = {"model": "text-embedding-3-small",
1116+
expected = {"model": "text-embedding-3-large",
11171117
"input": [query], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
11181118
"api_key": key,
11191119
"num_retries": 3}
@@ -1198,7 +1198,7 @@ async def test_gpt3_faq_embedding_predict_with_previous_bot_responses(self, mock
11981198

11991199
assert isinstance(time_elapsed, float) and time_elapsed > 0.0
12001200

1201-
expected = {"model": "text-embedding-3-small",
1201+
expected = {"model": "text-embedding-3-large",
12021202
"input": [query], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
12031203
"api_key": key,
12041204
"num_retries": 3}
@@ -1292,7 +1292,7 @@ async def test_gpt3_faq_embedding_predict_with_query_prompt(self, mock_embedding
12921292
'score_threshold': 0.70}
12931293
assert isinstance(time_elapsed, float) and time_elapsed > 0.0
12941294

1295-
expected = {"model": "text-embedding-3-small",
1295+
expected = {"model": "text-embedding-3-large",
12961296
"input": [query], 'metadata': {'user': user, 'bot': bot, 'invocation': None},
12971297
"api_key": key,
12981298
"num_retries": 3}
@@ -1926,7 +1926,7 @@ def test_rerank_embedding_raises_exception(self, mock_rerank):
19261926
# }
19271927
#
19281928
# mock_dense.assert_called_once_with(
1929-
# model="text-embedding-3-small",
1929+
# model="text-embedding-3-large",
19301930
# input=[text],
19311931
# metadata={'user': user, 'bot': bot, 'invocation': None},
19321932
# api_key=key,
@@ -1990,7 +1990,7 @@ def test_rerank_embedding_raises_exception(self, mock_rerank):
19901990
# }
19911991
#
19921992
# mock_dense.assert_called_once_with(
1993-
# model="text-embedding-3-small",
1993+
# model="text-embedding-3-large",
19941994
# input=texts,
19951995
# metadata={'user': user, 'bot': bot, 'invocation': None},
19961996
# api_key=key,
@@ -2043,7 +2043,7 @@ def test_rerank_embedding_raises_exception(self, mock_rerank):
20432043
# await processor.get_embedding(texts, user)
20442044
#
20452045
# mock_dense.assert_called_once_with(
2046-
# model="text-embedding-3-small",
2046+
# model="text-embedding-3-large",
20472047
# input=texts,
20482048
# metadata={'user': user, 'bot': bot, 'invocation': None},
20492049
# api_key=key,
@@ -2092,7 +2092,7 @@ def test_rerank_embedding_raises_exception(self, mock_rerank):
20922092
# await processor.get_embedding(texts, user)
20932093
#
20942094
# mock_dense.assert_called_once_with(
2095-
# model="text-embedding-3-small",
2095+
# model="text-embedding-3-large",
20962096
# input=texts,
20972097
# metadata={'user': user, 'bot': bot, 'invocation': None},
20982098
# api_key=key,
@@ -2141,7 +2141,7 @@ def test_rerank_embedding_raises_exception(self, mock_rerank):
21412141
# await processor.get_embedding(texts, user)
21422142
#
21432143
# mock_dense.assert_called_once_with(
2144-
# model="text-embedding-3-small",
2144+
# model="text-embedding-3-large",
21452145
# input=texts,
21462146
# metadata={'user': user, 'bot': bot, 'invocation': None},
21472147
# api_key=key,

0 commit comments

Comments
 (0)