Apiary Advisor

1"""Apiary Advisor: a small RAG + tool-calling agent built with Pixeltable.
2
3Hive-management guides are chunked into sentences, embedded with a custom batched UDF
4(feature hashing, no model download) and indexed with an EmbeddingIndex. A @pxt.query
5retrieves the best chunks; the Advisor table is an "agent table" whose computed columns
6retrieve, format context and draft an answer, and run.py turns that table into a reusable
7function with pxt.udf(table, return_value=...). Inspections use an expr_udf for mite load,
8and Questions call an LLM with tools (only when OPENAI_API_KEY is set).
9
10Layout: UDFs live in udfs.py; this file holds table models and @pxt.query functions.
11"""
12
13from __future__ import annotations
14
15import pixeltable as pxt
16import pixeltable.functions as pxtf
17from pixeltable.functions import openai
18
19from udfs import (
20    advisor_messages,
21    extractive_answer,
22    final_messages,
23    format_context,
24    hash_embed,
25    mite_load_pct,
26    varroa_action,
27    with_region,
28)
29
30TableModel = pxt.model_base()
31
32
33class Guides(TableModel, name='guides', has_default_idxs=False):
34    """One hive-management guide; source_uri may point at object storage or the web."""
35
36    guide_id = pxt.Column(type=pxt.String, primary_key=True)
37    title: pxt.String
38    topic: pxt.String
39    body: pxt.String
40    source_uri: pxt.String | None
41
42
43class GuideChunks(
44    TableModel,
45    name='guide_chunks',
46    base=Guides,
47    iterator=pxtf.string.string_splitter(Guides.body, separators='sentence'),
48):
49    """One row per sentence of a guide, searchable through a custom embedding index."""
50
51    n_chars = pxtf.string.len(text)  # type: ignore[name-defined]  # noqa: F821
52
53    __indexes__ = [
54        pxt.EmbeddingIndex(text, string_embed=hash_embed, metric='cosine', name='chunk_embed'),  # type: ignore[name-defined]  # noqa: F821
55    ]
56
57
58@pxt.query
59def search_guides(question: str, k: int = 3):
60    """Top-k guide sentences for a question (custom-embedding similarity)."""
61    sim = GuideChunks.text.similarity(string=question, idx='chunk_embed')
62    return (
63        GuideChunks.order_by(sim, asc=False)
64        .limit(k)
65        .select(GuideChunks.guide_id, GuideChunks.title, GuideChunks.text, score=sim)
66    )
67
68
69@pxt.query
70def guides_by_topic(topic: str):
71    """Guides for one topic, with a shareable link produced by functions.net.presigned_url."""
72    return (
73        Guides.where(Guides.topic == topic)
74        .select(Guides.guide_id, Guides.title, link=pxtf.net.presigned_url(Guides.source_uri, 3600))
75        .order_by(Guides.guide_id)
76    )
77
78
79class Advisor(TableModel, name='advisor', has_default_idxs=False):
80    """Agent table: question -> retrieval -> context -> offline extractive answer.
81
82    run.py wraps it as a function with pxt.udf(<advisor table>, return_value=...).
83    """
84
85    question: pxt.String
86    regional_question = with_region(question)
87    hits = search_guides(question, 3)
88    context = format_context(hits)
89    answer = extractive_answer(question, hits)
90
91
92class Inspections(TableModel, name='inspections', has_default_idxs=False):
93    """A hive inspection with a mite sample and a free-text note."""
94
95    inspection_id = pxt.Column(value=pxtf.uuid.uuid7(), primary_key=True)
96    hive_id: pxt.String
97    mites: pxt.Int
98    bees_sampled: pxt.Int
99    note: pxt.String
100
101    mite_pct = mite_load_pct(mites, bees_sampled)
102    action = varroa_action(mite_pct)
103
104    __indexes__ = [pxt.BtreeIndex(hive_id)]
105
106
107TOOLS = pxt.tools(varroa_action, search_guides)
108
109
110class Questions(TableModel, name='questions', has_default_idxs=False):
111    """Beekeeper questions answered by an LLM that may call varroa_action / search_guides."""
112
113    question: pxt.String
114    hits = search_guides(question, 3)
115    context = format_context(hits)
116    response = openai.chat_completions(
117        advisor_messages(question, context),
118        model='gpt-4o-mini',
119        model_kwargs={'max_tokens': 160, 'temperature': 0},
120        tools=TOOLS,
121    )
122    tool_output = openai.invoke_tools(TOOLS, response)
123    # Keep this model_kwargs literal different from the tool-calling one above: chat_completions
124    # adds `tools` to the dict it is given, and equal literals can end up shared between columns.
125    final = openai.chat_completions(
126        final_messages(question, context, tool_output),
127        model='gpt-4o-mini',
128        model_kwargs={'max_tokens': 200, 'temperature': 0},
129    )
130    answer = final.choices[0].message.content
131
132
133@pxt.query
134def hive_history(hive_id: str):
135    """Inspections for one hive, newest first (uses the hive_id btree index)."""
136    return (
137        Inspections.where(Inspections.hive_id == hive_id)
138        .select(Inspections.hive_id, Inspections.mites, Inspections.mite_pct, Inspections.action, Inspections.note)
139        .order_by(Inspections.inspection_id, asc=False)
140    )