isayahc commited on
Commit
a0edccd
·
1 Parent(s): 2fd3f0a

added the chroma search tool

Browse files
innovation_pathfinder_ai/structured_tools/structured_tools.py CHANGED
@@ -8,6 +8,7 @@ from langchain_community.utilities import GoogleSearchAPIWrapper
8
  from langchain_community.embeddings.sentence_transformer import (
9
  SentenceTransformerEmbeddings,
10
  )
 
11
  import arxiv
12
  import ast
13
 
@@ -87,6 +88,31 @@ def wikipedia_search(query: str) -> str:
87
  all_sources += create_wikipedia_urls_from_text(wikipedia_results)
88
  return wikipedia_results
89
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
90
 
91
  @tool
92
  def embed_arvix_paper(paper_id:str) -> None:
 
8
  from langchain_community.embeddings.sentence_transformer import (
9
  SentenceTransformerEmbeddings,
10
  )
11
+ from langchain_community.vectorstores import Chroma
12
  import arxiv
13
  import ast
14
 
 
88
  all_sources += create_wikipedia_urls_from_text(wikipedia_results)
89
  return wikipedia_results
90
 
91
+ @tool
92
+ def chroma_search(query:str) -> str:
93
+ """Search the Arxiv vector store for docmunets and relevent chunks"""
94
+ client = chromadb.PersistentClient(
95
+ # path=persist_directory,
96
+ )
97
+
98
+ collection_name="ArxivPapers"
99
+ #store using envar
100
+
101
+ embedding_function = SentenceTransformerEmbeddings(
102
+ model_name="all-MiniLM-L6-v2",
103
+ )
104
+
105
+ vector_db = Chroma(
106
+ client=client, # client for Chroma
107
+ collection_name=collection_name,
108
+ embedding_function=embedding_function,
109
+ )
110
+
111
+ retriever = vector_db.as_retriever()
112
+ docs = retriever.get_relevant_documents(query)
113
+
114
+ return docs.__str__()
115
+
116
 
117
  @tool
118
  def embed_arvix_paper(paper_id:str) -> None: