Skip to content

Commit e4b8afd

Browse files
authored
0.0.36 (#37)
Feature: client updates for retrieval gets Feature: Add processing time to object Fix: fix for pipeline run Fix: fix for weaviate class name
1 parent 4bf69b4 commit e4b8afd

4 files changed

Lines changed: 57 additions & 6 deletions

File tree

neumai/neumai/Client/NeumClient.py

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -202,6 +202,51 @@ def get_retrievals_by_pipeline_id(self, pipeline_id:str):
202202
except Exception as e:
203203
print(f"Retrievals fetch failed. Exception - {e}")
204204

205+
def get_retrievals_by_pipeline_id_user_id(self, pipeline_id:str, user_id:str):
206+
url = f"{self.endpoint}/retrievals/{pipeline_id}/user/{user_id}"
207+
208+
headers = {
209+
"accept": "application/json",
210+
"neum-api-key": self.api_key,
211+
"content-type": "application/json"
212+
}
213+
214+
try:
215+
response = requests.get(url, headers=headers)
216+
return json.loads(response.text)
217+
except Exception as e:
218+
print(f"Retrievals fetch failed. Exception - {e}")
219+
220+
def get_retrievals_by_user_id(self, user_id:str):
221+
url = f"{self.endpoint}/retrievals/user/{user_id}"
222+
223+
headers = {
224+
"accept": "application/json",
225+
"neum-api-key": self.api_key,
226+
"content-type": "application/json"
227+
}
228+
229+
try:
230+
response = requests.get(url, headers=headers)
231+
return json.loads(response.text)
232+
except Exception as e:
233+
print(f"Retrievals fetch failed. Exception - {e}")
234+
235+
def get_retrievals_by_file_id_user_id(self, pipeline_id:str, file_id:str, user_id:str):
236+
url = f"{self.endpoint}/retrievals/{pipeline_id}/files?file_id={file_id}&user_id={user_id}"
237+
238+
headers = {
239+
"accept": "application/json",
240+
"neum-api-key": self.api_key,
241+
"content-type": "application/json"
242+
}
243+
244+
try:
245+
response = requests.get(url, headers=headers)
246+
return json.loads(response.text)
247+
except Exception as e:
248+
print(f"Retrievals fetch failed. Exception - {e}")
249+
205250
def provide_retrieval_feedback(self, pipeline_id:str, retrieval_id:str, status:str):
206251
url = f"{self.endpoint}/retrievals/{pipeline_id}/{retrieval_id}"
207252

neumai/neumai/Pipelines/PipelineRun.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ class PipelineRun(BaseModel):
5858
last_updated: Optional[float] = None
5959
number_of_documents: Optional[int] = None
6060
finished_distributing: bool = False
61+
processing_time: Optional[float] = None
6162

6263
def set_id(self, id: str):
6364
self.id = id
@@ -76,6 +77,9 @@ def set_pipeline_id(self, pipeline_id: str):
7677

7778
def set_created(self, created: float):
7879
self.created = created
80+
81+
def set_processing_time(self, processing_time: float):
82+
self.processing_time = processing_time
7983

8084
def as_pipeline_run(dct:dict):
8185
if dct == None:
@@ -91,5 +95,6 @@ def as_pipeline_run(dct:dict):
9195
task_details=PipelineRunTaskDetails.as_pipeline_run_task_details(dct.get("task_details", None)),
9296
last_updated=dct.get("last_updated", None),
9397
number_of_documents=dct.get("number_of_documents",None),
94-
finished_distributing=dct.get("finished_distributing",None) # we could do something like in distributed tasks where we store the state of the DAG in the pipeline run object.. for now just doing finished_distributing
98+
finished_distributing=dct.get("finished_distributing",None),# we could do something like in distributed tasks where we store the state of the DAG in the pipeline run object.. for now just doing finished_distributing
99+
processing_time=dct.get("processing_time", None)
95100
)

neumai/neumai/SinkConnectors/WeaviateSink.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ def delete_vectors_with_file_id(self, file_id: str) -> bool:
105105
api_key = self.api_key
106106
url = self.url
107107
# Weaviate requires first letter to be capitalized
108-
class_name = self.class_name
108+
class_name = self.class_name.replace("-","_")
109109
class_name = _capitalize_first_letter(class_name)
110110
client = weaviate.Client(
111111
url=url,
@@ -128,7 +128,8 @@ def store(self, vectors_to_store:List[NeumVector]) -> Tuple[List, dict]:
128128
batch_size = self.batch_size
129129
is_dynamic_batch = self.is_dynamic_batch
130130
batch_connection_error_retries = self.batch_connection_error_retries
131-
class_name = self.class_name
131+
class_name = self.class_name.replace("-","_")
132+
class_name = _capitalize_first_letter(class_name)
132133
partial_failure = {'did_fail': False, 'latest_failure': None, 'number_of_failures': 0}
133134

134135
if 'https' not in url:
@@ -174,7 +175,7 @@ def search(self, vector: List[float], number_of_results: int, filter:dict={}) ->
174175
api_key = self.api_key
175176
url = self.url
176177
# Weaviate requires first letter to be capitalized
177-
class_name = self.class_name
178+
class_name = self.class_name.replace("-","_")
178179
class_name = _capitalize_first_letter(class_name)
179180
client = weaviate.Client(
180181
url=url,
@@ -215,7 +216,7 @@ def info(self) -> NeumSinkInfo:
215216
api_key = self.api_key
216217
url = self.url
217218

218-
class_name = self.class_name
219+
class_name = self.class_name.replace("-","_")
219220
class_name = _capitalize_first_letter(class_name)
220221
client = weaviate.Client(
221222
url=url,

neumai/pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[tool.poetry]
22
name = "neumai"
3-
version = "0.0.35"
3+
version = "0.0.36"
44
description = "Package containing connectors for Neum AI."
55
authors = ["David de Matheu <david@tryneum.com>"]
66
license = "MIT"

0 commit comments

Comments
 (0)