{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Knowledge Graph Engine v2 - Notebook Tutorial\n",
    "\n",
    "This notebook demonstrates how to use KG Engine v2 with Ollama (local LLM) and Neo4j in a development environment.\n",
    "\n",
    "## Prerequisites\n",
    "Make sure you have started the services:\n",
    "```bash\n",
    "docker-compose -f docker-compose.notebook.yml up -d\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup and Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Python path: /Users/dasein/.local/share/uv/python/cpython-3.12.9-macos-aarch64-none/lib/python312.zip\n",
      "Neo4j URI: neo4j+s://346d12e8.databases.neo4j.io\n",
      "LLM Provider: litellm\n",
      "Ollama Model: phi3:3.8b\n",
      "Ollama Base URL: http://localhost:11434/v1\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "import logging\n",
    "# Add src to Python path\n",
    "src_path = Path('../src').resolve()\n",
    "if str(src_path) not in sys.path:\n",
    "    sys.path.insert(0, str(src_path))\n",
    "\n",
    "# Load environment variables\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv('../.env.notebook')\n",
    "\n",
    "logging.getLogger().setLevel(logging.WARNING)\n",
    "# Disable tokenizer parallelism warning\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"] = \"false\"\n",
    "\n",
    "print(f\"Python path: {sys.path[0]}\")\n",
    "print(f\"Neo4j URI: {os.getenv('NEO4J_URI')}\")\n",
    "print(f\"LLM Provider: {os.getenv('LLM_PROVIDER', 'auto-detect')}\")\n",
    "print(f\"Ollama Model: {os.getenv('OLLAMA_MODEL', 'not set')}\")\n",
    "print(f\"Ollama Base URL: {os.getenv('OLLAMA_BASE_URL', 'not set')}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ KG Engine v2.2.0 loaded successfully\n"
     ]
    }
   ],
   "source": [
    "# Import KG Engine components\n",
    "from kg_engine import (\n",
    "    KnowledgeGraphEngineV2, \n",
    "    InputItem, \n",
    "    Neo4jConfig,\n",
    "    SearchType,\n",
    "    LLMClientFactory,\n",
    "    __version__\n",
    ")\n",
    "\n",
    "print(f\"✅ KG Engine v{__version__} loaded successfully\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Initialize the Engine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Neo4j connection successful\n"
     ]
    }
   ],
   "source": [
    "# Test Neo4j connection\n",
    "config = Neo4jConfig()\n",
    "if config.verify_connectivity():\n",
    "    print(\"✅ Neo4j connection successful\")\n",
    "else:\n",
    "    print(\"❌ Neo4j connection failed - check docker-compose services\")\n",
    "    raise ConnectionError(\"Neo4j not available\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "✅ Ollama connection successful\n",
      "Available models: phi3:3.8b, deepseek-coder:6.7b-base-q4_0\n",
      "✅ Selected model 'phi3:3.8b' is available\n"
     ]
    }
   ],
   "source": [
    "# Test Ollama connection\n",
    "import requests\n",
    "\n",
    "ollama_base = os.getenv('OLLAMA_BASE_URL', 'http://localhost:11434/v1')\n",
    "ollama_model = os.getenv('OLLAMA_MODEL', 'phi3:mini')\n",
    "\n",
    "try:\n",
    "    # Try to connect to Ollama API\n",
    "    response = requests.get(f\"{ollama_base.replace('/v1', '')}/api/tags\")\n",
    "    if response.status_code == 200:\n",
    "        models = [model['name'] for model in response.json().get('models', [])]\n",
    "        print(\"✅ Ollama connection successful\")\n",
    "        print(f\"Available models: {', '.join(models) if models else 'No models found'}\")\n",
    "        \n",
    "        if ollama_model in models:\n",
    "            print(f\"✅ Selected model '{ollama_model}' is available\")\n",
    "        else:\n",
    "            print(f\"⚠️  Selected model '{ollama_model}' not found. Available: {models}\")\n",
    "            print(\"   Run: ollama pull phi3:mini\")\n",
    "    else:\n",
    "        print(f\"❌ Ollama connection failed - status {response.status_code}\")\n",
    "except Exception as e:\n",
    "    print(f\"❌ Cannot connect to Ollama: {e}\")\n",
    "    print(\"   Make sure Ollama is running: docker-compose -f docker-compose.notebook.yml up -d\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🤖 LLM Interface initialized: gpt-4o via litellm\n",
      "🚀 Knowledge Graph Engine v2\n",
      "   - LLM interface: gpt-4o via litellm\n",
      "🚀 KG Engine initialized with litellm provider!\n",
      "   Model: gpt-4o\n",
      "   Base URL: https://litellm.chatcyber.ai\n"
     ]
    }
   ],
   "source": [
    "# Initialize KG Engine with LLMClientFactory\n",
    "try:\n",
    "    # Create LLM configuration from environment\n",
    "    llm_config = LLMClientFactory.create_from_env()\n",
    "    \n",
    "    # Initialize the engine\n",
    "    engine = KnowledgeGraphEngineV2(\n",
    "        llm_config=llm_config,\n",
    "        neo4j_config=config\n",
    "    )\n",
    "    \n",
    "    print(f\"🚀 KG Engine initialized with {llm_config.provider} provider!\")\n",
    "    print(f\"   Model: {llm_config.get_model_name()}\")\n",
    "    \n",
    "    # For Ollama, show the base URL\n",
    "    if hasattr(llm_config, 'base_url'):\n",
    "        print(f\"   Base URL: {llm_config.base_url}\")\n",
    "        \n",
    "except Exception as e:\n",
    "    print(f\"❌ Failed to initialize engine: {e}\")\n",
    "    print(\"\\nTroubleshooting:\")\n",
    "    print(\"1. Check if Ollama is running: curl http://localhost:11434/api/tags\")\n",
    "    print(\"2. Verify environment variables in .env.notebook\")\n",
    "    print(\"3. Ensure OLLAMA_MODEL is set correctly\")\n",
    "    raise"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### LLM Provider Configuration\n",
    "\n",
    "The LLMClientFactory automatically detects the provider based on environment variables:\n",
    "\n",
    "1. **Ollama** (current): Set `LLM_PROVIDER=ollama` or have `OLLAMA_MODEL` set\n",
    "2. **OpenAI**: Set `LLM_PROVIDER=openai` and `OPENAI_API_KEY`\n",
    "3. **LiteLLM**: Set `LLM_PROVIDER=litellm` and `LITELLM_BEARER_TOKEN`\n",
    "\n",
    "To switch providers, update the `.env.notebook` file and restart the kernel."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Basic Usage - Adding Knowledge"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Created 8 input items\n",
      "1. Alice works as a senior software engineer at Google\n",
      "2. Bob is a data scientist at Microsoft\n",
      "3. Alice lives in Mountain View, California\n",
      "4. Bob lives in Seattle, Washington\n",
      "5. Charlie is Alice's friend from Stanford University\n",
      "6. Google's headquarters is in Mountain View\n",
      "7. Alice graduated from Stanford in 2018\n",
      "8. Bob enjoys hiking and photography\n"
     ]
    }
   ],
   "source": [
    "# Create some sample knowledge\n",
    "sample_data = [\n",
    "    \"Alice works as a senior software engineer at Google\",\n",
    "    \"Bob is a data scientist at Microsoft\", \n",
    "    \"Alice lives in Mountain View, California\",\n",
    "    \"Bob lives in Seattle, Washington\",\n",
    "    \"Charlie is Alice's friend from Stanford University\",\n",
    "    \"Google's headquarters is in Mountain View\",\n",
    "    \"Alice graduated from Stanford in 2018\",\n",
    "    \"Bob enjoys hiking and photography\"\n",
    "]\n",
    "\n",
    "# Convert to InputItem objects\n",
    "input_items = [InputItem(description=text) for text in sample_data]\n",
    "\n",
    "print(f\"Created {len(input_items)} input items\")\n",
    "for i, item in enumerate(input_items, 1):\n",
    "    print(f\"{i}. {item.description}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔄 Processing input items...\n",
      "\n",
      "📊 Processing Results:\n",
      "   Items processed: 8\n",
      "   New edges created: 0\n",
      "   Updated edges: 0\n",
      "   Duplicates ignored: 12\n",
      "   Processing time: 79290.7ms\n"
     ]
    }
   ],
   "source": [
    "# Process the input items\n",
    "print(\"🔄 Processing input items...\")\n",
    "results = engine.process_input(input_items)\n",
    "\n",
    "print(\"\\n📊 Processing Results:\")\n",
    "print(f\"   Items processed: {results['processed_items']}\")\n",
    "print(f\"   New edges created: {results['new_edges']}\")\n",
    "print(f\"   Updated edges: {results['updated_edges']}\")\n",
    "print(f\"   Duplicates ignored: {results['duplicates_ignored']}\")\n",
    "print(f\"   Processing time: {results['processing_time_ms']:.1f}ms\")\n",
    "\n",
    "if results['errors']:\n",
    "    print(f\"\\n⚠️ Errors: {results['errors']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Searching the Knowledge Graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔍 Testing search queries:\n",
      "========================================\n"
     ]
    }
   ],
   "source": [
    "# Define some test queries\n",
    "test_queries = [\n",
    "    \"Who works at Google?\",\n",
    "    \"Where does Alice live?\",\n",
    "    \"What companies are mentioned?\",\n",
    "    \"Who are Alice's friends?\",\n",
    "    \"Where did Alice study?\",\n",
    "    \"What does Bob enjoy doing?\"\n",
    "]\n",
    "\n",
    "print(\"🔍 Testing search queries:\")\n",
    "print(\"=\" * 40)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "1. Query: 'Who works at Google?'\n",
      "Search type: SearchType.BOTH Who works at Google?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['Google'], relationships=['WORKS_AT', 'EMPLOYED_BY'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='list', temporal_context=None)\n",
      "Standardized relationships: ['SPEAKS', 'LEADS']\n",
      "Searching with entities: ['Google'] and relationships: ['SPEAKS', 'LEADS']\n",
      "Parsing query SEMANTIC...\n",
      "   Results found: 10\n",
      "   Answer: Alice currently works at Google. Bob also worked at Google, specifically on search algorithms.\n",
      "   Query time: 4959.9ms\n",
      "   Top results:\n",
      "     1. Alice VISITED Google (score: 0.82)\n",
      "     2. Bob Smith CONTINUED Google (score: 0.80)\n",
      "     3. TechCorp IS_ONE_OF ['Google', 'Microsoft', 'Amazon'] (score: 0.78)\n",
      "\n",
      "2. Query: 'Where does Alice live?'\n",
      "Search type: SearchType.BOTH Where does Alice live?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['Alice'], relationships=['LIVES_IN', 'RESIDES_IN'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='search', temporal_context=None)\n",
      "Standardized relationships: ['LIVES_IN', 'LIVES_IN']\n",
      "Searching with entities: ['Alice'] and relationships: ['LIVES_IN', 'LIVES_IN']\n",
      "Parsing query SEMANTIC...\n",
      "   Results found: 10\n",
      "   Answer: Alice lives in Mountain View, California.\n",
      "   Query time: 4322.1ms\n",
      "   Top results:\n",
      "     1. Alice LIVES_IN Mountain View, California (score: 1.00)\n",
      "     2. Alice VISITED Google (score: 0.81)\n",
      "     3. Charlie IS_ONE_OF Alice (score: 0.81)\n",
      "\n",
      "3. Query: 'What companies are mentioned?'\n",
      "Search type: SearchType.BOTH What companies are mentioned?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=[], relationships=['MENTIONED_IN'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='list', temporal_context=None)\n",
      "Standardized relationships: ['SPEAKS']\n",
      "Searching with entities: [] and relationships: ['SPEAKS']\n",
      "Parsing query SEMANTIC...\n",
      "   Results found: 10\n",
      "   Answer: The companies mentioned are Google, Microsoft, Amazon, Uber, Airbnb, and TechCorp.\n",
      "   Query time: 4184.7ms\n",
      "   Top results:\n",
      "     1. TechCorp IS_ONE_OF ['Google', 'Microsoft', 'Amazon'] (score: 0.73)\n",
      "     2. AI Assistant Platform LEADS businesses (score: 0.72)\n",
      "     3. Bob Smith BECAME TechCorp (score: 0.71)\n",
      "\n",
      "4. Query: 'Who are Alice's friends?'\n",
      "Search type: SearchType.BOTH Who are Alice's friends?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['Alice'], relationships=['HAS_FRIEND'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='list', temporal_context=None)\n",
      "Standardized relationships: ['HAS_ROLE']\n",
      "Searching with entities: ['Alice'] and relationships: ['HAS_ROLE']\n",
      "Parsing query SEMANTIC...\n",
      "   Results found: 10\n",
      "   Answer: Alice's friend is Charlie.\n",
      "   Query time: 4195.9ms\n",
      "   Top results:\n",
      "     1. Charlie IS_ONE_OF Alice (score: 0.90)\n",
      "     2. Bob Smith LEADS Alice to start TechCorp. (score: 0.82)\n",
      "     3. Alice SPEAKS senior software engineer (score: 0.81)\n",
      "\n",
      "5. Query: 'Where did Alice study?'\n",
      "Search type: SearchType.BOTH Where did Alice study?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['Alice'], relationships=['STUDIED_AT', 'ATTENDED'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='search', temporal_context=None)\n",
      "Standardized relationships: ['TEACHES', 'TEACHES']\n",
      "Searching with entities: ['Alice'] and relationships: ['TEACHES', 'TEACHES']\n",
      "Parsing query SEMANTIC...\n",
      "   Results found: 10\n",
      "   Answer: Alice studied at Stanford University, where she graduated with a degree in Computer Science in 2018.\n",
      "   Query time: 5341.7ms\n",
      "   Top results:\n",
      "     1. Alice TEACHES Stanford (score: 1.00)\n",
      "     2. Alice TEACHES Stanford University (score: 1.00)\n",
      "     3. Alice Johnson TEACHES Stanford University with a degree in Computer Science. (score: 0.83)\n",
      "\n",
      "6. Query: 'What does Bob enjoy doing?'\n",
      "Search type: SearchType.BOTH What does Bob enjoy doing?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['Bob'], relationships=['ENJOYS_DOING'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='list', temporal_context=None)\n",
      "Standardized relationships: ['SPEAKS']\n",
      "Searching with entities: ['Bob'] and relationships: ['SPEAKS']\n",
      "Parsing query SEMANTIC...\n",
      "   Results found: 10\n",
      "   Answer: Bob enjoys photography and hiking.\n",
      "   Query time: 3641.0ms\n",
      "   Top results:\n",
      "     1. Bob LIKES photography (score: 0.81)\n",
      "     2. Bob VISITED hiking (score: 0.81)\n",
      "     3. Bob LIVES_IN Seattle, Washington (score: 0.78)\n"
     ]
    }
   ],
   "source": [
    "# Test each query\n",
    "for i, query in enumerate(test_queries, 1):\n",
    "    print(f\"\\n{i}. Query: '{query}'\")\n",
    "    \n",
    "    try:\n",
    "        response = engine.search(query, search_type=SearchType.BOTH)\n",
    "        \n",
    "        print(f\"   Results found: {len(response.results)}\")\n",
    "        print(f\"   Answer: {response.answer}\")\n",
    "        print(f\"   Query time: {response.query_time_ms:.1f}ms\")\n",
    "        \n",
    "        # Show top results\n",
    "        if response.results:\n",
    "            print(\"   Top results:\")\n",
    "            for j, result in enumerate(response.results[:3], 1):\n",
    "                if result.triplet and result.triplet.edge:\n",
    "                    edge = result.triplet.edge\n",
    "                    print(f\"     {j}. {edge.subject} {edge.relationship} {edge.object} (score: {result.score:.2f})\")\n",
    "                \n",
    "    except Exception as e:\n",
    "        print(f\"   ❌ Error: {e}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔗 All relationships for 'Alice':\n",
      "1. Alice SPEAKS senior software engineer\n",
      "   Summary: Alice works as a senior software engineer\n",
      "   Confidence: 0.95\n",
      "\n",
      "2. Alice LIVES_IN Mountain View, California\n",
      "   Summary: Alice lives in Mountain View, California\n",
      "   Confidence: 0.95\n",
      "\n",
      "3. Alice BECAME Meta\n",
      "   Summary: Alice now works at Meta as a senior engineer\n",
      "   Confidence: 0.95\n",
      "\n",
      "4. Alice VISITED Google\n",
      "   Summary: Alice works at Google\n",
      "   Confidence: 0.95\n",
      "\n",
      "5. Alice TEACHES Stanford\n",
      "   Summary: Alice graduated from Stanford in 2018\n",
      "   Confidence: 0.95\n",
      "\n",
      "6. Alice IS_ONE_OF Charlie\n",
      "   Summary: Charlie is Alice's friend\n",
      "   Confidence: 0.9\n",
      "\n",
      "7. Alice TEACHES Stanford University\n",
      "   Summary: Alice attended Stanford University\n",
      "   Confidence: 0.85\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Get all relationships for a specific entity\n",
    "entity_name = \"Alice\"\n",
    "print(f\"🔗 All relationships for '{entity_name}':\")\n",
    "\n",
    "relations = engine.get_node_relations(entity_name, max_depth=1)\n",
    "\n",
    "for i, relation in enumerate(relations, 1):\n",
    "    if relation.triplet and relation.triplet.edge:\n",
    "        edge = relation.triplet.edge\n",
    "        print(f\"{i}. {edge.subject} {edge.relationship} {edge.object}\")\n",
    "        print(f\"   Summary: {edge.metadata.summary}\")\n",
    "        print(f\"   Confidence: {edge.metadata.confidence}\")\n",
    "        print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "📊 System Statistics:\n",
      "Graph Stats: {'total_entities': 207, 'total_edges': 201, 'active_edges': 201, 'obsolete_edges': 0, 'relationship_types': 45, 'relationships': ['RELATES_TO', 'BORN_IN', 'SPEAKS', 'HAS_ROLE', 'SPECIALIZES_IN', 'HAS_CHILDREN', 'LIVES_IN', 'PRACTICES', 'INTERESTED_IN', 'LIKES', 'LEADS', 'RULED', 'JOINED', 'CONTINUED', 'BECAME', 'CAPTURED', 'ORDERED', 'PROHIBITED', 'SUBJECTED_TO', 'INVADED']}\n",
      "Vector Stats: {'total_triplets': 201, 'active_triplets': 201, 'obsolete_triplets': 0, 'embedder_model': 'unknown'}\n",
      "Total Entities: 208\n",
      "Relationship Types: 45\n"
     ]
    }
   ],
   "source": [
    "# Get system statistics\n",
    "print(\"📊 System Statistics:\")\n",
    "stats = engine.get_stats()\n",
    "\n",
    "print(f\"Graph Stats: {stats.get('graph_stats', {})}\")\n",
    "print(f\"Vector Stats: {stats.get('vector_stats', {})}\")\n",
    "print(f\"Total Entities: {stats.get('entities', 0)}\")\n",
    "print(f\"Relationship Types: {len(stats.get('relationships', []))}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conflict Resolution Demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⚔️ Testing conflict resolution:\n",
      "   - Alice no longer works at Google\n",
      "   - Alice now works at Meta as a senior engineer\n",
      "   - Bob moved to Portland, Oregon in 2024\n",
      "\n",
      "🔄 Processing conflicting information...\n",
      "New edges: 0\n",
      "Updated edges: 0\n",
      "Obsoleted edges: 0\n"
     ]
    }
   ],
   "source": [
    "# Add conflicting information to demonstrate conflict resolution\n",
    "conflict_data = [\n",
    "    \"Alice no longer works at Google\",  # Negation\n",
    "    \"Alice now works at Meta as a senior engineer\",  # New job\n",
    "    \"Bob moved to Portland, Oregon in 2024\"  # Location change\n",
    "]\n",
    "\n",
    "conflict_items = [InputItem(description=text) for text in conflict_data]\n",
    "\n",
    "print(\"⚔️ Testing conflict resolution:\")\n",
    "for item in conflict_items:\n",
    "    print(f\"   - {item.description}\")\n",
    "\n",
    "print(\"\\n🔄 Processing conflicting information...\")\n",
    "conflict_results = engine.process_input(conflict_items)\n",
    "\n",
    "print(f\"New edges: {conflict_results['new_edges']}\")\n",
    "print(f\"Updated edges: {conflict_results['updated_edges']}\")\n",
    "print(f\"Obsoleted edges: {conflict_results['obsoleted_edges']}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "🔍 Testing queries after conflict resolution:\n",
      "Search type: SearchType.BOTH Where does Alice work now?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['Alice'], relationships=['WORKS_AT', 'EMPLOYED_BY'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='search', temporal_context=None)\n",
      "Standardized relationships: ['SPEAKS', 'LEADS']\n",
      "Searching with entities: ['Alice'] and relationships: ['SPEAKS', 'LEADS']\n",
      "Parsing query SEMANTIC...\n",
      "\n",
      "Q: Where does Alice work now?\n",
      "A: Alice now works at Meta as a senior engineer.\n",
      "Search type: SearchType.BOTH Where does Bob live?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['Bob'], relationships=['LIVES_IN', 'RESIDES_IN'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='search', temporal_context=None)\n",
      "Standardized relationships: ['LIVES_IN', 'LIVES_IN']\n",
      "Searching with entities: ['Bob'] and relationships: ['LIVES_IN', 'LIVES_IN']\n",
      "Parsing query SEMANTIC...\n",
      "\n",
      "Q: Where does Bob live?\n",
      "A: Bob currently lives in Portland, Oregon.\n",
      "Search type: SearchType.BOTH Who works at Google?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['Google'], relationships=['WORKS_AT', 'EMPLOYED_BY'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='list', temporal_context=None)\n",
      "Standardized relationships: ['SPEAKS', 'LEADS']\n",
      "Searching with entities: ['Google'] and relationships: ['SPEAKS', 'LEADS']\n",
      "Parsing query SEMANTIC...\n",
      "\n",
      "Q: Who works at Google?\n",
      "A: Alice currently works at Google, and Bob previously worked there on search algorithms.\n"
     ]
    }
   ],
   "source": [
    "# Test the updated information\n",
    "test_queries_after_conflict = [\n",
    "    \"Where does Alice work now?\",\n",
    "    \"Where does Bob live?\",\n",
    "    \"Who works at Google?\"\n",
    "]\n",
    "\n",
    "print(\"🔍 Testing queries after conflict resolution:\")\n",
    "for query in test_queries_after_conflict:\n",
    "    response = engine.search(query)\n",
    "    print(f\"\\nQ: {query}\")\n",
    "    print(f\"A: {response.answer}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Performance Testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "⏱️ Performance test: 'Who works at tech companies?' (3 trials)\n",
      "Search type: SearchType.BOTH Who works at tech companies?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['tech companies'], relationships=['WORKS_AT', 'EMPLOYED_BY'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='list', temporal_context=None)\n",
      "Standardized relationships: ['SPEAKS', 'LEADS']\n",
      "Searching with entities: ['tech companies'] and relationships: ['SPEAKS', 'LEADS']\n",
      "Parsing query SEMANTIC...\n",
      "Trial 1: 4432.7ms - 10 results\n",
      "Search type: SearchType.BOTH Who works at tech companies?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['tech companies'], relationships=['WORKS_AT', 'EMPLOYED_BY'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='list', temporal_context=None)\n",
      "Standardized relationships: ['SPEAKS', 'LEADS']\n",
      "Searching with entities: ['tech companies'] and relationships: ['SPEAKS', 'LEADS']\n",
      "Parsing query SEMANTIC...\n",
      "Trial 2: 4509.3ms - 10 results\n",
      "Search type: SearchType.BOTH Who works at tech companies?\n",
      "Parsing query DIRECT with LLM intuition...\n",
      "Parsed query: ParsedQuery(entities=['tech companies'], relationships=['WORKS_AT', 'EMPLOYED_BY'], search_type=<SearchType.DIRECT: 'direct'>, query_intent='list', temporal_context=None)\n",
      "Standardized relationships: ['SPEAKS', 'LEADS']\n",
      "Searching with entities: ['tech companies'] and relationships: ['SPEAKS', 'LEADS']\n",
      "Parsing query SEMANTIC...\n",
      "Trial 3: 5900.4ms - 10 results\n",
      "\n",
      "Average query time: 4947.5ms\n",
      "Provider: litellm\n",
      "Model: gpt-4o\n"
     ]
    }
   ],
   "source": [
    "import time\n",
    "\n",
    "# Test search performance\n",
    "search_query = \"Who works at tech companies?\"\n",
    "num_trials = 3\n",
    "\n",
    "print(f\"⏱️ Performance test: '{search_query}' ({num_trials} trials)\")\n",
    "\n",
    "times = []\n",
    "for i in range(num_trials):\n",
    "    start_time = time.time()\n",
    "    response = engine.search(search_query)\n",
    "    end_time = time.time()\n",
    "    \n",
    "    query_time = (end_time - start_time) * 1000\n",
    "    times.append(query_time)\n",
    "    \n",
    "    print(f\"Trial {i+1}: {query_time:.1f}ms - {len(response.results)} results\")\n",
    "\n",
    "avg_time = sum(times) / len(times)\n",
    "print(f\"\\nAverage query time: {avg_time:.1f}ms\")\n",
    "print(f\"Provider: {llm_config.provider}\")\n",
    "print(f\"Model: {llm_config.get_model_name()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean Up (Optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to clear all data\n",
    "# print(\"🧹 Clearing all data...\")\n",
    "# engine.clear_all_data()\n",
    "# print(\"✅ All data cleared\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This notebook demonstrated:\n",
    "1. ✅ Setting up KG Engine v2 with Ollama and Neo4j\n",
    "2. ✅ Processing natural language input\n",
    "3. ✅ Searching the knowledge graph with natural language queries\n",
    "4. ✅ Exploring entity relationships\n",
    "5. ✅ Conflict resolution with temporal tracking\n",
    "6. ✅ Performance testing\n",
    "\n",
    "The system is now ready for your knowledge management tasks! \n",
    "\n",
    "**Model Performance**: The `phi3:mini` model provides good balance of speed and accuracy for notebook environments. For production use, consider larger models or OpenAI GPT models."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
