mirror of
https://github.com/microsoft/graphrag.git
synced 2025-07-21 07:51:57 +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
14 lines
484 B
Python
14 lines
484 B
Python
# Copyright (c) 2024 Microsoft Corporation.
|
|
# Licensed under the MIT License
|
|
from pydantic import BaseModel
|
|
|
|
from graphrag.language_model.manager import ModelManager
|
|
from graphrag.language_model.protocol.base import ChatModel
|
|
|
|
|
|
def create_mock_llm(responses: list[str | BaseModel], name: str = "mock") -> ChatModel:
|
|
"""Creates a mock LLM that returns the given responses."""
|
|
return ModelManager().get_or_create_chat_model(
|
|
name, "mock_chat", responses=responses
|
|
)
|