"""OpenSearch persistence layer — the ABM engine's primary store. Aiven OpenSearch (v3.6, k-NN plugin enabled, cluster green) carries accounts, contacts, signal events and the vector memory index. Everything degrades gracefully: if OpenSearch is unreachable the service reports unhealthy in /health instead of crashing, so the Space can boot before the datastore. """ import asyncio import logging import uuid from datetime import datetime, timezone from typing import Any, Dict, List, Optional import httpx from ..config import settings logger = logging.getLogger("abm.os") def _now() -> str: return datetime.now(timezone.utc).isoformat() class OpenSearchStore: def __init__(self) -> None: self._client: Optional[httpx.AsyncClient] = None @property def configured(self) -> bool: return settings.os_configured def client(self) -> httpx.AsyncClient: if self._client is None: auth = (settings.os_user, settings.os_password) if settings.os_user else None self._client = httpx.AsyncClient( base_url=settings.os_url.rstrip("/"), auth=auth, timeout=httpx.Timeout(30.0), verify=True, ) return self._client async def close(self) -> None: if self._client is not None: await self._client.aclose() self._client = None # ── index bootstrap ────────────────────────────────────────────────── def _index(self, name: str) -> str: return f"{settings.os_index_prefix}-{name}" async def _put_index(self, name: str, body: dict) -> bool: res = await self.client().put(f"/{self._index(name)}", json=body) if res.status_code in (200, 201): return True if res.status_code == 400 and "resource_already_exists" in res.text: return True logger.warning("put_index %s failed: %s %s", name, res.status_code, res.text[:200]) return False async def bootstrap(self) -> dict: """Create the four indices if missing. Returns per-index status.""" dim = settings.embedding_dim knn_vector = { "type": "knn_vector", "dimension": dim, # nmslib is deprecated (cannot create new indexes on OpenSearch 3.x) "method": {"name": "hnsw", "space_type": "cosinesimil", "engine": "faiss"}, } results: Dict[str, bool] = {} results["accounts"] = await self._put_index("accounts", { "settings": {"index": {"knn": True}}, "mappings": {"properties": { "account_id": {"type": "keyword"}, "domain": {"type": "keyword"}, "name": {"type": "text", "fields": {"kw": {"type": "keyword"}}}, "industry": {"type": "keyword"}, "employee_count": {"type": "integer"}, "hq_country": {"type": "keyword"}, "city": {"type": "keyword"}, "website": {"type": "keyword"}, "phone": {"type": "keyword"}, "description": {"type": "text"}, "source": {"type": "keyword"}, "icp_score": {"type": "float"}, "icp_label": {"type": "keyword"}, "signals": {"type": "flat_object"}, "embedding": knn_vector, "updated_at": {"type": "date"}, }}, }) results["contacts"] = await self._put_index("contacts", { "mappings": {"properties": { "contact_id": {"type": "keyword"}, "account_id": {"type": "keyword"}, "email": {"type": "keyword"}, "full_name": {"type": "text"}, "title": {"type": "keyword"}, "linkedin_url": {"type": "keyword"}, "phone": {"type": "keyword"}, "source": {"type": "keyword"}, "updated_at": {"type": "date"}, }}, }) results["events"] = await self._put_index("events", { "mappings": {"properties": { "event_id": {"type": "keyword"}, "account_id": {"type": "keyword"}, "type": {"type": "keyword"}, "payload": {"type": "flat_object"}, "created_at": {"type": "date"}, }}, }) results["memory"] = await self._put_index("memory", { "settings": {"index": {"knn": True}}, "mappings": {"properties": { "doc_id": {"type": "keyword"}, "text": {"type": "text"}, "metadata": {"type": "flat_object"}, "embedding": knn_vector, "created_at": {"type": "date"}, }}, }) return results # ── generic helpers ────────────────────────────────────────────────── async def index_doc(self, index: str, doc: dict, doc_id: Optional[str] = None) -> str: _id = doc_id or doc.get("account_id") or doc.get("contact_id") or doc.get("event_id") or doc.get("doc_id") or str(uuid.uuid4()) res = await self.client().put( f"/{self._index(index)}/_doc/{_id}", json=doc, params={"refresh": "true"}, ) res.raise_for_status() return _id async def upsert_doc(self, index: str, doc: dict, doc_id: str) -> str: # Docs read via get_doc carry OpenSearch metadata keys (_id/_score); # writing them back is rejected as metadata fields. clean = {k: v for k, v in doc.items() if not str(k).startswith("_")} res = await self.client().put( f"/{self._index(index)}/_doc/{doc_id}", json=clean, params={"refresh": "true"}, ) res.raise_for_status() return doc_id async def get_doc(self, index: str, doc_id: str) -> Optional[dict]: res = await self.client().get(f"/{self._index(index)}/_doc/{doc_id}") if res.status_code == 404: return None res.raise_for_status() body = res.json() if not body.get("found"): return None return {"_id": body["_id"], **body["_source"]} async def delete_doc(self, index: str, doc_id: str) -> bool: res = await self.client().delete(f"/{self._index(index)}/_doc/{doc_id}") return res.status_code in (200, 404) async def search(self, index: str, body: dict) -> List[dict]: res = await self.client().post(f"/{self._index(index)}/_search", json=body) res.raise_for_status() hits = res.json().get("hits", {}).get("hits", []) return [{"_id": h["_id"], "_score": h.get("_score"), **h["_source"]} for h in hits] async def refresh(self, index: str) -> None: try: await self.client().post(f"/{self._index(index)}/_refresh") except Exception: pass async def count(self, index: str) -> int: res = await self.client().get(f"/{self._index(index)}/_count") if res.status_code != 200: return -1 return int(res.json().get("count", -1)) async def ping(self) -> bool: try: res = await self.client().get("/") return res.status_code == 200 except Exception: return False store = OpenSearchStore()