Spaces:
Running
Running
| """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 | |
| 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() | |