Merge pull request #708

Feat/add relevant chunk support to colang 2

Author: drazvan

Diff for: nemoguardrails/colang/v2_x/library/core.co

@@ -11,8 +11,11 @@ flow _user_said $text -> $event
   else
     match UtteranceUserAction.Finished() as $event
 
+  $text = $event.final_transcript
+
 flow _user_saying $text -> $event
   """The internal flow for all semantic 'user saying' flows."""
+
   if $text
     if is_regex($text)
       match UtteranceUserAction.TranscriptUpdated(interim_transcript=$text) as $event
@@ -22,10 +25,16 @@ flow _user_saying $text -> $event
   else
     match UtteranceUserAction.TranscriptUpdated() as $event
 
+  $text = $event.interim_transcript
+
 flow _user_said_something_unexpected -> $event
   """The internal flow for all semantic 'user said something unexpected' flows."""
+  global $last_user_message
   match UnhandledEvent(event="UtteranceUserActionFinished", loop_ids={$self.loop_id}) as $event
 
+  $text = $event.final_transcript
+  $last_user_message = $text
+
 @meta(user_action='user said "{$transcript}"')
 flow user said $text -> $transcript
   """Wait for a user to have said given text."""
@@ -73,6 +82,12 @@ flow user said something unexpected -> $event, $transcript
 
 flow _bot_say $text -> $action
   """The internal flow for all semantic level bot utterance flows."""
+  global $bot_message
+  global $last_bot_message
+
+  $bot_message = $text
+  $last_bot_message = $text
+
   await UtteranceBotAction(script=$text) as $action
 
 @meta(bot_action=True)

Diff for: nemoguardrails/colang/v2_x/library/guardrails.co

@@ -43,6 +43,7 @@ flow _user_saying $text -> $event
 flow _user_said_something_unexpected -> $event
   """Override core flow for when the user said something unexpected."""
   global $user_message
+  global $last_user_message
   match UnhandledEvent(event="UtteranceUserActionFinished", loop_ids={$self.loop_id}) as $event
 
   $text = $event.final_transcript

Diff for: nemoguardrails/colang/v2_x/library/llm.co

@@ -95,6 +95,13 @@ flow continuation on unhandled user utterance
 
   log 'start generating user intent and bot intent/action...'
   $action = 'user said "{$event.final_transcript}"'
+
+
+  # retrieve relevant chunks from KB if user_message is not empty
+
+  await RetrieveRelevantChunksAction()
+
+
   #await GenerateUserIntentAction(user_action=$action, max_example_flows=20) as $action_ref
   #$user_intent = $action_ref.return_value
   await GenerateUserIntentAndBotAction(user_action=$action, max_example_flows=20) as $action_ref
@@ -205,11 +212,18 @@ flow llm generate interaction continuation flow -> $flow_name
   activate polling llm request response
   # Generate continuation based current interaction history
 
+
+  # retrieve relevant chunks from KB if user_message is not empty
+  await RetrieveRelevantChunksAction()
+
+
   log 'start generating flow continuation...'
+
   $flow_info = await GenerateFlowContinuationAction(temperature=0.1)
   log "generated flow continuation: `{$flow_info}`"
   $exists = await CheckValidFlowExistsAction(flow_id=$flow_info.name)
 
+
   if $exists == False
     $flows = await AddFlowsAction(config=$flow_info.body)
     if len($flows) == 0

Diff for: nemoguardrails/llm/prompts/general.yml

@@ -146,6 +146,13 @@ prompts:
           # These are the most likely user intents:
           {{ examples }}
 
+          {% if context.relevant_chunks %}
+          # This is some additional context:
+          ```markdown
+          {{ context.relevant_chunks }}
+          ```
+          {% endif %}
+
           # This is the current conversation between the user and the bot:
           {{ history | colang }}
 
@@ -205,6 +212,14 @@ prompts:
           # This is the current conversation between the user and the bot:
           {{ history | colang }}
 
+          {% if context.relevant_chunks %}
+          # This is some additional context:
+          ```markdown
+          {{ context.relevant_chunks }}
+          ```
+          {% endif %}
+
+
           bot intent:
 
     - task: generate_flow_continuation_from_flow_nld

Diff for: tests/test_retrieve_relevant_chunks.py

@@ -0,0 +1,77 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from unittest.mock import MagicMock
+
+import pytest
+
+from nemoguardrails import LLMRails, RailsConfig
+from nemoguardrails.kb.kb import KnowledgeBase
+from tests.utils import TestChat
+
+config = RailsConfig.from_content(
+    """
+import llm
+import core
+
+flow main
+    activate llm continuation
+
+flow user express greeting
+   user said "hello"
+   or user said "hi"
+   or user said "how are you"
+
+flow bot express greeting
+   bot say "Hey!"
+
+flow greeting
+    user express greeting
+    bot express greeting
+""",
+    yaml_content="""
+    colang_version: 2.x
+    models: []
+    """,
+)
+
+
+def test_relevant_chunk_inserted_in_prompt():
+    mock_kb = MagicMock(spec=KnowledgeBase)
+
+    mock_kb.search_relevant_chunks.return_value = [
+        {"title": "Test Title", "body": "Test Body"}
+    ]
+
+    chat = TestChat(
+        config,
+        llm_completions=[
+            " user express greeting",
+            ' bot respond to aditional context\nbot action: "Hello is there anything else" ',
+        ],
+    )
+
+    rails = chat.app
+
+    rails.runtime.register_action_param("kb", mock_kb)
+
+    messages = [
+        {"role": "user", "content": "Hi!"},
+    ]
+
+    new_message = rails.generate(messages=messages)
+
+    info = rails.explain()
+    assert len(info.llm_calls) == 2
+    assert "Test Body" in info.llm_calls[1].prompt