2024-07-01 15:25:30 -06:00
|
|
|
# Copyright (c) 2024 Microsoft Corporation.
|
|
|
|
|
# Licensed under the MIT License
|
2025-02-20 08:56:20 -06:00
|
|
|
|
|
|
|
|
|
|
|
|
|
"""Tests for the GraphRAG LLM module."""
|
|
|
|
|
|
|
|
|
|
# Register MOCK providers
|
|
|
|
|
from graphrag.config.enums import ModelType
|
|
|
|
|
from graphrag.language_model.factory import ModelFactory
|
|
|
|
|
from tests.mock_provider import MockChatLLM, MockEmbeddingLLM
|
|
|
|
|
|
|
|
|
|
ModelFactory.register_chat(ModelType.MockChat, lambda **kwargs: MockChatLLM(**kwargs))
|
|
|
|
|
ModelFactory.register_embedding(
|
|
|
|
|
ModelType.MockEmbedding, lambda **kwargs: MockEmbeddingLLM(**kwargs)
|
|
|
|
|
)
|
