download llms.txt from docs when running cog init (#2504)

download llms.txt from docs when running cog init

Author: ethulia

pkg/cli/init-templates/pipeline/main.py

@@ -7,15 +7,18 @@
 flux_schnell = replicate.use("black-forest-labs/flux-schnell")
 claude = replicate.use("anthropic/claude-3.5-haiku")
 
+
 def run(
     prompt: str = Input(description="Describe the image to generate"),
-    seed: int = Input(description="A seed", default=0)
+    seed: int = Input(description="A seed", default=0),
 ) -> Path:
-    detailed_prompt = claude(prompt=f"""
+    detailed_prompt = claude(
+        prompt=f"""
     Generate a detailed prompt for a generative image model that will
     generate a high quality dynamic image based on the following
     theme: {prompt}
-    """)
+    """
+    )
     output_paths = flux_schnell(prompt=detailed_prompt, seed=seed)
 
     return Path(output_paths[0])

pkg/cli/init.go

@@ -3,8 +3,11 @@ package cli
 import (
 	"embed"
 	"fmt"
+	"io"
+	"net/http"
 	"os"
 	"path"
+	"time"
 
 	"github.com/spf13/cobra"
 
@@ -114,9 +117,28 @@ func processTemplateFile(fs embed.FS, templateDir, filename, cwd string) error {
 		return fmt.Errorf("Error creating directory %s: %w", dirPath, err)
 	}
 
-	content, err := fs.ReadFile(path.Join(templateDir, filename))
-	if err != nil {
-		return fmt.Errorf("Error reading %s: %w", filename, err)
+	var content []byte
+
+	// Special handling for AGENTS.md - try to download from Replicate docs
+	if filename == "AGENTS.md" {
+		downloadedContent, err := downloadAgentsFile()
+		if err != nil {
+			console.Infof("Failed to download AGENTS.md: %v", err)
+			console.Infof("Using template version instead...")
+			// Fall back to template version
+			content, err = fs.ReadFile(path.Join(templateDir, filename))
+			if err != nil {
+				return fmt.Errorf("Error reading template %s: %w", filename, err)
+			}
+		} else {
+			content = downloadedContent
+		}
+	} else {
+		// Regular template file processing
+		content, err = fs.ReadFile(path.Join(templateDir, filename))
+		if err != nil {
+			return fmt.Errorf("Error reading %s: %w", filename, err)
+		}
 	}
 
 	if err := os.WriteFile(filePath, content, 0o644); err != nil {
@@ -127,6 +149,31 @@ func processTemplateFile(fs embed.FS, templateDir, filename, cwd string) error {
 	return nil
 }
 
+func downloadAgentsFile() ([]byte, error) {
+	const agentsURL = "https://replicate.com/docs/reference/pipelines/llms.txt"
+
+	client := &http.Client{
+		Timeout: 10 * time.Second,
+	}
+
+	resp, err := client.Get(agentsURL)
+	if err != nil {
+		return nil, fmt.Errorf("%w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		return nil, fmt.Errorf("HTTP %d", resp.StatusCode)
+	}
+
+	content, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to read response body: %w", err)
+	}
+
+	return content, nil
+}
+
 func addPipelineInit(cmd *cobra.Command) {
 	const pipeline = "x-pipeline"
 	cmd.Flags().BoolVar(&pipelineTemplate, pipeline, false, "Initialize a pipeline template")