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

added tests

Browse source
This commit is contained in:
vasilije 2025-07-13 17:15:14 +02:00
parent 7055bd8a8f
commit 324a196552
1 changed files with 19 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
cognee/tests/test_falkordb.py
View file

@ -91,17 +91,17 @@ async def main():
from cognee.infrastructure.databases.vector import get_vector_engine from cognee.infrastructure.databases.vector import get_vector_engine
vector_engine = get_vector_engine() vector_engine = get_vector_engine()
# Try different search terms that are likely to be found as entities in the processed text # Try different search terms that are likely to be found as entities in the processed text
search_terms = ["GPT", "OpenAI", "language model", "LLM", "AI", "neural network"] search_terms = ["GPT", "OpenAI", "language model", "LLM", "AI", "neural network"]
search_results = [] search_results = []
for term in search_terms: for term in search_terms:
search_results = await vector_engine.search("Entity_name", term) search_results = await vector_engine.search("Entity_name", term)
if search_results: if search_results:
print(f"✅ Found {len(search_results)} results for '{term}'") print(f"✅ Found {len(search_results)} results for '{term}'")
break break
# If no entities found with common search terms, fallback to any Entity # If no entities found with common search terms, fallback to any Entity
if not search_results: if not search_results:
try: try:
@ -113,22 +113,26 @@ async def main():
entity_result = vector_engine.query(entity_query) entity_result = vector_engine.query(entity_query)
if entity_result.result_set: if entity_result.result_set:
# Convert to ScoredResult format # Convert to ScoredResult format
from cognee.infrastructure.databases.vector.models.ScoredResult import ScoredResult from cognee.infrastructure.databases.vector.models.ScoredResult import (
ScoredResult,
)
from cognee.shared.utils import parse_id from cognee.shared.utils import parse_id
node = entity_result.result_set[0][0] node = entity_result.result_set[0][0]
payload = dict(node.properties) if hasattr(node, 'properties') else {} payload = dict(node.properties) if hasattr(node, "properties") else {}
if 'text' not in payload and 'name' in payload: if "text" not in payload and "name" in payload:
payload['text'] = payload['name'] payload["text"] = payload["name"]
search_results = [ScoredResult( search_results = [
id=parse_id(payload.get('id', str(hash(str(node))))), ScoredResult(
score=1.0, id=parse_id(payload.get("id", str(hash(str(node))))),
payload=payload score=1.0,
)] payload=payload,
)
]
except Exception as e: except Exception as e:
print(f"❌ Error in fallback search: {e}") print(f"❌ Error in fallback search: {e}")
assert len(search_results) > 0, "No entities found in the vector database" assert len(search_results) > 0, "No entities found in the vector database"
random_node = search_results[0] random_node = search_results[0]
random_node_name = random_node.payload["text"] random_node_name = random_node.payload["text"]