Endpoint for streaming generated code

server/go.mod

@@ -20,4 +20,5 @@ require (
 	golang.org/x/net v0.32.0 // indirect
 	golang.org/x/sys v0.28.0 // indirect
 	golang.org/x/text v0.21.0 // indirect
+	golang.org/x/time v0.8.0 // indirect
 )

server/go.sum

@@ -37,6 +37,8 @@ golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA=
 golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo=
 golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ=
+golang.org/x/time v0.8.0 h1:9i3RxcPv3PZnitoVGMPDKZSq1xW1gK1Xy3ArNOGZfEg=
+golang.org/x/time v0.8.0/go.mod h1:3BpzKBy/shNhVucY/MWOyx10tF3SFh9QdLuxbVysPQM=
 gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
 gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=

server/main.go

@@ -2,25 +2,37 @@ package main
 
 import (
 	"context"
-	"encoding/json"
+	"fmt"
 	"log"
 	"net/http"
+	"strconv"
 
 	"github.com/labstack/echo/v4"
+	"github.com/labstack/echo/v4/middleware"
 	"github.com/tmc/langchaingo/llms"
 	"github.com/tmc/langchaingo/llms/ollama"
 )
 
-type GeneratedCode struct {
-	Code     string
-	Source   string
-	Language string
-}
-
 func main() {
 	ech := echo.New()
 
+	ech.Use(middleware.CORS())
 	ech.GET("/generate", func(ctx echo.Context) error {
+		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("llama3.1:8b"))
 
 		if err != nil {
@@ -29,26 +41,21 @@ func main() {
 
 		ollamaCtx := context.Background()
 		content := []llms.MessageContent{
-			llms.TextParts(llms.ChatMessageTypeHuman, `
+			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.`),
-				Generate 5 lines of code using a well known open source project
+			llms.TextParts(llms.ChatMessageTypeHuman, fmt.Sprintf(`
-				in the JavaScript programming language. I would like the result to be in
+				Generate max %d lines of code without any unncessary formatting from a well known open source project in the %s programming language. First line should always be a code comment in this format: Language/Project`, lines, lang)),
-				JSON format with the following keys: code for source code, source for project name and language for the used language.
-			`),
 		}
 
-		generated, err := llm.GenerateContent(ollamaCtx, content, llms.WithJSONMode())
+		if _, err := llm.GenerateContent(ollamaCtx, content, llms.WithStreamingFunc(func(streamCtx context.Context, chunk []byte) error {
-
+			ctx.Response().Write(chunk)
-		var generatedCode GeneratedCode
+			ctx.Response().Flush()
-
+			return nil
-		error := json.Unmarshal([]byte(generated.Choices[0].Content), &generatedCode)
+		})); err != nil {
-
+			log.Fatal(err)
-		if error != nil {
+			return ctx.String(http.StatusInternalServerError, err.Error())
-			log.Fatal(error)
 		}
 
-		return ctx.JSON(http.StatusOK,
+		return nil
-			generatedCode,
-		)
 	})
 
 	ech.Logger.Fatal(ech.Start(":5000"))