Create agent.py

agent.py

@@ -0,0 +1,131 @@
+from langgraph.graph import START, StateGraph, MessagesState
+from langgraph.prebuilt import tools_condition
+from langgraph.prebuilt import ToolNode
+
+from langchain_community.tools.tavily_search import TavilySearchResults
+from langchain_community.document_loaders import WikipediaLoader
+from langchain_community.document_loaders import ArxivLoader
+
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_core.tools import tool
+
+from langchain_google_genai import ChatGoogleGenerativeAI
+
+#load_dotenv()
+
+@tool
+def add(a: int, b: int) -> int:
+  """ Add a and b """
+  return a + b
+@tool
+def subtract(a: int,b: int) -> int:
+  """ Subract b from a """
+  return a - b
+@tool
+def multiply(a: int,b: int) -> int:
+  """ Multiply a and b """
+  return a * b
+@tool
+def divide(a: int,b: int) -> float:
+  """ Divide a by b """
+  if b == 0:
+    raise ValueError("Can't divide by 0.")
+  return a/b
+
+@tool
+def web_search(query: str) -> str:
+  """ Search for a query on web and return best 2 result."""
+
+  search_results = TavilySearchResults(max_results = 3).invoke(input=query)
+
+  formatted_search_results = "\n\n-----\n\n".join(
+      [
+          #f'<Result: source = "{result.metadata["source"]}", page = "{result.metadata.get("page","")}">\n{result.page_content}\n </Result>'
+          f'<Result: source = "{result.get("url", "")}", page = "{result.get("title","")}">\n{result.get("content","")}\n </Result>'
+          for result in search_results
+      ]
+  )
+  return {"web_results" : formatted_search_results}
+
+@tool
+def wikipedia_search(query: str) -> str:
+  """ Search for a query on wikipedia and return best 2 result."""
+
+  loader = WikipediaLoader(query=query, load_max_docs=2)
+  search_results = loader.load() # Now, just call load() without arguments
+
+  formatted_search_results = "\n\n-----\n\n".join(
+      [
+          # Each 'result' here is a Document object.
+          # Access metadata through .metadata and content through .page_content
+          f'<Result: source = "{result.metadata.get("source", "")}", page = "{result.metadata.get("title","")}">\n{result.page_content}\n </Result>'
+          for result in search_results
+      ]
+  )
+  return {"Wikipedia_results" : formatted_search_results}
+
+@tool
+def arxiv_search(query: str) -> str:
+  """ Search for a query on arxiv and return best 2 result."""
+
+  # Similar to WikipediaLoader, query and load_max_docs are passed during initialization
+  loader = ArxivLoader(query=query, load_max_docs=2)
+  search_results = loader.load() # Call load() without arguments
+
+  formatted_search_results = "\n\n-----\n\n".join(
+      [
+          f'<Result: source = "{result.metadata.get("source", "")}", page = "{result.metadata.get("title","")}">\n{result.page_content}\n </Result>'
+          for result in search_results
+      ]
+  )
+  return {"arxiv_results" : formatted_search_results}
+
+
+
+system_prompt = """You are a general AI assistant. I will ask you a question. Using your tools to report your thoughts, and finish your answer with the following template:
+FINAL ANSWER: [YOUR FINAL ANSWER]. YOUR FINAL ANSWER should be a number OR as few words as possible OR a comma separated list of numbers and/or strings.
+If you are asked for a number, don't use comma to write your number neither use units such as $ or percent sign unless specified otherwise.
+If you are asked for a string, don't use articles, neither abbreviations (e.g. for cities), and write the digits in plain text unless specified otherwise.
+If you are asked for a comma separated list, apply the above rules depending of whether the element to be put in the list is a number or a string."""
+#Using your tools to 추가하니 툴컬링 하게됨
+system_message = SystemMessage(content=system_prompt)
+
+tools = [
+    add,subtract,multiply,divide,web_search,wikipedia_search,arxiv_search
+]
+
+
+def build_graph(provider: str = "google"):
+  llm = ChatGoogleGenerativeAI(model="gemini-2.0-flash", temperature=0 , api_key = "")
+
+  llm_with_tools = llm.bind_tools(tools)
+
+  def assistant(state: MessagesState):
+    """ Use the tools to answer the query. you have add,subtract,multiply,divide,web_search,wikipedia_search,arxiv_search tools."""
+    return {"messages": state["messages"] + [llm_with_tools.invoke([system_message]+state["messages"])]}
+
+
+  builder = StateGraph(MessagesState)
+  builder.add_node("assistant", assistant)
+  builder.add_node("tools", ToolNode(tools))
+  builder.add_edge(START, "assistant")
+  builder.add_conditional_edges(
+      "assistant",
+      tools_condition
+  )
+  builder.add_edge("tools", "assistant")
+
+  return builder.compile()
+
+# test
+if __name__ == "__main__":
+    question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
+    # Build the graph
+    graph = build_graph(provider="google")
+    # Run the graph
+    messages = [HumanMessage(content=question)]
+    messages = graph.invoke({"messages": messages})
+    for m in messages["messages"]:
+        m.pretty_print()
+
+