nqtruong/Job_Knowledge_Graph

1import os2import yaml3from dotenv import load_dotenv4from langchain_google_genai import ChatGoogleGenerativeAI5from langchain_community.graphs import Neo4jGraph6from langchain_core.prompts.prompt import PromptTemplate7from langchain.chains import GraphCypherQAChain8from langchain_core.messages import SystemMessage, HumanMessage, AIMessage9 10 11 12def config():13    load_dotenv()14    15    # Set up Neo4J & Gemini API16    os.environ["NEO4J_URI"] = os.getenv("NEO4J_URI")17    os.environ["NEO4J_USERNAME"] = os.getenv("NEO4J_USERNAME")18    os.environ["NEO4J_PASSWORD"] = os.getenv("NEO4J_PASSWORD")19    os.environ["GOOGLE_API_KEY"] = os.getenv("GEMINI_API_KEY")20 21def load_prompt(filepath):22    with open(filepath, "r") as file:23        prompt = yaml.safe_load(file)24 25    return prompt26 27def init_():28    config()29    graph = Neo4jGraph()30    llm = ChatGoogleGenerativeAI(31        model= "gemini-1.5-flash-latest"32    )33 34    return graph, llm35 36def get_llm_response(query):37    # Connect to Neo4J Knowledge Graph38    knowledge_graph, llm_chat = init_()39    cypher_prompt = load_prompt("Agent/prompts/cypher_prompt.yaml")40    qa_prompt = load_prompt("Agent/prompts/qa_prompt.yaml")41 42    CYPHER_GENERATION_PROMPT = PromptTemplate(**cypher_prompt)43    QA_GENERATION_PROMPT = PromptTemplate(**qa_prompt)44 45    chain = GraphCypherQAChain.from_llm(46        llm_chat, graph=knowledge_graph, verbose=True,47        cypher_prompt= CYPHER_GENERATION_PROMPT,48        qa_prompt= QA_GENERATION_PROMPT49    )50 51    return chain.invoke({"query": query})["result"]52 53def llm_answer(message, history):54 55 56    try:57        response = get_llm_response(message["text"])58    except Exception:59        response = "Exception"60    except Error:61        response = "Error"62    return response63 64# if __name__ == "__main__":65#     message = "Have any company recruiting jobs about Machine Learning and coresponding job titles?"66#     history = [("What's your name?", "My name is Gemini")]67#     resp = llm_answer(message, history)68#     print(resp)