Upload RAG_prompt_enhancer.py

RAG_prompt_enhancer.py

@@ -0,0 +1,66 @@
+# Imports
+from chromadb import Client, Settings
+from langchain.vectorstores import Chroma
+from langchain.embeddings import SentenceTransformerEmbeddings
+import streamlit as st
+import requests
+
+# Vector Store setup
+def init_vector_store():
+    embeddings = SentenceTransformerEmbeddings('paraphrase-MiniLM-L6-v2')
+    client = Client(Settings(
+        persist_directory = "./chroma_db"
+    ))
+    return Chroma(
+        client=client,
+        embeddings=embeddings
+    )
+
+# Document processing
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.document_loaders import TextLoader, PyPDFLoader
+
+def process_documents(file_path):
+    # Determine loader based on file extension
+    loader = TextLoader() if file_path.endswith('.txt') else PyPDFLoader()
+
+    # Load and split documents
+    splitter = RecursiveCharacterTextSplitter(
+        chunk_size = 1000,
+        chunk_overlap = 100
+    )
+
+    docs = loader.load()
+    chunks = splitter.split_documents(docs)
+    return chunks
+
+# Prompt Template Management
+from langchain.prompts import PromptTemplate
+
+class PromptOptimizer:
+    def __init__(self):
+        self.base_template = PromptTemplate(
+            input_variables=["context", "prompt"],
+            template = "Use the following context to enhance the prompt provided."
+            template += "Context: {context}\n"
+            template += "Prompt: {prompt}\n"
+            template += "Generate an enhanced prompt that leverages the context provided\
+            while maintaining the original intent of the prompt."
+        )
+
+    def optimize_prompt(self, context, prompt):
+        return self.base_template.render(context=context, prompt=prompt)
+    
+# Streamlit frontend
+st.title("RAG-based Prompt Enhancer")
+
+# File upload
+uploaded_file = st.file_uploader("Choose a file")
+if uploaded_file:
+    files = {"file": uploaded_file}
+    response = requests.post("http://localhost:8000/upload", files=files)
+
+prompt = st.text_area("Enter a prompt you'd like to enhance:")
+if st.button("Enhance Prompt"):
+    st.write("Enhanced Prompt:")
+    st.write(response.json()["enhanced_prompt"])