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

Fix/ollama test cases (#1639)

Browse Source
This commit is contained in:
dbcontributions
2024-08-04 12:56:25 +05:30
committed by GitHub
parent 5aa7bedabe
commit b6cfd960d1
1 changed files with 15 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
tests/llms/test_ollama.py
View File

@@ -12,7 +12,6 @@ def mock_ollama_client():
mock_ollama.return_value = mock_client mock_ollama.return_value = mock_client
yield mock_client yield mock_client
@pytest.mark.skip(reason="Mock issue, need to be fixed")
def test_generate_response_without_tools(mock_ollama_client): def test_generate_response_without_tools(mock_ollama_client):
config = BaseLlmConfig(model="llama3.1:70b", temperature=0.7, max_tokens=100, top_p=1.0) config = BaseLlmConfig(model="llama3.1:70b", temperature=0.7, max_tokens=100, top_p=1.0)
llm = OllamaLLM(config) llm = OllamaLLM(config)
@@ -21,8 +20,9 @@ def test_generate_response_without_tools(mock_ollama_client):
{"role": "user", "content": "Hello, how are you?"} {"role": "user", "content": "Hello, how are you?"}
] ]
mock_response = Mock() mock_response = {
mock_response.message = {"content": "I'm doing well, thank you for asking!"} 'message': {"content": "I'm doing well, thank you for asking!"}
}
mock_ollama_client.chat.return_value = mock_response mock_ollama_client.chat.return_value = mock_response
response = llm.generate_response(messages) response = llm.generate_response(messages)
@@ -38,7 +38,6 @@ def test_generate_response_without_tools(mock_ollama_client):
) )
assert response == "I'm doing well, thank you for asking!" assert response == "I'm doing well, thank you for asking!"
@pytest.mark.skip(reason="Mock issue, need to be fixed")
def test_generate_response_with_tools(mock_ollama_client): def test_generate_response_with_tools(mock_ollama_client):
config = BaseLlmConfig(model="llama3.1:70b", temperature=0.7, max_tokens=100, top_p=1.0) config = BaseLlmConfig(model="llama3.1:70b", temperature=0.7, max_tokens=100, top_p=1.0)
llm = OllamaLLM(config) llm = OllamaLLM(config)
@@ -48,18 +47,20 @@ def test_generate_response_with_tools(mock_ollama_client):
] ]
tools = [ADD_MEMORY_TOOL] tools = [ADD_MEMORY_TOOL]
mock_response = Mock() mock_response = {
mock_message = {"content": "I've added the memory for you."} 'message': {
"content": "I've added the memory for you.",
mock_tool_call = { "tool_calls": [
"function": { {
"name": "add_memory", "function": {
"arguments": '{"data": "Today is a sunny day."}' "name": "add_memory",
"arguments": {"data": "Today is a sunny day."}
}
}
]
} }
} }
mock_message["tool_calls"] = [mock_tool_call]
mock_response.message = mock_message
mock_ollama_client.chat.return_value = mock_response mock_ollama_client.chat.return_value = mock_response
response = llm.generate_response(messages, tools=tools) response = llm.generate_response(messages, tools=tools)
@@ -79,3 +80,4 @@ def test_generate_response_with_tools(mock_ollama_client):
assert len(response["tool_calls"]) == 1 assert len(response["tool_calls"]) == 1
assert response["tool_calls"][0]["name"] == "add_memory" assert response["tool_calls"][0]["name"] == "add_memory"
assert response["tool_calls"][0]["arguments"] == {'data': 'Today is a sunny day.'} assert response["tool_calls"][0]["arguments"] == {'data': 'Today is a sunny day.'}