cwyodmodules 0.3.32__py3-none-any.whl → 0.3.33__py3-none-any.whl

cwyodmodules/batch/utilities/search/postgres_search_handler.py

@@ -6,134 +6,137 @@ from .search_handler_base import SearchHandlerBase
 from ..helpers.azure_postgres_helper import AzurePostgresHelper
 from ..common.source_document import SourceDocument
 
-from logging import getLogger
-from opentelemetry import trace, baggage
-from opentelemetry.propagate import extract
+from ...utilities.helpers.env_helper import EnvHelper
+from logging_config import logger
+env_helper: EnvHelper = EnvHelper()
+log_args = env_helper.LOG_ARGS
+log_result = env_helper.LOG_RESULT
 
 
 class AzurePostgresHandler(SearchHandlerBase):
 
     def __init__(self, env_helper):
         self.azure_postgres_helper = AzurePostgresHelper()
-        #self.logger = getLogger("__main__" + ".base_package")
-        self.logger = getLogger("__main__")
-        #self.tracer = trace.get_tracer("__main__" + ".base_package")
-        self.tracer = trace.get_tracer("__main__")
         super().__init__(env_helper)
 
+
+    @logger.trace_function(log_args=False, log_result=False)
     def query_search(self, question) -> List[SourceDocument]:
-        with self.tracer.start_as_current_span("query_search"):
-            self.logger.info(f"Starting query search for question: {question}")
-            user_input = question
-            query_embedding = self.azure_postgres_helper.llm_helper.generate_embeddings(
-                user_input
-            )
+        logger.info(f"Starting query search for question: {question}")
+        user_input = question
+        query_embedding = self.azure_postgres_helper.llm_helper.generate_embeddings(
+            user_input
+        )
 
-            embedding_array = np.array(query_embedding).tolist()
+        embedding_array = np.array(query_embedding).tolist()
 
-            search_results = self.azure_postgres_helper.get_vector_store(
-                embedding_array
-            )
+        search_results = self.azure_postgres_helper.get_vector_store(
+            embedding_array
+        )
 
-            source_documents = self._convert_to_source_documents(search_results)
-            self.logger.info(f"Found {len(source_documents)} source documents.")
-            return source_documents
+        source_documents = self._convert_to_source_documents(search_results)
+        logger.info(f"Found {len(source_documents)} source documents.")
+        return source_documents
 
+    @logger.trace_function(log_args=False, log_result=False)
     def _convert_to_source_documents(self, search_results) -> List[SourceDocument]:
-        with self.tracer.start_as_current_span("_convert_to_source_documents"):
-            source_documents = []
-            for source in search_results:
-                source_document = SourceDocument(
-                    id=source["id"],
-                    title=source["title"],
-                    chunk=source["chunk"],
-                    offset=source["offset"],
-                    page_number=source["page_number"],
-                    content=source["content"],
-                    source=source["source"],
-                )
-                source_documents.append(source_document)
-            return source_documents
+        source_documents = []
+        for source in search_results:
+            source_document = SourceDocument(
+                id=source["id"],
+                title=source["title"],
+                chunk=source["chunk"],
+                offset=source["offset"],
+                page_number=source["page_number"],
+                content=source["content"],
+                source=source["source"],
+            )
+            source_documents.append(source_document)
+        return source_documents
+
 
+    @logger.trace_function(log_args=log_args, log_result=False)
     def create_search_client(self):
-        with self.tracer.start_as_current_span("create_search_client"):
-            return self.azure_postgres_helper.get_search_client()
+        return self.azure_postgres_helper.get_search_client()
+
 
+    @logger.trace_function(log_args=False, log_result=False)
     def create_vector_store(self, documents_to_upload):
-        with self.tracer.start_as_current_span("create_vector_store"):
-            self.logger.info(
-                f"Creating vector store with {len(documents_to_upload)} documents."
-            )
-            return self.azure_postgres_helper.create_vector_store(documents_to_upload)
+        logger.info(
+            f"Creating vector store with {len(documents_to_upload)} documents."
+        )
+        return self.azure_postgres_helper.create_vector_store(documents_to_upload)
 
+    @logger.trace_function(log_args=log_args, log_result=False)
     def perform_search(self, filename):
-        with self.tracer.start_as_current_span("perform_search"):
-            self.logger.info(f"Performing search for filename: {filename}")
-            return self.azure_postgres_helper.perform_search(filename)
+        logger.info(f"Performing search for filename: {filename}")
+        return self.azure_postgres_helper.perform_search(filename)
 
+    @logger.trace_function(log_args=False, log_result=False)
     def process_results(self, results):
-        with self.tracer.start_as_current_span("process_results"):
-            if results is None:
-                self.logger.info("No results to process.")
-                return []
-            data = [
-                [json.loads(result["metadata"]).get("chunk", i), result["content"]]
-                for i, result in enumerate(results)
-            ]
-            self.logger.info(f"Processed {len(data)} results.")
-            return data
-
+        if results is None:
+            logger.info("No results to process.")
+            return []
+        data = [
+            [json.loads(result["metadata"]).get("chunk", i), result["content"]]
+            for i, result in enumerate(results)
+        ]
+        logger.info(f"Processed {len(data)} results.")
+        return data
+
+    @logger.trace_function(log_args=log_args, log_result=False)
     def get_files(self):
-        with self.tracer.start_as_current_span("get_files"):
-            results = self.azure_postgres_helper.get_files()
-            if results is None or len(results) == 0:
-                self.logger.info("No files found.")
-                return []
-            self.logger.info(f"Found {len(results)} files.")
-            return results
+        results = self.azure_postgres_helper.get_files()
+        if results is None or len(results) == 0:
+            logger.info("No files found.")
+            return []
+        logger.info(f"Found {len(results)} files.")
+        return results
+
 
+    @logger.trace_function(log_args=False, log_result=False)
     def output_results(self, results):
-        with self.tracer.start_as_current_span("output_results"):
-            files = {}
-            for result in results:
-                id = result["id"]
-                filename = result["title"]
-                if filename in files:
-                    files[filename].append(id)
-                else:
-                    files[filename] = [id]
-
-            return files
+        files = {}
+        for result in results:
+            id = result["id"]
+            filename = result["title"]
+            if filename in files:
+                files[filename].append(id)
+            else:
+                files[filename] = [id]
 
+        return files
+
+
+    @logger.trace_function(log_args=log_args, log_result=log_result)
     def delete_files(self, files):
-        with self.tracer.start_as_current_span("delete_files"):
-            ids_to_delete = []
-            files_to_delete = []
+        ids_to_delete = []
+        files_to_delete = []
 
-            for filename, ids in files.items():
-                files_to_delete.append(filename)
-                ids_to_delete += [{"id": id} for id in ids]
-            self.azure_postgres_helper.delete_documents(ids_to_delete)
+        for filename, ids in files.items():
+            files_to_delete.append(filename)
+            ids_to_delete += [{"id": id} for id in ids]
+        self.azure_postgres_helper.delete_documents(ids_to_delete)
 
-            return ", ".join(files_to_delete)
+        return ", ".join(files_to_delete)
 
+    @logger.trace_function(log_args=log_args, log_result=False)
     def search_by_blob_url(self, blob_url):
-        with self.tracer.start_as_current_span("search_by_blob_url"):
-            self.logger.info(f"Searching by blob URL: {blob_url}")
-            return self.azure_postgres_helper.search_by_blob_url(blob_url)
+        logger.info(f"Searching by blob URL: {blob_url}")
+        return self.azure_postgres_helper.search_by_blob_url(blob_url)
 
+    @logger.trace_function(log_args=log_args, log_result=log_result)
     def delete_from_index(self, blob_url) -> None:
-        with self.tracer.start_as_current_span("delete_from_index"):
-            self.logger.info(f"Deleting from index for blob URL: {blob_url}")
-            documents = self.search_by_blob_url(blob_url)
-            if documents is None or len(documents) == 0:
-                self.logger.info("No documents found for blob URL.")
-                return
-            files_to_delete = self.output_results(documents)
-            self.delete_files(files_to_delete)
-
+        logger.info(f"Deleting from index for blob URL: {blob_url}")
+        documents = self.search_by_blob_url(blob_url)
+        if documents is None or len(documents) == 0:
+            logger.info("No documents found for blob URL.")
+            return
+        files_to_delete = self.output_results(documents)
+        self.delete_files(files_to_delete)
+
+    @logger.trace_function(log_args=log_args, log_result=log_result)
     def get_unique_files(self):
-        with self.tracer.start_as_current_span("get_unique_files"):
-            results = self.azure_postgres_helper.get_unique_files()
-            unique_titles = [row["title"] for row in results]
-            return unique_titles
+        results = self.azure_postgres_helper.get_unique_files()
+        unique_titles = [row["title"] for row in results]
+        return unique_titles

cwyodmodules/batch/utilities/search/postgres_search_handler_light_rag.py

@@ -7,8 +7,11 @@ from ..helpers.azure_postgres_helper import AzurePostgresHelper
 from ..helpers.lightrag_helper import LightRAGHelper
 from ..common.source_document import SourceDocument
 
-from logging import getLogger
-from opentelemetry import trace
+from ...utilities.helpers.env_helper import EnvHelper
+from logging_config import logger
+env_helper: EnvHelper = EnvHelper()
+log_args = env_helper.LOG_ARGS
+log_result = env_helper.LOG_RESULT
 
 
 class AzurePostgresHandler(SearchHandlerBase):
@@ -16,132 +19,128 @@ class AzurePostgresHandler(SearchHandlerBase):
     def __init__(self, env_helper):
         self.azure_postgres_helper = AzurePostgresHelper()
         self.lightrag_helper = LightRAGHelper()
-        #self.logger = getLogger("__main__" + ".base_package")
-        self.logger = getLogger("__main__")
-        #self.tracer = trace.get_tracer("__main__" + ".base_package")
-        self.tracer = trace.get_tracer("__main__")
         super().__init__(env_helper)
 
+    @logger.trace_function(log_args=False, log_result=False)
     def query_search(self, question) -> List[SourceDocument]:
-        with self.tracer.start_as_current_span("query_search"):
-            self.logger.info(f"Starting query search for question: {question}")
-            user_input = question
-            query_embedding = self.azure_postgres_helper.llm_helper.generate_embeddings(
-                user_input
-            )
+        logger.info(f"Starting query search for question: {question}")
+        user_input = question
+        query_embedding = self.azure_postgres_helper.llm_helper.generate_embeddings(
+            user_input
+        )
 
-            embedding_array = np.array(query_embedding).tolist()
+        embedding_array = np.array(query_embedding).tolist()
 
-            search_results = self.azure_postgres_helper.get_vector_store(
-                embedding_array
-            )
+        search_results = self.azure_postgres_helper.get_vector_store(
+            embedding_array
+        )
 
-            source_documents = self._convert_to_source_documents(search_results)
-            self.logger.info(f"Found {len(source_documents)} source documents.")
-            return source_documents
+        source_documents = self._convert_to_source_documents(search_results)
+        logger.info(f"Found {len(source_documents)} source documents.")
+        return source_documents
 
+    @logger.trace_function(log_args=False, log_result=False)
     def _convert_to_source_documents(self, search_results) -> List[SourceDocument]:
-        with self.tracer.start_as_current_span("_convert_to_source_documents"):
-            source_documents = []
-            for source in search_results:
-                source_document = SourceDocument(
-                    id=source["id"],
-                    title=source["title"],
-                    chunk=source["chunk"],
-                    offset=source["offset"],
-                    page_number=source["page_number"],
-                    content=source["content"],
-                    source=source["source"],
-                )
-                source_documents.append(source_document)
-            return source_documents
+        source_documents = []
+        for source in search_results:
+            source_document = SourceDocument(
+                id=source["id"],
+                title=source["title"],
+                chunk=source["chunk"],
+                offset=source["offset"],
+                page_number=source["page_number"],
+                content=source["content"],
+                source=source["source"],
+            )
+            source_documents.append(source_document)
+        return source_documents
 
+    @logger.trace_function(log_args=log_args, log_result=False)
     def create_search_client(self):
-        with self.tracer.start_as_current_span("create_search_client"):
-            return self.azure_postgres_helper.get_search_client()
+        return self.azure_postgres_helper.get_search_client()
 
+    @logger.trace_function(log_args=False, log_result=False)
     def create_vector_store(self, documents_to_upload):
-        with self.tracer.start_as_current_span("create_vector_store"):
-            self.logger.info(
-                f"Creating vector store with {len(documents_to_upload)} documents."
-            )
-            return self.azure_postgres_helper.create_vector_store(documents_to_upload)
+        logger.info(
+            f"Creating vector store with {len(documents_to_upload)} documents."
+        )
+        return self.azure_postgres_helper.create_vector_store(documents_to_upload)
 
+    @logger.trace_function(log_args=log_args, log_result=False)
     def perform_search(self, filename):
-        with self.tracer.start_as_current_span("perform_search"):
-            self.logger.info(f"Performing search for filename: {filename}")
-            return self.azure_postgres_helper.perform_search(filename)
+        logger.info(f"Performing search for filename: {filename}")
+        return self.azure_postgres_helper.perform_search(filename)
 
+    @logger.trace_function(log_args=False, log_result=False)
     def process_results(self, results):
-        with self.tracer.start_as_current_span("process_results"):
-            if results is None:
-                self.logger.info("No results to process.")
-                return []
-            data = [
-                [json.loads(result["metadata"]).get("chunk", i), result["content"]]
-                for i, result in enumerate(results)
-            ]
-            self.logger.info(f"Processed {len(data)} results.")
-            return data
-
+        if results is None:
+            logger.info("No results to process.")
+            return []
+        data = [
+            [json.loads(result["metadata"]).get("chunk", i), result["content"]]
+            for i, result in enumerate(results)
+        ]
+        logger.info(f"Processed {len(data)} results.")
+        return data
+
+    @logger.trace_function(log_args=log_args, log_result=False)
     def get_files(self):
-        with self.tracer.start_as_current_span("get_files"):
-            results = self.azure_postgres_helper.get_files()
-            if results is None or len(results) == 0:
-                self.logger.info("No files found.")
-                return []
-            self.logger.info(f"Found {len(results)} files.")
-            return results
-
+        results = self.azure_postgres_helper.get_files()
+        if results is None or len(results) == 0:
+            logger.info("No files found.")
+            return []
+        logger.info(f"Found {len(results)} files.")
+        return results
+
+    @logger.trace_function(log_args=False, log_result=False)
     def output_results(self, results):
-        with self.tracer.start_as_current_span("output_results"):
-            files = {}
-            for result in results:
-                id = result["id"]
-                filename = result["title"]
-                if filename in files:
-                    files[filename].append(id)
-                else:
-                    files[filename] = [id]
-
-            return files
-
+        files = {}
+        for result in results:
+            id = result["id"]
+            filename = result["title"]
+            if filename in files:
+                files[filename].append(id)
+            else:
+                files[filename] = [id]
+
+        return files
+
+    @logger.trace_function(log_args=log_args, log_result=log_result)
     def delete_files(self, files):
-        with self.tracer.start_as_current_span("delete_files"):
-            ids_to_delete = []
-            files_to_delete = []
-
-            for filename, ids in files.items():
-                files_to_delete.append(filename)
-                ids_to_delete += [{"id": id} for id in ids]
-            self.azure_postgres_helper.delete_documents(ids_to_delete)
+        ids_to_delete = []
+        files_to_delete = []
+        for filename, ids in files.items():
+            files_to_delete.append(filename)
+            ids_to_delete += [{"id": id} for id in ids]
+        self.azure_postgres_helper.delete_documents(ids_to_delete)
 
-            return ", ".join(files_to_delete)
+        return ", ".join(files_to_delete)
 
+    @logger.trace_function(log_args=log_args, log_result=False)
     def search_by_blob_url(self, blob_url):
-        with self.tracer.start_as_current_span("search_by_blob_url"):
-            self.logger.info(f"Searching by blob URL: {blob_url}")
-            return self.azure_postgres_helper.search_by_blob_url(blob_url)
+        logger.info(f"Searching by blob URL: {blob_url}")
+        return self.azure_postgres_helper.search_by_blob_url(blob_url)
 
+    @logger.trace_function(log_args=log_args, log_result=log_result)
     def delete_from_index(self, blob_url) -> None:
-        with self.tracer.start_as_current_span("delete_from_index"):
-            self.logger.info(f"Deleting from index for blob URL: {blob_url}")
-            documents = self.search_by_blob_url(blob_url)
-            if documents is None or len(documents) == 0:
-                self.logger.info("No documents found for blob URL.")
-                return
-            files_to_delete = self.output_results(documents)
-            self.delete_files(files_to_delete)
-
+        logger.info(f"Deleting from index for blob URL: {blob_url}")
+        documents = self.search_by_blob_url(blob_url)
+        if documents is None or len(documents) == 0:
+            logger.info("No documents found for blob URL.")
+            return
+        files_to_delete = self.output_results(documents)
+        self.delete_files(files_to_delete)
+
+    @logger.trace_function(log_args=log_args, log_result=False)
     def get_unique_files(self):
-        with self.tracer.start_as_current_span("get_unique_files"):
-            results = self.azure_postgres_helper.get_unique_files()
-            unique_titles = [row["title"] for row in results]
-            return unique_titles
+        results = self.azure_postgres_helper.get_unique_files()
+        unique_titles = [row["title"] for row in results]
+        return unique_titles
 
+
+    @logger.trace_function(log_args=False, log_result=False) 
     def store_vector_and_text(self, documents_to_store):
-        with self.tracer.start_as_current_span("store_vector_and_text"):
-            self.logger.info(
-                f"Storing {len(documents_to_store)} documents with LightRAG."
-            )
-            self.lightrag_helper.store_documents(documents_to_store)
+        logger.info(
+            f"Storing {len(documents_to_store)} documents with LightRAG."
+        )
+        self.lightrag_helper.store_documents(documents_to_store)

cwyodmodules/batch/utilities/search/search.py

@@ -8,15 +8,12 @@ from ..search.search_handler_base import SearchHandlerBase
 from ..common.source_document import SourceDocument
 from ..helpers.env_helper import EnvHelper
 
-from logging import getLogger
-from opentelemetry import trace, baggage
-from opentelemetry.propagate import extract
 from ..search.lightrag_search_handler import LightRAGSearchHandler
 
-# logger = getLogger("__main__" + ".base_package")
-logger = getLogger("__main__")
-# tracer = trace.get_tracer("__main__" + ".base_package")
-tracer = trace.get_tracer("__main__")
+from logging_config import logger
+env_helper: EnvHelper = EnvHelper()
+log_args = env_helper.LOG_ARGS
+log_result = env_helper.LOG_RESULT
 
 
 class Search:
@@ -27,6 +24,7 @@ class Search:
     """
 
     @staticmethod
+    @logger.trace_function(log_args=False, log_result=False)
     def get_search_handler(env_helper: EnvHelper) -> SearchHandlerBase:
         """
         Determines and returns the appropriate search handler based on the
@@ -44,23 +42,23 @@ class Search:
             Since the full workflow for PostgreSQL indexing is not yet complete,
             you can comment out the condition checking for PostgreSQL.
         """
-        with tracer.start_as_current_span("get_search_handler"):
-            logger.info("Getting search handler...")
-            # if env_helper.DATABASE_TYPE == DatabaseType.POSTGRESQL.value:
-            #     logger.info("Using LightRAGSearchHandler.")
-            #     return LightRAGSearchHandler(env_helper)
-            if env_helper.DATABASE_TYPE == DatabaseType.POSTGRESQL.value:
-                logger.info("Using AzurePostgresHandler.")
-                return AzurePostgresHandler(env_helper)
+        logger.info("Getting search handler...")
+        # if env_helper.DATABASE_TYPE == DatabaseType.POSTGRESQL.value:
+        #     logger.info("Using LightRAGSearchHandler.")
+        #     return LightRAGSearchHandler(env_helper)
+        if env_helper.DATABASE_TYPE == DatabaseType.POSTGRESQL.value:
+            logger.info("Using AzurePostgresHandler.")
+            return AzurePostgresHandler(env_helper)
+        else:
+            if env_helper.AZURE_SEARCH_USE_INTEGRATED_VECTORIZATION:
+                logger.info("Using IntegratedVectorizationSearchHandler.")
+                return IntegratedVectorizationSearchHandler(env_helper)
             else:
-                if env_helper.AZURE_SEARCH_USE_INTEGRATED_VECTORIZATION:
-                    logger.info("Using IntegratedVectorizationSearchHandler.")
-                    return IntegratedVectorizationSearchHandler(env_helper)
-                else:
-                    logger.info("Using AzureSearchHandler.")
-                    return AzureSearchHandler(env_helper)
+                logger.info("Using AzureSearchHandler.")
+                return AzureSearchHandler(env_helper)
 
     @staticmethod
+    @logger.trace_function(log_args=False, log_result=False)
     def get_source_documents(
         search_handler: SearchHandlerBase, question: str
     ) -> list[SourceDocument]:
@@ -78,6 +76,5 @@ class Search:
             list[SourceDocument]: A list of SourceDocument instances that match the
                                   search query.
         """
-        with tracer.start_as_current_span("get_source_documents"):
-            logger.info(f"Getting source documents for question: {question}")
-            return search_handler.query_search(question)
+        logger.info(f"Getting source documents for question: {question}")
+        return search_handler.query_search(question)