Skip to content

Commit

Permalink
community: Add keep_alive parameter to control how long the model w… (
Browse files Browse the repository at this point in the history
#19005)

Add `keep_alive` parameter to control how long the model will stay
loaded into memory with Ollama。

---------

Co-authored-by: Bagatur <baskaryan@gmail.com>
  • Loading branch information
2 people authored and hinthornw committed Apr 26, 2024
1 parent fe88710 commit e513b50
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 1 deletion.
15 changes: 14 additions & 1 deletion libs/community/langchain_community/llms/ollama.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import json
from typing import Any, AsyncIterator, Dict, Iterator, List, Mapping, Optional
from typing import Any, AsyncIterator, Dict, Iterator, List, Mapping, Optional, Union

import aiohttp
import requests
Expand Down Expand Up @@ -111,6 +111,18 @@ class _OllamaCommon(BaseLanguageModel):
timeout: Optional[int] = None
"""Timeout for the request stream"""

keep_alive: Optional[Union[int, str]] = None
"""How long the model will stay loaded into memory.
The parameter (Default: 5 minutes) can be set to:
1. a duration string in Golang (such as "10m" or "24h");
2. a number in seconds (such as 3600);
3. any negative number which will keep the model loaded \
in memory (e.g. -1 or "-1m");
4. 0 which will unload the model immediately after generating a response;
See the [Ollama documents](https://github.com/ollama/ollama/blob/main/docs/faq.md#how-do-i-keep-a-model-loaded-in-memory-or-make-it-unload-immediately)"""

headers: Optional[dict] = None
"""Additional headers to pass to endpoint (e.g. Authorization, Referer).
This is useful when Ollama is hosted on cloud services that require
Expand Down Expand Up @@ -141,6 +153,7 @@ def _default_params(self) -> Dict[str, Any]:
},
"system": self.system,
"template": self.template,
"keep_alive": self.keep_alive,
}

@property
Expand Down
3 changes: 3 additions & 0 deletions libs/community/tests/unit_tests/llms/test_ollama.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,7 @@ def mock_post(url, headers, json, stream, timeout): # type: ignore[no-untyped-d
"prompt": "Test prompt",
"system": "Test system prompt",
"template": None,
"keep_alive": None,
}
assert stream is True
assert timeout == 300
Expand Down Expand Up @@ -147,6 +148,7 @@ def mock_post(url, headers, json, stream, timeout): # type: ignore[no-untyped-d
"prompt": "Test prompt",
"system": None,
"template": None,
"keep_alive": None,
}
assert stream is True
assert timeout == 300
Expand Down Expand Up @@ -178,6 +180,7 @@ def mock_post(url, headers, json, stream, timeout): # type: ignore[no-untyped-d
"prompt": "Test prompt",
"system": None,
"template": None,
"keep_alive": None,
}
assert stream is True
assert timeout == 300
Expand Down

0 comments on commit e513b50

Please sign in to comment.