mirror of
https://github.com/hazemKrimi/touch-programming.git
synced 2026-05-01 18:20:26 +00:00
Server directory structure and streaming optimization
This commit is contained in:
@@ -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
@@ -1,4 +1,4 @@
|
||||
module experiment/oss-code-generator-ai-powered
|
||||
module touch-programming.hazemkrimi.tech
|
||||
|
||||
go 1.23.4
|
||||
|
||||
|
||||
@@ -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
@@ -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)))
|
||||
}
|
||||
Reference in New Issue
Block a user