chore: limit tool output

Kujtim Hoxha created

Change summary

internal/llm/tools/tools.go | 75 ++++++++++++++++++++++++++++++++++++++
1 file changed, 74 insertions(+), 1 deletion(-)

Detailed changes

internal/llm/tools/tools.go 🔗

@@ -3,6 +3,8 @@ package tools
 import (
 	"context"
 	"encoding/json"
+	"fmt"
+	"strings"
 )
 
 type ToolInfo struct {
@@ -25,6 +27,10 @@ const (
 
 	SessionIDContextKey sessionIDContextKey = "session_id"
 	MessageIDContextKey messageIDContextKey = "message_id"
+
+	MaxResponseWidth  = 3000
+	MaxResponseHeight = 5000
+	MaxResponseChars  = 50000
 )
 
 type ToolResponse struct {
@@ -37,10 +43,77 @@ type ToolResponse struct {
 func NewTextResponse(content string) ToolResponse {
 	return ToolResponse{
 		Type:    ToolResponseTypeText,
-		Content: content,
+		Content: truncateContent(content),
 	}
 }
 
+func truncateContent(content string) string {
+	if len(content) <= MaxResponseChars {
+		return truncateWidthAndHeight(content)
+	}
+
+	truncated := content[:MaxResponseChars]
+
+	if lastNewline := strings.LastIndex(truncated, "\n"); lastNewline > MaxResponseChars/2 {
+		truncated = truncated[:lastNewline]
+	}
+
+	truncated += "\n\n... [Content truncated due to length] ..."
+
+	return truncateWidthAndHeight(truncated)
+}
+
+func truncateWidthAndHeight(content string) string {
+	lines := strings.Split(content, "\n")
+
+	heightTruncated := false
+	if len(lines) > MaxResponseHeight {
+		keepLines := MaxResponseHeight - 3
+		firstHalf := keepLines / 2
+		secondHalf := keepLines - firstHalf
+
+		truncatedLines := make([]string, 0, MaxResponseHeight)
+		truncatedLines = append(truncatedLines, lines[:firstHalf]...)
+		truncatedLines = append(truncatedLines, "")
+		truncatedLines = append(truncatedLines, fmt.Sprintf("... [%d lines truncated] ...", len(lines)-keepLines))
+		truncatedLines = append(truncatedLines, "")
+		truncatedLines = append(truncatedLines, lines[len(lines)-secondHalf:]...)
+
+		lines = truncatedLines
+		heightTruncated = true
+	}
+
+	widthTruncated := false
+	for i, line := range lines {
+		if len(line) > MaxResponseWidth {
+			if MaxResponseWidth > 20 {
+				keepChars := MaxResponseWidth - 10
+				firstHalf := keepChars / 2
+				secondHalf := keepChars - firstHalf
+				lines[i] = line[:firstHalf] + " ... " + line[len(line)-secondHalf:]
+			} else {
+				lines[i] = line[:MaxResponseWidth]
+			}
+			widthTruncated = true
+		}
+	}
+
+	result := strings.Join(lines, "\n")
+
+	if heightTruncated || widthTruncated {
+		notices := []string{}
+		if heightTruncated {
+			notices = append(notices, "height")
+		}
+		if widthTruncated {
+			notices = append(notices, "width")
+		}
+		result += fmt.Sprintf("\n\n[Note: Content truncated by %s to fit response limits]", strings.Join(notices, " and "))
+	}
+
+	return result
+}
+
 func WithResponseMetadata(response ToolResponse, metadata any) ToolResponse {
 	if metadata != nil {
 		metadataBytes, err := json.Marshal(metadata)