Add default lm configs to verb tests

dworthen · dworthen · commit 73eb5c8b79ff · 2025-01-10T09:48:15.000-08:00
diff --git a/tests/unit/config/test_config.py b/tests/unit/config/test_config.py
@@ -36,6 +36,10 @@ def test_missing_default_embedding_model_config() -> None:
         create_graphrag_config({"models": models_config})
 
 
+def test_skip_validation() -> None:
+    create_graphrag_config(skip_validation=True)
+
+
 def test_missing_openai_required_api_key() -> None:
     model_config_missing_api_key = {
         defs.DEFAULT_CHAT_MODEL_ID: {
diff --git a/tests/verbs/test_compute_communities.py b/tests/verbs/test_compute_communities.py
@@ -7,6 +7,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     compare_outputs,
     create_test_context,
     load_test_table,
@@ -20,7 +21,9 @@ async def test_compute_communities():
         storage=["base_relationship_edges"],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
 
     await run_workflow(
         config,
diff --git a/tests/verbs/test_create_base_text_units.py b/tests/verbs/test_create_base_text_units.py
@@ -7,6 +7,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     compare_outputs,
     create_test_context,
     load_test_table,
@@ -18,7 +19,9 @@ async def test_create_base_text_units():
 
     context = await create_test_context()
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     # test data was created with 4o, so we need to match the encoding for chunks to be identical
     config.chunks.encoding_model = "o200k_base"
 
diff --git a/tests/verbs/test_create_final_communities.py b/tests/verbs/test_create_final_communities.py
@@ -10,6 +10,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     compare_outputs,
     create_test_context,
     load_test_table,
@@ -27,7 +28,9 @@ async def test_create_final_communities():
         ],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
 
     await run_workflow(
         config,
diff --git a/tests/verbs/test_create_final_community_reports.py b/tests/verbs/test_create_final_community_reports.py
@@ -19,6 +19,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     compare_outputs,
     create_test_context,
     load_test_table,
@@ -61,7 +62,9 @@ async def test_create_final_community_reports():
         ]
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     config.community_reports.strategy = {
         "type": "graph_intelligence",
         "llm": MOCK_LLM_CONFIG,
@@ -96,7 +99,9 @@ async def test_create_final_community_reports_missing_llm_throws():
         ]
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     config.community_reports.strategy = {
         "type": "graph_intelligence",
     }
diff --git a/tests/verbs/test_create_final_covariates.py b/tests/verbs/test_create_final_covariates.py
@@ -15,6 +15,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     create_test_context,
     load_test_table,
 )
@@ -36,7 +37,9 @@ async def test_create_final_covariates():
         storage=["create_base_text_units"],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     config.claim_extraction.strategy = {
         "type": "graph_intelligence",
         "llm": MOCK_LLM_CONFIG,
@@ -85,7 +88,9 @@ async def test_create_final_covariates_missing_llm_throws():
         storage=["create_base_text_units"],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     config.claim_extraction.strategy = {
         "type": "graph_intelligence",
         "claim_description": "description",
diff --git a/tests/verbs/test_create_final_documents.py b/tests/verbs/test_create_final_documents.py
@@ -10,6 +10,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     compare_outputs,
     create_test_context,
     load_test_table,
@@ -23,7 +24,9 @@ async def test_create_final_documents():
         storage=["create_base_text_units"],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
 
     await run_workflow(
         config,
@@ -43,7 +46,9 @@ async def test_create_final_documents_with_attribute_columns():
         storage=["create_base_text_units"],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     config.input.document_attribute_columns = ["title"]
 
     await run_workflow(
diff --git a/tests/verbs/test_create_final_entities.py b/tests/verbs/test_create_final_entities.py
@@ -10,6 +10,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     compare_outputs,
     create_test_context,
     load_test_table,
@@ -23,7 +24,9 @@ async def test_create_final_entities():
         storage=["base_entity_nodes"],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
 
     await run_workflow(
         config,
diff --git a/tests/verbs/test_create_final_nodes.py b/tests/verbs/test_create_final_nodes.py
@@ -10,6 +10,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     compare_outputs,
     create_test_context,
     load_test_table,
@@ -27,7 +28,9 @@ async def test_create_final_nodes():
         ],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
 
     await run_workflow(
         config,
diff --git a/tests/verbs/test_create_final_relationships.py b/tests/verbs/test_create_final_relationships.py
@@ -11,6 +11,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     compare_outputs,
     create_test_context,
     load_test_table,
@@ -24,7 +25,9 @@ async def test_create_final_relationships():
         storage=["base_relationship_edges"],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
 
     await run_workflow(
         config,
diff --git a/tests/verbs/test_create_final_text_units.py b/tests/verbs/test_create_final_text_units.py
@@ -10,6 +10,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     compare_outputs,
     create_test_context,
     load_test_table,
@@ -28,7 +29,9 @@ async def test_create_final_text_units():
         ],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     config.claim_extraction.enabled = True
 
     await run_workflow(
@@ -54,7 +57,9 @@ async def test_create_final_text_units_no_covariates():
         ],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     config.claim_extraction.enabled = False
 
     await run_workflow(
diff --git a/tests/verbs/test_extract_graph.py b/tests/verbs/test_extract_graph.py
@@ -12,6 +12,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     create_test_context,
     load_test_table,
 )
@@ -55,7 +56,9 @@ async def test_extract_graph():
         storage=["create_base_text_units"],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     config.entity_extraction.strategy = {
         "type": "graph_intelligence",
         "llm": MOCK_LLM_ENTITY_CONFIG,
@@ -97,7 +100,9 @@ async def test_extract_graph_missing_llm_throws():
         storage=["create_base_text_units"],
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     config.entity_extraction.strategy = {
         "type": "graph_intelligence",
         "llm": MOCK_LLM_ENTITY_CONFIG,
diff --git a/tests/verbs/test_generate_text_embeddings.py b/tests/verbs/test_generate_text_embeddings.py
@@ -13,6 +13,7 @@
 from graphrag.utils.storage import load_table_from_storage
 
 from .util import (
+    DEFAULT_MODEL_CONFIG,
     create_test_context,
 )
 
@@ -28,7 +29,9 @@ async def test_generate_text_embeddings():
         ]
     )
 
-    config = create_graphrag_config(skip_validation=True)
+    config = create_graphrag_config(
+        {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True
+    )
     config.embeddings.strategy = {
         "type": "mock",
     }
diff --git a/tests/verbs/util.py b/tests/verbs/util.py
@@ -4,12 +4,32 @@
 import pandas as pd
 from pandas.testing import assert_series_equal
 
+import graphrag.config.defaults as defs
 from graphrag.index.context import PipelineRunContext
 from graphrag.index.run.utils import create_run_context
 from graphrag.utils.storage import write_table_to_storage
 
 pd.set_option("display.max_columns", None)
 
+FAKE_API_KEY = "NOT_AN_API_KEY"
+
+DEFAULT_CHAT_MODEL_CONFIG = {
+    "api_key": FAKE_API_KEY,
+    "type": defs.LLM_TYPE.value,
+    "model": defs.LLM_MODEL,
+}
+
+DEFAULT_EMBEDDING_MODEL_CONFIG = {
+    "api_key": FAKE_API_KEY,
+    "type": defs.EMBEDDING_TYPE.value,
+    "model": defs.EMBEDDING_MODEL,
+}
+
+DEFAULT_MODEL_CONFIG = {
+    defs.DEFAULT_CHAT_MODEL_ID: DEFAULT_CHAT_MODEL_CONFIG,
+    defs.DEFAULT_EMBEDDING_MODEL_ID: DEFAULT_EMBEDDING_MODEL_CONFIG,
+}
+
 
 async def create_test_context(storage: list[str] | None = None) -> PipelineRunContext:
     """Create a test context with tables loaded into storage storage."""

Original file line number	Diff line number	Diff line change
`@@ -19,6 +19,7 @@`
`19`	`19`	`from graphrag.utils.storage import load_table_from_storage`
`20`	`20`
`21`	`21`	`from .util import (`
	`22`	`+ DEFAULT_MODEL_CONFIG,`
`22`	`23`	`compare_outputs,`
`23`	`24`	`create_test_context,`
`24`	`25`	`load_test_table,`
`@@ -61,7 +62,9 @@ async def test_create_final_community_reports():`
`61`	`62`	`]`
`62`	`63`	`)`
`63`	`64`
`64`		`- config = create_graphrag_config(skip_validation=True)`
	`65`	`+ config = create_graphrag_config(`
	`66`	`+ {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True`
	`67`	`+ )`
`65`	`68`	`config.community_reports.strategy = {`
`66`	`69`	`"type": "graph_intelligence",`
`67`	`70`	`"llm": MOCK_LLM_CONFIG,`
`@@ -96,7 +99,9 @@ async def test_create_final_community_reports_missing_llm_throws():`
`96`	`99`	`]`
`97`	`100`	`)`
`98`	`101`
`99`		`- config = create_graphrag_config(skip_validation=True)`
	`102`	`+ config = create_graphrag_config(`
	`103`	`+ {"models": DEFAULT_MODEL_CONFIG}, skip_validation=True`
	`104`	`+ )`
`100`	`105`	`config.community_reports.strategy = {`
`101`	`106`	`"type": "graph_intelligence",`
`102`	`107`	`}`