Skip to content

Commit 7ccc8cd

Browse files
dartpaindartpain
authored
Merge pull request #1855 from siiddhantt/refactor/ddg-brave-tools
refactor: ddg and brave tools with sources fix
2 parents f0908af + 0d48159 commit 7ccc8cd

File tree

15 files changed

+310
-505
lines changed

15 files changed

+310
-505
lines changed

‎application/agents/classic_agent.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99

1010
class ClassicAgent(BaseAgent):
11-
"""A simplified classic agent with clear execution flow.
11+
"""A simplified agent with clear execution flow.
1212
1313
Usage:
1414
1. Processes a query through retrieval

‎application/agents/tools/brave.py‎

Lines changed: 45 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -25,27 +25,35 @@ def execute_action(self, action_name, **kwargs):
2525
else:
2626
raise ValueError(f"Unknown action: {action_name}")
2727

28-
def _web_search(self, query, country="ALL", search_lang="en", count=10,
29-
offset=0, safesearch="off", freshness=None,
30-
result_filter=None, extra_snippets=False, summary=False):
28+
def _web_search(
29+
self,
30+
query,
31+
country="ALL",
32+
search_lang="en",
33+
count=10,
34+
offset=0,
35+
safesearch="off",
36+
freshness=None,
37+
result_filter=None,
38+
extra_snippets=False,
39+
summary=False,
40+
):
3141
"""
3242
Performs a web search using the Brave Search API.
3343
"""
3444
print(f"Performing Brave web search for: {query}")
35-
45+
3646
url = f"{self.base_url}/web/search"
37-
38-
# Build query parameters
47+
3948
params = {
4049
"q": query,
4150
"country": country,
4251
"search_lang": search_lang,
4352
"count": min(count, 20),
4453
"offset": min(offset, 9),
45-
"safesearch": safesearch
54+
"safesearch": safesearch,
4655
}
47-
48-
# Add optional parameters only if they have values
56+
4957
if freshness:
5058
params["freshness"] = freshness
5159
if result_filter:
@@ -54,68 +62,69 @@ def _web_search(self, query, country="ALL", search_lang="en", count=10,
5462
params["extra_snippets"] = 1
5563
if summary:
5664
params["summary"] = 1
57-
58-
# Set up headers
5965
headers = {
6066
"Accept": "application/json",
6167
"Accept-Encoding": "gzip",
62-
"X-Subscription-Token": self.token
68+
"X-Subscription-Token": self.token,
6369
}
64-
65-
# Make the request
70+
6671
response = requests.get(url, params=params, headers=headers)
67-
72+
6873
if response.status_code == 200:
6974
return {
7075
"status_code": response.status_code,
7176
"results": response.json(),
72-
"message": "Search completed successfully."
77+
"message": "Search completed successfully.",
7378
}
7479
else:
7580
return {
7681
"status_code": response.status_code,
77-
"message": f"Search failed with status code: {response.status_code}."
82+
"message": f"Search failed with status code: {response.status_code}.",
7883
}
79-
80-
def _image_search(self, query, country="ALL", search_lang="en", count=5,
81-
safesearch="off", spellcheck=False):
84+
85+
def _image_search(
86+
self,
87+
query,
88+
country="ALL",
89+
search_lang="en",
90+
count=5,
91+
safesearch="off",
92+
spellcheck=False,
93+
):
8294
"""
8395
Performs an image search using the Brave Search API.
8496
"""
8597
print(f"Performing Brave image search for: {query}")
86-
98+
8799
url = f"{self.base_url}/images/search"
88-
89-
# Build query parameters
100+
90101
params = {
91102
"q": query,
92103
"country": country,
93104
"search_lang": search_lang,
94105
"count": min(count, 100), # API max is 100
95106
"safesearch": safesearch,
96-
"spellcheck": 1 if spellcheck else 0
107+
"spellcheck": 1 if spellcheck else 0,
97108
}
98-
99-
# Set up headers
109+
100110
headers = {
101111
"Accept": "application/json",
102112
"Accept-Encoding": "gzip",
103-
"X-Subscription-Token": self.token
113+
"X-Subscription-Token": self.token,
104114
}
105-
106-
# Make the request
115+
107116
response = requests.get(url, params=params, headers=headers)
108-
117+
109118
if response.status_code == 200:
110119
return {
111120
"status_code": response.status_code,
112121
"results": response.json(),
113-
"message": "Image search completed successfully."
122+
"message": "Image search completed successfully.",
114123
}
115124
else:
116125
return {
117126
"status_code": response.status_code,
118-
"message": f"Image search failed with status code: {response.status_code}."
127+
"message": f"Image search failed with status code: {response.status_code}.",
119128
}
120129

121130
def get_actions_metadata(self):
@@ -130,42 +139,14 @@ def get_actions_metadata(self):
130139
"type": "string",
131140
"description": "The search query (max 400 characters, 50 words)",
132141
},
133-
# "country": {
134-
# "type": "string",
135-
# "description": "The 2-character country code (default: US)",
136-
# },
137142
"search_lang": {
138143
"type": "string",
139144
"description": "The search language preference (default: en)",
140145
},
141-
# "count": {
142-
# "type": "integer",
143-
# "description": "Number of results to return (max 20, default: 10)",
144-
# },
145-
# "offset": {
146-
# "type": "integer",
147-
# "description": "Pagination offset (max 9, default: 0)",
148-
# },
149-
# "safesearch": {
150-
# "type": "string",
151-
# "description": "Filter level for adult content (off, moderate, strict)",
152-
# },
153146
"freshness": {
154147
"type": "string",
155148
"description": "Time filter for results (pd: last 24h, pw: last week, pm: last month, py: last year)",
156149
},
157-
# "result_filter": {
158-
# "type": "string",
159-
# "description": "Comma-delimited list of result types to include",
160-
# },
161-
# "extra_snippets": {
162-
# "type": "boolean",
163-
# "description": "Get additional excerpts from result pages",
164-
# },
165-
# "summary": {
166-
# "type": "boolean",
167-
# "description": "Enable summary generation in search results",
168-
# }
169150
},
170151
"required": ["query"],
171152
"additionalProperties": False,
@@ -181,37 +162,21 @@ def get_actions_metadata(self):
181162
"type": "string",
182163
"description": "The search query (max 400 characters, 50 words)",
183164
},
184-
# "country": {
185-
# "type": "string",
186-
# "description": "The 2-character country code (default: US)",
187-
# },
188-
# "search_lang": {
189-
# "type": "string",
190-
# "description": "The search language preference (default: en)",
191-
# },
192165
"count": {
193166
"type": "integer",
194167
"description": "Number of results to return (max 100, default: 5)",
195168
},
196-
# "safesearch": {
197-
# "type": "string",
198-
# "description": "Filter level for adult content (off, strict). Default: strict",
199-
# },
200-
# "spellcheck": {
201-
# "type": "boolean",
202-
# "description": "Whether to spellcheck provided query (default: true)",
203-
# }
204169
},
205170
"required": ["query"],
206171
"additionalProperties": False,
207172
},
208-
}
173+
},
209174
]
210175

211176
def get_config_requirements(self):
212177
return {
213178
"token": {
214-
"type": "string",
215-
"description": "Brave Search API key for authentication"
179+
"type": "string",
180+
"description": "Brave Search API key for authentication",
216181
},
217-
}
182+
}

‎application/agents/tools/duckduckgo.py‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
from application.agents.tools.base import Tool
2+
from duckduckgo_search import DDGS
3+
4+
5+
class DuckDuckGoSearchTool(Tool):
6+
"""
7+
DuckDuckGo Search
8+
A tool for performing web and image searches using DuckDuckGo.
9+
"""
10+
11+
def __init__(self, config):
12+
self.config = config
13+
14+
def execute_action(self, action_name, **kwargs):
15+
actions = {
16+
"ddg_web_search": self._web_search,
17+
"ddg_image_search": self._image_search,
18+
}
19+
20+
if action_name in actions:
21+
return actions[action_name](**kwargs)
22+
else:
23+
raise ValueError(f"Unknown action: {action_name}")
24+
25+
def _web_search(
26+
self,
27+
query,
28+
max_results=5,
29+
):
30+
print(f"Performing DuckDuckGo web search for: {query}")
31+
32+
try:
33+
results = DDGS().text(
34+
query,
35+
max_results=max_results,
36+
)
37+
38+
return {
39+
"status_code": 200,
40+
"results": results,
41+
"message": "Web search completed successfully.",
42+
}
43+
except Exception as e:
44+
return {
45+
"status_code": 500,
46+
"message": f"Web search failed: {str(e)}",
47+
}
48+
49+
def _image_search(
50+
self,
51+
query,
52+
max_results=5,
53+
):
54+
print(f"Performing DuckDuckGo image search for: {query}")
55+
56+
try:
57+
results = DDGS().images(
58+
keywords=query,
59+
max_results=max_results,
60+
)
61+
62+
return {
63+
"status_code": 200,
64+
"results": results,
65+
"message": "Image search completed successfully.",
66+
}
67+
except Exception as e:
68+
return {
69+
"status_code": 500,
70+
"message": f"Image search failed: {str(e)}",
71+
}
72+
73+
def get_actions_metadata(self):
74+
return [
75+
{
76+
"name": "ddg_web_search",
77+
"description": "Perform a web search using DuckDuckGo.",
78+
"parameters": {
79+
"type": "object",
80+
"properties": {
81+
"query": {
82+
"type": "string",
83+
"description": "Search query",
84+
},
85+
"max_results": {
86+
"type": "integer",
87+
"description": "Number of results to return (default: 5)",
88+
},
89+
},
90+
"required": ["query"],
91+
},
92+
},
93+
{
94+
"name": "ddg_image_search",
95+
"description": "Perform an image search using DuckDuckGo.",
96+
"parameters": {
97+
"type": "object",
98+
"properties": {
99+
"query": {
100+
"type": "string",
101+
"description": "Search query",
102+
},
103+
"max_results": {
104+
"type": "integer",
105+
"description": "Number of results to return (default: 5, max: 50)",
106+
},
107+
},
108+
"required": ["query"],
109+
},
110+
},
111+
]
112+
113+
def get_config_requirements(self):
114+
return {}

‎application/api/user/routes.py‎

Lines changed: 0 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -879,29 +879,6 @@ def get(self):
879879
"syncFrequency": index.get("sync_frequency", ""),
880880
}
881881
)
882-
if "duckduck_search" in settings.RETRIEVERS_ENABLED:
883-
data.append(
884-
{
885-
"name": "DuckDuckGo Search",
886-
"date": "duckduck_search",
887-
"model": settings.EMBEDDINGS_NAME,
888-
"location": "custom",
889-
"tokens": "",
890-
"retriever": "duckduck_search",
891-
}
892-
)
893-
if "brave_search" in settings.RETRIEVERS_ENABLED:
894-
data.append(
895-
{
896-
"name": "Brave Search",
897-
"language": "en",
898-
"date": "brave_search",
899-
"model": settings.EMBEDDINGS_NAME,
900-
"location": "custom",
901-
"tokens": "",
902-
"retriever": "brave_search",
903-
}
904-
)
905882
except Exception as err:
906883
current_app.logger.error(f"Error retrieving sources: {err}", exc_info=True)
907884
return make_response(jsonify({"success": False}), 400)

‎application/core/settings.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,7 +33,7 @@ class Settings(BaseSettings):
3333
VECTOR_STORE: str = (
3434
"faiss" # "faiss" or "elasticsearch" or "qdrant" or "milvus" or "lancedb"
3535
)
36-
RETRIEVERS_ENABLED: list = ["classic_rag", "duckduck_search"] # also brave_search
36+
RETRIEVERS_ENABLED: list = ["classic_rag"]
3737
AGENT_NAME: str = "classic"
3838
FALLBACK_LLM_PROVIDER: Optional[str] = None # provider for fallback llm
3939
FALLBACK_LLM_NAME: Optional[str] = None # model name for fallback llm
@@ -99,7 +99,6 @@ class Settings(BaseSettings):
9999
LANCEDB_TABLE_NAME: Optional[str] = (
100100
"docsgpts" # Name of the table to use for storing vectors
101101
)
102-
BRAVE_SEARCH_API_KEY: Optional[str] = None
103102

104103
FLASK_DEBUG_MODE: bool = False
105104
STORAGE_TYPE: str = "local" # local or s3

0 commit comments

Comments
 (0)