graphrag/tests/verbs/test_create_final_documents.py
gaudyb 17658c5df8
New workflow to generate embeddings in a single workflow (#1296)
* New workflow to generate embeddings in a single workflow

* New workflow to generate embeddings in a single workflow

* version change

* clean tests without any embeddings references

* clean tests without any embeddings references

* remove code

* feedback implemented

* changes in logic

* feedback implemented

* store in table bug fixed

* smoke test for generate_text_embeddings workflow

* smoke test fix

* add generate_text_embeddings to the list of transient workflows

* smoke tests

* fix

* ruff formatting updates

* fix

* smoke test fixed

* smoke test fixed

* fix lancedb import

* smoke test fix

* ignore sorting

* smoke test fixed

* smoke test fixed

* check smoke test

* smoke test fixed

* change config for vector store

* format fix

* vector store changes

* revert debug profile back to empty filepath

* merge conflict solved

* merge conflict solved

* format fixed

* format fixed

* fix return dataframe

* snapshot fix

* format fix

* embeddings param implemented

* validation fixes

* fix map

* fix map

* fix properties

* config updates

* smoke test fixed

* settings change

* Update collection config and rework back-compat

* Repalce . with - for embedding store

---------

Co-authored-by: Alonso Guevara <alonsog@microsoft.com>
Co-authored-by: Josh Bradley <joshbradley@microsoft.com>
Co-authored-by: Nathan Evans <github@talkswithnumbers.com>
2024-11-01 15:01:35 -07:00

61 lines
1.5 KiB
Python

# Copyright (c) 2024 Microsoft Corporation.
# Licensed under the MIT License
from graphrag.index.workflows.v1.create_final_documents import (
build_steps,
workflow_name,
)
from .util import (
compare_outputs,
get_config_for_workflow,
get_workflow_output,
load_expected,
load_input_tables,
)
async def test_create_final_documents():
input_tables = load_input_tables([
"workflow:create_final_text_units",
])
expected = load_expected(workflow_name)
config = get_config_for_workflow(workflow_name)
steps = build_steps(config)
actual = await get_workflow_output(
input_tables,
{
"steps": steps,
},
)
compare_outputs(actual, expected)
async def test_create_final_documents_with_attribute_columns():
input_tables = load_input_tables(["workflow:create_final_text_units"])
expected = load_expected(workflow_name)
config = get_config_for_workflow(workflow_name)
config["document_attribute_columns"] = ["title"]
steps = build_steps(config)
actual = await get_workflow_output(
input_tables,
{
"steps": steps,
},
)
# we should have dropped "title" and added "attributes"
# our test dataframe does not have attributes, so we'll assert without it
# and separately confirm it is in the output
compare_outputs(actual, expected, columns=["id", "text_unit_ids", "raw_content"])
assert len(actual.columns) == 4
assert "attributes" in actual.columns