Use dict of ModelConfig instead of list

Author: dworthen

graphrag/config/init_content.py

@@ -13,7 +13,7 @@
 ## There are a number of settings to tune the threading and token limits for LLM calls - check the docs.
 
 models:
-  - id: {defs.DEFAULT_CHAT_MODEL_ID}
+  {defs.DEFAULT_CHAT_MODEL_ID}:
     api_key: ${{GRAPHRAG_API_KEY}} # set this in the generated .env file
     type: {defs.LLM_TYPE.value} # or azure_openai_chat
     model: {defs.LLM_MODEL}
@@ -26,7 +26,7 @@
     # api_version: 2024-02-15-preview
     # organization: <organization_id>
     # deployment_name: <azure_model_deployment_name>
-  - id: {defs.DEFAULT_EMBEDDING_MODEL_ID}
+  {defs.DEFAULT_EMBEDDING_MODEL_ID}:
     api_key: ${{GRAPHRAG_API_KEY}}
     type: {defs.EMBEDDING_TYPE.value} # or azure_openai_embedding
     model: {defs.EMBEDDING_MODEL}

graphrag/config/models/graph_rag_config.py

@@ -45,8 +45,8 @@ def __str__(self):
         description="The root directory for the configuration.", default="."
     )
 
-    models: list[ModelConfig] = Field(
-        description="Available language model configurations.", default=[]
+    models: dict[str, ModelConfig] = Field(
+        description="Available language model configurations.", default={}
     )
 
     reporting: ReportingConfig = Field(
@@ -179,8 +179,7 @@ def get_model_config(self, model_id: str) -> ModelConfig:
         ValueError
             If the model ID is not found in the configuration.
         """
-        for model in self.models:
-            if model.id == model_id:
-                return model
-        err_msg = f"Model ID {model_id} not found in configuration."
-        raise ValueError(err_msg)
+        if model_id not in self.models:
+            err_msg = f"Model ID {model_id} not found in configuration."
+            raise ValueError(err_msg)
+        return self.models[model_id]

graphrag/config/models/model_config.py

@@ -13,9 +13,6 @@
 class ModelConfig(BaseModel):
     """Language model configuration."""
 
-    id: str = Field(
-        description="A user friendly ID of the model used elsewhere in GraphRAGConfig for selecting models."
-    )
     api_key: str | None = Field(
         description="The API key to use for the LLM service.",
         default=None,