code: import os
  import re
  from uuid import UUID
  from typing import Any, Dict, List, Optional, Union
  
  import asyncio
  import langchain
  import streamlit as st
  
  from langchain.schema import LLMResult
  from langchain.chat_models import ChatOpenAI
  from langchain.agents import Tool
  from langchain.agents import AgentType
  from langchain.llms import OpenAI
  
  from langchain.agents import initialize_agent
  from langchain.memory import ConversationBufferMemory
  from langchain.callbacks.base import AsyncCallbackHandler, BaseCallbackHandler
  
  from streamlit.delta_generator import DeltaGenerator
  
  from chains import doc_retriever
  from chains.conversational_retrieval_over_code import ConversationalRetrievalCodeChain
  from chains.parser import parse_code
  
  langchain.debug = st.secrets["langchain"]["debug"]
  python_script = os.path.join(os.getcwd(), "langchain", "generated_script.py")
  
  
  class AsyncHandler(AsyncCallbackHandler):
      """Async callback handler that can be used to handle callbacks from langchain."""
  
      def __init__(self, message_placeholder: DeltaGenerator) -> None:
          super().__init__()
          self.message_placeholder = message_placeholder
          self.code_block = False
          self.code_extracted = False
          self.full_response = ""
  
      async def on_llm_start(
          self, serialized: Dict[str, Any], prompts: List[str], **kwargs: Any
      ) -> None:
          """Run when chain starts running."""
          message = ""
          for chunk in "⌛Processing".split():
              message += chunk + " "
              await asyncio.sleep(0.05)
              # Add a blinking cursor to simulate typing
              self.message_placeholder.info(message + "▌")
  
      async def on_llm_new_token(
          self,
          token: str,
          *,
          run_id: UUID,
          parent_run_id: Optional[UUID] = None,
          **kwargs: Any,
      ) -> None:
          # Detect if the token is a code block, to print it all at once
          self.full_response += token
  
          if not self.code_extracted:
              if "`" in token and not self.code_block:
                  self.code_block = True
  
              if self.code_block and self.full_response.count("`") == 6:
                  # We have a full code block, print it now
                  code, explanation = parse_code(self.full_response)
                  container = self.message_placeholder.container()
                  ex = container.expander("Code")
                  ex.code(code)
                  container.markdown(explanation)
                  self.code_block = False
                  self.code_extracted = True
  
          if self.code_extracted:
              code, explanation = parse_code(self.full_response)
              container = self.message_placeholder.container()
              ex = container.expander("Code")
              ex.code(code)
              container.markdown(explanation + "▌")
  
      async def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:
          """Run when chain ends running."""
          self.code_extracted = False
          self.full_response = ""
  
  
  class Handler(BaseCallbackHandler):
      """Async callback handler that can be used to handle callbacks from langchain."""
  
      def __init__(self, message_placeholder: DeltaGenerator) -> None:
          super().__init__()
          self.message_placeholder = message_placeholder
          self.code_block = False
          self.code_extracted = False
          self.full_response = ""
  
      def on_chain_start(
          self,
          serialized: Dict[str, Any],
          inputs: Dict[str, Any],
          *,
          run_id: UUID,
          parent_run_id: Union[UUID, None] = None,
          tags: Union[List[str], None] = None,
          **kwargs: Any,
      ) -> Any:
          """Run when chain starts running."""
          return super().on_chain_start(
              serialized,
              inputs,
              run_id=run_id,
              parent_run_id=parent_run_id,
              tags=tags,
              **kwargs,
          )
  
      def on_llm_new_token(
          self,
          token: str,
          *,
          run_id: UUID,
          parent_run_id: Optional[UUID] = None,
          **kwargs: Any,
      ) -> None:
          # Detect if the token is a code block, to print it all at once
          self.full_response += token
  
          if not self.code_extracted:
              if "`" in token and not self.code_block:
                  self.code_block = True
  
              if self.code_block and self.full_response.count("`") >= 6:
                  # We have a full code block, print it now
                  self.message_placeholder.markdown(self.full_response)
                  self.code_block = False
                  self.code_extracted = True
  
          if self.code_extracted:
              message = ""
              code, explain = parse(self.full_response)
  
              if code:
                  message = f"```python\n{code}\n```\n"
              if explain:
                  message += f"{explain}"
  
              if message != "":
                  # Add a blinking cursor to simulate typing
                  self.message_placeholder.markdown(message + "▌")
          return super().on_llm_new_token(
              token, run_id=run_id, parent_run_id=parent_run_id, **kwargs
          )
  
      def on_chain_end(
          self,
          outputs: Dict[str, Any],
          *,
          run_id: UUID,
          parent_run_id: Union[UUID, None] = None,
          **kwargs: Any,
      ) -> Any:
          """Run when chain ends running."""
          self.code_extracted = False
          self.full_response = ""
          return super().on_chain_end(
              outputs, run_id=run_id, parent_run_id=parent_run_id, **kwargs
          )
  
  
  def load_conversation_chain(
      message_placeholder: DeltaGenerator, openai_api_key: str
  ) -> ConversationalRetrievalCodeChain:
      if openai_api_key is None:
          raise ValueError("OpenAI API key is required to load the chain.")
      llm = ChatOpenAI(
          model_name="gpt-3.5-turbo-16k",
          temperature=0,
          openai_api_key=openai_api_key,
          streaming=True,
          callbacks=[Handler(message_placeholder)],
      )
      condense_question_llm = ChatOpenAI(
          model_name="gpt-3.5-turbo", temperature=0, openai_api_key=openai_api_key
      )
      missing_imports_llm = ChatOpenAI(
          model_name="gpt-3.5-turbo",
          temperature=0,
          openai_api_key=openai_api_key,
          verbose=False,
      )
      retriever = doc_retriever.load_streamlit_doc_retriever(
          st.secrets["openai_api_key"],
          chroma_server_host=st.secrets["chroma"]["host"],
          chroma_server_port=st.secrets["chroma"]["port"],
          mode="docker",
      )
      qa_over_streamlit_code = ConversationalRetrievalCodeChain.from_llm(
          llm=llm,
          retriever=retriever,
          condense_question_llm=condense_question_llm,
          return_source_documents=True,
          missing_imports_llm=missing_imports_llm,
          return_revision_request=True,
          verbose=False,
      )
      return qa_over_streamlit_code
  
  
  def load_agent():
      if st.secrets["openai_api_key"] is None:
          st.error("OpenAI API key is missing! Please add it to your secrets.")
          st.stop()
      doc_chain = doc_retriever.load_streamlit_doc_chain(
          OpenAI(
              temperature=0, max_tokens=2000, openai_api_key=st.secrets["openai_api_key"]
          )
      )
  
      tools = [
          Tool(
              name="Streamlit up to date source code",
              func=doc_chain.run,
              description="useful for when you need to answer questions about the "
              "streamlit Python API. Input should be a fully formed question.",
          ),
      ]
  
      model_name = "text-davinci-003"
      memory = ConversationBufferMemory(memory_key="chat_history")
      llm = OpenAI(
          openai_api_key=st.secrets["openai_api_key"],
          max_tokens=2000,
          temperature=0,
          model_name=model_name,
      )
      agent_chain = initialize_agent(
          tools,
          llm,
          agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
          verbose=False,
          memory=memory,
      )
      return agent_chain
  
  
  def load_chat_agent():
      if st.secrets["openai_api_key"] is None:
          st.error("OpenAI API key is missing! Please add it to your secrets.")
          st.stop()
      doc_chain = doc_retriever.load_streamlit_doc_chain(
          OpenAI(
              temperature=0, max_tokens=2000, openai_api_key=st.secrets["openai_api_key"]
          )
      )
  
      tools = [
          Tool(
              name="Streamlit up to date source code",
              func=doc_chain.run,
              description="useful for when you need to answer questions about the streamlit "
              "Python API. Input should be a fully formed question.",
          ),
      ]
  
      model_name = "text-davinci-003"
      memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
      llm = ChatOpenAI(
          openai_api_key=st.secrets["openai_api_key"],
          max_tokens=2000,
          temperature=0,
          model_name=model_name,
      )
      agent_chain = initialize_agent(
          tools,
          llm,
          agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
          verbose=False,
          memory=memory,
      )
      return agent_chain
  
  
  # https://regex101.com/r/fHlyKq/1
  parse_code_regex = r"(```python(.*?)```)?(.*?)$"
  
  
  def parse(output):
      python_code = None
      explain_code = None
      python_code_match = re.search(parse_code_regex, output, re.DOTALL)
      if python_code_match:
          python_code = python_code_match.group(2)
          explain_code = python_code_match.group(3)
          if python_code == "None":
              python_code = None
      return python_code, explain_code
  
apis: langchain.agents.initialize_agent langchain.chat_models.ChatOpenAI langchain.llms.OpenAI langchain.memory.ConversationBufferMemory langchain.agents.Tool
extract_api: [((815, 826), 'os.getcwd', 'os.getcwd', ([], {}), '()\n', (824, 826), False, 'import os\n'), ((6031, 6120), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'temperature': '(0)', 'openai_api_key': 'openai_api_key'}), "(model_name='gpt-3.5-turbo', temperature=0, openai_api_key=\n    openai_api_key)\n", (6041, 6120), False, 'from langchain.chat_models import ChatOpenAI\n'), ((6156, 6260), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'model_name': '"""gpt-3.5-turbo"""', 'temperature': '(0)', 'openai_api_key': 'openai_api_key', 'verbose': '(False)'}), "(model_name='gpt-3.5-turbo', temperature=0, openai_api_key=\n    openai_api_key, verbose=False)\n", (6166, 6260), False, 'from langchain.chat_models import ChatOpenAI\n'), ((6311, 6505), 'chains.doc_retriever.load_streamlit_doc_retriever', 'doc_retriever.load_streamlit_doc_retriever', (["st.secrets['openai_api_key']"], {'chroma_server_host': "st.secrets['chroma']['host']", 'chroma_server_port': "st.secrets['chroma']['port']", 'mode': '"""docker"""'}), "(st.secrets['openai_api_key'],\n    chroma_server_host=st.secrets['chroma']['host'], chroma_server_port=st.\n    secrets['chroma']['port'], mode='docker')\n", (6353, 6505), False, 'from chains import doc_retriever\n'), ((6565, 6811), 'chains.conversational_retrieval_over_code.ConversationalRetrievalCodeChain.from_llm', 'ConversationalRetrievalCodeChain.from_llm', ([], {'llm': 'llm', 'retriever': 'retriever', 'condense_question_llm': 'condense_question_llm', 'return_source_documents': '(True)', 'missing_imports_llm': 'missing_imports_llm', 'return_revision_request': '(True)', 'verbose': '(False)'}), '(llm=llm, retriever=retriever,\n    condense_question_llm=condense_question_llm, return_source_documents=\n    True, missing_imports_llm=missing_imports_llm, return_revision_request=\n    True, verbose=False)\n', (6606, 6811), False, 'from chains.conversational_retrieval_over_code import ConversationalRetrievalCodeChain\n'), ((7573, 7624), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""'}), "(memory_key='chat_history')\n", (7597, 7624), False, 'from langchain.memory import ConversationBufferMemory\n'), ((7635, 7745), 'langchain.llms.OpenAI', 'OpenAI', ([], {'openai_api_key': "st.secrets['openai_api_key']", 'max_tokens': '(2000)', 'temperature': '(0)', 'model_name': 'model_name'}), "(openai_api_key=st.secrets['openai_api_key'], max_tokens=2000,\n    temperature=0, model_name=model_name)\n", (7641, 7745), False, 'from langchain.llms import OpenAI\n'), ((7799, 7912), 'langchain.agents.initialize_agent', 'initialize_agent', (['tools', 'llm'], {'agent': 'AgentType.CONVERSATIONAL_REACT_DESCRIPTION', 'verbose': '(False)', 'memory': 'memory'}), '(tools, llm, agent=AgentType.\n    CONVERSATIONAL_REACT_DESCRIPTION, verbose=False, memory=memory)\n', (7815, 7912), False, 'from langchain.agents import initialize_agent\n'), ((8661, 8734), 'langchain.memory.ConversationBufferMemory', 'ConversationBufferMemory', ([], {'memory_key': '"""chat_history"""', 'return_messages': '(True)'}), "(memory_key='chat_history', return_messages=True)\n", (8685, 8734), False, 'from langchain.memory import ConversationBufferMemory\n'), ((8745, 8859), 'langchain.chat_models.ChatOpenAI', 'ChatOpenAI', ([], {'openai_api_key': "st.secrets['openai_api_key']", 'max_tokens': '(2000)', 'temperature': '(0)', 'model_name': 'model_name'}), "(openai_api_key=st.secrets['openai_api_key'], max_tokens=2000,\n    temperature=0, model_name=model_name)\n", (8755, 8859), False, 'from langchain.chat_models import ChatOpenAI\n'), ((8913, 9031), 'langchain.agents.initialize_agent', 'initialize_agent', (['tools', 'llm'], {'agent': 'AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION', 'verbose': '(False)', 'memory': 'memory'}), '(tools, llm, agent=AgentType.\n    CHAT_CONVERSATIONAL_REACT_DESCRIPTION, verbose=False, memory=memory)\n', (8929, 9031), False, 'from langchain.agents import initialize_agent\n'), ((9274, 9320), 're.search', 're.search', (['parse_code_regex', 'output', 're.DOTALL'], {}), '(parse_code_regex, output, re.DOTALL)\n', (9283, 9320), False, 'import re\n'), ((6968, 7037), 'streamlit.error', 'st.error', (['"""OpenAI API key is missing! Please add it to your secrets."""'], {}), "('OpenAI API key is missing! Please add it to your secrets.')\n", (6976, 7037), True, 'import streamlit as st\n'), ((7046, 7055), 'streamlit.stop', 'st.stop', ([], {}), '()\n', (7053, 7055), True, 'import streamlit as st\n'), ((7120, 7208), 'langchain.llms.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'max_tokens': '(2000)', 'openai_api_key': "st.secrets['openai_api_key']"}), "(temperature=0, max_tokens=2000, openai_api_key=st.secrets[\n    'openai_api_key'])\n", (7126, 7208), False, 'from langchain.llms import OpenAI\n'), ((7255, 7467), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Streamlit up to date source code"""', 'func': 'doc_chain.run', 'description': '"""useful for when you need to answer questions about the streamlit Python API. Input should be a fully formed question."""'}), "(name='Streamlit up to date source code', func=doc_chain.run,\n    description=\n    'useful for when you need to answer questions about the streamlit Python API. Input should be a fully formed question.'\n    )\n", (7259, 7467), False, 'from langchain.agents import Tool\n'), ((8056, 8125), 'streamlit.error', 'st.error', (['"""OpenAI API key is missing! Please add it to your secrets."""'], {}), "('OpenAI API key is missing! Please add it to your secrets.')\n", (8064, 8125), True, 'import streamlit as st\n'), ((8134, 8143), 'streamlit.stop', 'st.stop', ([], {}), '()\n', (8141, 8143), True, 'import streamlit as st\n'), ((8208, 8296), 'langchain.llms.OpenAI', 'OpenAI', ([], {'temperature': '(0)', 'max_tokens': '(2000)', 'openai_api_key': "st.secrets['openai_api_key']"}), "(temperature=0, max_tokens=2000, openai_api_key=st.secrets[\n    'openai_api_key'])\n", (8214, 8296), False, 'from langchain.llms import OpenAI\n'), ((8343, 8555), 'langchain.agents.Tool', 'Tool', ([], {'name': '"""Streamlit up to date source code"""', 'func': 'doc_chain.run', 'description': '"""useful for when you need to answer questions about the streamlit Python API. Input should be a fully formed question."""'}), "(name='Streamlit up to date source code', func=doc_chain.run,\n    description=\n    'useful for when you need to answer questions about the streamlit Python API. Input should be a fully formed question.'\n    )\n", (8347, 8555), False, 'from langchain.agents import Tool\n'), ((2617, 2647), 'chains.parser.parse_code', 'parse_code', (['self.full_response'], {}), '(self.full_response)\n', (2627, 2647), False, 'from chains.parser import parse_code\n'), ((1529, 1548), 'asyncio.sleep', 'asyncio.sleep', (['(0.05)'], {}), '(0.05)\n', (1542, 1548), False, 'import asyncio\n'), ((2247, 2277), 'chains.parser.parse_code', 'parse_code', (['self.full_response'], {}), '(self.full_response)\n', (2257, 2277), False, 'from chains.parser import parse_code\n')]