update how prompts get saved in cosmosdb

Author: jgbradley1

backend/src/api/index.py

@@ -49,8 +49,8 @@ async def schedule_indexing_job(
     storage_name: str,
     index_name: str,
     entity_extraction_prompt: UploadFile | None = None,
-    community_report_prompt: UploadFile | None = None,
-    summarize_descriptions_prompt: UploadFile | None = None,
+    entity_summarization_prompt: UploadFile | None = None,
+    community_summarization_prompt: UploadFile | None = None,
 ):
     azure_client_manager = AzureClientManager()
     blob_service_client = azure_client_manager.get_blob_service_client()
@@ -80,14 +80,14 @@ async def schedule_indexing_job(
         if entity_extraction_prompt
         else None
     )
-    community_report_prompt_content = (
-        community_report_prompt.file.read().decode("utf-8")
-        if community_report_prompt
+    entity_summarization_prompt_content = (
+        entity_summarization_prompt.file.read().decode("utf-8")
+        if entity_summarization_prompt
         else None
     )
-    summarize_descriptions_prompt_content = (
-        summarize_descriptions_prompt.file.read().decode("utf-8")
-        if summarize_descriptions_prompt
+    community_summarization_prompt_content = (
+        community_summarization_prompt.file.read().decode("utf-8")
+        if community_summarization_prompt
         else None
     )
 
@@ -116,9 +116,9 @@ async def schedule_indexing_job(
             existing_job._failed_workflows
         ) = []
         existing_job._entity_extraction_prompt = entity_extraction_prompt_content
-        existing_job._community_report_prompt = community_report_prompt_content
-        existing_job._summarize_descriptions_prompt = (
-            summarize_descriptions_prompt_content
+        existing_job._entity_summarization_prompt = entity_summarization_prompt_content
+        existing_job._community_summarization_prompt = (
+            community_summarization_prompt_content
         )
         existing_job._epoch_request_time = int(time())
         existing_job.update_db()
@@ -128,8 +128,8 @@ async def schedule_indexing_job(
             human_readable_index_name=index_name,
             human_readable_storage_name=storage_name,
             entity_extraction_prompt=entity_extraction_prompt_content,
-            community_report_prompt=community_report_prompt_content,
-            summarize_descriptions_prompt=summarize_descriptions_prompt_content,
+            entity_summarization_prompt=entity_summarization_prompt_content,
+            community_summarization_prompt=community_summarization_prompt_content,
             status=PipelineJobState.SCHEDULED,
         )
 

backend/src/indexer/__init__.py

@@ -12,18 +12,20 @@
 from graphrag.callbacks.workflow_callbacks import WorkflowCallbacks
 from graphrag.config.create_graphrag_config import create_graphrag_config
 from graphrag.index.create_pipeline_config import create_pipeline_config
+from graphrag.index.typing import PipelineRunResult
 
-from src.logger import (
+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
+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
 
 
 def start_indexing_job(index_name: str):
+    return 0
     print("Start indexing job...")
     # get sanitized name
     sanitized_index_name = sanitize_name(index_name)
@@ -73,20 +75,20 @@ def start_indexing_job(index_name: str):
     else:
         data.pop("entity_extraction")
 
-    # set prompt for summarize descriptions
-    if pipeline_job.summarize_descriptions_prompt:
-        fname = "summarize-descriptions-prompt.txt"
+    # set prompt for entity summarization
+    if pipeline_job.entity_summarization_prompt:
+        fname = "entity-summarization-prompt.txt"
         with open(fname, "w") as outfile:
-            outfile.write(pipeline_job.summarize_descriptions_prompt)
+            outfile.write(pipeline_job.entity_summarization_prompt)
         data["summarize_descriptions"]["prompt"] = fname
     else:
         data.pop("summarize_descriptions")
 
-    # set prompt for community report
-    if pipeline_job.community_report_prompt:
-        fname = "community-report-prompt.txt"
+    # set prompt for community summarization
+    if pipeline_job.community_summarization_prompt:
+        fname = "community-summarization-prompt.txt"
         with open(fname, "w") as outfile:
-            outfile.write(pipeline_job.community_report_prompt)
+            outfile.write(pipeline_job.community_summarization_prompt)
         data["community_reports"]["prompt"] = fname
     else:
         data.pop("community_reports")
@@ -101,7 +103,7 @@ def start_indexing_job(index_name: str):
     pipeline_job.failed_workflows = []
     pipeline_config = create_pipeline_config(parameters)
     for workflow in pipeline_config.workflows:
-        pipeline_job.all_workflows.append(workflow.name)
+        pipeline_job.all_workflows = pipeline_job.all_workflows.append(workflow.name)
 
     # create new loggers/callbacks just for this job
     print("Creating generic loggers...")
@@ -118,16 +120,27 @@ def start_indexing_job(index_name: str):
     # run the pipeline
     try:
         print("Building index...")
-        asyncio.run(
+        pipeline_results: list[PipelineRunResult] = asyncio.run(
             api.build_index(
                 config=parameters,
                 callbacks=[logger, pipeline_job_updater],
             )
         )
-        print("Index building complete")
-        # if job is done, check if any pipeline steps failed
+
+        # once indexing job is done, check if any pipeline steps failed
+        for result in pipeline_results:
+            if result.errors:
+                pipeline_job.failed_workflows = pipeline_job.failed_workflows.append(
+                    result.workflow
+                )
+            else:
+                pipeline_job.completed_workflows = (
+                    pipeline_job.completed_workflows.append(result.workflow)
+                )
+        print("Indexing complete")
+
         if len(pipeline_job.failed_workflows) > 0:
-            print("Indexing pipeline encountered error.")
+            print("Indexing pipeline encountered errors.")
             pipeline_job.status = PipelineJobState.FAILED
             logger.error(
                 message=f"Indexing pipeline encountered error for index'{index_name}'.",
@@ -158,12 +171,10 @@ def start_indexing_job(index_name: str):
             exit(1)  # signal to AKS that indexing job failed
     except Exception as e:
         pipeline_job.status = PipelineJobState.FAILED
-        # update failed state in cosmos db
         error_details = {
             "index": index_name,
             "storage_name": storage_name,
         }
-        # log error in local index directory logs
         logger.error(
             message=f"Indexing pipeline failed for index '{index_name}'.",
             cause=e,
@@ -177,8 +188,4 @@ def start_indexing_job(index_name: str):
     parser.add_argument("-i", "--index-name", required=True)
     args = parser.parse_args()
 
-    asyncio.run(
-        start_indexing_job(
-            index_name=args.index_name,
-        )
-    )
+    start_indexing_job(index_name=args.index_name)

backend/src/indexer/indexer.py

@@ -13,11 +13,11 @@ llm:
   api_version: $GRAPHRAG_API_VERSION
   model: $GRAPHRAG_LLM_MODEL
   deployment_name: $GRAPHRAG_LLM_DEPLOYMENT_NAME
-  cognitive_services_endpoint: $GRAPHRAG_COGNITIVE_SERVICES_ENDPOINT
+  cognitive_services_endpoint: $COGNITIVE_SERVICES_AUDIENCE
   model_supports_json: True
   tokens_per_minute: 80_000
   requests_per_minute: 480
-  concurrent_requests: 50
+  concurrent_requests: 25
   max_retries: 250
   max_retry_wait: 60.0
   sleep_on_rate_limit_recommendation: True
@@ -43,7 +43,7 @@ embeddings:
     batch_size: 10
     model: $GRAPHRAG_EMBEDDING_MODEL
     deployment_name: $GRAPHRAG_EMBEDDING_DEPLOYMENT_NAME
-    cognitive_services_endpoint: $GRAPHRAG_COGNITIVE_SERVICES_ENDPOINT
+    cognitive_services_endpoint: $COGNITIVE_SERVICES_AUDIENCE
     tokens_per_minute: 350_000
     requests_per_minute: 2_100
 

backend/src/indexer/settings.yaml

@@ -21,14 +21,12 @@ def __init__(self, pipeline_job: PipelineJob):
         """
         self._pipeline_job = pipeline_job
 
-    def on_workflow_start(self, name: str, instance: object) -> None:
+    def workflow_start(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow starts."""
-        # if we are not already running, set the status to running
-        if self._pipeline_job.status != PipelineJobState.RUNNING:
-            self._pipeline_job.status = PipelineJobState.RUNNING
+        self._pipeline_job.status = PipelineJobState.RUNNING
         self._pipeline_job.progress = f"Workflow {name} started."
 
-    def on_workflow_end(self, name: str, instance: object) -> None:
+    def workflow_end(self, name: str, instance: object) -> None:
         """Execute this callback when a workflow ends."""
         self._pipeline_job.completed_workflows.append(name)
         self._pipeline_job.update_db()

backend/src/logger/pipeline_job_updater.py

@@ -27,16 +27,19 @@ class PipelineJob:
     _sanitized_index_name: str = field(default=None, init=False)
     _human_readable_storage_name: str = field(default=None, init=False)
     _sanitized_storage_name: str = field(default=None, init=False)
-    _entity_extraction_prompt: str = field(default=None, init=False)
-    _community_report_prompt: str = field(default=None, init=False)
-    _summarize_descriptions_prompt: str = field(default=None, init=False)
+
     _all_workflows: List[str] = field(default_factory=list, init=False)
     _completed_workflows: List[str] = field(default_factory=list, init=False)
     _failed_workflows: List[str] = field(default_factory=list, init=False)
+
     _status: PipelineJobState = field(default=None, init=False)
     _percent_complete: float = field(default=0, init=False)
     _progress: str = field(default="", init=False)
 
+    _entity_extraction_prompt: str = field(default=None, init=False)
+    _entity_summarization_prompt: str = field(default=None, init=False)
+    _community_summarization_prompt: str = field(default=None, init=False)
+
     @staticmethod
     def _jobs_container():
         azure_storage_client = AzureClientManager()
@@ -51,8 +54,8 @@ def create_item(
         human_readable_index_name: str,
         human_readable_storage_name: str,
         entity_extraction_prompt: str | None = None,
-        community_report_prompt: str | None = None,
-        summarize_descriptions_prompt: str | None = None,
+        entity_summarization_prompt: str | None = None,
+        community_summarization_prompt: str | None = None,
         **kwargs,
     ) -> "PipelineJob":
         """
@@ -95,18 +98,21 @@ def create_item(
         instance._sanitized_index_name = sanitize_name(human_readable_index_name)
         instance._human_readable_storage_name = human_readable_storage_name
         instance._sanitized_storage_name = sanitize_name(human_readable_storage_name)
-        instance._entity_extraction_prompt = entity_extraction_prompt
-        instance._community_report_prompt = community_report_prompt
-        instance._summarize_descriptions_prompt = summarize_descriptions_prompt
+
         instance._all_workflows = kwargs.get("all_workflows", [])
         instance._completed_workflows = kwargs.get("completed_workflows", [])
         instance._failed_workflows = kwargs.get("failed_workflows", [])
+
         instance._status = PipelineJobState(
             kwargs.get("status", PipelineJobState.SCHEDULED.value)
         )
         instance._percent_complete = kwargs.get("percent_complete", 0.0)
         instance._progress = kwargs.get("progress", "")
 
+        instance._entity_extraction_prompt = entity_extraction_prompt
+        instance._entity_summarization_prompt = entity_summarization_prompt
+        instance._community_summarization_prompt = community_summarization_prompt
+
         # Create the item in the database
         instance.update_db()
         return instance
@@ -140,17 +146,22 @@ def load_item(cls, id: str) -> "PipelineJob":
             "human_readable_storage_name"
         )
         instance._sanitized_storage_name = db_item.get("sanitized_storage_name")
-        instance._entity_extraction_prompt = db_item.get("entity_extraction_prompt")
-        instance._community_report_prompt = db_item.get("community_report_prompt")
-        instance._summarize_descriptions_prompt = db_item.get(
-            "summarize_descriptions_prompt"
-        )
+
         instance._all_workflows = db_item.get("all_workflows", [])
         instance._completed_workflows = db_item.get("completed_workflows", [])
         instance._failed_workflows = db_item.get("failed_workflows", [])
+
         instance._status = PipelineJobState(db_item.get("status"))
         instance._percent_complete = db_item.get("percent_complete", 0.0)
         instance._progress = db_item.get("progress", "")
+
+        instance._entity_extraction_prompt = db_item.get("entity_extraction_prompt")
+        instance._entity_summarization_prompt = db_item.get(
+            "entity_summarization_prompt"
+        )
+        instance._community_summarization_prompt = db_item.get(
+            "community_summarization_prompt"
+        )
         return instance
 
     @staticmethod
@@ -191,10 +202,12 @@ def dump_model(self) -> dict:
         }
         if self._entity_extraction_prompt:
             model["entity_extraction_prompt"] = self._entity_extraction_prompt
-        if self._community_report_prompt:
-            model["community_report_prompt"] = self._community_report_prompt
-        if self._summarize_descriptions_prompt:
-            model["summarize_descriptions_prompt"] = self._summarize_descriptions_prompt
+        if self._entity_summarization_prompt:
+            model["entity_summarization_prompt"] = self._entity_summarization_prompt
+        if self._community_summarization_prompt:
+            model["community_summarization_prompt"] = (
+                self._community_summarization_prompt
+            )
         return model
 
     def update_db(self):
@@ -268,21 +281,23 @@ def entity_extraction_prompt(self, entity_extraction_prompt: str) -> None:
         self.update_db()
 
     @property
-    def community_report_prompt(self) -> str:
-        return self._community_report_prompt
+    def entity_summarization_prompt(self) -> str:
+        return self._entity_summarization_prompt
 
-    @community_report_prompt.setter
-    def community_report_prompt(self, community_report_prompt: str) -> None:
-        self._community_report_prompt = community_report_prompt
+    @entity_summarization_prompt.setter
+    def entity_summarization_prompt(self, entity_summarization_prompt: str) -> None:
+        self._entity_summarization_prompt = entity_summarization_prompt
         self.update_db()
 
     @property
-    def summarize_descriptions_prompt(self) -> str:
-        return self._summarize_descriptions_prompt
-
-    @summarize_descriptions_prompt.setter
-    def summarize_descriptions_prompt(self, summarize_descriptions_prompt: str) -> None:
-        self._summarize_descriptions_prompt = summarize_descriptions_prompt
+    def community_summarization_prompt(self) -> str:
+        return self._community_summarization_prompt
+
+    @community_summarization_prompt.setter
+    def community_summarization_prompt(
+        self, community_summarization_prompt: str
+    ) -> None:
+        self._community_summarization_prompt = community_summarization_prompt
         self.update_db()
 
     @property

backend/src/utils/pipeline.py

@@ -38,7 +38,7 @@ graphragConfig:
   COSMOS_URI_ENDPOINT: ""
   GRAPHRAG_API_BASE: ""
   GRAPHRAG_API_VERSION: ""
-  GRAPHRAG_COGNITIVE_SERVICES_ENDPOINT: "https://cognitiveservices.azure.com/.default"
+  COGNITIVE_SERVICES_AUDIENCE: "https://cognitiveservices.azure.com/.default"
   GRAPHRAG_LLM_MODEL: ""
   GRAPHRAG_LLM_DEPLOYMENT_NAME: ""
   GRAPHRAG_EMBEDDING_MODEL: ""