kolzchut-ragbot 1.6.1__tar.gz → 1.7.2__tar.gz

{kolzchut_ragbot-1.6.1 → kolzchut_ragbot-1.7.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kolzchut-ragbot
-Version: 1.6.1
+Version: 1.7.2
 Summary: A search engine using machine learning models and Elasticsearch for advanced document retrieval.
 Home-page: https://github.com/shmuelrob/rag-bot
 Author: Shmuel Robinov

{kolzchut_ragbot-1.6.1 → kolzchut_ragbot-1.7.2}/kolzchut_ragbot/engine.py

@@ -170,20 +170,21 @@ class Engine:
 
         return fused_list
 
-    def search_documents(self, query: str, top_k: int):
+    def search_documents(self, query: str, top_k: int,retrieval_size: int, max_documents_from_same_page:int):
         """
         Searches for documents based on the query and returns the top_k results.
 
         Args:
             query (str): The query string.
             top_k (int): The number of top documents to return.
-
+            retrieval_size (int, optional): The number of documents to fetch from each model.
+            max_documents_from_same_page (int, optional): The maximum number of documents (paragraphs acutually) to return from the same page.
         Returns:
             list: A list of top k documents.
         """
         query_embeddings = {f"{semantic_model}": self.models[semantic_model].encode(query) for semantic_model in
                             definitions.models.keys()}
-        all_docs_by_model = self.elastic_model.search(query_embeddings)
+        all_docs_by_model = self.elastic_model.search(embedded_search=query_embeddings, size=retrieval_size)
         all_docs = []
         ids_for_fusion = []
         all_docs_and_scores = {}
@@ -202,20 +203,20 @@ class Engine:
         print(f"\nFusing {len(ids_for_fusion)} results\n")
         fused_ids = self.reciprocal_rank_fusion(ids_for_fusion, k=top_k)
         top_k_documents = []
-        top_titles = []
+        count_per_id = {}
 
-        for fused_id in fused_ids:
+        for fused_id in fused_ids[:top_k]:
             for doc in all_docs:
-                if doc["_source"]["page_id"] == fused_id and doc["_source"]["title"] not in top_titles:
+                if doc["_source"]["page_id"] == fused_id:
+                    count = count_per_id.get(fused_id, 0)
+                    if count >= max_documents_from_same_page:
+                        break;
                     top_k_documents.append(doc["_source"])
-                    top_titles.append(doc["_source"]["title"])
-                    break
-            if len(top_titles) >= top_k:
-                break
+                    count_per_id[fused_id] = count + 1
 
         return top_k_documents, all_docs_and_scores
 
-    def answer_query(self, query, top_k: int, model, additional_document: dict = None, send_complete_pages_to_llm: bool = False):
+    def answer_query(self, query, top_k: int, model, additional_document: dict = None, send_complete_pages_to_llm: bool = False, retrieval_size: int = 50, max_documents_from_same_page:int=3):
         """
         Answers a query using the top_k documents and the specified model.
 
@@ -225,23 +226,27 @@ class Engine:
             model: The model to use for answering the query.
             additional_document (dict, optional): An additional document to include in the search. Default is None.
             send_complete_pages_to_llm (bool, optional): Whether to send complete pages to the
+            retrieval_size(int, optional): The number of documents to fetch from each model. Default is 50.
+            max_documents_from_same_page(int, optional): The maximum number of documents (paragraphs acutually) to return from the same page. Default is 3.
 
         Returns:
             tuple: A tuple containing the top k documents, the answer, and the stats.
         """
         before_retrieval = time.perf_counter()
-        top_k_documents, all_docs_and_scores = self.search_documents(query, top_k)
+        top_k_documents, all_docs_and_scores = self.search_documents(query=query, top_k=top_k, retrieval_size=retrieval_size,max_documents_from_same_page=max_documents_from_same_page)
 
         if send_complete_pages_to_llm:
             top_k_documents = [self.transform_document_to_full_page(doc) for doc in top_k_documents]
 
+        top_k_documents_and_additional_document = top_k_documents.copy()
+
         if additional_document:
-            top_k_documents.append(additional_document)
+            top_k_documents_and_additional_document.append(additional_document)
 
         retrieval_time = round(time.perf_counter() - before_retrieval, 4)
         print(f"retrieval time: {retrieval_time}")
 
-        gpt_answer, gpt_elapsed, tokens, request_params = self.llms_client.answer(query, top_k_documents)
+        gpt_answer, gpt_elapsed, tokens, request_params = self.llms_client.answer(query, top_k_documents_and_additional_document)
         stats = {
             "retrieval_time": retrieval_time,
             "gpt_model": model,

{kolzchut_ragbot-1.6.1 → kolzchut_ragbot-1.7.2}/kolzchut_ragbot.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kolzchut-ragbot
-Version: 1.6.1
+Version: 1.7.2
 Summary: A search engine using machine learning models and Elasticsearch for advanced document retrieval.
 Home-page: https://github.com/shmuelrob/rag-bot
 Author: Shmuel Robinov

{kolzchut_ragbot-1.6.1 → kolzchut_ragbot-1.7.2}/pyproject.toml

@@ -14,7 +14,7 @@ push            = false
 
 [tool.poetry]
 name = "ragbot"
-version = "1.6.1"
+version = "1.7.2"
 description = ""
 authors = ["Your Name <your.email@example.com>"]
 

{kolzchut_ragbot-1.6.1 → kolzchut_ragbot-1.7.2}/setup.py

@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='kolzchut-ragbot',
-    version='1.6.1',
+    version='1.7.2',
     author='Shmuel Robinov',
     author_email='shmuel_robinov@webiks.com',
     description='A search engine using machine learning models and Elasticsearch for advanced document retrieval.',

{kolzchut_ragbot-1.6.1 → kolzchut_ragbot-1.7.2}/test/test_engine.py

@@ -111,7 +111,7 @@ class TestEngine:
         mock_reciprocal_rank_fusion.return_value = [3, 2, 4, 1, 5]
         engine = build_test_engine(es_model, llm_client)
 
-        result, all_docs_and_scores = engine.search_documents("test query", 5)
+        result, all_docs_and_scores = engine.search_documents("test query", 5,50, 1)
 
         # Just verify the method completes without error
         assert result is not None
@@ -136,7 +136,9 @@ class TestEngine:
             ], {})
 
             llm_client.answer.return_value = ('answer', 0.5, 100, {})
-            actual_top_k_documents, actual_gpt_answer, actual_stats, all_docs_and_score, request_params = engine.answer_query("test query", 5, 'gpt-4o')
+            # Make retrieval_time deterministic (0)
+            with patch('time.perf_counter', side_effect=[100.0, 100.0]):
+                actual_top_k_documents, actual_gpt_answer, actual_stats, all_docs_and_score, request_params = engine.answer_query("test query", 5, 'gpt-4o')
 
             expected_top_k_documents = [
                 {'page_id': 3, 'title': 'title3'},