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

82
cognee/tests/test_delete_default_graph.py
View file

@ -41,34 +41,22 @@ 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(
),
],
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",
),
],
),
KnowledgeGraph(
nodes=[ nodes=[
Node(id="Marie", name="Marie", type="Person", description="Marie is a person"), Node(id="Marie", name="Marie", type="Person", description="Marie is a person"),
Node( Node(
@ -90,14 +78,42 @@ async def main(mock_create_structured_output: AsyncMock):
target_node_id="Apple", target_node_id="Apple",
relationship_name="works_for", relationship_name="works_for",
), ),
Edge(source_node_id="Marie", target_node_id="MacOS", relationship_name="works_on"), Edge(
source_node_id="Marie", target_node_id="MacOS", relationship_name="works_on"
),
], ],
)
if "John" in text_input and response_model == KnowledgeGraph:
return KnowledgeGraph(
nodes=[
Node(id="John", name="John", type="Person", description="John is a person"),
Node(
id="Apple",
name="Apple",
type="Company",
description="Apple is a company",
), ),
SummarizedContent(summary="Summary of John's work.", description="Summary of John's work."), Node(
SummarizedContent( id="Food for Hungry",
summary="Summary of Marie's work.", description="Summary of Marie's work." 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'"