replypier/backend/app/agents/chat_agent.py

"""
Copyright 2024-2026 ReplyPier

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import traceback
from typing import Optional, Type
from agno.agent import Agent
from app.utils.agno_utils import create_model
from app.core.logger import get_logger
from app.channels.constants import is_widget_channel
from app.tools.knowledge_search_byagent import KnowledgeSearchByAgent
from app.tools.mcp_manager import ChatAgentMCPMixin
from app.database import get_db, SessionLocal, engine
from app.agents.choice_prompt import choice_instructions
from app.channels.interactive_limits import InteractiveLimits
from app.agents.encrypted_storage import EncryptedPostgresAgentStorage
from app.agents.history_memory import BoundedHistoryMemory, run_with_overflow_recovery
from app.repositories.chat import ChatRepository
from app.repositories.session_to_agent import SessionToAgentRepository
from app.services.widget_identity import trusts_identity_in_session
from app.models.session_to_agent import SessionStatus
from app.models.schemas.chat import ChannelChatResponse, ChatResponse, TransferReasonType, EndChatReasonType
from app.core.config import settings
from app.agents.structured_output import (
    build_groq_json_tool,
    lenient_json_load,
    salvage_groq_json_error,
)
from app.agents.transfer_agent import get_agent_availability_response
from app.agents.guardrail_policy import (
    GuardrailContext,
    apply_guardrail_policy,
    guardrail_scope_prompt,
    wrap_operator_block,
)
from app.utils.guardrail_runtime import BLOCK_REPLY, Surface, check_inbound, check_output
from app.services.notifications import ChatNotificationEvent, notify_chat_event
from app.models.user import User, user_groups
from datetime import datetime
from app.repositories.jira import JiraRepository
from app.tools.jira_toolkit import JiraTools
from app.tools.shopify_toolkit import ShopifyTools
from app.utils.response_parser import GENERIC_ERROR_MESSAGE, parse_response_content
from app.repositories.agent_shopify_config_repository import AgentShopifyConfigRepository
import re
import asyncio
import json

logger = get_logger(__name__)

# ---------------------------------------------------------------------------
# Groq structured-output tool
#
# Groq's API rejects response_format (JSON mode) whenever tools are present
# ("json mode cannot be combined with tool/function calling"), so agno's
# structured-output path degrades to prompt-only JSON — which GPT-OSS/Llama drift
# away from (emitting the fields as prose), losing end_chat and lead capture.
#
# GPT-OSS's native structured-output convention on Groq IS a tool call named
# `json`. So instead of response_format we register a real `json` tool whose
# parameters are the ChatResponse fields and mark it stop_after_tool_call — the
# model reliably calls it (even right after a knowledge search) and we read the
# final ChatResponse straight from its validated arguments. OpenAI/Anthropic/etc.
# keep using agno's native response_model path untouched.
# ---------------------------------------------------------------------------
_GROQ_JSON_TOOL_SCHEMA = {
    "type": "object",
    "properties": {
        "message": {"type": "string", "description": "The reply text shown to the visitor."},
        "transfer_to_human": {"type": ["boolean", "null"]},
        "transfer_reason": {"type": ["string", "null"],
                            "enum": [*[t.value for t in TransferReasonType], None]},
        "transfer_description": {"type": ["string", "null"]},
        "end_chat": {"type": ["boolean", "null"], "description": "true when the conversation is over (goodbye / task complete)."},
        "end_chat_reason": {"type": ["string", "null"],
                            "enum": [*[e.value for e in EndChatReasonType], None]},
        "end_chat_description": {"type": ["string", "null"]},
        "request_rating": {"type": ["boolean", "null"]},
        "create_ticket": {"type": ["boolean", "null"]},
        "request_lead_capture": {"type": ["boolean", "null"], "description": "true once a valid email is collected (and consent, if required)."},
        "lead_email": {"type": ["string", "null"]},
        "lead_name": {"type": ["string", "null"]},
        "lead_company": {"type": ["string", "null"]},
        "lead_phone": {"type": ["string", "null"]},
        "lead_summary": {"type": ["string", "null"]},
        "lead_consent": {"type": ["boolean", "null"]},
        "request_contact": {"type": ["boolean", "null"]},
    },
    "required": ["message"],
    "additionalProperties": False,
}

_GROQ_JSON_INSTRUCTION = (
    "\n\nCRITICAL OUTPUT RULE: You MUST end every single turn by calling the `json` tool "
    "with your final structured reply. Never write the reply as plain text and never put the "
    "JSON in your message — always deliver it through the `json` tool call, after any searching. "
    "Put the visitor-facing reply in `message`. Set `end_chat`=true when the conversation is "
    "ending. Set `request_lead_capture`=true and fill `lead_email` as soon as the visitor shares "
    "a valid email (and `lead_consent`=true once they agree to be contacted)."
)


# Added to the tool schema only for agents offered choices (see ChannelChatResponse)
_GROQ_QUICK_REPLIES_PROPERTY = {
    "type": ["array", "null"], "items": {"type": "string"},
    "description": ChannelChatResponse.model_fields["quick_replies"].description,
}


def build_groq_response_tool(capture: dict, offer_choices: bool = False):
    """Return an agno `json` tool that captures the final structured turn into `capture`.

    `capture` is mutated in place with the model's tool-call arguments; the caller reads
    it after `agent.arun()` and builds the ChatResponse from it.
    """
    schema = _GROQ_JSON_TOOL_SCHEMA
    if offer_choices:
        schema = {**schema, "properties": {**schema["properties"],
                                           "quick_replies": _GROQ_QUICK_REPLIES_PROPERTY}}
    return build_groq_json_tool(
        capture,
        schema,
        description="Return your final structured reply to the visitor. Call this exactly once to end the turn, after any searching.",
    )


# Shared implementations live in app.agents.structured_output (also used by the
# FAQ generator); keep the private aliases so existing call sites and tests
# don't churn.
_lenient_json_load = lenient_json_load
_salvage_groq_json_error = salvage_groq_json_error


def _build_chat_response_from_capture(capture: dict,
                                      model: Type[ChatResponse] = ChatResponse) -> ChatResponse:
    """Build a ChatResponse from the Groq `json` tool arguments.

    Drops None values so ChatResponse's own field defaults apply (the model may
    emit explicit nulls for optional fields), then lets pydantic validate/coerce
    (e.g. enum reasons). Falls back to a plain message on validation failure.
    """
    cleaned = {k: v for k, v in capture.items() if v is not None}
    try:
        return model(**cleaned)
    except Exception as e:
        logger.error(f"Groq json-tool args failed ChatResponse validation: {e}; args={cleaned}")
        return ChatResponse(message=str(cleaned.get("message") or "").strip() or "No response generated")


def drop_stale_choices(response_content: ChatResponse, model_message: str) -> None:
    """Choices belong to the reply the model wrote: not to a canned
    replacement (the empty-turn fallback or the output guardrail swapped the
    message), nor to a goodbye."""
    if not isinstance(response_content, ChannelChatResponse):
        return
    if response_content.message != model_message or response_content.end_chat:
        response_content.quick_replies = None


_EMPTY_TURN_FALLBACK = (
    "I'm sorry, I didn't quite catch that. Could you rephrase or give me a bit more detail?"
)


def ensure_nonempty_message(response_content: ChatResponse) -> ChatResponse:
    """Guarantee a user-facing message for a turn that produced none.

    A model can end a turn with no message and no action — e.g. Groq stopping
    after a tool call without emitting the final structured response. Downstream,
    the widget only emits when message is non-empty (widget_chat), so an empty
    turn is silently dropped and the typing indicator hangs forever. When there
    is genuinely nothing to say and nothing to do, substitute a graceful reply
    so the turn always completes. Turns that carry an action (transfer, end,
    rating, ticket, contact/lead capture, shopify) are left untouched — their
    own handlers own the message.
    """
    if (response_content.message or "").strip():
        return response_content
    has_action = any([
        response_content.transfer_to_human,
        response_content.end_chat,
        response_content.request_rating,
        response_content.create_ticket,
        getattr(response_content, 'request_contact', False),
        getattr(response_content, 'request_lead_capture', False),
        getattr(response_content, 'shopify_output', None),
    ])
    if not has_action:
        logger.warning("Empty model turn with no action — using fallback reply so the chat unblocks")
        response_content.message = _EMPTY_TURN_FALLBACK
    return response_content


def _salvage_groq_answer_text(response) -> str | None:
    """The last assistant text produced during the run.

    When Groq ends a turn without calling the `json` tool, its answer often
    still exists as the final assistant message even though response.content is
    empty (reasoning models leave `content` blank). Recover that text; skip the
    tool-call turns (empty content) and never surface reasoning_content, which
    is the model's internal chain-of-thought, not a visitor reply.
    """
    messages = getattr(response, 'messages', None) or []
    for msg in reversed(messages):
        if getattr(msg, 'role', None) != 'assistant':
            continue
        content = getattr(msg, 'content', None)
        if isinstance(content, str) and content.strip():
            return content.strip()
    return None


def recover_groq_no_capture(response) -> ChatResponse:
    """Groq path only: the model finished without calling the `json` tool, so
    the structured capture is empty. Recover the assistant's own answer text
    (which parse_response_content misses because it reads only response.content).
    Falls back to the normal parse — and then ensure_nonempty_message — when
    there is genuinely nothing to recover. Never reached for providers using
    agno's native structured output (OpenAI/Anthropic)."""
    salvaged = _salvage_groq_answer_text(response)
    if salvaged:
        logger.warning("Groq ended without the `json` tool; recovered the assistant's answer text")
        return parse_response_content(salvaged)
    logger.warning("Groq ended without the `json` tool and produced no answer text")
    return parse_response_content(response)


# Add a function to remove URLs from message content
def remove_urls_from_message(message: str) -> str:
    """Remove URLs from message text, but preserve markdown image URLs"""
    if not message:
        return message
    
    # Don't remove URLs from markdown images: ![alt](url)
    # We'll replace other URLs but skip those in image markdown
    
    # Pattern to match markdown images: ![...](url)
    image_pattern = r'!\[[^\]]*\]\(([^)]+)\)'
    
    # Find all markdown images and temporarily replace them with placeholders
    images = []
    def save_image(match):
        images.append(match.group(0))
        return f'__IMAGE_PLACEHOLDER_{len(images)-1}__'
    
    message = re.sub(image_pattern, save_image, message)
    
    # Now remove other URLs
    url_pattern = r'https?://[^\s\)\]"]+'
    message = re.sub(url_pattern, '[link removed]', message)
    
    # Restore markdown images
    for i, image in enumerate(images):
        message = message.replace(f'__IMAGE_PLACEHOLDER_{i}__', image)
    
    return message

def enrich_shopify_response(response_content: ChatResponse, session_id: str, fallback_cache_key: str = None) -> ChatResponse:
    """
    Enrich ChatResponse by converting ShopifyOutputDataLLM to ShopifyOutputData with full product data from Redis.

    This function:
    1. Takes ChatResponse with shopify_output (ShopifyOutputDataLLM - no products)
    2. Retrieves full products from Redis using product_cache_key
    3. Converts to ShopifyOutputData (with products) for socket/frontend

    The cache key is preferred from the LLM's structured shopify_output, but falls back
    to `fallback_cache_key` (the deterministic key the toolkit recorded when a product
    tool actually ran this turn). This makes product attachment reliable even when the
    LLM forgets to emit shopify_output / echo the key — the common failure mode.

    Args:
        response_content: The ChatResponse object from the LLM
        session_id: The current session ID (for logging)
        fallback_cache_key: Redis key recorded by the toolkit this turn, used when the
            LLM did not provide a usable product_cache_key.

    Returns:
        Enriched ChatResponse with shopify_output converted to ShopifyOutputData
    """
    from app.core.redis import get_redis
    from app.models.schemas.chat import ShopifyOutputData

    # Prefer the toolkit's recorded key — it is authoritative for the products actually
    # fetched this turn. The LLM's echoed key is only used when no tool ran this turn
    # (e.g. paginating a prior result). This makes attachment reliable regardless of
    # whether the LLM remembered to emit shopify_output.
    llm_product_ids = getattr(response_content.shopify_output, 'product_ids', None) if response_content.shopify_output else None
    if fallback_cache_key:
        cache_key = fallback_cache_key
        if not (response_content.shopify_output and getattr(response_content.shopify_output, 'product_cache_key', None)):
            logger.info(f"LLM omitted shopify_output; attaching products via toolkit fallback key for session {session_id}")
    elif response_content.shopify_output and getattr(response_content.shopify_output, 'product_cache_key', None):
        cache_key = response_content.shopify_output.product_cache_key
    else:
        cache_key = None

    if not cache_key:
        return response_content

    logger.debug(f"Enriching response with cache key: {cache_key}")

    try:
        redis_client = get_redis()
        if not redis_client:
            logger.warning("Redis client not available for enrichment, using cached data as-is")
            return response_content

        # Retrieve full product data from Redis
        cached_data = redis_client.get(cache_key)
        if not cached_data:
            logger.warning(f"Cache key {cache_key} not found or expired")
            return response_content

        # Parse cached data
        product_data = json.loads(cached_data)

        # Convert ShopifyOutputDataLLM to ShopifyOutputData with ALL fields from Redis
        # LLM only provides cache_key + product_ids, everything else comes from Redis.
        # When the LLM omitted shopify_output, derive product_ids from the cached products.
        pageInfo = product_data.get("pageInfo", {})
        products = product_data.get("products", [])
        product_ids = llm_product_ids or [p.get("id") for p in products if p.get("id")]

        enriched_output = ShopifyOutputData(
            products=products,
            search_query=product_data.get("search_query"),
            search_type=product_data.get("search_type"),
            total_count=product_data.get("total_count"),
            has_more=pageInfo.get("hasNextPage", False),
            shop_domain=product_data.get("shop_domain"),
            product_cache_key=cache_key,
            product_ids=product_ids
        )

        # Replace LLM output with enriched output (type change: ShopifyOutputDataLLM → ShopifyOutputData)
        response_content.shopify_output = enriched_output
        logger.debug(f"Enriched response with {len(enriched_output.products)} products")

    except Exception as e:
        logger.error(f"Failed to enrich response from Redis: {str(e)}")
        traceback.print_exc()
        # Continue with non-enriched response

    return response_content

class ChatAgent(ChatAgentMCPMixin):
    def __init__(self, api_key: str, model_name: str = "gpt-4o-mini", model_type: str = "OPENAI", org_id: str = None, agent_id: str = None, customer_id: str = None, session_id: str = None, custom_system_prompt: str = None, transfer_to_human: bool | None = None, mcp_tools: list = None, source: str = None, channel: str = None, extra_context: str = None, base_url: str = None, choice_limits: Optional[InteractiveLimits] = None):
        # NOTE: `source` is a knowledge-base document-name filter (see
        # KnowledgeSearchByAgent), NOT the messaging channel. `channel` is the
        # messaging channel tag ('web', 'telegram', ...) and must never be
        # passed as `source`.
        self.channel = channel or 'web'
        # Row id of the reply the last get_response stored (see get_response)
        self.last_bot_message_id = None
        # Channels that can show tappable choices get a reply schema that
        # offers them (ChannelChatResponse); everyone else's is unchanged.
        self._response_model = ChannelChatResponse if choice_limits is not None else ChatResponse
        # Initialize knowledge search tool if org_id and agent_id provided
        logger.debug(f"Initializing chat agent for agent_id: {agent_id} and org_id: {org_id} and source: {source}")
        tools = []
        knowledge_tool_prompt = ""  # Initialize to empty string
        self.knowledge_tool = None  # Holds the KnowledgeSearchByAgent for citation collection

        if org_id and agent_id:
            logger.debug(f"Initializing knowledge search tool for agent_id: {agent_id} and org_id: {org_id} and source: {source}")
            knowledge_tool = KnowledgeSearchByAgent(
                agent_id=agent_id, org_id=org_id, source=source)
            self.knowledge_tool = knowledge_tool
            tools.append(knowledge_tool)
            
            # Base knowledge tool prompt
            knowledge_tool_prompt = """
            You have access to the knowledge search tool. You can use this tool to search for information about the customer's query on product, services, policies, etc. Only use the tool if required, dont use it for general greeting. Dont hallucinate information. For all other queries other than general always search tools before answering.

            **IMPORTANT - Include URLs in Responses:**
            When tools return information that includes URLs, product links, documentation links, or reference URLs, you MUST include these URLs in your response to the customer. URLs provide valuable references and allow customers to access more detailed information. Always preserve and share URLs that are relevant to the customer's query."""

            # For non-Groq models, add the search limit instruction
            # For Groq, skip this to avoid discouraging tool usage
            if model_type.upper() != 'GROQ':
                knowledge_tool_prompt += """

            IMPORTANT: If you attempt to search for information but cannot find relevant results after a few tries, or if you've already searched multiple times without success, respond with a helpful message like "I apologize, but I don't have specific information about that in our knowledge base at the moment. Is there anything else I can help you with?" Do not keep searching indefinitely."""
            

        # Get template instructions and Jira config in a single optimized query
        # Use context manager for database operations
        with SessionLocal() as db:
            jira_repo = JiraRepository(db)
            if agent_id:
                self.agent_data = jira_repo.get_agent_with_jira_config(agent_id)
            else:
                self.agent_data = None
            
            # Check if Shopify is enabled for this agent while we have the db session
            shopify_config = None
            if agent_id and org_id and session_id:
                try:
                    shopify_config_repo = AgentShopifyConfigRepository(db)
                    shopify_config = shopify_config_repo.get_agent_shopify_config(agent_id)
                except Exception as e:
                    logger.error(f"Failed to get Shopify config: {e}")
                    shopify_config = None

            # Native AI ticketing takes precedence over the Jira toolkit when
            # the org's plan allows it AND this agent has ticketing switched on
            # (per-agent toggle, default on). Jira stays available for manual
            # escalation from the dashboard.
            self.ticketing_enabled = False
            if agent_id and org_id and session_id:
                try:
                    from app.services.ticket_access import ticketing_allowed
                    agent_toggle_on = getattr(self.agent_data, "ticketing_enabled", True) if self.agent_data else True
                    self.ticketing_enabled = agent_toggle_on and ticketing_allowed(db, org_id)
                except Exception as e:
                    logger.error(f"Failed to check ticketing access: {e}")

            # Load lead-capture config (prompt-driven, like transfer_to_human: a toggle;
            # the agent collects details conversationally and reports structured output).
            # Extract plain values while the session is open so they survive the block.
            # Known identity of an already-identified customer (authenticated
            # generate-token visitor, or one whose email/name we captured
            # earlier). Injected into the ticket instructions so the agent
            # doesn't ask an identified customer to repeat their email/name.
            self.known_customer_email = None
            self.known_customer_name = None
            if customer_id:
                try:
                    from app.repositories.customer import CustomerRepository
                    cust = CustomerRepository(db).get_by_id(customer_id)
                    if cust and trusts_identity_in_session(db, cust, session_id):
                        if not CustomerRepository.is_placeholder_email(cust.email):
                            self.known_customer_email = cust.email
                        if (cust.full_name or "").strip():
                            self.known_customer_name = cust.full_name.strip()
                except Exception as e:
                    logger.error(f"Failed to load customer identity: {e}")

            self.lead_capture_enabled = False
            self.lead_capture_fields = []
            self.lead_capture_require_consent = True
            self.lead_capture_guidance = None
            # Whether this visitor's lead was already captured on this agent. When true we
            # stop prompting for details (no nagging) and let the chat close normally.
            self.lead_already_captured = False
            if agent_id:
                try:
                    from app.repositories.lead_capture import LeadCaptureConfigRepository
                    lcc = LeadCaptureConfigRepository(db).get_by_agent(agent_id)
                    if lcc:
                        self.lead_capture_enabled = bool(lcc.enabled)
                        self.lead_capture_fields = lcc.fields or []
                        self.lead_capture_require_consent = bool(lcc.require_consent)
                        self.lead_capture_guidance = lcc.guidance
                        if self.lead_capture_enabled and customer_id:
                            from app.services.lead_capture import has_captured_lead
                            self.lead_already_captured = has_captured_lead(db, customer_id, agent_id)
                except Exception as e:
                    logger.error(f"Failed to load lead capture config: {e}")

            # Guardrail context: plain scalars captured while the session is
            # still open, so prompt assembly never touches the detached
            # `organization` relationship after this block closes.
            organization = getattr(self.agent_data, "organization", None) if self.agent_data else None
            agent_type = getattr(self.agent_data, "agent_type", None) if self.agent_data else None
            # Which surface guardrail events are attributed to. Derived once:
            # `channel` is 'web' for the widget and the channel name otherwise.
            self._guardrail_surface = Surface.WIDGET if self.channel == "web" else Surface.CHANNEL
            self._guardrail_ctx = GuardrailContext(
                org_name=getattr(organization, "name", None),
                domain=getattr(organization, "domain", None),
                agent_type=agent_type.value if hasattr(agent_type, "value") else agent_type,
                description=getattr(self.agent_data, "description", None) if self.agent_data else None,
                topic_scope=getattr(self.agent_data, "topic_scope", None) if self.agent_data else None,
                guardrail_prompt=getattr(self.agent_data, "guardrail_prompt", None) if self.agent_data else None,
                guardrail_enabled=getattr(self.agent_data, "guardrail_enabled", True) if self.agent_data else True,
                org_id=str(org_id) if org_id else None,
                agent_id=str(agent_id) if agent_id else None,
            )

        self.api_key = api_key
        self.model_name = model_name
        self.model_type = model_type
        self.base_url = base_url
        self.jira_instructions_added = False
        self.shopify_instructions_added = False
        self.mcp_instructions_added = False
        self.org_id = org_id
        self.agent_id = agent_id
        self.customer_id = customer_id
        self.session_id = session_id
        self.mcp_tools = mcp_tools or []
        
        # Determine transfer_to_human setting - use parameter if provided, otherwise use agent data
        if transfer_to_human is not None:
            self.transfer_to_human = transfer_to_human
        else:
            self.transfer_to_human = self.agent_data.transfer_to_human if self.agent_data else False

        # Initialize tools
        self.tools = []
        
        # Ticket tools: an agent's explicit Jira config wins (existing orgs
        # keep their Jira flow unchanged); native AI ticketing is the default
        # for everyone else. Never both — two create-ticket functions would
        # confuse the model.
        if self.agent_id and self.org_id and self.session_id and not self.transfer_to_human and self.agent_data and self.agent_data.jira_enabled:
            try:
                self.jira_tools = JiraTools(
                    agent_id=self.agent_id,
                    org_id=self.org_id,
                    session_id=self.session_id
                )
                self.tools.append(self.jira_tools)
            except Exception as e:
                logger.error(f"Failed to initialize Jira tools: {e}")
        elif self.agent_id and self.org_id and self.session_id and not self.transfer_to_human and self.ticketing_enabled:
            try:
                from app.tools.ticket_toolkit import TicketTools
                self.ticket_tools = TicketTools(
                    agent_id=self.agent_id,
                    org_id=self.org_id,
                    session_id=self.session_id
                )
                self.tools.append(self.ticket_tools)
            except Exception as e:
                logger.error(f"Failed to initialize ticket tools: {e}")
        
        # Add Shopify tools if agent has Shopify enabled
        if self.agent_id and self.org_id and self.session_id and not self.transfer_to_human and shopify_config and shopify_config.enabled:
            try:
                self.shopify_tools = ShopifyTools(
                    agent_id=self.agent_id,
                    org_id=self.org_id,
                    session_id=self.session_id
                )
                self.tools.append(self.shopify_tools)
            except Exception as e:
                logger.error(f"Failed to initialize Shopify tools: {e}")

        # Add MCP tools if provided
        if self.mcp_tools:
            self.tools.extend(self.mcp_tools)
            logger.debug(f"Added {len(self.mcp_tools)} MCP tools to agent")

        if self.agent_data:
            # Define end chat instructions to avoid long lines
            end_chat_with_rating = (
                "You should end the chat and request a rating ONLY when you are confident that: "
                "1) The customer's issue has been fully resolved and they have confirmed this, "
                "2) The customer explicitly requests to end the chat, "
                "3) There's a clear confirmation or acknowledgment from the customer that their needs have been met, "
                "4) The conversation has reached a natural conclusion after resolving the customer's query, or "
                "5) The requested task has been completed and confirmed by the customer. "
                "DO NOT end the chat just because the customer says \"thank you\" or \"thanks\" - "
                "this is often just politeness and not an indication that they want to end the conversation. "
                "Always check the conversation history to confirm the issue has been properly addressed before ending the chat."
            )
            
            end_chat_without_rating = (
                "You should end the chat ONLY when: "
                "1) The customer's issue has been fully resolved and they have confirmed this, "
                "2) The customer explicitly requests to end the chat, "
                "3) There's a clear confirmation or acknowledgment from the customer that their needs have been met, "
                "4) The conversation has reached a natural conclusion after resolving the customer's query, or "
                "5) The requested task has been completed and confirmed by the customer. "
                "DO NOT end the chat just because the customer says \"thank you\" or \"thanks\" - "
                "this is often just politeness and not an indication that they want to end the conversation. "
                "Always check the conversation history to confirm the issue has been properly addressed before ending the chat. Also generate a response in message field for end chat. e.g: Thank you for your time. Have a great day!"
            )
            
            # Build system message. Operator-authored text (workflow prompt or
            # dashboard instructions) is fenced so the platform policy block —
            # prepended at the end of assembly — can demote it below the
            # code-owned rules. Everything appended after this stays unfenced:
            # it is our own text and keeps full priority.
            system_message = ""
            if custom_system_prompt:
                # Use custom system prompt from workflow
                system_message = wrap_operator_block(custom_system_prompt)
            elif self.agent_data.instructions:
                system_message = wrap_operator_block("\n".join(self.agent_data.instructions)) + knowledge_tool_prompt
                # Grounding rule: scopes the agent by what its knowledge base
                # can support rather than by a topic list. Sits beside
                # knowledge_tool_prompt, outside the operator fence, so a
                # tenant can neither edit nor delete it.
                # NOT applied to the custom_system_prompt (workflow) branch:
                # a workflow is explicitly composed, so its scope belongs in an
                # explicit guardrails node the builder adds, not injected here.
                system_message += guardrail_scope_prompt(self._guardrail_ctx)
            
            # Add concise response instruction for better performance
            system_message += """
            
Keep your responses concise and focused. Provide clear, actionable information in 2-4 sentences unless a detailed explanation is specifically requested. Avoid unnecessary elaboration.

**CRITICAL: Tool Usage Guidelines:**
- If you need information from the user to complete a task, ASK them directly. DO NOT repeatedly call tools hoping to find the information.
- If a tool returns an error or indicates missing information, STOP calling tools and respond to the user.
- DO NOT call the same tool multiple times with the same parameters if it failed the first time.
- DO NOT call tools in a loop. If you've tried a few tools and haven't found what you need, ask the user for help.

**CRITICAL: Accuracy & Grounding (never invent facts):**
- NEVER make up or guess URLs, domain names, email addresses, phone numbers, prices, plan names, dates, or any other specific detail. Do not "complete" or "correct" a domain or link from memory.
- Only state a URL, contact detail, price, or fact if it appears in the knowledge base, tool results, or your configuration. Reproduce it exactly as written — do not alter the spelling, domain (e.g. .com vs .club), or path.
- If you don't have a specific detail from those sources, say you don't have it and offer to connect the visitor with the team, rather than providing a plausible-looking guess."""


            
            # Add transfer instructions if enabled
            if self.transfer_to_human:
                system_message += """
                You have the ability to transfer this conversation to a human agent if needed. You should transfer the conversation if:
                1. You are unable to answer the customer's question or solve their problem
                2. The customer explicitly asks to speak to a human
                3. The customer is expressing frustration with your responses
                4. The customer's request requires human judgment or decision-making
                5. The customer's issue is complex and would benefit from human expertise
                6. The customer needs to perform an action that you cannot assist with
                
                To transfer to a human, set transfer_to_human to true in your response and provide a transfer_reason and transfer_description.
                """
            else:
                # Scoped to THIS conversation on purpose. Stated as a flat
                # "transfer is disabled", the model reads it as a fact about the
                # business and answers "we cannot hand chats to a human" — which
                # for an agent whose subject is its own product turns a setting
                # on one agent into a false claim about what the product does.
                system_message += """
                You cannot hand this conversation to a human agent — that is switched off for this chat.
                Do not offer it, and do not set transfer_to_human. This is a setting on you, not a fact
                about the business or its product: never tell the customer that the business has no human
                support or that the product cannot transfer chats. If they ask to speak to a person, say
                you cannot put them through from here, then help them yourself or open a ticket. If they
                ask what the product can do, answer from the knowledge base, not from this setting.
                """

            # Add lead-capture instructions (enabled = a toggle, like transfer_to_human).
            # The agent collects details conversationally and reports them as structured
            # output (lead_data / lead_summary / request_lead_capture). No form, no triggers.
            # Only prompt while the lead is still pending — once captured, don't nag.
            lead_pending = self.lead_capture_enabled and not self.lead_already_captured
            if lead_pending:
                # Standard contact fields map to explicit scalar response fields
                # (lead_email/lead_name/lead_company/lead_phone) — these populate reliably
                # under strict structured outputs, unlike a free-form dict. Split the
                # configured fields into required vs optional and describe how to report each.
                STANDARD_FIELD_SLOT = {
                    'email': 'lead_email', 'name': 'lead_name',
                    'company': 'lead_company', 'phone': 'lead_phone',
                }
                required_labels, optional_labels = [], []
                for f in self.lead_capture_fields:
                    if not f.get('enabled', True):
                        continue
                    label = f.get('label') or f.get('key')
                    if not label:
                        continue
                    options = [str(o) for o in (f.get('options') or []) if str(o).strip()]
                    if options:
                        label = f"{label} (one of: {', '.join(options)})"
                    (required_labels if f.get('required') else optional_labels).append(label)
                # Email is always the minimum needed to record a lead.
                if not required_labels:
                    required_labels = ['email']
                lc_prompt = (
                    "\n\nLEAD CAPTURE (IMPORTANT): This is a lead-generation agent — collecting the visitor's "
                    "contact details is a primary goal. Help and deliver value first, then PROACTIVELY ask for "
                    "their details at a natural moment (a great time is right after you have answered their "
                    "question or given them something useful). Ask conversationally, one detail at a time — never "
                    "on the very first message, never mid-answer."
                )
                lc_prompt += " Collect these details, asking naturally one at a time. REQUIRED: " + \
                    ", ".join(required_labels) + "."
                if optional_labels:
                    lc_prompt += (
                        " ALSO TRY TO COLLECT (optional — they enrich the lead; genuinely ask for each one, but "
                        "do not insist if the visitor skips or declines): " + ", ".join(optional_labels) + "."
                        " Ask for these optional details EARLY, while collecting — do not skip straight to "
                        "recording after only getting the email. Ask for each optional field AT MOST ONCE."
                    )
                lc_prompt += (
                    " As you learn each standard detail, set the matching response field: email in lead_email, "
                    "name in lead_name, company in lead_company, phone in lead_phone. Only fill a field with what "
                    "the visitor ACTUALLY told you — never guess or infer a value (e.g. do not use the company as "
                    "the name); leave a field empty if they did not give it."
                    " EMAIL VALIDATION: a valid email must contain an '@' and a domain with a dot (e.g. "
                    "jane@acme.com). If what the visitor gives is NOT a valid email (e.g. 'arun.com', a bare "
                    "domain, or just a name), do NOT accept it, do NOT set lead_email, and do NOT claim you "
                    "recorded it — politely point out it looks incomplete and ask again for a full email address."
                )
                # The visitor may already be known - they passed the email gate, or the
                # embedding app identified them with a token. Without this the agent asks
                # them to type an address we are already storing and then tells them
                # "Email: Not provided", which is the same "we have this already" wart
                # the ticket instructions below avoid.
                known_lead_bits = []
                if self.known_customer_email:
                    known_lead_bits.append(
                        f'their email is "{self.known_customer_email}" - put exactly that in lead_email'
                    )
                if self.known_customer_name:
                    known_lead_bits.append(
                        f'their name is "{self.known_customer_name}" - put exactly that in lead_name'
                    )
                if known_lead_bits:
                    lc_prompt += (
                        " ALREADY KNOWN: " + "; ".join(known_lead_bits) + "."
                        " Do NOT ask for these and do NOT ask the visitor to confirm or repeat them - you"
                        " already have them, so count them as collected when deciding whether you can record"
                        " the lead. If the visitor volunteers a different value, use theirs instead."
                    )

                # The record trigger is: a valid email + every REQUIRED field + consent.
                # Required fields (beyond email) genuinely gate recording; optional ones
                # never do. Build the trigger text from the required fields so marking a
                # field "required" in the UI actually makes the agent insist on it.
                extra_required = [l for l in required_labels if l.strip().lower() != 'email']
                trigger = "a valid email (in lead_email)"
                if extra_required:
                    trigger += " AND the required details (" + ", ".join(extra_required) + ")"
                if self.lead_capture_require_consent:
                    lc_prompt += (
                        " ORDER: ask for the email and the other details first, and ask for consent LAST. "
                        " CONSENT REQUIRED: before recording you MUST get the visitor's explicit agreement to be "
                        "contacted (a clear yes). If the visitor clearly agrees to be contacted (e.g. 'yes', 'yes "
                        "you can contact me', 'sure, go ahead'), treat that as consent immediately — set "
                        "lead_consent=true and do NOT ask for consent again."
                    )
                    trigger += " AND consent"
                lc_prompt += (
                    " RECORD NOW — the MOMENT you have " + trigger + ", you MUST in that SAME response "
                    "(1) make sure lead_email holds the exact email the visitor gave, (2) set request_lead_capture "
                    "to true, (3) set lead_consent to true, and (4) write a short lead_summary qualifying this "
                    "lead. Do this even if OPTIONAL fields are still missing — NEVER keep asking for an optional "
                    "field once you have " + trigger + ", and never ask for the same field twice. Required fields "
                    "above DO gate recording: keep asking for a required field until you have it (or the visitor "
                    "clearly refuses). CRITICAL: request_lead_capture=true is INVALID unless lead_email is set. "
                    "Confirm back what you captured (e.g. 'Great — I'll have someone reach out at jane@acme.com'). "
                    "Make ONE genuine attempt overall; if they decline to share details, respect it and keep helping."
                )
                if self.lead_capture_guidance:
                    lc_prompt += " Additional guidance from the business: " + self.lead_capture_guidance
                system_message += lc_prompt

            # Add end chat instructions. While a lead-capture attempt is still pending, the
            # normal "end on natural conclusion / thank you" rules would let the model bail
            # before ever asking, so replace them with a gated version that forces the ask
            # first. Once the lead is captured (or capture is off), use the normal rules.
            if lead_pending:
                system_message += (
                    "\nEND CHAT (lead-capture pending): You have NOT yet collected this visitor's contact "
                    "details, so do NOT set end_chat=true yet — not even if the conversation seems to be "
                    "wrapping up or the visitor says \"thank you\", \"thanks\", \"bye\", or \"that's all\". "
                    "If the visitor is wrapping up and you have not asked yet, reply by asking for their "
                    "contact details now (see LEAD CAPTURE above) instead of ending. You may set end_chat=true "
                    "ONLY after you have recorded their details (request_lead_capture=true) OR they have clearly "
                    "declined to share them. When you do end, also generate a closing message in the message "
                    "field, e.g: Thank you for your time. Have a great day!"
                )
            elif self.agent_data.ask_for_rating:
                system_message += f"\n{end_chat_with_rating}"
            else:
                system_message += f"\n{end_chat_without_rating}"
            
            # Add native ticketing instructions when AI ticketing drives the
            # ticket tools (explicit agent Jira config takes precedence below)
            if self.ticketing_enabled and not self.transfer_to_human and not (self.agent_data and self.agent_data.jira_enabled):
                # For an already-identified customer, hand the agent their known
                # email/name so it passes them straight to create_ticket instead
                # of asking the customer to repeat details we already have.
                if self.known_customer_email or self.known_customer_name:
                    known_bits = []
                    if self.known_customer_name:
                        known_bits.append(f'name "{self.known_customer_name}"')
                    if self.known_customer_email:
                        known_bits.append(f'account email "{self.known_customer_email}"')
                    identity_instruction = (
                        "This customer is already identified — " + ", ".join(known_bits) + ". "
                        "Pass these directly as customer_name and customer_email when calling "
                        "create_ticket. Do NOT ask the customer to provide or confirm their "
                        "email or name again."
                    )
                else:
                    identity_instruction = (
                        "BEFORE creating a ticket, collect the customer's identity so the support "
                        "AI can look them up in the connected systems: their account email, and "
                        "their registered name. If they've already given the email or name in this "
                        "conversation, use those values. If not, ask for them in one short message "
                        "and wait for the reply before calling create_ticket. Then pass them as "
                        "customer_email and customer_name. Never invent an email or name; if the "
                        "customer declines, create the ticket without them."
                    )
                ticket_instructions = f"""
                You have access to native support-ticket tools:
                1. check_existing_ticket — check if this conversation already has a ticket (always call this first)
                2. create_ticket — open a support ticket that the team's AI investigator and humans will work on
                3. get_ticket_status — look up a ticket's current status

                Only create a ticket if:
                - The issue is a technical problem you cannot resolve from the knowledge base
                - The user explicitly asks for a ticket or escalation
                - You've tried to resolve the issue but were unable to do so
                - No ticket already exists for this conversation

                {identity_instruction}

                After creating a ticket, tell the customer their ticket number and that the team is investigating.
                Priorities are: urgent, high, medium, low.
                """
                system_message += "\n\n" + ticket_instructions
                self.jira_instructions_added = True
            # Add Jira instructions if Jira is enabled (and native ticketing is not)
            elif self.agent_data and self.agent_data.jira_enabled and not self.transfer_to_human:
                jira_instructions = """
                You have access to Jira integration tools. You can use these tools to:
                1. Create a Jira ticket for issues that need further attention
                2. Check if a ticket already exists for the current conversation
                3. Get the status of an existing ticket

                To create a ticket, you can either:
                - Use the create_jira_ticket function directly
                - Include the following fields in your response:
                - create_ticket: Set to true to create a ticket
                - ticket_summary: A brief summary of the issue (required if create_ticket is true)
                - ticket_description: A detailed description of the issue (required if create_ticket is true)
                - ticket_priority: The priority level of the ticket (optional, defaults to "Medium")

                Only create a ticket if:
                - The issue is complex and requires human intervention
                - The user explicitly requests to create a ticket
                - You've tried to resolve the issue but were unable to do so
                - No ticket already exists for this conversation
                """
                system_message += "\n\n" + jira_instructions
                self.jira_instructions_added = True
            
            # Add Shopify instructions if Shopify is enabled
            if shopify_config and shopify_config.enabled:
                # Simplified Shopify Instructions for faster LLM processing
                shopify_instructions = """
                You have access to Shopify tools for products and orders. Use `limit` of 8 for product searches.

                **CRITICAL RESPONSE RULES:**
                - Your message must ONLY contain simple conversational text (under 50 words)
                - NEVER include product details, prices, images, or URLs in your message
                - Examples: "Here are some options for you.", "I found several products.", "Your order has been shipped."

                **Product Search Filters:**
                Both `search_products` and `recommend_products` support these optional filters:
                - `min_price`: Minimum price (e.g., 100.00 for products >= $100)
                - `max_price`: Maximum price (e.g., 500.00 for products <= $500)
                - `vendor`: Brand/vendor name (e.g., 'Nike', 'Burton')

                Examples:
                - "snowboard under $600" → use `max_price=600`
                - "shoes between $50 and $200" → use `min_price=50, max_price=200`
                - "Nike products" → use `vendor='Nike'`

                **Product Tools Response:**
                When Shopify tools return products, you MUST include ONLY these 2 fields in your `shopify_output`:
                - `product_cache_key`: Copy this from the tool response (REQUIRED)
                - `product_ids`: Copy this array from the tool response (REQUIRED)

                DO NOT include any other fields (shop_domain, total_count, products, etc.) - backend will populate everything from cache

                **Order Tools:**
                - Before looking up an order, ask for BOTH the order number and the email address used at checkout. Order details are only shared when both match
                - Only use an email the customer typed in this chat; never guess one
                - Use customer-friendly language: "Your order has been shipped" not "FULFILLED"
                - Make tracking numbers clickable links
                """
                system_message += "\n\n" + shopify_instructions
                self.shopify_instructions_added = True

            # Add MCP tools instructions if MCP tools are available
            if self.mcp_tools:
                mcp_instructions = """
                You have access to MCP (Model Context Protocol) tools that provide additional capabilities.
                These tools allow you to interact with external systems and perform various operations.
                Use these tools when they can help answer the customer's questions or solve their problems.
                Always use the appropriate tool for the specific task at hand.
                """
                system_message += "\n\n" + mcp_instructions
                self.mcp_instructions_added = True
        else:
            system_message = [
                "You are a helpful customer service agent.",
            ]

        # Per-conversation context (e.g. the outbound template that opened this
        # thread). Appended last so it composes with the configured behaviour —
        # custom_system_prompt REPLACES the prompt, which is exactly what this
        # must never do. Same additive pattern as the Jira/Shopify sections.
        if extra_context:
            if isinstance(system_message, list):
                system_message = [*system_message, extra_context]
            else:
                system_message += "\n\n" + extra_context

        # Tappable choices, only where the channel can show them — every other
        # surface keeps an unchanged prompt.
        if choice_limits is not None:
            instructions = choice_instructions(choice_limits)
            if isinstance(system_message, list):
                system_message = [*system_message, instructions]
            else:
                system_message += "\n\n" + instructions

        # Platform guardrail policy — single choke point covering all three
        # branches above (workflow custom prompt, agent instructions, and the
        # no-agent default). Prepends the code-owned policy, appends the
        # precedence anchor, and always returns a str (which also fixes the
        # legacy list form crashing the Groq append below).
        system_message = apply_guardrail_policy(system_message, self._guardrail_ctx)

        # Initialize model with utility function
        base_max_tokens = 2000 if (self.shopify_instructions_added or self.mcp_instructions_added) else 1000
        # Groq's GPT-OSS/reasoning models spend output tokens on internal reasoning
        # BEFORE emitting the `json` tool call; too small a budget truncates the tool
        # arguments into invalid JSON (Groq 400). Give the Groq path extra headroom.
        if model_type.upper() == 'GROQ':
            base_max_tokens = max(base_max_tokens, 4000)
        model = create_model(
            model_type=model_type,
            api_key=api_key,
            model_name=model_name,
            max_tokens=base_max_tokens,
            base_url=base_url,
            # response_format={"type": "json_object"} if model_type.upper() != 'GROQ' else {"type": "text"}
        )

        # Use shared database engine to avoid connection leaks
        # Previously this created a new engine per ChatAgent instance, exhausting connections
        storage = EncryptedPostgresAgentStorage(table_name="agent_sessions", db_engine=engine)
        
       
        # Combine all tools
        all_tools = tools.copy()
        if hasattr(self, 'tools') and self.tools:
           all_tools.extend(self.tools)

        # Groq can't combine response_format (JSON mode) with tools, so its native
        # structured-output path degrades to unreliable prompt-only JSON. For Groq we
        # instead register a `json` tool (GPT-OSS's own structured-output convention)
        # and read the ChatResponse from its arguments; other providers keep agno's
        # native response_model path. See build_groq_response_tool above.
        self._groq_json_capture = {}
        self._use_groq_json_tool = model_type.upper() == 'GROQ'
        if self._use_groq_json_tool:
            all_tools.append(build_groq_response_tool(self._groq_json_capture, choice_limits is not None))
            system_message = (system_message or "") + _GROQ_JSON_INSTRUCTION
            response_model = None
            structured_outputs = False
        else:
            response_model = self._response_model
            structured_outputs = True

        self.agent = Agent(
           name=self.agent_data.name if self.agent_data else "Default Agent",
           session_id=session_id,
           model=model,
           tools=all_tools,
           instructions=system_message,
           agent_id=str(agent_id),
           storage=storage,
           memory=BoundedHistoryMemory(),
           add_history_to_messages=True,
           tool_call_limit=settings.AGENT_TOOL_CALL_LIMIT,
           num_history_responses=5,  # Reduced from 10 to 5 to minimize context size and improve speed
           read_chat_history=True,
           markdown=False,
           debug_mode=settings.ENVIRONMENT == "development",
           user_id=str(customer_id),
           session_state={"status": "active"},
           response_model=response_model,
           structured_outputs=structured_outputs,
           system_message_role="system",
           user_message_role="user",
           show_tool_calls=settings.ENVIRONMENT == "development"
          )

    async def _arun(self, message: str, session_id: str):
        """One agent run under AGENT_RUN_TIMEOUT, with the context-overflow
        retry. The retry is a fresh turn, so what the failed attempt searched
        must be searchable again."""
        knowledge_tool = getattr(self, "knowledge_tool", None)
        return await run_with_overflow_recovery(
            self.agent,
            message,
            session_id,
            timeout=settings.AGENT_RUN_TIMEOUT,
            before_retry=knowledge_tool.reset_turn if knowledge_tool is not None else None,
        )

    async def _get_llm_response_only(self, message: str, session_id: str = None, org_id: str = None, agent_id: str = None, customer_id: str = None) -> ChatResponse:
        """
        Get LLM response without storing messages in chat history.
        Used by workflow execution to avoid duplicate message storage.
        """
        try:
            # Update session and IDs if provided
            if session_id:
                self.session_id = session_id
            if org_id:
                self.org_id = org_id
            if agent_id:
                self.agent_id = agent_id
            if customer_id:
                self.customer_id = customer_id
                
            self.agent.session_id = session_id

            # Guardrail: COUNT-ONLY on the workflow path. This method persists
            # nothing and the workflow layer owns control flow, so a block here
            # would silently break node routing — signals are recorded and the
            # policy block in the prompt does the refusing.
            check_inbound(
                message,
                ctx=self._guardrail_ctx,
                surface=Surface.WORKFLOW,
                session_id=session_id,
                allow_block=False,
            )

            # Get AI response WITHOUT storing user message
            self._groq_json_capture.clear()
            if self.knowledge_tool is not None:
                self.knowledge_tool.reset_turn()
            try:
                response = await self._arun(message, session_id)
            except asyncio.TimeoutError:
                logger.warning(
                    f"Agent run timed out after {settings.AGENT_RUN_TIMEOUT}s and was cancelled "
                    f"(session_id={session_id}, agent_id={agent_id}, org_id={org_id})"
                )
                raise
            except Exception as arun_exc:
                # Groq only: a truncated `json` tool call is rejected as invalid JSON;
                # salvage the (mostly-complete) fields so the lead/end_chat survive.
                # Non-Groq providers re-raise unchanged — no behavior change for them.
                salvaged = _salvage_groq_json_error(arun_exc) if self._use_groq_json_tool else None
                # Empty salvage means nothing usable survived — re-raise so the normal
                # error reply is shown instead of a bare "No response generated".
                if not salvaged:
                    raise
                logger.warning("Groq json tool call unparseable (likely truncated); salvaged structured fields")
                response_content = _build_chat_response_from_capture(salvaged, self._response_model)
            else:
                # Groq path returns the structured turn via the `json` tool; everything
                # else parses agno's native structured output.
                if self._use_groq_json_tool and self._groq_json_capture:
                    response_content = _build_chat_response_from_capture(self._groq_json_capture, self._response_model)
                elif self._use_groq_json_tool:
                    # Groq finished without the `json` tool — recover its answer text.
                    response_content = recover_groq_no_capture(response)
                else:
                    response_content = parse_response_content(response, self._response_model)

            response_content = ensure_nonempty_message(response_content)
            logger.debug(f"Response content: {response_content}")

            # Output guardrail: replace a reply that leaked policy text,
            # count the model's own scope refusals.
            response_content.message, _ = check_output(
                response_content.message,
                ctx=self._guardrail_ctx,
                surface=Surface.WORKFLOW,
                session_id=session_id,
            )

            # Enrich Shopify response with full product data from Redis
            response_content = enrich_shopify_response(response_content, session_id, fallback_cache_key=self._shopify_fallback_cache_key())

            # If shopify_output has products, remove URLs from message
            # (URLs should only be removed when products are being displayed separately)
            if response_content.shopify_output and hasattr(response_content.shopify_output, 'products') and response_content.shopify_output.products:
                response_content.message = remove_urls_from_message(response_content.message)
                logger.debug(f"Cleaned message for Shopify output: {response_content.message}")

            # Don't handle end chat or transfer here - let workflow handle it
            # Don't store any messages - let workflow handle storage

            return response_content

        except Exception as e:
            traceback.print_exc()
            logger.error(f"Chat agent error: {str(e)}")
            error_message = GENERIC_ERROR_MESSAGE
            
            # Create error response without storing
            error_response = ChatResponse(
                message=error_message,
                transfer_to_human=False,
                transfer_reason=None,
                transfer_description=None,
                end_chat=False,
                end_chat_reason=None,
                end_chat_description=None,
                request_rating=False,
                create_ticket=False,
                shopify_output=None
            )
            
            return error_response

    async def _handle_end_chat(self, response_content: ChatResponse, session_id: str, db, force_rating: bool | None = None) -> ChatResponse:
        """
        Handle end chat logic including session updates and rating requests.
        
        Args:
            response_content: The chat response content
            session_id: The session ID
            db: Database session
            force_rating: Optional parameter to override agent's ask_for_rating setting.
                         If None, uses agent's default setting.
                         If True, forces rating request.
                         If False, disables rating request.
            
        Returns:
            Updated ChatResponse object
        """
        session_repo = SessionToAgentRepository(db)
        
        # Determine if rating should be requested
        if not is_widget_channel(self.channel):
            # Rating is a widget-only feature: external channels have no
            # rating UI, so asking there leaves the customer a dead-end prompt.
            # Shopify counts as the widget — see WIDGET_CHANNELS.
            should_request_rating = False
        elif force_rating is not None:
            # Use the forced setting from workflow configuration
            should_request_rating = force_rating
        else:
            # Use agent's default setting
            should_request_rating = self.agent_data and self.agent_data.ask_for_rating
            
        response_content.request_rating = should_request_rating

        session_repo.update_session(
            session_id,
            {
                "status": SessionStatus.CLOSED,
                "end_chat_reason": response_content.end_chat_reason.value if response_content.end_chat_reason else None,
                "end_chat_description": response_content.end_chat_description,
                "closed_at": datetime.now()
            }
        )

        # Add rating request to the message if enabled
        if should_request_rating:
            rating_message = "\n\nThank you for chatting with us! Would you please take a moment to rate your experience? Your feedback helps us improve our service."
            response_content.message += rating_message
            
        return response_content

    async def _handle_transfer(self, response_content: ChatResponse, session_id: str, org_id: str, agent_id: str, customer_id: str, db, chat_repo: ChatRepository, transfer_group_id: str = None) -> ChatResponse:
        """
        Handle transfer to human logic including session updates, notifications, and availability checks.
        
        Args:
            response_content: The chat response content (can be None for workflow transfers)
            session_id: The session ID
            org_id: Organization ID
            agent_id: Agent ID
            customer_id: Customer ID
            db: Database session
            chat_repo: Chat repository instance
            transfer_group_id: Optional specific group ID to transfer to (for workflow transfers)
            
        Returns:
            Updated ChatResponse object
        """
        from app.models.schemas.chat import TransferReasonType
        
        # Determine transfer source and group
        if transfer_group_id:
            logger.debug(f"Transfer group ID: {transfer_group_id}")
            # Workflow transfer - use provided group ID and transfer details from LLM response
            group_id = transfer_group_id
            # Use transfer reason/description from LLM response if available, otherwise fallback
            if response_content and response_content.transfer_reason:
                transfer_reason = response_content.transfer_reason.value
                transfer_description = response_content.transfer_description or "Transfer requested by workflow"
            else:
                transfer_reason = TransferReasonType.KNOWLEDGE_GAP.value
                transfer_description = "Transfer requested by workflow"
            notification_message = "A chat has been transferred to your group via workflow."
            is_workflow_transfer = True
        else:
            # Agent transfer - use agent's default group
            if not (self.agent_data and hasattr(self.agent_data, 'groups') and self.agent_data.groups):
                raise ValueError("No groups available for transfer")
            group_id = self.agent_data.groups[0].id
            transfer_reason = response_content.transfer_reason.value if response_content.transfer_reason else None
            transfer_description = response_content.transfer_description
            notification_message = f"A chat has been transferred to your group. Reason: {transfer_reason or 'Not specified'}"
            is_workflow_transfer = False
        
        # Get chat history
        chat_history = await chat_repo.get_session_history(session_id)
        
        # Update session with transfer details
        session_repo = SessionToAgentRepository(db)
        session_repo.update_session(
            session_id, 
            {
                "status": "TRANSFERRED",
                "transfer_reason": transfer_reason,
                "transfer_description": transfer_description,
                "group_id": group_id
            }
        )
        
        # Notify the target group's members who haven't muted transfers
        users = db.query(User).join(user_groups).filter(user_groups.c.group_id == group_id).all()

        await notify_chat_event(
            db=db,
            user_ids=[user.id for user in users],
            event=ChatNotificationEvent.CHAT_TRANSFER,
            title="New Chat Transfer",
            message=notification_message,
            metadata={
                "session_id": session_id,
                "transfer_reason": transfer_reason,
                "transfer_description": transfer_description
            }
        )

        # Get availability-based response
        availability_response = await get_agent_availability_response(
            agent=self.agent_data,
            customer_id=customer_id,
            chat_history=chat_history,
            db=db,
            api_key=self.api_key,
            model_name=self.model_name,
            model_type=self.model_type,
            base_url=self.base_url,
            session_id=session_id,
            transfer_group_id=transfer_group_id if is_workflow_transfer else None
        )
        
        # Create ChatResponse object
        updated_response = ChatResponse(
            message=availability_response["message"],
            transfer_to_human=availability_response["transfer_to_human"],
            transfer_reason=availability_response.get("transfer_reason"),
            transfer_description=availability_response.get("transfer_description"),
            end_chat=False,
            end_chat_reason=None,
            end_chat_description=None,
            request_rating=False,
            create_ticket=False,
            shopify_output=None
        )

        # Signal the widget to collect the visitor's contact details (handoff happened,
        # whether or not a live agent was available).
        updated_response.request_contact = True

        # Prepare message attributes
        attributes = {
            "transfer_to_human": updated_response.transfer_to_human,
            "transfer_reason": updated_response.transfer_reason.value if updated_response.transfer_reason else None,
            "transfer_description": updated_response.transfer_description,
            "end_chat": updated_response.end_chat,
            "end_chat_reason": updated_response.end_chat_reason.value if updated_response.end_chat_reason else None,
            "end_chat_description": updated_response.end_chat_description,
            "request_rating": updated_response.request_rating,
            "shopify_output": updated_response.shopify_output
        }
        
        # Add workflow-specific attributes
        if is_workflow_transfer:
            attributes["workflow_transfer"] = True
            attributes["transfer_group_id"] = transfer_group_id
        
        # Store transfer response
        self._store_bot_message(chat_repo, {
            "message": updated_response.message,
            "message_type": "bot",
            "session_id": session_id,
            "organization_id": org_id,
            "agent_id": agent_id,
            "customer_id": customer_id,
            "attributes": attributes
        })

        return updated_response

    def _shopify_fallback_cache_key(self):
        """Redis key for products a Shopify tool fetched during THIS turn, else None.

        Used to attach products when the model omits shopify_output. Gated on the
        toolkit having actually cached this turn so a previous turn's still-live
        cache can never be pinned onto an unrelated reply.
        """
        tools = getattr(self, 'shopify_tools', None)
        if tools and getattr(tools, 'has_cached_products', False):
            return tools.product_cache_key
        return None

    async def handle_workflow_transfer(self, session_id: str, org_id: str, agent_id: str, customer_id: str, transfer_group_id: str, db, chat_repo: ChatRepository, llm_response: ChatResponse = None) -> ChatResponse:
        """
        Handle transfer to human from workflow with specific group ID.
        This is a convenience wrapper around _handle_transfer for workflow transfers.
        
        Args:
            session_id: The session ID
            org_id: Organization ID
            agent_id: Agent ID
            customer_id: Customer ID
            transfer_group_id: The specific group ID to transfer to
            db: Database session
            chat_repo: Chat repository instance
            llm_response: The LLM response containing transfer reason and description
            
        Returns:
            ChatResponse object with transfer response
        """
        return await self._handle_transfer(
            response_content=llm_response,  # Pass the LLM response to get transfer reason/description
            session_id=session_id,
            org_id=org_id,
            agent_id=agent_id,
            customer_id=customer_id,
            db=db,
            chat_repo=chat_repo,
            transfer_group_id=transfer_group_id
        )

    def _store_bot_message(self, chat_repo: ChatRepository, message_data: dict) -> None:
        """Persist a reply and remember its row (see get_response). A ChatAgent
        serves a single turn, so the last one stored is this turn's reply."""
        self.last_bot_message_id = chat_repo.create_message(message_data).id

    async def get_response(self, message: str, session_id: str = None, org_id: str = None, agent_id: str = None, customer_id: str = None, user_attributes: dict = None) -> ChatResponse:
        """
        Get a response from the agent.

        `user_attributes` are stored on the customer's message row (e.g. the
        button a channel customer tapped). The stored reply's row id is left
        on `last_bot_message_id` so a caller can annotate that exact row.
        """
        self.last_bot_message_id = None
        try:
            # Update session and IDs if provided
            if session_id:
                self.session_id = session_id
            if org_id:
                self.org_id = org_id
            if agent_id:
                self.agent_id = agent_id
            if customer_id:
                self.customer_id = customer_id
                
            # Pre-LLM guardrail: strong injection signals are counted, and (in
            # blocking modes) stop the turn before inference. Checked before
            # the insert so the verdict rides on the user row's attributes.
            guardrail_verdict = check_inbound(
                message,
                ctx=self._guardrail_ctx,
                surface=self._guardrail_surface,
                session_id=session_id,
            )

            # Use context manager for database operations
            with SessionLocal() as db:
                chat_repo = ChatRepository(db)

                self.agent.session_id = session_id

                # Create user message
                chat_repo.create_message({
                    "message": message,
                    "message_type": "user",
                    "session_id": session_id,
                    "organization_id": org_id,
                    "agent_id": agent_id,
                    "customer_id": customer_id,
                    "attributes": {**(user_attributes or {}), **guardrail_verdict.as_attributes()}
                })

                if guardrail_verdict.block:
                    # Blocked pre-inference: reply with the canned line, store
                    # it like any bot turn, and never call the model.
                    blocked_response = ChatResponse(
                        message=guardrail_verdict.reply,
                        transfer_to_human=False,
                        transfer_reason=None,
                        transfer_description=None,
                        end_chat=False,
                        end_chat_reason=None,
                        end_chat_description=None,
                        request_rating=False,
                        create_ticket=False,
                        shopify_output=None
                    )
                    self._store_bot_message(chat_repo, {
                        "message": guardrail_verdict.reply,
                        "message_type": "bot",
                        "session_id": session_id,
                        "organization_id": org_id,
                        "agent_id": agent_id,
                        "customer_id": customer_id,
                        "attributes": guardrail_verdict.as_attributes()
                    })
                    return blocked_response

                # Reset citation collection and shown chunks for this turn
                if self.knowledge_tool is not None:
                    self.knowledge_tool.reset_turn()

                # Get AI response
                self._groq_json_capture.clear()
                _salvaged_content = None
                try:
                    response = await self._arun(message, session_id)
                except asyncio.TimeoutError:
                    logger.warning(
                        f"Agent run timed out after {settings.AGENT_RUN_TIMEOUT}s and was cancelled "
                        f"(session_id={session_id}, agent_id={agent_id}, org_id={org_id})"
                    )
                    raise
                except Exception as arun_exc:
                    # Groq only: salvage a truncated `json` tool call (see get_response).
                    salvaged = _salvage_groq_json_error(arun_exc) if self._use_groq_json_tool else None
                    # Empty salvage → nothing usable; re-raise for the normal error reply.
                    if not salvaged:
                        raise
                    logger.warning("Groq json tool call unparseable (likely truncated); salvaged structured fields")
                    _salvaged_content = _build_chat_response_from_capture(salvaged, self._response_model)
                    response = None

                # Groq path returns the structured turn via the `json` tool; everything
                # else parses agno's native structured output.
                if _salvaged_content is not None:
                    response_content = _salvaged_content
                elif self._use_groq_json_tool and self._groq_json_capture:
                    response_content = _build_chat_response_from_capture(self._groq_json_capture, self._response_model)
                elif self._use_groq_json_tool:
                    # Groq finished without the `json` tool — recover its answer text.
                    response_content = recover_groq_no_capture(response)
                else:
                    response_content = parse_response_content(response, self._response_model)

                # Attach knowledge-base citations gathered during this turn (overrides any
                # value the LLM may have produced — this field is system-managed).
                if self.knowledge_tool is not None and self.knowledge_tool.collected_sources:
                    from app.models.schemas.chat import SourceRef
                    response_content.sources = [
                        SourceRef(**s) for s in self.knowledge_tool.collected_sources
                    ]
                else:
                    response_content.sources = None

                # request_contact is system-managed (set only by the transfer handler);
                # clear anything the LLM may have produced.
                response_content.request_contact = False

                # What the model itself wrote — before an empty-turn fallback or
                # the guardrail can replace it (see drop_stale_choices)
                model_message = response_content.message
                response_content = ensure_nonempty_message(response_content)
                logger.debug(f"Response content: {response_content}")

                # Output guardrail: replace a reply that leaked policy text,
                # count the model's own scope refusals.
                response_content.message, _ = check_output(
                    response_content.message,
                    ctx=self._guardrail_ctx,
                    surface=self._guardrail_surface,
                    session_id=session_id,
                )
                drop_stale_choices(response_content, model_message)

                # Enrich Shopify response with full product data from Redis
                response_content = enrich_shopify_response(response_content, session_id, fallback_cache_key=self._shopify_fallback_cache_key())

                # If shopify_output has products, remove URLs from message
                # (URLs should only be removed when products are being displayed separately)
                if response_content.shopify_output and hasattr(response_content.shopify_output, 'products') and response_content.shopify_output.products:
                    response_content.message = remove_urls_from_message(response_content.message)
                    logger.debug(f"Cleaned message for Shopify output: {response_content.message}")
                
                # Handle end chat and rating request
                if response_content.end_chat:
                    response_content = await self._handle_end_chat(response_content, session_id, db)

                # Handle transfer 
                if self.agent_data and self.transfer_to_human and response_content.transfer_to_human and hasattr(self.agent_data, 'groups') and self.agent_data.groups:
                    response_content = await self._handle_transfer(
                        response_content=response_content,
                        session_id=session_id,
                        org_id=org_id,
                        agent_id=agent_id,
                        customer_id=customer_id,
                        db=db,
                        chat_repo=chat_repo,
                        transfer_group_id=None  # Use agent's default group for regular transfers
                    )
                    return response_content

                # Reaching here means no live transfer happened — either transfer is
                # disabled for this agent, the model didn't actually request one, or the
                # agent has no groups to route to. Clear any stray transfer flag the model
                # may have set so the widget does not prompt the visitor for contact
                # details or tell them it will connect them to a teammate.
                if response_content.transfer_to_human:
                    logger.debug(
                        "Clearing model-set transfer_to_human flag (no transfer performed)")
                    response_content.transfer_to_human = False
                    response_content.transfer_reason = None
                    response_content.transfer_description = None

                # Store AI response with all attributes
                attributes = {
                    "transfer_to_human": response_content.transfer_to_human,
                    "transfer_reason": response_content.transfer_reason.value if response_content.transfer_reason else None,
                    "transfer_description": response_content.transfer_description,
                    "end_chat": response_content.end_chat,
                    "end_chat_reason": response_content.end_chat_reason.value if response_content.end_chat_reason else None,
                    "end_chat_description": response_content.end_chat_description,
                    "request_rating": response_content.request_rating,
                    "shopify_output": response_content.shopify_output
                }

                # Persist citations so reloaded history can render them too
                if response_content.sources:
                    attributes["sources"] = [s.model_dump() for s in response_content.sources]

                # Add ticket attributes if present
                if response_content.create_ticket:
                    attributes.update({
                        "create_ticket": response_content.create_ticket,
                        "ticket_summary": response_content.ticket_summary,
                        "ticket_description": response_content.ticket_description,
                        "integration_type": response_content.integration_type,
                        "ticket_id": response_content.ticket_id,
                        "ticket_status": response_content.ticket_status,
                        "ticket_priority": response_content.ticket_priority
                    })
                
                
                self._store_bot_message(chat_repo, {
                    "message": response_content.message,
                    "message_type": "bot",
                    "session_id": session_id,
                    "organization_id": org_id,
                    "agent_id": agent_id,
                    "customer_id": customer_id,
                    "attributes": attributes
                })

                return response_content

        except Exception as e:
            traceback.print_exc()
            logger.error(f"Chat agent error: {str(e)}")
            error_message = GENERIC_ERROR_MESSAGE
            
            # Create error response
            error_response = ChatResponse(
                message=error_message,
                transfer_to_human=False,
                transfer_reason=None,
                transfer_description=None,
                end_chat=False,
                end_chat_reason=None,
                end_chat_description=None,
                request_rating=False,
                create_ticket=False,
                shopify_output=None
            )
            
            # Store error message
            try:
                with SessionLocal() as db:
                    chat_repo = ChatRepository(db)
                    self._store_bot_message(chat_repo, {
                        "message": error_message,
                        "message_type": "bot",
                        "session_id": session_id,
                        "organization_id": org_id,
                        "agent_id": agent_id,
                        "customer_id": customer_id,
                        "attributes": {"error": str(e)}
                    })
            except Exception as store_error:
                logger.error(f"Failed to store error message: {str(store_error)}")
            
            return error_response

    @staticmethod
    async def test_api_key(api_key: str, model_type: str, model_name: str, base_url: str = None) -> bool:
        """Test if the API key is valid for the given model type.

        Args:
            api_key: The API key to test
            model_type: The type of model (OPENAI, ANTHROPIC, etc.)
            model_name: The name of the model
            base_url: Custom API endpoint, used by OPENAI_COMPATIBLE

        Returns:
            bool: True if the API key is valid

        Raises:
            ValueError: If the model type is not supported
        """
        try:
            from app.utils.agno_utils import test_model_api_key
            return await test_model_api_key(api_key, model_type, model_name, base_url=base_url)
        except Exception as e:
            traceback.print_exc()
            logger.error(f"Error testing API key: {str(e)}")
            return False