arc53--docsgpt/application/agents/tools/internal_search.py

import json
import logging
from typing import Dict, List, Optional

from application.agents.tools.base import Tool
from application.core.settings import settings
from application.retriever.dispatcher import build_dispatcher
from application.retriever.retriever_creator import RetrieverCreator

logger = logging.getLogger(__name__)


class InternalSearchTool(Tool):
    """Wraps the ClassicRAG retriever as an LLM-callable tool.

    Instead of pre-fetching docs into the prompt, the LLM decides
    when and what to search. Supports multiple searches per session.

    Optional capabilities (enabled when sources have directory_structure):
    - path_filter on search: restrict results to a specific file/folder
    - list_files action: browse the file/folder structure
    """

    internal = True

    def __init__(self, config: Dict):
        self.config = config
        self.retrieved_docs: List[Dict] = []
        self._retriever = None
        self._directory_structure: Optional[Dict] = None
        self._dir_structure_loaded = False

    def _get_retriever(self):
        if self._retriever is None:
            retriever_kwargs = dict(
                source=self.config.get("source", {}),
                chat_history=[],
                prompt="",
                chunks=int(self.config.get("chunks", 2)),
                doc_token_limit=int(self.config.get("doc_token_limit", 50000)),
                model_id=self.config.get("model_id", "docsgpt-local"),
                model_user_id=self.config.get("model_user_id"),
                user_api_key=self.config.get("user_api_key"),
                agent_id=self.config.get("agent_id"),
                llm_name=self.config.get("llm_name", settings.LLM_PROVIDER),
                api_key=self.config.get("api_key", settings.API_KEY),
                decoded_token=self.config.get("decoded_token"),
                request_id=self.config.get("request_id"),
            )

            def _legacy_classic():
                return RetrieverCreator.create_retriever(
                    self.config.get("retriever_name", "classic"),
                    **retriever_kwargs,
                )

            # Dispatch per-source so on-demand agentic search honours the same
            # per-source config as pre-fetch; kill-switch falls back to legacy.
            self._retriever = build_dispatcher(
                _legacy_classic,
                sources=self.config.get("sources") or [],
                **retriever_kwargs,
            )
        return self._retriever

    def _get_directory_structure(self) -> Optional[Dict]:
        """Load directory structure from Postgres for the configured sources."""
        if self._dir_structure_loaded:
            return self._directory_structure

        self._dir_structure_loaded = True
        source = self.config.get("source", {})
        active_docs = source.get("active_docs", [])
        if not active_docs:
            return None

        try:
            # Per-operation session: this tool runs inside the answer
            # generator hot path, so we open a short-lived read
            # connection for the batch lookup and release immediately.
            from application.storage.db.repositories.sources import (
                SourcesRepository,
            )
            from application.storage.db.session import db_readonly

            if isinstance(active_docs, str):
                active_docs = [active_docs]

            decoded_token = self.config.get("decoded_token") or {}
            # Resolve the agent's sources as their OWNER: for a team-shared
            # agent run by a member, the sources belong to the owner, so using
            # the member's sub would 404. ``source_owner_id`` is the agent owner
            # (set at config-build time); fall back to the BYOM model_user_id,
            # then the invoker. Running the agent already authorized these
            # sources.
            user_id = (
                self.config.get("source_owner_id")
                or self.config.get("model_user_id")
                or (decoded_token.get("sub") if decoded_token else None)
            )

            merged_structure = {}
            with db_readonly() as conn:
                repo = SourcesRepository(conn)
                for doc_id in active_docs:
                    try:
                        source_doc = repo.get_any(str(doc_id), user_id) if user_id else None
                        if not source_doc:
                            continue
                        dir_str = source_doc.get("directory_structure")
                        if dir_str:
                            if isinstance(dir_str, str):
                                dir_str = json.loads(dir_str)
                            source_name = source_doc.get("name", doc_id)
                            if len(active_docs) > 1:
                                merged_structure[source_name] = dir_str
                            else:
                                merged_structure = dir_str
                    except Exception as e:
                        logger.debug(f"Could not load dir structure for {doc_id}: {e}")

            self._directory_structure = merged_structure if merged_structure else None
        except Exception as e:
            logger.debug(f"Failed to load directory structures: {e}")

        return self._directory_structure

    def execute_action(self, action_name: str, **kwargs):
        if action_name == "search":
            return self._execute_search(**kwargs)
        elif action_name == "list_files":
            return self._execute_list_files(**kwargs)
        return f"Unknown action: {action_name}"

    def _execute_search(self, **kwargs) -> str:
        query = kwargs.get("query", "")
        path_filter = kwargs.get("path_filter", "")

        if not query:
            return "Error: 'query' parameter is required."

        try:
            retriever = self._get_retriever()
            docs = retriever.search(query)
        except Exception as e:
            logger.error(f"Internal search failed: {e}", exc_info=True)
            return "Search failed: an internal error occurred."

        if not docs:
            return "No documents found matching your query."

        # Apply path filter if specified
        if path_filter:
            path_lower = path_filter.lower()
            docs = [
                d
                for d in docs
                if path_lower in d.get("source", "").lower()
                or path_lower in d.get("filename", "").lower()
                or path_lower in d.get("title", "").lower()
            ]
            if not docs:
                return f"No documents found matching query '{query}' in path '{path_filter}'."

        # Accumulate for source tracking
        for doc in docs:
            if doc not in self.retrieved_docs:
                self.retrieved_docs.append(doc)

        # Format results for the LLM
        formatted = []
        for i, doc in enumerate(docs, 1):
            title = doc.get("title", "Untitled")
            text = doc.get("text", "")
            source = doc.get("source", "Unknown")
            filename = doc.get("filename", "")
            header = filename or title
            formatted.append(f"[{i}] {header} (source: {source})\n{text}")

        return "\n\n---\n\n".join(formatted)

    def _execute_list_files(self, **kwargs) -> str:
        path = kwargs.get("path", "")
        dir_structure = self._get_directory_structure()

        if not dir_structure:
            return "No file structure available for the current sources."

        # Navigate to the requested path
        current = dir_structure
        if path:
            for part in path.strip("/").split("/"):
                if not part:
                    continue
                if isinstance(current, dict) and part in current:
                    current = current[part]
                else:
                    return f"Path '{path}' not found in the file structure."

        # Format the structure for the LLM
        return self._format_structure(current, path or "/")

    def _format_structure(self, node: Dict, current_path: str) -> str:
        if not isinstance(node, dict):
            return f"'{current_path}' is a file, not a directory."

        lines = [f"File structure at '{current_path}':\n"]
        folders = []
        files = []

        for name, value in sorted(node.items()):
            if isinstance(value, dict):
                # Check if it's a file metadata dict or a folder
                if "type" in value or "size_bytes" in value or "token_count" in value:
                    # It's a file with metadata
                    size = value.get("token_count", "")
                    ftype = value.get("type", "")
                    info_parts = []
                    if ftype:
                        info_parts.append(ftype)
                    if size:
                        info_parts.append(f"{size} tokens")
                    info = f" ({', '.join(info_parts)})" if info_parts else ""
                    files.append(f"  {name}{info}")
                else:
                    # It's a folder
                    count = self._count_files(value)
                    folders.append(f"  {name}/ ({count} items)")
            else:
                files.append(f"  {name}")

        if folders:
            lines.append("Folders:")
            lines.extend(folders)
        if files:
            lines.append("Files:")
            lines.extend(files)
        if not folders and not files:
            lines.append("  (empty)")

        return "\n".join(lines)

    def _count_files(self, node: Dict) -> int:
        count = 0
        for value in node.values():
            if isinstance(value, dict):
                if "type" in value or "size_bytes" in value or "token_count" in value:
                    count += 1
                else:
                    count += self._count_files(value)
            else:
                count += 1
        return count

    def get_actions_metadata(self):
        actions = [
            {
                "name": "search",
                "description": (
                    "Search the user's uploaded documents and knowledge base. "
                    "Use this before answering questions about their content. "
                    "Results include each document's source title — cite those "
                    "titles in your answer. You can call this multiple times "
                    "with different phrasings to improve coverage."
                ),
                "parameters": {
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query. Be specific and focused.",
                            "filled_by_llm": True,
                            "required": True,
                        },
                    }
                },
            }
        ]

        # Add path_filter and list_files only if directory structure exists
        has_structure = self.config.get("has_directory_structure", False)
        if has_structure:
            actions[0]["parameters"]["properties"]["path_filter"] = {
                "type": "string",
                "description": (
                    "Optional: filter results to a specific file or folder path. "
                    "Use list_files first to see available paths."
                ),
                "filled_by_llm": True,
                "required": False,
            }
            actions.append(
                {
                    "name": "list_files",
                    "description": (
                        "Browse the file and folder structure of the knowledge base. "
                        "Use this to see what files are available before searching. "
                        "Optionally provide a path to browse a specific folder."
                    ),
                    "parameters": {
                        "properties": {
                            "path": {
                                "type": "string",
                                "description": "Optional: folder path to browse. Leave empty for root.",
                                "filled_by_llm": True,
                                "required": False,
                            }
                        }
                    },
                }
            )

        return actions

    def get_config_requirements(self):
        return {}


# Constants for building synthetic tools_dict entries
INTERNAL_TOOL_ID = "internal"


def build_internal_tool_entry(has_directory_structure: bool = False) -> Dict:
    """Build the tools_dict entry for InternalSearchTool.

    Dynamically includes list_files and path_filter based on
    whether the sources have directory structure.
    """
    search_params = {
        "properties": {
            "query": {
                "type": "string",
                "description": "The search query. Be specific and focused.",
                "filled_by_llm": True,
                "required": True,
            }
        }
    }

    actions = [
        {
            "name": "search",
            "description": (
                "Search the user's uploaded documents and knowledge base. "
                "Use this to find relevant information before answering questions. "
                "You can call this multiple times with different queries."
            ),
            "active": True,
            "parameters": search_params,
        }
    ]

    if has_directory_structure:
        search_params["properties"]["path_filter"] = {
            "type": "string",
            "description": (
                "Optional: filter results to a specific file or folder path. "
                "Use list_files first to see available paths."
            ),
            "filled_by_llm": True,
            "required": False,
        }
        actions.append(
            {
                "name": "list_files",
                "description": (
                    "Browse the file and folder structure of the knowledge base. "
                    "Use this to see what files are available before searching. "
                    "Optionally provide a path to browse a specific folder."
                ),
                "active": True,
                "parameters": {
                    "properties": {
                        "path": {
                            "type": "string",
                            "description": "Optional: folder path to browse. Leave empty for root.",
                            "filled_by_llm": True,
                            "required": False,
                        }
                    }
                },
            }
        )

    return {"name": "internal_search", "actions": actions}


# Keep backward compat
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 a ``directory_structure`` row."""
    active_docs = source.get("active_docs", [])
    if not active_docs:
        return False

    try:
        # TODO(pg-cutover): SourcesRepository.get_any requires ``user_id``
        # scoping, but callers in the agent build path don't always
        # thread the decoded token through here. Use a direct
        # short-lived SQL lookup instead of the repo until the call
        # sites are updated to propagate user context.
        from sqlalchemy import text as _text

        from application.storage.db.session import db_readonly

        if isinstance(active_docs, str):
            active_docs = [active_docs]

        with db_readonly() as conn:
            for doc_id in active_docs:
                try:
                    value = str(doc_id)
                    if len(value) == 36 and "-" in value:
                        row = conn.execute(
                            _text(
                                "SELECT directory_structure FROM sources "
                                "WHERE id = CAST(:id AS uuid)"
                            ),
                            {"id": value},
                        ).fetchone()
                    else:
                        row = conn.execute(
                            _text(
                                "SELECT directory_structure FROM sources "
                                "WHERE legacy_mongo_id = :lid"
                            ),
                            {"lid": value},
                        ).fetchone()
                    if row is not None and row[0]:
                        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)
    # The executor resolves a tool row by ``id``; the internal tool is synthetic
    # (no DB row), so stamp its sentinel id or _get_or_load_tool drops it with
    # ``tool_missing_row_id``.
    internal_entry["id"] = INTERNAL_TOOL_ID
    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",
    chunks: int = 2,
    doc_token_limit: int = 50000,
    sources: Optional[List[Dict]] = None,
    model_id: str = "docsgpt-local",
    model_user_id: Optional[str] = None,
    source_owner_id: Optional[str] = None,
    user_api_key: Optional[str] = None,
    agent_id: Optional[str] = None,
    llm_name: str = None,
    api_key: str = None,
    decoded_token: Optional[Dict] = None,
    request_id: Optional[str] = None,
    has_directory_structure: bool = False,
) -> Dict:
    """Build the config dict for InternalSearchTool."""
    return {
        "source": source,
        "retriever_name": retriever_name,
        "chunks": chunks,
        "doc_token_limit": doc_token_limit,
        # Per-source list threaded through to the Dispatcher in _get_retriever.
        "sources": sources or [],
        "model_id": model_id,
        "model_user_id": model_user_id,
        # The agent owner — the sources belong to them, so directory-structure
        # resolution uses this (a team member running a shared agent has a
        # different sub). Independent of the BYOM ``model_user_id``.
        "source_owner_id": source_owner_id,
        "user_api_key": user_api_key,
        "agent_id": agent_id,
        "llm_name": llm_name or settings.LLM_PROVIDER,
        "api_key": api_key or settings.API_KEY,
        "decoded_token": decoded_token,
        "request_id": request_id,
        "has_directory_structure": has_directory_structure,
    }