{
 "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/use_cases/biomedical/01_Drug_Discovery_Pipeline.ipynb)\n",
    "\n",
    "# Drug Discovery Pipeline - Vector Similarity Search\n",
    "\n",
    "## Overview\n",
    "\n",
    "This notebook demonstrates a **complete drug discovery pipeline** using Semantica's modular architecture. We'll use individual modules directly to build a comprehensive system for drug-target interaction prediction using vector similarity search and knowledge graphs.\n",
    "\n",
    "### Key Features\n",
    "\n",
    "- **Modular Architecture**: Uses Semantica modules directly (`NERExtractor`, `GraphBuilder`, `EmbeddingGenerator`, `VectorStore`)\n",
    "- **Multiple Data Sources**: Ingests from 15+ PubMed RSS feeds, preprint servers, and journal feeds\n",
    "- **Vector Similarity Search**: Emphasizes embeddings and vector similarity for drug-target interaction prediction\n",
    "- **Entity Extraction**: Extracts drug compounds, proteins, targets, enzymes, and receptors\n",
    "- **Knowledge Graph**: Builds structured drug-target relationship graphs\n",
    "- **GraphRAG**: Hybrid vector + graph retrieval for enhanced querying\n",
    "\n",
    "### What You'll Learn\n",
    "\n",
    "- How to use Semantica modules directly (avoiding the core orchestrator)\n",
    "- How to ingest biomedical data from multiple sources\n",
    "- How to extract entities using `NERExtractor`\n",
    "- How to extract relationships using `RelationExtractor`\n",
    "- How to generate embeddings with `EmbeddingGenerator`\n",
    "- How to build knowledge graphs with `GraphBuilder`\n",
    "- How to perform similarity search with `VectorStore`\n",
    "- How to use GraphRAG with `AgentContext` for hybrid retrieval\n",
    "\n",
    "### Pipeline Flow\n",
    "\n",
    "```mermaid\n",
    "graph LR\n",
    "    A[Data Ingestion] --> B[Text Processing]\n",
    "    B --> C[Entity Extraction]\n",
    "    C --> D[Relationship Extraction]\n",
    "    D --> E[Deduplication]\n",
    "    E --> F[Embedding Generation]\n",
    "    F --> G[Vector Store]\n",
    "    G --> H[Knowledge Graph]\n",
    "    H --> I[Similarity Search]\n",
    "    H --> J[GraphRAG Queries]\n",
    "    I --> K[Visualization]\n",
    "    J --> K\n",
    "```\n",
    "\n",
    "### Data Sources\n",
    "\n",
    "**PubMed RSS Feeds:**\n",
    "- Drug Discovery, Drug Target Interaction, Pharmacokinetics, Pharmacodynamics\n",
    "- Clinical Trials, Protein Targets, Drug Repurposing, Molecular Docking\n",
    "- ADME, Drug Metabolism, Drug Safety, Precision Medicine\n",
    "- Biomarkers, Drug Resistance, Combinatorial Therapy\n",
    "\n",
    "**Preprint Servers:**\n",
    "- BioRxiv (Pharmacology & Toxicology, Drug Discovery)\n",
    "- MedRxiv (Clinical Trials)\n",
    "- ChemRxiv\n",
    "\n",
    "**Journal RSS Feeds:**\n",
    "- Nature (Drug Discovery, Pharmacology)\n",
    "- Science Translational Medicine\n",
    "- Cell Chemical Biology\n",
    "- Journal of Medicinal Chemistry\n",
    "- Drug Discovery Today\n",
    "- Trends in Pharmacological Sciences\n",
    "\n",
    "\n",
    "---\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Installation\n",
    "\n",
    "Install Semantica and required dependencies:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -qU semantica networkx matplotlib plotly pandas faiss-cpu beautifulsoup4 groq sentence-transformers scikit-learn\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Configuration & Setup\n",
    "\n",
    "Set up environment variables and configuration constants.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.environ[\"GROQ_API_KEY\"] = os.getenv(\"GROQ_API_KEY\", \"\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "EMBEDDING_DIMENSION = 384\n",
    "EMBEDDING_MODEL = \"all-MiniLM-L6-v2\"\n",
    "CHUNK_SIZE = 1000\n",
    "CHUNK_OVERLAP = 200\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ingesting Biomedical Data from Multiple Sources\n",
    "\n",
    "Ingest data from comprehensive biomedical sources including PubMed RSS feeds, preprint servers, and journal feeds.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.ingest import FeedIngestor, FileIngestor\n",
    "import os\n",
    "from contextlib import redirect_stderr\n",
    "from io import StringIO\n",
    "\n",
    "os.makedirs(\"data\", exist_ok=True)\n",
    "\n",
    "feed_sources = [\n",
    "    # Nature Feeds\n",
    "    (\"Nature - Drug Discovery\", \"https://www.nature.com/subjects/drug-discovery.rss\"),\n",
    "    (\"Nature - Pharmacology\", \"https://www.nature.com/subjects/pharmacology.rss\"),\n",
    "    (\"Nature Reviews Drug Discovery\", \"https://www.nature.com/nrd.rss\"),\n",
    "    \n",
    "    # FDA & Government Sources\n",
    "    (\"FDA MedWatch\", \"https://www.fda.gov/AboutFDA/ContactFDA/StayInformed/RSSFeeds/MedWatch/rss.xml\"),\n",
    "    (\"NCI News\", \"https://www.cancer.gov/syndication/rss\"),\n",
    "    \n",
    "    # Drug Information & News\n",
    "    (\"Drugs.com - MedNews\", \"https://www.drugs.com/rss/mednews.xml\"),\n",
    "    (\"Drugs.com - FDA Alerts\", \"https://www.drugs.com/rss/fda-alerts.xml\"),\n",
    "    (\"Drugs.com - Clinical Trials\", \"https://www.drugs.com/rss/clinical-trials.xml\"),\n",
    "    \n",
    "    # Medical News\n",
    "    (\"Labroots Health & Medicine\", \"http://www.labroots.com/rss/trending/health-and-medicine\"),\n",
    "    (\"Biology News Net\", \"https://www.biologynews.net/rss.php\"),\n",
    "    \n",
    "    # Open Access Journals\n",
    "    (\"PLOS ONE - Medicine\", \"https://journals.plos.org/plosone/feed/atom\"),\n",
    "    (\"PLOS Biology\", \"https://journals.plos.org/plosbiology/feed/atom\"),\n",
    "    (\"PLOS Medicine\", \"https://journals.plos.org/plosmedicine/feed/atom\"),\n",
    "    \n",
    "    # Preprint Servers\n",
    "    (\"arXiv - q-bio\", \"http://arxiv.org/rss/q-bio\"),\n",
    "    (\"arXiv - q-bio.BM\", \"http://arxiv.org/rss/q-bio.BM\"),\n",
    "]\n",
    "\n",
    "feed_ingestor = FeedIngestor()\n",
    "all_documents = []\n",
    "\n",
    "print(f\"Ingesting from {len(feed_sources)} feed sources...\")\n",
    "for i, (feed_name, feed_url) in enumerate(feed_sources, 1):\n",
    "    try:\n",
    "        with redirect_stderr(StringIO()):\n",
    "            feed_data = feed_ingestor.ingest_feed(feed_url, validate=False)\n",
    "        \n",
    "        feed_count = 0\n",
    "        for item in feed_data.items:\n",
    "            if not item.content:\n",
    "                item.content = item.description or item.title or \"\"\n",
    "            if item.content:\n",
    "                if not hasattr(item, 'metadata'):\n",
    "                    item.metadata = {}\n",
    "                item.metadata['source'] = feed_name\n",
    "                all_documents.append(item)\n",
    "                feed_count += 1\n",
    "        \n",
    "        if feed_count > 0:\n",
    "            print(f\"  [{i}/{len(feed_sources)}] {feed_name}: {feed_count} documents\")\n",
    "    except Exception:\n",
    "        continue\n",
    "\n",
    "if not all_documents:\n",
    "    sample_drug_data = \"\"\"\n",
    "    Aspirin (acetylsalicylic acid) is a medication used to reduce pain, fever, or inflammation. \n",
    "    It targets cyclooxygenase enzymes COX-1 and COX-2. Aspirin is commonly used for cardiovascular protection.\n",
    "    Ibuprofen is a nonsteroidal anti-inflammatory drug (NSAID) that targets COX-1 and COX-2 enzymes.\n",
    "    Metformin is an antidiabetic medication that targets AMP-activated protein kinase (AMPK).\n",
    "    Insulin targets the insulin receptor (INSR) to regulate glucose metabolism.\n",
    "    Warfarin is an anticoagulant that targets vitamin K epoxide reductase complex subunit 1 (VKORC1).\n",
    "    Atorvastatin is a statin medication that targets HMG-CoA reductase.\n",
    "    \"\"\"\n",
    "    \n",
    "    with open(\"data/sample_drugs.txt\", \"w\") as f:\n",
    "        f.write(sample_drug_data)\n",
    "    \n",
    "    file_ingestor = FileIngestor()\n",
    "    all_documents = file_ingestor.ingest(\"data/sample_drugs.txt\")\n",
    "\n",
    "documents = all_documents\n",
    "print(f\"Ingested {len(documents)} documents\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Normalizing and Chunking Documents\n",
    "\n",
    "Clean and normalize text, then split into chunks using entity-aware chunking to preserve drug/protein entity boundaries.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.normalize import TextNormalizer\n",
    "from semantica.split import TextSplitter\n",
    "\n",
    "normalizer = TextNormalizer()\n",
    "splitter = TextSplitter(\n",
    "    method=\"entity_aware\",\n",
    "    ner_method=\"spacy\",\n",
    "    chunk_size=CHUNK_SIZE,\n",
    "    chunk_overlap=CHUNK_OVERLAP\n",
    ")\n",
    "\n",
    "print(f\"Normalizing {len(documents)} documents...\")\n",
    "normalized_documents = []\n",
    "for i, doc in enumerate(documents, 1):\n",
    "    normalized_text = normalizer.normalize(\n",
    "        doc.content if hasattr(doc, 'content') else str(doc),\n",
    "        clean_html=True,\n",
    "        normalize_entities=True,\n",
    "        remove_extra_whitespace=True,\n",
    "        lowercase=False\n",
    "    )\n",
    "    normalized_documents.append(normalized_text)\n",
    "    if i % 50 == 0 or i == len(documents):\n",
    "        print(f\"  Normalized {i}/{len(documents)} documents...\")\n",
    "\n",
    "print(f\"Chunking {len(normalized_documents)} documents...\")\n",
    "chunked_documents = []\n",
    "for i, doc_text in enumerate(normalized_documents, 1):\n",
    "    try:\n",
    "        with redirect_stderr(StringIO()):\n",
    "            chunks = splitter.split(doc_text)\n",
    "        chunked_documents.extend(chunks)\n",
    "    except Exception:\n",
    "        simple_splitter = TextSplitter(method=\"recursive\", chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)\n",
    "        chunks = simple_splitter.split(doc_text)\n",
    "        chunked_documents.extend(chunks)\n",
    "    if i % 50 == 0 or i == len(normalized_documents):\n",
    "        print(f\"  Chunked {i}/{len(normalized_documents)} documents ({len(chunked_documents)} chunks so far)\")\n",
    "\n",
    "print(f\"Created {len(chunked_documents)} chunks from {len(normalized_documents)} documents\")\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.semantic_extract import NERExtractor\n",
    "\n",
    "# Using spaCy ML method (similar to NER cell)\n",
    "entity_extractor = NERExtractor(method=\"ml\", model=\"en_core_web_sm\")\n",
    "\n",
    "all_entities = []\n",
    "print(f\"Extracting entities from {len(chunked_documents)} chunks...\")\n",
    "\n",
    "for i, chunk in enumerate(chunked_documents, 1):\n",
    "    chunk_text = chunk.text if hasattr(chunk, 'text') else str(chunk)\n",
    "    try:\n",
    "        entities = entity_extractor.extract_entities(chunk_text)\n",
    "        all_entities.extend(entities)\n",
    "    except Exception:\n",
    "        continue\n",
    "    \n",
    "    if i % 20 == 0 or i == len(chunked_documents):\n",
    "        remaining = len(chunked_documents) - i\n",
    "        print(f\"  Processed {i}/{len(chunked_documents)} chunks ({len(all_entities)} entities found, {remaining} remaining)\")\n",
    "\n",
    "# Filter entities - spaCy returns standard types (PERSON, ORG, PRODUCT, etc.)\n",
    "# Map to biomedical categories based on context\n",
    "drugs = [e for e in all_entities if e.label == \"PRODUCT\" or (e.label == \"ORG\" and any(kw in e.text.lower() for kw in [\"drug\", \"pharma\", \"medication\"]))]\n",
    "proteins = [e for e in all_entities if e.label == \"ORG\" or (e.label == \"PRODUCT\" and any(kw in e.text.lower() for kw in [\"protein\", \"enzyme\", \"receptor\", \"kinase\", \"target\"]))]\n",
    "\n",
    "print(f\"Extracted {len(drugs)} drugs and {len(proteins)} proteins\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Extracting Drug-Target Relationships\n",
    "\n",
    "Extract relationships between drugs and proteins to understand drug-target interactions.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.semantic_extract import RelationExtractor\n",
    "\n",
    "# Using spaCy dependency parsing (similar to NER cell)\n",
    "relation_extractor = RelationExtractor(method=\"dependency\", model=\"en_core_web_sm\")\n",
    "\n",
    "all_relationships = []\n",
    "print(f\"Extracting relationships from {len(chunked_documents)} chunks...\")\n",
    "\n",
    "for i, chunk in enumerate(chunked_documents, 1):\n",
    "    chunk_text = chunk.text if hasattr(chunk, 'text') else str(chunk)\n",
    "    try:\n",
    "        relationships = relation_extractor.extract_relations(\n",
    "            chunk_text,\n",
    "            entities=all_entities,\n",
    "            relation_types=[\"targets\", \"inhibits\", \"activates\", \"binds_to\", \"interacts_with\"]\n",
    "        )\n",
    "        all_relationships.extend(relationships)\n",
    "    except Exception:\n",
    "        continue\n",
    "    \n",
    "    if i % 20 == 0 or i == len(chunked_documents):\n",
    "        print(f\"  Processed {i}/{len(chunked_documents)} chunks ({len(all_relationships)} relationships found)\")\n",
    "\n",
    "print(f\"Extracted {len(all_relationships)} relationships\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Resolving Duplicate Entities\n",
    "\n",
    "Detect and merge duplicate entities to ensure data quality and consistency.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conflict Detection and Resolution\n",
    "\n",
    "Detect and resolve conflicts in drug-target relationships from multiple research sources.\n",
    "\n",
    "- **Detection Method**: Relationship conflict detection identifies discrepancies in drug-target interactions across sources\n",
    "- **Resolution Strategy**: Credibility-weighted resolution prioritizes higher-credibility sources (e.g., Nature journals over arXiv preprints)\n",
    "- **Use Case**: Handles conflicting information when multiple sources report different drug-target relationships\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.conflicts import ConflictDetector, ConflictResolver\n",
    "\n",
    "detector = ConflictDetector()\n",
    "resolver = ConflictResolver(default_strategy=\"credibility_weighted\")\n",
    "\n",
    "# Convert to dict format for conflict detection\n",
    "entities = [\n",
    "    {\n",
    "        \"id\": ent.text if hasattr(ent, 'text') else str(ent),\n",
    "        \"name\": ent.text if hasattr(ent, 'text') else str(ent),\n",
    "        \"type\": ent.label if hasattr(ent, 'label') else \"ENTITY\",\n",
    "        \"confidence\": getattr(ent, 'confidence', 1.0),\n",
    "        \"source\": ent.metadata.get(\"source\", \"unknown\") if hasattr(ent, 'metadata') and ent.metadata else \"unknown\"\n",
    "    }\n",
    "    for ent in all_entities if hasattr(ent, 'text') or hasattr(ent, 'label')\n",
    "]\n",
    "\n",
    "relationships = [\n",
    "    {\n",
    "        \"id\": f\"{rel.subject.text}_{rel.object.text}_{rel.predicate}\",\n",
    "        \"source_id\": rel.subject.text,\n",
    "        \"target_id\": rel.object.text,\n",
    "        \"type\": rel.predicate,\n",
    "        \"confidence\": getattr(rel, 'confidence', 1.0),\n",
    "        \"source\": rel.metadata.get(\"source\", \"unknown\") if hasattr(rel, 'metadata') and rel.metadata else \"unknown\"\n",
    "    }\n",
    "    for rel in all_relationships if hasattr(rel, 'subject')\n",
    "]\n",
    "\n",
    "# Detect and resolve conflicts\n",
    "print(f\"Detecting conflicts in {len(entities)} entities, {len(relationships)} relationships...\")\n",
    "entity_conflicts = detector.detect_conflicts(entities)\n",
    "relationship_conflicts = detector.detect_relationship_conflicts(relationships)\n",
    "print(f\"Detected {len(entity_conflicts)} entity conflicts, {len(relationship_conflicts)} relationship conflicts\")\n",
    "\n",
    "# Resolve conflicts\n",
    "if entity_conflicts:\n",
    "    resolver.resolve_conflicts(entity_conflicts, strategy=\"credibility_weighted\")\n",
    "    print(f\"Resolved {len(entity_conflicts)} entity conflicts\")\n",
    "\n",
    "if relationship_conflicts:\n",
    "    resolver.resolve_conflicts(relationship_conflicts, strategy=\"credibility_weighted\")\n",
    "    print(f\"Resolved {len(relationship_conflicts)} relationship conflicts\")\n",
    "\n",
    "# GraphBuilder will use resolve_conflicts=True to apply resolutions automatically\n",
    "print(\"Conflicts resolved. GraphBuilder will use cleaned data.\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating Vector Embeddings\n",
    "\n",
    "Generate embeddings for drugs and proteins to enable similarity search.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.embeddings import EmbeddingGenerator\n",
    "from semantica.vector_store import VectorStore\n",
    "\n",
    "embedding_gen = EmbeddingGenerator(\n",
    "    provider=\"sentence_transformers\",\n",
    "    model=EMBEDDING_MODEL\n",
    ")\n",
    "\n",
    "vector_store = VectorStore(backend=\"faiss\", dimension=EMBEDDING_DIMENSION)\n",
    "\n",
    "print(f\"Generating embeddings for {len(drugs)} drugs and {len(proteins)} proteins...\")\n",
    "drug_texts = [d.text for d in drugs]\n",
    "drug_embeddings = embedding_gen.generate_embeddings(drug_texts)\n",
    "\n",
    "protein_texts = [p.text for p in proteins]\n",
    "protein_embeddings = embedding_gen.generate_embeddings(protein_texts)\n",
    "\n",
    "print(f\"Generated {len(drug_embeddings)} drug embeddings and {len(protein_embeddings)} protein embeddings\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Populating Vector Database\n",
    "\n",
    "Store drug and protein embeddings in the vector database with metadata for efficient similarity search.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"Storing {len(drug_embeddings)} drug vectors and {len(protein_embeddings)} protein vectors...\")\n",
    "drug_ids = vector_store.store_vectors(\n",
    "    vectors=drug_embeddings,\n",
    "    metadata=[{\"type\": \"drug\", \"name\": d.text, \"label\": d.label} for d in drugs]\n",
    ")\n",
    "\n",
    "protein_ids = vector_store.store_vectors(\n",
    "    vectors=protein_embeddings,\n",
    "    metadata=[{\"type\": \"protein\", \"name\": p.text, \"label\": p.label} for p in proteins]\n",
    ")\n",
    "\n",
    "print(f\"Stored {len(drug_ids)} drug vectors and {len(protein_ids)} protein vectors\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Building Drug-Target Knowledge Graph\n",
    "\n",
    "Construct a knowledge graph from extracted entities and relationships to enable graph-based reasoning.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.kg import GraphBuilder\n",
    "\n",
    "graph_builder = GraphBuilder()\n",
    "\n",
    "print(f\"Building graph from {len(all_entities)} entities, {len(all_relationships)} relationships...\")\n",
    "kg = graph_builder.build({\n",
    "    \"entities\": all_entities,\n",
    "    \"relationships\": all_relationships\n",
    "})\n",
    "\n",
    "entities_count = len(kg.get('entities', []))\n",
    "relationships_count = len(kg.get('relationships', []))\n",
    "print(f\"Graph: {entities_count} entities, {relationships_count} relationships\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Finding Similar Drugs via Vector Search\n",
    "\n",
    "Use vector similarity search to find drugs similar to a query drug based on their embeddings.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_drug = \"Aspirin\"\n",
    "query_embedding = embedding_gen.generate_embeddings([query_drug])[0]\n",
    "similar_drugs = vector_store.search_vectors(query_embedding, k=5)\n",
    "\n",
    "print(f\"Drugs similar to '{query_drug}':\")\n",
    "for i, result in enumerate(similar_drugs, 1):\n",
    "    metadata = result.get('metadata', {})\n",
    "    name = metadata.get('name', 'Unknown') if metadata else 'Unknown'\n",
    "    score = result.get('score', 0.0)\n",
    "    print(f\"{i}. {name} (similarity: {score:.3f})\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## GraphRAG: Hybrid Vector + Graph Retrieval\n",
    "\n",
    "Use GraphRAG to combine vector similarity search with knowledge graph traversal for enhanced retrieval and reasoning.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.context import AgentContext, ContextRetriever\n",
    "\n",
    "# Option 1: Use AgentContext (high-level, recommended)\n",
    "context = AgentContext(\n",
    "    vector_store=vector_store, \n",
    "    knowledge_graph=kg,\n",
    "    hybrid_alpha=0.6,\n",
    "    max_expansion_hops=2\n",
    ")\n",
    "\n",
    "# Option 2: Use ContextRetriever directly (more control)\n",
    "retriever = ContextRetriever(\n",
    "    vector_store=vector_store,\n",
    "    knowledge_graph=kg,\n",
    "    hybrid_alpha=0.6,\n",
    "    max_expansion_hops=2\n",
    ")\n",
    "\n",
    "# GraphRAG query using AgentContext\n",
    "query = \"What drugs target COX enzymes?\"\n",
    "results = context.retrieve(\n",
    "    query,\n",
    "    max_results=10,\n",
    "    use_graph=True,\n",
    "    expand_graph=True,\n",
    "    include_entities=True,\n",
    "    include_relationships=True\n",
    ")\n",
    "\n",
    "\n",
    "print(f\"Query: '{query}'\")\n",
    "print(f\"Retrieved {len(results)} results:\\n\")\n",
    "for i, result in enumerate(results[:5], 1):\n",
    "    print(f\"{i}. Score: {result.get('score', 0):.3f}\")\n",
    "    if result.get('content'):\n",
    "        print(f\"   {result['content'][:250]}\")\n",
    "    if result.get('related_entities'):\n",
    "        entities = result['related_entities']\n",
    "        names = [e.get('name', e.get('id', '')) for e in entities[:3]]\n",
    "        print(f\"   Entities: {', '.join(names)}\" + (f\" (+{len(entities)-3})\" if len(entities) > 3 else \"\"))\n",
    "    if result.get('related_relationships'):\n",
    "        print(f\"   Relationships: {len(result['related_relationships'])}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Visualizing the Knowledge Graph\n",
    "\n",
    "Generate an interactive visualization of the drug-target knowledge graph.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.visualization import KGVisualizer\n",
    "\n",
    "# Display interactive Plotly graph directly in notebook\n",
    "visualizer = KGVisualizer(layout=\"force\", node_size=20)\n",
    "fig = visualizer.visualize_network(kg, output=\"interactive\")\n",
    "\n",
    "# Display the figure (Plotly will show it automatically in notebook)\n",
    "fig.show() if fig else None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exporting Results\n",
    "\n",
    "Export the knowledge graph to various formats for further analysis or integration with other tools.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from semantica.export import GraphExporter\n",
    "\n",
    "exporter = GraphExporter()\n",
    "exporter.export(kg, output_path=\"drug_target_kg.json\", format=\"json\")\n",
    "exporter.export(kg, output_path=\"drug_target_kg.graphml\", format=\"graphml\")\n",
    "\n",
    "print(\"Exported knowledge graph to JSON and GraphML formats\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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": 2
}
