feat: make storage backend configuragble

ChenZiHong-Gavin · ChenZiHong-Gavin · commit 86760e99f128 · 2025-12-16T15:01:38.000+08:00
diff --git a/examples/extract/extract_schema_guided/schema_guided_extraction_config.yaml b/examples/extract/extract_schema_guided/schema_guided_extraction_config.yaml
@@ -1,5 +1,6 @@
 global_params:
   working_dir: cache
+  kv_backend: rocksdb # key-value store backend, support: rocksdb, json_kv
 
 nodes:
   - id: read
diff --git a/examples/generate/generate_aggregated_qa/aggregated_config.yaml b/examples/generate/generate_aggregated_qa/aggregated_config.yaml
@@ -1,5 +1,7 @@
 global_params:
   working_dir: cache
+  graph_backend: kuzu # graph database backend, support: kuzu, networkx
+  kv_backend: rocksdb # key-value store backend, support: rocksdb, json_kv
 
 nodes:
   - id: read_files # id is unique in the pipeline, and can be referenced by other steps
diff --git a/examples/generate/generate_atomic_qa/atomic_config.yaml b/examples/generate/generate_atomic_qa/atomic_config.yaml
@@ -1,5 +1,7 @@
 global_params:
   working_dir: cache
+  graph_backend: kuzu # graph database backend, support: kuzu, networkx
+  kv_backend: rocksdb # key-value store backend, support: rocksdb, json_kv
 
 nodes:
   - id: read
diff --git a/examples/generate/generate_cot_qa/cot_config.yaml b/examples/generate/generate_cot_qa/cot_config.yaml
@@ -1,5 +1,7 @@
 global_params:
   working_dir: cache
+  graph_backend: kuzu # graph database backend, support: kuzu, networkx
+  kv_backend: rocksdb # key-value store backend, support: rocksdb, json_kv
 
 nodes:
   - id: read
diff --git a/examples/generate/generate_multi_hop_qa/multi_hop_config.yaml b/examples/generate/generate_multi_hop_qa/multi_hop_config.yaml
@@ -1,5 +1,7 @@
 global_params:
   working_dir: cache
+  graph_backend: kuzu # graph database backend, support: kuzu, networkx
+  kv_backend: rocksdb # key-value store backend, support: rocksdb, json_kv
 
 nodes:
   - id: read
diff --git a/examples/generate/generate_vqa/vqa_config.yaml b/examples/generate/generate_vqa/vqa_config.yaml
@@ -1,5 +1,7 @@
 global_params:
   working_dir: cache
+  graph_backend: kuzu # graph database backend, support: kuzu, networkx
+  kv_backend: rocksdb # key-value store backend, support: rocksdb, json_kv
 
 nodes:
   - id: read
diff --git a/graphgen/models/llm/local/vllm_wrapper.py b/graphgen/models/llm/local/vllm_wrapper.py
@@ -1,13 +1,16 @@
-import uuid
 import math
+import uuid
 from typing import Any, List, Optional
+
 from graphgen.bases.base_llm_wrapper import BaseLLMWrapper
 from graphgen.bases.datatypes import Token
 
+
 class VLLMWrapper(BaseLLMWrapper):
     """
     Async inference backend based on vLLM.
     """
+
     def __init__(
         self,
         model: str,
@@ -33,7 +36,7 @@ def __init__(
             tensor_parallel_size=tensor_parallel_size,
             gpu_memory_utilization=gpu_memory_utilization,
             trust_remote_code=kwargs.get("trust_remote_code", True),
-            disable_log_stats=False, 
+            disable_log_stats=False,
         )
         self.engine = AsyncLLMEngine.from_engine_args(engine_args)
         self.temperature = temperature
@@ -96,7 +99,11 @@ async def generate_topk_per_token(
         async for request_output in result_generator:
             final_output = request_output
 
-        if not final_output or not final_output.outputs or not final_output.outputs[0].logprobs:
+        if (
+            not final_output
+            or not final_output.outputs
+            or not final_output.outputs[0].logprobs
+        ):
             return []
 
         top_logprobs = final_output.outputs[0].logprobs[0]
diff --git a/graphgen/operators/build_kg/build_kg_service.py b/graphgen/operators/build_kg/build_kg_service.py
@@ -12,11 +12,11 @@
 
 
 class BuildKGService(BaseOperator):
-    def __init__(self, working_dir: str = "cache"):
+    def __init__(self, working_dir: str = "cache", graph_backend: str = "kuzu"):
         super().__init__(working_dir=working_dir, op_name="build_kg_service")
         self.llm_client: BaseLLMWrapper = init_llm("synthesizer")
         self.graph_storage: BaseGraphStorage = init_storage(
-            backend="kuzu", working_dir=working_dir, namespace="graph"
+            backend=graph_backend, working_dir=working_dir, namespace="graph"
         )
 
     def process(self, batch: pd.DataFrame) -> pd.DataFrame:
diff --git a/graphgen/operators/chunk/chunk_service.py b/graphgen/operators/chunk/chunk_service.py
@@ -42,12 +42,14 @@ def split_chunks(text: str, language: str = "en", **kwargs) -> list:
 
 
 class ChunkService(BaseOperator):
-    def __init__(self, working_dir: str = "cache", **chunk_kwargs):
+    def __init__(
+        self, working_dir: str = "cache", kv_backend: str = "rocksdb", **chunk_kwargs
+    ):
         super().__init__(working_dir=working_dir, op_name="chunk_service")
         tokenizer_model = os.getenv("TOKENIZER_MODEL", "cl100k_base")
         self.tokenizer_instance: Tokenizer = Tokenizer(model_name=tokenizer_model)
         self.chunk_storage = init_storage(
-            backend="rocksdb",
+            backend=kv_backend,
             working_dir=working_dir,
             namespace="chunk",
         )
diff --git a/graphgen/operators/judge/judge_service.py b/graphgen/operators/judge/judge_service.py
@@ -11,11 +11,11 @@
 class JudgeService(BaseOperator):
     """Service for judging graph edges and nodes using a trainee LLM."""
 
-    def __init__(self, working_dir: str = "cache"):
+    def __init__(self, working_dir: str = "cache", graph_backend: str = "kuzu"):
         super().__init__(working_dir=working_dir, op_name="judge_service")
         self.llm_client: BaseLLMWrapper = init_llm("trainee")
         self.graph_storage: BaseGraphStorage = init_storage(
-            backend="kuzu",
+            backend=graph_backend,
             working_dir=working_dir,
             namespace="graph",
         )
diff --git a/graphgen/operators/partition/partition_service.py b/graphgen/operators/partition/partition_service.py
@@ -18,15 +18,21 @@
 
 
 class PartitionService(BaseOperator):
-    def __init__(self, working_dir: str = "cache", **partition_kwargs):
+    def __init__(
+        self,
+        working_dir: str = "cache",
+        graph_backend: str = "kuzu",
+        kv_backend: str = "rocksdb",
+        **partition_kwargs,
+    ):
         super().__init__(working_dir=working_dir, op_name="partition_service")
         self.kg_instance: BaseGraphStorage = init_storage(
-            backend="kuzu",
+            backend=graph_backend,
             working_dir=working_dir,
             namespace="graph",
         )
         self.chunk_storage: BaseKVStorage = init_storage(
-            backend="rocksdb",
+            backend=kv_backend,
             working_dir=working_dir,
             namespace="chunk",
         )
diff --git a/graphgen/operators/quiz/quiz_service.py b/graphgen/operators/quiz/quiz_service.py
@@ -12,18 +12,20 @@ class QuizService(BaseOperator):
     def __init__(
         self,
         working_dir: str = "cache",
+        graph_backend: str = "kuzu",
+        kv_backend: str = "rocksdb",
         quiz_samples: int = 1,
         concurrency_limit: int = 200,
     ):
         super().__init__(working_dir=working_dir, op_name="quiz_service")
         self.quiz_samples = quiz_samples
         self.llm_client: BaseLLMWrapper = init_llm("synthesizer")
         self.graph_storage: BaseGraphStorage = init_storage(
-            backend="kuzu", working_dir=working_dir, namespace="graph"
+            backend=graph_backend, working_dir=working_dir, namespace="graph"
         )
         # { _quiz_id: { "description": str, "quizzes": List[Tuple[str, str]] } }
         self.quiz_storage: BaseKVStorage = init_storage(
-            backend="rocksdb", working_dir=working_dir, namespace="quiz"
+            backend=kv_backend, working_dir=working_dir, namespace="quiz"
         )
         self.generator = QuizGenerator(self.llm_client)
         self.concurrency_limit = concurrency_limit
