Jason Thistlethwaite
73 lines
2.8 KiB
Python
73 lines
2.8 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import urllib.request
|
|
from typing import Any, Dict, Optional
|
|
|
|
from .app import build_services
|
|
from .models import SearchQuery, search_response
|
|
|
|
|
|
class SemanticIndexClient:
|
|
def __init__(
|
|
self,
|
|
base_url: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
search_service: Optional[Any] = None,
|
|
) -> None:
|
|
self.base_url = base_url.rstrip("/") if base_url else None
|
|
self.api_key = api_key
|
|
self.search_service = search_service
|
|
|
|
@classmethod
|
|
def local(cls) -> "SemanticIndexClient":
|
|
return cls(search_service=build_services()["search"])
|
|
|
|
def search(self, query: str, **filters: Any) -> Dict[str, Any]:
|
|
if self.base_url:
|
|
return self._post_json("/search", {"query": query, **filters})
|
|
search_service = self.search_service or build_services()["search"]
|
|
search_query = SearchQuery(
|
|
text=query,
|
|
source=filters.get("source"),
|
|
project_id=filters.get("project_id"),
|
|
project_identifier=filters.get("project_identifier"),
|
|
doc_type=filters.get("doc_type"),
|
|
issue_id=filters.get("issue_id"),
|
|
contact_id=filters.get("contact_id"),
|
|
contact_email=filters.get("contact_email"),
|
|
date_from=filters.get("date_from"),
|
|
date_to=filters.get("date_to"),
|
|
limit=int(filters.get("limit", 10)),
|
|
include_snippets=bool(filters.get("include_snippets", True)),
|
|
)
|
|
return search_response(search_query, search_service.search(search_query))
|
|
|
|
def get_document(self, document_id: str) -> Dict[str, Any]:
|
|
if self.base_url:
|
|
return self._get_json(f"/documents/{document_id}")
|
|
search_service = self.search_service or build_services()["search"]
|
|
return search_service.get_document(document_id) or {"error": "not_found", "id": document_id}
|
|
|
|
def _post_json(self, path: str, payload: Dict[str, Any]) -> Dict[str, Any]:
|
|
data = json.dumps(payload).encode("utf-8")
|
|
request = urllib.request.Request(
|
|
f"{self.base_url}{path}",
|
|
data=data,
|
|
headers=self._headers(),
|
|
method="POST",
|
|
)
|
|
with urllib.request.urlopen(request, timeout=60) as response:
|
|
return json.loads(response.read().decode("utf-8"))
|
|
|
|
def _get_json(self, path: str) -> Dict[str, Any]:
|
|
request = urllib.request.Request(f"{self.base_url}{path}", headers=self._headers())
|
|
with urllib.request.urlopen(request, timeout=60) as response:
|
|
return json.loads(response.read().decode("utf-8"))
|
|
|
|
def _headers(self) -> Dict[str, str]:
|
|
headers = {"Content-Type": "application/json"}
|
|
if self.api_key:
|
|
headers["Authorization"] = f"Bearer {self.api_key}"
|
|
return headers
|