Skip to content

Commit d0f71bc

Browse files
authored
client: load OLLAMA_API_KEY on init (#583)
1 parent b22c5fd commit d0f71bc

File tree

3 files changed

+111
-15
lines changed

3 files changed

+111
-15
lines changed

examples/web-search-crawl.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,11 @@
55
# "ollama",
66
# ]
77
# ///
8-
import os
98
from typing import Union
109

1110
from rich import print
1211

13-
from ollama import Client, WebCrawlResponse, WebSearchResponse
12+
from ollama import WebCrawlResponse, WebSearchResponse, chat, web_crawl, web_search
1413

1514

1615
def format_tool_results(results: Union[WebSearchResponse, WebCrawlResponse]):
@@ -49,15 +48,17 @@ def format_tool_results(results: Union[WebSearchResponse, WebCrawlResponse]):
4948
return '\n'.join(output).rstrip()
5049

5150

52-
client = Client(headers={'Authorization': (os.getenv('OLLAMA_API_KEY'))})
53-
available_tools = {'web_search': client.web_search, 'web_crawl': client.web_crawl}
51+
# Set OLLAMA_API_KEY in the environment variable or use the headers parameter to set the authorization header
52+
# client = Client(headers={'Authorization': 'Bearer <OLLAMA_API_KEY>'})
53+
54+
available_tools = {'web_search': web_search, 'web_crawl': web_crawl}
5455

5556
query = "ollama's new engine"
5657
print('Query: ', query)
5758

5859
messages = [{'role': 'user', 'content': query}]
5960
while True:
60-
response = client.chat(model='qwen3', messages=messages, tools=[client.web_search, client.web_crawl], think=True)
61+
response = chat(model='qwen3', messages=messages, tools=[web_search, web_crawl], think=True)
6162
if response.message.thinking:
6263
print('Thinking: ')
6364
print(response.message.thinking + '\n\n')

ollama/_client.py

Lines changed: 25 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -94,20 +94,25 @@ def __init__(
9494
`kwargs` are passed to the httpx client.
9595
"""
9696

97+
headers = {
98+
k.lower(): v
99+
for k, v in {
100+
**(headers or {}),
101+
'Content-Type': 'application/json',
102+
'Accept': 'application/json',
103+
'User-Agent': f'ollama-python/{__version__} ({platform.machine()} {platform.system().lower()}) Python/{platform.python_version()}',
104+
}.items()
105+
if v is not None
106+
}
107+
api_key = os.getenv('OLLAMA_API_KEY', None)
108+
if not headers.get('authorization') and api_key:
109+
headers['authorization'] = f'Bearer {api_key}'
110+
97111
self._client = client(
98112
base_url=_parse_host(host or os.getenv('OLLAMA_HOST')),
99113
follow_redirects=follow_redirects,
100114
timeout=timeout,
101-
# Lowercase all headers to ensure override
102-
headers={
103-
k.lower(): v
104-
for k, v in {
105-
**(headers or {}),
106-
'Content-Type': 'application/json',
107-
'Accept': 'application/json',
108-
'User-Agent': f'ollama-python/{__version__} ({platform.machine()} {platform.system().lower()}) Python/{platform.python_version()}',
109-
}.items()
110-
},
115+
headers=headers,
111116
**kwargs,
112117
)
113118

@@ -638,7 +643,12 @@ def web_search(self, queries: Sequence[str], max_results: int = 3) -> WebSearchR
638643
639644
Returns:
640645
WebSearchResponse with the search results
646+
Raises:
647+
ValueError: If OLLAMA_API_KEY environment variable is not set
641648
"""
649+
if not self._client.headers.get('authorization', '').startswith('Bearer '):
650+
raise ValueError('Authorization header with Bearer token is required for web search')
651+
642652
return self._request(
643653
WebSearchResponse,
644654
'POST',
@@ -658,7 +668,12 @@ def web_crawl(self, urls: Sequence[str]) -> WebCrawlResponse:
658668
659669
Returns:
660670
WebCrawlResponse with the crawl results
671+
Raises:
672+
ValueError: If OLLAMA_API_KEY environment variable is not set
661673
"""
674+
if not self._client.headers.get('authorization', '').startswith('Bearer '):
675+
raise ValueError('Authorization header with Bearer token is required for web fetch')
676+
662677
return self._request(
663678
WebCrawlResponse,
664679
'POST',

tests/test_client.py

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1195,3 +1195,83 @@ async def test_arbitrary_roles_accepted_in_message_request_async(monkeypatch: py
11951195
client = AsyncClient()
11961196

11971197
await client.chat(model='llama3.1', messages=[{'role': 'somerandomrole', 'content': "I'm ok with you adding any role message now!"}, {'role': 'user', 'content': 'Hello world!'}])
1198+
1199+
1200+
def test_client_web_search_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch):
1201+
monkeypatch.delenv('OLLAMA_API_KEY', raising=False)
1202+
1203+
client = Client()
1204+
1205+
with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web search'):
1206+
client.web_search(['test query'])
1207+
1208+
1209+
def test_client_web_crawl_requires_bearer_auth_header(monkeypatch: pytest.MonkeyPatch):
1210+
monkeypatch.delenv('OLLAMA_API_KEY', raising=False)
1211+
1212+
client = Client()
1213+
1214+
with pytest.raises(ValueError, match='Authorization header with Bearer token is required for web fetch'):
1215+
client.web_crawl(['https://example.com'])
1216+
1217+
1218+
def _mock_request_web_search(self, cls, method, url, json=None, **kwargs):
1219+
assert method == 'POST'
1220+
assert url == 'https://ollama.com/api/web_search'
1221+
assert json is not None and 'queries' in json and 'max_results' in json
1222+
return httpxResponse(status_code=200, content='{"results": {}, "success": true}')
1223+
1224+
1225+
def _mock_request_web_crawl(self, cls, method, url, json=None, **kwargs):
1226+
assert method == 'POST'
1227+
assert url == 'https://ollama.com/api/web_crawl'
1228+
assert json is not None and 'urls' in json
1229+
return httpxResponse(status_code=200, content='{"results": {}, "success": true}')
1230+
1231+
1232+
def test_client_web_search_with_env_api_key(monkeypatch: pytest.MonkeyPatch):
1233+
monkeypatch.setenv('OLLAMA_API_KEY', 'test-key')
1234+
monkeypatch.setattr(Client, '_request', _mock_request_web_search)
1235+
1236+
client = Client()
1237+
client.web_search(['what is ollama?'], max_results=2)
1238+
1239+
1240+
def test_client_web_crawl_with_env_api_key(monkeypatch: pytest.MonkeyPatch):
1241+
monkeypatch.setenv('OLLAMA_API_KEY', 'test-key')
1242+
monkeypatch.setattr(Client, '_request', _mock_request_web_crawl)
1243+
1244+
client = Client()
1245+
client.web_crawl(['https://example.com'])
1246+
1247+
1248+
def test_client_web_search_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch):
1249+
monkeypatch.delenv('OLLAMA_API_KEY', raising=False)
1250+
monkeypatch.setattr(Client, '_request', _mock_request_web_search)
1251+
1252+
client = Client(headers={'Authorization': 'Bearer custom-token'})
1253+
client.web_search(['what is ollama?'], max_results=1)
1254+
1255+
1256+
def test_client_web_crawl_with_explicit_bearer_header(monkeypatch: pytest.MonkeyPatch):
1257+
monkeypatch.delenv('OLLAMA_API_KEY', raising=False)
1258+
monkeypatch.setattr(Client, '_request', _mock_request_web_crawl)
1259+
1260+
client = Client(headers={'Authorization': 'Bearer custom-token'})
1261+
client.web_crawl(['https://example.com'])
1262+
1263+
1264+
def test_client_bearer_header_from_env(monkeypatch: pytest.MonkeyPatch):
1265+
monkeypatch.setenv('OLLAMA_API_KEY', 'env-token')
1266+
1267+
client = Client()
1268+
assert client._client.headers['authorization'] == 'Bearer env-token'
1269+
1270+
1271+
def test_client_explicit_bearer_header_overrides_env(monkeypatch: pytest.MonkeyPatch):
1272+
monkeypatch.setenv('OLLAMA_API_KEY', 'env-token')
1273+
monkeypatch.setattr(Client, '_request', _mock_request_web_search)
1274+
1275+
client = Client(headers={'Authorization': 'Bearer explicit-token'})
1276+
assert client._client.headers['authorization'] == 'Bearer explicit-token'
1277+
client.web_search(['override check'])

0 commit comments

Comments
 (0)