feat(llm):improve some RAG function UT(tests)

hugegraph-llm/run_tests.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+Test runner script for HugeGraph-LLM.
+This script sets up the environment and runs the tests.
+"""
+
+import os
+import sys
+import argparse
+import subprocess
+import nltk
+from pathlib import Path
+
+
+def setup_environment():
+    """Set up the environment for testing."""
+    # Add the project root to the Python path
+    project_root = os.path.dirname(os.path.abspath(__file__))
+    sys.path.insert(0, project_root)
+
+    # Download NLTK resources if needed
+    try:
+        nltk.data.find('corpora/stopwords')
+    except LookupError:
+        print("Downloading NLTK stopwords...")
+        nltk.download('stopwords', quiet=True)
+
+    # Set environment variable to skip external service tests by default
+    if 'HUGEGRAPH_LLM_SKIP_EXTERNAL_TESTS' not in os.environ:
+        os.environ['HUGEGRAPH_LLM_SKIP_EXTERNAL_TESTS'] = 'true'
+
+    # Create logs directory if it doesn't exist
+    logs_dir = os.path.join(project_root, 'logs')
+    os.makedirs(logs_dir, exist_ok=True)
+
+
+def run_tests(args):
+    """Run the tests with the specified arguments."""
+    # Construct the pytest command
+    cmd = ['pytest']
+
+    # Add verbosity
+    if args.verbose:
+        cmd.append('-v')
+
+    # Add coverage if requested
+    if args.coverage:
+        cmd.extend(['--cov=src/hugegraph_llm', '--cov-report=term', '--cov-report=html:coverage_html'])
+
+    # Add test pattern if specified
+    if args.pattern:
+        cmd.append(args.pattern)
+    else:
+        cmd.append('src/tests')
+
+    # Print the command being run
+    print(f"Running: {' '.join(cmd)}")
+
+    # Run the tests
+    result = subprocess.run(cmd)
+    return result.returncode
+
+
+def main():
+    """Parse arguments and run tests."""
+    parser = argparse.ArgumentParser(description='Run HugeGraph-LLM tests')
+    parser.add_argument('-v', '--verbose', action='store_true', help='Enable verbose output')
+    parser.add_argument('-c', '--coverage', action='store_true', help='Generate coverage report')
+    parser.add_argument('-p', '--pattern', help='Test pattern to run (e.g., src/tests/models)')
+    parser.add_argument('--external', action='store_true', help='Run tests that require external services')
+
+    args = parser.parse_args()
+
+    # Set up the environment
+    setup_environment()
+
+    # Configure external tests
+    if args.external:
+        os.environ['HUGEGRAPH_LLM_SKIP_EXTERNAL_TESTS'] = 'false'
+        print("Running tests including those that require external services")
+    else:
+        print("Skipping tests that require external services (use --external to include them)")
+
+    # Run the tests
+    return run_tests(args)
+
+
+if __name__ == '__main__':
+    sys.exit(main()) 

hugegraph-llm/src/tests/conftest.py

@@ -0,0 +1,47 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import os
+import sys
+import pytest
+import nltk
+
+# 获取项目根目录
+project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "../.."))
+# 添加到 Python 路径
+sys.path.insert(0, project_root)
+
+# 添加 src 目录到 Python 路径
+src_path = os.path.join(project_root, "src")
+sys.path.insert(0, src_path)
+
+# 下载 NLTK 资源
+def download_nltk_resources():
+    try:
+        nltk.data.find("corpora/stopwords")
+    except LookupError:
+        print("下载 NLTK stopwords 资源...")
+        nltk.download('stopwords', quiet=True)
+
+# 在测试开始前下载 NLTK 资源
+download_nltk_resources()
+
+# 设置环境变量，跳过外部服务测试
+os.environ['SKIP_EXTERNAL_SERVICES'] = 'true'
+
+# 打印当前 Python 路径，用于调试
+print("Python path:", sys.path) 

hugegraph-llm/src/tests/data/documents/sample.txt

@@ -0,0 +1,6 @@
+Alice is 25 years old and works as a software engineer at TechCorp.
+Bob is 30 years old and is a data scientist at DataInc.
+Alice and Bob are colleagues and they collaborate on AI projects.
+They are working on a knowledge graph project that uses natural language processing.
+The project aims to extract structured information from unstructured text.
+TechCorp and DataInc are partner companies in the technology sector. 

hugegraph-llm/src/tests/data/kg/schema.json

@@ -0,0 +1,42 @@
+{
+  "vertices": [
+    {
+      "vertex_label": "person",
+      "properties": ["name", "age", "occupation"]
+    },
+    {
+      "vertex_label": "company",
+      "properties": ["name", "industry"]
+    },
+    {
+      "vertex_label": "project",
+      "properties": ["name", "technology"]
+    }
+  ],
+  "edges": [
+    {
+      "edge_label": "works_at",
+      "source_vertex_label": "person",
+      "target_vertex_label": "company",
+      "properties": []
+    },
+    {
+      "edge_label": "colleague",
+      "source_vertex_label": "person",
+      "target_vertex_label": "person",
+      "properties": []
+    },
+    {
+      "edge_label": "works_on",
+      "source_vertex_label": "person",
+      "target_vertex_label": "project",
+      "properties": []
+    },
+    {
+      "edge_label": "partner",
+      "source_vertex_label": "company",
+      "target_vertex_label": "company",
+      "properties": []
+    }
+  ]
+} 

hugegraph-llm/src/tests/data/prompts/test_prompts.yaml

@@ -0,0 +1,36 @@
+rag_prompt:
+  system: |
+    You are a helpful assistant that answers questions based on the provided context.
+    Use only the information from the context to answer the question.
+    If you don't know the answer, say "I don't know" or "I don't have enough information".
+  user: |
+    Context:
+    {context}
+
+    Question:
+    {query}
+
+    Answer:
+
+kg_extraction_prompt:
+  system: |
+    You are a knowledge graph extraction assistant. Your task is to extract entities and relationships from the given text according to the provided schema.
+    Output the extracted information in a structured format that can be used to build a knowledge graph.
+  user: |
+    Text:
+    {text}
+
+    Schema:
+    {schema}
+
+    Extract entities and relationships from the text according to the schema:
+
+summarization_prompt:
+  system: |
+    You are a summarization assistant. Your task is to create a concise summary of the provided text.
+    The summary should capture the main points and key information.
+  user: |
+    Text:
+    {text}
+
+    Please provide a concise summary: 

hugegraph-llm/src/tests/document/test_document.py

@@ -0,0 +1,54 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import unittest
+import importlib
+
+
+class TestDocumentModule(unittest.TestCase):
+    def test_import_document_module(self):
+        """Test that the document module can be imported."""
+        try:
+            import hugegraph_llm.document
+            self.assertTrue(True)
+        except ImportError:
+            self.fail("Failed to import hugegraph_llm.document module")
+
+    def test_import_chunk_split(self):
+        """Test that the chunk_split module can be imported."""
+        try:
+            from hugegraph_llm.document import chunk_split
+            self.assertTrue(True)
+        except ImportError:
+            self.fail("Failed to import chunk_split module")
+
+    def test_chunk_splitter_class_exists(self):
+        """Test that the ChunkSplitter class exists in the chunk_split module."""
+        try:
+            from hugegraph_llm.document.chunk_split import ChunkSplitter
+            self.assertTrue(True)
+        except ImportError:
+            self.fail("ChunkSplitter class not found in chunk_split module")
+
+    def test_module_reload(self):
+        """Test that the document module can be reloaded."""
+        try:
+            import hugegraph_llm.document
+            importlib.reload(hugegraph_llm.document)
+            self.assertTrue(True)
+        except Exception as e:
+            self.fail(f"Failed to reload document module: {e}")