refactor: use SystemMessage and HumanMessage everywhere

Author: lpm0073

models/examples/prompt.py

@@ -2,6 +2,8 @@
 """Sales Support Model (hsr)"""
 import argparse
 
+from langchain.schema import HumanMessage, SystemMessage
+
 from models.hybrid_search_retreiver import HybridSearchRetriever
 
 
@@ -10,9 +12,11 @@
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="hsr examples")
-    parser.add_argument("system_prompt", type=str, help="A system prompt to send to the model.")
-    parser.add_argument("human_prompt", type=str, help="A human prompt to send to the model.")
+    parser.add_argument("system_message", type=str, help="A system prompt to send to the model.")
+    parser.add_argument("human_message", type=str, help="A human prompt to send to the model.")
     args = parser.parse_args()
 
-    result = hsr.cached_chat_request(args.system_prompt, args.human_prompt)
-    print(result)
+    system_message = SystemMessage(text=args.system_message)
+    human_message = HumanMessage(text=args.human_message)
+    result = hsr.cached_chat_request(system_message=system_message, human_message=human_message)
+    print(result.content)

models/examples/rag.py

@@ -2,6 +2,8 @@
 """Sales Support Model (hsr) Retrieval Augmented Generation (RAG)"""
 import argparse
 
+from langchain.schema import HumanMessage
+
 from models.hybrid_search_retreiver import HybridSearchRetriever
 
 
@@ -12,5 +14,6 @@
     parser.add_argument("prompt", type=str, help="A question about the PDF contents")
     args = parser.parse_args()
 
-    result = hsr.rag(prompt=args.prompt)
+    human_message = HumanMessage(text=args.prompt)
+    result = hsr.rag(human_message=human_message)
     print(result)

models/tests/test_examples.py

@@ -6,6 +6,7 @@
 from unittest.mock import MagicMock, patch
 
 import pytest  # pylint: disable=unused-import
+from langchain.schema import HumanMessage, SystemMessage
 
 from models.examples.prompt import hsr as prompt_hrs
 from models.examples.rag import hsr as rag_hsr
@@ -14,7 +15,7 @@
 from models.prompt_templates import NetecPromptTemplates
 
 
-HUMAN_PROMPT = 'return the word "SUCCESS" in upper case.'
+HUMAN_MESSAGE = 'return the word "SUCCESS" in upper case.'
 
 
 class TestExamples:
@@ -25,44 +26,47 @@ def test_prompt(self, mock_parse_args):
         """Test prompt example."""
         mock_args = MagicMock()
         mock_args.system_prompt = "you are a helpful assistant"
-        mock_args.human_prompt = HUMAN_PROMPT
+        mock_args.human_prompt = HUMAN_MESSAGE
         mock_parse_args.return_value = mock_args
 
-        result = prompt_hrs.cached_chat_request(mock_args.system_prompt, mock_args.human_prompt)
-        assert result == "SUCCESS"
+        system_message = SystemMessage(content="you are a helpful assistant")
+        human_message = HumanMessage(content=HUMAN_MESSAGE)
+        result = prompt_hrs.cached_chat_request(system_message=system_message, human_message=human_message)
+        assert result.content == "SUCCESS"
 
     @patch("argparse.ArgumentParser.parse_args")
     def test_rag(self, mock_parse_args):
         """Test RAG example."""
         mock_args = MagicMock()
-        mock_args.human_prompt = HUMAN_PROMPT
+        mock_args.human_message = HUMAN_MESSAGE
         mock_parse_args.return_value = mock_args
 
-        result = rag_hsr.rag(mock_args.human_prompt)
+        human_message = HumanMessage(content=mock_args.human_message)
+        result = rag_hsr.rag(human_message=human_message)
         assert result == "SUCCESS"
 
     @patch("argparse.ArgumentParser.parse_args")
     def test_training_services(self, mock_parse_args):
         """Test training services templates."""
         mock_args = MagicMock()
-        mock_args.human_prompt = HUMAN_PROMPT
+        mock_args.human_message = HUMAN_MESSAGE
         mock_parse_args.return_value = mock_args
 
         templates = NetecPromptTemplates()
         prompt = templates.training_services
 
-        result = training_services_hsr.prompt_with_template(prompt=prompt, concept=mock_args.human_prompt)
+        result = training_services_hsr.prompt_with_template(prompt=prompt, concept=mock_args.human_message)
         assert "SUCCESS" in result
 
     @patch("argparse.ArgumentParser.parse_args")
     def test_oracle_training_services(self, mock_parse_args):
         """Test oracle training services."""
         mock_args = MagicMock()
-        mock_args.human_prompt = HUMAN_PROMPT
+        mock_args.human_message = HUMAN_MESSAGE
         mock_parse_args.return_value = mock_args
 
         templates = NetecPromptTemplates()
         prompt = templates.oracle_training_services
 
-        result = training_services_oracle_hsr.prompt_with_template(prompt=prompt, concept=mock_args.human_prompt)
+        result = training_services_oracle_hsr.prompt_with_template(prompt=prompt, concept=mock_args.human_message)
         assert "SUCCESS" in result

models/tests/test_openai.py

@@ -19,4 +19,4 @@ def test_03_test_openai_connectivity(self):
         retval = hsr.cached_chat_request(
             "your are a helpful assistant", "please return the value 'CORRECT' in all upper case."
         )
-        assert retval == "CORRECT"
+        assert retval.content == "CORRECT"