feat: support_customize_entity_types (#1403)

* feat: support_customize_entity_types

* feat: support_customize_entity_types

Author: earayu

aperag/api/components/schemas/collection.yaml

@@ -109,6 +109,23 @@ collectionSource:
           type: string
 
 
+knowledgeGraphConfig:
+  type: object
+  description: Configuration for knowledge graph generation
+  properties:
+    language:
+      type: string
+      description: Language for entity extraction and query responses
+      default: "English"
+      example: "English"
+    entity_types:
+      type: array
+      items:
+        type: string
+      description: List of entity types to extract during graph indexing
+      default: ["organization", "person", "geo", "event", "product", "technology", "date", "category"]
+      example: ["organization", "person", "geo", "event"]
+
 collectionConfig:
   type: object
   properties:
@@ -136,6 +153,12 @@ collectionConfig:
       type: boolean
       description: Whether to enable vision index
       default: false
+    knowledge_graph_config:
+      allOf:
+        - $ref: '#/knowledgeGraphConfig'
+      default:
+        language: "English"
+        entity_types: ["organization", "person", "geo", "event", "product", "technology", "date", "category"]
     embedding:
       $ref: './model.yaml#/modelSpec'
     completion:

aperag/graph/lightrag/lightrag.py

@@ -229,9 +229,14 @@ class LightRAG:
     # Extensions
     # ---
 
-    addon_params: dict[str, Any] = field(
-        default_factory=lambda: {"language": get_env_value("SUMMARY_LANGUAGE", "English", str)}
-    )
+    language: str = field(default="English")
+    """Language for entity extraction and query responses."""
+
+    entity_types: list[str] = field(default_factory=lambda: PROMPTS["DEFAULT_ENTITY_TYPES"])
+    """List of entity types to extract during graph indexing."""
+
+    example_number: int | None = field(default=None)
+    """Number of examples to use in prompts. If None, uses all available examples."""
 
     # Storages Management
     # ---
@@ -544,7 +549,7 @@ async def _process_component_with_semaphore(task_data):
                     tokenizer=self.tokenizer,
                     llm_model_max_token_size=self.llm_model_max_token_size,
                     summary_to_max_tokens=self.summary_to_max_tokens,
-                    addon_params=self.addon_params or PROMPTS["DEFAULT_LANGUAGE"],
+                    language=self.language,
                     force_llm_summary_on_merge=self.force_llm_summary_on_merge,
                     lightrag_logger=self.lightrag_logger,
                 )
@@ -759,7 +764,9 @@ async def aprocess_graph_indexing(
                 chunks,
                 use_llm_func=self.llm_model_func,
                 entity_extract_max_gleaning=self.entity_extract_max_gleaning,
-                addon_params=self.addon_params,
+                language=self.language,
+                entity_types=self.entity_types,
+                example_number=self.example_number,
                 llm_model_max_async=self.llm_model_max_async,
                 lightrag_logger=self.lightrag_logger,
             )
@@ -809,7 +816,8 @@ async def aquery_context(
             param,
             self.tokenizer,
             self.llm_model_func,
-            self.addon_params,
+            language=self.language,
+            example_number=self.example_number,
             chunks_vdb=self.chunks_vdb,
         )
 
@@ -898,7 +906,8 @@ async def aquery(
                 param,
                 self.tokenizer,
                 self.llm_model_func,
-                self.addon_params,
+                language=self.language,
+                example_number=self.example_number,
                 system_prompt=system_prompt,
                 chunks_vdb=self.chunks_vdb,
             )

aperag/graph/lightrag/operate.py

@@ -501,7 +501,7 @@ async def merge_nodes_and_edges(
     tokenizer,
     llm_model_max_token_size,
     summary_to_max_tokens,
-    addon_params,
+    language: str,
     force_llm_summary_on_merge,
     lightrag_logger: LightRAGLogger,
 ) -> dict[str, int]:
@@ -516,7 +516,7 @@ async def merge_nodes_and_edges(
         tokenizer,
         llm_model_max_token_size,
         summary_to_max_tokens,
-        addon_params,
+        language,
         force_llm_summary_on_merge,
         lightrag_logger,
     )
@@ -532,15 +532,12 @@ async def _merge_nodes_and_edges_impl(
     tokenizer,
     llm_model_max_token_size,
     summary_to_max_tokens,
-    addon_params,
+    language: str,
     force_llm_summary_on_merge,
     lightrag_logger: LightRAGLogger,
 ) -> dict[str, int]:
     """Internal implementation of merge_nodes_and_edges with fine-grained locking"""
 
-    # Extract language from addon_params
-    language = addon_params.get("language", "English")
-
     # Collect all nodes and edges from all chunks
     all_nodes = defaultdict(list)
     all_edges = defaultdict(list)
@@ -572,7 +569,7 @@ async def _merge_nodes_and_edges_impl(
                 tokenizer,
                 llm_model_max_token_size,
                 summary_to_max_tokens,
-                language,  # Pass language instead of addon_params
+                language,
                 force_llm_summary_on_merge,
                 lightrag_logger,
                 workspace,
@@ -613,7 +610,7 @@ async def _merge_nodes_and_edges_impl(
                 tokenizer,
                 llm_model_max_token_size,
                 summary_to_max_tokens,
-                language,  # Pass language instead of addon_params
+                language,
                 force_llm_summary_on_merge,
                 lightrag_logger,
                 workspace,
@@ -644,15 +641,13 @@ async def extract_entities(
     chunks: dict[str, TextChunkSchema],
     use_llm_func: callable,
     entity_extract_max_gleaning: int,
-    addon_params: dict,
+    language: str,
+    entity_types: list[str],
+    example_number: int | None,
     llm_model_max_async: int,
     lightrag_logger: LightRAGLogger,
 ) -> list:
     ordered_chunks = list(chunks.items())
-    # add language and example number params to prompt
-    language = addon_params.get("language", PROMPTS["DEFAULT_LANGUAGE"])
-    entity_types = addon_params.get("entity_types", PROMPTS["DEFAULT_ENTITY_TYPES"])
-    example_number = addon_params.get("example_number", None)
     if example_number and example_number < len(PROMPTS["entity_extraction_examples"]):
         examples = "\n".join(PROMPTS["entity_extraction_examples"][: int(example_number)])
     else:
@@ -824,7 +819,8 @@ async def build_query_context(
     query_param: QueryParam,
     tokenizer: Tokenizer,
     llm_model_func: callable,
-    addon_params: dict,
+    language: str,
+    example_number: int | None,
     chunks_vdb: BaseVectorStorage = None,
 ):
     if query_param.model_func:
@@ -833,7 +829,7 @@ async def build_query_context(
         use_model_func = llm_model_func
 
     hl_keywords, ll_keywords = await get_keywords_from_query(
-        query, query_param, tokenizer, use_model_func, addon_params
+        query, query_param, tokenizer, use_model_func, language, example_number
     )
 
     logger.debug(f"High-level keywords: {hl_keywords}")
@@ -882,7 +878,8 @@ async def kg_query(
     query_param: QueryParam,
     tokenizer: Tokenizer,
     llm_model_func: callable,
-    addon_params: dict,
+    language: str,
+    example_number: int | None,
     system_prompt: str | None = None,
     chunks_vdb: BaseVectorStorage = None,
 ) -> str | AsyncIterator[str]:
@@ -901,7 +898,8 @@ async def kg_query(
         query_param,
         tokenizer,
         llm_model_func,
-        addon_params,
+        language,
+        example_number,
         chunks_vdb,
     )
 
@@ -980,7 +978,8 @@ async def get_keywords_from_query(
     query_param: QueryParam,
     tokenizer: Tokenizer,
     llm_model_func: callable,
-    addon_params: dict,
+    language: str,
+    example_number: int | None,
 ) -> tuple[list[str], list[str]]:
     """
     Retrieves high-level and low-level keywords for RAG operations.
@@ -996,7 +995,9 @@ async def get_keywords_from_query(
         return query_param.hl_keywords, query_param.ll_keywords
 
     # Extract keywords using extract_keywords_only function which already supports conversation history
-    hl_keywords, ll_keywords = await extract_keywords_only(query, query_param, tokenizer, llm_model_func, addon_params)
+    hl_keywords, ll_keywords = await extract_keywords_only(
+        query, query_param, tokenizer, llm_model_func, language, example_number
+    )
     return hl_keywords, ll_keywords
 
 
@@ -1005,20 +1006,19 @@ async def extract_keywords_only(
     param: QueryParam,
     tokenizer: Tokenizer,
     llm_model_func: callable,
-    addon_params: dict,
+    language: str,
+    example_number: int | None,
 ) -> tuple[list[str], list[str]]:
     """
     Extract high-level and low-level keywords from the given 'text' using the LLM.
     This method does NOT build the final RAG context or provide a final answer.
     It ONLY extracts keywords (hl_keywords, ll_keywords).
     """
     # 2. Build the examples
-    example_number = addon_params.get("example_number", None)
     if example_number and example_number < len(PROMPTS["keywords_extraction_examples"]):
         examples = "\n".join(PROMPTS["keywords_extraction_examples"][: int(example_number)])
     else:
         examples = "\n".join(PROMPTS["keywords_extraction_examples"])
-    language = addon_params.get("language", PROMPTS["DEFAULT_LANGUAGE"])
 
     # 3. Process conversation history
     history_context = ""

aperag/graph/lightrag/prompt.py

@@ -39,7 +39,6 @@
 
 PROMPTS: dict[str, Any] = {}
 
-PROMPTS["DEFAULT_LANGUAGE"] = "English"
 PROMPTS["DEFAULT_TUPLE_DELIMITER"] = "<|>"
 PROMPTS["DEFAULT_RECORD_DELIMITER"] = "##"
 PROMPTS["DEFAULT_COMPLETION_DELIMITER"] = "<|COMPLETE|>"

aperag/graph/lightrag_manager.py

@@ -22,6 +22,7 @@
 from aperag.db.models import Collection
 from aperag.db.ops import db_ops
 from aperag.graph.lightrag import LightRAG
+from aperag.graph.lightrag.prompt import PROMPTS
 from aperag.graph.lightrag.utils import EmbeddingFunc
 from aperag.llm.embed.base_embedding import get_collection_embedding_service_sync
 from aperag.llm.llm_error_types import (
@@ -46,7 +47,6 @@ class LightRAGConfig:
     SUMMARY_TO_MAX_TOKENS = 2000
     FORCE_LLM_SUMMARY_ON_MERGE = 10
     EMBEDDING_MAX_TOKEN_SIZE = 8192
-    # DEFAULT_LANGUAGE = "Simplified Chinese"
     DEFAULT_LANGUAGE = "The same language like input text"
 
 
@@ -76,6 +76,19 @@ async def create_lightrag_instance(collection: Collection) -> LightRAG:
         # Configure storage backends
         await _configure_storage_backends(kv_storage, vector_storage, graph_storage)
 
+        # Parse knowledge graph config from collection config
+        from aperag.schema.utils import parseCollectionConfig
+
+        config = parseCollectionConfig(collection.config)
+        kg_config = config.knowledge_graph_config
+        language = LightRAGConfig.DEFAULT_LANGUAGE
+        entity_types = PROMPTS["DEFAULT_ENTITY_TYPES"]
+        if kg_config:
+            if kg_config.language:
+                language = kg_config.language
+            if kg_config.entity_types:
+                entity_types = kg_config.entity_types
+
         # Create LightRAG instance
         rag = LightRAG(
             workspace=collection_id,
@@ -93,7 +106,8 @@ async def create_lightrag_instance(collection: Collection) -> LightRAG:
             entity_extract_max_gleaning=LightRAGConfig.ENTITY_EXTRACT_MAX_GLEANING,
             summary_to_max_tokens=LightRAGConfig.SUMMARY_TO_MAX_TOKENS,
             force_llm_summary_on_merge=LightRAGConfig.FORCE_LLM_SUMMARY_ON_MERGE,
-            addon_params={"language": LightRAGConfig.DEFAULT_LANGUAGE},
+            language=language,
+            entity_types=entity_types,
             kv_storage=kv_storage,
             vector_storage=vector_storage,
             graph_storage=graph_storage,

aperag/schema/view_models.py

@@ -14,7 +14,7 @@
 
 # generated by datamodel-codegen:
 #   filename:  openapi.merged.yaml
-#   timestamp: 2025-11-11T02:52:34+00:00
+#   timestamp: 2025-11-11T06:17:00+00:00
 
 from __future__ import annotations
 
@@ -66,6 +66,32 @@ class ModelSpec(BaseModel):
     )
 
 
+class KnowledgeGraphConfig(BaseModel):
+    """
+    Configuration for knowledge graph generation
+    """
+
+    language: Optional[str] = Field(
+        'English',
+        description='Language for entity extraction and query responses',
+        examples=['English'],
+    )
+    entity_types: Optional[list[str]] = Field(
+        [
+            'organization',
+            'person',
+            'geo',
+            'event',
+            'product',
+            'technology',
+            'date',
+            'category',
+        ],
+        description='List of entity types to extract during graph indexing',
+        examples=[['organization', 'person', 'geo', 'event']],
+    )
+
+
 class CollectionConfig(BaseModel):
     source: Optional[str] = Field(
         None, description='Source system identifier', examples=['system']
@@ -85,6 +111,23 @@ class CollectionConfig(BaseModel):
     enable_vision: Optional[bool] = Field(
         False, description='Whether to enable vision index'
     )
+    knowledge_graph_config: Optional[KnowledgeGraphConfig] = Field(
+        default_factory=lambda: KnowledgeGraphConfig.model_validate(
+            {
+                'language': 'English',
+                'entity_types': [
+                    'organization',
+                    'person',
+                    'geo',
+                    'event',
+                    'product',
+                    'technology',
+                    'date',
+                    'category',
+                ],
+            }
+        )
+    )
     embedding: Optional[ModelSpec] = None
     completion: Optional[ModelSpec] = None
     path: Optional[str] = Field(None, description='Path for local and ftp sources')

web/src/api/models/collection-config.ts

@@ -13,6 +13,9 @@
  */
 
 
+// May contain unused imports in some cases
+// @ts-ignore
+import type { KnowledgeGraphConfig } from './knowledge-graph-config';
 // May contain unused imports in some cases
 // @ts-ignore
 import type { ModelSpec } from './model-spec';
@@ -59,6 +62,12 @@ export interface CollectionConfig {
      * @memberof CollectionConfig
      */
     'enable_vision'?: boolean;
+    /**
+     * 
+     * @type {KnowledgeGraphConfig}
+     * @memberof CollectionConfig
+     */
+    'knowledge_graph_config'?: KnowledgeGraphConfig;
     /**
      * 
      * @type {ModelSpec}