Skip to content

Commit 0d04b99

Browse files
Copilotsawka
andauthored
Add Google AI file summarization package (#2455)
- [x] Create new directory pkg/aiusechat/google - [x] Implement SummarizeFile function with: - Context parameter for timeout - File validation (images, PDFs, text files only) - Use gemini-2.5-flash-lite model - Configurable API URL and prompt as constants - Return (string, usage, error) - [x] Define Google-specific usage struct - [x] Test the implementation (all tests pass) - [x] Verify with existing linting and build - [x] Run CodeQL security check (no issues found) - [x] Revert unintended tsunami demo dependency changes ## Summary Successfully implemented a new Google AI package at `pkg/aiusechat/google` with: 1. **SummarizeFile function** - A simple request-response API (not streaming, not SSE) - Takes context for timeout - Validates file types (images, PDFs, text only) - Enforces file size limits matching wshcmd-ai.go - Uses gemini-2.5-flash-lite model - Returns (summary string, usage stats, error) 2. **GoogleUsage struct** - Tracks token consumption: - PromptTokenCount - CachedContentTokenCount - CandidatesTokenCount - TotalTokenCount 3. **Configurable constants**: - GoogleAPIURL (for reference) - SummarizePrompt (customizable prompt) - SummarizeModel (gemini-2.5-flash-lite) 4. **Comprehensive tests** - 41.7% coverage with all tests passing 5. **Security verified** - No CodeQL alerts 6. **Package documentation** - doc.go with usage examples Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: sawka <2722291+sawka@users.noreply.github.com>
1 parent 2619c85 commit 0d04b99

File tree

4 files changed

+558
-0
lines changed

4 files changed

+558
-0
lines changed
Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
// Copyright 2025, Command Line Inc.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package main
5+
6+
import (
7+
"context"
8+
"flag"
9+
"fmt"
10+
"os"
11+
"time"
12+
13+
"github.com/wavetermdev/waveterm/pkg/aiusechat/google"
14+
)
15+
16+
func printUsage() {
17+
fmt.Println("Usage: go run main-testsummarize.go [--help] [--mode MODE] <filename>")
18+
fmt.Println("Examples:")
19+
fmt.Println(" go run main-testsummarize.go README.md")
20+
fmt.Println(" go run main-testsummarize.go --mode useful /path/to/image.png")
21+
fmt.Println(" go run main-testsummarize.go -m publiccode document.pdf")
22+
fmt.Println("")
23+
fmt.Println("Supported file types:")
24+
fmt.Println(" - Text files (up to 200KB)")
25+
fmt.Println(" - Images (up to 7MB)")
26+
fmt.Println(" - PDFs (up to 5MB)")
27+
fmt.Println("")
28+
fmt.Println("Flags:")
29+
fmt.Println(" --mode, -m Summarization mode (default: quick)")
30+
fmt.Println(" Options: quick, useful, publiccode, htmlcontent, htmlfull")
31+
fmt.Println("")
32+
fmt.Println("Environment variables:")
33+
fmt.Println(" GOOGLE_APIKEY (required)")
34+
}
35+
36+
func main() {
37+
var showHelp bool
38+
var mode string
39+
flag.BoolVar(&showHelp, "help", false, "Show usage information")
40+
flag.StringVar(&mode, "mode", "quick", "Summarization mode")
41+
flag.StringVar(&mode, "m", "quick", "Summarization mode (shorthand)")
42+
flag.Parse()
43+
44+
if showHelp {
45+
printUsage()
46+
os.Exit(0)
47+
}
48+
49+
apiKey := os.Getenv("GOOGLE_APIKEY")
50+
if apiKey == "" {
51+
fmt.Println("Error: GOOGLE_APIKEY environment variable not set")
52+
printUsage()
53+
os.Exit(1)
54+
}
55+
56+
args := flag.Args()
57+
if len(args) == 0 {
58+
fmt.Println("Error: filename required")
59+
printUsage()
60+
os.Exit(1)
61+
}
62+
63+
filename := args[0]
64+
65+
// Check if file exists
66+
if _, err := os.Stat(filename); os.IsNotExist(err) {
67+
fmt.Printf("Error: file '%s' does not exist\n", filename)
68+
os.Exit(1)
69+
}
70+
71+
ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
72+
defer cancel()
73+
74+
fmt.Printf("Summarizing file: %s\n", filename)
75+
fmt.Printf("Model: %s\n", google.SummarizeModel)
76+
fmt.Printf("Mode: %s\n", mode)
77+
78+
startTime := time.Now()
79+
summary, usage, err := google.SummarizeFile(ctx, filename, google.SummarizeOpts{
80+
APIKey: apiKey,
81+
Mode: mode,
82+
})
83+
latency := time.Since(startTime)
84+
85+
fmt.Printf("Latency: %d ms\n", latency.Milliseconds())
86+
fmt.Println("===")
87+
if err != nil {
88+
fmt.Printf("Error: %v\n", err)
89+
os.Exit(1)
90+
}
91+
92+
fmt.Println("\nSummary:")
93+
fmt.Println("---")
94+
fmt.Println(summary)
95+
fmt.Println("---")
96+
97+
if usage != nil {
98+
fmt.Println("\nUsage Statistics:")
99+
fmt.Printf(" Prompt tokens: %d\n", usage.PromptTokenCount)
100+
fmt.Printf(" Cached tokens: %d\n", usage.CachedContentTokenCount)
101+
fmt.Printf(" Response tokens: %d\n", usage.CandidatesTokenCount)
102+
fmt.Printf(" Total tokens: %d\n", usage.TotalTokenCount)
103+
}
104+
}

pkg/aiusechat/google/doc.go

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
// Copyright 2025, Command Line Inc.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
// Package google provides Google Generative AI integration for WaveTerm.
5+
//
6+
// This package implements file summarization using Google's Gemini models.
7+
// Unlike other AI provider implementations in the aiusechat package, this
8+
// package does NOT implement full SSE streaming. It uses a simple
9+
// request-response API for file summarization.
10+
//
11+
// # Supported File Types
12+
//
13+
// The package supports the same file types as defined in wshcmd-ai.go:
14+
// - Images (PNG, JPEG, etc.): up to 7MB
15+
// - PDFs: up to 5MB
16+
// - Text files: up to 200KB
17+
//
18+
// Binary files are rejected unless they are recognized as images or PDFs.
19+
//
20+
// # Usage
21+
//
22+
// To summarize a file:
23+
//
24+
// ctx := context.Background()
25+
// summary, usage, err := google.SummarizeFile(ctx, "/path/to/file.txt", google.SummarizeOpts{
26+
// APIKey: "YOUR_API_KEY",
27+
// Mode: google.ModeQuickSummary,
28+
// })
29+
// if err != nil {
30+
// log.Fatal(err)
31+
// }
32+
// fmt.Println("Summary:", summary)
33+
// fmt.Printf("Tokens used: %d\n", usage.TotalTokenCount)
34+
//
35+
// # Configuration
36+
//
37+
// The summarization behavior can be customized by modifying the constants:
38+
// - SummarizeModel: The Gemini model to use (default: "gemini-2.5-flash-lite")
39+
// - SummarizePrompt: The prompt sent to the model
40+
// - GoogleAPIURL: The base URL for the API (for reference, not currently used by the SDK)
41+
package google

0 commit comments

Comments
 (0)