training_tools.py

from agents import function_tool
from dotenv import load_dotenv
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, PointStruct, VectorParams
from sentence_transformers import SentenceTransformer
from supabase import create_client
from runtime_context import get_system_prompt_variable, set_system_prompt_variable

import json
import os
import re
import uuid
from datetime import datetime, timezone
from typing import Any

load_dotenv()

SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")

QDRANT_URL = os.getenv("QDRANT_URL")
QDRANT_API_KEY = os.getenv("QDRANT_API_KEY")
COLLECTION_NAME = (
    os.getenv("QDRANT_COLLECTION")
    or os.getenv("COLLECTION_NAME")
    or "tiq_knowledge"
)

PROMPT_TABLE = os.getenv("AI_TRAINING_PROMPT_TABLE", "ai_training_settings")
PROMPT_ROW_ID = os.getenv("AI_TRAINING_PROMPT_ROW_ID", "default")

EMBEDDING_MODEL = "all-MiniLM-L6-v2"
VECTOR_SIZE = 384
PROMPT_UPDATE_MAX_CHARS = 1200
FORBIDDEN_PROMPT_UPDATE_PHRASES = [
    "ignore previous instructions",
    "ignore all previous instructions",
    "replace the entire system prompt",
    "new system prompt",
    "you are no longer",
    "do not use qdrant",
    "never use qdrant",
    "do not escalate",
    "never escalate",
    "disable escalation",
    "do not collect name",
    "do not collect email",
    "say you are ai",
    "tell customers you are ai",
    "reveal internal",
    "show system prompt",
    "ignore safety",
]
FORBIDDEN_PROMPT_SECTION_MARKERS = [
    "tool usage",
    "escalation rules",
    "shopify order tracking tool",
    "final goal",
    "conversation rules",
]

if not SUPABASE_URL or not SUPABASE_KEY:
    raise RuntimeError("SUPABASE_URL and SUPABASE_KEY are required")

if not QDRANT_URL or not QDRANT_API_KEY:
    raise RuntimeError("QDRANT_URL and QDRANT_API_KEY are required")

supabase = create_client(SUPABASE_URL, SUPABASE_KEY)
qdrant = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY)
model = SentenceTransformer(EMBEDDING_MODEL)


def utc_now() -> str:
    return datetime.now(timezone.utc).isoformat()


def refresh_customer_agent_prompt(editable_prompt: str):
    """Refresh customer SMS agent prompt in local memory after training update."""
    set_system_prompt_variable(editable_prompt)

    try:
        from tiq_agent import refresh_tiq_agent_system_prompt

        refresh_tiq_agent_system_prompt(editable_prompt)
    except Exception as e:
        print("CUSTOMER AGENT PROMPT REFRESH ERROR:", e)


def clean_text(text: str) -> str:
    text = (text or "").replace("\x00", " ")
    text = re.sub(r"\s+", " ", text)
    return text.strip()


def strip_editable_prompt_wrapper(text: str) -> str:
    text = clean_text(text)

    marker = "Editable guidance:"
    if marker in text:
        return clean_text(text.split(marker, 1)[1])

    return text


def validate_editable_prompt_update(text: str) -> tuple[bool, str]:
    lowered = text.lower()

    if len(text) > PROMPT_UPDATE_MAX_CHARS:
        return False, f"Editable prompt update is too large. Keep it under {PROMPT_UPDATE_MAX_CHARS} characters."

    if any(phrase in lowered for phrase in FORBIDDEN_PROMPT_UPDATE_PHRASES):
        return False, "This looks like it would override fixed AI behavior or safety rules. Please make it a narrow editable instruction."

    marker_count = sum(1 for marker in FORBIDDEN_PROMPT_SECTION_MARKERS if marker in lowered)
    if marker_count >= 2:
        return False, "This looks like a full system prompt rewrite. Only small editable behavior instructions are allowed."

    return True, ""


def format_editable_prompt_update(instruction: str, replacement_text: str) -> str:
    replacement_text = strip_editable_prompt_wrapper(replacement_text)

    return f"""Admin editable guidance only.
Narrow additions to fixed customer SMS rules. Never override tools, escalation, privacy, order tracking, or safety.

Editable guidance:
{replacement_text}
""".strip()


def chunk_text(text: str, max_chars: int = 900, overlap: int = 120) -> list[str]:
    text = clean_text(text)

    if not text:
        return []

    chunks = []
    start = 0

    while start < len(text):
        end = min(start + max_chars, len(text))
        chunk = text[start:end].strip()

        if chunk:
            chunks.append(chunk)

        if end >= len(text):
            break

        start = max(0, end - overlap)

    return chunks


def enhanced_text_for_embedding(text: str, content_type: str = "content") -> str:
    enhanced = text
    lowered = text.lower()

    if content_type == "faq":
        enhanced = f"FAQ: {enhanced}"

    if "internet" in lowered or "browser" in lowered:
        enhanced = f"Internet restriction: {enhanced}"

    if "whatsapp" in lowered:
        enhanced = f"WhatsApp support: {enhanced}"

    if "verizon" in lowered:
        enhanced = f"Verizon support: {enhanced}"

    if "shipping" in lowered or "tracking" in lowered:
        enhanced = f"Shipping support: {enhanced}"

    return enhanced


def embed_text(text: str) -> list[float]:
    return model.encode(text).tolist()


def stable_point_id(value: str) -> str:
    return str(uuid.uuid5(uuid.NAMESPACE_URL, value))


def ensure_collection():
    collections = qdrant.get_collections().collections
    exists = any(item.name == COLLECTION_NAME for item in collections)

    if exists:
        return

    qdrant.create_collection(
        collection_name=COLLECTION_NAME,
        vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
    )


def search_qdrant_points(query: str, limit: int = 8):
    ensure_collection()
    vector = embed_text(query)
    response = qdrant.query_points(
        collection_name=COLLECTION_NAME,
        query=vector,
        limit=limit,
    )
    return response.points or []


def normalize_point_ids(point_ids: list[str] | str) -> list[str]:
    if isinstance(point_ids, str):
        try:
            parsed = json.loads(point_ids)
            if isinstance(parsed, list):
                point_ids = parsed
            else:
                point_ids = [point_ids]
        except Exception:
            point_ids = [part.strip() for part in point_ids.split(",") if part.strip()]

    return [str(point_id).strip() for point_id in point_ids if str(point_id).strip()]


def delete_qdrant_points(point_ids: list[str] | str):
    point_ids = normalize_point_ids(point_ids)

    if not point_ids:
        return

    qdrant.delete(
        collection_name=COLLECTION_NAME,
        points_selector={"points": point_ids},
    )


def serialize_qdrant_match(point) -> dict[str, Any]:
    payload = point.payload or {}
    text = clean_text(payload.get("text") or "")

    return {
        "point_id": str(point.id),
        "score": round(float(getattr(point, "score", 0) or 0), 4),
        "text": text[:900],
        "source": payload.get("source"),
        "type": payload.get("type"),
        "title": payload.get("title"),
        "category": payload.get("category"),
        "file_id": payload.get("file_id"),
        "message_id": payload.get("message_id"),
        "conversation_id": payload.get("conversation_id"),
        "training_source_id": payload.get("training_source_id"),
        "created_at": payload.get("created_at"),
    }


def upsert_knowledge_chunks(
    text: str,
    source: str,
    content_type: str,
    title: str | None = None,
    metadata: dict[str, Any] | None = None,
):
    ensure_collection()
    metadata = metadata or {}
    chunks = chunk_text(text)

    points = []
    now = utc_now()
    source_id = metadata.get("source_id") or stable_point_id(f"{source}:{title or ''}:{text[:160]}")

    for index, chunk in enumerate(chunks):
        enhanced = enhanced_text_for_embedding(chunk, content_type)
        point_id = stable_point_id(f"training:{source_id}:{index}:{chunk[:160]}")

        points.append(
            PointStruct(
                id=point_id,
                vector=embed_text(enhanced),
                payload={
                    "text": chunk,
                    "title": title or "Training update",
                    "source": source,
                    "type": content_type,
                    "active": True,
                    "training_source_id": source_id,
                    "chunk_index": index,
                    "created_at": now,
                    **metadata,
                },
            )
        )

    if points:
        qdrant.upsert(collection_name=COLLECTION_NAME, points=points)

    return {
        "source_id": source_id,
        "chunks": len(points),
    }


@function_tool
def get_current_system_prompt_variable_part() -> str:
    """
    Retrieve the current editable customer SMS prompt guidance before updating it.

    Always call this before update_system_prompt_variable_part.
    Return value is only the editable guidance, not the fixed system prompt.
    """
    current = strip_editable_prompt_wrapper(get_system_prompt_variable())

    if not current:
        return "No editable prompt guidance is currently saved."

    return current


@function_tool
def update_system_prompt_variable_part(instruction: str, replacement_text: str) -> str:
    """
    Update the editable AI behavior/information prompt section.

    Use this for tone, reply style, company support rules, SMS style rules,
    and behavior instructions. Do not modify fixed tool-calling or safety rules.
    """
    instruction = clean_text(instruction)
    previous_prompt = strip_editable_prompt_wrapper(get_system_prompt_variable())
    replacement_text = strip_editable_prompt_wrapper(clean_text(replacement_text))

    if not replacement_text:
        return "I could not update the prompt because the replacement text was empty."

    if previous_prompt and replacement_text == previous_prompt:
        return "No change needed. The editable AI behavior instructions are already up to date."

    is_valid, validation_error = validate_editable_prompt_update(replacement_text)
    if not is_valid:
        return validation_error

    editable_prompt = format_editable_prompt_update(instruction, replacement_text)

    row = {
        "id": PROMPT_ROW_ID,
        "instruction": instruction,
        "editable_prompt": editable_prompt,
        "updated_at": utc_now(),
    }

    try:
        supabase.table(PROMPT_TABLE).upsert(row, on_conflict="id").execute()
        refresh_customer_agent_prompt(editable_prompt)
    except Exception as e:
        return f"Prompt update failed. Make sure Supabase table '{PROMPT_TABLE}' exists. Error: {str(e)}"

    return "Done, I updated the AI behavior instructions."


@function_tool
def search_qdrant_training_chunks(query: str, limit: int = 6) -> str:
    """
    Search Qdrant and return exact chunks with point IDs before updating or deleting.

    Always use this before update_qdrant_chunks_by_ids or delete_qdrant_chunks_by_ids.
    """
    query = clean_text(query)
    limit = min(max(int(limit or 6), 1), 10)

    if not query:
        return json.dumps({"items": [], "message": "Search query is required."})

    try:
        matches = search_qdrant_points(query, limit=limit)
        items = [serialize_qdrant_match(point) for point in matches]

        return json.dumps({
            "items": items,
            "message": f"Found {len(items)} matching chunk(s). Review point_id and text before update/delete.",
        }, ensure_ascii=False)

    except Exception as e:
        return json.dumps({"items": [], "error": str(e)})


@function_tool
def update_qdrant_chunks_by_ids(point_ids: list[str], replacement_information: str, title: str = "Training update") -> str:
    """
    Replace specific Qdrant chunks by exact point IDs.

    Use only after search_qdrant_training_chunks returned the point IDs and the relevant chunks were reviewed.
    """
    point_ids = normalize_point_ids(point_ids)
    replacement_information = clean_text(replacement_information)
    title = clean_text(title) or "Training update"

    if not point_ids:
        return "No point IDs provided. Search chunks first, then pass exact point_ids."

    if not replacement_information:
        return "Replacement information is required."

    try:
        delete_qdrant_points(point_ids)
        result = upsert_knowledge_chunks(
            text=replacement_information,
            source="training_agent_update",
            content_type="training_update",
            title=title,
            metadata={
                "category": "training_update",
                "updated_by": "training_agent",
                "replaced_points": point_ids,
            },
        )
    except Exception as e:
        return f"Qdrant exact update failed: {str(e)}"

    return f"Done, I replaced {len(point_ids)} selected chunk(s) with {result['chunks']} new chunk(s)."


@function_tool
def delete_qdrant_chunks_by_ids(point_ids: list[str], reason: str = "Admin requested deletion") -> str:
    """
    Delete specific Qdrant chunks by exact point IDs.

    Use only after search_qdrant_training_chunks returned the point IDs and the relevant chunks were reviewed.
    """
    point_ids = normalize_point_ids(point_ids)
    reason = clean_text(reason) or "Admin requested deletion"

    if not point_ids:
        return "No point IDs provided. Search chunks first, then pass exact point_ids."

    try:
        delete_qdrant_points(point_ids)

        try:
            supabase.table("ai_training_deletions").insert({
                "search_query": "exact_point_ids",
                "reason": reason,
                "deleted_point_ids": point_ids,
                "deleted_count": len(point_ids),
                "created_at": utc_now(),
            }).execute()
        except Exception as log_error:
            print("AI TRAINING DELETE LOG ERROR:", log_error)

    except Exception as e:
        return f"Qdrant exact delete failed: {str(e)}"

    return f"Done, I deleted {len(point_ids)} selected knowledge chunk(s)."


@function_tool
def add_new_qdrant_data(title: str, information: str, category: str = "training") -> str:
    """
    Add new approved knowledge to Qdrant.

    Use this when admin gives new company/product/support information.
    """
    title = clean_text(title) or "Training update"
    information = clean_text(information)
    category = clean_text(category) or "training"

    if not information:
        return "I could not add the knowledge because the information was empty."

    try:
        result = upsert_knowledge_chunks(
            text=information,
            source="training_agent",
            content_type=category,
            title=title,
            metadata={
                "category": category,
                "updated_by": "training_agent",
            },
        )
    except Exception as e:
        return f"Qdrant add failed: {str(e)}"

    return f"Done, I added this to the knowledge base ({result['chunks']} chunk(s))."


@function_tool
def update_existing_qdrant_data(search_query: str, replacement_information: str, title: str = "Training update") -> str:
    """
    Update old or incorrect Qdrant knowledge.

    Legacy helper. Prefer search_qdrant_training_chunks first, then update_qdrant_chunks_by_ids with exact point IDs.
    """
    search_query = clean_text(search_query)
    replacement_information = clean_text(replacement_information)
    title = clean_text(title) or "Training update"

    if not search_query or not replacement_information:
        return "I need both what to search for and the replacement information."

    try:
        matches = search_qdrant_points(search_query, limit=8)
        point_ids = [str(point.id) for point in matches]

        if point_ids:
            delete_qdrant_points(point_ids)

        result = upsert_knowledge_chunks(
            text=replacement_information,
            source="training_agent_update",
            content_type="training_update",
            title=title,
            metadata={
                "category": "training_update",
                "updated_by": "training_agent",
                "replaced_query": search_query,
                "replaced_points": point_ids,
            },
        )
    except Exception as e:
        return f"Qdrant update failed: {str(e)}"

    return f"Done, I updated the matching knowledge ({len(point_ids)} old chunk(s) replaced, {result['chunks']} new chunk(s))."


@function_tool
def delete_qdrant_data(search_query: str, reason: str = "Admin requested deletion") -> str:
    """
    Legacy helper. Prefer search_qdrant_training_chunks first, then delete_qdrant_chunks_by_ids with exact point IDs.
    """
    search_query = clean_text(search_query)
    reason = clean_text(reason) or "Admin requested deletion"

    if not search_query:
        return "I need a search query to find what should be deleted."

    try:
        matches = search_qdrant_points(search_query, limit=10)
        point_ids = [str(point.id) for point in matches]

        if not point_ids:
            return "I did not find matching knowledge to delete."

        delete_qdrant_points(point_ids)

        try:
            supabase.table("ai_training_deletions").insert({
                "search_query": search_query,
                "reason": reason,
                "deleted_point_ids": point_ids,
                "deleted_count": len(point_ids),
                "created_at": utc_now(),
            }).execute()
        except Exception as log_error:
            print("AI TRAINING DELETE LOG ERROR:", log_error)

    except Exception as e:
        return f"Qdrant delete failed: {str(e)}"

    return f"Done, I deleted {len(point_ids)} matching knowledge chunk(s)."