Server directory structure and streaming optimization

server/main.go

@@ -1,151 +0,0 @@
-package main
-
-import (
-	"context"
-	"fmt"
-	"log"
-	"net/http"
-	"os"
-	"strconv"
-	"strings"
-
-	"github.com/joho/godotenv"
-	"github.com/labstack/echo/v4"
-	"github.com/labstack/echo/v4/middleware"
-	"github.com/tmc/langchaingo/llms"
-	"github.com/tmc/langchaingo/llms/ollama"
-)
-
-type CodeBlockParser struct {
-	channel      chan []byte
-	language     string
-	maxLines     int
-	currentLines int
-}
-
-func NewCodeBlockParser(lang string, lines int) *CodeBlockParser {
-	return &CodeBlockParser{
-		channel:      make(chan []byte),
-		language:     lang,
-		maxLines:     lines,
-		currentLines: 0,
-	}
-}
-
-func (parser *CodeBlockParser) ParseStream(chunk []byte) {
-	text := string(chunk)
-
-	if !strings.Contains(parser.language, text) && !strings.Contains(text, "```") && !strings.Contains(text, "``") {
-		if strings.Contains(text, "\n") {
-			parser.currentLines += 1
-		}
-
-		if parser.currentLines == parser.maxLines-1 {
-			indexOfNewLine := strings.Index(text, "\n")
-
-			if indexOfNewLine > -1 {
-				parser.channel <- []byte(text[:indexOfNewLine])
-				return
-			}
-
-			parser.channel <- nil
-			return
-		}
-
-		if parser.currentLines >= parser.maxLines {
-			parser.channel <- nil
-			return
-		}
-
-		parser.channel <- chunk
-		return
-	} else {
-		parser.channel <- nil
-		return
-	}
-}
-
-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())
-	ech.GET("/generate", func(ctx echo.Context) error {
-		// TODO: Make lines an environment variable and tweak it along with the prompt to get a suitable number of lines for the challenge to be fun and not too hard but still challenging.
-		lines, err := strconv.Atoi(ctx.QueryParam("lines"))
-
-		if err != nil {
-			return ctx.String(http.StatusBadRequest, "Lines param is not provided or incorrect!")
-		}
-
-		lang := ctx.QueryParam("lang")
-
-		if lang == "" {
-			return ctx.String(http.StatusBadRequest, "Lang param is not provided or incorrect!")
-		}
-
-		ctx.Response().Header().Set(echo.HeaderContentType, echo.MIMEApplicationJSON)
-		ctx.Response().WriteHeader(http.StatusOK)
-
-		llm, err := ollama.New(ollama.WithModel(fmt.Sprintf("%s", LLM_MODEL)))
-
-		if err != nil {
-			log.Fatal(err)
-		}
-
-		parser := NewCodeBlockParser(lang, lines)
-		ollamaCtx := context.Background()
-		prompt := []llms.MessageContent{
-			llms.TextParts(llms.ChatMessageTypeHuman, fmt.Sprintf(`
-				You must only generate code without any text descriptions or code comments and use spaces instead of tabs for spacing. 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, prompt, llms.WithStreamingFunc(func(streamCtx context.Context, chunk []byte) error {
-			go parser.ParseStream(chunk)
-
-			select {
-			case chunk := <-parser.channel:
-				if len(chunk) > 0 {
-					ctx.Response().Write(chunk)
-					ctx.Response().Flush()
-				}
-
-				if parser.currentLines == parser.maxLines {
-					cnx, _, err := ctx.Response().Hijack()
-
-					if err != nil {
-						log.Fatal(err)
-						return ctx.String(http.StatusInternalServerError, err.Error())
-					}
-
-					cnx.Close()
-				}
-			}
-			return nil
-		})); err != nil {
-			log.Fatal(err)
-			return ctx.String(http.StatusInternalServerError, err.Error())
-		}
-
-		defer close(parser.channel)
-		return nil
-	})
-
-	ech.Logger.Fatal(ech.Start(fmt.Sprintf(":%s", PORT)))
-}