Spaces:
Running
Running
File size: 7,450 Bytes
fb0ec21 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 | """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
@property
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()
|