{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Lesson 13 - Agent Memory with Cognee Knowledge Graphs" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Setup\n", "\n", "This notebook demonstrates how to build an intelligent **coding assistant** with persistent memory using [**Cognee**](https://www.cognee.ai/) knowledge graphs and the **Microsoft Agent Framework** (MAF).\n", "\n", "Cognee transforms unstructured text into a structured, queryable knowledge graph backed by vector embeddings — giving your agent a rich, relationship-aware long-term memory.\n", "\n", "### What You'll Learn\n", "1. **Build Knowledge Graphs**: Transform developer profiles and best practices into structured, queryable knowledge.\n", "2. **Integrate Cognee with MAF**: Use `@tool` functions to let an MAF agent query Cognee's knowledge graph.\n", "3. **Session-Aware Conversations**: Maintain context across multiple questions in the same session.\n", "4. **Long-Term Memory**: Persist important knowledge across sessions and retrieve it in new conversations.\n", "\n", "### Prerequisites\n", "- Python 3.9+\n", "- Redis running locally (`docker run -d -p 6379:6379 redis`) for session management\n", "- An LLM API key (e.g. OpenAI) — set `LLM_API_KEY` in `.env`\n", "- `CACHING=true` in `.env` (required for Cognee sessions)\n", "- A Microsoft Foundry project with a deployed chat model\n", "- `AZURE_AI_PROJECT_ENDPOINT` and `AZURE_AI_MODEL_DEPLOYMENT_NAME` in `.env`\n", "- Azure CLI authenticated (`az login`)\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "%pip install agent-framework azure-ai-projects azure-identity \"cognee[redis]==0.4.0\" -q" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "from pathlib import Path\n", "from typing import Annotated\n", "\n", "from dotenv import load_dotenv\n", "\n", "load_dotenv()\n", "\n", "os.environ[\"LLM_API_KEY\"] = os.getenv(\"LLM_API_KEY\", \"\")\n", "os.environ[\"CACHING\"] = os.getenv(\"CACHING\", \"true\")\n", "\n", "import cognee\n", "from cognee.modules.search.types import SearchType\n", "\n", "from agent_framework import tool\n", "from agent_framework.foundry import FoundryChatClient\n", "from azure.identity import AzureCliCredential\n", "\n", "print(f\"Cognee version: {cognee.__version__}\")\n", "print(f\"CACHING: {os.environ.get('CACHING')}\")\n" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "provider = FoundryChatClient(\n", " project_endpoint=os.environ[\"AZURE_AI_PROJECT_ENDPOINT\"],\n", " model=os.environ[\"AZURE_AI_MODEL_DEPLOYMENT_NAME\"],\n", " credential=AzureCliCredential(),\n", ")\n", "\n", "print(\"✅ FoundryChatClient created\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Types of Agent Memory\n", "\n", "This notebook explores the same three memory types from the main Lesson 13 notebook, but uses Cognee as the long-term memory backend:\n", "\n", "| Memory Type | Mechanism | Lifetime |\n", "|---|---|---|\n", "| **Working** | `agent.create_session()` (MAF) | Single conversation |\n", "| **Short-term** | Cognee session cache (Redis) | Single session |\n", "| **Long-term** | Cognee knowledge graph + vectors | Permanent |\n", "\n", "### Cognee's Memory Architecture\n", "```\n", "┌──────────────────────────┐\n", "│ Raw Data │ (developer profiles, docs, conversations)\n", "└───────────┬──────────────┘\n", " │ cognee.add() + cognee.cognify()\n", " ▼\n", "┌──────────────────────────────────────────┐\n", "│ Knowledge Graph + Vector Embeddings │\n", "└───────────┬──────────────────────────────┘\n", " │ cognee.search()\n", " ▼\n", "┌──────────────────┐ ┌────────────────┐\n", "│ MAF Agent │──────▶│ @tool funcs │\n", "│ (AgentSession) │ │ wrapping │\n", "│ │ │ cognee.search │\n", "└──────────────────┘ └────────────────┘\n", "```" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Prepare Cognee Storage" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "DATA_ROOT = Path('.data_storage').resolve()\n", "SYSTEM_ROOT = Path('.cognee_system').resolve()\n", "\n", "DATA_ROOT.mkdir(parents=True, exist_ok=True)\n", "SYSTEM_ROOT.mkdir(parents=True, exist_ok=True)\n", "\n", "cognee.config.data_root_directory(str(DATA_ROOT))\n", "cognee.config.system_root_directory(str(SYSTEM_ROOT))\n", "\n", "await cognee.prune.prune_data()\n", "await cognee.prune.prune_system(metadata=True)\n", "print(\"✅ Cognee storage configured and reset\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Part 1 — Building the Knowledge Base\n", "\n", "We ingest three types of data to create a comprehensive knowledge base for our coding assistant:\n", "\n", "1. **Developer Profile** — personal expertise and technical background\n", "2. **Python Best Practices** — the Zen of Python with practical guidelines\n", "3. **Historical Conversations** — past Q&A sessions between developers and AI assistants" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "developer_intro = (\n", " \"Hi, I'm an AI/Backend engineer. \"\n", " \"I build FastAPI services with Pydantic, heavy asyncio/aiohttp pipelines, \"\n", " \"and production testing via pytest-asyncio. \"\n", " \"I've shipped low-latency APIs on AWS, Azure, and GoogleCloud.\"\n", ")\n", "\n", "python_zen_principles = \"\"\"\n", "# The Zen of Python: Practical Guide\n", "\n", "## Key Principles With Guidance\n", "\n", "### 1. Beautiful is better than ugly\n", "Prefer descriptive names, clear structure, and consistent formatting.\n", "\n", "### 2. Explicit is better than implicit\n", "Be clear about behavior, imports, and types.\n", "\n", "### 3. Simple is better than complex\n", "Choose straightforward solutions first.\n", "\n", "### 4. Flat is better than nested\n", "Use early returns to reduce indentation.\n", "\n", "## Modern Python Tie-ins\n", "- Type hints reinforce explicitness\n", "- Context managers enforce safe resource handling\n", "- Dataclasses improve readability for data containers\n", "\"\"\"\n", "\n", "human_agent_conversations = \"\"\"\n", "\"conversations\": [\n", " {\n", " \"topic\": \"async/await patterns\",\n", " \"user_query\": \"I'm building a web scraper that needs to handle thousands of URLs concurrently. What's the best way to structure this with asyncio?\",\n", " \"assistant_response\": \"Use asyncio with aiohttp, a semaphore to cap concurrency, TCPConnector for connection pooling, and context managers for session lifecycle.\"\n", " },\n", " {\n", " \"topic\": \"dataclass vs pydantic\",\n", " \"user_query\": \"When should I use dataclasses vs Pydantic models?\",\n", " \"assistant_response\": \"For API input/output, prefer Pydantic: runtime validation, type coercion, JSON serialization. Integrates cleanly with FastAPI.\"\n", " },\n", " {\n", " \"topic\": \"testing patterns\",\n", " \"user_query\": \"What's the best approach for pytest with async functions?\",\n", " \"assistant_response\": \"Use pytest-asyncio, async fixtures, and an isolated test database or mocks to reliably test async code.\"\n", " },\n", " {\n", " \"topic\": \"error handling and logging\",\n", " \"user_query\": \"What's the best approach for production-ready error management?\",\n", " \"assistant_response\": \"Centralized error handling with custom exceptions, structured logging, and FastAPI middleware.\"\n", " }\n", " ]\n", "\"\"\"\n", "\n", "print(\"✅ Data sources prepared\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "await cognee.add(developer_intro, node_set=[\"developer_data\"])\n", "await cognee.add(human_agent_conversations, node_set=[\"developer_data\"])\n", "await cognee.add(python_zen_principles, node_set=[\"principles_data\"])\n", "\n", "await cognee.cognify()\n", "print(\"✅ Knowledge graph built\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Visualize the Knowledge Graph\n", "\n", "Cognee can render an interactive HTML visualization of the entities and relationships it extracted." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "from cognee import visualize_graph\n", "\n", "await visualize_graph('./cognee_graph.html')\n", "print(\"📊 Graph saved to cognee_graph.html — open it in a browser to explore.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Enrich Memory with Memify\n", "\n", "`memify()` analyzes the knowledge graph and generates intelligent rules — identifying patterns, best practices, and relationships between concepts." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "await cognee.memify()\n", "print(\"✅ Memory enriched with memify\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Part 2 — MAF Agent with Cognee Tools\n", "\n", "Now we create an MAF agent that can query Cognee's knowledge graph via `@tool` functions. This lets the agent leverage the full power of graph-aware semantic search while maintaining conversational context through sessions." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "@tool(approval_mode=\"never_require\")\n", "async def search_knowledge(\n", " query: Annotated[str, \"Natural-language question to search the knowledge graph\"],\n", ") -> str:\n", " \"\"\"Search the Cognee knowledge graph for relevant developer knowledge, best practices, and past conversations.\"\"\"\n", " results = await cognee.search(\n", " query_text=query,\n", " query_type=SearchType.GRAPH_COMPLETION,\n", " )\n", " if not results:\n", " return \"No relevant knowledge found.\"\n", " return str(results)\n", "\n", "\n", "@tool(approval_mode=\"never_require\")\n", "async def search_principles(\n", " query: Annotated[str, \"Question about Python principles or best practices\"],\n", ") -> str:\n", " \"\"\"Search only the Python principles subset of the knowledge graph.\"\"\"\n", " from cognee.modules.engine.models.node_set import NodeSet\n", " results = await cognee.search(\n", " query_text=query,\n", " query_type=SearchType.GRAPH_COMPLETION,\n", " node_type=NodeSet,\n", " node_name=[\"principles_data\"],\n", " )\n", " if not results:\n", " return \"No relevant principles found.\"\n", " return str(results)\n", "\n", "\n", "print(\"✅ Cognee tools defined: search_knowledge, search_principles\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "coding_agent = provider.as_agent(\n", " name=\"CodingAssistant\",\n", " instructions=(\n", " \"You are an expert coding assistant with access to a knowledge graph \"\n", " \"containing developer profiles, Python best practices, and past conversations.\\n\\n\"\n", " \"WORKFLOW:\\n\"\n", " \"1. Use search_knowledge() to find relevant information from the full knowledge graph.\\n\"\n", " \"2. Use search_principles() when the question is specifically about Python best practices.\\n\"\n", " \"3. Combine retrieved knowledge with your own expertise to give comprehensive answers.\\n\"\n", " \"4. Reference the developer's known tech stack (FastAPI, asyncio, Pydantic) when relevant.\"\n", " ),\n", ")\n", "\n", "print(\"✅ CodingAssistant agent created\")\n" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Working Memory with Sessions\n", "\n", "The `AgentSession` (created via `agent.create_session()`) provides working memory within a session. The agent can refer back to earlier messages while also querying Cognee's long-term knowledge graph." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "session = coding_agent.create_session()\n", "\n", "response = await coding_agent.run(\n", " \"How does my AsyncWebScraper implementation align with Python's design principles?\",\n", " session=session,\n", ")\n", "print(\"🤖 Agent:\", response)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = await coding_agent.run(\n", " \"Based on what you just said, when should I pick dataclasses versus Pydantic for this work?\",\n", " session=session,\n", ")\n", "print(\"🤖 Agent:\", response)\n", "print(\"\\n💡 The agent combined working memory (previous answer) with Cognee's knowledge graph.\")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## New Session — Long-Term Memory Persists\n", "\n", "Starting a fresh session clears working memory, but the Cognee knowledge graph is still available. The agent can retrieve the same long-term knowledge in a completely new conversation." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "session_2 = coding_agent.create_session()\n", "\n", "response = await coding_agent.run(\n", " \"What logging guidance should I follow for incident reviews?\",\n", " session=session_2,\n", ")\n", "print(\"🤖 Agent:\", response)\n", "print(\"\\n💡 New session, but the agent still has access to the full Cognee knowledge graph.\")" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "response = await coding_agent.run(\n", " \"How should variables be named according to Python best practices?\",\n", " session=session_2,\n", ")\n", "print(\"🤖 Agent:\", response)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Summary\n", "\n", "In this notebook you built a coding assistant that combines **MAF's working memory** (`agent.create_session()`) with **Cognee's long-term knowledge graph**.\n", "\n", "### What You've Learned\n", "1. **Knowledge graph construction**: Cognee ingests unstructured text and builds a graph + vector memory.\n", "2. **Graph enrichment with memify**: Derived facts and richer relationships on top of your existing graph.\n", "3. **MAF + Cognee integration**: `@tool` functions let an MAF agent query Cognee's graph naturally.\n", "4. **Working memory + long-term memory**: `AgentSession` (via `agent.create_session()`) provides session context while Cognee provides persistent knowledge.\n", "5. **Filtered search with NodeSets**: Target specific subsets of the knowledge graph (e.g. only principles).\n", "\n", "### Key Takeaways\n", "- **Cognee** turns raw text into structured, relationship-aware memory — more powerful than a flat vector store.\n", "- **`@tool` functions** bridge MAF agents and external knowledge systems cleanly.\n", "- **`AgentSession`** (via `agent.create_session()`) keeps per-conversation context separate from long-lived knowledge.\n", "- The same knowledge graph serves multiple sessions and agents.\n", "\n", "### Real-World Applications\n", "- **Developer copilots**: Code review, incident analysis, architecture assistants\n", "- **Customer-facing copilots**: Support agents over product docs, FAQs, and CRM notes\n", "- **Internal expert copilots**: Policy, legal, or security assistants reasoning over guidelines\n", "- **Unified data layers**: Combine structured and unstructured data into one queryable graph\n", "\n", "### Next Steps\n", "- Experiment with temporal awareness in Cognee\n", "- Define an OWL ontology for domain-specific graph quality\n", "- Add user feedback loops to improve retrieval over time\n", "- Scale to multi-agent systems sharing the same Cognee memory layer" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.12.0" } }, "nbformat": 4, "nbformat_minor": 5 }