73 lines
2.1 KiB
Python
73 lines
2.1 KiB
Python
import cognee
|
|
import asyncio
|
|
|
|
|
|
from cognee.shared.logging_utils import get_logger, ERROR
|
|
from cognee.modules.metrics.operations import get_pipeline_run_metrics
|
|
from cognee.modules.engine.models.Entity import Entity
|
|
from cognee.api.v1.search import SearchType
|
|
|
|
job_1 = """
|
|
Natural language processing (NLP) is an interdisciplinary
|
|
subfield of computer science and information retrieval.
|
|
"""
|
|
|
|
|
|
async def main(enable_steps):
|
|
# Step 1: Reset data and system state
|
|
if enable_steps.get("prune_data"):
|
|
await cognee.prune.prune_data()
|
|
print("Data pruned.")
|
|
|
|
if enable_steps.get("prune_system"):
|
|
await cognee.prune.prune_system(metadata=True)
|
|
print("System pruned.")
|
|
|
|
# Step 2: Add text
|
|
if enable_steps.get("add_text"):
|
|
text_list = [job_1]
|
|
for text in text_list:
|
|
await cognee.add(text)
|
|
print(f"Added text: {text[:35]}...")
|
|
|
|
# Step 3: Create knowledge graph
|
|
if enable_steps.get("cognify"):
|
|
pipeline_run = await cognee.cognify()
|
|
print("Knowledge graph created.")
|
|
|
|
# Step 4: Calculate descriptive metrics
|
|
if enable_steps.get("graph_metrics"):
|
|
await get_pipeline_run_metrics(pipeline_run, include_optional=True)
|
|
print("Descriptive graph metrics saved to database.")
|
|
|
|
# Step 5: Query insights
|
|
if enable_steps.get("retriever"):
|
|
search_results = await cognee.search(
|
|
query_type=SearchType.GRAPH_COMPLETION,
|
|
query_text="What is computer science?",
|
|
node_type=Entity,
|
|
node_name=["computer science"],
|
|
)
|
|
print(search_results)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logger = get_logger(level=ERROR)
|
|
|
|
rebuild_kg = True
|
|
retrieve = True
|
|
steps_to_enable = {
|
|
"prune_data": rebuild_kg,
|
|
"prune_system": rebuild_kg,
|
|
"add_text": rebuild_kg,
|
|
"cognify": rebuild_kg,
|
|
"graph_metrics": rebuild_kg,
|
|
"retriever": retrieve,
|
|
}
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
try:
|
|
loop.run_until_complete(main(steps_to_enable))
|
|
finally:
|
|
loop.run_until_complete(loop.shutdown_asyncgens())
|