coloring fixes

helpers.py

@@ -8,7 +8,7 @@ def highlight_code(code: str, language: str = 'py') -> None:
 	try:
 		lexer = get_lexer_by_name(language)
 	except ValueError:
-		debugMod.log("Warning: Language not recognized. Printing without highlighting.")
+		# debugMod.log("Warning: Language not recognized. Printing without highlighting.")
 		return code
 
 	formatter = TerminalFormatter()

queries.py

@@ -133,6 +133,8 @@ def call_ollama_and_print(task_type, prompt, silent=False):
 	in_code_block = False
 	code_lang = None
 	first_chunk = True
+	code_buffer = ""
+	prev_highlighted = ""  # initialize before processing stream
 
 	for chunk in stream:
 		if first_chunk:
@@ -141,31 +143,62 @@ def call_ollama_and_print(task_type, prompt, silent=False):
 			print("\nAI Response: ", end="", flush=True)
 		content: str = chunk.get('message', {}).get('content', '')
 
-		if content == '```' or re.match('```.*', content):
+		debugMod.log(content)
+
+		# Detect code block start/end
+		if content.startswith('```') or re.match(r'^```[a-zA-Z]*$', content):
 			if in_code_block:
 				in_code_block = False
-				print()
-				buffer += content
+				highlighted = highlight_code(code_buffer, code_lang)
+
+				# Overwrite the previously highlighted code before exiting the block
+				if prev_highlighted:
+					for _ in prev_highlighted.splitlines():
+						print("\033[F\033[2K", end='')  # Move up and clear line
+
+				print(f'{highlighted}\n```\n', flush=True)
+				code_buffer = ""
 				code_lang = None
+				prev_highlighted = ""
 			else:
 				in_code_block = True
+				print('\n```')
+
 				code_lang = content.replace('```', '').strip()
 				if (len(code_lang) == 0):
 					code_lang = "TODO"
+				else:
+					debugMod.log(f'detected language: {code_lang}')
 
 		elif code_lang == "TODO":
-			# last chunk was the backticks, now is lang
+			# Last chunk was the backticks, now is lang
 			splitVal = content.strip().split()
 			code_lang = splitVal[0]
+			debugMod.log(f'detected language: {code_lang}')
 
 			if (len(splitVal) > 1 and len(splitVal[1]) > 0):
 				hcode = highlight_code(splitVal[1], code_lang)
 				print(hcode, end="", flush=True)
-				buffer += hcode
+				code_buffer += content
+
+		elif in_code_block:
+			code_buffer += content
+			highlighted = highlight_code(code_buffer, code_lang)
+
+			for _ in prev_highlighted.splitlines():
+				print("\033[F\033[2K", end='')  # Move cursor up and clear each line
+
+			if not content.endswith('\n'):
+				print(f'{highlighted}', end='', flush=True)
+				prev_highlighted = highlighted
+			else:
+				print(f'{highlighted}', end='\n', flush=True)
+				prev_highlighted = highlighted + '\n'
 
 		else:
-			buffer += content
-			print(content, end="", flush=True)
+			debugMod.log('in normal for', content)
+			# Normal text handling
+			print(content, end='', flush=True)
 
 	print()  # Newline after streaming
 	debugMod.log("RAG query response received")