reranker.py 916 B

123456789101112131415161718192021222324252627282930
  1. import httpx
  2. from .config import RERANKER_URL
  3. def rerank_results(query: str, items: list, top_k: int) -> list:
  4. """Rerank mem0 search results via local reranker, fallback to original order."""
  5. if not items:
  6. return items
  7. documents = [r.get("memory", "") for r in items]
  8. try:
  9. response = httpx.post(
  10. RERANKER_URL,
  11. json={"query": query, "documents": documents, "top_k": top_k},
  12. timeout=5.0,
  13. )
  14. response.raise_for_status()
  15. reranked = response.json()["results"]
  16. except Exception as exc:
  17. print(f"[reranker] unavailable, skipping: {exc}")
  18. return items[:top_k]
  19. text_to_meta = {r.get("memory", ""): r for r in items}
  20. merged = []
  21. for item in reranked:
  22. meta = text_to_meta.get(item["text"])
  23. if meta:
  24. merged.append({**meta, "rerank_score": item["score"]})
  25. return merged