105 lines
3.2 KiB
Python
105 lines
3.2 KiB
Python
"""
|
||
Search Agent 核心模块 - A2A版本
|
||
|
||
基于LiteLLM和A2A协议的搜索Agent实现
|
||
"""
|
||
import sys
|
||
import os
|
||
|
||
# 先导入当前目录的config
|
||
from typing import Optional
|
||
from loguru import logger
|
||
from config import LiteLLMConfig, AgentConfig, get_config
|
||
|
||
# 然后添加search_agent目录到Python路径
|
||
search_agent_dir = os.path.join(os.path.dirname(__file__), 'search_agent')
|
||
if search_agent_dir not in sys.path:
|
||
sys.path.insert(0, search_agent_dir)
|
||
|
||
from search_agent.config import Config
|
||
from search_agent.agent.search_agent import SearchAgent as CoreSearchAgent
|
||
|
||
|
||
class SearchAgentWrapper:
|
||
"""
|
||
Search Agent包装器
|
||
|
||
用于适配A2A框架，将SearchAgent包装为可配置的Agent实例
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
litellm_config: Optional[LiteLLMConfig] = None,
|
||
agent_config: Optional[AgentConfig] = None,
|
||
api_key: Optional[str] = None,
|
||
model: Optional[str] = None
|
||
):
|
||
"""
|
||
初始化Search Agent
|
||
|
||
Args:
|
||
litellm_config: LiteLLM配置对象
|
||
agent_config: Agent配置对象
|
||
api_key: LiteLLM API密钥（可选，优先使用）
|
||
model: 模型名称（可选，优先使用）
|
||
"""
|
||
# 获取配置
|
||
if not litellm_config:
|
||
llm_config, _, _ = get_config(api_key=api_key, model=model)
|
||
else:
|
||
llm_config = litellm_config
|
||
|
||
if not agent_config:
|
||
_, agent_config, _ = get_config(api_key=api_key, model=model)
|
||
|
||
self.litellm_config = llm_config
|
||
self.agent_config = agent_config
|
||
|
||
# 验证配置
|
||
self.litellm_config.validate()
|
||
|
||
# 创建SearchAgent配置（使用litellm的base_url和api_key）
|
||
# 需要从环境变量获取其他配置
|
||
serper_api_key = os.getenv("SERPER_API_KEY", "")
|
||
jina_api_key = os.getenv("JINA_API_KEY", "")
|
||
|
||
self.search_config = Config(
|
||
llm_base_url=llm_config.base_url,
|
||
llm_api_key=llm_config.api_key,
|
||
llm_model=llm_config.model,
|
||
serper_api_key=serper_api_key,
|
||
jina_api_key=jina_api_key,
|
||
max_iterations=int(os.getenv("MAX_ITERATIONS", "3")),
|
||
max_results_per_query=int(os.getenv("MAX_RESULTS_PER_QUERY", "10")),
|
||
content_max_length=int(os.getenv("CONTENT_MAX_LENGTH", "5000")),
|
||
log_level=os.getenv("LOG_LEVEL", "INFO"),
|
||
timeout=int(os.getenv("TIMEOUT", "30"))
|
||
)
|
||
|
||
# 创建SearchAgent实例
|
||
self.agent = CoreSearchAgent(self.search_config)
|
||
|
||
logger.info(
|
||
"SearchAgent初始化完成",
|
||
agent_name=self.agent_config.name,
|
||
model=self.litellm_config.model,
|
||
base_url=self.litellm_config.base_url
|
||
)
|
||
|
||
async def search(self, query: str):
|
||
"""
|
||
执行搜索
|
||
|
||
Args:
|
||
query: 搜索查询
|
||
|
||
Returns:
|
||
AgentResponse对象
|
||
"""
|
||
return await self.agent.search(query)
|
||
|
||
async def close(self):
|
||
"""关闭资源（SearchAgent不需要特殊清理）"""
|
||
pass
|
||
|