cognee-integration-google-adk 0.1.0

Integrate Cognee's memory system seamlessly with your Google ADK agent using easy-to-use tools.

Cognee-Integration-Google-ADK

A powerful integration between Cognee and Google ADK that provides intelligent knowledge management and retrieval capabilities for AI agents.

Overview

cognee-integration-google-adk combines Cognee's advanced knowledge storage and retrieval system with Google's Agent Development Kit (ADK). This integration allows you to build AI agents that can efficiently store, search, and retrieve information from a persistent knowledge base.

Features

• Smart Knowledge Storage: Add and persist information using Cognee's advanced indexing
• Semantic Search: Retrieve relevant information using natural language queries
• Session Management: Support for user-specific data isolation
• Google ADK Integration: Seamless integration with Google's Agent Development Kit
• Async Support: Built with async/await for high-performance applications
• Long-Running Tools: Optimized for Google ADK's long-running tool capabilities
• Thread-Safe: Queue-based processing for concurrent operations

Installation

pip install cognee-integration-google-adk

Quick Start

import asyncio
from dotenv import load_dotenv
import cognee
from google.adk.agents import Agent
from google.adk.runners import InMemoryRunner
from cognee_integration_google_adk import add_tool, search_tool

load_dotenv()

async def main():
    # Initialize Cognee (optional - for data management)
    await cognee.prune.prune_data()
    await cognee.prune.prune_system(metadata=True)
    
    # Create an agent with memory capabilities
    agent = Agent(
        model="gemini-2.0-flash",
        name="research_analyst",
        description="You are an expert research analyst with access to a comprehensive knowledge base.",
        instruction="You are an expert research analyst with access to a comprehensive knowledge base.",
        tools=[add_tool, search_tool],
    )
    
    runner = InMemoryRunner(agent=agent)
    
    # Use the agent to store information
    events = await runner.run_debug(
        "Remember that our company signed a contract with HealthBridge Systems "
        "in the healthcare industry, starting Feb 2023, ending Jan 2026, worth £2.4M"
    )
    
    # Print agent response
    for event in events:
        if event.is_final_response() and event.content:
            for part in event.content.parts:
                if part.text:
                    print(part.text)
    
    # Query the stored information
    events = await runner.run_debug(
        "What contracts do we have in the healthcare industry?"
    )
    
    for event in events:
        if event.is_final_response() and event.content:
            for part in event.content.parts:
                if part.text:
                    print(part.text)

if __name__ == "__main__":
    asyncio.run(main())

Available Tools

Basic Tools

from cognee_integration_google_adk import add_tool, search_tool

# add_tool: Store information in the knowledge base
# search_tool: Search and retrieve previously stored information

Sessionized Tools

For multi-user applications, use sessionized tools to isolate data between users:

from cognee_integration_google_adk import get_sessionized_cognee_tools

# Get tools for a specific user session
add_tool, search_tool = get_sessionized_cognee_tools("user-123")

# Auto-generate a session ID
add_tool, search_tool = get_sessionized_cognee_tools()

Session Management

cognee-integration-google-adk supports user-specific sessions to isolate data between different users or contexts:

import asyncio
from google.adk.agents import Agent
from google.adk.runners import InMemoryRunner
from cognee_integration_google_adk import get_sessionized_cognee_tools

async def main():
    # Each user gets their own isolated session
    user1_add, user1_search = get_sessionized_cognee_tools("user-123")
    user2_add, user2_search = get_sessionized_cognee_tools("user-456")
    
    # Create separate agents for each user
    agent1 = Agent(
        model="gemini-2.0-flash",
        name="assistant_1",
        description="Assistant for user 1",
        instruction="You are a helpful assistant.",
        tools=[user1_add, user1_search]
    )
    
    agent2 = Agent(
        model="gemini-2.0-flash",
        name="assistant_2",
        description="Assistant for user 2",
        instruction="You are a helpful assistant.",
        tools=[user2_add, user2_search]
    )
    
    runner1 = InMemoryRunner(agent=agent1)
    runner2 = InMemoryRunner(agent=agent2)
    
    # Each agent works with isolated data
    await runner1.run_debug("Remember: I like pizza")
    await runner2.run_debug("Remember: I like sushi")

if __name__ == "__main__":
    asyncio.run(main())

Tool Reference

add_tool(data: str, node_set: Optional[List[str]] = None)

Store information in the knowledge base for later retrieval.

Parameters:

• data (str): The text or information you want to store
• node_set (Optional[List[str]]): Additional node set identifiers for organization

Returns: Confirmation message

Example:

agent = Agent(
    model="gemini-2.0-flash",
    name="data_manager",
    description="Data management specialist",
    instruction="You manage our knowledge base.",
    tools=[add_tool]
)

runner = InMemoryRunner(agent=agent)
await runner.run_debug(
    "Store this: Our Q4 revenue was $2.5M with 15% growth"
)

search_tool(query_text: str, node_set: Optional[List[str]] = None)

Search and retrieve previously stored information from the knowledge base.

Parameters:

• query_text (str): Natural language search query
• node_set (Optional[List[str]]): Additional node set identifiers for scoping the search

Returns: List of relevant search results

Example:

agent = Agent(
    model="gemini-2.0-flash",
    name="research_assistant",
    description="Research specialist",
    instruction="You help users find information quickly.",
    tools=[search_tool]
)

runner = InMemoryRunner(agent=agent)
await runner.run_debug("What was our Q4 revenue?")

get_sessionized_cognee_tools(session_id: Optional[str] = None)

Returns cognee tools with optional user-specific sessionization.

Parameters:

• session_id (Optional[str]): User identifier for data isolation. If not provided, a random session ID is auto-generated.

Returns: (add_tool, search_tool) - A tuple of sessionized tools

Example:

# With explicit session ID
add_tool, search_tool = get_sessionized_cognee_tools("user-123")

# Auto-generate session ID
add_tool, search_tool = get_sessionized_cognee_tools()

Configuration

Environment Variables

Create a .env file in your project root:

# OpenAI API key (used by Cognee for LLM operations)
LLM_API_KEY=your-openai-api-key-here

# Google API key (used by Google ADK for Gemini models)
GOOGLE_API_KEY=your-google-api-key-here

Cognee Configuration (Optional)

You can customize Cognee's data and system directories:

from cognee.api.v1.config import config
import os

config.data_root_directory(
    os.path.join(os.path.dirname(__file__), ".cognee/data_storage")
)

config.system_root_directory(
    os.path.join(os.path.dirname(__file__), ".cognee/system")
)

Examples

Check out the examples/ directory for comprehensive usage examples:

• examples/tools_example.py: Basic usage with add and search tools
• examples/sessionized_tools_example.py: Multi-user session management with visualization

Advanced Usage

Pre-loading Data

You can pre-load data into Cognee before creating agents:

import asyncio
import cognee
from cognee_integration_google_adk import search_tool
from google.adk.agents import Agent
from google.adk.runners import InMemoryRunner

async def main():
    # Pre-load data
    await cognee.add("Important company information here...")
    await cognee.add("More data to remember...")
    await cognee.cognify()  # Process and index the data
    
    # Now create an agent that can search this data
    agent = Agent(
        model="gemini-2.0-flash",
        name="analyst",
        description="Analyst with access to company knowledge base",
        instruction="You have access to our company knowledge base.",
        tools=[search_tool]
    )
    
    runner = InMemoryRunner(agent=agent)
    events = await runner.run_debug("What information do we have?")
    
    for event in events:
        if event.is_final_response() and event.content:
            for part in event.content.parts:
                if part.text:
                    print(part.text)

if __name__ == "__main__":
    asyncio.run(main())

Data Management

import asyncio
import cognee

async def reset_knowledge_base():
    """Clear all data and reset the knowledge base"""
    await cognee.prune.prune_data()
    await cognee.prune.prune_system(metadata=True)

async def visualize_knowledge_graph():
    """Generate a visualization of the knowledge graph"""
    await cognee.visualize_graph("graph.html")

Working with Multiple Agents

import asyncio
from google.adk.agents import Agent
from google.adk.runners import InMemoryRunner
from cognee_integration_google_adk import add_tool, search_tool

async def main():
    # Create a data entry agent
    data_agent = Agent(
        model="gemini-2.0-flash",
        name="data_collector",
        description="Collects and stores information",
        instruction="You collect and store important information.",
        tools=[add_tool]
    )
    
    # Create a research agent
    research_agent = Agent(
        model="gemini-2.0-flash",
        name="researcher",
        description="Searches and analyzes stored information",
        instruction="You search and analyze information from the knowledge base.",
        tools=[search_tool]
    )
    
    data_runner = InMemoryRunner(agent=data_agent)
    research_runner = InMemoryRunner(agent=research_agent)
    
    # Store data
    await data_runner.run_debug(
        "Store this: Project Alpha launched in Q1 2024 with $5M budget"
    )
    
    # Search data
    events = await research_runner.run_debug(
        "When did Project Alpha launch and what was the budget?"
    )
    
    for event in events:
        if event.is_final_response() and event.content:
            for part in event.content.parts:
                if part.text:
                    print(part.text)

if __name__ == "__main__":
    asyncio.run(main())

Requirements

• Python 3.10+
• Google API key (for Gemini models via Google ADK)
• OpenAI API key (or other LLM provider supported by Cognee)
• Dependencies automatically managed via pyproject.toml