first try to chainlit app with RAQA, WaB and Chains

.gitignore

@@ -3,6 +3,10 @@ __pycache__/
 *.py[cod]
 *$py.class
 
+# project
+cache/
+wandb/
+
 # C extensions
 *.so
 

app.py

@@ -7,30 +7,37 @@ import chainlit as cl  # importing chainlit for our app
 from chainlit.prompt import Prompt, PromptMessage  # importing prompt tools
 from chainlit.playground.providers import ChatOpenAI  # importing ChatOpenAI tools
 from dotenv import load_dotenv
-from aimakerspace.text_utils import PDFFileLoader, CharacterTextSplitter
-from aimakerspace.vectordatabase import VectorDatabase
+import arxiv
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+import pinecone
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.embeddings import CacheBackedEmbeddings
+from langchain.storage import LocalFileStore
+from utils.store import index_documents
+from utils.chain import create_chain
+from langchain.vectorstores import Pinecone
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts import ChatPromptTemplate
+from langchain.prompts import PromptTemplate
+from operator import itemgetter
+from langchain.schema.runnable import RunnableSequence
+from langchain.schema import format_document
+from langchain.schema.output_parser import StrOutputParser
+from langchain.prompts.prompt import PromptTemplate
+from pprint import pprint
+from langchain_core.documents.base import Document
+from langchain_core.vectorstores import VectorStoreRetriever
+import langchain
+from langchain.cache import InMemoryCache
 
 load_dotenv()
-
-# ChatOpenAI Templates
-system_template = """You are a Wizzard and everything you say is a spell!
-"""
-
-user_template = """{input}
-Wizzard, think through your response step by step.
-"""
-
-assistant_template = """Use the following context, if any, to help you
-answer the user's input, if the answer is not in the context say you don't
-know the answer.
-CONTEXT:
-===============
-{context}
-===============
-
-Spell away Wizzard!
-"""
-
+YOUR_API_KEY = os.environ["PINECONE_API_KEY"]
+YOUR_ENV = os.environ["PINECONE_ENV"]
+INDEX_NAME= 'arxiv-paper-index'
+WANDB_API_KEY=os.environ["WANDB_API_KEY"]
+WANDB_PROJECT=os.environ["WANDB_PROJECT"]
+first_run = False
 
 
 @cl.on_chat_start  # marks a function that will be executed at the start of a user session
@@ -38,113 +45,203 @@ async def start_chat():
     settings = {
         "model": "gpt-3.5-turbo",
         "temperature": 0,
-        "max_tokens": 500,
-        "top_p": 1,
-        "frequency_penalty": 0,
-        "presence_penalty": 0,
+        "max_tokens": 500
     }
 
-    cl.user_session.set("settings", settings)
-
-    files = None
-    while files is None:
-        files = await cl.AskFileMessage(
-            content="Please upload a PDF file to begin",
-            accept=["application/pdf"],
-            max_files=10,
-            max_size_mb=10,
-            timeout=60
-        ).send()
-
-    # let the user know you are processing the file(s)
     await cl.Message(
-        content="Loading your files..."
+        content="What would you like to learn about today? 😊"
     ).send()
 
-    # decode the file
-    documents = PDFFileLoader(path="", files=files).load_documents()
+    # instantiate arXiv client (on start)
+    arxiv_client = arxiv.Client()
 
-    # split the text into chunks
-    chunks = CharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=200
-    ).split_texts(documents)
+    # create an embedder through a cache interface (locally) (on start)
+    store = LocalFileStore("./cache/")
 
-    print(chunks[0])
+    core_embeddings_model = OpenAIEmbeddings(
+        api_key=os.environ['OPENAI_API_KEY']
+    )
 
-    # create a vector store
-    # let the user know you are processing the document(s)
-    await cl.Message(
-        content="Creating vector store"
-    ).send()
+    embedder = CacheBackedEmbeddings.from_bytes_store(
+        underlying_embeddings=core_embeddings_model,
+        document_embedding_cache=store,
+        namespace=core_embeddings_model.model
+    )
 
-    vector_db = VectorDatabase()
-    vector_db = await vector_db.abuild_from_list(chunks)
+    # instantiate pinecone (on start)
+    pinecone.init(
+        api_key=YOUR_API_KEY,
+        environment=YOUR_ENV
+    )
 
-    await cl.Message(
-        content="Done. Ask away!"
-    ).send()
+    if INDEX_NAME not in pinecone.list_indexes():
+        pinecone.create_index(
+            name=INDEX_NAME,
+            metric='cosine',
+            dimension=1536
+        )
+    index = pinecone.GRPCIndex(INDEX_NAME)
     
-    cl.user_session.set("vector_db", vector_db)
+    # setup your ChatOpenAI model (on start)
+    llm = ChatOpenAI(
+        model=settings['model'],
+        temperature=settings['temperature'],
+        max_tokens=settings['max_tokens'],
+        api_key=os.environ["OPENAI_API_KEY"],
+        streaming=True
+    )
+
+    # create a prompt cache (locally) (on start)
+    langchain.llm_cache = InMemoryCache()
+    
+    # log data in WaB (on start)
+    os.environ["LANGCHAIN_WANDB_TRACING"] = "true"
+
+    tools = {
+        "arxiv_client": arxiv_client,
+        "index": index,
+        "embedder": embedder,
+        "llm": llm
+    }
+    cl.user_session.set("tools", tools)
+    cl.user_session.set("settings", settings)
+    cl.user_session.set("first_run", False)
 
 
 @cl.on_message  # marks a function that should be run each time the chatbot receives a message from a user
 async def main(message: cl.Message):
-    vector_db = cl.user_session.get("vector_db")
     settings = cl.user_session.get("settings")
-
-    client = AsyncOpenAI()
-
-    print(message.content)
-
-    results_list = vector_db.search_by_text(query_text=message.content, k=3, return_as_text=True)
-    if results_list:
-        results_string = "\n\n".join(results_list)
-    else:
-        results_string = ""
-
-    prompt = Prompt(
-        provider=ChatOpenAI.id,
-        messages=[
-            PromptMessage(
-                role="system",
-                template=system_template,
-                formatted=system_template,
-            ),
-            PromptMessage(
-                role="user",
-                template=user_template,
-                formatted=user_template.format(input=message.content),
-            ),
-            PromptMessage(
-                role="assistant",
-                template=assistant_template,
-                formatted=assistant_template.format(context=results_string)
-            )
-        ],
-        inputs={
-            "input": message.content,
-            "context": results_string
-            },
-        settings=settings,
-    )
-
-    print([m.to_openai() for m in prompt.messages])
-
+    tools = cl.user_session.get("tools")
+    first_run = cl.user_session.get("first_run")
+
+    if not first_run:
+
+        arxiv_client: arxiv.Client  = tools['arxiv_client']
+        index: pinecone.GRPCIndex = tools['index']
+        embedder: CacheBackedEmbeddings = tools['embedder']
+        llm: ChatOpenAI = tools['llm']
+
+        # using query search for ArXiv documents (on message)
+        
+        search = arxiv.Search(
+            query = message.content,
+            max_results = 10,
+            sort_by = arxiv.SortCriterion.Relevance
+        )
+        paper_urls = []
+        
+        sys_message = cl.Message(content="")
+        await sys_message.send() # renders a loader
+        for result in arxiv_client.results(search):
+            paper_urls.append(result.pdf_url)
+        sys_message.content = """
+        I found some papers, let me study them real quick to help 
+        you learn, don't worry it'll be a few seconds 😉"""
+        await sys_message.update()
+        await sys_message.send()
+
+        sys_message = cl.Message(content="")
+        await sys_message.send() # renders a loader
+        # load them and split them (on message)
+        docs = []
+        for paper_url in paper_urls:
+            try:
+                loader = PyPDFLoader(paper_url)
+                docs.append(loader.load())
+            except:
+                print(f"Error loading {paper_url}")
+
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size = 400,
+            chunk_overlap = 30,
+            length_function = len
+        )
+
+        # create an index using pinecone (on message)
+        index_documents(docs, text_splitter, embedder, index)
+        sys_message.content = "Done studying :)"
+        await sys_message.update()
+        await sys_message.send()
+
+        text_field = "source_document"
+        index = pinecone.Index(INDEX_NAME)
+        vectorstore = Pinecone(
+            index=index,
+            embedding=embedder.embed_query,
+            text_key=text_field
+        )
+        retriever: VectorStoreRetriever = vectorstore.as_retriever()
+
+        # create the chain (on message)
+        retrieval_augmented_qa_chain: RunnableSequence = create_chain(retriever=retriever, llm=llm)
+        
+        # message.content = await cl.AskUserMessage(
+        #     content="Ask away"
+        # ).send()
+
+    # run
     msg = cl.Message(content="")
-
-    # Call OpenAI
-    async for stream_resp in await client.chat.completions.create(
-        messages=[m.to_openai() for m in prompt.messages], stream=True, **settings
-    ):
-        token = stream_resp.choices[0].delta.content
-        if not token:
-            token = ""
-        await msg.stream_token(token)
-
-    # Update the prompt object with the completion
-    prompt.completion = msg.content
-    msg.prompt = prompt
-
-    # Send and close the message stream
+    for chunk in retrieval_augmented_qa_chain.stream({"question": f"{message.content}"}):
+        pprint(chunk)
+        if res:= chunk.get('response'):
+            await msg.stream_token(res.content)
     await msg.send()
+    cl.user_session.set("first_run", True)
+    # first_run = True
+
+    
+    # client = AsyncOpenAI()
+
+    # print(message.content)
+
+    # results_list = vector_db.search_by_text(query_text=message.content, k=3, return_as_text=True)
+    # if results_list:
+    #     results_string = "\n\n".join(results_list)
+    # else:
+    #     results_string = ""
+
+    # prompt = Prompt(
+    #     provider=ChatOpenAI.id,
+    #     messages=[
+    #         PromptMessage(
+    #             role="system",
+    #             template=system_template,
+    #             formatted=system_template,
+    #         ),
+    #         PromptMessage(
+    #             role="user",
+    #             template=user_template,
+    #             formatted=user_template.format(input=message.content),
+    #         ),
+    #         PromptMessage(
+    #             role="assistant",
+    #             template=assistant_template,
+    #             formatted=assistant_template.format(context=results_string)
+    #         )
+    #     ],
+    #     inputs={
+    #         "input": message.content,
+    #         "context": results_string
+    #         },
+    #     settings=settings,
+    # )
+
+    # print([m.to_openai() for m in prompt.messages])
+
+    # msg = cl.Message(content="")
+
+    # # Call OpenAI
+    # async for stream_resp in await client.chat.completions.create(
+    #     messages=[m.to_openai() for m in prompt.messages], stream=True, **settings
+    # ):
+    #     token = stream_resp.choices[0].delta.content
+    #     if not token:
+    #         token = ""
+    #     await msg.stream_token(token)
+
+    # # Update the prompt object with the completion
+    # prompt.completion = msg.content
+    # msg.prompt = prompt
+
+    # # Send and close the message stream
+    # await msg.send()

requirements.txt

@@ -4,8 +4,6 @@ openai==1.3.5
 tiktoken==0.5.1
 python-dotenv==1.0.0
 numpy==1.25.2
-pandas
-scikit-learn
-matplotlib
-plotly
-pdfminer.six
+langchain
+pinecone-client[grpc]
+pypdf

utils/chain.py

@@ -0,0 +1,71 @@
+from operator import itemgetter
+from langchain_core.vectorstores import VectorStoreRetriever
+from langchain.schema.runnable import RunnableLambda, RunnableParallel, RunnableSequence
+from langchain.chat_models import ChatOpenAI
+from langchain.prompts import PromptTemplate
+from langchain_core.documents import Document
+from langchain_core.messages.ai import AIMessage
+
+
+template = """
+You are a helpful assistant, your job is to answer the user's question using the relevant context.
+CONTEXT
+=========
+{context}
+=========
+
+User question: {question}
+"""
+prompt = PromptTemplate.from_template(template=template)
+
+
+def to_doc(input: AIMessage) -> list[Document]:
+    return [Document(page_content="LLM", metadata={'chunk': 1.0, 'page_number': 1.0, 'text':input.content})]
+
+def merge_docs(a: dict[str, list[Document]]) -> list[Document]:
+    merged_docs = []
+    for key,value in a.items():
+        merged_docs.extend(value)
+    return merged_docs
+
+
+
+def create_chain(**kwargs) -> RunnableSequence:
+    """
+    Requires retriever, llm and prompt
+    """
+    
+    retriever: VectorStoreRetriever = kwargs["retriever"]
+    llm:ChatOpenAI = kwargs.get("llm", None)
+    
+
+    if not isinstance(retriever, VectorStoreRetriever):
+        raise ValueError
+    if not isinstance(llm, ChatOpenAI):
+        raise ValueError
+
+    docs_chain = (itemgetter("question") | retriever).with_config(config={"run_name": "docs"})
+    self_knowledge_chain = (itemgetter("question") | llm | to_doc).with_config(config={"run_name": "self knowledge"})
+    response_chain = (prompt | llm).with_config(config={"run_name": "response"})
+    merge_docs_link = RunnableLambda(merge_docs).with_config(config={"run_name": "merge docs"})
+    context_chain = (
+        RunnableParallel(
+            {
+                "docs": docs_chain,
+                "self_knowledge": self_knowledge_chain
+            }
+        ).with_config(config={"run_name": "parallel context"})
+        | merge_docs_link
+    )
+
+    retrieval_augmented_qa_chain = (
+        RunnableParallel({
+            "question": itemgetter("question"),
+            "context": context_chain
+        })
+        | RunnableParallel({
+            "response": response_chain,
+            "context": itemgetter("context"),
+        })
+    )
+    return retrieval_augmented_qa_chain

utils/store.py

@@ -0,0 +1,44 @@
+from tqdm.auto import tqdm
+from langchain.embeddings import CacheBackedEmbeddings
+from uuid import uuid4
+from langchain_core.documents import Document
+from typing import List
+from langchain.text_splitter import TextSplitter
+from pinecone import GRPCIndex
+
+BATCH_LIMIT = 100
+
+def index_documents(
+        docs: List[Document], 
+        text_splitter: TextSplitter, 
+        embedder: CacheBackedEmbeddings, 
+        index: GRPCIndex) -> None:
+    
+    texts = []
+    metadatas = []
+
+    for i in tqdm(range(len(docs))):
+        for doc in docs[i]:
+            metadata = {
+                'source_document' : doc.metadata["source"],
+                'page_number' : doc.metadata["page"]
+            }
+
+            record_texts = text_splitter.split_text(doc.page_content)
+
+            record_metadatas = [{
+                "chunk": j, "text": text, **metadata
+            } for j, text in enumerate(record_texts)]
+            texts.extend(record_texts)
+            metadatas.extend(record_metadatas)
+            if len(texts) >= BATCH_LIMIT:
+                ids = [str(uuid4()) for _ in range(len(texts))]
+                embeds = embedder.embed_documents(texts)
+                index.upsert(vectors=zip(ids, embeds, metadatas))
+                texts = []
+                metadatas = []
+
+    if len(texts) > 0:
+        ids = [str(uuid4()) for _ in range(len(texts))]
+        embeds = embedder.embed_documents(texts)
+        index.upsert(vectors=zip(ids, embeds, metadatas))