llms: fix memory and goroutine leaks in GoogleAI/Vertex and OpenAI streaming (#1364)

* llms: fix memory leaks by adding Close methods and improving goroutine handling

Add Close methods to GoogleAI and Vertex clients to properly clean up resources
and prevent memory leaks from underlying gRPC connections. Update examples to
demonstrate proper client cleanup with defer statements.

Improve goroutine management in OpenAI streaming response handling with proper
context cancellation and non-blocking channel operations to prevent goroutine
leaks in error scenarios.

* examples: add replace directives for testing memory leak fixes

Add temporary replace directives to GoogleAI and Vertex example go.mod files
to test the Close() method changes from the previous commit. These directives
allow the examples to use the local version of the library with the memory
leak fixes before they're included in the next release.

Author: tmc

examples/googleai-completion-example/go.mod

@@ -4,6 +4,10 @@ go 1.24.3
 
 require github.com/tmc/langchaingo v0.1.14-pre.0
 
+// Temporary replace directive for testing Close() method changes
+// TODO: Remove after next release includes the Close() methods
+replace github.com/tmc/langchaingo => ../..
+
 require (
 	cloud.google.com/go v0.116.0 // indirect
 	cloud.google.com/go/ai v0.7.0 // indirect

examples/googleai-completion-example/googleai-completion-example.go

@@ -18,6 +18,7 @@ func main() {
 	if err != nil {
 		log.Fatal(err)
 	}
+	defer llm.Close() // Clean up client when done
 
 	prompt := "Who was the second person to walk on the moon?"
 	answer, err := llms.GenerateFromSinglePrompt(ctx, llm, prompt)

examples/vertex-completion-example/go.mod

@@ -4,6 +4,10 @@ go 1.24.3
 
 require github.com/tmc/langchaingo v0.1.14-pre.0
 
+// Temporary replace directive for testing Close() method changes
+// TODO: Remove after next release includes the Close() methods
+replace github.com/tmc/langchaingo => ../..
+
 require (
 	cloud.google.com/go v0.116.0 // indirect
 	cloud.google.com/go/ai v0.7.0 // indirect

examples/vertex-completion-example/vertex-completion-example.go

@@ -30,6 +30,7 @@ func main() {
 	if err != nil {
 		log.Fatal(err)
 	}
+	defer llm.Close() // Clean up client when done
 
 	prompt := "Who was the second person to walk on the moon?"
 	answer, err := llms.GenerateFromSinglePrompt(ctx, llm, prompt)

llms/googleai/new.go

@@ -39,3 +39,13 @@ func New(ctx context.Context, opts ...Option) (*GoogleAI, error) {
 	gi.client = client
 	return gi, nil
 }
+
+// Close closes the underlying genai client.
+// This should be called when the GoogleAI instance is no longer needed
+// to prevent memory leaks from the underlying gRPC connections.
+func (g *GoogleAI) Close() error {
+	if g.client != nil {
+		return g.client.Close()
+	}
+	return nil
+}

llms/googleai/vertex/new.go

@@ -58,3 +58,15 @@ func New(ctx context.Context, opts ...googleai.Option) (*Vertex, error) {
 	}
 	return v, nil
 }
+
+// Close closes the underlying genai and palm clients.
+// This should be called when the Vertex instance is no longer needed
+// to prevent memory leaks from the underlying gRPC connections.
+func (v *Vertex) Close() error {
+	var err error
+	if v.client != nil {
+		err = v.client.Close()
+	}
+	// Note: palmClient doesn't have a Close method based on the codebase
+	return err
+}

llms/openai/internal/openaiclient/chat.go

@@ -438,9 +438,21 @@ func parseStreamingChatResponse(ctx context.Context, r *http.Response, payload *
 ) { //nolint:cyclop,lll
 	scanner := bufio.NewScanner(r.Body)
 	responseChan := make(chan StreamedChatResponsePayload)
+	
+	// Create a context that can be cancelled to stop the goroutine
+	readerCtx, cancelReader := context.WithCancel(ctx)
+	defer cancelReader()
+	
 	go func() {
 		defer close(responseChan)
 		for scanner.Scan() {
+			// Check if context is cancelled
+			select {
+			case <-readerCtx.Done():
+				return
+			default:
+			}
+			
 			line := scanner.Text()
 			if line == "" {
 				continue
@@ -459,15 +471,25 @@ func parseStreamingChatResponse(ctx context.Context, r *http.Response, payload *
 				// This allows the stream to continue processing valid JSON chunks
 				continue
 			}
-			responseChan <- streamPayload
+			
+			// Non-blocking send with context check
+			select {
+			case <-readerCtx.Done():
+				return
+			case responseChan <- streamPayload:
+			}
 		}
 		if err := scanner.Err(); err != nil {
-			responseChan <- StreamedChatResponsePayload{Error: fmt.Errorf("error reading streaming response: %w", err)}
+			select {
+			case <-readerCtx.Done():
+				return
+			case responseChan <- StreamedChatResponsePayload{Error: fmt.Errorf("error reading streaming response: %w", err)}:
+			}
 			return
 		}
 	}()
 	// Combine response
-	return combineStreamingChatResponse(ctx, payload, responseChan)
+	return combineStreamingChatResponse(readerCtx, payload, responseChan)
 }
 
 func combineStreamingChatResponse(