57 lines
1.9 KiB
Python
57 lines
1.9 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,
|
|
"showMatchesPosition": True,
|
|
})
|
|
results = []
|
|
for hit in response.get("hits", []):
|
|
formatted = hit.get("_formatted", {})
|
|
line_number = None
|
|
matches = hit.get("_matchesPosition", {}).get("content", [])
|
|
if matches and hit.get("content"):
|
|
start = matches[0]["start"]
|
|
line_number = hit["content"][:start].count("\n") + 1
|
|
results.append({
|
|
"key": hit["key"],
|
|
"title": hit["title"],
|
|
"path": hit["path"],
|
|
"snippet": formatted.get("content", hit.get("content", "")),
|
|
"lineNumber": line_number,
|
|
})
|
|
return results
|