From 754bc1f7d8b78fc269ab5649e003882cced679d0 Mon Sep 17 00:00:00 2001
From: Hazem Krimi <me@hazemkrimi.tech>
Date: Mon, 13 Jan 2025 19:59:00 +0100
Subject: [PATCH] Endpoint for streaming generated code

---
 server/go.mod  |  1 +
 server/go.sum  |  2 ++
 server/main.go | 53 ++++++++++++++++++++++++++++----------------------
 3 files changed, 33 insertions(+), 23 deletions(-)

diff --git a/server/go.mod b/server/go.mod
index ccd4094..eb8b02d 100644
--- a/server/go.mod
+++ b/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
 )
diff --git a/server/go.sum b/server/go.sum
index 6aaee76..c13f682 100644
--- a/server/go.sum
+++ b/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=
diff --git a/server/main.go b/server/main.go
index c4f07cc..4bbdd68 100644
--- a/server/main.go
+++ b/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, `
-				Generate 5 lines of code using a well known open source project
-				in the JavaScript programming language. I would like the result to be in
-				JSON format with the following keys: code for source code, source for project name and language for the used language.
-			`),
+			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.`),
+			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. First line should always be a code comment in this format: Language/Project`, lines, lang)),
 		}
 
-		generated, err := llm.GenerateContent(ollamaCtx, content, llms.WithJSONMode())
-
-		var generatedCode GeneratedCode
-
-		error := json.Unmarshal([]byte(generated.Choices[0].Content), &generatedCode)
-
-		if error != nil {
-			log.Fatal(error)
+		if _, err := llm.GenerateContent(ollamaCtx, content, llms.WithStreamingFunc(func(streamCtx context.Context, chunk []byte) error {
+			ctx.Response().Write(chunk)
+			ctx.Response().Flush()
+			return nil
+		})); err != nil {
+			log.Fatal(err)
+			return ctx.String(http.StatusInternalServerError, err.Error())
 		}
 
-		return ctx.JSON(http.StatusOK,
-			generatedCode,
-		)
+		return nil
 	})
 
 	ech.Logger.Fatal(ech.Start(":5000"))