Create rag.py

rag.py

@@ -0,0 +1,44 @@
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.document_loaders import TextLoader
+
+# We'll assume you have a documentation text file. If not, we can use some sample Python docs.
+# Let's create a sample if the file doesn't exist, or load it.
+
+def load_documents():
+    # Load the documents from a file (or multiple files)
+    # For demonstration, we'll create a sample document if it doesn't exist.
+    doc_path = "python_docs.txt"
+    if not os.path.exists(doc_path):
+        # Create a sample documentation about Python functions
+        with open(doc_path, 'w') as f:
+            f.write("""
+            Functions in Python are defined using the def keyword.
+            For example: def hello_world(): print("Hello, world!")
+            Functions can take parameters and return values.
+            """)
+    loader = TextLoader(doc_path)
+    documents = loader.load()
+    return documents
+
+def create_vector_store(documents):
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
+    texts = text_splitter.split_documents(documents)
+    embeddings = OpenAIEmbeddings()
+    vectorstore = Chroma.from_documents(documents=texts, embedding=embeddings)
+    return vectorstore
+
+def retrieve_relevant_docs(vectorstore, query, k=3):
+    """
+    Retrieve relevant documents for the query.
+    """
+    docs = vectorstore.similarity_search(query, k=k)
+    return "\n".join([doc.page_content for doc in docs])
+
+# Initialize the vector store once (for performance)
+documents = load_documents()
+vectorstore = create_vector_store(documents)
+
+def get_rag_context(query):
+    return retrieve_relevant_docs(vectorstore, query)