diff --git a/src/config/config.go b/src/config/config.go index 0201ad6..2e03218 100644 --- a/src/config/config.go +++ b/src/config/config.go @@ -94,6 +94,23 @@ type Config struct { // flow with explanation + full post content + tool inputs combined). MaxOutputTokens int64 `envconfig:"MAX_OUTPUT_TOKENS" default:"64000"` + // PostAssistantPlanner (CON-128) enables the hybrid model split for the + // Post Assistant: the orchestration/routing loop runs on the cheap + // PlanningModelID (Haiku) while the actual copywriting is delegated to a + // Sonnet (ModelID) editPost write-tool. Default on. Set to false to force + // the whole assistant back onto the proven single-Sonnet path (loop on + // ModelID, no editPost tool, inline content) — the instant rollback lever + // if Haiku routing regresses. Model ids stay tunable via MODEL_ID / + // PLANNING_MODEL_ID regardless. + PostAssistantPlanner bool `envconfig:"POST_ASSISTANT_PLANNER" default:"true"` + + // PostAssistantPlannerMaxOutputTokens caps the Haiku planner turn's output + // (CON-128). The planner only emits a short envelope (explanation + action + // + saveVersion + versionNote) plus tool inputs, so a small cap is plenty; + // the full post is produced by the writer sub-call under MaxOutputTokens. + // 0 falls back to a sensible default (8192). + PostAssistantPlannerMaxOutputTokens int64 `envconfig:"POST_ASSISTANT_PLANNER_MAX_OUTPUT_TOKENS" default:"8192"` + // Content-plan batching. The flow generates posts in K-sized batches in // parallel; the defaults are sized so a 64K-output Sonnet call comfortably // returns 30 posts with headroom, and so an account with default tier diff --git a/src/genkit/flows/post_assistant/edit_tool_test.go b/src/genkit/flows/post_assistant/edit_tool_test.go new file mode 100644 index 0000000..9059112 --- /dev/null +++ b/src/genkit/flows/post_assistant/edit_tool_test.go @@ -0,0 +1,96 @@ +package post_assistant + +import ( + "context" + "strings" + "testing" +) + +// The editPost tool must reject an empty instruction before it ever spins up +// the writer sub-call, and must not record an edit result. +func TestToolEditPost_RequiresInstruction(t *testing.T) { + st := &requestState{postID: "p1"} + ctx := withRequestState(context.Background(), st) + if _, err := toolEditPost(ctx, EditPostInput{Instruction: " "}); err == nil { + t.Fatal("expected an error for an empty instruction") + } + if st.editResult != nil { + t.Fatal("editResult must stay nil when no writer ran") + } +} + +// runWriter is a no-op in the legacy state (no genkit instance / provider / +// writer system prompt) — the guard keeps a stray call from panicking and +// clearly reports the writer is unavailable. +func TestRunWriter_Unavailable(t *testing.T) { + st := &requestState{postID: "p1"} // g / provider / writerSystem all zero + if _, err := runWriter(context.Background(), st, "shorten it", false); err == nil { + t.Fatal("expected an error when the writer is unavailable") + } +} + +// A well-formed editPost call whose writer is unavailable surfaces a wrapped +// write-content error and leaves editResult unset, so the runner never +// finalises a bogus "edited" turn. +func TestToolEditPost_WriterUnavailable(t *testing.T) { + st := &requestState{postID: "p1"} // writerSystem empty → runWriter errors + ctx := withRequestState(context.Background(), st) + _, err := toolEditPost(ctx, EditPostInput{Instruction: "make it punchier"}) + if err == nil { + t.Fatal("expected an error when content writing is unavailable") + } + if !strings.Contains(err.Error(), "write content") { + t.Fatalf("expected a wrapped write-content error, got: %v", err) + } + if st.editResult != nil { + t.Fatal("editResult must stay nil on writer failure") + } +} + +// The writer must receive the full retrieved excerpts as source material — +// alongside the unchanged, verbatim instruction — so an asset-grounded edit +// grounds on the retrieved text rather than the short preview (CON-128). +func TestComposeWriterInstruction_IncludesRetrievedExcerpts(t *testing.T) { + instruction := "Add a section on goroutine scheduling from the whitepaper." + excerpts := []retrievedExcerpt{ + {AssetID: "asset1", ChunkID: "c1", Content: "Goroutines are multiplexed onto OS threads by the Go runtime scheduler."}, + } + + out := composeWriterInstruction(instruction, excerpts) + + if !strings.HasPrefix(out, instruction) { + t.Fatalf("the verbatim instruction must lead the writer prompt; got: %q", out) + } + if !strings.Contains(out, "multiplexed onto OS threads") { + t.Fatalf("the retrieved excerpt content must reach the writer as source material; got: %q", out) + } + if !strings.Contains(out, "Source material") { + t.Fatalf("excerpts should be labelled as source material; got: %q", out) + } + + // No retrieval this turn → the instruction is passed through untouched. + if got := composeWriterInstruction(instruction, nil); got != instruction { + t.Fatalf("with no excerpts the instruction must be unchanged; got: %q", got) + } +} + +// Asset-retrieval tools may return the same chunk more than once across a turn; +// captureExcerpts must record each chunk once so the writer prompt isn't padded +// with duplicates. +func TestCaptureExcerpts_DedupesByChunkID(t *testing.T) { + st := &requestState{} + out := &ChunksOutput{Chunks: []ChunkContent{ + {ID: "c1", Content: "alpha"}, + {ID: "c2", Content: "beta"}, + }} + + st.captureExcerpts("a1", out) + st.captureExcerpts("a1", out) // same chunks retrieved again + + if len(st.retrieved) != 2 { + t.Fatalf("expected 2 deduped excerpts, got %d", len(st.retrieved)) + } + if st.retrieved[0].Content != "alpha" || st.retrieved[1].Content != "beta" { + t.Fatalf("captured excerpts lost their content: %+v", st.retrieved) + } +} diff --git a/src/genkit/flows/post_assistant/flow.go b/src/genkit/flows/post_assistant/flow.go index fae6c65..d4c2680 100644 --- a/src/genkit/flows/post_assistant/flow.go +++ b/src/genkit/flows/post_assistant/flow.go @@ -5,6 +5,7 @@ import ( "embed" "fmt" "log/slog" + "strings" "text/template" "github.com/firebase/genkit/go/core" @@ -38,9 +39,27 @@ func InitPostAssistant(g *genkit.Genkit, cfg PostAssistantFlowConfig, repos Post return fmt.Errorf("parse post_assistant.tmpl: %w", err) } systemTmpl := tmpl.Lookup("system") + plannerTmpl := tmpl.Lookup("planner") + writerTmpl := tmpl.Lookup("writer") contextTmpl := tmpl.Lookup("context") - if systemTmpl == nil || contextTmpl == nil { - return fmt.Errorf("post_assistant.tmpl must define both {{define \"system\"}} and {{define \"context\"}} blocks") + if systemTmpl == nil || plannerTmpl == nil || writerTmpl == nil || contextTmpl == nil { + return fmt.Errorf("post_assistant.tmpl must define \"system\", \"planner\", \"writer\", and \"context\" blocks") + } + + // CON-128: in the hybrid path the orchestration loop runs on the planner + // system prompt (routing + the editPost tool) and delegates copywriting to + // the Sonnet writer; the legacy path keeps the single system prompt that + // writes content inline. Resolve which system prompt the loop uses, and + // render the (static) writer instructions once for the writer sub-call. + activeSystemTmpl := systemTmpl + var writerInstructions string + if cfg.PlannerEnabled { + activeSystemTmpl = plannerTmpl + wi, err := renderTemplate(writerTmpl, contextTemplateData{}) + if err != nil { + return fmt.Errorf("render writer instructions: %w", err) + } + writerInstructions = strings.TrimSpace(wi) } tools := defineTools(g) @@ -53,12 +72,12 @@ func InitPostAssistant(g *genkit.Genkit, cfg PostAssistantFlowConfig, repos Post PostAssistantFlow = genkit.DefineFlow(g, "postAssistant", func(ctx context.Context, req PostAssistantRequest) (*PostAssistantResponse, error) { - return runPostAssistant(ctx, g, req, cfg, repos, systemTmpl, contextTmpl, tools, nil) + return runPostAssistant(ctx, g, req, cfg, repos, activeSystemTmpl, contextTmpl, writerInstructions, tools, nil) }, ) postAssistantRunner = func(ctx context.Context, req PostAssistantRequest, onEvent OnEventFunc) (*PostAssistantResponse, error) { - return runPostAssistant(ctx, g, req, cfg, repos, systemTmpl, contextTmpl, tools, onEvent) + return runPostAssistant(ctx, g, req, cfg, repos, activeSystemTmpl, contextTmpl, writerInstructions, tools, onEvent) } return nil diff --git a/src/genkit/flows/post_assistant/prewarm.go b/src/genkit/flows/post_assistant/prewarm.go index 7760444..11cd2b3 100644 --- a/src/genkit/flows/post_assistant/prewarm.go +++ b/src/genkit/flows/post_assistant/prewarm.go @@ -28,15 +28,26 @@ func prewarmToolCache(g *genkit.Genkit, cfg PostAssistantFlowConfig, t *toolSet) ctx, cancel := context.WithTimeout(context.Background(), 90*time.Second) defer cancel() + // CON-128: warm the grammar the real loop uses — in the hybrid path that's + // the planner tool set (incl. editPost) on the planning model; in the legacy + // path it's the base tool set on the generation model. The writer sub-call + // carries no tools, so there is no separate grammar to warm for it. + role := llm.RoleGeneration + tools := []ai.ToolRef{ + t.listAssets, t.getAssetChunks, t.searchAssetChunks, t.getCurrentContent, + t.clonePost, t.restoreVersion, t.schedulePost, t.createNote, + } + if cfg.PlannerEnabled { + role = llm.RolePlanning + tools = append(tools, t.editPost) + } + start := time.Now() _, err := genkit.Generate(ctx, g, - ai.WithModelName(cfg.Provider.Ref(llm.RoleGeneration)), + ai.WithModelName(cfg.Provider.Ref(role)), ai.WithSystem("warmup"), ai.WithPrompt("warmup"), - ai.WithTools( - t.listAssets, t.getAssetChunks, t.searchAssetChunks, t.getCurrentContent, - t.clonePost, t.restoreVersion, t.schedulePost, t.createNote, - ), + ai.WithTools(tools...), ai.WithMaxTurns(1), cfg.Provider.CallConfig(1), // max_tokens: 1 — grammar compiles during prep ) diff --git a/src/genkit/flows/post_assistant/prompts/post_assistant.tmpl b/src/genkit/flows/post_assistant/prompts/post_assistant.tmpl index 2545602..8dbc727 100644 --- a/src/genkit/flows/post_assistant/prompts/post_assistant.tmpl +++ b/src/genkit/flows/post_assistant/prompts/post_assistant.tmpl @@ -85,6 +85,97 @@ The updatedContent field MUST contain the complete post content formatted as Mar Respond ONLY with the JSON object. No markdown fences, no extra text.{{end}} +{{define "planner"}}You are the Post Assistant — an AI editor that helps users enhance post content for marketing campaigns. Your job is to ROUTE each request: understand what the user wants, then either answer directly or call the right tool. You do NOT write post copy yourself — a dedicated writer produces all content generation through the **editPost** tool. + +## Your role +- Route each request: change the post (via **editPost**), answer a question, or run an action tool (clone, restore, schedule, note). +- Pass the user's intent to editPost faithfully — do NOT paraphrase, pre-write, or summarise the copy; forward the instruction so the writer preserves the campaign's voice, tone, and phase objectives. +- Stay within scope: you help with post content and answer questions about the attached assets, the campaign, the current post, or your own recent changes. + +## Modes + +**1. Edit mode** — `action: "edited"`. Use when the user asks for ANY change to the post: rephrase, expand, shorten, adjust tone, restructure, incorporate asset information, or write/draft the post from a **draft_thesis** outline. Call the **editPost** tool: put the user's instruction VERBATIM in `instruction` and set `mode` to `edit` (or `draft` when writing the post body from a draft_thesis). Do NOT write the content yourself — the writer produces it and the server applies it. After editPost returns, put a short summary of what changed in `explanation`, and set `saveVersion`/`versionNote` per the versioning rules below. `action` becomes "edited" automatically. +- **Asset-grounded edits** ("incorporate what the whitepaper says about X", "add the stats from the report"): FIRST retrieve the relevant material with `getAssetChunks` or `searchAssetChunks`, THEN call editPost in the same turn. Every excerpt you retrieve this turn is automatically handed to the writer as source material, so it works from the full retrieved text — not the short asset preview. Do not paste the excerpts into `instruction`; just retrieve them first. + +**2. Answer mode** — `action: "declined"`. Use when the user asks a question or makes a request that does NOT change the post. Put your answer (or the brief reason you can't help) in `explanation` and do NOT call editPost. Examples: +- "What does the asset say about goroutines?" +- "Summarize the key messages of this campaign." +- "How long is the current post?" +- "Why did you choose that phrasing?" +- Anything off-topic ("what's the weather?") — explain briefly it's out of scope. + +**3. Clone mode** — for "clone this post" / "duplicate this" / "clone this for ". Call the **clonePost** tool. Cloning never changes the source, and you never write the clone's copy yourself. +- **Same platform** ("clone this post"): call clonePost with no targetPlatform — the server copies the content verbatim. +- **Different platform** ("clone this for Threads"): set `targetPlatform` to the platform name — the server's writer adapts the copy to that platform's conventions (length, tone, formatting, links). If the user adds a steer ("...and keep it short"), pass it in `instruction`. Do NOT pass adapted `content` — leave the writing to the server. +- If the platform name is unrecognized, the tool returns an error listing the available platforms — relay that and ask the user which one (answer mode); do not guess. +After the tool returns, `action` becomes "cloned"; leave any content alone and put a short confirmation (mentioning the new draft) in `explanation`. + +**4. Restore mode** — for "restore version N" / "restore v2" / "restore the previous version" / "undo the last change" / "go back". Call the **restoreVersion** tool. Do NOT edit the post yourself — the tool swaps the content for you and saves a new version (non-destructive; the history is preserved). After the tool returns, set `action` to `"restored"` and put a short confirmation in `explanation` mentioning which version you restored. +- **Explicit version** ("restore v2"): call restoreVersion with `versionNumber: 2`. Use the numbers shown in "Version history" below. +- **Relative** ("restore the previous version" / "undo" / "go back"): call restoreVersion with `relative: "previous"`. +- If the version doesn't exist or there's nothing earlier to restore, the tool returns an error — relay it and ask the user (answer mode); do NOT guess another number. +- If the tool reports it was a no-op (the post already matches that version), tell the user there was nothing to restore. + +**5. Schedule mode** — for "schedule this for tomorrow at 9am" / "publish next Monday morning" / "post this on Friday 8am". This is a **two-step, confirm-then-commit** flow. A **"Scheduling context"** block is injected with the user's message carrying the current time, the workspace timezone, the post's status, its auto/manual publishing mode, and a readiness summary — use it. +- **Turn 1 — confirm (do NOT call the tool yet):** resolve the relative time to an absolute time in the workspace timezone, then reply with `action: "declined"` echoing the **resolved absolute time** AND whether it will **auto-publish** or be **scheduled for manual publishing**, and ask the user to confirm. State any assumed time explicitly (e.g. "morning" → 09:00). Do not call schedulePost on this turn. +- **Turn 2 — commit:** only after the user confirms ("yes", "go ahead"), call the **schedulePost** tool with `scheduledAt` (ISO-8601 with the timezone offset) and `allowPromote: true` when the post is still a draft. After it returns, set `action` to `"scheduled"` and confirm in `explanation` (mention the time + mode). +- If the readiness summary says the draft is NOT publishable, decline and tell the user exactly what to fix — do NOT schedule. +- If the time is in the past or genuinely ambiguous, ask instead of guessing. If the post has no platform, ask the user to set one first. If the tool returns an error, relay it and ask the user (answer mode); do not retry blindly. + +**6. Notes mode** — for "make an image prompt" / "write a Nano Banana prompt" / "jot a note" / "save this as a note", or any request to produce a side artifact that should NOT go into the post body. Call the **createNote** tool with `type` `image_prompt` (for an image-generation prompt) or `note` (for anything else) and the artifact in `body`. +- **Note only**: after the tool returns, set `action` to `"noted"` and confirm in `explanation`. +- **Edit AND note in one turn**: if the user asks to change the copy AND capture a note (e.g. "tighten the intro and add an image prompt"), call **both** editPost (mode `edit`) AND createNote — both take effect. The note is captured separately from the body; the action becomes "edited". +- A **draft_thesis** note in the context below is the bullet-point outline for this post: to write/draft the post from it, call editPost with mode `draft`. Do not create draft_thesis notes yourself — createNote only accepts image_prompt or note. + +Never write prose outside of the JSON envelope. Even for purely informational answers, your full response must be the JSON object — the `explanation` field carries your prose. + +## Tools +- **editPost**: Generate or revise the post content — this is how ALL content changes happen. Pass the user's instruction verbatim and a `mode` (`edit` for a change to existing content, `draft` to write the post from a draft_thesis outline). You never write the post body yourself. +- **listAssets**: See which assets are attached to this post (ID, name, type, chunk count). +- **getAssetChunks**: Retrieve the full text of specific asset chunks. Use when the user references a specific asset or you need detailed content. +- **searchAssetChunks**: Semantic search over an asset's chunks for a topic or keyword. +- **getCurrentContent**: Re-read the latest post content if needed during long conversations. +- **clonePost**: Duplicate the current post as a new draft (see Clone mode). Set targetPlatform to clone for a different platform; the server adapts the copy. +- **restoreVersion**: Roll the post back to an earlier saved version (see Restore mode). Pass versionNumber (e.g. 2) or relative: "previous". +- **schedulePost**: Schedule the post for publishing at a confirmed absolute time (see Schedule mode). Only call it AFTER the user confirms the resolved time and the auto/manual mode. +- **createNote**: Save a side artifact (image prompt or free-form note) attached to the post instead of putting it in the body (see Notes mode). + +**Important**: Only call the asset-retrieval tools (listAssets, getAssetChunks, searchAssetChunks) when the user explicitly references assets or asks to incorporate external content — for an ordinary rephrase/shorten/expand/tone change, just call editPost with the instruction; the writer already has the current content and campaign context. When you DO retrieve asset content for an edit, retrieve it BEFORE calling editPost in the same turn — whatever you pulled is passed to the writer as source material so it grounds on the full excerpts, not the preview. The action tools (clonePost, restoreVersion, schedulePost, createNote) are driven by their own modes above. + +## Versioning rules +Set saveVersion to true for: +- Structural changes (added/removed paragraphs, reorganized flow) +- New information incorporated from assets or user-provided context +- Significant tone or audience shift +- Major rewrites or content pivots + +Set saveVersion to false for: +- Minor rephrasing or word-level tweaks +- Typo corrections +- Small formatting adjustments +- Iterative refinements within the same editing intent + +## Response format +You MUST respond with a single JSON object matching this exact schema. You do NOT emit the post content — the editPost writer produces it: +{ + "explanation": "edit mode: what you changed. answer mode: your answer to the user's question or the reason you can't help.", + "action": "edited", "declined", or "noted" (the tool-driven modes above also set cloned / restored / scheduled), + "saveVersion": true or false, + "versionNote": "short note describing the version (only in edit mode when saveVersion is true)" +} + +Respond ONLY with the JSON object. No markdown fences, no extra text.{{end}} + +{{define "writer"}}You are the Post Assistant's copywriter. You write and revise post content for marketing campaigns. You are given the campaign context (voice, tone, persona, key messages, phase) and the current post content below, followed by the change instruction — which may include a **## Source material** section carrying excerpts retrieved from the attached assets. + +Rules: +- Apply the instruction to produce the FULL updated post content. +- When a **## Source material** section is present, treat those retrieved excerpts as the authoritative source for asset-based facts and incorporate them; prefer them over any shorter asset previews in the campaign context. +- Preserve the campaign's voice, tone guidelines, and phase objectives across every edit. +- When the instruction asks to adapt/clone the post for a named platform, rewrite it to that platform's conventions (length, tone, formatting, links). +- When the instruction asks to write/draft the post from a draft_thesis outline, expand the outline into a complete post; do not dump the outline verbatim. +- Output ONLY the post content as Markdown — no preamble, no explanation, no JSON, no code fences. Use standard Markdown: headings (#, ##, ###), bold (**text**), italic (*text*), bullet lists (- item), numbered lists (1. item), links ([text](url)), blockquotes (> text), and code blocks. Always return the complete post, not just the changed parts.{{end}} + {{define "context"}}## Campaign context **Campaign**: {{.CampaignName}} diff --git a/src/genkit/flows/post_assistant/run.go b/src/genkit/flows/post_assistant/run.go index 7e8c2a0..7d12c75 100644 --- a/src/genkit/flows/post_assistant/run.go +++ b/src/genkit/flows/post_assistant/run.go @@ -41,6 +41,7 @@ func runPostAssistant( cfg PostAssistantFlowConfig, repos PostAssistantRepos, systemTmpl, contextTmpl *template.Template, + writerInstructions string, tools *toolSet, onEvent OnEventFunc, ) (out *PostAssistantResponse, retErr error) { @@ -167,6 +168,20 @@ func runPostAssistant( actor: post.CreatedBy, platforms: platforms, onEvent: onEvent, + // Writer support (CON-128): the editPost tool + clonePost adaptation + // run a nested Sonnet generation off this state. + g: g, + provider: cfg.Provider, + recorder: cfg.Recorder, + writerMaxTokens: cfg.MaxOutputTokens, + } + // The writer's system prompt is the copywriter instructions plus the same + // campaign/post context block the planner sees; only assembled in the + // hybrid path (writerSystem stays empty in the legacy path, which disables + // the writer helpers). Injecting the context here keeps the writer aware of + // the campaign voice + current content without a second DB read. + if cfg.PlannerEnabled { + st.writerSystem = writerInstructions + "\n\n" + actx.ContextBlock } ctx = withRequestState(ctx, st) @@ -183,27 +198,46 @@ func runPostAssistant( } // ── Call model ─────────────────────────────────────────────────────────── - // Assistant responses are short (description + explanation), so cap - // output tokens well below the content-plan default. - maxTokens := cfg.MaxOutputTokens - if maxTokens == 0 { - maxTokens = 64000 + // CON-128: in the hybrid path the orchestration loop routes on the cheap + // planning model and delegates all copywriting to the Sonnet editPost + // write-tool; the loop itself only emits a short envelope, so it takes a + // small output cap. The legacy path keeps the single generation-model call + // that writes the full post inline, so it needs the generous output budget. + planner := cfg.PlannerEnabled + loopRole := llm.RoleGeneration + loopMaxTokens := cfg.MaxOutputTokens + if loopMaxTokens == 0 { + loopMaxTokens = 64000 + } + if planner { + loopRole = llm.RolePlanning + loopMaxTokens = cfg.PlannerMaxOutputTokens + if loopMaxTokens == 0 { + loopMaxTokens = 8192 + } } maxTurns := cfg.MaxTurns if maxTurns == 0 { maxTurns = 8 } - modelName := cfg.Provider.Ref(llm.RoleGeneration) + modelName := cfg.Provider.Ref(loopRole) // System + context block forms the stable cached prefix. systemBlock := actx.SystemPrompt + "\n\n" + actx.ContextBlock - // Set up an incremental JSON scanner that watches the two string-valued - // fields whose deltas we surface to the client. The scanner decodes - // JSON escapes as they arrive, so the client never sees raw \n / \uXXXX. + // Set up an incremental JSON scanner that watches the string-valued fields + // whose deltas we surface to the client. The scanner decodes JSON escapes + // as they arrive, so the client never sees raw \n / \uXXXX. In the hybrid + // path the planner never emits updatedContent — the editPost writer sub-call + // streams content_delta itself — so only explanation is watched. (Values() + // still returns every top-level field for response assembly regardless.) + watchFields := []string{"explanation", "updatedContent"} + if planner { + watchFields = []string{"explanation"} + } scanner := jsonstream.New( - []string{"explanation", "updatedContent"}, + watchFields, func(key, delta string) { switch key { case "explanation": @@ -269,15 +303,22 @@ func runPostAssistant( // discarding the full response. Dropping the constraint lets us do the // parse ourselves with a tolerant preprocessor below. Format discipline // is enforced via the prompt, which is already explicit. + // Tool set: the editPost write-tool is attached only in the hybrid path; + // the legacy loop writes content inline and never routes through it. + toolRefs := []ai.ToolRef{tools.listAssets, tools.getAssetChunks, tools.searchAssetChunks, tools.getCurrentContent, tools.clonePost, tools.restoreVersion, tools.schedulePost, tools.createNote} + if planner { + toolRefs = append(toolRefs, tools.editPost) + } + resp, err := genkit.Generate(ctx, g, ai.WithModelName(modelName), ai.WithSystem(systemBlock), ai.WithMessages(history...), ai.WithPrompt(prompt), - ai.WithTools(tools.listAssets, tools.getAssetChunks, tools.searchAssetChunks, tools.getCurrentContent, tools.clonePost, tools.restoreVersion, tools.schedulePost, tools.createNote), + ai.WithTools(toolRefs...), ai.WithMaxTurns(maxTurns), ai.WithStreaming(streamCb), - cfg.Provider.CallConfig(maxTokens), + cfg.Provider.CallConfig(loopMaxTokens), ) if err != nil { slog.ErrorContext(ctx, "model call failed", logging.AttrComponent, "genkit.post_assistant", "post_id", req.PostID, "duration_ms", time.Since(start).Milliseconds(), logging.AttrError, err) @@ -293,13 +334,13 @@ func runPostAssistant( if resp.Usage != nil { outputTokens = int64(resp.Usage.OutputTokens) } - slog.WarnContext(ctx, "response truncated at max tokens", logging.AttrComponent, "genkit.post_assistant", "post_id", req.PostID, "output_tokens", outputTokens, "cap", maxTokens) + slog.WarnContext(ctx, "response truncated at max tokens", logging.AttrComponent, "genkit.post_assistant", "post_id", req.PostID, "output_tokens", outputTokens, "cap", loopMaxTokens) } if resp.Usage != nil { slog.InfoContext(ctx, "tokens", logging.AttrComponent, "genkit.post_assistant", "post_id", req.PostID, "input", resp.Usage.InputTokens, "output", resp.Usage.OutputTokens, "total", resp.Usage.InputTokens+resp.Usage.OutputTokens) } - cfg.Recorder.RecordResp(ctx, cfg.Provider.Vendor(), cfg.Provider.Model(llm.RoleGeneration), "post_assistant", resp) + cfg.Recorder.RecordResp(ctx, cfg.Provider.Vendor(), cfg.Provider.Model(loopRole), "post_assistant", resp) // ── Assemble response from scanner ─────────────────────────────────────── // The scanner has been processing every chunk in the streaming callback @@ -396,6 +437,17 @@ func runPostAssistant( } } + // ── Edit handling (CON-128) ────────────────────────────────────────────── + // In the hybrid path the editPost tool ran the Sonnet writer this turn; its + // content is authoritative and the planner never emits it. The planner + // supplies the metadata (action / saveVersion / versionNote) in its JSON + // envelope, already parsed into result above. Runs before the note handling + // so an edit-and-note turn is finalised as "edited" with the note attached. + if st.editResult != nil { + result.Action = "edited" + result.UpdatedContent = st.editResult.Content + } + // ── Note handling (CON-188) ────────────────────────────────────────────── // The createNote tool persisted its notes at call time (origin=assistant). // Notes are additive: a turn may create notes on their own or alongside an @@ -432,6 +484,29 @@ func runPostAssistant( } } + // Hybrid safety (CON-128): only the editPost writer may produce post copy, + // so an "edited" turn is legitimate ONLY when the writer actually ran + // (st.editResult set). Without it the planner either applied no edit, or + // emitted its own inline content in violation of the split — either way we + // must not persist that content (an empty body would wipe the post; a + // planner-written body would leak Haiku prose past the writer). Discard any + // such content and downgrade: to a notes-only turn if notes were captured + // this turn, otherwise to an answer, dropping the version snapshot. The + // legacy path can't hit this — there content comes from the same call. + if planner && result.Action == "edited" && st.editResult == nil { + slog.WarnContext(ctx, "planner claimed an edit without invoking editPost; discarding any inline content", logging.AttrComponent, "genkit.post_assistant", "post_id", req.PostID) + result.UpdatedContent = "" + if len(st.noteResults) > 0 { + result.Action = "noted" + } else { + result.Action = "declined" + } + result.SaveVersion = false + if result.Explanation == "" { + result.Explanation = "I couldn't apply that edit — could you rephrase what you'd like changed?" + } + } + // Pure-prose recovery: occasionally the model ignores the JSON // envelope entirely and answers in plain prose (often when the user // asks an informational question). Salvage the raw text as the diff --git a/src/genkit/flows/post_assistant/tools.go b/src/genkit/flows/post_assistant/tools.go index 25e4efc..b0c5ac5 100644 --- a/src/genkit/flows/post_assistant/tools.go +++ b/src/genkit/flows/post_assistant/tools.go @@ -17,6 +17,8 @@ import ( "github.com/ogen-app/ogen/src/post_actions/clone" "github.com/ogen-app/ogen/src/post_actions/restore" "github.com/ogen-app/ogen/src/post_actions/schedule" + "github.com/ogen-app/ogen/src/usage" + "github.com/ogen-app/ogen/src/vendors/llm" ) // Per-turn token budget for tool-retrieved chunks. @@ -68,6 +70,64 @@ type requestState struct { // generation to finalise the response. noteSvc *notes.Service noteResults []*models.PostNote + + // Writer support (CON-128). In the hybrid path the planner loop runs on + // the cheap planning model and delegates all copywriting to the Sonnet + // writer via the editPost tool (and clonePost adaptation). g runs the + // nested writer generation; provider/recorder resolve the generation model + // and meter it; writerSystem is the writer prompt + context block assembled + // per-request; writerMaxTokens caps the writer output. editResult holds the + // content the writer produced this turn, read by the runner to finalise an + // "edited" response — the content never flows back through the planner model. + g *genkit.Genkit + provider *llm.Provider + recorder *usage.Recorder + writerSystem string + writerMaxTokens int64 + editResult *editResult + + // retrieved holds the asset excerpts the planner pulled via the + // asset-retrieval tools this turn (CON-128). The writer's context block + // carries only short asset previews, so an asset-grounded edit must be fed + // the full retrieved text — captured here and passed to the writer as + // source material (see composeWriterInstruction). + retrieved []retrievedExcerpt +} + +// editResult is the internal outcome of the editPost write-tool (CON-128): the +// full post content the Sonnet writer produced this turn. +type editResult struct { + Content string +} + +// retrievedExcerpt is a chunk the planner pulled via an asset-retrieval tool +// this turn, captured so the writer sub-call can ground on the full retrieved +// text rather than the short preview in the context block. +type retrievedExcerpt struct { + AssetID string + ChunkID string + Content string +} + +// captureExcerpts records the chunks an asset-retrieval tool returned so the +// writer can use them as source material. Deduped by chunk ID across calls, so +// re-retrieving the same chunk in a later turn-step doesn't duplicate it. +func (st *requestState) captureExcerpts(assetID string, out *ChunksOutput) { + if out == nil { + return + } + for _, c := range out.Chunks { + dup := false + for _, e := range st.retrieved { + if e.ChunkID == c.ID { + dup = true + break + } + } + if !dup { + st.retrieved = append(st.retrieved, retrievedExcerpt{AssetID: assetID, ChunkID: c.ID, Content: c.Content}) + } + } } func withRequestState(ctx context.Context, s *requestState) context.Context { @@ -117,10 +177,28 @@ type SearchChunksInput struct { type ClonePostInput struct { TargetPlatform string `json:"targetPlatform,omitempty" jsonschema:"description=Platform name or ID for the clone (e.g. Threads). Omit to keep the source's platform."` TargetPostType string `json:"targetPostType,omitempty" jsonschema:"description=Post-type slug for the target platform (e.g. text-post). Omit to inherit or default."` - Content string `json:"content,omitempty" jsonschema:"description=Full Markdown content for the clone. For a cross-platform clone provide content adapted to the target platform; omit to copy the source content verbatim."` + Content string `json:"content,omitempty" jsonschema:"description=Full Markdown content for the clone. Normally leave empty — for a cross-platform clone the server's writer adapts the copy; provide content only to override it entirely. Omit to copy the source content verbatim (same-platform clone)."` + Instruction string `json:"instruction,omitempty" jsonschema:"description=Optional steer for a cross-platform adaptation (e.g. 'keep it short'), forwarded to the writer. Ignored for a verbatim same-platform clone."` Title string `json:"title,omitempty" jsonschema:"description=Title for the clone. Omit to apply the default naming."` } +// EditPostInput is the input for the editPost tool (CON-128). The planner +// forwards the user's change request verbatim; the Sonnet writer produces the +// full updated post from it. +type EditPostInput struct { + Instruction string `json:"instruction" jsonschema:"description=The user's change request, forwarded verbatim (e.g. 'shorten the intro' or 'write the post from the draft thesis'). Do not paraphrase or pre-write the content."` + Mode string `json:"mode,omitempty" jsonschema:"description=edit for a change to existing content; draft to write the post body from a draft_thesis outline.,enum=edit,enum=draft"` +} + +// EditPostOutput is the compact receipt returned to the planner after the +// writer runs. The full content deliberately does NOT come back through the +// planner model — it streams to the client and the runner finalises it from +// requestState — so the cheap planner stays cheap and can't mangle the copy. +type EditPostOutput struct { + OK bool `json:"ok"` + Chars int `json:"chars"` +} + // ClonePostOutput is returned to the model after a clone is created. type ClonePostOutput struct { NewPostID string `json:"newPostId"` @@ -186,6 +264,9 @@ type toolSet struct { restoreVersion ai.ToolRef schedulePost ai.ToolRef createNote ai.ToolRef + // editPost (CON-128) is the Sonnet write-tool used only in the hybrid + // planner path; the legacy single-Sonnet loop does not attach it. + editPost ai.ToolRef } func defineTools(g *genkit.Genkit) *toolSet { @@ -219,8 +300,10 @@ func defineTools(g *genkit.Genkit) *toolSet { clonePost := genkit.DefineTool(g, "clonePost", "Duplicates the current post as a new draft in the same campaign. "+ - "To clone for another platform, set targetPlatform and provide content adapted to that platform. "+ - "Omit content for a verbatim copy. Returns the new draft's id.", + "For a cross-platform clone, set targetPlatform and OMIT content — the server's writer "+ + "adapts the copy to the target platform (pass an optional instruction to steer it). "+ + "content is an explicit override only: provide it to set the clone's body yourself; omit it "+ + "for a verbatim same-platform copy. Returns the new draft's id.", func(ctx *ai.ToolContext, in ClonePostInput) (*ClonePostOutput, error) { return toolClonePost(ctx, in) }, @@ -258,6 +341,16 @@ func defineTools(g *genkit.Genkit) *toolSet { }, ) + editPost := genkit.DefineTool(g, "editPost", + "Generates or revises the current post's content — this is how ALL content changes happen. "+ + "Pass the user's change request verbatim in instruction (e.g. 'shorten the intro', "+ + "'make it punchier', 'write the post from the draft thesis') and a mode (edit or draft). "+ + "The content is produced and applied by the server; you do not write it yourself.", + func(ctx *ai.ToolContext, in EditPostInput) (*EditPostOutput, error) { + return toolEditPost(ctx, in) + }, + ) + return &toolSet{ listAssets: list, getAssetChunks: getChunks, @@ -267,6 +360,7 @@ func defineTools(g *genkit.Genkit) *toolSet { restoreVersion: restoreVersion, schedulePost: schedulePost, createNote: createNote, + editPost: editPost, } } @@ -312,7 +406,9 @@ func toolGetAssetChunks(ctx context.Context, in GetChunksInput) (*ChunksOutput, return nil, fmt.Errorf("fetch chunks: %w", err) } - return packChunks(chunks), nil + out := packChunks(chunks) + st.captureExcerpts(in.AssetID, out) + return out, nil } func toolSearchAssetChunks(ctx context.Context, in SearchChunksInput) (*ChunksOutput, error) { @@ -340,7 +436,9 @@ func toolSearchAssetChunks(ctx context.Context, in SearchChunksInput) (*ChunksOu return nil, fmt.Errorf("search chunks: %w", err) } - return packChunks(ranked), nil + out := packChunks(ranked) + st.captureExcerpts(in.AssetID, out) + return out, nil } func toolGetCurrentContent(ctx context.Context) (string, error) { @@ -367,8 +465,32 @@ func toolClonePost(ctx context.Context, in ClonePostInput) (*ClonePostOutput, er opts.TargetPlatformID = id } opts.TargetPostType = in.TargetPostType - if in.Content != "" { - opts.ContentOverride = &in.Content + + // CON-128: in the hybrid path the planner does not write copy, so a + // cross-platform clone's adapted content is produced by the Sonnet writer + // here rather than passed in. An explicit Content override still wins + // (legacy path / deliberate override); a verbatim same-platform clone (no + // targetPlatform) skips the writer entirely. writerSystem is empty in the + // legacy path, so this branch is inert there. + content := in.Content + if content == "" && in.TargetPlatform != "" && st.writerSystem != "" { + instr := fmt.Sprintf("Adapt the current post for the %s platform, following its conventions (length, tone, formatting, links). Output the full adapted post.", in.TargetPlatform) + if s := strings.TrimSpace(in.Instruction); s != "" { + instr += " " + s + } + adapted, err := runWriter(ctx, st, instr, false) + if err != nil { + return nil, fmt.Errorf("adapt content for %s: %w", in.TargetPlatform, err) + } + if adapted == "" { + // The writer returned no content (not an error) — fail rather than + // silently clone the source verbatim onto the target platform. + return nil, fmt.Errorf("adapt content for %s: the writer produced no content", in.TargetPlatform) + } + content = adapted + } + if content != "" { + opts.ContentOverride = &content } if in.Title != "" { opts.TitleOverride = &in.Title @@ -500,6 +622,113 @@ func toolCreateNote(ctx context.Context, in CreateNoteInput) (*CreateNoteOutput, return &CreateNoteOutput{ID: note.ID, Type: string(note.Type)}, nil } +// toolEditPost is the editPost write-tool (CON-128). In the hybrid path the +// planner (Haiku) calls it for any content change, forwarding the user's +// instruction verbatim; the Sonnet writer produces the full post, which +// streams to the client and is stashed on requestState for the runner. Only a +// compact receipt goes back to the planner. +func toolEditPost(ctx context.Context, in EditPostInput) (*EditPostOutput, error) { + st := getRequestState(ctx) + instruction := strings.TrimSpace(in.Instruction) + if instruction == "" { + return nil, fmt.Errorf("instruction is required to edit the post") + } + + content, err := runWriter(ctx, st, instruction, true) + if err != nil { + return nil, fmt.Errorf("write content: %w", err) + } + if content == "" { + return nil, fmt.Errorf("the writer produced no content") + } + st.editResult = &editResult{Content: content} + return &EditPostOutput{OK: true, Chars: len(content)}, nil +} + +// runWriter executes the Sonnet copywriting sub-call (CON-128). When stream is +// true it fans the generated Markdown out to the client as content_delta events +// (the editPost path, feeding the live editor); for clone adaptation it stays +// silent (the copy lands in a new draft, not the open editor). It returns the +// full generated post content. Usage is metered under the post_assistant_edit +// flow so the writer's Sonnet spend is attributable separately from the cheap +// planner loop. +func runWriter(ctx context.Context, st *requestState, instruction string, stream bool) (string, error) { + if st.g == nil || st.provider == nil || st.writerSystem == "" { + return "", fmt.Errorf("content writing is not available") + } + + var buf strings.Builder + streamCb := func(_ context.Context, chunk *ai.ModelResponseChunk) error { + if chunk == nil || chunk.Aggregated { + return nil + } + for _, part := range chunk.Content { + if part.IsText() { + buf.WriteString(part.Text) + if stream { + emit(st.onEvent, SSEEventContentDelta, DeltaEventPayload{Delta: part.Text}) + } + } + } + return nil + } + + maxTokens := st.writerMaxTokens + if maxTokens == 0 { + maxTokens = 64000 + } + + resp, err := genkit.Generate(ctx, st.g, + ai.WithModelName(st.provider.Ref(llm.RoleGeneration)), + ai.WithSystem(st.writerSystem), + ai.WithPrompt(composeWriterInstruction(instruction, st.retrieved)), + ai.WithStreaming(streamCb), + st.provider.CallConfig(maxTokens), + ) + if err != nil { + return "", err + } + if st.recorder != nil { + st.recorder.RecordResp(ctx, st.provider.Vendor(), st.provider.Model(llm.RoleGeneration), "post_assistant_edit", resp) + } + + content := strings.TrimSpace(buf.String()) + if content == "" { + content = strings.TrimSpace(resp.Text()) + } + return content, nil +} + +// composeWriterInstruction assembles the writer's prompt: the user's verbatim +// instruction, then any asset excerpts the planner retrieved this turn as +// clearly-delimited source material (CON-128). Appending the excerpts rather +// than folding them into the instruction keeps the instruction unchanged while +// giving an asset-grounded edit the full retrieved text — not just the short +// preview in the context block. Returns the instruction untouched when nothing +// was retrieved. +func composeWriterInstruction(instruction string, excerpts []retrievedExcerpt) string { + if len(excerpts) == 0 { + return instruction + } + var b strings.Builder + b.WriteString(instruction) + b.WriteString("\n\n## Source material (retrieved from the attached assets)\n") + b.WriteString("Use these retrieved excerpts as the source of truth for any asset-based facts; prefer them over the shorter previews in the context above.\n") + for _, e := range excerpts { + b.WriteString("\n---\n") + if e.AssetID != "" { + b.WriteString("Asset " + e.AssetID) + if e.ChunkID != "" { + b.WriteString(", chunk " + e.ChunkID) + } + b.WriteString(":\n") + } + b.WriteString(e.Content) + b.WriteString("\n") + } + return b.String() +} + // parseScheduledAt accepts the ISO-8601 instant the model resolved. It // requires an explicit timezone (offset or Z) so the stored UTC instant // is unambiguous; a naive local datetime is rejected with guidance the diff --git a/src/genkit/flows/post_assistant/types.go b/src/genkit/flows/post_assistant/types.go index 1aa45aa..c35b35f 100644 --- a/src/genkit/flows/post_assistant/types.go +++ b/src/genkit/flows/post_assistant/types.go @@ -130,7 +130,18 @@ type PostAssistantFlowConfig struct { // covers realistic asset-incorporation scenarios (browse + search a // few assets + read chunks + respond). MaxTurns int - Embedder ai.Embedder // nil = semantic search unavailable + // PlannerEnabled turns on the CON-128 hybrid split: the orchestration + // loop runs on the planning model (Haiku) and the actual copywriting is + // delegated to the Sonnet editPost write-tool. False = the legacy + // single-Sonnet path (loop on the generation model, no editPost tool, + // content written inline) — the instant-rollback fallback. + PlannerEnabled bool + // PlannerMaxOutputTokens caps the planner turn's output when + // PlannerEnabled. The planner only emits a short envelope + tool inputs, + // so a small cap is plenty; the writer sub-call uses MaxOutputTokens. + // 0 falls back to 8192. + PlannerMaxOutputTokens int64 + Embedder ai.Embedder // nil = semantic search unavailable // Hub is the event broker used to publish "operation finalised" // events on success/failure. nil = silent (no events emitted). Hub eventhub.Hub diff --git a/src/integration/post_assistant_test.go b/src/integration/post_assistant_test.go index b6a12ab..6c820c1 100644 --- a/src/integration/post_assistant_test.go +++ b/src/integration/post_assistant_test.go @@ -17,6 +17,7 @@ import ( "github.com/ogen-app/ogen/src/genkit/flows/post_assistant" "github.com/ogen-app/ogen/src/models" "github.com/ogen-app/ogen/src/repository" + "github.com/ogen-app/ogen/src/vendors/llm" ) var _ = Describe("Post assistant flow", Ordered, func() { @@ -97,7 +98,21 @@ var _ = Describe("Post assistant flow", Ordered, func() { if modelID == "" { modelID = "claude-haiku-4-5-20251001" } - flowCfg := post_assistant.PostAssistantFlowConfig{ModelID: modelID} + planningModelID := os.Getenv("PLANNING_MODEL_ID") + if planningModelID == "" { + planningModelID = "claude-haiku-4-5-20251001" + } + // CON-128: the flow resolves its loop model (and the editPost writer + // model) through the Provider — without one, cfg.Provider.Ref would + // panic. Exercise the hybrid planner path by default (the shipped + // default); set POST_ASSISTANT_PLANNER=false to run the legacy + // single-Sonnet path instead, so the same suite covers both. + provider := llm.NewProvider(modelID, modelID, planningModelID) + flowCfg := post_assistant.PostAssistantFlowConfig{ + Provider: provider, + ModelID: modelID, + PlannerEnabled: os.Getenv("POST_ASSISTANT_PLANNER") != "false", + } repos := post_assistant.PostAssistantRepos{ Posts: postRepo, Assets: assetRepo, @@ -248,6 +263,7 @@ A concise post about Go.`) var ( gotExplanationDelta bool + gotContentDelta bool gotComplete bool gotError bool eventOrder []post_assistant.SSEEventKind @@ -257,6 +273,8 @@ A concise post about Go.`) switch name { case post_assistant.SSEEventExplanationDelta: gotExplanationDelta = true + case post_assistant.SSEEventContentDelta: + gotContentDelta = true case post_assistant.SSEEventComplete: gotComplete = true case post_assistant.SSEEventError: @@ -272,6 +290,10 @@ A concise post about Go.`) Expect(gotError).To(BeFalse(), "error event should not fire on the happy path") Expect(gotExplanationDelta).To(BeTrue(), "explanation_delta should fire as the model writes the explanation") + // content_delta must stream the edited copy: in the hybrid path this + // proves the editPost writer sub-call streams through to the client; + // in the legacy path it is the inline updatedContent stream (CON-128). + Expect(gotContentDelta).To(BeTrue(), "content_delta should stream the rewritten post content on an edit") Expect(gotComplete).To(BeTrue(), "complete event signals the canonical final response") // complete should be the last event emitted. diff --git a/src/server/post_assistant.go b/src/server/post_assistant.go index 6d1f3d3..62b48b4 100644 --- a/src/server/post_assistant.go +++ b/src/server/post_assistant.go @@ -46,6 +46,11 @@ func initPostAssistant( RestoreService: restoreSvc, ScheduleService: scheduleSvc, NoteService: noteSvc, + // CON-128: hybrid model split — the planner loop runs on the planning + // model and delegates copywriting to the Sonnet editPost write-tool. + // Off reverts to the legacy single-Sonnet path. + PlannerEnabled: cfg.PostAssistantPlanner, + PlannerMaxOutputTokens: cfg.PostAssistantPlannerMaxOutputTokens, // CON-112: pre-warm the strict-tool grammar cache at boot when we're also // stabilizing tool order (otherwise the warmed key wouldn't match). PrewarmTools: cfg.AnthropicStableToolOrder,