From ec1e3fa81006db0f41a5f3cc7088e948812b7f76 Mon Sep 17 00:00:00 2001
From: Hazem Krimi <me@hazemkrimi.tech>
Date: Tue, 21 Jan 2025 18:17:50 +0100
Subject: [PATCH] Tweaks

---
 client/src/App.jsx |  2 +-
 server/main.go     | 15 ++++++---------
 2 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/client/src/App.jsx b/client/src/App.jsx
index abf158d..eb5acce 100644
--- a/client/src/App.jsx
+++ b/client/src/App.jsx
@@ -13,7 +13,7 @@ function App() {
 		(async function() {
 			setCode('');
 
-			const response = await fetch(`${import.meta.env.VITE_API_URL}/generate?lang=go&lines=10`);
+			const response = await fetch(`${import.meta.env.VITE_API_URL}/generate?lang=ocaml&lines=20`);
 			const reader = response.body.getReader();
 			const decoder = new TextDecoder();
 
diff --git a/server/main.go b/server/main.go
index c9df8a0..4fe83d6 100644
--- a/server/main.go
+++ b/server/main.go
@@ -20,7 +20,6 @@ import (
 type CodeBlockParser struct {
 	channel                  chan []byte
 	processedChunks          int
-	processedLines           int
 	removedStartingBackticks bool
 	removedLanguageName      bool
 	removedEndingBackticks   bool
@@ -30,7 +29,6 @@ func NewCodeBlockParser() *CodeBlockParser {
 	return &CodeBlockParser{
 		channel:                  make(chan []byte),
 		processedChunks:          0,
-		processedLines:           0,
 		removedStartingBackticks: false,
 		removedLanguageName:      false,
 		removedEndingBackticks:   false,
@@ -38,10 +36,6 @@ func NewCodeBlockParser() *CodeBlockParser {
 }
 
 func (p *CodeBlockParser) ParseStream(chunk []byte, language string) {
-	if strings.Contains(string(chunk), "\n") {
-		p.processedLines += 1
-	}
-
 	if !p.removedStartingBackticks {
 		if bytes.Contains(chunk, []byte("```")) {
 			p.removedStartingBackticks = true
@@ -62,11 +56,15 @@ func (p *CodeBlockParser) ParseStream(chunk []byte, language string) {
 
 	if p.removedStartingBackticks && !p.removedEndingBackticks {
 		if bytes.Contains(chunk, []byte("```")) {
-			p.removedEndingBackticks = true
 			chunk = nil
+			p.removedEndingBackticks = true
 		}
 	}
 
+	if p.removedEndingBackticks {
+		chunk = nil
+	}
+
 	p.processedChunks += 1
 	p.channel <- chunk
 }
@@ -117,9 +115,8 @@ func main() {
 		parser := NewCodeBlockParser()
 		ollamaCtx := context.Background()
 		prompt := []llms.MessageContent{
-			llms.TextParts(llms.ChatMessageTypeSystem, `You must only generate code without any descriptions or formatting like markdown code fences with backticks. Also don't include any code comments and use spaces instead of tabs for spacing. Most importantly respect the number of lines provided you get asked to generate!`),
 			llms.TextParts(llms.ChatMessageTypeHuman, fmt.Sprintf(`
-				Generate exactly %d lines of code from a well known open source project in the %s programming language.`, lines, lang)),
+				You must only generate code without any descriptions or code comments or formatting like markdown code fences with backticks. Use spaces instead of tabs for spacing. Generate accurately according to the number of lines you get provided. Generate exactly between %d and %d lines of code from a well known open source project in the %s programming language.`, lines/2, lines, lang)),
 		}
 
 		if _, err := llm.GenerateContent(ollamaCtx, prompt, llms.WithStreamingFunc(func(streamCtx context.Context, chunk []byte) error {