Update tools.py

tools.py

@@ -10,64 +10,65 @@ from io import BytesIO
 import chainlit as cl
 
 
-def get_image_name():
-    """
-    We need to keep track of images we generate, so we can reference them later
-    and display them correctly to our users.
-    """
-    image_count = cl.user_session.get("image_count")
-    if image_count is None:
-        image_count = 0
-    else:
-        image_count += 1
-
-    cl.user_session.set("image_count", image_count)
-
-    return f"image-{image_count}"
-
-
-def _generate_image(prompt: str):
-    """
-    This function is used to generate an image from a text prompt using
-    DALL-E 3.
-
-    We use the OpenAI API to generate the image, and then store it in our
-    user session so we can reference it later.
-    """
-    client = OpenAI()
-
-    response = client.images.generate(
-        model="dall-e-3",
-        prompt=prompt,
-        size="1024x1024",
-        quality="standard",
-        n=1,
+import os
+import openai
+from langchain.chat_models import ChatOpenAI
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.vectorstores import Chroma
+from langchain.chains.question_answering import load_qa_chain
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.document_loaders import UnstructuredPDFLoader
+
+# OpenAI API Key Setup
+openai.api_key = os.environ["OPENAI_API_KEY"]
+
+# Define our RAG tool function
+def rag(query):
+    # Load The Goal PDF
+    loader = UnstructuredPDFLoader("data/The Goal - A Process of Ongoing Improvement (Third Revised Edition).pdf") # , mode="elements"
+    docs = loader.load()
+    
+    # Split Text Chunks
+    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
+    splits = text_splitter.split_documents(docs)
+    
+    # Embed Chunks into Chroma Vector Store
+    vectorstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings())
+    retriever = vectorstore.as_retriever()
+    
+    # Use RAG Prompt Template
+    prompt = hub.pull("rlm/rag-prompt")
+    llm = ChatOpenAI(model_name="gpt-4-1106-preview", temperature=0) # or gpt-3.5-turbo
+    
+    
+    def format_docs(docs):
+        return "\n\n".join(doc.page_content for doc in docs)
+    
+    
+    rag_chain = (
+        {"context": retriever | format_docs, "question": RunnablePassthrough()}
+        | prompt
+        | llm
+        | StrOutputParser()
     )
 
-    image_payload = requests.get(response.data[0].url, stream=True)
-
-    image_bytes = BytesIO(image_payload.content)
-
-    print(type(image_bytes))
-
-    name = get_image_name()
-    cl.user_session.set(name, image_bytes.getvalue())
-    cl.user_session.set("generated_image", name)
-    return name
-
+    response = ""
+    for chunk in rag_chain.stream(query): #e.g. "What is a Bottleneck Constraint?"
+        cl.user_session(chunk, end="", flush=True)
+        response += f"\n{chunk}"
+    
+    # rag_chain.invoke("What is a Bottleneck Constraint?")
 
-def generate_image(prompt: str):
-    image_name = _generate_image(prompt)
-    return f"Here is {image_name}."
+    return response
 
 
-# this is our tool - which is what allows our agent to generate images in the first place!
+# this is our tool - which is what allows our agent to access RAG agent
 # the `description` field is of utmost imporance as it is what the LLM "brain" uses to determine
 # which tool to use for a given input.
-generate_image_format = '{{"prompt": "prompt"}}'
-generate_image_tool = Tool.from_function(
-    func=generate_image,
-    name="GenerateImage",
-    description=f"Useful to create an image from a text prompt. Input should be a single string strictly in the following JSON format: {generate_image_format}",
+rag_format = '{{"prompt": "prompt"}}'
+rag_tool = Tool.from_function(
+    func=rag,
+    name="RAG",
+    description=f"Useful for retrieving contextual information about the PDF to answer user questions. Input should be a single string strictly in the following JSON format: {generate_image_format}",
     return_direct=True,
 )