LamiaYT commited on
Commit
c45ce5d
·
1 Parent(s): 89373ab
Files changed (1) hide show
  1. agent.py +21 -0
agent.py CHANGED
@@ -24,28 +24,34 @@ from langchain.schema import Document
24
 
25
  @tool
26
  def multiply(a: int, b: int) -> int:
 
27
  return a * b
28
 
29
  @tool
30
  def add(a: int, b: int) -> int:
 
31
  return a + b
32
 
33
  @tool
34
  def subtract(a: int, b: int) -> int:
 
35
  return a - b
36
 
37
  @tool
38
  def divide(a: int, b: int) -> float:
 
39
  if b == 0:
40
  raise ValueError("Cannot divide by zero.")
41
  return a / b
42
 
43
  @tool
44
  def modulus(a: int, b: int) -> int:
 
45
  return a % b
46
 
47
  @tool
48
  def wiki_search(query: str) -> str:
 
49
  search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
50
  formatted = "\n\n---\n\n".join(
51
  [
@@ -57,6 +63,7 @@ def wiki_search(query: str) -> str:
57
 
58
  @tool
59
  def web_search(query: str) -> str:
 
60
  search_docs = TavilySearchResults(max_results=3).invoke(query=query)
61
  formatted = "\n\n---\n\n".join(
62
  [
@@ -68,6 +75,7 @@ def web_search(query: str) -> str:
68
 
69
  @tool
70
  def arvix_search(query: str) -> str:
 
71
  search_docs = ArxivLoader(query=query, load_max_docs=3).load()
72
  formatted = "\n\n---\n\n".join(
73
  [
@@ -77,6 +85,19 @@ def arvix_search(query: str) -> str:
77
  )
78
  return {"arvix_results": formatted}
79
 
 
 
 
 
 
 
 
 
 
 
 
 
 
80
  # ---- Embedding & Vector Store ----
81
 
82
  embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")
 
24
 
25
  @tool
26
  def multiply(a: int, b: int) -> int:
27
+ """Multiply two integers and return the result."""
28
  return a * b
29
 
30
  @tool
31
  def add(a: int, b: int) -> int:
32
+ """Add two integers and return the result."""
33
  return a + b
34
 
35
  @tool
36
  def subtract(a: int, b: int) -> int:
37
+ """Subtract second integer from the first and return the result."""
38
  return a - b
39
 
40
  @tool
41
  def divide(a: int, b: int) -> float:
42
+ """Divide first integer by second and return the result as a float."""
43
  if b == 0:
44
  raise ValueError("Cannot divide by zero.")
45
  return a / b
46
 
47
  @tool
48
  def modulus(a: int, b: int) -> int:
49
+ """Return the remainder when first integer is divided by second."""
50
  return a % b
51
 
52
  @tool
53
  def wiki_search(query: str) -> str:
54
+ """Search Wikipedia for the given query and return formatted documents."""
55
  search_docs = WikipediaLoader(query=query, load_max_docs=2).load()
56
  formatted = "\n\n---\n\n".join(
57
  [
 
63
 
64
  @tool
65
  def web_search(query: str) -> str:
66
+ """Search the web using Tavily API for the given query."""
67
  search_docs = TavilySearchResults(max_results=3).invoke(query=query)
68
  formatted = "\n\n---\n\n".join(
69
  [
 
75
 
76
  @tool
77
  def arvix_search(query: str) -> str:
78
+ """Search Arxiv for academic papers related to the query."""
79
  search_docs = ArxivLoader(query=query, load_max_docs=3).load()
80
  formatted = "\n\n---\n\n".join(
81
  [
 
85
  )
86
  return {"arvix_results": formatted}
87
 
88
+ @tool
89
+ def similar_question_search(query: str) -> str:
90
+ """Find and return top 3 most similar past questions from vector store."""
91
+ matched_docs = vector_store.similarity_search(query, 3)
92
+ formatted = "\n\n---\n\n".join(
93
+ [
94
+ f'<Document source="{doc.metadata["source"]}" page="{doc.metadata.get("page", "")}"/>\n{doc.page_content[:1000]}\n</Document>'
95
+ for doc in matched_docs
96
+ ]
97
+ )
98
+ return {"similar_questions": formatted}
99
+
100
+
101
  # ---- Embedding & Vector Store ----
102
 
103
  embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-mpnet-base-v2")