Update chat and embedding type field for litellm use and future migration away from fnllm.

Author: dworthen

graphrag/config/enums.py

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

graphrag/config/models/language_model_config.py

@@ -113,8 +113,7 @@ def _validate_model_provider(self) -> None:
             If the model provider is not recognized.
         """
         if (
-            self.type == ModelType.LitellmChat
-            or self.type == ModelType.LitellmEmbedding
+            self.type == ModelType.Chat or self.type == ModelType.Embedding
         ) and self.model_provider.strip() == "":
             msg = "Model provider must be specified when using Litellm."
             raise KeyError(msg)
@@ -145,8 +144,8 @@ def _validate_encoding_model(self) -> None:
             If the model name is not recognized.
         """
         if (
-            self.type != ModelType.LitellmChat
-            and self.type != ModelType.LitellmEmbedding
+            self.type != ModelType.Chat
+            and self.type != ModelType.Embedding
             and self.encoding_model.strip() == ""
         ):
             self.encoding_model = tiktoken.encoding_name_for_model(self.model)

graphrag/language_model/factory.py

@@ -109,9 +109,7 @@ def is_supported_model(cls, model_type: str) -> bool:
 ModelFactory.register_chat(
     ModelType.OpenAIChat.value, lambda **kwargs: OpenAIChatFNLLM(**kwargs)
 )
-ModelFactory.register_chat(
-    ModelType.LitellmChat, lambda **kwargs: LitellmChatModel(**kwargs)
-)
+ModelFactory.register_chat(ModelType.Chat, lambda **kwargs: LitellmChatModel(**kwargs))
 
 ModelFactory.register_embedding(
     ModelType.AzureOpenAIEmbedding.value,
@@ -121,5 +119,5 @@ def is_supported_model(cls, model_type: str) -> bool:
     ModelType.OpenAIEmbedding.value, lambda **kwargs: OpenAIEmbeddingFNLLM(**kwargs)
 )
 ModelFactory.register_embedding(
-    ModelType.LitellmEmbedding, lambda **kwargs: LitellmEmbeddingModel(**kwargs)
+    ModelType.Embedding, lambda **kwargs: LitellmEmbeddingModel(**kwargs)
 )