Add versioned prompt profiles, output gates, bounded generation concurrency and env-gated LLM client patches

MANIFEST.in

@@ -1,3 +1,4 @@
 include requirements.txt
 include graphgen/configs/*.yaml
 recursive-include graphgen/configs *.yaml
+recursive-include graphgen/templates/profiles *.json

graphgen/bases/base_generator.py

@@ -1,23 +1,45 @@
 from abc import ABC, abstractmethod
-from typing import Any
+from typing import Any, Optional
 
 from graphgen.bases.base_llm_wrapper import BaseLLMWrapper
 
 
 class BaseGenerator(ABC):
-    """
-    Generate QAs based on given prompts.
-    """
+    """Generate QAs based on given prompts."""
+
+    TEMPLATE_KEY: Optional[str] = None
 
     def __init__(self, llm_client: BaseLLMWrapper):
         self.llm_client = llm_client
+        self._profile_templates: Optional[dict] = None
+
+    def apply_prompt_profile(self, profile: str) -> None:
+        """Attach a versioned prompt profile to this generator instance."""
+        from graphgen.templates.prompt_profiles import load_profile_template
+
+        if self.TEMPLATE_KEY is None:
+            raise ValueError(
+                f"prompt_profile_unsupported_generator:{type(self).__name__}"
+            )
+        self._profile_templates = load_profile_template(profile, self.TEMPLATE_KEY)
+
+    def template(self, language: str, key: Optional[str] = None) -> str:
+        """Resolve the effective template for this generator."""
+        from graphgen.templates.prompt_profiles import OFFICIAL_TEMPLATES
+
+        if self.TEMPLATE_KEY is None:
+            raise ValueError(f"prompt_template_key_missing:{type(self).__name__}")
+        if self._profile_templates is not None:
+            node = self._profile_templates[language]
+        else:
+            node = OFFICIAL_TEMPLATES[self.TEMPLATE_KEY][language]
+        return node[key] if key is not None else node
 
-    @staticmethod
     @abstractmethod
     def build_prompt(
-        batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
+        self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
     ) -> str:
-        """Build prompt for LLM based on the given batch"""
+        """Build prompt for LLM based on the given batch."""
 
     @staticmethod
     @abstractmethod

graphgen/common/init_llm.py

@@ -129,19 +129,27 @@ def create_llm(
             )
 
         actor_name = f"Actor_LLM_{model_type}"
+        # Ray async actors default to max_concurrency=1: all workers' LLM
+        # requests serialize inside a single actor, which becomes the
+        # generation-throughput bottleneck. GRAPHGEN_LLM_ACTOR_MAX_CONCURRENCY
+        # raises it (async methods are concurrency-safe); unset keeps the
+        # official default behavior.
+        max_concurrency = int(os.environ.get("GRAPHGEN_LLM_ACTOR_MAX_CONCURRENCY", "1"))
+        actor_options: Dict[str, Any] = {
+            "name": actor_name,
+            "num_gpus": float(config.pop("num_gpus", 0)),
+            "get_if_exists": True,
+        }
+        if max_concurrency > 1:
+            actor_options["max_concurrency"] = max_concurrency
         try:
             actor_handle = ray.get_actor(actor_name)
             print(f"Using existing Ray actor: {actor_name}")
         except ValueError:
             print(f"Creating Ray actor for LLM {model_type} with backend {backend}.")
-            num_gpus = float(config.pop("num_gpus", 0))
             actor_handle = (
                 ray.remote(LLMServiceActor)
-                .options(
-                    name=actor_name,
-                    num_gpus=num_gpus,
-                    get_if_exists=True,
-                )
+                .options(**actor_options)
                 .remote(backend, config)
             )
 

graphgen/models/extractor/schema_guided_extractor.py

@@ -1,4 +1,5 @@
 import json
+import os
 
 from graphgen.bases import BaseExtractor, BaseLLMWrapper, Chunk
 from graphgen.templates import SCHEMA_GUIDED_EXTRACTION_PROMPT
@@ -63,6 +64,20 @@ async def extract(self, chunk: Chunk) -> dict:
 
         prompt = self.build_prompt(text)
         response = await self.llm_client.generate_answer(prompt)
+
+        # Strict-prompt extraction backends occasionally hit EOS as the first
+        # token. When GRAPHGEN_EMPTY_EXTRACTION_PREFILL=1 and the first
+        # response is empty, retry once with a deterministic assistant prefill
+        # that seeds the opening of the expected JSON output; non-empty
+        # responses are never retried or modified.
+        if (
+            not response.strip()
+            and os.environ.get("GRAPHGEN_EMPTY_EXTRACTION_PREFILL", "0") == "1"
+        ):
+            logger.warning("Empty extraction response; retrying with assistant prefill")
+            response = await self.llm_client.generate_answer(
+                prompt, assistant_prefill="{"
+            )
         try:
             extracted_info = json.loads(response)
             # Ensure all required keys are present

graphgen/models/generator/aggregated_generator.py

@@ -2,7 +2,6 @@
 from typing import Any, Optional
 
 from graphgen.bases import BaseGenerator
-from graphgen.templates import AGGREGATED_GENERATION_PROMPT
 from graphgen.utils import detect_main_language, logger
 
 
@@ -14,9 +13,10 @@ class AggregatedGenerator(BaseGenerator):
     2. question generation: Generate relevant questions based on the rephrased text.
     """
 
-    @staticmethod
+    TEMPLATE_KEY = "aggregated"
+
     def build_prompt(
-        batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
+        self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
     ) -> str:
         """
         Build prompts for REPHRASE.
@@ -51,7 +51,7 @@ def build_prompt(
         #                 for index, text in enumerate(original_text)
         #             ]
         #         )
-        prompt = AGGREGATED_GENERATION_PROMPT[language]["ANSWER_REPHRASING"].format(
+        prompt = self.template(language, "ANSWER_REPHRASING").format(
             entities=entities_str, relationships=relations_str
         )
         return prompt
@@ -73,17 +73,14 @@ def parse_rephrased_text(response: str) -> Optional[str]:
             return None
         return rephrased_text.strip('"').strip("'")
 
-    @staticmethod
-    def _build_prompt_for_question_generation(answer: str) -> str:
+    def _build_prompt_for_question_generation(self, answer: str) -> str:
         """
         Build prompts for QUESTION GENERATION.
         :param answer:
         :return:
         """
         language = detect_main_language(answer)
-        prompt = AGGREGATED_GENERATION_PROMPT[language]["QUESTION_GENERATION"].format(
-            answer=answer
-        )
+        prompt = self.template(language, "QUESTION_GENERATION").format(answer=answer)
         return prompt
 
     @staticmethod

graphgen/models/generator/atomic_generator.py

@@ -2,14 +2,14 @@
 from typing import Any
 
 from graphgen.bases import BaseGenerator
-from graphgen.templates import ATOMIC_GENERATION_PROMPT
 from graphgen.utils import detect_main_language, logger
 
 
 class AtomicGenerator(BaseGenerator):
-    @staticmethod
+    TEMPLATE_KEY = "atomic"
+
     def build_prompt(
-        batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
+        self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
     ) -> str:
         nodes, edges = batch
         context = ""
@@ -19,7 +19,7 @@ def build_prompt(
             context += f"- {edge[0]} - {edge[1]}: {edge[2]['description']}\n"
         language = detect_main_language(context)
 
-        prompt = ATOMIC_GENERATION_PROMPT[language].format(context=context)
+        prompt = self.template(language).format(context=context)
         return prompt
 
     @staticmethod

graphgen/models/generator/cot_generator.py

@@ -2,14 +2,14 @@
 from typing import Any
 
 from graphgen.bases import BaseGenerator
-from graphgen.templates import COT_GENERATION_PROMPT
 from graphgen.utils import detect_main_language, logger
 
 
 class CoTGenerator(BaseGenerator):
-    @staticmethod
+    TEMPLATE_KEY = "cot"
+
     def build_prompt(
-        batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
+        self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
     ) -> str:
         """
         Build prompts for COT Template Design.
@@ -30,13 +30,13 @@ def build_prompt(
             ]
         )
         language = detect_main_language(entities_str + relationships_str)
-        prompt = COT_GENERATION_PROMPT[language]["COT_TEMPLATE_DESIGN"].format(
+        prompt = self.template(language, "COT_TEMPLATE_DESIGN").format(
             entities=entities_str, relationships=relationships_str
         )
         return prompt
 
-    @staticmethod
     def build_prompt_for_cot_generation(
+        self,
         batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]],
         question: str,
         reasoning_path: str,
@@ -58,7 +58,7 @@ def build_prompt_for_cot_generation(
             ]
         )
         language = detect_main_language(entities_str + relationships_str)
-        prompt = COT_GENERATION_PROMPT[language]["COT_GENERATION"].format(
+        prompt = self.template(language, "COT_GENERATION").format(
             entities=entities_str,
             relationships=relationships_str,
             question=question,

graphgen/models/generator/fill_in_blank_generator.py

@@ -2,11 +2,12 @@
 from typing import Any
 
 from graphgen.bases import BaseGenerator
-from graphgen.templates import FILL_IN_BLANK_GENERATION_PROMPT
 from graphgen.utils import detect_main_language, logger
 
 
 class FillInBlankGenerator(BaseGenerator):
+    TEMPLATE_KEY = "fill_in_blank"
+
     def __init__(self, llm_client, num_of_questions) -> None:
         super().__init__(llm_client)
         self.num_of_questions = num_of_questions
@@ -72,7 +73,6 @@ def parse_response(response: str) -> list[dict]:
 
         return qa_pairs
 
-    # pylint: disable=W0221
     def build_prompt(
         self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
     ) -> str:
@@ -92,7 +92,7 @@ def build_prompt(
         )
         context = entities_str + "\n" + relationships_str
         language = detect_main_language(entities_str + relationships_str)
-        prompt = FILL_IN_BLANK_GENERATION_PROMPT[language].format(
+        prompt = self.template(language).format(
             context=context,
             num_of_questions=self.num_of_questions,
         )

graphgen/models/generator/masked_fill_in_blank_generator.py

@@ -3,7 +3,6 @@
 from typing import Any, Optional
 
 from graphgen.bases import BaseGenerator
-from graphgen.templates import AGGREGATED_GENERATION_PROMPT
 from graphgen.utils import detect_main_language, logger
 
 random.seed(42)
@@ -16,9 +15,10 @@ class MaskedFillInBlankGenerator(BaseGenerator):
     2. mask: Randomly select a node from the input nodes, and then mask the name of the node in the rephrased text.
     """
 
-    @staticmethod
+    TEMPLATE_KEY = "masked_fill_in_blank"
+
     def build_prompt(
-        batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
+        self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
     ) -> str:
         """
         Build prompts for REPHRASE.
@@ -53,7 +53,7 @@ def build_prompt(
         #                 for index, text in enumerate(original_text)
         #             ]
         #         )
-        prompt = AGGREGATED_GENERATION_PROMPT[language]["ANSWER_REPHRASING"].format(
+        prompt = self.template(language, "ANSWER_REPHRASING").format(
             entities=entities_str, relationships=relations_str
         )
         return prompt

graphgen/models/generator/multi_answer_generator.py

@@ -2,11 +2,12 @@
 from typing import Any
 
 from graphgen.bases import BaseGenerator
-from graphgen.templates import MAQ_GENERATION_PROMPT
 from graphgen.utils import detect_main_language, logger
 
 
 class MultiAnswerGenerator(BaseGenerator):
+    TEMPLATE_KEY = "multi_answer"
+
     def __init__(self, llm_client, num_of_questions) -> None:
         super().__init__(llm_client)
         self.num_of_questions = num_of_questions
@@ -93,7 +94,6 @@ def parse_response(response: str) -> list[dict]:
 
         return qa_pairs
 
-    # pylint: disable=W0221
     def build_prompt(
         self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
     ) -> str:
@@ -113,7 +113,7 @@ def build_prompt(
         )
         context = entities_str + "\n" + relationships_str
         language = detect_main_language(entities_str + relationships_str)
-        prompt = MAQ_GENERATION_PROMPT[language].format(
+        prompt = self.template(language).format(
             context=context,
             num_of_questions=self.num_of_questions,
         )

graphgen/models/generator/multi_choice_generator.py

@@ -2,11 +2,12 @@
 from typing import Any
 
 from graphgen.bases import BaseGenerator
-from graphgen.templates import MCQ_GENERATION_PROMPT
 from graphgen.utils import detect_main_language, logger
 
 
 class MultiChoiceGenerator(BaseGenerator):
+    TEMPLATE_KEY = "multi_choice"
+
     def __init__(self, llm_client, num_of_questions) -> None:
         super().__init__(llm_client)
         self.num_of_questions = num_of_questions
@@ -91,7 +92,6 @@ def parse_response(response: str) -> list[dict]:
 
         return qa_pairs
 
-    # pylint: disable=W0221
     def build_prompt(
         self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
     ) -> str:
@@ -111,7 +111,7 @@ def build_prompt(
         )
         context = entities_str + "\n" + relationships_str
         language = detect_main_language(entities_str + relationships_str)
-        prompt = MCQ_GENERATION_PROMPT[language].format(
+        prompt = self.template(language).format(
             context=context,
             num_of_questions=self.num_of_questions,
         )

graphgen/models/generator/multi_hop_generator.py

@@ -2,14 +2,14 @@
 from typing import Any
 
 from graphgen.bases import BaseGenerator
-from graphgen.templates import MULTI_HOP_GENERATION_PROMPT
 from graphgen.utils import detect_main_language, logger
 
 
 class MultiHopGenerator(BaseGenerator):
-    @staticmethod
+    TEMPLATE_KEY = "multi_hop"
+
     def build_prompt(
-        batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
+        self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
     ) -> str:
         nodes, edges = batch
         entities_str = "\n".join(
@@ -26,7 +26,7 @@ def build_prompt(
             ]
         )
         language = detect_main_language(entities_str + relationships_str)
-        prompt = MULTI_HOP_GENERATION_PROMPT[language].format(
+        prompt = self.template(language).format(
             entities=entities_str, relationships=relationships_str
         )
         return prompt

graphgen/models/generator/quiz_generator.py

@@ -10,9 +10,8 @@ class QuizGenerator(BaseGenerator):
     Quiz Generator rephrases given descriptions to create quiz questions.
     """
 
-    @staticmethod
     def build_prompt(
-        batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
+        self, batch: tuple[list[tuple[str, dict]], list[tuple[Any, Any, dict]]]
     ) -> str:
         """
         Build prompt for rephrasing the description.