import os
from typing import List, Union

import streamlit as st
import graphviz
from langchain.chat_models import ChatOpenAI
from langchain.graphs import Neo4jGraph
from langchain.schema import HumanMessage, AIMessage

from cypher_chain import CYPHER_QA_PROMPT

st.title("VC Chatbot")

from cypher_chain import CustomCypherChain

url = st.secrets["NEO4J_URI"]
username = st.secrets["NEO4J_USERNAME"]
password = st.secrets["NEO4J_PASSWORD"]
database = st.secrets["NEO4J_DATABASE"]

# Langchain x Neo4j connections
graph = Neo4jGraph(username=username, password=password, url=url, database=database)

graph_search = None

# Session state
if "generated" not in st.session_state:
    st.session_state["generated"] = []

if "user_input" not in st.session_state:
    st.session_state["user_input"] = []

if "viz_data" not in st.session_state:
    st.session_state["viz_data"] = []

if "database" not in st.session_state:
    st.session_state["database"] = []

if "cypher" not in st.session_state:
    st.session_state["cypher"] = []


def generate_context(
    prompt: str, context_data: str = "generated"
) -> List[Union[AIMessage, HumanMessage]]:
    context = []
    # If any history exists
    if st.session_state["generated"]:
        # Add the last three exchanges
        size = len(st.session_state["generated"])
        for i in range(max(size - 3, 0), size):
            context.append(HumanMessage(content=st.session_state["user_input"][i]))
            context.append(AIMessage(content=st.session_state[context_data][i]))
    # Add the latest user prompt
    context.append(HumanMessage(content=str(prompt)))
    return context


def dynamic_response_tabs(i):
    tabs_to_add = ["💬Chat"]
    data_check = {
        "🔍Cypher": st.session_state["cypher"][i],
        "🗃️Database results": st.session_state["database"][i],
        "🕸️Visualization": st.session_state["viz_data"][i]
        and st.session_state["viz_data"][i][0],
    }

    for tab_name, has_data in data_check.items():
        if has_data:
            tabs_to_add.append(tab_name)

    with st.chat_message("user"):
        st.write(st.session_state["user_input"][i])

    with st.chat_message("assistant"):
        selected_tabs = st.tabs(tabs_to_add)

        with selected_tabs[0]:
            st.write(st.session_state["generated"][i])
        if len(selected_tabs) > 1:
            with selected_tabs[1]:
                st.code(st.session_state["cypher"][i], language="cypher")
        if len(selected_tabs) > 2:
            with selected_tabs[2]:
                st.write(st.session_state["database"][i])
        if len(selected_tabs) > 3:
            with selected_tabs[3]:
                graph_object = graphviz.Digraph()
                for final_entity in st.session_state["viz_data"][i][1]:
                    graph_object.node(
                        final_entity, fillcolor="lightblue", style="filled"
                    )
                for record in st.session_state["viz_data"][i][0]:
                    graph_object.edge(
                        record["source"], record["target"], label=record["type"]
                    )
                st.graphviz_chart(graph_object)


def get_text() -> str:
    input_text = st.chat_input("Who is the CEO of Neo4j?")
    if not openai_api_key.startswith("sk-"):
        st.warning("Please enter your OpenAI API key!", icon="⚠")
    else:
        return input_text


openai_api_key = st.sidebar.text_input("OpenAI API Key", type="password")
os.environ["OPENAI_API_KEY"] = openai_api_key
if openai_api_key:
    graph_search = CustomCypherChain.from_llm(
        cypher_llm=ChatOpenAI(temperature=0.0, model_name="gpt-4"),
        qa_llm=ChatOpenAI(temperature=0.0),
        graph=graph,
        qa_prompt=CYPHER_QA_PROMPT,
    )
st.sidebar.markdown(
    """
## Example questions

* What do you know about Neo4j organization?
* Who is Neo4j CEO?
* How is Emil Eifrem connected to Magnus Christerson?
* Which company has the most subsidiaries?
* Who are the competitors of databricks?
* How many levels of subsidiaries does Blackstone has?
* Are there any news regarding employee satisfaction?
* What are the latest news around companies where Emil Eifrem is CEO?
* Are there any news about new partnerships mentioning Neo4j?

You can also ask follow up questions as we use a conversational LLM under the hood.

Code is available on [GitHub](https://github.com/tomasonjo/streamlit-neo4j-hackathon)
"""
)

user_input = get_text()

if user_input:
    with st.spinner("Processing"):
        context = generate_context(user_input)
        output = graph_search({"query": user_input, "chat_history": context})

        st.session_state.user_input.append(user_input)
        st.session_state.generated.append(output["result"])
        st.session_state.viz_data.append(output["viz_data"])
        st.session_state.database.append(output["database"])
        st.session_state.cypher.append(output["cypher"])

if st.session_state["generated"]:
    size = len(st.session_state["generated"])
    # Display only the last three exchanges
    for i in range(max(size - 3, 0), size):
        dynamic_response_tabs(i)

src/app.py
tomasonjo/streamlit-neo4j-hackathon
import os

from typing import Any, Dict, List, Optional, Tuple
from langchain.chains.graph_qa.cypher import extract_cypher
from langchain.chains.openai_functions import create_structured_output_chain
from langchain.schema import SystemMessage
from langchain.prompts import ChatPromptTemplate, PromptTemplate

from langchain.chains import GraphCypherQAChain
from langchain.callbacks.manager import CallbackManagerForChainRun

try:
    from pydantic.v1.main import BaseModel, Field
except ImportError:
    from pydantic.main import BaseModel, Field

from cypher_validator import CypherQueryCorrector, Schema


def remove_entities(doc):
    """
    Replace named entities in the given text with their corresponding entity labels.

    Parameters:
    - doc (Spacy Document): processed SpaCy document of the input text.

    Returns:
    - str: The modified text with named entities replaced by their entity labels.

    Example:
    >>> replace_entities_with_labels("Apple is looking at buying U.K. startup for $1 billion.")
    'ORG is looking at buying GPE startup for MONEY .'
    """
    # Initialize an empty list to store the new tokens
    new_tokens = []
    # Keep track of the end index of the last entity
    last_end = 0

    # Iterate through entities, replacing them with their entity label
    for ent in doc.ents:
        # Add the tokens that come before this entity
        new_tokens.extend([token.text for token in doc[last_end : ent.start]])
        # Replace the entity with its label
        new_tokens.append(f"{ent.label_}")
        # Update the last entity end index
        last_end = ent.end

    # Add any remaining tokens after the last entity
    new_tokens.extend([token.text for token in doc[last_end:]])
    # Join the new tokens into a single string
    new_text = " ".join(new_tokens)
    return new_text


AVAILABLE_RELATIONSHIPS = [
    Schema("Person", "HAS_PARENT", "Person"),
    Schema("Person", "HAS_CHILD", "Person"),
    Schema("Organization", "HAS_SUPPLIER", "Organization"),
    Schema("Organization", "IN_CITY", "City"),
    Schema("Organization", "HAS_CATEGORY", "IndustryCategory"),
    Schema("Organization", "HAS_CEO", "Person"),
    Schema("Organization", "HAS_SUBSIDIARY", "Organization"),
    Schema("Organization", "HAS_COMPETITOR", "Organization"),
    Schema("Organization", "HAS_BOARD_MEMBER", "Person"),
    Schema("Organization", "HAS_INVESTOR", "Organization"),
    Schema("Organization", "HAS_INVESTOR", "Person"),
    Schema("City", "IN_COUNTRY", "Country"),
    Schema("Article", "HAS_CHUNK", "Chunk"),
    Schema("Article", "MENTIONS", "Organization")
]

CYPHER_SYSTEM_TEMPLATE = """
Purpose:
Your role is to convert user questions concerning data in a Neo4j database into accurate Cypher queries.
"""

cypher_query_corrector = CypherQueryCorrector(AVAILABLE_RELATIONSHIPS)

CYPHER_QA_TEMPLATE = """You are an assistant that helps to form nice and human understandable answers.
The information part contains the provided information that you must use to construct an answer.
The provided information is authoritative, you must never doubt it or try to use your internal knowledge to correct it.
Make the answer sound as a response to the question. Do not mention that you based the result on the given information.
If the provided information is empty, say that you don't know the answer.
Even if the question doesn't provide full person or organization names, you should use the full names from the provided
information to construct an answer.
Information:
{context}

Question: {question}
Helpful Answer:"""
CYPHER_QA_PROMPT = PromptTemplate(
    input_variables=["context", "question"], template=CYPHER_QA_TEMPLATE
)


class Entities(BaseModel):
    """Identifying information about entities."""

    name: List[str] = Field(
        ...,
        description="All the person, organization, or business entities that appear in the text",
    )


class CustomCypherChain(GraphCypherQAChain):
    def process_entities(self, text: str) -> List[str]:
        prompt = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    "You are extracting organization and person entities from the text.",
                ),
                (
                    "human",
                    "Use the given format to extract information from the following input: {input}",
                ),
            ]
        )

        entity_chain = create_structured_output_chain(
            Entities, self.qa_chain.llm, prompt
        )
        entities = entity_chain.run(text)
        print(entities)
        return entities.name

    def get_viz_data(self, entities: List[str]) -> List[Tuple[str, str]]:
        viz_query = """
        MATCH (n:Person|Organization) WHERE n.name IN $entities
        CALL {
            WITH n
            MATCH (n)-[r:!MENTIONS]->(m)
            WHERE m.name IS NOT NULL
            RETURN n.name AS source, type(r) AS type, m.name AS target
            LIMIT 5
            UNION
            WITH n
            MATCH (n)<-[r:!MENTIONS]-(m)
            WHERE m.name IS NOT NULL
            RETURN n.name AS target, type(r) AS type, m.name AS source
            LIMIT 5
        }
        RETURN source, type, target LIMIT 20
        """
        results = self.graph.query(viz_query, {"entities": entities})
        return results

    def find_entity_match(self, entity: str, k: int = 3) -> List[str]:
        fts_query = """
        CALL db.index.fulltext.queryNodes('entity', $entity + "*", {limit:$k})
        YIELD node,score
        RETURN node.name AS result
        """

        return [
            el["result"]
            for el in self.graph.query(
                fts_query, {"entity": "AND ".join(entity.split()), "k": k}
            )
        ]

    def generate_system_message(
        self, relevant_entities: str = "", fewshot_examples: str = ""
    ) -> SystemMessage:
        system_message = CYPHER_SYSTEM_TEMPLATE
        system_message += (
            f"Database Schema: Please refer to the provided database schema {self.graph_schema} for reference. "
            "Guidelines: Relationships & Properties: Utilize only the relationship types "
            "and properties specified in the provided schema. Do not introduce new ones.\n"
        )
        if relevant_entities:
            system_message += (
                f"Entity Substitution: If the question mentions specific entities, replace them in the query with corresponding entities from "
                f"the given list. Given list of entities is: {relevant_entities}\n"
                "Example: If the list contains John: ['John Goodman'], replace 'John' in the query with 'John Goodman'."
                "Flexible Queries: Ensure your Cypher queries can capture all relevant entities.\n"
                "Correct: MATCH (p:Person) WHERE p.name IN ['John Goodman', 'John Stockton'] MATCH (p)<-[:BOARD_MEMBER]-(o:Organization)"
            )
        if fewshot_examples:
            system_message += (
                f"Example Queries: Please refer to the provided example queries for constructing Cypher statements:\n"
                f"{fewshot_examples}\n"
            )

        system_message += (
            "Informative Responses: Ensure your responses contain sufficient information to answer the question for anyone lacking context.\n"
            "Include details about both persons and organizations, if applicable. \n"
            "Text Search: To search text, avoid using the CONTAINS clause. "
            "Instead, employ apoc.ml.openai.embedding and gds.similarity.cosine functions or db.index.vector.queryNodes.\n"
            "Note: If you use apoc.ml.openai.embedding, it must be the first operation in your query.\n"
            "Result Chunking: When returning text, always provide exactly three chunks—no more, no less.\n"
            "By following these guidelines, you'll ensure the generated Cypher queries accurately reflect the database entities and relationships."
            "Solve it step by step."
        )
        return SystemMessage(content=system_message)

    def get_fewshot_examples(self, question):
        results = self.graph.query(
            """
        CALL apoc.ml.openai.embedding([$question], $openai_api_key)
                                    YIELD embedding                             
        CALL db.index.vector.queryNodes('fewshot', 3, embedding)
                                    YIELD node, score
        RETURN node.Question AS question, node.Cypher as cypher
                                    """,
            {"question": question, "openai_api_key": os.environ["OPENAI_API_KEY"]},
        )

        fewshot = "\n".join([f"#{el['question']}\n{el['cypher']}" for el in results])
        print("-" * 30)
        print(fewshot)
        return fewshot

    def _call(
        self,
        inputs: Dict[str, Any],
        run_manager: Optional[CallbackManagerForChainRun] = None,
    ) -> Dict[str, Any]:

        _run_manager = run_manager or CallbackManagerForChainRun.get_noop_manager()
        callbacks = _run_manager.get_child()
        question = inputs[self.input_key]
        chat_history = inputs["chat_history"]
        # Extract mentioned people and organizations and match them to database values
        entities = self.process_entities(question)
        print(f"NER found: {entities}")
        relevant_entities = dict()
        for entity in entities:
            relevant_entities[entity] = self.find_entity_match(entity)
        print(f"Relevant entities are: {relevant_entities}")

        # Get few-shot examples using vector search
        fewshots = self.get_fewshot_examples(question)

        system = self.generate_system_message(str(relevant_entities), fewshots)
        generated_cypher = self.cypher_generation_chain.llm.predict_messages(
            [system] + chat_history
        )
        print(generated_cypher.content)
        generated_cypher = extract_cypher(generated_cypher.content)
        validated_cypher = cypher_query_corrector(
            generated_cypher
        )
        print(validated_cypher)
        # If Cypher statement wasn't generated
        # Usually happens when LLM decides it can't answer
        if not "RETURN" in validated_cypher:
            chain_result: Dict[str, Any] = {
                self.output_key: validated_cypher,
                "viz_data": (None, None),
                "database": None,
                "cypher": None,
            }
            return chain_result

        # Retrieve and limit the number of results
        context = self.graph.query(
            validated_cypher, {"openai_api_key": os.environ["OPENAI_API_KEY"]}
        )[: self.top_k]

        result = self.qa_chain(
            {"question": question, "context": context}, callbacks=callbacks
        )
        final_result = result[self.qa_chain.output_key]

        final_entities = self.process_entities(final_result)
        if final_entities:
            viz_data = self.get_viz_data(final_entities)
        else:
            viz_data = None

        chain_result: Dict[str, Any] = {
            self.output_key: final_result,
            "viz_data": (viz_data, final_entities),
            "database": context,
            "cypher": validated_cypher,
        }
        return chain_result

src/cypher_chain.py
tomasonjo/streamlit-neo4j-hackathon
import re
from collections import namedtuple
from typing import Any, Dict, List, Optional, Tuple

Schema = namedtuple("Schema", ["left_node", "relation", "right_node"])


class CypherQueryCorrector:
    """
    Used to correct relationship direction in generated Cypher statements.
    This code is copied from the winner's submission to the Cypher competition:
    https://github.com/sakusaku-rich/cypher-direction-competition
    """

    property_pattern = re.compile(r"\{.+?\}")
    node_pattern = re.compile(r"\(.+?\)")
    path_pattern = re.compile(r"\(.*\).*-.*-.*\(.*\)")
    node_relation_node_pattern = re.compile(
        r"(\()+(?P<left_node>[^()]*?)\)(?P<relation>.*?)\((?P<right_node>[^()]*?)(\))+"
    )
    relation_type_pattern = re.compile(r":(?P<relation_type>.+?)?(\{.+\})?]")

    def __init__(self, schemas: List[Schema]):
        """
        Args:
            schemas: list of schemas
        """
        self.schemas = schemas

    def clean_node(self, node: str) -> str:
        """
        Args:
            node: node in string format
        """
        node = re.sub(self.property_pattern, "", node)
        node = node.replace("(", "")
        node = node.replace(")", "")
        node = node.strip()
        return node

    def detect_node_variables(self, query: str) -> Dict[str, List[str]]:
        """
        Args:
            query: cypher query
        """
        nodes = re.findall(self.node_pattern, query)
        nodes = [self.clean_node(node) for node in nodes]
        res: Dict[str, Any] = {}
        for node in nodes:
            parts = node.split(":")
            if parts == "":
                continue
            variable = parts[0]
            if variable not in res:
                res[variable] = []
            res[variable] += parts[1:]
        return res

    def extract_paths(self, query: str) -> "List[str]":
        """
        Args:
            query: cypher query
        """
        return re.findall(self.path_pattern, query)

    def judge_direction(self, relation: str) -> str:
        """
        Args:
            relation: relation in string format
        """
        direction = "BIDIRECTIONAL"
        if relation[0] == "<":
            direction = "INCOMING"
        if relation[-1] == ">":
            direction = "OUTGOING"
        return direction

    def extract_node_variable(self, part: str) -> Optional[str]:
        """
        Args:
            part: node in string format
        """
        part = part.lstrip("(").rstrip(")")
        idx = part.find(":")
        if idx != -1:
            part = part[:idx]
        return None if part == "" else part

    def detect_labels(
        self, str_node: str, node_variable_dict: Dict[str, Any]
    ) -> List[str]:
        """
        Args:
            str_node: node in string format
            node_variable_dict: dictionary of node variables
        """
        splitted_node = str_node.split(":")
        variable = splitted_node[0]
        labels = []
        if variable in node_variable_dict:
            labels = node_variable_dict[variable]
        elif variable == "" and len(splitted_node) > 1:
            labels = splitted_node[1:]
        return labels

    def verify_schema(
        self,
        from_node_labels: List[str],
        relation_types: List[str],
        to_node_labels: List[str],
    ) -> bool:
        """
        Args:
            from_node_labels: labels of the from node
            relation_type: type of the relation
            to_node_labels: labels of the to node
        """
        valid_schemas = self.schemas
        if from_node_labels != []:
            from_node_labels = [label.strip("`") for label in from_node_labels]
            valid_schemas = [
                schema for schema in valid_schemas if schema[0] in from_node_labels
            ]
        if to_node_labels != []:
            to_node_labels = [label.strip("`") for label in to_node_labels]
            valid_schemas = [
                schema for schema in valid_schemas if schema[2] in to_node_labels
            ]
        if relation_types != []:
            relation_types = [type.strip("`") for type in relation_types]
            valid_schemas = [
                schema for schema in valid_schemas if schema[1] in relation_types
            ]
        return valid_schemas != []

    def detect_relation_types(self, str_relation: str) -> Tuple[str, List[str]]:
        """
        Args:
            str_relation: relation in string format
        """
        relation_direction = self.judge_direction(str_relation)
        relation_type = self.relation_type_pattern.search(str_relation)
        if relation_type is None or relation_type.group("relation_type") is None:
            return relation_direction, []
        relation_types = [
            t.strip().strip("!")
            for t in relation_type.group("relation_type").split("|")
        ]
        return relation_direction, relation_types

    def correct_query(self, query: str) -> str:
        """
        Args:
            query: cypher query
        """
        node_variable_dict = self.detect_node_variables(query)
        paths = self.extract_paths(query)
        for path in paths:
            original_path = path
            start_idx = 0
            while start_idx < len(path):
                match_res = re.match(self.node_relation_node_pattern, path[start_idx:])
                if match_res is None:
                    break
                start_idx += match_res.start()
                match_dict = match_res.groupdict()
                left_node_labels = self.detect_labels(
                    match_dict["left_node"], node_variable_dict
                )
                right_node_labels = self.detect_labels(
                    match_dict["right_node"], node_variable_dict
                )
                end_idx = (
                    start_idx
                    + 4
                    + len(match_dict["left_node"])
                    + len(match_dict["relation"])
                    + len(match_dict["right_node"])
                )
                original_partial_path = original_path[start_idx : end_idx + 1]
                relation_direction, relation_types = self.detect_relation_types(
                    match_dict["relation"]
                )

                if relation_types != [] and "".join(relation_types).find("*") != -1:
                    start_idx += (
                        len(match_dict["left_node"]) + len(match_dict["relation"]) + 2
                    )
                    continue

                if relation_direction == "OUTGOING":
                    is_legal = self.verify_schema(
                        left_node_labels, relation_types, right_node_labels
                    )
                    if not is_legal:
                        is_legal = self.verify_schema(
                            right_node_labels, relation_types, left_node_labels
                        )
                        if is_legal:
                            corrected_relation = "<" + match_dict["relation"][:-1]
                            corrected_partial_path = original_partial_path.replace(
                                match_dict["relation"], corrected_relation
                            )
                            query = query.replace(
                                original_partial_path, corrected_partial_path
                            )
                        else:
                            return ""
                elif relation_direction == "INCOMING":
                    is_legal = self.verify_schema(
                        right_node_labels, relation_types, left_node_labels
                    )
                    if not is_legal:
                        is_legal = self.verify_schema(
                            left_node_labels, relation_types, right_node_labels
                        )
                        if is_legal:
                            corrected_relation = match_dict["relation"][1:] + ">"
                            corrected_partial_path = original_partial_path.replace(
                                match_dict["relation"], corrected_relation
                            )
                            query = query.replace(
                                original_partial_path, corrected_partial_path
                            )
                        else:
                            return ""
                else:
                    is_legal = self.verify_schema(
                        left_node_labels, relation_types, right_node_labels
                    )
                    is_legal |= self.verify_schema(
                        right_node_labels, relation_types, left_node_labels
                    )
                    if not is_legal:
                        return ""

                start_idx += (
                    len(match_dict["left_node"]) + len(match_dict["relation"]) + 2
                )
        return query

    def __call__(self, query: str) -> str:
        """Correct the query to make it valid. If
        Args:
            query: cypher query
        """
        return self.correct_query(query)
src/cypher_validator.py
tomasonjo/streamlit-neo4j-hackathon