test: 补充知识处理、检索与 Agent 回归用例

覆盖切块、向量检索、工具调用和历史预算；模型相关用例使用桩。

backend/tests/knowledge/test_optimized_pgvector.py

@@ -0,0 +1,205 @@
+"""
+Copyright 2024-2026 ChatterMate
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+
+import pytest
+from unittest.mock import MagicMock, patch
+import numpy as np
+
+from agno.document import Document
+from agno.vectordb.pgvector.pgvector import PgVector
+from app.knowledge.optimized_pgvector import OptimizedPgVector
+
+
+class TestOptimizedPgVector:
+    """Test cases for OptimizedPgVector"""
+    
+    @pytest.fixture
+    def mock_embedder(self):
+        """Create a mock embedder"""
+        mock = MagicMock()
+        mock.get_embedding.return_value = np.array([0.1, 0.2, 0.3])
+        mock.dimensions = 3
+        return mock
+    
+    @pytest.fixture
+    def mock_session(self):
+        """Create a mock database session"""
+        mock = MagicMock()
+        # Configure session to return a context manager
+        mock.return_value.__enter__.return_value = mock
+        return mock
+    
+    @patch('sqlalchemy.engine.create_engine')
+    def test_init(self, mock_create_engine, mock_embedder):
+        """Test initialization"""
+        # Arrange
+        mock_engine = MagicMock()
+        mock_create_engine.return_value = mock_engine
+        
+        # Act
+        db = OptimizedPgVector(
+            table_name="test_table",
+            db_url="postgresql+psycopg://user:pass@localhost/testdb",
+            embedder=mock_embedder
+        )
+        
+        # Assert
+        assert db.table_name == "test_table"
+        assert db.embedder == mock_embedder
+        assert db.dimensions == 3
+    
+    @patch('sqlalchemy.engine.create_engine')
+    def test_upsert_with_pre_embedded_documents(self, mock_create_engine, mock_embedder, mock_session):
+        """Test upserting pre-embedded documents"""
+        # Arrange
+        mock_engine = MagicMock()
+        mock_create_engine.return_value = mock_engine
+        
+        db = OptimizedPgVector(
+            table_name="test_table",
+            db_url="postgresql+psycopg://user:pass@localhost/testdb",
+            embedder=mock_embedder
+        )
+        db.Session = mock_session
+        db._clean_content = lambda x: x  # Mock the content cleaning function
+        
+        # Create a document with an existing embedding
+        pre_embedded_doc = Document(
+            id="test1",
+            name="Test Document 1",
+            content="This is a test document",
+            embedding=np.array([0.5, 0.6, 0.7])
+        )
+        
+        # Create a document without an embedding
+        unembedded_doc = Document(
+            id="test2",
+            name="Test Document 2",
+            content="This is another test document"
+        )
+        
+        # Mock the document's embed method
+        with patch.object(Document, 'embed') as mock_embed:
+            # Act
+            with patch('sqlalchemy.dialects.postgresql.insert') as mock_insert:
+                mock_insert.return_value.on_conflict_do_update.return_value = "upsert_stmt"
+                db.upsert([pre_embedded_doc, unembedded_doc])
+            
+            # Assert
+            # The document's embed method should be called for the unembedded document
+            mock_embed.assert_called_once_with(embedder=mock_embedder)
+            
+            # Both documents should be passed to the upsert statement
+            assert mock_session.execute.call_count == 1
+        
+    @patch('sqlalchemy.engine.create_engine')
+    def test_insert_with_pre_embedded_documents(self, mock_create_engine, mock_embedder, mock_session):
+        """Test inserting pre-embedded documents"""
+        # Arrange
+        mock_engine = MagicMock()
+        mock_create_engine.return_value = mock_engine
+        
+        db = OptimizedPgVector(
+            table_name="test_table",
+            db_url="postgresql+psycopg://user:pass@localhost/testdb",
+            embedder=mock_embedder
+        )
+        db.Session = mock_session
+        db._clean_content = lambda x: x  # Mock the content cleaning function
+        
+        # Create a document with an existing embedding
+        pre_embedded_doc = Document(
+            id="test1",
+            name="Test Document 1",
+            content="This is a test document",
+            embedding=np.array([0.5, 0.6, 0.7])
+        )
+        
+        # Create a document without an embedding
+        unembedded_doc = Document(
+            id="test2",
+            name="Test Document 2",
+            content="This is another test document"
+        )
+        
+        # Mock the document's embed method
+        with patch.object(Document, 'embed') as mock_embed:
+            # Act
+            with patch('sqlalchemy.dialects.postgresql.insert') as mock_insert:
+                db.insert([pre_embedded_doc, unembedded_doc])
+            
+            # Assert
+            # The document's embed method should be called for the unembedded document
+            mock_embed.assert_called_once_with(embedder=mock_embedder)
+            
+            # Both documents should be passed to the insert statement
+            assert mock_session.execute.call_count == 1
+
+
+if __name__ == "__main__":
+    pytest.main(["-v", "test_optimized_pgvector.py"])
+
+
+class _PgError(Exception):
+    """Stand-in for a psycopg error, which carries the SQLSTATE as .sqlstate."""
+
+    def __init__(self, sqlstate):
+        super().__init__(sqlstate)
+        self.sqlstate = sqlstate
+
+
+class TestConcurrentTableCreate:
+    """The per-org vector table is created by whichever queue item gets there first.
+
+    agno's PgVector.create() does table_exists() and then CREATE TABLE, which is
+    not atomic. The processor runs two items at once, so an organisation whose
+    first two sources are queued together had both workers create the table and
+    the loser died with
+    `duplicate key value violates unique constraint "pg_type_typname_nsp_index"`,
+    losing that source (production queue item 372, goldmaidsteam.com).
+    """
+
+    def _vector_db(self):
+        return OptimizedPgVector.__new__(OptimizedPgVector)
+
+    @pytest.mark.parametrize("sqlstate", ["42P07", "23505"])
+    def test_losing_the_create_race_is_not_an_error(self, sqlstate):
+        db = self._vector_db()
+        db.table_name = "d_org"
+        db.schema = "ai"
+
+        with patch.object(PgVector, "create", side_effect=_PgError(sqlstate)):
+            db.create()  # the other worker created it; that is the table we wanted
+
+    def test_a_real_create_failure_still_raises(self):
+        db = self._vector_db()
+        db.table_name = "d_org"
+        db.schema = "ai"
+
+        with patch.object(PgVector, "create", side_effect=_PgError("42501")):
+            with pytest.raises(_PgError):
+                db.create()
+
+    def test_duplicate_detected_through_sqlalchemy_wrapping(self):
+        """SQLAlchemy wraps the driver error; the SQLSTATE lives on .orig."""
+        wrapped = RuntimeError("duplicate")
+        wrapped.orig = _PgError("23505")
+
+        with patch.object(PgVector, "create", side_effect=wrapped):
+            db = self._vector_db()
+            db.table_name = "d_org"
+            db.schema = "ai"
+            db.create()