Update app.py

app.py

@@ -1,47 +1,129 @@
 import chainlit as cl
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.document_loaders.csv_loader import CSVLoader
-from langchain.embeddings import CacheBackedEmbeddings
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.vectorstores import FAISS
-from langchain.chains import RetrievalQA
-from langchain.chat_models import ChatOpenAI
-from langchain.storage import LocalFileStore
-from langchain.prompts.chat import (
-    ChatPromptTemplate,
-    SystemMessagePromptTemplate,
-    HumanMessagePromptTemplate,
+from llama_index import ServiceContext
+from llama_index.node_parser.simple import SimpleNodeParser
+from llama_index.langchain_helpers.text_splitter import TokenTextSplitter
+from llama_index.llms import OpenAI
+from llama_index.embeddings.openai import OpenAIEmbedding
+from llama_index import VectorStoreIndex
+from llama_index.vector_stores import ChromaVectorStore
+from llama_index.storage.storage_context import StorageContext
+import chromadb
+from llama_index.readers.wikipedia import WikipediaReader
+from llama_index.tools import FunctionTool
+from llama_index.vector_stores.types import (
+    VectorStoreInfo,
+    MetadataInfo,
+    ExactMatchFilter,
+    MetadataFilters,
 )
-import chainlit as cl
+from llama_index.retrievers import VectorIndexRetriever
+from llama_index.query_engine import RetrieverQueryEngine
+
+from typing import List, Tuple, Any
+from pydantic import BaseModel, Field
+from llama_index.agent import OpenAIAgent
+
 
 text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
 
-system_template = """
-Use the following pieces of context to answer the user's question.
-Please respond as if you were Ken from the movie Barbie. Ken is a well-meaning but naive character who loves to Beach. He talks like a typical Californian Beach Bro, but he doesn't use the word "Dude" so much.
-If you don't know the answer, just say that you don't know, don't try to make up an answer.
-You can make inferences based on the context as long as it still faithfully represents the feedback.
+embed_model = OpenAIEmbedding()
+chunk_size = 1000
+llm = OpenAI(
+    temperature=0, 
+    model="gpt-3.5-turbo", 
+    streaming=True
+)
+
+service_context = ServiceContext.from_defaults(
+    llm=llm, 
+    chunk_size=chunk_size, 
+    embed_model=embed_model
+)
+
+text_splitter = TokenTextSplitter(
+    chunk_size=chunk_size
+)
+
+node_parser = SimpleNodeParser(
+    text_splitter=text_splitter
+)
 
-Example of your response should be:
+chroma_client = chromadb.Client()
+chroma_collection = chroma_client.create_collection("wikipedia_barbie_opp")
 
-```
-The answer is foo
-```
+vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
+storage_context = StorageContext.from_defaults(vector_store=vector_store)
+wiki_vector_index = VectorStoreIndex([], storage_context=storage_context, service_context=service_context)
 
-Begin!
-----------------
-{context}"""
+movie_list = ["Barbie (film)", "Oppenheimer (film)"]
 
-messages = [
-    SystemMessagePromptTemplate.from_template(system_template),
-    HumanMessagePromptTemplate.from_template("{question}"),
-]
-prompt = ChatPromptTemplate(messages=messages)
-chain_type_kwargs = {"prompt": prompt}
+wiki_docs = WikipediaReader().load_data(pages=movie_list, auto_suggest=False)
+
+top_k = 3
+vector_store_info = VectorStoreInfo(
+    content_info="semantic information about movies",
+    metadata_info=[MetadataInfo(
+        name="title",
+        type="str",
+        description="title of the movie, one of [Barbie (film), Oppenheimer (film)]",
+    )]
+)
+
+class AutoRetrieveModel(BaseModel):
+    query: str = Field(..., description="natural language query string")
+    filter_key_list: List[str] = Field(
+        ..., description="List of metadata filter field names"
+    )
+    filter_value_list: List[str] = Field(
+        ...,
+        description=(
+            "List of metadata filter field values (corresponding to names specified in filter_key_list)"
+        )
+    )
+
+def auto_retrieve_fn(
+    query: str, filter_key_list: List[str], filter_value_list: List[str]
+):
+    """Auto retrieval function.
+
+    Performs auto-retrieval from a vector database, and then applies a set of filters.
+
+    """
+    query = query or "Query"
+
+    exact_match_filters = [
+        ExactMatchFilter(key=k, value=v)
+        for k, v in zip(filter_key_list, filter_value_list)
+    ]
+    retriever = VectorIndexRetriever(
+        wiki_vector_index, filters=MetadataFilters(filters=exact_match_filters), top_k=top_k
+    )
+    query_engine = RetrieverQueryEngine.from_args(retriever)
+
+    response = query_engine.query(query)
+    return str(response)
+
+description = f"""\
+Use this tool to look up semantic information about films.
+The vector database schema is given below:
+{vector_store_info.json()}
+"""
+
+auto_retrieve_tool = FunctionTool.from_defaults(
+    fn=auto_retrieve_fn,
+    name="auto_retrieve_tool",
+    description=description,
+    fn_schema=AutoRetrieveModel,
+)
+
+
+agent = OpenAIAgent.from_tools(
+    [auto_retrieve_tool], llm=llm, verbose=True
+)
 
 @cl.author_rename
 def rename(orig_author: str):
-    rename_dict = {"RetrievalQA": "Consulting The Kens"}
+    rename_dict = {"RetrievalQA": "Consulting The Llamaindex Tools"}
     return rename_dict.get(orig_author, orig_author)
 
 @cl.on_chat_start
@@ -49,17 +131,13 @@ async def init():
     msg = cl.Message(content=f"Building Index...")
     await msg.send()
 
-    # build FAISS index from csv
-    loader = CSVLoader(file_path="./data/barbie.csv", source_column="Review_Url")
-    data = loader.load()
-    documents = text_splitter.transform_documents(data)
-    store = LocalFileStore("./cache/")
-    core_embeddings_model = OpenAIEmbeddings()
-    embedder = CacheBackedEmbeddings.from_bytes_store(
-        core_embeddings_model, store, namespace=core_embeddings_model.model
-    )
-    # make async docsearch
-    docsearch = await cl.make_async(FAISS.from_documents)(documents, embedder)
+    for movie, wiki_doc in zip(movie_list, wiki_docs):
+    nodes = node_parser.get_nodes_from_documents([wiki_doc])
+    for node in nodes:
+        node.metadata = {'title' : movie}
+    wiki_vector_index.insert_nodes(nodes)
+
+    
 
     chain = RetrievalQA.from_chain_type(
         ChatOpenAI(model="gpt-3.5-turbo", temperature=0, streaming=True),