Fix slicing of doc

h2oai · Oct 26, 2024 · 95761c2 · 95761c2
1 parent 6348f31
commit 95761c2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 4 deletions.
diff --git a/src/utils_langchain.py b/src/utils_langchain.py
@@ -130,10 +130,11 @@ def _huggingface_tokenizer_length(text: str) -> int:
 
 
 def select_docs_with_score(docs_with_score, top_k_docs, one_doc_size):
-    if top_k_docs > 0:
+    if one_doc_size is not None and len(docs_with_score) > 0:
+        doc1 = Document(page_content=docs_with_score[0][0].page_content[:one_doc_size], metadata=docs_with_score[0][0].metadata)
+        docs_with_score = [(doc1, docs_with_score[0][1])]
+    elif top_k_docs > 0:
         docs_with_score = docs_with_score[:top_k_docs]
-    elif one_doc_size is not None:
-        docs_with_score = [(docs_with_score[0][:one_doc_size], docs_with_score[0][1])]
     else:
         # do nothing
         pass

diff --git a/src/version.py b/src/version.py
@@ -1 +1 @@
-__version__ = "2f91e7bd7e48904bed9ac1e50d5eca1ca6da541f"
+__version__ = "6348f316022b004dd125d2e465df07e654b4287b"