Skip to content

Commit 3880ccb

Browse files
alcholiclgalcholiclg
authored
Feat/Generalize Search Engine Integration and Workflow Enhancements (#687)
1 parent eee92de commit 3880ccb

16 files changed

Lines changed: 977 additions & 46 deletions

File tree

‎conf.yaml‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
# Search Engine, Supported values: EXA, SERPAPI, ARXIV
2+
# SEARCH_ENGINE:
3+
# engine: exa
4+
# exa_api_key: $EXA_API_KEY
5+
6+
# SEARCH_ENGINE:
7+
# engine: serpapi
8+
# serpapi_api_key: $SERPAPI_API_KEY
9+
# provider: google
10+
11+
# SEARCH_ENGINE:
12+
# engine: arxiv

‎ms_agent/tools/exa/search.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
from exa_py import Exa
55
from ms_agent.tools.exa.schema import ExaSearchRequest, ExaSearchResult
6+
from ms_agent.tools.search.search_base import SearchEngineType
67

78

89
class ExaSearch:
@@ -13,6 +14,7 @@ def __init__(self, api_key: str = None):
1314
assert api_key, 'EXA_API_KEY must be set either as an argument or as an environment variable'
1415

1516
self.client = Exa(api_key=api_key)
17+
self.engine_type = SearchEngineType.EXA
1618

1719
def search(self, search_request: ExaSearchRequest) -> ExaSearchResult:
1820
"""

‎ms_agent/tools/search/arxiv/__init__.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
# flake8: noqa
2+
from ms_agent.tools.search.arxiv.schema import (ArxivSearchRequest,
3+
ArxivSearchResult)
4+
from ms_agent.tools.search.arxiv.search import ArxivSearch

‎ms_agent/tools/search/arxiv/schema.py‎

Lines changed: 140 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,140 @@
1+
# flake8: noqa
2+
from dataclasses import dataclass, field
3+
from typing import Any, Dict, Generator, List, Optional
4+
5+
import arxiv
6+
import json
7+
from arxiv import SortCriterion, SortOrder
8+
from ms_agent.tools.search.search_base import (BaseResult, SearchRequest,
9+
SearchResponse, SearchResult)
10+
11+
12+
class ArxivSearchRequest(SearchRequest):
13+
"""
14+
A class representing a search request to ArXiv.
15+
"""
16+
17+
def __init__(self,
18+
query: str = None,
19+
num_results: Optional[int] = 10,
20+
sort_strategy: SortCriterion = SortCriterion.Relevance,
21+
sort_order: SortOrder = SortOrder.Descending):
22+
"""
23+
Initialize ArxivSearchRequest with search parameters.
24+
25+
Args:
26+
query: The search query string
27+
num_results: Number of results to return, default is 10
28+
sort_strategy: The strategy to sort results, default is relevance
29+
sort_order: The order of sorting, default is descending
30+
"""
31+
super().__init__(query=query, num_results=num_results)
32+
self.sort_strategy = sort_strategy
33+
self.sort_order = sort_order
34+
35+
def to_dict(self) -> Dict[str, Any]:
36+
"""
37+
Convert the request parameters to a dictionary.
38+
39+
Returns:
40+
Dict[str, Any]: The parameters as a dictionary
41+
"""
42+
return {
43+
'query': self.query,
44+
'max_results': self.num_results,
45+
'sort_by': self.sort_strategy,
46+
'sort_order': self.sort_order
47+
}
48+
49+
def to_json(self) -> Dict[str, Any]:
50+
"""
51+
Convert the request parameters to a JSON string.
52+
53+
Returns:
54+
Dict[str, Any]: The parameters as a JSON string
55+
"""
56+
return json.dumps(
57+
{
58+
'query': self.query,
59+
'max_results': self.num_results,
60+
'sort_strategy': self.sort_strategy.value,
61+
'sort_order': self.sort_order.value
62+
},
63+
ensure_ascii=False)
64+
65+
66+
class ArxivSearchResult(SearchResult):
67+
"""ArXiv search result implementation."""
68+
69+
def __init__(self,
70+
query: str,
71+
arguments: Dict[str, Any] = None,
72+
response: List['arxiv.Result'] = None):
73+
"""
74+
Initialize ArxivSearchResult.
75+
76+
Args:
77+
query: The original search query string
78+
arguments: The arguments used for the search
79+
response: The raw results returned by the search
80+
"""
81+
super().__init__(query, arguments, response)
82+
self.arguments = self._process_arguments()
83+
self.response = self._process_results()
84+
85+
def _process_results(self) -> SearchResponse:
86+
"""
87+
Process the raw results into a standardized format.
88+
89+
Returns:
90+
SearchResponse: Processed search results
91+
"""
92+
if isinstance(self.response, Generator):
93+
self.response = list(self.response)
94+
95+
if not self.response:
96+
print(
97+
'***Warning: No search results found. This may happen because '
98+
'Arxiv\'s search functionality relies on precise metadata matching (e.g., title, '
99+
'author, abstract keywords) rather than the full-text indexing and complex '
100+
'ranking algorithms used by search engines like Google, or the semantic search '
101+
'capabilities of some neural search engines. The search query rewritten by the '
102+
'model may not align perfectly with Arxiv\'s metadata-driven engine. For a more '
103+
'robust and stable search experience, consider configuring an advanced search '
104+
'provider (such as Exa, SerpApi, etc.) in the `conf.yaml` file.'
105+
)
106+
return SearchResponse(results=[])
107+
108+
processed = []
109+
for res in self.response:
110+
if not isinstance(res, arxiv.Result):
111+
print(
112+
f'***Warning: Result {res} is not an instance of arxiv.Result.'
113+
)
114+
continue
115+
116+
processed.append(
117+
BaseResult(
118+
url=getattr(res, 'pdf_url', None)
119+
or getattr(res, 'entry_id', None),
120+
id=getattr(res, 'entry_id', None),
121+
title=getattr(res, 'title', None),
122+
highlights=None,
123+
highlight_scores=None,
124+
summary=getattr(res, 'summary', None),
125+
markdown=None))
126+
127+
return SearchResponse(results=processed)
128+
129+
def _process_arguments(self) -> Dict[str, Any]:
130+
"""Process the search arguments to be JSON serializable."""
131+
return {
132+
'query':
133+
self.query,
134+
'max_results':
135+
self.arguments.get('max_results', None),
136+
'sort_strategy':
137+
self.arguments.get('sort_strategy', SortCriterion.Relevance).value,
138+
'sort_order':
139+
self.arguments.get('sort_order', SortOrder.Descending).value
140+
}

‎ms_agent/tools/search/arxiv/search.py‎

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
# flake8: noqa
2+
import os
3+
4+
import arxiv
5+
from ms_agent.tools.search.arxiv.schema import (ArxivSearchRequest,
6+
ArxivSearchResult)
7+
from ms_agent.tools.search.search_base import SearchEngine, SearchEngineType
8+
9+
10+
class ArxivSearch(SearchEngine):
11+
"""
12+
A class to perform web searches using the arxiv service.
13+
"""
14+
15+
def __init__(self):
16+
17+
self.client = arxiv.Client()
18+
self.engine_type = SearchEngineType.ARXIV
19+
20+
def search(self, search_request: ArxivSearchRequest) -> ArxivSearchResult:
21+
"""Perform a search using arxiv and return the results."""
22+
search_args: dict = search_request.to_dict()
23+
24+
try:
25+
response = list(
26+
self.client.results(search=arxiv.Search(**search_args)))
27+
search_result = ArxivSearchResult(
28+
query=search_request.query,
29+
arguments=search_args,
30+
response=response)
31+
except Exception as e:
32+
raise RuntimeError(f'Failed to perform search: {e}') from e
33+
34+
return search_result

‎ms_agent/tools/search/search_base.py‎

Lines changed: 154 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,154 @@
1+
# flake8: noqa
2+
import enum
3+
import os
4+
from abc import ABC, abstractmethod
5+
from dataclasses import dataclass
6+
from typing import Any, Dict, Generic, List, Optional, TypeVar
7+
8+
import json
9+
10+
T = TypeVar('T')
11+
12+
13+
class SearchEngineType(enum.Enum):
14+
EXA = 'exa'
15+
SERPAPI = 'serpapi'
16+
ARXIV = 'arxiv'
17+
18+
19+
@dataclass
20+
class BaseResult:
21+
"""A class representing the base fields of a search result.
22+
23+
Attributes:
24+
url (str): The URL of the search result.
25+
id (str): The temporary ID for the document.
26+
title (str): The title of the search result.
27+
highlights (Optional[List[str]]): Highlights from the search result.
28+
highlight_scores (Optional[List[float]]): Scores for the highlights.
29+
summary (Optional[str]): A summary of the search result.
30+
markdown (Optional[str]): Markdown content of the search result.
31+
"""
32+
33+
url: Optional[str] = None
34+
id: Optional[str] = None
35+
title: Optional[str] = None
36+
highlights: Optional[List[str]] = None
37+
highlight_scores: Optional[List[float]] = None
38+
summary: Optional[str] = None
39+
markdown: Optional[str] = None
40+
41+
42+
@dataclass
43+
class SearchResponse(Generic[T]):
44+
"""Base class for search responses."""
45+
46+
# A list of search results.
47+
results: List[T]
48+
49+
50+
class SearchRequest(ABC):
51+
"""Abstract base class for search requests."""
52+
53+
def __init__(self, query: str, num_results: Optional[int] = 10):
54+
"""
55+
Initialize SearchRequest with search parameters.
56+
57+
Args:
58+
query: The search query string
59+
num_results: Number of results to return, default is 10
60+
"""
61+
self.query = query
62+
self.num_results = num_results
63+
64+
@abstractmethod
65+
def to_dict(self) -> Dict[str, Any]:
66+
"""Convert the request parameters to a dictionary."""
67+
pass
68+
69+
def to_json(self) -> str:
70+
"""
71+
Convert the request parameters to a JSON string.
72+
73+
Returns:
74+
str: The parameters as a JSON string
75+
"""
76+
return json.dumps(self.to_dict(), ensure_ascii=False)
77+
78+
79+
class SearchResult(ABC):
80+
"""Base class for search results."""
81+
82+
def __init__(self,
83+
query: str,
84+
arguments: Optional[Dict[str, Any]] = None,
85+
response: Any = None):
86+
"""
87+
Initialize SearchResult.
88+
89+
Args:
90+
query: The original search query string
91+
arguments: The arguments used for the search
92+
response: The raw results returned by the search
93+
"""
94+
self.query = query
95+
self.arguments = arguments
96+
self.response = response
97+
98+
@abstractmethod
99+
def _process_results(self) -> SearchResponse:
100+
"""
101+
Process the raw results into a standardized format.
102+
103+
Returns:
104+
SearchResponse: Processed search results
105+
"""
106+
pass
107+
108+
def to_list(self) -> List[Dict[str, Any]]:
109+
"""
110+
Convert the search results to a list of dictionaries.
111+
"""
112+
113+
if not self.response or not self.response.results:
114+
print('***Warning: No search results found.')
115+
return []
116+
117+
if not self.query:
118+
print('***Warning: No query provided for search results.')
119+
return []
120+
121+
res_list: List[Dict[str, Any]] = []
122+
for res in self.response.results:
123+
res_list.append({
124+
'url': res.url,
125+
'id': res.id,
126+
'title': res.title,
127+
'highlights': res.highlights,
128+
'highlight_scores': res.highlight_scores,
129+
'summary': res.summary,
130+
'markdown': res.markdown,
131+
})
132+
133+
return res_list
134+
135+
@staticmethod
136+
def load_from_disk(file_path: str) -> List[Dict[str, Any]]:
137+
"""Load search results from a JSON file."""
138+
if not os.path.exists(file_path):
139+
return []
140+
141+
with open(file_path, 'r', encoding='utf-8') as f:
142+
data = json.load(f)
143+
print(f'Search results loaded from {file_path}')
144+
145+
return data
146+
147+
148+
class SearchEngine(ABC):
149+
"""Abstract base class for search engines."""
150+
151+
@abstractmethod
152+
def search(self, search_request: SearchRequest) -> SearchResult:
153+
"""Perform a search and return results."""
154+
pass

0 commit comments

Comments
 (0)