From da3385b2abfdb20caca347188bb228b0fb8fb457 Mon Sep 17 00:00:00 2001
From: bilge-ince <bilge.ince@enterprisedb.com>
Date: Wed, 31 Jul 2024 14:44:18 +0100
Subject: [PATCH] parameterize embedding model and remove streamlit library

---
 .env-example            |  1 +
 embedding.py            |  6 +++---
 streamlit/chatgptui.py  | 44 -----------------------------------------
 streamlit/llamaindex.py |  0
 4 files changed, 4 insertions(+), 47 deletions(-)
 delete mode 100644 streamlit/chatgptui.py
 delete mode 100644 streamlit/llamaindex.py

diff --git a/.env-example b/.env-example
index 1c58177..54f98d0 100644
--- a/.env-example
+++ b/.env-example
@@ -6,6 +6,7 @@ DB_HOST=localhost
 DB_PORT=5432
 
 # MODEl
+AIDB_MODEL_NAME=all-MiniLM-L6-v2
 MODEL_NAME=mistralai/Mistral-7B-Instruct-v0.2
 TOKENIZER_NAME=mistralai/Mistral-7B-Instruct-v0.2
 HUGGING_FACE_ACCESS_TOKEN=
diff --git a/embedding.py b/embedding.py
index 7d1328b..01ecdee 100644
--- a/embedding.py
+++ b/embedding.py
@@ -1,4 +1,4 @@
-# importing all the required modules
+import os
 import PyPDF2
 from db import get_connection
 
@@ -6,12 +6,12 @@ def generate_embeddings():
     conn = get_connection()
     cursor = conn.cursor()
     
-    cursor.execute("""
+    cursor.execute(f"""
                     SELECT aidb.create_pg_retriever(
                         'documents_embeddings',
                         'public',
                         'id',
-                        'all-MiniLM-L6-v2',
+                        '{os.getenv("EMBEDDING_MODEL")}',
                         'text',
                         'documents',
                         ARRAY['id', 'doc_fragment'],
diff --git a/streamlit/chatgptui.py b/streamlit/chatgptui.py
deleted file mode 100644
index c560796..0000000
--- a/streamlit/chatgptui.py
+++ /dev/null
@@ -1,44 +0,0 @@
-import streamlit as st
-from openai import OpenAI
-
-st.title("A simple RAG application using Postgres and pgvector")
-
-# Set OpenAI API key from Streamlit secrets
-client = OpenAI (api_key=st.secrets["OPENAI_API_KEY"])
-
-# Set a default model
-if "openai_model" not in st.session_state:
-    st.session_state["openai_model"] = "gpt-3.5-turbo"
-
-# Initialize chat history
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-# Display chat messages from history on app rerun
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# Accept user input
-if prompt := st.chat_input("Ask me anything"):
-    # Add user message to chat history
-    st.session_state.messages.append({"role": "user", "content": prompt})
-    # Display user message in chat message container
-    with st.chat_message("user"):
-        st.markdown(prompt)
-
-    # Display assistant response in chat message container
-    with st.chat_message("assistant"):
-        stream = client.chat.completions.create(
-            model=st.session_state["openai_model"],
-            messages=[
-                {"role": m["role"], "content": m["content"]}
-                for m in st.session_state.messages
-            ],
-            stream=True,
-        )
-        response = st.write_stream(stream)
-    st.session_state.messages.append({"role": "assistant", "content": response})
-
-
-
diff --git a/streamlit/llamaindex.py b/streamlit/llamaindex.py
deleted file mode 100644
index e69de29..0000000