refactor: move model family

Kujtim Hoxha created 2 days ago

Change summary

internal/agent/prompt/model_family.go      | 2 +-
internal/agent/prompt/model_family_test.go | 2 +-
internal/agent/prompt/prompt.go            | 6 +++---
internal/agent/prompts.go                  | 3 +--
4 files changed, 6 insertions(+), 7 deletions(-)

Detailed changes

internal/agent/model_family.go → internal/agent/prompt/model_family.go

@@ -1,4 +1,4 @@
-package agent
+package prompt
 
 import "strings"
 

internal/agent/model_family_test.go → internal/agent/prompt/model_family_test.go

@@ -1,4 +1,4 @@
-package agent
+package prompt
 
 import (
 	"testing"

internal/agent/prompt/prompt.go

@@ -24,13 +24,13 @@ type Prompt struct {
 	now         func() time.Time
 	platform    string
 	workingDir  string
-	modelFamily string
+	modelFamily ModelFamily
 }
 
 type PromptDat struct {
 	Provider      string
 	Model         string
-	ModelFamily   string
+	ModelFamily   ModelFamily
 	Config        config.Config
 	WorkingDir    string
 	IsGitRepo     bool
@@ -66,7 +66,7 @@ func WithWorkingDir(workingDir string) Option {
 	}
 }
 
-func WithModelFamily(modelFamily string) Option {
+func WithModelFamily(modelFamily ModelFamily) Option {
 	return func(p *Prompt) {
 		p.modelFamily = modelFamily
 	}

internal/agent/prompts.go

@@ -18,8 +18,7 @@ var taskPromptTmpl []byte
 var initializePromptTmpl []byte
 
 func coderPrompt(modelName string, opts ...prompt.Option) (*prompt.Prompt, error) {
-	family := DetectModelFamily(modelName)
-	opts = append(opts, prompt.WithModelFamily(string(family)))
+	opts = append(opts, prompt.WithModelFamily(prompt.DetectModelFamily(modelName)))
 	return prompt.NewPrompt("coder", string(coderPromptTmpl), opts...)
 }