style: apply ruff lint fixes and code formatting (#26)

- Fix import sorting (I001) across all modules
- Remove unused imports (F401) and variables (F841)
- Fix bare except clauses (E722) in url_fetcher.py
- Remove duplicate DataType import (F811) in milvus.py
- Remove trailing whitespace in docstrings (W293)
- Fix module-level import ordering (E402)
- Remove unused num_select parameter from LLMReranker.rerank()
- Apply consistent code formatting with ruff format

Signed-off-by: Cheney Zhang <chen.zhang@zilliz.com>

Author: zc277584121

src/vector_graph_rag/__init__.py

@@ -6,17 +6,17 @@
 """
 
 from vector_graph_rag.config import Settings
-from vector_graph_rag.models import Document, Triplet, Entity, Relation, Passage
-from vector_graph_rag.llm.extractor import TripletExtractor
-from vector_graph_rag.storage.embeddings import EmbeddingModel
-from vector_graph_rag.storage.milvus import MilvusStore
 from vector_graph_rag.graph.builder import GraphBuilder
-from vector_graph_rag.graph.retriever import GraphRetriever
-from vector_graph_rag.graph.knowledge_graph import SubGraph
 from vector_graph_rag.graph.graph import Graph
+from vector_graph_rag.graph.knowledge_graph import SubGraph
+from vector_graph_rag.graph.retriever import GraphRetriever
+from vector_graph_rag.llm.cache import LLMCache, get_llm_cache
+from vector_graph_rag.llm.extractor import TripletExtractor
 from vector_graph_rag.llm.reranker import LLMReranker
+from vector_graph_rag.models import Document, Entity, Passage, Relation, Triplet
 from vector_graph_rag.rag import VectorGraphRAG, create_rag
-from vector_graph_rag.llm.cache import LLMCache, get_llm_cache
+from vector_graph_rag.storage.embeddings import EmbeddingModel
+from vector_graph_rag.storage.milvus import MilvusStore
 
 __version__ = "0.1.3"
 

src/vector_graph_rag/api/app.py

@@ -3,9 +3,10 @@
 """
 
 import os
-from typing import Optional, Dict, Any
-from pydantic_settings import BaseSettings
+from typing import Any, Dict, Optional
+
 from pydantic import Field
+from pydantic_settings import BaseSettings
 
 
 class Settings(BaseSettings):
@@ -21,9 +22,7 @@ class Settings(BaseSettings):
         default_factory=lambda: os.getenv("OPENAI_API_KEY"),
         description="OpenAI API key for LLM and embeddings",
     )
-    openai_base_url: Optional[str] = Field(
-        default=None, description="Custom OpenAI API base URL"
-    )
+    openai_base_url: Optional[str] = Field(default=None, description="Custom OpenAI API base URL")
 
     # Model Settings
     llm_model: str = Field(
@@ -80,12 +79,8 @@ class Settings(BaseSettings):
     )
 
     # Retrieval Settings
-    entity_top_k: int = Field(
-        default=20, description="Number of top entities to retrieve"
-    )
-    relation_top_k: int = Field(
-        default=20, description="Number of top relations to retrieve"
-    )
+    entity_top_k: int = Field(default=20, description="Number of top entities to retrieve")
+    relation_top_k: int = Field(default=20, description="Number of top relations to retrieve")
     entity_similarity_threshold: float = Field(
         default=0.9,
         description="Similarity threshold for entity retrieval (keep if score > threshold)",
@@ -101,31 +96,23 @@ class Settings(BaseSettings):
         default=1000,
         description="Maximum number of expanded relations. If exceeded, use eviction strategy to filter by similarity.",
     )
-    final_top_k: int = Field(
-        default=3, description="Number of final passages to return"
-    )
+    final_top_k: int = Field(default=3, description="Number of final passages to return")
 
     # LLM Settings
-    llm_temperature: float = Field(
-        default=0.0, description="Temperature for LLM generation"
-    )
-    llm_max_retries: int = Field(
-        default=3, description="Maximum retries for LLM API calls"
-    )
-    use_llm_cache: bool = Field(
-        default=True, description="Whether to use LLM response caching"
-    )
+    llm_temperature: float = Field(default=0.0, description="Temperature for LLM generation")
+    llm_max_retries: int = Field(default=3, description="Maximum retries for LLM API calls")
+    use_llm_cache: bool = Field(default=True, description="Whether to use LLM response caching")
 
     # Processing Settings
-    batch_size: int = Field(
-        default=32, description="Batch size for embedding and insertion"
-    )
+    batch_size: int = Field(default=32, description="Batch size for embedding and insertion")
 
     # NER Cache Settings
     ner_cache_dir: Optional[str] = Field(
         default_factory=lambda: os.path.join(
             os.path.dirname(os.path.dirname(os.path.dirname(__file__))),  # -> vector-graph-rag/
-            "evaluation", "data", "ner_cache"
+            "evaluation",
+            "data",
+            "ner_cache",
         ),
         description="Directory containing NER cache TSV files (HippoRAG format). "
         "Files should be named {dataset}_queries.named_entity_output.tsv",

src/vector_graph_rag/config.py

@@ -3,14 +3,14 @@
 """
 
 from vector_graph_rag.graph.builder import GraphBuilder
-from vector_graph_rag.graph.retriever import GraphRetriever, RetrievalResult
+from vector_graph_rag.graph.graph import Graph
 from vector_graph_rag.graph.knowledge_graph import (
-    SubGraph,
     GraphEntity,
-    GraphRelation,
     GraphPassage,
+    GraphRelation,
+    SubGraph,
 )
-from vector_graph_rag.graph.graph import Graph
+from vector_graph_rag.graph.retriever import GraphRetriever, RetrievalResult
 
 __all__ = [
     "Graph",

src/vector_graph_rag/graph/__init__.py

@@ -4,17 +4,17 @@
 
 import uuid
 from collections import defaultdict
-from typing import List, Dict, Optional
+from typing import Dict, List, Optional
 
+from vector_graph_rag.config import Settings, get_settings
+from vector_graph_rag.llm.extractor import processing_phrases
 from vector_graph_rag.models import (
     Document,
     Entity,
+    ExtractionResult,
     Relation,
     Triplet,
-    ExtractionResult,
 )
-from vector_graph_rag.config import Settings, get_settings
-from vector_graph_rag.llm.extractor import processing_phrases
 
 
 def generate_id() -> str:
@@ -171,10 +171,7 @@ def build_from_documents(self, documents: List[Document]) -> ExtractionResult:
         self._process_documents(documents)
 
         # Build result
-        entities = [
-            Entity(id=eid, name=self.entities[eid])
-            for eid in self.entity_ids
-        ]
+        entities = [Entity(id=eid, name=self.entities[eid]) for eid in self.entity_ids]
 
         relations = []
         for rid in self.relation_ids:

src/vector_graph_rag/graph/builder.py

@@ -11,14 +11,14 @@
 - Internal operations automatically handle Entity/Relation creation and linking
 """
 
-from typing import List, Optional, Dict, Any
+from typing import Any, Dict, List, Optional
 
 from vector_graph_rag.config import Settings, get_settings
-from vector_graph_rag.storage.milvus import MilvusStore, generate_id
-from vector_graph_rag.storage.embeddings import EmbeddingModel
-from vector_graph_rag.models import Entity, Relation, Passage, Triplet
 from vector_graph_rag.graph.knowledge_graph import SubGraph
 from vector_graph_rag.llm.extractor import processing_phrases
+from vector_graph_rag.models import Entity, Passage, Relation, Triplet
+from vector_graph_rag.storage.embeddings import EmbeddingModel
+from vector_graph_rag.storage.milvus import MilvusStore, generate_id
 
 
 class Graph:
@@ -128,8 +128,12 @@ def _create_entity(
                 existing = self._store._get_entities_by_ids([existing_id])
                 if existing:
                     current = existing[0]
-                    new_relation_ids = list(set(current.get("relation_ids", []) + (relation_ids or [])))
-                    new_passage_ids = list(set(current.get("passage_ids", []) + (passage_ids or [])))
+                    new_relation_ids = list(
+                        set(current.get("relation_ids", []) + (relation_ids or []))
+                    )
+                    new_passage_ids = list(
+                        set(current.get("passage_ids", []) + (passage_ids or []))
+                    )
                     self._store._update_entity(
                         existing_id,
                         relation_ids=new_relation_ids,
@@ -326,8 +330,12 @@ def _create_relation(
         relation_id = id or generate_id()
 
         # Create entities (if not exist) and get their IDs
-        subject_id = self._create_entity(subject, relation_ids=[relation_id], passage_ids=passage_ids)
-        object_id = self._create_entity(object_, relation_ids=[relation_id], passage_ids=passage_ids)
+        subject_id = self._create_entity(
+            subject, relation_ids=[relation_id], passage_ids=passage_ids
+        )
+        object_id = self._create_entity(
+            object_, relation_ids=[relation_id], passage_ids=passage_ids
+        )
 
         # Generate embedding
         embedding = self._embedding_model.embed(relation_text)
@@ -413,8 +421,12 @@ def _update_relation(
         # Build new text if any triplet field changed
         new_text = None
         if subject or predicate or object_:
-            new_subject = self._normalize_entity_name(subject) if subject else data.get("subject", "")
-            new_predicate = processing_phrases(predicate) if predicate else data.get("predicate", "")
+            new_subject = (
+                self._normalize_entity_name(subject) if subject else data.get("subject", "")
+            )
+            new_predicate = (
+                processing_phrases(predicate) if predicate else data.get("predicate", "")
+            )
             new_object = self._normalize_entity_name(object_) if object_ else data.get("object", "")
             new_text = f"{new_subject} {new_predicate} {new_object}"
 
@@ -453,15 +465,19 @@ def _delete_relation(self, relation_id: str) -> bool:
             entities = self._store._get_entities_by_ids([eid])
             if entities:
                 e_data = entities[0]
-                new_relation_ids = [rid for rid in e_data.get("relation_ids", []) if rid != relation_id]
+                new_relation_ids = [
+                    rid for rid in e_data.get("relation_ids", []) if rid != relation_id
+                ]
                 self._store._update_entity(eid, relation_ids=new_relation_ids)
 
         # Update related passages (remove this relation from relation_ids)
         for pid in passage_ids:
             passages = self._store.get_passages_by_ids([pid])
             if passages:
                 p_data = passages[0]
-                new_relation_ids = [rid for rid in p_data.get("relation_ids", []) if rid != relation_id]
+                new_relation_ids = [
+                    rid for rid in p_data.get("relation_ids", []) if rid != relation_id
+                ]
                 self._store.update_passage(pid, relation_ids=new_relation_ids)
 
         # Delete the relation
@@ -528,7 +544,9 @@ def create_passage(
                 relation_ids.append(relation_id)
 
                 # Get entity IDs for this triplet
-                subject_id = self._entity_name_to_id.get(self._normalize_entity_name(triplet.subject))
+                subject_id = self._entity_name_to_id.get(
+                    self._normalize_entity_name(triplet.subject)
+                )
                 object_id = self._entity_name_to_id.get(self._normalize_entity_name(triplet.object))
 
                 if subject_id and subject_id not in entity_ids:
@@ -596,12 +614,14 @@ def search_passages(
         passages = []
         for r in results:
             data = r["entity"]
-            passages.append(Passage(
-                id=data["id"],
-                text=data["text"],
-                entity_ids=data.get("entity_ids", []),
-                relation_ids=data.get("relation_ids", []),
-            ))
+            passages.append(
+                Passage(
+                    id=data["id"],
+                    text=data["text"],
+                    entity_ids=data.get("entity_ids", []),
+                    relation_ids=data.get("relation_ids", []),
+                )
+            )
 
         return passages
 
@@ -657,15 +677,19 @@ def delete_passage(self, passage_id: str) -> bool:
             entities = self._store._get_entities_by_ids([eid])
             if entities:
                 e_data = entities[0]
-                new_passage_ids = [pid for pid in e_data.get("passage_ids", []) if pid != passage_id]
+                new_passage_ids = [
+                    pid for pid in e_data.get("passage_ids", []) if pid != passage_id
+                ]
                 self._store._update_entity(eid, passage_ids=new_passage_ids)
 
         # Update related relations (remove this passage from passage_ids)
         for rid in relation_ids:
             relations = self._store._get_relations_by_ids([rid])
             if relations:
                 r_data = relations[0]
-                new_passage_ids = [pid for pid in r_data.get("passage_ids", []) if pid != passage_id]
+                new_passage_ids = [
+                    pid for pid in r_data.get("passage_ids", []) if pid != passage_id
+                ]
                 self._store._update_relation(rid, passage_ids=new_passage_ids)
 
         # Delete the passage

src/vector_graph_rag/graph/graph.py

@@ -6,8 +6,9 @@
 """
 
 from __future__ import annotations
-from typing import List, Dict, Set, Optional, Any, TYPE_CHECKING
+
 from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set
 
 if TYPE_CHECKING:
     from vector_graph_rag.storage.milvus import MilvusStore
@@ -401,7 +402,15 @@ def _fetch_relations(self, relation_ids: List[str]) -> None:
         results = self._store.client.query(
             collection_name=self._store.relation_collection,
             filter=filter_expr,
-            output_fields=["id", "text", "entity_ids", "passage_ids", "subject", "predicate", "object"],
+            output_fields=[
+                "id",
+                "text",
+                "entity_ids",
+                "passage_ids",
+                "subject",
+                "predicate",
+                "object",
+            ],
         )
 
         for r in results:
@@ -476,23 +485,17 @@ def passage_ids(self) -> Set[str]:
     @property
     def entities(self) -> List[GraphEntity]:
         """Get all entity objects in this subgraph."""
-        return [
-            self._entities[eid] for eid in self._entity_ids if eid in self._entities
-        ]
+        return [self._entities[eid] for eid in self._entity_ids if eid in self._entities]
 
     @property
     def relations(self) -> List[GraphRelation]:
         """Get all relation objects in this subgraph."""
-        return [
-            self._relations[rid] for rid in self._relation_ids if rid in self._relations
-        ]
+        return [self._relations[rid] for rid in self._relation_ids if rid in self._relations]
 
     @property
     def passages(self) -> List[GraphPassage]:
         """Get all passage objects in this subgraph."""
-        return [
-            self._passages[pid] for pid in self._passage_ids if pid in self._passages
-        ]
+        return [self._passages[pid] for pid in self._passage_ids if pid in self._passages]
 
     @property
     def entity_names(self) -> List[str]:

src/vector_graph_rag/graph/knowledge_graph.py

@@ -3,17 +3,17 @@
 """
 
 import logging
-from typing import List, Dict, Any, Optional, Tuple
 from dataclasses import dataclass, field
-
-logger = logging.getLogger(__name__)
+from typing import List, Optional, Tuple
 
 from vector_graph_rag.config import Settings, get_settings
-from vector_graph_rag.storage.embeddings import EmbeddingModel
-from vector_graph_rag.storage.milvus import MilvusStore
 from vector_graph_rag.graph.builder import GraphBuilder
 from vector_graph_rag.graph.knowledge_graph import SubGraph
 from vector_graph_rag.llm.extractor import EntityExtractor
+from vector_graph_rag.storage.embeddings import EmbeddingModel
+from vector_graph_rag.storage.milvus import MilvusStore
+
+logger = logging.getLogger(__name__)
 
 
 @dataclass
@@ -93,9 +93,7 @@ def __init__(
         self.graph_builder = graph_builder
 
         self.embedding_model = embedding_model or EmbeddingModel(settings=self.settings)
-        self.entity_extractor = entity_extractor or EntityExtractor(
-            settings=self.settings
-        )
+        self.entity_extractor = entity_extractor or EntityExtractor(settings=self.settings)
 
     def _extract_query_entities(self, query: str) -> List[str]:
         """Extract named entities from the query."""
@@ -338,9 +336,7 @@ def retrieve(
         )
 
         # Expand subgraph
-        subgraph = self._expand_subgraph(
-            entity_ids, relation_ids, degree=expansion_degree
-        )
+        subgraph = self._expand_subgraph(entity_ids, relation_ids, degree=expansion_degree)
 
         # Apply eviction strategy if needed
         threshold = relation_number_threshold or self.settings.relation_number_threshold