kolzchut-ragbot 1.0.1__tar.gz

kolzchut_ragbot-1.0.1/PKG-INFO

@@ -0,0 +1,67 @@
+Metadata-Version: 2.2
+Name: kolzchut-ragbot
+Version: 1.0.1
+Summary: A search engine using machine learning models and Elasticsearch for advanced document retrieval.
+Home-page: https://github.com/shmuelrob/ragbot
+Author: Shmuel Robinov
+Author-email: shmuel_robinov@webiks.com
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: elasticsearch==8.14.0
+Requires-Dist: sentence-transformers==3.0.1
+Requires-Dist: torch==2.3.1
+Requires-Dist: transformers==4.42.3
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# **Webiks-Hebrew-RAGbot**
+
+## **Overview**
+
+This project is a search engine that uses machine learning models and Elasticsearch to provide advanced document retrieval.
+You can use [Webiks-Hebrew-RAGbot-Demo](https://github.com/NNLP-IL/Webiks-Hebrew-RAGbot-Demo) to demonstrate the engine's document retrieval abilities
+
+## **Features**
+
+Document representation and validation
+Document embedding and indexing in Elasticsearch
+Advanced search using machine learning model
+Integration with LLM (Large Language Model) client for query answering
+
+## **Installation**
+
+1.  Clone the repository:
+   
+`git clone https://github.com/NNLP-IL/Webiks-Hebrew-RAGbot.git`
+
+`cd Webiks-Hebrew-RAGbot`
+
+2.  Create a virtual environment and activate it:  
+
+`python -m venv venv`
+
+`source venv/bin/activate`
+
+On Windows use `\venv\\Scripts\\activate\`
+
+3.  Install the required dependencies:  
+
+`pip install -r requirements.txt`
+
+## **Configuration**
+
+Set the following environment variables:  
+
+ES\_EMBEDDING\_INDEX: The name of the Elasticsearch index for embeddings.
+
+TOKENIZER\_LOCATION: The location of the tokenizer model.

kolzchut_ragbot-1.0.1/README.md

@@ -0,0 +1,41 @@
+# **Webiks-Hebrew-RAGbot**
+
+## **Overview**
+
+This project is a search engine that uses machine learning models and Elasticsearch to provide advanced document retrieval.
+You can use [Webiks-Hebrew-RAGbot-Demo](https://github.com/NNLP-IL/Webiks-Hebrew-RAGbot-Demo) to demonstrate the engine's document retrieval abilities
+
+## **Features**
+
+Document representation and validation
+Document embedding and indexing in Elasticsearch
+Advanced search using machine learning model
+Integration with LLM (Large Language Model) client for query answering
+
+## **Installation**
+
+1.  Clone the repository:
+   
+`git clone https://github.com/NNLP-IL/Webiks-Hebrew-RAGbot.git`
+
+`cd Webiks-Hebrew-RAGbot`
+
+2.  Create a virtual environment and activate it:  
+
+`python -m venv venv`
+
+`source venv/bin/activate`
+
+On Windows use `\venv\\Scripts\\activate\`
+
+3.  Install the required dependencies:  
+
+`pip install -r requirements.txt`
+
+## **Configuration**
+
+Set the following environment variables:  
+
+ES\_EMBEDDING\_INDEX: The name of the Elasticsearch index for embeddings.
+
+TOKENIZER\_LOCATION: The location of the tokenizer model.

kolzchut_ragbot-1.0.1/kolzchut_ragbot.egg-info/PKG-INFO

@@ -0,0 +1,67 @@
+Metadata-Version: 2.2
+Name: kolzchut-ragbot
+Version: 1.0.1
+Summary: A search engine using machine learning models and Elasticsearch for advanced document retrieval.
+Home-page: https://github.com/shmuelrob/ragbot
+Author: Shmuel Robinov
+Author-email: shmuel_robinov@webiks.com
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: elasticsearch==8.14.0
+Requires-Dist: sentence-transformers==3.0.1
+Requires-Dist: torch==2.3.1
+Requires-Dist: transformers==4.42.3
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# **Webiks-Hebrew-RAGbot**
+
+## **Overview**
+
+This project is a search engine that uses machine learning models and Elasticsearch to provide advanced document retrieval.
+You can use [Webiks-Hebrew-RAGbot-Demo](https://github.com/NNLP-IL/Webiks-Hebrew-RAGbot-Demo) to demonstrate the engine's document retrieval abilities
+
+## **Features**
+
+Document representation and validation
+Document embedding and indexing in Elasticsearch
+Advanced search using machine learning model
+Integration with LLM (Large Language Model) client for query answering
+
+## **Installation**
+
+1.  Clone the repository:
+   
+`git clone https://github.com/NNLP-IL/Webiks-Hebrew-RAGbot.git`
+
+`cd Webiks-Hebrew-RAGbot`
+
+2.  Create a virtual environment and activate it:  
+
+`python -m venv venv`
+
+`source venv/bin/activate`
+
+On Windows use `\venv\\Scripts\\activate\`
+
+3.  Install the required dependencies:  
+
+`pip install -r requirements.txt`
+
+## **Configuration**
+
+Set the following environment variables:  
+
+ES\_EMBEDDING\_INDEX: The name of the Elasticsearch index for embeddings.
+
+TOKENIZER\_LOCATION: The location of the tokenizer model.

kolzchut_ragbot-1.0.1/kolzchut_ragbot.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+README.md
+pyproject.toml
+setup.py
+kolzchut_ragbot.egg-info/PKG-INFO
+kolzchut_ragbot.egg-info/SOURCES.txt
+kolzchut_ragbot.egg-info/dependency_links.txt
+kolzchut_ragbot.egg-info/requires.txt
+kolzchut_ragbot.egg-info/top_level.txt
+test/test_configs.py
+test/test_document.py
+test/test_engine.py
+test/test_model.py

kolzchut_ragbot-1.0.1/kolzchut_ragbot.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

kolzchut_ragbot-1.0.1/kolzchut_ragbot.egg-info/requires.txt

@@ -0,0 +1,4 @@
+elasticsearch==8.14.0
+sentence-transformers==3.0.1
+torch==2.3.1
+transformers==4.42.3

kolzchut_ragbot-1.0.1/kolzchut_ragbot.egg-info/top_level.txt

@@ -0,0 +1 @@
+

kolzchut_ragbot-1.0.1/pyproject.toml

@@ -0,0 +1,33 @@
+# pyproject.toml
+
+[build-system]
+requires = ["poetry-core>=1.0.0"]
+build-backend = "poetry.core.masonry.api"
+
+[tool.bumpver]
+current_version = "1.0.53"
+version_pattern = "MAJOR.MINOR.PATCH"
+commit_message  = "Bump version {old_version} -> {new_version}"
+commit          = true
+tag             = true
+push            = false
+
+[tool.poetry]
+name = "ragbot"
+version = "1.0.53"
+description = ""
+authors = ["Your Name <your.email@example.com>"]
+
+[tool.poetry.dependencies]
+python = "^3.7"
+elasticsearch = ">=8.10.1"
+tomli = { version = "*", python = "<3.11" }
+transformers = ">=4.41.0"
+sentence_transformers = ">=2.2.2"
+torch = ">=2.1.2"
+
+[tool.bumpver.file_patterns]
+"pyproject.toml" = [
+    'current_version = "{version}"',
+]
+

kolzchut_ragbot-1.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

kolzchut_ragbot-1.0.1/setup.py

@@ -0,0 +1,25 @@
+from setuptools import setup, find_packages
+
+setup(
+    name='kolzchut-ragbot',
+    version='1.0.1',
+    author='Shmuel Robinov',
+    author_email='shmuel_robinov@webiks.com',
+    description='A search engine using machine learning models and Elasticsearch for advanced document retrieval.',
+    long_description=open('README.md').read(),
+    long_description_content_type='text/markdown',
+    url='https://github.com/shmuelrob/ragbot',
+    packages=find_packages(),
+    install_requires=[
+        'elasticsearch==8.14.0',
+        'sentence-transformers==3.0.1',
+        'torch==2.3.1',
+        'transformers==4.42.3'
+    ],
+    classifiers=[
+        'Programming Language :: Python :: 3',
+        'License :: OSI Approved :: MIT License',
+        'Operating System :: OS Independent',
+    ],
+    python_requires='>=3.10',
+)

kolzchut_ragbot-1.0.1/test/test_configs.py

@@ -0,0 +1,42 @@
+bad_conf_1 = '''
+{
+  "identifier_field": "no_such_field",
+  "saved_fields": {
+    "title": "text",
+    "doc_id": "integer",
+    "content": "text"
+  },
+  "models": {
+    "me5_large-v10": "content"
+  }
+}
+'''
+
+bad_conf_2 = '''
+{
+  "identifier_field": "doc_id",
+    "saved_fields": {
+    "title": "text",
+    "doc_id": "integer",
+    "content": "text"
+  },
+  "models": {
+    "me5_large-v10": "content",
+    "OMG": "no_such_field"
+  }
+}
+'''
+
+good_conf = '''
+{
+  "identifier_field": "doc_id",
+    "saved_fields": {
+    "title": "text",
+    "doc_id": "integer",
+    "content": "text"
+  },
+  "models": {
+    "me5_large-v10": "content"
+  }
+}
+'''

kolzchut_ragbot-1.0.1/test/test_document.py

@@ -0,0 +1,54 @@
+import unittest
+from unittest.mock import patch, mock_open
+import json
+import sys
+import os
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))
+
+from ragbot.Document import initialize_definitions
+
+
+class TestDocument(unittest.TestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        pass
+
+    @patch('builtins.open', new_callable=mock_open, read_data=json.dumps({
+        "saved_fields": {"doc_id": "doc_id"},
+        "models": {"model1": "doc_id"},
+        "identifier_field": "doc_id"
+    }))
+    @patch('os.getenv', return_value='example-conf.json')
+    def test_initialize_definitions_good(self, mock_getenv, mock_open):
+        definitions = initialize_definitions()
+        self.assertEqual(definitions.saved_fields, {"doc_id": "doc_id"})
+        self.assertEqual(definitions.models, {"model1": "doc_id"})
+        self.assertEqual(definitions.identifier, "doc_id")
+
+    @patch('builtins.open', new_callable=mock_open, read_data=json.dumps({
+        "saved_fields": {"field1": {"field_name": "field1", "required": True}},
+        "models": {"model1": "field1"},
+        "identifier_field": "no_such_field"
+    }))
+    @patch('os.getenv', return_value='example-conf.json')
+    def test_initialize_definitions_bad_identifier(self, mock_getenv, mock_open):
+        with self.assertRaises(ValueError) as context:
+            initialize_definitions()
+        self.assertIn("identifier_field must be one of the saved fields", str(context.exception))
+
+    @patch('builtins.open', new_callable=mock_open, read_data=json.dumps({
+        "saved_fields": {"field1": {"field_name": "field1", "required": True}},
+        "models": "no_such_field",
+        "identifier_field": "doc_id"
+    }))
+    @patch('os.getenv', return_value='example-conf.json')
+    def test_initialize_definitions_bad_model(self, mock_getenv, mock_open):
+        with self.assertRaises(ValueError) as context:
+            initialize_definitions()
+        self.assertIn("must be one of the saved fields", str(context.exception))
+
+
+if __name__ == '__main__':
+    unittest.main()

kolzchut_ragbot-1.0.1/test/test_engine.py

@@ -0,0 +1,172 @@
+import unittest
+import importlib
+from unittest.mock import patch, MagicMock, ANY
+import sys
+import os
+
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))
+import ragbot.engine
+importlib.reload(ragbot.engine)
+def build_test_engine(es_model, llm_client):
+        reranker_model = MagicMock()
+        reranker_tokenizer = MagicMock()
+        models = MagicMock()
+        return ragbot.engine.Engine(llms_client=llm_client, es_client=es_model, models=models,
+                                      reranker_model=reranker_model, reranker_tokenizer=reranker_tokenizer)
+
+
+class TestEngine(unittest.TestCase):
+
+    @patch('ragbot.llm_client.LLMClient')
+    @patch('ragbot.model.Model')
+    @patch('elasticsearch.Elasticsearch')
+    def test_update_title_summary(self, Elasticsearch, Model, LLMClient):
+        list_of_docs = [
+            {"doc_id": 1, 'title': 'title1', 'summary': 'summary1', 'content': 'content1'},
+            {"doc_id": 2, 'title': 'title2', 'summary': 'summary2', 'content': 'content2'}
+        ]
+
+        es_client = Elasticsearch()
+        es_model = Model(es_client)
+        llm_client = LLMClient()
+        engine = build_test_engine(es_model, llm_client)
+        engine.update_docs(list_of_docs=list_of_docs, embed_only_fields=['title', 'summary'], delete_existing=False)
+        es_model.create_or_update_documents.assert_called_once_with(list_of_docs, False)
+
+    @patch('ragbot.llm_client.LLMClient')
+    @patch('ragbot.model.Model')
+    @patch('elasticsearch.Elasticsearch')
+    def test_update_content_without_delete(self, Elasticsearch, Model, LLMClient):
+        list_of_docs = [
+            {"doc_id": 1, 'title': 'title1', 'summary': 'summary1', 'content': 'content1'},
+            {"doc_id": 2, 'title': 'title2', 'summary': 'summary2', 'content': 'content2'}
+        ]
+        es_client = Elasticsearch()
+        es_model = Model(es_client)
+        llm_client = LLMClient()
+        engine = build_test_engine(es_model, llm_client)
+        engine.update_docs(list_of_docs, embed_only_fields=['content'], delete_existing=False)
+        es_model.create_or_update_documents.assert_called_once_with(list_of_docs, False)
+
+    @patch('ragbot.llm_client.LLMClient')
+    @patch('ragbot.model.Model')
+    @patch('elasticsearch.Elasticsearch')
+    def test_update_content_with_delete(self, Elasticsearch, Model, LLMClient):
+        list_of_docs = [
+            {"doc_id": 1, 'title': 'title1', 'summary': 'summary1', 'content': 'content1'},
+            {"doc_id": 2, 'title': 'title2', 'summary': 'summary2', 'content': 'content2'}
+        ]
+        reranker_model = MagicMock()
+        reranker_tokenizer = MagicMock()
+        llm_client = LLMClient()
+        models = MagicMock()
+        es_client = Elasticsearch()
+        es_model = Model(es_client)
+        engine = ragbot.engine.Engine(llms_client=llm_client, es_client=es_model, models=models,
+                                      reranker_model=reranker_model, reranker_tokenizer=reranker_tokenizer)
+        engine.update_docs(list_of_docs, embed_only_fields=['content'], delete_existing=True)
+        es_model.create_or_update_documents.assert_called_once_with(list_of_docs, True)
+
+    @patch('ragbot.llm_client.LLMClient')
+    @patch('ragbot.model.Model')
+    @patch('elasticsearch.Elasticsearch')
+    def test_reciprocal_rank_fusion(self, Elasticsearch, Model, LLMClient):
+        es_client = Elasticsearch()
+        es_model = Model(es_client)
+        llm_client = LLMClient()
+        engine = build_test_engine(es_model, llm_client)
+
+        ranking_lists = [
+            [1, 2, 3],
+            [2, 3, 4],
+            [3, 4, 5]
+        ]
+        expected_fused_list = [3,2,4,1,5]
+        fused_list = engine.reciprocal_rank_fusion(ranking_lists)
+        self.assertEqual(fused_list, expected_fused_list)
+
+    @patch.object(ragbot.engine.Engine, 'reciprocal_rank_fusion')
+    @patch('ragbot.llm_client.LLMClient')
+    @patch('ragbot.model.Model')
+    @patch('elasticsearch.Elasticsearch')
+    def test_search_documents(self, Elasticsearch, Model, LLMClient, mock_reciprocal_rank_fusion):
+        llm_client = LLMClient()
+        models = MagicMock()
+        es_client = Elasticsearch()
+        es_model = Model(es_client)
+        es_model.search.return_value = {
+                "title":[
+                    {'_source': {'page_id': 1, 'title': 'title1'}},
+                    {'_source': {'page_id': 2, 'title': 'title2'}},
+                    {'_source': {'page_id': 3, 'title': 'title3'}},
+                ],
+                "summary":[
+                    {'_source': {'page_id': 2, 'title': 'title2'}},
+                    {'_source': {'page_id': 3, 'title': 'title3'}},
+                    {'_source': {'page_id': 4, 'title': 'title4'}},
+                ],
+                "content":[
+                    {'_source': {'page_id': 3, 'title': 'title3'}},
+                    {'_source': {'page_id': 4, 'title': 'title4'}},
+                    {'_source': {'page_id': 5, 'title': 'title5'}}
+                ]
+        }
+        mock_reciprocal_rank_fusion.return_value = [3, 2, 4, 1, 5]
+        engine = build_test_engine(es_model, llm_client)
+
+        result = engine.search_documents("test query", 5)
+
+        es_model.search.assert_called_once()
+        self.assertEqual(mock_reciprocal_rank_fusion.return_value, [3, 2, 4, 1, 5])
+        self.assertEqual([
+            {'page_id': 3, 'title': 'title3'},
+            {'page_id': 2, 'title': 'title2'},
+            {'page_id': 4, 'title': 'title4'},
+            {'page_id': 1, 'title': 'title1'},
+            {'page_id': 5, 'title': 'title5'}
+        ], result)
+
+    @patch('ragbot.llm_client.LLMClient')
+    @patch('ragbot.model.Model')
+    @patch('elasticsearch.Elasticsearch')
+    def test_answer_query(self, Elasticsearch, Model, LLMClient):
+        es_client = Elasticsearch()
+        es_model = Model(es_client)
+        llm_client = LLMClient()
+        engine = build_test_engine(es_model, llm_client)
+
+        with patch.object(ragbot.engine.Engine, 'search_documents') as mock_search_documents:
+            mock_search_documents.return_value = [
+                {'page_id': 3, 'title': 'title3'},
+                {'page_id': 2, 'title': 'title2'},
+                {'page_id': 4, 'title': 'title4'},
+                {'page_id': 1, 'title': 'title1'},
+                {'page_id': 5, 'title': 'title5'}
+            ]
+
+            llm_client.answer.return_value = ('answer', 0.5, 100)
+            actual_top_k_documents, actual_gpt_answer, actual_stats = engine.answer_query("test query", 5, 'gpt-4o')
+
+            expected_top_k_documents = [
+                {'page_id': 3, 'title': 'title3'},
+                {'page_id': 2, 'title': 'title2'},
+                {'page_id': 4, 'title': 'title4'},
+                {'page_id': 1, 'title': 'title1'},
+                {'page_id': 5, 'title': 'title5'}
+            ]
+            expected_gpt_answer = llm_client.answer.return_value[0]
+            expected_stats = {
+                "retrieval_time": 0,
+                "gpt_model": 'gpt-4o',
+                "gpt_time": llm_client.answer.return_value[1],
+                "tokens": llm_client.answer.return_value[2]
+            }
+
+            self.assertEqual(expected_top_k_documents, actual_top_k_documents)
+            self.assertEqual(expected_gpt_answer, actual_gpt_answer)
+            self.assertEqual(expected_stats, actual_stats)
+
+
+if __name__ == '__main__':
+    unittest.main()
+

kolzchut_ragbot-1.0.1/test/test_model.py

@@ -0,0 +1,153 @@
+import unittest
+from unittest.mock import patch, ANY
+import sys
+import os
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))
+from ragbot.model import Model, index_from_page_id
+
+search_json = {
+   "script_score": {
+       "query": {
+           "exists": {
+               "field": 'content-me5_large-v10'
+           }
+       },
+       "script": {
+           "source": f"cosineSimilarity(params.query_vector, 'content-me5_large-v10') + 1.0",
+           "params": {
+               "Query_vector": [0.0, 0.11, ]
+           }
+       }
+   }
+}
+
+
+class TestModel(unittest.TestCase):
+
+    @patch('elasticsearch.Elasticsearch')
+    def test_create_index(self, EsMock):
+        es_mock = EsMock()
+        model = Model(es_mock)
+        expected_index_mapping = {
+            'properties': {
+                'last_update': {'type': 'date'},
+                'me5_large-v10_content_vectors': {'type': 'dense_vector', 'dims': 1024},
+                'title': {'type': 'text'}, 'doc_id': {'type': 'integer'}, 'content': {'type': 'text'}}
+        }
+        es_mock.indices.exists.return_value = False
+        model.create_index()
+        self.assertTrue(es_mock.indices.create.called)
+        _, kwargs = es_mock.indices.create.call_args
+        self.assertEqual(expected_index_mapping, kwargs.get("mappings"))
+
+    @patch('elasticsearch.Elasticsearch')
+    def test_create_index_false(self, EsMock):
+        es_mock = EsMock()
+        model = Model(es_mock)
+        es_mock.indices.exists.return_value = True
+        model.create_index()
+        self.assertTrue(es_mock.indices.create.not_called)
+
+    @patch('elasticsearch.Elasticsearch')
+    def test_create_or_update_document_no_delete(self, EsMock):
+        es_mock = EsMock()
+        model = Model(es_mock)
+        es_mock.search.return_value = {"hits": {"hits": []}}
+        new_doc = {"doc_id": 1, "title": "title", "content": "content"}
+        model.create_or_update_documents([new_doc], True)
+
+        es_mock.search.assert_called_with(
+            index=index_from_page_id(1),
+            body={
+                "query": {
+                    "term": {
+                        "doc_id": {"value": 1}
+                    }
+                }
+            })
+        self.assertTrue(es_mock.delete.call_count == 0)
+        es_mock.index.called_once()
+
+    @patch('elasticsearch.Elasticsearch')
+    def test_create_or_update_document_but_delete(self, EsMock):
+        es_mock = EsMock()
+        model = Model(es_mock)
+        es_mock.search.return_value = {"hits": {"hits": [{"_id":"1","doc_id": 1, "title": "title", "content": "content"}]}}
+        new_doc = {"doc_id": 1, "title": "edited", "content": "edited"}
+        model.create_or_update_documents([new_doc], True)
+
+        es_mock.search.assert_called_with(
+            index=index_from_page_id(1),
+            body={
+                "query": {
+                    "term": {
+                        "doc_id": {"value": 1}
+                    }
+                }
+            })
+        self.assertTrue(es_mock.delete.call_count == 1)
+        es_mock.index.called_once()
+
+    @patch('elasticsearch.Elasticsearch')
+    def test_create_or_update_document_delete_false(self, EsMock):
+        es_mock = EsMock()
+        model = Model(es_mock)
+        es_mock.search.return_value = {
+            "hits": {"hits": [{"_id": "1", "doc_id": 1, "title": "title", "content": "content"}]}}
+        new_doc = {"doc_id": 1, "title": "edited", "content": "edited"}
+        model.create_or_update_documents([new_doc], False)
+
+        self.assertEqual(0, es_mock.search.call_count)
+        self.assertEqual(0, es_mock.delete.call_count)
+        es_mock.index.called_once()
+
+    @patch('elasticsearch.Elasticsearch')
+    def test_search(self, EsMock):
+        es_mock = EsMock()
+        model = Model(es_mock)
+        es_mock.search.return_value = {
+            "hits": {
+                "hits": [
+                    {"_id": "1", "_source": {"field": "value1"}},
+                    {"_id": "2", "_source": {"field": "value2"}}
+                ]
+            }
+        }
+
+        embedded_search = {
+            "me5_large-v10": [0.1, 0.2, 0.3],
+            "model2": [0.4, 0.5, 0.6]
+        }
+
+        results = model.search(embedded_search, size=2)
+
+        expected_query = {
+            "size": 2,
+            "query": {
+                "script_score": {
+                    "query": {
+                        "exists": {
+                            "field": "doc_id_me5_large-v10_vectors"
+                        }
+                    },
+                    "script": {
+                        "source": "cosineSimilarity(params.query_vector, 'doc_id_me5_large-v10_vectors') + 1.0",
+                        "params": {
+                            "query_vector": embedded_search["me5_large-v10"]
+                        }
+                    }
+                }
+            }
+        }
+
+        expected_results = {
+            "doc_id": [
+                {"_id": "1", "_source": {"field": "value1"}},
+                {"_id": "2", "_source": {"field": "value2"}}
+            ]
+        }
+        self.assertEqual(results["content"], expected_results["doc_id"])
+
+
+if __name__ == '__main__':
+    unittest.main()