mirror of
https://github.com/microsoft/graphrag.git
synced 2025-06-26 23:19:58 +00:00
7bdeaee94a
* 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
16 lines
506 B
Python
16 lines
506 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 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)
|
|
)
|