feat: (0.6)New UI (#1855)

Co-authored-by: 夏姜 <wenfengjiang.jwf@digital-engine.com>
Co-authored-by: aries_ckt <916701291@qq.com>
Co-authored-by: wb-lh513319 <wb-lh513319@alibaba-inc.com>
Co-authored-by: csunny <cfqsunny@163.com>

dbgpt/serve/rag/retriever/knowledge_space.py

@@ -5,8 +5,12 @@ from dbgpt.component import ComponentType
 from dbgpt.configs.model_config import EMBEDDING_MODEL_CONFIG
 from dbgpt.core import Chunk
 from dbgpt.rag.embedding.embedding_factory import EmbeddingFactory
+from dbgpt.rag.retriever import EmbeddingRetriever, QueryRewrite, Ranker
 from dbgpt.rag.retriever.base import BaseRetriever
 from dbgpt.serve.rag.connector import VectorStoreConnector
+from dbgpt.serve.rag.models.models import KnowledgeSpaceDao
+from dbgpt.serve.rag.retriever.qa_retriever import QARetriever
+from dbgpt.serve.rag.retriever.retriever_chain import RetrieverChain
 from dbgpt.storage.vector_store.filters import MetadataFilters
 from dbgpt.util.executor_utils import ExecutorFactory, blocking_func_to_async
 
@@ -18,18 +22,24 @@ class KnowledgeSpaceRetriever(BaseRetriever):
 
     def __init__(
         self,
-        space_name: str = None,
+        space_id: str = None,
         top_k: Optional[int] = 4,
+        query_rewrite: Optional[QueryRewrite] = None,
+        rerank: Optional[Ranker] = None,
     ):
         """
         Args:
-            space_name (str): knowledge space name
+            space_id (str): knowledge space name
             top_k (Optional[int]): top k
+            query_rewrite: (Optional[QueryRewrite]) query rewrite
+            rerank: (Optional[Ranker]) rerank
         """
-        if space_name is None:
-            raise ValueError("space_name is required")
-        self._space_name = space_name
+        if space_id is None:
+            raise ValueError("space_id is required")
+        self._space_id = space_id
         self._top_k = top_k
+        self._query_rewrite = query_rewrite
+        self._rerank = rerank
         embedding_factory = CFG.SYSTEM_APP.get_component(
             "embedding_factory", EmbeddingFactory
         )
@@ -38,7 +48,9 @@ class KnowledgeSpaceRetriever(BaseRetriever):
         )
         from dbgpt.storage.vector_store.base import VectorStoreConfig
 
-        config = VectorStoreConfig(name=self._space_name, embedding_fn=embedding_fn)
+        space_dao = KnowledgeSpaceDao()
+        space = space_dao.get_one({"id": space_id})
+        config = VectorStoreConfig(name=space.name, embedding_fn=embedding_fn)
         self._vector_store_connector = VectorStoreConnector(
             vector_store_type=CFG.VECTOR_STORE_TYPE,
             vector_store_config=config,
@@ -47,6 +59,19 @@ class KnowledgeSpaceRetriever(BaseRetriever):
             ComponentType.EXECUTOR_DEFAULT, ExecutorFactory
         ).create()
 
+        self._retriever_chain = RetrieverChain(
+            retrievers=[
+                QARetriever(space_id=space_id, top_k=top_k, embedding_fn=embedding_fn),
+                EmbeddingRetriever(
+                    index_store=self._vector_store_connector.index_client,
+                    top_k=top_k,
+                    query_rewrite=self._query_rewrite,
+                    rerank=self._rerank,
+                ),
+            ],
+            executor=self._executor,
+        )
+
     def _retrieve(
         self, query: str, filters: Optional[MetadataFilters] = None
     ) -> List[Chunk]:
@@ -59,9 +84,7 @@ class KnowledgeSpaceRetriever(BaseRetriever):
         Return:
             List[Chunk]: list of chunks
         """
-        candidates = self._vector_store_connector.similar_search(
-            doc=query, topk=self._top_k, filters=filters
-        )
+        candidates = self._retriever_chain.retrieve(query=query, filters=filters)
         return candidates
 
     def _retrieve_with_score(
@@ -80,13 +103,10 @@ class KnowledgeSpaceRetriever(BaseRetriever):
         Return:
             List[Chunk]: list of chunks with score
         """
-        candidates_with_score = self._vector_store_connector.similar_search_with_scores(
-            doc=query,
-            topk=self._top_k,
-            score_threshold=score_threshold,
-            filters=filters,
+        candidates_with_scores = self._retriever_chain.retrieve_with_scores(
+            query, score_threshold, filters
         )
-        return candidates_with_score
+        return candidates_with_scores
 
     async def _aretrieve(
         self, query: str, filters: Optional[MetadataFilters] = None

dbgpt/serve/rag/retriever/qa_retriever.py

@@ -0,0 +1,216 @@
+import ast
+import json
+import logging
+from typing import Any, List, Optional
+
+from dbgpt._private.config import Config
+from dbgpt.app.knowledge.chunk_db import DocumentChunkDao, DocumentChunkEntity
+from dbgpt.app.knowledge.document_db import KnowledgeDocumentDao
+from dbgpt.component import ComponentType
+from dbgpt.core import Chunk
+from dbgpt.rag.retriever.base import BaseRetriever
+from dbgpt.serve.rag.models.models import KnowledgeSpaceDao, KnowledgeSpaceEntity
+from dbgpt.storage.vector_store.filters import MetadataFilters
+from dbgpt.util.executor_utils import ExecutorFactory, blocking_func_to_async
+from dbgpt.util.similarity_util import calculate_cosine_similarity
+from dbgpt.util.string_utils import remove_trailing_punctuation
+
+CFG = Config()
+CHUNK_PAGE_SIZE = 1000
+logger = logging.getLogger(__name__)
+
+
+class QARetriever(BaseRetriever):
+    """Document QA retriever."""
+
+    def __init__(
+        self,
+        space_id: str = None,
+        top_k: Optional[int] = 4,
+        embedding_fn: Optional[Any] = 4,
+        lambda_value: Optional[float] = 1e-5,
+    ):
+        """
+        Args:
+            space_id (str): knowledge space name
+            top_k (Optional[int]): top k
+        """
+        if space_id is None:
+            raise ValueError("space_id is required")
+        self._top_k = top_k
+        self._lambda_value = lambda_value
+        self._space_dao = KnowledgeSpaceDao()
+        self._document_dao = KnowledgeDocumentDao()
+        self._chunk_dao = DocumentChunkDao()
+        self._embedding_fn = embedding_fn
+
+        space = self._space_dao.get_one({"id": space_id})
+        if not space:
+            raise ValueError("space not found")
+        self.documents = self._document_dao.get_list({"space": space.name})
+        self._executor = CFG.SYSTEM_APP.get_component(
+            ComponentType.EXECUTOR_DEFAULT, ExecutorFactory
+        ).create()
+
+    def _retrieve(
+        self, query: str, filters: Optional[MetadataFilters] = None
+    ) -> List[Chunk]:
+        """Retrieve knowledge chunks.
+        Args:
+            query (str): query text
+            filters: (Optional[MetadataFilters]) metadata filters.
+        Return:
+            List[Chunk]: list of chunks
+        """
+        query = remove_trailing_punctuation(query)
+        candidate_results = []
+        for doc in self.documents:
+            if doc.questions:
+                questions = json.loads(doc.questions)
+                if query in questions:
+                    chunks = self._chunk_dao.get_document_chunks(
+                        DocumentChunkEntity(document_id=doc.id),
+                        page_size=CHUNK_PAGE_SIZE,
+                    )
+                    candidates = [
+                        Chunk(
+                            content=chunk.content,
+                            metadata=ast.literal_eval(chunk.meta_info),
+                            retriever=self.name(),
+                            score=0.0,
+                        )
+                        for chunk in chunks
+                    ]
+                    candidate_results.extend(
+                        self._cosine_similarity_rerank(candidates, query)
+                    )
+        return candidate_results
+
+    def _retrieve_with_score(
+        self,
+        query: str,
+        score_threshold: float,
+        filters: Optional[MetadataFilters] = None,
+        lambda_value: Optional[float] = None,
+    ) -> List[Chunk]:
+        """Retrieve knowledge chunks with score.
+        Args:
+            query (str): query text
+            score_threshold (float): score threshold
+            filters: (Optional[MetadataFilters]) metadata filters.
+        Return:
+            List[Chunk]: list of chunks with score
+        """
+        query = remove_trailing_punctuation(query)
+        candidate_results = []
+        doc_ids = [doc.id for doc in self.documents]
+        query_param = DocumentChunkEntity()
+        chunks = self._chunk_dao.get_chunks_with_questions(
+            query=query_param, document_ids=doc_ids
+        )
+        for chunk in chunks:
+            if chunk.questions:
+                questions = json.loads(chunk.questions)
+                if query in questions:
+                    logger.info(f"qa chunk hit:{chunk}, question:{query}")
+                    candidate_results.append(
+                        Chunk(
+                            content=chunk.content,
+                            chunk_id=str(chunk.id),
+                            metadata={"prop_field": ast.literal_eval(chunk.meta_info)},
+                            retriever=self.name(),
+                            score=1.0,
+                        )
+                    )
+        if len(candidate_results) > 0:
+            return self._cosine_similarity_rerank(candidate_results, query)
+
+        for doc in self.documents:
+            if doc.questions:
+                questions = json.loads(doc.questions)
+                if query in questions:
+                    logger.info(f"qa document hit:{doc}, question:{query}")
+                    chunks = self._chunk_dao.get_document_chunks(
+                        DocumentChunkEntity(document_id=doc.id),
+                        page_size=CHUNK_PAGE_SIZE,
+                    )
+                    candidates_with_scores = [
+                        Chunk(
+                            content=chunk.content,
+                            chunk_id=str(chunk.id),
+                            metadata={"prop_field": ast.literal_eval(chunk.meta_info)},
+                            retriever=self.name(),
+                            score=1.0,
+                        )
+                        for chunk in chunks
+                    ]
+                    candidate_results.extend(
+                        self._cosine_similarity_rerank(candidates_with_scores, query)
+                    )
+        return candidate_results
+
+    async def _aretrieve(
+        self, query: str, filters: Optional[MetadataFilters] = None
+    ) -> List[Chunk]:
+        """Retrieve knowledge chunks.
+        Args:
+            query (str): query text
+            filters: (Optional[MetadataFilters]) metadata filters.
+        Return:
+            List[Chunk]: list of chunks
+        """
+        candidates = await blocking_func_to_async(
+            self._executor, self._retrieve, query, filters
+        )
+        return candidates
+
+    async def _aretrieve_with_score(
+        self,
+        query: str,
+        score_threshold: float,
+        filters: Optional[MetadataFilters] = None,
+    ) -> List[Chunk]:
+        """Retrieve knowledge chunks with score.
+        Args:
+            query (str): query text
+            score_threshold (float): score threshold
+            filters: (Optional[MetadataFilters]) metadata filters.
+        Return:
+            List[Chunk]: list of chunks with score
+        """
+        candidates_with_score = await blocking_func_to_async(
+            self._executor, self._retrieve_with_score, query, score_threshold, filters
+        )
+        return candidates_with_score
+
+    def _cosine_similarity_rerank(
+        self, candidates_with_scores: List[Chunk], query: str
+    ) -> List[Chunk]:
+        """Rerank candidates using cosine similarity."""
+        if len(candidates_with_scores) > self._top_k:
+            for candidate in candidates_with_scores:
+                similarity = calculate_cosine_similarity(
+                    embeddings=self._embedding_fn,
+                    prediction=query,
+                    contexts=[candidate.content],
+                )
+                score = float(similarity.mean())
+                candidate.score = score
+            candidates_with_scores.sort(key=lambda x: x.score, reverse=True)
+            candidates_with_scores = candidates_with_scores[: self._top_k]
+            candidates_with_scores = [
+                Chunk(
+                    content=candidate.content,
+                    chunk_id=candidate.chunk_id,
+                    metadata=candidate.metadata,
+                    retriever=self.name(),
+                    score=1.0,
+                )
+                for candidate in candidates_with_scores
+            ]
+        return candidates_with_scores
+
+    @classmethod
+    def name(cls):
+        """Return retriever name."""
+        return "qa_retriever"

dbgpt/serve/rag/retriever/retriever_chain.py

@@ -0,0 +1,91 @@
+from concurrent.futures import Executor, ThreadPoolExecutor
+from typing import List, Optional
+
+from dbgpt.core import Chunk
+from dbgpt.rag.retriever.base import BaseRetriever
+from dbgpt.storage.vector_store.filters import MetadataFilters
+from dbgpt.util.executor_utils import blocking_func_to_async
+
+
+class RetrieverChain(BaseRetriever):
+    """Retriever chain class."""
+
+    def __init__(
+        self,
+        retrievers: Optional[List[BaseRetriever]] = None,
+        executor: Optional[Executor] = None,
+    ):
+        """Create retriever chain instance."""
+        self._retrievers = retrievers or []
+        self._executor = executor or ThreadPoolExecutor()
+
+    def _retrieve(
+        self, query: str, filters: Optional[MetadataFilters] = None
+    ) -> List[Chunk]:
+        """Retrieve knowledge chunks.
+        Args:
+            query (str): query text
+            filters: (Optional[MetadataFilters]) metadata filters.
+        Return:
+            List[Chunk]: list of chunks
+        """
+        for retriever in self._retrievers:
+            candidates = retriever.retrieve(query, filters)
+            if candidates:
+                return candidates
+        return []
+
+    async def _aretrieve(
+        self, query: str, filters: Optional[MetadataFilters] = None
+    ) -> List[Chunk]:
+        """Retrieve knowledge chunks.
+        Args:
+            query (str): query text
+            filters: (Optional[MetadataFilters]) metadata filters.
+        Return:
+            List[Chunk]: list of chunks
+        """
+        candidates = await blocking_func_to_async(
+            self._executor, self._retrieve, query, filters
+        )
+        return candidates
+
+    def _retrieve_with_score(
+        self,
+        query: str,
+        score_threshold: float,
+        filters: Optional[MetadataFilters] = None,
+    ) -> List[Chunk]:
+        """Retrieve knowledge chunks.
+        Args:
+            query (str): query text
+            filters: (Optional[MetadataFilters]) metadata filters.
+        Return:
+            List[Chunk]: list of chunks
+        """
+        for retriever in self._retrievers:
+            candidates_with_scores = retriever.retrieve_with_scores(
+                query=query, score_threshold=score_threshold, filters=filters
+            )
+            if candidates_with_scores:
+                return candidates_with_scores
+        return []
+
+    async def _aretrieve_with_score(
+        self,
+        query: str,
+        score_threshold: float,
+        filters: Optional[MetadataFilters] = None,
+    ) -> List[Chunk]:
+        """Retrieve knowledge chunks with score.
+        Args:
+            query (str): query text
+            score_threshold (float): score threshold
+            filters: (Optional[MetadataFilters]) metadata filters.
+        Return:
+            List[Chunk]: list of chunks with score
+        """
+        candidates_with_score = await blocking_func_to_async(
+            self._executor, self._retrieve_with_score, query, score_threshold, filters
+        )
+        return candidates_with_score