refreshed ui, added support for response streaming, and added new database.

.streamlit/secrets.toml

@@ -3,4 +3,5 @@ LANGCHAIN_ENDPOINT="https://api.smith.langchain.com"
 LANGCHAIN_API_KEY="ls__3382b1f40a7f4eefa6959cb2b03dd687"
 LANGCHAIN_PROJECT="ConsultAI v1"
 OPENAI_API_KEY = "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"
-APIFY_CLIENT_KEY = "apify_api_GhFIqZgUf2BGqO46OdBcQOyk2rekQt0ns3Wv"
+APIFY_CLIENT_KEY = "apify_api_GhFIqZgUf2BGqO46OdBcQOyk2rekQt0ns3Wv"
+DEV_DATASET = "https://huggingface.co/datasets/aniruddhdoki/Atticus_DEV"

app.py

@@ -1,142 +1,93 @@
-import os
-import openai
 import streamlit as st
-from streamlit_chat import message
-from utils.db import Database
+import openai
+from supabase import client, create_client
 
-st.title("ATTICUS")
-st.subheader("I know everything")
+openai.api_key="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"
+supabase = create_client(
+  "https://djfytlaeuxwbcztcwrzt.supabase.co",
+  "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6ImRqZnl0bGFldXh3YmN6dGN3cnp0Iiwicm9sZSI6ImFub24iLCJpYXQiOjE3MDA1MTM4NTEsImV4cCI6MjAxNjA4OTg1MX0.G0ni6ZeuU1NxnoLhXyY99agZz6TkufxFc5-LAat2sk4",
+)
+
+st.title("Atticus")
 
 if "messages" not in st.session_state:
-  st.session_state["messages"] = []
+  st.session_state.messages = []
 
 if "history" not in st.session_state:
-  st.session_state["history"] = []
-
-if "input_disabled" not in st.session_state:
-  st.session_state["input_disabled"] = False
-
-
-openai.api_key = st.secrets.get('OPENAI_API_KEY')
-if not openai.api_key:
-  openai_api_key = st.text_input(
-    "Enter your OpenAI API key here:", 
-    disabled=st.session_state["input_disabled"], 
-    type="password"
-  )
-  st.write("Please enter your OpenAI API key above")
-  st.stop()
-
-@st.cache_resource
-def initialize_db():
-  return Database()
-db = initialize_db()
-
-def ingest(uploaded_files):
-  db.add_files(uploaded_files)
-
-def disable_input_cb():
-  st.session_state["input_disabled"] = True
-
-def enable_input_cb():
-  st.session_state["input_disabled"] = False
-
-with st.sidebar:
-  st.write('Upload knowledge here. Can digest .pdf files. May take a while to process')
-  uploaded_files = st.file_uploader(
-    "Upload Files",
-    accept_multiple_files=True, 
-    type=['pdf'],
-    disabled=st.session_state["input_disabled"]
-  )
-
-  if uploaded_files:
-    if st.button("Ingest PDFs", on_click=disable_input_cb):
-      with st.spinner("Analyzing... DON'T INTERACT WITH THE PAGE UNTIL IT'S DONE!"):
-        disable_input_cb()
-        ingest(uploaded_files)
-        st.success("Done!")
-        st.balloons()
-        enable_input_cb()
-        st.experimental_rerun()
-
-  st.subheader('current files in db')
-  for file in db.view_db().get('files', {}).get('documents', []):
-    st.write(file)
-
-def get_text():
-  input_text = st.text_input(
-    "You: ",
-    help="Ask your questions here!",
-    key="input",
-    disabled=st.session_state["input_disabled"]
-  )
-  return input_text
-
-user_input = get_text()
-
-def generate_response(query):
-  search = db.similarity_search(query, 5)
-  sources = """"""
-  print(search)
-  for metadata in search.get('metadatas', [])[0]:
-    if metadata.get('page_num', False):
-      sources += f"\nHarvard CBE Consultants"
-    else:
-      sources += f"\n{metadata['link']}"
-
-  context = '\n'.join([' '.join(document) for document in search['documents']])
-  print(context)
-  template = f"""
-  You are a smart, witty, hipster, know-it-all young generative AI consultant out of Harvard University. You will answer a question given the following context:
-
-  Here's the context: {context}
-
-  Here's the question: {query}
-
-  Now generate an answer to the question using the context provided.
-  """
-  st.session_state["history"].append({
-    'role': 'user', 
-    'content': template.format(context=context, query=query),
-  })
-  st.session_state["messages"].append("You: " + query)
-
-  response = openai.ChatCompletion.create(
-    model='gpt-3.5-turbo-0613',
-    messages=st.session_state["history"]
-  )
-  formatted_response = f"""
-  Assistant: {response.get('choices', [])[0].get('message', {}).get('content', 'Response generated but error displaying. Please try again')}
-
-  Sources: {sources}
+  st.session_state.history = []
+
+for message in st.session_state.history:
+  with st.chat_message(message["role"]):
+    st.markdown(message["content"])
+
+prompt = st.chat_input("I literally know everything about financial AI. Ask me anything.")
+
+# def example_prompt_cb(text):
+#   global prompt
+#   prompt = text
+
+# if not prompt:
+#   st.button(
+#     "What are top financial companies thinking about when it comes to Artificial Intelligence and large language models?",
+#     on_click=example_prompt_cb,
+#     args=["What are top financial companies thinking about when it comes to Artificial Intelligence and large language models?"]
+#   )
+
+
+if prompt:
+  # embed question
+  embedding = openai.embeddings.create(
+    model='text-embedding-ada-002',
+    input=prompt,
+    encoding_format='float'
+  ).data[0].embedding
+
+  # perform similarity search
+  data = supabase.rpc("match_documents", {
+    "match_count": 10, 
+    "query_embedding": embedding
+  }).execute().data
+
+  # construct prompt
+  context = '\n'.join([doc.get('content', '') for doc in data])
+  sources = '\n'.join(set([doc.get('metadata', {}).get('source', '') for doc in data]))
+  query = """
+  You are a smart, witty, hipster, know-it-all young finance and generative AI consultant out of Harvard University. Be snarky, know-it-all, hipster, and above all, be cool. Give as much information as possible and be as helpful as possible. Be detailed. Give as much detail as possible. Produce your answer in MARKDOWN format. Only answer the question asked. Only answer about finance.
+
+  LLMS = Large Language Models
+
+  Answer the following question given the context below. Do not use any information outside of this context. If you don't have the answer in the given context, say you do not know. DO NOT PRODUCE URLS.
+  
+  Here is the question: {prompt}
+
+  Here is the context: {context}
   """
-  return formatted_response
-
-if user_input:
-  res = generate_response(user_input)
-  st.session_state["history"].append({
-    'role': 'assistant',
-    'content': res
-  })
-  st.session_state["messages"].append("Assistant: " + res)
-
-print(st.session_state["messages"])
-
-for msg in st.session_state["messages"]:
-  st.write(msg)
-
-with st.expander("debug db"):
-  if st.button(
-    "empty db. this is very dangerous. pls don't.", 
-    on_click=lambda: disable_input_cb,
-    disabled=st.session_state["input_disabled"],
-  ):
-    db.empty()
-    st.write("emptied db")
-    enable_input_cb()
-    st.experimental_rerun()
-
-  st.write(db.view_db())
-
-db._save_db()
+  query = query.format(prompt=prompt, context=context)
+
+  # generate response
+  st.session_state.history.append({"role": "user", "content": prompt})
+  st.session_state.messages.append({"role": "user", "content": prompt})
+  with st.chat_message("user"):
+    st.markdown(prompt)
+
+  with st.chat_message("assistant"):
+    message_placeholder = st.empty()
+    full_response = ""
+    for response in openai.chat.completions.create(
+      model='gpt-3.5-turbo-16k-0613',
+      messages=[
+          {"role": m["role"], "content": m["content"]}
+          for m in st.session_state.messages
+      ],
+      stream=True,
+    ):
+      full_response += (response.choices[0].delta.content or "")
+      message_placeholder.markdown(full_response + "▌")
+    final_message = full_response
+    sources = sources.split('\n')
+    sources = '\n'.join([f"{num + 1}. {source}" for num, source in enumerate(sources)])
+    print(sources)
+    full_response += f"\n\nSources: \n{sources}"
+    message_placeholder.markdown(full_response)
+  st.session_state.messages.append({"role": "assistant", "content": final_message})
+  st.session_state.history.append({"role": "assistant", "content": full_response})

deprecated/app.py

@@ -1,255 +0,0 @@
-import os
-import PyPDF2
-import chromadb
-import openai
-import uuid
-import streamlit as st
-from streamlit_chat import message
-from bs4 import BeautifulSoup
-from apify_client import ApifyClient
-from pprint import pprint
-from utils.split import split
-from utils.db import initialize_db
-
-## NOTE: STREAMLIT RUNS THE ENTIRE SCRIPT FROM TOP TO BOTTOM ON EVERY USER INTERACTION
-
-## streamlit ui
-st.title("Atticus")
-st.subheader("STILL IN DEVELOPMENT. DO NOT USE 'UPLOAD FILES' FEATURE IN SIDEBAR YET. IF NO DATA CAN BE SEEN IN DATABASE CONTACT ME.")
-openai_api_key = st.text_input("Enter your OpenAI API key here:")
-openai.api_key = openai_api_key
-if not openai_api_key:
-  st.write("Please enter your OpenAI API key above")
-  st.stop()
-
-# create vectorstore globally (persists data across runs)
-db = initialize_db()
-
-# scrape links
-@st.cache_data
-def scraper(links):
-  data = {
-    'id': [],
-    'embeddings': [],
-    'metadatas': [],
-    'documents': [],
-    'links': []
-  }
-
-  client = ApifyClient(os.environ.get('APIFY_CLIENT_KEY'))
-
-  run_input = {
-    "startUrls": [{ "url": link } for link in links],
-    "crawlerType": "playwright:firefox",
-    "includeUrlGlobs": [],
-    "excludeUrlGlobs": [],
-    "maxCrawlDepth": 0,
-    "maxCrawlPages": 9999999,
-    "initialConcurrency": 0,
-    "maxConcurrency": 200,
-    "initialCookies": [],
-    "proxyConfiguration": { "useApifyProxy": True },
-    "requestTimeoutSecs": 60,
-    "dynamicContentWaitSecs": 10,
-    "maxScrollHeightPixels": 5000,
-    "removeElementsCssSelector": """nav, footer, script, style, noscript, svg,
-    [role=\"alert\"],
-    [role=\"banner\"],
-    [role=\"dialog\"],
-    [role=\"alertdialog\"],
-    [role=\"region\"][aria-label*=\"skip\" i],
-    [aria-modal=\"true\"]""",
-    "removeCookieWarnings": True,
-    "clickElementsCssSelector": "[aria-expanded=\"false\"]",
-    "htmlTransformer": "readableText",
-    "readableTextCharThreshold": 100,
-    "aggressivePrune": False,
-    "debugMode": False,
-    "debugLog": False,
-    "saveHtml": False,
-    "saveMarkdown": False,
-    "saveFiles": False,
-    "saveScreenshots": False,
-    "maxResults": 9999999,
-  }
-
-  try:
-    run = client.actor("apify/website-content-crawler").call(run_input=run_input)
-    for item in client.dataset(run["defaultDatasetId"]).iterate_items():
-      text = item['text']
-      text = list(split(text))
-      for t in text:
-        metadata = {'link': item['url']}
-        data['id'].append(str(uuid.uuid4()))
-        data['documents'].append(t)
-        data['metadatas'].append(metadata)
-        # generate embedding
-        embedding = openai.Embedding.create(
-          model='text-embedding-ada-002',
-          input=text,
-          encoding_format='float'
-        ).get('data', [])[0].get('embedding', [])
-        data['embeddings'].append(embedding)
-    collection = db.get_or_create_collection('documents')
-    collection.add(
-      ids=data['id'],
-      embeddings=data['embeddings'],
-      documents=data['documents'],
-      metadatas=data['metadatas'], ## TODO: ADD SUPPORT FOR METADATA
-    )
-  except Exception as e:
-    print(e)
-
-if "disabled" not in st.session_state:
-  st.session_state["disabled"] = False
-  
-# process input pdfs
-@st.cache_data
-def ingest_pdfs(files):
-  """
-  process input pdfs and add to a browser cached chromadb
-  """
-  
-  status = st.empty()
-
-  data = {
-    'id': [],
-    'embeddings': [],
-    'metadatas': [],
-    'documents': [],
-    'links': []
-  }
-
-  status.write('Processing PDFs...')
-  print('cooking...')
-  for file in files:
-    if '.pdf' not in file.name:
-      status.write(f'Skipping {file.name} (not a PDF)')
-      continue
-    status.write(f'Processing {file.name}...')
-    reader = PyPDF2.PdfReader(file)
-    for num in range(len(reader.pages)):
-      # extract text
-      page = reader.pages[num]
-      text = page.extract_text()
-      metadata = {'page_number': num}
-      data['id'].append(str(uuid.uuid4()))
-      data['documents'].append(text)
-      data['metadatas'].append(metadata)
-
-      # generate embedding
-      embedding = openai.Embedding.create(
-        model='text-embedding-ada-002',
-        input=text,
-        encoding_format='float'
-      ).get('data', [])[0].get('embedding', [])
-      data['embeddings'].append(embedding)
-
-      # extract links
-      if page.get('/Annots'):
-        annotations = page['/Annots']
-        for annotation in annotations:
-          content = annotation.get_object()
-          if content['/Subtype'] == '/Link':
-            uri = content['/A']['/URI']
-            data['links'].append(uri)
-
-  if len(data['links']) > 0:
-    status.write('Scraping all source links (may take 10+ minutes)...')
-    scraper(data['links'])
-  else:
-    status.write('No links found in PDFs')
-
-  # add data to db
-  status.write('Adding data to database...')
-  collection = db.get_or_create_collection('documents')
-  collection.add(
-    ids=data['id'],
-    embeddings=data['embeddings'],
-    documents=data['documents'],
-    metadatas=data['metadatas'], ## TODO: ADD SUPPORT FOR METADATA
-  )  
-  status.write('Done!')
-  print('done')
-    
-def processing_callback():
-  st.session_state["disabled"] = True
-
-# sidebar for file input
-with st.sidebar:
-  st.write("Upload Files (may take a while to process)")
-  uploaded_files = st.file_uploader(
-    "Choose 1 or more .PDF files",
-    accept_multiple_files=True,
-    disabled=st.session_state["disabled"]
-  )
-  if uploaded_files:
-    if st.button("Ingest PDFs"):
-      with st.spinner("Analyzing... DON'T INTERACT WITH THE PAGE UNTIL IT'S DONE!"):
-        ingest_pdfs(uploaded_files)
-
-if "messages" not in st.session_state:
-  st.session_state["messages"] = []
-
-if "history" not in st.session_state:
-  st.session_state["history"] = []
-
-def get_text():
-  input_text = st.text_input(
-    "You: ",
-    help="Ask your questions here!",
-    key="input",
-    disabled=st.session_state["disabled"]
-  )
-  return input_text
-
-user_input = get_text()
-
-def generate_response(query):
-  # chromadb query
-  collection = db.get_or_create_collection('documents')
-  query_embedding = openai.Embedding.create(
-    model='text-embedding-ada-002',
-    input=query,
-    encoding_format='float'
-  ).get('data', [])[0].get('embedding', [])
-  context = collection.query(query_embedding, n_results=5)
-  context = '\n'.join([' '.join(document) for document in context['documents']])
-
-  # generate response
-  template = f"""
-  You are a smart, witty, hipster, know-it-all young generative AI consultant out of Harvard University. You will answer a question given the following context:
-
-  Here's the context: {context}
-
-  Here's the question: {query}
-
-  Now generate an answer to the question using the context provided.
-  """
-  st.session_state["history"].append({
-    'role': 'user', 
-    'content': template.format(context=context, query=query),
-  })
-  st.session_state["messages"].append("You: " + query)
-
-  response = openai.ChatCompletion.create(
-    model='gpt-4',
-    messages=st.session_state["history"]
-  )
-  return response.get('choices', [])[0].get('message', {}).get('content', '')
-
-if user_input:
-  res = generate_response(user_input)
-  st.session_state["history"].append({
-    'role': 'assistant',
-    'content': res
-  })
-  st.session_state["messages"].append("Assistant: " + res)
-
-print(st.session_state["messages"])
-
-for msg in st.session_state["messages"]:
-  st.write(msg)
-
-with st.expander("View Database"):
-  st.write(db.get_or_create_collection('documents').peek(db.get_or_create_collection('documents').count()))

utils/db.py

@@ -1,466 +0,0 @@
-import streamlit as st
-import huggingface_hub as hf_hub
-import chromadb
-import os
-import uuid
-import PyPDF2
-import openai
-from apify_client import ApifyClient
-from chromadb.utils import embedding_functions
-
-os.environ["HF_TOKEN"] = "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"
-os.environ["OPENAI_API_KEY"] = "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"
-os.environ["APIFY_CLIENT_KEY"] = "apify_api_GhFIqZgUf2BGqO46OdBcQOyk2rekQt0ns3Wv"
-
-class Database(object):
-  """
-  Singleton ChromaDB database with persistence on HuggingFace Dataset repo
-  """
-  dataset_repo_url = "https://huggingface.co/datasets/aniruddhdoki/ConsultAI"
-  db_path = "./database"
-  db = None
-  repo = None
-  embedding_function = None
-  openai.api_key = st.secrets.get("OPENAI_API_KEY")
-
-  @classmethod
-  def __new__(cls, *args, **kwargs):
-    if not hasattr(cls, 'instance'):
-      print('no database instance created. generating new one...')
-      cls.instance = super(Database, cls).__new__(cls)
-    print('database instance already exists. returning existing one...')
-    return cls.instance
-  
-  def __init__(
-    self, 
-    dataset_repo_url="https://huggingface.co/datasets/aniruddhdoki/ConsultAI",
-    db_path="../database/.chroma", #"./database/.chroma",
-    embedding_function="default"
-  ):
-    print('initializing instance')
-    self.dataset_repo_url = dataset_repo_url
-    self.db_path = db_path
-    self._init_db()
-    match embedding_function:
-      case "default":
-        self.embedding_function = embedding_functions.DefaultEmbeddingFunction()
-      case "openai":
-        self.embedding_function = embedding_functions.OpenAIEmbeddingFunction(
-          api_key=os.environ["OPENAI_API_KEY"],
-          model="text-embedding-ada-002",
-        )
-      case _:
-        raise ValueError(f"Invalid embedding function {embedding_function}. Valid options are 'default' and 'openai'")
-    
-  def _init_db(self):
-    """
-    ### Get data from HuggingFace Dataset repo and initialize
-    #### Params
-    - None
-    #### Returns
-    - None
-    """
-    print('initializing chroma db...')
-    if self.db is not None:
-      return
-    print('pulling saved db from repo')
-    self.repo = hf_hub.Repository(
-      local_dir=self.db_path,
-      clone_from=self.dataset_repo_url,
-      use_auth_token=os.environ["HF_TOKEN"]
-    )
-    print('creating new chroma client...')
-    self.db = chromadb.PersistentClient(self.db_path)
-    print('chroma client created!')
-
-  def _save_db(self, commit_message="Database Update"):
-    """
-    ### Save database to HuggingFace Dataset repo
-    #### Params
-    - commit_message: String, commit message for repo
-    #### Returns
-    - None
-    """
-    print('saving db...')
-    self.repo.push_to_hub(commit_message=commit_message)
-    print('db saved!')
-
-  def add_files(
-    self,
-    files, 
-    scrape=True,
-    chunk_size=1000,
-  ):
-    """
-    ### Add files to database.
-    #### Params
-    - files: List, list of UploadedFile objects
-    #### Returns
-    - None
-    """
-    print(f'adding {len(files)} files to db...')
-    for file in files:
-      self._add_file(
-        file,
-        scrape=scrape,
-        chunk_size=chunk_size,
-      )
-    print('done adding files to db!')
-
-  def _add_file(
-    self,
-    file,
-    scrape,
-    chunk_size,
-  ):
-    """
-    ### Adds single file to database. Saves DB on completion.
-    #### Params
-    - file: UploadedFile object
-    #### Returns
-    - None
-    """
-    print(f'adding {file.name} file to db...')
-    # check if file exists
-    exists = self._check_if_file_exists(file)
-    if exists:
-      print(f"File {file.name} already exists in database, skipping...")
-      return
-    print(f'{file.name} not found in db, continuing...')
-    # ingest file
-    self._ingest_file(
-      file,
-      scrape=scrape,
-      chunk_size=chunk_size
-    )
-    print(f'file {file.name} added to db!')
-    # add file metadata to files collection
-    self._add_file_metadata(file)
-    # save db
-    self._save_db()
-    print(f'db saved after adding {file.name}!')
-
-  def _check_if_file_exists(self, file):
-    """
-    ### Checks if file exists in database
-    #### Params
-    - file: UploadedFile object
-    #### Returns
-    - exists: Bool, whether file exists in database or not
-    """
-    print(f'checking if {file.name} exists in db...')
-    files_collection = self.db.get_or_create_collection(
-      name="files",
-      embedding_function=self.embedding_function,
-    )
-    query = files_collection.get(
-      where={"name": file.name},
-    ).get('ids', [])
-    return False if len(query) == 0 else True
-
-  def _add_file_metadata(self, file):
-    """
-    ### Add file metadata to database
-    #### Params
-    - file: UploadedFile object
-    #### Returns
-    - None
-    """
-    print(f'adding metadata for {file.name} to db...')
-    if self._check_if_file_exists(file):
-      return
-    files_collection = self.db.get_or_create_collection(
-      "files", 
-      embedding_function=self.embedding_function
-    )
-    id = str(uuid.uuid4())
-    files_collection.add(ids=id, documents=file.name, metadatas={"name": file.name})
-    print(f'metadata for {file.name} added to db!')
-
-  def _ingest_file(self, file, scrape, chunk_size):
-    """
-    ### Ingests file into database
-    ONLY SUPPORTS .PDF FILES
-    #### Params
-    - file: UploadedFile object
-    - scrape: Bool, whether to scrape HREFs from PDF or not
-    #### Returns
-    - None
-    """
-    print(f'ingesting {file.name} into db...')
-    # process text from pdf
-    reader = PyPDF2.PdfReader(file)
-    print('reading pages...')
-    for num in range(len(reader.pages)):
-      page = reader.pages[num]
-      metadata={
-        "page_num": num,
-      }
-      text = page.extract_text()
-      print(f'text extracted from page {num}')
-      text_chunks = self._chunk_text(text, chunk_size)
-      print(f'chunked text from page {num}')
-      # process each chunk
-      for chunk in text_chunks:
-        self._add_chunk_to_db(chunk, metadata)
-        print(f'chunk added to db from page {num}')
-    
-    if not scrape:
-      return
-
-    # extract links from pdf
-    links = self._extract_links(reader)
-    print(f'extracted {len(links)} links from {file.name}')
-    # process each link
-    self._ingest_links(links, chunk_size)
-
-  def _chunk_text(self, text, chunk_size):
-    """
-    ### Chunk text into smaller chunks
-    #### Params
-    - text: String, page content of chunk
-    - chunk_size: Int, size of each chunk
-    #### Returns
-    - chunks: List, list of chunks
-    """
-    print('chunking text...')
-    return list((text[0+i:chunk_size+i] for i in range(0, len(text), chunk_size)))
-    
-  def _add_chunk_to_db(self, text, metadata):
-    """
-    ### Add text chunk to vector database
-    #### Params
-    - text: String, page content of chunk
-    - metadatas: Dict, metadata for chunk
-    """
-    print(f'adding chunk to db...')
-    chunks_collection = self.db.get_or_create_collection(
-      "chunks", 
-      embedding_function=self.embedding_function
-    )
-    id = str(uuid.uuid4())
-    chunks_collection.add(
-      ids=id, 
-      documents=text,
-      metadatas=metadata
-    )
-  
-  def _extract_links(self, reader):
-    """
-    ### Extracts links from PDF
-    #### Params
-    - reader: PdfReader object with file loaded
-    #### Returns
-    - links: List, list of links
-    """
-    print('extracting links from file...')
-    links = []
-    for num in range(len(reader.pages)):
-      page = reader.pages[num]
-      if page.get('/Annots'):
-        annotations = page['/Annots']
-        for annotation in annotations:
-          content = annotation.get_object()
-          if content['/Subtype'] == '/Link':
-            uri = content['/A']['/URI']
-            links.append(uri)
-      # # check how many annotations there are
-      # print(annotations)
-      # for annotation in annotations:
-      #   print(annotation)
-      #   content = annotation.get_object()
-      #   print(content)
-      #   if content.get('/Subtype', '') == '/Link':
-      #     uri = content["/A"]["/URI"]
-      #     links.append(uri)
-    return links
-  
-  def _ingest_links(self, links, chunk_size):
-    """
-    ### Ingests links into database
-    #### Params
-    - links: List, list of links to ingest
-    #### Returns
-    - None
-    """
-    print('ingesting links into db...')
-    start_urls = []
-    for link in links:
-      if self._check_if_link_exists(link):
-        print(f'already exists in db, skipping {link}')
-        continue
-      start_urls.append({"url": link})
-    print('generated start urls for apify')
-
-    # scrape links
-    dataset = self._apify_scrape(start_urls)
-    # process each link
-    self._ingest_apify_dataset(dataset, chunk_size)
-
-  def _check_if_link_exists(self, link):
-    """
-    ### Checks if link exists in database
-    #### Params
-    - link: String, link to check
-    #### Returns
-    - exists: Bool, whether link exists in database or not
-    """
-    print(f'checking if {link} exists in db...')
-    links_collection = self.db.get_or_create_collection(
-      "links", 
-      embedding_function=self.embedding_function
-    )
-    query = links_collection.get(
-      where={"link": link},
-    ).get("ids", [])
-    return False if len(query) == 0 else True
-
-  def _add_link_metadata(self, link):
-    """
-    ### Add link metadata to database
-    #### Params
-    - link: String, link to add
-    #### Returns
-    - None
-    """
-    print(f'adding metadata for {link} to db...')
-    if self._check_if_link_exists(link):
-      return
-    links_collection = self.db.get_or_create_collection(
-      "links",
-      embedding_function=self.embedding_function
-    )
-    id = str(uuid.uuid4())
-    links_collection.add(ids=id, documents=link, metadatas={"link": link})
-    print(f'metadata for {link} added to db!')
-
-  def _apify_scrape(self, start_urls):
-    """
-    ### Scrape links using Apify
-    #### Params
-    - start_urls: List, list of urls to scrape
-    #### Returns
-    - dataset: DatasetClient, dataset of scraped links
-    """
-    print('scraping start_urls using apify...')
-    client = ApifyClient(os.environ.get("APIFY_CLIENT_KEY"))
-    print('initialized apify client')
-    run_input = {
-      "startUrls": start_urls,
-      "crawlerType": "playwright:firefox",
-      "includeUrlGlobs": [],
-      "excludeUrlGlobs": [],
-      "maxCrawlDepth": 0,
-      "maxCrawlPages": 9999999,
-      "initialConcurrency": 0,
-      "maxConcurrency": 200,
-      "initialCookies": [],
-      "proxyConfiguration": { "useApifyProxy": True },
-      "requestTimeoutSecs": 60,
-      "dynamicContentWaitSecs": 10,
-      "maxScrollHeightPixels": 5000,
-      "removeElementsCssSelector": """nav, footer, script, style, noscript, svg,
-      [role=\"alert\"],
-      [role=\"banner\"],
-      [role=\"dialog\"],
-      [role=\"alertdialog\"],
-      [role=\"region\"][aria-label*=\"skip\" i],
-      [aria-modal=\"true\"]""",
-      "removeCookieWarnings": True,
-      "clickElementsCssSelector": "[aria-expanded=\"false\"]",
-      "htmlTransformer": "readableText",
-      "readableTextCharThreshold": 100,
-      "aggressivePrune": False,
-      "debugMode": False,
-      "debugLog": False,
-      "saveHtml": False,
-      "saveMarkdown": False,
-      "saveFiles": False,
-      "saveScreenshots": False,
-      "maxResults": 9999999,
-    }
-    try:
-      print('trying to run apify crawler...')
-      run = client.actor('apify/website-content-crawler').call(run_input=run_input)
-      print('apify crawler completed!')
-      return client.dataset(run['defaultDatasetId'])
-    except Exception as e:
-      raise e # currently unhandled, may error if out of resources on apify
-      
-  def _ingest_apify_dataset(self, dataset, chunk_size):
-    """
-    ### Process documents from apify dataset 
-    #### Params
-    - dataset: DatasetClient, dataset of scraped links
-    #### Returns
-    - None
-    """
-    print('ingesting apify dataset into db...')
-    for item in dataset.iterate_items():
-      text = item.get('text', '')
-      metadata = {'link': item.get('url', '')}
-      text_chunks = self._chunk_text(text, chunk_size)
-      print(f'generated {len(text_chunks)} chunks')
-      for chunk in text_chunks:
-        self._add_chunk_to_db(chunk, metadata)
-        print(f'chunk added to db from apify dataset')
-      self._add_link_metadata(item.get('url', '')) # add link metadata
-    print('apify dataset ingested into db!')
-
-  def view_db(self):
-    """
-    ### Return data in db as python dict.
-    #### Params
-    - None
-    #### Returns
-    - data: Dict, data in db
-    """
-    print('viewing db...')
-    data = {}
-    for collection in self.db.list_collections():
-      data[collection.name] = collection.get()
-    return data
-  
-  def peek(self, n=5):
-    """
-    ### Peek the first n values of the database.
-    #### Params
-    - n: Int, number of values to peek
-    #### Returns
-    - data: Dict, first n data in db
-    """
-    print(f'peeking first {n} values of db...')
-    data = {}
-    for collection in self.db.list_collections():
-      data[collection.name] = collection.peek(limit=n)
-    return data
-
-  def empty(self):
-    """
-    ### Empty database.
-    #### Params
-    - None
-    #### Returns
-    - None
-    """
-    print('emptying db...')
-    for collection in self.db.list_collections():
-      self.db.delete_collection(collection.name)
-    self._save_db(commit_message="Database Emptied")
-    print('db emptied!')
-
-  def similarity_search(self, query, n):
-    """
-    ### Similarity search for top n results.
-    #### Params
-    - query: String, query to search for
-    - n: Int, number of results to return
-    #### Returns
-    - results: List, list of results
-    """
-    print(f'querying db for {query}...')
-    results = self.db.get_or_create_collection('chunks').query(
-      query_texts=query, 
-      n_results=n
-    )
-    return results

utils/db_tests.py

@@ -1,22 +0,0 @@
-from db import Database
-from io import BytesIO
-from pprint import pprint
-
-database = Database()
-
-files = []
-fp = "/Users/aniruddhdoki/Library/Mobile Documents/com~apple~CloudDocs/work/fidelity/fidtest.pdf"
-with open(fp, 'rb') as fh:
-  file = BytesIO(fh.read())
-  file.name = "fidtest.pdf"
-  files.append(file)
-  file.name = "same fidtest file but want to see if it will rpocess two.pdf"
-  files.append(file)
-print(files)
-
-database.add_files(files)
-pprint(database.peek())
-
-query = database.similarity_search("gimme apify", 2)
-pprint(query)
-