mini refactors

Author: dartpain

application/agents/agentic_agent.py

@@ -4,8 +4,7 @@
 from application.agents.base import BaseAgent
 from application.agents.tools.internal_search import (
     INTERNAL_TOOL_ID,
-    build_internal_tool_config,
-    build_internal_tool_entry,
+    add_internal_search_tool,
 )
 from application.logging import LogContext
 
@@ -32,24 +31,8 @@ def __init__(
     def _gen_inner(
         self, query: str, log_context: LogContext
     ) -> Generator[Dict, None, None]:
-        # 1. Get user tools (same as ClassicAgent)
         tools_dict = self.tool_executor.get_tools()
-
-        # 2. Add internal search as a synthetic tool (only if sources are configured)
-        source = self.retriever_config.get("source", {})
-        has_sources = bool(source.get("active_docs"))
-        if self.retriever_config and has_sources:
-            has_dir = _sources_have_directory_structure(source)
-            internal_entry = build_internal_tool_entry(
-                has_directory_structure=has_dir
-            )
-            internal_entry["config"] = build_internal_tool_config(
-                **self.retriever_config,
-                has_directory_structure=has_dir,
-            )
-            tools_dict[INTERNAL_TOOL_ID] = internal_entry
-
-        # 3. Prepare all tools for the LLM
+        add_internal_search_tool(tools_dict, self.retriever_config)
         self._prepare_tools(tools_dict)
 
         # 4. Build messages (prompt has NO pre-fetched docs)
@@ -78,40 +61,3 @@ def _collect_internal_sources(self):
         tool = self.tool_executor._loaded_tools.get(cache_key)
         if tool and hasattr(tool, "retrieved_docs") and tool.retrieved_docs:
             self.retrieved_docs = tool.retrieved_docs
-
-
-def _sources_have_directory_structure(source: Dict) -> bool:
-    """Check if any of the active sources have directory_structure in MongoDB."""
-    active_docs = source.get("active_docs", [])
-    if not active_docs:
-        return False
-
-    try:
-        from bson.objectid import ObjectId
-        from application.core.mongo_db import MongoDB
-
-        mongo = MongoDB.get_client()
-        db = mongo[settings.MONGO_DB_NAME]
-        sources_collection = db["sources"]
-
-        if isinstance(active_docs, str):
-            active_docs = [active_docs]
-
-        for doc_id in active_docs:
-            try:
-                source_doc = sources_collection.find_one(
-                    {"_id": ObjectId(doc_id)},
-                    {"directory_structure": 1},
-                )
-                if source_doc and source_doc.get("directory_structure"):
-                    return True
-            except Exception:
-                continue
-    except Exception as e:
-        logger.debug(f"Could not check directory structure: {e}")
-
-    return False
-
-
-# Import settings at module level for _sources_have_directory_structure
-from application.core.settings import settings  # noqa: E402

application/agents/research_agent.py

@@ -6,11 +6,9 @@
 
 from application.agents.base import BaseAgent
 from application.agents.tool_executor import ToolExecutor
-from application.agents.agentic_agent import _sources_have_directory_structure
 from application.agents.tools.internal_search import (
     INTERNAL_TOOL_ID,
-    build_internal_tool_config,
-    build_internal_tool_entry,
+    add_internal_search_tool,
 )
 from application.agents.tools.think import THINK_TOOL_ENTRY, THINK_TOOL_ID
 from application.logging import LogContext
@@ -130,6 +128,7 @@ def __init__(
         self.citations = CitationManager()
         self._start_time: float = 0
         self._tokens_used: int = 0
+        self._last_token_snapshot: int = 0
 
     # ------------------------------------------------------------------
     # Budget & timeout helpers
@@ -153,7 +152,9 @@ def _is_over_budget(self) -> bool:
     def _snapshot_llm_tokens(self) -> int:
         """Read current token usage from LLM and return delta since last snapshot."""
         current = self.llm.token_usage.get("prompt_tokens", 0) + self.llm.token_usage.get("generated_tokens", 0)
-        return current
+        delta = current - self._last_token_snapshot
+        self._last_token_snapshot = current
+        return delta
 
     # ------------------------------------------------------------------
     # Main orchestration
@@ -272,21 +273,7 @@ def _setup_tools(self) -> Dict:
         """Build tools_dict with user tools + internal search + think."""
         tools_dict = self.tool_executor.get_tools()
 
-        # Only add internal search if sources are configured
-        source = self.retriever_config.get("source", {})
-        has_sources = bool(source.get("active_docs"))
-        if self.retriever_config and has_sources:
-            has_dir = _sources_have_directory_structure(source)
-            internal_entry = build_internal_tool_entry(
-                has_directory_structure=has_dir
-            )
-            internal_entry["config"] = build_internal_tool_config(
-                **self.retriever_config,
-                has_directory_structure=has_dir,
-            )
-            tools_dict[INTERNAL_TOOL_ID] = internal_entry
-        elif self.retriever_config and not has_sources:
-            logger.info("ResearchAgent: No sources configured, skipping internal_search tool")
+        add_internal_search_tool(tools_dict, self.retriever_config)
 
         think_entry = dict(THINK_TOOL_ENTRY)
         think_entry["config"] = {}
@@ -580,7 +567,14 @@ def _execute_step_tools_with_refinement(
             call_id = None
             while True:
                 try:
-                    next(gen)
+                    event = next(gen)
+                    # Log tool_call status events instead of discarding them
+                    if isinstance(event, dict) and event.get("type") == "tool_call":
+                        logger.debug(
+                            "Tool %s status: %s",
+                            event.get("data", {}).get("action_name", ""),
+                            event.get("data", {}).get("status", ""),
+                        )
                 except StopIteration as e:
                     result, call_id = e.value
                     break

application/agents/tools/internal_search.py

@@ -354,6 +354,59 @@ def build_internal_tool_entry(has_directory_structure: bool = False) -> Dict:
 INTERNAL_TOOL_ENTRY = build_internal_tool_entry(has_directory_structure=False)
 
 
+def sources_have_directory_structure(source: Dict) -> bool:
+    """Check if any of the active sources have directory_structure in MongoDB."""
+    active_docs = source.get("active_docs", [])
+    if not active_docs:
+        return False
+
+    try:
+        from bson.objectid import ObjectId
+        from application.core.mongo_db import MongoDB
+
+        mongo = MongoDB.get_client()
+        db = mongo[settings.MONGO_DB_NAME]
+        sources_collection = db["sources"]
+
+        if isinstance(active_docs, str):
+            active_docs = [active_docs]
+
+        for doc_id in active_docs:
+            try:
+                source_doc = sources_collection.find_one(
+                    {"_id": ObjectId(doc_id)},
+                    {"directory_structure": 1},
+                )
+                if source_doc and source_doc.get("directory_structure"):
+                    return True
+            except Exception:
+                continue
+    except Exception as e:
+        logger.debug(f"Could not check directory structure: {e}")
+
+    return False
+
+
+def add_internal_search_tool(tools_dict: Dict, retriever_config: Dict) -> None:
+    """Add the internal search tool to tools_dict if sources are configured.
+
+    Shared by AgenticAgent and ResearchAgent to avoid duplicate setup logic.
+    Mutates tools_dict in place.
+    """
+    source = retriever_config.get("source", {})
+    has_sources = bool(source.get("active_docs"))
+    if not retriever_config or not has_sources:
+        return
+
+    has_dir = sources_have_directory_structure(source)
+    internal_entry = build_internal_tool_entry(has_directory_structure=has_dir)
+    internal_entry["config"] = build_internal_tool_config(
+        **retriever_config,
+        has_directory_structure=has_dir,
+    )
+    tools_dict[INTERNAL_TOOL_ID] = internal_entry
+
+
 def build_internal_tool_config(
     source: Dict,
     retriever_name: str = "classic",

application/agents/workflows/node_agent.py

@@ -2,8 +2,10 @@
 
 from typing import Any, Dict, List, Optional, Type
 
+from application.agents.agentic_agent import AgenticAgent
 from application.agents.base import BaseAgent
 from application.agents.classic_agent import ClassicAgent
+from application.agents.research_agent import ResearchAgent
 from application.agents.workflows.schemas import AgentType
 
 
@@ -35,7 +37,8 @@ def _get_tools(self, api_key: str = None) -> Dict[str, Dict[str, Any]]:
         return filtered_tools
 
 
-class WorkflowNodeClassicAgent(ToolFilterMixin, ClassicAgent):
+class _WorkflowNodeMixin:
+    """Common __init__ for all workflow node agents."""
 
     def __init__(
         self,
@@ -56,11 +59,25 @@ def __init__(
         self._allowed_tool_ids = tool_ids or []
 
 
+class WorkflowNodeClassicAgent(ToolFilterMixin, _WorkflowNodeMixin, ClassicAgent):
+    pass
+
+
+class WorkflowNodeAgenticAgent(ToolFilterMixin, _WorkflowNodeMixin, AgenticAgent):
+    pass
+
+
+class WorkflowNodeResearchAgent(ToolFilterMixin, _WorkflowNodeMixin, ResearchAgent):
+    pass
+
+
 class WorkflowNodeAgentFactory:
 
     _agents: Dict[AgentType, Type[BaseAgent]] = {
         AgentType.CLASSIC: WorkflowNodeClassicAgent,
         AgentType.REACT: WorkflowNodeClassicAgent,  # backwards compat
+        AgentType.AGENTIC: WorkflowNodeAgenticAgent,
+        AgentType.RESEARCH: WorkflowNodeResearchAgent,
     }
 
     @classmethod

application/agents/workflows/schemas.py

@@ -18,6 +18,8 @@ class NodeType(str, Enum):
 class AgentType(str, Enum):
     CLASSIC = "classic"
     REACT = "react"
+    AGENTIC = "agentic"
+    RESEARCH = "research"
 
 
 class ExecutionStatus(str, Enum):

application/agents/workflows/workflow_engine.py

@@ -7,6 +7,7 @@
 from application.agents.workflows.node_agent import WorkflowNodeAgentFactory
 from application.agents.workflows.schemas import (
     AgentNodeConfig,
+    AgentType,
     ConditionNodeConfig,
     ExecutionStatus,
     NodeExecutionLog,
@@ -223,18 +224,32 @@ def _execute_agent_node(
                     f'Model "{node_model_id}" does not support structured output for node "{node.title}"'
                 )
 
-        node_agent = WorkflowNodeAgentFactory.create(
-            agent_type=node_config.agent_type,
-            endpoint=self.agent.endpoint,
-            llm_name=node_llm_name,
-            model_id=node_model_id,
-            api_key=node_api_key,
-            tool_ids=node_config.tools,
-            prompt=node_config.system_prompt,
-            chat_history=self.agent.chat_history,
-            decoded_token=self.agent.decoded_token,
-            json_schema=node_json_schema,
-        )
+        factory_kwargs = {
+            "agent_type": node_config.agent_type,
+            "endpoint": self.agent.endpoint,
+            "llm_name": node_llm_name,
+            "model_id": node_model_id,
+            "api_key": node_api_key,
+            "tool_ids": node_config.tools,
+            "prompt": node_config.system_prompt,
+            "chat_history": self.agent.chat_history,
+            "decoded_token": self.agent.decoded_token,
+            "json_schema": node_json_schema,
+        }
+
+        # Agentic/research agents need retriever_config for on-demand search
+        if node_config.agent_type in (AgentType.AGENTIC, AgentType.RESEARCH):
+            factory_kwargs["retriever_config"] = {
+                "source": {"active_docs": node_config.sources} if node_config.sources else {},
+                "retriever_name": node_config.retriever or "classic",
+                "chunks": int(node_config.chunks) if node_config.chunks else 2,
+                "model_id": node_model_id,
+                "llm_name": node_llm_name,
+                "api_key": node_api_key,
+                "decoded_token": self.agent.decoded_token,
+            }
+
+        node_agent = WorkflowNodeAgentFactory.create(**factory_kwargs)
 
         full_response_parts: List[str] = []
         structured_response_parts: List[str] = []

application/prompts/agentic/creative.txt

@@ -13,4 +13,4 @@ Use the search_internal tool to find relevant information before answering quest
 You may search multiple times with different queries if needed.
 Do not guess when documents are available — search first, then answer based on what you find.
 If no relevant documents are found, use your general knowledge and tool capabilities.
-Allow yourself to be very creative and use your imagination.
+Allow yourself to be very creative and use your imagination.

application/prompts/agentic/default.txt

@@ -12,4 +12,4 @@ You have access to a search tool that searches the user's uploaded documents and
 Use the search_internal tool to find relevant information before answering questions.
 You may search multiple times with different queries if needed.
 Do not guess when documents are available — search first, then answer based on what you find.
-If no relevant documents are found, use your general knowledge and tool capabilities.
+If no relevant documents are found, use your general knowledge and tool capabilities.

application/prompts/agentic/strict.txt

@@ -13,4 +13,4 @@ Use the search_internal tool to find relevant information before answering quest
 You may search multiple times with different queries if needed.
 You MUST search before answering any factual question. Do not guess or use general knowledge when documents are available.
 If you dont have enough information from the search results or tools, answer "I don't know" or "I don't have enough information".
-Never make up information or provide false information!
+Never make up information or provide false information!

application/prompts/research/clarification.txt

@@ -20,4 +20,4 @@ You MUST respond with ONLY a valid JSON object (no markdown, no code fences):
     "needs_clarification": true or false,
     "questions": ["question 1", "question 2"] (only if needs_clarification is true, max 3 questions),
     "reason": "brief explanation of why clarification is or isn't needed"
-}
+}