nqtruong/Job_Knowledge_Graph

1from langchain.agents import Tool, AgentType, initialize_agent2from langchain.memory import ConversationBufferMemory3from langchain_google_genai import ChatGoogleGenerativeAI4from langchain.agents import AgentExecutor5from langchain import hub6from langchain.agents.format_scratchpad import format_log_to_str7from langchain.agents.output_parsers import ReActSingleInputOutputParser8from langchain.tools.render import render_text_description9import os10from Agent.tools.kg_search import lookup_kg11from Agent.tools.tavily_search_v2 import tavily_search, tavily_qna_search12 13from dotenv import load_dotenv14from langchain.agents import Tool15from langchain_core.prompts import PromptTemplate16 17load_dotenv()18os.environ["GOOGLE_API_KEY"] = os.getenv("GEMINI_API_KEY")19llm = ChatGoogleGenerativeAI(20    model= "gemini-1.5-flash-latest",21    temperature = 022)23 24 25 26kg_query = Tool(27    name = 'Query Knowledge Graph',28    func = lookup_kg,29    description='Useful for when you need to answer questions about job posts.'30)31 32 33web_search = Tool(34    name = 'Web Search',35    func = tavily_qna_search,36    description = "Useful for when you need to search for external information."37)38 39tools = [kg_query, web_search]40 41 42with open("Agent/prompts/react_prompt_v2.txt", "r") as file:43    react_template = file.read()44 45react_prompt = PromptTemplate(46    input_variables = ["tools", "tool_names", "input", "agent_scratchpad", "chat_history"],47    template = react_template48)49 50prompt = react_prompt.partial(51    tools = render_text_description(tools),52    tool_names = ", ".join([t.name for t in tools]),53)54 55llm_with_stop = llm.bind(stop=["\nObservation"])56 57agent = (58    {59        "input": lambda x: x["input"],60        "agent_scratchpad": lambda x: format_log_to_str(x["intermediate_steps"]),61        "chat_history": lambda x: x["chat_history"],62    }63    | prompt64    | llm_with_stop65    | ReActSingleInputOutputParser()66)67 68memory = ConversationBufferMemory(memory_key="chat_history")69 70agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True, memory=memory)71 72 73def get_react_agent(memory):74    agent_executor = AgentExecutor(75        agent=agent,76        tools=tools,77        verbose=True,78        memory=memory79    )80 81    return agent_executor82 83 84 85# if __name__ == "__main__":86#     while True:87#         try:88#             question = input("> ")89#             result = agent_executor.invoke({90#                 "input": question91#             })92#         except:93#             break