naas-abi-core 1.4.1__py3-none-any.whl

naas_abi_core/services/agent/test_agent_memory.py

@@ -0,0 +1,278 @@
+"""Tests for Agent memory configuration with PostgreSQL support."""
+
+import os
+from unittest.mock import ANY, MagicMock, patch
+
+import pytest
+from langgraph.checkpoint.base import BaseCheckpointSaver
+from langgraph.checkpoint.memory import MemorySaver
+from naas_abi_core.services.agent.Agent import (
+    Agent,
+    AgentSharedState,
+    create_checkpointer,
+)
+
+
+class TestCreateCheckpointer:
+    """Test the create_checkpointer function that detects PostgreSQL configuration."""
+
+    def test_create_checkpointer_without_postgres_url(self):
+        """Test that MemorySaver is returned when POSTGRES_URL is not set."""
+        with patch.dict(os.environ, {}, clear=True):
+            checkpointer = create_checkpointer()
+            assert isinstance(checkpointer, MemorySaver)
+
+    def test_create_checkpointer_with_postgres_url(self):
+        """Test that PostgresSaver is returned when POSTGRES_URL is set."""
+        test_url = "postgresql://user:pass@localhost:5432/testdb"
+
+        # Mock the PostgresSaver class, psycopg connection, and instance
+        mock_postgres_saver = MagicMock()
+        mock_postgres_saver.setup = MagicMock()
+        mock_postgres_class = MagicMock(return_value=mock_postgres_saver)
+        mock_connection = MagicMock()
+        mock_connection_connect = MagicMock(return_value=mock_connection)
+
+        with patch.dict(os.environ, {"POSTGRES_URL": test_url}):
+            with patch(
+                "langgraph.checkpoint.postgres.PostgresSaver", mock_postgres_class
+            ):
+                with patch("psycopg.Connection.connect", mock_connection_connect):
+                    checkpointer = create_checkpointer()
+
+                    # Verify Connection.connect was called with proper parameters
+                    mock_connection_connect.assert_called_once_with(
+                        test_url,
+                        autocommit=True,
+                        prepare_threshold=0,
+                        row_factory=ANY,  # dict_row import is mocked
+                    )
+                    # Verify PostgresSaver constructor was called with connection
+                    mock_postgres_class.assert_called_once_with(mock_connection)
+                    # Verify setup() was called
+                    mock_postgres_saver.setup.assert_called_once()
+                    assert checkpointer == mock_postgres_saver
+
+    def test_create_checkpointer_postgres_import_error(self):
+        """Test fallback to MemorySaver when PostgresSaver import fails."""
+        test_url = "postgresql://user:pass@localhost:5432/testdb"
+
+        with patch.dict(os.environ, {"POSTGRES_URL": test_url}):
+            # Simulate ImportError when trying to import PostgresSaver
+            with patch(
+                "builtins.__import__",
+                side_effect=ImportError(
+                    "No module named 'langgraph.checkpoint.postgres'"
+                ),
+            ):
+                checkpointer = create_checkpointer()
+                assert isinstance(checkpointer, MemorySaver)
+
+    def test_create_checkpointer_postgres_connection_error(self):
+        """Test fallback to MemorySaver when PostgreSQL connection fails."""
+        test_url = "postgresql://user:pass@localhost:5432/testdb"
+
+        # Mock Connection.connect to raise an exception
+        mock_connection_connect = MagicMock(side_effect=Exception("Connection failed"))
+
+        with patch.dict(os.environ, {"POSTGRES_URL": test_url}):
+            with patch("psycopg.Connection.connect", mock_connection_connect):
+                checkpointer = create_checkpointer()
+                assert isinstance(checkpointer, MemorySaver)
+
+    def test_create_checkpointer_postgres_setup_error(self):
+        """Test fallback to MemorySaver when PostgreSQL setup fails."""
+        test_url = "postgresql://user:pass@localhost:5432/testdb"
+
+        # Mock successful connection but setup failure
+        mock_postgres_saver = MagicMock()
+        mock_postgres_saver.setup.side_effect = Exception("Table creation failed")
+        mock_postgres_class = MagicMock(return_value=mock_postgres_saver)
+        mock_connection = MagicMock()
+        mock_connection_connect = MagicMock(return_value=mock_connection)
+
+        with patch.dict(os.environ, {"POSTGRES_URL": test_url}):
+            with patch(
+                "langgraph.checkpoint.postgres.PostgresSaver", mock_postgres_class
+            ):
+                with patch("psycopg.Connection.connect", mock_connection_connect):
+                    checkpointer = create_checkpointer()
+                    assert isinstance(checkpointer, MemorySaver)
+
+
+class TestAgentMemoryConfiguration:
+    """Test Agent class memory configuration."""
+
+    @pytest.fixture
+    def mock_chat_model(self):
+        """Create a mock chat model."""
+        model = MagicMock()
+        model.bind_tools = MagicMock(return_value=model)
+        return model
+
+    def test_agent_uses_provided_memory(self, mock_chat_model):
+        """Test that Agent uses explicitly provided memory."""
+        custom_memory = MagicMock(spec=BaseCheckpointSaver)
+
+        agent = Agent(
+            name="test_agent",
+            description="Test agent",
+            chat_model=mock_chat_model,
+            memory=custom_memory,
+        )
+
+        assert agent._checkpointer == custom_memory
+
+    def test_agent_creates_memory_when_none_provided(self, mock_chat_model):
+        """Test that Agent creates memory based on environment when None is provided."""
+        with patch("abi.services.agent.Agent.create_checkpointer") as mock_create:
+            mock_checkpointer = MagicMock(spec=BaseCheckpointSaver)
+            mock_create.return_value = mock_checkpointer
+
+            agent = Agent(
+                name="test_agent",
+                description="Test agent",
+                chat_model=mock_chat_model,
+                memory=None,
+            )
+
+            mock_create.assert_called_once()
+            assert agent._checkpointer == mock_checkpointer
+
+    def test_agent_default_memory_creation(self, mock_chat_model):
+        """Test that Agent creates memory automatically when not provided."""
+        with patch("abi.services.agent.Agent.create_checkpointer") as mock_create:
+            mock_checkpointer = MagicMock(spec=BaseCheckpointSaver)
+            mock_create.return_value = mock_checkpointer
+
+            # Don't provide memory parameter at all
+            agent = Agent(
+                name="test_agent", description="Test agent", chat_model=mock_chat_model
+            )
+
+            mock_create.assert_called_once()
+            assert agent._checkpointer == mock_checkpointer
+
+
+class TestAgentPostgresIntegration:
+    """Integration tests for Agent with PostgreSQL checkpointer."""
+
+    @pytest.fixture
+    def mock_chat_model(self):
+        """Create a mock chat model that simulates real behavior."""
+        model = MagicMock()
+        model.bind_tools = MagicMock(return_value=model)
+
+        # Mock invoke to return a proper message
+        from langchain_core.messages import AIMessage
+
+        model.invoke = MagicMock(return_value=AIMessage(content="Test response"))
+
+        return model
+
+    @pytest.mark.integration
+    def test_agent_with_postgres_preserves_state(self, mock_chat_model):
+        """Test that Agent with PostgreSQL preserves conversation state."""
+        # This test would require a running PostgreSQL instance
+        # It's marked as integration test and can be skipped in CI
+
+        test_url = "postgresql://abi_user:abi_password@localhost:5432/abi_memory"
+
+        with patch.dict(os.environ, {"POSTGRES_URL": test_url}):
+            # First agent instance
+            agent1 = Agent(
+                name="test_agent",
+                description="Test agent",
+                chat_model=mock_chat_model,
+                state=AgentSharedState(thread_id="123"),
+            )
+
+            # Simulate some conversation
+            # In a real test, we would invoke the agent and check state persistence
+
+            # Second agent instance with same thread_id should have access to same state
+            agent2 = Agent(
+                name="test_agent",
+                description="Test agent",
+                chat_model=mock_chat_model,
+                state=AgentSharedState(thread_id="123"),
+            )
+
+            # Both agents should share the same checkpointer type
+            assert type(agent1._checkpointer) is type(agent2._checkpointer)
+
+    def test_agent_duplicate_preserves_memory_config(self, mock_chat_model):
+        """Test that duplicating an agent preserves memory configuration."""
+        custom_memory = MagicMock(spec=BaseCheckpointSaver)
+
+        original_agent = Agent(
+            name="test_agent",
+            description="Test agent",
+            chat_model=mock_chat_model,
+            memory=custom_memory,
+        )
+
+        duplicated_agent = original_agent.duplicate()
+
+        # The duplicated agent should use the same memory type
+        assert duplicated_agent._checkpointer == original_agent._checkpointer
+
+    def test_agent_thread_id_passed_as_string(self, mock_chat_model):
+        """Test that thread_id is passed as string to graph config."""
+        # Mock the graph.stream method to capture the config
+        mock_graph = MagicMock()
+        mock_graph.stream.return_value = iter([("source", {"messages": []})])
+
+        agent = Agent(
+            name="test_agent", description="Test agent", chat_model=mock_chat_model
+        )
+        agent.graph = mock_graph
+
+        # Set a specific thread_id
+        agent._state.set_thread_id("123")
+
+        # Call stream which should pass thread_id as string directly
+        list(agent.stream("test prompt"))
+
+        # Verify that graph.stream was called with thread_id as string
+        mock_graph.stream.assert_called_once()
+        call_args = mock_graph.stream.call_args
+        config = call_args[1]["config"]  # kwargs
+
+        assert config["configurable"]["thread_id"] == "123"
+        assert isinstance(config["configurable"]["thread_id"], str)
+
+    def test_agent_thread_id_increment(self, mock_chat_model):
+        """Test that thread_id increments correctly with string type."""
+        # Create agent with explicit thread_id to avoid test interference
+        from naas_abi_core.services.agent.Agent import AgentSharedState
+
+        state = AgentSharedState(thread_id="1")
+
+        agent = Agent(
+            name="test_agent",
+            description="Test agent",
+            chat_model=mock_chat_model,
+            state=state,
+        )
+
+        # Default thread_id should be "1"
+        assert agent._state.thread_id == "1"
+
+        # Test increment
+        agent.reset()
+        assert agent._state.thread_id == "2"
+
+        # Test multiple increments
+        agent.reset()
+        agent.reset()
+        assert agent._state.thread_id == "4"
+
+        # Test setting a custom string thread_id and incrementing
+        agent._state.set_thread_id("100")
+        agent.reset()
+        assert agent._state.thread_id == "101"
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, "-v"])

naas_abi_core/services/agent/test_postgres_integration.py

@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+"""Integration test to verify PostgreSQL memory works with Agent.
+
+This test requires PostgreSQL to be running. It can be started with:
+    make dev-up
+
+Run this test with:
+    POSTGRES_URL=postgresql://abi_user:abi_password@localhost:5432/abi_memory python test_postgres_integration.py
+"""
+
+import os
+import sys
+from unittest.mock import MagicMock
+
+from langchain_core.messages import AIMessage
+
+# Add parent directory to path to import Agent
+sys.path.insert(
+    0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
+)
+
+from naas_abi_core.services.agent.Agent import (
+    Agent,
+    AgentSharedState,
+    create_checkpointer,
+)
+
+
+def test_postgres_checkpointer():
+    """Test that PostgreSQL checkpointer is created when POSTGRES_URL is set."""
+    postgres_url = os.getenv("POSTGRES_URL")
+
+    if not postgres_url:
+        print("⚠️  POSTGRES_URL not set. Skipping PostgreSQL integration test.")
+        print("   To run this test, start PostgreSQL with: make dev-up")
+        print(
+            "   Then set POSTGRES_URL=postgresql://abi_user:abi_password@localhost:5432/abi_memory"
+        )
+        return False
+
+    print(f"✓ POSTGRES_URL found: {postgres_url}")
+
+    try:
+        checkpointer = create_checkpointer()
+        print(f"✓ Created checkpointer: {type(checkpointer).__name__}")
+
+        # Check if it's a PostgreSQL checkpointer
+        if "PostgresSaver" in type(checkpointer).__name__:
+            print("✓ PostgreSQL checkpointer successfully created!")
+            print(
+                "✓ Connection configured with autocommit=True, prepare_threshold=0, row_factory=dict_row"
+            )
+            print("✓ Database tables initialized with setup()")
+            return True
+        else:
+            print(f"⚠️  Expected PostgresSaver but got {type(checkpointer).__name__}")
+            return False
+
+    except Exception as e:
+        print(f"✗ Error creating checkpointer: {e}")
+        return False
+
+
+def test_agent_with_postgres():
+    """Test that Agent can use PostgreSQL for memory persistence."""
+    postgres_url = os.getenv("POSTGRES_URL")
+
+    if not postgres_url:
+        print("⚠️  Skipping Agent PostgreSQL test (POSTGRES_URL not set)")
+        return False
+
+    try:
+        # Create a mock chat model
+        mock_model = MagicMock()
+        mock_model.bind_tools = MagicMock(return_value=mock_model)
+        mock_model.invoke = MagicMock(return_value=AIMessage(content="Test response"))
+
+        # Create agent with a specific thread ID
+        thread_id = 12345
+        agent = Agent(
+            name="test_agent",
+            description="Test agent for PostgreSQL integration",
+            chat_model=mock_model,
+            state=AgentSharedState(thread_id=thread_id),
+        )
+
+        print(f"✓ Created agent with thread_id={thread_id}")
+        print(f"✓ Agent using checkpointer: {type(agent._checkpointer).__name__}")
+
+        # Test that the agent's checkpointer is PostgreSQL-based
+        if "PostgresSaver" in type(agent._checkpointer).__name__:
+            print("✓ Agent successfully using PostgreSQL for memory!")
+            return True
+        else:
+            print(
+                f"⚠️  Agent using {type(agent._checkpointer).__name__} instead of PostgresSaver"
+            )
+            return False
+
+    except Exception as e:
+        print(f"✗ Error creating agent with PostgreSQL: {e}")
+        import traceback
+
+        traceback.print_exc()
+        return False
+
+
+def main():
+    """Run all integration tests."""
+    print("=" * 60)
+    print("PostgreSQL Integration Tests for Agent Memory")
+    print("=" * 60)
+    print()
+
+    tests_passed = 0
+    tests_total = 2
+
+    # Test 1: Checkpointer creation
+    print("Test 1: PostgreSQL Checkpointer Creation")
+    print("-" * 40)
+    if test_postgres_checkpointer():
+        tests_passed += 1
+    print()
+
+    # Test 2: Agent with PostgreSQL
+    print("Test 2: Agent with PostgreSQL Memory")
+    print("-" * 40)
+    if test_agent_with_postgres():
+        tests_passed += 1
+    print()
+
+    # Summary
+    print("=" * 60)
+    if tests_passed == tests_total:
+        print(f"✅ All tests passed! ({tests_passed}/{tests_total})")
+    else:
+        print(f"⚠️  {tests_passed}/{tests_total} tests passed")
+    print("=" * 60)
+
+    return tests_passed == tests_total
+
+
+if __name__ == "__main__":
+    success = main()
+    sys.exit(0 if success else 1)

naas_abi_core/services/cache/CacheFactory.py

@@ -0,0 +1,31 @@
+import os
+
+from naas_abi_core.services.cache.adapters.secondary.CacheFSAdapter import (
+    CacheFSAdapter,
+)
+from naas_abi_core.services.cache.CacheService import CacheService
+from naas_abi_core.utils.Storage import NoStorageFolderFound, find_storage_folder
+
+
+class CacheFactory:
+    @staticmethod
+    def CacheFS_find_storage(
+        subpath: str = "cache", needle: str = "storage"
+    ) -> CacheService:
+        if not subpath.startswith("cache"):
+            subpath = os.path.join("cache", subpath)
+
+        try:
+            return CacheService(
+                CacheFSAdapter(
+                    os.path.join(find_storage_folder(os.getcwd(), needle), subpath)
+                )
+            )
+        except NoStorageFolderFound as _:
+            # Create a "storage" folder for the cache
+            os.makedirs(os.path.join(os.getcwd(), "storage"), exist_ok=True)
+            return CacheService(
+                CacheFSAdapter(
+                    os.path.join(find_storage_folder(os.getcwd(), needle), subpath)
+                )
+            )

naas_abi_core/services/cache/CachePort.py

@@ -0,0 +1,63 @@
+from pydantic import BaseModel, Field
+import datetime
+from enum import Enum
+from typing import Any
+
+class CacheNotFoundError(Exception):
+    pass
+
+class CacheExpiredError(Exception):
+    pass
+
+class DataType(str, Enum):
+    TEXT = "text"
+    JSON = "json"
+    BINARY = "binary"
+    PICKLE = "pickle"
+
+class CachedData(BaseModel):
+    key: str
+    data: Any
+    data_type: DataType
+    created_at: str = Field(default_factory=lambda: datetime.datetime.now().isoformat())
+
+class ICacheAdapter:
+    
+    def get(self, key: str) -> CachedData:
+        raise NotImplementedError("Not implemented")
+    
+    def set(self, key: str, value: CachedData) -> None:
+        raise NotImplementedError("Not implemented")
+    
+    def delete(self, key: str) -> None:
+        raise NotImplementedError("Not implemented")
+    
+    def exists(self, key: str) -> bool:
+        raise NotImplementedError("Not implemented")
+
+class ICacheService:
+    
+    adapter: ICacheAdapter
+    
+    def __init__(self, adapter: ICacheAdapter):
+        self.adapter = adapter
+    
+    def exists(self, key: str) -> bool:
+        raise NotImplementedError("Not implemented")
+    
+    def get(self, key: str, ttl: datetime.timedelta | None = None) -> Any:
+        raise NotImplementedError("Not implemented")
+    
+    def set_text(self, key: str, value: str) -> None:
+        raise NotImplementedError("Not implemented")
+    
+    def set_json(self, key: str, value: dict) -> None:
+        raise NotImplementedError("Not implemented")
+    
+    def set_binary(self, key: str, value: bytes) -> None:
+        raise NotImplementedError("Not implemented")
+
+    def set_pickle(self, key: str, value: Any) -> None:
+        raise NotImplementedError("Not implemented")
+    
+