lam-ho commited on
Commit
bef8270
·
verified ·
1 Parent(s): bbbcb72

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +3 -3
agent.py CHANGED
@@ -1,5 +1,5 @@
1
  from smolagents import HfApiModel, CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, Tool, LiteLLMModel
2
- from langchain_community.tools.tavily_search import TavilySearchResults
3
  from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
4
  import os, time, math, pandas, numpy
5
 
@@ -173,7 +173,7 @@ class TavilySearchTool(Tool):
173
  output_type = "string"
174
 
175
  def forward(self, query: str) -> str:
176
- documents = TavilySearchResults(max_results=3).invoke(query=query)
177
  condensed_docs = "\n\n---\n\n".join(
178
  [
179
  f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
@@ -204,7 +204,7 @@ class ArvixSearchTool(Tool):
204
  documents = ArxivLoader(query=query, load_max_docs=3).load()
205
  condensed_docs = "\n\n---\n\n".join(
206
  [
207
- f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
208
  for doc in documents
209
  ])
210
  return {"arvix_search_results": condensed_docs}
 
1
  from smolagents import HfApiModel, CodeAgent, DuckDuckGoSearchTool, WikipediaSearchTool, Tool, LiteLLMModel
2
+ from langchain_tavily import TavilySearch
3
  from langchain_community.document_loaders import WikipediaLoader, ArxivLoader
4
  import os, time, math, pandas, numpy
5
 
 
173
  output_type = "string"
174
 
175
  def forward(self, query: str) -> str:
176
+ documents = TavilySearch(max_results=3).invoke(query=query)
177
  condensed_docs = "\n\n---\n\n".join(
178
  [
179
  f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
 
204
  documents = ArxivLoader(query=query, load_max_docs=3).load()
205
  condensed_docs = "\n\n---\n\n".join(
206
  [
207
+ f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.content[:1000]}\n</Document>'
208
  for doc in documents
209
  ])
210
  return {"arvix_search_results": condensed_docs}