refactor: RAG Refactor (#985)

Co-authored-by: Aralhi <xiaoping0501@gmail.com>
Co-authored-by: csunny <cfqsunny@163.com>

dbgpt/rag/retriever/tests/test_embedding.py

@@ -0,0 +1,39 @@
+from unittest.mock import MagicMock
+import pytest
+
+from dbgpt.rag.chunk import Chunk
+from dbgpt.rag.retriever.embedding import EmbeddingRetriever
+
+
+@pytest.fixture
+def top_k():
+    return 4
+
+
+@pytest.fixture
+def query():
+    return "test query"
+
+
+@pytest.fixture
+def mock_vector_store_connector():
+    return MagicMock()
+
+
+@pytest.fixture
+def embedding_retriever(top_k, mock_vector_store_connector):
+    return EmbeddingRetriever(
+        top_k=top_k,
+        query_rewrite=False,
+        vector_store_connector=mock_vector_store_connector,
+    )
+
+
+def test_retrieve(query, top_k, mock_vector_store_connector, embedding_retriever):
+    expected_chunks = [Chunk() for _ in range(top_k)]
+    mock_vector_store_connector.similar_search.return_value = expected_chunks
+
+    retrieved_chunks = embedding_retriever._retrieve(query)
+
+    mock_vector_store_connector.similar_search.assert_called_once_with(query, top_k)
+    assert len(retrieved_chunks) == top_k