format code with black

Author: xionghuaidong

kag/common/conf.py

@@ -207,7 +207,8 @@ def update_conf(self, configs: dict):
 """
 KAG_QA_TASK_CONFIG stores per-task configuration and should be cleaned up after use.
 """
-KAG_QA_TASK_CONFIG = knext.common.cache.LinkCache(maxsize=100, ttl=300)
+KAG_QA_TASK_CONFIG = knext.common.cache.LinkCache(maxsize=100, ttl=300)
+
 
 class KAGConfigAccessor:
  @staticmethod
@@ -248,7 +249,11 @@ def init_env(config_file: str = None):
  project_id = os.getenv(KAGConstants.ENV_KAG_PROJECT_ID)
  host_addr = os.getenv(KAGConstants.ENV_KAG_PROJECT_HOST_ADDR)
  prod = False
- if project_id is not None and host_addr is not None and not validate_config_file(config_file):
+ if (
+ project_id is not None
+ and host_addr is not None
+ and not validate_config_file(config_file)
+ ):
  prod = True
  global KAG_CONFIG
  KAG_CONFIG.initialize(prod, config_file)

kag/common/tools/algorithm_tool/chunk_retriever/atomic_query_chunk_retriever.py

@@ -256,15 +256,11 @@ def invoke(self, task: Task, **kwargs) -> RetrieverOutput:
 
  chunks = chunks + query_text_related_chunks
 
- out = RetrieverOutput(
- retriever_method=self.name, chunks=chunks
- )
+ out = RetrieverOutput(retriever_method=self.name, chunks=chunks)
  return out
  except Exception as e:
  logger.error(f"run calculate_sim_scores failed, info: {e}", exc_info=True)
- return RetrieverOutput(
- retriever_method=self.name, err_msg=str(e)
- )
+ return RetrieverOutput(retriever_method=self.name, err_msg=str(e))
 
  def schema(self):
  return {

kag/common/tools/algorithm_tool/chunk_retriever/outline_chunk_retriever.py

@@ -156,23 +156,17 @@ def invoke(self, task, **kwargs) -> RetrieverOutput:
  )
 
  # to retrieve output
- out = RetrieverOutput(
- retriever_method=self.name, chunks=chunks
- )
+ out = RetrieverOutput(retriever_method=self.name, chunks=chunks)
  chunk_cached_by_query_map.put(query, out)
  return out
 
  except Exception as e:
  logger.error(f"run calculate_sim_scores failed, info: {e}", exc_info=True)
- return RetrieverOutput(
- retriever_method=self.name, err_msg=str(e)
- )
+ return RetrieverOutput(retriever_method=self.name, err_msg=str(e))
 
  @property
  def input_indices(self):
  return ["Outline"]
 
  def schema(self):
- return {
- "name": "outline_chunk_retriever"
- }
+ return {"name": "outline_chunk_retriever"}

kag/common/tools/algorithm_tool/chunk_retriever/ppr_chunk_retriever.py

@@ -309,9 +309,7 @@ def invoke(self, task, **kwargs) -> RetrieverOutput:
  properties=node,
  )
  )
- return RetrieverOutput(
- retriever_method=self.name, chunks=matched_docs
- )
+ return RetrieverOutput(retriever_method=self.name, chunks=matched_docs)
 
  def schema(self):
  return {

kag/common/tools/algorithm_tool/chunk_retriever/rc_retriever.py

@@ -52,13 +52,12 @@ def invoke(self, task, **kwargs) -> RetrieverOutput:
  output.retriever_method = self.name
  except Exception as e:
  logger.error(e, exc_info=True)
- output = RetrieverOutput(
- retriever_method=self.name, err_msg=f"{task} {e}"
- )
+ output = RetrieverOutput(retriever_method=self.name, err_msg=f"{task} {e}")
  logger.debug(
  f"{self.schema().get('name', '')} `{task.arguments['query']}` Retrieved chunks num: {len(output.chunks)} cost={time.time() - start_time}"
  )
  return output
+
  def schema(self):
  return {
  "name": "kg_rc_retriever",

kag/common/tools/algorithm_tool/chunk_retriever/summary_chunk_retriever.py

@@ -21,7 +21,8 @@
  VectorizeModelABC,
  ChunkData,
  RetrieverOutput,
- EntityData, Task,
+ EntityData,
+ Task,
 )
 from kag.interface.solver.model.schema_utils import SchemaUtils
 from kag.common.config import LogicFormConfiguration
@@ -36,13 +37,13 @@
 @RetrieverABC.register("summary_chunk_retriever")
 class SummaryChunkRetriever(RetrieverABC):
  def __init__(
-  self,
-  vectorize_model: VectorizeModelABC = None,
-  search_api: SearchApiABC = None,
-  graph_api: GraphApiABC = None,
-  top_k: int = 10,
-  score_threshold=0.85,
-  **kwargs,
+ self,
+ vectorize_model: VectorizeModelABC = None,
+ search_api: SearchApiABC = None,
+ graph_api: GraphApiABC = None,
+ top_k: int = 10,
+ score_threshold=0.85,
+ **kwargs,
  ):
  super().__init__(top_k, **kwargs)
  self.vectorize_model = vectorize_model or VectorizeModelABC.from_config(
@@ -69,12 +70,12 @@ def _get_summaries(self, query, top_k) -> List[str]:
 
  # recall top_k summaries
  top_k_summaries = self.search_api.search_vector(
-  label=self.schema_helper.get_label_within_prefix("Summary"),
-  property_key="content",
-  query_vector=query_vector,
-  topk=top_k,
-  ef_search=top_k * 3,
-  )
+ label=self.schema_helper.get_label_within_prefix("Summary"),
+ property_key="content",
+ query_vector=query_vector,
+ topk=top_k,
+ ef_search=top_k * 3,
+ )
  for item in top_k_summaries:
  topk_summary_ids.append(item["node"]["id"])
 
@@ -172,23 +173,17 @@ def invoke(self, task: Task, **kwargs) -> RetrieverOutput:
  )
 
  # to retrieve output
- out = RetrieverOutput(
- chunks=chunks, retriever_method=self.name
- )
+ out = RetrieverOutput(chunks=chunks, retriever_method=self.name)
  chunk_cached_by_query_map.put(query, out)
  return out
 
  except Exception as e:
  logger.error(f"run calculate_sim_scores failed, info: {e}", exc_info=True)
- return RetrieverOutput(
- retriever_method=self.name, err_msg=str(e)
- )
+ return RetrieverOutput(retriever_method=self.name, err_msg=str(e))
 
  @property
  def input_indices(self):
  return ["Summary"]
 
  def schema(self):
- return {
- "name": "summary_chunk_retriever"
- }
+ return {"name": "summary_chunk_retriever"}

kag/common/tools/algorithm_tool/chunk_retriever/table_retriever.py

@@ -153,23 +153,17 @@ def invoke(self, task, **kwargs) -> RetrieverOutput:
  chunks = self.get_related_chunks(topk_table_ids)
 
  # to retrieve output
- out = RetrieverOutput(
- retriever_method=self.name, chunks=chunks
- )
+ out = RetrieverOutput(retriever_method=self.name, chunks=chunks)
  chunk_cached_by_query_map.put(query, out)
  return out
 
  except Exception as e:
  logger.error(f"run calculate_sim_scores failed, info: {e}", exc_info=True)
- return RetrieverOutput(
- retriever_method=self.name, err_msg=str(e)
- )
+ return RetrieverOutput(retriever_method=self.name, err_msg=str(e))
 
  @property
  def input_indices(self):
  return ["Outline"]
 
  def schema(self):
- return {
- "name": "table_retriever"
- }
+ return {"name": "table_retriever"}

kag/common/tools/algorithm_tool/chunk_retriever/text_chunk_retriever.py

@@ -48,14 +48,10 @@ def invoke(self, task, **kwargs) -> RetrieverOutput:
  score=score,
  )
  )
- return RetrieverOutput(
- chunks=chunks, retriever_method=self.name
- )
+ return RetrieverOutput(chunks=chunks, retriever_method=self.name)
  except Exception as e:
  logger.error(f"run calculate_sim_scores failed, info: {e}", exc_info=True)
- return RetrieverOutput(
- retriever_method=self.name, err_msg=str(e)
- )
+ return RetrieverOutput(retriever_method=self.name, err_msg=str(e))
 
  def schema(self):
  return {

kag/common/tools/algorithm_tool/chunk_retriever/vector_chunk_retriever.py

@@ -88,11 +88,11 @@ def invoke(self, task, **kwargs) -> RetrieverOutput:
  score = item.get("score", 0.0)
  if score >= self.score_threshold:
  chunk = ChunkData(
-  content=item["node"].get("content", ""),
-  title=item["node"]["name"],
-  chunk_id=item["node"]["id"],
-  score=score,
-  )
+ content=item["node"].get("content", ""),
+ title=item["node"]["name"],
+ chunk_id=item["node"]["id"],
+ score=score,
+ )
  if chunk.chunk_id not in merged:
  merged[chunk.chunk_id] = score
  if merged[chunk.chunk_id] < score:
@@ -111,17 +111,13 @@ def invoke(self, task, **kwargs) -> RetrieverOutput:
  score=score,
  )
  )
- out = RetrieverOutput(
- chunks=chunks, retriever_method=self.name
- )
+ out = RetrieverOutput(chunks=chunks, retriever_method=self.name)
  chunk_cached_by_query_map.put(query, out)
  return out
 
  except Exception as e:
  logger.error(f"run calculate_sim_scores failed, info: {e}", exc_info=True)
- return RetrieverOutput(
- retriever_method=self.name, err_msg=str(e)
- )
+ return RetrieverOutput(retriever_method=self.name, err_msg=str(e))
 
  def schema(self):
  return {

kag/common/tools/algorithm_tool/graph_retriever/entity_linking.py

@@ -222,7 +222,9 @@ def invoke(self, query, name, type_name, topk_k=None, **kwargs) -> List[EntityDa
  # 6. Final filtering and sorting
  if not topk_k:
  topk_k = self.top_k
- recognition_threshold = kwargs.get("recognition_threshold", self.recognition_threshold)
+ recognition_threshold = kwargs.get(
+ "recognition_threshold", self.recognition_threshold
+ )
  retdata = []
  if name is None:
  return retdata
@@ -260,7 +262,7 @@ def invoke(self, query, name, type_name, topk_k=None, **kwargs) -> List[EntityDa
  else:
  break
 
- return retdata[: topk_k]
+ return retdata[:topk_k]
  except Exception as e:
  logger.error(
  f"Error in entity_linking {query} name={name} type={type_name}: {e}",