Weights & Biases

GoodMem retrieval for Weights & Biases Weave — traced retrieval ops, an evaluable weave.Model, and retrieval scorers.

Preview· Python

Overview

GoodMem is a memory layer for AI agents that handles embedding, vector search, reranking, and LLM-powered answering server-side. This integration brings GoodMem retrieval into Weights & Biases Weave: GoodMemRetriever runs every search as a traced Weave op, so a RAG call shows up in the Weave UI with its query, each hit's score and what kind of score it is, and any degradation the server reported. GoodMemRetrievalModel wraps the same retriever as a weave.Model, and the bundled scorers — RecallAtK, MRR, FactRecall, and RetrievalHealth — let weave.Evaluation compare retrieval configurations side by side.

Installation

pip install goodmem-wandb

The pip package is goodmem-wandb; the import name is goodmem_wandb. Installing it pulls in weave and the goodmem SDK as dependencies.

Configuration

Credentials are read from environment variables by default:

Variable	Description
`GOODMEM_BASE_URL`	Base URL of your GoodMem server (no `/v1` suffix), e.g. `https://your-goodmem-server.example.com`.
`GOODMEM_API_KEY`	API key sent as `X-API-Key` to authenticate requests.

They can also be passed as the constructor keywords base_url and api_key, alongside verify_ssl (default True; only turn it off for a local server with a self-signed certificate), timeout, and client (a caller-owned goodmem SDK client, which the integration never closes).

Credentials are deliberately not Weave fields. Weave publishes an object's pydantic fields verbatim to the trace server, and its redaction helper does not run on that path — a field named api_key on a weave.Model would be uploaded to W&B in plaintext. This package keeps the connection on a private attribute instead, so what Weave versions and publishes is the retrieval configuration only.

Quick start

import weave
from goodmem_wandb import GoodMemRetriever

weave.init("my-project")

retriever = GoodMemRetriever(space_name="docs")   # credentials from the environment
result = retriever.search("How do I rotate an API key?")

for hit in result["hits"]:
    print(hit["score"], hit["score_kind"], hit["chunk_text"][:80])

search() returns a dict with query, hits (each hit carries chunk_id, chunk_text, memory_id, space_id, source, score, score_kind, and metadata, in the server's order), score_kind, statuses, partial, abstract_reply, and space_ids. partial is True when the server reported a real problem during the retrieval — an empty hits with partial=True is a failed search, not a miss. score_kind is "vector" or "reranker" and reports what the server actually returned: the two are different scales (vector scores are frequently negative), so results keep the server's order and are never re-sorted client-side.

Local development. If you are running GoodMem locally with a self-signed certificate, use https://localhost:8080 and disable TLS verification with the verify_ssl constructor keyword. Keep TLS verification enabled for any deployed server.

Retriever configuration

Beyond the space selectors (space_name, space_id, or space_ids), GoodMemRetriever and GoodMemRetrievalModel accept:

Field	Description
`limit`	Maximum number of hits to return (default 5).
`fetch_k`	Candidates to retrieve before reranking; defaults to `limit`.
`reranker_id`	UUID of a reranker model that reorders matched chunks by relevance.
`min_score`	Client-side threshold applied only to reranker scores; reranker scales are model-dependent, so calibrate it for the reranker in use.
`filter`	A GoodMem filter expression applied to every space searched.
`metadata_filter`	Field/value pairs AND-ed into the filter; each value is compared as its own type (`str` as TEXT, `bool` as BOOLEAN, `int`/`float` as NUMERIC).
`llm_id`	UUID of an LLM that generates a contextual answer (`abstract_reply`) alongside the chunks.
`llm_temperature`	Sampling temperature for the LLM post-processor.
`embedder_id` / `create_space`	With `space_name`: reuse an existing space whose embedder matches, or create it when `create_space=True`.
`trace_chunk_text`	Set `False` for a sensitive corpus: chunk ids, scores, and statuses are still traced, but the retrieved text is left out of the Weave trace.

Evaluate a retrieval configuration

Wrap the retriever as a weave.Model and score it with weave.Evaluation. Changing any field on the model versions it, so two configurations are directly comparable in the Weave UI:

import asyncio

import weave
from goodmem_wandb import GoodMemRetrievalModel, RecallAtK, MRR, FactRecall, RetrievalHealth

weave.init("my-project")

dataset = [
    {
        "question": "How do I rotate an API key?",
        "expected_memory_ids": ["01a0..."],
        "expected_text": "rotate the key from the console",
    },
]

baseline = GoodMemRetrievalModel(space_name="docs", limit=5)
reranked = GoodMemRetrievalModel(space_name="docs", limit=5, reranker_id="...")

evaluation = weave.Evaluation(
    dataset=dataset,
    scorers=[RecallAtK(k=5), MRR(), FactRecall(), RetrievalHealth()],
)
asyncio.run(evaluation.evaluate(baseline))
asyncio.run(evaluation.evaluate(reranked))   # compare the two in the Weave UI

Evaluation.evaluate is a coroutine: call it with asyncio.run (or await it in a notebook), or nothing runs.

Scorer	Measures
`RecallAtK(k=5)`	Fraction of `expected_memory_ids` in the top k distinct memories (several chunks of one memory count as one document).
`MRR()`	Reciprocal rank of the first expected memory; `0.0` if none was retrieved.
`FactRecall()`	Whether the expected text actually appears in the retrieved chunks — survives re-chunking and re-embedding, unlike an id-based metric.
`RetrievalHealth()`	Whether the retrieval was complete, so a silently degraded run is visible as its own metric rather than only as a recall drop.

None of the scorers use the raw relevance number, because that number does not mean the same thing between two configurations.

Learn more

The full README, examples, and API reference are available on the PyPI package page.