Spaces:
Sleeping
Sleeping
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:
|