46 lines
No EOL
1.3 KiB
Python
46 lines
No EOL
1.3 KiB
Python
#!/usr/bin/env python
|
|
"""
|
|
Script to test if Cognee tools are working correctly.
|
|
Run this script to test if the tools are correctly importing and functioning.
|
|
"""
|
|
|
|
import os
|
|
import cognee
|
|
from src.latest_ai_development.tools import CogneeAdd, CogneeSearch
|
|
|
|
# Set COGNEE_API_KEY if not already set
|
|
if "LLM_API_KEY" not in os.environ:
|
|
openai_api_key = os.environ.get("OPENAI_API_KEY")
|
|
if openai_api_key:
|
|
os.environ["LLM_API_KEY"] = openai_api_key
|
|
|
|
|
|
def test_tools():
|
|
"""Test the CogneeAdd and CogneeSearch tools."""
|
|
print("Testing Cognee tools...")
|
|
|
|
print("\nTesting CogneeAdd tool...")
|
|
add_tool = CogneeAdd()
|
|
test_input = (
|
|
"This is a test text to add to Cognee memory. It contains information about AI LLMs."
|
|
)
|
|
node_set = ["AI", "LLMs"]
|
|
try:
|
|
result = add_tool._run(context=test_input, node_set=node_set)
|
|
print(f"CogneeAdd result: {result}")
|
|
except Exception as e:
|
|
print(f"Error testing CogneeAdd: {str(e)}")
|
|
|
|
print("\nTesting CogneeSearch tool...")
|
|
search_tool = CogneeSearch()
|
|
search_query = "AI LLMs"
|
|
node_set = ["AI"]
|
|
try:
|
|
result = search_tool._run(query_text=search_query, node_set=node_set)
|
|
print(f"CogneeSearch result: {result}")
|
|
except Exception as e:
|
|
print(f"Error testing CogneeSearch: {str(e)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
test_tools() |