Skip to content

Commit 8c02efd

Browse files
gotsysdbagotsysdba
authored
Minor fixes to 1.3.0 with tests (#317)
1 parent 66e78cc commit 8c02efd

File tree

6 files changed

+347
-7
lines changed

6 files changed

+347
-7
lines changed

‎opentofu/modules/vm/templates/cloudinit-compute.tpl‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,10 +11,7 @@ users:
1111

1212
package_update: false
1313
packages:
14-
- python36-oci-cli
1514
- python3.11
16-
- jre
17-
- sqlcl
1815

1916
write_files:
2017
- path: /etc/systemd/system/ai-optimizer.service
@@ -77,8 +74,9 @@ write_files:
7774

7875
# Install Models
7976
if ${install_ollama}; then
80-
ollama pull llama3.1
81-
ollama pull mxbai-embed-large
77+
echo "Pulling Ollama Models"
78+
ollama pull llama3.1 > /dev/null 2>&1
79+
ollama pull mxbai-embed-large > /dev/null 2>&1
8280
fi
8381

8482
# Wait for python modules to finish

‎opentofu/modules/vm/templates/cloudinit-database.tpl‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,12 @@
33
# All rights reserved. The Universal Permissive License (UPL), Version 1.0 as shown at http://oss.oracle.com/licenses/upl
44
# spell-checker: disable
55

6+
package_update: false
7+
packages:
8+
- python36-oci-cli
9+
- jre
10+
- sqlcl
11+
612
write_files:
713
- path: /tmp/db_priv_sql.sh
814
permissions: '0755'

‎src/client/utils/st_common.py‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -258,6 +258,11 @@ def _update_set_tool():
258258
state.client_settings["vector_search"]["enabled"] = state.selected_tool == "Vector Search"
259259
state.client_settings["selectai"]["enabled"] = state.selected_tool == "SelectAI"
260260

261+
if state.client_settings["vector_search"]["enabled"]:
262+
switch_prompt("sys", "Vector Search Example")
263+
else:
264+
switch_prompt("sys", "Basic Example")
265+
261266
disable_selectai = not is_db_configured()
262267
disable_vector_search = not is_db_configured()
263268

‎src/server/api/v1/chat.py‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ async def chat_stream(
6666
response_model=list[schema.ChatMessage],
6767
)
6868
async def chat_history_clean(client: schema.ClientIdType = Header(default="server")) -> list[ChatMessage]:
69-
"""Delete all Chat History"""
69+
"""Delete all Chat History and clear RAG document context"""
7070
agent: CompiledStateGraph = chatbot.chatbot_graph
7171
try:
7272
_ = agent.update_state(
@@ -75,7 +75,13 @@ async def chat_history_clean(client: schema.ClientIdType = Header(default="serve
7575
"thread_id": client,
7676
}
7777
),
78-
values={"messages": RemoveMessage(id=REMOVE_ALL_MESSAGES)},
78+
values={
79+
"messages": RemoveMessage(id=REMOVE_ALL_MESSAGES),
80+
"cleaned_messages": [],
81+
"context_input": "",
82+
"documents": {},
83+
"final_response": {},
84+
},
7985
)
8086
return [ChatMessage(content="As requested, I've forgotten our conversation.", role="system")]
8187
except KeyError:

‎tests/client/content/test_chatbot.py‎

Lines changed: 260 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@
55
# spell-checker: disable
66
# pylint: disable=import-error
77

8+
from unittest.mock import patch
9+
810

911
#############################################################################
1012
# Test Streamlit UI
@@ -22,3 +24,261 @@ def test_disabled(self, app_server, app_test):
2224
at.error[0].value == "No language models are configured and/or enabled. Disabling Client."
2325
and at.error[0].icon == "🛑"
2426
)
27+
28+
29+
#############################################################################
30+
# Test Prompt Switching Functions
31+
#############################################################################
32+
class TestPromptSwitching:
33+
"""Test automatic prompt switching based on tool selection"""
34+
35+
def test_switch_prompt_to_vector_search(self, app_server, app_test):
36+
"""Test that selecting Vector Search switches to Vector Search Example prompt"""
37+
from client.utils.st_common import switch_prompt
38+
39+
assert app_server is not None
40+
at = app_test(TestStreamlit.ST_FILE)
41+
at.run()
42+
43+
# Setup: ensure we're not starting with Vector Search Example or Custom
44+
at.session_state.client_settings["prompts"]["sys"] = "Basic Example"
45+
46+
# Mock streamlit's info to track if prompt was switched
47+
with patch("client.utils.st_common.state", at.session_state):
48+
with patch("client.utils.st_common.st") as mock_st:
49+
# Act: Switch to Vector Search prompt
50+
switch_prompt("sys", "Vector Search Example")
51+
52+
# Assert: Prompt was updated
53+
assert at.session_state.client_settings["prompts"]["sys"] == "Vector Search Example"
54+
# Assert: User was notified via st.info
55+
mock_st.info.assert_called_once()
56+
assert "Vector Search Example" in str(mock_st.info.call_args)
57+
58+
def test_switch_prompt_to_basic_example(self, app_server, app_test):
59+
"""Test that disabling Vector Search switches to Basic Example prompt"""
60+
from client.utils.st_common import switch_prompt
61+
62+
assert app_server is not None
63+
at = app_test(TestStreamlit.ST_FILE)
64+
at.run()
65+
66+
# Setup: Start with Vector Search Example
67+
at.session_state.client_settings["prompts"]["sys"] = "Vector Search Example"
68+
69+
# Mock streamlit's state and info
70+
with patch("client.utils.st_common.state", at.session_state):
71+
with patch("client.utils.st_common.st") as mock_st:
72+
# Act: Switch to Basic Example
73+
switch_prompt("sys", "Basic Example")
74+
75+
# Assert: Prompt was updated
76+
assert at.session_state.client_settings["prompts"]["sys"] == "Basic Example"
77+
# Assert: User was notified
78+
mock_st.info.assert_called_once()
79+
assert "Basic Example" in str(mock_st.info.call_args)
80+
81+
def test_switch_prompt_does_not_override_custom(self, app_server, app_test):
82+
"""Test that automatic switching respects Custom prompt selection"""
83+
from client.utils.st_common import switch_prompt
84+
85+
assert app_server is not None
86+
at = app_test(TestStreamlit.ST_FILE)
87+
at.run()
88+
89+
# Setup: User has selected Custom prompt
90+
at.session_state.client_settings["prompts"]["sys"] = "Custom"
91+
92+
# Mock streamlit's state and info
93+
with patch("client.utils.st_common.state", at.session_state):
94+
with patch("client.utils.st_common.st") as mock_st:
95+
# Act: Attempt to switch to Vector Search Example
96+
switch_prompt("sys", "Vector Search Example")
97+
98+
# Assert: Prompt remains Custom (not overridden)
99+
assert at.session_state.client_settings["prompts"]["sys"] == "Custom"
100+
# Assert: User was NOT notified (no switching occurred)
101+
mock_st.info.assert_not_called()
102+
103+
def test_switch_prompt_does_not_switch_if_already_set(self, app_server, app_test):
104+
"""Test that switching to the same prompt doesn't trigger notification"""
105+
from client.utils.st_common import switch_prompt
106+
107+
assert app_server is not None
108+
at = app_test(TestStreamlit.ST_FILE)
109+
at.run()
110+
111+
# Setup: Already on Vector Search Example
112+
at.session_state.client_settings["prompts"]["sys"] = "Vector Search Example"
113+
114+
# Mock streamlit's state and info
115+
with patch("client.utils.st_common.state", at.session_state):
116+
with patch("client.utils.st_common.st") as mock_st:
117+
# Act: Try to switch to Vector Search Example again
118+
switch_prompt("sys", "Vector Search Example")
119+
120+
# Assert: Prompt remains the same
121+
assert at.session_state.client_settings["prompts"]["sys"] == "Vector Search Example"
122+
# Assert: User was NOT notified (no change occurred)
123+
mock_st.info.assert_not_called()
124+
125+
def test_vector_search_tool_enables_vector_search_prompt(self, app_server, app_test):
126+
"""Test that selecting Vector Search tool enables Vector Search and switches prompt"""
127+
assert app_server is not None
128+
at = app_test(TestStreamlit.ST_FILE)
129+
at.run()
130+
131+
# Setup: Start with None tool selected and Basic Example prompt
132+
at.session_state.selected_tool = "None"
133+
at.session_state.client_settings["prompts"]["sys"] = "Basic Example"
134+
at.session_state.client_settings["vector_search"] = {"enabled": False}
135+
at.session_state.client_settings["selectai"] = {"enabled": False}
136+
137+
# Simulate selecting Vector Search tool
138+
# This would trigger the _update_set_tool callback in tools_sidebar()
139+
at.session_state.selected_tool = "Vector Search"
140+
141+
# Mock the switch_prompt behavior
142+
with patch("client.utils.st_common.state", at.session_state):
143+
with patch("client.utils.st_common.st"):
144+
# Import and call the function that would be triggered
145+
from client.utils.st_common import switch_prompt
146+
147+
# Simulate what _update_set_tool does
148+
at.session_state.client_settings["vector_search"]["enabled"] = (
149+
at.session_state.selected_tool == "Vector Search"
150+
)
151+
at.session_state.client_settings["selectai"]["enabled"] = (
152+
at.session_state.selected_tool == "SelectAI"
153+
)
154+
155+
# Apply prompt switching logic
156+
if at.session_state.client_settings["vector_search"]["enabled"]:
157+
switch_prompt("sys", "Vector Search Example")
158+
else:
159+
switch_prompt("sys", "Basic Example")
160+
161+
# Assert: Vector Search is enabled
162+
assert at.session_state.client_settings["vector_search"]["enabled"] is True
163+
assert at.session_state.client_settings["selectai"]["enabled"] is False
164+
# Assert: Prompt switched to Vector Search Example
165+
assert at.session_state.client_settings["prompts"]["sys"] == "Vector Search Example"
166+
167+
def test_selectai_tool_uses_basic_prompt(self, app_server, app_test):
168+
"""Test that selecting SelectAI tool uses Basic Example prompt"""
169+
assert app_server is not None
170+
at = app_test(TestStreamlit.ST_FILE)
171+
at.run()
172+
173+
# Setup: Start with Vector Search selected
174+
at.session_state.selected_tool = "Vector Search"
175+
at.session_state.client_settings["prompts"]["sys"] = "Vector Search Example"
176+
at.session_state.client_settings["vector_search"] = {"enabled": True}
177+
at.session_state.client_settings["selectai"] = {"enabled": False}
178+
179+
# Simulate selecting SelectAI tool
180+
at.session_state.selected_tool = "SelectAI"
181+
182+
# Mock the switch_prompt behavior
183+
with patch("client.utils.st_common.state", at.session_state):
184+
with patch("client.utils.st_common.st"):
185+
from client.utils.st_common import switch_prompt
186+
187+
# Simulate what _update_set_tool does
188+
at.session_state.client_settings["vector_search"]["enabled"] = (
189+
at.session_state.selected_tool == "Vector Search"
190+
)
191+
at.session_state.client_settings["selectai"]["enabled"] = (
192+
at.session_state.selected_tool == "SelectAI"
193+
)
194+
195+
# Apply prompt switching logic
196+
if at.session_state.client_settings["vector_search"]["enabled"]:
197+
switch_prompt("sys", "Vector Search Example")
198+
else:
199+
switch_prompt("sys", "Basic Example")
200+
201+
# Assert: SelectAI is enabled, Vector Search is disabled
202+
assert at.session_state.client_settings["vector_search"]["enabled"] is False
203+
assert at.session_state.client_settings["selectai"]["enabled"] is True
204+
# Assert: Prompt switched to Basic Example
205+
assert at.session_state.client_settings["prompts"]["sys"] == "Basic Example"
206+
207+
def test_none_tool_uses_basic_prompt(self, app_server, app_test):
208+
"""Test that selecting None tool uses Basic Example prompt"""
209+
assert app_server is not None
210+
at = app_test(TestStreamlit.ST_FILE)
211+
at.run()
212+
213+
# Setup: Start with Vector Search selected
214+
at.session_state.selected_tool = "Vector Search"
215+
at.session_state.client_settings["prompts"]["sys"] = "Vector Search Example"
216+
at.session_state.client_settings["vector_search"] = {"enabled": True}
217+
at.session_state.client_settings["selectai"] = {"enabled": False}
218+
219+
# Simulate selecting None tool
220+
at.session_state.selected_tool = "None"
221+
222+
# Mock the switch_prompt behavior
223+
with patch("client.utils.st_common.state", at.session_state):
224+
with patch("client.utils.st_common.st"):
225+
from client.utils.st_common import switch_prompt
226+
227+
# Simulate what _update_set_tool does
228+
at.session_state.client_settings["vector_search"]["enabled"] = (
229+
at.session_state.selected_tool == "Vector Search"
230+
)
231+
at.session_state.client_settings["selectai"]["enabled"] = (
232+
at.session_state.selected_tool == "SelectAI"
233+
)
234+
235+
# Apply prompt switching logic
236+
if at.session_state.client_settings["vector_search"]["enabled"]:
237+
switch_prompt("sys", "Vector Search Example")
238+
else:
239+
switch_prompt("sys", "Basic Example")
240+
241+
# Assert: Both tools are disabled
242+
assert at.session_state.client_settings["vector_search"]["enabled"] is False
243+
assert at.session_state.client_settings["selectai"]["enabled"] is False
244+
# Assert: Prompt switched to Basic Example
245+
assert at.session_state.client_settings["prompts"]["sys"] == "Basic Example"
246+
247+
def test_custom_prompt_not_overridden_by_tool_selection(self, app_server, app_test):
248+
"""Test that Custom prompt is not overridden when switching tools"""
249+
assert app_server is not None
250+
at = app_test(TestStreamlit.ST_FILE)
251+
at.run()
252+
253+
# Setup: User has Custom prompt selected
254+
at.session_state.selected_tool = "None"
255+
at.session_state.client_settings["prompts"]["sys"] = "Custom"
256+
at.session_state.client_settings["vector_search"] = {"enabled": False}
257+
at.session_state.client_settings["selectai"] = {"enabled": False}
258+
259+
# Simulate selecting Vector Search tool
260+
at.session_state.selected_tool = "Vector Search"
261+
262+
# Mock the switch_prompt behavior
263+
with patch("client.utils.st_common.state", at.session_state):
264+
with patch("client.utils.st_common.st"):
265+
from client.utils.st_common import switch_prompt
266+
267+
# Simulate what _update_set_tool does
268+
at.session_state.client_settings["vector_search"]["enabled"] = (
269+
at.session_state.selected_tool == "Vector Search"
270+
)
271+
at.session_state.client_settings["selectai"]["enabled"] = (
272+
at.session_state.selected_tool == "SelectAI"
273+
)
274+
275+
# Apply prompt switching logic
276+
if at.session_state.client_settings["vector_search"]["enabled"]:
277+
switch_prompt("sys", "Vector Search Example")
278+
else:
279+
switch_prompt("sys", "Basic Example")
280+
281+
# Assert: Vector Search is enabled
282+
assert at.session_state.client_settings["vector_search"]["enabled"] is True
283+
# Assert: Prompt remains Custom (not overridden)
284+
assert at.session_state.client_settings["prompts"]["sys"] == "Custom"

0 commit comments

Comments
 (0)