mirror of
https://github.com/danswer-ai/danswer.git
synced 2025-05-14 05:40:16 +02:00
28 lines
844 B
Python
28 lines
844 B
Python
import requests
|
|
|
|
from danswer.document_index.vespa.index import DOCUMENT_ID_ENDPOINT
|
|
|
|
|
|
class TestVespaClient:
|
|
def __init__(self, index_name: str):
|
|
self.index_name = index_name
|
|
self.vespa_document_url = DOCUMENT_ID_ENDPOINT.format(index_name=index_name)
|
|
|
|
def get_documents_by_id(
|
|
self, document_ids: list[str], wanted_doc_count: int = 1_000
|
|
) -> dict:
|
|
selection = " or ".join(
|
|
f"{self.index_name}.document_id=='{document_id}'"
|
|
for document_id in document_ids
|
|
)
|
|
params = {
|
|
"selection": selection,
|
|
"wantedDocumentCount": wanted_doc_count,
|
|
}
|
|
response = requests.get(
|
|
self.vespa_document_url,
|
|
params=params, # type: ignore
|
|
)
|
|
response.raise_for_status()
|
|
return response.json()
|