fix(rag): Fix db schema aretriever bug (#1755)

This commit is contained in:
Fangyin Cheng
2024-07-30 10:30:42 +08:00
committed by GitHub
parent 55c8b39e2e
commit 25d7d94b89
3 changed files with 18 additions and 30 deletions

View File

@@ -167,7 +167,7 @@ class DBSchemaRetriever(BaseRetriever):
result_candidates = await run_async_tasks(
tasks=candidates, concurrency_limit=1
)
return result_candidates
return cast(List[Chunk], reduce(lambda x, y: x + y, result_candidates))
else:
from dbgpt.rag.summary.rdbms_db_summary import ( # noqa: F401
_parse_db_summary,
@@ -177,7 +177,9 @@ class DBSchemaRetriever(BaseRetriever):
tasks=[self._aparse_db_summary(root_tracer.get_current_span_id())],
concurrency_limit=1,
)
return [Chunk(content=table_summary) for table_summary in table_summaries]
return [
Chunk(content=table_summary) for table_summary in table_summaries[0]
]
async def _aretrieve_with_score(
self,