Phase 2 Semantic Extraction (Ingester) Implementation Plan¶

func TestGraphPushExtractorAndFileSHAsShape(t *testing.T) {
    p := contract.GraphPush{
        Repo:      "tatara-cli",
        Commit:    "abc123",
        Extractor: "semantic",
        Files:     []string{"cmd/root.go"},
        Entities:  []contract.Entity{},
        Edges:     []contract.Edge{},
        FileSHAs:  map[string]string{"cmd/root.go": "deadbeef"},
    }
    b, err := json.Marshal(p)
    require.NoError(t, err)
    var got map[string]any
    require.NoError(t, json.Unmarshal(b, &got))
    require.ElementsMatch(t,
        []string{"repo", "commit", "extractor", "files", "entities", "edges", "file_shas"},
        keys(got))
    require.Equal(t, "semantic", got["extractor"])
    shas := got["file_shas"].(map[string]any)
    require.Equal(t, "deadbeef", shas["cmd/root.go"])
}

func TestGraphPushExtractorAndFileSHAsOmitEmpty(t *testing.T) {
    p := contract.GraphPush{Repo: "r", Files: []string{"a.go"}, Entities: []contract.Entity{}, Edges: []contract.Edge{}}
    b, err := json.Marshal(p)
    require.NoError(t, err)
    var got map[string]any
    require.NoError(t, json.Unmarshal(b, &got))
    _, hasExtractor := got["extractor"]
    _, hasSHAs := got["file_shas"]
    require.False(t, hasExtractor, "extractor must be absent when empty")
    require.False(t, hasSHAs, "file_shas must be absent when empty")
}

func TestSemanticMissesRequestShape(t *testing.T) {
    req := contract.SemanticMissesRequest{
        Repo: "r",
        Files: []contract.FileSHA{
            {Path: "a.go", ContentSHA: "sha1"},
            {Path: "b.go", ContentSHA: "sha2"},
        },
    }
    b, err := json.Marshal(req)
    require.NoError(t, err)
    var got map[string]any
    require.NoError(t, json.Unmarshal(b, &got))
    require.ElementsMatch(t, []string{"repo", "files"}, keys(got))
    first := got["files"].([]any)[0].(map[string]any)
    require.ElementsMatch(t, []string{"path", "content_sha"}, keys(first))
}

// GraphPush is one /code-graph:bulk request. Extractor tags the origin of every
// row in this push ("" or "ast" for the AST extractor, "semantic" for the LLM
// stage); reconcile scopes its per-src_file deletes by it so the two origins do
// not clobber each other. FileSHAs (path->content_sha) is set on the semantic
// push to update the server's extraction cache.
type GraphPush struct {
    Repo       string            `json:"repo"`
    Commit     string            `json:"commit,omitempty"`
    Extractor  string            `json:"extractor,omitempty"`
    Files      []string          `json:"files"`
    Entities   []Entity          `json:"entities"`
    Edges      []Edge            `json:"edges"`
    Symbols    []SymbolRow       `json:"symbols,omitempty"`
    Hyperedges []Hyperedge       `json:"hyperedges,omitempty"`
    FileSHAs   map[string]string `json:"file_shas,omitempty"`
}

// Extractor origin tags written onto every row of a GraphPush.
const (
    ExtractorAST      = "ast"
    ExtractorSemantic = "semantic"
)

// FileSHA pairs a repo-relative path with the sha256 of its working-tree content.
type FileSHA struct {
    Path       string `json:"path"`
    ContentSHA string `json:"content_sha"`
}

// SemanticMissesRequest is the POST /code-graph/semantic-misses body: the set of
// analyzed files with their current content_sha. The server returns the subset
// whose stored content_sha differs or is absent (cache miss -> needs extraction).
type SemanticMissesRequest struct {
    Repo  string    `json:"repo"`
    Files []FileSHA `json:"files"`
}

package semantic

import (
    "strings"
    "testing"

    "github.com/stretchr/testify/require"
)

func TestPromptEmbedsExtractionSpecVerbatim(t *testing.T) {
    // Anchor lines from graphify's extraction-spec prompt block must be present verbatim.
    require.Contains(t, extractionSpec, "You are a graphify extraction subagent. Read the files listed and extract a knowledge graph fragment.")
    require.Contains(t, extractionSpec, "Output ONLY valid JSON matching the schema below - no explanation, no markdown fences, no preamble.")
    require.Contains(t, extractionSpec, "- EXTRACTED: relationship explicit in source (import, call, citation, \"see §3.2\")")
    require.Contains(t, extractionSpec, "Maximum 3 hyperedges per chunk.")
    require.Contains(t, extractionSpec, "confidence_score is REQUIRED on every edge - never omit it, never use 0.5 as a default:")
    require.Contains(t, extractionSpec, "\"nodes\":[{\"id\":\"session_validatetoken\"")
}

func TestBuildPromptSubstitutesPlaceholders(t *testing.T) {
    got := BuildPrompt(PromptVars{
        FileList:    "- a.go\n- b.go",
        ChunkNum:    2,
        TotalChunks: 5,
        ChunkPath:   "/tmp/chunk-2.json",
    })
    require.Contains(t, got, "Files (chunk 2 of 5):")
    require.Contains(t, got, "- a.go\n- b.go")
    require.Contains(t, got, "/tmp/chunk-2.json")
    // Placeholders must be fully consumed.
    require.NotContains(t, got, "FILE_LIST")
    require.NotContains(t, got, "CHUNK_NUM")
    require.NotContains(t, got, "TOTAL_CHUNKS")
    require.NotContains(t, got, "CHUNK_PATH")
    // DEEP_MODE is off: the literal token must not leak into the prompt.
    require.NotContains(t, got, "DEEP_MODE (if --mode deep was given)")
}

func TestBuildPromptKeepsSchemaIntact(t *testing.T) {
    got := BuildPrompt(PromptVars{FileList: "- a.go", ChunkNum: 1, TotalChunks: 1, ChunkPath: "/tmp/c.json"})
    require.Contains(t, got, "Generate the extraction JSON matching this schema exactly:")
    require.Contains(t, got, "\"semantically_similar_to\"")
}

// Package semantic runs the LLM extraction stage: it chunks analyzed files,
// builds the graphify extraction prompt, and maps the returned JSON fragment to
// contract graph types. It is best-effort: callers log and skip on any failure.
package semantic

import (
    _ "embed"
    "strconv"
    "strings"
)

//go:embed extraction_spec.txt
var extractionSpec string

// PromptVars are the placeholder substitutions for the extraction prompt.
type PromptVars struct {
    FileList    string
    ChunkNum    int
    TotalChunks int
    ChunkPath   string
}

// BuildPrompt returns the verbatim extraction-spec prompt with FILE_LIST,
// CHUNK_NUM, TOTAL_CHUNKS, and CHUNK_PATH substituted. DEEP_MODE is always off.
func BuildPrompt(v PromptVars) string {
    r := strings.NewReplacer(
        "CHUNK_NUM", strconv.Itoa(v.ChunkNum),
        "TOTAL_CHUNKS", strconv.Itoa(v.TotalChunks),
        "FILE_LIST", v.FileList,
        "CHUNK_PATH", v.ChunkPath,
    )
    return r.Replace(extractionSpec)
}

package semantic

import (
    "testing"

    "github.com/stretchr/testify/require"
)

func TestChunkGroupsByFileCount(t *testing.T) {
    files := make([]LoadedFile, 0, 20)
    for i := 0; i < 20; i++ {
        files = append(files, LoadedFile{Path: "f" + string(rune('a'+i)) + ".go", Content: "tiny"})
    }
    chunks := Chunk(files, ChunkBudget{MaxFiles: 8, MaxBytes: 1 << 20})
    require.Len(t, chunks, 3) // 8 + 8 + 4
    require.Len(t, chunks[0].Files, 8)
    require.Len(t, chunks[1].Files, 8)
    require.Len(t, chunks[2].Files, 4)
}

func TestChunkSplitsOnByteBudget(t *testing.T) {
    big := make([]byte, 600)
    for i := range big {
        big[i] = 'x'
    }
    files := []LoadedFile{
        {Path: "a.go", Content: string(big)},
        {Path: "b.go", Content: string(big)},
        {Path: "c.go", Content: string(big)},
    }
    // MaxBytes 1000 admits one 600-byte file per chunk (a second would exceed it).
    chunks := Chunk(files, ChunkBudget{MaxFiles: 8, MaxBytes: 1000})
    require.Len(t, chunks, 3)
    for _, c := range chunks {
        require.Len(t, c.Files, 1)
    }
}

func TestChunkOversizeFileGetsItsOwnChunk(t *testing.T) {
    big := make([]byte, 2000)
    files := []LoadedFile{
        {Path: "small.go", Content: "x"},
        {Path: "huge.go", Content: string(big)}, // exceeds MaxBytes alone
    }
    chunks := Chunk(files, ChunkBudget{MaxFiles: 8, MaxBytes: 1000})
    require.Len(t, chunks, 2)
    require.Equal(t, "small.go", chunks[0].Files[0].Path)
    require.Equal(t, "huge.go", chunks[1].Files[0].Path)
}

func TestChunkEmptyInputYieldsNoChunks(t *testing.T) {
    require.Empty(t, Chunk(nil, ChunkBudget{MaxFiles: 8, MaxBytes: 1000}))
}

func TestDefaultChunkBudget(t *testing.T) {
    b := DefaultChunkBudget()
    require.Equal(t, 8, b.MaxFiles)
    require.Greater(t, b.MaxBytes, 0)
}

package semantic

// LoadedFile is one analyzed file with its content, ready to be chunked.
type LoadedFile struct {
    Path    string
    Content string
}

// FileChunk is a token/size-bounded group of files extracted together.
type FileChunk struct {
    Files []LoadedFile
}

// ChunkBudget bounds a single chunk: at most MaxFiles files and MaxBytes of
// content. A single file larger than MaxBytes still gets its own chunk.
type ChunkBudget struct {
    MaxFiles int
    MaxBytes int
}

// DefaultChunkBudget is ~8 files and ~48KB of content per chunk (a rough proxy
// for the gpt-4o-mini context budget; ~4 bytes/token).
func DefaultChunkBudget() ChunkBudget {
    return ChunkBudget{MaxFiles: 8, MaxBytes: 48 * 1024}
}

// Chunk groups files greedily, closing a chunk when adding the next file would
// exceed either bound. Order is preserved.
func Chunk(files []LoadedFile, b ChunkBudget) []FileChunk {
    if b.MaxFiles <= 0 {
        b.MaxFiles = 8
    }
    if b.MaxBytes <= 0 {
        b.MaxBytes = 48 * 1024
    }
    var chunks []FileChunk
    var cur []LoadedFile
    curBytes := 0
    flush := func() {
        if len(cur) > 0 {
            chunks = append(chunks, FileChunk{Files: cur})
            cur = nil
            curBytes = 0
        }
    }
    for _, f := range files {
        n := len(f.Content)
        if len(cur) > 0 && (len(cur) >= b.MaxFiles || curBytes+n > b.MaxBytes) {
            flush()
        }
        cur = append(cur, f)
        curBytes += n
    }
    flush()
    return chunks
}

package semantic

import (
    "testing"

    "github.com/stretchr/testify/require"

    "github.com/szymonrychu/tatara-memory-repo-ingester/internal/contract"
)

const sampleFragment = `{
  "nodes": [
    {"id":"auth_session_validatetoken","label":"Validate Token","file_type":"code","source_file":"auth/session.go"},
    {"id":"retry_backoff","label":"Retry Backoff","file_type":"concept","source_file":"http/client.go"},
    {"id":"why_jitter","label":"Why Jitter","file_type":"rationale","source_file":"http/client.go"}
  ],
  "edges": [
    {"source":"auth_session_validatetoken","target":"retry_backoff","relation":"semantically_similar_to","confidence":"INFERRED","confidence_score":0.85,"source_file":"http/client.go","weight":1.0},
    {"source":"why_jitter","target":"retry_backoff","relation":"rationale_for","confidence":"EXTRACTED","confidence_score":1.0,"source_file":"http/client.go","weight":1.0},
    {"source":"a","target":"b","relation":"calls","confidence":"AMBIGUOUS","confidence_score":0.2,"source_file":"http/client.go"}
  ],
  "hyperedges": [
    {"id":"auth_flow","label":"Auth Flow","nodes":["auth_session_validatetoken","retry_backoff","why_jitter"],"relation":"participate_in","confidence":"INFERRED","confidence_score":0.75,"source_file":"auth/session.go"}
  ],
  "input_tokens": 0,
  "output_tokens": 0
}`

func TestParseFragmentConceptNodeIDs(t *testing.T) {
    res, err := ParseFragment("myrepo", []byte(sampleFragment))
    require.NoError(t, err)
    // code file_type nodes are NOT emitted as entities (they reference AST nodes).
    // concept/rationale file_type nodes ARE emitted with deterministic ids.
    byType := map[string]contract.Entity{}
    for _, e := range res.Entities {
        byType[e.Type] = e
    }
    require.Contains(t, byType, contract.EntityConcept)
    require.Contains(t, byType, contract.EntityRationale)
    require.Equal(t, "concept:myrepo:retry-backoff", byType[contract.EntityConcept].ID)
    require.Equal(t, "concept:myrepo:why-jitter", byType[contract.EntityRationale].ID)
    for _, e := range res.Entities {
        require.Equal(t, "http/client.go", e.FilePath)
    }
}

func TestParseFragmentEdgesMapToContract(t *testing.T) {
    res, err := ParseFragment("myrepo", []byte(sampleFragment))
    require.NoError(t, err)
    require.Len(t, res.Edges, 3)
    var sim contract.Edge
    for _, e := range res.Edges {
        if e.Relation == contract.RelSemanticallySimilar {
            sim = e
        }
    }
    require.Equal(t, contract.RelSemanticallySimilar, sim.Relation)
    require.InDelta(t, 0.85, sim.ConfidenceScore, 1e-9)
    require.Equal(t, contract.TierInferred, sim.ConfidenceTier)
    require.Equal(t, "http/client.go", sim.SrcFile)
}

func TestParseFragmentConfidenceTiers(t *testing.T) {
    res, err := ParseFragment("myrepo", []byte(sampleFragment))
    require.NoError(t, err)
    tiers := map[string]string{}
    for _, e := range res.Edges {
        tiers[e.Relation] = e.ConfidenceTier
    }
    require.Equal(t, contract.TierInferred, tiers[contract.RelSemanticallySimilar]) // 0.85
    require.Equal(t, contract.TierExtracted, tiers[contract.RelRationaleFor])       // 1.0
    require.Equal(t, contract.TierAmbiguous, tiers[contract.RelCalls])              // 0.2
}

func TestParseFragmentHyperedges(t *testing.T) {
    res, err := ParseFragment("myrepo", []byte(sampleFragment))
    require.NoError(t, err)
    require.Len(t, res.Hyperedges, 1)
    h := res.Hyperedges[0]
    require.Equal(t, "Auth Flow", h.Label)
    require.Equal(t, "participate_in", h.Relation)
    require.InDelta(t, 0.75, h.ConfidenceScore, 1e-9)
    require.Equal(t, "auth/session.go", h.SrcFile)
    require.Len(t, h.Members, 3)
    require.Equal(t, "he:myrepo:auth/session.go:auth_flow", h.ID)
}

func TestParseFragmentHyperedgesCappedAtThree(t *testing.T) {
    frag := `{"nodes":[],"edges":[],"hyperedges":[
      {"id":"h1","label":"a","nodes":["x","y","z"],"relation":"form","confidence_score":0.7,"source_file":"f.go"},
      {"id":"h2","label":"b","nodes":["x","y","z"],"relation":"form","confidence_score":0.7,"source_file":"f.go"},
      {"id":"h3","label":"c","nodes":["x","y","z"],"relation":"form","confidence_score":0.7,"source_file":"f.go"},
      {"id":"h4","label":"d","nodes":["x","y","z"],"relation":"form","confidence_score":0.7,"source_file":"f.go"}
    ]}`
    res, err := ParseFragment("r", []byte(frag))
    require.NoError(t, err)
    require.Len(t, res.Hyperedges, 3)
}

func TestParseFragmentRejectsMalformedJSON(t *testing.T) {
    _, err := ParseFragment("r", []byte("not json"))
    require.Error(t, err)
}

func TestParseFragmentStripsCodeFences(t *testing.T) {
    wrapped := "```json\n" + sampleFragment + "\n```"
    res, err := ParseFragment("myrepo", []byte(wrapped))
    require.NoError(t, err)
    require.Len(t, res.Edges, 3)
}

func TestSlugLabel(t *testing.T) {
    require.Equal(t, "retry-backoff", slugLabel("Retry Backoff"))
    require.Equal(t, "why-jitter", slugLabel("Why Jitter!"))
    require.Equal(t, "a-b-c", slugLabel("a  b   c"))
    require.Equal(t, "leadingtrailing", slugLabel("  LeadingTrailing  "))
}

package semantic

import (
    "encoding/json"
    "fmt"
    "strings"

    "github.com/szymonrychu/tatara-memory-repo-ingester/internal/analyze"
    "github.com/szymonrychu/tatara-memory-repo-ingester/internal/contract"
)

// rawFragment mirrors graphify's extraction JSON schema.
type rawFragment struct {
    Nodes []struct {
        ID         string `json:"id"`
        Label      string `json:"label"`
        FileType   string `json:"file_type"`
        SourceFile string `json:"source_file"`
        SourceURL  string `json:"source_url"`
        CapturedAt string `json:"captured_at"`
        Author     string `json:"author"`
    } `json:"nodes"`
    Edges []struct {
        Source          string  `json:"source"`
        Target          string  `json:"target"`
        Relation        string  `json:"relation"`
        Confidence      string  `json:"confidence"`
        ConfidenceScore float64 `json:"confidence_score"`
        SourceFile      string  `json:"source_file"`
    } `json:"edges"`
    Hyperedges []struct {
        ID              string   `json:"id"`
        Label           string   `json:"label"`
        Nodes           []string `json:"nodes"`
        Relation        string   `json:"relation"`
        ConfidenceScore float64  `json:"confidence_score"`
        SourceFile      string   `json:"source_file"`
    } `json:"hyperedges"`
}

// maxHyperedgesPerChunk matches the extraction-spec cap.
const maxHyperedgesPerChunk = 3

// ParseFragment maps a graphify extraction JSON fragment to contract graph
// types for one repo. Concept/rationale nodes become Entities with deterministic
// ids (concept:<repo>:<slug>); code/document/paper/image nodes reference AST
// entity ids and are not re-emitted. Edges carry the semantic relation, score,
// and tier. Hyperedges are capped at 3 with deterministic ids.
func ParseFragment(repo string, body []byte) (analyze.Result, error) {
    var f rawFragment
    if err := json.Unmarshal(stripFences(body), &f); err != nil {
        return analyze.Result{}, fmt.Errorf("parse extraction fragment: %w", err)
    }
    var res analyze.Result
    for _, n := range f.Nodes {
        typ := entityTypeFor(n.FileType)
        if typ == "" {
            continue // code/document/paper/image: references an AST node, not re-emitted
        }
        res.Entities = append(res.Entities, contract.Entity{
            ID:         conceptID(repo, n.Label),
            Name:       n.Label,
            Type:       typ,
            FilePath:   n.SourceFile,
            SourceURL:  n.SourceURL,
            Author:     n.Author,
            CapturedAt: n.CapturedAt,
        })
    }
    for _, e := range f.Edges {
        res.Edges = append(res.Edges, contract.Edge{
            From:            e.Source,
            To:              e.Target,
            Relation:        e.Relation,
            SrcFile:         e.SourceFile,
            ConfidenceScore: e.ConfidenceScore,
            ConfidenceTier:  contract.TierForScore(e.ConfidenceScore),
        })
    }
    for i, h := range f.Hyperedges {
        if i >= maxHyperedgesPerChunk {
            break
        }
        res.Hyperedges = append(res.Hyperedges, contract.Hyperedge{
            ID:              fmt.Sprintf("he:%s:%s:%s", repo, h.SourceFile, h.ID),
            Label:           h.Label,
            Relation:        h.Relation,
            ConfidenceScore: h.ConfidenceScore,
            SrcFile:         h.SourceFile,
            Members:         h.Nodes,
        })
    }
    return res, nil
}

// entityTypeFor maps a graphify file_type to a contract concept/rationale entity
// type, or "" for node types that reference an existing AST entity.
func entityTypeFor(fileType string) string {
    switch fileType {
    case "concept":
        return contract.EntityConcept
    case "rationale":
        return contract.EntityRationale
    default:
        return ""
    }
}

// conceptID is the deterministic id for a concept/rationale node: a slug of the
// label scoped to the repo. Re-extraction of the same label upserts, not dupes.
func conceptID(repo, label string) string {
    return "concept:" + repo + ":" + slugLabel(label)
}

// slugLabel lowercases a label and collapses runs of non-[a-z0-9] into single
// hyphens, trimming leading/trailing hyphens.
func slugLabel(label string) string {
    var b strings.Builder
    prevHyphen := false
    for _, r := range strings.ToLower(label) {
        switch {
        case (r >= 'a' && r <= 'z') || (r >= '0' && r <= '9'):
            b.WriteRune(r)
            prevHyphen = false
        default:
            if !prevHyphen {
                b.WriteByte('-')
                prevHyphen = true
            }
        }
    }
    return strings.Trim(b.String(), "-")
}

// stripFences removes a leading ```json / ``` fence and a trailing ``` fence if
// the model wrapped its JSON despite instructions to the contrary.
func stripFences(body []byte) []byte {
    s := strings.TrimSpace(string(body))
    s = strings.TrimPrefix(s, "```json")
    s = strings.TrimPrefix(s, "```")
    s = strings.TrimSuffix(s, "```")
    return []byte(strings.TrimSpace(s))
}

package llm

import (
    "context"
    "encoding/json"
    "io"
    "net/http"
    "net/http/httptest"
    "sync/atomic"
    "testing"

    "github.com/stretchr/testify/require"
)

func TestCompleteSendsJSONModeRequest(t *testing.T) {
    var gotBody map[string]any
    var gotAuth string
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        require.Equal(t, "/chat/completions", r.URL.Path)
        gotAuth = r.Header.Get("Authorization")
        b, _ := io.ReadAll(r.Body)
        _ = json.Unmarshal(b, &gotBody)
        w.WriteHeader(200)
        _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"{\"ok\":true}"}}]}`))
    }))
    defer srv.Close()

    c := New(Config{APIKey: "sk-test", Model: "gpt-4o-mini", BaseURL: srv.URL}, http.DefaultClient)
    out, err := c.Complete(context.Background(), "do the thing")
    require.NoError(t, err)
    require.Equal(t, `{"ok":true}`, out)
    require.Equal(t, "Bearer sk-test", gotAuth)
    require.Equal(t, "gpt-4o-mini", gotBody["model"])
    rf := gotBody["response_format"].(map[string]any)
    require.Equal(t, "json_object", rf["type"])
    msgs := gotBody["messages"].([]any)
    require.GreaterOrEqual(t, len(msgs), 1)
    last := msgs[len(msgs)-1].(map[string]any)
    require.Equal(t, "user", last["role"])
    require.Equal(t, "do the thing", last["content"])
}

func TestCompleteRetriesOnce5xxThenSucceeds(t *testing.T) {
    var calls int32
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if atomic.AddInt32(&calls, 1) == 1 {
            w.WriteHeader(503)
            _, _ = w.Write([]byte(`{"error":"overloaded"}`))
            return
        }
        w.WriteHeader(200)
        _, _ = w.Write([]byte(`{"choices":[{"message":{"content":"{\"ok\":1}"}}]}`))
    }))
    defer srv.Close()

    c := New(Config{APIKey: "k", Model: "m", BaseURL: srv.URL}, http.DefaultClient)
    out, err := c.Complete(context.Background(), "x")
    require.NoError(t, err)
    require.Equal(t, `{"ok":1}`, out)
    require.Equal(t, int32(2), atomic.LoadInt32(&calls))
}

func TestCompleteRetriesOnce429ThenFails(t *testing.T) {
    var calls int32
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        atomic.AddInt32(&calls, 1)
        w.WriteHeader(429)
        _, _ = w.Write([]byte(`{"error":"rate limited"}`))
    }))
    defer srv.Close()

    c := New(Config{APIKey: "k", Model: "m", BaseURL: srv.URL}, http.DefaultClient)
    _, err := c.Complete(context.Background(), "x")
    require.Error(t, err)
    require.Equal(t, int32(2), atomic.LoadInt32(&calls), "one initial + one retry only")
}

func TestCompleteDoesNotRetryOn400(t *testing.T) {
    var calls int32
    srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        atomic.AddInt32(&calls, 1)
        w.WriteHeader(400)
        _, _ = w.Write([]byte(`{"error":"bad request"}`))
    }))
    defer srv.Close()

    c := New(Config{APIKey: "k", Model: "m", BaseURL: srv.URL}, http.DefaultClient)
    _, err := c.Complete(context.Background(), "x")
    require.Error(t, err)
    require.Equal(t, int32(1), atomic.LoadInt32(&calls), "4xx (non-429) must not retry")
}

func TestConfigFromEnvDefaults(t *testing.T) {
    cfg := ConfigFromEnv(func(k string) string {
        switch k {
        case "OPENAI_API_KEY":
            return "sk-xyz"
        default:
            return ""
        }
    })
    require.Equal(t, "sk-xyz", cfg.APIKey)
    require.Equal(t, "gpt-4o-mini", cfg.Model)
    require.Equal(t, "https://api.openai.com/v1", cfg.BaseURL)
}

func TestConfigFromEnvOverrides(t *testing.T) {
    cfg := ConfigFromEnv(func(k string) string {
        switch k {
        case "OPENAI_API_KEY":
            return "sk-1"
        case "SEMANTIC_MODEL":
            return "gpt-4o"
        case "OPENAI_BASE_URL":
            return "http://localhost:1234/v1/"
        default:
            return ""
        }
    })
    require.Equal(t, "gpt-4o", cfg.Model)
    require.Equal(t, "http://localhost:1234/v1", cfg.BaseURL, "trailing slash trimmed")
}

// Package llm is a minimal OpenAI chat/completions client used by the semantic
// extraction stage. It requests JSON mode and retries once on transient errors.
package llm

import (
    "bytes"
    "context"
    "encoding/json"
    "fmt"
    "io"
    "net/http"
    "strings"
    "time"
)

// Config holds the OpenAI client configuration.
type Config struct {
    APIKey  string
    Model   string
    BaseURL string
}

// ConfigFromEnv reads OPENAI_API_KEY, SEMANTIC_MODEL (default gpt-4o-mini), and
// OPENAI_BASE_URL (default https://api.openai.com/v1, trailing slash trimmed).
func ConfigFromEnv(getenv func(string) string) Config {
    model := getenv("SEMANTIC_MODEL")
    if model == "" {
        model = "gpt-4o-mini"
    }
    base := getenv("OPENAI_BASE_URL")
    if base == "" {
        base = "https://api.openai.com/v1"
    }
    return Config{
        APIKey:  getenv("OPENAI_API_KEY"),
        Model:   model,
        BaseURL: strings.TrimRight(base, "/"),
    }
}

// Client posts chat/completions to an OpenAI-compatible endpoint.
type Client struct {
    cfg  Config
    http *http.Client
}

// New constructs an OpenAI client.
func New(cfg Config, hc *http.Client) *Client {
    return &Client{cfg: cfg, http: hc}
}

type chatRequest struct {
    Model          string            `json:"model"`
    Messages       []chatMessage     `json:"messages"`
    ResponseFormat map[string]string `json:"response_format"`
    Temperature    float64           `json:"temperature"`
}

type chatMessage struct {
    Role    string `json:"role"`
    Content string `json:"content"`
}

type chatResponse struct {
    Choices []struct {
        Message struct {
            Content string `json:"content"`
        } `json:"message"`
    } `json:"choices"`
}

// Complete sends a single user prompt in JSON mode and returns the message
// content. It retries once on 429/5xx with a short backoff.
func (c *Client) Complete(ctx context.Context, prompt string) (string, error) {
    reqBody := chatRequest{
        Model:          c.cfg.Model,
        Messages:       []chatMessage{{Role: "user", Content: prompt}},
        ResponseFormat: map[string]string{"type": "json_object"},
        Temperature:    0,
    }
    b, err := json.Marshal(reqBody)
    if err != nil {
        return "", fmt.Errorf("marshal chat request: %w", err)
    }

    var lastErr error
    for attempt := 0; attempt < 2; attempt++ {
        if attempt > 0 {
            select {
            case <-ctx.Done():
                return "", ctx.Err()
            case <-time.After(500 * time.Millisecond):
            }
        }
        content, retry, err := c.try(ctx, b)
        if err == nil {
            return content, nil
        }
        lastErr = err
        if !retry {
            return "", err
        }
    }
    return "", lastErr
}

// try performs one request. retry is true only for transient (429/5xx) failures.
func (c *Client) try(ctx context.Context, body []byte) (content string, retry bool, err error) {
    req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.cfg.BaseURL+"/chat/completions", bytes.NewReader(body))
    if err != nil {
        return "", false, fmt.Errorf("request: %w", err)
    }
    req.Header.Set("Content-Type", "application/json")
    req.Header.Set("Authorization", "Bearer "+c.cfg.APIKey)

    resp, err := c.http.Do(req)
    if err != nil {
        return "", true, fmt.Errorf("call openai: %w", err)
    }
    defer func() { _ = resp.Body.Close() }()

    if resp.StatusCode != http.StatusOK {
        b, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
        transient := resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode >= 500
        return "", transient, fmt.Errorf("openai status %d: %s", resp.StatusCode, string(b))
    }

    var cr chatResponse
    if err := json.NewDecoder(resp.Body).Decode(&cr); err != nil {
        return "", false, fmt.Errorf("decode openai response: %w", err)
    }
    if len(cr.Choices) == 0 {
        return "", false, fmt.Errorf("openai response has no choices")
    }
    return cr.Choices[0].Message.Content, false, nil
}