klas/t6_mem0
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Download Ollama model if not present (#1397)

Browse Source
This commit is contained in:
Dev Khant
2024-06-09 12:13:03 +05:30
committed by GitHub
parent 6fa946557f
commit 73e53aaff1
3 changed files with 13 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
tests/llm/test_ollama.py
View File

@@ -11,6 +11,7 @@ def ollama_llm_config():
def test_get_llm_model_answer(ollama_llm_config, mocker):
mocker.patch("embedchain.llm.ollama.Client.list", return_value={"models": [{"name": "llama2"}]})
mocker.patch("embedchain.llm.ollama.OllamaLlm._get_answer", return_value="Test answer")
llm = OllamaLlm(ollama_llm_config)
@@ -20,6 +21,7 @@ def test_get_llm_model_answer(ollama_llm_config, mocker):
def test_get_answer_mocked_ollama(ollama_llm_config, mocker):
mocker.patch("embedchain.llm.ollama.Client.list", return_value={"models": [{"name": "llama2"}]})
mocked_ollama = mocker.patch("embedchain.llm.ollama.Ollama")
mock_instance = mocked_ollama.return_value
mock_instance.invoke.return_value = "Mocked answer"