kateforsberg commited on
Commit
ccb4963
·
1 Parent(s): 25d18e3

updated vector store retrieval

Browse files
Files changed (1) hide show
  1. uw_programmatic/base_machine.py +12 -5
uw_programmatic/base_machine.py CHANGED
@@ -298,11 +298,18 @@ class UWBaseMachine(StateMachine):
298
  all_kbs = {}
299
  if response.status_code == 200:
300
  data = response.json()
301
- for kb in data["knowledge_bases"]:
302
- name = kb["name"]
303
- kb_id = kb["knowledge_base_id"]
304
- if "KB_section" in name:
305
- all_kbs[name] = kb_id
 
 
 
 
 
 
 
306
  else:
307
  raise ValueError(response.status_code)
308
  self.kb_ids = all_kbs
 
298
  all_kbs = {}
299
  if response.status_code == 200:
300
  data = response.json()
301
+ next_page = data["pagination"]["next_page"]
302
+ while next_page is not None:
303
+ for kb in data["knowledge_bases"]:
304
+ name = kb["name"]
305
+ kb_id = kb["knowledge_base_id"]
306
+ if "KB_section" in name:
307
+ all_kbs[name] = kb_id
308
+ page_url = kb_url + f"?page={next_page}"
309
+ response = requests.get(url=page_url, headers=headers)
310
+ response.raise_for_status()
311
+ data = response.json()
312
+ next_page = data["pagination"]["next_page"]
313
  else:
314
  raise ValueError(response.status_code)
315
  self.kb_ids = all_kbs