import os
import langchain
from langchain import (
    agents,
    prompts,
    chains,
    llms
)


class BOAgent:
    def __init__(
            self,
            tools,
            memory,
            model="text-davinci-003",
            temp=0.1,
            max_steps=30,
    ):
        self.openai_key = os.getenv("OPENAI_API_KEY")
        self.memory = memory

        # Initialize LLM
        if model.startswith("gpt-3.5-turbo") or model.startswith("gpt-4"):
            self.llm = langchain.chat_models.ChatOpenAI(
                temperature=temp,
                openai_api_key=self.openai_key,
                model_name=model,
            )
        else:
            self.llm = langchain.OpenAI(
                temperature=temp,
                openai_api_key=self.openai_key,
                model_name=model
            )

        # Initialize agent
        self.agent = agents.initialize_agent(
            tools,
            self.llm,
            agent="conversational-react-description",
            verbose=True,
            max_iterations=max_steps,
            memory=self.memory
        )

    def run(self, prompt):
        return self.agent.run(prompt)

langchain.agents.initialize_agent langchain.OpenAI langchain.chat_models.ChatOpenAI
[((310, 337), 'os.getenv', 'os.getenv', (['"""OPENAI_API_KEY"""'], {}), "('OPENAI_API_KEY')\n", (319, 337), False, 'import os\n'), ((888, 1040), 'langchain.agents.initialize_agent', 'agents.initialize_agent', (['tools', 'self.llm'], {'agent': '"""conversational-react-description"""', 'verbose': '(True)', 'max_iterations': 'max_steps', 'memory': 'self.memory'}), "(tools, self.llm, agent=\n    'conversational-react-description', verbose=True, max_iterations=\n    max_steps, memory=self.memory)\n", (911, 1040), False, 'from langchain import agents, prompts, chains, llms\n'), ((491, 596), 'langchain.chat_models.ChatOpenAI', 'langchain.chat_models.ChatOpenAI', ([], {'temperature': 'temp', 'openai_api_key': 'self.openai_key', 'model_name': 'model'}), '(temperature=temp, openai_api_key=self.\n    openai_key, model_name=model)\n', (523, 596), False, 'import langchain\n'), ((692, 780), 'langchain.OpenAI', 'langchain.OpenAI', ([], {'temperature': 'temp', 'openai_api_key': 'self.openai_key', 'model_name': 'model'}), '(temperature=temp, openai_api_key=self.openai_key,\n    model_name=model)\n', (708, 780), False, 'import langchain\n')]