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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kolzchut-ragbot
-Version: 1.5.0
+Version: 1.6.1
 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.5.0 → kolzchut_ragbot-1.6.1}/kolzchut_ragbot/engine.py

@@ -7,7 +7,7 @@ from .model import es_client_factory
 from .Document import factory
 from sentence_transformers import SentenceTransformer
 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
 
@@ -215,7 +215,7 @@ class Engine:
 
         return top_k_documents, all_docs_and_scores
 
-    def answer_query(self, query, top_k: int, model):
+    def answer_query(self, query, top_k: int, model, additional_document: dict = None, send_complete_pages_to_llm: bool = False):
         """
         Answers a query using the top_k documents and the specified model.
 
@@ -223,6 +223,8 @@ class Engine:
             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
 
         Returns:
             tuple: A tuple containing the top k documents, the answer, and the stats.
@@ -230,18 +232,59 @@ class Engine:
         before_retrieval = time.perf_counter()
         top_k_documents, all_docs_and_scores = self.search_documents(query, top_k)
 
+        if send_complete_pages_to_llm:
+            top_k_documents = [self.transform_document_to_full_page(doc) for doc in top_k_documents]
+
+        if additional_document:
+            top_k_documents.append(additional_document)
 
         retrieval_time = round(time.perf_counter() - before_retrieval, 4)
         print(f"retrieval time: {retrieval_time}")
 
-        gpt_answer, gpt_elapsed, tokens = self.llms_client.answer(query, top_k_documents)
+        gpt_answer, gpt_elapsed, tokens, request_params = self.llms_client.answer(query, top_k_documents)
         stats = {
             "retrieval_time": retrieval_time,
             "gpt_model": model,
             "gpt_time": gpt_elapsed,
             "tokens": tokens
         }
-        return top_k_documents, gpt_answer, stats, all_docs_and_scores
+        return top_k_documents, gpt_answer, stats, all_docs_and_scores, request_params
+
+    def transform_document_to_full_page(self, document: dict) -> dict:
+        """
+        Adds the full page content to the document by retrieving it from Elasticsearch.
+
+        Args:
+            document (dict): The document to which the full page content will be added.
+
+        Returns:
+            dict: The updated document with the full page content added.
+        """
+        if not document.get("page_id"):
+            return document
+        full_document = self.get_full_document_by_page_id(document["page_id"])
+        if full_document and full_document.get("content"):
+            document["content"] = full_document["content"]
+        return document
+
+    def get_full_document_by_page_id(self, page_id: int) -> dict | None:
+        """
+        Retrieves a unified document instance for a given page_id by searching all indices in Elasticsearch.
+
+        Args:
+            page_id (int): The page ID to search for.
+
+        Returns:
+            dict | None: A single dict representing the united document (with metadata and concatenated content),
+                        or None if no documents are found.
+        """
+        es_client = self.elastic_model.es_client
+        indices = es_client.indices.get_alias(index="*").keys()
+        parts_of_documents = find_page_id_in_all_indices(page_id=page_id, es_client=es_client, indices=indices)
+        if not parts_of_documents:
+            return None
+        full_document = unite_docs_to_single_instance(parts_of_documents)
+        return full_document
 
 
 engine = None

kolzchut_ragbot-1.6.1/kolzchut_ragbot/get_full_documents_utilities.py

@@ -0,0 +1,45 @@
+from elasticsearch import Elasticsearch
+from typing import List, Dict, Any
+
+
+def find_page_id_in_all_indices(page_id: int, es_client: Elasticsearch, indices: List[str]) -> List[Dict[str, Any]]:
+    """
+    Search all provided indices using the given es_client for documents with the given page_id.
+    Returns a list of all matching documents' _source fields.
+    """
+    all_docs = []
+    fixed_indicies =  ["_".join(index.split("_")[1:]) for index in indices]
+    for index in fixed_indicies:
+        resp = es_client.search(index=index, body={
+            "query": {
+                "term": {"page_id": page_id}
+            }
+        }, size=100)
+        hits = resp.get('hits', {}).get('hits', [])
+        all_docs.extend([doc['_source'] for doc in hits])
+    return all_docs
+
+
+def unite_docs_to_single_instance(docs: List[Dict[str, Any]]) -> Dict[str, Any] | None:
+    """
+    Unites a list of Elasticsearch document dicts (with the same page_id) into a single instance dict.
+    - Takes metadata fields (page_id, title, url, articleType, articleContentArea, etc.) from the first document.
+    - Concatenates all 'content' fields with a line break between them.
+    - Returns a single dict representing the united document, or None if docs is empty.
+    """
+    if not docs:
+        return None
+    first = docs[0]
+    united_content = '\n'.join([doc.get('content', '') for doc in docs if doc.get('content')])
+    instance = {
+        "page_id": first.get("page_id"),
+        "title": first.get("title"),
+        "url": first.get("url"),
+        "link": first.get("url").split("/")[-1] if first.get("url") else None,
+        "articleType": first.get("articleType"),
+        "articleContentArea": first.get("articleContentArea"),
+        "summary": first.get("summary"),
+        "categories": first.get("categories"),
+        "content": united_content
+    }
+    return instance

{kolzchut_ragbot-1.5.0 → kolzchut_ragbot-1.6.1}/kolzchut_ragbot/llm_client.py

@@ -7,5 +7,5 @@ class LLMClient(ABC):
     def __init__(self):
         self.field_for_answer = definitions.field_for_llm
     @abstractmethod
-    def answer(self, _question, _top_k_docs) -> tuple[str, float, int]:
+    def answer(self, _question, _top_k_docs) -> tuple[str, float, int, dict]:
         raise NotImplementedError

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kolzchut-ragbot
-Version: 1.5.0
+Version: 1.6.1
 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.5.0 → kolzchut_ragbot-1.6.1}/kolzchut_ragbot.egg-info/SOURCES.txt

@@ -6,6 +6,7 @@ kolzchut_ragbot/IntegrateService.py
 kolzchut_ragbot/__init__.py
 kolzchut_ragbot/config.py
 kolzchut_ragbot/engine.py
+kolzchut_ragbot/get_full_documents_utilities.py
 kolzchut_ragbot/llm_client.py
 kolzchut_ragbot/model.py
 kolzchut_ragbot.egg-info/PKG-INFO

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

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

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

@@ -7,7 +7,7 @@ from setuptools import setup, find_packages
 
 setup(
     name='kolzchut-ragbot',
-    version='1.5.0',
+    version='1.6.1',
     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.5.0 → kolzchut_ragbot-1.6.1}/test/test_engine.py

@@ -135,8 +135,8 @@ class TestEngine:
                 {'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')
+            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')
 
             expected_top_k_documents = [
                 {'page_id': 3, 'title': 'title3'},
@@ -156,4 +156,3 @@ class TestEngine:
             assert expected_top_k_documents == actual_top_k_documents
             assert expected_gpt_answer == actual_gpt_answer
             assert expected_stats == actual_stats
-