Parthiban007 commited on
Commit
a11eeb8
·
verified ·
1 Parent(s): 81917a3

Upload 3 files

Browse files
Files changed (3) hide show
  1. .env +2 -0
  2. .gitignore +1 -0
  3. llm_agent.py +174 -0
.env ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ GOOGLE_API_KEY = "AIzaSyDB2kx888D7HUuya3vvb0wvHzDUP2BeC7M"
2
+ TAVILY_API_KEY = "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"
.gitignore ADDED
@@ -0,0 +1 @@
 
 
1
+ .env
llm_agent.py ADDED
@@ -0,0 +1,174 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ from dotenv import load_dotenv
3
+ from langgraph.graph import START, StateGraph, MessagesState
4
+ from langgraph.prebuilt import tools_condition
5
+ from langgraph.prebuilt import ToolNode
6
+ from langchain_google_genai import ChatGoogleGenerativeAI,GoogleGenerativeAIEmbeddings
7
+ from langchain_community.document_loaders import WikipediaLoader
8
+ from langchain_community.document_loaders import ArxivLoader
9
+ from langchain_community.tools.tavily_search import TavilySearchResults
10
+ from langchain_community.vectorstores import FAISS
11
+ from langchain_core.messages import SystemMessage,HumanMessage
12
+ from langchain_core.tools import tool
13
+
14
+
15
+ load_dotenv()
16
+
17
+ os.environ["GOOGLE_API_KEY"] = os.getenv("GOOGLE_API_KEY")
18
+
19
+ @tool
20
+ def add(a:int,b:int)->int:
21
+ """Add two Numbers
22
+ Args:
23
+ a:int
24
+ b:int
25
+
26
+ """
27
+ return a+b
28
+
29
+ @tool
30
+ def subtract(a:int,b:int)->int:
31
+ """Subtract two numbers
32
+ Args:
33
+ a:int
34
+ b:int
35
+ """
36
+ return a-b
37
+
38
+ @tool
39
+ def multiply(a:int,b:int)->int:
40
+ """Multiply Two Numbers
41
+ Args:
42
+ a:int
43
+ b:int
44
+ """
45
+ return a*b
46
+
47
+ @tool
48
+ def divide(a:int,b:int)->int:
49
+ """Divide two numbers
50
+ Args:
51
+ a:int
52
+ b:int
53
+ """
54
+ if b==0:
55
+ raise ValueError("Cannot Divide by Zero")
56
+ return a//b
57
+
58
+ @tool
59
+ def modulus(a:int,b:int)->int:
60
+ """Modulus of the two numbers
61
+ Args:
62
+ a:int
63
+ b:int
64
+ """
65
+ return a%b
66
+
67
+
68
+ @tool
69
+ def wiki_search(query:str)->str:
70
+ """Search Wikipedia for a query and return maximum 2 results
71
+ Args:
72
+ query: The Search Query : str
73
+ """
74
+ print(query)
75
+
76
+ search_docs = WikipediaLoader(query=query,load_max_docs=2).load()
77
+
78
+ return {"wiki_results": search_docs}
79
+
80
+
81
+ @tool
82
+ def web_search(query:str)->str:
83
+ """ search Tavily for a query and return maximum 3 results
84
+
85
+ Args:
86
+ query: The Search Query
87
+ """
88
+ search_docs = TavilySearchResults(max_results=3).invoke(input=query)
89
+
90
+ return {"web_results": search_docs}
91
+
92
+ @tool
93
+ def arxiv_search(query:str)->str:
94
+ """"search Arxiv for a query and return maximum 3 results
95
+ Args:
96
+ query: search query
97
+ """
98
+
99
+ search_docs = ArxivLoader(query=query,load_max_docs = 3).load()
100
+
101
+ return {"arxiv_resutls": search_docs}
102
+
103
+
104
+ system_prompt = """
105
+ You are a helpful assistant tasked with answering questions using a set of tools.
106
+ Now, I will ask you a question. Report your thoughts, and finish your answer with the following template:
107
+ FINAL ANSWER: [YOUR FINAL ANSWER].
108
+ 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.
109
+ Your answer should only start with "FINAL ANSWER: ", then follows with the answer.
110
+ """
111
+
112
+ sys_msg = SystemMessage(content=system_prompt)
113
+
114
+
115
+ embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
116
+
117
+
118
+
119
+ tools = [
120
+ add,
121
+ subtract,
122
+ multiply,
123
+ divide,
124
+ modulus,
125
+ wiki_search,
126
+ web_search,
127
+ arxiv_search,
128
+ ]
129
+
130
+
131
+
132
+ def build_graph():
133
+ llm = ChatGoogleGenerativeAI(model = "gemini-2.0-flash")
134
+ print(tools)
135
+ # Bind tools to LLM
136
+ llm_with_tools = llm.bind_tools(tools)
137
+
138
+ # Node
139
+ def assistant(state: MessagesState):
140
+ """Assistant node"""
141
+ return {"messages": [llm_with_tools.invoke(state["messages"])]}
142
+
143
+ # def retriever(state: MessagesState):
144
+ # """Retriever node"""
145
+ # similar_question = vector_store.similarity_search(state["messages"][0].content)
146
+ # example_msg = HumanMessage(
147
+ # content=f"Here I provide a similar question and answer for reference: \n\n{similar_question}",
148
+ # )
149
+ # return {"messages": [sys_msg] + state["messages"] + [example_msg]}
150
+
151
+ builder = StateGraph(MessagesState)
152
+ # builder.add_node("retriever", retriever)
153
+ builder.add_node("assistant", assistant)
154
+ builder.add_node("tools", ToolNode(tools))
155
+ builder.add_edge(START, "assistant")
156
+ # builder.add_edge("retriever", "assistant")
157
+ builder.add_conditional_edges(
158
+ "assistant",
159
+ tools_condition,
160
+ )
161
+ builder.add_edge("tools", "assistant")
162
+
163
+ # Compile graph
164
+ return builder.compile()
165
+
166
+
167
+ if __name__ == "__main__":
168
+ question = "When was a picture of St. Thomas Aquinas first added to the Wikipedia page on the Principle of double effect?"
169
+ graph = build_graph()
170
+ # Run the graph
171
+ messages = [HumanMessage(content=question)]
172
+ messages = graph.invoke({"messages": messages})
173
+ for m in messages["messages"]:
174
+ m.pretty_print()