vectara-agentic 0.4.0__py3-none-any.whl → 0.4.2__py3-none-any.whl

tests/test_agent_memory_consistency.py

@@ -0,0 +1,229 @@
+# Suppress external dependency warnings before any other imports
+import warnings
+
+warnings.simplefilter("ignore", DeprecationWarning)
+
+import unittest
+import threading
+
+from vectara_agentic.agent import Agent, AgentType
+from vectara_agentic.agent_config import AgentConfig
+from vectara_agentic.types import ModelProvider, AgentConfigType
+from vectara_agentic.tools import ToolsFactory
+from llama_index.core.llms import ChatMessage, MessageRole
+from conftest import mult, STANDARD_TEST_TOPIC, STANDARD_TEST_INSTRUCTIONS
+
+
+ARIZE_LOCK = threading.Lock()
+
+
+class TestAgentMemoryConsistency(unittest.TestCase):
+    """Test memory consistency behavior for main/fallback agent switching"""
+
+    def setUp(self):
+        """Set up test fixtures"""
+        self.tools = [ToolsFactory().create_tool(mult)]
+        self.topic = STANDARD_TEST_TOPIC
+        self.custom_instructions = STANDARD_TEST_INSTRUCTIONS
+
+        # Main agent config
+        self.main_config = AgentConfig(
+            agent_type=AgentType.FUNCTION_CALLING,
+            main_llm_provider=ModelProvider.ANTHROPIC,
+        )
+
+        # Fallback agent config
+        self.fallback_config = AgentConfig(
+            agent_type=AgentType.REACT, main_llm_provider=ModelProvider.ANTHROPIC
+        )
+
+        self.session_id = "test-memory-consistency-123"
+
+    def test_agent_recreation_on_config_switch(self):
+        """Test that agent instances are properly recreated when switching configurations"""
+        agent = Agent(
+            tools=self.tools,
+            topic=self.topic,
+            custom_instructions=self.custom_instructions,
+            agent_config=self.main_config,
+            fallback_agent_config=self.fallback_config,
+            session_id=self.session_id,
+        )
+
+        # Load main agent first
+        original_main_agent = agent.agent
+        self.assertIsNotNone(original_main_agent)
+        self.assertEqual(agent.agent_config_type, AgentConfigType.DEFAULT)
+
+        # Switch to fallback - should clear fallback agent instance for recreation
+        agent._switch_agent_config()
+        self.assertEqual(agent.agent_config_type, AgentConfigType.FALLBACK)
+        self.assertIsNone(agent._fallback_agent)  # Should be cleared for recreation
+
+        # Load fallback agent - should be new instance
+        new_fallback_agent = agent.fallback_agent
+        self.assertIsNotNone(new_fallback_agent)
+
+        # Switch back to main - should clear main agent instance for recreation
+        agent._switch_agent_config()
+        self.assertEqual(agent.agent_config_type, AgentConfigType.DEFAULT)
+        self.assertIsNone(agent._agent)  # Should be cleared for recreation
+
+        # Load main agent again - should be new instance
+        recreated_main_agent = agent.agent
+        self.assertIsNotNone(recreated_main_agent)
+        self.assertIsNot(recreated_main_agent, original_main_agent)
+
+    def test_memory_persistence_across_config_switches(self):
+        """Test that Agent memory persists correctly when switching configurations"""
+        agent = Agent(
+            tools=self.tools,
+            topic=self.topic,
+            custom_instructions=self.custom_instructions,
+            agent_config=self.main_config,
+            fallback_agent_config=self.fallback_config,
+            session_id=self.session_id,
+        )
+
+        # Add initial memory
+        initial_messages = [
+            ChatMessage(role=MessageRole.USER, content="Initial question"),
+            ChatMessage(role=MessageRole.ASSISTANT, content="Initial response"),
+        ]
+        agent.memory.put_messages(initial_messages)
+
+        # Verify initial memory
+        self.assertEqual(len(agent.memory.get()), 2)
+        self.assertEqual(agent.memory.get()[0].content, "Initial question")
+
+        # Switch to fallback configuration
+        agent._switch_agent_config()
+        self.assertEqual(agent.agent_config_type, AgentConfigType.FALLBACK)
+
+        # Memory should persist at the Agent level
+        self.assertEqual(len(agent.memory.get()), 2)
+        self.assertEqual(agent.memory.get()[0].content, "Initial question")
+
+        # Add more memory while in fallback mode
+        fallback_messages = [
+            ChatMessage(role=MessageRole.USER, content="Fallback question"),
+            ChatMessage(role=MessageRole.ASSISTANT, content="Fallback response"),
+        ]
+        agent.memory.put_messages(fallback_messages)
+
+        # Verify combined memory
+        self.assertEqual(len(agent.memory.get()), 4)
+        self.assertEqual(agent.memory.get()[2].content, "Fallback question")
+
+        # Switch back to main configuration
+        agent._switch_agent_config()
+        self.assertEqual(agent.agent_config_type, AgentConfigType.DEFAULT)
+
+        # All memory should still be present
+        self.assertEqual(len(agent.memory.get()), 4)
+        self.assertEqual(agent.memory.get()[0].content, "Initial question")
+        self.assertEqual(agent.memory.get()[2].content, "Fallback question")
+
+    def test_clear_memory_resets_agent_instances(self):
+        """Test that clearing memory properly resets agent instances"""
+        agent = Agent(
+            tools=self.tools,
+            topic=self.topic,
+            custom_instructions=self.custom_instructions,
+            agent_config=self.main_config,
+            fallback_agent_config=self.fallback_config,
+            session_id=self.session_id,
+        )
+
+        # Add memory and load both agents
+        test_messages = [
+            ChatMessage(role=MessageRole.USER, content="Test question"),
+            ChatMessage(role=MessageRole.ASSISTANT, content="Test response"),
+        ]
+        agent.memory.put_messages(test_messages)
+
+        # Load both agents
+        _ = agent.agent
+        _ = agent.fallback_agent
+
+        # Verify memory exists
+        self.assertEqual(len(agent.memory.get()), 2)
+
+        # Clear memory
+        agent.clear_memory()
+
+        # Verify memory is cleared
+        self.assertEqual(len(agent.memory.get()), 0)
+
+        # Verify agent instances were reset
+        self.assertIsNone(agent._agent)
+        self.assertIsNone(agent._fallback_agent)
+
+    def test_session_id_consistency(self):
+        """Test that session_id remains consistent throughout agent lifecycle"""
+        agent = Agent(
+            tools=self.tools,
+            topic=self.topic,
+            custom_instructions=self.custom_instructions,
+            agent_config=self.main_config,
+            fallback_agent_config=self.fallback_config,
+            session_id=self.session_id,
+        )
+
+        # Verify initial session_id
+        self.assertEqual(agent.session_id, self.session_id)
+        self.assertEqual(agent.memory.chat_store_key, self.session_id)
+
+        # Switch configurations multiple times
+        agent._switch_agent_config()
+        self.assertEqual(agent.session_id, self.session_id)
+        self.assertEqual(agent.memory.chat_store_key, self.session_id)
+
+        agent._switch_agent_config()
+        self.assertEqual(agent.session_id, self.session_id)
+        self.assertEqual(agent.memory.chat_store_key, self.session_id)
+
+        # Clear memory
+        agent.clear_memory()
+        self.assertEqual(agent.session_id, self.session_id)
+        self.assertEqual(agent.memory.chat_store_key, self.session_id)
+
+    def test_serialization_preserves_consistency(self):
+        """Test that serialization/deserialization preserves memory consistency behavior"""
+        agent = Agent(
+            tools=self.tools,
+            topic=self.topic,
+            custom_instructions=self.custom_instructions,
+            agent_config=self.main_config,
+            fallback_agent_config=self.fallback_config,
+            session_id=self.session_id,
+        )
+
+        # Add memory and switch configurations
+        test_messages = [
+            ChatMessage(role=MessageRole.USER, content="Serialization test"),
+            ChatMessage(role=MessageRole.ASSISTANT, content="Serialization response"),
+        ]
+        agent.memory.put_messages(test_messages)
+        agent._switch_agent_config()  # Switch to fallback
+
+        # Serialize and deserialize
+        serialized_data = agent.dumps()
+        restored_agent = Agent.loads(serialized_data)
+
+        # Verify restored agent has same memory (config type resets to DEFAULT on deserialization)
+        self.assertEqual(restored_agent.session_id, self.session_id)
+        self.assertEqual(len(restored_agent.memory.get()), 2)
+        self.assertEqual(restored_agent.memory.get()[0].content, "Serialization test")
+        self.assertEqual(
+            restored_agent.agent_config_type, AgentConfigType.DEFAULT
+        )  # Resets to default
+
+        # Verify memory consistency behavior is preserved
+        restored_agent._switch_agent_config()  # Switch to fallback
+        self.assertEqual(restored_agent.agent_config_type, AgentConfigType.FALLBACK)
+        self.assertEqual(len(restored_agent.memory.get()), 2)  # Memory should persist
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_agent_type.py

@@ -5,6 +5,10 @@ warnings.simplefilter("ignore", DeprecationWarning)
 
 import unittest
 
+import sys
+import os
+sys.path.insert(0, os.path.dirname(__file__))
+
 from conftest import (
     AgentTestMixin,
     mult,

tests/test_bedrock.py

@@ -3,43 +3,58 @@ import warnings
 warnings.simplefilter("ignore", DeprecationWarning)
 
 import unittest
+import threading
 
-from vectara_agentic.agent import Agent, AgentType
-from vectara_agentic.agent_config import AgentConfig
+from vectara_agentic.agent import Agent
 from vectara_agentic.tools import ToolsFactory
-from vectara_agentic.types import ModelProvider
 
 import nest_asyncio
 nest_asyncio.apply()
 
-
-def mult(x: float, y: float) -> float:
-    "Multiply two numbers"
-    return x * y
-
-
-fc_config_bedrock = AgentConfig(
-    agent_type=AgentType.FUNCTION_CALLING,
-    main_llm_provider=ModelProvider.BEDROCK,
-    tool_llm_provider=ModelProvider.BEDROCK,
-)
-
-class TestBedrock(unittest.TestCase):
-
-    def test_multiturn(self):
-        tools = [ToolsFactory().create_tool(mult)]
-        topic = "AI topic"
-        instructions = "Always do as your father tells you, if your mother agrees!"
-        agent = Agent(
-            tools=tools,
-            topic=topic,
-            custom_instructions=instructions,
-        )
-
-        agent.chat("What is 5 times 10. Only give the answer, nothing else")
-        agent.chat("what is 3 times 7. Only give the answer, nothing else")
-        res = agent.chat("multiply the results of the last two questions. Output only the answer.")
-        self.assertEqual(res.response, "1050")
+from conftest import mult, fc_config_bedrock, STANDARD_TEST_TOPIC, STANDARD_TEST_INSTRUCTIONS
+
+ARIZE_LOCK = threading.Lock()
+
+class TestBedrock(unittest.IsolatedAsyncioTestCase):
+
+    async def test_multiturn(self):
+        with ARIZE_LOCK:
+            tools = [ToolsFactory().create_tool(mult)]
+            agent = Agent(
+                tools=tools,
+                topic=STANDARD_TEST_TOPIC,
+                custom_instructions=STANDARD_TEST_INSTRUCTIONS,
+                agent_config=fc_config_bedrock,
+            )
+
+            # First calculation: 5 * 10 = 50
+            stream1 = await agent.astream_chat(
+                "What is 5 times 10. Only give the answer, nothing else"
+            )
+            # Consume the stream
+            async for chunk in stream1.async_response_gen():
+                pass
+            _ = await stream1.aget_response()
+
+            # Second calculation: 3 * 7 = 21
+            stream2 = await agent.astream_chat(
+                "what is 3 times 7. Only give the answer, nothing else"
+            )
+            # Consume the stream
+            async for chunk in stream2.async_response_gen():
+                pass
+            _ = await stream2.aget_response()
+
+            # Final calculation: 50 * 21 = 1050
+            stream3 = await agent.astream_chat(
+                "multiply the results of the last two questions. Output only the answer."
+            )
+            # Consume the stream
+            async for chunk in stream3.async_response_gen():
+                pass
+            response3 = await stream3.aget_response()
+
+            self.assertEqual(response3.response, "1050")
 
 
 if __name__ == "__main__":

tests/test_fallback.py

@@ -54,7 +54,7 @@ class TestFallback(unittest.TestCase):
         config = AgentConfig(
             agent_type=AgentType.REACT,
             main_llm_provider=ModelProvider.PRIVATE,
-            main_llm_model_name="gpt-4o",
+            main_llm_model_name="gpt-4.1-mini",
             private_llm_api_base=f"http://127.0.0.1:{FLASK_PORT}/v1",
             private_llm_api_key="TEST_API_KEY",
         )

tests/test_gemini.py

@@ -4,15 +4,15 @@ warnings.simplefilter("ignore", DeprecationWarning)
 
 import unittest
 
-from vectara_agentic.agent import Agent, AgentType
-from vectara_agentic.agent_config import AgentConfig
-from vectara_agentic.types import ModelProvider
+from vectara_agentic.agent import Agent
 from vectara_agentic.tools import ToolsFactory
 
 
 import nest_asyncio
 nest_asyncio.apply()
 
+from conftest import mult, fc_config_gemini, STANDARD_TEST_TOPIC, STANDARD_TEST_INSTRUCTIONS
+
 tickers = {
     "C": "Citigroup",
     "COF": "Capital One",
@@ -34,11 +34,6 @@ tickers = {
 years = list(range(2015, 2025))
 
 
-def mult(x: float, y: float) -> float:
-    "Multiply two numbers"
-    return x * y
-
-
 def get_company_info() -> list[str]:
     """
     Returns a dictionary of companies you can query about. Always check this before using any other tool.
@@ -56,23 +51,15 @@ def get_valid_years() -> list[str]:
     return years
 
 
-fc_config_gemini = AgentConfig(
-    agent_type=AgentType.FUNCTION_CALLING,
-    main_llm_provider=ModelProvider.GEMINI,
-    tool_llm_provider=ModelProvider.GEMINI,
-)
-
 class TestGEMINI(unittest.TestCase):
     def test_gemini(self):
         tools = [ToolsFactory().create_tool(mult)]
-        topic = "AI topic"
-        instructions = "Always do as your father tells you, if your mother agrees!"
 
         agent = Agent(
             agent_config=fc_config_gemini,
             tools=tools,
-            topic=topic,
-            custom_instructions=instructions,
+            topic=STANDARD_TEST_TOPIC,
+            custom_instructions=STANDARD_TEST_INSTRUCTIONS,
         )
         _ = agent.chat("What is 5 times 10. Only give the answer, nothing else")
         _ = agent.chat("what is 3 times 7. Only give the answer, nothing else")
@@ -81,14 +68,12 @@ class TestGEMINI(unittest.TestCase):
 
     def test_gemini_single_prompt(self):
         tools = [ToolsFactory().create_tool(mult)]
-        topic = "AI topic"
-        instructions = "Always do as your father tells you, if your mother agrees!"
 
         agent = Agent(
             agent_config=fc_config_gemini,
             tools=tools,
-            topic=topic,
-            custom_instructions=instructions,
+            topic=STANDARD_TEST_TOPIC,
+            custom_instructions=STANDARD_TEST_INSTRUCTIONS,
         )
         res = agent.chat("First, multiply 5 by 10. Then, multiply 3 by 7. Finally, multiply the results of the first two calculations.")
         self.assertIn("1050", res.response)

tests/test_groq.py

@@ -3,43 +3,58 @@ import warnings
 warnings.simplefilter("ignore", DeprecationWarning)
 
 import unittest
+import threading
 
-from vectara_agentic.agent import Agent, AgentType
-from vectara_agentic.agent_config import AgentConfig
+from vectara_agentic.agent import Agent
 from vectara_agentic.tools import ToolsFactory
-from vectara_agentic.types import ModelProvider
 
 import nest_asyncio
 nest_asyncio.apply()
 
-def mult(x: float, y: float) -> float:
-    "Multiply two numbers"
-    return x * y
-
-
-fc_config_groq = AgentConfig(
-    agent_type=AgentType.FUNCTION_CALLING,
-    main_llm_provider=ModelProvider.GROQ,
-    tool_llm_provider=ModelProvider.GROQ,
-)
-
-
-class TestGROQ(unittest.TestCase):
-
-    def test_multiturn(self):
-        tools = [ToolsFactory().create_tool(mult)]
-        topic = "AI topic"
-        instructions = "Always do as your father tells you, if your mother agrees!"
-        agent = Agent(
-            tools=tools,
-            topic=topic,
-            custom_instructions=instructions,
-        )
-
-        agent.chat("What is 5 times 10. Only give the answer, nothing else")
-        agent.chat("what is 3 times 7. Only give the answer, nothing else")
-        res = agent.chat("multiply the results of the last two questions. Output only the answer.")
-        self.assertEqual(res.response, "1050")
+from conftest import mult, fc_config_groq, STANDARD_TEST_TOPIC, STANDARD_TEST_INSTRUCTIONS
+
+ARIZE_LOCK = threading.Lock()
+
+class TestGROQ(unittest.IsolatedAsyncioTestCase):
+
+    async def test_multiturn(self):
+        with ARIZE_LOCK:
+            tools = [ToolsFactory().create_tool(mult)]
+            agent = Agent(
+                tools=tools,
+                topic=STANDARD_TEST_TOPIC,
+                custom_instructions=STANDARD_TEST_INSTRUCTIONS,
+                agent_config=fc_config_groq,
+            )
+
+            # First calculation: 5 * 10 = 50
+            stream1 = await agent.astream_chat(
+                "What is 5 times 10. Only give the answer, nothing else"
+            )
+            # Consume the stream
+            async for chunk in stream1.async_response_gen():
+                pass
+            _ = await stream1.aget_response()
+
+            # Second calculation: 3 * 7 = 21
+            stream2 = await agent.astream_chat(
+                "what is 3 times 7. Only give the answer, nothing else"
+            )
+            # Consume the stream
+            async for chunk in stream2.async_response_gen():
+                pass
+            _ = await stream2.aget_response()
+
+            # Final calculation: 50 * 21 = 1050
+            stream3 = await agent.astream_chat(
+                "multiply the results of the last two questions. Output only the answer."
+            )
+            # Consume the stream
+            async for chunk in stream3.async_response_gen():
+                pass
+            response3 = await stream3.aget_response()
+
+            self.assertEqual(response3.response, "1050")
 
 
 if __name__ == "__main__":

tests/test_private_llm.py

@@ -54,7 +54,7 @@ class TestPrivateLLM(unittest.TestCase):
         config = AgentConfig(
             agent_type=AgentType.FUNCTION_CALLING,
             main_llm_provider=ModelProvider.PRIVATE,
-            main_llm_model_name="gpt-4.1",
+            main_llm_model_name="gpt-4.1-mini",
             private_llm_api_base=f"http://127.0.0.1:{FLASK_PORT}/v1",
             private_llm_api_key="TEST_API_KEY",
         )

tests/test_serialization.py

@@ -14,8 +14,7 @@ from vectara_agentic.tools import ToolsFactory
 from llama_index.core.utilities.sql_wrapper import SQLDatabase
 from sqlalchemy import create_engine
 
-def mult(x: float, y: float) -> float:
-    return x * y
+from conftest import mult, STANDARD_TEST_TOPIC, STANDARD_TEST_INSTRUCTIONS
 
 
 ARIZE_LOCK = threading.Lock()
@@ -44,12 +43,10 @@ class TestAgentSerialization(unittest.TestCase):
             )
 
             tools = [ToolsFactory().create_tool(mult)] + ToolsFactory().standard_tools() + db_tools
-            topic = "AI topic"
-            instructions = "Always do as your father tells you, if your mother agrees!"
             agent = Agent(
                 tools=tools,
-                topic=topic,
-                custom_instructions=instructions,
+                topic=STANDARD_TEST_TOPIC,
+                custom_instructions=STANDARD_TEST_INSTRUCTIONS,
                 agent_config=config
             )