Upload folder using huggingface_hub

.github/workflows/update_space.yml

@@ -1,28 +1,28 @@
-name: Run Python script
-
-on:
-  push:
-    branches:
-      - main
-
-jobs:
-  build:
-    runs-on: ubuntu-latest
-
-    steps:
-    - name: Checkout
-      uses: actions/checkout@v2
-
-    - name: Set up Python
-      uses: actions/setup-python@v2
-      with:
-        python-version: '3.9'
-
-    - name: Install Gradio
-      run: python -m pip install gradio
-
-    - name: Log in to Hugging Face
-      run: python -c 'import huggingface_hub; huggingface_hub.login(token="${{ secrets.hf_token }}")'
-
-    - name: Deploy to Spaces
-      run: gradio deploy
+name: Run Python script
+
+on:
+  push:
+    branches:
+      - main
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+
+    steps:
+    - name: Checkout
+      uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: '3.9'
+
+    - name: Install Gradio
+      run: python -m pip install gradio
+
+    - name: Log in to Hugging Face
+      run: python -c 'import huggingface_hub; huggingface_hub.login(token="${{ secrets.hf_token }}")'
+
+    - name: Deploy to Spaces
+      run: gradio deploy

README.md

@@ -3,4 +3,5 @@ title: sanatan_ai
 app_file: app.py
 sdk: gradio
 sdk_version: 5.38.0
+python_version: 3.12
 ---

app.py

@@ -1,74 +1,74 @@
-import os
-import gradio as gr
-from config import SanatanConfig
-from drive_downloader import ZipDownloader
-from graph_helper import generate_graph
-import uuid
-import logging
-from dotenv import load_dotenv
-
-logging.basicConfig()
-logger = logging.getLogger()
-logger.setLevel(logging.INFO)
-
-graph = generate_graph()
-
-def init():
-    load_dotenv(override=True)
-
-    downloader = ZipDownloader(
-        service_account_json=os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON")
-    )
-    # 1. Download zip file
-    zip_path = downloader.download_zip_from_drive(
-        file_id=os.getenv("CHROMADB_FILE_ID"),
-        output_path=SanatanConfig.dbStorePath,
-    )
-
-    # 2. Extract it
-    downloader.unzip(zip_path, extract_to="./")
-
-def init_session():
-    # Generate a unique ID per browser session
-    return str(uuid.uuid4())
-
-
-def chat(message, history, thread_id):
-    print("received chat message for thread:", thread_id)
-    config = {"configurable": {"thread_id": thread_id}}
-    response = graph.invoke(
-        {"messages": [{"role": "user", "content": message}]}, config=config
-    )
-    return response["messages"][-1].content
-
-
-thread_id = gr.State(init_session)
-supported_scriptures = "\n - ".join(
-    [scripture["title"] for scripture in SanatanConfig.scriptures]
-)
-intro_messages = [
-    {"role" : "assistant","content" : f"Namaskaram! I can assist with the following scriptures: \n - {supported_scriptures}"},
-    {"role" : "assistant","content" : """
-    #### You can ask me questions like:
-      - How do all scriptures describe the form of Vishnu?
-      - What did Krishna teach in Gita?
-      - How did Arjun react on seeing the Vishwarupa form of Krishna?
-      - Give some names of Vishnu from Sahasranamam.
-      """}
-]
-chatbot = gr.Chatbot(
-    elem_id="chatbot",
-    avatar_images=("assets/avatar_user.png", "assets/adiyen_bot.png"),
-    value=intro_messages,
-    label="Sanatan-Bot",
-    show_copy_button=True,
-    show_copy_all_button=True,
-    type="messages",
-    height=800
-)
-chatInterface = gr.ChatInterface(title="Sanatan-AI",
-    fn=chat, additional_inputs=[thread_id], chatbot=chatbot
-)
-# initializze database
-init()
-chatInterface.launch()
+import os
+import gradio as gr
+from config import SanatanConfig
+from drive_downloader import ZipDownloader
+from graph_helper import generate_graph
+import uuid
+import logging
+from dotenv import load_dotenv
+
+logging.basicConfig()
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+graph = generate_graph()
+
+def init():
+    load_dotenv(override=True)
+
+    downloader = ZipDownloader(
+        service_account_json=os.getenv("GOOGLE_SERVICE_ACCOUNT_JSON")
+    )
+    # 1. Download zip file
+    zip_path = downloader.download_zip_from_drive(
+        file_id=os.getenv("CHROMADB_FILE_ID"),
+        output_path=SanatanConfig.dbStorePath,
+    )
+
+    # 2. Extract it
+    downloader.unzip(zip_path, extract_to="./")
+
+def init_session():
+    # Generate a unique ID per browser session
+    return str(uuid.uuid4())
+
+
+def chat(message, history, thread_id):
+    print("received chat message for thread:", thread_id)
+    config = {"configurable": {"thread_id": thread_id}}
+    response = graph.invoke(
+        {"messages": [{"role": "user", "content": message}]}, config=config
+    )
+    return response["messages"][-1].content
+
+
+thread_id = gr.State(init_session)
+supported_scriptures = "\n - ".join(
+    [scripture["title"] for scripture in SanatanConfig.scriptures]
+)
+intro_messages = [
+    {"role" : "assistant","content" : f"Namaskaram! I can assist with the following scriptures: \n - {supported_scriptures}"},
+    {"role" : "assistant","content" : """
+    #### You can ask me questions like:
+      - How do all scriptures describe the form of Vishnu?
+      - What did Krishna teach in Gita?
+      - How did Arjun react on seeing the Vishwarupa form of Krishna?
+      - Give some names of Vishnu from Sahasranamam.
+      """}
+]
+chatbot = gr.Chatbot(
+    elem_id="chatbot",
+    avatar_images=("assets/avatar_user.png", "assets/adiyen_bot.png"),
+    value=intro_messages,
+    label="Sanatan-Bot",
+    show_copy_button=True,
+    show_copy_all_button=True,
+    type="messages",
+    height=800
+)
+chatInterface = gr.ChatInterface(title="Sanatan-AI",
+    fn=chat, additional_inputs=[thread_id], chatbot=chatbot
+)
+# initializze database
+init()
+chatInterface.launch()

config.py

@@ -1,108 +1,108 @@
-class SanatanConfig:
-    # shuklaYajurVedamPdfPath: str = "./data/shukla-yajur-veda.pdf"
-    # shuklaYajurVedamSmallPdfPath: str = "./data/shukla-yajur-veda-small.pdf"
-    # vishnuPuranamPdfPath = "./data/vishnu_puranam.pdf"
-    # datastores = [{"name": "sanskrit_001", "dbStorePath": "./chromadb-store"}, {"name": "nalayiram", "dbStorePath": "./chromadb-store-4000"}]
-    dbStorePath: str = "./chromadb-store"
-    # shuklaYajurVedamCollectionName: str = "shukla_yajur_vedam"
-    # vishnuPuranamCollectionName: str = "vishnu_puranam"
-    # shuklaYajurVedamOutputDir = "./output/shukla_yajur_vedam"
-    # vishnuPuranamOutputDir = "./output/vishnu_puranam"
-    scriptures = [
-        {
-            "name": "vishnu_puranam",
-            "title": "Sri Vishnu Puranam",
-            "output_dir": "./output/vishnu_puranam",
-            "collection_name": "vishnu_puranam",
-            "pdf_path": "./data/vishnu_puranam.pdf",
-            "language": "san+eng",
-            "example_labels": [
-                "Vishnu's form",
-                "About the five elements",
-                "About Garuda",
-                "Weapons of Vishnu",
-            ],
-            "examples": [
-                "describe Vishnu's form",
-                "five elements and their significance",
-                "What is the significance of Garuda? Show some verses that describe him.",
-                "What weapons does Vishnu hold?",
-            ],
-        },
-        {
-            "name": "shukla_yajur_vedam",
-            "title": "Shukla Yajur Vedam",
-            "output_dir": "./output/shukla_yajur_vedam",
-            "collection_name": "shukla_yajur_vedam",
-            "pdf_path": "./data/shukla-yajur-veda.pdf",
-            "language": "san+eng",
-            "example_labels": [
-                "About Vedam",
-                "About the five elements",
-                "About Brahma",
-            ],
-            "examples": [
-                "Gist of Shukla Yajur Vedam. Give me some sanskrit verses.",
-                "What is the significance of fire and water. show some sanskrit verses",
-                "Brahma",
-            ],
-        },
-        {
-            "name": "bhagavat_gita",
-            "title": "Bhagavat Gita",
-            "output_dir": "./output/bhagavat_gita",
-            "collection_name": "bhagavat_gita",
-            "pdf_path": "./data/bhagavat_gita.pdf",
-            "language": "san+eng",
-            "example_labels": [
-                "About Arjuna",
-                "About Karma",
-                "About birth and death",
-                "About the battle field",
-                "About Krishna's form",
-            ],
-            "examples": [
-                "Show some verses where Krishna advises Arjuna",
-                "What does Krishna say about Karma",
-                "What does Krishna say about birth and death",
-                "describe the battle field",
-                "Vishwarupa",
-            ],
-        },
-        {
-            "name": "valmiki_ramayanam",
-            "title": "Valmiki Ramayanam",
-            "output_dir": "./output/valmiki_ramayanam",
-            "collection_name": "valmiki_ramayanam",
-            "pdf_path": "./data/valmiki_ramayanam.pdf",
-            "language": "san+eng",
-            "example_labels": [
-                "About Jatayu",
-                "About Hanuman",
-                "About Vali",
-                "About Sita",
-                "About Ravana",
-            ],
-            "examples": [
-                "What is the significance of Jatayu? show some sanskrit verses to support the argument",
-                "Show some verses where Hanuman is mentioned",
-                "How did Rama kill Vali",
-                "How was Sita abducted",
-                "How did Rama kill Ravana?",
-            ],
-        },
-        {
-            "name": "vishnu_sahasranamam",
-            "title": "Vishnu Sahasranamam",
-            "output_dir": "./output/vishnu_sahasranamam",
-            "collection_name": "vishnu_sahasranamam",
-            "pdf_path": "./data/vishnu_sahasranamam.pdf",
-            "language": "san+eng",
-            "example_labels": ["Vanamali", "1000 names", "Sanskrit text search"],
-            "examples": [
-                "Vanamali",
-                "Show some of the 1000 names of Vishnu along with their meaning",
-                "show the verse that begins with शुक्लाम्बरधरं",
-            ],
-        },
-    ]
+class SanatanConfig:
+    # shuklaYajurVedamPdfPath: str = "./data/shukla-yajur-veda.pdf"
+    # shuklaYajurVedamSmallPdfPath: str = "./data/shukla-yajur-veda-small.pdf"
+    # vishnuPuranamPdfPath = "./data/vishnu_puranam.pdf"
+    # datastores = [{"name": "sanskrit_001", "dbStorePath": "./chromadb-store"}, {"name": "nalayiram", "dbStorePath": "./chromadb-store-4000"}]
+    dbStorePath: str = "./chromadb-store"
+    # shuklaYajurVedamCollectionName: str = "shukla_yajur_vedam"
+    # vishnuPuranamCollectionName: str = "vishnu_puranam"
+    # shuklaYajurVedamOutputDir = "./output/shukla_yajur_vedam"
+    # vishnuPuranamOutputDir = "./output/vishnu_puranam"
+    scriptures = [
+        {
+            "name": "vishnu_puranam",
+            "title": "Sri Vishnu Puranam",
+            "output_dir": "./output/vishnu_puranam",
+            "collection_name": "vishnu_puranam",
+            "pdf_path": "./data/vishnu_puranam.pdf",
+            "language": "san+eng",
+            "example_labels": [
+                "Vishnu's form",
+                "About the five elements",
+                "About Garuda",
+                "Weapons of Vishnu",
+            ],
+            "examples": [
+                "describe Vishnu's form",
+                "five elements and their significance",
+                "What is the significance of Garuda? Show some verses that describe him.",
+                "What weapons does Vishnu hold?",
+            ],
+        },
+        {
+            "name": "shukla_yajur_vedam",
+            "title": "Shukla Yajur Vedam",
+            "output_dir": "./output/shukla_yajur_vedam",
+            "collection_name": "shukla_yajur_vedam",
+            "pdf_path": "./data/shukla-yajur-veda.pdf",
+            "language": "san+eng",
+            "example_labels": [
+                "About Vedam",
+                "About the five elements",
+                "About Brahma",
+            ],
+            "examples": [
+                "Gist of Shukla Yajur Vedam. Give me some sanskrit verses.",
+                "What is the significance of fire and water. show some sanskrit verses",
+                "Brahma",
+            ],
+        },
+        {
+            "name": "bhagavat_gita",
+            "title": "Bhagavat Gita",
+            "output_dir": "./output/bhagavat_gita",
+            "collection_name": "bhagavat_gita",
+            "pdf_path": "./data/bhagavat_gita.pdf",
+            "language": "san+eng",
+            "example_labels": [
+                "About Arjuna",
+                "About Karma",
+                "About birth and death",
+                "About the battle field",
+                "About Krishna's form",
+            ],
+            "examples": [
+                "Show some verses where Krishna advises Arjuna",
+                "What does Krishna say about Karma",
+                "What does Krishna say about birth and death",
+                "describe the battle field",
+                "Vishwarupa",
+            ],
+        },
+        {
+            "name": "valmiki_ramayanam",
+            "title": "Valmiki Ramayanam",
+            "output_dir": "./output/valmiki_ramayanam",
+            "collection_name": "valmiki_ramayanam",
+            "pdf_path": "./data/valmiki_ramayanam.pdf",
+            "language": "san+eng",
+            "example_labels": [
+                "About Jatayu",
+                "About Hanuman",
+                "About Vali",
+                "About Sita",
+                "About Ravana",
+            ],
+            "examples": [
+                "What is the significance of Jatayu? show some sanskrit verses to support the argument",
+                "Show some verses where Hanuman is mentioned",
+                "How did Rama kill Vali",
+                "How was Sita abducted",
+                "How did Rama kill Ravana?",
+            ],
+        },
+        {
+            "name": "vishnu_sahasranamam",
+            "title": "Vishnu Sahasranamam",
+            "output_dir": "./output/vishnu_sahasranamam",
+            "collection_name": "vishnu_sahasranamam",
+            "pdf_path": "./data/vishnu_sahasranamam.pdf",
+            "language": "san+eng",
+            "example_labels": ["Vanamali", "1000 names", "Sanskrit text search"],
+            "examples": [
+                "Vanamali",
+                "Show some of the 1000 names of Vishnu along with their meaning",
+                "show the verse that begins with शुक्लाम्बरधरं",
+            ],
+        },
+    ]

db.py

@@ -1,35 +1,35 @@
-import chromadb
-
-from config import SanatanConfig
-from embeddings import get_embedding
-import logging
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-
-class SanatanDatabase:
-    def __init__(self) -> None:
-        self.chroma_client = chromadb.PersistentClient(path=SanatanConfig.dbStorePath)
-
-    def does_data_exist(self, collection_name: str) -> bool:
-        collection = self.chroma_client.get_or_create_collection(name=collection_name)
-        num_rows = collection.count()
-        logger.info("num_rows in %s = %d", collection_name, num_rows)
-        return num_rows > 0
-
-    def load(self, collection_name: str, ids, documents, embeddings, metadatas):
-        collection = self.chroma_client.get_or_create_collection(name=collection_name)
-        collection.add(
-            ids=ids,
-            documents=documents,
-            embeddings=embeddings,
-            metadatas=metadatas,
-        )
-
-    def search(self, collection_name: str, query: str, n_results=2):
-        collection = self.chroma_client.get_or_create_collection(name=collection_name)
-        response = collection.query(
-            query_embeddings=[get_embedding(query)], n_results=n_results
-        )
-        return response
+import chromadb
+
+from config import SanatanConfig
+from embeddings import get_embedding
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+class SanatanDatabase:
+    def __init__(self) -> None:
+        self.chroma_client = chromadb.PersistentClient(path=SanatanConfig.dbStorePath)
+
+    def does_data_exist(self, collection_name: str) -> bool:
+        collection = self.chroma_client.get_or_create_collection(name=collection_name)
+        num_rows = collection.count()
+        logger.info("num_rows in %s = %d", collection_name, num_rows)
+        return num_rows > 0
+
+    def load(self, collection_name: str, ids, documents, embeddings, metadatas):
+        collection = self.chroma_client.get_or_create_collection(name=collection_name)
+        collection.add(
+            ids=ids,
+            documents=documents,
+            embeddings=embeddings,
+            metadatas=metadatas,
+        )
+
+    def search(self, collection_name: str, query: str, n_results=2):
+        collection = self.chroma_client.get_or_create_collection(name=collection_name)
+        response = collection.query(
+            query_embeddings=[get_embedding(query)], n_results=n_results
+        )
+        return response

drive_downloader.py

@@ -1,44 +1,44 @@
-import os
-import json
-import io
-import zipfile
-
-from google.oauth2.service_account import Credentials
-from googleapiclient.discovery import build
-from googleapiclient.http import MediaIoBaseDownload
-
-
-class ZipDownloader:
-    def __init__(self, service_account_json: str):
-        self.creds = Credentials.from_service_account_info(json.loads(service_account_json))
-        self.service = build("drive", "v3", credentials=self.creds)
-
-    def download_zip_from_drive(self, file_id: str, output_path: str) -> str:
-        """
-        Downloads a ZIP file from Google Drive and saves it locally.
-        Returns the path to the downloaded ZIP.
-        """
-        request = self.service.files().get_media(fileId=file_id)
-        local_zip_path = os.path.join(output_path, "downloaded.zip")
-
-        os.makedirs(output_path, exist_ok=True)
-        fh = io.FileIO(local_zip_path, 'wb')
-        downloader = MediaIoBaseDownload(fh, request)
-
-        print(f"⬇️ Downloading ZIP file from Drive ID: {file_id}")
-        done = False
-        while not done:
-            status, done = downloader.next_chunk()
-            print(f"   ⏬ Progress: {int(status.progress() * 100)}%")
-
-        print(f"✅ ZIP downloaded to: {local_zip_path}")
-        return local_zip_path
-
-    def unzip(self, zip_path: str, extract_to: str):
-        """
-        Unzips the downloaded ZIP file to a specified directory.
-        """
-        print(f"📂 Extracting ZIP: {zip_path} -> {extract_to}")
-        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
-            zip_ref.extractall(extract_to)
-        print("✅ Extraction complete.")
+import os
+import json
+import io
+import zipfile
+
+from google.oauth2.service_account import Credentials
+from googleapiclient.discovery import build
+from googleapiclient.http import MediaIoBaseDownload
+
+
+class ZipDownloader:
+    def __init__(self, service_account_json: str):
+        self.creds = Credentials.from_service_account_info(json.loads(service_account_json))
+        self.service = build("drive", "v3", credentials=self.creds)
+
+    def download_zip_from_drive(self, file_id: str, output_path: str) -> str:
+        """
+        Downloads a ZIP file from Google Drive and saves it locally.
+        Returns the path to the downloaded ZIP.
+        """
+        request = self.service.files().get_media(fileId=file_id)
+        local_zip_path = os.path.join(output_path, "downloaded.zip")
+
+        os.makedirs(output_path, exist_ok=True)
+        fh = io.FileIO(local_zip_path, 'wb')
+        downloader = MediaIoBaseDownload(fh, request)
+
+        print(f"⬇️ Downloading ZIP file from Drive ID: {file_id}")
+        done = False
+        while not done:
+            status, done = downloader.next_chunk()
+            print(f"   ⏬ Progress: {int(status.progress() * 100)}%")
+
+        print(f"✅ ZIP downloaded to: {local_zip_path}")
+        return local_zip_path
+
+    def unzip(self, zip_path: str, extract_to: str):
+        """
+        Unzips the downloaded ZIP file to a specified directory.
+        """
+        print(f"📂 Extracting ZIP: {zip_path} -> {extract_to}")
+        with zipfile.ZipFile(zip_path, 'r') as zip_ref:
+            zip_ref.extractall(extract_to)
+        print("✅ Extraction complete.")

embeddings.py

@@ -1,9 +1,9 @@
-
-from sentence_transformers import SentenceTransformer
-
-# Step 1: Load SentenceTransformer model
-# model = SentenceTransformer("all-MiniLM-L6-v2")
-model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
-
-def get_embedding(text: str) -> list:
-    return model.encode(text).tolist()
+
+from sentence_transformers import SentenceTransformer
+
+# Step 1: Load SentenceTransformer model
+# model = SentenceTransformer("all-MiniLM-L6-v2")
+model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")
+
+def get_embedding(text: str) -> list:
+    return model.encode(text).tolist()

graph_helper.py

@@ -1,103 +1,103 @@
-from typing import Annotated, TypedDict
-from langgraph.graph import StateGraph, START, END
-from langgraph.checkpoint.memory import MemorySaver
-from langgraph.graph.message import add_messages
-from langchain_openai import ChatOpenAI
-from langgraph.graph.state import CompiledStateGraph
-from tools import (
-    tool_format_scripture_answer,
-    tool_search_db,
-    tool_search_web,
-    tool_push,
-)
-from langgraph.prebuilt import ToolNode, tools_condition
-from langchain_core.messages import SystemMessage, ToolMessage, HumanMessage
-import logging
-
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-
-class ChatState(TypedDict):
-    messages: Annotated[list[str], add_messages]
-
-
-def generate_graph() -> CompiledStateGraph:
-    memory = MemorySaver()
-    tools = [tool_search_web, tool_push, tool_search_db, tool_format_scripture_answer]
-    llm = ChatOpenAI(model="gpt-4o-mini").bind_tools(tools)
-
-    def chatNode(state: ChatState) -> ChatState:
-        # logger.info("messages before LLM: %s", str(state["messages"]))
-        response = llm.invoke(state["messages"])
-        # return {"messages": [response]}
-        return {"messages": state["messages"] + [response]}
-
-    def init_system_prompt_node(state: ChatState) -> ChatState:
-        messages = state["messages"] or []
-
-        # Check if system prompts were already added
-        already_has_prompt = any(
-            isinstance(m, SystemMessage) and "format_scripture_answer" in m.content
-            for m in messages
-        )
-
-        if not already_has_prompt:
-            messages += [
-                SystemMessage(
-                    content="⚠️ Do NOT summarize or compress the output from the `query` tool. It will be passed directly to `format_scripture_answer` tool that formats the answer **AS IS**. DO NOT REMOVE SANSKRIT TEXTS"
-                ),
-                SystemMessage(
-                    content="You MUST call the `format_scripture_answer` tool if the user question is about scripture content and the `query` tool has returned a result."
-                ),
-                SystemMessage(
-                    content="If the user's question is about any scripture content (even if multiple scriptures), you must use the `tool_search_db`. Only use `tool_search_web` for general non-scriptural questions."
-                ),
-            ]
-
-        return {"messages": messages}
-
-    # def chatNode(state: ChatState) -> ChatState:
-    #     messages = state["messages"]
-    #     system_prompt = None
-    #     new_messages = []
-
-    #     for m in messages:
-    #         if isinstance(m, ToolMessage):
-    #             print("m.name = ", m.name)
-    #             if m.name == "format_scripture_answer":
-    #                 system_prompt = m.content
-    #         else:
-    #             new_messages.append(m)
-
-    #     if system_prompt:
-    #         full_history = [
-    #             SystemMessage(content=system_prompt),
-    #             SystemMessage(
-    #                 content="⚠️ Do NOT summarize or compress the context from the query tool. It will be passed directly to another tool that formats the answer."
-    #             ),
-    #             SystemMessage(
-    #                 content="You MUST call the `format_scripture_answer` tool if the user question is about scripture content and the query tool has returned a result."
-    #             ),
-    #         ] + new_messages
-    #     else:
-    #         full_history = messages
-
-    #     # 🔍 Debug log (optional)
-    #     # print("\n🧠 LLM Full History:")
-    #     # for m in full_history:
-    #     #    print(f"- {m.type.upper()}: {m.content[:100]}...\n")
-
-    #     ai_response = llm.invoke(full_history)
-    #     return {"messages": messages + [ai_response]}
-
-    graph = StateGraph(ChatState)
-    graph.add_node("init", init_system_prompt_node)
-    graph.add_node("chat", chatNode)
-    graph.add_node("tools", ToolNode(tools))
-    graph.add_edge(START, "init")
-    graph.add_edge("init", "chat")
-    graph.add_conditional_edges("chat", tools_condition, "tools")
-    graph.add_edge("tools", "chat")
-
-    return graph.compile(checkpointer=memory)
+from typing import Annotated, TypedDict
+from langgraph.graph import StateGraph, START, END
+from langgraph.checkpoint.memory import MemorySaver
+from langgraph.graph.message import add_messages
+from langchain_openai import ChatOpenAI
+from langgraph.graph.state import CompiledStateGraph
+from tools import (
+    tool_format_scripture_answer,
+    tool_search_db,
+    tool_search_web,
+    tool_push,
+)
+from langgraph.prebuilt import ToolNode, tools_condition
+from langchain_core.messages import SystemMessage, ToolMessage, HumanMessage
+import logging
+
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+
+class ChatState(TypedDict):
+    messages: Annotated[list[str], add_messages]
+
+
+def generate_graph() -> CompiledStateGraph:
+    memory = MemorySaver()
+    tools = [tool_search_web, tool_push, tool_search_db, tool_format_scripture_answer]
+    llm = ChatOpenAI(model="gpt-4o-mini").bind_tools(tools)
+
+    def chatNode(state: ChatState) -> ChatState:
+        # logger.info("messages before LLM: %s", str(state["messages"]))
+        response = llm.invoke(state["messages"])
+        # return {"messages": [response]}
+        return {"messages": state["messages"] + [response]}
+
+    def init_system_prompt_node(state: ChatState) -> ChatState:
+        messages = state["messages"] or []
+
+        # Check if system prompts were already added
+        already_has_prompt = any(
+            isinstance(m, SystemMessage) and "format_scripture_answer" in m.content
+            for m in messages
+        )
+
+        if not already_has_prompt:
+            messages += [
+                SystemMessage(
+                    content="⚠️ Do NOT summarize or compress the output from the `query` tool. It will be passed directly to `format_scripture_answer` tool that formats the answer **AS IS**. DO NOT REMOVE SANSKRIT TEXTS"
+                ),
+                SystemMessage(
+                    content="You MUST call the `format_scripture_answer` tool if the user question is about scripture content and the `query` tool has returned a result."
+                ),
+                SystemMessage(
+                    content="If the user's question is about any scripture content (even if multiple scriptures), you must use the `tool_search_db`. Only use `tool_search_web` for general non-scriptural questions."
+                ),
+            ]
+
+        return {"messages": messages}
+
+    # def chatNode(state: ChatState) -> ChatState:
+    #     messages = state["messages"]
+    #     system_prompt = None
+    #     new_messages = []
+
+    #     for m in messages:
+    #         if isinstance(m, ToolMessage):
+    #             print("m.name = ", m.name)
+    #             if m.name == "format_scripture_answer":
+    #                 system_prompt = m.content
+    #         else:
+    #             new_messages.append(m)
+
+    #     if system_prompt:
+    #         full_history = [
+    #             SystemMessage(content=system_prompt),
+    #             SystemMessage(
+    #                 content="⚠️ Do NOT summarize or compress the context from the query tool. It will be passed directly to another tool that formats the answer."
+    #             ),
+    #             SystemMessage(
+    #                 content="You MUST call the `format_scripture_answer` tool if the user question is about scripture content and the query tool has returned a result."
+    #             ),
+    #         ] + new_messages
+    #     else:
+    #         full_history = messages
+
+    #     # 🔍 Debug log (optional)
+    #     # print("\n🧠 LLM Full History:")
+    #     # for m in full_history:
+    #     #    print(f"- {m.type.upper()}: {m.content[:100]}...\n")
+
+    #     ai_response = llm.invoke(full_history)
+    #     return {"messages": messages + [ai_response]}
+
+    graph = StateGraph(ChatState)
+    graph.add_node("init", init_system_prompt_node)
+    graph.add_node("chat", chatNode)
+    graph.add_node("tools", ToolNode(tools))
+    graph.add_edge(START, "init")
+    graph.add_edge("init", "chat")
+    graph.add_conditional_edges("chat", tools_condition, "tools")
+    graph.add_edge("tools", "chat")
+
+    return graph.compile(checkpointer=memory)

push_notifications_helper.py

@@ -1,24 +1,24 @@
-import os
-import requests
-from dotenv import load_dotenv
-import logging
-
-logger = logging.getLogger()
-logger.setLevel(logging.INFO)
-
-load_dotenv(override=True)
-
-pushover_token = os.getenv("PUSHOVER_TOKEN")
-pushover_user = os.getenv("PUSHOVER_USER")
-pushover_url = "https://api.pushover.net/1/messages.json"
-
-
-def push(text: str):
-    """Send a push notification to the user"""
-    logger.info("Sending a push notification for %s", text)
-    response = requests.post(
-        pushover_url,
-        data={"token": pushover_token, "user": pushover_user, "message": text},
-    )
-    logger.info("response = %s", response.json())
-    logger.info("Sent notification")
+import os
+import requests
+from dotenv import load_dotenv
+import logging
+
+logger = logging.getLogger()
+logger.setLevel(logging.INFO)
+
+load_dotenv(override=True)
+
+pushover_token = os.getenv("PUSHOVER_TOKEN")
+pushover_user = os.getenv("PUSHOVER_USER")
+pushover_url = "https://api.pushover.net/1/messages.json"
+
+
+def push(text: str):
+    """Send a push notification to the user"""
+    logger.info("Sending a push notification for %s", text)
+    response = requests.post(
+        pushover_url,
+        data={"token": pushover_token, "user": pushover_user, "message": text},
+    )
+    logger.info("response = %s", response.json())
+    logger.info("Sent notification")

sanatan_assistant.py

@@ -1,98 +1,98 @@
-import logging
-from typing import Literal
-from dotenv import load_dotenv
-from config import SanatanConfig
-from db import SanatanDatabase
-
-load_dotenv(override=True)
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-sanatanDatabase = SanatanDatabase()
-allowedCollections = Literal[
-    *[scripture["collection_name"] for scripture in SanatanConfig.scriptures]
-]
-
-
-def format_scripture_answer(scripture_title: allowedCollections, question: str, query_tool_output: str):
-    """
-    Use this tool to generate a custom system prompt based on the scripture title, question, and query_tool_output.
-
-    This is especially useful when the user has asked a question about a scripture, and the relevant context has been fetched using the `query` tool.
-
-    The generated prompt will guide the assistant to respond using only that scripture’s content, with a clear format including Sanskrit verses, English explanations, and source chapters.
-    """
-
-    prompt = f"""You are a knowledgeable assistant on the scripture *{scripture_title}*, well-versed in both **Sanskrit** and **English**.
-
-You must answer the question using **only** the content from *{scripture_title}* provided in the context below.  
-- Do **not** bring in information from **any other scripture or source**, or from prior knowledge, even if the answer seems obvious or well-known.  
-- Do **not** quote any Sanskrit verses unless they appear **explicitly** in the provided context.  
-- Do **not** use verse numbers or line references unless clearly mentioned in the context.  
-- If the answer cannot be found in the context, clearly say:  
-  **"I do not have enough information from the {scripture_title} to answer this."**
-
-If the answer is not directly stated in the verses but is present in explanatory notes within the context, you may interpret — but **explicitly mention that it is an interpretation**.
-
-If the user query is not small talk, use the following response format (in Markdown):
-
-### 🧾 Answer  
-- Present the explanation in clear, concise **English**.  
-- If it is an interpretation, say so explicitly.
-
-### 🔍 Scripture  
-- Mention the scripture from which the references were taken.
-
-### 🔍 Chapter  
-- Mention the chapter(s) from which the references were taken.  
-- Only mention chapters if they are explicitly part of the context.
-
-### 📜 Supporting Sanskrit Verse(s)  
-- Quote **only the most relevant** Sanskrit verse(s) from the context.  
-- Do **not** add verses from memory or outside the context.  
-- Format each verse clearly, one per line.  
-- **Avoid transliteration unless no Devanagari is available.**  
-- Do **not** provide English text in this section.
-
-### 🔍 English Translation  
-- Provide the **corresponding English meaning** for each Sanskrit verse shown.  
-- Keep it readable and aligned with the verse above.
-
-### Notes  
-- Bullet any extra points or cross-references from explanatory notes **only if present in the context**.  
-- Do **not** include anything that is not supported or implied in the context.
-
-**Question:**  
-{question}
-
----
-
-**Context:**  
-{query_tool_output}
-
----
-
-Respond in **Markdown** format only. Ensure Sanskrit verses are always clearly shown and translated. If a section does not apply (e.g. no verses), you may omit it.
-"""
-
-    return prompt
-
-
-
-def query(collection_name: allowedCollections, query: str, n_results=5):
-    """
-    Search a scripture collection.
-
-    Parameters:
-    - collection_name (str): The name of the scripture collection to search. ...
-    - query (str): The search query.
-    - n_results (int): Number of results to return. Default is 5.
-
-    Returns:
-    - A list of matching results.
-    """
-    logger.info("Searching collection [%s] for [%s]", collection_name, query)
-    response = sanatanDatabase.search(
-        collection_name=collection_name, query=query, n_results=n_results
-    )
-    return "\n\n".join(response["documents"][0])
+import logging
+from typing import Literal
+from dotenv import load_dotenv
+from config import SanatanConfig
+from db import SanatanDatabase
+
+load_dotenv(override=True)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+sanatanDatabase = SanatanDatabase()
+allowedCollections = Literal[
+    *[scripture["collection_name"] for scripture in SanatanConfig.scriptures]
+]
+
+
+def format_scripture_answer(scripture_title: allowedCollections, question: str, query_tool_output: str):
+    """
+    Use this tool to generate a custom system prompt based on the scripture title, question, and query_tool_output.
+
+    This is especially useful when the user has asked a question about a scripture, and the relevant context has been fetched using the `query` tool.
+
+    The generated prompt will guide the assistant to respond using only that scripture’s content, with a clear format including Sanskrit verses, English explanations, and source chapters.
+    """
+
+    prompt = f"""You are a knowledgeable assistant on the scripture *{scripture_title}*, well-versed in both **Sanskrit** and **English**.
+
+You must answer the question using **only** the content from *{scripture_title}* provided in the context below.  
+- Do **not** bring in information from **any other scripture or source**, or from prior knowledge, even if the answer seems obvious or well-known.  
+- Do **not** quote any Sanskrit verses unless they appear **explicitly** in the provided context.  
+- Do **not** use verse numbers or line references unless clearly mentioned in the context.  
+- If the answer cannot be found in the context, clearly say:  
+  **"I do not have enough information from the {scripture_title} to answer this."**
+
+If the answer is not directly stated in the verses but is present in explanatory notes within the context, you may interpret — but **explicitly mention that it is an interpretation**.
+
+If the user query is not small talk, use the following response format (in Markdown):
+
+### 🧾 Answer  
+- Present the explanation in clear, concise **English**.  
+- If it is an interpretation, say so explicitly.
+
+### 🔍 Scripture  
+- Mention the scripture from which the references were taken.
+
+### 🔍 Chapter  
+- Mention the chapter(s) from which the references were taken.  
+- Only mention chapters if they are explicitly part of the context.
+
+### 📜 Supporting Sanskrit Verse(s)  
+- Quote **only the most relevant** Sanskrit verse(s) from the context.  
+- Do **not** add verses from memory or outside the context.  
+- Format each verse clearly, one per line.  
+- **Avoid transliteration unless no Devanagari is available.**  
+- Do **not** provide English text in this section.
+
+### 🔍 English Translation  
+- Provide the **corresponding English meaning** for each Sanskrit verse shown.  
+- Keep it readable and aligned with the verse above.
+
+### Notes  
+- Bullet any extra points or cross-references from explanatory notes **only if present in the context**.  
+- Do **not** include anything that is not supported or implied in the context.
+
+**Question:**  
+{question}
+
+---
+
+**Context:**  
+{query_tool_output}
+
+---
+
+Respond in **Markdown** format only. Ensure Sanskrit verses are always clearly shown and translated. If a section does not apply (e.g. no verses), you may omit it.
+"""
+
+    return prompt
+
+
+
+def query(collection_name: allowedCollections, query: str, n_results=5):
+    """
+    Search a scripture collection.
+
+    Parameters:
+    - collection_name (str): The name of the scripture collection to search. ...
+    - query (str): The search query.
+    - n_results (int): Number of results to return. Default is 5.
+
+    Returns:
+    - A list of matching results.
+    """
+    logger.info("Searching collection [%s] for [%s]", collection_name, query)
+    response = sanatanDatabase.search(
+        collection_name=collection_name, query=query, n_results=n_results
+    )
+    return "\n\n".join(response["documents"][0])

serperdev_helper.py

@@ -1,14 +1,14 @@
-from langchain_community.utilities import GoogleSerperAPIWrapper
-from dotenv import load_dotenv
-import logging
-
-load_dotenv(override=True)
-logger = logging.getLogger(__name__)
-logger.setLevel(logging.INFO)
-
-serper = GoogleSerperAPIWrapper()
-
-def search(query: str):
-    """Search the web for any given query"""
-    logger.info("Searching the web for %s", query)
+from langchain_community.utilities import GoogleSerperAPIWrapper
+from dotenv import load_dotenv
+import logging
+
+load_dotenv(override=True)
+logger = logging.getLogger(__name__)
+logger.setLevel(logging.INFO)
+
+serper = GoogleSerperAPIWrapper()
+
+def search(query: str):
+    """Search the web for any given query"""
+    logger.info("Searching the web for %s", query)
     return serper.run(query)

tools.py

@@ -1,28 +1,28 @@
-from langchain.agents import Tool
-from langchain_core.tools import StructuredTool
-
-from config import SanatanConfig
-from push_notifications_helper import push
-from serperdev_helper import search as search_web
-from sanatan_assistant import format_scripture_answer, query
-
-tool_push = Tool(
-    name="push", description="Send a push notification to the user", func=push
-)
-
-allowed_collections = [s["collection_name"] for s in SanatanConfig.scriptures]
-
-tool_search_db = StructuredTool.from_function(
-    query,
-    description=(
-        "Search within a specific scripture collection. "
-        f"The collection_name must be one of: {', '.join(allowed_collections)}."
-        " Use this to find relevant scripture verses or explanations."
-    ),
-)
-
-tool_search_web = Tool(
-    name="search_web", description="Search the web for information", func=search_web
-)
-
-tool_format_scripture_answer = StructuredTool.from_function(format_scripture_answer)
+from langchain.agents import Tool
+from langchain_core.tools import StructuredTool
+
+from config import SanatanConfig
+from push_notifications_helper import push
+from serperdev_helper import search as search_web
+from sanatan_assistant import format_scripture_answer, query
+
+tool_push = Tool(
+    name="push", description="Send a push notification to the user", func=push
+)
+
+allowed_collections = [s["collection_name"] for s in SanatanConfig.scriptures]
+
+tool_search_db = StructuredTool.from_function(
+    query,
+    description=(
+        "Search within a specific scripture collection. "
+        f"The collection_name must be one of: {', '.join(allowed_collections)}."
+        " Use this to find relevant scripture verses or explanations."
+    ),
+)
+
+tool_search_web = Tool(
+    name="search_web", description="Search the web for information", func=search_web
+)
+
+tool_format_scripture_answer = StructuredTool.from_function(format_scripture_answer)