Spaces:
Running
Running
Z User
feat: ABM Engine v1 — OpenSearch store + Apify enrichment + ICP scoring + k-NN memory + optional Neo4j mirror
fb0ec21 | """Optional graph mirror — BetterDB (managed Neo4j) or direct Aura. | |
| The ABM graph (accounts, contacts, enrichment events) mirrors into Neo4j when | |
| a connection is available; the engine runs fully without it (OpenSearch is the | |
| source of truth). Two transports, auto-detected at boot: | |
| 1. BETTERDB_URL + BETTERDB_TOKEN — managed Neo4j behind betterdb.com. The | |
| HTTP query API is probed at a few known paths; BetterDB answers "waiting" | |
| until the owner connects their instance in the console, in which case the | |
| mirror stays disabled and /health reports it honestly. | |
| 2. NEO4J_URI (+user/password) — a direct Neo4j/Aura endpoint via the official | |
| driver (bolt+ssc for Aura). | |
| """ | |
| import logging | |
| from datetime import datetime, timezone | |
| from typing import Any, Dict, List, Optional | |
| import httpx | |
| from ..config import settings | |
| logger = logging.getLogger("abm.graph") | |
| class GraphStore: | |
| def __init__(self) -> None: | |
| self._mode: Optional[str] = None # "betterdb" | "neo4j" | None | |
| self._driver = None | |
| self._http: Optional[httpx.AsyncClient] = None | |
| self._query_path: Optional[str] = None | |
| def mode(self) -> Optional[str]: | |
| return self._mode | |
| async def connect(self) -> bool: | |
| if settings.neo4j_uri: | |
| try: | |
| from neo4j import GraphDatabase # optional dependency | |
| scheme = "neo4j+ssc" if "+ssc" not in settings.neo4j_uri else None | |
| uri = settings.neo4j_uri | |
| if uri.startswith("neo4j://"): | |
| uri = uri.replace("neo4j://", "neo4j+ssc://", 1) | |
| elif scheme and not uri.startswith("neo4j+"): | |
| uri = f"{scheme}://{uri}" | |
| self._driver = GraphDatabase.driver( | |
| uri, auth=(settings.neo4j_user or "neo4j", settings.neo4j_password or "") | |
| ) | |
| await asyncio_wrap(lambda: self._driver.verify_connectivity()) | |
| self._mode = "neo4j" | |
| logger.info("graph: direct Neo4j connected (%s)", uri) | |
| return True | |
| except Exception as exc: | |
| logger.warning("graph: direct Neo4j connect failed: %s", exc) | |
| self._driver = None | |
| if settings.betterdb_url and settings.betterdb_token: | |
| self._http = httpx.AsyncClient( | |
| base_url=settings.betterdb_url.rstrip("/"), | |
| headers={"Authorization": f"Bearer {settings.betterdb_token}"}, | |
| timeout=httpx.Timeout(20.0), | |
| ) | |
| for path in ( | |
| "/api/neo4j/query/v2", | |
| "/neo4j/query/v2", | |
| "/api/db/neo4j/query/v2", | |
| "/db/neo4j/query/v2", | |
| "/api/query", | |
| ): | |
| try: | |
| res = await self._http.post( | |
| path, json={"statement": "RETURN 1 AS ok"} | |
| ) | |
| if res.status_code == 200: | |
| self._query_path = path | |
| self._mode = "betterdb" | |
| logger.info("graph: BetterDB query API live at %s", path) | |
| return True | |
| if res.status_code in (401, 402, 403): | |
| logger.warning("graph: BetterDB answered %s at %s — token %s", | |
| res.status_code, path, "rejected") | |
| except Exception: | |
| continue | |
| logger.info("graph: BetterDB reachable but no query API answered yet " | |
| "(owner must connect the instance in the console)") | |
| return False | |
| async def close(self) -> None: | |
| if self._http is not None: | |
| await self._http.aclose() | |
| if self._driver is not None: | |
| self._driver.close() | |
| async def run(self, cypher: str, params: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]: | |
| """Run a write/query statement; returns rows (best-effort, logged).""" | |
| if self._mode == "betterdb" and self._query_path: | |
| try: | |
| res = await self._http.post( | |
| self._query_path, json={"statement": cypher, "parameters": params or {}} | |
| ) | |
| if res.status_code != 200: | |
| logger.warning("graph query %s: %s", res.status_code, res.text[:150]) | |
| return [] | |
| body = res.json() | |
| return body.get("records") or body.get("data") or [] | |
| except Exception as exc: | |
| logger.warning("graph query failed: %s", exc) | |
| return [] | |
| if self._mode == "neo4j" and self._driver is not None: | |
| def _run(): | |
| with self._driver.session() as session: | |
| result = session.run(cypher, params or {}) | |
| return [r.data() for r in result] | |
| try: | |
| return await asyncio_wrap(_run) | |
| except Exception as exc: | |
| logger.warning("graph query failed: %s", exc) | |
| return [] | |
| return [] | |
| # ── mirror API (idempotent MERGEs) ─────────────────────────────────── | |
| async def upsert_account(self, account: Dict[str, Any]) -> None: | |
| if self._mode is None: | |
| return | |
| await self.run( | |
| """ | |
| MERGE (a:Account {account_id: $account_id}) | |
| SET a.name = $name, a.domain = $domain, a.industry = $industry, | |
| a.hq_country = $hq_country, a.icp_score = $icp_score, | |
| a.updated_at = $updated_at | |
| """, | |
| { | |
| "account_id": account.get("account_id"), | |
| "name": account.get("name"), | |
| "domain": account.get("domain"), | |
| "industry": account.get("industry"), | |
| "hq_country": account.get("hq_country"), | |
| "icp_score": account.get("icp_score"), | |
| "updated_at": datetime.now(timezone.utc).isoformat(), | |
| }, | |
| ) | |
| async def upsert_contact(self, account_id: str, contact: Dict[str, Any]) -> None: | |
| if self._mode is None: | |
| return | |
| await self.run( | |
| """ | |
| MATCH (a:Account {account_id: $account_id}) | |
| MERGE (c:Contact {contact_id: $contact_id}) | |
| SET c.full_name = $full_name, c.title = $title, c.email = $email, | |
| c.linkedin_url = $linkedin_url | |
| MERGE (c)-[:WORKS_AT]->(a) | |
| """, | |
| { | |
| "account_id": account_id, | |
| "contact_id": contact.get("contact_id"), | |
| "full_name": contact.get("full_name"), | |
| "title": contact.get("title"), | |
| "email": contact.get("email"), | |
| "linkedin_url": contact.get("linkedin_url"), | |
| }, | |
| ) | |
| async def overview(self) -> Dict[str, Any]: | |
| if self._mode is None: | |
| return {"mode": None, "status": "not connected"} | |
| rows = await self.run( | |
| "MATCH (a:Account) WITH count(a) AS accounts " | |
| "OPTIONAL MATCH (c:Contact) RETURN accounts, count(c) AS contacts" | |
| ) | |
| return {"mode": self._mode, "status": "connected", "counts": rows or []} | |
| import asyncio as _asyncio | |
| async def asyncio_wrap(fn): | |
| return await _asyncio.to_thread(fn) | |
| graph = GraphStore() | |