systex_phase2/ai_agent.py

from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate

from langchain.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

# graph usage
from pprint import pprint
from typing import List
from langchain_core.documents import Document
from typing_extensions import TypedDict
from langgraph.graph import END, StateGraph, START
from langgraph.pregel import RetryPolicy

# supabase db
from langchain_community.utilities import SQLDatabase
import os
from dotenv import load_dotenv
load_dotenv()
URI: str =  os.environ.get('SUPABASE_URI')
db = SQLDatabase.from_uri(URI)

# LLM
# local_llm = "llama3.1:8b-instruct-fp16"
local_llm = "llama3-groq-tool-use:latest"
llm_json = ChatOllama(model=local_llm, format="json", temperature=0)
llm = ChatOllama(model=local_llm, temperature=0)

# RAG usage
from faiss_index import create_faiss_retriever, faiss_query
retriever = create_faiss_retriever()

# text-to-sql usage
from text_to_sql2 import run, get_query, query_to_nl, table_description


def faiss_query(question: str, docs, llm, multi_query: bool = False) -> str:
    
    context = docs
    
    system_prompt: str = "你是一個來自台灣的AI助理，樂於以台灣人的立場幫助使用者，會用繁體中文回答問題。"
    template = """
    <|begin_of_text|>
    
    <|start_header_id|>system<|end_header_id|>
    你是一個來自台灣的ESG的AI助理，請用繁體中文回答問題 \n
    You should not mention anything about "根據提供的文件內容" or other similar terms.
    Use five sentences maximum and keep the answer concise.
    如果你不知道答案請回答："很抱歉，目前我無法回答您的問題，請將您的詢問發送至 test@systex.com 以便獲得更進一步的幫助，謝謝。"
    勿回答無關資訊
    <|eot_id|>
    
    <|start_header_id|>user<|end_header_id|>
    Answer the following question based on this context:

    {context}

    Question: {question}
    用繁體中文
    <|eot_id|>
    
    <|start_header_id|>assistant<|end_header_id|>
    """
    prompt = ChatPromptTemplate.from_template(
        system_prompt + "\n\n" +
        template
    )

    rag_chain = prompt | llm | StrOutputParser()
    return rag_chain.invoke({"context": context, "question": question})

### Hallucination Grader

def Hallucination_Grader():
    # Prompt
    prompt = PromptTemplate(
        template=""" <|begin_of_text|><|start_header_id|>system<|end_header_id|> 
        You are a grader assessing whether an answer is grounded in / supported by a set of facts. 
        Give 'yes' or 'no' score to indicate whether the answer is grounded in / supported by a set of facts. 
        Provide 'yes' or 'no' score as a JSON with a single key 'score' and no preamble or explanation. 
        Return the a JSON with a single key 'score' and no premable or explanation. 
        <|eot_id|><|start_header_id|>user<|end_header_id|>
        Here are the facts:
        \n ------- \n
        {documents} 
        \n ------- \n
        Here is the answer: {generation} 
        Provide 'yes' or 'no' score as a JSON with a single key 'score' and no premable or explanation.
        <|eot_id|><|start_header_id|>assistant<|end_header_id|>""",
        input_variables=["generation", "documents"],
    )

    hallucination_grader = prompt | llm_json | JsonOutputParser()
    
    return hallucination_grader

### Answer Grader

def Answer_Grader():
    # Prompt
    prompt = PromptTemplate(
        template="""
        <|begin_of_text|><|start_header_id|>system<|end_header_id|> You are a grader assessing whether an 
        answer is useful to resolve a question. Give a binary score 'yes' or 'no' to indicate whether the answer is 
        useful to resolve a question. Provide the binary score as a JSON with a single key 'score' and no preamble or explanation.
        <|eot_id|><|start_header_id|>user<|end_header_id|> Here is the answer:
        \n ------- \n
        {generation} 
        \n ------- \n
        Here is the question: {question} 
        <|eot_id|><|start_header_id|>assistant<|end_header_id|>""",
        input_variables=["generation", "question"],
    )

    answer_grader = prompt | llm_json | JsonOutputParser()
    
    return answer_grader

# Text-to-SQL
def run_text_to_sql(question: str):
    selected_table = ['104_112碳排放公開及建準資料', '水電使用量(GHG)', '水電使用量(ISO)']
    # question = "建準去年的固定燃燒總排放量是多少?"
    query, result, answer = run(db, question, selected_table, llm)
    
    return  answer, query

def _get_query(question: str):
    selected_table = ['104_112碳排放公開及建準資料', '水電使用量(GHG)', '水電使用量(ISO)']
    query = get_query(db, question, selected_table, llm)
    return  query

def _query_to_nl(question: str, query: str):
    answer = query_to_nl(db, question, query, llm)
    return  answer


### SQL Grader

def SQL_Grader():
    prompt = PromptTemplate(
        template="""<|begin_of_text|><|start_header_id|>system<|end_header_id|> 
        You are a SQL query grader assessing correctness of PostgreSQL query to a user question. 
        Based on following database description, you need to grade whether the PostgreSQL query exactly matches the user question.
        
        Here is database description:
        {table_info}
        
        You need to check that each where statement is correctly filtered out what user question need.
        
        For example, if user question is "建準去年的固定燃燒總排放量是多少?", and the PostgreSQL query is 
        "SELECT SUM("排放量(公噸CO2e)") AS "下游租賃總排放量"
        FROM "104_112碳排放公開及建準資料"
        WHERE "事業名稱" like '%建準%'
        AND "排放源" = '下游租賃'
        AND "盤查標準" = 'GHG'
        AND "年度" = EXTRACT(YEAR FROM CURRENT_DATE)-1;"
        For the above example, we can find that user asked for "固定燃燒", but the PostgreSQL query gives "排放源" = '下游租賃' in WHERE statement, which means the PostgreSQL query is incorrect for the user question.
        
        Another example like "建準去年的固定燃燒總排放量是多少?", and the PostgreSQL query is 
        "SELECT SUM("排放量(公噸CO2e)") AS "固定燃燒總排放量"
        FROM "104_112碳排放公開及建準資料"
        WHERE "事業名稱" like '%台積電%'
        AND "排放源" = '固定燃燒'
        AND "盤查標準" = 'GHG'
        AND "年度" = EXTRACT(YEAR FROM CURRENT_DATE)-1;"
        For the above example, we can find that user asked for "建準", but the PostgreSQL query gives "事業名稱" like '%台積電%' in WHERE statement, which means the PostgreSQL query is incorrect for the user question.
        
        and so on. You need to strictly examine whether the sql PostgreSQL query matches the user question.
        
        If the PostgreSQL query do not exactly matches the user question, grade it as incorrect. 
        You need to strictly examine whether the sql PostgreSQL query matches the user question.
        Give a binary score 'yes' or 'no' score to indicate whether the PostgreSQL query is correct to the question. \n
        Provide the binary score as a JSON with a single key 'score' and no premable or explanation.
        <|eot_id|>
        
        <|start_header_id|>user<|end_header_id|>
        Here is the PostgreSQL query: \n\n {sql_query} \n\n
        Here is the user question: {question} \n <|eot_id|><|start_header_id|>assistant<|end_header_id|>
        """,
        input_variables=["table_info", "question", "sql_query"],
    )

    sql_query_grader = prompt | llm_json | JsonOutputParser()
    
    return sql_query_grader

### Router
def Router():
    prompt = PromptTemplate(
        template="""<|begin_of_text|><|start_header_id|>system<|end_header_id|> 
        You are an expert at routing a user question to a vectorstore or company private data. 
        Use company private data for questions about the informations about a company's greenhouse gas emissions data.
        Otherwise, use the vectorstore for questions on ESG field knowledge or news about ESG. 
        You do not need to be stringent with the keywords in the question related to these topics. 
        Give a binary choice 'company_private_data' or 'vectorstore' based on the question. 
        Return the a JSON with a single key 'datasource' and no premable or explanation. 
        
        Question to route: {question} 
        <|eot_id|><|start_header_id|>assistant<|end_header_id|>""",
        input_variables=["question"],
    )

    question_router = prompt | llm_json | JsonOutputParser()
    
    return question_router

class GraphState(TypedDict):
    """
    Represents the state of our graph.

    Attributes:
        question: question
        generation: LLM generation
        company_private_data: whether to search company private data
        documents: list of documents
    """

    question: str
    generation: str
    documents: List[str]
    retry: int
    sql_query: str
    
# Node
def retrieve_and_generation(state):
    """
    Retrieve documents from vectorstore

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, documents, that contains retrieved documents, and generation, genrating by LLM
    """
    print("---RETRIEVE---")
    question = state["question"]

    # Retrieval
    # documents = retriever.invoke(question)
    # TODO: correct Retrieval function
    documents = retriever.get_relevant_documents(question, k=30)
    # docs_documents = "\n\n".join(doc.page_content for doc in documents)
    # print(documents)
    generation = faiss_query(question, documents, llm)
    return {"documents": documents, "question": question, "generation": generation}

def company_private_data_get_sql_query(state):
    """
    Get PostgreSQL query according to question

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): return generated PostgreSQL query and record retry times
    """
    print("---SQL QUERY---")
    question = state["question"]
    
    if state["retry"]:
        retry = state["retry"]
        retry += 1
    else: 
        retry = 0
    # print("RETRY: ", retry)
    
    sql_query = _get_query(question)
    
    return {"sql_query": sql_query, "question": question, "retry": retry}
    
def company_private_data_search(state):
    """
    Execute PostgreSQL query and convert to nature language.

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): Appended sql results to state
    """

    print("---SQL TO NL---")
    # print(state)
    question = state["question"]
    sql_query = state["sql_query"]
    generation = _query_to_nl(question, sql_query)
    
    # generation = [company_private_data_result]
    
    return {"sql_query": sql_query, "question": question, "generation": generation}

### Conditional edge


def route_question(state):
    """
    Route question to web search or RAG.

    Args:
        state (dict): The current graph state

    Returns:
        str: Next node to call
    """

    print("---ROUTE QUESTION---")
    question = state["question"]
    # print(question)
    question_router = Router()
    source = question_router.invoke({"question": question})
    # print(source)
    print(source["datasource"])
    if source["datasource"] == "company_private_data":
        print("---ROUTE QUESTION TO TEXT-TO-SQL---")
        return "company_private_data"
    elif source["datasource"] == "vectorstore":
        print("---ROUTE QUESTION TO RAG---")
        return "vectorstore"
    
def grade_generation_v_documents_and_question(state):
    """
    Determines whether the generation is grounded in the document and answers question.

    Args:
        state (dict): The current graph state

    Returns:
        str: Decision for next node to call
    """

    print("---CHECK HALLUCINATIONS---")
    question = state["question"]
    documents = state["documents"]
    generation = state["generation"]

    
    # print(docs_documents)
    # print(generation)
    hallucination_grader = Hallucination_Grader()
    score = hallucination_grader.invoke(
        {"documents": documents, "generation": generation}
    )
    # print(score)
    grade = score["score"]

    # Check hallucination
    if grade in ["yes", "true", 1, "1"]:
        print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
        # Check question-answering
        print("---GRADE GENERATION vs QUESTION---")
        answer_grader = Answer_Grader()
        score = answer_grader.invoke({"question": question, "generation": generation})
        grade = score["score"]
        if grade in ["yes", "true", 1, "1"]:
            print("---DECISION: GENERATION ADDRESSES QUESTION---")
            return "useful"
        else:
            print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")
            return "not useful"
    else:
        pprint("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")
        return "not supported"
    
def grade_sql_query(state):
    """
    Determines whether the Postgresql query are correct to the question

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): Decision for retry or continue
    """

    print("---CHECK SQL CORRECTNESS TO QUESTION---")
    question = state["question"]
    sql_query = state["sql_query"]
    retry = state["retry"]

    # Score each doc
    sql_query_grader = SQL_Grader()
    score = sql_query_grader.invoke({"table_info": table_description(), "question": question, "sql_query": sql_query})
    grade = score["score"]
    # Document relevant
    if grade in ["yes", "true", 1, "1"]:
        print("---GRADE: CORRECT SQL QUERY---")
        return "correct"
    elif retry >= 5:
        print("---GRADE: INCORRECT SQL QUERY AND REACH RETRY LIMIT---")
        return "failed"
    else:
        print("---GRADE: INCORRECT SQL QUERY---")
        return "incorrect"

def build_graph():
    workflow = StateGraph(GraphState)

    # Define the nodes
    workflow.add_node("company_private_data_query", company_private_data_get_sql_query, retry=RetryPolicy(max_attempts=5))  # web search
    workflow.add_node("company_private_data_search", company_private_data_search, retry=RetryPolicy(max_attempts=5))  # web search
    workflow.add_node("retrieve_and_generation", retrieve_and_generation, retry=RetryPolicy(max_attempts=5))  # retrieve
    
    workflow.add_conditional_edges(
        START,
        route_question,
        {
            "company_private_data": "company_private_data_query",
            "vectorstore": "retrieve_and_generation",
        },
    )

    workflow.add_conditional_edges(
        "retrieve_and_generation",
        grade_generation_v_documents_and_question,
        {
            "not supported": "retrieve_and_generation",
            "useful": END,
            "not useful": "retrieve_and_generation",
        },
    )
    workflow.add_conditional_edges(
        "company_private_data_query",
        grade_sql_query,
        {
            "correct": "company_private_data_search",
            "incorrect": "company_private_data_query",
            "failed": END
            
        },
    )
    workflow.add_edge("company_private_data_search", END)

    app = workflow.compile()    
    
    return app

def main():
    app = build_graph()
    #建準去年的類別一排放量?
    inputs = {"question": "溫室氣體是什麼"}
    for output in app.stream(inputs, {"recursion_limit": 10}):
        for key, value in output.items():
            pprint(f"Finished running: {key}:")
    pprint(value["generation"])
    
    return value["generation"]

if __name__ == "__main__":
    result = main()
    print("------------------------------------------------------")
    print(result)