Generate diagram based on prompt

app.py

@@ -0,0 +1,55 @@
+import os
+import subprocess
+import time
+# import tempfile
+from virtualenv_mgr import VirtualenvManager
+
+import streamlit as st
+from langchain.llms import OpenAI
+from langchain.prompts import PromptTemplate
+from langchain.chains import LLMChain, SequentialChain
+from langchain.memory import ConversationBufferMemory
+
+os.environ['OPENAI_API_KEY'] = st.secrets["openai_api_key"]
+
+st.title("dAIgramGen : The AI-based Diagram Generator")
+
+daig_name = st.text_input("Diagram name")
+daig_name_lower = daig_name.lower()
+daig_name_underscore = daig_name_lower.replace(" ", "_")
+daig_python_filename = daig_name_underscore + ".py"
+daig_png_filename = daig_name_underscore + ".png"
+
+daig_prompt = st.text_area("Prompt for diagrams")
+
+daig_template = PromptTemplate(
+    input_variables = ['topic', 'name'],
+    template = 'Using the Python diagrams library generate the code for {topic} and name the diagram {name}. Just return the code. No explaination of the code needed'
+)
+
+daig_memory = ConversationBufferMemory(input_key='topic', memory_key='daig_history')
+
+llm = OpenAI(temperature=0.9)
+daig_chain = LLMChain(llm=llm, prompt=daig_template, verbose=True, output_key='code', memory=daig_memory)
+
+if daig_prompt and daig_name:
+	code = daig_chain.run({'topic':daig_prompt, 'name': daig_name})
+	st.subheader('Generated Code')
+	st.code(code)
+	venv = VirtualenvManager()
+	venv.create_env()
+	venv.add_dependency("diagrams")
+	
+	if os.path.exists(daig_python_filename):
+		os.remove(daig_python_filename)
+	f = open(daig_python_filename, "w")
+	f.write(code)
+	f.flush()
+	process = venv.run_code(f.name)
+
+	time.sleep(5)
+	if os.path.exists(daig_png_filename):
+		st.image(daig_png_filename)
+    
+	with st.expander('Code History'):
+		st.info(daig_memory.buffer)

virtualenv_mgr.py

@@ -0,0 +1,43 @@
+import logging
+import string
+import random
+import os
+import subprocess
+from virtualenv import cli_run
+
+logger = logging.getLogger(__name__)
+
+RANDOM_NAME_LENGTH = 16
+
+
+class VirtualenvManager:
+    def __init__(self, name: str = "", base_path="/tmp") -> None:
+        if not name:
+            name = ""
+            for _ in range(RANDOM_NAME_LENGTH):
+                population = string.ascii_letters + string.digits
+                char = random.sample(population, k=1)
+                name += char[0]
+        self.name = name
+        self.path = os.path.join(base_path, name)
+        self.python_interpreter = os.path.join(self.path, "bin/python3")
+        self.dependencies = []
+
+    def add_dependency(self, dependency):
+        logger.info("Adding dependency '%s' ", dependency)
+        self.dependencies.append(dependency)
+
+    def create_env(self):
+        logger.info("Creating virtualenv at path '%s' ", self.path)
+        cli_run([self.path], setup_logging=False)
+
+    def run_code(self, fname):
+        logger.info("Installing dependencies")
+        process = subprocess.run(
+            [self.python_interpreter, "-m", "pip", "install"] + self.dependencies,
+            capture_output=True,
+        )
+        process = subprocess.run(
+            [self.python_interpreter, fname], capture_output=True,
+        )
+        return process