gmakstutis/cognee
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: improve llm output mock

Browse source
This commit is contained in:
Boris Arzentar 2025-11-04 14:48:06 +01:00
parent e29e53df43
commit ba4e5c3c69
No known key found for this signature in database
GPG key ID: D5CC274C784807B7
1 changed files with 73 additions and 57 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

130
cognee/tests/test_delete_default_graph.py
View file

@ -41,63 +41,79 @@ async def main(mock_create_structured_output: AsyncMock):
assert not await vector_engine.has_collection("TextSummary_text") assert not await vector_engine.has_collection("TextSummary_text")
assert not await vector_engine.has_collection("TextDocument_text") assert not await vector_engine.has_collection("TextDocument_text")
mock_create_structured_output.side_effect = [ def mock_llm_output(text_input: str, system_prompt: str, response_model):
"", # For LLM connection test if text_input == "test": # LLM connection test
KnowledgeGraph( return "test"
nodes=[
Node(id="John", name="John", type="Person", description="John is a person"), if "John" in text_input and response_model == SummarizedContent:
Node( return SummarizedContent(
id="Apple", summary="Summary of John's work.", description="Summary of John's work."
name="Apple", )
type="Company",
description="Apple is a company", if "Marie" in text_input and response_model == SummarizedContent:
), return SummarizedContent(
Node( summary="Summary of Marie's work.", description="Summary of Marie's work."
id="Food for Hungry", )
name="Food for Hungry",
type="Non-profit organization", if "Marie" in text_input and response_model == KnowledgeGraph:
description="Food for Hungry is a non-profit organization", return KnowledgeGraph(
), nodes=[
], Node(id="Marie", name="Marie", type="Person", description="Marie is a person"),
edges=[ Node(
Edge(source_node_id="John", target_node_id="Apple", relationship_name="works_for"), id="Apple",
Edge( name="Apple",
source_node_id="John", type="Company",
target_node_id="Food for Hungry", description="Apple is a company",
relationship_name="works_for", ),
), Node(
], id="MacOS",
), name="MacOS",
KnowledgeGraph( type="Product",
nodes=[ description="MacOS is Apple's operating system",
Node(id="Marie", name="Marie", type="Person", description="Marie is a person"), ),
Node( ],
id="Apple", edges=[
name="Apple", Edge(
type="Company", source_node_id="Marie",
description="Apple is a company", target_node_id="Apple",
), relationship_name="works_for",
Node( ),
id="MacOS", Edge(
name="MacOS", source_node_id="Marie", target_node_id="MacOS", relationship_name="works_on"
type="Product", ),
description="MacOS is Apple's operating system", ],
), )
],
edges=[ if "John" in text_input and response_model == KnowledgeGraph:
Edge( return KnowledgeGraph(
source_node_id="Marie", nodes=[
target_node_id="Apple", Node(id="John", name="John", type="Person", description="John is a person"),
relationship_name="works_for", Node(
), id="Apple",
Edge(source_node_id="Marie", target_node_id="MacOS", relationship_name="works_on"), name="Apple",
], type="Company",
), description="Apple is a company",
SummarizedContent(summary="Summary of John's work.", description="Summary of John's work."), ),
SummarizedContent( Node(
summary="Summary of Marie's work.", description="Summary of Marie's work." id="Food for Hungry",
), name="Food for Hungry",
] type="Non-profit organization",
description="Food for Hungry is a non-profit organization",
),
],
edges=[
Edge(
source_node_id="John", target_node_id="Apple", relationship_name="works_for"
),
Edge(
source_node_id="John",
target_node_id="Food for Hungry",
relationship_name="works_for",
),
],
)
mock_create_structured_output.side_effect = mock_llm_output
add_john_result = await cognee.add( add_john_result = await cognee.add(
"John works for Apple. He is also affiliated with a non-profit organization called 'Food for Hungry'" "John works for Apple. He is also affiliated with a non-profit organization called 'Food for Hungry'"