Implement NL2SQL tooling

.gitignore

@@ -21,7 +21,7 @@ sbin/**
 !opentofu/examples/manual-test.sh
 !src/entrypoint.sh
 !src/client/spring_ai/templates/env.sh
-tests/db_startup_temp/**
+test*/db_startup_temp
 
 ##############################################################################
 # Environment (PyVen, IDE, etc.)

.pylintrc

@@ -52,7 +52,7 @@ ignore=CVS,.venv
 # ignore-list. The regex matches against paths and can be in Posix or Windows
 # format. Because '\\' represents the directory delimiter on Windows systems,
 # it can't be used as an escape character.
-ignore-paths=.*[/\\]wip[/\\].*,src/client/mcp,docs/themes/relearn,docs/public,docs/static/demoware
+ignore-paths=.*[/\\]wip[/\\].*,src/client/mcp,docs/themes/relearn,docs/public,docs/static/demoware,src/server/agents/chatbot.py
 
 # Files or directories matching the regular expression patterns are skipped.
 # The regex matches against base names, not paths. The default value ignores

pyproject.toml

@@ -32,7 +32,7 @@ server = [
     "langchain-aimlapi==0.1.0",
     "langchain-cohere==0.4.6",
     "langchain-community==0.3.31",
-    "langchain-fireworks==0.3.0",
+    #"langchain-fireworks==0.3.0",
     "langchain-google-genai==2.1.12",
     "langchain-ibm==0.3.20",
     "langchain-mcp-adapters==0.1.13",
@@ -43,7 +43,7 @@ server = [
     "langchain-openai==0.3.35",
     "langchain-together==0.3.1",
     "langgraph==1.0.1",
-    "litellm==1.80.0",
+    "litellm==1.80.7",
     "llama-index==0.14.8",
     "lxml==6.0.2",
     "matplotlib==3.10.7",
@@ -70,6 +70,7 @@ test = [
     "pytest",
     "pytest-asyncio",
     "pytest-cov",
+    "types-jsonschema",
     "yamllint"
 ]
 

src/client/content/chatbot.py

@@ -16,7 +16,7 @@
 from streamlit import session_state as state
 
 from client.content.config.tabs.models import get_models
-from client.utils import st_common, api_call, client
+from client.utils import st_common, api_call, client, vs_options
 from client.utils.st_footer import render_chat_footer
 from common import logging_config
 
@@ -26,15 +26,15 @@
 #############################################################################
 # Functions
 #############################################################################
-def show_vector_search_refs(context):
+def show_vector_search_refs(context, vs_metadata=None):
     """When Vector Search Content Found, show the references"""
     st.markdown("**References:**")
     ref_src = set()
     ref_cols = st.columns([3, 3, 3])
     # Create a button in each column
-    for i, (ref_col, chunk) in enumerate(zip(ref_cols, context[0])):
+    for i, (ref_col, chunk) in enumerate(zip(ref_cols, context["documents"])):
         with ref_col.popover(f"Reference: {i + 1}"):
-            chunk = context[0][i]
+            chunk = context["documents"][i]
             logger.debug("Chunk Content: %s", chunk)
             st.subheader("Reference Text", divider="red")
             st.markdown(chunk["page_content"])
@@ -46,9 +46,32 @@ def show_vector_search_refs(context):
             except KeyError:
                 logger.error("Chunk Metadata NOT FOUND!!")
 
-    for link in ref_src:
-        st.markdown("- " + link)
-    st.markdown(f"**Notes:** Vector Search Query - {context[1]}")
+    # Display Vector Search details in expander
+    if vs_metadata or ref_src:
+        with st.expander("Vector Search Details", expanded=False):
+            if ref_src:
+                st.markdown("**Source Documents:**")
+                for link in ref_src:
+                    st.markdown(f"- {link}")
+
+            if vs_metadata and vs_metadata.get("searched_tables"):
+                st.markdown("**Tables Searched:**")
+                for table in vs_metadata["searched_tables"]:
+                    st.markdown(f"- {table}")
+
+            if vs_metadata and vs_metadata.get("context_input"):
+                st.markdown(f"**Search Query:** {vs_metadata.get('context_input')}")
+            elif context.get("context_input"):
+                st.markdown(f"**Search Query:** {context.get('context_input')}")
+
+
+def show_token_usage(token_usage):
+    """Display token usage for AI responses using caption"""
+    if token_usage:
+        prompt_tokens = token_usage.get("prompt_tokens", 0)
+        completion_tokens = token_usage.get("completion_tokens", 0)
+        total_tokens = token_usage.get("total_tokens", 0)
+        st.caption(f"Token usage: {prompt_tokens} prompt + {completion_tokens} completion = {total_tokens} total")
 
 
 def setup_sidebar():
@@ -62,7 +85,7 @@ def setup_sidebar():
     st_common.tools_sidebar()
     st_common.history_sidebar()
     st_common.ll_sidebar()
-    st_common.vector_search_sidebar()
+    vs_options.vector_search_sidebar()
 
     if not state.enable_client:
         st.stop()
@@ -80,22 +103,33 @@ def create_client():
 
 
 def display_chat_history(history):
-    """Display chat history messages"""
+    """Display chat history messages with metadata"""
     st.chat_message("ai").write("Hello, how can I help you?")
     vector_search_refs = []
 
     for message in history or []:
         if not message["content"]:
             continue
 
-        if message["role"] == "tool" and message["name"] == "oraclevs_tool":
+        if message["role"] == "tool" and message["name"] == "optimizer_vs-retriever":
             vector_search_refs = json.loads(message["content"])
 
         elif message["role"] in ("ai", "assistant"):
             with st.chat_message("ai"):
                 st.markdown(message["content"])
+
+                # Extract metadata from response_metadata
+                response_metadata = message.get("response_metadata", {})
+                vs_metadata = response_metadata.get("vs_metadata", {})
+                token_usage = response_metadata.get("token_usage", {})
+
+                # Show token usage immediately after message
+                if token_usage:
+                    show_token_usage(token_usage)
+
+                # Show vector search references if available
                 if vector_search_refs:
-                    show_vector_search_refs(vector_search_refs)
+                    show_vector_search_refs(vector_search_refs, vs_metadata)
                     vector_search_refs = []
 
         elif message["role"] in ("human", "user"):
@@ -131,9 +165,32 @@ async def handle_chat_input(user_client):
         try:
             message_placeholder = st.chat_message("ai").empty()
             full_answer = ""
-            async for chunk in user_client.stream(message=human_request.text, image_b64=file_b64):
-                full_answer += chunk
-                message_placeholder.markdown(full_answer)
+
+            # Animated thinking indicator
+            async def animate_thinking():
+                """Animate the thinking indicator with increasing dots"""
+                dots = 0
+                while True:
+                    message_placeholder.markdown(f"🤔 Thinking{'.' * (dots % 4)}")
+                    dots += 1
+                    await asyncio.sleep(0.5)  # Update every 500ms
+
+            # Start the thinking animation
+            thinking_task = asyncio.create_task(animate_thinking())
+
+            try:
+                async for chunk in user_client.stream(message=human_request.text, image_b64=file_b64):
+                    # Cancel thinking animation on first chunk
+                    if thinking_task and not thinking_task.done():
+                        thinking_task.cancel()
+                        thinking_task = None
+                    full_answer += chunk
+                    message_placeholder.markdown(full_answer)
+            finally:
+                # Ensure thinking task is cancelled
+                if thinking_task and not thinking_task.done():
+                    thinking_task.cancel()
+
             st.rerun()
         except (ConnectionError, TimeoutError, api_call.ApiError) as ex:
             logger.exception("Error during chat streaming: %s", ex)

src/client/content/config/tabs/settings.py

@@ -363,17 +363,23 @@ def spring_ai_conf_check(ll_model: dict, embed_model: dict) -> str:
 def spring_ai_obaas(src_dir, file_name, provider, ll_config, embed_config):
     """Get the system prompt for SpringAI export"""
 
+    ## FUTURE FEATURE:
     # Determine which system prompt would be active based on tools_enabled
-    tools_enabled = state.client_settings.get("tools_enabled", [])
+    # tools_enabled = state.client_settings.get("tools_enabled", [])
 
     # Select prompt name based on tools configuration
-    if not tools_enabled:
-        prompt_name = "optimizer_basic-default"
-        if state.client_settings["vector_search"]["enabled"]:
-            prompt_name = "optimizer_vs-no-tools-default"
+    # if not tools_enabled:
+    #     prompt_name = "optimizer_basic-default"
+    #     if state.client_settings["vector_search"]["enabled"]:
+    #         prompt_name = "optimizer_vs-no-tools-default"
+    # else:
+    #     # Tools are enabled, use tools-default prompt
+    #     prompt_name = "optimizer_tools-default"
+    ## Legacy Feature:
+    if "Vector Search" in state.client_settings.get("tools_enabled", []):
+        prompt_name = "optimizer_vs-no-tools-default"
     else:
-        # Tools are enabled, use tools-default prompt
-        prompt_name = "optimizer_tools-default"
+        prompt_name = "optimizer_basic-default"
 
     # Find the prompt in configs
     sys_prompt_obj = next((item for item in state.prompt_configs if item["name"] == prompt_name), None)

src/client/content/testbed.py

@@ -17,7 +17,7 @@
 
 from client.content.config.tabs.models import get_models
 
-from client.utils import st_common, api_call
+from client.utils import st_common, api_call, vs_options
 
 from common import logging_config
 
@@ -85,9 +85,12 @@ def create_gauge(value):
     st.dataframe(ll_settings_reversed, hide_index=True)
     if report["settings"]["testbed"]["judge_model"]:
         st.markdown(f"**Judge Model**: {report['settings']['testbed']['judge_model']}")
-    # if discovery; then list out the tables that were discovered (MCP implementation)
-    # if report["settings"]["vector_search"].get("discovery"):
-    if report["settings"]["vector_search"]["enabled"]:
+    # Backward compatibility
+    try:
+        vs_enabled = report["settings"]["vector_search"]["enabled"]
+    except KeyError:
+        vs_enabled = "Vector Search" in report["settings"]["tools_enabled"]
+    if vs_enabled:
         st.subheader("Vector Search Settings")
         st.markdown(f"""**Database**: {report["settings"]["database"]["alias"]};
             **Vector Store**: {report["settings"]["vector_search"]["vector_store"]}
@@ -100,6 +103,8 @@ def create_gauge(value):
         if report["settings"]["vector_search"]["search_type"] == "Similarity":
             embed_settings.drop(["score_threshold", "fetch_k", "lambda_mult"], axis=1, inplace=True)
         st.dataframe(embed_settings, hide_index=True)
+        # if discovery; then list out the tables that were discovered (MCP implementation)
+        # if report["settings"]["vector_search"].get("discovery"):
     else:
         st.markdown("**Evaluated without Vector Search**")
 
@@ -493,7 +498,7 @@ def render_evaluation_ui(available_ll_models: list) -> None:
     st.info("Use the sidebar settings for chatbot evaluation parameters", icon="⬅️")
     st_common.tools_sidebar()
     st_common.ll_sidebar()
-    st_common.vector_search_sidebar()
+    vs_options.vector_search_sidebar()
     st.write("Choose a model to judge the correctness of the chatbot answer, then start evaluation.")
     col_left, col_center, _ = st.columns([4, 3, 3])
 
@@ -510,20 +515,13 @@ def render_evaluation_ui(available_ll_models: list) -> None:
         on_change=st_common.update_client_settings("testbed"),
     )
 
-    # Check if vector search is enabled but no vector store is selected
-    evaluation_disabled = False
-    if state.client_settings.get("vector_search", {}).get("enabled", False):
-        # If vector search is enabled, check if a vector store is selected
-        if not state.client_settings.get("vector_search", {}).get("vector_store"):
-            evaluation_disabled = True
-
     if col_center.button(
         "Start Evaluation",
         type="primary",
         key="evaluate_button",
         help="Evaluation will automatically save the TestSet to the Database",
         on_click=qa_update_db,
-        disabled=evaluation_disabled,
+        disabled=not state.get("enable_client", True),
     ):
         with st.spinner("Starting Q&A evaluation... please be patient.", show_time=True):
             st_common.clear_state_key("testbed_evaluations")