File size: 3,091 Bytes
e6ce96e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
"""
search_endpoints tool — semantic search over endpoint embeddings from browser_agent.

Returns top-3 endpoint schemas (full text) for a natural language query.
"""

from __future__ import annotations

import os
import numpy as np


def search_endpoints(query: str, episode_store: dict) -> list[str]:
    """
    Semantic search over endpoint embeddings built by browser_agent.

    Args:
        query: Natural language query (e.g. "create guest cart", "add item to cart")
        episode_store: Mutable dict containing embeddings from browser_agent.

    Returns:
        List of up to 3 endpoint schema text strings.
    """
    chunks: list[str] = episode_store.get("endpoint_chunks", [])
    embeddings = episode_store.get("endpoint_embeddings")
    spec_entries: list[dict] = episode_store.get("spec_entries", [])

    if not chunks:
        return ["No endpoint index available. Call browser_agent(task, url) first."]

    # If no embeddings, use keyword fallback directly
    if embeddings is None or (hasattr(embeddings, '__len__') and len(embeddings) == 0):
        query_lower = query.lower()
        query_terms = query_lower.split()
        matches = []
        for chunk in chunks:
            if any(term in chunk.lower() for term in query_terms):
                matches.append(chunk)
        return matches[:3] if matches else chunks[:3]

    try:
        model, model_name = _get_embedding_model()

        if hasattr(model, "encode_query"):
            q_emb = model.encode_query([query], show_progress_bar=False)
        else:
            q_emb = model.encode([query], show_progress_bar=False)

        if not isinstance(q_emb, np.ndarray):
            q_emb = np.array(q_emb)

        # Normalize query embedding
        norm = np.linalg.norm(q_emb, axis=1, keepdims=True)
        if norm[0, 0] > 0:
            q_emb = q_emb / norm

        # Cosine similarity (embeddings already normalized)
        scores = (embeddings @ q_emb.T).flatten()
        top_k = min(3, len(scores))
        top_indices = np.argsort(scores)[::-1][:top_k]

        results = []
        for idx in top_indices:
            results.append(chunks[int(idx)])
        return results

    except Exception as e:
        # Fallback: keyword match
        print(f"[search_endpoints] Embedding search failed: {e}. Using keyword fallback.", flush=True)
        query_lower = query.lower()
        matches = []
        for chunk in chunks:
            if any(word in chunk.lower() for word in query_lower.split()):
                matches.append(chunk)
        return matches[:3] if matches else chunks[:3]


# Lazy-load the model (shared with browser_agent)
_embedding_model = None
_embedding_model_name = None


def _get_embedding_model():
    global _embedding_model, _embedding_model_name
    if _embedding_model is not None:
        return _embedding_model, _embedding_model_name

    # Re-use browser_agent's loader
    from .browser_agent import _get_embedding_model as _ba_get
    model, name = _ba_get()
    _embedding_model = model
    _embedding_model_name = name
    return model, name