Feat/add tool - cognee

.python-version

@@ -0,0 +1 @@
+3.12.0

agentstack/_tools/cognee/__init__.py

@@ -0,0 +1,128 @@
+"""
+Implementation of the cognee for AgentStack.
+These functions wrap cognee's asynchronous methods and expose them
+as synchronous functions with typed parameters & docstrings for use by AI agents.
+"""
+
+import cognee
+import asyncio
+from typing import List
+from cognee.api.v1.search import SearchType
+
+
+def prune_data(metadata: bool = False) -> str:
+    """
+    Prune the cognee data. If metadata is True, also prune system metadata.
+
+    :param metadata: Whether to prune system metadata as well.
+    :return: A confirmation message.
+    """
+
+    async def _prune():
+        await cognee.prune.prune_data()
+        if metadata:
+            await cognee.prune.prune_system(metadata=True)
+        return "Data pruned successfully."
+
+    return asyncio.run(_prune())
+
+
+def add_text(text: str) -> str:
+    """
+    Add text to cognee's knowledge system for future 'cognify' operations.
+
+    :param text: The text to add.
+    :return: A confirmation message.
+    """
+
+    async def _add():
+        await cognee.add(text)
+        return "Text added successfully."
+
+    return asyncio.run(_add())
+
+
+def cognify() -> str:
+    """
+    Run cognee's 'cognify' pipeline to build the knowledge graph,
+    summaries, and other metadata from previously added text.
+
+    :return: A confirmation message.
+    """
+
+    async def _cognify():
+        await cognee.cognify()
+        return "Cognify process complete."
+
+    return asyncio.run(_cognify())
+
+
+def search_insights(query_text: str) -> str:
+    """
+    Perform an INSIGHTS search on the knowledge graph for the given query text.
+
+    :param query_text: The query to search for.
+    :return: The search results as a (stringified) list of matches.
+    """
+
+    async def _search():
+        results = await cognee.search(SearchType.INSIGHTS, query_text=query_text)
+        return str(results)
+
+    return asyncio.run(_search())
+
+def search_summaries(query_text: str) -> str:
+    """
+    Perform a SUMMARIES search on the knowledge graph for the given query text.
+
+    :param query_text: The query to search for.
+    :return: The search results as a (stringified) list of matches.
+    """
+
+    async def _search():
+        results = await cognee.search(SearchType.SUMMARIES, query_text=query_text)
+        return str(results)
+
+    return asyncio.run(_search())
+
+def search_chunks(query_text: str) -> str:
+    """
+    Perform a CHUNKS search on the knowledge graph for the given query text.
+
+    :param query_text: The query to search for.
+    :return: The search results as a (stringified) list of matches.
+    """
+
+    async def _search():
+        results = await cognee.search(SearchType.CHUNKS, query_text=query_text)
+        return str(results)
+
+    return asyncio.run(_search())
+
+def search_completion(query_text: str) -> str:
+    """
+    Perform a COMPLETION search on the knowledge graph for the given query text.
+
+    :param query_text: The query to search for.
+    :return: The search results as a (stringified) list of matches.
+    """
+
+    async def _search():
+        results = await cognee.search(SearchType.COMPLETION, query_text=query_text)
+        return str(results)
+
+    return asyncio.run(_search())
+
+def search_graph_completion(query_text: str) -> str:
+    """
+    Perform a GRAPH_COMPLETION search on the knowledge graph for the given query text.
+
+    :param query_text: The query to search for.
+    :return: The search results as a (stringified) list of matches.
+    """
+
+    async def _search():
+        results = await cognee.search(SearchType.GRAPH_COMPLETION, query_text=query_text)
+        return str(results)
+
+    return asyncio.run(_search())

agentstack/_tools/cognee/config.json

@@ -0,0 +1,64 @@
+{
+    "name": "cognee",
+    "category": "Memory", 
+    "tools": [
+        "prune_data",
+        "add_text",
+        "cognify",
+        "search_insights",
+        "search_summaries",
+        "search_chunks",
+        "search_completion",
+        "search_graph_completion"
+    ],
+    "url": "https://github.com/topoteretes/cognee",
+    "tools_bundled": true,
+    "cta": "Cognee installed! Please set your cognee env variables.",
+
+    "env": {
+        "ENV": "local",
+        "TOKENIZERS_PARALLELISM": "false",
+        "LLM_API_KEY": "",
+        "LLM_MODEL": "openai/gpt-4o-mini",
+        "LLM_PROVIDER": "openai",
+        "LLM_ENDPOINT": "",
+        "LLM_API_VERSION": "",
+        "GRAPHISTRY_USERNAME": "",
+        "GRAPHISTRY_PASSWORD": "",
+        "SENTRY_REPORTING_URL": "",
+        "EMBEDDING_PROVIDER": "openai",
+        "EMBEDDING_API_KEY": "",
+        "EMBEDDING_MODEL": "openai/text-embedding-3-large",
+        "EMBEDDING_ENDPOINT": "",
+        "EMBEDDING_API_VERSION": "",
+        "EMBEDDING_DIMENSIONS": 3072,
+        "EMBEDDING_MAX_TOKENS": 8191,
+        "_comment1": "# neo4j or networkx", 
+        "GRAPH_DATABASE_PROVIDER": "networkx",
+        "_comment2": "# Not needed if using networkx",
+        "GRAPH_DATABASE_URL": "",
+        "GRAPH_DATABASE_USERNAME": "",
+        "GRAPH_DATABASE_PASSWORD": "",
+        "_comment3": "# qdrant, pgvector, weaviate, milvus or lancedb",
+        "VECTOR_DB_PROVIDER": "lancedb",
+        "_comment4": "# Not needed if using lancedb or pgvector",
+        "VECTOR_DB_URL": "",
+        "VECTOR_DB_KEY": "",
+        "_comment5": "Relational Database provider sqlite or postgres",
+        "DB_PROVIDER": "sqlite",
+        "_comment6": "# Database name",
+        "DB_NAME": "cognee_db",
+        "_comment7": "# Postgres specific parameters (Only if Postgres or PGVector is used)", 
+        "DB_HOST": "127.0.0.1",
+        "DB_PORT": 5432,
+        "DB_USERNAME": "cognee",
+        "DB_PASSWORD": "cognee"
+      },
+
+    "packages": [
+        "cognee"
+    ],
+    "post_install": "Now, you can start cognifying!",
+    "post_remove": "Cognee is removed!"
+  }
+