Ame42/GrapeFruit

1from langchain_core.runnables import (2    RunnableBranch,3    RunnableLambda,4    RunnableParallel,5    RunnablePassthrough,6)7from langchain_core.prompts import ChatPromptTemplate8from langchain_core.prompts.prompt import PromptTemplate9from langchain_core.pydantic_v1 import BaseModel, Field10from typing import Tuple, List, Optional11from langchain_core.messages import AIMessage, HumanMessage12from langchain_core.output_parsers import StrOutputParser13import os14from langchain_community.graphs import Neo4jGraph15from langchain.document_loaders import WikipediaLoader16from langchain.text_splitter import TokenTextSplitter17from langchain_openai import ChatOpenAI18from langchain_experimental.graph_transformers import LLMGraphTransformer19from neo4j import GraphDatabase20from yfiles_jupyter_graphs import GraphWidget21from langchain_community.vectorstores import Neo4jVector22from langchain_openai import OpenAIEmbeddings23from langchain_community.vectorstores.neo4j_vector import remove_lucene_chars24from langchain_core.runnables import ConfigurableField, RunnableParallel, RunnablePassthrough25from langchain_community.document_loaders import PyPDFLoader26from langchain_community.vectorstores import FAISS27from langchain_openai import OpenAIEmbeddings28 29import os30 31 32# Extract entities from text33class Entities(BaseModel):34    """Identifying information about entities."""35 36    names: List[str] = Field(37        ...,38        description="All the person, organization, or business entities that "39        "appear in the text",40    )41 42 43def _format_chat_history(chat_history: List[Tuple[str, str]]) -> List:44    buffer = []45    for human, ai in chat_history:46        buffer.append(HumanMessage(content=human))47        buffer.append(AIMessage(content=ai))48    return buffer49 50 51def generate_full_text_query(input: str) -> str:52    """53    Generate a full-text search query for a given input string.54 55    This function constructs a query string suitable for a full-text search.56    It processes the input string by splitting it into words and appending a57    similarity threshold (~2 changed characters) to each word, then combines58    them using the AND operator. Useful for mapping entities from user questions59    to database values, and allows for some misspelings.60    """61    full_text_query = ""62    words = [el for el in remove_lucene_chars(input).split() if el]63    for word in words[:-1]:64        full_text_query += f" {word}~2 AND"65    full_text_query += f" {words[-1]}~2"66    return full_text_query.strip()67 68 69# Fulltext index query70def structured_retriever(question: str) -> str:71    """72    Collects the neighborhood of entities mentioned73    in the question74    """75    result = ""76    entities = entity_chain.invoke({"question": question})77    for entity in entities.names:78        response = graph.query(79            """CALL db.index.fulltext.queryNodes('entity', $query, {limit:2})80            YIELD node,score81            CALL {82              WITH node83              MATCH (node)-[r:!MENTIONS]->(neighbor)84              RETURN node.id + ' - ' + type(r) + ' -> ' + neighbor.id AS output85              UNION ALL86              WITH node87              MATCH (node)<-[r:!MENTIONS]-(neighbor)88              RETURN neighbor.id + ' - ' + type(r) + ' -> ' +  node.id AS output89            }90            RETURN output LIMIT 5091            """,92            {"query": generate_full_text_query(entity)},93        )94        result += "\n".join([el['output'] for el in response])95    return result96 97 98#Final retrieval99def retriever(question: str):100    print(f"Search query: {question}")101    structured_data = structured_retriever(question)102    unstructured_data = [el.page_content for el in vector_index.similarity_search(question)]103    final_data = f"""Structured data:104{structured_data}105Unstructured data:106{"#Document ". join(unstructured_data)}107    """108    return final_data109 110 111def query(query):112    chain = (RunnableParallel(113            {114                "context": _search_query | retriever,115                "question": RunnablePassthrough(),116            }117        )118        | prompt119        | llm120        | StrOutputParser()121    )122    return chain.invoke(query)123 124 125os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_API_KEY')126os.environ["NEO4J_URI"] = os.getenv('neo4j_url')127os.environ["NEO4J_USERNAME"] = os.getenv('neo4j_username')128os.environ["NEO4J_PASSWORD"] = os.getenv('neo4j_password')129 130graph = Neo4jGraph()131llm=ChatOpenAI(temperature=0, model_name="gpt-3.5-turbo-0125", openai_api_key = os.getenv('OPENAI_API_KEY')) # gpt-4-0125-preview occasionally has issues132 133vector_index = Neo4jVector.from_existing_graph(134    OpenAIEmbeddings(),135    search_type="hybrid",136    node_label="Document",137    text_node_properties=["text"],138    embedding_node_property="embedding"139)140 141URI = os.getenv('neo4j_url')142AUTH = (os.getenv('neo4j_username'), os.getenv('neo4j_password'))143 144graph_driver = GraphDatabase.driver(URI, auth=AUTH)145 146with graph_driver as driver:147    driver.verify_connectivity()148 149# Retriever150graph.query(151    "CREATE FULLTEXT INDEX entity IF NOT EXISTS FOR (e:__Entity__) ON EACH [e.id]")152 153mes_prompt = ChatPromptTemplate.from_messages(154    [155        (156            "system",157            "You are extracting organization and person entities from the text.",158        ),159        (160            "human",161            "Use the given format to extract information from the following "162            "input: {question}",163        ),164    ]165)166 167entity_chain = mes_prompt | llm.with_structured_output(Entities)168 169_template = """Given the following conversation and a follow up question, rephrase the follow up question to be a standalone question,170in its original language.171Chat History:172{chat_history}173Follow Up Input: {question}174Standalone question:"""  # noqa: E501175CONDENSE_QUESTION_PROMPT = PromptTemplate.from_template(_template)176 177template = """Answer the question based only on the following context:178{context}179 180Question: {question}181Use natural language and be concise.182Answer:"""183 184prompt = ChatPromptTemplate.from_template(template)185 186_search_query = RunnableBranch(187    # If input includes chat_history, we condense it with the follow-up question188    (189        RunnableLambda(lambda x: bool(x.get("chat_history"))).with_config(190            run_name="HasChatHistoryCheck"191        ),  # Condense follow-up question and chat into a standalone_question192        RunnablePassthrough.assign(193            chat_history=lambda x: _format_chat_history(x["chat_history"])194        )195        | CONDENSE_QUESTION_PROMPT196        | ChatOpenAI(temperature=0)197        | StrOutputParser(),198    ),199    # Else, we have no chat history, so just pass through the question200    RunnableLambda(lambda x : x["question"]),201)202