"""
Base agent class for LangGraph-based agents.
Provides common functionality for all agents including:
- LLM initialization
- Tool management
- Graph building with standard structure
- Common node implementations (_llm_call, _tool_node, _should_continue)
"""
import json
import logging
from abc import ABC, abstractmethod
from typing import Any, Literal
from langchain.chat_models import init_chat_model
from langchain_core.messages import AIMessage, AnyMessage, SystemMessage, ToolMessage
from langgraph.graph import END, START, StateGraph
from config import config
from src.checkpoint import get_checkpointer
from src.models.state import MessagesState
from src.utils.prompts import _is_eval_mode, strip_suggested_prompts
from src.utils.response_recovery import recover_local_response
logger = logging.getLogger(__name__)
# Marks the ToolMessage ``_force_end_node`` appends to close a call the forced
# stop abandoned. The supervisor's idle-loop detector asks "did any tool run
# during this delegation" by looking for ToolMessages -- and this one means the
# exact opposite, so it has to be distinguishable from a real tool result.
FORCED_STOP_MARKER = "engiai_forced_stop"
[docs]
class BaseAgent(ABC):
"""Abstract base class for LangGraph agents with common functionality."""
[docs]
def __init__(
self,
model_name: str | None = None,
tools: list | None = None,
require_confirmation: bool = False,
temperature: float | None = None,
seed: int | None = None,
):
"""Initialize the base agent.
Args:
model_name: Name of the LLM model to use (defaults to config.llm_model)
tools: List of LangChain tools for this agent
require_confirmation: Whether to require user confirmation before tool execution
temperature: Model temperature (defaults to config.llm_temperature)
seed: Random seed for model (defaults to config.llm_seed, None for non-deterministic)
"""
self.model_name = model_name or config.llm_model
self.temperature = (
temperature if temperature is not None else config.llm_temperature
)
self.seed = seed if seed is not None else config.llm_seed
logger.debug(
f"Initializing {self.__class__.__name__} with temperature={self.temperature}, seed={self.seed}"
)
# Build kwargs for model initialization
model_kwargs: dict = {"temperature": self.temperature}
# Add seed if provided
if self.seed is not None:
model_kwargs["seed"] = self.seed
# Add Ollama-specific configuration if using an Ollama model
if self.model_name.startswith("ollama:"):
model_kwargs["num_ctx"] = config.ollama_num_ctx
logger.debug(
f"Ollama model detected, using num_ctx={config.ollama_num_ctx}"
)
self.llm = init_chat_model(self.model_name, **model_kwargs)
self.require_confirmation = require_confirmation
# Initialize tools
self.tools = tools or self._create_tools()
self.tools_by_name = {tool.name: tool for tool in self.tools}
self.llm_with_tools = self.llm.bind_tools(self.tools)
# Build the agent graph
self.agent = self._build_agent()
@abstractmethod
def _create_tools(self) -> list:
"""Create the list of tools for this agent.
Returns:
List of LangChain tools
Note:
Subclasses must implement this method to define their specific tools.
"""
@abstractmethod
def _get_system_prompt(self) -> str:
"""Get the system prompt for this agent.
Returns:
System prompt string
Note:
Subclasses must implement this method to define their system prompt.
"""
def _llm_call(self, state: MessagesState) -> dict:
"""LLM decides whether to call a tool or not.
Args:
state: Current conversation state
Returns:
Updated state with LLM response
"""
prompt = self._get_system_prompt()
if _is_eval_mode():
prompt = strip_suggested_prompts(prompt)
messages: list[AnyMessage] = [SystemMessage(content=prompt)] + state["messages"]
response = self._invoke_llm_with_tools(messages)
# Some locally served agent models have learned the common ``web_fetch``
# tool name even when only EngiAI's ``web_search`` tool is bound. Repair
# that narrow alias before it reaches the graph. Searching for the URL
# plus the model's extraction prompt provides useful page snippets while
# preserving the keyless search tool's existing network safeguards.
if (
isinstance(response, AIMessage)
and response.tool_calls
and "web_search" in self.tools_by_name
):
for tool_call in response.tool_calls:
if tool_call.get("name") != "web_fetch":
continue
args = tool_call.get("args", {})
query_parts = [
str(args.get(key, "")).strip()
for key in ("prompt", "query", "url")
if str(args.get(key, "")).strip()
]
tool_call["name"] = "web_search"
tool_call["args"] = {"query": " ".join(query_parts)}
logger.warning(
"%s: repaired hallucinated web_fetch call to web_search",
self.__class__.__name__,
)
return {
"messages": [response],
"llm_calls": state.get("llm_calls", 0) + 1,
}
def _invoke_llm_with_tools(self, messages: list[AnyMessage]) -> AIMessage:
"""Recover an invisible local answer without replaying tool calls."""
response = self.llm_with_tools.invoke(messages)
return recover_local_response(
response,
messages,
model_name=self.model_name,
answer_model=self.llm,
allow_tools=bool(self.tools),
)
def _tool_node(self, state: MessagesState) -> dict:
"""Execute tool calls from the LLM.
Args:
state: Current conversation state
Returns:
Updated state with tool results
"""
result = []
last_message = state["messages"][-1]
# Only AIMessage has tool_calls
if isinstance(last_message, AIMessage) and last_message.tool_calls:
for tool_call in last_message.tool_calls:
try:
tool = self.tools_by_name[tool_call["name"]]
observation = tool.invoke(tool_call["args"])
result.append(
ToolMessage(
content=str(observation),
tool_call_id=tool_call["id"],
name=tool_call["name"],
)
)
except Exception as e:
# CRITICAL: Always return a ToolMessage, even for errors
# Otherwise OpenAI API will fail with "tool_call_id did not have response"
#
# Logged as well as returned: the agent sees only the
# formatted string and usually just retries, so without this
# a tool failing every time is indistinguishable from the
# model choosing not to use it.
# warning, not exception: the retry loop exists to absorb
# these (malformed arguments from the model are routine), and
# now that the eval harness configures logging an ERROR-level
# traceback per occurrence would drown its stderr. exc_info
# keeps the traceback for when it is a real failure.
logger.warning(
"Tool %s raised during execution: %s",
tool_call["name"],
e,
exc_info=True,
)
error_content = (
f"❌ Error executing tool '{tool_call['name']}': {e!s}"
)
result.append(
ToolMessage(
content=error_content,
tool_call_id=tool_call["id"],
name=tool_call["name"],
)
)
return {"messages": result}
def _should_force_end(self, state: MessagesState) -> bool: # noqa: ARG002
"""Hook: force the agent loop to END regardless of any pending tool calls.
Returns ``False`` by default (no forced termination). Subclasses override it
to add an agent-specific stop condition -- e.g. ``HPCAgent`` ends the loop a
few turns after the per-trial submission budget is exhausted, so a
self-correcting agent stops churning instead of running until the recursion
limit (which would otherwise kill the whole trial).
"""
return False
def _force_end_node(self, state: MessagesState) -> dict:
"""Answer any tool calls the forced stop is cutting off, then end.
``_should_force_end`` deliberately preempts a pending tool call --
``HPCAgent`` inspects the call itself to decide whether to allow it --
which leaves an AIMessage carrying ``tool_calls`` that nothing responds
to. That history is invalid: OpenAI rejects an assistant message whose
``tool_call_id``s have no matching tool response, so the *next* request
on that thread fails rather than the current one.
The state reducer is ``operator.add``, so the message cannot be replaced
in place -- the only way to close the pair is to append a response, the
same way ``_tool_node`` always answers even on error.
The response carries ``FORCED_STOP_MARKER`` because appending a
ToolMessage otherwise tells the supervisor a tool *ran*: its idle-loop
detector reads "any ToolMessage" as a productive delegation, and would
re-route to an agent that had just given up.
"""
last_message = state["messages"][-1]
if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
return {"messages": []}
logger.warning(
"%s: forced stop with %d unanswered tool call(s) %s — closing them",
self.__class__.__name__,
len(last_message.tool_calls),
[tc.get("name") for tc in last_message.tool_calls],
)
return {
"messages": [
ToolMessage(
status="error",
additional_kwargs={FORCED_STOP_MARKER: True},
content=(
"Not executed — the agent stopped before this call ran. "
"Do not assume any effect from it."
),
tool_call_id=tool_call["id"],
name=tool_call["name"],
)
for tool_call in last_message.tool_calls
]
}
def _should_continue(
self, state: MessagesState
) -> Literal["tool_node", "force_end", "__end__"]:
"""Decide whether to continue to tool execution or end.
Args:
state: Current conversation state
Returns:
Next node to execute ("tool_node", "force_end" or "__end__")
"""
messages = state["messages"]
last_message = messages[-1]
# Agent-specific hard stop (default no-op). Lets a subclass terminate a
# runaway loop -- e.g. churning after the submission budget is exhausted.
if self._should_force_end(state):
# Route via force_end when the stop is cutting a tool call off, so
# the abandoned call still gets a response. Plain END otherwise.
if isinstance(last_message, AIMessage) and last_message.tool_calls:
return "force_end"
return "__end__"
# If the LLM makes a tool call, then perform an action
if isinstance(last_message, AIMessage) and last_message.tool_calls:
# Guard: some models (e.g. Gemini) hallucinate tool calls even when
# bind_tools([]) was used and no tool definitions were provided.
# Routing to tool_node would just produce KeyError ToolMessages and
# cause an infinite loop. End immediately instead.
if not self.tools:
logger.warning(
"%s: LLM emitted tool calls but agent has no tools — ignoring",
self.__class__.__name__,
)
return "__end__"
return "tool_node"
# Otherwise, we stop (reply to the user)
return "__end__"
def _after_tools(
self, state: MessagesState
) -> Literal["llm_call", "clarification_response", "__end__"]:
"""Route after tool execution.
If the current batch of tool results contains a successful
ask_human_for_clarification call, route to clarification_response
which synthesizes an AIMessage from the tool result — no additional
LLM round-trip needed. This avoids a ~30-60s unnecessary LLM call.
Args:
state: Current conversation state
Returns:
Next node to execute
"""
messages = state["messages"]
# Check only the current batch of ToolMessages (trailing ToolMessages
# at the end of the message list, before any non-ToolMessage).
for message in reversed(messages):
if not isinstance(message, ToolMessage):
break
if message.name != "ask_human_for_clarification":
continue
try:
content = message.content
if not isinstance(content, str):
continue
payload = json.loads(content)
if payload.get("success") is True:
return "clarification_response"
except (json.JSONDecodeError, AttributeError):
pass
return "llm_call"
def _clarification_response(self, state: MessagesState) -> dict:
"""Synthesize an AIMessage from the clarification tool result.
Extracts the question from the last successful
ask_human_for_clarification ToolMessage and returns it as a proper
AIMessage so the user sees the question. No LLM call needed.
"""
for message in reversed(state["messages"]):
if not isinstance(message, ToolMessage):
break
if message.name != "ask_human_for_clarification":
continue
try:
content = message.content
if isinstance(content, str):
payload = json.loads(content)
question = payload.get("question", "")
if question:
return {"messages": [AIMessage(content=question)]}
except (json.JSONDecodeError, AttributeError):
pass
return {"messages": []}
def _build_agent(self) -> Any:
"""Build the agent workflow graph.
Returns:
Compiled agent graph
"""
# Build workflow
agent_builder = StateGraph(MessagesState)
# Add nodes
agent_builder.add_node("llm_call", self._llm_call)
agent_builder.add_node("tool_node", self._tool_node)
agent_builder.add_node("clarification_response", self._clarification_response)
agent_builder.add_node("force_end", self._force_end_node)
# Add edges to connect nodes
agent_builder.add_edge(START, "llm_call")
agent_builder.add_conditional_edges(
"llm_call", self._should_continue, ["tool_node", "force_end", END]
)
agent_builder.add_edge("force_end", END)
agent_builder.add_conditional_edges(
"tool_node",
self._after_tools,
["llm_call", "clarification_response", END],
)
agent_builder.add_edge("clarification_response", END)
# Compile with persistent checkpointer for conversation memory
checkpointer = get_checkpointer()
# Use interrupt_before for human-in-the-loop confirmation if required
if self.require_confirmation:
return agent_builder.compile(
checkpointer=checkpointer,
interrupt_before=["tool_node"],
)
else:
return agent_builder.compile(checkpointer=checkpointer)
[docs]
def invoke(self, state: MessagesState, config: dict | None = None) -> MessagesState:
"""Invoke the agent with a given state.
Args:
state: Input conversation state
config: Optional configuration (e.g., thread_id for conversation tracking)
Returns:
Updated conversation state
"""
return self.agent.invoke(state, config) # type: ignore[return-value]
