fix: typo lancedb (#760)

This commit is contained in:
kan_cin 2025-07-01 17:11:22 +07:00 committed by GitHub
parent ffe766f24d
commit ec1f6abdc4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 10 additions and 4 deletions

View File

@ -1,6 +1,7 @@
from __future__ import annotations from __future__ import annotations
import os import os
from decouple import config from decouple import config
from kotaemon.base import Document, Param from kotaemon.base import Document, Param
@ -25,7 +26,9 @@ class CohereReranking(BaseReranking):
required=True, required=True,
) )
base_url: str = Param( base_url: str = Param(
None, help="Rerank API base url. Default is https://api.cohere.com", required=False None,
help="Rerank API base url. Default is https://api.cohere.com",
required=False,
) )
def run(self, documents: list[Document], query: str) -> list[Document]: def run(self, documents: list[Document], query: str) -> list[Document]:
@ -42,7 +45,9 @@ class CohereReranking(BaseReranking):
print("Cohere API key not found. Skipping rerankings.") print("Cohere API key not found. Skipping rerankings.")
return documents return documents
cohere_client = cohere.Client(self.cohere_api_key, base_url=self.base_url or os.getenv("CO_API_URL")) cohere_client = cohere.Client(
self.cohere_api_key, base_url=self.base_url or os.getenv("CO_API_URL")
)
compressed_docs: list[Document] = [] compressed_docs: list[Document] = []
if not documents: # to avoid empty api call if not documents: # to avoid empty api call

View File

@ -114,10 +114,11 @@ class LanceDBDocumentStore(BaseDocumentStore):
except (ValueError, FileNotFoundError): except (ValueError, FileNotFoundError):
docs = [] docs = []
# return the documents using the order of original ids (which were ordered by score) # return the documents using the order of original
# ids (which were ordered by score)
doc_dict = { doc_dict = {
doc["id"]: Document( doc["id"]: Document(
d_=doc["id"], id_=doc["id"],
text=doc["text"] if doc["text"] else "<empty>", text=doc["text"] if doc["text"] else "<empty>",
metadata=json.loads(doc["attributes"]), metadata=json.loads(doc["attributes"]),
) )