app.py

import os
import configparser

from typing import List, Union, Optional, Any, Dict
import re
import sys
import time
import json
import asyncio
import aiohttp
import requests
import threading

import pandas as pd
from langchain import SerpAPIWrapper, LLMChain
from langchain.agents import Tool, AgentType, AgentExecutor, LLMSingleActionAgent, AgentOutputParser
from langchain.callbacks.streaming_stdout_final_only import FinalStreamingStdOutCallbackHandler
from langchain.chat_models import ChatOpenAI
from langchain.chains import LLMChain, SimpleSequentialChain
from langchain.chains.query_constructor.base import AttributeInfo
from langchain.document_loaders import DataFrameLoader, SeleniumURLLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.indexes import VectorstoreIndexCreator
from langchain.prompts import PromptTemplate, StringPromptTemplate, load_prompt, BaseChatPromptTemplate
from langchain.llms import OpenAI
from langchain.retrievers.self_query.base import SelfQueryRetriever
from langchain.schema import AgentAction, AgentFinish, HumanMessage
from langchain.vectorstores import Chroma

import gradio as gr

# config = configparser.ConfigParser()
# config.read('./secrets.ini')
# openai_api_key = config['OPENAI']['OPENAI_API_KEY']
# serper_api_key = config['SERPER']['SERPER_API_KEY']
# serp_api_key = config['SERPAPI']['SERPAPI_API_KEY']
# kakao_api_key = config['KAKAO_MAP']['KAKAO_API_KEY']
# huggingface_token = config['HUGGINGFACE']['HUGGINGFACE_TOKEN']


# os.environ.update({'OPENAI_API_KEY': openai_api_key})
# os.environ.update({'SERPER_API_KEY': serper_api_key})
# os.environ.update({'SERPAPI_API_KEY': serp_api_key})
# os.environ.update({'KAKAO_API_KEY': kakao_api_key})
# os.environ.update({'HUGGINGFACE_TOKEN': huggingface_token})

huggingface_token = os.getenv('HUGGINGFACE_TOKEN')
kakao_api_key = os.getenv('KAKAO_API_KEY')

### Load wine database json
df = pd.read_json('./data/unified_wine_data.json', encoding='utf-8', lines=True)
### Prepare Langchain Tool
#### Tool1: Wine database 1
df['page_content'] = ''
columns = ['name', 'pairing']
for column in columns:
    if column != 'page_content':
        df['page_content'] += column + ':' + df[column].astype(str) + ','
columns = ['rating', 'price', 'body', 'sweetness', 'alcohol', 'acidity', 'tannin']
for idx in df.index:
    for column in columns:
        if type(df[column][idx]) == str:
            df[column][idx] = df[column][idx].replace(',', '')
        df[column][idx] = float(df[column][idx]) if df[column][idx] != '' else -1
loader =DataFrameLoader(data_frame=df, page_content_column='page_content')
docs = loader.load()
embeddings = OpenAIEmbeddings()
# 아래는 wine database1에 metadata_field Attribute이다. 아래를 기준으로 서치를 진행하게 된다.
metadata_field_info = [
    AttributeInfo(
        name="body",
        description="1-5 rating for the body of wine",
        type="int",
    ),
    AttributeInfo(
        name="tannin",
        description="1-5 rating for the tannin of wine",
        type="int",
    ),
    AttributeInfo(
        name="sweetness",
        description="1-5 rating for the sweetness of wine",
        type="int",
    ),
    AttributeInfo(
        name="alcohol",
        description="1-5 rating for the alcohol of wine",
        type="int",
    ),
    AttributeInfo(
        name="price",
        description="The price of the wine",
        type="int",
    ),
    AttributeInfo(
        name="rating", 
        description="1-5 rating for the wine", 
        type="float"
    ),
    AttributeInfo(
        name="wine_type", 
        description="The type of wine. It can be '레드', '로제', '스파클링', '화이트', '디저트', '주정강화'", 
        type="string"
    ),
    AttributeInfo(
        name="country", 
        description="The country of wine. It can be '기타 신대륙', '기타구대륙', '뉴질랜드', '독일', '미국', '스페인', '아르헨티나', '이탈리아', '칠레', '포루투칼', '프랑스', '호주'", 
        type="float"
    ),
]

wine_vectorstore = Chroma.from_documents(docs, embeddings)
document_content_description = "A database of wines. 'name' and 'pairing' must be included in the query, and 'Body', 'Tannin', 'Sweetness', 'Alcohol', 'Price', 'Rating', 'Wine_Type', and 'Country' can be included in the filter. query and filter must be form of 'key: value'. For example, query: 'name: 돔페리뇽, pairing:육류'."
llm = OpenAI(temperature=0)
wine_retriever = SelfQueryRetriever.from_llm(
    llm, wine_vectorstore, document_content_description, metadata_field_info, verbose=True
)  # Added missing closing parenthesis

#### Tool2: Wine bar database
df = pd.read_json('./data/wine_bar.json', encoding='utf-8', lines=True)
df['page_content'] = ''
columns = ['summary']
for column in columns:
    if column != 'page_content':
        df['page_content'] += df[column].astype(str) + ','
df = df.drop(columns=['review'])

loader =DataFrameLoader(data_frame=df, page_content_column='page_content')
docs = loader.load()
embeddings = OpenAIEmbeddings()
wine_bar_vectorstore = Chroma.from_documents(docs, embeddings)
wine_bar_vectorstore.similarity_search_with_score('여자친구랑 갈만한 와인바', k=5)
metadata_field_info = [
    AttributeInfo(
        name="name",
        description="The name of the wine bar",
        type="str",
    ),
    AttributeInfo(
        name="rating", 
        description="1-5 rating for the wine bar", 
        type="float"
    ),
    AttributeInfo(
        name="district",
        description="The district of the wine bar.",
        type="str",
    ),
]

document_content_description = "Database of a winebar"
llm = OpenAI(temperature=0)
wine_bar_retriever = SelfQueryRetriever.from_llm(
    llm, wine_bar_vectorstore, document_content_description, metadata_field_info=metadata_field_info, verbose=True
)  # Added missing closing parenthesis
#### Tool3: Search in Google
search = SerpAPIWrapper()
#### Tool4: Kakao Map API


class KakaoMap:
    def __init__(self):
        self.url = 'https://dapi.kakao.com/v2/local/search/keyword.json' 
        self.headers = {"Authorization": f"KakaoAK {kakao_api_key}"}
        
    async def arun(self, query):
        async with aiohttp.ClientSession() as session:
            params = {'query': query,'page': 1}
            async with session.get(self.url, params=params, headers=self.headers) as response:
                places = await response.json()
                address = places['documents'][0]['address_name']
                if not address.split()[0].startswith('서울'):
                    return {'district': 'not in seoul'}
                else:
                    return {'district': address.split()[1]}
                
    def run(self, query):
        params = {'query': query,'page': 1}
        places = requests.get(self.url, params=params, headers=self.headers).json()
        address = places['documents'][0]['address_name']
        if not address.split()[0].startswith('서울'):
            return {'district': 'not in seoul'}
        else:
            return {'district': address.split()[1]}
kakao_map = KakaoMap()
tools = [
    Tool(
        name="Wine database",
        func=wine_retriever.get_relevant_documents,
        coroutine=wine_retriever.aget_relevant_documents,
        description="""
Database about the wines in wine store.
You can search wines with the following attributes:
- price: The price range of the wine. You have to specify greater than and less than.
- rating: 1-5 rating float for the wine. You have to specify greater than and less than.
- wine_type: The type of wine. It can be '레드', '로제', '스파클링', '화이트', '디저트', '주정강화'
- name: The name of wine.
- pairing: The food pairing of wine.
The form of Action Input must be 'key1: value1, key2: value2, ...'. For example, to search for wines with a rating of less than 3 points, a price range of 50000원 or more, and a meat pairing, enter 'rating: gt 0 lt 3, price: gt 50000, pairing: 고기'.
--------------------------------------------------
You can get the following attributes:
- url: Wine purchase site URL.
- vivino_link: Vivino link of wine.
- flavor_description
- site_name: Wine purchase site name.
- name: The name of wine in korean.
- en_name: The name of wine in english.
- price: The price of wine in 원.
- rating: 1-5 vivino rating.
- wine_type: The type of wine.
- pairing: The food pairing of wine.
- pickup_location: Offline stores where you can purchase wine
- img_url
- country
- body
- tannin
- sweetness
- acidity
- alcohol
- grape
The form of Desired Outcome must be 'key1, key2, ...'. For example to get the name and price of wine, enter 'name, price'.
"""
    ),
    Tool(
        name = "Wine bar database",
        func=wine_bar_retriever.get_relevant_documents,
        coroutine=wine_bar_retriever.aget_relevant_documents,
        description="Database about the winebars in Seoul. It should be the first thing you use when looking for information about a wine bar."
"""
- query: The query of winebar. You can search wines with review data like mood or something.
- name: The name of winebar.
- price: The average price point of a wine bar.
- rating: 1-5 rating float for the wine bar. 
- district: The district of wine bar. Input district must be korean. For example, if you want to search for wines in Gangnam, enter 'district: 강남구'
The form of Action Input must be 'key1: value1, key2: value2, ...'. 
--------------------------------------------------
You can get the following attributes:
- name: The name of winebar.
- url: Wine purchase site URL.
- rating: 1-5 망고플레이트(맛집검색 앱) rating.
- summary: Summarized information about wine bars
- address
- phone
- parking
- opening_hours
- menu
- holidays
- img_url
The form of Desired Outcome must be 'key1, key2, ...'. For example to get the name and price of wine, enter 'name, price'.
"""
    ),
    Tool(
        name = "Search",
        func=search.run,
        coroutine=search.arun,
        description="Useful for when you need to ask with search. Search in English only."
    ),
    Tool(
        name = "Map",
        func=kakao_map.run,
        coroutine=kakao_map.arun,
        description="The tool used to draw a district for a region. When looking for wine bars, you can use this before applying filters based on location. The query must be in Korean. You can get the following attribute: district."
    ),
]
template = """
Your role is a chatbot that asks customers questions about wine and makes recommendations.
Never forget your name is "이우선".
Keep your responses in short length to retain the user's attention unless you describe the wine for recommendations.
Be sure to actively empathize and respond to your users.
Only generate one response at a time! When you are done generating, end with '<END_OF_TURN>' to give the user a chance to respond.
Responses should be in Korean.

Complete the objective as best you can. You have access to the following tools:

{tools}

Use the following format:
Thought: you should always think about what to do
Action: the action to take, should be one of [{tool_names}]
Action Input: the input to the action
Desired Outcome: the desired outcome from the action (optional)
Observation: the result of the action
... (this Thought/Action/Action Input/Observation can repeat N times)
Thought: I now know the final answer
이우선: the final response to the user

You must respond according to the conversation stage within the triple backticks and conversation history within in '======'.

Current conversation stage: 
```{conversation_stage}```

Conversation history: 
=======
{conversation_history}
=======

Last user saying: {input}
{agent_scratchpad}
"""

conversation_stages_dict = {
    "1": "Introduction: Start the conversation by introducing yourself. Maintain politeness, respect, and a professional tone.",
    "2": "Needs Analysis: Identify the customer's needs to make wine recommendations. Note that the wine database tools are not available. You ask about the occasion the customer will enjoy the wine, what they eat with it, and their desired price point. Ask only ONE question at a time.",
    "3": "Checking Price Range: Asking the customer's preferred price point. Again, remember that the tool for this is not available. But if you know the customer's perferences and price range, then search for the three most suitable wines with tool and recommend them product cards in a list format with a Vivino link, price, rating, wine type, flavor description, and image.",
    "4": "Wine Recommendation: Propose the three most suitable wines based on the customer's needs and price range. Before the recommendation, you should have identified the occasion the customer will enjoy the wine, what they will eat with it, and their desired price point. Each wine recommendation should form of product cards in a list format with a Vivino link, price, rating, wine type, flavor description, and image. Use only wines available in the database for recommendations. If there are no suitable wines in the database, inform the customer. After making a recommendation, inquire whether the customer likes the suggested wine.",
    "5": "Sales: If the customer approves of the recommended wine, provide a detailed description. Supply a product card in a list format with a Vivino link, price, rating, wine type, flavor description, and image.",
    "6": "Location Suggestions: Recommend wine bars based on the customer's location and occasion. Before making a recommendation, always use the map tool to find the district of the customer's preferred location. Then use the wine bar database tool to find a suitable wine bar. Provide form of product cards in a list format with the wine bar's name, url, rating, address, menu, opening_hours, holidays, phone, summary, and image with img_urls.",
    "7": "Concluding the Conversation: Respond appropriately to the customer's comments to wrap up the conversation.",
    "8": "Questions and Answers: This stage involves answering customer's inquiries. Use the search tool or wine database tool to provide specific answers where possible. Describe answer as detailed as possible",
    "9": "Other Situations: Use this step when the situation does not fit into any of the steps between 1 and 8."
}

# Set up a prompt template
class CustomPromptTemplate(StringPromptTemplate):
    # The template to use
    template: str
    # The list of tools available
    tools: List[Tool]
    
    def format(self, **kwargs) -> str:
        stage_number = kwargs.pop("stage_number")
        kwargs["conversation_stage"] = conversation_stages_dict[stage_number]
        # Get the intermediate steps (AgentAction, Observation tuples)
        # Format them in a particular way
        intermediate_steps = kwargs.pop("intermediate_steps")
        thoughts = ""
        special_chars = "()'[]{}"
        for action, observation in intermediate_steps:
            thoughts += action.log

            if ('Desired Outcome: ' in action.log) and (('Action: Wine database' in action.log) or ('Action: Wine bar database' in action.log)):
                regex = r"Desired Outcome:(.*)"
                match = re.search(regex, action.log, re.DOTALL)
                if not match:
                    raise ValueError(f"Could not parse Desired Outcome: `{action.log}`")
                desired_outcome_keys = [key.strip() for key in match.group(1).split(',')]

                pattern = re.compile(r'metadata=\{(.*?)\}')
                matches = pattern.findall(f'{observation}')
                documents = ['{'+f'{match}'+'}' for match in matches]
                
                pattern = re.compile(r"'(\w+)':\s*('[^']+'|\b[^\s,]+\b)")
                output=[]

                for doc in documents:
                    # Extract key-value pairs from the document string
                    matches = pattern.findall(doc)

                    # Convert matches to a dictionary
                    doc_dict = dict(matches)

                    # Create a new dictionary containing only the desired keys
                    item_dict = {}
                    for key in desired_outcome_keys:
                        value = doc_dict.get(key, "")
                        for c in special_chars:
                            value = value.replace(c, "")                        
                        item_dict[key] = value
                    output.append(item_dict)
                
                observation = ','.join([str(i) for i in output])

            thoughts += f"\nObservation: {observation}\nThought: "
        # Set the agent_scratchpad variable to that value
        kwargs["agent_scratchpad"] = thoughts
        # Create a tools variable from the list of tools provided
        kwargs["tools"] = "\n".join([f"{tool.name}: {tool.description}" for tool in self.tools])
        # Create a list of tool names for the tools provided
        kwargs["tool_names"] = ", ".join([tool.name for tool in self.tools])
        return self.template.format(**kwargs)


prompt = CustomPromptTemplate(
    template=template,
    tools=tools,
    input_variables=["input", "intermediate_steps", "conversation_history", "stage_number"]
)
class CustomOutputParser(AgentOutputParser):
    
    def parse(self, llm_output: str) -> Union[AgentAction, AgentFinish]:
        # Check if agent should finish
        if "이우선: " in llm_output:
            return AgentFinish(
                # Return values is generally always a dictionary with a single `output` key
                # It is not recommended to try anything else at the moment :)
                return_values={"output": llm_output.split("이우선: ")[-1].strip()},
                log=llm_output,
            )
        # Parse out the action and action input
        regex = r"Action\s*\d*\s*:(.*?)\nAction\s*\d*\s*Input\s*\d*\s*:[\s]*(.*?)\n"
        match = re.search(regex, llm_output, re.DOTALL)
        if not match:
            raise ValueError(f"Could not parse LLM output: `{llm_output}`")
        action = match.group(1).strip()
        action_input = match.group(2)
        # desired_outcome = match.group(3).strip() if match.group(3) else None

        # Return the action and action input
        return AgentAction(tool=action, tool_input=action_input.strip(" ").strip('"'), log=llm_output)

output_parser = CustomOutputParser()

### Gradio

class CustomStreamingStdOutCallbackHandler(FinalStreamingStdOutCallbackHandler):
    """Callback handler for streaming in agents.
    Only works with agents using LLMs that support streaming.

    The output will be streamed until "<END" is reached.
    """
    def __init__(
        self,
        *,
        answer_prefix_tokens: Optional[List[str]] = None,
        end_prefix_tokens: str = "<END",
        strip_tokens: bool = True,
        stream_prefix: bool = False,
        sender: str
    ) -> None:
        """Instantiate EofStreamingStdOutCallbackHandler.

        Args:
            answer_prefix_tokens: Token sequence that prefixes the anwer.
                Default is ["Final", "Answer", ":"]
            end_of_file_token: Token that signals end of file.
                Default is "END"
            strip_tokens: Ignore white spaces and new lines when comparing
                answer_prefix_tokens to last tokens? (to determine if answer has been
                reached)
            stream_prefix: Should answer prefix itself also be streamed?
        """
        super().__init__(answer_prefix_tokens=answer_prefix_tokens, strip_tokens=strip_tokens, stream_prefix=stream_prefix)
        self.end_prefix_tokens = end_prefix_tokens
        self.end_reached = False
        self.sender = sender

    def append_to_last_tokens(self, token: str) -> None:
        self.last_tokens.append(token)
        self.last_tokens_stripped.append(token.strip())
        if len(self.last_tokens) > 5:
            self.last_tokens.pop(0)
            self.last_tokens_stripped.pop(0)

    def on_llm_start(
        self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
    ) -> None:
        """Run when LLM starts running."""
        self.answer_reached = False
        self.end_reached = False

    def check_if_answer_reached(self) -> bool:
        if self.strip_tokens:
            return ''.join(self.last_tokens_stripped) in self.answer_prefix_tokens_stripped
        else:
            unfied_last_tokens = ''.join(self.last_tokens)
            try:
                unfied_last_tokens.index(self.answer_prefix_tokens)
                return True
            except:
                return False
            
    def check_if_end_reached(self) -> bool:
        if self.strip_tokens:
            return ''.join(self.last_tokens_stripped) in self.answer_prefix_tokens_stripped
        else:
            unfied_last_tokens = ''.join(self.last_tokens)
            try:
                unfied_last_tokens.index(self.end_prefix_tokens)
                self.sender[1] = True
                return True
            except:
                # try:
                #     unfied_last_tokens.index('Action Input')
                #     self.sender[1] = False
                #     return False
                # except:
                #     return False
                return False

    def on_llm_new_token(self, token: str, **kwargs: Any) -> None:
        """Run on new LLM token. Only available when streaming is enabled."""
        # Remember the last n tokens, where n = len(answer_prefix_tokens)
        self.append_to_last_tokens(token)
        
        # Check if the last n tokens match the answer_prefix_tokens list ...
        if not self.answer_reached and self.check_if_answer_reached():
            self.answer_reached = True
            if self.stream_prefix:
                for t in self.last_tokens:
                    sys.stdout.write(t)
                sys.stdout.flush()
            return
        
        if not self.end_reached and self.check_if_end_reached():
            self.end_reached = True

        if self.end_reached:
            pass
        elif self.answer_reached:
            if self.last_tokens[-2] == ":":
                pass
            else:
                self.sender[0] += self.last_tokens[-2]

class UnifiedAgent:
    def __init__(self):

        tools = [
    Tool(
        name="Wine database",
        func=wine_retriever.get_relevant_documents,
        coroutine=wine_retriever.aget_relevant_documents,
        description="""
Database about the wines in wine store.
You can search wines with the following attributes:
- price: The price range of the wine. You have to specify greater than and less than.
- rating: 1-5 rating float for the wine. You have to specify greater than and less than.
- wine_type: The type of wine. It can be '레드', '로제', '스파클링', '화이트', '디저트', '주정강화'
- name: The name of wine.
- pairing: The food pairing of wine.
The form of Action Input must be 'key1: value1, key2: value2, ...'. For example, to search for wines with a rating of less than 3 points, a price range of 50000원 or more, and a meat pairing, enter 'rating: gt 0 lt 3, price: gt 50000, pairing: 고기'.
--------------------------------------------------
You can get the following attributes:
- url: Wine purchase site URL.
- vivino_link: Vivino link of wine.
- flavor_description
- site_name: Wine purchase site name.
- name: The name of wine in korean.
- en_name: The name of wine in english.
- price: The price of wine in 원.
- rating: 1-5 vivino rating.
- wine_type: The type of wine.
- pairing: The food pairing of wine.
- pickup_location: Offline stores where you can purchase wine
- img_url
- country
- body
- tannin
- sweetness
- acidity
- alcohol
- grape
The form of Desired Outcome must be 'key1, key2, ...'. For example to get the name and price of wine, enter 'name, price'.
"""
    ),
    Tool(
        name = "Wine bar database",
        func=wine_bar_retriever.get_relevant_documents,
        coroutine=wine_bar_retriever.aget_relevant_documents,
        description="Database about the winebars in Seoul. It should be the first thing you use when looking for information about a wine bar."
"""
- query: The query of winebar. You can search wines with review data like mood or something.
- name: The name of winebar.
- price: The average price point of a wine bar.
- rating: 1-5 rating float for the wine bar. 
- district: The district of wine bar. Input district must be korean. For example, if you want to search for wines in Gangnam, enter 'district: 강남구'
The form of Action Input must be 'key1: value1, key2: value2, ...'. 
--------------------------------------------------
You can get the following attributes:
- name: The name of winebar.
- url: Wine purchase site URL.
- rating: 1-5 망고플레이트(맛집검색 앱) rating.
- summary: Summarized information about wine bars
- address
- phone
- parking
- opening_hours
- menu
- holidays
- img_url
The form of Desired Outcome must be 'key1, key2, ...'. For example to get the name and price of wine, enter 'name, price'.
"""
    ),
    Tool(
        name = "Search",
        func=search.run,
        coroutine=search.arun,
        description="Useful for when you need to ask with search. Search in English only."
    ),
    Tool(
        name = "Map",
        func=kakao_map.run,
        coroutine=kakao_map.arun,
        description="The tool used to draw a district for a region. When looking for wine bars, you can use this before applying filters based on location. The query must be in Korean. You can get the following attribute: district."
    ),
]

        llm_chain = LLMChain(llm=ChatOpenAI(model='gpt-4', temperature=0.5, streaming=True), prompt=prompt, verbose=True,)

        tool_names = [tool.name for tool in tools]
        agent = LLMSingleActionAgent(
            llm_chain=llm_chain, 
            output_parser=output_parser,
            stop=["\nObservation:"], 
            allowed_tools=tool_names
        )
        agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, verbose=False)

        self.agent_executor = agent_executor

    async def arun(self, sender, *args, **kwargs):
        resp = await self.agent_executor.arun(kwargs, callbacks=[CustomStreamingStdOutCallbackHandler(answer_prefix_tokens='이우선:', end_prefix_tokens='<END', strip_tokens=False, sender=sender)])
        return resp


class UnifiedChain:
    def __init__(self):
        stage_analyzer_inception_prompt = load_prompt("./templates/stage_analyzer_inception_prompt_template.json")
        llm = ChatOpenAI(model='gpt-3.5-turbo', temperature=0.0)
        stage_analyzer_chain = LLMChain(
            llm=llm,
            prompt=stage_analyzer_inception_prompt, 
            verbose=False, 
            output_key="stage_number")
        
        user_response_prompt = load_prompt("./templates/user_response_prompt.json")
        # 랭체인 모델 선언, 랭체인은 언어모델과 프롬프트로 구성됩니다.
        llm = ChatOpenAI(model='gpt-3.5-turbo', temperature=0.5)
        user_response_chain = LLMChain(
            llm=llm,
            prompt=user_response_prompt, 
            verbose=False, # 과정을 출력할지
            output_key="user_responses"
        )

        self.stage_analyzer_chain = stage_analyzer_chain
        self.user_response_chain = user_response_chain

    async def arun_stage_analyzer_chain(self, *args, **kwargs):
        resp = await self.stage_analyzer_chain.arun(kwargs)
        return resp
    
    async def arun_user_response_chain(self, *args, **kwargs):
        resp = await self.user_response_chain.arun(kwargs)
        return resp
    
unified_chain = UnifiedChain()
unified_agent = UnifiedAgent()

# logging
# callback = gr.CSVLogger()
hf_writer = gr.HuggingFaceDatasetSaver(huggingface_token, "chatwine-korean")


with gr.Blocks(css='#chatbot .overflow-y-auto{height:750px}') as demo:
    
    with gr.Row():
        gr.HTML("""<div style="text-align: center; max-width: 500px; margin: 0 auto;">
            <div>
                <h1>ChatWine</h1>
            </div>
            <p style="margin-bottom: 10px; font-size: 94%">
                LinkedIn <a href="https://www.linkedin.com/company/audrey-ai/about/">Audrey.ai</a>
            </p>
        </div>""")
    
    chatbot = gr.Chatbot()

    with gr.Row():
        with gr.Column(scale=0.85):
            msg = gr.Textbox()
        with gr.Column(scale=0.15, min_width=0):
            submit_btn = gr.Button("전송")

    user_response_examples = gr.Dataset(samples=[["이번 주에 친구들과 모임이 있는데, 훌륭한 와인 한 병을 추천해줄래?"], ["입문자에게 좋은 와인을 추천해줄래?"], ["연인과 가기 좋은 와인바를 알려줘"]], components=[msg], type="index")
    clear_btn = gr.ClearButton([msg, chatbot])

    dev_mod = True
    cur_stage = gr.Textbox(visible=dev_mod, interactive=False, label='current_stage')
    stage_hist = gr.Textbox(visible=dev_mod, value="stage history: ", interactive=False, label='stage history')
    chat_hist = gr.Textbox(visible=dev_mod, interactive=False, label='chatting_history')
    response_examples_text = gr.Textbox(visible=dev_mod, interactive=False, value="이번 주에 친구들과 모임이 있는데, 훌륭한 와인 한 병을 추천해줄래?|입문자에게 좋은 와인을 추천해줄래?|연인과 가기 좋은 와인바를 알려줘", label='response_examples')
    btn = gr.Button("Flag", visible=dev_mod)
    hf_writer.setup(components=[chat_hist, stage_hist, response_examples_text], flagging_dir="chatwine-korean")

    def click_flag_btn(*args):
        hf_writer.flag(flag_data=[*args])

    def clean(*args):
        return gr.Dataset.update(samples=[["이번 주에 친구들과 모임이 있는데, 훌륭한 와인 한 병을 추천해줄래?"], ["입문자에게 좋은 와인을 추천해줄래?"], ["연인과 가기 좋은 와인바를 알려줘"]]), "", "stage history: ", "", "이번 주에 친구들과 모임이 있는데, 훌륭한 와인 한 병을 추천해줄래?|입문자에게 좋은 와인을 추천해줄래?|연인과 가기 좋은 와인바를 알려줘"

    def load_example(response_text, input_idx):
        response_examples = []
        for user_response_example in response_text.split('|'):
            response_examples.append([user_response_example])
        return response_examples[input_idx][0]

    async def agent_run(agent_exec, inp, sender):
        sender[0] = ""
        await agent_exec.arun(inp)

    def user_chat(user_message, chat_history_list, chat_history):
        return (chat_history_list + [[user_message, None]], chat_history + f"User: {user_message} <END_OF_TURN>\n", [])

    async def bot_stage_pred(user_response, chat_history, stage_history):
        pre_chat_history = '<END_OF_TURN>'.join(chat_history.split('<END_OF_TURN>')[:-2])
        if pre_chat_history != '':
            pre_chat_history += '<END_OF_TURN>'
        # stage_number = unified_chain.stage_analyzer_chain.run({'conversation_history': pre_chat_history, 'stage_history': stage_history.replace('stage history: ', ''), 'last_user_saying':user_response+' <END_OF_TURN>\n'})
        stage_number = await unified_chain.arun_stage_analyzer_chain(conversation_history=pre_chat_history, stage_history= stage_history.replace('stage history: ', ''), last_user_saying=user_response+' <END_OF_TURN>\n')
        stage_number = stage_number[-1]
        stage_history += stage_number if stage_history == "stage history: " else ", " + stage_number

        return stage_number, stage_history

    async def bot_chat(user_response, chat_history, chat_history_list, current_stage): # stream output by yielding
        
        pre_chat_history = '<END_OF_TURN>'.join(chat_history.split('<END_OF_TURN>')[:-2])
        if pre_chat_history != '':
            pre_chat_history += '<END_OF_TURN>'

        sender = ["", False]        
        task = asyncio.create_task(unified_agent.arun(sender = sender, input=user_response+' <END_OF_TURN>\n', conversation_history=pre_chat_history, stage_number= current_stage))
        await asyncio.sleep(0)

        while(sender[1] == False):
            await asyncio.sleep(0.2)
            chat_history_list[-1][1] = sender[0]
            yield chat_history_list, chat_history + f"이우선: {sender[0]}<END_OF_TURN>\n"

        chat_history_list[-1][1] = sender[0]
        yield chat_history_list, chat_history + f"이우선: {sender[0]}<END_OF_TURN>\n"

    async def bot_response_pred(chat_history):
        response_examples = []
        pre_chat_history = '<END_OF_TURN>'.join(chat_history.split('<END_OF_TURN>')[-3:])
        out = await unified_chain.arun_user_response_chain(conversation_history=pre_chat_history)
        for user_response_example in out.split('|'):
            response_examples.append([user_response_example])
        return [response_examples, out, ""]
    
    # btn.click(lambda *args: hf_writer.flag(args), [msg, chat_hist, stage_hist, response_examples_text], None, preprocess=False)

    msg.submit(
        user_chat, [msg, chatbot, chat_hist], [chatbot, chat_hist, user_response_examples], queue=False
    ).then(
        bot_stage_pred, [msg, chat_hist, stage_hist], [cur_stage, stage_hist], queue=False
    ).then(
        bot_chat, [msg, chat_hist, chatbot, cur_stage], [chatbot, chat_hist]
    ).then(
        bot_response_pred, chat_hist, [user_response_examples, response_examples_text, msg]
    ).then(
        click_flag_btn, [chat_hist, stage_hist, response_examples_text], None
    )



    submit_btn.click(
        user_chat, [msg, chatbot, chat_hist], [chatbot, chat_hist, user_response_examples], queue=False
    ).then(
        bot_stage_pred, [msg, chat_hist, stage_hist], [cur_stage, stage_hist], queue=False
    ).then(
        bot_chat, [msg, chat_hist, chatbot, cur_stage], [chatbot, chat_hist]
    ).then(
        bot_response_pred, chat_hist, [user_response_examples, response_examples_text, msg]
    ).then(
        click_flag_btn, [chat_hist, stage_hist, response_examples_text], None
    )



    clear_btn.click(
        clean, 
        inputs=[user_response_examples, cur_stage, stage_hist, chat_hist, response_examples_text], 
        outputs=[user_response_examples, cur_stage, stage_hist, chat_hist, response_examples_text], 
        queue=False)
    user_response_examples.click(load_example, inputs=[response_examples_text, user_response_examples], outputs=[msg], queue=False)
    btn.click(lambda *args: hf_writer.flag(args), [chat_hist, stage_hist, response_examples_text], None, preprocess=False)
demo.queue(concurrency_count=100)
demo.launch()