Shaukat39 commited on
Commit
d7a1a8c
·
verified ·
1 Parent(s): 15479b8

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +30 -0
agent.py CHANGED
@@ -70,19 +70,49 @@ def modulus(a: int, b: int) -> str:
70
  # ------------------ Retrieval Tools ------------------
71
  @tool
72
  def wiki_search(query: str) -> str:
 
 
 
 
 
 
 
 
 
73
  docs = WikipediaLoader(query=query, load_max_docs=2).load()
74
  return "\n\n---\n\n".join(doc.page_content for doc in docs)
75
 
 
76
  @tool
77
  def web_search(query: str) -> str:
 
 
 
 
 
 
 
 
 
78
  docs = TavilySearchResults(max_results=3).invoke(query)
79
  return "\n\n---\n\n".join(doc.page_content for doc in docs)
80
 
 
81
  @tool
82
  def arvix_search(query: str) -> str:
 
 
 
 
 
 
 
 
 
83
  docs = ArxivLoader(query=query, load_max_docs=3).load()
84
  return "\n\n---\n\n".join(doc.page_content[:1000] for doc in docs)
85
 
 
86
  # ------------------ System Setup ------------------
87
  with open("system_prompt.txt", "r", encoding="utf-8") as f:
88
  system_prompt = f.read()
 
70
  # ------------------ Retrieval Tools ------------------
71
  @tool
72
  def wiki_search(query: str) -> str:
73
+ """
74
+ Search Wikipedia for a given query and return the content of up to two matching articles.
75
+
76
+ Args:
77
+ query: A string query to search on Wikipedia.
78
+
79
+ Returns:
80
+ A string containing the content from up to two relevant Wikipedia articles, separated by dividers.
81
+ """
82
  docs = WikipediaLoader(query=query, load_max_docs=2).load()
83
  return "\n\n---\n\n".join(doc.page_content for doc in docs)
84
 
85
+
86
  @tool
87
  def web_search(query: str) -> str:
88
+ """
89
+ Perform a web search using Tavily and return content from the top three results.
90
+
91
+ Args:
92
+ query: A string query representing the web search topic.
93
+
94
+ Returns:
95
+ A string of up to three relevant result contents, separated by dividers.
96
+ """
97
  docs = TavilySearchResults(max_results=3).invoke(query)
98
  return "\n\n---\n\n".join(doc.page_content for doc in docs)
99
 
100
+
101
  @tool
102
  def arvix_search(query: str) -> str:
103
+ """
104
+ Search arXiv for academic papers matching the query and return excerpts from up to three results.
105
+
106
+ Args:
107
+ query: A string query to search on arXiv.
108
+
109
+ Returns:
110
+ A string containing excerpts (first 1000 characters) from up to three relevant papers.
111
+ """
112
  docs = ArxivLoader(query=query, load_max_docs=3).load()
113
  return "\n\n---\n\n".join(doc.page_content[:1000] for doc in docs)
114
 
115
+
116
  # ------------------ System Setup ------------------
117
  with open("system_prompt.txt", "r", encoding="utf-8") as f:
118
  system_prompt = f.read()