Coool2 commited on
Commit
6242081
·
verified ·
1 Parent(s): 0c9e128

Update agent.py

Browse files
Files changed (1) hide show
  1. agent.py +6 -47
agent.py CHANGED
@@ -242,9 +242,9 @@ class IntelligentSourceRouter:
242
  def __init__(self):
243
  # Initialize ArXiv and DuckDuckGo as LlamaIndex tools
244
  self.arxiv_tool = ArxivToolSpec().to_tool_list()[0]
245
- self.duckduckgo_tool = DuckDuckGoSearchToolSpec().to_tool_list()[0]
246
 
247
- def detect_intent_and_route(self, query: str) -> str:
248
  # Use your LLM to decide between arxiv and web_search
249
  intent_prompt = f"""
250
  Analyze this query and determine if it's scientific research or general information:
@@ -260,10 +260,10 @@ class IntelligentSourceRouter:
260
  results = [f"**Query**: {query}", f"**Selected Source**: {selected_source}", "="*50]
261
  try:
262
  if selected_source == 'arxiv':
263
- result = self.arxiv_tool.call(query=query, max_results=3)
264
  results.append(f"**ArXiv Research:**\n{result}")
265
  else:
266
- result = self.duckduckgo_tool.call(query=query, max_results=5)
267
  # Format results if needed
268
  if isinstance(result, list):
269
  formatted = []
@@ -277,55 +277,14 @@ class IntelligentSourceRouter:
277
  results.append(f"**Search failed**: {str(e)}")
278
  return "\n\n".join(results)
279
 
280
- class IntelligentSourceRouter:
281
- def __init__(self):
282
- # Initialize Arxiv and DuckDuckGo tools
283
- self.arxiv_tool = ArxivToolSpec().to_tool_list()[0]
284
- self.duckduckgo_tool = DuckDuckGoSearchToolSpec().to_tool_list()[0]
285
-
286
- def detect_intent_and_extract_content(self, query: str, max_results: int = 3) -> str:
287
- # Use your LLM to decide between arxiv and web_search
288
- intent_prompt = f"""
289
- Analyze this query and determine if it's scientific research or general information:
290
- Query: "{query}"
291
- Choose ONE source:
292
- - arxiv: For scientific research, academic papers, technical studies, algorithms, experiments
293
- - web_search: For all other information (current events, general facts, weather, how-to guides, etc.)
294
- Respond with ONLY "arxiv" or "web_search".
295
- """
296
- response = proj_llm.complete(intent_prompt)
297
- selected_source = response.text.strip().lower()
298
-
299
- results = [f"**Query**: {query}", f"**Selected Source**: {selected_source}", "="*50]
300
- try:
301
- if selected_source == 'arxiv':
302
- # Extract abstracts and paper summaries (deep content)
303
- arxiv_results = self.arxiv_tool.call(query=query, max_results=max_results)
304
- results.append(f"**Extracted ArXiv Content:**\n{arxiv_results}")
305
- else:
306
- # DuckDuckGo returns a list of dicts with 'href', 'title', 'body'
307
- web_results = self.duckduckgo_tool.call(query=query, max_results=max_results)
308
- if isinstance(web_results, list):
309
- formatted = []
310
- for i, r in enumerate(web_results, 1):
311
- formatted.append(
312
- f"{i}. **{r.get('title', '')}**\n URL: {r.get('href', '')}\n {r.get('body', '')}"
313
- )
314
- web_content = "\n".join(formatted)
315
- else:
316
- web_content = str(web_results)
317
- results.append(f"**Extracted Web Content:**\n{web_content}")
318
- except Exception as e:
319
- results.append(f"**Extraction failed**: {str(e)}")
320
- return "\n\n".join(results)
321
 
322
  # Initialize router
323
  intelligent_router = IntelligentSourceRouter()
324
 
325
  # Create enhanced research tool
326
- def enhanced_smart_research_tool(query: str, task_context: str = "", max_results: int = 3) -> str:
327
  full_query = f"{query} {task_context}".strip()
328
- return intelligent_router.detect_intent_and_extract_content(full_query, max_results=max_results)
329
 
330
  research_tool = FunctionTool.from_defaults(
331
  fn=enhanced_smart_research_tool,
 
242
  def __init__(self):
243
  # Initialize ArXiv and DuckDuckGo as LlamaIndex tools
244
  self.arxiv_tool = ArxivToolSpec().to_tool_list()[0]
245
+ self.duckduckgo_tool = DuckDuckGoSearchToolSpec().to_tool_list()[1]
246
 
247
+ def detect_intent_and_extract_content(self, query: str, max_results = 3) -> str:
248
  # Use your LLM to decide between arxiv and web_search
249
  intent_prompt = f"""
250
  Analyze this query and determine if it's scientific research or general information:
 
260
  results = [f"**Query**: {query}", f"**Selected Source**: {selected_source}", "="*50]
261
  try:
262
  if selected_source == 'arxiv':
263
+ result = self.arxiv_tool.call(query=query)
264
  results.append(f"**ArXiv Research:**\n{result}")
265
  else:
266
+ result = self.duckduckgo_tool.call(query=query, max_results=max_results)
267
  # Format results if needed
268
  if isinstance(result, list):
269
  formatted = []
 
277
  results.append(f"**Search failed**: {str(e)}")
278
  return "\n\n".join(results)
279
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
280
 
281
  # Initialize router
282
  intelligent_router = IntelligentSourceRouter()
283
 
284
  # Create enhanced research tool
285
+ def enhanced_smart_research_tool(query: str, task_context: str = "") -> str:
286
  full_query = f"{query} {task_context}".strip()
287
+ return intelligent_router.detect_intent_and_extract_content(full_query)
288
 
289
  research_tool = FunctionTool.from_defaults(
290
  fn=enhanced_smart_research_tool,