{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/Hawksight-AI/semantica/blob/main/cookbook/introduction/19_Context_Module.ipynb)\n",
    "\n",
    "# Context Engineering Module\n",
    "\n",
    "## Overview\n",
    "\n",
    "This notebook provides a comprehensive guide to Semantica's **Context Engineering Module** - a powerful system for building context graphs, managing agent memory, retrieving context, and linking entities. You'll learn how to use the new synchronous Architecture 2.0 features, including hierarchical memory with token management.\n",
    "\n",
    "**Documentation**: [API Reference](https://semantica.readthedocs.io/reference/context/)\n",
    "\n",
    "### Learning Objectives\n",
    "\n",
    "- **Hierarchical Memory**: Manage short-term (token-buffered) and long-term (vector-stored) memory\n",
    "- **Context Graph**: Build and query dynamic knowledge graphs\n",
    "- **Hybrid Retrieval**: Combine vector search, graph traversal, and keyword matching\n",
    "- **Entity Linking**: Resolve entities across conversations\n",
    "- **Configuration**: Customize behavior via YAML or environment variables\n",
    "\n",
    "---\n",
    "\n",
    "## Installation\n",
    "\n",
    "```bash\n",
    "pip install semantica\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setup: Create a mock vector store for demonstration\n",
    "from typing import List, Dict, Any, Optional\n",
    "from semantica.context import VectorStore\n",
    "\n",
    "class MockVectorStore(VectorStore):\n",
    "    def __init__(self):\n",
    "        self.items = {}\n",
    "        self.counter = 0\n",
    "        \n",
    "    def add(self, texts: List[str], metadata: Optional[List[Dict[str, Any]]] = None, **kwargs) -> List[str]:\n",
    "        ids = []\n",
    "        for i, text in enumerate(texts):\n",
    "            id_ = f\"id_{self.counter}\"\n",
    "            self.items[id_] = {\"text\": text, \"metadata\": metadata[i] if metadata else {}}\n",
    "            ids.append(id_)\n",
    "            self.counter += 1\n",
    "        return ids\n",
    "        \n",
    "    def search(self, query: str, limit: int = 5, **kwargs) -> List[Dict[str, Any]]:\n",
    "        # Simple keyword match for mock\n",
    "        results = []\n",
    "        for id_, item in self.items.items():\n",
    "            if any(w.lower() in item[\"text\"].lower() for w in query.split()):\n",
    "                results.append({\n",
    "                    \"id\": id_,\n",
    "                    \"content\": item[\"text\"],\n",
    "                    \"score\": 0.9,\n",
    "                    \"metadata\": item[\"metadata\"]\n",
    "                })\n",
    "        return results[:limit]\n",
    "        \n",
    "    def delete(self, ids: List[str], **kwargs) -> bool:\n",
    "        for id_ in ids:\n",
    "            self.items.pop(id_, None)\n",
    "        return True\n",
    "\n",
    "vs = MockVectorStore()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. High-Level Interface: AgentContext\n",
    "\n",
    "The `AgentContext` class is the easiest way to get started. It unifies vector storage, knowledge graphs, and memory management."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context import AgentContext, ContextGraph\n",
    "\n",
    "# Initialize with vector store and a new in-memory knowledge graph\n",
    "kg = ContextGraph()\n",
    "context = AgentContext(\n",
    "    vector_store=vs,\n",
    "    knowledge_graph=kg,\n",
    "    token_limit=2000,      # Max tokens in short-term memory\n",
    "    short_term_limit=10    # Max items in short-term memory\n",
    ")\n",
    "\n",
    "# Store a memory (automatically goes to short-term and long-term)\n",
    "context.store(\n",
    "    \"The user, Alice, is a data scientist interested in Python.\",\n",
    "    conversation_id=\"conv_1\",\n",
    "    user_id=\"alice_01\"\n",
    ")\n",
    "\n",
    "# Retrieve context (automatically uses hybrid retrieval)\n",
    "results = context.retrieve(\"What does Alice do?\")\n",
    "\n",
    "for res in results:\n",
    "    print(f\"Found: {res['content']} (Score: {res['score']})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Hierarchical Memory Management\n",
    "\n",
    "Semantica uses a two-tier memory system:\n",
    "1. **Short-Term Memory**: A fast, in-memory buffer limited by tokens (to fit in LLM context windows) and item count.\n",
    "2. **Long-Term Memory**: Persistent storage backed by the vector store.\n",
    "\n",
    "Let's observe how the token limit works."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context import AgentMemory\n",
    "\n",
    "# Initialize memory with strict limits for demonstration\n",
    "memory = AgentMemory(\n",
    "    vector_store=vs,\n",
    "    token_limit=50,       # Very small token limit\n",
    "    short_term_limit=5    # Max 5 items\n",
    ")\n",
    "\n",
    "# Add memories\n",
    "for i in range(10):\n",
    "    memory.store(f\"Memory item {i}: This is a sentence with some tokens.\")\n",
    "    print(f\"Added item {i}. Short-term size: {len(memory.short_term_memory)}\")\n",
    "\n",
    "print(\"\\nFinal short-term memory content:\")\n",
    "for item in memory.short_term_memory:\n",
    "    print(f\"- {item.content}\")\n",
    "    \n",
    "# Notice that older items are pruned to respect the token limit and item count."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Context Graph & GraphRAG\n",
    "\n",
    "The `ContextGraph` allows you to structure information as nodes and edges, enabling \"GraphRAG\" - retrieving information based on relationships rather than just semantic similarity."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context import ContextGraph\n",
    "\n",
    "graph = ContextGraph()\n",
    "\n",
    "# Manually building a graph\n",
    "graph.add_node(\"n1\", \"person\", \"Alice\")\n",
    "graph.add_node(\"n2\", \"language\", \"Python\")\n",
    "graph.add_node(\"n3\", \"library\", \"Semantica\")\n",
    "\n",
    "graph.add_edge(\"n1\", \"n2\", \"uses\")\n",
    "graph.add_edge(\"n2\", \"n3\", \"powers\")\n",
    "\n",
    "# Query the graph\n",
    "neighbors = graph.get_neighbors(\"n2\", hops=1)\n",
    "print(\"Neighbors of Python:\", neighbors)\n",
    "\n",
    "# Using the graph in AgentContext\n",
    "context = AgentContext(vector_store=vs, knowledge_graph=graph)\n",
    "\n",
    "# Retrieve with graph expansion\n",
    "results = context.retrieve(\n",
    "    \"Alice\",\n",
    "    use_graph=True,\n",
    "    expand_graph=True  # Will pull in 'Python' because Alice uses it\n",
    ")\n",
    "\n",
    "print(\"\\nGraph-enhanced Retrieval:\")\n",
    "for res in results:\n",
    "    print(f\"- {res['content']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Entity Linking\n",
    "\n",
    "The `EntityLinker` helps ensure that \"Alice\", \"Alice Smith\", and \"she\" (in context) refer to the same entity ID."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context import EntityLinker\n",
    "\n",
    "linker = EntityLinker()\n",
    "\n",
    "# Generate a canonical URI\n",
    "uri = linker.generate_uri(\"Python Programming Language\")\n",
    "print(f\"Canonical URI: {uri}\")\n",
    "\n",
    "# Check similarity\n",
    "score = linker._calculate_text_similarity(\"Python\", \"Python Lang\")\n",
    "print(f\"Similarity Score: {score}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Configuration\n",
    "\n",
    "You can configure the context module using the `config` object or environment variables."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context import config\n",
    "\n",
    "# Set global configuration\n",
    "config.context_config.set(\"token_limit\", 4096)\n",
    "config.context_config.set(\"retention_days\", 30)\n",
    "\n",
    "print(f\"Current Token Limit: {config.context_config.get('token_limit')}\")"
   ]
  }
 ],
 "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.8.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
