# Copyright 2023 Lei Zhang
#
# 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
from typing import List

from langchain.agents import initialize_agent, AgentType
from langchain.chains import RetrievalQA
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import TextLoader, WebBaseLoader
from langchain.embeddings import OpenAIEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.tools import Tool
from langchain.vectorstores import Chroma

from langchain_plantuml import diagram
from langchain_plantuml.core.plantuml_callback_handler import (
    BasePlantUMLCallbackHandler,
)
from dotenv import load_dotenv

load_dotenv()


# Define an Agent
class MyAgent:
    def __init__(self):
        llm = ChatOpenAI(model_name="gpt-3.5-turbo-0613")

        """Create the state_of_union Vectorstore"""
        current_path = os.path.abspath(os.path.dirname(__file__))
        doc_path = os.path.join(current_path, "state_of_the_union.txt")
        loader = TextLoader(doc_path)
        documents = loader.load()
        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
        texts = text_splitter.split_documents(documents)
        embeddings = OpenAIEmbeddings()
        docsearch = Chroma.from_documents(
            texts, embeddings, collection_name="state-of-union"
        )
        state_of_union = RetrievalQA.from_chain_type(
            llm=llm, chain_type="stuff", retriever=docsearch.as_retriever()
        )

        """Create the ruff Vectorstore"""
        loader = WebBaseLoader("https://beta.ruff.rs/docs/faq/")
        docs = loader.load()
        ruff_texts = text_splitter.split_documents(docs)
        ruff_db = Chroma.from_documents(ruff_texts, embeddings, collection_name="ruff")
        ruff = RetrievalQA.from_chain_type(
            llm=llm, chain_type="stuff", retriever=ruff_db.as_retriever()
        )

        """Create the Agent"""
        tools = [
            Tool(
                name="State of Union QA System",
                func=state_of_union.run,
                description="useful for when you need to answer questions about the most recent state of the union address. Input should be a fully formed question.",
            ),
            Tool(
                name="Ruff QA System",
                func=ruff.run,
                description="useful for when you need to answer questions about ruff (a python linter). Input should be a fully formed question.",
            ),
        ]
        self.agent = initialize_agent(
            tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION
        )

    def run(self, question: str, callbacks: List[BasePlantUMLCallbackHandler]):
        self.agent.run(question, callbacks=callbacks)


# Run the Agent
agent = MyAgent()
activity_diagram = diagram.activity_diagram_callback(note_max_length=2000)
sequence_diagram = diagram.sequence_diagram_callback(note_max_length=2000)
question = "What did biden say about ketanji brown jackson in the state of the union address?"
try:
    agent.run(question=question, callbacks=[activity_diagram, sequence_diagram])
finally:
    activity_diagram.save_uml_content("example_2_activity-plantuml.puml")
    sequence_diagram.save_uml_content("example_2_sequence-plantuml.puml")

langchain.text_splitter.CharacterTextSplitter langchain.agents.initialize_agent langchain_plantuml.diagram.activity_diagram_callback langchain.document_loaders.TextLoader langchain.tools.Tool langchain.chat_models.ChatOpenAI langchain_plantuml.diagram.sequence_diagram_callback langchain.document_loaders.WebBaseLoader langchain.vectorstores.Chroma.from_documents langchain.embeddings.OpenAIEmbeddings
[((1171, 1184), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (1182, 1184), False, 'from dotenv import load_dotenv\n'), ((3316, 3371), 'langchain_plantuml.diagram.activity_diagram_callback', 'diagram.activity_diagram_callback', ([], {'note_max_length': '(2000)'}), '(note_max_length=2000)\n', (3349, 3371), False, 'from langchain_plantuml import diagram\n'), ((3391, 3446), 'langchain_plantuml.diagram.sequence_diagram_callback', 'diagram.sequence_diagram_callback', ([], {'note_max_length': '(2000)'}), '(note_max_length=2000)\n', (3424, 3446), False, 'from langchain_plantuml import diagram\n'), ((1258, 1301), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo-0613"""'}), "(model_name='gpt-3.5-turbo-0613')\n", (1268, 1301), False, 'from langchain.chat_models import ChatOpenAI\n'), ((1440, 1492), 'os.path.join', 'os.path.join', (['current_path', '"""state_of_the_union.txt"""'], {}), "(current_path, 'state_of_the_union.txt')\n", (1452, 1492), False, 'import os\n'), ((1510, 1530), 'langchain.document_loaders.TextLoader', 'TextLoader', (['doc_path'], {}), '(doc_path)\n', (1520, 1530), False, 'from langchain.document_loaders import TextLoader, WebBaseLoader\n'), ((1589, 1644), 'langchain.text_splitter.CharacterTextSplitter', 'CharacterTextSplitter', ([], {'chunk_size': '(1000)', 'chunk_overlap': '(0)'}), '(chunk_size=1000, chunk_overlap=0)\n', (1610, 1644), False, 'from langchain.text_splitter import CharacterTextSplitter\n'), ((1723, 1741), 'langchain.embeddings.OpenAIEmbeddings', 'OpenAIEmbeddings', ([], {}), '()\n', (1739, 1741), False, 'from langchain.embeddings import OpenAIEmbeddings\n'), ((1762, 1836), 'langchain.vectorstores.Chroma.from_documents', 'Chroma.from_documents', (['texts', 'embeddings'], {'collection_name': '"""state-of-union"""'}), "(texts, embeddings, collection_name='state-of-union')\n", (1783, 1836), False, 'from langchain.vectorstores import Chroma\n'), ((2059, 2106), 'langchain.document_loaders.WebBaseLoader', 'WebBaseLoader', (['"""https://beta.ruff.rs/docs/faq/"""'], {}), "('https://beta.ruff.rs/docs/faq/')\n", (2072, 2106), False, 'from langchain.document_loaders import TextLoader, WebBaseLoader\n'), ((2211, 2280), 'langchain.vectorstores.Chroma.from_documents', 'Chroma.from_documents', (['ruff_texts', 'embeddings'], {'collection_name': '"""ruff"""'}), "(ruff_texts, embeddings, collection_name='ruff')\n", (2232, 2280), False, 'from langchain.vectorstores import Chroma\n'), ((3030, 3103), 'langchain.agents.initialize_agent', 'initialize_agent', (['tools', 'llm'], {'agent': 'AgentType.ZERO_SHOT_REACT_DESCRIPTION'}), '(tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION)\n', (3046, 3103), False, 'from langchain.agents import initialize_agent, AgentType\n'), ((1394, 1419), 'os.path.dirname', 'os.path.dirname', (['__file__'], {}), '(__file__)\n', (1409, 1419), False, 'import os\n'), ((2471, 2694), 'langchain.tools.Tool', 'Tool', ([], {'name': '"""State of Union QA System"""', 'func': 'state_of_union.run', 'description': '"""useful for when you need to answer questions about the most recent state of the union address. Input should be a fully formed question."""'}), "(name='State of Union QA System', func=state_of_union.run, description=\n    'useful for when you need to answer questions about the most recent state of the union address. Input should be a fully formed question.'\n    )\n", (2475, 2694), False, 'from langchain.tools import Tool\n'), ((2761, 2944), 'langchain.tools.Tool', 'Tool', ([], {'name': '"""Ruff QA System"""', 'func': 'ruff.run', 'description': '"""useful for when you need to answer questions about ruff (a python linter). Input should be a fully formed question."""'}), "(name='Ruff QA System', func=ruff.run, description=\n    'useful for when you need to answer questions about ruff (a python linter). Input should be a fully formed question.'\n    )\n", (2765, 2944), False, 'from langchain.tools import Tool\n')]