import langchain
from dotenv import load_dotenv
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from rmrkl import ChatZeroShotAgent, RetryAgentExecutor

from .prompt import FORMAT_INSTRUCTIONS, QUESTION_PROMPT, SUFFIX
from .tools import make_tools, Doc, Text,search_texts, load_texts

import time
load_dotenv()


def _make_llm(model, temp, verbose):
    if model.startswith("gpt-3.5-turbo") or model.startswith("gpt-4"):
        llm = langchain.chat_models.ChatOpenAI(
            temperature=temp,
            model_name=model,
            request_timeout=1000,
            streaming=True if verbose else False,
            callbacks=[StreamingStdOutCallbackHandler()] if verbose else [None],
        )
    elif model.startswith("text-"):
        llm = langchain.OpenAI(
            temperature=temp,
            model_name=model,
            streaming=True if verbose else False,
            callbacks=[StreamingStdOutCallbackHandler()] if verbose else [None],
        )
    else:
        raise ValueError(f"Invalid model name: {model}")
    return llm


class HVACAgent:
    def __init__(
        self,
        tools=None,
        model="gpt-4",
        tools_model="gpt-3.5-turbo",
        temp=0.1,
        max_iterations=40,
        verbose=True,
    ):
        self.llm = _make_llm(model, temp, verbose)
        if tools is None:
            tools_llm = _make_llm(tools_model, temp, verbose)
            tools = make_tools(tools_llm, verbose=verbose)

        # Initialize agent
        self.agent_executor = RetryAgentExecutor.from_agent_and_tools(
            tools=tools,
            agent=ChatZeroShotAgent.from_llm_and_tools(
                self.llm,
                tools=tools,
                suffix=SUFFIX,
                format_instructions=FORMAT_INSTRUCTIONS,
                question_prompt=QUESTION_PROMPT,
            ),
            verbose=True,
            max_iterations=max_iterations,
            return_intermediate_steps=True,
        )

    def run(self, prompt):
        #wait three seconds
        time.sleep(3)
        outputs = self.agent_executor({"input": prompt})
        # Parse long output (with intermediate steps)
        intermed = outputs["intermediate_steps"]

        final = ""
        for step in intermed:
            final += f"Thought: {step[0].log}\n" f"Observation: {step[1]}\n"
        final += f"Final Answer: {outputs['output']}"

        return final
langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler
[((329, 342), 'dotenv.load_dotenv', 'load_dotenv', ([], {}), '()\n', (340, 342), False, 'from dotenv import load_dotenv\n'), ((2064, 2077), 'time.sleep', 'time.sleep', (['(3)'], {}), '(3)\n', (2074, 2077), False, 'import time\n'), ((1632, 1784), 'rmrkl.ChatZeroShotAgent.from_llm_and_tools', 'ChatZeroShotAgent.from_llm_and_tools', (['self.llm'], {'tools': 'tools', 'suffix': 'SUFFIX', 'format_instructions': 'FORMAT_INSTRUCTIONS', 'question_prompt': 'QUESTION_PROMPT'}), '(self.llm, tools=tools, suffix=SUFFIX,\n    format_instructions=FORMAT_INSTRUCTIONS, question_prompt=QUESTION_PROMPT)\n', (1668, 1784), False, 'from rmrkl import ChatZeroShotAgent, RetryAgentExecutor\n'), ((668, 700), 'langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler', 'StreamingStdOutCallbackHandler', ([], {}), '()\n', (698, 700), False, 'from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n'), ((937, 969), 'langchain.callbacks.streaming_stdout.StreamingStdOutCallbackHandler', 'StreamingStdOutCallbackHandler', ([], {}), '()\n', (967, 969), False, 'from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler\n')]