import os
import json
from typing import List
from dotenv import load_dotenv
from pydantic import BaseModel, Field
from supabase.client import Client, create_client
from langchain.chat_models import ChatOpenAI
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.tools import StructuredTool
from langchain.chains.openai_functions import create_structured_output_chain
from langchain.prompts import (
    ChatPromptTemplate,
    HumanMessagePromptTemplate,
    SystemMessagePromptTemplate,
)
import langchain

load_dotenv()

# Set debug to True to see A LOT of details of langchain's inner workings
# langchain.debug = True

# The name of the table in Supabase, where the vectors are stored
matchVectorFunctionName = "match_embeddings"

# Create the supabase client
SUPABASE_URL = os.getenv("SUPABASE_URL")
SUPABASE_KEY = os.getenv("SUPABASE_KEY")
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)


class ToolInputSchema(BaseModel):
    question: str = Field(..., description="A fully formed question.")


class KnowledgeAnswer(BaseModel):
    answer: str = Field(..., description="The answer to the question.")
    sources: List[str] = Field(
        ...,
        description="The sources which contributed to the answer.",
    )


llm = ChatOpenAI(model_name="gpt-3.5-turbo-16k", temperature=0.3)

prompt_msgs = [
    SystemMessagePromptTemplate.from_template(
        """You're an elite algorithm, answering queries based solely on given context. If the context lacks the answer, state ignorance. If you are not 100% sure tell the user.

        Context:
        {context}"""
    ),
    HumanMessagePromptTemplate.from_template("{question}"),
]
prompt = ChatPromptTemplate.from_messages(prompt_msgs)

chain = create_structured_output_chain(KnowledgeAnswer, llm, prompt)


def get_answer(question: str) -> str:
    try:
        vectors = OpenAIEmbeddings().embed_documents([question])
        embeddings = supabase.rpc(
            matchVectorFunctionName, dict(query_embedding=vectors[0], match_count=7)
        ).execute()

        print(f"⚡ Retrieved {len(embeddings.data)} vectors from Supabase:")
        for entry in embeddings.data:
            print("🔖 Title:", entry["metadata"]["title"])
            print("🌐 Source:", entry["metadata"]["source"])
            print("📊 Similarity:", entry["similarity"])
            print("📄 Content:", entry["content"].replace("\n", " ")[:100] + "...")
            print("-" * 50)

        result = chain.run(context=json.dumps(embeddings.data), question=question)
        print("📝 Result of knowledge extraction chain:", result)

        return f"""Answer: {result.answer}
        Sources: {json.dumps(result.sources)}
        """

    except Exception as e:
        print(e)
        return "The wiki knowledgebase is currently not available. We are working on it. Tell the user to use the wiki directly. https://www.defichainwiki.com/"


description = """Use this if you need to answer any question about DeFiChain which does not require live-data. Make sure to include the source of the answer in your response."""

wikiTool = StructuredTool(
    name="defichain_wiki_knowledge",
    description=description,
    func=get_answer,
    args_schema=ToolInputSchema,
)


if __name__ == "__main__":
    while True:
        question = input(
            "Ask something, that can be answered using information from DeFiChainWiki: "
        )
        print("✅", get_answer(question))

langchain.chains.openai_functions.create_structured_output_chain langchain.tools.StructuredTool langchain.prompts.HumanMessagePromptTemplate.from_template langchain.chat_models.ChatOpenAI langchain.prompts.ChatPromptTemplate.from_messages langchain.prompts.SystemMessagePromptTemplate.from_template langchain.embeddings.openai.OpenAIEmbeddings
[((528, 541), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (539, 541), False, 'from dotenv import load_dotenv\n'), ((799, 824), 'os.getenv', 'os.getenv', (['"""SUPABASE_URL"""'], {}), "('SUPABASE_URL')\n", (808, 824), False, 'import os\n'), ((840, 865), 'os.getenv', 'os.getenv', (['"""SUPABASE_KEY"""'], {}), "('SUPABASE_KEY')\n", (849, 865), False, 'import os\n'), ((885, 926), 'supabase.client.create_client', 'create_client', (['SUPABASE_URL', 'SUPABASE_KEY'], {}), '(SUPABASE_URL, SUPABASE_KEY)\n', (898, 926), False, 'from supabase.client import Client, create_client\n'), ((1269, 1328), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo-16k"""', 'temperature': '(0.3)'}), "(model_name='gpt-3.5-turbo-16k', temperature=0.3)\n", (1279, 1328), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1687, 1732), 'langchain.prompts.ChatPromptTemplate.from_messages', 'ChatPromptTemplate.from_messages', (['prompt_msgs'], {}), '(prompt_msgs)\n', (1719, 1732), False, 'from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate\n'), ((1742, 1802), 'langchain.chains.openai_functions.create_structured_output_chain', 'create_structured_output_chain', (['KnowledgeAnswer', 'llm', 'prompt'], {}), '(KnowledgeAnswer, llm, prompt)\n', (1772, 1802), False, 'from langchain.chains.openai_functions import create_structured_output_chain\n'), ((3106, 3228), 'langchain.tools.StructuredTool', 'StructuredTool', ([], {'name': '"""defichain_wiki_knowledge"""', 'description': 'description', 'func': 'get_answer', 'args_schema': 'ToolInputSchema'}), "(name='defichain_wiki_knowledge', description=description,\n    func=get_answer, args_schema=ToolInputSchema)\n", (3120, 3228), False, 'from langchain.tools import StructuredTool\n'), ((983, 1033), 'pydantic.Field', 'Field', (['...'], {'description': '"""A fully formed question."""'}), "(..., description='A fully formed question.')\n", (988, 1033), False, 'from pydantic import BaseModel, Field\n'), ((1088, 1141), 'pydantic.Field', 'Field', (['...'], {'description': '"""The answer to the question."""'}), "(..., description='The answer to the question.')\n", (1093, 1141), False, 'from pydantic import BaseModel, Field\n'), ((1167, 1237), 'pydantic.Field', 'Field', (['...'], {'description': '"""The sources which contributed to the answer."""'}), "(..., description='The sources which contributed to the answer.')\n", (1172, 1237), False, 'from pydantic import BaseModel, Field\n'), ((1350, 1610), 'langchain.prompts.SystemMessagePromptTemplate.from_template', 'SystemMessagePromptTemplate.from_template', (['"""You\'re an elite algorithm, answering queries based solely on given context. If the context lacks the answer, state ignorance. If you are not 100% sure tell the user.\n\n        Context:\n        {context}"""'], {}), '(\n    """You\'re an elite algorithm, answering queries based solely on given context. If the context lacks the answer, state ignorance. If you are not 100% sure tell the user.\n\n        Context:\n        {context}"""\n    )\n', (1391, 1610), False, 'from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate\n'), ((1620, 1674), 'langchain.prompts.HumanMessagePromptTemplate.from_template', 'HumanMessagePromptTemplate.from_template', (['"""{question}"""'], {}), "('{question}')\n", (1660, 1674), False, 'from langchain.prompts import ChatPromptTemplate, HumanMessagePromptTemplate, SystemMessagePromptTemplate\n'), ((1870, 1888), 'langchain.embeddings.openai.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (1886, 1888), False, 'from langchain.embeddings.openai import OpenAIEmbeddings\n'), ((2493, 2520), 'json.dumps', 'json.dumps', (['embeddings.data'], {}), '(embeddings.data)\n', (2503, 2520), False, 'import json\n'), ((2684, 2710), 'json.dumps', 'json.dumps', (['result.sources'], {}), '(result.sources)\n', (2694, 2710), False, 'import json\n')]