kolzchut-ragbot 1.4.2__py3-none-any.whl → 1.5.0__py3-none-any.whl

kolzchut_ragbot/Document.py

@@ -1,101 +1,101 @@
-import json
-import os
-
-DEFINITIONS_FILE = os.getenv("DOCUMENT_DEFINITION_CONFIG", "example-conf.json")
-
-
-class DocumentFieldDefinition:
-    """
-    Represents the definition of a document field.
-
-    Attributes:
-        field_name (str): The name of the field.
-        required (bool): Indicates if the field is required. Default is False.
-    """
-    def __init__(self, field_name: str, required: bool = False):
-        self.field_name = field_name
-        self.required = required
-
-
-class DocumentDefinitions:
-    """
-    Represents the definitions for a document.
-
-    Attributes:
-        saved_fields (dict[str, DocumentFieldDefinition]): A dictionary of saved fields.
-        models (dict[str, str]): A dictionary of models.
-        identifier (str): The identifier field.
-        field_for_llm (str, optional): The field for LLM. Default is None.
-    """
-    def __init__(self, saved_fields: dict[str, DocumentFieldDefinition], models: dict[str, str],
-                 identifier: str, field_for_llm: str = None):
-        self.saved_fields = saved_fields
-        self.models = models
-        self.identifier = identifier
-        self.field_for_llm = field_for_llm
-
-
-def initialize_definitions():
-    """
-    Initializes the document definitions by reading the configuration file.
-
-    Raises:
-        ValueError: If the identifier field is not one of the saved fields or if any model field is not one of the saved fields.
-
-    Returns:
-        DocumentDefinitions: The initialized document definitions.
-    """
-    with open(DEFINITIONS_FILE, 'r', encoding='utf-8') as f:
-        definitions = json.load(f)
-
-        saved_fields = definitions['saved_fields']
-        models = definitions['models']
-        identifier_field = definitions['identifier_field']
-        field_for_llm = definitions.get('field_for_llm', None)
-        if identifier_field not in saved_fields.keys():
-            raise ValueError("identifier_field must be one of the saved fields, check the configuration file")
-
-        for embedded_field in models.values():
-            if embedded_field not in saved_fields.keys():
-                raise ValueError(f"{embedded_field} must be one of the saved fields {saved_fields.keys()}, check the configuration file")
-
-        return DocumentDefinitions(saved_fields, models, identifier_field, field_for_llm)
-
-
-definitions_singleton = None
-
-
-def factory():
-    """
-    Factory method to get the singleton instance of DocumentDefinitions.
-
-    Returns:
-        DocumentDefinitions: The singleton instance of document definitions.
-    """
-    global definitions_singleton
-    if definitions_singleton is None:
-        definitions_singleton = initialize_definitions()
-    return definitions_singleton
-
-
-class Document:
-    """
-    Represents a document.
-
-    Attributes:
-        page_id (str): The ID of the page.
-        fields (dict): The fields of the document.
-
-    Raises:
-        ValueError: If the fields do not match the required fields or if a required field is missing.
-    """
-    def __init__(self, page_id: str, fields: dict):
-        definitions = factory()
-        self.page_id = page_id
-        if fields.keys() != definitions.saved_fields.keys():
-            raise ValueError("fields do not match the required fields")
-        for defined_field in definitions.saved_fields.values():
-            if defined_field.required and defined_field.field_name not in fields:
-                raise ValueError(f"field {defined_field.field_name} is required")
-            if defined_field.field_name in fields:
-                setattr(self, defined_field.field_name, fields[defined_field.field_name])
+import json
+import os
+
+DEFINITIONS_FILE = os.getenv("DOCUMENT_DEFINITION_CONFIG", "example-conf.json")
+
+
+class DocumentFieldDefinition:
+    """
+    Represents the definition of a document field.
+
+    Attributes:
+        field_name (str): The name of the field.
+        required (bool): Indicates if the field is required. Default is False.
+    """
+    def __init__(self, field_name: str, required: bool = False):
+        self.field_name = field_name
+        self.required = required
+
+
+class DocumentDefinitions:
+    """
+    Represents the definitions for a document.
+
+    Attributes:
+        saved_fields (dict[str, DocumentFieldDefinition]): A dictionary of saved fields.
+        models (dict[str, str]): A dictionary of models.
+        identifier (str): The identifier field.
+        field_for_llm (str, optional): The field for LLM. Default is None.
+    """
+    def __init__(self, saved_fields: dict[str, DocumentFieldDefinition], models: dict[str, str],
+                 identifier: str, field_for_llm: str = None):
+        self.saved_fields = saved_fields
+        self.models = models
+        self.identifier = identifier
+        self.field_for_llm = field_for_llm
+
+
+def initialize_definitions():
+    """
+    Initializes the document definitions by reading the configuration file.
+
+    Raises:
+        ValueError: If the identifier field is not one of the saved fields or if any model field is not one of the saved fields.
+
+    Returns:
+        DocumentDefinitions: The initialized document definitions.
+    """
+    with open(DEFINITIONS_FILE, 'r', encoding='utf-8') as f:
+        definitions = json.load(f)
+
+        saved_fields = definitions['saved_fields']
+        models = definitions['models']
+        identifier_field = definitions['identifier_field']
+        field_for_llm = definitions.get('field_for_llm', None)
+        if identifier_field not in saved_fields.keys():
+            raise ValueError("identifier_field must be one of the saved fields, check the configuration file")
+
+        for embedded_field in models.values():
+            if embedded_field not in saved_fields.keys():
+                raise ValueError(f"{embedded_field} must be one of the saved fields {saved_fields.keys()}, check the configuration file")
+
+        return DocumentDefinitions(saved_fields, models, identifier_field, field_for_llm)
+
+
+definitions_singleton = None
+
+
+def factory():
+    """
+    Factory method to get the singleton instance of DocumentDefinitions.
+
+    Returns:
+        DocumentDefinitions: The singleton instance of document definitions.
+    """
+    global definitions_singleton
+    if definitions_singleton is None:
+        definitions_singleton = initialize_definitions()
+    return definitions_singleton
+
+
+class Document:
+    """
+    Represents a document.
+
+    Attributes:
+        page_id (str): The ID of the page.
+        fields (dict): The fields of the document.
+
+    Raises:
+        ValueError: If the fields do not match the required fields or if a required field is missing.
+    """
+    def __init__(self, page_id: str, fields: dict):
+        definitions = factory()
+        self.page_id = page_id
+        if fields.keys() != definitions.saved_fields.keys():
+            raise ValueError("fields do not match the required fields")
+        for defined_field in definitions.saved_fields.values():
+            if defined_field.required and defined_field.field_name not in fields:
+                raise ValueError(f"field {defined_field.field_name} is required")
+            if defined_field.field_name in fields:
+                setattr(self, defined_field.field_name, fields[defined_field.field_name])

kolzchut_ragbot/IntegrateService.py

@@ -1,4 +1,4 @@
-class IntegrateService:
-
-    def on_update_docs(self, _docs):
-        raise NotImplementedError
+class IntegrateService:
+
+    def on_update_docs(self, _docs):
+        raise NotImplementedError

kolzchut_ragbot/__init__.py

@@ -1,2 +0,0 @@
-__version__ = "1.0.53"
-

kolzchut_ragbot/config.py

@@ -1,5 +1,5 @@
-import os
-
-EMBEDDING_INDEX = os.getenv("ES_EMBEDDING_INDEX", "embeddings")
-MODELS_LOCATION = os.getenv("MODELS_LOCATION", "models")
-
+import os
+
+EMBEDDING_INDEX = os.getenv("ES_EMBEDDING_INDEX", "embeddings")
+MODELS_LOCATION = os.getenv("MODELS_LOCATION", "models")
+

kolzchut_ragbot/engine.py

@@ -1,246 +1,254 @@
-import time
-from collections import defaultdict
-from datetime import datetime
-from .llm_client import LLMClient
-from . import config
-from .model import es_client_factory
-from .Document import factory
-from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
-
-import torch
-import os
-
-device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
-definitions = factory()
-
-
-class Engine:
-    """
-    Engine class for handling document search and retrieval using Elasticsearch and LLMs.
-
-    Attributes:
-        llms_client (LLMClient): The LLM client instance.
-        elastic_model (Model): The Elasticsearch model instance.
-        models (dict): A dictionary of SentenceTransformer models.
-        reranker_tokenizer (AutoTokenizer): The tokenizer for the reranker model.
-        reranker_model (AutoModelForSequenceClassification): The reranker model.
-        identifier_field (str): The identifier field for documents.
-
-    Methods:
-        rerank_with_me5(query, documents, k=5):
-            Reranks documents based on the query using the reranker model.
-
-        update_docs(list_of_docs, embed_only_fields=None, delete_existing=False):
-            Updates or creates documents in the Elasticsearch index.
-
-        reciprocal_rank_fusion(ranking_lists, k=60, weights=None):
-            Performs Reciprocal Rank Fusion on a list of ranking lists.
-
-        search_documents(query, top_k):
-            Searches for documents based on the query and returns the top_k results.
-
-        answer_query(query, top_k, model):
-            Answers a query using the top_k documents and the specified model.
-    """
-
-    def __init__(self, llms_client: LLMClient, elastic_model=None, models=None, reranker_tokenizer=None,
-                 reranker_model=None, es_client=None):
-        """
-        Initializes the Engine instance.
-
-        Args:
-            llms_client (LLMClient): The LLM client instance.
-            elastic_model (Model, optional): The Elasticsearch model instance. Default is None.
-            models (dict, optional): A dictionary of SentenceTransformer models. Default is None.
-            reranker_tokenizer (AutoTokenizer, optional): The tokenizer for the reranker model. Default is None.
-            reranker_model (AutoModelForSequenceClassification, optional): The reranker model. Default is None.
-            es_client (optional): The Elasticsearch client instance. Default is None.
-        """
-        if elastic_model is None:
-            self.elastic_model = es_client_factory(es_client)
-        else:
-            self.elastic_model = elastic_model
-
-        self.llms_client = llms_client
-
-        self.identifier_field = factory().identifier
-
-        if models is None:
-            self.models = {f"{model_name}": SentenceTransformer(config.MODELS_LOCATION + "/" + model_name).to(device)
-                           for model_name in definitions.models.keys()}
-        else:
-            self.models = models
-        for model in self.models.values():
-            model.eval()
-
-        if reranker_tokenizer is None:
-            self.reranker_tokenizer = AutoTokenizer.from_pretrained(os.getenv("TOKENIZER_LOCATION"))
-        else:
-            self.reranker_tokenizer = reranker_tokenizer
-
-        if reranker_model is None:
-            self.reranker_model = AutoModelForSequenceClassification.from_pretrained(os.getenv("TOKENIZER_LOCATION"))
-        else:
-            self.reranker_model = reranker_model
-        self.reranker_model.eval()
-
-
-    def rerank_with_me5(self, query, documents, k=5):
-        """
-        Reranks documents based on the query using the reranker model.
-
-        Args:
-            query (str): The query string.
-            documents (list): A list of documents to be reranked.
-            k (int, optional): The number of top documents to return. Default is 5.
-
-        Returns:
-            list: A list of top k reranked documents.
-        """
-        pairs = [(query, doc) for doc in set(documents)]
-        inputs = self.reranker_tokenizer(pairs, return_tensors='pt', padding=True, truncation=True, max_length=512)
-
-        # Make predictions
-        with torch.no_grad():
-            outputs = self.reranker_model(**inputs)
-
-        scores = outputs.logits.squeeze()
-
-        if scores.ndim > 1:
-            scores = scores[:, 1]  # Assuming binary classification and index 1 is the relevance score
-
-        sorted_indices = torch.argsort(scores, descending=True)
-        # If there is only one document, return it to avoid torch error
-        if len(sorted_indices) == 1:
-            return [pairs[0][1]]
-        # Sort documents by their highest score
-        sorted_docs = [pairs[i][1] for i in sorted_indices]
-        return sorted_docs[:k]
-
-    def update_docs(self, list_of_docs: list[dict], embed_only_fields=None, delete_existing=False):
-        """
-        Updates or creates documents in the Elasticsearch index.
-
-        Args:
-            list_of_docs (list[dict]): A list of dictionaries representing the documents to be indexed.
-            embed_only_fields (list, optional): A list of fields to be embedded. Default is None.
-            delete_existing (bool, optional): Whether to delete existing documents. Default is False.
-        """
-        embed_only_fields = embed_only_fields or definitions.models.values()
-        for doc in list_of_docs:
-            for semantic_model, field in definitions.models.items():
-                if field in doc.keys() and field in embed_only_fields:
-                    content_vectors = self.models[semantic_model].encode(doc[field])
-                    doc[f'{field}_{semantic_model}_vectors'] = content_vectors
-
-            doc['last_update'] = datetime.now()
-        self.elastic_model.create_or_update_documents(list_of_docs, delete_existing)
-
-    def reciprocal_rank_fusion(self, ranking_lists, k=60, weights=None):
-        """
-        Performs Reciprocal Rank Fusion on a list of ranking lists.
-
-        Args:
-        :param ranking_lists: List of ranking lists, where each ranking list is a list of documents returned by a model.
-        :param k: The parameter for the reciprocal rank calculation (default is 60).
-        :param: weights: Optional. Weights for each ranking list.
-
-        Returns:
-            list: A fused ranking list of documents.
-        """
-        scores = defaultdict(float)
-
-        for list_index, rank_list in enumerate(ranking_lists):
-            for rank, identifier in enumerate(rank_list):
-                # Reciprocal rank score
-                w = weights[list_index] if weights else 1
-                scores[identifier] += w / (k + rank + 1)
-
-        # Sort the documents by their cumulative scores in descending order
-        fused_list = sorted(scores, key=scores.get, reverse=True)
-
-        return fused_list
-
-    def search_documents(self, query: str, top_k: 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.
-
-        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 = []
-        ids_for_fusion = []
-        all_docs_and_scores = {}
-
-        for key, values in all_docs_by_model.items():
-            print(f"\nFound {len(values)} documents for model\n")
-            model_ids = []
-            scores_for_model = []
-
-            for doc in values:
-                model_ids.append(doc["_source"]["page_id"])
-                all_docs.append(doc)
-                scores_for_model.append({"doc": doc["_source"]["title"], "score": doc["_score"]})
-            ids_for_fusion.append(model_ids)
-            all_docs_and_scores[f'{key}'] = scores_for_model
-        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 = []
-
-        for fused_id in fused_ids:
-            for doc in all_docs:
-                if doc["_source"]["page_id"] == fused_id and doc["_source"]["title"] not in top_titles:
-                    top_k_documents.append(doc["_source"])
-                    top_titles.append(doc["_source"]["title"])
-                    break
-            if len(top_titles) >= top_k:
-                break
-
-        return top_k_documents, all_docs_and_scores
-
-    def answer_query(self, query, top_k: int, model):
-        """
-        Answers a query using the top_k documents and the specified model.
-
-        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.
-
-        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)
-
-
-        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)
-        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
-
-
-engine = None
-
-
-def engine_factory(llms_client: LLMClient, es_client=None):
-    global engine
-    if engine is None:
-        engine = Engine(llms_client=llms_client, es_client=es_client)
-    return engine
+import time
+from collections import defaultdict
+from datetime import datetime
+from .llm_client import LLMClient
+from . import config
+from .model import es_client_factory
+from .Document import factory
+from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+
+import torch
+import os
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+definitions = factory()
+
+
+class Engine:
+    """
+    Engine class for handling document search and retrieval using Elasticsearch and LLMs.
+
+    Attributes:
+        llms_client (LLMClient): The LLM client instance.
+        elastic_model (Model): The Elasticsearch model instance.
+        models (dict): A dictionary of SentenceTransformer models.
+        reranker_tokenizer (AutoTokenizer): The tokenizer for the reranker model.
+        reranker_model (AutoModelForSequenceClassification): The reranker model.
+        identifier_field (str): The identifier field for documents.
+
+    Methods:
+        rerank_with_me5(query, documents, k=5):
+            Reranks documents based on the query using the reranker model.
+
+        update_docs(list_of_docs, embed_only_fields=None, delete_existing=False):
+            Updates or creates documents in the Elasticsearch index.
+
+        reciprocal_rank_fusion(ranking_lists, k=60, weights=None):
+            Performs Reciprocal Rank Fusion on a list of ranking lists.
+
+        search_documents(query, top_k):
+            Searches for documents based on the query and returns the top_k results.
+
+        answer_query(query, top_k, model):
+            Answers a query using the top_k documents and the specified model.
+    """
+
+    def __init__(self, llms_client: LLMClient, elastic_model=None, models=None, reranker_tokenizer=None,
+                 reranker_model=None, es_client=None):
+        """
+        Initializes the Engine instance.
+
+        Args:
+            llms_client (LLMClient): The LLM client instance.
+            elastic_model (Model, optional): The Elasticsearch model instance. Default is None.
+            models (dict, optional): A dictionary of SentenceTransformer models. Default is None.
+            reranker_tokenizer (AutoTokenizer, optional): The tokenizer for the reranker model. Default is None.
+            reranker_model (AutoModelForSequenceClassification, optional): The reranker model. Default is None.
+            es_client (optional): The Elasticsearch client instance. Default is None.
+        """
+        if elastic_model is None:
+            self.elastic_model = es_client_factory(es_client)
+        else:
+            self.elastic_model = elastic_model
+
+        self.llms_client = llms_client
+
+        self.identifier_field = factory().identifier
+
+        if models is None:
+            self.models = {f"{model_name}": SentenceTransformer(config.MODELS_LOCATION + "/" + model_name).to(device)
+                           for model_name in definitions.models.keys()}
+        else:
+            self.models = models
+        for model in self.models.values():
+            model.eval()
+
+        if reranker_tokenizer is None:
+            self.reranker_tokenizer = AutoTokenizer.from_pretrained(os.getenv("TOKENIZER_LOCATION"))
+        else:
+            self.reranker_tokenizer = reranker_tokenizer
+
+        if reranker_model is None:
+            self.reranker_model = AutoModelForSequenceClassification.from_pretrained(os.getenv("TOKENIZER_LOCATION"))
+        else:
+            self.reranker_model = reranker_model
+        self.reranker_model.eval()
+
+    def change_llm(self, llms_client: LLMClient):
+        """
+        Changes the LLM client for the Engine instance.
+
+        Args:
+            llms_client (LLMClient): The new LLM client instance.
+        """
+        self.llms_client = llms_client
+
+    def rerank_with_me5(self, query, documents, k=5):
+        """
+        Reranks documents based on the query using the reranker model.
+
+        Args:
+            query (str): The query string.
+            documents (list): A list of documents to be reranked.
+            k (int, optional): The number of top documents to return. Default is 5.
+
+        Returns:
+            list: A list of top k reranked documents.
+        """
+        pairs = [(query, doc) for doc in set(documents)]
+        inputs = self.reranker_tokenizer(pairs, return_tensors='pt', padding=True, truncation=True, max_length=512)
+
+        # Make predictions
+        with torch.no_grad():
+            outputs = self.reranker_model(**inputs)
+
+        scores = outputs.logits.squeeze()
+
+        if scores.ndim > 1:
+            scores = scores[:, 1]  # Assuming binary classification and index 1 is the relevance score
+
+        sorted_indices = torch.argsort(scores, descending=True)
+        # If there is only one document, return it to avoid torch error
+        if len(sorted_indices) == 1:
+            return [pairs[0][1]]
+        # Sort documents by their highest score
+        sorted_docs = [pairs[i][1] for i in sorted_indices]
+        return sorted_docs[:k]
+
+    def update_docs(self, list_of_docs: list[dict], embed_only_fields=None, delete_existing=False):
+        """
+        Updates or creates documents in the Elasticsearch index.
+
+        Args:
+            list_of_docs (list[dict]): A list of dictionaries representing the documents to be indexed.
+            embed_only_fields (list, optional): A list of fields to be embedded. Default is None.
+            delete_existing (bool, optional): Whether to delete existing documents. Default is False.
+        """
+        embed_only_fields = embed_only_fields or definitions.models.values()
+        for doc in list_of_docs:
+            for semantic_model, field in definitions.models.items():
+                if field in doc.keys() and field in embed_only_fields:
+                    content_vectors = self.models[semantic_model].encode(doc[field])
+                    doc[f'{field}_{semantic_model}_vectors'] = content_vectors
+
+            doc['last_update'] = datetime.now()
+        self.elastic_model.create_or_update_documents(list_of_docs, delete_existing)
+
+    def reciprocal_rank_fusion(self, ranking_lists, k=60, weights=None):
+        """
+        Performs Reciprocal Rank Fusion on a list of ranking lists.
+
+        Args:
+        :param ranking_lists: List of ranking lists, where each ranking list is a list of documents returned by a model.
+        :param k: The parameter for the reciprocal rank calculation (default is 60).
+        :param: weights: Optional. Weights for each ranking list.
+
+        Returns:
+            list: A fused ranking list of documents.
+        """
+        scores = defaultdict(float)
+
+        for list_index, rank_list in enumerate(ranking_lists):
+            for rank, identifier in enumerate(rank_list):
+                # Reciprocal rank score
+                w = weights[list_index] if weights else 1
+                scores[identifier] += w / (k + rank + 1)
+
+        # Sort the documents by their cumulative scores in descending order
+        fused_list = sorted(scores, key=scores.get, reverse=True)
+
+        return fused_list
+
+    def search_documents(self, query: str, top_k: 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.
+
+        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 = []
+        ids_for_fusion = []
+        all_docs_and_scores = {}
+
+        for key, values in all_docs_by_model.items():
+            print(f"\nFound {len(values)} documents for model\n")
+            model_ids = []
+            scores_for_model = []
+
+            for doc in values:
+                model_ids.append(doc["_source"]["page_id"])
+                all_docs.append(doc)
+                scores_for_model.append({"doc": doc["_source"]["title"], "score": doc["_score"]})
+            ids_for_fusion.append(model_ids)
+            all_docs_and_scores[f'{key}'] = scores_for_model
+        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 = []
+
+        for fused_id in fused_ids:
+            for doc in all_docs:
+                if doc["_source"]["page_id"] == fused_id and doc["_source"]["title"] not in top_titles:
+                    top_k_documents.append(doc["_source"])
+                    top_titles.append(doc["_source"]["title"])
+                    break
+            if len(top_titles) >= top_k:
+                break
+
+        return top_k_documents, all_docs_and_scores
+
+    def answer_query(self, query, top_k: int, model):
+        """
+        Answers a query using the top_k documents and the specified model.
+
+        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.
+
+        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)
+
+
+        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)
+        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
+
+
+engine = None
+
+
+def engine_factory(llms_client: LLMClient, es_client=None):
+    global engine
+    if engine is None:
+        engine = Engine(llms_client=llms_client, es_client=es_client)
+    return engine

kolzchut_ragbot/llm_client.py

@@ -1,11 +1,11 @@
-from .Document import factory
-from abc import ABC, abstractmethod
-definitions = factory()
-
-class LLMClient(ABC):
-    @abstractmethod
-    def __init__(self):
-        self.field_for_answer = definitions.field_for_llm
-    @abstractmethod
-    def answer(self, _question, _top_k_docs) -> tuple[str, float, int]:
-        raise NotImplementedError
+from .Document import factory
+from abc import ABC, abstractmethod
+definitions = factory()
+
+class LLMClient(ABC):
+    @abstractmethod
+    def __init__(self):
+        self.field_for_answer = definitions.field_for_llm
+    @abstractmethod
+    def answer(self, _question, _top_k_docs) -> tuple[str, float, int]:
+        raise NotImplementedError

kolzchut_ragbot/model.py

@@ -1,182 +1,182 @@
-import datetime
-import logging
-import os
-from .Document import factory as definitions_factory
-
-definitions_singleton = definitions_factory()
-EMBEDDING_INDEX = os.getenv("ES_EMBEDDING_INDEX", "embeddings")
-semantic_models = definitions_singleton.models.keys()
-
-
-def index_from_page_id(page_id: int):
-    """
-    Generates an index name based on the page ID.
-
-    Args:
-        page_id (int): The ID of the page.
-
-    Returns:
-        str: The generated index name.
-    """
-    index_postfix = round(page_id / 1000)
-    return EMBEDDING_INDEX + "_" + str(index_postfix)
-
-
-def create_mapping():
-    """
-    Creates a mapping for the model in Elasticsearch.
-    """
-    vector_fields = {f'{semantic_model}_{name}_vectors': {"type": "dense_vector", "dims": 1024}
-                     for name, semantic_model in definitions_singleton.models.items()}
-
-    data_fields = {}
-    for field in definitions_singleton.saved_fields.keys():
-        field_type = definitions_singleton.saved_fields[field]
-        field_mapping = {"type": field_type}
-        if field_type == "date":
-            field_mapping["format"] = "yyyyMMddHHmmss"
-        data_fields[f"{field}"] = field_mapping
-
-    mappings = {
-        "properties": {
-            "last_update": {
-                "type": "date",
-            },
-            **vector_fields,
-            **data_fields,
-        }
-    }
-    return mappings
-
-
-class Model:
-    """
-    Represents the model for creating, updating, and searching documents in Elasticsearch.
-
-    Attributes:
-        custom_result_selection_function (callable): A custom function for selecting search results.
-        es_client: The Elasticsearch client instance.
-
-    Methods:
-        create_index():
-            Creates an index for the model in Elasticsearch.
-
-        create_or_update_documents(paragraphs_dicts: list[dict], update=False):
-            Creates or updates documents in the Elasticsearch index.
-
-        search(embedded_search: dict[str, list[float]], size=50) -> dict[str, list[dict]]:
-            Searches for similar documents using cosine similarity.
-    """
-
-    custom_result_selection_function = None
-
-    def __init__(self, es_client, custom_result_selection_function=None):
-        """
-        Initializes the Model instance.
-
-        Args:
-            es_client: The Elasticsearch client instance.
-            custom_result_selection_function (callable, optional): A custom function for selecting search results.
-        """
-        self.es_client = es_client
-        if custom_result_selection_function is not None:
-            self.custom_result_selection_function = custom_result_selection_function
-
-    def create_index(self, index_name):
-        """
-        Creates an index for the model in Elasticsearch.
-        """
-        mapping = create_mapping()
-        if not self.es_client.indices.exists(index=index_name):
-            self.es_client.indices.create(
-                index=index_name,
-                mappings=mapping
-            )
-
-    def create_or_update_documents(self, paragraphs_dicts: list[dict], update=False):
-        """
-        Creates or updates documents in the Elasticsearch index.
-
-        Args:
-            paragraphs_dicts (list[dict]): A list of dictionaries representing the paragraphs to be indexed.
-            update (bool, optional): Whether to update existing documents. Default is False.
-        """
-
-        identifier = definitions_singleton.identifier
-        print(f"Creating or updating documents in the index, {len(paragraphs_dicts)} paragraphs\n")
-        # Identify the doc from the first paragraph - all paragraphs should have the same doc_id
-        doc_id = paragraphs_dicts[0][identifier]
-        index = index_from_page_id(int(doc_id))
-
-        if update:
-            try:
-                query = {
-                    "query": {
-                        "match": {
-                            f"{identifier}": doc_id
-                        }
-                    }
-                }
-                self.es_client.delete_by_query(index=f"{EMBEDDING_INDEX}*", body=query)
-
-            except Exception as e:
-                logging.error(f"Error while searching for existing document: {e}")
-        self.create_index(index)
-        for i, doc_dict in enumerate(paragraphs_dicts):
-            print(f"saving paragraph {i + 1} / {len(paragraphs_dicts)}")
-            doc = {
-                "last_update": datetime.datetime.now(),
-                **doc_dict
-            }
-
-            self.es_client.index(index=index, body=doc)
-
-    def search(self, embedded_search: dict[str, list[float]], size=50) -> dict[str, list[dict]]:
-        """
-        Searches for similar documents using cosine similarity.
-
-        Args:
-            embedded_search (dict[str, list[float]]): A dictionary containing the embedded search vectors.
-            size (int, optional): The number of search results to return. Default is 50.
-
-        Returns:
-            dict[str, list[dict]]: A dictionary containing the search results.
-        """
-        results = {}
-        for semantic_model, field in definitions_singleton.models.items():
-            results[field] = [] if field not in results.keys() else results[field]
-            body = {
-                "script_score": {
-                    "query": {
-                        "exists": {
-                            "field": f'{field}_{semantic_model}_vectors'
-                        }
-                    },
-                    "script": {
-                        "source": f"cosineSimilarity(params.query_vector, '{field}_{semantic_model}_vectors') + 1.0",
-                        "params": {
-                            "query_vector": embedded_search[semantic_model]
-                        }
-                    }
-                }
-            }
-            print(f"Searching for {field} using {semantic_model} on index {EMBEDDING_INDEX}\n")
-            field_results = self.es_client.search(
-                index=EMBEDDING_INDEX + "*",
-                body={
-                    "size": size,
-                    "query": body
-                })
-            results[field] = results[field] + field_results["hits"]["hits"]
-
-        return results
-
-
-model = None
-
-
-def es_client_factory(es_client) -> Model:
-    global model
-    if model is None:
-        model = Model(es_client)
-    return model
+import datetime
+import logging
+import os
+from .Document import factory as definitions_factory
+
+definitions_singleton = definitions_factory()
+EMBEDDING_INDEX = os.getenv("ES_EMBEDDING_INDEX", "embeddings")
+semantic_models = definitions_singleton.models.keys()
+
+
+def index_from_page_id(page_id: int):
+    """
+    Generates an index name based on the page ID.
+
+    Args:
+        page_id (int): The ID of the page.
+
+    Returns:
+        str: The generated index name.
+    """
+    index_postfix = round(page_id / 1000)
+    return EMBEDDING_INDEX + "_" + str(index_postfix)
+
+
+def create_mapping():
+    """
+    Creates a mapping for the model in Elasticsearch.
+    """
+    vector_fields = {f'{semantic_model}_{name}_vectors': {"type": "dense_vector", "dims": 1024}
+                     for name, semantic_model in definitions_singleton.models.items()}
+
+    data_fields = {}
+    for field in definitions_singleton.saved_fields.keys():
+        field_type = definitions_singleton.saved_fields[field]
+        field_mapping = {"type": field_type}
+        if field_type == "date":
+            field_mapping["format"] = "yyyyMMddHHmmss"
+        data_fields[f"{field}"] = field_mapping
+
+    mappings = {
+        "properties": {
+            "last_update": {
+                "type": "date",
+            },
+            **vector_fields,
+            **data_fields,
+        }
+    }
+    return mappings
+
+
+class Model:
+    """
+    Represents the model for creating, updating, and searching documents in Elasticsearch.
+
+    Attributes:
+        custom_result_selection_function (callable): A custom function for selecting search results.
+        es_client: The Elasticsearch client instance.
+
+    Methods:
+        create_index():
+            Creates an index for the model in Elasticsearch.
+
+        create_or_update_documents(paragraphs_dicts: list[dict], update=False):
+            Creates or updates documents in the Elasticsearch index.
+
+        search(embedded_search: dict[str, list[float]], size=50) -> dict[str, list[dict]]:
+            Searches for similar documents using cosine similarity.
+    """
+
+    custom_result_selection_function = None
+
+    def __init__(self, es_client, custom_result_selection_function=None):
+        """
+        Initializes the Model instance.
+
+        Args:
+            es_client: The Elasticsearch client instance.
+            custom_result_selection_function (callable, optional): A custom function for selecting search results.
+        """
+        self.es_client = es_client
+        if custom_result_selection_function is not None:
+            self.custom_result_selection_function = custom_result_selection_function
+
+    def create_index(self, index_name):
+        """
+        Creates an index for the model in Elasticsearch.
+        """
+        mapping = create_mapping()
+        if not self.es_client.indices.exists(index=index_name):
+            self.es_client.indices.create(
+                index=index_name,
+                mappings=mapping
+            )
+
+    def create_or_update_documents(self, paragraphs_dicts: list[dict], update=False):
+        """
+        Creates or updates documents in the Elasticsearch index.
+
+        Args:
+            paragraphs_dicts (list[dict]): A list of dictionaries representing the paragraphs to be indexed.
+            update (bool, optional): Whether to update existing documents. Default is False.
+        """
+
+        identifier = definitions_singleton.identifier
+        print(f"Creating or updating documents in the index, {len(paragraphs_dicts)} paragraphs\n")
+        # Identify the doc from the first paragraph - all paragraphs should have the same doc_id
+        doc_id = paragraphs_dicts[0][identifier]
+        index = index_from_page_id(int(doc_id))
+
+        if update:
+            try:
+                query = {
+                    "query": {
+                        "match": {
+                            f"{identifier}": doc_id
+                        }
+                    }
+                }
+                self.es_client.delete_by_query(index=f"{EMBEDDING_INDEX}*", body=query)
+
+            except Exception as e:
+                logging.error(f"Error while searching for existing document: {e}")
+        self.create_index(index)
+        for i, doc_dict in enumerate(paragraphs_dicts):
+            print(f"saving paragraph {i + 1} / {len(paragraphs_dicts)}")
+            doc = {
+                "last_update": datetime.datetime.now(),
+                **doc_dict
+            }
+
+            self.es_client.index(index=index, body=doc)
+
+    def search(self, embedded_search: dict[str, list[float]], size=50) -> dict[str, list[dict]]:
+        """
+        Searches for similar documents using cosine similarity.
+
+        Args:
+            embedded_search (dict[str, list[float]]): A dictionary containing the embedded search vectors.
+            size (int, optional): The number of search results to return. Default is 50.
+
+        Returns:
+            dict[str, list[dict]]: A dictionary containing the search results.
+        """
+        results = {}
+        for semantic_model, field in definitions_singleton.models.items():
+            results[field] = [] if field not in results.keys() else results[field]
+            body = {
+                "script_score": {
+                    "query": {
+                        "exists": {
+                            "field": f'{field}_{semantic_model}_vectors'
+                        }
+                    },
+                    "script": {
+                        "source": f"cosineSimilarity(params.query_vector, '{field}_{semantic_model}_vectors') + 1.0",
+                        "params": {
+                            "query_vector": embedded_search[semantic_model]
+                        }
+                    }
+                }
+            }
+            print(f"Searching for {field} using {semantic_model} on index {EMBEDDING_INDEX}\n")
+            field_results = self.es_client.search(
+                index=EMBEDDING_INDEX + "*",
+                body={
+                    "size": size,
+                    "query": body
+                })
+            results[field] = results[field] + field_results["hits"]["hits"]
+
+        return results
+
+
+model = None
+
+
+def es_client_factory(es_client) -> Model:
+    global model
+    if model is None:
+        model = Model(es_client)
+    return model

kolzchut_ragbot-1.5.0.dist-info/METADATA

@@ -0,0 +1,76 @@
+Metadata-Version: 2.4
+Name: kolzchut-ragbot
+Version: 1.5.0
+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
+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.17.1
+Requires-Dist: sentence-transformers==3.4.1
+Requires-Dist: torch==2.6.0
+Requires-Dist: transformers==4.48.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
+
+# kolzchut-ragbot
+
+## Overview
+
+This project is a search engine that uses machine learning models and Elasticsearch to provide advanced document retrieval.
+You can use [kolzchut-ragbot](https://github.com/shmuelrob/rag-bot) 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
+
+### From PyPI
+
+```bash
+pip install kolzchut-ragbot
+```
+
+### From Source
+
+1. Clone the repository:
+   
+   ```bash
+   git clone https://github.com/shmuelrob/rag-bot.git
+   cd rag-bot
+   ```
+
+2. Create a virtual environment and activate it:  
+
+   ```bash
+   python -m venv venv
+   source venv/bin/activate  # On Windows use: venv\Scripts\activate
+   ```
+
+3. Install the required dependencies:  
+
+   ```bash
+   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.5.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+kolzchut_ragbot/Document.py,sha256=5OyBBTZyAJFM_1Pjs3SUC-_s5zEJ5U6wjhw12_FFkdE,3621
+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=tULiiCdqLoXydxx7VrZKChDgbe4ygSoEft-k76j3_t0,10158
+kolzchut_ragbot/llm_client.py,sha256=Frp7CL0OIlQA6ltohrGWedI6uD6MpGg6TbpZTBE0qIo,341
+kolzchut_ragbot/model.py,sha256=HCi3r4YztPknnbgTOA7I-GVaqxn8CzrTeLFkEg-7fg0,6320
+kolzchut_ragbot-1.5.0.dist-info/METADATA,sha256=nzWcZsfn23nVWANcWMGoQvQXJ-idymUKg0ZheMbWp1c,1999
+kolzchut_ragbot-1.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+kolzchut_ragbot-1.5.0.dist-info/top_level.txt,sha256=NTZoY4GGw3v_7jm0MgcdHw8simoZ78PsR7Meqmkgd_Q,16
+kolzchut_ragbot-1.5.0.dist-info/RECORD,,

{kolzchut_ragbot-1.4.2.dist-info → kolzchut_ragbot-1.5.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

kolzchut_ragbot-1.4.2.dist-info/METADATA

@@ -1,67 +0,0 @@
-Metadata-Version: 2.4
-Name: kolzchut-ragbot
-Version: 1.4.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
-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.17.1
-Requires-Dist: sentence-transformers==3.4.1
-Requires-Dist: torch==2.6.0
-Requires-Dist: transformers==4.48.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.4.2.dist-info/RECORD

@@ -1,11 +0,0 @@
-kolzchut_ragbot/Document.py,sha256=ySawnD06HA0zHjHp4Y_CPjMMZqLp8onaEgd1dGP5sbs,3722
-kolzchut_ragbot/IntegrateService.py,sha256=CqB9vW6W5oj6Ig3crEa6hXqwro21z97UaG9ngxFTzYs,100
-kolzchut_ragbot/__init__.py,sha256=KKAc2xjCl5Aui2Cj0FWyvJ51nmnFv7MspLMqOYb-QHA,26
-kolzchut_ragbot/config.py,sha256=pcKVJVJ8P2YximjTrmVlrocHXSmzmNu_DFzNoPLa22E,138
-kolzchut_ragbot/engine.py,sha256=V8WUWyqvBWbGt-rRRf8G6BEyD-4GjsmtJrxBb6aPon8,10154
-kolzchut_ragbot/llm_client.py,sha256=q_cUZq645P7i1PliYzpJRTWlsoSECVIhE-y9wU5eRtQ,352
-kolzchut_ragbot/model.py,sha256=M7i9B-zzwa-ATblY-5c7gmbkOXKwS8wWmYMP8l0HE40,6502
-kolzchut_ragbot-1.4.2.dist-info/METADATA,sha256=YctMSApfbXBBvx7d0APy28lDWmulO3cfHJKjKOZ_LYQ,2024
-kolzchut_ragbot-1.4.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-kolzchut_ragbot-1.4.2.dist-info/top_level.txt,sha256=NTZoY4GGw3v_7jm0MgcdHw8simoZ78PsR7Meqmkgd_Q,16
-kolzchut_ragbot-1.4.2.dist-info/RECORD,,