feat: adds copilot enabled context refinement for ephemeral contexts too

1x-eng · Nov 8, 2024 · 0b61130 · 0b61130
1 parent 9a49e76
commit 0b61130
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 15 deletions.
diff --git a/pkg/context/manager.go b/pkg/context/manager.go
@@ -21,13 +21,15 @@ type ContextManager struct {
 	au                 aurora.Aurora
 	presenter          *ui.ContextPresenter
 	currentContextFile string
+	llmClient          *llm.PerplexityAI
 }
 
-func NewContextManager(contextDir string, au aurora.Aurora, theme *ui.Theme) *ContextManager {
+func NewContextManager(contextDir string, au aurora.Aurora, theme *ui.Theme, llmClient *llm.PerplexityAI) *ContextManager {
 	return &ContextManager{
 		contextDir: contextDir,
 		au:         au,
 		presenter:  ui.NewContextPresenter(theme),
+		llmClient:  llmClient,
 	}
 }
 
@@ -57,13 +59,7 @@ func (cm *ContextManager) GetSessionContext(llmClient *llm.PerplexityAI) (string
 		return "", err
 	}
 
-	// After getting the context (either from file or input), offer refinement
-	refinedContext, err := cm.RefineContext(context, llmClient)
-	if err != nil {
-		return "", fmt.Errorf("failed to refine context: %w", err)
-	}
-
-	return refinedContext, nil
+	return context, nil
 }
 
 func (cm *ContextManager) getContextFromFile() (string, error) {
@@ -167,19 +163,25 @@ func (cm *ContextManager) getContextFromInput() (string, error) {
 
 	context := strings.Join(lines, "\n")
 
+	refinedContext, err := cm.RefineContext(context, cm.llmClient)
+	if err != nil {
+		fmt.Println(cm.au.Red("\nError during context refinement. Proceeding with original context."))
+		refinedContext = context
+	}
+
 	var saveContext bool
 	prompt := &survey.Confirm{
 		Message: "Would you like to save this context for future sessions?",
 	}
 	survey.AskOne(prompt, &saveContext)
 
 	if saveContext {
-		if err := cm.saveContext(context); err != nil {
+		if err := cm.saveContext(refinedContext); err != nil {
 			fmt.Println(cm.au.Red("Failed to save context:"), err)
 		}
 	}
 
-	return context, nil
+	return refinedContext, nil
 }
 
 func (cm *ContextManager) saveContext(context string) error {

diff --git a/pkg/pomodoro/pom.go b/pkg/pomodoro/pom.go
@@ -79,7 +79,7 @@ func (p *TomatickMemento) StartCycle() {
 	if p.cycleCount == 0 {
 		displayWelcomeMessage(p.auroraInstance)
 
-		contextManager := context.NewContextManager(p.cfg.ContextDir, p.auroraInstance, p.theme)
+		contextManager := context.NewContextManager(p.cfg.ContextDir, p.auroraInstance, p.theme, p.llmClient)
 
 		sessionContext, err := contextManager.GetSessionContext(p.llmClient)
 		if err != nil {
@@ -614,7 +614,7 @@ func displayWelcomeMessage(au aurora.Aurora) {
 	   ██║   ██║   ██║██╔████╔██║███████║   ██║   ██║██║     █████╔╝ 
 	   ██║   ██║   ██║██║╚██╔╝██║██╔══██║   ██║   ██║██║     ██╔═██╗ 
 	   ██║   ╚██████╔╝██║ ╚═╝ ██║██║  ██║   ██║   ██║╚██████╗██║  ██╗
-	   ╚═╝    ╚═══╝ ╚═╝     ╚═╝╚═╝  ╚═╝   ╚═╝   ╚═╝ ╚═════╝╚═╝  ╚═╝
+		╚═╝    ╚═══╝  ╚═╝     ╚═╝╚═╝  ╚═╝   ╚═╝   ╚═╝ ╚═════╝╚═╝  ╚═╝
 	`
 	welcomeText := `
 	🌟 Your Productivity Partner 🌟

diff --git a/pkg/ui/context.go b/pkg/ui/context.go
@@ -193,8 +193,8 @@ func (cp *ContextPresenter) PresentRefinedContext(refinedContext string) string
 
 	cleanContext := strings.TrimSpace(refinedContext)
 	if strings.HasPrefix(strings.ToLower(cleanContext), "context refinement complete") {
-		if idx := strings.Index(cleanContext, "•"); idx != -1 {
-			cleanContext = cleanContext[idx:]
+		if idx := strings.Index(strings.ToLower(cleanContext), "here's the comprehensive breakdown:"); idx != -1 {
+			cleanContext = strings.TrimSpace(cleanContext[idx+len("here's the comprehensive breakdown:"):])
 		}
 	}
 
@@ -206,7 +206,8 @@ Please review the refined context above:
 • Does it capture all key points?
 • Are there any missing details?`
 
-	sb.WriteString("\n" + cp.theme.Styles.SystemInstruction.Render(reviewPrompt))
+	sb.WriteString("\n" + cp.theme.Styles.InfoText.Render(reviewPrompt) + "\n")
+
 	return sb.String()
 }