fix

Author: yanchaomei

hugegraph-llm/src/hugegraph_llm/operators/index_op/build_gremlin_example_index.py

@@ -37,17 +37,17 @@ def __init__(self, embedding: BaseEmbedding, examples: List[Dict[str, str]]):
 
     def run(self, context: Dict[str, Any]) -> Dict[str, Any]:
         embed_dim = 0
-        
+
         if len(self.examples) > 0:
             # Use the new async parallel embedding approach from upstream
             queries = [example["query"] for example in self.examples]
             # TODO: refactor function chain async to avoid blocking
             examples_embedding = asyncio.run(get_embeddings_parallel(self.embedding, queries))
             embed_dim = len(examples_embedding[0])
-            
+
             vector_index = VectorIndex(embed_dim)
             vector_index.add(examples_embedding, self.examples)
             vector_index.to_index_file(self.index_dir, self.filename_prefix)
-            
+
         context["embed_dim"] = embed_dim
         return context

hugegraph-llm/src/hugegraph_llm/operators/llm_op/keyword_extract.py

@@ -75,6 +75,40 @@ def run(self, context: Dict[str, Any]) -> Dict[str, Any]:
         context["call_count"] = context.get("call_count", 0) + 1
         return context
 
+    async def arun(self, context: Dict[str, Any]) -> Dict[str, Any]:
+        """异步版本的关键词提取"""
+        if self._query is None:
+            self._query = context.get("query")
+            assert self._query is not None, "No query for keywords extraction."
+        else:
+            context["query"] = self._query
+
+        if self._llm is None:
+            self._llm = LLMs().get_extract_llm()
+            assert isinstance(self._llm, BaseLLM), "Invalid LLM Object."
+
+        self._language = context.get("language", self._language).lower()
+        self._max_keywords = context.get("max_keywords", self._max_keywords)
+
+        prompt_run = f"{self._extract_template.format(question=self._query, max_keywords=self._max_keywords)}"
+        start_time = time.perf_counter()
+        
+        # 关键改动：使用异步LLM调用
+        response = await self._llm.agenerate(prompt=prompt_run)
+        
+        end_time = time.perf_counter()
+        log.debug("Keyword extraction time: %.2f seconds", end_time - start_time)
+
+        keywords = self._extract_keywords_from_response(
+            response=response, lowercase=False, start_token="KEYWORDS:"
+        )
+        keywords = {k.replace("'", "") for k in keywords}
+        context["keywords"] = list(keywords)
+        log.info("User Query: %s\nKeywords: %s", self._query, context["keywords"])
+
+        context["call_count"] = context.get("call_count", 0) + 1
+        return context
+
     def _extract_keywords_from_response(
             self,
             response: str,

hugegraph-llm/src/tests/operators/index_op/test_build_gremlin_example_index.py

@@ -51,11 +51,11 @@ def setUp(self):
         self.patcher2 = patch("hugegraph_llm.operators.index_op.build_gremlin_example_index.get_index_folder_name")
         self.mock_get_index_folder_name = self.patcher2.start()
         self.mock_get_index_folder_name.return_value = "hugegraph"
-        
+
         self.patcher3 = patch("hugegraph_llm.operators.index_op.build_gremlin_example_index.get_filename_prefix")
         self.mock_get_filename_prefix = self.patcher3.start()
         self.mock_get_filename_prefix.return_value = "test_prefix"
-        
+
         self.patcher4 = patch("hugegraph_llm.operators.index_op.build_gremlin_example_index.get_embeddings_parallel")
         self.mock_get_embeddings_parallel = self.patcher4.start()
         self.mock_get_embeddings_parallel.return_value = [[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]]
@@ -128,7 +128,7 @@ def test_run_with_empty_examples(self):
 
         # The run method should handle empty examples gracefully
         result = builder.run(context)
-        
+
         # Should return embed_dim as 0 for empty examples
         self.assertEqual(result["embed_dim"], 0)
         self.assertEqual(result["test"], "value")  # Original context should be preserved