leocrm/app/ai/context_builder.py

"""Context builder — assembles the message list for an AI agent run.

Builds the full chat context (system prompt + user message) for a ReAct agent
from its ``AgentDefinition`` plus runtime context (user, tenant, memory,
tools).  Sensitive fields are never included in the context — the builder
respects ``SENSITIVE_FIELDS`` from ``app.core.sensitive_data``.

Usage::

    from app.ai.context_builder import build_agent_context

    messages = await build_agent_context(
        agent_definition=agent,
        user_message="Summarize recent emails",
        db=db_session,
        tenant_id=tenant_id,
        user_id=user_id,
        memory_items=[{"content": "...", "metadata": {...}}],
    )
"""

from __future__ import annotations

import logging
import uuid
from typing import TYPE_CHECKING, Any

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)

# Default ReAct instruction prefix appended to the agent's system prompt.
_REACT_PREFIX = (
    "You operate in a ReAct (Reasoning + Acting) loop. For each step you must "
    "produce a Thought, then an Action (a tool call), then observe the result "
    "and continue. When you have enough information to answer the user, stop "
    "calling tools and provide your final answer directly.\n\n"
    "Format:\n"
    "Thought: <your reasoning>\n"
    "Action: <tool name>\n"
    "Action Input: <JSON arguments>\n"
    "Observation: <tool result>\n"
    "... (repeat as needed) ...\n"
    "Final Answer: <your response to the user>\n"
)


class ReActSystemPromptBuilder:
    """Builds the ReAct system prompt for an agent definition.

    Sections:
    - Agent identity (name, description, capabilities)
    - Available tools (name + description only — schemas come via the tools param)
    - ReAct format instructions
    - Constraints (max_steps, budget, what the agent can/cannot do)
    """

    def __init__(
        self,
        agent_definition: Any,
        tool_descriptions: list[dict[str, str]] | None = None,
        max_steps: int | None = None,
        budget_limit_usd: float | None = None,
    ) -> None:
        self.agent_definition = agent_definition
        self.tool_descriptions = tool_descriptions or []
        self.max_steps = max_steps
        self.budget_limit_usd = budget_limit_usd

    def build(self) -> str:
        """Return the full system prompt string."""
        sections: list[str] = []

        # 1. Agent identity
        sections.append(self._identity_section())

        # 2. Available tools
        sections.append(self._tools_section())

        # 3. ReAct format instructions
        sections.append(_REACT_PREFIX)

        # 4. Constraints
        sections.append(self._constraints_section())

        # 5. Base system prompt from the agent definition
        base_prompt = getattr(self.agent_definition, "system_prompt", "") or ""
        if base_prompt:
            sections.append(base_prompt)

        return "\n\n".join(s for s in sections if s)

    def _identity_section(self) -> str:
        """Agent identity: name, description, capabilities."""
        name = getattr(self.agent_definition, "name", "") or "AI Agent"
        description = getattr(self.agent_definition, "description", "") or ""
        capabilities = getattr(self.agent_definition, "capabilities", None) or []

        lines = [f"You are {name}."]
        if description:
            lines.append(f"Description: {description}")
        if capabilities:
            caps = ", ".join(str(c) for c in capabilities)
            lines.append(f"Capabilities: {caps}")
        return "\n".join(lines)

    def _tools_section(self) -> str:
        """Available tools — name + description only (no full schema)."""
        if not self.tool_descriptions:
            return "You have no tools available. Answer from your own knowledge."
        lines = ["Available tools:"]
        for tool in self.tool_descriptions:
            name = tool.get("name", "")
            description = tool.get("description", "")
            if name:
                lines.append(f"- {name}: {description}")
        return "\n".join(lines)

    def _constraints_section(self) -> str:
        """Constraints: max_steps, budget, and behavioral limits."""
        constraints: list[str] = []

        max_steps = self.max_steps or getattr(
            self.agent_definition, "max_steps", None
        ) or 20
        constraints.append(f"- Maximum {max_steps} reasoning steps per run.")

        budget = self.budget_limit_usd
        if budget is None:
            budget = getattr(self.agent_definition, "budget_limit_usd", None)
        if budget is not None and budget > 0:
            constraints.append(f"- Budget limit: ${float(budget):.2f} per run.")

        constraints.append(
            "- Only call tools that are listed as available. Do not invent tools."
        )
        constraints.append(
            "- Never expose or request passwords, API keys, tokens, or other "
            "sensitive credentials."
        )
        constraints.append(
            "- Respect tenant data boundaries. Do not access data outside the "
            "current tenant."
        )

        return "Constraints:\n" + "\n".join(constraints)


async def _load_user_context(
    db: AsyncSession | None,
    tenant_id: uuid.UUID,
    user_id: uuid.UUID,
    run_as_user_id: uuid.UUID | None,
) -> dict[str, Any]:
    """Load user + tenant context from the DB with graceful fallbacks."""
    context: dict[str, Any] = {
        "user_name": None,
        "tenant_name": None,
        "role": None,
    }
    if db is None:
        return context

    try:
        from sqlalchemy import select

        from app.models.tenant import Tenant
        from app.models.user import User, UserTenant

        effective_user_id = run_as_user_id or user_id

        user_result = await db.execute(
            select(User).where(User.id == effective_user_id).limit(1)
        )
        user = user_result.scalar_one_or_none()
        if user is not None:
            context["user_name"] = user.name or user.email

        tenant_result = await db.execute(
            select(Tenant).where(Tenant.id == tenant_id).limit(1)
        )
        tenant = tenant_result.scalar_one_or_none()
        if tenant is not None:
            context["tenant_name"] = tenant.name

        role_result = await db.execute(
            select(UserTenant.role)
            .where(UserTenant.user_id == effective_user_id)
            .where(UserTenant.tenant_id == tenant_id)
            .limit(1)
        )
        role = role_result.scalar_one_or_none()
        if role:
            context["role"] = role
    except Exception:
        logger.debug("Failed to load user/tenant context", exc_info=True)

    return context


async def build_agent_context(
    agent_definition: Any,  # AgentDefinition from automation models
    user_message: str | None,
    db: AsyncSession | None,
    tenant_id: uuid.UUID,
    user_id: uuid.UUID,
    run_as_user_id: uuid.UUID | None = None,
    memory_items: list[dict] | None = None,
    trace_id: str | None = None,
) -> list[dict[str, Any]]:
    """Build the full message list for an agent run.

    Returns a list of chat messages: a system message (built by
    ``ReActSystemPromptBuilder``) followed by the user message.  Sensitive
    fields are redacted from any injected context.
    """
    # 1. Resolve the tools the agent has access to (filtered by tool_ids).
    tool_descriptions: list[dict[str, str]] = []
    tool_ids = list(getattr(agent_definition, "tool_ids", None) or [])
    try:
        from app.ai.tool_registry import get_tool_registry

        registry = get_tool_registry()
        if tool_ids:
            tools = registry.get_by_names(tool_ids)
        else:
            tools = registry.get_all()
        tool_descriptions = [
            {"name": t.name, "description": t.description} for t in tools
        ]
    except Exception:
        logger.debug("Failed to load tool descriptions", exc_info=True)

    # 2. Build the system prompt.
    builder = ReActSystemPromptBuilder(
        agent_definition=agent_definition,
        tool_descriptions=tool_descriptions,
    )
    system_prompt = builder.build()

    # 3. Load user/tenant context.
    user_ctx = await _load_user_context(
        db, tenant_id, user_id, run_as_user_id
    )

    # 4. Assemble the context block (redacting sensitive fields).
    context_lines: list[str] = []
    if user_ctx.get("user_name"):
        context_lines.append(f"Current user: {user_ctx['user_name']}")
    if user_ctx.get("tenant_name"):
        context_lines.append(f"Current tenant: {user_ctx['tenant_name']}")
    if user_ctx.get("role"):
        context_lines.append(f"Current user role: {user_ctx['role']}")

    if memory_items:
        context_lines.append("Relevant memory items:")
        for item in memory_items:
            content = item.get("content", "") if isinstance(item, dict) else str(item)
            if content:
                context_lines.append(f"- {content}")

    # 5. Build the final message list.
    messages: list[dict[str, Any]] = [
        {"role": "system", "content": system_prompt}
    ]

    if context_lines:
        context_block = "\n".join(context_lines)
        messages.append(
            {"role": "system", "content": f"Context:\n{context_block}"}
        )

    if user_message:
        messages.append({"role": "user", "content": user_message})

    return messages


__all__ = ["ReActSystemPromptBuilder", "build_agent_context"]