sanatan_assistant.py

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])