Add base code

internal/llm/text_rewriter.go

@@ -0,0 +1,99 @@
+/*
+Copyright © 2024 Matteo Schiff <matteo@underdesk.net>
+
+*/
+
+package llm
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"log"
+	"os"
+	"time"
+
+	"git.underdesk.net/Matte23/transcriber/internal/utils"
+	"github.com/sashabaranov/go-openai"
+)
+
+func RewriteText(client *openai.Client, paragraphs []ParagraphItem, finalFile string) {
+	outputFile, err := os.Create(finalFile)
+	if err != nil {
+		panic(err)
+	}
+	// close output file on exit and check for its returned error
+	defer func() {
+		if err := outputFile.Close(); err != nil {
+			panic(err)
+		}
+	}()
+
+	systemPrompt := openai.ChatCompletionMessage{
+		Role: openai.ChatMessageRoleSystem,
+		Content: `You are provided with:
+
+    A transcription paragraph to rewrite, aimed at improving clarity, grammar, and flow while preserving the original meaning and details.
+    A JSON structure of the document that lists section and subsection titles, showing how this paragraph fits within the larger structure.
+    The current built document, containing all previously written sections up to this point.
+
+Your task: Rewrite the transcription paragraph using clear and polished language, while keeping all key information intact. Format the paragraph using Markdown for improved readability (e.g., bold for emphasis, bullet points if applicable, etc.). Do not add any new information or leave out any critical details. Focus solely on rewriting the paragraph provided in the transcription, and do not add headers, titles, extra context, or explanations beyond the paragraph's Markdown-formatted text.
+
+Only write the Markdown-formatted paragraph text in your response.`,
+	}
+
+	structureJson, _ := json.Marshal(paragraphs)
+
+	documentStructure := openai.ChatCompletionMessage{
+		Role:    openai.ChatMessageRoleUser,
+		Content: "Document structure is: " + string(structureJson),
+	}
+
+	currentText := ""
+
+	log.Println("Starting rewriting text")
+	for _, currentMessage := range paragraphs {
+		startTime := time.Now()
+		heading := ""
+		switch currentMessage.Type {
+		case "heading1":
+			heading += "# "
+		case "heading2":
+			heading += "## "
+		default:
+			heading += "### "
+		}
+		heading += currentMessage.Title + "\n"
+		currentText += heading
+		outputFile.Write([]byte(heading))
+
+		currentTranscription := openai.ChatCompletionMessage{
+			Role:    openai.ChatMessageRoleUser,
+			Content: "Paragraph to rewrite is: '" + currentMessage.Content + "'",
+		}
+
+		currentDocument := openai.ChatCompletionMessage{
+			Role:    openai.ChatMessageRoleUser,
+			Content: "Current rewritten document is: \n\n" + currentText,
+		}
+
+		resp1, err := client.CreateChatCompletion(
+			context.Background(),
+			openai.ChatCompletionRequest{
+				Model:    "mistral-small-instruct",
+				Messages: []openai.ChatCompletionMessage{systemPrompt, documentStructure, currentDocument, currentTranscription},
+			},
+		)
+
+		if err != nil {
+			fmt.Printf("LLM process error: %v\n", err)
+			return
+		}
+		result := resp1.Choices[0].Message.Content
+
+		outputFile.Write([]byte(result + "\n\n"))
+		currentText += result + "\n\n"
+
+		utils.MeasureTime(startTime, "Text rewrite iteration for "+currentMessage.Title)
+	}
+}