DocHub/backend/app/search_client.py

50 lines
1.6 KiB
Python

import base64
def _encode_key(key: str) -> str:
return base64.urlsafe_b64encode(key.encode("utf-8")).decode("ascii").rstrip("=")
class SearchClient:
def __init__(self, client, index_name: str):
self._client = client
self._index_name = index_name
def index_document(self, key: str, path: str, title: str, content: str) -> None:
index = self._client.index(self._index_name)
index.add_documents([{
"id": _encode_key(key),
"key": key,
"path": path,
"title": title,
"content": content,
}])
def delete_document(self, key: str) -> None:
index = self._client.index(self._index_name)
index.delete_document(_encode_key(key))
def delete_all_documents(self) -> None:
index = self._client.index(self._index_name)
index.delete_all_documents()
def search(self, query: str) -> list[dict]:
index = self._client.index(self._index_name)
response = index.search(query, {
"attributesToHighlight": ["content"],
"highlightPreTag": "<em>",
"highlightPostTag": "</em>",
"attributesToCrop": ["content"],
"cropLength": 30,
})
results = []
for hit in response.get("hits", []):
formatted = hit.get("_formatted", {})
results.append({
"key": hit["key"],
"title": hit["title"],
"path": hit["path"],
"snippet": formatted.get("content", hit.get("content", "")),
})
return results