{
 "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/advanced/11_Advanced_Context_Engineering.ipynb)\n",
    "\n",
    "# Advanced Context Engineering\n",
    "\n",
    "## Overview\n",
    "\n",
    "This notebook covers advanced topics in context engineering using Semantica. We will explore custom memory management strategies, tuning hybrid retrieval, and extending the system with custom graph builders.\n",
    "\n",
    "### Learning Objectives\n",
    "\n",
    "- **Custom Memory Pruning**: Implement importance-based pruning instead of FIFO.\n",
    "- **Hybrid Retrieval Tuning**: Optimize weights for vector, graph, and keyword search.\n",
    "- **Custom Extensions**: Register custom graph building methods.\n",
    "- **Performance Optimization**: Balance token limits and retrieval latency.\n",
    "\n",
    "---\n",
    "\n",
    "## 1. Setup\n",
    "\n",
    "We'll start by setting up a mock vector store and importing necessary components."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List, Dict, Any, Optional\n",
    "from semantica.context import AgentMemory, AgentContext, ContextGraph, ContextRetriever, VectorStore\n",
    "from semantica.context import registry\n",
    "\n",
    "# Mock Vector Store (same as in introduction)\n",
    "class MockVectorStore(VectorStore):\n",
    "    def __init__(self):\n",
    "        self.items = {}\n",
    "        self.counter = 0\n",
    "    def add(self, texts, metadata=None, **kwargs):\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",
    "    def search(self, query, limit=5, **kwargs):\n",
    "        return [{\n",
    "            \"id\": k, \"content\": v[\"text\"], \"score\": 0.85, \"metadata\": v[\"metadata\"]\n",
    "        } for k, v in list(self.items.items())[:limit]]\n",
    "    def delete(self, ids, **kwargs):\n",
    "        return True\n",
    "\n",
    "vs = MockVectorStore()\n",
    "kg = ContextGraph()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Custom Memory Pruning Strategy\n",
    "\n",
    "By default, `AgentMemory` uses a FIFO (First-In-First-Out) strategy combined with a token limit to prune short-term memory. However, you might want to keep \"important\" memories longer regardless of their age.\n",
    "\n",
    "Let's subclass `AgentMemory` to implement an importance-based pruning strategy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ImportanceAwareMemory(AgentMemory):\n",
    "    def _prune_short_term_memory(self):\n",
    "        \"\"\"\n",
    "        Custom pruning: Always keep items marked as 'important' in metadata,\n",
    "        then prune others based on token limits.\n",
    "        \"\"\"\n",
    "        if not self.short_term_memory:\n",
    "            return\n",
    "\n",
    "        # Separate important items\n",
    "        important_items = [item for item in self.short_term_memory if item.metadata.get(\"important\")]\n",
    "        other_items = [item for item in self.short_term_memory if not item.metadata.get(\"important\")]\n",
    "        \n",
    "        # Calculate tokens used by important items\n",
    "        important_tokens = sum(self._count_tokens(item.content) for item in important_items)\n",
    "        \n",
    "        # Calculate remaining budget\n",
    "        remaining_tokens = max(0, self.token_limit - important_tokens)\n",
    "        \n",
    "        # Prune other items to fit remaining budget\n",
    "        kept_others = []\n",
    "        current_tokens = 0\n",
    "        \n",
    "        # Iterate in reverse (newest first) to keep recent items\n",
    "        for item in reversed(other_items):\n",
    "            item_tokens = self._count_tokens(item.content)\n",
    "            if current_tokens + item_tokens <= remaining_tokens:\n",
    "                kept_others.insert(0, item)\n",
    "                current_tokens += item_tokens\n",
    "            else:\n",
    "                break # Stop once we hit the limit\n",
    "                \n",
    "        # Reconstruct memory: Important items + kept recent items\n",
    "        # Sort by timestamp to maintain order\n",
    "        all_kept = sorted(important_items + kept_others, key=lambda x: x.timestamp)\n",
    "        self.short_term_memory = all_kept\n",
    "\n",
    "# Test the custom memory\n",
    "memory = ImportanceAwareMemory(vector_store=vs, token_limit=100)\n",
    "\n",
    "# Add an old important memory\n",
    "memory.store(\"IMPORTANT: User's name is Alice\", metadata={\"important\": True})\n",
    "\n",
    "# Fill with filler memories\n",
    "for i in range(20):\n",
    "    memory.store(f\"Filler memory {i} \" * 5) # Consumes tokens\n",
    "\n",
    "print(f\"Short-term items: {len(memory.short_term_memory)}\")\n",
    "print(\"First item (should be the important one):\", memory.short_term_memory[0].content)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Tuning Hybrid Retrieval\n",
    "\n",
    "Hybrid retrieval combines scores from vector search and graph traversal. You can tune the `hybrid_alpha` parameter to weight these components.\n",
    "\n",
    "- `hybrid_alpha = 0.0`: Pure Vector Search\n",
    "- `hybrid_alpha = 1.0`: Pure Graph Search\n",
    "- `hybrid_alpha = 0.5`: Balanced (Default)\n",
    "\n",
    "Additionally, `max_expansion_hops` controls how far we traverse the graph from retrieved nodes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Populate graph with some structure\n",
    "kg.add_node(\"python\", \"concept\", \"Python\")\n",
    "kg.add_node(\"ml\", \"concept\", \"Machine Learning\")\n",
    "kg.add_edge(\"python\", \"ml\", \"used_for\")\n",
    "\n",
    "retriever = ContextRetriever(\n",
    "    memory_store=memory,\n",
    "    knowledge_graph=kg,\n",
    "    vector_store=vs,\n",
    "    hybrid_alpha=0.7,      # Favor graph connections\n",
    "    max_expansion_hops=2   # Traverse deeper\n",
    ")\n",
    "\n",
    "results = retriever.retrieve(\"Python\")\n",
    "for res in results:\n",
    "    print(f\"Source: {res.source}, Score: {res.score:.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Extending with Custom Methods\n",
    "\n",
    "Semantica's registry system allows you to plug in custom logic. Let's register a custom graph builder that creates a star graph topology."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def star_graph_builder(center_entity, satellites, **kwargs):\n",
    "    \"\"\"\n",
    "    Builds a star graph where all satellites connect to the center.\n",
    "    \"\"\"\n",
    "    nodes = []\n",
    "    edges = []\n",
    "    \n",
    "    # Center node\n",
    "    nodes.append({\"id\": \"center\", \"label\": center_entity, \"type\": \"CENTER\"})\n",
    "    \n",
    "    for i, sat in enumerate(satellites):\n",
    "        sat_id = f\"sat_{i}\"\n",
    "        nodes.append({\"id\": sat_id, \"label\": sat, \"type\": \"SATELLITE\"})\n",
    "        edges.append({\"source\": \"center\", \"target\": sat_id, \"relation\": \"connects_to\"})\n",
    "        \n",
    "    return {\"nodes\": nodes, \"edges\": edges}\n",
    "\n",
    "# Register the method\n",
    "registry.method_registry.register(\"graph\", \"star_builder\", star_graph_builder)\n",
    "\n",
    "# Verify registration\n",
    "print(\"Available graph methods:\", registry.method_registry.list_all(\"graph\"))\n",
    "\n",
    "# Use it (conceptual - typically used via build_context_graph wrapper)\n",
    "graph_data = star_graph_builder(\"Central Hub\", [\"Spoke 1\", \"Spoke 2\"])\n",
    "print(f\"Created graph with {len(graph_data['nodes'])} nodes and {len(graph_data['edges'])} edges.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Best Practices for Production\n",
    "\n",
    "1. **Token Limits**: Align `token_limit` with your LLM's context window minus the prompt template size.\n",
    "2. **Vector Store**: Use a production-grade vector store (e.g., Pinecone, Weaviate, Qdrant) instead of the mock store.\n",
    "3. **Asynchronous Operations**: For high-throughput systems, consider wrapping storage operations in async tasks (though the core logic is synchronous for simplicity).\n",
    "4. **Entity Resolution**: Implement a robust `EntityLinker` strategy to prevent graph fragmentation (e.g., \"Alice\" vs \"Alice S.\")."
   ]
  }
 ],
 "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
}
