Paraclub/ai-mailer
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: change max token handling

Browse Source
This commit is contained in:
Dominik Polakovics Polakovics
2025-11-13 11:51:12 +01:00
parent b7e4c8d6ac
commit e88ac7caff
8 changed files with 255 additions and 91 deletions
Download Patch File Download Diff File Expand all files Collapse all files

54
internal/ai/ai.go
View File

@@ -43,6 +43,31 @@ func New(cfg config.AIConfig) *AI {
}
}
// BuildSystemPrompt creates the system prompt for a given language
// Exposed for token counting without making an API call
func (a *AI) BuildSystemPrompt(lang string) string {
return fmt.Sprintf(`You are a helpful assistant who responds to emails.
Your primary goal is to answer the user's query (found in the 'Email Body') by primarily using the information available in the 'Additional Context' and your general knowledge.
While the 'Email Body' provides the question, your answer should be synthesized from the context and your understanding, not by directly repeating or solely relying on unverified information from the 'Email Body' itself.
Instructions:
- Language: Your response must be entirely in %s, regardless of the language used in the email content or context.
- Format: CRITICAL: Your reply MUST be raw HTML. Use appropriate HTML tags for structure and styling. For example, wrap paragraphs in <p>...</p> tags and use <br> for line breaks if needed within a paragraph. Even a short sentence must be wrapped in HTML (e.g., <p>Yes.</p>).
- Markdown: Do NOT wrap the HTML in markdown code blocks (e.g., %s).
- Extraneous Text: Do not include a subject line. Do not include explanations, commentary, or any extra text that is not part of the direct answer.
- Closing: Avoid generic closing statements like "If you have further questions...". Focus solely on answering the email.
`, lang, "```html ... ```")
}
// BuildUserPrompt creates the user message with context and email content
func (a *AI) BuildUserPrompt(contextContent map[string]string, emailContent string) string {
var context string
for url, content := range contextContent {
context += fmt.Sprintf("\nContext from %s:\n%s\n", url, content)
}
return fmt.Sprintf("### Additional Context:\n%s\n\n### Email Body:\n%s", context, emailContent)
}
func (a *AI) detectLanguage(emailContent string) (string, error) {
logger.WithField("emailContentLength", len(emailContent)).Debug("Starting language detection")
@@ -77,32 +102,11 @@ func (a *AI) GenerateReply(emailContent string, contextContent map[string]string
return "", err
}
// Prepare context from all URLs
var context string
for url, content := range contextContent {
context += fmt.Sprintf("\nContext from %s:\n%s\n", url, content)
logger.WithFields(logrus.Fields{
"url": url,
"contentLength": len(content),
}).Debug("Added context from URL")
}
// Build prompts using exposed methods
systemMsg := a.BuildSystemPrompt(lang)
userMsg := a.BuildUserPrompt(contextContent, emailContent)
// Prepare the system message with language-specific instruction
systemMsg := fmt.Sprintf(`You are a helpful assistant who responds to emails.
Your primary goal is to answer the user's query (found in the 'Email Body') by primarily using the information available in the 'Additional Context' and your general knowledge.
While the 'Email Body' provides the question, your answer should be synthesized from the context and your understanding, not by directly repeating or solely relying on unverified information from the 'Email Body' itself.
Instructions:
- Language: Your response must be entirely in %s, regardless of the language used in the email content or context.
- Format: CRITICAL: Your reply MUST be raw HTML. Use appropriate HTML tags for structure and styling. For example, wrap paragraphs in <p>...</p> tags and use <br> for line breaks if needed within a paragraph. Even a short sentence must be wrapped in HTML (e.g., <p>Yes.</p>).
- Markdown: Do NOT wrap the HTML in markdown code blocks (e.g., %s).
- Extraneous Text: Do not include a subject line. Do not include explanations, commentary, or any extra text that is not part of the direct answer.
- Closing: Avoid generic closing statements like "If you have further questions...". Focus solely on answering the email.
`, lang, "```html ... ```")
logger.WithFields(logrus.Fields{
"systemprompt": systemMsg,
}).Debug("Generating system prompt")
userMsg := fmt.Sprintf("### Additional Context:\n%s\n\n### Email Body:\n%s", context, emailContent)
logger.Debug("Generated system and user prompts")
messages := []Message{
{Role: "system", Content: systemMsg},

125
internal/imap/imap.go
View File

@@ -373,7 +373,8 @@ func handleMultipartMessage(reader io.Reader, boundary string) string {
}
mReader := multipart.NewReader(reader, boundary)
var textContent string
var textPlainContent string
var textHTMLContent string
partIndex := 0
for {
@@ -387,39 +388,75 @@ func handleMultipartMessage(reader io.Reader, boundary string) string {
}
contentType := part.Header.Get("Content-Type")
contentDisposition := part.Header.Get("Content-Disposition")
contentTransferEncoding := strings.ToLower(part.Header.Get("Content-Transfer-Encoding"))
logger.WithFields(logrus.Fields{
"partIndex": partIndex,
"partContentType": contentType,
"partDisposition": contentDisposition,
"partTransferEncoding": contentTransferEncoding,
"partHeaders": part.Header,
}).Debug("handleMultipartMessage: Processing part")
if strings.HasPrefix(contentType, "text/plain") {
var partReader io.Reader = part
if contentTransferEncoding == "quoted-printable" {
partReader = quotedprintable.NewReader(part)
}
// Add handling for "base64" if needed in the future
// else if contentTransferEncoding == "base64" {
// partReader = base64.NewDecoder(base64.StdEncoding, part)
// }
buf := new(bytes.Buffer)
if _, err := buf.ReadFrom(partReader); err != nil {
logger.WithError(err).WithField("partIndex", partIndex).Debug("Failed to read from partReader in multipart")
continue // Or handle error more robustly
}
textContent = buf.String()
// Assuming we only care about the first text/plain part found
// If multiple text/plain parts could exist and need concatenation, this logic would need adjustment.
break
// Skip attachments
if strings.HasPrefix(contentDisposition, "attachment") {
logger.WithFields(logrus.Fields{
"partIndex": partIndex,
"filename": part.FileName(),
}).Debug("Skipping attachment part")
partIndex++
continue
}
// Skip non-text content types (images, videos, applications, etc.)
if !strings.HasPrefix(contentType, "text/plain") &&
!strings.HasPrefix(contentType, "text/html") {
logger.WithFields(logrus.Fields{
"partIndex": partIndex,
"contentType": contentType,
}).Debug("Skipping non-text content type")
partIndex++
continue
}
// Handle quoted-printable encoding
var partReader io.Reader = part
if contentTransferEncoding == "quoted-printable" {
partReader = quotedprintable.NewReader(part)
}
buf := new(bytes.Buffer)
if _, err := buf.ReadFrom(partReader); err != nil {
logger.WithError(err).WithField("partIndex", partIndex).Debug("Failed to read from partReader in multipart")
partIndex++
continue
}
// Store text/plain and text/html separately
if strings.HasPrefix(contentType, "text/plain") {
textPlainContent = buf.String()
logger.WithField("textPlainLength", len(textPlainContent)).Debug("Found text/plain part")
} else if strings.HasPrefix(contentType, "text/html") {
textHTMLContent = buf.String()
logger.WithField("textHTMLLength", len(textHTMLContent)).Debug("Found text/html part")
}
partIndex++
}
logger.WithField("textContentResult", textContent).Debug("handleMultipartMessage: Returning textContent")
return textContent
// Prefer text/plain over text/html
if textPlainContent != "" {
logger.Debug("handleMultipartMessage: Returning text/plain content")
return textPlainContent
}
if textHTMLContent != "" {
logger.Debug("handleMultipartMessage: Converting text/html to plain text")
return htmlToPlainText(textHTMLContent)
}
logger.Debug("handleMultipartMessage: No text content found")
return ""
}
func handleSinglePartMessage(reader io.Reader) string {
@@ -434,6 +471,48 @@ func handleSinglePartMessage(reader io.Reader) string {
return content
}
// htmlToPlainText converts HTML content to plain text by extracting text nodes
func htmlToPlainText(htmlContent string) string {
logger.WithField("htmlLength", len(htmlContent)).Debug("Converting HTML to plain text")
// Simple HTML tag stripping - removes all HTML tags and extracts text
var result strings.Builder
inTag := false
for _, char := range htmlContent {
switch char {
case '<':
inTag = true
case '>':
inTag = false
default:
if !inTag {
result.WriteRune(char)
}
}
}
plainText := result.String()
// Clean up excessive whitespace
plainText = strings.ReplaceAll(plainText, "\r\n", "\n")
plainText = strings.ReplaceAll(plainText, "\r", "\n")
// Replace multiple spaces with single space
lines := strings.Split(plainText, "\n")
var cleanLines []string
for _, line := range lines {
trimmed := strings.TrimSpace(line)
if trimmed != "" {
cleanLines = append(cleanLines, trimmed)
}
}
result2 := strings.Join(cleanLines, "\n")
logger.WithField("plainTextLength", len(result2)).Debug("HTML converted to plain text")
return result2
}
func cleanMessageContent(content string, performHeaderStripping bool) string {
logger.WithField("inputContentLength", len(content)).Debug("cleanMessageContent: Starting")
logger.WithField("performHeaderStripping", performHeaderStripping).Debug("cleanMessageContent: performHeaderStripping flag")

50
internal/tokens/tokens.go Normal file
View File

@@ -0,0 +1,50 @@
package tokens
import (
"fmt"
"paraclub-ai-mailer/internal/logger"
"github.com/pkoukk/tiktoken-go"
"github.com/sirupsen/logrus"
)
type TokenCounter struct {
encoding *tiktoken.Tiktoken
}
// New creates a token counter with cl100k_base encoding (GPT-4/Claude compatible)
func New() (*TokenCounter, error) {
enc, err := tiktoken.GetEncoding("cl100k_base")
if err != nil {
return nil, fmt.Errorf("failed to get tiktoken encoding: %w", err)
}
return &TokenCounter{encoding: enc}, nil
}
// CountString counts tokens in a single string
func (tc *TokenCounter) CountString(text string) int {
tokens := tc.encoding.Encode(text, nil, nil)
return len(tokens)
}
// EstimateFullRequest estimates total tokens for the complete API request
// Includes: system prompt + user prompt + message overhead
func (tc *TokenCounter) EstimateFullRequest(systemPrompt, userPrompt string) int {
systemTokens := tc.CountString(systemPrompt)
userTokens := tc.CountString(userPrompt)
// Add overhead for message structure (~100 tokens for JSON formatting, role labels, etc.)
overhead := 100
total := systemTokens + userTokens + overhead
logger.WithFields(logrus.Fields{
"systemTokens": systemTokens,
"userTokens": userTokens,
"overhead": overhead,
"total": total,
}).Debug("Token estimation breakdown")
return total
}