VinogradovAI's picture
Update agent.py
ff8f4bc verified
# agent.py
import argparse
from langchain_openai.chat_models import ChatOpenAI
from langchain.agents import initialize_agent
from langchain.tools import Tool
from config import config
from tools.gaia_tool import GaiaQATool
from tools.wiki_tool import WikiSearchTool
from tools.youtube_species_tool import YouTubeSpeciesTool
from tools.youtube_transcription_tool import YouTubeTranscriptionTool
from tools.excel_tool import ExcelTool
from tools.codegen_tool import CodeGenTool
from tools.web_search_tool import WebSearchTool
from tools.string_tool import StringTool
from tools.chess_tool import ChessTool
from tools.group_tool import GroupTool
from tools.grocery_tool import GroceryTool
def main() -> None:
parser = argparse.ArgumentParser(description="Многофункциональный QA агент")
parser.add_argument(
"--input", required=True,
help="Вопрос, URL или команда для инструмента"
)
args = parser.parse_args()
user_input = args.input
# Инициализируем LLM
llm = ChatOpenAI(
model_name=config.OPENAI_MODEL,
openai_api_key=config.OPENAI_API_KEY,
temperature=0.2
)
# Создаём и сохраняем инстансы инструментов
gaia = GaiaQATool()
wiki = WikiSearchTool()
species = YouTubeSpeciesTool()
transcript = YouTubeTranscriptionTool()
excel = ExcelTool()
codegen = CodeGenTool()
web = WebSearchTool()
string_tool = StringTool()
chess = ChessTool()
group = GroupTool()
grocery = GroceryTool()
# Собираем список Tool(...) из инстансов
tools = [
Tool(name=gaia.name, func=gaia._run, description=gaia.description),
Tool(name=wiki.name, func=wiki._run, description=wiki.description),
Tool(name=species.name, func=species._run, description=species.description),
Tool(name=transcript.name, func=transcript._run, description=transcript.description),
Tool(name=excel.name, func=excel._run, description=excel.description),
Tool(name=codegen.name, func=codegen._run, description=codegen.description),
Tool(name=web.name, func=web._run, description=web.description),
Tool(name=string_tool.name, func=string_tool._run, description=string_tool.description),
Tool(name=chess.name, func=chess._run, description=chess.description),
Tool(name=group.name, func=group._run, description=group.description),
Tool(name=grocery.name, func=grocery._run, description=grocery.description),
]
# Инициализируем агента в режиме zero-shot-react-description
agent = initialize_agent(
tools=tools,
llm=llm,
agent="zero-shot-react-description",
verbose=True,
max_iterations=30,
early_stopping_method="generate", # или "generate"
handle_parsing_errors=True # при ошибках парсинга пробует ещё раз
)
# Запускаем и получаем ответ
answer = agent.invoke(user_input)
print("\nОтвет:\n", answer)
if __name__ == "__main__":
main()