ktluege commited on
Commit
292f5f0
Β·
verified Β·
1 Parent(s): 442e89c

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +5 -29
agent.py CHANGED
@@ -17,55 +17,31 @@ load_dotenv()
17
  # ----- TOOLS -----
18
  @tool
19
  def multiply(a: int, b: int) -> int:
 
20
  return a * b
21
 
22
  @tool
23
  def add(a: int, b: int) -> int:
 
24
  return a + b
25
 
26
  @tool
27
  def subtract(a: int, b: int) -> int:
 
28
  return a - b
29
 
30
  @tool
31
  def divide(a: int, b: int) -> float:
 
32
  if b == 0:
33
  raise ValueError("Cannot divide by zero.")
34
  return a / b
35
 
36
  @tool
37
  def modulus(a: int, b: int) -> int:
 
38
  return a % b
39
 
40
- @tool
41
- def wiki_search(query: str) -> str:
42
- search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
43
- formatted_search_docs = "\n\n---\n\n".join(
44
- [
45
- f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
46
- for doc in search_docs
47
- ])
48
- return {"wiki_results": formatted_search_docs}
49
-
50
- @tool
51
- def web_search(query: str) -> str:
52
- search_docs = TavilySearchResults(max_results=3).invoke(query=query)
53
- formatted_search_docs = "\n\n---\n\n".join(
54
- [
55
- f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content}\n</Document>'
56
- for doc in search_docs
57
- ])
58
- return {"web_results": formatted_search_docs}
59
-
60
- @tool
61
- def arvix_search(query: str) -> str:
62
- search_docs = ArxivLoader(query=query, load_max_docs=3).load()
63
- formatted_search_docs = "\n\n---\n\n".join(
64
- [
65
- f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
66
- for doc in search_docs
67
- ])
68
- return {"arvix_results": formatted_search_docs}
69
 
70
  tools = [
71
  multiply, add, subtract, divide, modulus,
 
17
  # ----- TOOLS -----
18
  @tool
19
  def multiply(a: int, b: int) -> int:
20
+ """Multiply two numbers."""
21
  return a * b
22
 
23
  @tool
24
  def add(a: int, b: int) -> int:
25
+ """Add two numbers."""
26
  return a + b
27
 
28
  @tool
29
  def subtract(a: int, b: int) -> int:
30
+ """Subtract two numbers."""
31
  return a - b
32
 
33
  @tool
34
  def divide(a: int, b: int) -> float:
35
+ """Divide two numbers."""
36
  if b == 0:
37
  raise ValueError("Cannot divide by zero.")
38
  return a / b
39
 
40
  @tool
41
  def modulus(a: int, b: int) -> int:
42
+ """Get the modulus of two numbers."""
43
  return a % b
44
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
45
 
46
  tools = [
47
  multiply, add, subtract, divide, modulus,