KB injection tuning: pgvector substrate + retrieval telemetry + title-first auto-inject #74

Merged
bvandeusen merged 4 commits from dev into main 2026-06-22 20:56:15 -04:00
22 changed files with 1037 additions and 32 deletions
+4 -2
View File
@@ -165,7 +165,9 @@ jobs:
SECRET_KEY: ci_integration_placeholder SECRET_KEY: ci_integration_placeholder
services: services:
postgres: postgres:
image: postgres:16-alpine # pgvector image so `alembic upgrade head` can run migration 0067
# (CREATE EXTENSION vector). PG17 — matches the prod/quickstart image.
image: pgvector/pgvector:pg17
env: env:
POSTGRES_USER: scribe POSTGRES_USER: scribe
POSTGRES_PASSWORD: ci_integration POSTGRES_PASSWORD: ci_integration
@@ -189,7 +191,7 @@ jobs:
set -eux set -eux
echo "=== container landscape (diagnostic for the name filter) ===" echo "=== container landscape (diagnostic for the name filter) ==="
docker ps -a --format '{{.ID}} {{.Image}} -> {{.Names}}' docker ps -a --format '{{.ID}} {{.Image}} -> {{.Names}}'
PG=$(docker ps --filter "name=integration" --filter "ancestor=postgres:16-alpine" -q | head -n1) PG=$(docker ps --filter "name=integration" --filter "ancestor=pgvector/pgvector:pg17" -q | head -n1)
test -n "$PG" test -n "$PG"
PG_IP=$(docker inspect -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$PG") PG_IP=$(docker inspect -f '{{range .NetworkSettings.Networks}}{{.IPAddress}}{{end}}' "$PG")
test -n "$PG_IP" test -n "$PG_IP"
@@ -0,0 +1,73 @@
"""pgvector: note_embeddings.embedding JSONB -> vector(384) + HNSW index
Revision ID: 0067
Revises: 0066
Create Date: 2026-06-22
Moves semantic search off the full-table Python cosine scan onto a native
pgvector column so ranking + top-k run as an indexed `ORDER BY embedding <=> :q
LIMIT k` in Postgres (see services/embeddings.semantic_search_notes).
Requires a Postgres image that bundles the `vector` extension — the stack moved
from postgres:16-alpine to pgvector/pgvector:pg16 in the same change (compose +
CI). `CREATE EXTENSION IF NOT EXISTS vector` below is the in-db half.
Embeddings are DERIVED data (regenerated from note text by
backfill_note_embeddings at startup), so this migration is free to drop any row
it can't cleanly convert: only rows whose stored JSONB array is exactly 384-dim
are carried over (guarding against stale vectors from an earlier model — the
same mixed-dim hazard _cosine_similarity defended against). Dropped rows are
re-embedded on next boot.
"""
from alembic import op
revision = "0067"
down_revision = "0066"
branch_labels = None
depends_on = None
def upgrade() -> None:
op.execute("CREATE EXTENSION IF NOT EXISTS vector")
# New native-vector column, populated only from cleanly-convertible rows.
# A JSONB array like [0.1, 0.2, ...] renders to text that is exactly
# pgvector's input literal, so (embedding::text)::vector is a direct cast.
op.execute("ALTER TABLE note_embeddings ADD COLUMN embedding_vec vector(384)")
op.execute(
"""
UPDATE note_embeddings
SET embedding_vec = (embedding::text)::vector
WHERE jsonb_array_length(embedding) = 384
"""
)
# Stale-dim rows (couldn't convert) are derived data — drop and let the
# startup backfill regenerate them at the current dimension.
op.execute("DELETE FROM note_embeddings WHERE embedding_vec IS NULL")
op.execute("ALTER TABLE note_embeddings ALTER COLUMN embedding_vec SET NOT NULL")
op.execute("ALTER TABLE note_embeddings DROP COLUMN embedding")
op.execute("ALTER TABLE note_embeddings RENAME COLUMN embedding_vec TO embedding")
# HNSW index for cosine distance — matches Vector.cosine_distance (`<=>`).
op.execute(
"""
CREATE INDEX ix_note_embeddings_embedding_hnsw
ON note_embeddings
USING hnsw (embedding vector_cosine_ops)
"""
)
def downgrade() -> None:
# Back to JSONB. pgvector renders a vector to a text literal that is a valid
# JSON array, so the reverse cast is symmetric. The `vector` extension is
# intentionally left installed (other objects may depend on it; dropping an
# extension is the riskier, rarely-wanted direction).
op.execute("DROP INDEX IF EXISTS ix_note_embeddings_embedding_hnsw")
op.execute("ALTER TABLE note_embeddings ADD COLUMN embedding_json jsonb")
op.execute("UPDATE note_embeddings SET embedding_json = (embedding::text)::jsonb")
op.execute("ALTER TABLE note_embeddings ALTER COLUMN embedding_json SET NOT NULL")
op.execute("ALTER TABLE note_embeddings DROP COLUMN embedding")
op.execute("ALTER TABLE note_embeddings RENAME COLUMN embedding_json TO embedding")
+60
View File
@@ -0,0 +1,60 @@
"""retrieval_logs: per-call semantic-retrieval telemetry for KB-injection tuning
Revision ID: 0068
Revises: 0067
Create Date: 2026-06-22
One row per semantic-retrieval call (MCP search tool, REST search route, and —
once it lands — the title-first auto-inject path). Captures the effective query
params and the score distribution of the results so the similarity threshold
and top-k can be tuned from real usage. FK-free on user_id (mirrors app_logs):
telemetry should outlive the row it describes.
"""
from alembic import op
import sqlalchemy as sa
from sqlalchemy.dialects.postgresql import JSONB
revision = "0068"
down_revision = "0067"
branch_labels = None
depends_on = None
def upgrade() -> None:
op.create_table(
"retrieval_logs",
sa.Column("id", sa.Integer(), primary_key=True),
sa.Column(
"created_at", sa.DateTime(timezone=True), nullable=False,
server_default=sa.text("now()"),
),
sa.Column("user_id", sa.Integer(), nullable=True),
sa.Column("source", sa.Text(), nullable=False),
sa.Column("query", sa.Text(), nullable=True),
sa.Column("threshold", sa.Float(), nullable=True),
sa.Column("limit_n", sa.Integer(), nullable=True),
sa.Column("project_id", sa.Integer(), nullable=True),
sa.Column("is_task", sa.Boolean(), nullable=True),
sa.Column("result_count", sa.Integer(), nullable=False, server_default="0"),
sa.Column("top_score", sa.Float(), nullable=True),
sa.Column("min_score", sa.Float(), nullable=True),
sa.Column("result_ids", JSONB(), nullable=False, server_default=sa.text("'[]'::jsonb")),
sa.Column("duration_ms", sa.Float(), nullable=True),
)
op.create_index("ix_retrieval_logs_created_at", "retrieval_logs", ["created_at"])
op.create_index("ix_retrieval_logs_user_id", "retrieval_logs", ["user_id"])
op.create_index("ix_retrieval_logs_source", "retrieval_logs", ["source"])
op.create_index(
"ix_retrieval_logs_source_created_at",
"retrieval_logs",
["source", sa.text("created_at DESC")],
)
def downgrade() -> None:
op.drop_index("ix_retrieval_logs_source_created_at", table_name="retrieval_logs")
op.drop_index("ix_retrieval_logs_source", table_name="retrieval_logs")
op.drop_index("ix_retrieval_logs_user_id", table_name="retrieval_logs")
op.drop_index("ix_retrieval_logs_created_at", table_name="retrieval_logs")
op.drop_table("retrieval_logs")
+5 -1
View File
@@ -21,7 +21,11 @@ services:
max_attempts: 5 max_attempts: 5
db: db:
image: postgres:16-alpine # pgvector image (Debian/glibc, PG17) — bundles the `vector` extension that
# migration 0067 enables. Moved off postgres:16-alpine via logical
# dump/restore (which doubles as the PG16->PG17 major upgrade); see the
# TRANSITION runbook in the PR.
image: pgvector/pgvector:pg17
stop_grace_period: 120s stop_grace_period: 120s
volumes: volumes:
- pgdata:/var/lib/postgresql/data - pgdata:/var/lib/postgresql/data
+2 -1
View File
@@ -35,7 +35,8 @@ services:
start_period: 30s start_period: 30s
db: db:
image: postgres:16-alpine # pgvector image (PG17) — bundles the `vector` extension (migration 0067).
image: pgvector/pgvector:pg17
stop_grace_period: 120s stop_grace_period: 120s
volumes: volumes:
- pgdata:/var/lib/postgresql/data - pgdata:/var/lib/postgresql/data
+88
View File
@@ -17,6 +17,13 @@ const timezoneSaved = ref(false);
const trashRetentionDays = ref("90"); const trashRetentionDays = ref("90");
const savingRetention = ref(false); const savingRetention = ref(false);
const retentionSaved = ref(false); const retentionSaved = ref(false);
// Knowledge auto-inject (per-user). Defaults mirror the backend
// (services/plugin_context: enabled, threshold 0.55, top-k 3).
const kbInjectEnabled = ref(true);
const kbInjectThreshold = ref("0.55");
const kbInjectTopK = ref("3");
const savingKbInject = ref(false);
const kbInjectSaved = ref(false);
// think_enabled setting removed 2026-05-23. The chat+curator architecture // think_enabled setting removed 2026-05-23. The chat+curator architecture
// has tools=[] on the chat model; think on a no-tools conversational pass // has tools=[] on the chat model; think on a no-tools conversational pass
@@ -56,6 +63,28 @@ async function saveRetention() {
savingRetention.value = false; savingRetention.value = false;
} }
} }
async function saveKbInject() {
const t = Math.min(1, Math.max(0, Number(kbInjectThreshold.value) || 0));
const k = Math.min(10, Math.max(1, Math.floor(Number(kbInjectTopK.value) || 1)));
kbInjectThreshold.value = String(t);
kbInjectTopK.value = String(k);
savingKbInject.value = true;
kbInjectSaved.value = false;
try {
await apiPut('/api/settings', {
kb_autoinject_enabled: kbInjectEnabled.value ? 'true' : 'false',
kb_autoinject_threshold: String(t),
kb_autoinject_top_k: String(k),
});
kbInjectSaved.value = true;
setTimeout(() => (kbInjectSaved.value = false), 2000);
} catch {
toastStore.show('Failed to save auto-inject settings', 'error');
} finally {
savingKbInject.value = false;
}
}
const newEmail = ref(""); const newEmail = ref("");
const emailPassword = ref(""); const emailPassword = ref("");
const changingEmail = ref(false); const changingEmail = ref(false);
@@ -435,6 +464,13 @@ onMounted(async () => {
const allSettings = await apiGet<Record<string, string>>("/api/settings"); const allSettings = await apiGet<Record<string, string>>("/api/settings");
userTimezone.value = allSettings.user_timezone ?? ""; userTimezone.value = allSettings.user_timezone ?? "";
trashRetentionDays.value = allSettings.trash_retention_days ?? "90"; trashRetentionDays.value = allSettings.trash_retention_days ?? "90";
kbInjectEnabled.value = allSettings.kb_autoinject_enabled !== "false";
if (allSettings.kb_autoinject_threshold !== undefined) {
kbInjectThreshold.value = allSettings.kb_autoinject_threshold;
}
if (allSettings.kb_autoinject_top_k !== undefined) {
kbInjectTopK.value = allSettings.kb_autoinject_top_k;
}
if (allSettings.notify_task_reminders !== undefined) { if (allSettings.notify_task_reminders !== undefined) {
notifyTaskReminders.value = allSettings.notify_task_reminders !== "false"; notifyTaskReminders.value = allSettings.notify_task_reminders !== "false";
} }
@@ -1165,6 +1201,58 @@ function formatUserDate(iso: string): string {
</div> </div>
</section> </section>
<section class="settings-section full-width">
<h2>Knowledge auto-inject</h2>
<p class="section-desc">
When enabled, the Scribe plugin quietly surfaces the titles of your most
relevant notes on each prompt — never their full text — so Claude can pull
one in with <code>get_note(id)</code> only when it helps. Titles only, each
note at most once per session, and nothing is shown unless it clears the
confidence bar below.
</p>
<div class="checkbox-field">
<label>
<input type="checkbox" v-model="kbInjectEnabled" />
Surface relevant note titles each prompt
</label>
<p class="field-hint">Off = notes reach context only when Claude searches for them.</p>
</div>
<div class="field">
<label for="kb-inject-threshold">Confidence threshold (01)</label>
<input
id="kb-inject-threshold"
v-model="kbInjectThreshold"
type="number"
min="0"
max="1"
step="0.05"
class="input"
style="max-width: 8rem"
/>
<p class="field-hint">Minimum similarity to surface a note. Higher = stricter (fewer, more certain). Deliberately above the 0.45 used for searches you trigger yourself.</p>
</div>
<div class="field">
<label for="kb-inject-topk">Max notes per prompt</label>
<input
id="kb-inject-topk"
v-model="kbInjectTopK"
type="number"
min="1"
max="10"
step="1"
class="input"
style="max-width: 8rem"
/>
<p class="field-hint">Ceiling on titles surfaced at once (110).</p>
</div>
<div class="actions">
<button class="btn-save" @click="saveKbInject" :disabled="savingKbInject">
{{ savingKbInject ? 'Saving' : 'Save' }}
</button>
<span v-if="kbInjectSaved" class="saved-msg">Saved!</span>
</div>
</section>
</div> </div>
<!-- ── Account ── --> <!-- ── Account ── -->
+10
View File
@@ -13,6 +13,16 @@
} }
] ]
} }
],
"UserPromptSubmit": [
{
"hooks": [
{
"type": "command",
"command": "bash \"${CLAUDE_PLUGIN_ROOT}/hooks/scribe_autoinject.sh\""
}
]
}
] ]
} }
} }
+86
View File
@@ -0,0 +1,86 @@
#!/usr/bin/env bash
# Scribe plugin — UserPromptSubmit push channel (knowledge auto-inject, Path A).
#
# On each user prompt, asks the operator's Scribe instance for a TITLE-FIRST
# awareness hint: the few notes that clear the per-user auto-inject gates
# (high-confidence threshold, margin gate, session dedup, top-k). Titles + ids
# only — never bodies; the agent calls get_note(id) to pull anything it judges
# relevant. Most turns inject nothing.
#
# Best-effort enrichment ONLY: unlike the SessionStart channel there is no
# static floor here. If the instance is unconfigured/unreachable, or anything
# fails, the hook stays SILENT and exits 0 — it must never block a prompt.
#
# Config (same as scribe_session_context.sh), exported to the hook by Claude Code:
# CLAUDE_PLUGIN_OPTION_api_endpoint base URL, no trailing slash
# CLAUDE_PLUGIN_OPTION_api_token fmcp_ API key (sensitive)
# SCRIBE_URL / SCRIBE_TOKEN override for the settings.json dogfooding path.
#
# Session dedup: each surfaced note id is remembered in a per-session file so a
# note is injected at most once per session. Passed back as exclude_ids.
set -uo pipefail
command -v jq >/dev/null 2>&1 || exit 0
command -v curl >/dev/null 2>&1 || exit 0
# UserPromptSubmit delivers a JSON event on stdin: { prompt, session_id, cwd, ... }
event=$(cat 2>/dev/null || true)
prompt=$(printf '%s' "$event" | jq -r '.prompt // empty' 2>/dev/null) || prompt=""
session_id=$(printf '%s' "$event" | jq -r '.session_id // empty' 2>/dev/null) || session_id=""
event_cwd=$(printf '%s' "$event" | jq -r '.cwd // empty' 2>/dev/null) || event_cwd=""
# Nothing to retrieve against.
[ -n "$prompt" ] || exit 0
url=${SCRIBE_URL:-${CLAUDE_PLUGIN_OPTION_api_endpoint:-}}
token=${SCRIBE_TOKEN:-${CLAUDE_PLUGIN_OPTION_api_token:-}}
# Guard against an unexpanded ${...} placeholder arriving as a literal.
case "$url" in *'${'*) url="" ;; esac
case "$token" in *'${'*) token="" ;; esac
# Unconfigured install → silent (auto-inject is pure enrichment).
[ -n "$url" ] && [ -n "$token" ] || exit 0
# Cap the query length — a giant prompt makes a giant URL for no extra signal.
q=$(printf '%s' "$prompt" | cut -c1-2000)
q_enc=$(printf '%s' "$q" | jq -rR '@uri' 2>/dev/null) || exit 0
# Resolve the working repo's remote so the server can scope to the bound project.
repo_dir=${event_cwd:-${CLAUDE_PROJECT_DIR:-$PWD}}
repo=$(git -C "$repo_dir" remote get-url origin 2>/dev/null || true)
repo_q=""
if [ -n "$repo" ]; then
enc=$(printf '%s' "$repo" | jq -rR '@uri' 2>/dev/null) || enc=""
[ -n "$enc" ] && repo_q="&repo=${enc}"
fi
# Per-session dedup: ids already injected this session are skipped.
state_dir="${TMPDIR:-/tmp}/scribe-autoinject"
mkdir -p "$state_dir" 2>/dev/null || true
idfile=""
exclude_q=""
if [ -n "$session_id" ]; then
# session_id is an opaque token from Claude Code; keep only filename-safe chars.
safe_sid=$(printf '%s' "$session_id" | tr -c 'A-Za-z0-9._-' '_')
idfile="$state_dir/${safe_sid}.ids"
if [ -f "$idfile" ]; then
seen=$(tr '\n' ',' < "$idfile" 2>/dev/null | sed 's/,$//')
[ -n "$seen" ] && exclude_q="&exclude_ids=${seen}"
fi
fi
body=$(curl -fsS --max-time 5 \
-H "Authorization: Bearer ${token}" \
"${url%/}/api/plugin/retrieve?q=${q_enc}${repo_q}${exclude_q}" 2>/dev/null) || exit 0
[ -n "$body" ] || exit 0
context=$(printf '%s' "$body" | jq -r '.context // empty' 2>/dev/null) || exit 0
[ -n "$context" ] || exit 0
# Remember the surfaced ids so they aren't injected again this session.
if [ -n "$idfile" ]; then
printf '%s' "$body" | jq -r '.note_ids[]? // empty' 2>/dev/null >> "$idfile" || true
fi
jq -n --arg c "$context" \
'{hookSpecificOutput: {hookEventName: "UserPromptSubmit", additionalContext: $c}}'
exit 0
+9 -4
View File
@@ -12,10 +12,15 @@ for the operator's work, and as your own working memory across sessions.
recent notes in one shot. recent notes in one shot.
**While you work:** **While you work:**
- **Recall before acting** — `search` Scribe for related prior work before - **Recall before acting** — before you answer anything about the operator's
answering a question about the operator's work, starting a task, or work or start a task, `search` Scribe first; assume a related note, task, or
re-deriving a decision. Assume a related note, task, or decision already decision already exists. Concretely, reach for recall whenever a request
exists. touches the operator's projects, people, places, prior decisions, or existing
work: check for an existing task before opening a new one, and for a prior
note/decision before re-deriving one. When a project is in scope (you entered
one), pass its id to `search` so results stay scoped to it. Treating Scribe as
the first place you look — not just somewhere you write — is what makes it a
trustworthy record.
- **Record as you go** — track work as Scribe tasks and log progress with - **Record as you go** — track work as Scribe tasks and log progress with
`add_task_log`. Always log when you **complete a task** and when you **hit or `add_task_log`. Always log when you **complete a task** and when you **hit or
discover a problem** — so changes of direction are captured, not just discover a problem** — so changes of direction are captured, not just
+1
View File
@@ -21,6 +21,7 @@ dependencies = [
"APScheduler>=3.10,<4.0", "APScheduler>=3.10,<4.0",
"mcp[cli]>=1.0", "mcp[cli]>=1.0",
"fastembed>=0.4", "fastembed>=0.4",
"pgvector>=0.3",
] ]
[project.optional-dependencies] [project.optional-dependencies]
+11 -1
View File
@@ -7,8 +7,11 @@ working. Differences from fable-mcp:
""" """
from __future__ import annotations from __future__ import annotations
import time
from scribe.mcp._context import current_user_id from scribe.mcp._context import current_user_id
from scribe.services.embeddings import semantic_search_notes from scribe.services.embeddings import DEFAULT_SIMILARITY_THRESHOLD, semantic_search_notes
from scribe.services.retrieval_telemetry import record_retrieval
async def search( async def search(
@@ -43,10 +46,17 @@ async def search(
uid = current_user_id() uid = current_user_id()
limit = max(1, min(limit, 50)) limit = max(1, min(limit, 50))
is_task = {"note": False, "task": True}.get(content_type) # None => any is_task = {"note": False, "task": True}.get(content_type) # None => any
t0 = time.perf_counter()
raw = await semantic_search_notes( raw = await semantic_search_notes(
uid, q, limit=limit, is_task=is_task, uid, q, limit=limit, is_task=is_task,
project_id=project_id or None, project_id=project_id or None,
) )
record_retrieval(
user_id=uid, source="mcp_search", query=q,
threshold=DEFAULT_SIMILARITY_THRESHOLD, limit=limit,
project_id=project_id or None, is_task=is_task, results=raw,
duration_ms=(time.perf_counter() - t0) * 1000.0,
)
return { return {
"results": [ "results": [
{ {
+1
View File
@@ -26,6 +26,7 @@ from scribe.models.app_log import AppLog # noqa: E402, F401
from scribe.models.password_reset import PasswordResetToken # noqa: E402, F401 from scribe.models.password_reset import PasswordResetToken # noqa: E402, F401
from scribe.models.invitation import InvitationToken # noqa: E402, F401 from scribe.models.invitation import InvitationToken # noqa: E402, F401
from scribe.models.embedding import NoteEmbedding # noqa: E402, F401 from scribe.models.embedding import NoteEmbedding # noqa: E402, F401
from scribe.models.retrieval_log import RetrievalLog # noqa: E402, F401
from scribe.models.project import Project # noqa: E402, F401 from scribe.models.project import Project # noqa: E402, F401
from scribe.models.event import Event # noqa: E402, F401 from scribe.models.event import Event # noqa: E402, F401
from scribe.models.milestone import Milestone # noqa: E402, F401 from scribe.models.milestone import Milestone # noqa: E402, F401
+8 -2
View File
@@ -1,11 +1,17 @@
from datetime import datetime, timezone from datetime import datetime, timezone
from pgvector.sqlalchemy import Vector
from sqlalchemy import DateTime, ForeignKey, Integer from sqlalchemy import DateTime, ForeignKey, Integer
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column from sqlalchemy.orm import Mapped, mapped_column
from scribe.models import Base from scribe.models import Base
# bge-small-en-v1.5 produces 384-dim unit-normalized vectors. The column is a
# native pgvector `vector(384)` (see migration 0067) so similarity search runs
# as an indexed `ORDER BY embedding <=> :q LIMIT k` in Postgres rather than a
# full-table Python cosine scan.
EMBEDDING_DIM = 384
class NoteEmbedding(Base): class NoteEmbedding(Base):
"""Stores the embedding vector for a note, used for semantic search.""" """Stores the embedding vector for a note, used for semantic search."""
@@ -18,7 +24,7 @@ class NoteEmbedding(Base):
primary_key=True, primary_key=True,
) )
user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True) user_id: Mapped[int] = mapped_column(Integer, nullable=False, index=True)
embedding: Mapped[list] = mapped_column(JSONB, nullable=False) embedding: Mapped[list] = mapped_column(Vector(EMBEDDING_DIM), nullable=False)
updated_at: Mapped[datetime] = mapped_column( updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), DateTime(timezone=True),
default=lambda: datetime.now(timezone.utc), default=lambda: datetime.now(timezone.utc),
+70
View File
@@ -0,0 +1,70 @@
from datetime import datetime, timezone
from sqlalchemy import Boolean, DateTime, Float, Index, Integer, Text
from sqlalchemy.dialects.postgresql import JSONB
from sqlalchemy.orm import Mapped, mapped_column
from scribe.models import Base
class RetrievalLog(Base):
"""One row per semantic-retrieval call, for KB-injection tuning.
Captures what a query asked for, what came back, and the score
distribution of the results — the empirical basis for tuning the
similarity threshold and top-k per surface. `result_ids` holds the ranked
hits (id + score + rank) so a later pass can correlate "what we surfaced"
against "what the agent then fetched/referenced".
Deliberately FK-free on user_id (mirrors AppLog): telemetry should outlive
the row it describes, and a deleted user shouldn't cascade away history.
"""
__tablename__ = "retrieval_logs"
id: Mapped[int] = mapped_column(primary_key=True)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), default=lambda: datetime.now(timezone.utc)
)
user_id: Mapped[int | None] = mapped_column(Integer, nullable=True)
# Retrieval surface: 'mcp_search' | 'rest_search' | 'auto_inject' | ...
source: Mapped[str] = mapped_column(Text, nullable=False)
query: Mapped[str | None] = mapped_column(Text, nullable=True)
# Effective parameters actually used for this call.
threshold: Mapped[float | None] = mapped_column(Float, nullable=True)
limit_n: Mapped[int | None] = mapped_column(Integer, nullable=True)
project_id: Mapped[int | None] = mapped_column(Integer, nullable=True)
# The content-type filter as passed to semantic_search_notes: True=tasks,
# False=notes, NULL=any.
is_task: Mapped[bool | None] = mapped_column(Boolean, nullable=True)
result_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
top_score: Mapped[float | None] = mapped_column(Float, nullable=True)
min_score: Mapped[float | None] = mapped_column(Float, nullable=True)
# [{"id": int, "score": float, "rank": int}, ...], highest-first.
result_ids: Mapped[list] = mapped_column(JSONB, nullable=False, default=list)
duration_ms: Mapped[float | None] = mapped_column(Float, nullable=True)
__table_args__ = (
Index("ix_retrieval_logs_created_at", "created_at"),
Index("ix_retrieval_logs_user_id", "user_id"),
Index("ix_retrieval_logs_source", "source"),
Index("ix_retrieval_logs_source_created_at", "source", created_at.desc()),
)
def to_dict(self) -> dict:
return {
"id": self.id,
"created_at": self.created_at.isoformat() if self.created_at else None,
"user_id": self.user_id,
"source": self.source,
"query": self.query,
"threshold": self.threshold,
"limit_n": self.limit_n,
"project_id": self.project_id,
"is_task": self.is_task,
"result_count": self.result_count,
"top_score": self.top_score,
"min_score": self.min_score,
"result_ids": self.result_ids,
"duration_ms": self.duration_ms,
}
+42
View File
@@ -57,6 +57,48 @@ async def session_context():
return jsonify(result) return jsonify(result)
@plugin_bp.get("/retrieve")
@login_required
async def autoinject_retrieve():
"""Title-first knowledge auto-inject for the plugin's UserPromptSubmit hook.
Given the user's prompt (`q`), returns a compact awareness hint — note
titles + scores only, never bodies — for the top hits that clear the
per-user gates (see services.plugin_context.build_autoinject_hint). Returns
empty context (most of the time) when disabled or nothing is relevant.
Query:
q (str) — the user's prompt to retrieve against.
repo (optional) — working repo remote; resolved to the bound project
to scope the search (mirrors /context). Unbound or
absent → searches all the user's notes.
project_id (opt) — explicit project scope override (ad-hoc/testing).
exclude_ids (opt) — comma-separated note ids already injected this
session; skipped so each note injects at most once.
"""
q = (request.args.get("q") or "").strip()
try:
project_id = int(request.args.get("project_id", 0) or 0)
except (TypeError, ValueError):
project_id = 0
repo = (request.args.get("repo") or "").strip()
if repo and not project_id:
resolved = await repo_bindings_svc.resolve_project(g.user.id, repo)
if resolved:
project_id = resolved
exclude_ids = [
int(p) for p in (request.args.get("exclude_ids") or "").split(",")
if p.strip().isdigit()
]
result = await plugin_ctx_svc.build_autoinject_hint(
g.user.id, q, project_id=project_id, exclude_ids=exclude_ids
)
return jsonify(result)
@plugin_bp.get("/processes") @plugin_bp.get("/processes")
@login_required @login_required
async def process_manifest(): async def process_manifest():
+15 -1
View File
@@ -1,7 +1,14 @@
import time
from quart import Blueprint, jsonify, request from quart import Blueprint, jsonify, request
from scribe.auth import login_required, get_current_user_id from scribe.auth import login_required, get_current_user_id
from scribe.services.embeddings import semantic_search_notes from scribe.services.embeddings import semantic_search_notes
from scribe.services.retrieval_telemetry import record_retrieval
# This route searches with a looser floor than the MCP tool default — it powers
# an interactive feed where loosely-related hits still have value.
_REST_SEARCH_THRESHOLD = 0.3
search_bp = Blueprint("search", __name__, url_prefix="/api/search") search_bp = Blueprint("search", __name__, url_prefix="/api/search")
@@ -27,8 +34,15 @@ async def search_route():
limit = min(request.args.get("limit", 10, type=int), 50) limit = min(request.args.get("limit", 10, type=int), 50)
is_task = _content_type_to_is_task(content_type) is_task = _content_type_to_is_task(content_type)
t0 = time.perf_counter()
results = await semantic_search_notes( results = await semantic_search_notes(
uid, q, limit=limit, is_task=is_task, threshold=0.3 uid, q, limit=limit, is_task=is_task, threshold=_REST_SEARCH_THRESHOLD
)
record_retrieval(
user_id=uid, source="rest_search", query=q,
threshold=_REST_SEARCH_THRESHOLD, limit=limit,
project_id=None, is_task=is_task, results=results,
duration_ms=(time.perf_counter() - t0) * 1000.0,
) )
return jsonify({ return jsonify({
"results": [ "results": [
+23 -20
View File
@@ -28,6 +28,10 @@ logger = logging.getLogger(__name__)
# loosely-related results that pad the sidebar without adding real value. # loosely-related results that pad the sidebar without adding real value.
_SIMILARITY_THRESHOLD = 0.45 _SIMILARITY_THRESHOLD = 0.45
# Public alias so callers (and telemetry) can record the effective default
# threshold without reaching for the underscored name.
DEFAULT_SIMILARITY_THRESHOLD = _SIMILARITY_THRESHOLD
_MODEL_NAME = "BAAI/bge-small-en-v1.5" _MODEL_NAME = "BAAI/bge-small-en-v1.5"
_CACHE_DIR = os.environ.get("FASTEMBED_CACHE_DIR", "/data/fastembed-cache") _CACHE_DIR = os.environ.get("FASTEMBED_CACHE_DIR", "/data/fastembed-cache")
@@ -115,6 +119,14 @@ async def semantic_search_notes(
Scores are cosine similarities in [-1, 1]; only notes at or above Scores are cosine similarities in [-1, 1]; only notes at or above
*threshold* are returned, sorted highest-first. *threshold* are returned, sorted highest-first.
Ranking and the top-k cut happen in Postgres via pgvector's cosine-distance
operator (`<=>`, exposed as ``Vector.cosine_distance``) backed by the HNSW
index from migration 0067 — so this is an indexed ``ORDER BY ... LIMIT k``
rather than a full-table scan. Cosine distance is ``1 - cosine_similarity``,
so a similarity floor of *threshold* is a distance ceiling of
``1 - threshold`` and similarity is recovered as ``1 - distance``.
Returns an empty list if the embedder is unavailable or on any error. Returns an empty list if the embedder is unavailable or on any error.
""" """
if not query or not query.strip(): if not query or not query.strip():
@@ -125,10 +137,17 @@ async def semantic_search_notes(
logger.debug("Semantic search skipped — embedder unavailable") logger.debug("Semantic search skipped — embedder unavailable")
return [] return []
# Distance ceiling equivalent to the similarity floor. Clamp to the valid
# cosine-distance range [0, 2] so a threshold of, say, -1 doesn't produce a
# nonsensical ceiling.
max_distance = min(2.0, max(0.0, 1.0 - threshold))
distance = NoteEmbedding.embedding.cosine_distance(query_vec)
try: try:
async with async_session() as session: async with async_session() as session:
stmt = ( stmt = (
select(NoteEmbedding, Note) select(Note, distance.label("distance"))
.select_from(NoteEmbedding)
.join(Note, NoteEmbedding.note_id == Note.id) .join(Note, NoteEmbedding.note_id == Note.id)
.where(NoteEmbedding.user_id == user_id, Note.deleted_at.is_(None)) .where(NoteEmbedding.user_id == user_id, Note.deleted_at.is_(None))
) )
@@ -142,30 +161,14 @@ async def semantic_search_notes(
stmt = stmt.where(Note.status.is_(None)) stmt = stmt.where(Note.status.is_(None))
if exclude_ids: if exclude_ids:
stmt = stmt.where(NoteEmbedding.note_id.notin_(exclude_ids)) stmt = stmt.where(NoteEmbedding.note_id.notin_(exclude_ids))
stmt = stmt.where(distance <= max_distance).order_by(distance.asc()).limit(limit)
rows = list((await session.execute(stmt)).all()) rows = list((await session.execute(stmt)).all())
except Exception: except Exception:
logger.warning("Failed to query note embeddings", exc_info=True) logger.warning("Failed to query note embeddings", exc_info=True)
return [] return []
if not rows: # Recover similarity (1 - distance) and preserve the highest-first contract.
return [] return [(1.0 - float(dist), note) for note, dist in rows]
def _score() -> list[tuple[float, Note]]:
out: list[tuple[float, Note]] = []
for ne, note in rows:
try:
sim = _cosine_similarity(query_vec, ne.embedding)
except Exception:
continue
if sim >= threshold:
out.append((sim, note))
out.sort(key=lambda x: x[0], reverse=True)
return out[:limit]
# Offload the O(rows) cosine scoring off the event loop so a large corpus
# doesn't stall other requests while ranking. Results are unchanged; the
# real scaling fix (ORDER BY / LIMIT in pgvector) is a separate effort.
return await asyncio.to_thread(_score)
async def backfill_note_embeddings() -> None: async def backfill_note_embeddings() -> None:
+115
View File
@@ -15,6 +15,7 @@ index alone already steers behavior.
from __future__ import annotations from __future__ import annotations
import re import re
import time
from sqlalchemy import select from sqlalchemy import select
@@ -24,6 +25,9 @@ from scribe.services import knowledge as knowledge_svc
from scribe.services import notes as notes_svc from scribe.services import notes as notes_svc
from scribe.services import projects as projects_svc from scribe.services import projects as projects_svc
from scribe.services import rulebooks as rulebooks_svc from scribe.services import rulebooks as rulebooks_svc
from scribe.services.embeddings import semantic_search_notes
from scribe.services.retrieval_telemetry import record_retrieval
from scribe.services.settings import get_setting
# Defensive cap below Claude Code's 10k additionalContext limit. # Defensive cap below Claude Code's 10k additionalContext limit.
_MAX_CHARS = 9000 _MAX_CHARS = 9000
@@ -31,6 +35,28 @@ _MAX_CHARS = 9000
# Max chars of a Process body to fold into the auto-surface description. # Max chars of a Process body to fold into the auto-surface description.
_PROC_PREVIEW_CHARS = 200 _PROC_PREVIEW_CHARS = 200
# --- Knowledge auto-inject (Path A: per-turn awareness push) -----------------
# Per-user settings (keys live in the generic settings table). The threshold is
# deliberately STRICTER than the pull-search default (embeddings
# DEFAULT_SIMILARITY_THRESHOLD = 0.45): an unsolicited per-turn inject must clear
# a higher bar than a search the agent chose to run. Defaults start conservative
# and are meant to be tuned from retrieval_logs (source='auto_inject') once data
# accrues — they're exposed in the Settings UI, no restart needed.
AUTOINJECT_ENABLED_KEY = "kb_autoinject_enabled"
AUTOINJECT_THRESHOLD_KEY = "kb_autoinject_threshold"
AUTOINJECT_TOP_K_KEY = "kb_autoinject_top_k"
AUTOINJECT_DEFAULT_ENABLED = True
AUTOINJECT_DEFAULT_THRESHOLD = 0.55
AUTOINJECT_DEFAULT_TOP_K = 3
# Margin gate: drop any hit more than this far below the top hit's score, so a
# single strong match doesn't drag in a wall of barely-passing neighbours.
_AUTOINJECT_BAND = 0.10
# Hard ceiling on top-k regardless of the user's setting — this is an
# awareness menu (titles only), never a content dump.
_AUTOINJECT_MAX_TOP_K = 10
def _slugify(text: str) -> str: def _slugify(text: str) -> str:
"""kebab-case slug for a skill directory name (a-z0-9 + single hyphens).""" """kebab-case slug for a skill directory name (a-z0-9 + single hyphens)."""
@@ -82,6 +108,95 @@ async def build_process_manifest(user_id: int) -> dict:
return {"processes": procs, "total": len(procs)} return {"processes": procs, "total": len(procs)}
async def get_autoinject_config(user_id: int) -> dict:
"""Resolve a user's auto-inject settings, falling back to the defaults.
Returns {"enabled": bool, "threshold": float, "top_k": int}, clamped to
sane ranges (threshold to [0,1]; top_k to [1, _AUTOINJECT_MAX_TOP_K]).
"""
enabled_raw = await get_setting(
user_id, AUTOINJECT_ENABLED_KEY,
"true" if AUTOINJECT_DEFAULT_ENABLED else "false",
)
enabled = enabled_raw.strip().lower() in ("true", "1", "yes", "on")
try:
threshold = float(await get_setting(
user_id, AUTOINJECT_THRESHOLD_KEY, str(AUTOINJECT_DEFAULT_THRESHOLD)))
except (TypeError, ValueError):
threshold = AUTOINJECT_DEFAULT_THRESHOLD
threshold = min(1.0, max(0.0, threshold))
try:
top_k = int(float(await get_setting(
user_id, AUTOINJECT_TOP_K_KEY, str(AUTOINJECT_DEFAULT_TOP_K))))
except (TypeError, ValueError):
top_k = AUTOINJECT_DEFAULT_TOP_K
top_k = min(_AUTOINJECT_MAX_TOP_K, max(1, top_k))
return {"enabled": enabled, "threshold": threshold, "top_k": top_k}
async def build_autoinject_hint(
user_id: int,
query: str,
project_id: int = 0,
exclude_ids: list[int] | None = None,
) -> dict:
"""Title-first awareness hint for the plugin's UserPromptSubmit hook.
The four anti-bloat gates (see the module + milestone-93 design):
1. high-confidence threshold (stricter than pull) — set per-user;
2. margin gate — keep only hits within _AUTOINJECT_BAND of the top score;
3. session dedup — caller passes already-injected ids as `exclude_ids`;
4. title-first payload — id + title + score only, never bodies.
Disabled, blank-query, or nothing-clears-the-gates all return empty context,
so most turns inject nothing.
Returns {"context": str, "note_ids": list[int], "config": dict}. Every
retrieval (even empty) is logged to retrieval_logs as source='auto_inject'
so the threshold can be tuned from data.
"""
cfg = await get_autoinject_config(user_id)
empty = {"context": "", "note_ids": [], "config": cfg}
q = (query or "").strip()
if not cfg["enabled"] or not q:
return empty
t0 = time.perf_counter()
hits = await semantic_search_notes(
user_id, q,
limit=cfg["top_k"],
threshold=cfg["threshold"],
project_id=(project_id or None),
exclude_ids=set(exclude_ids or []),
)
record_retrieval(
user_id=user_id, source="auto_inject", query=q,
threshold=cfg["threshold"], limit=cfg["top_k"],
project_id=(project_id or None), is_task=None, results=hits,
duration_ms=(time.perf_counter() - t0) * 1000.0,
)
if not hits:
return empty
# Margin gate: keep only hits close to the strongest one.
top_score = hits[0][0]
kept = [(s, n) for s, n in hits if s >= top_score - _AUTOINJECT_BAND]
lines = [
"> Possibly relevant from your Scribe notes — call `get_note(id)` to "
"open any in full (titles only; injected once per session):",
]
note_ids: list[int] = []
for score, note in kept:
note_ids.append(int(note.id))
title = (note.title or "(untitled)").replace("\n", " ").strip()
lines.append(f"> - #{note.id} \"{title}\" ({score:.2f})")
return {"context": "\n".join(lines), "note_ids": note_ids, "config": cfg}
async def _topic_titles(topic_ids: set[int]) -> dict[int, str]: async def _topic_titles(topic_ids: set[int]) -> dict[int, str]:
"""Map topic_id -> title for the given ids (live topics only).""" """Map topic_id -> title for the given ids (live topics only)."""
if not topic_ids: if not topic_ids:
+115
View File
@@ -0,0 +1,115 @@
"""Retrieval telemetry — one RetrievalLog row per semantic-retrieval call.
This is the empirical basis for KB-injection tuning: it records what each query
asked for, the score distribution of what came back, and the effective params,
so the similarity threshold and top-k can be tuned from data rather than guessed.
Design notes:
- Fire-and-forget, mirroring upsert_note_embedding: `record_retrieval` extracts
the primitives it needs SYNCHRONOUSLY (while the caller's Note objects are
still valid) and schedules the DB insert as a background task, so logging
never adds latency to — or can break — the search response.
- Result objects are reduced to {id, score, rank} before scheduling; the
background writer touches only plain data, never a possibly-detached ORM row.
- Every failure path is swallowed: telemetry must never take down retrieval.
"""
from __future__ import annotations
import asyncio
import logging
from scribe.models import async_session
from scribe.models.note import Note
from scribe.models.retrieval_log import RetrievalLog
logger = logging.getLogger(__name__)
def _build_payload(
*,
user_id: int | None,
source: str,
query: str | None,
threshold: float | None,
limit: int | None,
project_id: int | None,
is_task: bool | None,
results: list[tuple[float, Note]],
duration_ms: float | None,
) -> dict:
"""Reduce a retrieval call to a flat, JSON-safe RetrievalLog payload.
Pure and synchronous (no DB, no event loop) so it is unit-testable and safe
to run inline before scheduling the write. `results` is the
`(score, Note)` list from semantic_search_notes, already highest-first.
"""
items = [
{"id": int(note.id), "score": round(float(score), 5), "rank": rank}
for rank, (score, note) in enumerate(results)
]
scores = [it["score"] for it in items]
return {
"user_id": user_id,
"source": source,
"query": query,
"threshold": threshold,
"limit_n": limit,
"project_id": project_id,
"is_task": is_task,
"result_count": len(items),
"top_score": (scores[0] if scores else None),
"min_score": (scores[-1] if scores else None),
"result_ids": items,
"duration_ms": (round(duration_ms, 2) if duration_ms is not None else None),
}
async def _insert_retrieval_log(payload: dict) -> None:
"""Persist one RetrievalLog row. Best-effort: all errors are swallowed."""
try:
async with async_session() as session:
session.add(RetrievalLog(**payload))
await session.commit()
except Exception:
logger.debug("retrieval telemetry write skipped", exc_info=True)
def record_retrieval(
*,
user_id: int | None,
source: str,
query: str | None,
threshold: float | None,
limit: int | None,
project_id: int | None,
is_task: bool | None,
results: list[tuple[float, Note]],
duration_ms: float | None = None,
) -> None:
"""Fire-and-forget: record one retrieval call.
Builds the payload inline (synchronously) then schedules the insert so the
caller returns immediately. Never raises — telemetry must not affect search.
"""
try:
payload = _build_payload(
user_id=user_id,
source=source,
query=query,
threshold=threshold,
limit=limit,
project_id=project_id,
is_task=is_task,
results=results,
duration_ms=duration_ms,
)
except Exception:
logger.debug("retrieval telemetry payload build failed", exc_info=True)
return
try:
asyncio.get_running_loop().create_task(_insert_retrieval_log(payload))
except RuntimeError:
# No running loop (e.g. called from sync context outside the app) —
# skip rather than block. The app paths always run on the loop.
logger.debug("retrieval telemetry skipped — no running event loop")
+101
View File
@@ -0,0 +1,101 @@
"""Real-Postgres integration test for pgvector semantic search.
Runs only in the CI integration lane (real Postgres + `vector` extension +
schema built by `alembic upgrade head`, which includes migration 0067). This
exercises what the unit mocks cannot: the native `vector(384)` column, the
`<=>` cosine-distance operator behind `Vector.cosine_distance`, the HNSW index,
and the distance->similarity recovery in `semantic_search_notes`.
The embedder itself is stubbed (get_embedding is patched) so the test does not
depend on downloading the fastembed model — only the Postgres/pgvector path is
under test.
"""
from unittest.mock import AsyncMock, patch
import pytest
import pytest_asyncio
from sqlalchemy import delete
from scribe.models import async_session, engine
from scribe.models.embedding import EMBEDDING_DIM, NoteEmbedding
from scribe.models.note import Note
from scribe.models.user import User
from scribe.services.embeddings import semantic_search_notes
pytestmark = pytest.mark.integration
def _vec(*nonzero_first):
"""A 384-dim vector with the given leading values, zero-padded."""
v = list(nonzero_first) + [0.0] * (EMBEDDING_DIM - len(nonzero_first))
return v[:EMBEDDING_DIM]
@pytest_asyncio.fixture(autouse=True)
async def _dispose_engine():
# Per-loop pool: dispose after each test (see test_integration_db_maintenance).
yield
await engine.dispose()
@pytest_asyncio.fixture
async def seeded():
"""Insert a user + a near and a far note with hand-crafted embeddings.
Returns (user_id, near_note_id, far_note_id). Cleaned up after the test.
"""
async with async_session() as s:
user = User(username="pgvec_itest")
s.add(user)
await s.flush()
near = Note(user_id=user.id, title="near", body="near body")
far = Note(user_id=user.id, title="far", body="far body")
s.add_all([near, far])
await s.flush()
# query vector will be [1,0,0,...]; near ~ identical (sim≈1.0),
# far is orthogonal (sim≈0.0 -> filtered by the default threshold).
s.add(NoteEmbedding(note_id=near.id, user_id=user.id, embedding=_vec(1.0)))
s.add(NoteEmbedding(note_id=far.id, user_id=user.id, embedding=_vec(0.0, 1.0)))
await s.commit()
ids = (user.id, near.id, far.id)
yield ids
user_id = ids[0]
async with async_session() as s:
await s.execute(delete(NoteEmbedding).where(NoteEmbedding.user_id == user_id))
await s.execute(delete(Note).where(Note.user_id == user_id))
await s.execute(delete(User).where(User.id == user_id))
await s.commit()
@pytest.mark.asyncio
async def test_semantic_search_ranks_and_thresholds_via_pgvector(seeded):
user_id, near_id, far_id = seeded
with patch(
"scribe.services.embeddings.get_embedding",
AsyncMock(return_value=_vec(1.0)),
):
results = await semantic_search_notes(user_id=user_id, query="anything", limit=10)
ids = [note.id for _score, note in results]
# Near note returned and ranked first; far (orthogonal, sim≈0) excluded by
# the default 0.45 similarity threshold.
assert near_id in ids
assert far_id not in ids
assert ids[0] == near_id
top_score = results[0][0]
assert top_score == pytest.approx(1.0, abs=1e-3)
@pytest.mark.asyncio
async def test_low_threshold_lets_orthogonal_through(seeded):
user_id, near_id, far_id = seeded
with patch(
"scribe.services.embeddings.get_embedding",
AsyncMock(return_value=_vec(1.0)),
):
results = await semantic_search_notes(
user_id=user_id, query="anything", limit=10, threshold=-1.0,
)
ids = [note.id for _score, note in results]
# With the floor dropped, both come back and near still ranks above far.
assert ids.index(near_id) < ids.index(far_id)
+86
View File
@@ -10,6 +10,92 @@ def _rule(rid, title, topic_id):
return r return r
def _note(nid, title):
n = MagicMock()
n.id, n.title = nid, title
return n
# ─── knowledge auto-inject (Path A) ──────────────────────────────────────────
@pytest.mark.asyncio
async def test_get_autoinject_config_defaults_and_clamps():
from scribe.services import plugin_context as pc
# No settings stored → defaults.
with patch.object(pc, "get_setting", AsyncMock(side_effect=lambda uid, k, d: d)):
cfg = await pc.get_autoinject_config(1)
assert cfg == {
"enabled": pc.AUTOINJECT_DEFAULT_ENABLED,
"threshold": pc.AUTOINJECT_DEFAULT_THRESHOLD,
"top_k": pc.AUTOINJECT_DEFAULT_TOP_K,
}
# Out-of-range values are clamped; top_k capped at the hard ceiling.
stored = {
pc.AUTOINJECT_ENABLED_KEY: "false",
pc.AUTOINJECT_THRESHOLD_KEY: "5",
pc.AUTOINJECT_TOP_K_KEY: "999",
}
with patch.object(pc, "get_setting",
AsyncMock(side_effect=lambda uid, k, d: stored.get(k, d))):
cfg = await pc.get_autoinject_config(1)
assert cfg["enabled"] is False
assert cfg["threshold"] == 1.0
assert cfg["top_k"] == pc._AUTOINJECT_MAX_TOP_K
@pytest.mark.asyncio
async def test_build_autoinject_hint_disabled_returns_empty_and_skips_search():
from scribe.services import plugin_context as pc
search = AsyncMock()
with patch.object(pc, "get_autoinject_config",
AsyncMock(return_value={"enabled": False, "threshold": 0.55, "top_k": 3})), \
patch.object(pc, "semantic_search_notes", search), \
patch.object(pc, "record_retrieval", MagicMock()):
out = await pc.build_autoinject_hint(1, "anything")
assert out["context"] == "" and out["note_ids"] == []
search.assert_not_called() # disabled → no retrieval at all
@pytest.mark.asyncio
async def test_build_autoinject_hint_titles_only_with_margin_gate():
from scribe.services import plugin_context as pc
# top=0.80; 0.74 within band (0.10), 0.61 outside → dropped.
hits = [(0.80, _note(11, "Pool sizing decision")),
(0.74, _note(22, "run_maintenance thresholds")),
(0.61, _note(33, "unrelated-ish"))]
rec = MagicMock()
with patch.object(pc, "get_autoinject_config",
AsyncMock(return_value={"enabled": True, "threshold": 0.55, "top_k": 3})), \
patch.object(pc, "semantic_search_notes", AsyncMock(return_value=hits)), \
patch.object(pc, "record_retrieval", rec):
out = await pc.build_autoinject_hint(1, "postgres pool", project_id=2,
exclude_ids=[99])
# Margin gate kept the top two, dropped the straggler.
assert out["note_ids"] == [11, 22]
assert '#11 "Pool sizing decision" (0.80)' in out["context"]
assert "#33" not in out["context"]
# Title-first: no body text, ever.
assert "get_note(id)" in out["context"]
# Telemetry fired with the auto_inject source and the full candidate set.
rec.assert_called_once()
assert rec.call_args.kwargs["source"] == "auto_inject"
@pytest.mark.asyncio
async def test_build_autoinject_hint_blank_query_returns_empty():
from scribe.services import plugin_context as pc
search = AsyncMock()
with patch.object(pc, "get_autoinject_config",
AsyncMock(return_value={"enabled": True, "threshold": 0.55, "top_k": 3})), \
patch.object(pc, "semantic_search_notes", search):
out = await pc.build_autoinject_hint(1, " ")
assert out["context"] == ""
search.assert_not_called()
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_build_session_context_renders_titles_grouped_by_topic(): async def test_build_session_context_renders_titles_grouped_by_topic():
rules = [ rules = [
+112
View File
@@ -0,0 +1,112 @@
"""Tests for services.retrieval_telemetry.
_build_payload is pure (no DB, no loop) and gets unit coverage. The persistence
path (_insert_retrieval_log + the RetrievalLog model / JSONB roundtrip) is an
integration test against real Postgres.
"""
from types import SimpleNamespace
import pytest
import pytest_asyncio
from scribe.services.retrieval_telemetry import (
_build_payload,
record_retrieval,
)
def _note(nid):
"""Minimal stand-in — _build_payload only reads .id."""
return SimpleNamespace(id=nid)
# ─── _build_payload (pure) ───────────────────────────────────────────────────
def test_build_payload_ranks_and_score_bounds():
results = [(0.91, _note(11)), (0.72, _note(22)), (0.55, _note(33))]
p = _build_payload(
user_id=7, source="mcp_search", query="hello", threshold=0.45,
limit=10, project_id=3, is_task=None, results=results, duration_ms=12.345,
)
assert p["result_count"] == 3
assert p["top_score"] == 0.91
assert p["min_score"] == 0.55
assert [it["rank"] for it in p["result_ids"]] == [0, 1, 2]
assert [it["id"] for it in p["result_ids"]] == [11, 22, 33]
assert p["duration_ms"] == 12.35 # rounded to 2dp
assert p["user_id"] == 7 and p["project_id"] == 3 and p["threshold"] == 0.45
def test_build_payload_empty_results():
p = _build_payload(
user_id=1, source="rest_search", query="x", threshold=0.3,
limit=5, project_id=None, is_task=False, results=[], duration_ms=None,
)
assert p["result_count"] == 0
assert p["top_score"] is None and p["min_score"] is None
assert p["result_ids"] == []
assert p["duration_ms"] is None
def test_build_payload_rounds_scores_to_5dp():
p = _build_payload(
user_id=1, source="mcp_search", query="q", threshold=0.45,
limit=1, project_id=None, is_task=None,
results=[(0.123456789, _note(1))], duration_ms=0.0,
)
assert p["result_ids"][0]["score"] == 0.12346
def test_record_retrieval_without_event_loop_is_safe():
"""Called from a sync context (no running loop) it must swallow and return,
never raise — telemetry can't be allowed to break a caller."""
# No event loop running in this plain sync test.
assert record_retrieval(
user_id=1, source="mcp_search", query="q", threshold=0.45,
limit=10, project_id=None, is_task=None,
results=[(0.9, _note(1))],
) is None
# ─── persistence (integration) ───────────────────────────────────────────────
@pytest_asyncio.fixture
async def _dispose_engine():
from scribe.models import engine
yield
await engine.dispose()
@pytest.mark.integration
@pytest.mark.asyncio
async def test_insert_retrieval_log_roundtrip(_dispose_engine):
from sqlalchemy import delete, select
from scribe.models import async_session
from scribe.models.retrieval_log import RetrievalLog
from scribe.services.retrieval_telemetry import _insert_retrieval_log
payload = _build_payload(
user_id=990001, source="mcp_search", query="pgvector tuning",
threshold=0.45, limit=10, project_id=None, is_task=None,
results=[(0.88, _note(501)), (0.61, _note(502))], duration_ms=9.9,
)
await _insert_retrieval_log(payload)
async with async_session() as s:
row = (
await s.execute(
select(RetrievalLog).where(RetrievalLog.user_id == 990001)
)
).scalars().first()
assert row is not None
assert row.source == "mcp_search"
assert row.result_count == 2
assert row.top_score == 0.88
# JSONB roundtrips as a list of dicts with the expected shape.
assert row.result_ids[0] == {"id": 501, "score": 0.88, "rank": 0}
assert row.created_at is not None # server_default now()
await s.execute(delete(RetrievalLog).where(RetrievalLog.user_id == 990001))
await s.commit()