Skip to content

Commit 839be27

Browse files
authored
"extras" must be resolved (#248)
* "extras" must be resolved * Add a test case * Ruff * Move test to e2e tests * mypy + comments * Fix test
1 parent 5340073 commit 839be27

File tree

3 files changed

+28
-5
lines changed

3 files changed

+28
-5
lines changed

src/neo4j_graphrag/experimental/pipeline/config/pipeline_config.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,9 @@
2727
LLMType,
2828
Neo4jDriverType,
2929
)
30+
from neo4j_graphrag.experimental.pipeline.config.param_resolver import (
31+
ParamConfig,
32+
)
3033
from neo4j_graphrag.experimental.pipeline.config.types import PipelineType
3134
from neo4j_graphrag.experimental.pipeline.types import (
3235
ComponentDefinition,
@@ -48,7 +51,7 @@ class AbstractPipelineConfig(AbstractConfig):
4851
llm_config: dict[str, LLMType] = {}
4952
embedder_config: dict[str, EmbedderType] = {}
5053
# extra parameters values that can be used in different places of the config file
51-
extras: dict[str, Any] = {}
54+
extras: dict[str, ParamConfig] = {}
5255

5356
DEFAULT_NAME: ClassVar[str] = "default"
5457
"""Name of the default item in dict

tests/e2e/data/config_files/simple_kg_pipeline_config.json

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,12 @@
11
{
22
"version_": "1",
33
"template_": "SimpleKGPipeline",
4+
"extras": {
5+
"openai_api_key": {
6+
"resolver_": "ENV",
7+
"var_": "MY_OPENAI_KEY"
8+
}
9+
},
410
"neo4j_config": {
511
"params_": {
612
"uri": {
@@ -36,8 +42,8 @@
3642
"class_": "OpenAIEmbeddings",
3743
"params_": {
3844
"api_key": {
39-
"resolver_": "ENV",
40-
"var_": "OPENAI_API_KEY"
45+
"resolver_": "CONFIG_KEY",
46+
"key_": "extras.openai_api_key"
4147
}
4248
}
4349
},

tests/e2e/experimental/pipeline/config/test_pipeline_runner_e2e.py

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -129,19 +129,33 @@ async def test_simple_kg_pipeline_from_json_config(
129129
os.environ["NEO4J_USER"] = "neo4j"
130130
os.environ["NEO4J_PASSWORD"] = "password"
131131
os.environ["OPENAI_API_KEY"] = "sk-my-secret-key"
132+
os.environ["MY_OPENAI_KEY"] = "my-openai-key"
132133

133134
runner = PipelineRunner.from_config_file(
134135
"tests/e2e/data/config_files/simple_kg_pipeline_config.json"
135136
)
137+
138+
# check extras and API keys are handled as expected
139+
config = runner.config
140+
assert config is not None
141+
# extras must be resolved:
142+
assert config._global_data["extras"] == {"openai_api_key": "my-openai-key"}
143+
# API key for LLM is read from env vars (see config file)
144+
default_llm = config._global_data["llm_config"]["default"]
145+
assert default_llm.client.api_key == "sk-my-secret-key"
146+
# API key for embedder is read from extras (see config file)
147+
default_embedder = config._global_data["embedder_config"]["default"]
148+
assert default_embedder.client.api_key == "my-openai-key"
149+
150+
# then run pipeline and check results
136151
res = await runner.run({"file_path": "tests/e2e/data/documents/harry_potter.pdf"})
137152
assert isinstance(res, PipelineResult)
138-
# print(await runner.pipeline.store.get_result_for_component(res.run_id, "splitter"))
139153
assert res.result["resolver"] == {
140154
"number_of_nodes_to_resolve": 3,
141155
"number_of_created_nodes": 3,
142156
}
143157
nodes = driver.execute_query("MATCH (n) RETURN n")
144-
# 1 chunk + 1 document + 3 nodes
158+
# 1 chunk + 1 document + 3 __Entity__ nodes
145159
assert len(nodes.records) == 5
146160

147161

0 commit comments

Comments
 (0)