Add mypy & ruff configs, fix type issues, add assertions, and initial vector DB CRUD with tests

hugegraph-llm/pyproject.toml

@@ -55,14 +55,17 @@ dependencies = [
     "rich~=13.9.4",
     "apscheduler~=3.10.4",
     "litellm~=1.61.13",
-    "hugegraph-python"
+    "hugegraph-python",
+    "pymilvus==2.5.9",
+    "qdrant-client==1.14.2",
 ]
 [project.urls]
 homepage = "https://hugegraph.apache.org/"
 repository = "https://github.com/apache/incubator-hugegraph-ai"
 documentation = "https://hugegraph.apache.org/docs/quickstart/hugegraph-ai/"
 "Bug Tracker" = "https://github.com/apache/incubator-hugegraph-ai/issues"
 
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
@@ -76,16 +79,27 @@ default = true
 packages = ["src/hugegraph_llm"]
 
 [tool.hatch.build.targets.sdist]
-include = [
-    "src/hugegraph_llm",
-    "README.md",
-    "LICENSE",
-    "NOTICE",
-    "MANIFEST.in"
-]
+include = ["src/hugegraph_llm", "README.md", "LICENSE", "NOTICE", "MANIFEST.in"]
 
 [tool.hatch.metadata]
 allow-direct-references = true
 
+[tool.mypy]
+disable_error_code = ["import-untyped"]
+check_untyped_defs = true
+disallow_untyped_defs = false
+
+
+[tool.ruff]
+line-length = 120
+indent-width = 4
+extend-exclude = []
+extend-select = ["I"]
+
+[tool.ruff.format]
+quote-style = "preserve"
+indent-style = "space"
+line-ending = "auto"
+
 [tool.uv.sources]
-hugegraph-python = { path = "../hugegraph-python-client/", editable = true }
+hugegraph-python = { path = "../hugegraph-python-client/", editable = true }

hugegraph-llm/requirements.txt

@@ -18,3 +18,5 @@ openpyxl~=3.1.5
 pydantic-settings~=2.6.1
 apscheduler~=3.10.4
 litellm~=1.61.13
+pymilvus ~=2.5.9
+qdrant-client ~=1.14.2

hugegraph-llm/src/hugegraph_llm/config/__init__.py

@@ -16,21 +16,24 @@
 # under the License.
 
 
-__all__ = ["huge_settings", "admin_settings", "llm_settings", "resource_path"]
+__all__ = ["huge_settings", "admin_settings", "llm_settings", "resource_path", "index_settings"]
 
 import os
 
-from .prompt_config import PromptConfig
-from .hugegraph_config import HugeGraphConfig
+from hugegraph_llm.config.index_config import IndexConfig
+
 from .admin_config import AdminConfig
+from .hugegraph_config import HugeGraphConfig
 from .llm_config import LLMConfig
+from .prompt_config import PromptConfig
 
 prompt = PromptConfig()
 prompt.ensure_yaml_file_exists()
 
 huge_settings = HugeGraphConfig()
 admin_settings = AdminConfig()
 llm_settings = LLMConfig()
+index_settings = IndexConfig()
 
 package_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
 resource_path = os.path.join(package_path, "resources")

hugegraph-llm/src/hugegraph_llm/config/generate.py

@@ -18,7 +18,7 @@
 
 import argparse
 
-from hugegraph_llm.config import huge_settings, admin_settings, llm_settings, PromptConfig
+from hugegraph_llm.config import PromptConfig, admin_settings, huge_settings, index_settings, llm_settings
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Generate hugegraph-llm config file")
@@ -28,4 +28,5 @@
         huge_settings.generate_env()
         admin_settings.generate_env()
         llm_settings.generate_env()
+        index_settings.generate_env()
         PromptConfig().generate_yaml_file()

hugegraph-llm/src/hugegraph_llm/config/hugegraph_config.py

@@ -16,27 +16,29 @@
 # under the License.
 
 from typing import Optional
+
 from .models import BaseConfig
 
 
 class HugeGraphConfig(BaseConfig):
     """HugeGraph settings"""
+
     # graph server config
-    graph_url: Optional[str] = "127.0.0.1:8080"
-    graph_name: Optional[str] = "hugegraph"
-    graph_user: Optional[str] = "admin"
-    graph_pwd: Optional[str] = "xxx"
+    graph_url: str = "127.0.0.1:8080"
+    graph_name: str = "hugegraph"
+    graph_user: str = "admin"
+    graph_pwd: str = "xxx"
     graph_space: Optional[str] = None
 
     # graph query config
-    limit_property: Optional[str] = "False"
-    max_graph_path: Optional[int] = 10
-    max_graph_items: Optional[int] = 30
-    edge_limit_pre_label: Optional[int] = 8
+    limit_property: str = "False"
+    max_graph_path: int = 10
+    max_graph_items: int = 30
+    edge_limit_pre_label: int = 8
 
     # vector config
-    vector_dis_threshold: Optional[float] = 0.9
-    topk_per_keyword: Optional[int] = 1
+    vector_dis_threshold: float = 0.9
+    topk_per_keyword: int = 1
 
     # rerank config
-    topk_return_results: Optional[int] = 20
+    topk_return_results: int = 20

hugegraph-llm/src/hugegraph_llm/config/index_config.py

@@ -0,0 +1,36 @@
+# 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
+from typing import Optional
+
+from .models import BaseConfig
+
+
+class IndexConfig(BaseConfig):
+    """LLM settings"""
+
+    qdrant_host: Optional[str] = os.environ.get("QDRANT_HOST", None)
+    qdrant_port: int = int(os.environ.get("QDRANT_PORT", "6333"))
+    qdrant_api_key: Optional[str] = os.environ.get("QDRANT_API_KEY") if os.environ.get("QDRANT_API_KEY") else None
+
+    milvus_host: Optional[str] = os.environ.get("MILVUS_HOST", None)
+    milvus_port: int = int(os.environ.get("MILVUS_PORT", "19530"))
+    milvus_user: str = os.environ.get("MILVUS_USER", "")
+    milvus_password: str = os.environ.get("MILVUS_PASSWORD", "")
+
+    now_vector_index: str = 'Faiss'

hugegraph-llm/src/hugegraph_llm/config/llm_config.py

@@ -17,7 +17,7 @@
 
 
 import os
-from typing import Optional, Literal
+from typing import Literal, Optional
 
 from .models import BaseConfig
 
@@ -31,38 +31,43 @@ class LLMConfig(BaseConfig):
     embedding_type: Optional[Literal["openai", "litellm", "ollama/local", "qianfan_wenxin"]] = "openai"
     reranker_type: Optional[Literal["cohere", "siliconflow"]] = None
     # 1. OpenAI settings
-    openai_chat_api_base: Optional[str] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
-    openai_chat_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY")
-    openai_chat_language_model: Optional[str] = "gpt-4.1-mini"
-    openai_extract_api_base: Optional[str] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
-    openai_extract_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY")
-    openai_extract_language_model: Optional[str] = "gpt-4.1-mini"
-    openai_text2gql_api_base: Optional[str] = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
-    openai_text2gql_api_key: Optional[str] = os.environ.get("OPENAI_API_KEY")
-    openai_text2gql_language_model: Optional[str] = "gpt-4.1-mini"
-    openai_embedding_api_base: Optional[str] = os.environ.get("OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1")
-    openai_embedding_api_key: Optional[str] = os.environ.get("OPENAI_EMBEDDING_API_KEY")
-    openai_embedding_model: Optional[str] = "text-embedding-3-small"
+    openai_chat_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
+    openai_chat_api_key: str | None = os.environ.get("OPENAI_API_KEY")
+    openai_chat_language_model: str = "gpt-4.1-mini"
+    openai_extract_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
+    openai_extract_api_key: str | None = os.environ.get("OPENAI_API_KEY")
+    openai_extract_language_model: str = "gpt-4.1-mini"
+    openai_text2gql_api_base: str = os.environ.get("OPENAI_BASE_URL", "https://api.openai.com/v1")
+    openai_text2gql_api_key: str | None = os.environ.get("OPENAI_API_KEY")
+    openai_text2gql_language_model: str = "gpt-4.1-mini"
+    openai_embedding_api_base: str = os.environ.get("OPENAI_EMBEDDING_BASE_URL", "https://api.openai.com/v1")
+    openai_embedding_api_key: str | None = os.environ.get("OPENAI_EMBEDDING_API_KEY")
+    openai_embedding_model: str = "text-embedding-3-small"
+    openai_embedding_model_dim: int = 1536
     openai_chat_tokens: int = 8192
     openai_extract_tokens: int = 256
     openai_text2gql_tokens: int = 4096
     # 2. Rerank settings
-    cohere_base_url: Optional[str] = os.environ.get("CO_API_URL", "https://api.cohere.com/v1/rerank")
+    cohere_base_url: str = os.environ.get("CO_API_URL", "https://api.cohere.com/v1/rerank")
     reranker_api_key: Optional[str] = None
     reranker_model: Optional[str] = None
     # 3. Ollama settings
-    ollama_chat_host: Optional[str] = "127.0.0.1"
-    ollama_chat_port: Optional[int] = 11434
-    ollama_chat_language_model: Optional[str] = None
-    ollama_extract_host: Optional[str] = "127.0.0.1"
-    ollama_extract_port: Optional[int] = 11434
-    ollama_extract_language_model: Optional[str] = None
-    ollama_text2gql_host: Optional[str] = "127.0.0.1"
-    ollama_text2gql_port: Optional[int] = 11434
-    ollama_text2gql_language_model: Optional[str] = None
-    ollama_embedding_host: Optional[str] = "127.0.0.1"
-    ollama_embedding_port: Optional[int] = 11434
-    ollama_embedding_model: Optional[str] = None
+    ollama_chat_host: str = "127.0.0.1"
+    ollama_chat_port: int = 11434
+    ollama_chat_language_model: str | None = None
+    ollama_extract_host: str = "127.0.0.1"
+    ollama_extract_port: int = 11434
+    ollama_extract_language_model: str | None = None
+    ollama_text2gql_host: str = "127.0.0.1"
+    ollama_text2gql_port: int = 11434
+    ollama_text2gql_language_model: str | None = None
+    ollama_embedding_host: str = "127.0.0.1"
+    ollama_embedding_port: int = int(11434)
+    ollama_embedding_model: str = 'quentinz/bge-large-zh-v1.5'
+    ollama_embedding_model_dim: Optional[int] = (
+        int(os.getenv("OLLAMA_EMBEDDING_MODEL_DIM")) if os.getenv("OLLAMA_EMBEDDING_MODEL_DIM") else None  # type:ignore
+    )
+
     # 4. QianFan/WenXin settings
     # TODO: update to one token key mode
     qianfan_chat_api_key: Optional[str] = None
@@ -77,27 +82,30 @@ class LLMConfig(BaseConfig):
     qianfan_embedding_api_key: Optional[str] = None
     qianfan_embedding_secret_key: Optional[str] = None
     # 4.1 URL settings
-    qianfan_url_prefix: Optional[str] = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop"
-    qianfan_chat_url: Optional[str] = qianfan_url_prefix + "/chat/"
-    qianfan_chat_language_model: Optional[str] = "ERNIE-Speed-128K"
-    qianfan_extract_language_model: Optional[str] = "ERNIE-Speed-128K"
-    qianfan_text2gql_language_model: Optional[str] = "ERNIE-Speed-128K"
-    qianfan_embed_url: Optional[str] = qianfan_url_prefix + "/embeddings/"
+    qianfan_url_prefix: str = "https://aip.baidubce.com/rpc/2.0/ai_custom/v1/wenxinworkshop"
+    qianfan_chat_url: str = qianfan_url_prefix + "/chat/"
+    qianfan_chat_language_model: str = "ERNIE-Speed-128K"
+    qianfan_extract_language_model: str = "ERNIE-Speed-128K"
+    qianfan_text2gql_language_model: str = "ERNIE-Speed-128K"
+    qianfan_embed_url: str = qianfan_url_prefix + "/embeddings/"
+    qianfan_embedding_model_dim: int = 384
+
     # refer https://cloud.baidu.com/doc/WENXINWORKSHOP/s/alj562vvu to get more details
-    qianfan_embedding_model: Optional[str] = "embedding-v1"
+    qianfan_embedding_model: str = "embedding-v1"
     # 5. LiteLLM settings
     litellm_chat_api_key: Optional[str] = None
     litellm_chat_api_base: Optional[str] = None
-    litellm_chat_language_model: Optional[str] = "openai/gpt-4.1-mini"
+    litellm_chat_language_model: str = "openai/gpt-4.1-mini"
     litellm_chat_tokens: int = 8192
     litellm_extract_api_key: Optional[str] = None
     litellm_extract_api_base: Optional[str] = None
-    litellm_extract_language_model: Optional[str] = "openai/gpt-4.1-mini"
+    litellm_extract_language_model: str = "openai/gpt-4.1-mini"
     litellm_extract_tokens: int = 256
     litellm_text2gql_api_key: Optional[str] = None
     litellm_text2gql_api_base: Optional[str] = None
-    litellm_text2gql_language_model: Optional[str] = "openai/gpt-4.1-mini"
+    litellm_text2gql_language_model: str = "openai/gpt-4.1-mini"
     litellm_text2gql_tokens: int = 4096
     litellm_embedding_api_key: Optional[str] = None
     litellm_embedding_api_base: Optional[str] = None
-    litellm_embedding_model: Optional[str] = "openai/text-embedding-3-small"
+    litellm_embedding_model: str = "openai/text-embedding-3-small"
+    litellm_embedding_model_dim: int = 1536

hugegraph-llm/src/hugegraph_llm/config/models/base_config.py

@@ -96,8 +96,12 @@ def _sync_env_to_object(self, env_config, config_dict):
                 obj_value_str = str(obj_value) if obj_value is not None else ""
 
                 if env_value != obj_value_str:
-                    log.info("Update configuration from the file: %s=%s (Original value: %s)",
-                             env_key, env_value, obj_value_str)
+                    log.info(
+                        "Update configuration from the file: %s=%s (Original value: %s)",
+                        env_key,
+                        env_value,
+                        obj_value_str,
+                    )
                     # Update the object attribute (using lowercase key)
                     setattr(self, env_key.lower(), env_value)
 
@@ -106,8 +110,7 @@ def _sync_object_to_env(self, env_config, config_dict):
         for obj_key, obj_value in config_dict.items():
             if obj_key not in env_config:
                 obj_value_str = str(obj_value) if obj_value is not None else ""
-                log.info("Add configuration items to the environment variable file: %s=%s",
-                         obj_key, obj_value)
+                log.info("Add configuration items to the environment variable file: %s=%s", obj_key, obj_value)
                 # Add to .env
                 set_key(env_path, obj_key, obj_value_str, quote_mode="never")
 

hugegraph-llm/src/hugegraph_llm/config/models/base_prompt_config.py

@@ -55,12 +55,12 @@ def save_to_yaml(self):
         indented_gremlin_prompt = "\n".join([f"  {line}" for line in self.gremlin_generate_prompt.splitlines()])
         indented_example_prompt = "\n".join([f"    {line}" for line in self.extract_graph_prompt.splitlines()])
         indented_question = "\n".join([f"    {line}" for line in self.default_question.splitlines()])
-        indented_custom_related_information = (
-            "\n".join([f"    {line}" for line in self.custom_rerank_info.splitlines()])
+        indented_custom_related_information = "\n".join(
+            [f"    {line}" for line in self.custom_rerank_info.splitlines()]
         )
         indented_default_answer_template = "\n".join([f"    {line}" for line in self.answer_prompt.splitlines()])
-        indented_keywords_extract_template = (
-            "\n".join([f"    {line}" for line in self.keywords_extract_prompt.splitlines()])
+        indented_keywords_extract_template = "\n".join(
+            [f"    {line}" for line in self.keywords_extract_prompt.splitlines()]
         )
         indented_doc_input_text = "\n".join([f"  {line}" for line in self.doc_input_text.splitlines()])
 

hugegraph-llm/src/hugegraph_llm/demo/rag_demo/app.py

@@ -108,7 +108,7 @@ def init_rag_ui() -> gr.Interface:
         def refresh_ui_config_prompt() -> tuple:
             # we can use its __init__() for in-place reload
             # settings.from_env()
-            huge_settings.__init__()  # pylint: disable=C2801
+            huge_settings.__init__()  # type: ignore[misc] # pylint: disable=C2801
             prompt.ensure_yaml_file_exists()
             return (
                 huge_settings.graph_url,
@@ -125,7 +125,7 @@ def refresh_ui_config_prompt() -> tuple:
                 prompt.custom_rerank_info,
                 prompt.default_question,
                 huge_settings.graph_name,
-                prompt.gremlin_generate_prompt
+                prompt.gremlin_generate_prompt,
             )
 
         hugegraph_llm_ui.load(  # pylint: disable=E1101
@@ -157,6 +157,7 @@ def create_app():
     # we don't need to manually check the env now
     # settings.check_env()
     prompt.update_yaml_file()
+    assert admin_settings.enable_login
     auth_enabled = admin_settings.enable_login.lower() == "true"
     log.info("(Status) Authentication is %s now.", "enabled" if auth_enabled else "disabled")
     api_auth = APIRouter(dependencies=[Depends(authenticate)] if auth_enabled else [])