INTPYTHON-629 Take 2 - Encouraged LLM to find as many entities and relationships as it could (#148)

Author: caseyclements

libs/langchain-mongodb/langchain_mongodb/graphrag/graph.py

@@ -98,7 +98,7 @@ def __init__(
         entity_extraction_model: BaseChatModel,
         entity_prompt: Optional[ChatPromptTemplate] = None,
         query_prompt: Optional[ChatPromptTemplate] = None,
-        max_depth: int = 2,
+        max_depth: int = 3,
         allowed_entity_types: Optional[List[str]] = None,
         allowed_relationship_types: Optional[List[str]] = None,
         entity_examples: Optional[str] = None,
@@ -235,7 +235,7 @@ def from_connection_string(
         entity_extraction_model: BaseChatModel,
         entity_prompt: ChatPromptTemplate = prompts.entity_prompt,
         query_prompt: ChatPromptTemplate = prompts.query_prompt,
-        max_depth: int = 2,
+        max_depth: int = 3,
         allowed_entity_types: Optional[List[str]] = None,
         allowed_relationship_types: Optional[List[str]] = None,
         entity_examples: Optional[str] = None,

libs/langchain-mongodb/langchain_mongodb/graphrag/prompts.py

@@ -6,17 +6,9 @@
 
 ENTITY_EXTRACTION_INSTRUCTIONS = """
 ## Overview
-You are a meticulous analyst tasked with extracting information from unstructured text
+You are a meticulous analyst tasked with identifying potential entities from unstructured text
 to build a knowledge graph in a structured json format of entities (nodes) and their relationships (edges).
-The graph will be stored in a MongoDB Collection and traversed using $graphLookup
-from starting points of entity nodes matching names found in a query, and follow their relationships.
-
-Use the following as guidelines.
-
-- Simplicity: The graph should have as few entities and relationship types as needed to convey the information in the input.
-- Consistency: Connections can only be made if entities and relationships use consistent naming.
-- Generality: The graph should be useful for describing the concepts in not just this document but other similar documents.
-- Accuracy: Do not add any information that is not explicitly mentioned in the text.
+**Include as many entities and relationships as you can.**
 
 INPUT: You will be provided a text document.
 OUTPUT:
@@ -84,21 +76,19 @@
 
 
 NAME_EXTRACTION_INSTRUCTIONS = """
-You are a meticulous analyst tasked with extracting information from documents to form
-knowledge graphs of entities (nodes) and their relationships (edges).
-
-You will be provided a short document (query) from which you infer the entity names.
-You need not think about relationships between the entities. You only need names.
+You are an analyst tasked with identifying potential entities in text documents.
+You will be provided a short document from which you infer entity names.
+Identify as many as possible.
 
 Provide your response as a valid JSON Array of entity names
 or human-readable identifiers, found in the text.
 
 **Allowed Entity Types**:
-- Extract ONLY entities whose `type` matches one of the following: {allowed_entity_types}.
-- NOTE: If this list is empty, ANY `type` is permitted.
+- By default, all types are permitted.
+- If a non-empty list is provided, extract ONLY entities whose `type` matches one of the following: [{allowed_entity_types}].
 
 ### Examples of Exclusions:
-- If `allowed_entity_types` is `["Person", "Organization"]`, and the text mentions "Event" or "Location",
+- If `allowed_entity_types` is `["Person", "Organization"]`, and the text mentions an "Event" or "Location",
   these entities must **NOT** be included in the output.
 
  ## Examples:

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

@@ -4,7 +4,6 @@
 from typing import Generator
 
 import pytest
-from flaky import flaky  # type:ignore[import-untyped]
 from langchain_core.documents import Document
 from langchain_core.language_models.chat_models import BaseChatModel
 from langchain_core.messages import AIMessage
@@ -69,7 +68,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.
@@ -125,8 +124,6 @@ def graph_store(collection, entity_extraction_model, documents) -> MongoDBGraphS
     store = MongoDBGraphStore(
         collection=collection,
         entity_extraction_model=entity_extraction_model,
-        entity_prompt=entity_prompt,
-        query_prompt=query_prompt,
     )
     bulkwrite_results = store.add_documents(documents)
     assert len(bulkwrite_results) == len(documents)
@@ -145,7 +142,6 @@ def test_add_docs_store(graph_store):
     assert 4 <= len(extracted_entities) < 8
 
 
-@flaky(max_runs=3, min_passes=2)
 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 +153,6 @@ def test_extract_entities_from_empty_string_names(graph_store):
     assert len(no_names) == 0
 
 
-@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 +163,6 @@ def test_related_entities(graph_store):
     assert len(no_entities) == 0
 
 
-@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 +186,13 @@ def test_additional_entity_examples(entity_extraction_model, entity_example, doc
     assert len(new_entities) >= 2
 
 
-@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=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 +201,6 @@ def test_similarity_search(graph_store, query_connection):
     assert any("attributes" in d.keys() for d in docs)
 
 
-@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 +254,6 @@ def test_validator(documents, entity_extraction_model):
     client.close()
 
 
-@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 +281,6 @@ def test_allowed_entity_types(documents, entity_extraction_model):
     all([len(e["relationships"].get("attributes", [])) == 0 for e in entities])
 
 
-@flaky(max_runs=3, min_passes=1)
 def test_allowed_relationship_types(documents, entity_extraction_model):
     # drop collection
     client = MongoClient(CONNECTION_STRING)