HippoRAG-PDF-v2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Implementation of HippoRAG using only APIs to search over a PDF file\n",
    "\n",
    "#### V2  \n",
    "- Change the prompts to use \"entities\" in place of \"named entities\", this seems to work better for Brazillian Portugues\n",
    "- Add node specificity calculation on ppr_graphdb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### PDF extraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dotenv import load_dotenv\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "from langchain_community.document_loaders import PyPDFLoader\n",
    "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
    "from langchain_openai import ChatOpenAI\n",
    "import numpy as np\n",
    "import faiss\n",
    "import pickle\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain.output_parsers.json import SimpleJsonOutputParser\n",
    "from neo4j import GraphDatabase\n",
    "import os\n",
    "load_dotenv()\n",
    "\n",
    "from uuid import uuid4\n",
    "unique_id = uuid4().hex[0:8]\n",
    "os.environ[\"LANGCHAIN_TRACING_V2\"] = \"true\"\n",
    "os.environ[\"LANGCHAIN_PROJECT\"] = f\"HippoRag-API-v2.1 - {unique_id}\"\n",
    "os.environ[\"LANGCHAIN_ENDPOINT\"] = \"https://api.smith.langchain.com\"\n",
    "# os.environ[\"LANGCHAIN_API_KEY\"] = \"LANGCHAIN_API_KEY\"  # Update to your API key\n",
    "\n",
    "llm_model = \"gpt-3.5-turbo-1106\" # 1106 works better extracting more entities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract text from the PDF\n",
    "\n",
    "pdf_file_path = \"/Users/viewitpro/Downloads/HIDROPONIA.pdf\"\n",
    "# pdf_file_path = \"/Users/viewitpro/Downloads/MANUAL_DUSTER.pdf\"\n",
    "\n",
    "pdf_loader = PyPDFLoader(pdf_file_path)\n",
    "pdf_pages = pdf_loader.load_and_split()\n",
    "pdf_full_text = \"\\n\".join([d.page_content for d in pdf_pages])\n",
    "\n",
    "text_splitter = RecursiveCharacterTextSplitter(\n",
    "    chunk_size=5000,\n",
    "    chunk_overlap=0,\n",
    "    length_function=len,\n",
    "    is_separator_regex=False,\n",
    ")\n",
    "\n",
    "passages_doc_arr = text_splitter.create_documents([pdf_full_text])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### RippoRAG INDEXING \n",
    "- Extract name entities from each passage\n",
    "\n",
    "- Create Triplets\n",
    "- Create P matrix N x P with number of times n appears in p \n",
    "- Encode Entities (OpenAI text-embedding-3-small with 256 vector size)\n",
    "- Find similar entities to create E' relations triplets \n",
    "- Add to GraphDB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract name entities from each passage\n",
    "\n",
    "\n",
    "# changed the prompt to include \"in the same language as the paragraph\"\n",
    "extract_entities_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"\"\"Your task is to extract entities from the given paragraph, in the same language as the paragraph.\n",
    "Respond with a JSON list of entities.\"\"\"),\n",
    "        (\"human\", \"\"\"Paragraph:\n",
    "```\n",
    "Radio City\n",
    "Radio City is India's first private FM radio station and was started on 3 July 2001.\n",
    "It plays Hindi, English and regional songs.\n",
    "Radio City recently forayed into New Media in May 2008 with the launch of a music portal - PlanetRadiocity.com that offers music related news, videos, songs, and other music-related features.\n",
    "```\"\"\"),\n",
    "        (\"ai\", \"\"\"{{\"entities\":\n",
    "    [\"Radio City\", \"India\", \"3 July 2001\", \"Hindi\", \"English\", \"May 2008\", \"PlanetRadiocity.com\"]\n",
    "}}\"\"\"),\n",
    "        (\"human\", \"\"\"Paragraph:```\n",
    "{passage_text}\n",
    "```\"\"\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "extract_triplets_prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"\"\"Your task is to construct an RDF (Resource Description Framework) graph from the given passages and entity lists. \n",
    "Respond with a JSON list of triples, with each triple representing a relationship in the RDF graph. \n",
    "\n",
    "Pay attention to the following requirements:\n",
    "- Each triple should contain at least one, but preferably two, of the named entities in the list for each passage.\n",
    "- Clearly resolve pronouns to their specific names to maintain clarity.\n",
    "\"\"\"),\n",
    "        (\"human\", \"\"\"Convert the paragraph into a JSON dict, it has a named entity list and a triple list.\n",
    "Paragraph:\n",
    "```\n",
    "Radio City\n",
    "Radio City is India's first private FM radio station and was started on 3 July 2001.\n",
    "It plays Hindi, English and regional songs.\n",
    "Radio City recently forayed into New Media in May 2008 with the launch of a music portal - PlanetRadiocity.com that offers music related news, videos, songs, and other music-related features.\n",
    "```\n",
    "\n",
    "{{\"entities\":\n",
    "    [\"Radio City\", \"India\", \"3 July 2001\", \"Hindi\", \"English\", \"May 2008\", \"PlanetRadiocity.com\"]\n",
    "}}\"\"\"),\n",
    "        (\"ai\", \"\"\"{{\"triples\": [\n",
    "            [\"Radio City\", \"located in\", \"India\"],\n",
    "            [\"Radio City\", \"is\", \"private FM radio station\"],\n",
    "            [\"Radio City\", \"started on\", \"3 July 2001\"],\n",
    "            [\"Radio City\", \"plays songs in\", \"Hindi\"],\n",
    "            [\"Radio City\", \"plays songs in\", \"English\"]\n",
    "            [\"Radio City\", \"forayed into\", \"New Media\"],\n",
    "            [\"Radio City\", \"launched\", \"PlanetRadiocity.com\"],\n",
    "            [\"PlanetRadiocity.com\", \"launched in\", \"May 2008\"],\n",
    "            [\"PlanetRadiocity.com\", \"is\", \"music portal\"],\n",
    "            [\"PlanetRadiocity.com\", \"offers\", \"news\"],\n",
    "            [\"PlanetRadiocity.com\", \"offers\", \"videos\"],\n",
    "            [\"PlanetRadiocity.com\", \"offers\", \"songs\"]\n",
    "    ]\n",
    "}}\"\"\"),\n",
    "        (\"human\", \"\"\"Convert the paragraph into a JSON dict, it has a entity list and a triple list.\n",
    "Paragraph:\n",
    "```\n",
    "{passage_text}\n",
    "```\n",
    "\n",
    "{entities}\"\"\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# passages_arr = [ { \"id\":idx, \"text\":passage.page_content } for idx, passage in enumerate(passages_doc_arr[:4])]\n",
    "passages_arr = [ { \"id\":idx, \"text\":passage.page_content } for idx, passage in enumerate(passages_doc_arr)]\n",
    "\n",
    "for passage in passages_arr:\n",
    "    passage[\"named_entities\"] =[]\n",
    "    passage[\"triples\"] = []\n",
    "    try:\n",
    "        json_output_parser = SimpleJsonOutputParser()\n",
    "        chain_entities = extract_entities_prompt | ChatOpenAI(model=llm_model, temperature=0.0) | json_output_parser\n",
    "        named_entities = chain_entities.invoke({\"passage_text\": passage[\"text\"]})\n",
    "        passage[\"named_entities\"] = named_entities[\"entities\"]\n",
    "\n",
    "        chain_triples = extract_triplets_prompt | ChatOpenAI(model=llm_model, temperature=0.0) | json_output_parser\n",
    "        triples = chain_triples.invoke({\"passage_text\": passage[\"text\"], \"entities\": named_entities})\n",
    "        passage[\"triples\"] = triples[\"triples\"]\n",
    "    except Exception as e:\n",
    "        print(f\"Error processing passage: {e}\")\n",
    "        continue\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "# save the passages to a file\n",
    "pickle.dump(passages_arr, open(\"passages_arr.pkl\", \"wb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the passages from a file\n",
    "passages_arr = pickle.load(open(\"passages_arr.pkl\", \"rb\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add to the passage named_entities entities created from the triples\n",
    "for passage in passages_arr:\n",
    "    named_entities = passage[\"named_entities\"]\n",
    "    named_entities = [entity.lower() for entity in named_entities]\n",
    "    for triple in passage[\"triples\"]:\n",
    "        if len(triple) != 3:\n",
    "            continue\n",
    "        named_entities.extend([triple[0].lower(), triple[2].lower()])\n",
    "    \n",
    "    passage[\"named_entities\"] = list(set(named_entities)) # remove duplicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a unique id for each named entity\n",
    "named_entities_dict = {}\n",
    "named_entities_next_id = 0\n",
    "\n",
    "for passage in passages_arr:\n",
    "    named_entities = passage[\"named_entities\"]\n",
    "    for idx, named_entity in enumerate(named_entities):\n",
    "        if named_entity not in named_entities_dict:\n",
    "            named_entities_dict[named_entity] = named_entities_next_id\n",
    "            named_entities_next_id += 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create P matrix which contains the number\n",
    "# of times each named_entitie appears in each passage\n",
    "\n",
    "n_p_matrix = np.zeros((len(named_entities_dict), len(passages_arr)))\n",
    "for idx, passage in enumerate(passages_arr):\n",
    "    named_entities = passage[\"named_entities\"]\n",
    "    for named_entity in named_entities:\n",
    "        # if named_entity.lower() == \"modelo\":\n",
    "        #     print(f\"Named entities: {named_entities}\")\n",
    "        named_entity_id = named_entities_dict[named_entity.lower()]\n",
    "        # get number of times named_entity appears in the passage\n",
    "        n_p_matrix[named_entity_id][idx] = passage[\"text\"].lower().count(named_entity.lower())\n",
    "\n",
    "\n",
    "# np.set_printoptions(edgeitems=10,linewidth=180)\n",
    "# print(n_p_matrix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Encode Entities\n",
    "\n",
    "entities = [key for key in named_entities_dict.keys()]\n",
    "\n",
    "# Embedding Entities using OpenAI\n",
    "openai_embeddings = OpenAIEmbeddings(model=\"text-embedding-3-small\", dimensions=256)\n",
    "entities_embeddings = openai_embeddings.embed_documents(entities)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find similar entities to create E' relations triplets  \n",
    "\n",
    "# Build the index \n",
    "faiss_index = faiss.IndexHNSWFlat(256, 32)\n",
    "\n",
    "# convert to numpy array\n",
    "vectors = np.array(entities_embeddings)\n",
    "# Add the vectors to the index\n",
    "faiss_index.add(vectors)\n",
    "\n",
    "# Perform the search\n",
    "k = 2\n",
    "distances, indices = faiss_index.search(vectors, k)\n",
    "\n",
    "filter_distance = 0.7\n",
    "# Create the similar enties list  \n",
    "similar_entities = []\n",
    "for idx, entity in enumerate(entities):\n",
    "    for i in range(k):\n",
    "        if indices[idx][i] != idx and distances[idx][i] <= filter_distance:\n",
    "\n",
    "            similar_entities.append(\n",
    "                {\n",
    "                    \"entity\": entity, \n",
    "                    \"similar_entity\": entities[indices[idx][i]]\n",
    "                }\n",
    "            )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# for passage in passages_arr:\n",
    "#     print(passage[\"named_entities\"])\n",
    "#     # print(passage[\"triples\"])\n",
    "#     print(\"\\n\\n\\n\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "URI = os.environ.get(\"NEO4J_URI\")\n",
    "AUTH = (os.environ.get(\"NEO4J_USER\"), os.environ.get(\"NEO4J_PASSWORD\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add to GraphDB \n",
    "\n",
    "with GraphDatabase.driver(URI, auth=AUTH) as driver:\n",
    "    with driver.session() as session:\n",
    "\n",
    "        # create all entities with embeddings and node id\n",
    "        def create_all_entities(tx, all_entities):\n",
    "            query = \"\"\"\n",
    "            UNWIND $all_entities AS ae\n",
    "            MERGE (a:Entity {name: ae.name, node_id: ae.node_id, embedding: ae.embedding})\n",
    "            \"\"\"\n",
    "            tx.run(query, all_entities=all_entities)\n",
    "        \n",
    "        all_entities = [{\"name\": entity, \"node_id\": named_entities_dict[entity], \"embedding\": entities_embeddings[named_entities_dict[entity]]} for entity in named_entities_dict.keys()]\n",
    "        session.execute_write(lambda tx: create_all_entities(tx, all_entities))\n",
    "\n",
    "\n",
    "        # create RELATE_TO relationships\n",
    "        def create_relateto_relationships(tx, triplets):\n",
    "            query = \"\"\"\n",
    "            UNWIND $triplets AS triplet\n",
    "            MATCH (a:Entity {name: triplet.subject}), (b:Entity {name: triplet.object})\n",
    "            MERGE (a)-[:RELATES_TO {type: triplet.predicate}]->(b)\n",
    "            \"\"\"\n",
    "            tx.run(query, triplets=triplets)\n",
    "\n",
    "        triplets = []\n",
    "        for passage in passages_arr:\n",
    "            for triple in passage[\"triples\"]:\n",
    "                if len(triple) != 3:\n",
    "                    continue\n",
    "                # print(triple)\n",
    "                predicate=triple[1].replace(\" \", \"_\").upper()\n",
    "                subject=triple[0].lower()\n",
    "                object=triple[2].lower()\n",
    "                triplets.append({\n",
    "                    \"subject\": subject, \n",
    "                    \"predicate\": predicate, \n",
    "                    \"object\": object,\n",
    "                    \"passageId_subject\": passage[\"id\"],\n",
    "                    \"passageId_object\": passage[\"id\"],\n",
    "                })\n",
    "        \n",
    "        session.execute_write(lambda tx: create_relateto_relationships(tx, triplets))\n",
    "\n",
    "        # create the vector index on  Neo4j\n",
    "        def create_vector_index(tx):\n",
    "            query = \"\"\"\n",
    "            CREATE VECTOR INDEX entitieDB IF NOT EXISTS\n",
    "            FOR (m:Entity)\n",
    "            ON m.embedding\n",
    "            OPTIONS {indexConfig: {\n",
    "                `vector.dimensions`: 256,\n",
    "                `vector.similarity_function`: 'cosine'\n",
    "            }}\n",
    "            \"\"\"\n",
    "            tx.run(query)\n",
    "        \n",
    "        session.execute_write(lambda tx: create_vector_index(tx))\n",
    "\n",
    "        # create similar entities graph\n",
    "        def create_similar_entities(tx, similar_entities):\n",
    "            query = \"\"\"\n",
    "            UNWIND $similar_entities AS se\n",
    "            MATCH (a:Entity {name: se.entity}), (b:Entity {name: se.similar_entity})\n",
    "            MERGE (a)-[:SIMILAR_TO]->(b)\n",
    "            \"\"\"\n",
    "            tx.run(query, similar_entities=similar_entities)\n",
    "        \n",
    "        session.execute_write(lambda tx: create_similar_entities(tx, similar_entities))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### RippoRAG RETRIEVAL\n",
    "- Extract name entities from query\n",
    "\n",
    "- Vector search on GraphDB to find similar Entities\n",
    "- PPR on GraphDB to find \"good\" nodes\n",
    "- Rank passages from nodes found \n",
    "- Send passages + query to LLM for answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Extract name entities from query\n",
    "\n",
    "def extract_entities_from_query(user_query):\n",
    "    # This prompt is a simpler version o the original, it works better for small paragraphs and less entities and\n",
    "    # in other languages like portuguese\n",
    "    extract_entities_custom_prompt = ChatPromptTemplate.from_messages(\n",
    "        [\n",
    "            (\"system\", \"\"\"Your task is to extract all entities from the given paragraph, in the same language as the paragraph.\n",
    "    Respond with a JSON list of entities like {{\"entities\":[\"entity1\", \"entity2\", ...]}}\"\"\"),\n",
    "            (\"human\", \"\"\"Paragraph:```\n",
    "    {passage_text}\n",
    "    ```\"\"\"),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    json_output_parser = SimpleJsonOutputParser()\n",
    "    chain_query_entities = extract_entities_custom_prompt | ChatOpenAI(model=llm_model, temperature=0.0) | json_output_parser\n",
    "    #chain_query_entities = extract_entities_prompt | ChatOpenAI(model=llm_model, temperature=0.0) | json_output_parser\n",
    "    query_entities = chain_query_entities.invoke({\"passage_text\": user_query})\n",
    "    query_entities[\"named_entities\"] = query_entities[\"entities\"] # change the name to named_entities\n",
    "\n",
    "    return query_entities\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Vector search on GraphDB to find similar Entities\n",
    "\n",
    "#https://neo4j.com/docs/cypher-manual/current/indexes/semantic-indexes/vector-indexes/\n",
    "\n",
    "def vector_search_graphdb(query_entitie, min_score=0.8):\n",
    "    openai_embeddings = OpenAIEmbeddings(model=\"text-embedding-3-small\", dimensions=256)\n",
    "    query_entitie_embedding = openai_embeddings.embed_query(query_entitie)\n",
    "    results = []\n",
    "    with GraphDatabase.driver(URI, auth=AUTH) as driver:\n",
    "        with driver.session() as session:\n",
    "            def search_vector(tx, query_embedding):\n",
    "                cypher_query = f\"\"\"\n",
    "                CALL db.index.vector.queryNodes('entitieDB', 3, {query_embedding}) YIELD node, score\n",
    "                RETURN node.node_id as id, node.name as name, score\n",
    "                \"\"\"\n",
    "                return tx.run(cypher_query).data()\n",
    "\n",
    "            db_result = session.execute_read(lambda tx: search_vector(tx, query_entitie_embedding))\n",
    "            \n",
    "            for result in db_result:\n",
    "                if result[\"score\"] >= min_score:\n",
    "                    results.append(result)\n",
    "\n",
    "            # return the SIMILAR_TO entities too  E'\n",
    "            def similar_search(tx, similar_entities):\n",
    "                cypher_query = f\"\"\"\n",
    "                MATCH (a:Entity)-[:SIMILAR_TO]->(b:Entity)\n",
    "                WHERE a.name IN $similar_entities\n",
    "                RETURN b.node_id AS id, b.name AS name\n",
    "                \"\"\"\n",
    "                return tx.run(cypher_query, similar_entities=similar_entities).data()\n",
    "\n",
    "            similar_entities = [ result[\"name\"] for result in results]\n",
    "            db_result = session.execute_read(lambda tx: similar_search(tx, similar_entities))\n",
    "            results.extend(db_result)\n",
    "\n",
    "    return results\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PPR on GraphDB to find \"good\" nodes\n",
    "\n",
    "# gds.pageRank only works on Neo4j 4.0 auraDS and its a payed feature on Neo4j 4.0 \n",
    "# for this reason I will export the graph from auraDB to igraph and use the igraph implementation of PageRank  \n",
    "import igraph as ig\n",
    "\n",
    "def ppr_graphdb(nodes):\n",
    "    # remove duplicated nodes\n",
    "    unique_data = {}\n",
    "    for item in nodes:\n",
    "        unique_data[item['id']] = item\n",
    "\n",
    "    unique_nodes = list(unique_data.values())\n",
    "\n",
    "    with GraphDatabase.driver(URI, auth=AUTH) as driver:\n",
    "        with driver.session() as session:\n",
    "            #get nodes and relations from graphDB\n",
    "            graph_nodes = session.execute_read(lambda tx: tx.run(\"MATCH (n:Entity) RETURN n.node_id AS node_id\").data())\n",
    "            relations = session.execute_read(lambda tx: tx.run(\"MATCH (a)-[r:RELATES_TO]->(b) RETURN a.node_id AS source, b.node_id AS target\").data())\n",
    "\n",
    "            # Create an igraph graph\n",
    "            g = ig.Graph(directed=True) \n",
    "            # add nodes\n",
    "            for gnode in graph_nodes:\n",
    "                # g.add_vertex(name=str(gnode[\"node_id\"]), labels=str(gnode[\"node_id\"]), passageId=gnode[\"passageId\"])\n",
    "                g.add_vertex(name=str(gnode[\"node_id\"]), labels=str(gnode[\"node_id\"]))\n",
    "\n",
    "            # add edges\n",
    "            g.add_edges([(str(rel[\"source\"]), str(rel[\"target\"])) for rel in relations])\n",
    "\n",
    "            # Personalized PageRank\n",
    "            personalization = [0] * len(g.vs)\n",
    "            # Set personalization vector \n",
    "            personalization_value  = 1.0 / len(nodes)\n",
    "            for node in unique_nodes:\n",
    "                idx = g.vs.find(name=str(node[\"id\"])).index\n",
    "                personalization[idx] += personalization_value \n",
    "\n",
    "                # calculate node specificity len(node_passages) ** -1\n",
    "                node_sum = np.sum(n_p_matrix[node[\"id\"]])\n",
    "                if node_sum == 0:\n",
    "                    print(f\"Node sum{node['id']} is zero\", node)\n",
    "                else:\n",
    "                    personalization[idx] *= node_sum ** -1\n",
    "\n",
    "\n",
    "            #https://igraph.org/python/api/0.9.11/igraph._igraph.GraphBase.html#personalized_pagerank\n",
    "            pagerank_scores = g.personalized_pagerank(damping=0.85, reset=personalization)\n",
    "\n",
    "    return pagerank_scores\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Query entities: {'entities': ['modelo'], 'named_entities': ['modelo']}\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[{'id': 100, 'name': 'modelo', 'score': 0.9999996423721313},\n",
       " {'id': 165, 'name': 'tipo', 'score': 0.8274419903755188},\n",
       " {'id': 165, 'name': 'tipo'},\n",
       " {'id': 100, 'name': 'modelo'}]"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create the final answer with the passages\n",
    "from langchain_core.output_parsers.string import StrOutputParser\n",
    "\n",
    "# User question\n",
    "user_query = \"Qual o modelo mais versátil?\"\n",
    "# user_query = \"Quais as diferenças entre os modelos de casas de vegetação, capela e dente-de-serra?\"\n",
    "\n",
    "# Extract name entities from query\n",
    "query_entities = extract_entities_from_query(user_query)\n",
    "print(\"Query entities:\",query_entities)\n",
    "\n",
    "# find the nodes on the graphDB that are similar to the named entities in the query\n",
    "base_nodes = []\n",
    "if len(query_entities[\"named_entities\"]) == 0:\n",
    "    print(\"No named entities found in the query\")\n",
    "else:\n",
    "    for query_entitie in query_entities[\"named_entities\"]:\n",
    "        results = vector_search_graphdb(query_entitie)\n",
    "        base_nodes.extend(results)\n",
    "\n",
    "base_nodes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "nodes_score = ppr_graphdb(base_nodes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import numpy as np\n",
    "# indices = np.where(np.array(nodes_score) > 0)[0]\n",
    "# print(indices)\n",
    "\n",
    "# # nodes_score\n",
    "# for idx in indices:\n",
    "#     print(idx, entities[idx])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.         3.28205335 0.         0.01096668 0.         0.\n",
      " 0.         0.         0.         0.         0.         0.\n",
      " 0.         0.         0.         0.         0.         0.\n",
      " 0.         0.         0.         0.         0.        ]\n",
      "[ 1  3 22 10  2  4  5  6  7  8  9 11 21 12 13 14 15 16 17 18 19 20  0]\n"
     ]
    }
   ],
   "source": [
    "# get the passages score and order\n",
    "\n",
    "def rank_passages(p_matrix, ppr_nodes_scores):\n",
    "    passages_scores = np.dot(p_matrix.T, np.array(ppr_nodes_scores))\n",
    "    print(passages_scores)\n",
    "    print(np.argsort(passages_scores)[::-1])\n",
    "    return passages_scores, np.argsort(passages_scores)[::-1]\n",
    "\n",
    "\n",
    "passages_score, passages_order = rank_passages(n_p_matrix, nodes_score)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'O modelo mais versátil é o \"teto em arco\", pois ele pode ser utilizado nas Regiões Sul, Sudeste, Centro-Oeste e Norte, fazendo-se as devidas adaptações e adicionando acessórios como janelas advectivas, teto zenital e sistemas de refrigeração ou aquecimento.'"
      ]
     },
     "execution_count": 49,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# create the context with the passages \n",
    "context = \"\"\n",
    "max_passages = 3\n",
    "passage_count = 0\n",
    "\n",
    "for idx in passages_order:\n",
    "    if passages_score[idx] > 0:\n",
    "        if passage_count > max_passages-1:\n",
    "            break\n",
    "        passage_count += 1\n",
    "        context += f\"{passages_arr[idx]['text']}\\n\\n\"\n",
    "\n",
    "\n",
    "template = \"\"\"Answer the question based only on the provided context, answer in the same languague as the question.\n",
    "{context}\n",
    "Question: {question}\n",
    "\"\"\"\n",
    "rag_prompt_text = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"Answer the question based only on the provided context, answer in the same languague as the question.\"),\n",
    "        (\"human\", \"\"\"Context:```\n",
    "{context}\n",
    "```\n",
    "\n",
    "Question: {question}\"\"\"),\n",
    "    ]\n",
    ")\n",
    "# Create the chain to answer the question, will use a stronger model like GPT-4o\n",
    "final_chain = rag_prompt_text | ChatOpenAI(model=\"gpt-4o\", temperature=0.0) | StrOutputParser()\n",
    "\n",
    "# Invoke the chain with a query\n",
    "answer = final_chain.invoke({\"context\": context, \"question\": user_query})\n",
    "\n",
    "answer"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ircot",
   "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
}