[pull] main from microsoft:main

.semversioner/2.6.0.json

@@ -0,0 +1,54 @@
+{
+  "changes": [
+    {
+      "description": "Add LiteLLM chat and embedding model providers.",
+      "type": "minor"
+    },
+    {
+      "description": "Add LoggerFactory and clean up related API.",
+      "type": "minor"
+    },
+    {
+      "description": "Add config for NLP async mode.",
+      "type": "minor"
+    },
+    {
+      "description": "Add optional input documents to indexing API.",
+      "type": "minor"
+    },
+    {
+      "description": "add customization to vector store",
+      "type": "minor"
+    },
+    {
+      "description": "Add gpt-5 support by updating fnllm dependency.",
+      "type": "patch"
+    },
+    {
+      "description": "Fix all human_readable_id fields to be 0-based.",
+      "type": "patch"
+    },
+    {
+      "description": "Fix multi-index search.",
+      "type": "patch"
+    },
+    {
+      "description": "Improve upon recent logging refactor",
+      "type": "patch"
+    },
+    {
+      "description": "Make cache, storage, and vector_store factories consistent with similar registration support",
+      "type": "patch"
+    },
+    {
+      "description": "Remove hard-coded community rate limiter.",
+      "type": "patch"
+    },
+    {
+      "description": "generate_text_embeddings only loads tables if embedding field is specified.",
+      "type": "patch"
+    }
+  ],
+  "created_at": "2025-09-22T21:44:51+00:00",
+  "version": "2.6.0"
+}

CHANGELOG.md

@@ -1,6 +1,21 @@
 # Changelog
 Note: version releases in the 0.x.y range may introduce breaking changes.
 
+## 2.6.0
+
+- minor: Add LiteLLM chat and embedding model providers.
+- minor: Add LoggerFactory and clean up related API.
+- minor: Add config for NLP async mode.
+- minor: Add optional input documents to indexing API.
+- minor: add customization to vector store
+- patch: Add gpt-5 support by updating fnllm dependency.
+- patch: Fix all human_readable_id fields to be 0-based.
+- patch: Fix multi-index search.
+- patch: Improve upon recent logging refactor
+- patch: Make cache, storage, and vector_store factories consistent with similar registration support
+- patch: Remove hard-coded community rate limiter.
+- patch: generate_text_embeddings only loads tables if embedding field is specified.
+
 ## 2.5.0
 
 - minor: Add additional context variable to build index signature for custom parameter bag

dictionary.txt

@@ -81,6 +81,7 @@ typer
 spacy
 kwargs
 ollama
+litellm
 
 # Library Methods
 iterrows
@@ -103,6 +104,8 @@ isin
 nocache
 nbconvert
 levelno
+acompletion
+aembedding
 
 # HTML
 nbsp

graphrag/api/prompt_tune.py

@@ -47,6 +47,7 @@
 from graphrag.prompt_tune.generator.persona import generate_persona
 from graphrag.prompt_tune.loader.input import load_docs_in_chunks
 from graphrag.prompt_tune.types import DocSelectionType
+from graphrag.tokenizer.get_tokenizer import get_tokenizer
 
 logger = logging.getLogger(__name__)
 
@@ -166,7 +167,7 @@ async def generate_indexing_prompts(
         examples=examples,
         language=language,
         json_mode=False,  # config.llm.model_supports_json should be used, but these prompts are used in non-json mode by the index engine
-        encoding_model=extract_graph_llm_settings.encoding_model,
+        tokenizer=get_tokenizer(model_config=extract_graph_llm_settings),
         max_token_count=max_tokens,
         min_examples_required=min_examples_required,
     )

graphrag/config/defaults.py

@@ -3,6 +3,7 @@
 
 """Common default configuration values."""
 
+from collections.abc import Callable
 from dataclasses import dataclass, field
 from pathlib import Path
 from typing import ClassVar, Literal
@@ -23,6 +24,25 @@
 from graphrag.index.operations.build_noun_graph.np_extractors.stop_words import (
     EN_STOP_WORDS,
 )
+from graphrag.language_model.providers.litellm.services.rate_limiter.rate_limiter import (
+    RateLimiter,
+)
+from graphrag.language_model.providers.litellm.services.rate_limiter.static_rate_limiter import (
+    StaticRateLimiter,
+)
+from graphrag.language_model.providers.litellm.services.retry.exponential_retry import (
+    ExponentialRetry,
+)
+from graphrag.language_model.providers.litellm.services.retry.incremental_wait_retry import (
+    IncrementalWaitRetry,
+)
+from graphrag.language_model.providers.litellm.services.retry.native_wait_retry import (
+    NativeRetry,
+)
+from graphrag.language_model.providers.litellm.services.retry.random_wait_retry import (
+    RandomWaitRetry,
+)
+from graphrag.language_model.providers.litellm.services.retry.retry import Retry
 
 DEFAULT_OUTPUT_BASE_DIR = "output"
 DEFAULT_CHAT_MODEL_ID = "default_chat_model"
@@ -39,6 +59,18 @@
 COGNITIVE_SERVICES_AUDIENCE = "https://cognitiveservices.azure.com/.default"
 
 
+DEFAULT_RETRY_SERVICES: dict[str, Callable[..., Retry]] = {
+    "native": NativeRetry,
+    "exponential_backoff": ExponentialRetry,
+    "random_wait": RandomWaitRetry,
+    "incremental_wait": IncrementalWaitRetry,
+}
+
+DEFAULT_RATE_LIMITER_SERVICES: dict[str, Callable[..., RateLimiter]] = {
+    "static": StaticRateLimiter,
+}
+
+
 @dataclass
 class BasicSearchDefaults:
     """Default values for basic search."""
@@ -275,6 +307,7 @@ class LanguageModelDefaults:
 
     api_key: None = None
     auth_type: ClassVar[AuthType] = AuthType.APIKey
+    model_provider: str | None = None
     encoding_model: str = ""
     max_tokens: int | None = None
     temperature: float = 0
@@ -294,6 +327,7 @@ class LanguageModelDefaults:
     model_supports_json: None = None
     tokens_per_minute: Literal["auto"] = "auto"
     requests_per_minute: Literal["auto"] = "auto"
+    rate_limit_strategy: str | None = "static"
     retry_strategy: str = "native"
     max_retries: int = 10
     max_retry_wait: float = 10.0

graphrag/config/enums.py

@@ -86,10 +86,12 @@ class ModelType(str, Enum):
     # Embeddings
     OpenAIEmbedding = "openai_embedding"
     AzureOpenAIEmbedding = "azure_openai_embedding"
+    Embedding = "embedding"
 
     # Chat Completion
     OpenAIChat = "openai_chat"
     AzureOpenAIChat = "azure_openai_chat"
+    Chat = "chat"
 
     # Debug
     MockChat = "mock_chat"

graphrag/config/models/graph_rag_config.py

@@ -37,6 +37,12 @@
 from graphrag.config.models.text_embedding_config import TextEmbeddingConfig
 from graphrag.config.models.umap_config import UmapConfig
 from graphrag.config.models.vector_store_config import VectorStoreConfig
+from graphrag.language_model.providers.litellm.services.rate_limiter.rate_limiter_factory import (
+    RateLimiterFactory,
+)
+from graphrag.language_model.providers.litellm.services.retry.retry_factory import (
+    RetryFactory,
+)
 
 
 class GraphRagConfig(BaseModel):
@@ -89,6 +95,47 @@ def _validate_models(self) -> None:
         if defs.DEFAULT_EMBEDDING_MODEL_ID not in self.models:
             raise LanguageModelConfigMissingError(defs.DEFAULT_EMBEDDING_MODEL_ID)
 
+    def _validate_retry_services(self) -> None:
+        """Validate the retry services configuration."""
+        retry_factory = RetryFactory()
+
+        for model_id, model in self.models.items():
+            if model.retry_strategy != "none":
+                if model.retry_strategy not in retry_factory:
+                    msg = f"Retry strategy '{model.retry_strategy}' for model '{model_id}' is not registered. Available strategies: {', '.join(retry_factory.keys())}"
+                    raise ValueError(msg)
+
+                _ = retry_factory.create(
+                    strategy=model.retry_strategy,
+                    max_attempts=model.max_retries,
+                    max_retry_wait=model.max_retry_wait,
+                )
+
+    def _validate_rate_limiter_services(self) -> None:
+        """Validate the rate limiter services configuration."""
+        rate_limiter_factory = RateLimiterFactory()
+
+        for model_id, model in self.models.items():
+            if model.rate_limit_strategy is not None:
+                if model.rate_limit_strategy not in rate_limiter_factory:
+                    msg = f"Rate Limiter strategy '{model.rate_limit_strategy}' for model '{model_id}' is not registered. Available strategies: {', '.join(rate_limiter_factory.keys())}"
+                    raise ValueError(msg)
+
+                rpm = (
+                    model.requests_per_minute
+                    if type(model.requests_per_minute) is int
+                    else None
+                )
+                tpm = (
+                    model.tokens_per_minute
+                    if type(model.tokens_per_minute) is int
+                    else None
+                )
+                if rpm is not None or tpm is not None:
+                    _ = rate_limiter_factory.create(
+                        strategy=model.rate_limit_strategy, rpm=rpm, tpm=tpm
+                    )
+
     input: InputConfig = Field(
         description="The input configuration.", default=InputConfig()
     )
@@ -300,6 +347,11 @@ def _validate_vector_store_db_uri(self) -> None:
                     raise ValueError(msg)
                 store.db_uri = str((Path(self.root_dir) / store.db_uri).resolve())
 
+    def _validate_factories(self) -> None:
+        """Validate the factories used in the configuration."""
+        self._validate_retry_services()
+        self._validate_rate_limiter_services()
+
     def get_language_model_config(self, model_id: str) -> LanguageModelConfig:
         """Get a model configuration by ID.
 
@@ -360,4 +412,5 @@ def _validate_model(self):
         self._validate_multi_output_base_dirs()
         self._validate_update_index_output_base_dir()
         self._validate_vector_store_db_uri()
+        self._validate_factories()
         return self