Merge pull request #12 from almutareb/add_document_to_chroma_tool

innovation_pathfinder_ai/structured_tools/structured_tools.py

@@ -5,8 +5,15 @@ from langchain_community.tools import WikipediaQueryRun
 from langchain_community.utilities import WikipediaAPIWrapper
 #from langchain.tools import Tool
 from langchain_community.utilities import GoogleSearchAPIWrapper
+from langchain_community.embeddings.sentence_transformer import (
+    SentenceTransformerEmbeddings,
+)
+from langchain_community.vectorstores import Chroma
 import arxiv
 import ast
+
+import chromadb
+
 # hacky and should be replaced with a database
 from innovation_pathfinder_ai.source_container.container import (
     all_sources
@@ -18,6 +25,15 @@ from innovation_pathfinder_ai.database.db_handler import (
     add_many
 )
 
+from innovation_pathfinder_ai.vector_store.chroma_vector_store import (
+    add_pdf_to_vector_store
+)
+from innovation_pathfinder_ai.utils.utils import (
+    create_wikipedia_urls_from_text, create_folder_if_not_exists,
+)
+import os
+# from innovation_pathfinder_ai.utils import create_wikipedia_urls_from_text
+
 @tool
 def arxiv_search(query: str) -> str:
     """Search arxiv database for scientific research papers and studies. This is your primary information source.
@@ -72,9 +88,71 @@ def wikipedia_search(query: str) -> str:
     api_wrapper = WikipediaAPIWrapper()
     wikipedia_search = WikipediaQueryRun(api_wrapper=api_wrapper)
     wikipedia_results = wikipedia_search.run(query)
-    formatted_summaries = format_wiki_summaries(wikipedia_results)
-    all_sources += formatted_summaries
-    parsed_summaries = parse_list_to_dicts(formatted_summaries)
-    add_many(parsed_summaries)
+    all_sources += create_wikipedia_urls_from_text(wikipedia_results)
+    return wikipedia_results
+
+@tool
+def chroma_search(query:str) -> str:
+    """Search the Arxiv vector store for docmunets and relevent chunks"""
+    client = chromadb.PersistentClient(
+    # path=persist_directory,
+    )
+    
+    collection_name="ArxivPapers"
+    #store using envar
+    
+    embedding_function = SentenceTransformerEmbeddings(
+        model_name="all-MiniLM-L6-v2",
+        )
+    
+    vector_db = Chroma(
+    client=client, # client for Chroma
+    collection_name=collection_name,
+    embedding_function=embedding_function,
+    )
+    
+    retriever = vector_db.as_retriever()
+    docs = retriever.get_relevant_documents(query)
+    
+    return docs.__str__()
+
+
+@tool
+def embed_arvix_paper(paper_id:str) -> None:
+    """Download a paper from axriv to download a paper please input 
+    the axriv id such as "1605.08386v1" This tool is named get_arxiv_paper
+    If you input "http://arxiv.org/abs/2312.02813", This will break the code. Also only do 
+    "2312.02813". In addition please download one paper at a time. Pleaase keep the inputs/output
+    free of additional information only have the id. 
+    """
+    # code from https://lukasschwab.me/arxiv.py/arxiv.html
+    paper = next(arxiv.Client().results(arxiv.Search(id_list=[paper_id])))
+    
+    number_without_period = paper_id.replace('.', '')
+    
+    pdf_file_name = f"{number_without_period}.pdf"
+    
+    pdf_directory = "./downloaded_papers"
+    create_folder_if_not_exists(pdf_directory)
+    
+    # Download the PDF to a specified directory with a custom filename.
+    paper.download_pdf(dirpath=pdf_directory, filename=f"{number_without_period}.pdf")
+    
+    client = chromadb.PersistentClient(
+    # path=persist_directory,
+    )
+    
+    collection_name="ArxivPapers"
+    #store using envar
+    
+    embedding_function = SentenceTransformerEmbeddings(
+        model_name="all-MiniLM-L6-v2",
+        )
+    
+    full_path = os.path.join(pdf_directory, pdf_file_name)
     
-    return wikipedia_results.__str__()
+    add_pdf_to_vector_store(
+        collection_name=collection_name,
+        pdf_file_location=full_path,
+    )
+    

innovation_pathfinder_ai/utils/utils.py

@@ -1,5 +1,6 @@
 import hashlib
 import datetime
+import os
 
 from innovation_pathfinder_ai.utils import logger
 
@@ -168,4 +169,17 @@ def hash_text(text: str) -> str:
 
 
 def convert_timestamp_to_datetime(timestamp: str) -> str:
-    return datetime.datetime.fromtimestamp(int(timestamp)).strftime("%Y-%m-%d %H:%M:%S")
+    return datetime.datetime.fromtimestamp(int(timestamp)).strftime("%Y-%m-%d %H:%M:%S")
+
+def create_folder_if_not_exists(folder_path: str) -> None:
+    """
+    Create a folder if it doesn't already exist.
+
+    Args:
+    - folder_path (str): The path of the folder to create.
+    """
+    if not os.path.exists(folder_path):
+        os.makedirs(folder_path)
+        print(f"Folder '{folder_path}' created.")
+    else:
+        print(f"Folder '{folder_path}' already exists.")