Reverting the tools commit (#2404)

tests/llms/test_azure_openai.py

@@ -20,10 +20,8 @@ def mock_openai_client():
         yield mock_client
 
 
-def test_generate_response(mock_openai_client):
-    config = BaseLlmConfig(
-        model=MODEL, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, top_p=TOP_P
-    )
+def test_generate_response_without_tools(mock_openai_client):
+    config = BaseLlmConfig(model=MODEL, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, top_p=TOP_P)
     llm = AzureOpenAILLM(config)
     messages = [
         {"role": "system", "content": "You are a helpful assistant."},
@@ -31,21 +29,67 @@ def test_generate_response(mock_openai_client):
     ]
 
     mock_response = Mock()
-    mock_response.choices = [
-        Mock(message=Mock(content="I'm doing well, thank you for asking!"))
-    ]
+    mock_response.choices = [Mock(message=Mock(content="I'm doing well, thank you for asking!"))]
     mock_openai_client.chat.completions.create.return_value = mock_response
 
     response = llm.generate_response(messages)
 
+    mock_openai_client.chat.completions.create.assert_called_once_with(
+        model=MODEL, messages=messages, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, top_p=TOP_P
+    )
+    assert response == "I'm doing well, thank you for asking!"
+
+
+def test_generate_response_with_tools(mock_openai_client):
+    config = BaseLlmConfig(model=MODEL, temperature=TEMPERATURE, max_tokens=MAX_TOKENS, top_p=TOP_P)
+    llm = AzureOpenAILLM(config)
+    messages = [
+        {"role": "system", "content": "You are a helpful assistant."},
+        {"role": "user", "content": "Add a new memory: Today is a sunny day."},
+    ]
+    tools = [
+        {
+            "type": "function",
+            "function": {
+                "name": "add_memory",
+                "description": "Add a memory",
+                "parameters": {
+                    "type": "object",
+                    "properties": {"data": {"type": "string", "description": "Data to add to memory"}},
+                    "required": ["data"],
+                },
+            },
+        }
+    ]
+
+    mock_response = Mock()
+    mock_message = Mock()
+    mock_message.content = "I've added the memory for you."
+
+    mock_tool_call = Mock()
+    mock_tool_call.function.name = "add_memory"
+    mock_tool_call.function.arguments = '{"data": "Today is a sunny day."}'
+
+    mock_message.tool_calls = [mock_tool_call]
+    mock_response.choices = [Mock(message=mock_message)]
+    mock_openai_client.chat.completions.create.return_value = mock_response
+
+    response = llm.generate_response(messages, tools=tools)
+
     mock_openai_client.chat.completions.create.assert_called_once_with(
         model=MODEL,
         messages=messages,
         temperature=TEMPERATURE,
         max_tokens=MAX_TOKENS,
         top_p=TOP_P,
+        tools=tools,
+        tool_choice="auto",
     )
-    assert response == "I'm doing well, thank you for asking!"
+
+    assert response["content"] == "I've added the memory for you."
+    assert len(response["tool_calls"]) == 1
+    assert response["tool_calls"][0]["name"] == "add_memory"
+    assert response["tool_calls"][0]["arguments"] == {"data": "Today is a sunny day."}
 
 
 @pytest.mark.parametrize(
@@ -84,6 +128,4 @@ def test_generate_with_http_proxies(default_headers):
             api_version=None,
             default_headers=default_headers,
         )
-        mock_http_client.assert_called_once_with(
-            proxies="http://testproxy.mem0.net:8000"
-        )
+        mock_http_client.assert_called_once_with(proxies="http://testproxy.mem0.net:8000")