2 files changed, 147 insertions, 0 deletions
diff --git a/generators/gemini/gemini.go b/generators/gemini/gemini.go
new file mode 100644
index 0000000..8b24475
--- /dev/null
+++ b/generators/gemini/gemini.go
@@ -0,0 +1,67 @@
+package gemini
+
+import (
+	"context"
+	"errors"
+	"thunderbird-ai-compose-server/config"
+
+	"github.com/google/generative-ai-go/genai"
+	"google.golang.org/api/option"
+)
+
+func GenerateResponse(prompt string) (string, error) {
+	ctx := context.Background()
+
+	client, err := genai.NewClient(ctx, option.WithAPIKey(config.Config.APIKey))
+	if err != nil {
+		return "", err
+	}
+
+	defer client.Close()
+
+	model := client.GenerativeModel(config.Config.Model)
+	model.SafetySettings = []*genai.SafetySetting{
+		{
+			Category:  genai.HarmCategoryHarassment,
+			Threshold: genai.HarmBlockNone,
+		},
+		{
+			Category:  genai.HarmCategoryHateSpeech,
+			Threshold: genai.HarmBlockNone,
+		},
+		{
+			Category:  genai.HarmCategorySexuallyExplicit,
+			Threshold: genai.HarmBlockNone,
+		},
+		{
+			Category:  genai.HarmCategoryDangerousContent,
+			Threshold: genai.HarmBlockNone,
+		},
+	}
+
+	resp, err := model.GenerateContent(ctx, genai.Text(prompt))
+	if err != nil {
+		return "", err
+	}
+
+	if len(resp.Candidates) == 0 {
+		return "", errors.New("no response candidates received")
+	}
+
+	candidate := resp.Candidates[0]
+
+	if candidate.Content == nil {
+		return "", errors.New("no content in the response candidate")
+	}
+
+	if len(candidate.Content.Parts) == 0 {
+		return "", errors.New("no parts in the response content")
+	}
+
+	if textPart, ok := candidate.Content.Parts[0].(genai.Text); ok {
+		result := string(textPart)
+		return result, nil
+	}
+
+	return "", errors.New("response part is not text")
+}
diff --git a/generators/prompt.go b/generators/prompt.go
new file mode 100644
index 0000000..8e7fbd1
--- /dev/null
+++ b/generators/prompt.go
@@ -0,0 +1,80 @@
+package generators
+
+import (
+	"strings"
+	"thunderbird-ai-compose-server/types"
+)
+
+const systemPrompt = `You are a professional email writing assistant integrated into the Thunderbird email client.
+Your goal is to help users compose high-quality emails based on their prompts and context. Follow these guidelines:
+1. Understand the user's prompt and the context of the email, including any provided details such as subject, recipients, and body content.
+2. Understand the writing style the user is going for, whether formal, informal, concise, or detailed.
+3. Emails can be in plain text or HTML format. Ensure the output matches the specified format. If HTML is requested, use appropriate HTML tags for formatting.
+4. You only need to generate the email body or the reply in case of a reply.
+5. If the user provides a draft or partial content, use it as a base and enhance it according to the prompt.
+6. Ensure the email is clear, coherent, and free of grammatical errors.
+7. If the prompt is vague or lacks details, make reasonable assumptions based on common email practices.
+8. Always maintain a professional tone unless the prompt specifies otherwise.
+9. You only need to output plain text or HTML content of the email body. Do not include any additional explanations or code blocks like ` + "```" + `html` + `, etc.
+10. You should not include HTML syntax like DOCTYPE or <head> or <html> or <body> tags in the output. Only include the content that goes inside the <body> tag if HTML is requested, like <p>, <br>, <b>, <i>, etc.
+
+Examples:
+- If the user wants to write a formal business email, ensure the tone is professional and polite.
+- For informal emails to friends or family, a casual and friendly tone is appropriate.
+- When replying to an email, ensure you address all points raised in the original message.
+
+Remember, your primary objective is to assist users in crafting effective emails that meet their needs and expectations.
+
+Following context is provided along with every request automatically:
+- Account details (email, name) of the user sending the email. You can use this to personalize the email.
+- The subject of the email. You can use this to ensure the email is relevant to the topic. The subject can also be empty.
+- The people the email is being sent to (To, Cc, Bcc). Further personalization can be done if needed using this information. This information can also be empty.
+- The body of the email in plain text. This can also be empty. If this is a reply, this will contain the original email being replied to.
+- The body of the email in HTML (if applicable). This can also be empty. If this is a reply, this will contain the HTML of the original email being replied to.
+- Identity ID of the user (this is inserted by Thunderbird, and can be ignored).
+- Whether the email is in HTML format or plain text. Use this to determine if you need to generate HTML or plain text.
+
+Use the above guidelines and context to generate high-quality emails that meet the user's requirements.
+`
+
+func BuildPrompt(payload types.Payload) string {
+	var prompt strings.Builder
+
+	prompt.WriteString(systemPrompt)
+	prompt.WriteString("\n\n")
+	prompt.WriteString("User Prompt:\n")
+	prompt.WriteString(payload.Prompt)
+	prompt.WriteString("\n\n")
+	prompt.WriteString("Email Context:\n")
+	if payload.Context.Compose.Subject != "" {
+		prompt.WriteString("Subject: " + payload.Context.Compose.Subject + "\n")
+	}
+	if len(payload.Context.Compose.To) > 0 {
+		prompt.WriteString("To: " + strings.Join(payload.Context.Compose.To, ", ") + "\n")
+	}
+	if len(payload.Context.Compose.Cc) > 0 {
+		prompt.WriteString("Cc: " + strings.Join(payload.Context.Compose.Cc, ", ") + "\n")
+	}
+	if len(payload.Context.Compose.Bcc) > 0 {
+		prompt.WriteString("Bcc: " + strings.Join(payload.Context.Compose.Bcc, ", ") + "\n")
+	}
+	if payload.Context.Account.Name != "" {
+		prompt.WriteString("Sender Name: " + payload.Context.Account.Name + "\n")
+	}
+	if payload.Context.Account.Email != "" {
+		prompt.WriteString("Sender Email: " + payload.Context.Account.Email + "\n")
+	}
+	if payload.Context.Compose.BodyPlain != "" {
+		prompt.WriteString("Body (Plain Text):\n" + payload.Context.Compose.BodyPlain + "\n")
+	}
+	if payload.Context.Compose.BodyHTML != "" {
+		prompt.WriteString("Body (HTML):\n" + payload.Context.Compose.BodyHTML + "\n")
+	}
+	if payload.Context.Compose.IsHTML {
+		prompt.WriteString("Is HTML: true. The email should be in HTML format.\n")
+	} else {
+		prompt.WriteString("Is HTML: false. The email should be in plain text format.\n")
+	}
+
+	return prompt.String()
+}