supratipb commited on
Commit
696927b
Β·
verified Β·
1 Parent(s): 04d91b0

Upload agent.py

Browse files
Files changed (1) hide show
  1. agent.py +26 -18
agent.py CHANGED
@@ -80,34 +80,42 @@ def modulus(a: int, b: int) -> int:
80
 
81
  @tool
82
  def wiki_search(query: str) -> str:
83
- """Search Wikipedia for a query and return maximum 2 results.
84
-
85
- Args:
86
- query: The search query."""
87
- search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
88
- formatted_search_docs = "\n\n---\n\n".join(
89
- [
90
- f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
91
- for doc in search_docs
92
- ])
93
- return {"wiki_results": formatted_search_docs}
 
 
 
 
94
 
95
 
96
 
97
  @tool
98
  def web_search(query: str) -> str:
99
- """Search Tavily for a query and return maximum 3 results.
100
-
101
- Args:
102
- query: The search query."""
103
  search_docs = TavilySearch(max_results=3).invoke(query=query)
 
 
 
104
  formatted_search_docs = "\n\n---\n\n".join(
105
  [
106
- f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
 
 
107
  for doc in search_docs
108
- ])
 
109
  return {"web_results": formatted_search_docs}
110
 
 
111
  @tool
112
  def arvix_search(query: str) -> str:
113
  """Search Arxiv for a query and return maximum 3 result.
@@ -177,7 +185,7 @@ tools = [
177
  ]
178
 
179
  # Build graph function
180
- def build_graph(provider: str = "openai"):
181
  """Build the graph"""
182
  # Load environment variables from .env file
183
  if provider == "openai":
 
80
 
81
  @tool
82
  def wiki_search(query: str) -> str:
83
+ """Search Wikipedia for a query and return up to 5 results summarized."""
84
+ search_docs = WikipediaLoader(query=query, load_max_docs=5).load()
85
+ if not search_docs:
86
+ return "No Wikipedia results found."
87
+
88
+ summaries = []
89
+ for doc in search_docs:
90
+ content = doc.page_content
91
+ # Truncate if too long or summarize with a simple heuristic or model
92
+ content = content[:1000] + ("..." if len(content) > 1000 else "")
93
+ summaries.append(
94
+ f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}">\n{content}\n</Document>'
95
+ )
96
+ return {"wiki_results": "\n\n---\n\n".join(summaries)}
97
+
98
 
99
 
100
 
101
  @tool
102
  def web_search(query: str) -> str:
103
+ """Search Tavily for a query and return maximum 3 results."""
 
 
 
104
  search_docs = TavilySearch(max_results=3).invoke(query=query)
105
+ if not search_docs:
106
+ return "No web search results found."
107
+
108
  formatted_search_docs = "\n\n---\n\n".join(
109
  [
110
+ f'<Document source="{doc.metadata.get("source", "unknown")}" page="{doc.metadata.get("page", "")}"/>\n'
111
+ f'{doc.page_content[:1000]}' # truncate to 1000 chars if needed
112
+ f'\n</Document>'
113
  for doc in search_docs
114
+ ]
115
+ )
116
  return {"web_results": formatted_search_docs}
117
 
118
+
119
  @tool
120
  def arvix_search(query: str) -> str:
121
  """Search Arxiv for a query and return maximum 3 result.
 
185
  ]
186
 
187
  # Build graph function
188
+ def build_graph(provider: str = "groq"):
189
  """Build the graph"""
190
  # Load environment variables from .env file
191
  if provider == "openai":