mirror of
https://github.com/microsoft/graphrag.git
synced 2025-06-26 23:19:58 +00:00

* Refactor config - Add new ModelConfig to represent LLM settings - Combines LLMParameters, ParallelizationParameters, encoding_model, and async_mode - Add top level models config that is a list of available LLM ModelConfigs - Remove LLMConfig inheritance and delete LLMConfig - Replace the inheritance with a model_id reference to the ModelConfig listed in the top level models config - Remove all fallbacks and hydration logic from create_graphrag_config - This removes the automatic env variable overrides - Support env variables within config files using Templating - This requires "$" to be escaped with extra "$" so ".*\\.txt$" becomes ".*\\.txt$$" - Update init content to initialize new config file with the ModelConfig structure * Use dict of ModelConfig instead of list * Add model validations and unit tests * Fix ruff checks * Add semversioner change * Fix unit tests * validate root_dir in pydantic model * Rename ModelConfig to LanguageModelConfig * Rename ModelConfigMissingError to LanguageModelConfigMissingError * Add validationg for unexpected API keys * Allow skipping pydantic validation for testing/mocking purposes. * Add default lm configs to verb tests * smoke test * remove config from flows to fix llm arg mapping * Fix embedding llm arg mapping * Remove timestamp from smoke test outputs * Remove unused "subworkflows" smoke test properties * Add models to smoke test configs * Update smoke test output path * Send logs to logs folder * Fix output path * Fix csv test file pattern * Update placeholder * Format * Instantiate default model configs * Fix unit tests for config defaults * Fix migration notebook * Remove create_pipeline_config * Remove several unused config models * Remove indexing embedding and input configs * Move embeddings function to config * Remove skip_workflows * Remove skip embeddings in favor of explicit naming * fix unit test spelling mistake * self.models[model_id] is already a language model. Remove redundant casting. * update validation errors to instruct users to rerun graphrag init * instantiate LanguageModelConfigs with validation * skip validation in unit tests * update verb tests to use default model settings instead of skipping validation * test using llm settings * cleanup verb tests * remove unsafe default model config * remove the ability to skip pydantic validation * remove None union types when default values are set * move vector_store from embeddings to top level of config and delete resolve_paths * update vector store settings * fix vector store and smoke tests * fix serializing vector_store settings * fix vector_store usage * fix vector_store type * support cli overrides for loading graphrag config * rename storage to output * Add --force flag to init * Remove run_id and resume, fix Drift config assignment * Ruff --------- Co-authored-by: Nathan Evans <github@talkswithnumbers.com> Co-authored-by: Alonso Guevara <alonsog@microsoft.com>
72 lines
2.2 KiB
Python
72 lines
2.2 KiB
Python
# Copyright (c) 2024 Microsoft Corporation.
|
|
# Licensed under the MIT License
|
|
|
|
from graphrag.callbacks.noop_workflow_callbacks import NoopWorkflowCallbacks
|
|
from graphrag.config.create_graphrag_config import create_graphrag_config
|
|
from graphrag.config.embeddings import (
|
|
all_embeddings,
|
|
)
|
|
from graphrag.config.enums import TextEmbeddingTarget
|
|
from graphrag.index.workflows.generate_text_embeddings import (
|
|
run_workflow,
|
|
)
|
|
from graphrag.utils.storage import load_table_from_storage
|
|
|
|
from .util import (
|
|
DEFAULT_MODEL_CONFIG,
|
|
create_test_context,
|
|
)
|
|
|
|
|
|
async def test_generate_text_embeddings():
|
|
context = await create_test_context(
|
|
storage=[
|
|
"create_final_documents",
|
|
"create_final_relationships",
|
|
"create_final_text_units",
|
|
"create_final_entities",
|
|
"create_final_community_reports",
|
|
]
|
|
)
|
|
|
|
config = create_graphrag_config({"models": DEFAULT_MODEL_CONFIG})
|
|
llm_settings = config.get_language_model_config(
|
|
config.embeddings.model_id
|
|
).model_dump()
|
|
|
|
config.embeddings.strategy = {
|
|
"type": "mock",
|
|
"llm": llm_settings,
|
|
}
|
|
config.embeddings.target = TextEmbeddingTarget.all
|
|
config.snapshots.embeddings = True
|
|
|
|
await run_workflow(
|
|
config,
|
|
context,
|
|
NoopWorkflowCallbacks(),
|
|
)
|
|
|
|
parquet_files = context.storage.keys()
|
|
|
|
for field in all_embeddings:
|
|
assert f"embeddings.{field}.parquet" in parquet_files
|
|
|
|
# entity description should always be here, let's assert its format
|
|
entity_description_embeddings = await load_table_from_storage(
|
|
"embeddings.entity.description", context.storage
|
|
)
|
|
|
|
assert len(entity_description_embeddings.columns) == 2
|
|
assert "id" in entity_description_embeddings.columns
|
|
assert "embedding" in entity_description_embeddings.columns
|
|
|
|
# every other embedding is optional but we've turned them all on, so check a random one
|
|
document_text_embeddings = await load_table_from_storage(
|
|
"embeddings.document.text", context.storage
|
|
)
|
|
|
|
assert len(document_text_embeddings.columns) == 2
|
|
assert "id" in document_text_embeddings.columns
|
|
assert "embedding" in document_text_embeddings.columns
|