From e5cb3b73c731363d6b64c8552b1a2bbe2e501550 Mon Sep 17 00:00:00 2001
From: leehk <leehongkai@gmail.com>
Date: Thu, 6 Mar 2025 11:10:11 +0800
Subject: [PATCH] update

---
 .vscode/settings.json                         |  7 ++
 .../.github/workflows/app-testing.yml         | 30 ++++++++
 app/streamlit/Chatbot.py                      |  1 -
 app/streamlit/tests/test_chatbot.py           | 72 ++++++++++---------
 4 files changed, 77 insertions(+), 33 deletions(-)
 create mode 100644 .vscode/settings.json
 create mode 100644 app/streamlit/.github/workflows/app-testing.yml

diff --git a/.vscode/settings.json b/.vscode/settings.json
new file mode 100644
index 0000000..d171c99
--- /dev/null
+++ b/.vscode/settings.json
@@ -0,0 +1,7 @@
+{
+    "python.testing.pytestArgs": [
+        "app"
+    ],
+    "python.testing.unittestEnabled": false,
+    "python.testing.pytestEnabled": true
+}
\ No newline at end of file
diff --git a/app/streamlit/.github/workflows/app-testing.yml b/app/streamlit/.github/workflows/app-testing.yml
new file mode 100644
index 0000000..36cf8a5
--- /dev/null
+++ b/app/streamlit/.github/workflows/app-testing.yml
@@ -0,0 +1,30 @@
+name: App testing
+
+on:
+  push:
+    branches: [ "main" ]
+  pull_request:
+    branches: [ "main" ]
+
+permissions:
+  contents: read
+
+jobs:
+  streamlit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5
+        with:
+          python-version: '3.11'
+      - uses: streamlit/streamlit-app-action@v0.0.3
+        with:
+          app-path: Chatbot.py
+          ruff: true
+          pytest-args: -v --junit-xml=test-results.xml
+      - if: always()
+        uses: pmeier/pytest-results-action@v0.6.0
+        with:
+          path: test-results.xml
+          summary: true
+          display-options: fEX
diff --git a/app/streamlit/Chatbot.py b/app/streamlit/Chatbot.py
index beb7f08..16266d4 100644
--- a/app/streamlit/Chatbot.py
+++ b/app/streamlit/Chatbot.py
@@ -1,5 +1,4 @@
 import os
-import subprocess
 import streamlit as st
 import chromadb
 from decouple import config
diff --git a/app/streamlit/tests/test_chatbot.py b/app/streamlit/tests/test_chatbot.py
index 97ad3d1..7a86629 100644
--- a/app/streamlit/tests/test_chatbot.py
+++ b/app/streamlit/tests/test_chatbot.py
@@ -1,39 +1,47 @@
+import os
 import pytest
-import streamlit as st
-from unittest.mock import patch
+import chromadb
+from langchain.prompts import PromptTemplate
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_deepseek import ChatDeepSeek
+from langchain_community.llms.moonshot import Moonshot
 
-# add app/streamlit to sys.path
 import sys
-sys.path.insert(0, "/Users/leehongkai/projects/aimingmed/aimingmed-ai/app/streamlit")
+sys.path.append(".")
+import streamlit as st
+import pytest
+from unittest.mock import patch
+from Chatbot import CHAT_MODEL_PROVIDER, INPUT_CHROMADB_LOCAL, COLLECTION_NAME, cot_template, answer_template
 
-from unittest.mock import patch, MagicMock
+@pytest.fixture(autouse=True)
+def mock_session_state():
+    with patch.object(st, "session_state", {"messages": []}):
+        yield
 
+def test_prompt_templates():
+    # Test that the prompt templates are correctly formatted
+    assert "documents_text" in cot_template
+    assert "question" in cot_template
+    assert "cot" in answer_template
+    assert "question" in answer_template
 
-def test_title():
-    with patch("streamlit.title") as mock_title, \
-         patch("streamlit.session_state", new_callable=MagicMock) as mock_session_state:
-        import Chatbot
-        st.session_state["messages"] = []
-        mock_title.assert_called_once_with("💬 RAG AI for Medical Guideline")
+def test_chromadb_connection():
+    # Test that the ChromaDB client is initialized correctly
+    chroma_client = chromadb.PersistentClient(path=INPUT_CHROMADB_LOCAL)
+    collection = chroma_client.get_collection(name=COLLECTION_NAME)
+    assert collection is not None
 
-def test_caption():
-    with patch("streamlit.caption") as mock_caption, \
-         patch("streamlit.session_state", new_callable=MagicMock) as mock_session_state:
-        import Chatbot
-        st.session_state["messages"] = []
-        mock_caption.assert_called()
-
-def test_chat_input():
-    with patch("streamlit.chat_input", return_value="test_prompt") as mock_chat_input, \
-         patch("streamlit.session_state", new_callable=MagicMock) as mock_session_state:
-        import Chatbot
-        st.session_state["messages"] = []
-        mock_chat_input.assert_called_once()
-
-def test_chat_message():
-    with patch("streamlit.chat_message") as mock_chat_message, \
-         patch("streamlit.session_state", new_callable=MagicMock) as mock_session_state:
-        with patch("streamlit.chat_input", return_value="test_prompt"):
-            import Chatbot
-        st.session_state["messages"] = []
-        mock_chat_message.assert_called()
\ No newline at end of file
+@pytest.mark.skipif(CHAT_MODEL_PROVIDER not in ["deepseek", "gemini", "moonshot"], reason="requires a valid CHAT_MODEL_PROVIDER")
+def test_llm_initialization():
+    # Test that the correct LLM is initialized based on the CHAT_MODEL_PROVIDER environment variable
+    if CHAT_MODEL_PROVIDER == "deepseek":
+        llm = ChatDeepSeek(model="deepseek-chat")
+        assert isinstance(llm, ChatDeepSeek)
+    elif CHAT_MODEL_PROVIDER == "gemini":
+        llm = ChatGoogleGenerativeAI(model="gemini-1.5-flash")
+        assert isinstance(llm, ChatGoogleGenerativeAI)
+    elif CHAT_MODEL_PROVIDER == "moonshot":
+        llm = Moonshot(model="moonshot-v1-128k")
+        assert isinstance(llm, Moonshot)
+        llm = Moonshot(model="moonshot-v1-128k")
+        assert isinstance(llm, Moonshot)
\ No newline at end of file