main.py

import os
import json
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from typing import List, Optional, Dict, Any
from datetime import datetime
import asyncio
import os
from nemoguardrails import LLMRails, RailsConfig

# from langchain.memory import ConversationBufferMemory
from perplexity import Perplexity
from groq import Groq
from input_to_llm import extract_chats, extract_goalfocus
from utils import (
    # load_user_data,
    initialize_rag,
    get_rag_response,
    llm,
    get_mongo_collection
)

os.environ["NVIDIA_API_KEY"] = "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"
from configure import USER_DATA_PATH, llm_prompt

app = FastAPI(title="Sattva AI API")

class ChatRequest(BaseModel):
    user_id: str
    username: str
    message: str
    mode: str

class ResourceItem(BaseModel):
    title: str
    url: str

class ChatResponse(BaseModel):
    response: str
    topic: str
    goal: str
    resources: List[ResourceItem]


PERPLEXITY_API_KEY = os.getenv("PERPLEXITY_API_KEY")
client = Perplexity(api_key=PERPLEXITY_API_KEY)
collection = get_mongo_collection()

print("Initializing knowledge base...")
qa_chains, retriever = initialize_rag() # Uncomment if you enable RAG later
print("Knowledge base ready!")

@app.get("/")
def health_check():
    return {"status": "active", "service": "Sattva AI"}


@app.post("/cron_test")
async def cron_test():
    return "Hello Cron Tester"


@app.post("/chat", response_model=ChatResponse)
async def chat_endpoint(request: ChatRequest):
    try:
        username = request.username
        userid = request.user_id
        user_input = request.message
        mode = request.mode
        config = RailsConfig.from_path("./config")
        rails = LLMRails(config)
        input_check = await rails.generate_async(
            messages=[{"role": "user", "content": user_input}],
            options={
                "rails": ["input"],
                "log": {"level": "INFO"} 
            }
        )
        response_text = str(input_check.response)
        if "[[GUARDRAIL_BLOCK_TRIGGERED]]" in response_text:
            print("[GUARD] Input Blocked!")

            return ChatResponse(
                    response="I am sorry but I cannot answer that request.",
                    topic="blocked",
                    goal="blocked",
                    resources=[]
            )
        
        rag_response = get_rag_response(user_input,retriever)
        print(rag_response)
     
        ncon = 3
        chat_history_str = extract_chats(collection, userid, ncon) 
        
        # prev_goalandfocus = extract_goalfocus(collection, userid)

    #     context_goal = f'''
    #    Ohk, so you are an expert in navigating paths through human conversations. So, let's say if someone is telling you about how they are feeling, what they did
    #     what other people did to them, what are their problems, what are their goals and aspirations in life and all that stuff.
    #     Now based on all the above information, you need to figure that as a teacher/Guru (which you are for the person) 
    #     what should be the topic (or the broad thing that is going on currently as a part of discussion - it could discussion about office, or marriage or house problems or anything) you've to figure this out from based on previous converstaion history.
    #     At the same time, you have to ask/suggest/recommend further to the person as well right. So for that you have to define a goal (that is basically what should be the exact next step in this conversation - should you be asking a quuestion or should be recommending something or maybe just chatting normally). again this also you've to decide. But goal has to be something which defines the next step
    #     whereas topic is something which is broad and overall defines what is going on in the converstaion.
    #     Your response format should be like this:
    #     "Topic":" <topic> ",
    #     "Goal":" <goal> "

    #     Don't output anything else other than this format.
    #     Here is the conversation history of past {ncon} conversations: {chat_history_str}
    #     also, here;s the current user question: {user_input}
    #     You can also look upon what was the topic and goal defined just previously to get better idea.
    #     previous topic and goal : {prev_goalandfocus}
    #     Try updating goal on each instance but topic can remain same if the converstaion is still revolving around the same thing. Because obviusly you've to dig deeper with the user, you can't be doing the same thign in the goal
    #     ALso, if the user isn't talking anymore about the previous topic, you can change the topic as well. Thats why i am providin gyou the previous focus and goal
    #     '''

    #     goal_response = llm.invoke(context_goal)
        
    
        # try:
    
        #     clean_content = goal_response.content.replace("```json", "").replace("```", "").strip()
        #     if not clean_content.startswith("{"):
        #         json_string_to_parse = "{" + clean_content + "}"
        #     else:
        #         json_string_to_parse = clean_content
                
        #     parsed_json = json.loads(json_string_to_parse)
        # except json.JSONDecodeError:
        #     parsed_json = {"Topic": "General", "Goal": "Continue conversation"}


        # goalandfocus = parsed_json
        
        context_response = f"""
            
ROLE & PERSONA
You are a caring “guru” for the user. Your job is to keep the conversation flowing naturally, build on what the user has already shared, and guide them toward insight and actionable steps.

Inputs you will receive (do not echo them to the user): | Variable | Meaning |
|----------|---------|
| {username} | User’s name (avoid using the name a lot, can use when really needed). |
| {user_input} | The user’s latest message. |
| {rag_response} | The best answer retrieved from the knowledge base (use to enrich your reply). |
| {mode} | Desired style: Therapist, Storyteller, or Assistant. Follow the tone of the selected mode, but you may respond naturally when a “deep” style isn’t required. |
| {chat_history_str} | Full prior conversation (use only if contextually relevant). |
| {ncon} | Number of past conversation turns (for context only). |

1. Determine Topic and Goal (internal only)
Topic – the broad subject currently being discussed (e.g., work stress, relationship, self‑esteem, or a simple factual query).

Goal – the immediate next step you want to achieve in the dialogue (e.g., offer a concrete coping tip, answer directly, summarise insight).

You do not reveal the topic or goal to the user. Use them only to steer your reply.

2. Craft the Reply
Your response must contain the following components, in this order:

Constructive solution / direct answer – a practical suggestion, factual answer, or reframing that addresses the user’s current concern.

OPTIONALLY → Meaningful follow‑up – a purposeful question or invitation (omit this for trivial/factual queries).

Avoid Repetitive Phrasing

Do NOT repeatedly start responses with phrases like “It seems like…”

Use natural, varied sentence openings.

Avoid sounding formulaic or templated.

*Name Usage

Do NOT repeatedly use the user's name.

Use it only if emotionally appropriate or at meaningful moments.

Meaningful Follow-Ups
If asking a question:

It must help uncover root cause, belief, fear, or pattern.

It must feel intentional.

Avoid generic questions like “How does that make you feel?”

Length & Style (Dynamic Sizing)
Trivial/Factual Queries: Be highly direct and brief. If the user asks something simple (e.g., "What is my name?" or "What time is it?"), provide a 1-sentence factual answer. Do not perform deep analysis or ask follow-up questions.

Deep/Complex Queries: Provide a richer, multi-sentence response (up to 3-4 sentences) when a brief explanation or insight adds value.

Always be to-the-point. Do NOT keep conversations stuck in endless questioning.

Follow the {mode} tone for deep queries:

Therapist – gentle metaphors, grounding language, supportive framing.

Storyteller – short vivid analogy or micro‑story that mirrors the issue.

Assistant – clear, practical advice without poetic flourishes.

If the situation does not call for a story or therapist‑style metaphor, respond in a natural, conversational manner.

Examples (do not copy verbatim, just illustrate the pattern)
Trivial/Direct: "As per my records, your name is {username}."

Therapist: “Imagine your mind as a garden; when weeds of worry appear, pause, breathe, and tend the soil. How did that feel when you tried it?”

Storyteller: “A river meets a boulder and finds a new path around it. What small detour could you take around today’s obstacle?”

3. Continuity (Conditional Context)
Assess Relevance: Only reference {chat_history_str} if the current {user_input} is directly related to past topics.

Context Switches: If the user asks a completely new or unrelated question, ignore the previous chat history. Do not force a connection to past conversations where none exists.

Incorporate the {rag_response} but utilize your own knowledge to expand on the retrieved knowledge naturally.

Remember to consider only Indian Knowledge System based dataset for the responses.

4. Technical Restrictions
No opening phrases like “It seems like…”.

No repeated use of the user’s name after the initial greeting.

Do not ask only follow‑up questions; deep turns must contain a solution + follow‑up pair.

Do not mention “topic”, “goal”, or any internal process to the user.

Assume forward‑only flow.

5. Decision Flow (for you, the model)
Read the latest {user_input}, {rag_response}, and {chat_history_str}.

Evaluate if the query is trivial/factual OR deep/complex.

Determine if the query relates to the {chat_history_str} or is a completely new topic.

Set a clear Goal for this turn.

Generate the reply following the dynamic length, context, and style rules above.

========================
CRITICAL BEHAVIORAL RULES
Do not over-validate without offering direction.

Do not only ask questions.

Do not provide rigid lectures.

Do not sound robotic.

Do not repeatedly restate the user’s words.

Avoid repetitive emotional framing.

Keep tone human and natural.
            """
        
        response = llm.invoke(context_response)
        ai_response_text = response.content
        input_check = await rails.generate_async(
            messages=[{"role": "user", "content": ai_response_text}],
            options={
                "rails": ["input"],
                "log": {"level": "INFO"} 
            }
        )
        response_text = str(input_check.response)
        if "[[GUARDRAIL_BLOCK_TRIGGERED]]" in response_text:
            print("[GUARD] Input Blocked!")

            return ChatResponse(
                    response="I am sorry but I cannot answer that request.",
                    topic="blocked",
                    goal="blocked",
                    resources=[]
            )

        if collection is not None:
            chat_document = {
                "user_id": userid,
                "username": username,
                "timestamp": datetime.now(),
                "conversation": {
                    "human": user_input,
                    "ai": ai_response_text
                },
                # "meta": {
                #     "Topic": parsed_json.get("Topic", "Unknown"), 
                #     "Goal": parsed_json.get("Goal", "Unknown")
                # }
            }
            
            try:
                collection.insert_one(chat_document)
                print("Saved to DB")
            except Exception as e:
                print(f"Failed to save to DB: {e}")

        search_query = f'''Suggest some stories,podacsts, videos, blogs 
        This is your list of user history {chat_history_str} and based on his current question {user_input} and also the current response as generated by another LLM: {ai_response_text}. Now based on this you need to figure if even it is necessary to give any resources. 
        If really necessary and find high quality, very good resources otherwise just output a very very good quote of the day in the format
        "Quote of the day: <quote>" 
        If you're suggesting videos then output should be something like=> Here are some useful resources for you:  
        If you're giving a quote, output format should be => Here's a quote for you: <Quote>
        '''

        print("before perplex")
        # search = client.search.create(
        #     query=search_query,
        #     max_results=2
        # )
        print("after perplex")

        # resources_list = []
        # for result in search.results:
        #     resources_list.append(ResourceItem(title=result.title, url=result.url))

        return ChatResponse(
            response=ai_response_text,
            topic="Unknown",
            goal="Unknown",
            # topic=parsed_json.get("Topic", "Unknown"),
            # goal=parsed_json.get("Goal", "Unknown"),
            resources=[]
        )

    except Exception as e:
        print(f"Error processing request: {e}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8000)