56 lines
1.8 KiB
Python
56 lines
1.8 KiB
Python
# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
|
|
# SPDX-License-Identifier: MIT
|
|
|
|
import logging
|
|
from typing import List, Optional
|
|
|
|
from langgraph.prebuilt import create_react_agent
|
|
|
|
from src.config.agents import AGENT_LLM_MAP
|
|
from src.llms.llm import get_llm_by_type
|
|
from src.prompts import apply_prompt_template
|
|
from src.agents.tool_interceptor import wrap_tools_with_interceptor
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
# Create agents using configured LLM types
|
|
def create_agent(
|
|
agent_name: str,
|
|
agent_type: str,
|
|
tools: list,
|
|
prompt_template: str,
|
|
pre_model_hook: callable = None,
|
|
interrupt_before_tools: Optional[List[str]] = None,
|
|
):
|
|
"""Factory function to create agents with consistent configuration.
|
|
|
|
Args:
|
|
agent_name: Name of the agent
|
|
agent_type: Type of agent (researcher, coder, etc.)
|
|
tools: List of tools available to the agent
|
|
prompt_template: Name of the prompt template to use
|
|
pre_model_hook: Optional hook to preprocess state before model invocation
|
|
interrupt_before_tools: Optional list of tool names to interrupt before execution
|
|
|
|
Returns:
|
|
A configured agent graph
|
|
"""
|
|
# Wrap tools with interrupt logic if specified
|
|
processed_tools = tools
|
|
if interrupt_before_tools:
|
|
logger.info(
|
|
f"Creating agent '{agent_name}' with tool-specific interrupts: {interrupt_before_tools}"
|
|
)
|
|
processed_tools = wrap_tools_with_interceptor(tools, interrupt_before_tools)
|
|
|
|
return create_react_agent(
|
|
name=agent_name,
|
|
model=get_llm_by_type(AGENT_LLM_MAP[agent_type]),
|
|
tools=processed_tools,
|
|
prompt=lambda state: apply_prompt_template(
|
|
prompt_template, state, locale=state.get("locale", "en-US")
|
|
),
|
|
pre_model_hook=pre_model_hook,
|
|
)
|