Spaces:
Running
Running
File size: 2,261 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 | """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)
|