查询 (Query)

import heavenbase as hb


class Article(hb.Entity):
    title = hb.field(hb.ShortText)
    body = hb.field(hb.LongText)
    tags = hb.field(hb.Array[hb.ShortText]).default([])
    embedding = hb.field(hb.Vector[2])


ws = hb.HeavenBase("query-demo", preset="debug")
ws.register(Article)
ws.upsert(
    Article,
    {
        "object_id": "a1",
        "title": "Agent memory",
        "body": "Agents need searchable memory.",
        "tags": ["agent"],
        "embedding": [1.0, 0.0],
    },
)

base = ws.query(Article).where(Article.body.match("agent"))
nearest = base.near(Article.embedding, [1.0, 0.0], top_k=3)
rows = nearest.select("title", "score").execute().rows()

frame = nearest.select("title", "score").execute()

print(frame.rows())
print(frame.ids)
print(frame.column("title"))

filtered = (
    ws.query(Article)
    .where((Article.body.match("agent")) & (Article.tags.array_contains("agent")))
    .execute()
    .rows()
)

query = ws.query_json(
    Article,
    {
        "filter": {
            "body": {"$match": "agent"},
            "tags": {"$array_contains": "agent"},
        },
        "select": ["title"],
        "limit": 5,
    },
)

print(query.execute().rows())

total = ws.query(Article).where(Article.body.match("agent")).count().execute().scalar()

stats = (
    ws.query(Article)
    .group_by(Article.title)
    .aggregate(total=Article.object_id.count())
    .execute()
    .rows()
)

plan = ws.query(Article).where(Article.body.match("agent")).explain()
step = plan["steps"][0]

print(step["backend"], step["handler_mode"])

object_id = ws.upsert(Article, {"object_id": "draft", "title": "Draft", "body": "Edit me.", "embedding": [0.0, 1.0]})
ws.set(object_id, {"title": "Published"}, entity=Article)
ws.delete((Article, object_id))