Phoenix21 commited on
Commit
54b13ed
·
1 Parent(s): 1b9109f

Updated web Search function

Browse files
Files changed (1) hide show
  1. pipeline.py +23 -18
pipeline.py CHANGED
@@ -26,6 +26,7 @@ from langchain.prompts import PromptTemplate
26
  from langchain.docstore.document import Document
27
  from langchain_core.caches import BaseCache
28
  from langchain_core.callbacks import Callbacks
 
29
 
30
  from chain.classification_chain import get_classification_chain
31
  from chain.refusal_chain import get_refusal_chain
@@ -220,32 +221,36 @@ def retrieve_websearch_result(query: str) -> Optional[str]:
220
  return web_search_cache.get(query.strip().lower())
221
 
222
  def do_web_search(query: str) -> str:
223
- """Performs web search if no cached result exists."""
224
  try:
225
  cached = retrieve_websearch_result(query)
226
  if cached:
227
  logger.info("Using cached web search result.")
228
  return cached
229
 
230
- logger.info("Performing a new web search for: '%s'", query)
231
- search_agent = ToolCallingAgent(
232
- tools=[DuckDuckGoSearchTool(), VisitWebpageTool()],
233
- model=HfApiModel(),
234
- name="search_agent",
235
- description="This is an agent that can do web search.",
236
- )
237
-
238
- manager_agent = CodeAgent(
239
- tools=[],
240
- model=HfApiModel(),
241
- managed_agents=[search_agent]
242
- )
 
 
 
243
 
244
- new_search_result = manager_agent.run(f"Search for information about: {query}")
245
- store_websearch_result(query, new_search_result)
246
- return str(new_search_result).strip()
 
247
  except Exception as e:
248
- logger.error(f"Web search failed: {e}")
249
  return ""
250
 
251
  def is_greeting(query: str) -> bool:
 
26
  from langchain.docstore.document import Document
27
  from langchain_core.caches import BaseCache
28
  from langchain_core.callbacks import Callbacks
29
+ from langchain_community.tools import TavilySearchResults
30
 
31
  from chain.classification_chain import get_classification_chain
32
  from chain.refusal_chain import get_refusal_chain
 
221
  return web_search_cache.get(query.strip().lower())
222
 
223
  def do_web_search(query: str) -> str:
224
+ """Performs web search using Tavily if no cached result exists."""
225
  try:
226
  cached = retrieve_websearch_result(query)
227
  if cached:
228
  logger.info("Using cached web search result.")
229
  return cached
230
 
231
+ logger.info("Performing a new Tavily web search for: '%s'", query)
232
+ #Intialize Tavily search tool
233
+ tavily_api_key = os.environ.get("TAVILY_API_KEY")
234
+ if not tavily_api_key:
235
+ logger.error("Tavily API key not found.")
236
+ return "Unable to perform web search API key not set"
237
+ #Create Tavily Search Tool
238
+ tavily_search=TavilySearchResults(api_key=tavily_api_key)
239
+ #Perform search
240
+ search_results = tavily_search.search(query, num_results=3)
241
+
242
+ result_text = "Web Search Results:\n\n"
243
+ for i, result in enumerate(search_results):
244
+ result_text += f"{i+1}. {result.get('title', 'No Title')}\n"
245
+ result_text += f" URL: {result.get('url', 'No URL')}\n"
246
+ result_text += f" {result.get('content', 'No content available')[:300]}...\n\n"
247
 
248
+ store_websearch_result(query, result_text)
249
+ return result_text.strip()
250
+
251
+
252
  except Exception as e:
253
+ logger.error(f"Tavily Web search failed: {e}")
254
  return ""
255
 
256
  def is_greeting(query: str) -> bool: