db123-ai
3895ca68c0
- internal/llama/template.go: ChatML/Llama3 template engine with stop tokens - internal/agent/reasoning.go: parse <think>/<reasoning> tags from LLM output - companion/promptbuilder.go: output structured BuildResult, add date/time context - agent/agent.go: wire templates + reasoning split into message loop - storage: add reasoning TEXT column to messages table + migration - channel/tui.go: collapsible reasoning (tab toggle), mouse scroll, polished dimmed style - config: CHAT_TEMPLATE env var (chatml|llama3|none) - plan.md: mark Phase 2 complete, add themes TODO
72 lines
1.6 KiB
Go
72 lines
1.6 KiB
Go
package llama
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
|
|
"git.db123.ir/db123/divacode/pkg/types"
|
|
)
|
|
|
|
type TemplateEngine struct {
|
|
format string
|
|
}
|
|
|
|
func NewTemplateEngine(format string) *TemplateEngine {
|
|
return &TemplateEngine{format: format}
|
|
}
|
|
|
|
func (te *TemplateEngine) Build(segments []types.PromptSegment) string {
|
|
switch te.format {
|
|
case "chatml":
|
|
return te.buildChatML(segments)
|
|
case "llama3":
|
|
return te.buildLlama3(segments)
|
|
default:
|
|
return te.buildRaw(segments)
|
|
}
|
|
}
|
|
|
|
func (te *TemplateEngine) StopTokens() []string {
|
|
switch te.format {
|
|
case "chatml":
|
|
return []string{"<|im_end|>"}
|
|
case "llama3":
|
|
return []string{"<|eot_id|>"}
|
|
default:
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func (te *TemplateEngine) buildChatML(segments []types.PromptSegment) string {
|
|
var b strings.Builder
|
|
for _, seg := range segments {
|
|
b.WriteString(fmt.Sprintf("<|im_start|>%s\n%s<|im_end|>\n", seg.Type, seg.Content))
|
|
}
|
|
b.WriteString("<|im_start|>assistant\n")
|
|
return b.String()
|
|
}
|
|
|
|
func (te *TemplateEngine) buildLlama3(segments []types.PromptSegment) string {
|
|
var b strings.Builder
|
|
b.WriteString("<|begin_of_text|>")
|
|
for _, seg := range segments {
|
|
b.WriteString(fmt.Sprintf("<|start_header_id|>%s<|end_header_id|>\n\n%s<|eot_id|>", seg.Type, seg.Content))
|
|
}
|
|
b.WriteString("<|start_header_id|>assistant<|end_header_id|>\n\n")
|
|
return b.String()
|
|
}
|
|
|
|
func (te *TemplateEngine) buildRaw(segments []types.PromptSegment) string {
|
|
var b strings.Builder
|
|
for _, seg := range segments {
|
|
switch seg.Type {
|
|
case types.SegSystem:
|
|
b.WriteString("[System]\n" + seg.Content + "\n\n")
|
|
case types.SegUser:
|
|
b.WriteString("User: " + seg.Content + "\n")
|
|
}
|
|
}
|
|
b.WriteString("You:")
|
|
return b.String()
|
|
}
|