HARvestGym / server /tools /search_endpoints.py
kdcyberdude's picture
Upload folder using huggingface_hub
e6ce96e verified
Raw
History Blame
3.09 kB
"""
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