mirror of https://github.com/hwchase17/langchain
community[patch]: OllamaEmbeddings - Pass headers to post request (#16880)
## Feature - Set additional headers in constructor - Headers will be sent in post request This feature is useful if deploying Ollama on a cloud service such as hugging face, which requires authentication tokens to be passed in the request header. ## Tests - Test if header is passed - Test if header is not passed Similar to https://github.com/langchain-ai/langchain/pull/15881 --------- Co-authored-by: Bagatur <baskaryan@gmail.com>pull/16805/head^2
parent
e0f137dbe0
commit
f51e6a35ba
@ -0,0 +1,61 @@
|
||||
import requests
|
||||
from pytest import MonkeyPatch
|
||||
|
||||
from langchain_community.embeddings.ollama import OllamaEmbeddings
|
||||
|
||||
|
||||
class MockResponse:
|
||||
status_code = 200
|
||||
|
||||
def json(self) -> dict:
|
||||
return {"embedding": [1, 2, 3]}
|
||||
|
||||
|
||||
def mock_response() -> MockResponse:
|
||||
return MockResponse()
|
||||
|
||||
|
||||
def test_pass_headers_if_provided(monkeypatch: MonkeyPatch) -> None:
|
||||
embedder = OllamaEmbeddings(
|
||||
base_url="https://ollama-hostname:8000",
|
||||
model="foo",
|
||||
headers={
|
||||
"Authorization": "Bearer TEST-TOKEN-VALUE",
|
||||
"Referer": "https://application-host",
|
||||
},
|
||||
)
|
||||
|
||||
def mock_post(url: str, headers: dict, json: str) -> MockResponse:
|
||||
assert url == "https://ollama-hostname:8000/api/embeddings"
|
||||
assert headers == {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": "Bearer TEST-TOKEN-VALUE",
|
||||
"Referer": "https://application-host",
|
||||
}
|
||||
assert json is not None
|
||||
|
||||
return mock_response()
|
||||
|
||||
monkeypatch.setattr(requests, "post", mock_post)
|
||||
|
||||
embedder.embed_query("Test prompt")
|
||||
|
||||
|
||||
def test_handle_if_headers_not_provided(monkeypatch: MonkeyPatch) -> None:
|
||||
embedder = OllamaEmbeddings(
|
||||
base_url="https://ollama-hostname:8000",
|
||||
model="foo",
|
||||
)
|
||||
|
||||
def mock_post(url: str, headers: dict, json: str) -> MockResponse:
|
||||
assert url == "https://ollama-hostname:8000/api/embeddings"
|
||||
assert headers == {
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
assert json is not None
|
||||
|
||||
return mock_response()
|
||||
|
||||
monkeypatch.setattr(requests, "post", mock_post)
|
||||
|
||||
embedder.embed_query("Test prompt")
|
Loading…
Reference in New Issue