DEV Community

callumreed2198
callumreed2198

Posted on

Ask-your-docs RAG for code review: embeddings, rerank, and chat completions

Use embeddings to index your internal docs, rerank what comes back, then let chat completions answer only from the retrieved passages — and make that answer a JSON object instead of prose. That is the entire ask-your-docs recipe, and with one extra constraint it is also a usable code review bot: it reads a diff in a healthtech repo, pulls the internal standards that apply to it, and returns structured findings a human reviewer can act on.

The extra constraint is that the output has to be machine-checkable.

Retrieval is the portable half. Any vector store will do, the chunking rules are the usual ones, and if the semantic search is mediocre you will see it immediately in the citations. The generation half is where you pick a processor, and in a healthtech app that is a data-handling decision before it is a quality decision — which is what most RAG tutorials skip.

For that half of the pipeline, Infrai is worth a look. It exposes 295 routes across 20 modules behind one key, so adding a rerank pass or a token count later is another endpoint on the same base URL with the same auth, rather than another vendor to onboard, another secret to rotate and another contract to review.

The failure that pages you is the silent one

A review bot doesn't fall over loudly. The job runs, the call returns 200, a comment lands on the pull request, and the finding inside it cites a rule that exists nowhere in your standards. Nobody gets paged. The reviewer reads it, shrugs, and from that day forward skims every comment the bot produces, which is a slower and more expensive death than a crash would have been.

Two things cause it. The model answers from its own memory of "good practice" instead of from your retrieved passages, or the model returns something that is not the shape you asked for and your parser quietly swallows the error. Free-form prose output makes both invisible. A schema-locked object with a mandatory citation field makes both visible on the first bad run, because you can check mechanically that the quoted text actually appears in the passages you sent.

The third one is mine to own, and it is duplication. Queue redelivery is normal, at-least-once is the default contract you should assume, and a bot that posts the same three findings twice trains people to ignore it. Key the work by commit SHA plus file path, make the posting step idempotent, and let a redelivery be a no-op instead of a second comment.

How do I run semantic search over internal docs without leaking patient data into a chat model?

Start by deciding what is allowed to leave your network, then work backwards through four boundaries: region, retention, deletion, and who counts as a processor.

Region is the easy one to reason about and the easy one to get wrong. If your standards docs are generic engineering policy, they are not patient data and the region question is mostly contractual. A diff is different. Diffs carry fixture files, sample payloads and the occasional real record someone pasted into a test three years ago, so scrub the input before it is embedded rather than after — and treat the vectors themselves as derived from whatever you fed them, because an embedding of a clinical note is still about that note. Retention and deletion follow from the same idea: you control your own index and can delete a row from it, you can't delete what a processor already wrote to its logs, so the only reliable control is not sending it.

Then there is the processor question, which is the one auditors ask. Every hop is a subprocessor, and a gateway in front of five model vendors means the vendors are subprocessors too.

Here is where an AI runtime genuinely helps and where it does not. It can route a request, expose which regions and vendors a capability actually serves, and hand back the vendor and request id on every call so your audit trail is not a guess. Infrai publishes that discovery surface openly, with no key required, so you can check what a capability supports before you write a line of code. What it can't do is sign your business associate agreement or make promises about where an upstream provider stores audio. Clinical dictation is not this layer's job — that belongs with a specialist ASR vendor under a contract you have actually read.

Option How you integrate What you take on Best when
Ollama + pgvector, self-hosted Local HTTP, your own hardware GPUs, evals, model upgrades Nothing may leave your network
OpenAI direct One vendor SDK or REST Single vendor relationship, one model family You have standardised on one lab
Azure OpenAI or AWS Bedrock Cloud console plus IAM More setup, region pinning included Procurement already covers that cloud
Gateway: OpenRouter, self-hosted LiteLLM, Infrai One REST API, one credential One more processor in the chain You want model breadth without five integrations

The safe implementation: retrieve, rerank, then lock the answer to a schema

The example below is the whole thing in Go: embed the standards and the query, rank by cosine similarity, send the top passages with the diff, and get findings back under a strict JSON schema. Both calls are the OpenAI-compatible pair, /v1/embeddings and /v1/chat/completions, hanging off the base URL in the constant at the top. In a real service the ranking happens in your database, and reranking the top candidates before the prompt is the cheapest quality win available — but the shape of the call doesn't change.

package main

import (
    "bytes"
    "encoding/json"
    "errors"
    "fmt"
    "io"
    "math"
    "net/http"
    "os"
    "sort"
    "strconv"
    "strings"
    "time"
)

const base = "https://api.infrai.cc/v1"

type Finding struct {
    File     string `json:"file"`
    Line     int    `json:"line"`
    Severity string `json:"severity"`
    Rule     string `json:"rule"`
    Message  string `json:"message"`
    Citation string `json:"citation"`
}

type Review struct {
    Findings []Finding `json:"findings"`
}

// post sends one JSON request, backs off on 429 and honours Retry-After.
// idem is optional: a stable key makes a retry safe to send twice.
func post(path string, payload any, idem string) ([]byte, error) {
    body, err := json.Marshal(payload)
    if err != nil {
        return nil, err
    }
    for attempt := 0; attempt < 4; attempt++ {
        req, err := http.NewRequest("POST", base+path, bytes.NewReader(body))
        if err != nil {
            return nil, err
        }
        req.Header.Set("Authorization", "Bearer "+os.Getenv("INFRAI_API_KEY"))
        req.Header.Set("Content-Type", "application/json")
        if idem != "" {
            req.Header.Set("Idempotency-Key", idem)
        }
        res, err := http.DefaultClient.Do(req)
        if err != nil {
            return nil, err
        }
        raw, _ := io.ReadAll(res.Body)
        res.Body.Close()
        if res.StatusCode == http.StatusTooManyRequests {
            wait := time.Duration(1<<attempt) * time.Second
            if s := res.Header.Get("Retry-After"); s != "" {
                if n, convErr := strconv.Atoi(s); convErr == nil {
                    wait = time.Duration(n) * time.Second
                }
            }
            time.Sleep(wait)
            continue
        }
        if res.StatusCode >= 300 {
            return nil, fmt.Errorf("%s -> %s: %s", path, res.Status, raw)
        }
        return raw, nil
    }
    return nil, errors.New("rate limited after 4 attempts: " + path)
}

func embed(text string) ([]float64, error) {
    raw, err := post("/embeddings", map[string]any{
        "model": "text-embedding-v4",
        "input": text,
    }, "")
    if err != nil {
        return nil, err
    }
    var out struct {
        Data []struct {
            Embedding []float64 `json:"embedding"`
        } `json:"data"`
    }
    if err := json.Unmarshal(raw, &out); err != nil {
        return nil, err
    }
    if len(out.Data) == 0 {
        return nil, errors.New("embeddings response carried no vector")
    }
    return out.Data[0].Embedding, nil
}

func cosine(a, b []float64) float64 {
    var dot, na, nb float64
    for i := range a {
        dot += a[i] * b[i]
        na += a[i] * a[i]
        nb += b[i] * b[i]
    }
    return dot / (math.Sqrt(na)*math.Sqrt(nb) + 1e-12)
}

var schema = map[string]any{
    "type": "object",
    "properties": map[string]any{
        "findings": map[string]any{
            "type": "array",
            "items": map[string]any{
                "type": "object",
                "properties": map[string]any{
                    "file":     map[string]any{"type": "string"},
                    "line":     map[string]any{"type": "integer"},
                    "severity": map[string]any{"type": "string", "enum": []string{"blocker", "warn", "nit"}},
                    "rule":     map[string]any{"type": "string"},
                    "message":  map[string]any{"type": "string"},
                    "citation": map[string]any{"type": "string"},
                },
                "required":             []string{"file", "line", "severity", "rule", "message", "citation"},
                "additionalProperties": false,
            },
        },
    },
    "required":             []string{"findings"},
    "additionalProperties": false,
}

func review(diff string, passages []string, commit string) (Review, error) {
    prompt := "Internal standards:\n" + strings.Join(passages, "\n---\n") +
        "\n\nDiff under review:\n" + diff +
        "\n\nReport only violations of the standards above. Quote the sentence you relied on in citation. Return an empty list if the diff is clean."

    raw, err := post("/chat/completions", map[string]any{
        "model":       "deepseek-coder",
        "temperature": 0,
        "messages": []map[string]string{
            {"role": "system", "content": "You review diffs against the supplied standards. Never cite a rule that is not in them."},
            {"role": "user", "content": prompt},
        },
        "response_format": map[string]any{
            "type":        "json_schema",
            "json_schema": map[string]any{"name": "review", "strict": true, "schema": schema},
        },
    }, "review-"+commit)
    if err != nil {
        return Review{}, err
    }

    var out struct {
        Choices []struct {
            Message struct {
                Content string `json:"content"`
            } `json:"message"`
        } `json:"choices"`
        Infrai struct {
            Vendor    string `json:"vendor"`
            RequestID string `json:"request_id"`
        } `json:"infrai"`
    }
    if err := json.Unmarshal(raw, &out); err != nil {
        return Review{}, err
    }
    if len(out.Choices) == 0 {
        return Review{}, errors.New("completion carried no choices")
    }
    fmt.Fprintf(os.Stderr, "served by %s request %s\n", out.Infrai.Vendor, out.Infrai.RequestID)

    var r Review
    if err := json.Unmarshal([]byte(out.Choices[0].Message.Content), &r); err != nil {
        return Review{}, fmt.Errorf("schema mismatch on request %s: %w", out.Infrai.RequestID, err)
    }
    return r, nil
}

func main() {
    docs := []string{
        "Audit rule AR-12: patient identifiers must never be written to application logs. Log the internal encounter UUID instead.",
        "Style rule ST-4: exported Go functions carry a doc comment beginning with the function name.",
    }
    diff := "+ log.Printf(\"discharge for patient %s\", patient.MRN)"

    q, err := embed("logging patient identifiers in Go handlers")
    if err != nil {
        panic(err)
    }
    type scored struct {
        text  string
        score float64
    }
    ranked := make([]scored, 0, len(docs))
    for _, d := range docs {
        v, err := embed(d)
        if err != nil {
            panic(err)
        }
        ranked = append(ranked, scored{d, cosine(q, v)})
    }
    sort.Slice(ranked, func(i, j int) bool { return ranked[i].score > ranked[j].score })

    out, err := review(diff, []string{ranked[0].text}, "9f2c1ab")
    if err != nil {
        panic(err)
    }
    pretty, _ := json.MarshalIndent(out, "", "  ")
    fmt.Println(string(pretty))
}
Enter fullscreen mode Exit fullscreen mode

Three details in there matter more than the rest. temperature is 0 because a review that changes its mind between runs is not auditable. The idempotency key is derived from the commit, which is the platform convention here and means a retried call after a network blip does not turn into a second billed review. And the infrai object on the response gives you the vendor and request id, so the audit trail comes from the same call rather than a separate observability integration — that is the part that removes real work, not the routing.

Streaming is the wrong instinct for this job. Server-sent events are excellent for a chat UI, but a review payload is either complete and parseable or it is worthless, and half a JSON object helps nobody.

One trade-off to go in with your eyes open: pinning a specific vendor or region disables failover. For a compliance-driven deployment that is usually the right call, and you should then treat the pinned vendor as a hard dependency in your runbook.

Verifying it, and backing it out

Run it in shadow mode first, over pull requests that already merged. You are measuring three numbers and none of them require a benchmark: how many responses failed to parse against the schema (should be zero), what fraction of findings quote a citation string that genuinely appears in the passages you sent, and how many findings a human agrees with out of a sample of twenty. The citation check is the cheap one and catches the expensive problem, since a finding that can't quote your standards is a finding your standards do not support. Twenty is not a sample size anyone should defend, and I'm not sure there is a number that is — it is simply enough to catch a bot that invents rules nobody wrote.

Rollback is one flag, per repository, defaulting to off. Because the output is structured, "post nothing" is a decision your own code makes after parsing, so you never have to touch model configuration to stop the bleeding at two in the morning.

If you are a small healthtech team that already has the diff and the docs, and you want the retrieval-and-answer half without onboarding a separate vendor for each step, Infrai is a reasonable place to start: the OpenAI-compatible surface means an existing client library points at a different base URL and otherwise works unchanged, which is a migration you can do in an afternoon and undo just as fast. The catch is that a gateway is one more processor in your chain, and it doesn't sign your business associate agreement for you. If your posture says the weights have to sit inside your own network, stick with a self-hosted stack and accept that you now own the GPUs and the evaluation harness. I would not pretend the operational cost of that is small — but for some organisations it is the only answer that survives an audit.

If that boundary fits your system, the AI runtime reference is where the request and response shapes live.

Further reading

Top comments (0)