347 lines
11 KiB
Go
347 lines
11 KiB
Go
package pipeline
|
||
|
||
import (
|
||
"context"
|
||
"crypto/sha256"
|
||
"encoding/hex"
|
||
"encoding/json"
|
||
"fmt"
|
||
"log/slog"
|
||
"os"
|
||
"strings"
|
||
"time"
|
||
|
||
"textmachine/backend/internal/config"
|
||
"textmachine/backend/internal/ledger"
|
||
"textmachine/backend/internal/llm"
|
||
"textmachine/backend/internal/obs"
|
||
"textmachine/backend/internal/store"
|
||
)
|
||
|
||
// runner.go: the Phase-0 mini-runner — one chunk through the configured stage
|
||
// list (C1: draft→edit) with full money discipline: reserve → call → settle+
|
||
// checkpoint (одна транзакция) → request_log. Resume: найденный чекпоинт =
|
||
// вызов не повторяется и не оплачивается. Гейты/эскалация/циклы по главам —
|
||
// Фаза 1, сюда же (код раннера).
|
||
|
||
// Runner executes a book's pipeline.
|
||
type Runner struct {
|
||
Book *config.Book
|
||
Models *config.Models
|
||
Pipeline *config.Pipeline
|
||
Store *store.Store
|
||
Pricer *ledger.Pricer
|
||
Log *slog.Logger
|
||
|
||
clients map[string]llm.LLMClient
|
||
templates map[string]*PromptTemplate
|
||
}
|
||
|
||
// NewRunner loads all three configs rooted at book.yaml and opens the project
|
||
// store. Caller owns Close.
|
||
func NewRunner(bookPath string, logger *slog.Logger) (*Runner, error) {
|
||
book, err := config.LoadBook(bookPath)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
models, err := config.LoadModels(book.ModelsFile)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
pipe, err := config.LoadPipeline(book.Pipeline, models)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
pricer, err := models.Prices()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
st, err := store.Open(book.ProjectDB)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
r := &Runner{
|
||
Book: book,
|
||
Models: models,
|
||
Pipeline: pipe,
|
||
Store: st,
|
||
Pricer: pricer,
|
||
Log: logger,
|
||
clients: map[string]llm.LLMClient{},
|
||
templates: map[string]*PromptTemplate{},
|
||
}
|
||
if err := r.loadTemplates(); err != nil {
|
||
st.Close()
|
||
return nil, err
|
||
}
|
||
return r, nil
|
||
}
|
||
|
||
func (r *Runner) Close() error { return r.Store.Close() }
|
||
|
||
func (r *Runner) loadTemplates() error {
|
||
for _, st := range r.Pipeline.Stages {
|
||
tpl, err := LoadPromptTemplate(st.Prompt)
|
||
if err != nil {
|
||
return err
|
||
}
|
||
r.templates[st.Name] = tpl
|
||
}
|
||
return nil
|
||
}
|
||
|
||
func (r *Runner) client(model string) (llm.LLMClient, error) {
|
||
if c, ok := r.clients[model]; ok {
|
||
return c, nil
|
||
}
|
||
c, err := BuildClient(r.Models, model, r.Log)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
r.clients[model] = c
|
||
return c, nil
|
||
}
|
||
|
||
// snapshotID materializes the job-context snapshot (§3.2): brief_hash, chunker
|
||
// version and the full stage plan (model, prompt version + CONTENT hash,
|
||
// sampling). Payload is rendered with a fixed field order — the id is a
|
||
// content hash, so identical context re-upserts idempotently.
|
||
func (r *Runner) snapshotID() (id, payload string, err error) {
|
||
type stageSnap struct {
|
||
Name string `json:"name"`
|
||
Role string `json:"role"`
|
||
Model string `json:"model"`
|
||
PromptVersion string `json:"prompt_version"`
|
||
PromptSHA256 string `json:"prompt_sha256"`
|
||
Temperature float64 `json:"temperature"`
|
||
Reasoning string `json:"reasoning"`
|
||
}
|
||
snap := struct {
|
||
BriefHash string `json:"brief_hash"`
|
||
ChunkerVersion string `json:"chunker_version"`
|
||
PipelineCore string `json:"pipeline_core"`
|
||
Stages []stageSnap `json:"stages"`
|
||
}{
|
||
BriefHash: r.Book.BriefHash(),
|
||
ChunkerVersion: chunkerVersion,
|
||
PipelineCore: r.Pipeline.Core,
|
||
}
|
||
for _, st := range r.Pipeline.Stages {
|
||
snap.Stages = append(snap.Stages, stageSnap{
|
||
Name: st.Name, Role: st.Role, Model: st.Model,
|
||
PromptVersion: st.PromptVersion, PromptSHA256: r.templates[st.Name].SHA256,
|
||
Temperature: st.Temperature, Reasoning: st.Reasoning,
|
||
})
|
||
}
|
||
data, err := json.Marshal(snap)
|
||
if err != nil {
|
||
return "", "", err
|
||
}
|
||
sum := sha256.Sum256(data)
|
||
return hex.EncodeToString(sum[:]), string(data), nil
|
||
}
|
||
|
||
// StageResult reports one executed stage.
|
||
type StageResult struct {
|
||
Stage string
|
||
Role string
|
||
Model string // фактически ответившая модель
|
||
FromResume bool // served from a checkpoint, no call made
|
||
Usage llm.Usage
|
||
CostUSD float64
|
||
LatencyMS int
|
||
FinishReason string
|
||
Text string
|
||
}
|
||
|
||
// RunResult is the whole chunk's outcome.
|
||
type RunResult struct {
|
||
BookID string
|
||
Chapter int
|
||
ChunkIdx int
|
||
Stages []StageResult
|
||
FinalText string
|
||
TotalUSD float64
|
||
}
|
||
|
||
// TranslateOneChunk runs the configured stages over the book's source file as
|
||
// a single chunk (Фаза 0). Chapter/chunk are fixed at 1/0 until the chunker
|
||
// lands in Phase 1.
|
||
func (r *Runner) TranslateOneChunk(ctx context.Context) (*RunResult, error) {
|
||
srcRaw, err := os.ReadFile(r.Book.SourceFile)
|
||
if err != nil {
|
||
return nil, fmt.Errorf("pipeline: read source: %w", err)
|
||
}
|
||
source := strings.TrimSpace(string(srcRaw))
|
||
if source == "" {
|
||
return nil, fmt.Errorf("pipeline: source file %s is empty", r.Book.SourceFile)
|
||
}
|
||
|
||
snapID, snapPayload, err := r.snapshotID()
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if err := r.Store.UpsertSnapshot(snapID, r.Book.BriefHash(), snapPayload); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
const chapter, chunkIdx = 1, 0
|
||
res := &RunResult{BookID: r.Book.BookID, Chapter: chapter, ChunkIdx: chunkIdx}
|
||
prev := ""
|
||
|
||
for _, st := range r.Pipeline.Stages {
|
||
sr, err := r.runStage(ctx, st, snapID, chapter, chunkIdx, source, prev)
|
||
if err != nil {
|
||
return res, err
|
||
}
|
||
res.Stages = append(res.Stages, *sr)
|
||
res.TotalUSD += sr.CostUSD
|
||
prev = sr.Text
|
||
}
|
||
res.FinalText = prev
|
||
return res, nil
|
||
}
|
||
|
||
func (r *Runner) runStage(ctx context.Context, st config.Stage, snapID string, chapter, chunkIdx int, source, prev string) (*StageResult, error) {
|
||
job, err := r.Store.EnsureJob(r.Book.BookID, chapter, st.Name, snapID)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
if err := r.Store.SetJobStatus(job.ID, "running"); err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
msgs, err := Messages(r.templates[st.Name], RenderVars{Book: r.Book, Text: source, Draft: prev})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// max_tokens: калибровка полигона — русский выход ≈1.9× входа в токенах;
|
||
// ratio из конфига, floor защищает короткие чанки. Входит в request-hash
|
||
// (меняет вызов), поэтому считается детерминированно от исходника.
|
||
inputEst := EstimateTokens(source)
|
||
maxTokens := int(float64(inputEst) * r.Pipeline.Defaults.MaxOutputRatio)
|
||
if maxTokens < r.Pipeline.Defaults.MinMaxTokens {
|
||
maxTokens = r.Pipeline.Defaults.MinMaxTokens
|
||
}
|
||
|
||
reqHash := RequestHash(r.Book.BookID, chapter, chunkIdx, st.Name, st.Role, st.Model,
|
||
st.Temperature, st.Reasoning, false, maxTokens, snapID, msgs)
|
||
|
||
ctx = obs.WithReqInfo(ctx, obs.ReqInfo{
|
||
TraceID: obs.TraceFromContext(ctx), Book: r.Book.BookID,
|
||
Chapter: chapter, Chunk: chunkIdx, Stage: st.Name, Role: st.Role,
|
||
})
|
||
|
||
// Resume path: a checkpoint means the call already happened and was paid —
|
||
// serve it, log the hit, never re-bill (приёмка: kill -9 теряет ≤1 вызов).
|
||
if cp, err := r.Store.GetCheckpoint(reqHash); err != nil {
|
||
return nil, err
|
||
} else if cp != nil {
|
||
var usage llm.Usage
|
||
_ = json.Unmarshal([]byte(cp.UsageJSON), &usage)
|
||
r.Store.LogRequest(ctx, r.Log, store.RequestLog{
|
||
BookID: r.Book.BookID, Chapter: chapter, ChunkIdx: chunkIdx,
|
||
Stage: st.Name, Role: st.Role, ModelRequested: st.Model, ModelActual: cp.ModelActual,
|
||
RequestHash: reqHash, TMHit: true, OK: true, FinishReason: cp.FinishReason,
|
||
})
|
||
if err := r.Store.SetJobStatus(job.ID, "done"); err != nil {
|
||
return nil, err
|
||
}
|
||
r.Log.InfoContext(ctx, "stage served from checkpoint", "stage", st.Name, "hash", reqHash[:12])
|
||
return &StageResult{Stage: st.Name, Role: st.Role, Model: cp.ModelActual,
|
||
FromResume: true, Usage: usage, CostUSD: 0, FinishReason: cp.FinishReason, Text: cp.ResponseText}, nil
|
||
}
|
||
|
||
price := r.Pricer.PriceFor(st.Model)
|
||
promptEst := 0
|
||
for _, m := range msgs {
|
||
promptEst += EstimateTokens(m.Content)
|
||
}
|
||
estimate := ledger.EstimateUSD(price, promptEst, maxTokens)
|
||
|
||
resv, verdict, err := r.Store.Reserve(r.Book.BookID, estimate, store.Ceilings{
|
||
BookUSD: r.Book.Ceilings.BookUSD, DayUSD: r.Book.Ceilings.DayUSD,
|
||
})
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
switch verdict {
|
||
case store.ReserveDeniedBook:
|
||
return nil, fmt.Errorf("pipeline: book USD ceiling reached (%.2f$) — raise ceilings.book_usd or stop", r.Book.Ceilings.BookUSD)
|
||
case store.ReserveDeniedDay:
|
||
return nil, fmt.Errorf("pipeline: daily USD ceiling reached (%.2f$)", r.Book.Ceilings.DayUSD)
|
||
}
|
||
|
||
client, err := r.client(st.Model)
|
||
if err != nil {
|
||
_ = r.Store.Release(resv)
|
||
return nil, err
|
||
}
|
||
|
||
start := time.Now()
|
||
resp, err := client.Complete(ctx, llm.LLMRequest{
|
||
Model: st.Model,
|
||
Messages: msgs,
|
||
MaxTokens: maxTokens,
|
||
Temperature: st.Temperature,
|
||
ReasoningEffort: st.Reasoning,
|
||
})
|
||
latency := int(time.Since(start).Milliseconds())
|
||
if err != nil {
|
||
// No 2xx ever arrived: nothing was billed by the provider — release the
|
||
// reservation and surface the failure with a telemetry row.
|
||
if rerr := r.Store.Release(resv); rerr != nil {
|
||
r.Log.ErrorContext(ctx, "release after failed call also failed", "err", rerr)
|
||
}
|
||
r.Store.LogRequest(ctx, r.Log, store.RequestLog{
|
||
BookID: r.Book.BookID, Chapter: chapter, ChunkIdx: chunkIdx,
|
||
Stage: st.Name, Role: st.Role, ModelRequested: st.Model,
|
||
RequestHash: reqHash, LatencyMS: latency, Err: err.Error(), OK: false,
|
||
})
|
||
_ = r.Store.SetJobStatus(job.ID, "failed")
|
||
return nil, fmt.Errorf("pipeline: stage %s call: %w", st.Name, err)
|
||
}
|
||
|
||
modelActual := resp.Model
|
||
if modelActual == "" {
|
||
modelActual = st.Model
|
||
}
|
||
cost := ledger.CostUSD(r.Pricer.PriceFor(modelActual), resp.Usage)
|
||
usageJSON, err := json.Marshal(resp.Usage)
|
||
if err != nil {
|
||
return nil, err
|
||
}
|
||
|
||
// Деньги: settle + сырой ответ — одна транзакция (§3.3).
|
||
if err := r.Store.SettleWithCheckpoint(resv, cost, store.Checkpoint{
|
||
RequestHash: reqHash, JobID: job.ID, ChunkIdx: chunkIdx, Stage: st.Name, Role: st.Role,
|
||
ModelRequested: st.Model, ModelActual: modelActual, ResponseText: resp.Text,
|
||
UsageJSON: string(usageJSON), CostUSD: cost, FinishReason: resp.FinishReason,
|
||
ProviderRequestID: resp.ProviderRequestID,
|
||
}); err != nil {
|
||
// The provider HAS billed this 2xx; failing to persist means the money
|
||
// state is behind reality — fail loud, never continue on top.
|
||
return nil, fmt.Errorf("pipeline: settle stage %s: %w", st.Name, err)
|
||
}
|
||
|
||
r.Store.LogRequest(ctx, r.Log, store.RequestLog{
|
||
BookID: r.Book.BookID, Chapter: chapter, ChunkIdx: chunkIdx,
|
||
Stage: st.Name, Role: st.Role, ModelRequested: st.Model, ModelActual: modelActual,
|
||
RequestHash: reqHash,
|
||
PromptTokens: resp.Usage.PromptTokens, CachedTokens: resp.Usage.CachedTokens,
|
||
CacheCreationTokens: resp.Usage.CacheCreationTokens,
|
||
CompletionTokens: resp.Usage.CompletionTokens, ReasoningTokens: resp.Usage.ReasoningTokens,
|
||
CostUSD: cost, LatencyMS: latency, FinishReason: resp.FinishReason, OK: true,
|
||
})
|
||
if err := r.Store.SetJobStatus(job.ID, "done"); err != nil {
|
||
return nil, err
|
||
}
|
||
r.Log.InfoContext(ctx, "stage completed", "stage", st.Name, "model", modelActual,
|
||
"cost_usd", fmt.Sprintf("%.6f", cost), "latency_ms", latency,
|
||
"prompt_tokens", resp.Usage.PromptTokens, "completion_tokens", resp.Usage.CompletionTokens)
|
||
|
||
return &StageResult{Stage: st.Name, Role: st.Role, Model: modelActual,
|
||
Usage: resp.Usage, CostUSD: cost, LatencyMS: latency,
|
||
FinishReason: resp.FinishReason, Text: resp.Text}, nil
|
||
}
|