Some improvements

server/main.go

@@ -5,8 +5,10 @@ import (
 	"fmt"
 	"log"
 	"net/http"
+	"os"
 	"strconv"
 
+	"github.com/joho/godotenv"
 	"github.com/labstack/echo/v4"
 	"github.com/labstack/echo/v4/middleware"
 	"github.com/tmc/langchaingo/llms"
@@ -14,6 +16,23 @@ import (
 )
 
 func main() {
+	err := godotenv.Load()
+
+	if err != nil {
+		log.Fatal("Error loading environment!")
+	}
+
+	LLM_MODEL := os.Getenv("LLM_MODEL")
+	PORT := os.Getenv("PORT")
+
+	if len(LLM_MODEL) == 0 {
+		log.Fatal("No LLM model specified in environment!")
+	}
+
+	if len(PORT) == 0 {
+		PORT = "8080"
+	}
+
 	ech := echo.New()
 
 	ech.Use(middleware.CORS())
@@ -33,17 +52,18 @@ func main() {
 		ctx.Response().Header().Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
 		ctx.Response().WriteHeader(http.StatusOK)
 
-		llm, err := ollama.New(ollama.WithModel("llama3.1:8b"))
+		llm, err := ollama.New(ollama.WithModel(fmt.Sprintf("%s", LLM_MODEL)))
 
 		if err != nil {
 			log.Fatal(err)
 		}
 
+		// TODO: Clean the prompt result of any unnecessary formatting or text
 		ollamaCtx := context.Background()
 		content := []llms.MessageContent{
-			llms.TextParts(llms.ChatMessageTypeSystem, `You are only a code generator. You must not respond with anything else but code and do not format with code fences. Always use spaces instead of tabs.`),
+			llms.TextParts(llms.ChatMessageTypeSystem, `You must only generate code without any descriptions. Also don't include any code comments and use spaces instead of tabs for spacing. Most importantly. Most importantly you must remove any markdown code fences that wrap the content!`),
 			llms.TextParts(llms.ChatMessageTypeHuman, fmt.Sprintf(`
-				Generate max %d lines of code without any unncessary formatting from a well known open source project in the %s programming language.`, lines, lang)),
+				Generate a maximum of %d lines of code from a well known open source project in the %s programming language.`, lines, lang)),
 		}
 
 		if _, err := llm.GenerateContent(ollamaCtx, content, llms.WithStreamingFunc(func(streamCtx context.Context, chunk []byte) error {
@@ -58,5 +78,5 @@ func main() {
 		return nil
 	})
 
-	ech.Logger.Fatal(ech.Start(":5000"))
+	ech.Logger.Fatal(ech.Start(fmt.Sprintf(":%s", PORT)))
 }