7d0a6ad · crush · git.iko.soy

git.iko.soy/crush

on main help sign in

refactor: adjustment based on fantasy api changes

7d0a6adAndrey Nering1 year ago

internal/agent/agent_tool.go+2−2

compare with main

unifiedsplit

@@ -25,7 +25,7 @@ const (

25 AgentToolName = "agent"

26)

28−func (c *coordinator) agentTool() (fantasy.AgentTool, error) {

28+func (c *coordinator) agentTool(ctx context.Context) (fantasy.AgentTool, error) {

29 agentCfg, ok := c.cfg.Agents[config.AgentTask]

30 if !ok {

31 return nil, errors.New("task agent not configured")

@@ -35,7 +35,7 @@ func (c *coordinator) agentTool() (fantasy.AgentTool, error) {

35 return nil, err

36 }

38− agent, err := c.buildAgent(prompt, agentCfg)

38+ agent, err := c.buildAgent(ctx, prompt, agentCfg)

39 if err != nil {

40 return nil, err

41 }

internal/agent/common_test.go+22−10

@@ -47,43 +47,55 @@ type modelPair struct {

47}

49func anthropicBuilder(model string) builderFunc {

50− return func(_ *testing.T, r *recorder.Recorder) (fantasy.LanguageModel, error) {

50+ return func(t *testing.T, r *recorder.Recorder) (fantasy.LanguageModel, error) {

51− provider := anthropic.New(

51+ provider, err := anthropic.New(

52 anthropic.WithAPIKey(os.Getenv("CRUSH_ANTHROPIC_API_KEY")),

53 anthropic.WithHTTPClient(&http.Client{Transport: r}),

54 )

55+ if err != nil {

56+ return nil, err

57+ }

55− return provider.LanguageModel(model)

58+ return provider.LanguageModel(t.Context(), model)

56 }

59 }

57}

60}

59func openaiBuilder(model string) builderFunc {

62func openaiBuilder(model string) builderFunc {

60− return func(_ *testing.T, r *recorder.Recorder) (fantasy.LanguageModel, error) {

63+ return func(t *testing.T, r *recorder.Recorder) (fantasy.LanguageModel, error) {

61− provider := openai.New(

64+ provider, err := openai.New(

62 openai.WithAPIKey(os.Getenv("CRUSH_OPENAI_API_KEY")),

65 openai.WithAPIKey(os.Getenv("CRUSH_OPENAI_API_KEY")),

63 openai.WithHTTPClient(&http.Client{Transport: r}),

66 openai.WithHTTPClient(&http.Client{Transport: r}),

64 )

67 )

68+ if err != nil {

69+ return nil, err

70+ }

65− return provider.LanguageModel(model)

71+ return provider.LanguageModel(t.Context(), model)

66 }

72 }

67}

73}

69func openRouterBuilder(model string) builderFunc {

75func openRouterBuilder(model string) builderFunc {

70 return func(t *testing.T, r *recorder.Recorder) (fantasy.LanguageModel, error) {

76 return func(t *testing.T, r *recorder.Recorder) (fantasy.LanguageModel, error) {

71− provider := openrouter.New(

77+ provider, err := openrouter.New(

72 openrouter.WithAPIKey(os.Getenv("CRUSH_OPENROUTER_API_KEY")),

78 openrouter.WithAPIKey(os.Getenv("CRUSH_OPENROUTER_API_KEY")),

73 openrouter.WithHTTPClient(&http.Client{Transport: r}),

79 openrouter.WithHTTPClient(&http.Client{Transport: r}),

74 )

80 )

81+ if err != nil {

82+ return nil, err

83+ }

75− return provider.LanguageModel(model)

84+ return provider.LanguageModel(t.Context(), model)

76 }

85 }

77}

86}

79func zAIBuilder(model string) builderFunc {

88func zAIBuilder(model string) builderFunc {

80 return func(t *testing.T, r *recorder.Recorder) (fantasy.LanguageModel, error) {

89 return func(t *testing.T, r *recorder.Recorder) (fantasy.LanguageModel, error) {

81− provider := openaicompat.New(

90+ provider, err := openaicompat.New(

82 openaicompat.WithBaseURL("https://api.z.ai/api/coding/paas/v4"),

91 openaicompat.WithBaseURL("https://api.z.ai/api/coding/paas/v4"),

83 openaicompat.WithAPIKey(os.Getenv("CRUSH_ZAI_API_KEY")),

92 openaicompat.WithAPIKey(os.Getenv("CRUSH_ZAI_API_KEY")),

84 openaicompat.WithHTTPClient(&http.Client{Transport: r}),

93 openaicompat.WithHTTPClient(&http.Client{Transport: r}),

85 )

94 )

95+ if err != nil {

96+ return nil, err

97+ }

86− return provider.LanguageModel(model)

98+ return provider.LanguageModel(t.Context(), model)

87 }

99 }

88}

100}

101

internal/agent/coordinator.go+29−30

@@ -45,7 +45,7 @@ type Coordinator interface {

45 ClearQueue(sessionID string)

46 Summarize(context.Context, string) error

47 Model() Model

48− UpdateModels() error

48+ UpdateModels(ctx context.Context) error

49}

51type coordinator struct {

@@ -61,6 +61,7 @@ type coordinator struct {

61}

63func NewCoordinator(

64+ ctx context.Context,

64 cfg *config.Config,

65 cfg *config.Config,

65 sessions session.Service,

66 sessions session.Service,

66 messages message.Service,

67 messages message.Service,

@@ -89,7 +90,7 @@ func NewCoordinator(

89 return nil, err

90 return nil, err

90 }

91 }

92− agent, err := c.buildAgent(prompt, agentCfg)

93+ agent, err := c.buildAgent(ctx, prompt, agentCfg)

93 if err != nil {

94 if err != nil {

94 return nil, err

95 return nil, err

95 }

96 }

@@ -255,8 +256,8 @@ func mergeCallOptions(model Model, tp catwalk.Type) (fantasy.ProviderOptions, *f

255 return modelOptions, temp, topP, topK, freqPenalty, presPenalty

256 return modelOptions, temp, topP, topK, freqPenalty, presPenalty

256}

257}

257

258

258−func (c *coordinator) buildAgent(prompt *prompt.Prompt, agent config.Agent) (SessionAgent, error) {

259+func (c *coordinator) buildAgent(ctx context.Context, prompt *prompt.Prompt, agent config.Agent) (SessionAgent, error) {

259− large, small, err := c.buildAgentModels()

260+ large, small, err := c.buildAgentModels(ctx)

260 if err != nil {

261 if err != nil {

261 return nil, err

262 return nil, err

262 }

263 }

@@ -266,17 +267,17 @@ func (c *coordinator) buildAgent(prompt *prompt.Prompt, agent config.Agent) (Ses

266 return nil, err

267 return nil, err

267 }

268 }

268

269

269− tools, err := c.buildTools(agent)

270+ tools, err := c.buildTools(ctx, agent)

270 if err != nil {

271 if err != nil {

271 return nil, err

272 return nil, err

272 }

273 }

273 return NewSessionAgent(SessionAgentOptions{large, small, systemPrompt, c.cfg.Options.DisableAutoSummarize, c.sessions, c.messages, tools}), nil

274 return NewSessionAgent(SessionAgentOptions{large, small, systemPrompt, c.cfg.Options.DisableAutoSummarize, c.sessions, c.messages, tools}), nil

274}

275}

275

276

276−func (c *coordinator) buildTools(agent config.Agent) ([]fantasy.AgentTool, error) {

277+func (c *coordinator) buildTools(ctx context.Context, agent config.Agent) ([]fantasy.AgentTool, error) {

277 var allTools []fantasy.AgentTool

278 var allTools []fantasy.AgentTool

278 if slices.Contains(agent.AllowedTools, AgentToolName) {

279 if slices.Contains(agent.AllowedTools, AgentToolName) {

279− agentTool, err := c.agentTool()

280+ agentTool, err := c.agentTool(ctx)

280 if err != nil {

281 if err != nil {

281 return nil, err

282 return nil, err

282 }

283 }

@@ -334,7 +335,7 @@ func (c *coordinator) buildTools(agent config.Agent) ([]fantasy.AgentTool, error

334}

335}

335

336

336// TODO: when we support multiple agents we need to change this so that we pass in the agent specific model config

337// TODO: when we support multiple agents we need to change this so that we pass in the agent specific model config

337−func (c *coordinator) buildAgentModels() (Model, Model, error) {

338+func (c *coordinator) buildAgentModels(ctx context.Context) (Model, Model, error) {

338 largeModelCfg, ok := c.cfg.Models[config.SelectedModelTypeLarge]

339 largeModelCfg, ok := c.cfg.Models[config.SelectedModelTypeLarge]

339 if !ok {

340 if !ok {

340 return Model{}, Model{}, errors.New("large model not selected")

341 return Model{}, Model{}, errors.New("large model not selected")

@@ -386,11 +387,11 @@ func (c *coordinator) buildAgentModels() (Model, Model, error) {

386 return Model{}, Model{}, errors.New("snall model not found in provider config")

387 return Model{}, Model{}, errors.New("snall model not found in provider config")

387 }

388 }

388

389

389− largeModel, err := largeProvider.LanguageModel(largeModelCfg.Model)

390+ largeModel, err := largeProvider.LanguageModel(ctx, largeModelCfg.Model)

390 if err != nil {

391 if err != nil {

391 return Model{}, Model{}, err

392 return Model{}, Model{}, err

392 }

393 }

393− smallModel, err := smallProvider.LanguageModel(smallModelCfg.Model)

394+ smallModel, err := smallProvider.LanguageModel(ctx, smallModelCfg.Model)

394 if err != nil {

395 if err != nil {

395 return Model{}, Model{}, err

396 return Model{}, Model{}, err

396 }

397 }

@@ -406,7 +407,7 @@ func (c *coordinator) buildAgentModels() (Model, Model, error) {

406 }, nil

407 }, nil

407}

408}

408

409

409−func (c *coordinator) buildAnthropicProvider(baseURL, apiKey string, headers map[string]string) fantasy.Provider {

410+func (c *coordinator) buildAnthropicProvider(baseURL, apiKey string, headers map[string]string) (fantasy.Provider, error) {

410 hasBearerAuth := false

411 hasBearerAuth := false

411 for key := range headers {

412 for key := range headers {

412 if strings.ToLower(key) == "authorization" {

413 if strings.ToLower(key) == "authorization" {

@@ -441,7 +442,7 @@ func (c *coordinator) buildAnthropicProvider(baseURL, apiKey string, headers map

441 return anthropic.New(opts...)

442 return anthropic.New(opts...)

442}

443}

443

444

444−func (c *coordinator) buildOpenaiProvider(baseURL, apiKey string, headers map[string]string) fantasy.Provider {

445+func (c *coordinator) buildOpenaiProvider(baseURL, apiKey string, headers map[string]string) (fantasy.Provider, error) {

445 opts := []openai.Option{

446 opts := []openai.Option{

446 openai.WithAPIKey(apiKey),

447 openai.WithAPIKey(apiKey),

447 openai.WithUseResponsesAPI(),

448 openai.WithUseResponsesAPI(),

@@ -459,7 +460,7 @@ func (c *coordinator) buildOpenaiProvider(baseURL, apiKey string, headers map[st

459 return openai.New(opts...)

460 return openai.New(opts...)

460}

461}

461

462

462−func (c *coordinator) buildOpenrouterProvider(_, apiKey string, headers map[string]string) fantasy.Provider {

463+func (c *coordinator) buildOpenrouterProvider(_, apiKey string, headers map[string]string) (fantasy.Provider, error) {

463 opts := []openrouter.Option{

464 opts := []openrouter.Option{

464 openrouter.WithAPIKey(apiKey),

465 openrouter.WithAPIKey(apiKey),

465 }

466 }

@@ -473,7 +474,7 @@ func (c *coordinator) buildOpenrouterProvider(_, apiKey string, headers map[stri

473 return openrouter.New(opts...)

474 return openrouter.New(opts...)

474}

475}

475

476

476−func (c *coordinator) buildOpenaiCompatProvider(baseURL, apiKey string, headers map[string]string) fantasy.Provider {

477+func (c *coordinator) buildOpenaiCompatProvider(baseURL, apiKey string, headers map[string]string) (fantasy.Provider, error) {

477 opts := []openaicompat.Option{

478 opts := []openaicompat.Option{

478 openaicompat.WithBaseURL(baseURL),

479 openaicompat.WithBaseURL(baseURL),

479 openaicompat.WithAPIKey(apiKey),

480 openaicompat.WithAPIKey(apiKey),

@@ -489,7 +490,7 @@ func (c *coordinator) buildOpenaiCompatProvider(baseURL, apiKey string, headers

489 return openaicompat.New(opts...)

490 return openaicompat.New(opts...)

490}

491}

491

492

492−func (c *coordinator) buildAzureProvider(baseURL, apiKey string, headers map[string]string, options map[string]string) fantasy.Provider {

493+func (c *coordinator) buildAzureProvider(baseURL, apiKey string, headers map[string]string, options map[string]string) (fantasy.Provider, error) {

493 opts := []azure.Option{

494 opts := []azure.Option{

494 azure.WithBaseURL(baseURL),

495 azure.WithBaseURL(baseURL),

495 azure.WithAPIKey(apiKey),

496 azure.WithAPIKey(apiKey),

@@ -511,7 +512,7 @@ func (c *coordinator) buildAzureProvider(baseURL, apiKey string, headers map[str

511 return azure.New(opts...)

512 return azure.New(opts...)

512}

513}

513

514

514−func (c *coordinator) buildGoogleProvider(baseURL, apiKey string, headers map[string]string) fantasy.Provider {

515+func (c *coordinator) buildGoogleProvider(baseURL, apiKey string, headers map[string]string) (fantasy.Provider, error) {

515 opts := []google.Option{

516 opts := []google.Option{

516 google.WithBaseURL(baseURL),

517 google.WithBaseURL(baseURL),

517 google.WithGeminiAPIKey(apiKey),

518 google.WithGeminiAPIKey(apiKey),

@@ -526,7 +527,7 @@ func (c *coordinator) buildGoogleProvider(baseURL, apiKey string, headers map[st

526 return google.New(opts...)

527 return google.New(opts...)

527}

528}

528

529

529−func (c *coordinator) buildGoogleVertexProvider(headers map[string]string, options map[string]string) fantasy.Provider {

530+func (c *coordinator) buildGoogleVertexProvider(headers map[string]string, options map[string]string) (fantasy.Provider, error) {

530 opts := []google.Option{}

531 opts := []google.Option{}

531 if c.cfg.Options.Debug {

532 if c.cfg.Options.Debug {

532 httpClient := log.NewHTTPClient()

533 httpClient := log.NewHTTPClient()

@@ -574,27 +575,25 @@ func (c *coordinator) buildProvider(providerCfg config.ProviderConfig, model con

574 // TODO: make sure we have

575 // TODO: make sure we have

575 apiKey, _ := c.cfg.Resolve(providerCfg.APIKey)

576 apiKey, _ := c.cfg.Resolve(providerCfg.APIKey)

576 baseURL, _ := c.cfg.Resolve(providerCfg.BaseURL)

577 baseURL, _ := c.cfg.Resolve(providerCfg.BaseURL)

577− var provider fantasy.Provider

578+

578 switch providerCfg.Type {

579 switch providerCfg.Type {

579 case openai.Name:

580 case openai.Name:

580− provider = c.buildOpenaiProvider(baseURL, apiKey, headers)

581+ return c.buildOpenaiProvider(baseURL, apiKey, headers)

581 case anthropic.Name:

582 case anthropic.Name:

582− provider = c.buildAnthropicProvider(baseURL, apiKey, headers)

583+ return c.buildAnthropicProvider(baseURL, apiKey, headers)

583 case openrouter.Name:

584 case openrouter.Name:

584− provider = c.buildOpenrouterProvider(baseURL, apiKey, headers)

585+ return c.buildOpenrouterProvider(baseURL, apiKey, headers)

585 case azure.Name:

586 case azure.Name:

586− provider = c.buildAzureProvider(baseURL, apiKey, headers, providerCfg.ExtraParams)

587+ return c.buildAzureProvider(baseURL, apiKey, headers, providerCfg.ExtraParams)

587 case google.Name:

588 case google.Name:

588− provider = c.buildGoogleProvider(baseURL, apiKey, headers)

589+ return c.buildGoogleProvider(baseURL, apiKey, headers)

589− // this is not in fantasy since its just the google provider with extra stuff

590 case "vertexai":

591− provider = c.buildGoogleVertexProvider(headers, providerCfg.ExtraParams)

591+ return c.buildGoogleVertexProvider(headers, providerCfg.ExtraParams)

592 case openaicompat.Name:

593− provider = c.buildOpenaiCompatProvider(baseURL, apiKey, headers)

593+ return c.buildOpenaiCompatProvider(baseURL, apiKey, headers)

594 default:

595 return nil, errors.New("provider type not supported")

596 }

597− return provider, nil

598}

597}

599

598

600func (c *coordinator) Cancel(sessionID string) {

599func (c *coordinator) Cancel(sessionID string) {

@@ -621,9 +620,9 @@ func (c *coordinator) Model() Model {

621 return c.currentAgent.Model()

620 return c.currentAgent.Model()

622}

621}

623

622

624−func (c *coordinator) UpdateModels() error {

623+func (c *coordinator) UpdateModels(ctx context.Context) error {

625 // build the models again so we make sure we get the latest config

624 // build the models again so we make sure we get the latest config

626− large, small, err := c.buildAgentModels()

625+ large, small, err := c.buildAgentModels(ctx)

627 if err != nil {

626 if err != nil {

628 return err

627 return err

629 }

628 }

@@ -634,7 +633,7 @@ func (c *coordinator) UpdateModels() error {

634 return errors.New("coder agent not configured")

633 return errors.New("coder agent not configured")

635 }

634 }

636

635

637− tools, err := c.buildTools(agentCfg)

636+ tools, err := c.buildTools(ctx, agentCfg)

638 if err != nil {

637 if err != nil {

639 return err

638 return err

640 }

639 }

internal/app/app.go+5−4

@@ -87,7 +87,7 @@ func New(ctx context.Context, conn *sql.DB, cfg *config.Config) (*App, error) {

88 // TODO: remove the concept of agent config, most likely.

89 if cfg.IsConfigured() {

90− if err := app.InitCoderAgent(); err != nil {

90+ if err := app.InitCoderAgent(ctx); err != nil {

91 return nil, fmt.Errorf("failed to initialize coder agent: %w", err)

92 }

93 } else {

@@ -207,8 +207,8 @@ func (app *App) RunNonInteractive(ctx context.Context, prompt string, quiet bool

207 }

208}

209

210−func (app *App) UpdateAgentModel() error {

210+func (app *App) UpdateAgentModel(ctx context.Context) error {

211− return app.AgentCoordinator.UpdateModels()

211+ return app.AgentCoordinator.UpdateModels(ctx)

212}

213

214func (app *App) setupEvents() {

@@ -262,13 +262,14 @@ func setupSubscriber[T any](

262 })

263}

264

265−func (app *App) InitCoderAgent() error {

265+func (app *App) InitCoderAgent(ctx context.Context) error {

266 coderAgentCfg := app.config.Agents[config.AgentCoder]

267 if coderAgentCfg.ID == "" {

268 return fmt.Errorf("coder agent configuration is missing")

269 }

270 var err error

271 app.AgentCoordinator, err = agent.NewCoordinator(

272+ ctx,

272 app.config,

273 app.config,

273 app.Sessions,

274 app.Sessions,

274 app.Messages,

275 app.Messages,

internal/tui/page/chat/chat.go+3−3

@@ -346,7 +346,7 @@ func (p *chatPage) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

346 p.splashFullScreen = true

347 return p, p.SetSize(p.width, p.height)

348 }

349− err := p.app.InitCoderAgent()

349+ err := p.app.InitCoderAgent(context.TODO())

350 if err != nil {

351 return p, util.ReportError(err)

352 }

@@ -538,7 +538,7 @@ func (p *chatPage) toggleThinking() tea.Cmd {

538 cfg.Models[agentCfg.Model] = currentModel

539

540 // Update the agent with the new configuration

541− if err := p.app.UpdateAgentModel(); err != nil {

541+ if err := p.app.UpdateAgentModel(context.TODO()); err != nil {

542 return util.InfoMsg{

543 Type: util.InfoTypeError,

544 Msg: "Failed to update thinking mode: " + err.Error(),

@@ -589,7 +589,7 @@ func (p *chatPage) handleReasoningEffortSelected(effort string) tea.Cmd {

589 }

590

591 // Update the agent with the new configuration

592− if err := p.app.UpdateAgentModel(); err != nil {

592+ if err := p.app.UpdateAgentModel(context.TODO()); err != nil {

593 return util.InfoMsg{

594 Type: util.InfoTypeError,

595 Msg: "Failed to update reasoning effort: " + err.Error(),

internal/tui/tui.go+1−1

@@ -204,7 +204,7 @@ func (a *appModel) Update(msg tea.Msg) (tea.Model, tea.Cmd) {

204 config.Get().UpdatePreferredModel(msg.ModelType, msg.Model)

205

206 // Update the agent with the new model/provider configuration

207− if err := a.app.UpdateAgentModel(); err != nil {

207+ if err := a.app.UpdateAgentModel(context.TODO()); err != nil {

208 return a, util.ReportError(fmt.Errorf("model changed to %s but failed to update agent: %v", msg.Model.Model, err))

209 }

210