Upload folder using huggingface_hub

app.py

@@ -5,7 +5,6 @@ import asyncio
 import logging
 import time
 import traceback
-import uuid
 from html import escape
 import gradio as gr
 from dotenv import load_dotenv
@@ -13,7 +12,12 @@ from langchain_core.messages.ai import AIMessageChunk, AIMessage
 from langchain_core.messages.system import SystemMessage
 from langchain_core.messages.tool import ToolMessage
 
-from chat_utils import chat_wrapper
+from chat_utils import (
+    MAX_MESSAGES_IN_CONVERSATION,
+    chat_wrapper,
+    init_session,
+    limited_chat_wrapper,
+)
 from config import SanatanConfig
 from db import SanatanDatabase
 from drive_downloader import ZipDownloader
@@ -55,16 +59,13 @@ def init():
         downloader.unzip(zip_path, extract_to="./")
 
 
-def init_session():
-    return str(uuid.uuid4())
-
-
 def render_message_with_tooltip(content: str, max_chars=200):
     short = escape(content[:max_chars]) + ("…" if len(content) > max_chars else "")
     return f"<div title='{escape(content)}'>{short}</div>"
 
 
 # UI Elements
+message_count = gr.State(0)
 thread_id = gr.State(init_session)
 
 supported_scriptures = "\n - ".join(
@@ -213,6 +214,7 @@ with gr.Blocks(
         "📜 Explore Ashwar": "From the same azhwar as the above pasuram, show me other pasurams",
         "🏛️ Another divya desam (same āzhwār)": "show pasuram from another divya desam by the same azhwar",
         "👤 Another āzhwār (same divya desam)": "show pasuram from the same divya desam by another azhwar",
+        "❓ Quiz": "Pick any pasuram. Frame a question to ask me related to that pasuram based on its explanatory notes and word by word meanings. Output ONLY the pasuram title, the verse number, the question you framed and the answer to that question.",
     }
 
     chatbot = gr.Chatbot(
@@ -233,15 +235,23 @@ with gr.Blocks(
     )
     chatInterface = gr.ChatInterface(
         title="Sanatan-AI",
-        fn=chat_wrapper,
-        additional_inputs=[thread_id, debug_checkbox, preferred_language],
+        fn=limited_chat_wrapper,
+        additional_inputs=[
+            thread_id,
+            debug_checkbox,
+            preferred_language,
+            message_count,
+        ],
         additional_inputs_accordion=additional_inputs,
+        additional_outputs=[thread_id, message_count],
         chatbot=chatbot,
         textbox=message_textbox,
         type="messages",
     )
 
     with gr.Column(visible=False) as followup_examples:
+        with gr.Row():
+            followup_count_textbox = gr.Markdown(container=False, show_label=False)
         with gr.Row():
             gr.Examples(
                 label="Quick Navigation Follow-ups",
@@ -286,3 +296,21 @@ with gr.Blocks(
         inputs=[gr.State(False)],
         outputs=[followup_examples],
     )
+
+    def update_followup_counter(count):
+        remaining_followups = MAX_MESSAGES_IN_CONVERSATION - count
+        no_more_followups = False
+        if remaining_followups > 1:
+            text = f"✨ `{remaining_followups}` more follow-ups to go."
+        elif remaining_followups == 1:
+            text = "🌟 Just one more follow-up to go!"
+        else:
+            text = "✅ That was the last follow-up."
+            no_more_followups = True
+        return gr.update(value=text), gr.update(visible=not no_more_followups)
+
+    message_count.change(
+        update_followup_counter,
+        inputs=[message_count],
+        outputs=[followup_count_textbox, followup_examples],
+    )

chat_utils.py

@@ -1,3 +1,4 @@
+import uuid
 import json
 import random
 import asyncio
@@ -58,6 +59,7 @@ thinking_verbs = [
 
 graph = generate_graph()
 
+
 def add_node_to_tree(
     node_tree: list[str], node_label: str, tooltip: str = "no arguments to show"
 ) -> list[str]:
@@ -114,6 +116,7 @@ def chat(debug_mode, message, history, thread_id, preferred_language="English"):
     )
     return response["messages"][-1].content
 
+
 async def chat_streaming(
     debug_mode: bool, message, history, thread_id, preferred_language="English"
 ):
@@ -284,4 +287,60 @@ async def chat_streaming(
             "\nhere is what I got so far ...\n"
             f"\n{streamed_response}"
         )
-        return
+        return
+
+
+def init_session():
+    return str(uuid.uuid4())
+
+
+MAX_MESSAGES_IN_CONVERSATION = 5
+
+
+async def limited_chat_wrapper(
+    message, history, thread_id, debug, preferred_language, count
+):
+    # increment **after processing the message**
+    count += 1
+
+    # warn before reset
+    if count == MAX_MESSAGES_IN_CONVERSATION - 1:
+        yield [
+            {
+                "role": "system",
+                "content": "⚠️ You are allowed to ask one more follow-up. The next question will be considered a new conversation. Please wait ... processing your request ...",
+            }
+        ], thread_id, count
+        await asyncio.sleep(1)
+
+    # reset
+    if count > MAX_MESSAGES_IN_CONVERSATION:
+        thread_id = init_session()
+        history = []
+        count = 1
+        yield [
+            {
+                "role": "system",
+                "content": "🔄 This is now considered a new question. Don't worry, your message shall still be processed! If I am giving irrelevant responses, you know why :-)",
+            }
+        ], thread_id, count
+        await asyncio.sleep(1)
+
+    # normal flow: stream from your original chat_wrapper
+    final_chunk = []
+    async for chunk in chat_wrapper(
+        message, history, thread_id, debug, preferred_language
+    ):
+        yield chunk, thread_id, count
+        final_chunk = chunk
+
+    # Simulating LLM Response
+    # for i in range(5):
+    #     final_chunk += [{
+    #             "role": "assistant",
+    #             "content": f"Simulated LLM output {i+1}",
+    #         }]
+    #     yield final_chunk, thread_id, count
+    #     await asyncio.sleep(0.25)
+
+    yield final_chunk, thread_id, count