mydomain_agent.py

from langgraph.graph import StateGraph, START, END
# from llm_initializer import initialize_llm, generate_prompt_phi4
from langgraph.graph import MessagesState
from langchain_core.messages import ToolMessage, HumanMessage, SystemMessage
from typing_extensions import Literal, TypedDict
from pydantic import BaseModel, Field
from pydantic import BaseModel, Field, validator
from typing import List, Optional, Dict, Any, TypedDict,Generic, TypeVar
import uuid
import io
import os
import PyPDF2
import re
import logging
import time
from docx import Document as dx
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import (
    DirectoryLoader,
    PyPDFLoader,
    TextLoader
)
import tempfile
import faiss
from langchain_community.vectorstores import FAISS
from langchain_core.prompts import PromptTemplate
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_huggingface import HuggingFaceEmbeddings
from langgraph.checkpoint.memory import MemorySaver
from langgraph.graph import StateGraph, END
from sqlalchemy import create_engine, Column, String, Integer, DateTime, ForeignKey, Text
from sqlalchemy.dialects.sqlite import JSON as SQLiteJSON
# from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, relationship
from sentence_transformers import SentenceTransformer
from huggingface_hub import login
from langchain_google_genai import ChatGoogleGenerativeAI
import datetime
from enum import Enum as PyEnum
from sqlalchemy.orm import DeclarativeBase
# from config import Config
from functools import lru_cache
from dotenv import load_dotenv

load_dotenv()
hf_token = os.getenv("hf_user_token")
login(hf_token)

T = TypeVar("T")
# --- 1. Database Setup ---
DATABASE_URL = "sqlite:///Db_domain_agent.db"
engine = create_engine(DATABASE_URL)
SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)

class Base(DeclarativeBase):
    pass

class FeedbackScore(PyEnum):
    POSITIVE = 1
    NEGATIVE = -1

class Telemetry(Base):
    __tablename__ = "telemetry_table"
    transaction_id = Column(String, primary_key=True)
    session_id = Column(String)
    user_question = Column(Text)
    response = Column(Text)
    context = Column(Text)
    model_name = Column(String)
    input_tokens = Column(Integer)
    output_tokens = Column(Integer)
    total_tokens = Column(Integer)
    latency = Column(Integer)
    dtcreatedon = Column(DateTime)
    
    feedback = relationship("Feedback", back_populates="telemetry_entry", uselist=False)

class Feedback(Base):
    __tablename__ = "feedback_table"
    id = Column(Integer, primary_key=True, autoincrement=True)
    telemetry_entry_id = Column(String, ForeignKey("telemetry_table.transaction_id"), nullable=False, unique=True)
    feedback_score = Column(Integer, nullable=False)
    feedback_text = Column(Text, nullable=True)
    user_query = Column(Text, nullable=False)
    llm_response = Column(Text, nullable=False)
    timestamp = Column(DateTime, default=datetime.datetime.now)
    
    telemetry_entry = relationship("Telemetry", back_populates="feedback")

class ConversationHistory(Base):
    __tablename__ = "conversation_history"
    session_id = Column(String, primary_key=True)
    messages = Column(SQLiteJSON, nullable=False)
    last_updated = Column(DateTime, default=datetime.datetime.now)

Base.metadata.create_all(bind=engine)
# --- 2. Initialize LLM and Embeddings ---
gak = os.getenv("Gapi_key")
llm = ChatGoogleGenerativeAI(model="gemini-2.5-flash-lite",google_api_key=gak)
# embedding_model = SentenceTransformer("ibm-granite/granite-embedding-english-r2")

# my_model_name = "gemma3:1b-it-qat"
# llm = ChatOllama(model=my_model_name)
embedding_model = HuggingFaceEmbeddings(
    model_name="ibm-granite/granite-embedding-english-r2",
    model_kwargs={'device': 'cpu'},
    encode_kwargs={'normalize_embeddings': False}
)

# --- 3. LangGraph State and Workflow ---
class GraphState(TypedDict):
    chat_history: List[Dict[str, Any]]
    retrieved_documents: List[str]
    user_question: str
    decision:str
    session_id: str
    telemetry_id: Optional[str] = None

class Route(BaseModel):
    step: Literal['HR Agent','Finance Agent','Legal Compliance Agent'] = Field(
        None, description="The next step in routing process"
    )
router = llm.with_structured_output(Route)

# class State(TypedDict):
#     input:str
#     decision:str
#     output:str

chathistory = {}

def retrieve_documents(state: GraphState):
    # global vectorstore_retriever
    # upload_documents()
    saved_vectorstore_index = FAISS.load_local('domain_index', embedding_model,allow_dangerous_deserialization=True)
    user_question = state["user_question"]
    # meta_filter = {'Domain':'HR'}
    if saved_vectorstore_index is None:
        raise ValueError("Knowledge base not loaded.")
    retrieved_docs = saved_vectorstore_index.as_retriever(search_type="mmr", search_kwargs={"k": 5})
    top_docs = retrieved_docs.invoke(user_question)
    print("Top Docs: ", top_docs)
    retrieved_docs_content = [doc.page_content if doc.page_content else doc for doc in top_docs]
    print("retrieved_documents List: ", retrieved_docs_content)
    return {"retrieved_documents": retrieved_docs_content}

def generate_response(user_question, retrieved_documents):
    print("Inside generate_response--------------")
    global llm
    global chathistory
    global agent_name
    # user_question = state["user_question"]
    # retrieved_documents = state["retrieved_documents"]
    
    formatted_chat_history = []
    for msg in chathistory["chat_history"]:
        if msg['role'] == 'user':
            formatted_chat_history.append(HumanMessage(content=msg['content']))
        elif msg['role'] == 'assistant':
            formatted_chat_history.append(AIMessage(content=msg['content']))
    
    if not retrieved_documents:
        response_content = "I couldn't find any relevant information in the uploaded documents for your question. Can you please rephrase or provide more context?"
        response_obj = AIMessage(content=response_content)
    else:
        context = "\n\n".join(retrieved_documents)
        template = """
            You are a helpful AI assistant. Answer the user's question based on the provided context {context} and the conversation history {chat_history}.
            If the answer is not in the context, state that you don't have enough information.
            Do not make up answers. Only use the given context and chat_history.
            Remove unwanted words like 'Response:' or 'Answer:' from answers.
            \n\nHere is the Question:\n{user_question}
        """
        rag_prompt = PromptTemplate(
            input_variables=["context", "chat_history", "user_question"],
            template=template
        )
        rag_chain = rag_prompt | llm
        time.sleep(3)
        response_obj = rag_chain.invoke({
            "context": [SystemMessage(content=context)],
            "chat_history": formatted_chat_history,
            "user_question": [HumanMessage(content=user_question)]
        })
    
    telemetry_data = response_obj.model_dump()
    input_tokens = telemetry_data.get('usage_metadata', {}).get('input_tokens', 0)
    output_tokens = telemetry_data.get('usage_metadata', {}).get('output_tokens', 0)
    total_tokens = telemetry_data.get('usage_metadata', {}).get('total_tokens', 0)
    model_name = telemetry_data.get('response_metadata', {}).get('model', 'unknown')
    total_duration = telemetry_data.get('response_metadata', {}).get('total_duration', 0)
    
    db = SessionLocal()
    transaction_id = str(uuid.uuid4())
    try:
        telemetry_record = Telemetry(
            transaction_id=transaction_id,
            session_id=chathistory.get("session_id"),
            user_question=user_question,
            response=response_obj.content,
            context="\n\n".join(retrieved_documents) if retrieved_documents else "No documents retrieved",
            model_name=model_name,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            total_tokens=total_tokens,
            latency=total_duration,
            dtcreatedon=datetime.datetime.now()
        )
        db.add(telemetry_record)
        
        new_messages = chathistory["chat_history"] + [
            {"role": "user", "content": user_question}, 
            {"role": "assistant", "content": response_obj.content, "telemetry_id": transaction_id}
        ]
        
        # --- FIX: Refactored Database Save Logic ---
        print(f"Saving conversation for session_id: {chathistory.get('session_id')}")
        conversation_entry = db.query(ConversationHistory).filter_by(session_id=chathistory.get("session_id")).first()
        if conversation_entry:
            print(f"Updating existing conversation for session_id: {chathistory.get('session_id')}")
            conversation_entry.messages = new_messages
            conversation_entry.last_updated = datetime.datetime.now()
        else:
            print(f"Creating new conversation for session_id: {chathistory.get('session_id')}")
            new_conversation_entry = ConversationHistory(
                session_id=chathistory.get("session_id"),
                messages=new_messages,
                last_updated=datetime.datetime.now()
            )
            db.add(new_conversation_entry)
        
        db.commit()
        print(f"Successfully saved conversation for session_id: {chathistory.get('session_id')}")

    except Exception as e:
        db.rollback()
        print(f"***CRITICAL ERROR***: Failed to save data to database. Error: {e}")
    finally:
        db.close()
    
    return {
        "chat_history": new_messages,
        "telemetry_id": transaction_id,
        "agent_name": agent_name
    }

agent_name = ""
def hr_agent(state:GraphState):
    """Answer the user question based on Human Resource(HR)"""
    global agent_name
    user_question = state["user_question"]
    retrieved_documents = state["retrieved_documents"]
    print("HR Agent")
    agent_name = "HR Agent"
    result = generate_response(user_question,retrieved_documents)
    # return {"output":result}
    return result

def finance_agent(state:GraphState):
    """Answer the user question based on Finance and Bank"""
    global agent_name
    user_question = state["user_question"]
    retrieved_documents = state["retrieved_documents"]
    print("Finance Agent")
    agent_name = "Finance Agent"
    result = generate_response(user_question,retrieved_documents)
    return result

def legals_agent(state:GraphState):
    """Answer the user question based on Legal Compliance"""
    global agent_name
    user_question = state["user_question"]
    retrieved_documents = state["retrieved_documents"]
    print("LC agent")
    agent_name = "Legal Compliance Agent"
    result = generate_response(user_question,retrieved_documents)
    # return {"output":result}
    return result

def llm_call_router(state:GraphState):
    decision = router.invoke(
        [
            SystemMessage(
                content="Route the user_question to HR Agent, Finance Agent, Legal Compliance Agent based on the user's request"
            ),
            HumanMessage(
                content=state['user_question']
            ),
        ]
    )
    return {"decision":decision.step}

def route_decision(state:GraphState):
    
    if state['decision'] == 'HR Agent':
        return "hr_agent"
    elif state['decision'] == 'Finance Agent':
        return "finance_agent"
    elif state['decision'] == 'Legal Compliance Agent':
        return "legals_agent"
    
router_builder = StateGraph(GraphState)

router_builder.add_node("retrieve", retrieve_documents)
router_builder.add_node("hr_agent", hr_agent)
router_builder.add_node("finance_agent", finance_agent)
router_builder.add_node("legals_agent", legals_agent)
router_builder.add_node("llm_call_router", llm_call_router)

# router_builder.add_node("generate", generate_response)
# router_builder.set_entry_point("retrieve")
# router_builder.add_edge("retrieve", "generate")
# router_builder.add_edge("generate", END)
# compiled_app = workflow.compile(checkpointer=memory)


router_builder.add_edge(START, "llm_call_router")
router_builder.add_conditional_edges(
    "llm_call_router",
    route_decision,
    {
        "hr_agent":"hr_agent",
        "finance_agent":"finance_agent",
        "legals_agent":"legals_agent",
    },
)
router_builder.set_entry_point("retrieve")
router_builder.add_edge("retrieve","llm_call_router")
router_builder.add_edge("hr_agent",END)
router_builder.add_edge("finance_agent",END)
router_builder.add_edge("legals_agent",END)

route_workflow = router_builder.compile()

# state = route_workflow.invoke({'input': "Write a poem about a wicked cat"})
# print(state['output'])



vectorstore_retriever = None
compiled_app = None
memory = MemorySaver()

# --- 4. LangGraph Nodes ---
# def load_documents(state:GraphState):
#     global selected_domain







# --- 5. API Models ---
class ChatHistoryEntry(BaseModel):
    role: str
    content: str
    telemetry_id: Optional[str] = None

class ChatRequest(BaseModel):
    user_question: str
    session_id: str
    chat_history: Optional[List[ChatHistoryEntry]] = Field(default_factory=list)

    @validator('user_question')
    def validate_prompt(cls, v):
        v = v.strip()
        if not v:
            raise ValueError('Question cannot be empty')
        return v

class ChatResponse(BaseModel):
    ai_response: str
    updated_chat_history: List[ChatHistoryEntry]
    telemetry_entry_id: str
    is_restricted: bool = False
    moderation_reason: Optional[str] = None

class FeedbackRequest(BaseModel):
    session_id: str
    telemetry_entry_id: str
    feedback_score: int
    feedback_text: Optional[str] = None

class ConversationSummary(BaseModel):
    session_id: str
    title: str


@lru_cache(maxsize=5)
def process_text(file):
    string_data = (file.read()).decode("utf-8")
    return string_data
    
@lru_cache(maxsize=5)
def process_pdf(file):
    pdf_bytes = io.BytesIO(file.read())
    reader = PyPDF2.PdfReader(pdf_bytes)
    pdf_text = "".join([page.extract_text() + "\n" for page in reader.pages])
    return pdf_text
    
@lru_cache(maxsize=5)
def process_docx(file):
    docx_bytes = io.BytesIO(file.read())
    docx_docs = dx(docx_bytes)
    docx_content = "\n".join([para.text for para in docx_docs.paragraphs])
    return docx_content
    

# @app.post("/upload-documents")
# def upload_documents(files):
def upload_documents():
    global vectorstore_retriever
    # saved_vectorstore_index = FAISS.load_local('domain_index', embedding_model,allow_dangerous_deserialization=True)
    try:
        saved_vectorstore_index = faiss.read_index("domain_index_sec.faiss")
        if saved_vectorstore_index:
            vectorstore_retriever = saved_vectorstore_index

            msg = f"Successfully loaded the knowledge base."
            return msg, True
    except Exception as e:
        print("unable to find index...", e)
        print("Creating new index.....")
        all_documents = []
        hr_loader = PyPDFLoader("D:\Pdf_data\Developments_in_HR_management_in_QAAs.pdf").load()
        hr_finance = PyPDFLoader("D:\Pdf_data\White Paper_QA Practice.pdf").load()
        hr_legal = PyPDFLoader("D:\Pdf_data\Legal-Aspects-Compliances.pdf").load()

        for doc in hr_loader:
            doc.metadata['Domain'] = 'HR'
            all_documents.append(doc)
        for doc in hr_finance:
            doc.metadata['Domain'] = 'Finance'
            all_documents.append(doc)
        for doc in hr_legal:
            doc.metadata['Domain'] = 'Legal'
            all_documents.append(doc)
        # for uploaded_file in files:
            # doc_loader = PyPDFLoader(uploaded_file)
            # all_documents.extend(doc_loader.load())
        
        if not all_documents:
            raise Exception(status_code=400, detail="No supported documents uploaded.")

        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
        text_chunks = text_splitter.split_documents(all_documents)
        print("text_chucks: ", text_chunks[:100])

        # processed_chunks_with_ids = []
        # for i, chunk in enumerate(text_chunks):
        #     # Generate a unique ID for each chunk
        #     # Option 1 (Recommended): Using UUID for global uniqueness
        #     # chunk_id = str(uuid.uuid4())
            
        #     # Option 2 (Alternative): Combining source file path with chunk index
        #     # This is good if you want IDs to be deterministic based on file/chunk.
        #     # You might need to make the file path more robust (e.g., hash it or normalize it).
        #     file_source = chunk.metadata.get('source', 'unknown_source')
        #     chunk_id = f"{file_source.replace('.','_')}_chunk_{i}"

        #     # Add the unique ID to the chunk's metadata
        #     # It's good practice to keep original metadata and just add your custom ID.
        #     chunk.metadata['doc_id'] = chunk_id
            
            
        #     processed_chunks_with_ids.append(chunk)
        # embeddings = [embedding_model.encode(doc_chunks.page_content, convert_to_numpy=True) for doc_chunks in processed_chunks_with_ids]
        
        print(f"Split {len(text_chunks)} chunks.")
        print(f"Assigned unique 'doc_id' to each chunk in metadata.")
        # dimension = 768
        # # hnsw_m = 32
        # # index = faiss.IndexHNSWFlat(dimension, hnsw_m, faiss.METRIC_INNER_PRODUCT)
        # index = faiss.IndexFlatL2(dimension)
        # vector_store = FAISS(
        #     embedding_function=embedding_model.embed_query,
        #     index=index,
        #     docstore= InMemoryDocstore(),
        #     index_to_docstore_id={}
        # )
        vectorstore = FAISS.from_documents(documents=text_chunks, embedding=embedding_model)
        # vectorstore.add_documents(text_chunks, ids = [cid.metadata['doc_id'] for cid in text_chunks])
        vectorstore.add_documents(text_chunks)
        # vectorstore_retriever = vectorstore.as_retriever(search_kwargs={'k': 5})
        faiss.write_index(vectorstore.index, "domain_index_sec.faiss")
        # vectorstore.save_local("domain_index")
        vectorstore_retriever = vectorstore
        if vectorstore:
            msg = f"Successfully loaded the knowledge base."
            return msg, True
        else:
            msg = f"Failed to process documents."
            return msg, False

# @app.post("/chat", response_model=ChatResponse)
def chat_with_rag(chatdata):
    global compiled_app
    global vectorstore_retriever
    global chathistory
    if vectorstore_retriever is None:
        raise Exception(status_code=400, detail="Knowledge base not loaded. Please upload documents first.")
    print(f"Received request: {chatdata}")
    # moderation_result = moderator.moderate_content(request.user_question)
    # if moderation_result["is_restricted"]:
    #     # Get appropriate response based on restriction type
    #     response_type = moderation_result.get("response_type", "general")
    #     response_text = Config.RESTRICTED_RESPONSES.get(
    #         response_type, 
    #         Config.RESTRICTED_RESPONSES["general"]
    #     )
        
    #     logger.warning(
    #         f"Restricted query: {request.prompt[:100]}... "
    #         f"Reason: {moderation_result['reason']}"
    #     )
        
    #     return ChatResponse(
    #         ai_response=response_text,
    #         updated_chat_history=[],
    #         telemetry_entry_id=request.session_id,
    #         is_restricted=True,
    #         moderation_reason=moderation_result["reason"],
    #     )
    print("✅ Question passed the RAI check.........")
    print("Received data from UI: ", chatdata)
    chathistory = chatdata
    initial_state = {
        # "chat_history": [msg.model_dump() for msg in chatdata.get('chat_history')],
        "chat_history": [msg for msg in chatdata.get('chat_history')],
        "retrieved_documents": [],
        "user_question": chatdata.get('user_question'),
        "session_id": chatdata.get('session_id')
    }

    try:
        config = {"configurable": {"thread_id": chatdata.get('session_id')}}
        final_state = route_workflow.invoke(initial_state, config=config)

        # chathistory = final_state
        print("chathistory inside chat_with_rag-----------------")
        print("Final State--- : ", final_state)

        ai_response_message = final_state["chat_history"][-1]["content"]
        updated_chat_history_dicts = final_state["chat_history"]
        agent_name = final_state.get("decision","No Agent")

        response_chat = ChatResponse(
            ai_response=ai_response_message,
            updated_chat_history=updated_chat_history_dicts,
            telemetry_entry_id=final_state.get("telemetry_id"),
            is_restricted=False,
        )
        
        return agent_name,response_chat.dict()
    except Exception as e:
        print(f"Internal Server Error: {e}")
        raise Exception(status_code=500, detail=f"An error occurred during chat processing: {e}")


def submit_feedback(feedbackdata):
    db = SessionLocal()
    try:
        telemetry_record = db.query(Telemetry).filter(
            Telemetry.transaction_id == feedbackdata['telemetry_entry_id'],
            Telemetry.session_id == feedbackdata['session_id']
        ).first()

        if not telemetry_record:
            raise Exception(status_code=404, detail="Telemetry entry not found or session ID mismatch.")

        existing_feedback = db.query(Feedback).filter(
            Feedback.telemetry_entry_id == feedbackdata['telemetry_entry_id']
        ).first()

        if existing_feedback:
            existing_feedback.feedback_score = feedbackdata['feedback_score']
            existing_feedback.feedback_text = feedbackdata['feedback_text']
            existing_feedback.timestamp = datetime.datetime.now()
        else:
            feedback_record = Feedback(
                telemetry_entry_id=feedbackdata['telemetry_entry_id'],
                feedback_score=feedbackdata['feedback_score'],
                feedback_text=feedbackdata['feedback_text'],
                user_query=telemetry_record.user_question,
                llm_response=telemetry_record.response,
                timestamp=datetime.datetime.now()
            )
            db.add(feedback_record)
            
        db.commit()

        return {"message": "Feedback submitted successfully."}

    except Exception as e:
        raise e
    except Exception as e:
        db.rollback()
        raise Exception(status_code=500, detail=f"An error occurred: {str(e)}")
    finally:
        db.close()

# @app.get("/conversations", response_model=List[ConversationSummary])
def get_conversations():
    db = SessionLocal()
    try:
        conversations = db.query(ConversationHistory).order_by(ConversationHistory.last_updated.desc()).all()
        summaries = []
        for conv in conversations:
            for msg in conv.messages:
                print(msg)
            first_user_message = next((msg for msg in conv.messages if msg["role"] == "user"), None)
            title = first_user_message.get("content") if first_user_message else "New Conversation"
            summaries.append({"session_id":conv.session_id, "title":title[:30] + "..." if len(title) > 30 else title})
        return summaries
    finally:
        db.close()

# @app.get("/conversations/{session_id}", response_model=List[ChatHistoryEntry])
def get_conversation_history(session_id: str):
    db = SessionLocal()
    try:
        conversation = db.query(ConversationHistory).filter(ConversationHistory.session_id == session_id).first()
        if not conversation:
            raise Exception(status_code=404, detail="Conversation not found.")
        return conversation.messages
    finally:
        db.close()




# if 'selected_model' not in st.session_state:
#     st.session_state.selected_model = ""
# @st.dialog("Choose a domain")
# def domain_modal():
#     domain = st.selectbox("Select a domain",["HR","Finance","Legal"])
#     st.session_state.selected_model = domain
#     if st.button("submit"):
#         st.rerun()

# domain_modal()
# print("Selected Domain: ",st.session_state['selected_model'])

# llm = initialize_llm()