Spaces:

sitcod3
/

abm-engine

Sleeping

Z User

feat: ABM Engine v1 — OpenSearch store + Apify enrichment + ICP scoring + k-NN memory + optional Neo4j mirror

fb0ec21 2 days ago

10.9 kB

"""FortyDN ABM Engine — outbound intelligence service.
Account graph + enrichment + ICP signals + vector memory, served as the
product backend for the ABM module. Components (all optional at boot):
  - Aiven OpenSearch : primary store, search, k-NN vector memory (required)
  - Apify            : enrichment actors (Google Maps / LinkedIn)
  - BetterDB / Aura  : optional Neo4j graph mirror
  - Upstash Redis    : optional distributed cache
  - Supabase JWKS    : user identity (same session JWT as the validator)
Endpoints
  GET  /health                 component report
  POST /accounts               upsert account
  GET  /accounts               search (?q=&industry=&country=&min_score=&limit=)
  GET  /accounts/{id}          account + contacts + events
  POST /enrich                 Apify enrichment run (google_maps | linkedin)
  POST /signals/score          rescore an account
  POST /memory                 remember a text document (embedded)
  POST /memory/batch           remember many
  POST /memory/search          k-NN recall
  GET  /graph/overview         Neo4j mirror status/counts
"""
import logging
from contextlib import asynccontextmanager
from typing import Any, Dict, List, Optional

from fastapi import FastAPI, HTTPException, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field

from .auth import authenticate
from .config import settings
from .services import accounts as accounts_svc
from .services import enrich as enrich_svc
from .services import memory as memory_svc
from .services import signals as signals_svc
from .stores import opensearch
from .stores.apify import apify_client as apify
from .stores.cache import cache
from .stores.embeddings import provider_name
from .stores.graph import graph

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("abm.main")


@asynccontextmanager
async def lifespan(app: FastAPI):
    logger.info("%s v%s starting", settings.app_name, settings.app_version)
    if settings.os_configured:
        try:
            results = await opensearch.store.bootstrap()
            logger.info("OpenSearch bootstrap: %s", results)
            graph_ok = await graph.connect()
            logger.info("graph mirror: %s", graph_ok)
        except Exception as exc:
            logger.warning("OpenSearch bootstrap failed (best-effort): %s", exc)
    else:
        logger.warning("OpenSearch not configured — set OS_URL/OS_USER/OS_PASSWORD")
    yield
    await opensearch.store.close()
    await graph.close()
    await apify.close()


app = FastAPI(title=settings.app_name, version=settings.app_version, lifespan=lifespan)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=False,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.middleware("http")
async def auth_gate(request: Request, call_next):
    """Paid-route gate: /health is open; everything under /api requires
    authentication. Enforcement off → anonymous demo access still works."""
    path = request.url.path
    if path == "/health" or path == "/" or not path.startswith("/api"):
        return await call_next(request)
    try:
        request.state.auth = await authenticate(request)
    except PermissionError as exc:
        message = str(exc)
        if message.startswith("402"):
            return JSONResponse(status_code=402, content={"detail": message.split(":", 1)[1]})
        return JSONResponse(status_code=401, content={"detail": message.split(":", 1)[-1]})
    return await call_next(request)


# All product routes live under /api/v1 (productApi.enginePath() convention).
api_prefix = "/api/v1"


class AccountIn(BaseModel):
    domain: Optional[str] = None
    name: Optional[str] = None
    industry: Optional[str] = None
    employee_count: Optional[int] = None
    hq_country: Optional[str] = None
    city: Optional[str] = None
    website: Optional[str] = None
    phone: Optional[str] = None
    address: Optional[str] = None
    description: Optional[str] = None
    source: Optional[str] = None


class EnrichIn(BaseModel):
    source: str = Field(default="google_maps", description="google_maps | linkedin")
    queries: List[str] = Field(..., min_length=1, max_length=20)
    limit: int = Field(default=20, ge=1, le=500)


class ScoreIn(BaseModel):
    account_id: str


class MemoryIn(BaseModel):
    text: str = Field(..., min_length=1, max_length=8000)
    metadata: Optional[Dict[str, Any]] = None


class MemoryBatchIn(BaseModel):
    texts: List[str] = Field(..., min_length=1, max_length=200)
    metadata: Optional[Dict[str, Any]] = None


class MemorySearchIn(BaseModel):
    query: str = Field(..., min_length=1)
    k: int = Field(default=8, ge=1, le=50)


# ── health ───────────────────────────────────────────────────────────────────

@app.get("/")
async def root():
    return {"app": settings.app_name, "version": settings.app_version, "docs": "/docs"}


@app.get("/api/health")
@app.get(f"{api_prefix}/health")
@app.get("/health")
async def health():
    os_ok = await opensearch.store.ping()
    apify_ok = await apify.ping()
    cache_ok = await cache.ping()
    components = {
        "opensearch": {"ok": os_ok, "configured": opensearch.store.configured},
        "apify": {"ok": apify_ok, "configured": apify.configured},
        "graph": {"mode": graph.mode},
        "cache": {"mode": cache.mode, "ok": cache_ok},
        "embeddings": {"provider": provider_name()},
        "enforcement": bool(settings.license_enforcement),
    }
    counts = {}
    if os_ok:
        try:
            counts = {
                "accounts": await opensearch.store.count("accounts"),
                "contacts": await opensearch.store.count("contacts"),
                "memory": await opensearch.store.count("memory"),
            }
        except Exception:
            pass
    core_ok = os_ok  # the store is the only hard dependency
    return {
        "status": "healthy" if core_ok else "degraded",
        "app": settings.app_name,
        "version": settings.app_version,
        "components": components,
        "counts": counts,
        "timestamp": __import__("time").time(),
    }


# ── accounts ────────────────────────────────────────────────────────────────

@app.post(f"{api_prefix}/accounts")
async def create_account(payload: AccountIn):
    try:
        doc = await accounts_svc.upsert_account(payload.model_dump(exclude_none=True))
    except ValueError as exc:
        raise HTTPException(400, str(exc))
    except Exception as exc:
        raise HTTPException(503, f"Datastore unavailable: {exc}")
    await graph.upsert_account(doc)
    return {"account": doc}


@app.get(f"{api_prefix}/accounts")
async def list_accounts(q: str = "", industry: str = "", country: str = "",
                        min_score: float = 0.0, limit: int = 25):
    try:
        hits = await accounts_svc.search_accounts(q, industry, country, min_score, limit)
    except Exception as exc:
        raise HTTPException(503, f"Datastore unavailable: {exc}")
    return {"total": len(hits), "accounts": hits}


@app.get(f"{api_prefix}/accounts/{{account_id}}")
async def read_account(account_id: str):
    result = await accounts_svc.get_account(account_id)
    if result is None:
        raise HTTPException(404, f"Account {account_id} not found")
    return result


# ── enrichment ──────────────────────────────────────────────────────────────

@app.post(f"{api_prefix}/enrich")
async def run_enrich(payload: EnrichIn):
    try:
        summary = await enrich_svc.enrich(payload.source, payload.queries, payload.limit)
    except ValueError as exc:
        raise HTTPException(400, str(exc))
    except RuntimeError as exc:
        raise HTTPException(502, str(exc))
    except Exception as exc:
        logger.exception("enrichment failed")
        raise HTTPException(500, f"Enrichment failed: {exc}")
    return summary


# ── signals ─────────────────────────────────────────────────────────────────

@app.post(f"{api_prefix}/signals/score")
async def score(payload: ScoreIn):
    account = await opensearch.store.get_doc("accounts", payload.account_id)
    if account is None:
        raise HTTPException(404, f"Account {payload.account_id} not found")
    try:
        res = await opensearch.store.client().post(
            f"/{settings.os_index_prefix}-contacts/_count",
            json={"query": {"term": {"account_id": payload.account_id}}},
        )
        contact_count = int(res.json().get("count", 0)) if res.status_code == 200 else 0
    except Exception:
        contact_count = 0
    scored = signals_svc.score_account(account, contact_count)
    doc = {**account, **scored}
    await opensearch.store.upsert_doc("accounts", doc, payload.account_id)
    await accounts_svc.record_event(payload.account_id, "rescored", scored)
    return scored


# ── vector memory ───────────────────────────────────────────────────────────

@app.post(f"{api_prefix}/memory")
async def memory_add(payload: MemoryIn):
    try:
        return await memory_svc.remember(payload.text, payload.metadata)
    except ValueError as exc:
        raise HTTPException(400, str(exc))
    except Exception as exc:
        raise HTTPException(503, f"Datastore unavailable: {exc}")


@app.post(f"{api_prefix}/memory/batch")
async def memory_add_batch(payload: MemoryBatchIn):
    try:
        return {"docs": await memory_svc.remember_batch(payload.texts, payload.metadata)}
    except Exception as exc:
        raise HTTPException(503, f"Datastore unavailable: {exc}")


@app.post(f"{api_prefix}/memory/search")
async def memory_search(payload: MemorySearchIn):
    try:
        hits = await memory_svc.recall(payload.query, payload.k)
    except Exception as exc:
        raise HTTPException(503, f"Datastore unavailable: {exc}")
    return {"total": len(hits), "results": hits}


# ── graph ───────────────────────────────────────────────────────────────────

@app.get(f"{api_prefix}/graph/overview")
async def graph_overview():
    return await graph.overview()

	"""FortyDN ABM Engine — outbound intelligence service.

	Account graph + enrichment + ICP signals + vector memory, served as the
	product backend for the ABM module. Components (all optional at boot):

	- Aiven OpenSearch : primary store, search, k-NN vector memory (required)
	- Apify : enrichment actors (Google Maps / LinkedIn)
	- BetterDB / Aura : optional Neo4j graph mirror
	- Upstash Redis : optional distributed cache
	- Supabase JWKS : user identity (same session JWT as the validator)

	Endpoints
	GET /health component report
	POST /accounts upsert account
	GET /accounts search (?q=&industry=&country=&min_score=&limit=)
	GET /accounts/{id} account + contacts + events
	POST /enrich Apify enrichment run (google_maps \| linkedin)
	POST /signals/score rescore an account
	POST /memory remember a text document (embedded)
	POST /memory/batch remember many
	POST /memory/search k-NN recall
	GET /graph/overview Neo4j mirror status/counts
	"""
	import logging
	from contextlib import asynccontextmanager
	from typing import Any, Dict, List, Optional

	from fastapi import FastAPI, HTTPException, Request
	from fastapi.middleware.cors import CORSMiddleware
	from fastapi.responses import JSONResponse
	from pydantic import BaseModel, Field

	from .auth import authenticate
	from .config import settings
	from .services import accounts as accounts_svc
	from .services import enrich as enrich_svc
	from .services import memory as memory_svc
	from .services import signals as signals_svc
	from .stores import opensearch
	from .stores.apify import apify_client as apify
	from .stores.cache import cache
	from .stores.embeddings import provider_name
	from .stores.graph import graph

	logging.basicConfig(level=logging.INFO)
	logger = logging.getLogger("abm.main")


	@asynccontextmanager
	async def lifespan(app: FastAPI):
	logger.info("%s v%s starting", settings.app_name, settings.app_version)
	if settings.os_configured:
	try:
	results = await opensearch.store.bootstrap()
	logger.info("OpenSearch bootstrap: %s", results)
	graph_ok = await graph.connect()
	logger.info("graph mirror: %s", graph_ok)
	except Exception as exc:
	logger.warning("OpenSearch bootstrap failed (best-effort): %s", exc)
	else:
	logger.warning("OpenSearch not configured — set OS_URL/OS_USER/OS_PASSWORD")
	yield
	await opensearch.store.close()
	await graph.close()
	await apify.close()


	app = FastAPI(title=settings.app_name, version=settings.app_version, lifespan=lifespan)

	app.add_middleware(
	CORSMiddleware,
	allow_origins=["*"],
	allow_credentials=False,
	allow_methods=["*"],
	allow_headers=["*"],
	)


	@app.middleware("http")
	async def auth_gate(request: Request, call_next):
	"""Paid-route gate: /health is open; everything under /api requires
	authentication. Enforcement off → anonymous demo access still works."""
	path = request.url.path
	if path == "/health" or path == "/" or not path.startswith("/api"):
	return await call_next(request)
	try:
	request.state.auth = await authenticate(request)
	except PermissionError as exc:
	message = str(exc)
	if message.startswith("402"):
	return JSONResponse(status_code=402, content={"detail": message.split(":", 1)[1]})
	return JSONResponse(status_code=401, content={"detail": message.split(":", 1)[-1]})
	return await call_next(request)


	# All product routes live under /api/v1 (productApi.enginePath() convention).
	api_prefix = "/api/v1"


	class AccountIn(BaseModel):
	domain: Optional[str] = None
	name: Optional[str] = None
	industry: Optional[str] = None
	employee_count: Optional[int] = None
	hq_country: Optional[str] = None
	city: Optional[str] = None
	website: Optional[str] = None
	phone: Optional[str] = None
	address: Optional[str] = None
	description: Optional[str] = None
	source: Optional[str] = None


	class EnrichIn(BaseModel):
	source: str = Field(default="google_maps", description="google_maps \| linkedin")
	queries: List[str] = Field(..., min_length=1, max_length=20)
	limit: int = Field(default=20, ge=1, le=500)


	class ScoreIn(BaseModel):
	account_id: str


	class MemoryIn(BaseModel):
	text: str = Field(..., min_length=1, max_length=8000)
	metadata: Optional[Dict[str, Any]] = None


	class MemoryBatchIn(BaseModel):
	texts: List[str] = Field(..., min_length=1, max_length=200)
	metadata: Optional[Dict[str, Any]] = None


	class MemorySearchIn(BaseModel):
	query: str = Field(..., min_length=1)
	k: int = Field(default=8, ge=1, le=50)


	# ── health ───────────────────────────────────────────────────────────────────

	@app.get("/")
	async def root():
	return {"app": settings.app_name, "version": settings.app_version, "docs": "/docs"}


	@app.get("/api/health")
	@app.get(f"{api_prefix}/health")
	@app.get("/health")
	async def health():
	os_ok = await opensearch.store.ping()
	apify_ok = await apify.ping()
	cache_ok = await cache.ping()
	components = {
	"opensearch": {"ok": os_ok, "configured": opensearch.store.configured},
	"apify": {"ok": apify_ok, "configured": apify.configured},
	"graph": {"mode": graph.mode},
	"cache": {"mode": cache.mode, "ok": cache_ok},
	"embeddings": {"provider": provider_name()},
	"enforcement": bool(settings.license_enforcement),
	}
	counts = {}
	if os_ok:
	try:
	counts = {
	"accounts": await opensearch.store.count("accounts"),
	"contacts": await opensearch.store.count("contacts"),
	"memory": await opensearch.store.count("memory"),
	}
	except Exception:
	pass
	core_ok = os_ok # the store is the only hard dependency
	return {
	"status": "healthy" if core_ok else "degraded",
	"app": settings.app_name,
	"version": settings.app_version,
	"components": components,
	"counts": counts,
	"timestamp": __import__("time").time(),
	}


	# ── accounts ────────────────────────────────────────────────────────────────

	@app.post(f"{api_prefix}/accounts")
	async def create_account(payload: AccountIn):
	try:
	doc = await accounts_svc.upsert_account(payload.model_dump(exclude_none=True))
	except ValueError as exc:
	raise HTTPException(400, str(exc))
	except Exception as exc:
	raise HTTPException(503, f"Datastore unavailable: {exc}")
	await graph.upsert_account(doc)
	return {"account": doc}


	@app.get(f"{api_prefix}/accounts")
	async def list_accounts(q: str = "", industry: str = "", country: str = "",
	min_score: float = 0.0, limit: int = 25):
	try:
	hits = await accounts_svc.search_accounts(q, industry, country, min_score, limit)
	except Exception as exc:
	raise HTTPException(503, f"Datastore unavailable: {exc}")
	return {"total": len(hits), "accounts": hits}


	@app.get(f"{api_prefix}/accounts/{{account_id}}")
	async def read_account(account_id: str):
	result = await accounts_svc.get_account(account_id)
	if result is None:
	raise HTTPException(404, f"Account {account_id} not found")
	return result


	# ── enrichment ──────────────────────────────────────────────────────────────

	@app.post(f"{api_prefix}/enrich")
	async def run_enrich(payload: EnrichIn):
	try:
	summary = await enrich_svc.enrich(payload.source, payload.queries, payload.limit)
	except ValueError as exc:
	raise HTTPException(400, str(exc))
	except RuntimeError as exc:
	raise HTTPException(502, str(exc))
	except Exception as exc:
	logger.exception("enrichment failed")
	raise HTTPException(500, f"Enrichment failed: {exc}")
	return summary


	# ── signals ─────────────────────────────────────────────────────────────────

	@app.post(f"{api_prefix}/signals/score")
	async def score(payload: ScoreIn):
	account = await opensearch.store.get_doc("accounts", payload.account_id)
	if account is None:
	raise HTTPException(404, f"Account {payload.account_id} not found")
	try:
	res = await opensearch.store.client().post(
	f"/{settings.os_index_prefix}-contacts/_count",
	json={"query": {"term": {"account_id": payload.account_id}}},
	)
	contact_count = int(res.json().get("count", 0)) if res.status_code == 200 else 0
	except Exception:
	contact_count = 0
	scored = signals_svc.score_account(account, contact_count)
	doc = {account, scored}
	await opensearch.store.upsert_doc("accounts", doc, payload.account_id)
	await accounts_svc.record_event(payload.account_id, "rescored", scored)
	return scored


	# ── vector memory ───────────────────────────────────────────────────────────

	@app.post(f"{api_prefix}/memory")
	async def memory_add(payload: MemoryIn):
	try:
	return await memory_svc.remember(payload.text, payload.metadata)
	except ValueError as exc:
	raise HTTPException(400, str(exc))
	except Exception as exc:
	raise HTTPException(503, f"Datastore unavailable: {exc}")


	@app.post(f"{api_prefix}/memory/batch")
	async def memory_add_batch(payload: MemoryBatchIn):
	try:
	return {"docs": await memory_svc.remember_batch(payload.texts, payload.metadata)}
	except Exception as exc:
	raise HTTPException(503, f"Datastore unavailable: {exc}")


	@app.post(f"{api_prefix}/memory/search")
	async def memory_search(payload: MemorySearchIn):
	try:
	hits = await memory_svc.recall(payload.query, payload.k)
	except Exception as exc:
	raise HTTPException(503, f"Datastore unavailable: {exc}")
	return {"total": len(hits), "results": hits}


	# ── graph ───────────────────────────────────────────────────────────────────

	@app.get(f"{api_prefix}/graph/overview")
	async def graph_overview():
	return await graph.overview()