Skip to content

fix(pgvector): tag similarity-search SQL as RETRIEVER + populate input#177

Open
SuhaniNagpal7 wants to merge 1 commit into
devfrom
fix/pgvector-retriever-attrs
Open

fix(pgvector): tag similarity-search SQL as RETRIEVER + populate input#177
SuhaniNagpal7 wants to merge 1 commit into
devfrom
fix/pgvector-retriever-attrs

Conversation

@SuhaniNagpal7
Copy link
Copy Markdown
Contributor

Summary

PgvectorInstrumentor wraps psycopg.cursor.execute() and executemany() and tags the span based on the SQL it sees: query for SELECT with a vector op, insert/update/delete otherwise. None of these currently set the FI canonical retriever keys, so a textbook RAG similarity-search query lands in Future AGI as Type: unknown with an empty Input panel.

This PR adds gen_ai.span.kind = RETRIEVER and input.value only when the detected operation is a similarity-search read (db.operation.name == \"query\"). INSERTs that happen to contain a vector operator (e.g. INSERT ... RETURNING id <=> '[...]') are deliberately not tagged — they're writes, not retrievals.

What changes

All in traceai_pgvector/_wrappers.py:

  • Optional fi_instrumentation.fi_types import with raw-string fallback.
  • New _is_pgvector_query_op(metadata) predicate — True iff the detected operation is \"query\".
  • New _truncate_sql(sql, limit=2000) helper.
  • ExecuteWrapper.__call__ and ExecuteManyWrapper.__call__ — when _is_pgvector_query_op(metadata):
    • gen_ai.span.kind = \"RETRIEVER\"
    • input.value = the SQL string (resolved via as_string(None) for psycopg Composed objects, else str()), truncated to 2000 chars
    • input.mime_type = text/plain

Output is deliberately not set. psycopg.execute() returns None (the rows live on the cursor, not the return value), so there's no clean way to capture results from this wrapping layer. The Output panel will remain empty for pgvector retriever spans — that's expected.

All pre-existing db.vector.* attrs preserved. INSERT/UPDATE/DELETE behavior is unchanged — they still get the existing pgvector technical attrs but without the RETRIEVER tag.

Verified

  • In-process attribute check on both positive and negative cases:
    • Similarity-search SQL → gen_ai.span.kind = \"RETRIEVER\", input.value populated
    • INSERT containing a vector op → gen_ai.span.kind NOT set (correctly stays None/unknown)
  • Real end-to-end ingest to Future AGI → confirmed via Future AGI MCP:
    • pgvector query (SELECT ... embedding <=> '[...]' ORDER BY dist LIMIT 5) → Type: Retriever, Input: full SQL string
    • pgvector insert (INSERT with a vector op) → Type: unknown ✅ (negative case correctly handled)

Out of scope

  • Output capture from psycopg results (would require wrapping cursor.fetchall() / iteration — separate concern).
  • Per-document retrieval.documents.N.* attrs (Tier 3).
  • Other 6 vector DBs get their own PRs.

The pgvector instrumentor wraps psycopg's cursor.execute() and
executemany() and tags the span based on the SQL it sees: "query"
for SELECT with a vector op, "insert"/"update"/"delete" for the
respective DML, "unknown" otherwise. Currently none of these set the
FI canonical retriever keys, so even a textbook RAG similarity-search
query lands in Future AGI as Type=unknown with an empty Input panel.

This PR adds the FI canonical kind + input.value, but ONLY when the
detected operation is a similarity-search read (db.operation.name ==
"query"). INSERT/UPDATE/DELETE that happen to contain a vector
operator (e.g. INSERT ... RETURNING id <=> '[...]') are deliberately
NOT tagged RETRIEVER — they're writes, not retrievals.

Changes (all in traceai_pgvector/_wrappers.py)

- Optional `fi_instrumentation.fi_types` import with raw-string fallback.
- New `_is_pgvector_query_op(metadata)` predicate — True iff the
  detected operation is "query".
- New `_truncate_sql(sql, limit=2000)` helper to cap the SQL string
  attached to input.value at a reasonable size for span attributes.
- `ExecuteWrapper.__call__` — when `_is_pgvector_query_op(metadata)`:
  - Set `gen_ai.span.kind = "RETRIEVER"`.
  - Set `input.value` to the SQL string (resolved via `as_string(None)`
    when a psycopg `Composed` object is passed, else `str()`), truncated.
  - Set `input.mime_type = text/plain`.
- `ExecuteManyWrapper.__call__` — same treatment for batched
  similarity-search reads.

Output.value is deliberately not set: psycopg's execute() returns
None (rows live on the cursor, not the return value), so there's no
clean way to capture results from this wrapping layer.

All pre-existing `db.vector.*` attrs preserved. INSERT/UPDATE/DELETE
behavior unchanged — they still get the existing pgvector attrs but
without the RETRIEVER tag, which is correct.

Verified end-to-end via Future AGI MCP:
  * `pgvector query` (SELECT ... <=> ... ORDER BY ... LIMIT) → Type=Retriever
  * `pgvector insert` (INSERT with a vector op) → Type=unknown ✅
@SuhaniNagpal7 SuhaniNagpal7 self-assigned this May 13, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant