refactor AzureClientManager code and update pytests

Author: jgbradley1

backend/manage-indexing-jobs.py

@@ -17,9 +17,9 @@
     config,
 )
 
-from src.api.azure_clients import AzureClientManager
 from src.logger.load_logger import load_pipeline_logger
 from src.typing.pipeline import PipelineJobState
+from src.utils.azure_clients import AzureClientManager
 from src.utils.common import sanitize_name
 from src.utils.pipeline import PipelineJob
 

backend/src/api/data.py

@@ -13,12 +13,12 @@
     UploadFile,
 )
 
-from src.api.azure_clients import AzureClientManager
 from src.logger.load_logger import load_pipeline_logger
 from src.typing.models import (
     BaseResponse,
     StorageNameList,
 )
+from src.utils.azure_clients import AzureClientManager
 from src.utils.common import (
     delete_blob_container,
     delete_cosmos_container_item,

backend/src/api/graph.py

@@ -7,8 +7,8 @@
 )
 from fastapi.responses import StreamingResponse
 
-from src.api.azure_clients import AzureClientManager
 from src.logger.load_logger import load_pipeline_logger
+from src.utils.azure_clients import AzureClientManager
 from src.utils.common import (
     sanitize_name,
     validate_index_file_exist,

backend/src/api/index.py

@@ -18,14 +18,14 @@
     config as kubernetes_config,
 )
 
-from src.api.azure_clients import AzureClientManager
 from src.logger.load_logger import load_pipeline_logger
 from src.typing.models import (
     BaseResponse,
     IndexNameList,
     IndexStatusResponse,
 )
 from src.typing.pipeline import PipelineJobState
+from src.utils.azure_clients import AzureClientManager
 from src.utils.common import (
     delete_blob_container,
     sanitize_name,

backend/src/api/prompt_tuning.py

@@ -13,8 +13,8 @@
 )
 from graphrag.config.create_graphrag_config import create_graphrag_config
 
-from src.api.azure_clients import AzureClientManager
 from src.logger.load_logger import load_pipeline_logger
+from src.utils.azure_clients import AzureClientManager
 from src.utils.common import sanitize_name
 
 prompt_tuning_route = APIRouter(prefix="/index/config", tags=["Index Configuration"])

backend/src/api/query.py

@@ -25,13 +25,13 @@
     VectorStoreSearchResult,
 )
 
-from src.api.azure_clients import AzureClientManager
 from src.logger.load_logger import load_pipeline_logger
 from src.typing.models import (
     GraphRequest,
     GraphResponse,
 )
 from src.typing.pipeline import PipelineJobState
+from src.utils.azure_clients import AzureClientManager
 from src.utils.common import (
     get_df,
     sanitize_name,

backend/src/api/query_streaming.py

@@ -21,10 +21,10 @@
 )
 from graphrag.config import create_graphrag_config
 
-from src.api.azure_clients import AzureClientManager
 from src.api.query import _is_index_complete
 from src.logger.load_logger import load_pipeline_logger
 from src.typing.models import GraphRequest
+from src.utils.azure_clients import AzureClientManager
 from src.utils.common import (
     get_df,
     sanitize_name,

backend/src/api/source.py

@@ -29,8 +29,8 @@
 COVARIATES_TABLE = "output/create_final_covariates.parquet"
 ENTITY_EMBEDDING_TABLE = "output/create_final_entities.parquet"
 RELATIONSHIPS_TABLE = "output/create_final_relationships.parquet"
-TEXT_UNITS_TABLE = "output/create_base_text_units.parquet"
-DOCUMENTS_TABLE = "output/create_base_documents.parquet"
+TEXT_UNITS_TABLE = "output/create_final_text_units.parquet"
+DOCUMENTS_TABLE = "output/create_final_documents.parquet"
 
 
 @source_route.get(
@@ -39,7 +39,7 @@
     response_model=ReportResponse,
     responses={200: {"model": ReportResponse}},
 )
-async def get_report_info(index_name: str, report_id: str):
+async def get_report_info(index_name: str, report_id: int):
     # check for existence of file the query relies on to validate the index is complete
     sanitized_index_name = sanitize_name(index_name)
     validate_index_file_exist(sanitized_index_name, COMMUNITY_REPORT_TABLE)
@@ -49,15 +49,15 @@ async def get_report_info(index_name: str, report_id: str):
             storage_options=pandas_storage_options(),
         )
         # check if report_id exists in the index
-        if not report_table["community"].isin([report_id]).any():
+        if not report_table["human_readable_id"].isin([report_id]).any():
             raise ValueError(f"Report '{report_id}' not found in index '{index_name}'.")
         # check if multiple reports with the same id exist (should not happen)
-        if len(report_table.loc[report_table["community"] == report_id]) > 1:
+        if len(report_table.loc[report_table["human_readable_id"] == report_id]) > 1:
             raise ValueError(
                 f"Multiple reports with id '{report_id}' found in index '{index_name}'."
             )
         report_content = report_table.loc[
-            report_table["community"] == report_id, "full_content"
+            report_table["human_readable_id"] == report_id, "full_content_json"
         ].to_numpy()[0]
         return ReportResponse(text=report_content)
     except Exception:
@@ -97,7 +97,7 @@ async def get_chunk_info(index_name: str, text_unit_id: str):
         text_units = text_units.explode("document_ids")
 
         # verify that text_unit_id exists in the index
-        if not text_units["chunk_id"].isin([text_unit_id]).any():
+        if not text_units["human_readable_id"].isin([text_unit_id]).any():
             raise ValueError(
                 f"Text unit '{text_unit_id}' not found in index '{index_name}'."
             )

backend/src/indexer/indexer.py

@@ -13,12 +13,12 @@
 from graphrag.config.create_graphrag_config import create_graphrag_config
 from graphrag.index.create_pipeline_config import create_pipeline_config
 
-from src.api.azure_clients import AzureClientManager
 from src.logger import (
     PipelineJobUpdater,
     load_pipeline_logger,
 )
 from src.typing.pipeline import PipelineJobState
+from src.utils.azure_clients import AzureClientManager
 from src.utils.common import sanitize_name
 from src.utils.pipeline import PipelineJob
 

backend/src/logger/__init__.py

@@ -6,7 +6,6 @@
 )
 from src.logger.console_workflow_callbacks import ConsoleWorkflowCallbacks
 from src.logger.load_logger import load_pipeline_logger
-from src.logger.logger_singleton import LoggerSingleton
 from src.logger.pipeline_job_updater import PipelineJobUpdater
 from src.logger.typing import (
     Logger,
@@ -18,7 +17,6 @@
     "Logger",
     "ApplicationInsightsWorkflowCallbacks",
     "ConsoleWorkflowCallbacks",
-    "LoggerSingleton",
     "PipelineAppInsightsReportingConfig",
     "PipelineJobUpdater",
     "PipelineReportingConfigTypes",