From b75ad9067a063e09e9769761f270f468e22e2c8a Mon Sep 17 00:00:00 2001
From: dwrz <dwrz@dwrz.net>
Date: Tue, 2 Jun 2026 12:02:49 +0000
Subject: [PATCH] Add llm client max iterations

---
 internal/llm/llm.go         | 24 +++++++++++++++++++-----
 internal/service/service.go |  4 ++--
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/internal/llm/llm.go b/internal/llm/llm.go
index bd685dc..985b59a 100644
--- a/internal/llm/llm.go
+++ b/internal/llm/llm.go
@@ -120,9 +120,15 @@ func (c *Client) DefaultModel() string {
 	return c.model
 }
 
+// ErrMaxIterations is returned when the agent loop exceeds the configured
+// maximum number of iterations.
+var ErrMaxIterations = fmt.Errorf("max iterations exceeded")
+
 // Query sends messages to the LLM using the specified model.
 // If model is empty, uses the default configured model.
 // If systemMessage is non-empty, it overrides the configured system message.
+// maxIterations caps the number of agent loop iterations; a value of 0
+// means unlimited (no cap).
 // Returns all messages generated during the query, including tool calls
 // and tool results. The final message is the last element in the slice.
 func (c *Client) Query(
@@ -130,6 +136,7 @@ func (c *Client) Query(
 	messages []openai.ChatCompletionMessage,
 	model string,
 	systemMessage string,
+	maxIterations int,
 ) ([]openai.ChatCompletionMessage, error) {
 	ctx, cancel := context.WithTimeout(ctx, c.timeout)
 	defer cancel()
@@ -160,8 +167,8 @@ func (c *Client) Query(
 	// Track messages generated during this query.
 	var generated []openai.ChatCompletionMessage
 
-	// Loop for tool calls.
-	for {
+	// Loop for tool calls. maxIterations == 0 means unlimited.
+	for i := 0; maxIterations == 0 || i < maxIterations; i++ {
 		req := openai.ChatCompletionRequest{
 			Model:    model,
 			Messages: messages,
@@ -236,6 +243,8 @@ func (c *Client) Query(
 		}
 		// Loop to get LLM's response after tool execution.
 	}
+
+	return generated, ErrMaxIterations
 }
 
 // StreamEvent wraps a ChatCompletionMessage produced during streaming.
@@ -248,12 +257,15 @@ type StreamEvent struct {
 // tool result) is sent to the events channel as it becomes available.
 // The channel is closed before returning.
 // If systemMessage is non-empty, it overrides the configured system message.
-// Returns all messages generated during the query.
+// maxIterations caps the number of agent loop iterations; a value of 0
+// means unlimited (no cap).
+// Returns ErrMaxIterations if the iteration limit is exceeded.
 func (c *Client) QueryStream(
 	ctx context.Context,
 	messages []openai.ChatCompletionMessage,
 	model string,
 	systemMessage string,
+	maxIterations int,
 	events chan<- StreamEvent,
 ) error {
 	defer close(events)
@@ -284,8 +296,8 @@ func (c *Client) QueryStream(
 		)
 	}
 
-	// Loop for tool calls.
-	for {
+	// Loop for tool calls. maxIterations == 0 means unlimited.
+	for i := 0; maxIterations == 0 || i < maxIterations; i++ {
 		req := openai.ChatCompletionRequest{
 			Model:    model,
 			Messages: messages,
@@ -423,4 +435,6 @@ func (c *Client) QueryStream(
 		}
 		// Loop to get LLM's response after tool execution.
 	}
+
+	return ErrMaxIterations
 }
diff --git a/internal/service/service.go b/internal/service/service.go
index 6acf622..93e5d4f 100644
--- a/internal/service/service.go
+++ b/internal/service/service.go
@@ -453,7 +453,7 @@ func (svc *Service) voice(w http.ResponseWriter, r *http.Request) {
 	if model == "" {
 		model = svc.llm.DefaultModel()
 	}
-	msgs, err := svc.llm.Query(ctx, messages, model, req.SystemMessage)
+	msgs, err := svc.llm.Query(ctx, messages, model, req.SystemMessage, 0)
 	if err != nil {
 		log.ErrorContext(
 			ctx,
@@ -718,7 +718,7 @@ func (svc *Service) voiceStream(w http.ResponseWriter, r *http.Request) {
 		llmErr error
 	)
 	go func() {
-		llmErr = svc.llm.QueryStream(ctx, messages, model, req.SystemMessage, events)
+		llmErr = svc.llm.QueryStream(ctx, messages, model, req.SystemMessage, 0, events)
 	}()
 
 	// Consume events and send as SSE.