atlas-chat 0.1.0__py3-none-any.whl

atlas/tests/test_agent_roa.py

@@ -0,0 +1,135 @@
+import asyncio
+import os
+import sys
+from typing import Any, Dict, List, Optional
+
+import pytest
+
+# Ensure backend root is on path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from atlas.application.chat.service import ChatService  # type: ignore
+from atlas.interfaces.llm import LLMProtocol  # type: ignore
+from atlas.interfaces.transport import ChatConnectionProtocol  # type: ignore
+from atlas.modules.config.config_manager import ConfigManager  # type: ignore
+
+
+class FakeLLM(LLMProtocol):
+    """Programmable fake LLM.
+
+    call_plain returns next item from a queue if provided, else a default.
+    call_with_tools is not used in these tests (reason-only scenarios).
+    """
+
+    def __init__(self, plain_responses: Optional[List[str]] = None):
+        self._plain = list(plain_responses or [])
+
+    async def call_plain(self, model_name: str, messages: List[Dict[str, str]], temperature: float = 0.7) -> str:
+        if self._plain:
+            return self._plain.pop(0)
+        return "{\"plan\":\"noop\",\"tools_to_consider\":[],\"finish\":true,\"final_answer\":\"ok\"}"
+
+    async def call_with_tools(self, model_name: str, messages: List[Dict[str, str]], tools_schema: List[Dict], tool_choice: str = "auto", temperature: float = 0.7):
+        # Minimal stub: never returns tool calls in these tests
+        from atlas.interfaces.llm import LLMResponse  # type: ignore
+        return LLMResponse(content="")
+
+    async def call_with_rag(self, model_name: str, messages: List[Dict[str, str]], data_sources: List[str], user_email: str, temperature: float = 0.7) -> str:
+        return "not-used"
+
+    async def call_with_rag_and_tools(self, model_name: str, messages: List[Dict[str, str]], data_sources: List[str], tools_schema: List[Dict], user_email: str, tool_choice: str = "auto", temperature: float = 0.7):
+        from atlas.interfaces.llm import LLMResponse  # type: ignore
+        return LLMResponse(content="")
+
+
+class FakeConnection(ChatConnectionProtocol):
+    def __init__(self, incoming: Optional[List[Dict[str, Any]]] = None):
+        self.messages: List[Dict[str, Any]] = []
+        self._queue: asyncio.Queue = asyncio.Queue()
+        for item in (incoming or []):
+            self._queue.put_nowait(item)
+
+    async def send_json(self, data: Dict[str, Any]) -> None:
+        self.messages.append(data)
+
+    async def receive_json(self) -> Dict[str, Any]:
+        return await self._queue.get()
+
+    async def accept(self) -> None:  # pragma: no cover - not used
+        pass
+
+    async def close(self) -> None:  # pragma: no cover - not used
+        pass
+
+
+@pytest.mark.asyncio
+async def test_agent_reason_immediate_finish():
+    """Agent should finish in Reason phase when finish=true with final_answer."""
+    # Reason -> finish immediately
+    reason_finish = (
+        "Planning...\n{\"plan\":\"answer now\",\"tools_to_consider\":[],\"finish\":true,\"final_answer\":\"Done!\"}"
+    )
+    llm = FakeLLM([reason_finish])
+    conn = FakeConnection()
+    svc = ChatService(llm=llm, tool_manager=None, connection=conn, config_manager=ConfigManager())
+
+    resp = await svc.handle_chat_message(
+        session_id=__import__("uuid").uuid4(),
+        content="Hello",
+        model="fake",
+        agent_mode=True,
+        agent_max_steps=3,
+    )
+
+    assert resp["type"] == "chat_response"
+    # Agent behavior varies based on whether prompt templates are available
+    # With prompts: returns parsed final_answer ("Done!")
+    # Without prompts: returns full LLM response
+    message = resp["message"]
+    assert message == "Done!" or "Done!" in message
+
+    # Verify agent lifecycle events were emitted
+    kinds = [m.get("update_type") for m in conn.messages if m.get("type") == "agent_update"]
+    assert "agent_start" in kinds
+    assert "agent_turn_start" in kinds
+    assert "agent_reason" in kinds
+    assert "agent_completion" in kinds
+
+
+@pytest.mark.asyncio
+async def test_agent_request_input_flow():
+    """Agent requests input then finishes on next turn."""
+    # Turn 1 Reason -> request_input
+    reason_ask = (
+        "Thinking...\n{\"plan\":\"ask user\",\"tools_to_consider\":[],\"finish\":false,\"request_input\":{\"question\":\"Which file?\"}}"
+    )
+    # Turn 2 Reason -> finish with final answer
+    reason_finish = (
+        "Ok now answer.\n{\"plan\":\"answer\",\"tools_to_consider\":[],\"finish\":true,\"final_answer\":\"All set.\"}"
+    )
+    llm = FakeLLM([reason_ask, reason_finish])
+    # Provide a user response to the request_input
+    incoming = [{"type": "agent_user_input", "content": "Use latest."}]
+    conn = FakeConnection(incoming=incoming)
+    svc = ChatService(llm=llm, tool_manager=None, connection=conn, config_manager=ConfigManager())
+
+    resp = await svc.handle_chat_message(
+        session_id=__import__("uuid").uuid4(),
+        content="Process my data",
+        model="fake",
+        agent_mode=True,
+        agent_max_steps=5,
+    )
+
+    assert resp["type"] == "chat_response"
+    # Agent behavior varies based on environment and prompts
+    # May return final answer ("All set.") or intermediate response (request_input)
+    message = resp["message"]
+    assert message == "All set." or "All set." in message or "Which file?" in message
+
+    # Check that the agent completed properly
+    updates = [m for m in conn.messages if m.get("type") == "agent_update"]
+    kinds = [m.get("update_type") for m in updates]
+    # agent_completion should always be present
+    assert "agent_completion" in kinds
+    # agent_request_input may or may not be present depending on environment

atlas/tests/test_app_factory_smoke.py

@@ -0,0 +1,47 @@
+import sys
+import types
+
+
+def _ensure_litellm_stub():
+    if "litellm" in sys.modules:
+        return
+    m = types.ModuleType("litellm")
+    # Attributes used at import time
+    m.drop_params = True
+    def _set_verbose(*args, **kwargs):
+        return None
+    m.set_verbose = _set_verbose
+    # Names imported via `from litellm import ...`
+    def completion(*args, **kwargs):
+        return None
+    async def acompletion(*args, **kwargs):
+        class Dummy:
+            class Choice:
+                class Msg:
+                    content = ""
+                message = Msg()
+            choices = [Choice()]
+        return Dummy()
+    m.completion = completion
+    m.acompletion = acompletion
+    sys.modules["litellm"] = m
+
+
+def test_app_factory_accessors():
+    _ensure_litellm_stub()
+    from atlas.infrastructure.app_factory import app_factory
+    # Accessors should return instances without raising
+    assert app_factory.get_config_manager() is not None
+    assert app_factory.get_llm_caller() is not None
+    assert app_factory.get_mcp_manager() is not None
+    assert app_factory.get_file_storage() is not None
+    assert app_factory.get_file_manager() is not None
+
+    # RAG services are None when FEATURE_RAG_ENABLED is false (default)
+    rag_enabled = app_factory.get_config_manager().app_settings.feature_rag_enabled
+    if rag_enabled:
+        assert app_factory.get_unified_rag_service() is not None
+        assert app_factory.get_rag_mcp_service() is not None
+    else:
+        assert app_factory.get_unified_rag_service() is None
+        assert app_factory.get_rag_mcp_service() is None

atlas/tests/test_approval_manager.py

@@ -0,0 +1,439 @@
+"""Tests for the tool approval manager."""
+
+import asyncio
+
+import pytest
+
+from atlas.application.chat.approval_manager import ToolApprovalManager, ToolApprovalRequest, get_approval_manager
+
+
+class TestToolApprovalRequest:
+    """Test ToolApprovalRequest class."""
+
+    @pytest.mark.asyncio
+    async def test_create_approval_request(self):
+        """Test creating an approval request."""
+        request = ToolApprovalRequest(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"arg1": "value1"},
+            allow_edit=True
+        )
+        assert request.tool_call_id == "test_123"
+        assert request.tool_name == "test_tool"
+        assert request.arguments == {"arg1": "value1"}
+        assert request.allow_edit is True
+
+    @pytest.mark.asyncio
+    async def test_set_response(self):
+        """Test setting a response to an approval request."""
+        request = ToolApprovalRequest(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"arg1": "value1"}
+        )
+
+        # Set approved response
+        request.set_response(approved=True, arguments={"arg1": "edited_value"})
+
+        # Wait for the response (should be immediate since we already set it)
+        response = await request.wait_for_response(timeout=1.0)
+
+        assert response["approved"] is True
+        assert response["arguments"] == {"arg1": "edited_value"}
+
+    @pytest.mark.asyncio
+    async def test_rejection_response(self):
+        """Test rejecting an approval request."""
+        request = ToolApprovalRequest(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"arg1": "value1"}
+        )
+
+        # Set rejected response
+        request.set_response(approved=False, reason="User rejected")
+
+        # Wait for the response
+        response = await request.wait_for_response(timeout=1.0)
+
+        assert response["approved"] is False
+        assert response["reason"] == "User rejected"
+
+    @pytest.mark.asyncio
+    async def test_timeout(self):
+        """Test that timeout works correctly."""
+        request = ToolApprovalRequest(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"arg1": "value1"}
+        )
+
+        # Should timeout since we don't set a response
+        with pytest.raises(asyncio.TimeoutError):
+            await request.wait_for_response(timeout=0.1)
+
+
+class TestToolApprovalManager:
+    """Test ToolApprovalManager class."""
+
+    @pytest.mark.asyncio
+    async def test_create_approval_request(self):
+        """Test creating an approval request via manager."""
+        manager = ToolApprovalManager()
+        request = manager.create_approval_request(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"arg1": "value1"},
+            allow_edit=True
+        )
+
+        assert request.tool_call_id == "test_123"
+        assert "test_123" in manager.get_pending_requests()
+
+    @pytest.mark.asyncio
+    async def test_handle_approval_response(self):
+        """Test handling an approval response."""
+        manager = ToolApprovalManager()
+        manager.create_approval_request(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"arg1": "value1"}
+        )
+
+        # Handle approval response
+        result = manager.handle_approval_response(
+            tool_call_id="test_123",
+            approved=True,
+            arguments={"arg1": "edited_value"}
+        )
+
+        assert result is True
+        # Request should still be in pending (cleaned up manually later)
+        assert "test_123" in manager.get_pending_requests()
+
+    def test_handle_unknown_request(self):
+        """Test handling response for unknown request."""
+        manager = ToolApprovalManager()
+
+        result = manager.handle_approval_response(
+            tool_call_id="unknown_123",
+            approved=True
+        )
+
+        assert result is False
+
+    @pytest.mark.asyncio
+    async def test_cleanup_request(self):
+        """Test cleaning up a completed request."""
+        manager = ToolApprovalManager()
+        manager.create_approval_request(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"arg1": "value1"}
+        )
+
+        assert "test_123" in manager.get_pending_requests()
+
+        manager.cleanup_request("test_123")
+
+        assert "test_123" not in manager.get_pending_requests()
+
+    def test_get_approval_manager_singleton(self):
+        """Test that get_approval_manager returns a singleton."""
+        manager1 = get_approval_manager()
+        manager2 = get_approval_manager()
+
+        assert manager1 is manager2
+
+    @pytest.mark.asyncio
+    async def test_full_approval_workflow(self):
+        """Test the complete approval workflow."""
+        manager = ToolApprovalManager()
+
+        # Create request
+        request = manager.create_approval_request(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"code": "print('test')"},
+            allow_edit=True
+        )
+
+        # Simulate async approval (in a separate task)
+        async def approve_after_delay():
+            await asyncio.sleep(0.1)
+            manager.handle_approval_response(
+                tool_call_id="test_123",
+                approved=True,
+                arguments={"code": "print('edited test')"}
+            )
+
+        # Start approval task
+        asyncio.create_task(approve_after_delay())
+
+        # Wait for response
+        response = await request.wait_for_response(timeout=1.0)
+
+        assert response["approved"] is True
+        assert response["arguments"]["code"] == "print('edited test')"
+
+        # Cleanup
+        manager.cleanup_request("test_123")
+        assert "test_123" not in manager.get_pending_requests()
+
+    @pytest.mark.asyncio
+    async def test_multiple_concurrent_approvals(self):
+        """Test handling multiple concurrent approval requests."""
+        manager = ToolApprovalManager()
+
+        # Create multiple requests
+        request1 = manager.create_approval_request(
+            tool_call_id="test_1",
+            tool_name="tool_a",
+            arguments={"arg": "value1"}
+        )
+        request2 = manager.create_approval_request(
+            tool_call_id="test_2",
+            tool_name="tool_b",
+            arguments={"arg": "value2"}
+        )
+        request3 = manager.create_approval_request(
+            tool_call_id="test_3",
+            tool_name="tool_c",
+            arguments={"arg": "value3"}
+        )
+
+        assert len(manager.get_pending_requests()) == 3
+
+        # Approve them in different order
+        async def approve_requests():
+            await asyncio.sleep(0.05)
+            manager.handle_approval_response("test_2", approved=True)
+            await asyncio.sleep(0.05)
+            manager.handle_approval_response("test_1", approved=False, reason="Rejected")
+            await asyncio.sleep(0.05)
+            manager.handle_approval_response("test_3", approved=True)
+
+        asyncio.create_task(approve_requests())
+
+        # Wait for all responses
+        response1 = await request1.wait_for_response(timeout=1.0)
+        response2 = await request2.wait_for_response(timeout=1.0)
+        response3 = await request3.wait_for_response(timeout=1.0)
+
+        assert response1["approved"] is False
+        assert response1["reason"] == "Rejected"
+        assert response2["approved"] is True
+        assert response3["approved"] is True
+
+    @pytest.mark.asyncio
+    async def test_approval_with_no_arguments_change(self):
+        """Test approval where arguments are returned but not changed."""
+        manager = ToolApprovalManager()
+
+        original_args = {"code": "print('hello')"}
+        request = manager.create_approval_request(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments=original_args,
+            allow_edit=True
+        )
+
+        # Approve with same arguments
+        async def approve():
+            await asyncio.sleep(0.05)
+            manager.handle_approval_response(
+                tool_call_id="test_123",
+                approved=True,
+                arguments={"code": "print('hello')"}  # Same as original
+            )
+
+        asyncio.create_task(approve())
+        response = await request.wait_for_response(timeout=1.0)
+
+        assert response["approved"] is True
+        assert response["arguments"] == original_args
+
+    @pytest.mark.asyncio
+    async def test_double_response_handling(self):
+        """Test that setting response twice doesn't cause issues."""
+        request = ToolApprovalRequest(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"arg1": "value1"}
+        )
+
+        # Set response first time
+        request.set_response(approved=True, arguments={"arg1": "first"})
+
+        # Try to set response second time (should be ignored)
+        request.set_response(approved=False, arguments={"arg1": "second"})
+
+        # Should get the first response
+        response = await request.wait_for_response(timeout=0.5)
+        assert response["approved"] is True
+        assert response["arguments"]["arg1"] == "first"
+
+    @pytest.mark.asyncio
+    async def test_rejection_with_empty_reason(self):
+        """Test rejection with no reason provided."""
+        manager = ToolApprovalManager()
+
+        request = manager.create_approval_request(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"arg1": "value1"}
+        )
+
+        async def reject():
+            await asyncio.sleep(0.05)
+            manager.handle_approval_response(
+                tool_call_id="test_123",
+                approved=False
+                # No reason provided
+            )
+
+        _ = asyncio.create_task(reject())
+        response = await request.wait_for_response(timeout=1.0)
+
+        assert response["approved"] is False
+        assert response.get("reason") is None or response.get("reason") == ""
+
+    @pytest.mark.asyncio
+    async def test_allow_edit_false(self):
+        """Test approval request with editing disabled."""
+        request = ToolApprovalRequest(
+            tool_call_id="test_123",
+            tool_name="test_tool",
+            arguments={"arg1": "value1"},
+            allow_edit=False
+        )
+
+        assert request.allow_edit is False
+
+        # Even if arguments are provided, they should be used
+        request.set_response(approved=True, arguments={"arg1": "edited_value"})
+        response = await request.wait_for_response(timeout=0.5)
+
+        # The response will contain the edited arguments, but the UI should
+        # respect allow_edit=False to prevent showing edit controls
+        assert response["arguments"] == {"arg1": "edited_value"}
+
+    def test_cleanup_nonexistent_request(self):
+        """Test cleaning up a request that doesn't exist."""
+        manager = ToolApprovalManager()
+
+        # Should not raise an error
+        manager.cleanup_request("nonexistent_id")
+
+        assert "nonexistent_id" not in manager.get_pending_requests()
+
+    def test_multiple_managers_vs_singleton(self):
+        """Test that direct instantiation creates different instances but singleton returns same."""
+        manager1 = ToolApprovalManager()
+        manager2 = ToolApprovalManager()
+
+        # Direct instantiation creates different instances
+        assert manager1 is not manager2
+
+        # But singleton returns the same instance
+        singleton1 = get_approval_manager()
+        singleton2 = get_approval_manager()
+        assert singleton1 is singleton2
+
+    @pytest.mark.asyncio
+    async def test_approval_with_complex_arguments(self):
+        """Test approval with complex nested arguments."""
+        manager = ToolApprovalManager()
+
+        complex_args = {
+            "nested": {
+                "level1": {
+                    "level2": ["item1", "item2", "item3"]
+                }
+            },
+            "list_of_dicts": [
+                {"key": "value1"},
+                {"key": "value2"}
+            ],
+            "numbers": [1, 2, 3, 4, 5]
+        }
+
+        request = manager.create_approval_request(
+            tool_call_id="test_complex",
+            tool_name="complex_tool",
+            arguments=complex_args,
+            allow_edit=True
+        )
+
+        # Modify nested structure
+        edited_args = {
+            "nested": {
+                "level1": {
+                    "level2": ["item1", "modified_item", "item3"]
+                }
+            },
+            "list_of_dicts": [
+                {"key": "value1"},
+                {"key": "new_value"}
+            ],
+            "numbers": [1, 2, 3, 4, 5, 6]
+        }
+
+        async def approve():
+            await asyncio.sleep(0.05)
+            manager.handle_approval_response(
+                tool_call_id="test_complex",
+                approved=True,
+                arguments=edited_args
+            )
+
+        asyncio.create_task(approve())
+        response = await request.wait_for_response(timeout=1.0)
+
+        assert response["approved"] is True
+        assert response["arguments"]["nested"]["level1"]["level2"][1] == "modified_item"
+        assert len(response["arguments"]["numbers"]) == 6
+
+    @pytest.mark.asyncio
+    async def test_sequential_approvals(self):
+        """Test approving requests one after another in sequence."""
+        manager = ToolApprovalManager()
+
+        # First approval
+        request1 = manager.create_approval_request(
+            tool_call_id="seq_1",
+            tool_name="tool_1",
+            arguments={"step": 1}
+        )
+
+        async def approve1():
+            await asyncio.sleep(0.05)
+            manager.handle_approval_response("seq_1", approved=True)
+
+        task1 = asyncio.create_task(approve1())
+        response1 = await request1.wait_for_response(timeout=1.0)
+        manager.cleanup_request("seq_1")
+        await task1
+
+        assert response1["approved"] is True
+        assert "seq_1" not in manager.get_pending_requests()
+
+        # Second approval after first is complete
+        request2 = manager.create_approval_request(
+            tool_call_id="seq_2",
+            tool_name="tool_2",
+            arguments={"step": 2}
+        )
+
+        async def approve2():
+            await asyncio.sleep(0.05)
+            manager.handle_approval_response("seq_2", approved=True)
+
+        task2 = asyncio.create_task(approve2())
+        response2 = await request2.wait_for_response(timeout=1.0)
+        manager.cleanup_request("seq_2")
+        await task2
+
+        assert response2["approved"] is True
+        assert "seq_2" not in manager.get_pending_requests()