Create Language Model Providers and Registry methods. Remove fnllm coupling (#1724)

* Base structure

* Add fnllm providers and Mock LLM

* Remove fnllm coupling, introduce llm providers

* Ruff + Tests fix

* Spellcheck

* Semver

* Format

* Default MockChat params

* Fix more tests

* Fix embedding smoke test

* Fix embeddings smoke test

* Fix MockEmbeddingLLM

* Rename LLM to model. Package organization

* Fix prompt tuning

* Oops

* Oops II

Author: AlonsoGuevara

.semversioner/next-release/minor-20250219193155428669.json

@@ -0,0 +1,4 @@
+{
+  "type": "minor",
+  "description": "Add LMM Manager and Factory, to support provider registration"
+}

graphrag/api/prompt_tune.py

@@ -16,7 +16,7 @@
 import graphrag.config.defaults as defs
 from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
 from graphrag.config.models.graph_rag_config import GraphRagConfig
-from graphrag.index.llm.load_llm import load_llm
+from graphrag.language_model.manager import ModelManager
 from graphrag.logger.print_progress import PrintProgressLogger
 from graphrag.prompt_tune.defaults import MAX_TOKEN_COUNT, PROMPT_TUNING_MODEL_ID
 from graphrag.prompt_tune.generator.community_report_rating import (
@@ -104,11 +104,12 @@ async def generate_indexing_prompts(
     if default_llm_settings.max_retries == -1:
         default_llm_settings.max_retries = min(len(doc_list), defs.LLM_MAX_RETRIES)
 
-    llm = load_llm(
-        "prompt_tuning",
-        default_llm_settings,
-        cache=None,
+    llm = ModelManager().register_chat(
+        name="prompt_tuning",
+        model_type=default_llm_settings.type,
+        config=default_llm_settings,
         callbacks=NoopWorkflowCallbacks(),
+        cache=None,
     )
 
     if not domain:

graphrag/cli/index.py

@@ -153,7 +153,7 @@ def _run_index(
             True,
         )
 
-    if skip_validation:
+    if not skip_validation:
         validate_config_names(progress_logger, config)
 
     info(f"Starting pipeline run. {dry_run=}", verbose)

graphrag/config/defaults.py

@@ -12,7 +12,7 @@
     ChunkStrategyType,
     InputFileType,
     InputType,
-    LLMType,
+    ModelType,
     NounPhraseExtractorType,
     OutputType,
     ReportingType,
@@ -30,7 +30,7 @@
 # LLM Parameters
 #
 LLM_FREQUENCY_PENALTY = 0.0
-LLM_TYPE = LLMType.OpenAIChat
+LLM_TYPE = ModelType.OpenAIChat
 LLM_MODEL = "gpt-4-turbo-preview"
 LLM_MAX_TOKENS = 4000
 LLM_TEMPERATURE = 0
@@ -48,7 +48,7 @@
 #
 # Text embedding
 #
-EMBEDDING_TYPE = LLMType.OpenAIEmbedding
+EMBEDDING_TYPE = ModelType.OpenAIEmbedding
 EMBEDDING_MODEL = "text-embedding-3-small"
 EMBEDDING_BATCH_SIZE = 16
 EMBEDDING_BATCH_MAX_TOKENS = 8191

graphrag/config/enums.py

@@ -98,7 +98,7 @@ def __repr__(self):
         return f'"{self.value}"'
 
 
-class LLMType(str, Enum):
+class ModelType(str, Enum):
     """LLMType enum class definition."""
 
     # Embeddings
@@ -110,7 +110,8 @@ class LLMType(str, Enum):
     AzureOpenAIChat = "azure_openai_chat"
 
     # Debug
-    StaticResponse = "static_response"
+    MockChat = "mock_chat"
+    MockEmbedding = "mock_embedding"
 
     def __repr__(self):
         """Get a string representation."""

graphrag/config/models/language_model_config.py

@@ -7,7 +7,7 @@
 from pydantic import BaseModel, Field, model_validator
 
 import graphrag.config.defaults as defs
-from graphrag.config.enums import AsyncType, AuthType, LLMType
+from graphrag.config.enums import AsyncType, AuthType, ModelType
 from graphrag.config.errors import (
     ApiKeyMissingError,
     AzureApiBaseMissingError,
@@ -71,12 +71,12 @@ def _validate_auth_type(self) -> None:
             If the Azure authentication type conflicts with the model being used.
         """
         if self.auth_type == AuthType.AzureManagedIdentity and (
-            self.type == LLMType.OpenAIChat or self.type == LLMType.OpenAIEmbedding
+            self.type == ModelType.OpenAIChat or self.type == ModelType.OpenAIEmbedding
         ):
             msg = f"auth_type of azure_managed_identity is not supported for model type {self.type.value}. Please rerun `graphrag init` and set the auth_type to api_key."
             raise ConflictingSettingsError(msg)
 
-    type: LLMType = Field(description="The type of LLM model to use.")
+    type: ModelType = Field(description="The type of LLM model to use.")
     model: str = Field(description="The LLM model to use.")
     encoding_model: str = Field(description="The encoding model to use", default="")
 
@@ -133,8 +133,8 @@ def _validate_api_base(self) -> None:
             If the API base is missing and is required.
         """
         if (
-            self.type == LLMType.AzureOpenAIChat
-            or self.type == LLMType.AzureOpenAIEmbedding
+            self.type == ModelType.AzureOpenAIChat
+            or self.type == ModelType.AzureOpenAIEmbedding
         ) and (self.api_base is None or self.api_base.strip() == ""):
             raise AzureApiBaseMissingError(self.type.value)
 
@@ -153,8 +153,8 @@ def _validate_api_version(self) -> None:
             If the API base is missing and is required.
         """
         if (
-            self.type == LLMType.AzureOpenAIChat
-            or self.type == LLMType.AzureOpenAIEmbedding
+            self.type == ModelType.AzureOpenAIChat
+            or self.type == ModelType.AzureOpenAIEmbedding
         ) and (self.api_version is None or self.api_version.strip() == ""):
             raise AzureApiVersionMissingError(self.type.value)
 
@@ -173,8 +173,8 @@ def _validate_deployment_name(self) -> None:
             If the deployment name is missing and is required.
         """
         if (
-            self.type == LLMType.AzureOpenAIChat
-            or self.type == LLMType.AzureOpenAIEmbedding
+            self.type == ModelType.AzureOpenAIChat
+            or self.type == ModelType.AzureOpenAIEmbedding
         ) and (self.deployment_name is None or self.deployment_name.strip() == ""):
             raise AzureDeploymentNameMissingError(self.type.value)
 

graphrag/model/__init__.py graphrag/data_model/__init__.pygraphrag/model/__init__.py renamed to graphrag/data_model/__init__.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from typing import Any
 
-from graphrag.model.named import Named
+from graphrag.data_model.named import Named
 
 
 @dataclass

graphrag/model/community.py graphrag/data_model/community.pygraphrag/model/community.py renamed to graphrag/data_model/community.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from typing import Any
 
-from graphrag.model.named import Named
+from graphrag.data_model.named import Named
 
 
 @dataclass

graphrag/model/community_report.py graphrag/data_model/community_report.pygraphrag/model/community_report.py renamed to graphrag/data_model/community_report.py

@@ -6,7 +6,7 @@
 from dataclasses import dataclass
 from typing import Any
 
-from graphrag.model.identified import Identified
+from graphrag.data_model.identified import Identified
 
 
 @dataclass