{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c21e9c8d",
   "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 Module — Practical Guide\n",
    "\n",
    "Semantica’s `context` module is the layer that makes an agent “stateful”. It combines:\n",
    "\n",
    "- **Memory** (short-term + long-term) via `AgentMemory`\n",
    "- **Graph context** via `ContextGraph`\n",
    "- **Hybrid retrieval** (vector + memory + graph) via `ContextRetriever`\n",
    "- **High-level UX** via `AgentContext` (recommended entry point)\n",
    "- **Entity linking** via `EntityLinker`\n",
    "- **Extensibility + config** via `registry` and `config`\n",
    "\n",
    "This notebook focuses on small, runnable examples and keeps imports scoped to each cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "257bad40",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -q semantica"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a48e0f10",
   "metadata": {},
   "source": [
    "## 1) Vector store (for long-term memory)\n",
    "\n",
    "The `VectorStore` can generate embeddings via its internal embedder. If no embedder is available in your environment, it falls back to random vectors so the API stays usable for demos."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c845a94",
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.vector_store import VectorStore\n",
    "\n",
    "vs = VectorStore(backend=\"inmemory\", dimension=384)\n",
    "\n",
    "if getattr(vs, \"embedder\", None) and hasattr(vs.embedder, \"set_text_model\"):\n",
    "    vs.embedder.set_text_model(method=\"fastembed\", model_name=\"BAAI/bge-small-en-v1.5\")\n",
    "\n",
    "vs.backend, vs.dimension"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1b1ba34",
   "metadata": {},
   "source": [
    "## 2) Quick start with `AgentContext` (recommended)\n",
    "\n",
    "`AgentContext` is the user-friendly interface that ties memory, vector store, and graph together. If you pass a `ContextGraph`, the system can do GraphRAG-style retrieval."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4d788b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context import AgentContext, ContextGraph\n",
    "\n",
    "kg = ContextGraph()\n",
    "context = AgentContext(vector_store=vs, knowledge_graph=kg)\n",
    "\n",
    "context.config"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "638bdbc8",
   "metadata": {},
   "source": [
    "## 3) Store and retrieve memory\n",
    "\n",
    "A single string is treated as a memory item. You can attach `conversation_id` and `user_id` through metadata-friendly parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d65eb00",
   "metadata": {},
   "outputs": [],
   "source": [
    "memory_id = context.store(\n",
    "    \"User prefers short answers about Python.\",\n",
    "    conversation_id=\"conv_1\",\n",
    "    user_id=\"user_1\",\n",
    "    metadata={\"type\": \"preference\"},\n",
    ")\n",
    "\n",
    "context.get_memory(memory_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c1be718",
   "metadata": {},
   "outputs": [],
   "source": [
    "context.store(\n",
    "    \"User is working on Semantica context module examples.\",\n",
    "    conversation_id=\"conv_1\",\n",
    "    user_id=\"user_1\",\n",
    "    metadata={\"type\": \"note\"},\n",
    ")\n",
    "\n",
    "context.retrieve(\"Python answers\", max_results=3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "485acf33",
   "metadata": {},
   "outputs": [],
   "source": [
    "context.conversation(\"conv_1\", max_items=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e43cddd",
   "metadata": {},
   "source": [
    "## 4) Export, save, load\n",
    "\n",
    "`AgentContext` includes simple persistence helpers. This example uses a temporary directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a264ef4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "export_json = context.export(conversation_id=\"conv_1\", format=\"json\")\n",
    "export_json[:300]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b62d1859",
   "metadata": {},
   "outputs": [],
   "source": [
    "import tempfile\n",
    "\n",
    "with tempfile.TemporaryDirectory() as d:\n",
    "    context.save(d)\n",
    "    context.load(d)\n",
    "\n",
    "context.conversation_summary(\"conv_1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b8bf553",
   "metadata": {},
   "source": [
    "## 5) Store documents and build a context graph\n",
    "\n",
    "If you store a list, `AgentContext.store(...)` treats it as documents. To keep this notebook lightweight and deterministic, we pass pre-extracted entities and relationships per document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72930ae7",
   "metadata": {},
   "outputs": [],
   "source": [
    "documents = [\n",
    "    {\n",
    "        \"id\": \"doc_1\",\n",
    "        \"content\": \"Python is used for machine learning.\",\n",
    "        \"metadata\": {\"source\": \"docs\"},\n",
    "        \"entities\": [\n",
    "            {\"id\": \"e_python\", \"text\": \"Python\", \"type\": \"PROGRAMMING_LANGUAGE\"},\n",
    "            {\"id\": \"e_ml\", \"text\": \"Machine Learning\", \"type\": \"CONCEPT\"},\n",
    "        ],\n",
    "        \"relationships\": [\n",
    "            {\n",
    "                \"source_id\": \"e_python\",\n",
    "                \"target_id\": \"e_ml\",\n",
    "                \"type\": \"used_for\",\n",
    "                \"confidence\": 0.9,\n",
    "            }\n",
    "        ],\n",
    "    },\n",
    "    {\n",
    "        \"id\": \"doc_2\",\n",
    "        \"content\": \"PyTorch is a machine learning framework.\",\n",
    "        \"metadata\": {\"source\": \"docs\"},\n",
    "        \"entities\": [\n",
    "            {\"id\": \"e_pytorch\", \"text\": \"PyTorch\", \"type\": \"FRAMEWORK\"},\n",
    "            {\"id\": \"e_ml\", \"text\": \"Machine Learning\", \"type\": \"CONCEPT\"},\n",
    "        ],\n",
    "        \"relationships\": [\n",
    "            {\n",
    "                \"source_id\": \"e_pytorch\",\n",
    "                \"target_id\": \"e_ml\",\n",
    "                \"type\": \"implements\",\n",
    "                \"confidence\": 0.95,\n",
    "            }\n",
    "        ],\n",
    "    },\n",
    "]\n",
    "\n",
    "stats = context.store(\n",
    "    documents,\n",
    "    extract_entities=False,\n",
    "    extract_relationships=False,\n",
    "    link_entities=True,\n",
    ")\n",
    "\n",
    "stats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24f8dba8",
   "metadata": {},
   "outputs": [],
   "source": [
    "kg.stats()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4671f2a",
   "metadata": {},
   "source": [
    "## 6) Explore the graph with `ContextGraph`\n",
    "\n",
    "The graph supports keyword querying and neighbor expansion."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df2e5fcd",
   "metadata": {},
   "outputs": [],
   "source": [
    "kg.query(\"machine learning\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0836feeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "kg.get_neighbors(\"e_python\", hops=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1be1adf1",
   "metadata": {},
   "source": [
    "## 7) Entity linking with `EntityLinker`\n",
    "\n",
    "`EntityLinker` assigns stable URIs and can link related or duplicate entities across sources."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24b011c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context import EntityLinker\n",
    "\n",
    "linker = EntityLinker(knowledge_graph={\"entities\": [{\"id\": \"e_py\", \"text\": \"Python\", \"type\": \"PROGRAMMING_LANGUAGE\"}]})\n",
    "\n",
    "entities = [\n",
    "    {\"id\": \"e1\", \"text\": \"Python\", \"type\": \"PROGRAMMING_LANGUAGE\"},\n",
    "    {\"id\": \"e2\", \"text\": \"PyTorch\", \"type\": \"FRAMEWORK\"},\n",
    "]\n",
    "\n",
    "linked = linker.link(\"Python and PyTorch\", entities=entities)\n",
    "[(e.entity_id, e.uri, len(e.linked_entities)) for e in linked]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a282de3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "linker.link_entities(\"e1\", \"e2\", link_type=\"related_to\", confidence=0.8)\n",
    "linker.get_entity_links(\"e1\")[:2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d11f87f",
   "metadata": {},
   "outputs": [],
   "source": [
    "linker.build_entity_web()[\"statistics\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "072efafd",
   "metadata": {},
   "source": [
    "## 8) Low-level building blocks: `AgentMemory` and `ContextRetriever`\n",
    "\n",
    "If you want more control than `AgentContext`, you can wire the parts directly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3791d6c6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context import AgentMemory, ContextRetriever\n",
    "\n",
    "memory = AgentMemory(vector_store=vs, knowledge_graph=kg, retention_policy=\"unlimited\")\n",
    "memory.store(\"Python powers Semantica.\", metadata={\"type\": \"fact\", \"conversation_id\": \"conv_2\"})\n",
    "\n",
    "retriever = ContextRetriever(memory_store=memory, knowledge_graph=kg, vector_store=vs)\n",
    "results = retriever.retrieve(\"Python Semantica\", max_results=5)\n",
    "\n",
    "[(r.content, r.source, round(r.score, 3)) for r in results]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92060402",
   "metadata": {},
   "source": [
    "## 9) Methods, registry, and configuration\n",
    "\n",
    "The `methods` layer exposes convenience functions, while `registry` lets you plug in your own implementations. `config` provides runtime configuration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "896e7001",
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context.config import context_config\n",
    "\n",
    "context_config.set(\"retention_policy\", \"7_days\")\n",
    "context_config.get(\"retention_policy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e925a2e0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context.methods import build_context_graph\n",
    "from semantica.context.registry import method_registry\n",
    "\n",
    "def custom_graph_method(entities, relationships, conversations=None, **kwargs):\n",
    "    return {\n",
    "        \"nodes\": [],\n",
    "        \"edges\": [],\n",
    "        \"statistics\": {\"node_count\": 0, \"edge_count\": 0},\n",
    "    }\n",
    "\n",
    "method_registry.register(\"graph\", \"custom_demo\", custom_graph_method)\n",
    "method_registry.list_all(\"graph\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21fa6cb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "build_context_graph(\n",
    "    entities=[{\"id\": \"e1\", \"text\": \"Python\", \"type\": \"PROGRAMMING_LANGUAGE\"}],\n",
    "    relationships=[{\"source_id\": \"e1\", \"target_id\": \"e2\", \"type\": \"related_to\"}],\n",
    "    method=\"custom_demo\",\n",
    ")"
   ]
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
