Spaces:
Running
Running
Z User
feat: ABM Engine v1 — OpenSearch store + Apify enrichment + ICP scoring + k-NN memory + optional Neo4j mirror
fb0ec21 | """Vector memory over the OpenSearch k-NN index.""" | |
| import json | |
| import uuid | |
| from datetime import datetime, timezone | |
| from typing import Any, Dict, List, Optional | |
| from ..stores import opensearch | |
| from ..stores.embeddings import embed, embed_batch, provider_name | |
| def _now() -> str: | |
| return datetime.now(timezone.utc).isoformat() | |
| def _memory_index() -> str: | |
| return f"{opensearch.settings.os_index_prefix}-memory" | |
| async def remember(text: str, metadata: Optional[dict] = None) -> dict: | |
| text = (text or "").strip() | |
| if not text: | |
| raise ValueError("text must not be empty") | |
| doc_id = str(uuid.uuid4()) | |
| doc = { | |
| "doc_id": doc_id, | |
| "text": text[:8000], | |
| "metadata": metadata or {}, | |
| "embedding": embed(text), | |
| "created_at": _now(), | |
| } | |
| await opensearch.store.index_doc("memory", doc) | |
| return {"doc_id": doc_id, "provider": provider_name()} | |
| async def remember_batch(texts: List[str], metadata: Optional[dict] = None) -> List[dict]: | |
| texts = [t.strip()[:8000] for t in texts if t and t.strip()] | |
| if not texts: | |
| return [] | |
| vectors = embed_batch(texts) | |
| docs = [] | |
| for text, vector in zip(texts, vectors): | |
| docs.append({ | |
| "doc_id": str(uuid.uuid4()), | |
| "text": text, | |
| "metadata": metadata or {}, | |
| "embedding": vector, | |
| "created_at": _now(), | |
| }) | |
| lines: List[str] = [] | |
| for d in docs: | |
| lines.append(json.dumps({"index": {"_index": _memory_index(), "_id": d["doc_id"]}})) | |
| lines.append(json.dumps(d)) | |
| res = await opensearch.store.client().post( | |
| "/_bulk", | |
| content=("\n".join(lines) + "\n").encode(), | |
| headers={"Content-Type": "application/x-ndjson"}, | |
| ) | |
| res.raise_for_status() | |
| await opensearch.store.refresh("memory") | |
| return [{"doc_id": d["doc_id"]} for d in docs] | |
| async def recall(query: str, k: int = 8) -> List[dict]: | |
| query = (query or "").strip() | |
| if not query: | |
| return [] | |
| k = max(1, min(k, 50)) | |
| body = { | |
| "size": k, | |
| "query": {"knn": {"embedding": {"vector": embed(query), "k": k}}}, | |
| "_source": ["text", "metadata", "created_at"], | |
| } | |
| return await opensearch.store.search("memory", body) | |