mirror of
https://github.com/microsoft/graphrag.git
synced 2026-01-14 09:07:20 +08:00
17 lines
536 B
Python
17 lines
536 B
Python
# Copyright (c) 2024 Microsoft Corporation.
|
|
# Licensed under the MIT License
|
|
|
|
|
|
"""Tests for the GraphRAG LLM module."""
|
|
|
|
# Register MOCK providers
|
|
from graphrag.config.enums import ModelType
|
|
from graphrag.language_model.factory import ChatModelFactory, EmbeddingModelFactory
|
|
|
|
from tests.mock_provider import MockChatLLM, MockEmbeddingLLM
|
|
|
|
ChatModelFactory().register(ModelType.MockChat, lambda **kwargs: MockChatLLM(**kwargs))
|
|
EmbeddingModelFactory().register(
|
|
ModelType.MockEmbedding, lambda **kwargs: MockEmbeddingLLM(**kwargs)
|
|
)
|