Skip to content

Commit

Permalink
Fix tmp dir not exists (#401)
Browse files Browse the repository at this point in the history
* Fix tmp dir not exists

* Update tests to make it more clear

* Add check if save path is not None
  • Loading branch information
thinkall authored Oct 24, 2023
1 parent 3991f86 commit 80954e4
Show file tree
Hide file tree
Showing 2 changed files with 5 additions and 2 deletions.
3 changes: 3 additions & 0 deletions autogen/retrieve_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -254,7 +254,10 @@ def get_files_from_dir(dir_path: Union[str, List[str]], types: list = TEXT_FORMA
def get_file_from_url(url: str, save_path: str = None):
"""Download a file from a URL."""
if save_path is None:
os.makedirs("/tmp/chromadb", exist_ok=True)
save_path = os.path.join("/tmp/chromadb", os.path.basename(url))
else:
os.makedirs(os.path.dirname(save_path), exist_ok=True)
with requests.get(url, stream=True) as r:
r.raise_for_status()
with open(save_path, "wb") as f:
Expand Down
4 changes: 2 additions & 2 deletions test/test_retrieve_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -140,7 +140,7 @@ def query_vector_db(
db = lancedb.connect(db_path)
table = db.open_table("my_table")
query = table.search(vector).where(f"documents LIKE '%{search_string}%'").limit(n_results).to_df()
return {"ids": query["id"].tolist(), "documents": query["documents"].tolist()}
return {"ids": [query["id"].tolist()], "documents": [query["documents"].tolist()]}

def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str = ""):
results = self.query_vector_db(
Expand All @@ -166,7 +166,7 @@ def retrieve_docs(self, problem: str, n_results: int = 20, search_string: str =

create_lancedb()
ragragproxyagent.retrieve_docs("This is a test document spark", n_results=10, search_string="spark")
assert ragragproxyagent._results["ids"] == [3, 1, 5]
assert ragragproxyagent._results["ids"] == [[3, 1, 5]]

def test_custom_text_split_function(self):
def custom_text_split_function(text):
Expand Down

0 comments on commit 80954e4

Please sign in to comment.