refactor AI response generation and enhance IMAP connection handling
This commit is contained in:
@@ -48,52 +48,63 @@ func (a *AI) GenerateReply(emailContent string, contextContent map[string]string
|
||||
}
|
||||
|
||||
// Prepare the system message and user message
|
||||
systemMsg := "You are a helpful assistant. Use the provided context to help answer the email professionally and accurately."
|
||||
userMsg := fmt.Sprintf("Using the following context:\n%s\n\nPlease generate a reply for this email:\n%s", context, emailContent)
|
||||
systemMsg := "You are a helpful assistant responding to emails. Analyze the language of the incoming email and translate the response to the same language. Format your response in HTML. Do not include any explanations or extra text - just write the email response directly in HTML format. Use appropriate HTML tags for formatting."
|
||||
userMsg := fmt.Sprintf("Using the following context:\n%s\n\nPlease analyze the language of and generate a reply in the same language for this email:\n%s", context, emailContent)
|
||||
|
||||
messages := []Message{
|
||||
{Role: "system", Content: systemMsg},
|
||||
{Role: "user", Content: userMsg},
|
||||
}
|
||||
|
||||
reqBody := OpenRouterRequest{
|
||||
Model: a.config.Model,
|
||||
Messages: messages,
|
||||
Temperature: a.config.Temperature,
|
||||
MaxTokens: a.config.MaxTokens,
|
||||
const maxRetries = 3
|
||||
var lastErr error
|
||||
|
||||
for attempt := 0; attempt < maxRetries; attempt++ {
|
||||
reqBody := OpenRouterRequest{
|
||||
Model: a.config.Model,
|
||||
Messages: messages,
|
||||
Temperature: a.config.Temperature,
|
||||
MaxTokens: a.config.MaxTokens,
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("POST", "https://openrouter.ai/api/v1/chat/completions", bytes.NewBuffer(jsonData))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", a.config.OpenRouterAPIKey))
|
||||
|
||||
resp, err := a.client.Do(req)
|
||||
if err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
lastErr = fmt.Errorf("OpenRouter API returned status code: %d", resp.StatusCode)
|
||||
continue
|
||||
}
|
||||
|
||||
var result OpenRouterResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
lastErr = err
|
||||
continue
|
||||
}
|
||||
|
||||
if len(result.Choices) == 0 || result.Choices[0].Message.Content == "" {
|
||||
lastErr = fmt.Errorf("empty response received")
|
||||
continue
|
||||
}
|
||||
|
||||
return result.Choices[0].Message.Content, nil
|
||||
}
|
||||
|
||||
jsonData, err := json.Marshal(reqBody)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
req, err := http.NewRequest("POST", "https://openrouter.ai/api/v1/chat/completions", bytes.NewBuffer(jsonData))
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", a.config.OpenRouterAPIKey))
|
||||
|
||||
resp, err := a.client.Do(req)
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
if resp.StatusCode != http.StatusOK {
|
||||
return "", fmt.Errorf("OpenRouter API returned status code: %d", resp.StatusCode)
|
||||
}
|
||||
|
||||
var result OpenRouterResponse
|
||||
if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
|
||||
return "", err
|
||||
}
|
||||
|
||||
if len(result.Choices) == 0 {
|
||||
return "", fmt.Errorf("no response generated")
|
||||
}
|
||||
|
||||
return result.Choices[0].Message.Content, nil
|
||||
return "", fmt.Errorf("failed after %d attempts, last error: %v", maxRetries, lastErr)
|
||||
}
|
||||
|
||||
@@ -1,8 +1,12 @@
|
||||
package imap
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"fmt"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
"net/mail"
|
||||
"paraclub-ai-mailer/config"
|
||||
"strings"
|
||||
"time"
|
||||
@@ -20,10 +24,19 @@ type IMAPClient struct {
|
||||
// MessageLiteral implements imap.Literal for draft messages
|
||||
type MessageLiteral struct {
|
||||
content []byte
|
||||
pos int
|
||||
}
|
||||
|
||||
func (m *MessageLiteral) Read(p []byte) (n int, err error) {
|
||||
return copy(p, m.content), io.EOF
|
||||
if m.pos >= len(m.content) {
|
||||
return 0, io.EOF
|
||||
}
|
||||
n = copy(p, m.content[m.pos:])
|
||||
m.pos += n
|
||||
if m.pos >= len(m.content) {
|
||||
err = io.EOF
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (m *MessageLiteral) Len() int {
|
||||
@@ -63,10 +76,63 @@ type Email struct {
|
||||
Body string
|
||||
}
|
||||
|
||||
func (ic *IMAPClient) FetchUnprocessedEmails() ([]Email, error) {
|
||||
_, err := ic.client.Select(ic.config.MailboxIn, false)
|
||||
func (ic *IMAPClient) ensureConnection() error {
|
||||
if ic.client == nil {
|
||||
return ic.reconnect()
|
||||
}
|
||||
|
||||
// Try to check connection by NOOPing
|
||||
if err := ic.client.Noop(); err != nil {
|
||||
return ic.reconnect()
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (ic *IMAPClient) reconnect() error {
|
||||
addr := fmt.Sprintf("%s:%d", ic.config.Server, ic.config.Port)
|
||||
var c *client.Client
|
||||
var err error
|
||||
|
||||
if ic.config.UseTLS {
|
||||
c, err = client.DialTLS(addr, nil)
|
||||
} else {
|
||||
c, err = client.Dial(addr)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return fmt.Errorf("failed to connect to IMAP server: %v", err)
|
||||
}
|
||||
|
||||
if err := c.Login(ic.config.Username, ic.config.Password); err != nil {
|
||||
c.Close()
|
||||
return fmt.Errorf("failed to login: %v", err)
|
||||
}
|
||||
|
||||
if ic.client != nil {
|
||||
ic.client.Close()
|
||||
}
|
||||
ic.client = c
|
||||
return nil
|
||||
}
|
||||
|
||||
func (ic *IMAPClient) FetchUnprocessedEmails() ([]Email, error) {
|
||||
if err := ic.ensureConnection(); err != nil {
|
||||
return nil, fmt.Errorf("failed to ensure IMAP connection: %v", err)
|
||||
}
|
||||
|
||||
// Make sure we have a valid mailbox name
|
||||
if ic.config.MailboxIn == "" {
|
||||
return nil, fmt.Errorf("mailbox_in not configured")
|
||||
}
|
||||
|
||||
// Select mailbox and get status
|
||||
mbox, err := ic.client.Select(ic.config.MailboxIn, false)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to select mailbox %s: %v", ic.config.MailboxIn, err)
|
||||
}
|
||||
|
||||
// If mailbox is empty, return early
|
||||
if mbox.Messages == 0 {
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
criteria := imap.NewSearchCriteria()
|
||||
@@ -74,7 +140,7 @@ func (ic *IMAPClient) FetchUnprocessedEmails() ([]Email, error) {
|
||||
|
||||
uids, err := ic.client.Search(criteria)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
return nil, fmt.Errorf("search failed: %v", err)
|
||||
}
|
||||
|
||||
if len(uids) == 0 {
|
||||
@@ -84,35 +150,54 @@ func (ic *IMAPClient) FetchUnprocessedEmails() ([]Email, error) {
|
||||
seqSet := new(imap.SeqSet)
|
||||
seqSet.AddNum(uids...)
|
||||
|
||||
messages := make(chan *imap.Message, 10)
|
||||
// Fetch both envelope and body
|
||||
section := &imap.BodySectionName{Peek: true}
|
||||
items := []imap.FetchItem{imap.FetchEnvelope, section.FetchItem()}
|
||||
|
||||
messages := make(chan *imap.Message, 10)
|
||||
done := make(chan error, 1)
|
||||
|
||||
go func() {
|
||||
done <- ic.client.Fetch(seqSet, []imap.FetchItem{section.FetchItem()}, messages)
|
||||
done <- ic.client.Fetch(seqSet, items, messages)
|
||||
}()
|
||||
|
||||
var emails []Email
|
||||
for msg := range messages {
|
||||
// Skip if message doesn't have an envelope
|
||||
if msg.Envelope == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// Skip if already processed
|
||||
if ic.processed[msg.Envelope.MessageId] {
|
||||
continue
|
||||
}
|
||||
|
||||
// Skip if no message ID
|
||||
if msg.Envelope.MessageId == "" {
|
||||
continue
|
||||
}
|
||||
|
||||
r := msg.GetBody(section)
|
||||
if r == nil {
|
||||
continue
|
||||
}
|
||||
|
||||
var bodyBuilder strings.Builder
|
||||
_, err := io.Copy(&bodyBuilder, r)
|
||||
if err != nil {
|
||||
if _, err := io.Copy(&bodyBuilder, r); err != nil {
|
||||
continue
|
||||
}
|
||||
|
||||
// Make sure we have a valid from address
|
||||
var from string
|
||||
if len(msg.Envelope.From) > 0 {
|
||||
from = msg.Envelope.From[0].Address()
|
||||
}
|
||||
|
||||
emails = append(emails, Email{
|
||||
ID: msg.Envelope.MessageId,
|
||||
Subject: msg.Envelope.Subject,
|
||||
From: msg.Envelope.From[0].Address(),
|
||||
From: from,
|
||||
Body: bodyBuilder.String(),
|
||||
})
|
||||
|
||||
@@ -120,42 +205,137 @@ func (ic *IMAPClient) FetchUnprocessedEmails() ([]Email, error) {
|
||||
}
|
||||
|
||||
if err := <-done; err != nil {
|
||||
return nil, err
|
||||
return nil, fmt.Errorf("fetch failed: %v", err)
|
||||
}
|
||||
|
||||
return emails, nil
|
||||
}
|
||||
|
||||
func (ic *IMAPClient) SaveDraft(email Email, response string) error {
|
||||
_, err := ic.client.Select(ic.config.DraftBox, false)
|
||||
if err != nil {
|
||||
return err
|
||||
if err := ic.ensureConnection(); err != nil {
|
||||
return fmt.Errorf("failed to ensure IMAP connection: %v", err)
|
||||
}
|
||||
|
||||
// Format the draft message
|
||||
_, err := ic.client.Select(ic.config.DraftBox, false)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to select draft box: %v", err)
|
||||
}
|
||||
|
||||
// Format the draft message with HTML response and original email headers + content
|
||||
draft := fmt.Sprintf("From: %s\r\n"+
|
||||
"To: %s\r\n"+
|
||||
"Subject: Re: %s\r\n"+
|
||||
"Content-Type: text/plain; charset=UTF-8\r\n"+
|
||||
"\r\n%s",
|
||||
"MIME-Version: 1.0\r\n"+
|
||||
"Content-Type: text/html; charset=UTF-8\r\n"+
|
||||
"\r\n"+
|
||||
"%s\r\n"+
|
||||
"<br><br>\r\n"+
|
||||
"<div style=\"border-top: 1px solid #B5C4DF; padding-top: 10px; margin-top: 10px;\">\r\n"+
|
||||
"<div style=\"font-family: Arial, sans-serif; color: #666666; margin-bottom: 10px;\">\r\n"+
|
||||
"<b>From:</b> %s<br>\r\n"+
|
||||
"<b>Subject:</b> %s<br>\r\n"+
|
||||
"<b>Date:</b> %s<br>\r\n"+
|
||||
"</div>\r\n"+
|
||||
"<div style=\"margin-left: 10px; padding-left: 10px; border-left: 2px solid #B5C4DF;\">\r\n"+
|
||||
"%s\r\n"+
|
||||
"</div></div>",
|
||||
ic.config.Username,
|
||||
email.From,
|
||||
email.Subject,
|
||||
response)
|
||||
response,
|
||||
email.From,
|
||||
email.Subject,
|
||||
time.Now().Format("Mon, 02 Jan 2006 15:04:05 -0700"),
|
||||
extractMessageContent(email.Body))
|
||||
|
||||
// Create literal message
|
||||
literal := &MessageLiteral{content: []byte(draft)}
|
||||
literal := &MessageLiteral{
|
||||
content: []byte(draft),
|
||||
pos: 0,
|
||||
}
|
||||
|
||||
// Save the draft
|
||||
flags := []string{"\\Draft"}
|
||||
if err := ic.client.Append(ic.config.DraftBox, flags, time.Now(), literal); err != nil {
|
||||
return err
|
||||
return fmt.Errorf("failed to append draft: %v", err)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// extractMessageContent attempts to extract just the message content
|
||||
// by removing email headers and MIME boundaries
|
||||
func extractMessageContent(body string) string {
|
||||
msg, err := mail.ReadMessage(strings.NewReader(body))
|
||||
if err != nil {
|
||||
// Fallback to previous method if parsing fails
|
||||
return fallbackExtractContent(body)
|
||||
}
|
||||
|
||||
mediaType, params, err := mime.ParseMediaType(msg.Header.Get("Content-Type"))
|
||||
if err != nil {
|
||||
return fallbackExtractContent(body)
|
||||
}
|
||||
|
||||
if strings.HasPrefix(mediaType, "multipart/") {
|
||||
boundary := params["boundary"]
|
||||
if boundary == "" {
|
||||
return fallbackExtractContent(body)
|
||||
}
|
||||
|
||||
reader := multipart.NewReader(msg.Body, boundary)
|
||||
var textContent string
|
||||
|
||||
for {
|
||||
part, err := reader.NextPart()
|
||||
if err == io.EOF {
|
||||
break
|
||||
}
|
||||
if err != nil {
|
||||
return fallbackExtractContent(body)
|
||||
}
|
||||
|
||||
// Check Content-Type of the part
|
||||
contentType := part.Header.Get("Content-Type")
|
||||
if strings.HasPrefix(contentType, "text/plain") {
|
||||
buf := new(bytes.Buffer)
|
||||
_, err := buf.ReadFrom(part)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
textContent = buf.String()
|
||||
break
|
||||
}
|
||||
}
|
||||
|
||||
if textContent != "" {
|
||||
return strings.TrimSpace(textContent)
|
||||
}
|
||||
}
|
||||
|
||||
// For non-multipart messages, just read the body
|
||||
buf := new(bytes.Buffer)
|
||||
_, err = buf.ReadFrom(msg.Body)
|
||||
if err != nil {
|
||||
return fallbackExtractContent(body)
|
||||
}
|
||||
|
||||
return strings.TrimSpace(buf.String())
|
||||
}
|
||||
|
||||
// fallbackExtractContent is the previous implementation used as fallback
|
||||
func fallbackExtractContent(body string) string {
|
||||
parts := strings.Split(body, "\r\n\r\n")
|
||||
if len(parts) > 1 {
|
||||
return strings.TrimSpace(strings.Join(parts[1:], "\r\n\r\n"))
|
||||
}
|
||||
return strings.TrimSpace(body)
|
||||
}
|
||||
|
||||
func (ic *IMAPClient) MarkAsProcessed(email Email) error {
|
||||
if err := ic.ensureConnection(); err != nil {
|
||||
return fmt.Errorf("failed to ensure IMAP connection: %v", err)
|
||||
}
|
||||
|
||||
_, err := ic.client.Select(ic.config.MailboxIn, false)
|
||||
if err != nil {
|
||||
return err
|
||||
|
||||
Reference in New Issue
Block a user