1+ """
2+ Implementation of the cognee for AgentStack.
3+ These functions wrap cognee's asynchronous methods and expose them
4+ as synchronous functions with typed parameters & docstrings for use by AI agents.
5+ """
6+
7+ import cognee
8+ import asyncio
9+ from typing import List
10+ from cognee .api .v1 .search import SearchType
11+
12+
13+ def prune_data (metadata : bool = False ) -> str :
14+ """
15+ Prune the cognee data. If metadata is True, also prune system metadata.
16+
17+ :param metadata: Whether to prune system metadata as well.
18+ :return: A confirmation message.
19+ """
20+
21+ async def _prune ():
22+ await cognee .prune .prune_data ()
23+ if metadata :
24+ await cognee .prune .prune_system (metadata = True )
25+ return "Data pruned successfully."
26+
27+ return asyncio .run (_prune ())
28+
29+
30+ def add_text (text : str ) -> str :
31+ """
32+ Add text to cognee's knowledge system for future 'cognify' operations.
33+
34+ :param text: The text to add.
35+ :return: A confirmation message.
36+ """
37+
38+ async def _add ():
39+ await cognee .add (text )
40+ return "Text added successfully."
41+
42+ return asyncio .run (_add ())
43+
44+
45+ def cognify () -> str :
46+ """
47+ Run cognee's 'cognify' pipeline to build the knowledge graph,
48+ summaries, and other metadata from previously added text.
49+
50+ :return: A confirmation message.
51+ """
52+
53+ async def _cognify ():
54+ await cognee .cognify ()
55+ return "Cognify process complete."
56+
57+ return asyncio .run (_cognify ())
58+
59+
60+ def search_insights (query_text : str ) -> str :
61+ """
62+ Perform an INSIGHTS search on the knowledge graph for the given query text.
63+
64+ :param query_text: The query to search for.
65+ :return: The search results as a (stringified) list of matches.
66+ """
67+
68+ async def _search ():
69+ results = await cognee .search (SearchType .INSIGHTS , query_text = query_text )
70+ return str (results )
71+
72+ return asyncio .run (_search ())
73+
74+ def search_summaries (query_text : str ) -> str :
75+ """
76+ Perform a SUMMARIES search on the knowledge graph for the given query text.
77+
78+ :param query_text: The query to search for.
79+ :return: The search results as a (stringified) list of matches.
80+ """
81+
82+ async def _search ():
83+ results = await cognee .search (SearchType .SUMMARIES , query_text = query_text )
84+ return str (results )
85+
86+ return asyncio .run (_search ())
87+
88+ def search_chunks (query_text : str ) -> str :
89+ """
90+ Perform a CHUNKS search on the knowledge graph for the given query text.
91+
92+ :param query_text: The query to search for.
93+ :return: The search results as a (stringified) list of matches.
94+ """
95+
96+ async def _search ():
97+ results = await cognee .search (SearchType .CHUNKS , query_text = query_text )
98+ return str (results )
99+
100+ return asyncio .run (_search ())
101+
102+ def search_completion (query_text : str ) -> str :
103+ """
104+ Perform a COMPLETION search on the knowledge graph for the given query text.
105+
106+ :param query_text: The query to search for.
107+ :return: The search results as a (stringified) list of matches.
108+ """
109+
110+ async def _search ():
111+ results = await cognee .search (SearchType .COMPLETION , query_text = query_text )
112+ return str (results )
113+
114+ return asyncio .run (_search ())
115+
116+ def search_graph_completion (query_text : str ) -> str :
117+ """
118+ Perform a GRAPH_COMPLETION search on the knowledge graph for the given query text.
119+
120+ :param query_text: The query to search for.
121+ :return: The search results as a (stringified) list of matches.
122+ """
123+
124+ async def _search ():
125+ results = await cognee .search (SearchType .GRAPH_COMPLETION , query_text = query_text )
126+ return str (results )
127+
128+ return asyncio .run (_search ())
0 commit comments