Update app.py

app.py

@@ -1,23 +1,20 @@
 import os
 import configparser
 
-# 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']
-# 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})
-
-from typing import List, Union
+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
@@ -28,37 +25,58 @@ from langchain.prompts import PromptTemplate, StringPromptTemplate, load_prompt,
 from langchain.llms import OpenAI
 from langchain.retrievers.self_query.base import SelfQueryRetriever
 from langchain.schema import AgentAction, AgentFinish, HumanMessage
-from langchain.vectorstores import DocArrayInMemorySearch, Chroma
-
-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-4', temperature=0.5)
-user_response_chain = LLMChain(
-    llm=llm,
-    prompt=user_response_prompt, 
-    verbose=False,
-    output_key="user_responses"
-)
+from langchain.vectorstores import Chroma
 
-df = pd.read_json('./data/unified_wine_data.json', encoding='utf-8', lines=True)
+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')
 
-loader =DataFrameLoader(data_frame=df, page_content_column='name')
+### 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",
@@ -91,57 +109,164 @@ metadata_field_info = [
     ),
 ]
 
-vectorstore = Chroma.from_documents(docs, embeddings)
-document_content_description = "Database of a wine"
+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)
-retriever = SelfQueryRetriever.from_llm(
-    llm, vectorstore, document_content_description, metadata_field_info, verbose=False
+wine_retriever = SelfQueryRetriever.from_llm(
+    llm, wine_vectorstore, document_content_description, metadata_field_info, verbose=True
 )  # Added missing closing parenthesis
 
-def search_with_url(query):
-    return SeleniumURLLoader(urls=[query]).load()
+#### 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'])
 
-index = VectorstoreIndexCreator(
-    vectorstore_cls=DocArrayInMemorySearch
-).from_loaders([loader])
+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=retriever.get_relevant_documents,
+        func=wine_retriever.get_relevant_documents,
+        coroutine=wine_retriever.aget_relevant_documents,
         description="""
-Database about the wines in wine store. You can get information such as the price of the wine, purchase URL, features, rating information, and more.
+Database about the wines in wine store.
 You can search wines with the following attributes:
-- body: 1-5 rating int for the body of wine. You have to specify greater than or less than. For example, if you want to search for wines with a body rating of less than 3, enter 'body: gt 0 lt 3'
-- price: The price range of the wine. Please enter the price range in the form of range. For example, if you want to search for wines that cost less than 20,000 won, enter 'price: gt 0 lt20000'
-- rating: 1-5 rating float for the wine. You have to specify greater than or less than. For example, if you want to search for wines with a rating of less than 3, enter 'rating: gt 0 lt 3'
+- 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. 입력할 때는 '와인 이름은 "비냐 조잘" 입니다' 이런 식으로 입력해주세요.
+- 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 = "Search specific wine with url",
-    #     func=search_with_url,
-    #     description="Search specific wine with url. Query must be url"
-    # ),
     Tool(
-        name = "Wine database 2",
-        func=index.query,
-        description="Database about the wines in wine store. You can use this tool if you're having trouble getting information from the wine database tool above. Query must be in String"
+        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. 
+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.
 
@@ -153,6 +278,7 @@ 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
@@ -173,14 +299,15 @@ Last user saying: {input}
 """
 
 conversation_stages_dict = {
-    "1": "Start: Start the conversation by introducing yourself. Be polite and respectful while maintaining a professional tone of conversation.",
-    "2": "Analyze: Figuring out the customer's preferences in order to make wine recommendations. Ask questions to figure out the preferences of your customer in order to make wine recommendations. Ask only one question at a time. The wine database tool is not available here.",
-    "3": "Product Recommendation: Recommend the right wine based on the user's preferences identified. Recommendations must be limited to wines in wine database, and you can use tools to do this. After making a wine recommendation, it asks if the user likes the wine you recommended.",
-    "4": "Sales: If the customer wants to get the wine you recommended, provides a link and image of wine and price of it. Link you provide must be open in new tab when clicked. Otherwise, it takes you back to the recommendation stage.",
-    "5": "Close: When you're done, say goodbye to the customer.",
-    "6": "Question and Answering: This is where you answer the customer's questions. To answer customer question, you can use the search tool or the wine database tool.",
-    "7": "Place Recommendation: Recommend wine bar based on location. you must be ask the location the customer wants before recommend the wine bar", 
-    "8": "Not in the given steps: This step is for when none of the steps between 1 and 7 apply.",
+    "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
@@ -197,8 +324,42 @@ class CustomPromptTemplate(StringPromptTemplate):
         # 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
@@ -208,12 +369,12 @@ class CustomPromptTemplate(StringPromptTemplate):
         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]:
@@ -226,39 +387,263 @@ class CustomOutputParser(AgentOutputParser):
                 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]*(.*)"
+        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()
 
-llm_chain = LLMChain(llm=ChatOpenAI(model='gpt-4', temperature=0.0), prompt=prompt, verbose=False,)
-
-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=True)
+### 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."
+    ),
+]
 
-import gradio as gr
+        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()
 
-# user_response, stage_history, conversation_history, pre_conversation_history = "", "", """""", """"""
+# logging
+# callback = gr.CSVLogger()
+hf_writer = gr.HuggingFaceDatasetSaver(huggingface_token, "chatwine-korean")
 
-stage_description = ""
-for key, value in conversation_stages_dict.items():
-    stage_description += f"{key}.{value}\n"
 
 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>
@@ -268,52 +653,118 @@ with gr.Blocks(css='#chatbot .overflow-y-auto{height:750px}') as demo:
                 LinkedIn <a href="https://www.linkedin.com/company/audrey-ai/about/">Audrey.ai</a>
             </p>
         </div>""")
- 
+    
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(label='User input')
-    samples = [["이번 주에 친구들과 모임이 있는데, 훌륭한 와인 한 병을 추천해줄래?"], ["입문자에게 좋은 와인을 추천해줄래?"], ["보르도와 부르고뉴 와인의 차이점은 뭐야?"]]
-    user_response_examples = gr.Dataset(samples=samples, components=[msg], type="index")
-    stage_history = gr.Textbox(value="stage history: ", interactive=False, label='stage history')
-    submit_btn = gr.Button("전송")
+
+    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])
-    stage_info = gr.Textbox(value=stage_description, interactive=False, label='stage description')
-
-    def load_example(example_id):
-        global samples
-        return samples[example_id][0]
-
-    def answer(user_response, chat_history, stage_history):
-        global samples
-        chat_history = chat_history or []
-        stage_history = stage_history or ""
-        pre_conversation_history = ""
-        for idx, chat in enumerate(chat_history):
-            pre_conversation_history += f"User: {chat[0]} <END_OF_TURN>\n"
-            pre_conversation_history += f"이우선: {chat[1]} <END_OF_TURN>\n"
-        conversation_history = pre_conversation_history + f"User: {user_response} <END_OF_TURN>\n"
-        stage_number = stage_analyzer_chain.run({'conversation_history': conversation_history, 'stage_history': stage_history})
+
+    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
-        response = agent_executor.run({'input':user_response, 'conversation_history': pre_conversation_history, 'stage_number': stage_number})
-        conversation_history += "이우선: " + response + "\n"
-        for line in conversation_history.split('\n'):
-            print(line)
-        response = response.split('<END_OF_TURN>')[0]
-        chat_history.append((user_response, response))
-        user_response_examples = []
-        for user_response_example in user_response_chain.run({'conversation_history': conversation_history}).split('|'):
-            user_response_examples.append([user_response_example])
-        samples = user_response_examples
- 
-        return "", chat_history, stage_history, gr.Dataset.update(samples=samples)
-
-    def clear(*args):
-        global samples
-        samples = [["이번 주에 친구들과 모임이 있는데, 훌륭한 와인 한 병을 추천해줄래?"], ["입문자에게 좋은 와인을 추천해줄래?"], ["보르도와 부르고뉴 와인의 차이점은 뭐야?"]]
-        return gr.Dataset.update(samples=samples), "stage history: "
-
-    clear_btn.click(fn=clear, inputs=[user_response_examples, stage_history], outputs=[user_response_examples, stage_history])
-    user_response_examples.click(load_example, inputs=[user_response_examples], outputs=[msg])      
-    submit_btn.click(answer, [msg, chatbot, stage_history], [msg, chatbot, stage_history, user_response_examples])
-    msg.submit(answer, [msg, chatbot, stage_history], [msg, chatbot, stage_history, user_response_examples])
-demo.launch()
+
+        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()