Phoenix21 commited on
Commit
90437f0
·
verified ·
1 Parent(s): d9f3ab5

updated web search function

Browse files
Files changed (1) hide show
  1. pipeline.py +20 -3
pipeline.py CHANGED
@@ -39,7 +39,7 @@ gemini_llm = ChatGoogleGenerativeAI(
39
  # Additional parameters or safety_settings can be added here if needed
40
  )
41
 
42
- web_gemini_llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))
43
 
44
  ################################################################################
45
  # Pydantic Models
@@ -315,9 +315,26 @@ def do_cached_web_search(query: str) -> str:
315
  # 2) If no suitable cached answer, do a new search
316
  try:
317
  print("DEBUG: Performing a new web search...")
 
318
  search_tool = DuckDuckGoSearchTool()
319
- search_agent = ManagedAgent(llm=web_gemini_llm, tools=[search_tool])
320
- new_search_result = search_agent.run(f"Search for information about: {query}")
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
321
 
322
  # 3) Store in cache for future reuse
323
  store_websearch_result(query, new_search_result)
 
39
  # Additional parameters or safety_settings can be added here if needed
40
  )
41
 
42
+ # web_gemini_llm = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))
43
 
44
  ################################################################################
45
  # Pydantic Models
 
315
  # 2) If no suitable cached answer, do a new search
316
  try:
317
  print("DEBUG: Performing a new web search...")
318
+ model = LiteLLMModel(model_id="gemini/gemini-pro", api_key=os.environ.get("GEMINI_API_KEY"))
319
  search_tool = DuckDuckGoSearchTool()
320
+ web_agent = CodeAgent(
321
+ tools=[search_tool],
322
+ model=model
323
+ )
324
+
325
+ managed_web_agent = ManagedAgent(
326
+ agent=web_agent,
327
+ name="web_search",
328
+ description="Runs a web search for you. Provide your query as an argument."
329
+ )
330
+
331
+ manager_agent = CodeAgent(
332
+ tools=[], # If you have additional tools for the manager, add them here
333
+ model=model,
334
+ managed_agents=[managed_web_agent]
335
+ )
336
+
337
+ new_search_result = manager_agent.run(f"Search for information about: {query}")
338
 
339
  # 3) Store in cache for future reuse
340
  store_websearch_result(query, new_search_result)