refactor: Separate app in modules (#9)

* refactor: Separate app in modules

* enhancement: Create a sidebar

app.py

@@ -1,45 +1,34 @@
 """ A simple example of Streamlit. """
-from datetime import datetime as Date
 import textwrap
 import os
 import tiktoken
-import chromadb
-from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
 import fitz
 import streamlit as st
 import openai
 from dotenv import load_dotenv
-from openai import OpenAI
+from src.chroma_client import ChromaDB
+import src.gui_messages as gm
+from src import settings
 
 load_dotenv()
 
-chroma_client = chromadb.PersistentClient(path="tmp/chroma")
-chroma_client.heartbeat()
-
-
-def api_message(api_key):
-    """Inform if the api key is set."""
-    if api_key is None:
-        return st.warning("Add your OpenAI API key")
-
-    return st.success("Your API key is setup ")
-
 
 def set_api_key():
     """Set the OpenAI API key."""
     openai.api_key = st.session_state.api_key
-    st.session_state.api_message = api_message(openai.api_key)
+    st.session_state.api_message = gm.api_message(openai.api_key)
 
 
 openai.api_key = os.getenv("OPENAI_API_KEY")
 
 if "api_message" not in st.session_state:
-    st.session_state.api_message = api_message(openai.api_key)
+    st.session_state.api_message = gm.api_message(openai.api_key)
 
-if os.getenv("OPENAI_API_KEY") is None:
-    message = st.session_state.api_message
+# Sidebar
+with st.sidebar:
+    st.write("## OpenAI API key")
     openai.api_key = st.text_input(
-        "Enter your OpenAI API key",
+        "Enter OpenAI API key",
         value="",
         type="password",
         key="api_key",
@@ -47,23 +36,13 @@ if os.getenv("OPENAI_API_KEY") is None:
         on_change=set_api_key,
         label_visibility="collapsed",
     )
-    st.write("You can find your API key at https://beta.openai.com/account/api-keys")
-    client = OpenAI(api_key=openai.api_key)
-    embedding_function = OpenAIEmbeddingFunction(
-        api_key=openai.api_key, model_name="text-embedding-ada-002"
-    )
-    collection = chroma_client.get_or_create_collection(
-        name="pdf-explainer", embedding_function=embedding_function
-    )
-else:
-    client = OpenAI()
-    embedding_function = OpenAIEmbeddingFunction(
-        api_key=openai.api_key, model_name="text-embedding-ada-002"
-    )
-    collection = chroma_client.get_or_create_collection(
-        name="pdf-explainer", embedding_function=embedding_function
+    st.write(
+        "You can find your API key at https://platform.openai.com/account/api-keys"
     )
 
+# Build settings
+chroma_db = ChromaDB(openai.api_key)
+openai_client, collection = settings.build(chroma_db)
 
 # Query ChromaDb
 query = st.text_input(
@@ -90,14 +69,10 @@ pdf = st.file_uploader("Upload a file", type="pdf")
 
 if pdf is not None:
     with fitz.open(stream=pdf.read(), filetype="pdf") as doc:  # open document
-        text = chr(12).join([page.get_text() for page in doc])
-        st.write(text[0:200])
-        if st.button("Add to collection"):
-            collection.add(
-                documents=[text],
-                metadatas=[{"source": pdf.name}],
-                ids=[pdf.name + str(Date.now())],
-            )
+        with st.spinner("Extracting text..."):
+            text = chr(12).join([page.get_text() for page in doc])
+        st.subheader("Text preview")
+        st.write(text[0:300] + "...")
         if st.button("Save chunks"):
             with st.spinner("Saving chunks..."):
                 chunks = textwrap.wrap(text, 3000)
@@ -105,7 +80,7 @@ if pdf is not None:
                     encoding = tiktoken.get_encoding("cl100k_base")
                     num_tokens = len(encoding.encode(chunk))
                     response = (
-                        client.embeddings.create(
+                        openai_client.embeddings.create(
                             input=chunk, model="text-embedding-ada-002"
                         )
                         .data[0]
@@ -125,6 +100,6 @@ if st.button("Chroma data collection"):
 
 if st.button("Delete Chroma Collection"):
     try:
-        chroma_client.delete_collection(collection.name)
+        chroma_db.client.delete_collection(collection.name)
     except AttributeError:
         st.error("Collection erased.")

src/chroma_client.py

@@ -0,0 +1,35 @@
+"""A client for ChromaDB."""
+import chromadb
+from chromadb.utils.embedding_functions import OpenAIEmbeddingFunction
+import streamlit as st
+
+
+class ChromaDB:
+    """A class for creating a client for ChromaDB."""
+
+    def __init__(self, api_key, path="tmp/chroma"):
+        """Initialize the client."""
+        self.client = chromadb.PersistentClient(path=path)
+        self.api_key = api_key
+        self.client.heartbeat()
+
+    def get_collection(self, name):
+        """Get a Chroma collection."""
+        try:
+            collection = self.client.get_collection(name=name)
+            return collection
+        except AttributeError:
+            return st.error("An error ocurred while getting the collection.")
+
+    def create_collection(self, name):
+        """Create a Chroma collection."""
+        try:
+            embedding_function = OpenAIEmbeddingFunction(
+                api_key=self.api_key, model_name="text-embedding-ada-002"
+            )
+            collection = self.client.get_or_create_collection(
+                name=name, embedding_function=embedding_function
+            )
+            return collection
+        except AttributeError:
+            return st.error("An error ocurred while creating the collection.")

src/gui_messages.py

@@ -0,0 +1,10 @@
+"""Streamlit GUI messages."""
+import streamlit as st
+
+
+def api_message(api_key):
+    """Inform if the api key is set."""
+    if api_key is None:
+        return st.warning("Add your OpenAI API key")
+
+    return st.success("Your API key is setup ")

src/openai_client.py

@@ -0,0 +1,13 @@
+"""OpenAI client creator."""
+import os
+from openai import OpenAI
+
+
+def create_client(api_key=None):
+    """Create an OpenAI client."""
+    if os.getenv("OPENAI_API_KEY"):
+        api_key = os.getenv("OPENAI_API_KEY")
+
+    client = OpenAI(api_key=api_key)
+
+    return client

src/settings.py

@@ -0,0 +1,10 @@
+"""Build settings for the app."""
+from src.openai_client import create_client
+
+
+def build(chroma_db):
+    """Build the app."""
+    openai_client = create_client(chroma_db.api_key)
+    collection = chroma_db.create_collection("pdf-explainer")
+
+    return openai_client, collection