Add Google AI file summarization package

Copilot · sawka · Copilot · commit 98c7fd3cf28d · 2025-10-17T19:26:37.000Z
Co-authored-by: sawka &lt;2722291+sawka@users.noreply.github.com&gt;
diff --git a/pkg/aiusechat/google/doc.go b/pkg/aiusechat/google/doc.go
@@ -0,0 +1,38 @@
+// Copyright 2025, Command Line Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+// Package google provides Google Generative AI integration for WaveTerm.
+//
+// This package implements file summarization using Google's Gemini models.
+// Unlike other AI provider implementations in the aiusechat package, this
+// package does NOT implement full SSE streaming. It uses a simple
+// request-response API for file summarization.
+//
+// # Supported File Types
+//
+// The package supports the same file types as defined in wshcmd-ai.go:
+//   - Images (PNG, JPEG, etc.): up to 7MB
+//   - PDFs: up to 5MB
+//   - Text files: up to 200KB
+//
+// Binary files are rejected unless they are recognized as images or PDFs.
+//
+// # Usage
+//
+// To summarize a file:
+//
+//	ctx := context.Background()
+//	summary, usage, err := google.SummarizeFile(ctx, "/path/to/file.txt", "YOUR_API_KEY")
+//	if err != nil {
+//	    log.Fatal(err)
+//	}
+//	fmt.Println("Summary:", summary)
+//	fmt.Printf("Tokens used: %d\n", usage.TotalTokenCount)
+//
+// # Configuration
+//
+// The summarization behavior can be customized by modifying the constants:
+//   - SummarizeModel: The Gemini model to use (default: "gemini-2.5-flash-lite")
+//   - SummarizePrompt: The prompt sent to the model
+//   - GoogleAPIURL: The base URL for the API (for reference, not currently used by the SDK)
+package google
diff --git a/pkg/aiusechat/google/google-summarize.go b/pkg/aiusechat/google/google-summarize.go
@@ -0,0 +1,151 @@
+// Copyright 2025, Command Line Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package google
+
+import (
+	"context"
+	"fmt"
+	"net/http"
+	"os"
+	"strings"
+
+	"github.com/google/generative-ai-go/genai"
+	"github.com/wavetermdev/waveterm/pkg/util/utilfn"
+	"google.golang.org/api/option"
+)
+
+const (
+	// GoogleAPIURL is the base URL for the Google Generative AI API
+	GoogleAPIURL = "https://generativelanguage.googleapis.com"
+
+	// SummarizePrompt is the prompt used for file summarization
+	SummarizePrompt = "Please provide a concise summary of this file. Include the main topics, key points, and any notable information."
+
+	// SummarizeModel is the model used for file summarization
+	SummarizeModel = "gemini-2.5-flash-lite"
+)
+
+// GoogleUsage represents token usage information from Google's Generative AI API
+type GoogleUsage struct {
+	PromptTokenCount        int32 `json:"prompt_token_count"`
+	CachedContentTokenCount int32 `json:"cached_content_token_count"`
+	CandidatesTokenCount    int32 `json:"candidates_token_count"`
+	TotalTokenCount         int32 `json:"total_token_count"`
+}
+
+func detectMimeType(data []byte) string {
+	mimeType := http.DetectContentType(data)
+	return strings.Split(mimeType, ";")[0]
+}
+
+func getMaxFileSize(mimeType string) (int, string) {
+	if mimeType == "application/pdf" {
+		return 5 * 1024 * 1024, "5MB"
+	}
+	if strings.HasPrefix(mimeType, "image/") {
+		return 7 * 1024 * 1024, "7MB"
+	}
+	return 200 * 1024, "200KB"
+}
+
+// SummarizeFile reads a file and generates a summary using Google's Generative AI.
+// It supports images, PDFs, and text files based on the limits defined in wshcmd-ai.go.
+// Returns the summary text, usage information, and any error encountered.
+func SummarizeFile(ctx context.Context, filename string, apiKey string) (string, *GoogleUsage, error) {
+	// Read the file
+	data, err := os.ReadFile(filename)
+	if err != nil {
+		return "", nil, fmt.Errorf("reading file: %w", err)
+	}
+
+	// Detect MIME type
+	mimeType := detectMimeType(data)
+
+	isPDF := mimeType == "application/pdf"
+	isImage := strings.HasPrefix(mimeType, "image/")
+
+	if !isPDF && !isImage {
+		mimeType = "text/plain"
+		if utilfn.ContainsBinaryData(data) {
+			return "", nil, fmt.Errorf("file contains binary data and cannot be summarized")
+		}
+	}
+
+	// Validate file size
+	maxSize, sizeStr := getMaxFileSize(mimeType)
+	if len(data) > maxSize {
+		return "", nil, fmt.Errorf("file exceeds maximum size of %s for %s files", sizeStr, mimeType)
+	}
+
+	// Create client
+	client, err := genai.NewClient(ctx, option.WithAPIKey(apiKey))
+	if err != nil {
+		return "", nil, fmt.Errorf("creating Google AI client: %w", err)
+	}
+	defer client.Close()
+
+	// Create model
+	model := client.GenerativeModel(SummarizeModel)
+
+	// Prepare the content parts
+	var parts []genai.Part
+
+	// Add the prompt
+	parts = append(parts, genai.Text(SummarizePrompt))
+
+	// Add the file content based on type
+	if isImage {
+		// For images, use Blob
+		parts = append(parts, genai.Blob{
+			MIMEType: mimeType,
+			Data:     data,
+		})
+	} else if isPDF {
+		// For PDFs, use Blob
+		parts = append(parts, genai.Blob{
+			MIMEType: mimeType,
+			Data:     data,
+		})
+	} else {
+		// For text files, convert to string
+		parts = append(parts, genai.Text(string(data)))
+	}
+
+	// Generate content
+	resp, err := model.GenerateContent(ctx, parts...)
+	if err != nil {
+		return "", nil, fmt.Errorf("generating content: %w", err)
+	}
+
+	// Check if we got any candidates
+	if len(resp.Candidates) == 0 {
+		return "", nil, fmt.Errorf("no response candidates returned")
+	}
+
+	// Extract the text from the first candidate
+	candidate := resp.Candidates[0]
+	if candidate.Content == nil || len(candidate.Content.Parts) == 0 {
+		return "", nil, fmt.Errorf("no content in response")
+	}
+
+	var summary strings.Builder
+	for _, part := range candidate.Content.Parts {
+		if textPart, ok := part.(genai.Text); ok {
+			summary.WriteString(string(textPart))
+		}
+	}
+
+	// Convert usage metadata
+	var usage *GoogleUsage
+	if resp.UsageMetadata != nil {
+		usage = &GoogleUsage{
+			PromptTokenCount:        resp.UsageMetadata.PromptTokenCount,
+			CachedContentTokenCount: resp.UsageMetadata.CachedContentTokenCount,
+			CandidatesTokenCount:    resp.UsageMetadata.CandidatesTokenCount,
+			TotalTokenCount:         resp.UsageMetadata.TotalTokenCount,
+		}
+	}
+
+	return summary.String(), usage, nil
+}
diff --git a/pkg/aiusechat/google/google-summarize_test.go b/pkg/aiusechat/google/google-summarize_test.go
@@ -0,0 +1,167 @@
+// Copyright 2025, Command Line Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package google
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+)
+
+func TestDetectMimeType(t *testing.T) {
+	tests := []struct {
+		name     string
+		data     []byte
+		expected string
+	}{
+		{
+			name:     "plain text",
+			data:     []byte("Hello, World!"),
+			expected: "text/plain",
+		},
+		{
+			name:     "empty file",
+			data:     []byte{},
+			expected: "text/plain",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := detectMimeType(tt.data)
+			if !containsMimeType(result, tt.expected) {
+				t.Errorf("detectMimeType() = %v, want to contain %v", result, tt.expected)
+			}
+		})
+	}
+}
+
+func containsMimeType(got, want string) bool {
+	// DetectContentType may return variations like "text/plain; charset=utf-8"
+	return got == want || (want == "text/plain" && got == "text/plain; charset=utf-8")
+}
+
+func TestGetMaxFileSize(t *testing.T) {
+	tests := []struct {
+		name         string
+		mimeType     string
+		expectedSize int
+		expectedStr  string
+	}{
+		{
+			name:         "PDF file",
+			mimeType:     "application/pdf",
+			expectedSize: 5 * 1024 * 1024,
+			expectedStr:  "5MB",
+		},
+		{
+			name:         "PNG image",
+			mimeType:     "image/png",
+			expectedSize: 7 * 1024 * 1024,
+			expectedStr:  "7MB",
+		},
+		{
+			name:         "JPEG image",
+			mimeType:     "image/jpeg",
+			expectedSize: 7 * 1024 * 1024,
+			expectedStr:  "7MB",
+		},
+		{
+			name:         "text file",
+			mimeType:     "text/plain",
+			expectedSize: 200 * 1024,
+			expectedStr:  "200KB",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			size, sizeStr := getMaxFileSize(tt.mimeType)
+			if size != tt.expectedSize {
+				t.Errorf("getMaxFileSize() size = %v, want %v", size, tt.expectedSize)
+			}
+			if sizeStr != tt.expectedStr {
+				t.Errorf("getMaxFileSize() sizeStr = %v, want %v", sizeStr, tt.expectedStr)
+			}
+		})
+	}
+}
+
+func TestSummarizeFile_FileNotFound(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	_, _, err := SummarizeFile(ctx, "/nonexistent/file.txt", "fake-api-key")
+	if err == nil {
+		t.Error("SummarizeFile() expected error for nonexistent file, got nil")
+	}
+}
+
+func TestSummarizeFile_BinaryFile(t *testing.T) {
+	// Create a temporary binary file
+	tmpDir := t.TempDir()
+	binFile := filepath.Join(tmpDir, "test.bin")
+
+	// Create binary data (not text, image, or PDF)
+	binaryData := []byte{0x00, 0x01, 0x02, 0x03, 0x7F, 0x80, 0xFF}
+	if err := os.WriteFile(binFile, binaryData, 0644); err != nil {
+		t.Fatalf("Failed to create test file: %v", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	_, _, err := SummarizeFile(ctx, binFile, "fake-api-key")
+	if err == nil {
+		t.Error("SummarizeFile() expected error for binary file, got nil")
+	}
+	if err != nil && !containsString(err.Error(), "binary data") {
+		t.Errorf("SummarizeFile() error = %v, want error containing 'binary data'", err)
+	}
+}
+
+func TestSummarizeFile_FileTooLarge(t *testing.T) {
+	// Create a temporary text file that exceeds the limit
+	tmpDir := t.TempDir()
+	textFile := filepath.Join(tmpDir, "large.txt")
+
+	// Create a file larger than 200KB (text file limit)
+	largeData := make([]byte, 201*1024)
+	for i := range largeData {
+		largeData[i] = 'a'
+	}
+	if err := os.WriteFile(textFile, largeData, 0644); err != nil {
+		t.Fatalf("Failed to create test file: %v", err)
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+
+	_, _, err := SummarizeFile(ctx, textFile, "fake-api-key")
+	if err == nil {
+		t.Error("SummarizeFile() expected error for file too large, got nil")
+	}
+	if err != nil && !containsString(err.Error(), "exceeds maximum size") {
+		t.Errorf("SummarizeFile() error = %v, want error containing 'exceeds maximum size'", err)
+	}
+}
+
+func containsString(s, substr string) bool {
+	return len(s) >= len(substr) && (s == substr || len(substr) == 0 ||
+		(len(s) > 0 && len(substr) > 0 && stringContains(s, substr)))
+}
+
+func stringContains(s, substr string) bool {
+	for i := 0; i <= len(s)-len(substr); i++ {
+		if s[i:i+len(substr)] == substr {
+			return true
+		}
+	}
+	return false
+}
+
+// Note: We don't test the actual API call without a real API key
+// Integration tests would require setting GOOGLE_API_KEY environment variable
