Skip to content

Commit e8b8511

Browse files
ddematheuddematheu
authored
0.0.34 (#31)
Fix: Website robust Fix: neum-tools dependency Fix: S3 pre-fix Feature: metadata filters in search for Weaviate, Pinecone, Supabase Feature: Client update with additional endpoints
1 parent 62ef674 commit e8b8511

13 files changed

Lines changed: 143 additions & 24 deletions

File tree

‎neumai/neumai/Chunkers/CustomChunker.py‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
from typing import List, Generator, Optional
22
from neumai.Chunkers.Chunker import Chunker
33
from neumai.Shared.NeumDocument import NeumDocument
4-
from neumai_tools.SemanticHelpers import semantic_chunking
54
from pydantic import Field
65
from neumai.Shared.Exceptions import CustomChunkerException
76

@@ -37,6 +36,10 @@ def optional_properties(self) -> List[str]:
3736
return ["batch_size"]
3837

3938
def chunk(self, documents:List[NeumDocument]) -> Generator[List[NeumDocument], None, None]:
39+
try:
40+
from neumai_tools.SemanticHelpers import semantic_chunking
41+
except ImportError:
42+
raise ImportError("You must run " "`pip install neumai-tools")
4043

4144
chunking_code_exec=self.code
4245
batch_size = self.batch_size
@@ -60,6 +63,10 @@ def chunk(self, documents:List[NeumDocument]) -> Generator[List[NeumDocument], N
6063
yield documents_to_embed
6164

6265
def config_validation(self) -> bool:
66+
try:
67+
from neumai_tools.SemanticHelpers import semantic_chunking
68+
except ImportError:
69+
raise ImportError("You must run " "`pip install neumai-tools")
6370
try:
6471
chunks = semantic_chunking(documents=[NeumDocument(id="test", content="test", metadata={})], chunking_code_exec=self.code)
6572
except Exception as e:

‎neumai/neumai/Client/NeumClient.py‎

Lines changed: 102 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,14 +56,114 @@ def trigger_pipeline(self, pipeline_id:str, sync_type:TriggerSyncTypeEnum):
5656
except Exception as e:
5757
print(f"Pipeline trigger failed. Exception - {e}")
5858

59-
def search_pipeline(self, pipeline_id:str, query:str, num_of_results:int = 3, track:bool = False):
59+
def search_pipeline(self, pipeline_id:str, query:str, num_of_results:int = 3, track:bool = False, filter:dict = {}, requested_by:str = None):
6060
url = f"{self.endpoint}/pipelines/{pipeline_id}/search"
6161

6262
payload = {
6363
"number_of_results": num_of_results,
6464
"query": query,
65-
"collect_retrieval":track
65+
"collect_retrieval":track,
66+
"requested_by":requested_by,
67+
"filter":filter
6668
}
69+
headers = {
70+
"accept": "application/json",
71+
"neum-api-key": self.api_key,
72+
"content-type": "application/json"
73+
}
74+
try:
75+
response = requests.post(url, json=payload, headers=headers)
76+
return json.loads(response.text)
77+
except Exception as e:
78+
print(f"Pipeline trigger failed. Exception - {e}")
79+
80+
def search_file(self, pipeline_id:str, file_id:str, query:str, num_of_results:int = 3, track:bool = False, requested_by:str = None):
81+
url = f"{self.endpoint}/pipelines/{pipeline_id}/files/search?file_id={file_id}"
82+
83+
payload = {
84+
"number_of_results": num_of_results,
85+
"query": query,
86+
"collect_retrieval":track,
87+
"requested_by":requested_by,
88+
}
89+
headers = {
90+
"accept": "application/json",
91+
"neum-api-key": self.api_key,
92+
"content-type": "application/json"
93+
}
94+
try:
95+
response = requests.post(url, json=payload, headers=headers)
96+
return json.loads(response.text)
97+
except Exception as e:
98+
print(f"Pipeline trigger failed. Exception - {e}")
99+
100+
def get_files(self, pipeline_id:str):
101+
url = f"{self.endpoint}/pipelines/{pipeline_id}/files"
102+
103+
headers = {
104+
"accept": "application/json",
105+
"neum-api-key": self.api_key,
106+
"content-type": "application/json"
107+
}
108+
109+
try:
110+
response = requests.get(url, headers=headers)
111+
return json.loads(response.text)
112+
except Exception as e:
113+
print(f"Pipeline trigger failed. Exception - {e}")
114+
115+
def get_file(self, pipeline_id:str, file_id:str):
116+
url = f"{self.endpoint}/pipelines/{pipeline_id}/files?file_id={file_id}"
117+
118+
headers = {
119+
"accept": "application/json",
120+
"neum-api-key": self.api_key,
121+
"content-type": "application/json"
122+
}
123+
124+
try:
125+
response = requests.get(url, headers=headers)
126+
return json.loads(response.text)
127+
except Exception as e:
128+
print(f"Pipeline trigger failed. Exception - {e}")
129+
130+
def get_retrievals_by_file_id(self, pipeline_id:str, file_id:str):
131+
url = f"{self.endpoint}/retrievals/{pipeline_id}/files?file_id={file_id}"
132+
133+
headers = {
134+
"accept": "application/json",
135+
"neum-api-key": self.api_key,
136+
"content-type": "application/json"
137+
}
138+
139+
try:
140+
response = requests.get(url, headers=headers)
141+
return json.loads(response.text)
142+
except Exception as e:
143+
print(f"Pipeline trigger failed. Exception - {e}")
144+
145+
def get_retrievals_by_pipeline_id(self, pipeline_id:str):
146+
url = f"{self.endpoint}/retrievals/{pipeline_id}"
147+
148+
headers = {
149+
"accept": "application/json",
150+
"neum-api-key": self.api_key,
151+
"content-type": "application/json"
152+
}
153+
154+
try:
155+
response = requests.get(url, headers=headers)
156+
return json.loads(response.text)
157+
except Exception as e:
158+
print(f"Pipeline trigger failed. Exception - {e}")
159+
160+
def provide_retrieval_feedback(self, pipeline_id:str, retrieval_id:str, status:str):
161+
url = f"{self.endpoint}/retrievals/{pipeline_id}/{retrieval_id}"
162+
163+
payload = {
164+
"status":status
165+
}
166+
67167
headers = {
68168
"accept": "application/json",
69169
"neum-api-key": self.api_key,

‎neumai/neumai/DataConnectors/S3Connector.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ class S3Connector(DataConnector):
3838

3939
bucket_name: str = Field(..., description="S3 bucket name.")
4040

41-
prefix: Optional[str] = Field(None, description="Optional prefix for S3 files.")
41+
prefix: Optional[str] = Field("", description="Optional prefix for S3 files.")
4242

4343
selector: Optional[Selector] = Field(Selector(to_embed=[], to_metadata=[]), description="Selector for data connector metadata")
4444

@@ -147,7 +147,6 @@ def connect_and_download(self, cloudFile:CloudFile) -> Generator[LocalFile, None
147147
def config_validation(self) -> bool:
148148
if not all(x in self.available_metadata for x in self.selector.to_metadata):
149149
raise ValueError("Invalid metadata values provided")
150-
151150
try:
152151
session = boto3.Session(
153152
aws_access_key_id=self.aws_key_id,

‎neumai/neumai/DataConnectors/WebsiteConnector.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,6 +78,9 @@ def connect_and_download(self, cloudFile:CloudFile) -> Generator[LocalFile, None
7878
soup = BeautifulSoup(response.content, 'html.parser')
7979
# Find the <body> element and extract its HTML content
8080
body = soup.find('body')
81+
# Some sites don't have a body, so instead just get all the text off it.
82+
if body == None:
83+
body = soup.get_text()
8184
body_html = str(body) # Convert the body tag to a string to get its HTML content
8285
# Create a temporary file and write the extracted HTML to it
8386
with tempfile.NamedTemporaryFile(delete=False, suffix=".html", mode='w', encoding="utf-8") as temp:

‎neumai/neumai/Pipelines/Pipeline.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -152,9 +152,9 @@ def run(self) -> int:
152152
except Exception as e:
153153
raise e
154154

155-
def search(self, query:str, number_of_results:int) -> List[NeumSearchResult]:
155+
def search(self, query:str, number_of_results:int, filter:dict={}) -> List[NeumSearchResult]:
156156
vector_for_query = self.embed.embed_query(query=query)
157-
matches = self.sink.search(vector=vector_for_query, number_of_results=number_of_results)
157+
matches = self.sink.search(vector=vector_for_query, number_of_results=number_of_results, filter=filter)
158158
return matches
159159

160160
# Todo standardize the model serialization as we are mixing FE and BE concepts into the SDK

‎neumai/neumai/Shared/NeumSearch.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,4 +7,4 @@ class NeumSearchResult(BaseModel):
77
id:str = Field(..., description="Search result vector ID")
88
metadata:dict = Field(...,description="Search result vector metadata")
99
score:float = Field(..., description="Search result similarity score")
10-
vector: Optional[List[float]] = Field(..., description="Search result vector")
10+
vector: Optional[List[float]] = Field(None, description="Search result vector")

‎neumai/neumai/SinkConnectors/PineconeSink.py‎

Lines changed: 8 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -88,20 +88,25 @@ def store(self, vectors_to_store:List[NeumVector]) -> int:
8888
raise PineconeInsertionException(f"Failed to store in Pinecone. Exception - {e}")
8989
return int(vectors_stored)
9090

91-
def search(self, vector: List[float], number_of_results:int) -> List[NeumSearchResult]:
91+
def search(self, vector: List[float], number_of_results:int, filter:dict = {}) -> List[NeumSearchResult]:
9292
import pinecone
9393
api_key = self.api_key
9494
environment = self.environment
9595
index = self.index
9696
namespace = self.namespace
9797
if environment == "gcp-starter": namespace = None # short-term fix given gcp-starter limitation
98-
9998
try:
10099
pinecone.init(
101100
api_key=api_key,
102101
environment=environment)
103102
index = pinecone.Index(index)
104-
results = index.query(vector=vector, top_k=number_of_results, namespace=namespace, include_values=False, include_metadata=True)["matches"]
103+
results = index.query(
104+
vector=vector,
105+
filter=filter,
106+
top_k=number_of_results,
107+
namespace=namespace,
108+
include_values=False,
109+
include_metadata=True)["matches"]
105110
except Exception as e:
106111
raise PineconeQueryException(f"Failed to query pinecone. Exception - {e}")
107112

‎neumai/neumai/SinkConnectors/QdrantSink.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ def store(self, vectors_to_store:List[NeumVector]) -> int:
8282
return len(points)
8383
raise QdrantInsertionException("Qdrant storing failed. Try again later.")
8484

85-
def search(self, vector: List[float], number_of_results: int) -> List:
85+
def search(self, vector: List[float], number_of_results: int, filter:dict = {}) -> List:
8686
url = self.url
8787
api_key = self.api_key
8888
collection_name = self.collection_name
@@ -96,7 +96,7 @@ def search(self, vector: List[float], number_of_results: int) -> List:
9696
collection_name=collection_name,
9797
query_vector=vector,
9898
with_payload= True,
99-
limit=number_of_results
99+
limit=number_of_results,
100100
)
101101
except Exception as e:
102102
raise QdrantQueryException(f"Failed to query Qdrant. Exception - {e}")

‎neumai/neumai/SinkConnectors/SingleStoreSink.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -108,7 +108,7 @@ def store(self, vectors_to_store:List[NeumVector]) -> int:
108108

109109
return len(vectors_to_store), None
110110

111-
def search(self, vector: List[float], number_of_results: int) -> List[NeumSearchResult]:
111+
def search(self, vector: List[float], number_of_results: int, filter:dict={}) -> List[NeumSearchResult]:
112112
url = self.url
113113
table = self.table
114114

‎neumai/neumai/SinkConnectors/SinkConnector.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,7 +32,7 @@ def store(self, vectors_to_store:List[NeumVector]) -> int:
3232
"""Store vectors with a given service"""
3333

3434
@abstractmethod
35-
def search(self, vector:List[float], number_of_results:int) -> List[NeumSearchResult]:
35+
def search(self, vector:List[float], number_of_results:int, filter:dict={}) -> List[NeumSearchResult]:
3636
"""Search vectors for a given service"""
3737

3838
@abstractmethod

0 commit comments

Comments
 (0)