tests/tool_search_endpoints.py

"""
Tool 1: search_endpoints — Semantic search over endpoint catalog.

Uses GEMMA embeddings (google/embeddinggemma-300m) for semantic search.
Falls back to keyword matching when GEMMA is not available (test mode).
"""

import json
import os
import re
import math
from collections import Counter

# ---------------------------------------------------------------------------
# Keyword-based fallback search (for testing without GEMMA model)
# Uses TF-IDF-like scoring
# ---------------------------------------------------------------------------


def _tokenize(text: str) -> list[str]:
    """Simple whitespace + punctuation tokenizer."""
    return re.findall(r'[a-zA-Z0-9_/{}]+', text.lower())


class KeywordSearchIndex:
    """Simple TF-IDF search index for testing without neural embeddings."""

    def __init__(self):
        self.documents: list[str] = []
        self.doc_tokens: list[list[str]] = []
        self.idf: dict[str, float] = {}

    def add_documents(self, docs: list[str]):
        self.documents = docs
        self.doc_tokens = [_tokenize(d) for d in docs]
        self._build_idf()

    def _build_idf(self):
        n = len(self.documents)
        df = Counter()
        for tokens in self.doc_tokens:
            for t in set(tokens):
                df[t] += 1
        self.idf = {t: math.log(n / (1 + count)) for t, count in df.items()}

    def search(self, query: str, top_k: int = 3) -> list[tuple[int, float, str]]:
        """Returns list of (index, score, document) tuples."""
        query_tokens = _tokenize(query)
        scores = []
        for i, doc_toks in enumerate(self.doc_tokens):
            tf = Counter(doc_toks)
            score = sum(
                (tf.get(qt, 0) / max(len(doc_toks), 1)) * self.idf.get(qt, 0)
                for qt in query_tokens
            )
            scores.append((i, score, self.documents[i]))

        scores.sort(key=lambda x: x[1], reverse=True)
        return scores[:top_k]


# ---------------------------------------------------------------------------
# Catalog loading
# ---------------------------------------------------------------------------

def load_catalog(catalog_path: str) -> list[dict]:
    """Load a ground truth catalog JSON file."""
    with open(catalog_path) as f:
        data = json.load(f)
    return data.get("endpoints", data if isinstance(data, list) else [])


def catalog_entry_to_text(entry: dict, app_name: str = "") -> str:
    """Convert a catalog endpoint to a searchable text document."""
    parts = []
    if app_name:
        parts.append(f"app: {app_name}")

    endpoint = entry.get("endpoint", "")
    parts.append(f"endpoint: {endpoint}")

    auth = entry.get("auth", "none")
    parts.append(f"auth: {auth}")

    # Query params
    qp = entry.get("query_params", {})
    if qp:
        param_strs = []
        for k, v in qp.items():
            if isinstance(v, dict):
                param_strs.append(f"{k} ({v.get('type', '?')}, source: {v.get('source', '?')})")
            else:
                param_strs.append(f"{k}: {v}")
        parts.append(f"query_params: {', '.join(param_strs)}")

    # Path params
    pp = entry.get("path_params", {})
    if pp:
        param_strs = []
        for k, v in pp.items():
            if isinstance(v, dict):
                src = v.get("source", "?")
                from_ep = v.get("from_endpoint", "")
                param_strs.append(f"{k} ({v.get('type', '?')}, source: {src}, from: {from_ep})")
            else:
                param_strs.append(f"{k}: {v}")
        parts.append(f"path_params: {', '.join(param_strs)}")

    # Body params
    bp = entry.get("body_params", entry.get("form_params", {}))
    if bp:
        param_strs = []
        for k, v in bp.items():
            if isinstance(v, dict):
                src = v.get("source", "?")
                from_ep = v.get("from_endpoint", "")
                notes = v.get("notes", "")
                param_strs.append(f"{k} ({v.get('type', '?')}, source: {src})")
            else:
                param_strs.append(f"{k}: {v}")
        parts.append(f"body_params: {', '.join(param_strs)}")

    # Response fields
    rkf = entry.get("response_key_fields", [])
    if rkf:
        parts.append(f"returns: {', '.join(str(f) for f in rkf)}")

    # Notes
    notes = entry.get("notes", "")
    if notes:
        parts.append(f"notes: {notes}")

    return " | ".join(parts)


# ---------------------------------------------------------------------------
# search_endpoints tool
# ---------------------------------------------------------------------------

class SearchEndpoints:
    """
    Tool 1 implementation.
    Loads catalog, builds search index, provides search interface.
    """

    def __init__(self):
        self.index = KeywordSearchIndex()
        self.raw_entries: list[dict] = []
        self.text_chunks: list[str] = []

    def load_catalog(self, catalog_path: str, app_name: str = ""):
        """Load a catalog and build the search index."""
        self.raw_entries = load_catalog(catalog_path)
        self.text_chunks = [catalog_entry_to_text(e, app_name) for e in self.raw_entries]
        self.index.add_documents(self.text_chunks)

    def load_from_browser_agent(self, text_chunks: list[str]):
        """Load text chunks produced by browser_agent Stage 4."""
        self.text_chunks = text_chunks
        self.index.add_documents(text_chunks)

    def search(self, query: str, top_k: int = 3) -> list[str]:
        """
        Search endpoints by natural language query.
        Returns top-k matching endpoint schema texts.
        """
        results = self.index.search(query, top_k)
        return [doc for _, _, doc in results]

    def search_with_scores(self, query: str, top_k: int = 3) -> list[tuple[float, str]]:
        """Search with scores for debugging."""
        results = self.index.search(query, top_k)
        return [(score, doc) for _, score, doc in results]


# ---------------------------------------------------------------------------
# Test
# ---------------------------------------------------------------------------

if __name__ == "__main__":
    print("=" * 70)
    print("TEST: search_endpoints with browser_agent output")
    print("=" * 70)

    # PRIMARY TEST: load from browser_agent output (this is the real data flow)
    # In production, search_endpoints searches GEMMA embeddings built by browser_agent
    # from HAR data. Here we test with keyword search as a fallback for GEMMA.
    print("\n--- Primary: load from browser_agent HAR output ---")
    from tool_browser_agent import extract_openapi_spec, spec_entry_to_text

    mock_har_path = os.path.join(os.path.dirname(__file__), "mock_data", "mock_har.json")
    with open(mock_har_path) as f:
        har_data = json.load(f)

    spec = extract_openapi_spec(har_data, "http://localhost:7770/")
    chunks = [spec_entry_to_text(e, "shopping") for e in spec]

    tool = SearchEndpoints()
    tool.load_from_browser_agent(chunks)

    print(f"\nLoaded {len(tool.text_chunks)} endpoint documents from browser_agent output\n")
    for i, chunk in enumerate(tool.text_chunks):
        print(f"  [{i}] {chunk[:100]}...")

    # Test queries against browser_agent output
    queries = [
        "find product by name get sku",
        "create guest cart",
        "add item to guest cart",
        "authenticate customer login",
        "shipping methods for cart",
        "get cart total",
        "list categories",
    ]

    print(f"\n--- Search Results (from browser_agent HAR output) ---\n")
    for q in queries:
        print(f"Query: \"{q}\"")
        results = tool.search_with_scores(q, top_k=3)
        for score, doc in results:
            # Extract just the endpoint name for display
            ep_match = re.search(r'endpoint: (\S+ \S+)', doc)
            ep_name = ep_match.group(1) if ep_match else doc[:60]
            print(f"  [{score:.3f}] {ep_name}")
        print()

    # SECONDARY TEST: catalog loading (used by judge for ground truth, NOT by search_endpoints)
    print("--- Secondary: catalog loading (for judge ground truth, not search_endpoints) ---")
    catalog_path = os.path.join(os.path.dirname(__file__), "mock_data", "mock_catalog.json")

    tool2 = SearchEndpoints()
    tool2.load_catalog(catalog_path, app_name="shopping")
    print(f"  Catalog loaded: {len(tool2.text_chunks)} endpoint documents (judge reference only)")

    results = tool2.search("add item to cart", top_k=1)
    print(f"  Query: 'add item to cart' → top result:")
    print(f"    {results[0][:120]}...")

    print("\n[PASS] search_endpoints tool tests completed successfully")