Skip to content
This repository was archived by the owner on Apr 23, 2026. It is now read-only.

Commit 68976e4

Browse files
author
reid_liu
committed
add --no-decoration option for chat
- add a new flag to chat without decoration Signed-off-by: reid_liu <guliu@redhat.com>
1 parent aaca968 commit 68976e4

3 files changed

Lines changed: 90 additions & 27 deletions

File tree

‎src/instructlab/cli/model/chat.py‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,6 +142,12 @@
142142
config_class="rag",
143143
config_sections="retriever",
144144
)
145+
@click.option(
146+
"-nd",
147+
"--no-decoration",
148+
is_flag=True,
149+
help="Disable decorations for chat responses.",
150+
)
145151
@click.pass_context
146152
@clickext.display_params
147153
def chat(
@@ -166,6 +172,7 @@ def chat(
166172
collection_name,
167173
embedding_model_path,
168174
top_k,
175+
no_decoration,
169176
):
170177
"""Runs a chat using the modified model"""
171178
chat_model(
@@ -189,6 +196,7 @@ def chat(
189196
collection_name,
190197
embedding_model_path,
191198
top_k,
199+
no_decoration,
192200
backend_type=ctx.obj.config.serve.server.backend_type,
193201
host=ctx.obj.config.serve.server.host,
194202
port=ctx.obj.config.serve.server.port,

‎src/instructlab/model/chat.py‎

Lines changed: 30 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ def __init__(
100100
max_ctx_size=None,
101101
temperature=1.0,
102102
backend_type="",
103+
box=True,
103104
):
104105
self.client = client
105106
self.retriever: DocumentStoreRetriever | None = retriever
@@ -112,6 +113,7 @@ def __init__(
112113
self.max_ctx_size = max_ctx_size
113114
self.temperature = temperature
114115
self.backend_type = backend_type
116+
self.box = box
115117

116118
self.console = Console()
117119

@@ -135,7 +137,10 @@ def _reset_session(self, hard=False):
135137
)
136138

137139
def _sys_print(self, *args, **kwargs):
138-
self.console.print(Panel(*args, title="system", **kwargs))
140+
if self.box:
141+
self.console.print(Panel(*args, title="system", **kwargs))
142+
else:
143+
self.console.print(*args)
139144

140145
def log_message(self, msg):
141146
if self.log_file:
@@ -146,19 +151,23 @@ def greet(self, help=False, new=False, session_name="new session"): # pylint: d
146151
side_info_str = (" (type `/h` for help)" if help else "") + (
147152
f" ({session_name})" if new else ""
148153
)
149-
self._sys_print(
150-
Markdown(
151-
f"Welcome to InstructLab Chat w/ **{self.model_name.upper()}**"
152-
+ side_info_str
153-
)
154+
message = (
155+
f"Welcome to InstructLab Chat w/ **{self.model_name.upper()}**"
156+
+ side_info_str
154157
)
158+
if self.box:
159+
self._sys_print(Markdown(message))
160+
else:
161+
self.console.print(message)
155162

156163
@property
157164
def model_name(self):
158165
return os.path.basename(os.path.normpath(self.model))
159166

160167
@property
161168
def _right_prompt(self):
169+
if not self.box:
170+
return None
162171
return FormattedText(
163172
[
164173
(
@@ -262,7 +271,10 @@ def __handle_replay(self, content, display_wrapper=lambda x: x):
262271
raise KeyboardInterrupt
263272

264273
def _handle_display(self, content):
265-
return self.__handle_replay(content, display_wrapper=lambda x: Panel(x)) # pylint: disable=unnecessary-lambda
274+
return self.__handle_replay(
275+
content,
276+
display_wrapper=lambda x: Panel(x) if self.box else x, # pylint: disable=unnecessary-lambda
277+
)
266278

267279
def _load_session_history(self, content=None):
268280
data = self.info["messages"]
@@ -274,7 +286,10 @@ def _load_session_history(self, content=None):
274286
"\n" + PROMPT_PREFIX + m["content"], style="dim grey0"
275287
)
276288
else:
277-
self.console.print(Panel(m["content"]), style="dim grey0")
289+
if self.box:
290+
self.console.print(Panel(m["content"]), style="dim grey0")
291+
else:
292+
self.console.print(m["content"], style="dim grey0")
278293

279294
def _handle_plain(self, content):
280295
return self.__handle_replay(content)
@@ -362,7 +377,6 @@ def start_prompt(
362377
self,
363378
logger, # pylint: disable=redefined-outer-name
364379
content=None,
365-
box=True,
366380
):
367381
handlers = {
368382
"/q": self._handle_quit,
@@ -520,7 +534,7 @@ def start_prompt(
520534
title=self.model_name,
521535
subtitle_align="right",
522536
)
523-
if box
537+
if self.box
524538
else response_content
525539
)
526540
subtitle = None
@@ -536,7 +550,7 @@ def start_prompt(
536550
if chunk_message.content:
537551
response_content.append(chunk_message.content)
538552

539-
if box:
553+
if self.box:
540554
panel.subtitle = f"elapsed {time.time() - start_time:.3f} seconds"
541555
subtitle = f"elapsed {time.time() - start_time:.3f} seconds"
542556

@@ -569,6 +583,7 @@ def chat_model(
569583
collection_name,
570584
embedding_model_path,
571585
top_k,
586+
no_decoration,
572587
backend_type,
573588
host,
574589
port,
@@ -723,6 +738,7 @@ def chat_model(
723738
top_k=top_k,
724739
backend_type=backend_type,
725740
params=params,
741+
no_decoration=no_decoration,
726742
)
727743
except ChatException as exc:
728744
print(f"{RED}Executing chat failed with: {exc}{RESET}")
@@ -752,6 +768,7 @@ def chat_cli(
752768
vi_mode,
753769
visible_overflow,
754770
params,
771+
no_decoration,
755772
):
756773
"""Starts a CLI-based chat with the server"""
757774
client = OpenAI(
@@ -834,6 +851,7 @@ def chat_cli(
834851
max_tokens=(max_tokens if max_tokens else max_tokens),
835852
max_ctx_size=max_ctx_size,
836853
backend_type=backend_type,
854+
box=not no_decoration,
837855
)
838856

839857
if not qq and session is None:
@@ -846,7 +864,7 @@ def chat_cli(
846864
if not qq:
847865
print(f"{PROMPT_PREFIX}{question}")
848866
try:
849-
ccb.start_prompt(logger, content=question, box=not qq)
867+
ccb.start_prompt(logger, content=question)
850868
except ChatException as exc:
851869
raise ChatException(f"API issue found while executing chat: {exc}") from exc
852870
except (ChatQuitException, KeyboardInterrupt, EOFError):

‎tests/test_model_chat.py‎

Lines changed: 52 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@
22
from unittest.mock import MagicMock
33
import contextlib
44
import logging
5-
import re
65

76
# Third Party
87
from click.testing import CliRunner
98
from rich.console import Console
9+
from rich.panel import Panel
1010
import pytest
1111

1212
# First Party
@@ -60,18 +60,20 @@ def test_retriever_is_called_when_present():
6060
retriever.augmented_context.assert_called_with(user_query=user_query)
6161

6262

63-
def handle_output(output):
64-
return re.sub(r"\s+", " ", output).strip()
65-
66-
67-
def test_list_contexts_output():
63+
def test_list_contexts_and_decoration():
6864
chatbot = ConsoleChatBot(model="/var/model/file", client=None, loaded={})
6965

70-
def mock_sys_print(output):
71-
mock_sys_print.output = output
66+
def mock_sys_print_output(*args, **kwargs):
67+
if chatbot.box:
68+
panel = Panel(*args, **kwargs)
69+
mock_sys_print_output.output = panel
70+
else:
71+
mock_sys_print_output.output = args[0]
7272

73-
chatbot._sys_print = mock_sys_print
73+
chatbot._sys_print = mock_sys_print_output
7474

75+
# Test when box=True
76+
chatbot.box = True
7577
mock_prompt_session = MagicMock()
7678
mock_prompt_session.prompt.return_value = "/lc"
7779
chatbot.input = mock_prompt_session
@@ -81,15 +83,50 @@ def mock_sys_print(output):
8183

8284
console = Console(force_terminal=False)
8385
with console.capture() as capture:
84-
console.print(mock_sys_print.output)
86+
console.print(mock_sys_print_output.output)
87+
88+
rendered_output = capture.get().strip()
89+
90+
expected_output_with_box = (
91+
"╭──────────────────────────────────────────────────────────────────────────────╮\n"
92+
"│ Available contexts: │\n"
93+
"│ │\n"
94+
"│ default: I am an advanced AI language model designed to assist you with a │\n"
95+
"│ wide range of tasks and provide helpful, clear, and accurate responses. My │\n"
96+
"│ primary role is to serve as a chat assistant, engaging in natural, │\n"
97+
"│ conversational dialogue, answering questions, generating ideas, and offering │\n"
98+
"│ support across various topics. │\n"
99+
"│ │\n"
100+
"│ cli_helper: You are an expert for command line interface and know all common │\n"
101+
"│ commands. Answer the command to execute as it without any explanation. │\n"
102+
"╰──────────────────────────────────────────────────────────────────────────────╯"
103+
)
104+
105+
assert rendered_output == expected_output_with_box
106+
107+
# Test when box=False
108+
chatbot.box = False
109+
mock_prompt_session = MagicMock()
110+
mock_prompt_session.prompt.return_value = "/lc"
111+
chatbot.input = mock_prompt_session
112+
113+
with contextlib.suppress(KeyboardInterrupt):
114+
chatbot.start_prompt(logger=None)
115+
116+
with console.capture() as capture:
117+
console.print(mock_sys_print_output.output)
85118

86119
rendered_output = capture.get().strip()
87120

88-
expected_output = (
89-
"Available contexts:\n\n"
90-
"default: I am an advanced AI language model designed to assist you with a wide range of tasks and provide helpful, clear, and accurate responses. My primary role is to serve as a chat assistant, engaging in natural, conversational dialogue, answering questions, generating ideas, and offering support across various topics.\n\n"
91-
"cli_helper: You are an expert for command line interface and know all common "
121+
expected_output_without_box = (
122+
"Available contexts: \n\n"
123+
"default: I am an advanced AI language model designed to assist you with a wide \n"
124+
"range of tasks and provide helpful, clear, and accurate responses. My primary \n"
125+
"role is to serve as a chat assistant, engaging in natural, conversational \n"
126+
"dialogue, answering questions, generating ideas, and offering support across \n"
127+
"various topics. \n\n"
128+
"cli_helper: You are an expert for command line interface and know all common \n"
92129
"commands. Answer the command to execute as it without any explanation."
93130
)
94131

95-
assert handle_output(rendered_output) == handle_output(expected_output)
132+
assert rendered_output == expected_output_without_box

0 commit comments

Comments
 (0)