mirror of
https://github.com/csunny/DB-GPT.git
synced 2025-10-31 14:56:04 +00:00
22 lines
743 B
Python
22 lines
743 B
Python
"""Datasource operator for RDBMS database."""
|
|
|
|
from typing import Any, List
|
|
|
|
from dbgpt.core.interface.operators.retriever import RetrieverOperator
|
|
from dbgpt.datasource.base import BaseConnector
|
|
from dbgpt.rag.summary.rdbms_db_summary import _parse_db_summary
|
|
|
|
|
|
class DatasourceRetrieverOperator(RetrieverOperator[Any, List[str]]):
|
|
"""The Datasource Retriever Operator."""
|
|
|
|
def __init__(self, connector: BaseConnector, **kwargs):
|
|
"""Create a new DatasourceRetrieverOperator."""
|
|
super().__init__(**kwargs)
|
|
self._connector = connector
|
|
|
|
def retrieve(self, input_value: Any) -> List[str]:
|
|
"""Retrieve the database summary."""
|
|
summary = _parse_db_summary(self._connector)
|
|
return summary
|