Skip to content

Commit 1174e95

Browse files
committed
Bumped graphrag retries to 10.
1 parent c83235a commit 1174e95

File tree

1 file changed

+8
-8
lines changed

1 file changed

+8
-8
lines changed

libs/langchain-mongodb/tests/integration_tests/test_graphrag.py

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -145,7 +145,7 @@ def test_add_docs_store(graph_store):
145145
assert 4 <= len(extracted_entities) < 8
146146

147147

148-
@flaky(max_runs=3, min_passes=2)
148+
@flaky(max_runs=10, min_passes=2)
149149
def test_extract_entity_names(graph_store, query_connection):
150150
query_entity_names = graph_store.extract_entity_names(query_connection)
151151
assert set(query_entity_names) == {"John Doe", "Jane Smith"}
@@ -157,7 +157,7 @@ def test_extract_entities_from_empty_string_names(graph_store):
157157
assert len(no_names) == 0
158158

159159

160-
@flaky(max_runs=3, min_passes=1)
160+
@flaky(max_runs=10, min_passes=1)
161161
def test_related_entities(graph_store):
162162
entity_names = ["John Doe", "Jane Smith"]
163163
related_entities = graph_store.related_entities(entity_names)
@@ -168,7 +168,7 @@ def test_related_entities(graph_store):
168168
assert len(no_entities) == 0
169169

170170

171-
@flaky(max_runs=3, min_passes=1)
171+
@flaky(max_runs=10, min_passes=1)
172172
def test_additional_entity_examples(entity_extraction_model, entity_example, documents):
173173
# First, create one client just to drop any existing collections
174174
client = MongoClient(CONNECTION_STRING)
@@ -192,15 +192,15 @@ def test_additional_entity_examples(entity_extraction_model, entity_example, doc
192192
assert len(new_entities) >= 2
193193

194194

195-
@flaky(max_runs=3, min_passes=1)
195+
@flaky(max_runs=10, min_passes=1)
196196
def test_chat_response(graph_store, query_connection):
197197
"""Displays querying an existing Knowledge Graph Database"""
198198
answer = graph_store.chat_response(query_connection)
199199
assert isinstance(answer, AIMessage)
200200
assert "acme corporation" in answer.content.lower()
201201

202202

203-
@flaky(max_runs=3, min_passes=1)
203+
@flaky(max_runs=10, min_passes=1)
204204
def test_similarity_search(graph_store, query_connection):
205205
docs = graph_store.similarity_search(query_connection)
206206
assert len(docs) >= 4
@@ -209,7 +209,7 @@ def test_similarity_search(graph_store, query_connection):
209209
assert any("attributes" in d.keys() for d in docs)
210210

211211

212-
@flaky(max_runs=3, min_passes=1)
212+
@flaky(max_runs=10, min_passes=1)
213213
def test_validator(documents, entity_extraction_model):
214214
# Case 1. No existing collection.
215215
client = MongoClient(CONNECTION_STRING)
@@ -263,7 +263,7 @@ def test_validator(documents, entity_extraction_model):
263263
client.close()
264264

265265

266-
@flaky(max_runs=3, min_passes=1)
266+
@flaky(max_runs=10, min_passes=1)
267267
def test_allowed_entity_types(documents, entity_extraction_model):
268268
"""Add allowed_entity_types. Use the validator to confirm behaviour."""
269269
allowed_entity_types = ["Person"]
@@ -291,7 +291,7 @@ def test_allowed_entity_types(documents, entity_extraction_model):
291291
all([len(e["relationships"].get("attributes", [])) == 0 for e in entities])
292292

293293

294-
@flaky(max_runs=3, min_passes=1)
294+
@flaky(max_runs=10, min_passes=1)
295295
def test_allowed_relationship_types(documents, entity_extraction_model):
296296
# drop collection
297297
client = MongoClient(CONNECTION_STRING)

0 commit comments

Comments
 (0)