github/DB-GPT
1
0
Fork 0
mirror of https://github.com/csunny/DB-GPT.git synced 2025-09-05 19:11:52 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(rag): Support RAG SDK (#1322)

Browse Source
This commit is contained in:
Fangyin Cheng
2024-03-22 15:36:57 +08:00
committed by GitHub
parent e65732d6e4
commit 8a17099dd2
69 changed files with 1332 additions and 558 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
dbgpt/serve/rag/operators/base.py
View File

@@ -1,23 +0,0 @@
from abc import abstractmethod
from dbgpt.core.awel import MapOperator
from dbgpt.core.awel.task.base import IN, OUT
class AssemblerOperator(MapOperator[IN, OUT]):
"""The Base Assembler Operator."""
async def map(self, input_value: IN) -> OUT:
"""Map input value to output value.
Args:
input_value (IN): The input value.
Returns:
OUT: The output value.
"""
return await self.blocking_func_to_async(self.assemble, input_value)
@abstractmethod
def assemble(self, input_value: IN) -> OUT:
"""assemble knowledge for input value."""

36
dbgpt/serve/rag/operators/db_schema.py
View File

@@ -1,36 +0,0 @@
from typing import Any, Optional
from dbgpt.core.awel.task.base import IN
from dbgpt.datasource.rdbms.base import RDBMSConnector
from dbgpt.serve.rag.assembler.db_schema import DBSchemaAssembler
from dbgpt.serve.rag.operators.base import AssemblerOperator
from dbgpt.storage.vector_store.connector import VectorStoreConnector
class DBSchemaAssemblerOperator(AssemblerOperator[Any, Any]):
"""The DBSchema Assembler Operator.
Args:
connection (RDBMSConnector): The connection.
chunk_parameters (Optional[ChunkParameters], optional): The chunk parameters. Defaults to None.
vector_store_connector (VectorStoreConnector, optional): The vector store connector. Defaults to None.
"""
def __init__(
self,
connection: RDBMSConnector = None,
vector_store_connector: Optional[VectorStoreConnector] = None,
**kwargs
):
self._connection = connection
self._vector_store_connector = vector_store_connector
self._assembler = DBSchemaAssembler.load_from_connection(
connection=self._connection,
vector_store_connector=self._vector_store_connector,
)
super().__init__(**kwargs)
def assemble(self, input_value: IN) -> Any:
"""assemble knowledge for input value."""
if self._vector_store_connector:
self._assembler.persist()
return self._assembler.get_chunks()

44
dbgpt/serve/rag/operators/embedding.py
View File

@@ -1,44 +0,0 @@
from typing import Any, Optional
from dbgpt.core.awel.task.base import IN
from dbgpt.rag.chunk_manager import ChunkParameters
from dbgpt.rag.knowledge.base import Knowledge
from dbgpt.serve.rag.assembler.embedding import EmbeddingAssembler
from dbgpt.serve.rag.operators.base import AssemblerOperator
from dbgpt.storage.vector_store.connector import VectorStoreConnector
class EmbeddingAssemblerOperator(AssemblerOperator[Any, Any]):
"""The Embedding Assembler Operator.
Args:
knowledge (Knowledge): The knowledge.
chunk_parameters (Optional[ChunkParameters], optional): The chunk parameters. Defaults to None.
vector_store_connector (VectorStoreConnector, optional): The vector store connector. Defaults to None.
"""
def __init__(
self,
chunk_parameters: Optional[ChunkParameters] = ChunkParameters(
chunk_strategy="CHUNK_BY_SIZE"
),
vector_store_connector: VectorStoreConnector = None,
**kwargs
):
"""
Args:
chunk_parameters (Optional[ChunkParameters], optional): The chunk parameters. Defaults to ChunkParameters(chunk_strategy="CHUNK_BY_SIZE").
vector_store_connector (VectorStoreConnector, optional): The vector store connector. Defaults to None.
"""
self._chunk_parameters = chunk_parameters
self._vector_store_connector = vector_store_connector
super().__init__(**kwargs)
def assemble(self, knowledge: IN) -> Any:
"""assemble knowledge for input value."""
assembler = EmbeddingAssembler.load_from_knowledge(
knowledge=knowledge,
chunk_parameters=self._chunk_parameters,
vector_store_connector=self._vector_store_connector,
)
assembler.persist()
return assembler.get_chunks()