kolzchut-ragbot 1.7.3__py3-none-any.whl → 1.7.5__py3-none-any.whl

kolzchut_ragbot/engine.py

@@ -10,6 +10,7 @@ from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from .get_full_documents_utilities import find_page_id_in_all_indices, unite_docs_to_single_instance
 import torch
 import os
+import asyncio
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 definitions = factory()
@@ -170,7 +171,7 @@ class Engine:
 
         return fused_list
 
-    def search_documents(self, query: str, top_k: int,retrieval_size: int, max_documents_from_same_page: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.
 
@@ -216,43 +217,129 @@ class Engine:
 
         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, retrieval_size: int = 50, max_documents_from_same_page:int=3):
+    def get_page_content_by_page_id(self, page_id: int) -> tuple:
         """
-        Answers a query using the top_k documents and the specified model.
+        Fetches the full content of a page and measures how long it takes.
 
         Args:
-            query (str): The query string.
-            top_k (int): The number of top documents to use for answering the query.
-            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.
+            page_id (int): The ID of the page to retrieve.
+
+        Returns:
+            tuple: (page_content, elapsed_time) where `page_content` is the retrieved content
+                   and `elapsed_time` is the time in seconds.
+        """
+        before_getting_additional_page = time.perf_counter()
+        additional_page_content = self.get_full_document_by_page_id(page_id)
+        after_getting_additional_page = time.perf_counter()
+        elapsed_time = after_getting_additional_page - before_getting_additional_page
+        return additional_page_content, elapsed_time
+
+    def retrieve_documents(self, query: str, top_k: int, retrieval_size: int,
+                           max_documents_from_same_page: int, send_complete_pages_to_llm: bool) -> tuple:
+        """
+        Retrieves documents matching a query and optionally converts them to full pages.
+
+        Args:
+            query (str): Search query.
+            top_k (int): Number of top documents to return.
+            retrieval_size (int): Number of documents to fetch from the source.
+            max_documents_from_same_page (int): Max documents from a single page.
+            send_complete_pages_to_llm (bool): If True, returns full page content.
 
         Returns:
-            tuple: A tuple containing the top k documents, the answer, and the stats.
+            tuple: (top_k_documents, all_docs_and_scores, retrieval_time)
         """
         before_retrieval = time.perf_counter()
-        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)
+        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()
+        retrieval_time = round(time.perf_counter() - before_retrieval, 4)
+        print(f"retrieval time: {retrieval_time}")
+
+        return top_k_documents, all_docs_and_scores, retrieval_time
+
+    async def answer_query(self, query: str, top_k: int, model, page_id: int | None = None,
+                           send_complete_pages_to_llm: bool = False, retrieval_size: int = 50,
+                           max_documents_from_same_page: int = 3) -> tuple:
+        """
+        Answers a query using top documents and an LLM model, optionally including a full page.
+
+        Args:
+            query (str): Query string.
+            top_k (int): Number of top documents to use.
+            model: LLM model to generate the answer.
+            page_id (int | None): Optional page to include.
+            send_complete_pages_to_llm (bool): If True, sends full pages to the LLM.
+            retrieval_size (int): Number of documents to fetch (default 50).
+            max_documents_from_same_page (int): Max documents from one page (default 3).
 
+        Returns:
+            tuple: (top_k_documents, gpt_answer, stats, all_docs_and_scores, request_params)
+        """
+        before_answer = time.perf_counter()
+
+        tasks = [
+            asyncio.to_thread(
+                self.retrieve_documents,
+                query, top_k, retrieval_size, max_documents_from_same_page, send_complete_pages_to_llm
+            )
+        ]
+
+        if page_id:
+            tasks.append(asyncio.to_thread(self.get_page_content_by_page_id, page_id))
+
+        results = await asyncio.gather(*tasks)
+
+        # Unpack results
+        top_k_documents, all_docs_and_scores, retrieval_time = results[0]
+        additional_document = None
+        additional_page_time = None
+        if page_id:
+            additional_document, additional_page_time = results[1]
+
+        # Combine documents
+        top_k_documents_and_additional_document = top_k_documents.copy()
+        # Remove documents with the same page_id as the additional_document before appending
         if additional_document:
+            additional_page_id = additional_document.get("page_id")
+            top_k_documents_and_additional_document = [
+                doc for doc in top_k_documents_and_additional_document
+                if doc.get("page_id") != additional_page_id
+            ]
             top_k_documents_and_additional_document.append(additional_document)
 
-        retrieval_time = round(time.perf_counter() - before_retrieval, 4)
-        print(f"retrieval time: {retrieval_time}")
+        # Query LLM
+        gpt_answer, gpt_elapsed, tokens, request_params = await asyncio.to_thread(
+            self.llms_client.answer,
+            query, top_k_documents_and_additional_document
+        )
+        after_answer = time.perf_counter()
+        answer_time = after_answer - before_answer
 
-        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,
             "gpt_time": gpt_elapsed,
-            "tokens": tokens
+            "tokens": tokens,
+            "answer_time": answer_time
         }
+
+        request_params['timers_ms'] = {
+            "answer_time": int(answer_time * 1000),
+            "retrieval_time": int(retrieval_time * 1000),
+            "llm_time": int(gpt_elapsed * 1000)
+        }
+
+        if additional_page_time:
+            request_params['timers_ms']['additional_page_time'] = int(additional_page_time * 1000)
+
         return top_k_documents, gpt_answer, stats, all_docs_and_scores, request_params
 
     def transform_document_to_full_page(self, document: dict) -> dict:

{kolzchut_ragbot-1.7.3.dist-info → kolzchut_ragbot-1.7.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kolzchut-ragbot
-Version: 1.7.3
+Version: 1.7.5
 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.7.3.dist-info → kolzchut_ragbot-1.7.5.dist-info}/RECORD

@@ -2,11 +2,11 @@ kolzchut_ragbot/Document.py,sha256=5OyBBTZyAJFM_1Pjs3SUC-_s5zEJ5U6wjhw12_FFkdE,3
 kolzchut_ragbot/IntegrateService.py,sha256=rcwUY2RkclCY3l8BGAmNbstdxhxwhLO9oA8BofqLyts,96
 kolzchut_ragbot/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kolzchut_ragbot/config.py,sha256=uILFvgn9W92-NRaKXYtaoQXpn3KOWKK8SZYRsIAa5Yw,133
-kolzchut_ragbot/engine.py,sha256=ttUxJRsX7p-sLuU4sC-W2GIt3pUClPchbckJAMBfT2k,13221
+kolzchut_ragbot/engine.py,sha256=-qNSnyq4A3rQZz_HU_5Djwc35Su0M4tAqNvqLlDQ_UU,16383
 kolzchut_ragbot/get_full_documents_utilities.py,sha256=YWljmGWM6h1ghLDCAUnDdhmn-0k6R_t7b1g7wSojzvg,1882
 kolzchut_ragbot/llm_client.py,sha256=JdDeOn2THpkOM2Mwe2DucTaYXul1fL2agIisBuHFtc8,347
 kolzchut_ragbot/model.py,sha256=HCi3r4YztPknnbgTOA7I-GVaqxn8CzrTeLFkEg-7fg0,6320
-kolzchut_ragbot-1.7.3.dist-info/METADATA,sha256=6weg2Y1MWvZ98nuG2laVp_aCVfxdetTrB5psXqCMBNs,1999
-kolzchut_ragbot-1.7.3.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-kolzchut_ragbot-1.7.3.dist-info/top_level.txt,sha256=NTZoY4GGw3v_7jm0MgcdHw8simoZ78PsR7Meqmkgd_Q,16
-kolzchut_ragbot-1.7.3.dist-info/RECORD,,
+kolzchut_ragbot-1.7.5.dist-info/METADATA,sha256=9jpBjTYYUzmuNuQ_t5L-fPjOJfKoCfBSEree9DnuazE,1999
+kolzchut_ragbot-1.7.5.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+kolzchut_ragbot-1.7.5.dist-info/top_level.txt,sha256=NTZoY4GGw3v_7jm0MgcdHw8simoZ78PsR7Meqmkgd_Q,16
+kolzchut_ragbot-1.7.5.dist-info/RECORD,,