Reduced retries. Make connection from Jane Smith to GreenTech explicit.

Author: caseyclements

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

@@ -69,7 +69,7 @@ def documents():
 GreenTech Ltd. Leads SolarGrid Initiative
 
 San Francisco, CA — GreenTech Ltd. has emerged as a leader in renewable energy projects with the SolarGrid Initiative,
-a collaboration with ACME Corporation. Jane Smith, the SolarGrid project’s Lead Manager, highlighted its ambitious goal:
+a collaboration with ACME Corporation. GreenTech's Jane Smith, the SolarGrid project’s Lead Manager, highlighted its ambitious goal:
 providing affordable solar energy to underserved communities.
 
 GreenTech, headquartered in San Francisco, has worked closely with ACME since their partnership began in May 2021.
@@ -145,7 +145,7 @@ def test_add_docs_store(graph_store):
     assert 4 <= len(extracted_entities) < 8
 
 
-@flaky(max_runs=10, min_passes=2)
+@flaky(max_runs=3, min_passes=1)
 def test_extract_entity_names(graph_store, query_connection):
     query_entity_names = graph_store.extract_entity_names(query_connection)
     assert set(query_entity_names) == {"John Doe", "Jane Smith"}
@@ -157,7 +157,7 @@ def test_extract_entities_from_empty_string_names(graph_store):
     assert len(no_names) == 0
 
 
-@flaky(max_runs=10, min_passes=1)
+@flaky(max_runs=3, min_passes=1)
 def test_related_entities(graph_store):
     entity_names = ["John Doe", "Jane Smith"]
     related_entities = graph_store.related_entities(entity_names)
@@ -168,7 +168,7 @@ def test_related_entities(graph_store):
     assert len(no_entities) == 0
 
 
-@flaky(max_runs=10, min_passes=1)
+@flaky(max_runs=3, min_passes=1)
 def test_additional_entity_examples(entity_extraction_model, entity_example, documents):
     # First, create one client just to drop any existing collections
     client = MongoClient(CONNECTION_STRING)
@@ -192,15 +192,15 @@ def test_additional_entity_examples(entity_extraction_model, entity_example, doc
     assert len(new_entities) >= 2
 
 
-@flaky(max_runs=10, min_passes=1)
+@flaky(max_runs=3, min_passes=1)
 def test_chat_response(graph_store, query_connection):
     """Displays querying an existing Knowledge Graph Database"""
     answer = graph_store.chat_response(query_connection)
     assert isinstance(answer, AIMessage)
     assert "acme corporation" in answer.content.lower()
 
 
-@flaky(max_runs=10, min_passes=1)
+@flaky(max_runs=3, min_passes=1)
 def test_similarity_search(graph_store, query_connection):
     docs = graph_store.similarity_search(query_connection)
     assert len(docs) >= 4
@@ -209,7 +209,7 @@ def test_similarity_search(graph_store, query_connection):
     assert any("attributes" in d.keys() for d in docs)
 
 
-@flaky(max_runs=10, min_passes=1)
+@flaky(max_runs=3, min_passes=1)
 def test_validator(documents, entity_extraction_model):
     # Case 1. No existing collection.
     client = MongoClient(CONNECTION_STRING)
@@ -263,7 +263,7 @@ def test_validator(documents, entity_extraction_model):
     client.close()
 
 
-@flaky(max_runs=10, min_passes=1)
+@flaky(max_runs=3, min_passes=1)
 def test_allowed_entity_types(documents, entity_extraction_model):
     """Add allowed_entity_types. Use the validator to confirm behaviour."""
     allowed_entity_types = ["Person"]
@@ -291,7 +291,7 @@ def test_allowed_entity_types(documents, entity_extraction_model):
     all([len(e["relationships"].get("attributes", [])) == 0 for e in entities])
 
 
-@flaky(max_runs=10, min_passes=1)
+@flaky(max_runs=3, min_passes=1)
 def test_allowed_relationship_types(documents, entity_extraction_model):
     # drop collection
     client = MongoClient(CONNECTION_STRING)