Server directory structure and streaming optimization

This commit is contained in:
2025-01-28 01:17:30 +01:00
parent 2674bc6c20
commit 40954e1f52
5 changed files with 176 additions and 153 deletions
+1 -1
View File
@@ -14,7 +14,7 @@ function App() {
(async function() {
setCode('');
const response = await fetch(`${import.meta.env.VITE_API_URL}/generate?lang=java&lines=15`);
const response = await fetch(`${import.meta.env.VITE_API_URL}/generate?lang=lua`);
const reader = response.body.getReader();
const decoder = new TextDecoder();
+32
View File
@@ -0,0 +1,32 @@
package main
import (
"fmt"
"log"
"os"
"github.com/joho/godotenv"
"github.com/labstack/echo/v4"
"github.com/labstack/echo/v4/middleware"
"touch-programming.hazemkrimi.tech/internal/handlers"
)
func main() {
err := godotenv.Load()
if err != nil {
log.Fatal("Error loading environment!")
}
PORT := os.Getenv("PORT")
if len(PORT) == 0 {
PORT = "8080"
}
ech := echo.New()
ech.Use(middleware.CORS())
ech.GET("/generate", handlers.Generate)
ech.Logger.Fatal(ech.Start(fmt.Sprintf(":%s", PORT)))
}
+1 -1
View File
@@ -1,4 +1,4 @@
module experiment/oss-code-generator-ai-powered
module touch-programming.hazemkrimi.tech
go 1.23.4
+142
View File
@@ -0,0 +1,142 @@
package handlers
import (
"context"
"fmt"
"log"
"net/http"
"os"
"strconv"
"strings"
"sync"
"github.com/labstack/echo/v4"
"github.com/tmc/langchaingo/llms"
"github.com/tmc/langchaingo/llms/ollama"
)
type ErrorMaxLinesReached struct{}
func (err *ErrorMaxLinesReached) Error() string {
return fmt.Sprintf("Reached maximum lines configured!")
}
type CodeBlockParser struct {
wg sync.WaitGroup
channel chan []byte
language string
maxLines int
currentLines int
backticksRemoved int
}
func NewCodeBlockParser(lang string, lines int) *CodeBlockParser {
return &CodeBlockParser{
channel: make(chan []byte),
language: lang,
maxLines: lines,
currentLines: 0,
backticksRemoved: 0,
}
}
func (parser *CodeBlockParser) ParseStream(chunk []byte, cancelCtx context.CancelFunc) {
defer parser.wg.Done()
text := string(chunk)
// TODO: After getting the project up and running optimize this to take into account for example open brackets or function ending statements that will be closed after the max lines gets reached.
if !strings.Contains(parser.language, text) && !strings.Contains(text, "```") && !strings.Contains(text, "``") {
if parser.currentLines == 0 && parser.backticksRemoved > 0 && text == "\n" {
parser.channel <- nil
return
}
if strings.Contains(text, "\n") {
parser.currentLines++
}
if parser.currentLines == parser.maxLines-1 {
indexOfNewLine := strings.Index(text, "\n")
if indexOfNewLine > -1 {
parser.channel <- []byte(text[:indexOfNewLine])
} else {
parser.channel <- nil
}
} else if parser.currentLines >= parser.maxLines {
cancelCtx()
} else {
parser.channel <- chunk
}
} else {
parser.channel <- nil
parser.backticksRemoved++
}
}
func Generate(ctx echo.Context) error {
LLM_MODEL := os.Getenv("LLM_MODEL")
if len(LLM_MODEL) == 0 {
return ctx.String(http.StatusInternalServerError, "No LLM model specified in environment!")
}
MAX_LINES, err := strconv.Atoi(os.Getenv("MAX_LINES"))
if err != nil {
return ctx.String(http.StatusInternalServerError, "Error setting max lines!")
}
lang := ctx.QueryParam("lang")
if len(lang) == 0 {
return ctx.String(http.StatusBadRequest, "Lang param is incorrect!")
}
llm, err := ollama.New(ollama.WithModel(LLM_MODEL))
if err != nil {
log.Println(err)
return ctx.String(http.StatusInternalServerError, "Error initializing LLM!")
}
parser := NewCodeBlockParser(lang, MAX_LINES)
ollamaCtx, cancelOllamaCtx := context.WithCancel(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.`, MAX_LINES, lang)),
}
ctx.Response().WriteHeader(http.StatusOK)
_, err = llm.GenerateContent(ollamaCtx, prompt, llms.WithStreamingFunc(func(streamCtx context.Context, chunk []byte) error {
parser.wg.Add(1)
go parser.ParseStream(chunk, cancelOllamaCtx)
select {
case chunk, ok := <-parser.channel:
if ok && len(chunk) > 0 {
ctx.Response().Write(chunk)
ctx.Response().Flush()
}
case <-ollamaCtx.Done():
return &ErrorMaxLinesReached{}
}
return nil
}))
parser.wg.Wait()
defer close(parser.channel)
if err != nil {
if _, ok := err.(*ErrorMaxLinesReached); ok {
return nil
}
return ctx.String(http.StatusInternalServerError, "Error generating code!")
}
return nil
}
-151
View File
@@ -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)))
}