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 10def config():11    load_dotenv()12 13    # Set up Neo4J & Gemini API14    os.environ["NEO4J_URI"] = os.getenv("NEO4J_URI")15    os.environ["NEO4J_USERNAME"] = os.getenv("NEO4J_USERNAME")16    os.environ["NEO4J_PASSWORD"] = os.getenv("NEO4J_PASSWORD")17    os.environ["GOOGLE_API_KEY"] = os.getenv("GEMINI_API_KEY")18 19def load_prompt(filepath):20    with open(filepath, "r") as file:21        prompt = yaml.safe_load(file)22 23    return prompt24 25def init_():26    config()27    graph = Neo4jGraph(enhanced_schema= True)28    llm = ChatGoogleGenerativeAI(29        model= "gemini-1.5-flash-latest",30        temperature = 031    )32 33    return graph, llm34 35def get_llm_response(query):36    # Connect to Neo4J Knowledge Graph37    knowledge_graph, llm_chat = init_()38    cypher_prompt = load_prompt("Agent/prompts/cypher_prompt.yaml")39    qa_prompt = load_prompt("Agent/prompts/qa_prompt.yaml")40 41    CYPHER_GENERATION_PROMPT = PromptTemplate(**cypher_prompt)42    QA_GENERATION_PROMPT = PromptTemplate(**qa_prompt)43 44    chain = GraphCypherQAChain.from_llm(45        llm_chat, graph=knowledge_graph, verbose=True,46        cypher_prompt= CYPHER_GENERATION_PROMPT,47        qa_prompt= QA_GENERATION_PROMPT48    )49 50    return chain.invoke({"query": query})["result"]51 52def llm_answer(message, history):53 54 55    try:56        response = get_llm_response(message["text"])57    except Exception:58        response = "Exception"59    except Error:60        response = "Error"61    return response62 63# if __name__ == "__main__":64#     message = "Have any company recruiting jobs about Machine Learning and coresponding job titles?"65#     history = [("What's your name?", "My name is Gemini")]66#     resp = llm_answer(message, history)67#     print(resp)