wrappers.py

import os
import openai
from langchain_openai import AzureChatOpenAI
from langchain.prompts import ChatPromptTemplate, PromptTemplate
from qdrant_client.http import models as rest
from qdrant_client import QdrantClient, models
from langchain_core.pydantic_v1 import BaseModel, Field
from langchain.memory import ConversationBufferMemory
from langchain.chains import LLMChain
from langchain_community.utilities import GoogleSearchAPIWrapper
from langchain_core.tools import Tool
from typing_extensions import TypedDict
from typing import List
from langchain.schema import Document
from pprint import pprint
from langgraph.graph import END, StateGraph

from db_operations import DatabaseOperations
# from initialize_db import QdrantClientInitializer
from embedding_loader import *

class Wrappers:

    def __init__(self, collection_name_manual, client, embeddings, LLM, db, CAR_ID, memory):
        self.collection_name_manual = collection_name_manual
        self.embeddings = embeddings
        self.client = client
        self.myLLM = LLM
        self.db = db
        self.CAR_ID = CAR_ID
        self.memory =  memory
        # self.memory =  ConversationBufferMemory(memory_key="history",
        #                           input_key="question")
    def translater(self):
        template = """You are a Turkish-English translator. Translate the input to English considering vehicle/car domain terms.

        Input: {question}

        Answer: """

        PROMPT = PromptTemplate.from_template(template)  # ----------------->  ChatPromptTemplate kullanmak daha sağlıklı

        translate_chain = LLMChain(llm=self.myLLM, prompt=PROMPT)
        # translate_chain = PROMPT | self.myLLM

        return translate_chain
    
    def retriever(self):

        retriever = self.db.as_retriever(search_kwargs={'k': 4}, filter=rest.Filter(
            must=[
                models.FieldCondition(key="car_id", match=models.MatchValue(value=self.CAR_ID))
                ]
              ))
        
        return retriever
    
    def grade_documents(self):
        class GradeDocuments(BaseModel):
            """Binary score for relevance check on retrieved documents."""

            binary_score: str = Field(description="Documents are relevant to the question, 'yes' or 'no'")

        # LLM with function call 
        structured_llm_grader_docs = self.myLLM.with_structured_output(GradeDocuments)

        # Prompt 
        system = """You are a grader assessing relevance of a retrieved document to a user question. \n
            Consider the following when making your assessment: \n
            - Does the document directly address the user's question? \n
            - Does it provide information or context that is pertinent to the question? \n
            - Does it discuss relevant risks, benefits, recommendations, or considerations related to the question? \n
            If the document contains keyword(s) or semantic meaning related or partially related to the question, grade it as relevant. \n
            Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""

        grade_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", system),
                ("human", "Retrieved document: \n\n {document} \n\n User question: {question}"),
            ]
        )

        retrieval_grader_relevance = grade_prompt | structured_llm_grader_docs

        return retrieval_grader_relevance
    
    def lead_check(self):
        
        class LeadCheck(BaseModel):
            """Binary score for service relevance check on question."""

            binary_score: str = Field(description="Services are relevant to the question, 'yes' or 'no'")

        # LLM with function call 
        structured_llm_grader_service = self.myLLM.with_structured_output(LeadCheck)

        # Prompt 
        system = """You are a grader assessing relevance of services to a user question. \n
            If the provided services related or partially related to the question, grade it as relevant. \n
            Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""

        lead_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", system),
                ("human", "Provided services: \n\n {hizmet_listesi} \n\n User question: {question}"),
            ]
        )

        service_grader_relevance = lead_prompt | structured_llm_grader_service
        
        return service_grader_relevance
    
    def main_prompt(self):

        prompt = ChatPromptTemplate.from_template(
            """
            You are an expert assistant named ARVI focused solely on car troubles and vehicle information. 
            Your goal is to provide accurate, helpful, and clear answers to any questions related to car issues, maintenance, repairs, specifications, and other vehicle-related topics.

            You are also designed to respond politely and appropriately to basic courteous interactions. Here are the guidelines:

            1. Car Troubles and Vehicle Information:

                - Always answer questions regarding car issues, diagnostics, repairs, maintenance, and vehicle specifications.
                - Provide detailed and practical advice that can help users resolve their car troubles or understand more about their vehicles based on context.
            
            2. References:

                - When answering a question, if you use specific information from the context, add the context's metadata (source and page) as references at the end of your response.
                - Do not repeat same reference.
                - Never include irrelevant context.

            The first information you have is the relevant text that is automatically extracted from the manual or through the web search. \n
            Context: {context} \n

            History: {history} \n

            Based on all the information provided, answer the following question briefly: \n
            Question: {question} \n

            Lead: {lead} \n
            Do not use your prior knowledge! \n
            If the question is too general, ask for specific information. \n
            Answer: Answer in Turkish\n
            References: 1: \n
                        2: \n
                        ...
                        """
        )

        # Chain

        # rag_chain = prompt | myLLM | StrOutputParser()
        rag_chain =LLMChain(llm=self.myLLM, prompt=prompt, memory=self.memory)

        return rag_chain
    
    def hallucination_grader(self):
        class GradeHallucinations(BaseModel):
            """Binary score for hallucination present in generation answer."""

            binary_score: str = Field(description="Don't consider calling external APIs for additional information. Answer is supported by the facts, 'yes' or 'no'.")
        
        # LLM with function call 
        structured_llm_grader_hallucination = self.myLLM.with_structured_output(GradeHallucinations)
        
        # Prompt 
        system = """You are a grader assessing whether an LLM generation is supported by a set of retrieved facts. \n 
            If the LLM generation is greetings sentences or say 'cannot answer the question", always consider it is a yes. \n
            For others, restrict yourself to give a binary score, either 'yes' or 'no'. If the answer is supported or partially supported by the set of facts, consider it a yes. \n
            Don't consider calling external APIs or prior knowledge for additional information as consistent with the facts."""

        hallucination_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", system),
                ("human", "Set of facts: \n\n {documents} \n\n LLM generation: {generation}"),
        ]
        )
        
        hallucination_grader = hallucination_prompt | structured_llm_grader_hallucination

        return hallucination_grader
    
    def answer_grader(self):
        class GradeAnswer(BaseModel):
            """Binary score to assess answer addresses question."""

            binary_score: str = Field(description="Answer addresses the question, 'yes' or 'no'")

        # LLM with function call 
        structured_llm_grader_answer = self.myLLM.with_structured_output(GradeAnswer)

        # Prompt 
        system = """You are a grader assessing whether an answer addresses / resolves a question \n 
            Give a binary score 'yes' or 'no'. Yes' means that the answer resolves the question. \n
            If the LLM generation is greetings sentences, always consider it is a yes.\n
            If the LLM generation said that 'I cannot answer that question", consider it is a yes."""

        answer_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", system),
                ("human", "User question: \n\n {question} \n\n LLM generation: {generation}"),
            ]
        )

        answer_grader = answer_prompt | structured_llm_grader_answer

        return answer_grader
    

    def web_search(self):
        os.environ["GOOGLE_CSE_ID"] = os.getenv('google_search_id')
        os.environ["GOOGLE_API_KEY"] = os.getenv('google_search_api')

        search = GoogleSearchAPIWrapper()

        def top3_results(query):
            return search.results(query, 3)

        web_search_tool = Tool(
            name="google_search",
            description="Search Google for recent results.",
            func=top3_results,
        )
    
        return web_search_tool
    

    def lagchain_graph(self):

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

            Attributes:
                question: question
                generation: LLM generation
                web_search: whether to add search
                documents: list of documents 
            """
            question : str
            generation : str
            web_search : str
            documents : List[str]
            iter_halucination: int
            lead: str

        ### Nodes

        def retrieve(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
            """
            print("---RETRIEVE from Vector Store DB---")
            question = state["question"]
            # Retrieval
            documents = self.retriever().invoke(question)
            return {"documents": documents, "question": question}

        def generate(state):
            """
            Generate answer using RAG on retrieved documents

            Args:
                state (dict): The current graph state

            Returns:
                state (dict): New key added to state, generation, that contains LLM generation
            """
            print("---GENERATE Answer---")
            question = state["question"]
            documents = state["documents"]
            lead = state["lead"]
            
            # RAG generation
            generation = self.main_prompt().invoke({"context": documents, "question": question, "history": self.memory, "lead": lead})
            return {"documents": documents, "question": question, "generation": generation}

        def history_router(state):

            question = state["question"]
            history_log = DatabaseOperations.question_history_search(client=self.client, 
                                            collection_name=self.collection_name_manual,
                                            car_id=self.CAR_ID, 
                                            question=question, embeddings=self.embeddings)
            
            if len(history_log) > 0:
                print("---ANSWER FROM HISTORY---")
                return 'question history'

            else: 
                print("---ANSWER FROM MANUAL---")
                return 'user manual'
            
        def generate_from_history(state):
            # sohbet devamlılığı için history'den konuşmayı memory'e ekle
            question = state["question"]
            history_log = DatabaseOperations.question_history_search(client=self.client, 
                                            collection_name=self.collection_name_manual,
                                            car_id=self.CAR_ID,  
                                            question=question, embeddings=self.embeddings)
            return {"generation": {"text": history_log[0].payload["answer"]}}

        def grade_documents(state):
            """
            Determines whether the retrieved documents are relevant to the question
            If any document is not relevant, we will set a flag to run web search

            Args:
                state (dict): The current graph state

            Returns:
                state (dict): Filtered out irrelevant documents and updated web_search state
            """

            print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
            question = state["question"]
            documents = state["documents"]
            # Score each doc
            filtered_docs = []
            web_search = "No"
            for d in documents:
                score = self.grade_documents().invoke({"question": question, "document": d.page_content})
                grade = score.binary_score
                # Document relevant
                if grade.lower() == "yes":
                    print("---GRADE: DOCUMENT RELEVANT---")
                    filtered_docs.append(d)
                # Document not relevant
                else:
                    print("---GRADE: DOCUMENT NOT RELEVANT---")
                    # We do not include the document in filtered_docs
                    # We set a flag to indicate that we want to run web search
                    # web_search = "Yes"
                    continue
            if filtered_docs == []:
                web_search = "Yes"
                    
            return {"documents": filtered_docs, "question": question, "web_search": web_search}

        def grade_service(state):
            hizmet_listesi = hizmet_listesi = {"Bakım": """Check-Up, Periyodik Bakım, Aks Değişimi, Amortisör Değişimi, Amortisör Takozu Değişimi, Baskı Balata Değişimi, Benzin Filtresi Değişimi,
            Debriyaj Balatası Değişimi, Direksiyon Kutusu Değişimi, Dizel Araç Bakımı, Egzoz Muayenesi, Fren Kaliperi Değişimi, El Freni Teli Değişimi,
            Fren Balatası Değişimi, Fren Disk Değişimi, Hava Filtresi Değişimi, Helezon Yay Değişimi, Kampana Fren Balatası Değişimi, 
            Kızdırma Bujisi Değişimi, Rot Başı Değişimi, Rot Kolu Değişimi, Rotil Değişimi, Silecek Değişimi, Süspansiyon, Triger Kayışı Değişimi,
            Triger Zinciri Değişimi, V Kayışı Değişimi, Yağ Filtresi Değişimi, Yakıt Filtresi Değişimi,""",
                            "Yağ ve Sıvılar": """Şanzıman Yağı Değişimi, Dizel Araçlarda Yağ Değişimi, Yağ Değişimi, Fren Hidrolik Değişimi, Antifriz Değişimi,""",
                            "Akü": """Akü Şarj Etme, Akü Değişimi""",
                            "Klima": """Oto Klima Kompresörü Tamiri, Oto Klima Tamiri, Araç Klima Temizliği, Araç Klima Bakteri Temizliği, Klima Gazı Dolumu, Klima Dezenfeksiyonu, Polen Filtresi Değişimi""",
                            "Elektrik": """Servis Uyarı Lambası Sıfırlama,Buji Kablosu Değişimi, Arıza Tespit, Göstergelerin Kontrolü, Far Ayarı ve Ampul Değişimi, Buji Değişimi, Sigorta Değişimi""",
                            "Lastik/ Jant": """Lastik Jant Satış, Lastik Değişimi, Balans Ayarı, Rot Ayarı, Rotasyon, Lastik Tamiri, Hava Kontrolü, Nitrojen Dolumu, Supap Değişimi, Lastik Saklama (Lastik Oteli), Jant Sökme Takma,""",
                            "Diğer": """Cam Tamiri""",
                            "Hibrit Araçlar": "Hibrit Araç Aküsü"}
            
            print("---CHECK SERVICE RELEVANCE TO QUESTION---")
            question = state["question"]
            score = self.lead_check().invoke({"hizmet_listesi": hizmet_listesi, "question": question})
            grade = score.binary_score
            lead = 'Always use "XXX Servisi" only instead of "Yetkili Servis" or "Yetkili Bayi".'
            
            # Document relevant
            if grade.lower() == "yes":
                print("---GRADE: SERVICE RELEVANT---")
                lead = 'Always use "Otobakım Servisi" only instead of "Toyota Servisi", "Yetkili Servis" or "Toyota Bayisi". At the end of the answer, offer the user to visit Otobakım shops for quick and affordable solutions, starting with the phrase: "If you are not sure how to proceed.".'
            # Document not relevant
            else:
                print("---GRADE: SERVICE NOT RELEVANT---")

            return {"lead": lead}
            
        def web_search(state):
            """
            Web search based based on the question

            Args:
                state (dict): The current graph state

            Returns:
                state (dict): Appended web results to documents
            """

            print("---WEB SEARCH. Append to vector store db---")
            question = state["question"]
            documents = state["documents"]
            translated_question = self.translater().invoke({"question": question})["text"]

            # Web search
            docs = self.web_search().invoke({"query": translated_question})
            # web_results = [d["snippet"] for d in docs]
            # source = [d["link"] for d in docs]
            # web_results = Document(page_content = web_results, metadata={'source': source})
            web_results = [Document(page_content=d["snippet"], metadata={'source': d["link"]}) for d in docs]
            # if documents is not None:
            #     documents.append(web_results)
            # else:
            #     documents = [web_results]
            if documents == []:
                documents = web_results
            print(documents)
            return {"documents": documents, "question": question}

            
        def decide_to_generate(state):
            """
            Determines whether to generate an answer, or add web search

            Args:
                state (dict): The current graph state

            Returns:
                str: Binary decision for next node to call
            """

            print("---ASSESS GRADED DOCUMENTS---")
            question = state["question"]
            web_search = state["web_search"]
            filtered_documents = state["documents"]

            if web_search == "Yes":
                # All documents have been filtered check_relevance
                # We will re-generate a new query
                print("---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, INCLUDE WEB SEARCH---")
                return "websearch"
            else:
                # We have relevant documents, so generate answer
                print("---DECISION: GENERATE---")
                return "generate"

        def hallucination_router(state):
            print("---QUESTION CHANGING---")
            question = 'You are hallucinating! Please change your answer by sticking to the context.'
            iter_halucination = state["iter_halucination"]
            iter_halucination += 1
            return {'question': question, "iter_halucination": iter_halucination}

        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"]
            iter_halucination = state["iter_halucination"]
            # print("Generation:", generation)
            score = self.hallucination_grader().invoke({"documents": documents, "generation": generation})
            grade = score.binary_score
            # Check hallucination
            if grade == "yes":
                print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
                # Check question-answering
                print("---GRADE GENERATION vs QUESTION---")
                score = self.answer_grader().invoke({"question": question,"generation": generation})
                grade = score.binary_score
                if grade == "yes":
                    print("---DECISION: GENERATION ADDRESSES QUESTION---")
                    return "useful"
                else:
                    print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")
                    return "not useful"
            else:
                if iter_halucination < 2:
                    pprint("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")
                    return "not supported"
                else:
                    return "not useful"
            
        def lead_check(answer):

            if "otobakım" in answer.lower():
                return 1
            else: 
                return 0
            

        def print_result(question, result):
            output_text = f"""### Question: 
            {question}
            ### Answer: 
            {result}
            """
            return(output_text)
        
        workflow = StateGraph(GraphState)

        # Define the nodes
        workflow.add_node("websearch", web_search) # web search # key: action to 0do
        workflow.add_node("retrieve", retrieve) # retrieve
        workflow.add_node("grade_documents", grade_documents) # grade documents
        workflow.add_node("generate", generate) # generatae
        workflow.add_node("hallucination_router", hallucination_router)
        workflow.add_node("grade_service", grade_service)

        # workflow.add_node("generate_from_history", generate_from_history)

        workflow.add_edge("grade_service", "retrieve")
        workflow.add_edge("websearch", "generate") #start -> end of node
        workflow.add_edge("retrieve", "grade_documents")
        workflow.add_edge("hallucination_router", "generate")


        # workflow.add_edge("generate_from_history", END)

        # Build graph
        # workflow.set_conditional_entry_point(
        #     history_router, # defined function
        #     {
        #         "question history": "generate_from_history", #returns of the function
        #         "user manual": "grade_service",   #returns of the function
        #     },
        # )
        workflow.set_entry_point(
            "grade_service")

        workflow.add_conditional_edges(
            "grade_documents", # start: node
            decide_to_generate, # defined function
            {
                "websearch": "websearch", #returns of the function
                "generate": "generate",   #returns of the function
            },
        )

        workflow.add_conditional_edges(
            "generate", # start: node
            grade_generation_v_documents_and_question, # defined function
            {
                "not supported": "hallucination_router", #returns of the function
                "not useful": END,               #returns of the function
                "useful": END,   #returns of the function
            },
        )

        # Compile
        app = workflow.compile()

        return app