notebooks/kg-construction/4-geocoding.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run 'shared.ipynb'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import googlemaps\n",
    "\n",
    "gmaps = googlemaps.Client(key=GOOGLE_MAPS_API_KEY)\n",
    "\n",
    "# Geocoding an address\n",
    "geocode_result = gmaps.geocode('1600 Amphitheatre Parkway, Mountain View, CA')\n",
    "\n",
    "geocode_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_city(address_component):\n",
    "    if (address_component['types'] == ['locality', 'political']\n",
    "        or address_component['types'] == ['administrative_area_level_3', 'political'] \n",
    "        or address_component['types'] == ['postal_town']\n",
    "        or address_component['types'] == ['political', 'sublocality', 'sublocality_level_1']\n",
    "        or address_component['types'] == ['neighborhood', 'political']\n",
    "        or address_component['types'] == ['locality', 'political']\n",
    "        ):\n",
    "          return True\n",
    "    return False\n",
    "\n",
    "def get_city(geocode):\n",
    "    found_city = None\n",
    "    administrative_area_level_3 = None\n",
    "    locality = None\n",
    "    postal_town = None\n",
    "    sublocality_level_1 = None\n",
    "    neighborhood = None\n",
    "\n",
    "    for address_component in geocode['address_components']:\n",
    "        match address_component['types']:\n",
    "            case ['administrative_area_level_3', 'political']:\n",
    "                administrative_area_level_3 = address_component\n",
    "                continue\n",
    "            case ['locality', 'political']:\n",
    "                locality = address_component\n",
    "                continue\n",
    "            case ['postal_town']:\n",
    "                postal_town = address_component\n",
    "                continue\n",
    "            case ['political', 'sublocality', 'sublocality_level_1']:\n",
    "                sublocality_level_1 = address_component\n",
    "                continue\n",
    "            case ['neighborhood', 'political']:\n",
    "                neighborhood = address_component\n",
    "                continue\n",
    "    possible_city = [locality, administrative_area_level_3, sublocality_level_1, postal_town, neighborhood]\n",
    "    found_city = list(filter(None, possible_city))\n",
    "    return found_city[0] if found_city else None\n",
    "\n",
    "def filter_state(address_component):\n",
    "    if address_component['types'] == ['administrative_area_level_1', 'political']:\n",
    "          return True  \n",
    "    return False\n",
    "\n",
    "def get_state(geocode):\n",
    "    state_list = list(filter(filter_state, geocode['address_components']))\n",
    "    if (state_list):\n",
    "        return state_list[0]\n",
    "    else:\n",
    "        return None\n",
    "    \n",
    "def filter_postalcode(address_component):\n",
    "    if address_component['types'] == ['postal_code']:\n",
    "          return True  \n",
    "    return False\n",
    "    \n",
    "def get_postalcode(geocode):\n",
    "    zipcode_list = list(filter(filter_postalcode, geocode['address_components']))\n",
    "    if (zipcode_list):\n",
    "        return zipcode_list[0]\n",
    "    else:\n",
    "        return None\n",
    "    \n",
    "def filter_country(address_component):\n",
    "    if address_component['types'] == ['country', 'political']:\n",
    "          return True  \n",
    "    return False\n",
    "\n",
    "def get_country(geocode):\n",
    "    country_list = list(filter(filter_country, geocode['address_components']))\n",
    "    if (country_list):\n",
    "        return country_list[0]\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "def get_location(geocode):\n",
    "    return geocode['geometry']['location']\n",
    "\n",
    "\n",
    "def long_name(address_component):\n",
    "    if address_component:\n",
    "        return address_component['long_name'] if address_component['long_name'] else None\n",
    "    else:\n",
    "        return None\n",
    "\n",
    "def print_address(name, address, city, state, postal, country):\n",
    "    print(f\"{name} is located at {address}\")\n",
    "    print(f\"\\tcomponents: {city}, {state} {postal}, {country}\")\n",
    "    \n",
    "# example_geocode = geocode_result[0]\n",
    "# print(\"city\", get_city(example_geocode))\n",
    "# print(\"state\", get_state(example_geocode))\n",
    "# print(\"zip\", get_postalcode(example_geocode))\n",
    "# print(\"country\", get_country(example_geocode))\n",
    "# print(\"location\", get_location(example_geocode))\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cypher_geoloc = gdb.execute_query(\"CALL apoc.spatial.geocode('1600 Amphitheatre Parkway, Mountain View, CA');\").records[0]\n",
    "\n",
    "print(f\"latitude: {cypher_geoloc['latitude']} longitude: {cypher_geoloc['longitude']}\")\n",
    "\n",
    "print(f\"address_components: {cypher_geoloc}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "gdb.execute_query(\"\"\"\n",
    "  CREATE POINT INDEX address_locations IF NOT EXISTS\n",
    "  FOR (n:Address) ON (n.location)\n",
    "\"\"\")\n",
    "\n",
    "gdb.execute_query(\"SHOW INDEXES\").records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "get_managers_cypher = \"\"\"\n",
    "  MATCH (mgr:Manager)\n",
    "  RETURN mgr { .cik, .name, .address, .location}\n",
    "\"\"\"\n",
    "\n",
    "manager_rows = gdb.execute_query(get_managers_cypher).records\n",
    "\n",
    "managers = list(map(lambda row: row['mgr'], manager_rows))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "managers[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add geocode to each manager\n",
    "\n",
    "for manager in managers:\n",
    "  if 'geocode' not in manager:\n",
    "    geocode_for_address = gmaps.geocode(manager['address'])\n",
    "    if len(geocode_for_address) > 0:\n",
    "      manager['geocode'] = geocode_for_address[0] # accept first result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create composite index on city, state for Addresses\n",
    "\n",
    "gdb.execute_query(\"\"\"\n",
    "CREATE INDEX composite_address_index IF NOT EXISTS\n",
    "FOR (n:Address) \n",
    "ON (n.city, n.state)\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_manager_location_cypher = \"\"\"\n",
    "  MATCH (mgr:Manager {cik: $managerCik})\n",
    "  SET mgr.location = point({latitude: $latitude, longitude: $longitude})\n",
    "  MERGE (addr:Address {city: $city, state: $state})\n",
    "   ON CREATE SET addr.country = $country\n",
    "   ON MATCH SET addr.location = point({latitude: $latitude, longitude: $longitude})\n",
    "  MERGE (mgr)-[:LOCATED_AT]->(addr)\n",
    "\"\"\"\n",
    "\n",
    "for manager in managers:\n",
    "  if 'geocode' not in manager:\n",
    "    continue\n",
    "\n",
    "  location = get_location(manager['geocode'])\n",
    "  city = get_city(manager['geocode'])\n",
    "  state = get_state(manager['geocode'])\n",
    "  country = get_country(manager['geocode'])\n",
    "  postal = get_postalcode(manager['geocode'])\n",
    "\n",
    "  cityOrState = city if city else state\n",
    "  stateOrCountry = state if state else country \n",
    "\n",
    "  print_address(manager['name'], manager['address'],\n",
    "      long_name(cityOrState), \n",
    "      long_name(stateOrCountry), \n",
    "      long_name(postal), \n",
    "      long_name(country)\n",
    "  )\n",
    "\n",
    "  if location and cityOrState and stateOrCountry:\n",
    "    gdb.execute_query(set_manager_location_cypher,\n",
    "      managerCik= manager['cik'],\n",
    "      latitude= location['lat'],\n",
    "      longitude= location['lng'],\n",
    "      city= long_name(cityOrState),\n",
    "      state= long_name(stateOrCountry),\n",
    "      country= long_name(country)\n",
    "    )\n",
    "  else:\n",
    "    print (f\"no location for {manager['name']} {manager['cik']} within {manager['geocode']}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Where is Blackrock located?\n",
    "gdb.execute_query(\"\"\"\n",
    "  CALL db.index.fulltext.queryNodes(\"fullTextManagerNames\", \"Blackrock\") YIELD node, score\n",
    "  MATCH p=(node)-[:LOCATED_AT]->(address:Address)\n",
    "  RETURN p\n",
    "\"\"\").records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# How many investment firms are at the same address as Blackrock?\n",
    "gdb.execute_query(\"\"\"\n",
    "  CALL db.index.fulltext.queryNodes(\"fullTextManagerNames\", \"Blackrock\") YIELD node, score\n",
    "  MATCH p=(node)-[:LOCATED_AT]->(address:Address)<-[:LOCATED_AT]-(other:Manager)\n",
    "  RETURN count(other) as numManagers\n",
    "\"\"\").records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Which state has the most investment firms?\n",
    "gdb.execute_query(\"\"\"\n",
    "  MATCH p=(:Manager)-[:LOCATED_AT]->(address:Address)\n",
    "  RETURN address.state as state, count(address.state) as numManagers\n",
    "    ORDER BY numManagers DESC\n",
    "    LIMIT 10\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# What are the cities in California with the most investment firms?\n",
    "gdb.execute_query(\"\"\"\n",
    "  MATCH p=(:Manager)-[:LOCATED_AT]->(address:Address)\n",
    "         WHERE address.state = 'California'\n",
    "  RETURN address.city as city, count(address.city) as numManagers\n",
    "    ORDER BY numManagers DESC\n",
    "    LIMIT 10\n",
    "\"\"\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# What are top investment firms in San Francisco?\n",
    "gdb.execute_query(\"\"\"\n",
    "  MATCH p=(mgr:Manager)-[:LOCATED_AT]->(address:Address),\n",
    "         (mgr)-[owns:OWNS_STOCK_IN]->(:Company)\n",
    "         WHERE address.city = $city\n",
    "            AND address.state = $state\n",
    "  RETURN mgr.name as city, sum(owns.value) as totalInvestmentValue\n",
    "    ORDER BY totalInvestmentValue DESC\n",
    "    LIMIT 10\n",
    "\"\"\", \n",
    "  city=\"San Francisco\", \n",
    "  state=\"California\"\n",
    ").records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain.prompts.prompt import PromptTemplate\n",
    "from langchain.chains import GraphCypherQAChain\n",
    "\n",
    "CYPHER_GENERATION_TEMPLATE = \"\"\"Task:Generate Cypher statement to query a graph database.\n",
    "Instructions:\n",
    "Use only the provided relationship types and properties in the schema.\n",
    "Do not use any other relationship types or properties that are not provided.\n",
    "Schema:\n",
    "{schema}\n",
    "Note: Do not include any explanations or apologies in your responses.\n",
    "Do not respond to any questions that might ask anything else than for you to construct a Cypher statement.\n",
    "Do not include any text except the generated Cypher statement.\n",
    "Examples: Here are a few examples of generated Cypher statements for particular questions:\n",
    "# What investment firms are in San Francisco?\n",
    "MATCH (mgr:Manager)-[:LOCATED_AT]->(a:Address)\n",
    "    WHERE a.city = 'San Francisco'\n",
    "RETURN mgr.name\n",
    "\n",
    "# What firms in San Francisco have the most investments?\n",
    "MATCH (mgr:Manager)-[:LOCATED_AT]->(a:Address),\n",
    "        (mgr:Manager)-[owns:OWNS_STOCK_IN]->(com:Company)\n",
    "    WHERE a.city = 'San Francisco'\n",
    "WITH mgr, sum(owns.value) as totalInvestmentValue\n",
    "RETURN mgr.name + \" owns $\" + apoc.number.format(totalInvestmentValue) + \" worth of shares.\"\n",
    "\n",
    "\n",
    "The question is:\n",
    "{question}\"\"\"\n",
    "\n",
    "CYPHER_GENERATION_PROMPT = PromptTemplate(\n",
    "    input_variables=[\"schema\", \"question\"], template=CYPHER_GENERATION_TEMPLATE\n",
    ")\n",
    "\n",
    "kg=Neo4jGraph(\n",
    "        url=NEO4J_URI, username=NEO4J_USERNAME, password=NEO4J_PASSWORD, database=NEO4J_DATABASE\n",
    "    )\n",
    "cypherChain = prettifyChain(GraphCypherQAChain.from_llm(\n",
    "    ChatOpenAI(temperature=0),\n",
    "    graph=kg,\n",
    "    verbose=True,\n",
    "    cypher_prompt=CYPHER_GENERATION_PROMPT,\n",
    "))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cypherChain(\"What investment firms are in San Francisco?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cypherChain(\"What investment firms are in New York?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cypherChain(\"What firms in San Francisco have the most investments?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "company_rows = gdb.execute_query(\"\"\"\n",
    "  MATCH (com:Company)\n",
    "  RETURN com { .cusip6, .name } as company\n",
    "\"\"\").records\n",
    "\n",
    "companies = list(map(lambda row: row['company'], company_rows))\n",
    "\n",
    "print(companies[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a langchain vector store from the existing Neo4j knowledge graph.\n",
    "neo4j_vector_store = Neo4jVector.from_existing_graph(\n",
    "    embedding=embeddings_api,\n",
    "    url=NEO4J_URI,\n",
    "    username=NEO4J_USERNAME,\n",
    "    password=NEO4J_PASSWORD,\n",
    "    index_name=VECTOR_INDEX_NAME,\n",
    "    node_label=VECTOR_NODE_LABEL,\n",
    "    text_node_properties=[VECTOR_SOURCE_PROPERTY],\n",
    "    embedding_node_property=VECTOR_EMBEDDING_PROPERTY,\n",
    ")\n",
    "\n",
    "# Create a retriever from the vector store\n",
    "retriever = neo4j_vector_store.as_retriever()\n",
    "\n",
    "# Create a chatbot Question & Answer chain from the retriever\n",
    "chain = RetrievalQAWithSourcesChain.from_chain_type(\n",
    "    chat_api, chain_type=\"stuff\", retriever=retriever\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chain(f\"Where is Fedex headquartrered?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_company_location_cypher = \"\"\"\n",
    "  MATCH (com:Company {cusip6: $cusip6})\n",
    "  SET com.address = $companyAddress\n",
    "  SET com.location = point({latitude: $latitude, longitude: $longitude})\n",
    "  MERGE (addr:Address {city: $city, state: $state} )\n",
    "  SET addr.country = $country\n",
    "  SET addr.location = point({latitude: $latitude, longitude: $longitude})\n",
    "  MERGE (com)-[:LOCATED_AT]->(addr)\n",
    "\"\"\"\n",
    "\n",
    "for company in companies:\n",
    "    result = chain(f\"Where is {company['name']} headquartered?\")\n",
    "    address_statement = result['answer']\n",
    "    address_geocodes = gmaps.geocode(address_statement)\n",
    "    if len(address_geocodes) > 0:\n",
    "        address_geocode = address_geocodes[0]\n",
    "\n",
    "        cusip6 = company['cusip6']\n",
    "\n",
    "        location = get_location(address_geocode)\n",
    "        city = get_city(address_geocode)\n",
    "        state = get_state(address_geocode)\n",
    "        postal = get_postalcode(address_geocode)\n",
    "        country = get_country(address_geocode)\n",
    "\n",
    "        cityOrState = city if city else state\n",
    "        stateOrCountry = state if state else country \n",
    "\n",
    "        print_address(company['name'], address_geocode['formatted_address'],\n",
    "            long_name(cityOrState), \n",
    "            long_name(stateOrCountry), \n",
    "            long_name(postal), \n",
    "            long_name(country)\n",
    "        )\n",
    "\n",
    "        if location and cityOrState and stateOrCountry:\n",
    "            kg.query(set_company_location_cypher, params={\n",
    "                \"companyAddress\": address_geocode['formatted_address'],\n",
    "                \"cusip6\": company['cusip6'],\n",
    "                \"latitude\": location['lat'],\n",
    "                \"longitude\": location['lng'],\n",
    "                \"city\": long_name(cityOrState),\n",
    "                \"state\": long_name(stateOrCountry),\n",
    "                \"country\": long_name(country)\n",
    "            })\n",
    "    else:\n",
    "        print(f\"no geocode found for {company['companyName']} at {address_statement}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Which state has the most public companies listed?\n",
    "gdb.execute_query(\"\"\"\n",
    "  MATCH p=(:Company)-[:LOCATED_AT]->(address:Address)\n",
    "  RETURN address.state as state, count(address.state) as numCompanies\n",
    "    ORDER BY numCompanies DESC\n",
    "\"\"\").records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Which city in California has the most companies listed?\n",
    "gdb.execute_query(\"\"\"\n",
    "  MATCH p=(:Company)-[:LOCATED_AT]->(address:Address)\n",
    "         WHERE address.state = 'California'\n",
    "  RETURN address.city as city, count(address.city) as numCompanies\n",
    "    ORDER BY numCompanies DESC\n",
    "\"\"\").records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gdb.execute_query(\"\"\"\n",
    "  CALL db.index.fulltext.queryNodes(\"fullTextCompanyNames\", $companyName) \n",
    "    YIELD node, score\n",
    "  WITH node as com\n",
    "  RETURN com.location, com.address\n",
    "\"\"\", companyName=\"Fedex\").records[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Which investment firms are near Fedex?\n",
    "gdb.execute_query(\"\"\"\n",
    "  CALL db.index.fulltext.queryNodes(\"fullTextCompanyNames\", $companyName) \n",
    "         YIELD node, score\n",
    "  WITH node as com\n",
    "  MATCH (mgr:Manager)\n",
    "    WHERE point.distance(mgr.location, com.location) < $withinMeters\n",
    "  WITH mgr, \n",
    "    toInteger(point.distance(mgr.location, com.location) / 1000) as distanceKm\n",
    "  RETURN mgr.name, mgr.address, \n",
    "        apoc.number.format(distanceKm) + \"km\" as distance\n",
    "\n",
    "\"\"\", companyName=\"Fedex\", withinMeters= 100 * 1000).records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Investment firms within an hour drive of Netapp?\n",
    "gdb.execute_query(\"\"\"\n",
    "  CALL db.index.fulltext.queryNodes(\"fullTextCompanyNames\", $companyName) \n",
    "         YIELD node, score\n",
    "  WITH node as com\n",
    "  MATCH (com)-[:LOCATED_AT]->(comAddr:Address),\n",
    "         (mgr:Manager)-[:LOCATED_AT]->(mgrAddr:Address)\n",
    "    WHERE point.distance(comAddr.location, mgrAddr.location) < $withinMeters\n",
    "  WITH mgr, mgrAddr,\n",
    "    toInteger(point.distance(comAddr.location, mgrAddr.location) / 1000) as distanceKm\n",
    "  RETURN count(mgr.name) as numManagers, mgrAddr.city + \", \" + mgrAddr.state as cityState,\n",
    "        apoc.number.format(distanceKm) + \"km\" as distance\n",
    "    ORDER BY distance ASC LIMIT 10\n",
    "\n",
    "\"\"\", companyName=\"Netapp\", withinMeters= 50 * 1000).records"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "kg.query(\"SHOW INDEXES\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "kg-query-patterns-fsqh8llA-py3.11",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}