kanshan/ollama-python
1
0
Fork 0
mirror of https://github.com/ollama/ollama-python.git synced 2026-01-13 21:57:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fixing empty header + ensuring security (#313)

Browse Source 
* Fixing empty header + ensuring security
This commit is contained in:
Parth Sareen 2024-11-07 13:12:22 -08:00 committed by Michael Yang
parent f25834217b
commit 72052188c3
2 changed files with 25 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
ollama/_client.py
View File

@ -90,11 +90,16 @@ class BaseClient:
base_url=_parse_host(host or os.getenv('OLLAMA_HOST')),
follow_redirects=follow_redirects,
timeout=timeout,
# Lowercase all headers to ensure override
headers={
'Content-Type': 'application/json',
'Accept': 'application/json',
'User-Agent': f'ollama-python/{__version__} ({platform.machine()} {platform.system().lower()}) Python/{platform.python_version()}',
}.update(headers or {}),
k.lower(): v
for k, v in {
**(headers or {}),
'Content-Type': 'application/json',
'Accept': 'application/json',
'User-Agent': f'ollama-python/{__version__} ({platform.machine()} {platform.system().lower()}) Python/{platform.python_version()}',
}.items()
},
**kwargs,
)

16
tests/test_client.py
View File

@ -968,3 +968,19 @@ async def test_async_client_copy(httpserver: HTTPServer):
client = AsyncClient(httpserver.url_for('/api/copy'))
response = await client.copy('dum', 'dummer')
assert response['status'] == 'success'
def test_headers():
client = Client()
assert client._client.headers['content-type'] == 'application/json'
assert client._client.headers['accept'] == 'application/json'
assert client._client.headers['user-agent'].startswith('ollama-python/')
client = Client(
headers={
'X-Custom': 'value',
'Content-Type': 'text/plain',
}
)
assert client._client.headers['x-custom'] == 'value'
assert client._client.headers['content-type'] == 'application/json'