Some Ruff fixes

Author: AlonsoGuevara

graphrag/index/flows/create_final_community_reports.py

@@ -116,15 +116,15 @@ async def create_final_community_reports(
 def _prep_nodes(input: pd.DataFrame) -> pd.DataFrame:
     """Prepare nodes by filtering, filling missing descriptions, and creating NODE_DETAILS."""
     # Filter rows where community is not -1
-    input = input[input[COMMUNITY_ID] != -1]
+    input = input.loc[input.loc[:,COMMUNITY_ID] != -1]
 
     # Fill missing values in NODE_DESCRIPTION
     input.loc[:, NODE_DESCRIPTION] = input.loc[:, NODE_DESCRIPTION].fillna(
         "No Description"
     )
 
     # Create NODE_DETAILS column
-    input.loc[:, NODE_DETAILS] = input.loc[
+    input[NODE_DETAILS] = input.loc[
         :, [NODE_ID, NODE_NAME, NODE_DESCRIPTION, NODE_DEGREE]
     ].to_dict(orient="records")
 
@@ -136,7 +136,7 @@ def _prep_edges(input: pd.DataFrame) -> pd.DataFrame:
     input.fillna(value={NODE_DESCRIPTION: "No Description"}, inplace=True)
 
     # Create EDGE_DETAILS column
-    input[EDGE_DETAILS] = input[
+    input[EDGE_DETAILS] = input.loc[:,
         [EDGE_ID, EDGE_SOURCE, EDGE_TARGET, EDGE_DESCRIPTION, EDGE_DEGREE]
     ].to_dict(orient="records")
 
@@ -148,7 +148,7 @@ def _prep_claims(input: pd.DataFrame) -> pd.DataFrame:
     input.fillna(value={NODE_DESCRIPTION: "No Description"}, inplace=True)
 
     # Create CLAIM_DETAILS column
-    input[CLAIM_DETAILS] = input[
+    input[CLAIM_DETAILS] = input.loc[:,
         [CLAIM_ID, CLAIM_SUBJECT, CLAIM_TYPE, CLAIM_STATUS, CLAIM_DESCRIPTION]
     ].to_dict(orient="records")
 

graphrag/index/graph/extractors/community_reports/prep_community_report_context.py

@@ -46,12 +46,12 @@ def prep_community_report_context(
 
     # Filter by community level
     level_context_df = local_context_df[
-        local_context_df[schemas.COMMUNITY_LEVEL] == level
+        local_context_df.loc[:,schemas.COMMUNITY_LEVEL] == level
     ]
 
     # Filter valid and invalid contexts using boolean logic
-    valid_context_df = level_context_df[~level_context_df[schemas.CONTEXT_EXCEED_FLAG]]
-    invalid_context_df = level_context_df[level_context_df[schemas.CONTEXT_EXCEED_FLAG]]
+    valid_context_df = level_context_df[~level_context_df.loc[:,schemas.CONTEXT_EXCEED_FLAG]]
+    invalid_context_df = level_context_df[level_context_df.loc[:, schemas.CONTEXT_EXCEED_FLAG]]
 
     # there is no report to substitute with, so we just trim the local context of the invalid context records
     # this case should only happen at the bottom level of the community hierarchy where there are no sub-communities
@@ -62,10 +62,10 @@ def prep_community_report_context(
         invalid_context_df[schemas.CONTEXT_STRING] = _sort_and_trim_context(
             invalid_context_df, max_tokens
         )
-        invalid_context_df.loc[schemas.CONTEXT_SIZE] = invalid_context_df[
+        invalid_context_df[schemas.CONTEXT_SIZE] = invalid_context_df[
             schemas.CONTEXT_STRING
         ].map(num_tokens)
-        invalid_context_df.loc[:, schemas.CONTEXT_EXCEED_FLAG] = 0
+        invalid_context_df[schemas.CONTEXT_EXCEED_FLAG] = 0
         return union(valid_context_df, invalid_context_df)
 
     level_context_df = _antijoin_reports(level_context_df, report_df)
@@ -85,7 +85,7 @@ def prep_community_report_context(
     )
 
     result = union(valid_context_df, community_df, remaining_df)
-    result.loc[schemas.CONTEXT_SIZE] = result[schemas.CONTEXT_STRING].map(num_tokens)
+    result.loc[schemas.CONTEXT_SIZE] = result.loc[:, schemas.CONTEXT_STRING].map(num_tokens)
 
     result[schemas.CONTEXT_EXCEED_FLAG] = 0
     return result