# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.



import os
import langchain
from langchain.llms import OpenAI
from langchain import PromptTemplate


from configparser import ConfigParser
from langchain.cache import InMemoryCache
from posttext.src.posttext import PostText
langchain.llm_cache = InMemoryCache()


class ViewEngine:

    def __init__(self, path):
        print(path)
        config = ConfigParser(comment_prefixes=None)
        config.read(os.path.join(path, 'config.ini'))
        self.posttext = PostText(config, path)
        self.llm = OpenAI()

    def verbalize(self, query: str, answer: str):
        template = "Question: {query}\n\nAnswer: {context}\n\n Can you summarize the answer in a single sentence (durations are in seconds)?"
        prompt = PromptTemplate(
            input_variables=["query", "context"],
            template=template)
        
        return self.llm(prompt.format(query=query, context=answer))


    def flatten(self, lst):
        """Flatten a nested list
        """
        flattened_list = []
        for item in lst:
            if isinstance(item, tuple):
                item = list(item)

            if isinstance(item, list):
                flattened_list.extend(self.flatten(item))
            else:
                flattened_list.append(item)
        return flattened_list

    def query(self, 
              query: str):
        """Answer a query using a View-based QA method.
        """
        formattedprompt, sqlquery_before, sqlquery, view_res, eng_answer, provenance_ids, retrieval_res = self.posttext.query(query)
        
        # print(provenance_ids)
        provenance_ids = self.flatten(provenance_ids)

        # sources = [tpl[1][1] for tpl in sources]
        return {"question": query, 
                "view_res": view_res,
                "eng_answer": eng_answer,
                "answer": eng_answer,
                "sources": provenance_ids,
                "sql": sqlquery,
                "sql_before": sqlquery_before}


if __name__ == '__main__':
    # engine = ViewEngine("public/digital_data/")
    engine = ViewEngine("personal-data/app_data/")
    # print(engine.query("How many cities did I visit when I travel to Japan?"))

langchain.cache.InMemoryCache langchain.llms.OpenAI langchain.PromptTemplate
[((847, 862), 'langchain.cache.InMemoryCache', 'InMemoryCache', ([], {}), '()\n', (860, 862), False, 'from langchain.cache import InMemoryCache\n'), ((951, 986), 'configparser.ConfigParser', 'ConfigParser', ([], {'comment_prefixes': 'None'}), '(comment_prefixes=None)\n', (963, 986), False, 'from configparser import ConfigParser\n'), ((1065, 1087), 'posttext.src.posttext.PostText', 'PostText', (['config', 'path'], {}), '(config, path)\n', (1073, 1087), False, 'from posttext.src.posttext import PostText\n'), ((1107, 1115), 'langchain.llms.OpenAI', 'OpenAI', ([], {}), '()\n', (1113, 1115), False, 'from langchain.llms import OpenAI\n'), ((1326, 1397), 'langchain.PromptTemplate', 'PromptTemplate', ([], {'input_variables': "['query', 'context']", 'template': 'template'}), "(input_variables=['query', 'context'], template=template)\n", (1340, 1397), False, 'from langchain import PromptTemplate\n'), ((1007, 1039), 'os.path.join', 'os.path.join', (['path', '"""config.ini"""'], {}), "(path, 'config.ini')\n", (1019, 1039), False, 'import os\n')]