Skip to content

Commit 2f4f027

Browse files
robinji0robinji0
committed
save schema embedding update
1 parent 454adcc commit 2f4f027

5 files changed

Lines changed: 195 additions & 8 deletions

File tree

‎chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/ChatController.java‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,7 @@
4545
import ai.chat2db.server.web.api.controller.ai.rest.client.RestAIClient;
4646
import ai.chat2db.server.web.api.http.GatewayClientService;
4747
import ai.chat2db.server.web.api.http.model.TableSchema;
48+
import ai.chat2db.server.web.api.http.request.TableSchemaRequest;
4849
import ai.chat2db.server.web.api.http.response.TableSchemaResponse;
4950
import ai.chat2db.server.web.api.util.ApplicationContextUtil;
5051
import ai.chat2db.server.web.api.controller.ai.openai.client.OpenAIClient;
@@ -528,7 +529,15 @@ public String queryDatabaseSchema(ChatQueryRequest queryRequest) throws IOExcept
528529
contentVector.add(response.getData().get(0).getEmbedding());
529530

530531
// search embedding
531-
DataResult<TableSchemaResponse> result = gatewayClientService.schemaVectorSearch(contentVector, queryRequest.getDataSourceId());
532+
TableSchemaRequest tableSchemaRequest = new TableSchemaRequest();
533+
tableSchemaRequest.setSchemaVector(contentVector);
534+
tableSchemaRequest.setDataSourceId(queryRequest.getDataSourceId());
535+
String databaseName = StringUtils.isNotBlank(queryRequest.getDatabaseName()) ? queryRequest.getDatabaseName() : queryRequest.getSchemaName();
536+
if (Objects.isNull(databaseName)) {
537+
databaseName = "";
538+
}
539+
tableSchemaRequest.setDatabaseName(databaseName);
540+
DataResult<TableSchemaResponse> result = gatewayClientService.schemaVectorSearch(tableSchemaRequest);
532541

533542
List<String> schemas = Lists.newArrayList();
534543
if (CollectionUtils.isNotEmpty(result.getData().getTableSchemas())) {

‎chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/EmbeddingController.java‎

Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
package ai.chat2db.server.web.api.controller.ai;
2+
3+
import ai.chat2db.server.domain.api.param.ShowCreateTableParam;
4+
import ai.chat2db.server.domain.api.param.TablePageQueryParam;
5+
import ai.chat2db.server.domain.api.param.TableSelector;
6+
import ai.chat2db.server.domain.api.service.TableService;
7+
import ai.chat2db.server.tools.base.wrapper.result.ActionResult;
8+
import ai.chat2db.server.tools.base.wrapper.result.DataResult;
9+
import ai.chat2db.server.tools.base.wrapper.result.PageResult;
10+
import ai.chat2db.server.tools.common.exception.ParamBusinessException;
11+
import ai.chat2db.server.web.api.aspect.ConnectionInfoAspect;
12+
import ai.chat2db.server.web.api.controller.ai.fastchat.embeddings.FastChatEmbeddingResponse;
13+
import ai.chat2db.server.web.api.controller.rdb.converter.RdbWebConverter;
14+
import ai.chat2db.server.web.api.controller.rdb.request.TableBriefQueryRequest;
15+
import ai.chat2db.server.web.api.http.GatewayClientService;
16+
import ai.chat2db.server.web.api.http.request.TableSchemaRequest;
17+
import ai.chat2db.spi.model.Table;
18+
import com.google.common.collect.Lists;
19+
import jakarta.annotation.Resource;
20+
import jakarta.validation.Valid;
21+
import lombok.extern.slf4j.Slf4j;
22+
import org.apache.commons.collections4.CollectionUtils;
23+
import org.apache.commons.lang3.StringUtils;
24+
import org.springframework.beans.factory.annotation.Autowired;
25+
import org.springframework.web.bind.annotation.CrossOrigin;
26+
import org.springframework.web.bind.annotation.PostMapping;
27+
import org.springframework.web.bind.annotation.RequestMapping;
28+
import org.springframework.web.bind.annotation.RestController;
29+
30+
import java.io.IOException;
31+
import java.math.BigDecimal;
32+
import java.util.ArrayList;
33+
import java.util.List;
34+
import java.util.Objects;
35+
36+
/**
37+
* @author moji
38+
*/
39+
@RestController
40+
@ConnectionInfoAspect
41+
@RequestMapping("/api/ai/embedding")
42+
@Slf4j
43+
public class EmbeddingController extends ChatController {
44+
45+
46+
@Resource
47+
private GatewayClientService gatewayClientService;
48+
49+
@Autowired
50+
private RdbWebConverter rdbWebConverter;
51+
52+
@Autowired
53+
private TableService tableService;
54+
55+
/**
56+
* save knowledge embeddings from pdf file
57+
*
58+
* @param request
59+
* @return
60+
* @throws IOException
61+
*/
62+
@PostMapping("/datasource")
63+
@CrossOrigin
64+
public ActionResult embeddings(@Valid TableBriefQueryRequest request)
65+
throws Exception {
66+
67+
// query tables
68+
request.setPageSize(1000);
69+
TablePageQueryParam queryParam = rdbWebConverter.tablePageRequest2param(request);
70+
TableSelector tableSelector = new TableSelector();
71+
tableSelector.setColumnList(false);
72+
tableSelector.setIndexList(false);
73+
PageResult<Table> tableDTOPageResult = tableService.pageQuery(queryParam, tableSelector);
74+
75+
List<Table> tables = tableDTOPageResult.getData();
76+
if (CollectionUtils.isEmpty(tables)) {
77+
return ActionResult.isSuccess();
78+
}
79+
80+
String tableName = tables.get(0).getName();
81+
String tableSchema = queryTableDdl(tableName, request);
82+
83+
if (StringUtils.isBlank(tableSchema)) {
84+
throw new ParamBusinessException("tableSchema is empty");
85+
}
86+
87+
// save first table embedding
88+
TableSchemaRequest tableSchemaRequest = new TableSchemaRequest();
89+
tableSchemaRequest.setDataSourceId(request.getDataSourceId());
90+
tableSchemaRequest.setDeleteBeforeInsert(true);
91+
String databaseName = StringUtils.isNotBlank(request.getDatabaseName()) ? request.getDatabaseName() : request.getSchemaName();
92+
if (Objects.isNull(databaseName)) {
93+
databaseName = "";
94+
}
95+
tableSchemaRequest.setDatabaseName(databaseName);
96+
97+
saveTableEmbedding(tableSchema, tableSchemaRequest);
98+
99+
// save other table embedding
100+
tableSchemaRequest.setDeleteBeforeInsert(false);
101+
for (int i = 1; i < tables.size(); i++) {
102+
tableName = tables.get(i).getName();
103+
tableSchema = queryTableDdl(tableName, request);
104+
if (StringUtils.isBlank(tableSchema)) {
105+
continue;
106+
}
107+
saveTableEmbedding(tableSchema, tableSchemaRequest);
108+
}
109+
110+
// query all the tables
111+
Long totalTableCount = tableDTOPageResult.getTotal();
112+
Integer pageSize = queryParam.getPageSize();
113+
if (pageSize < totalTableCount) {
114+
for (int i = 2; i < totalTableCount/pageSize + 1; i++) {
115+
queryParam.setPageNo(i);
116+
tableDTOPageResult = tableService.pageQuery(queryParam, tableSelector);
117+
tables = tableDTOPageResult.getData();
118+
for (Table table : tables) {
119+
tableName = table.getName();
120+
tableSchema = queryTableDdl(tableName, request);
121+
if (StringUtils.isBlank(tableSchema)) {
122+
continue;
123+
}
124+
saveTableEmbedding(tableSchema, tableSchemaRequest);
125+
}
126+
}
127+
}
128+
129+
return ActionResult.isSuccess();
130+
}
131+
132+
/**
133+
* save table embedding
134+
*
135+
* @param tableSchema
136+
* @param tableSchemaRequest
137+
* @throws Exception
138+
*/
139+
private void saveTableEmbedding(String tableSchema, TableSchemaRequest tableSchemaRequest) throws Exception{
140+
List<String> schemaList = Lists.newArrayList(tableSchema);
141+
tableSchemaRequest.setSchemaList(schemaList);
142+
143+
List<List<BigDecimal>> contentVector = new ArrayList<>();
144+
for(String str : schemaList){
145+
// request embedding
146+
FastChatEmbeddingResponse response = distributeAIEmbedding(str);
147+
if(response == null){
148+
continue;
149+
}
150+
contentVector.add(response.getData().get(0).getEmbedding());
151+
}
152+
tableSchemaRequest.setSchemaVector(contentVector);
153+
154+
// save table embedding
155+
gatewayClientService.schemaVectorSave(tableSchemaRequest);
156+
}
157+
158+
/**
159+
* query table schema
160+
*
161+
* @param tableName
162+
* @param request
163+
* @return
164+
*/
165+
private String queryTableDdl(String tableName, TableBriefQueryRequest request) {
166+
ShowCreateTableParam param = new ShowCreateTableParam();
167+
param.setTableName(tableName);
168+
param.setDataSourceId(request.getDataSourceId());
169+
param.setDatabaseName(request.getDatabaseName());
170+
param.setSchemaName(request.getSchemaName());
171+
DataResult<String> tableSchema = tableService.showCreateTable(param);
172+
return tableSchema.getData();
173+
}
174+
175+
}

‎chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/controller/ai/KnowledgeController.java‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -102,7 +102,9 @@ public SseEmitter search(ChatQueryRequest queryRequest, @RequestHeader Map<Strin
102102
contentVector.add(response.getData().get(0).getEmbedding());
103103

104104
// search embedding
105-
DataResult<KnowledgeResponse> result = gatewayClientService.knowledgeVectorSearch(contentVector);
105+
KnowledgeRequest knowledgeRequest = new KnowledgeRequest();
106+
knowledgeRequest.setContentVector(contentVector);
107+
DataResult<KnowledgeResponse> result = gatewayClientService.knowledgeVectorSearch(knowledgeRequest);
106108
queryRequest.setPromptType(PromptType.TEXT_GENERATION.getCode());
107109
String prompt = queryRequest.getMessage();
108110
if (CollectionUtils.isNotEmpty(result.getData().getKnowledgeList())) {

‎chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/http/GatewayClientService.java‎

Lines changed: 3 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -81,15 +81,14 @@ public interface GatewayClientService {
8181
* @return
8282
*/
8383
@Get("/api/milvus/knowledge/search")
84-
DataResult<KnowledgeResponse> knowledgeVectorSearch(List<List<BigDecimal>> searchVectors);
84+
DataResult<KnowledgeResponse> knowledgeVectorSearch(KnowledgeRequest searchVectors);
8585

8686
/**
8787
* save table schema vector
8888
*
89-
* @param searchVectors
90-
* @param datasourceId
89+
* @param request
9190
* @return
9291
*/
9392
@Get("/api/milvus/schema/search")
94-
DataResult<TableSchemaResponse> schemaVectorSearch(List<List<BigDecimal>> searchVectors, Long datasourceId);
93+
DataResult<TableSchemaResponse> schemaVectorSearch(TableSchemaRequest request);
9594
}

‎chat2db-server/chat2db-server-web/chat2db-server-web-api/src/main/java/ai/chat2db/server/web/api/http/request/TableSchemaRequest.java‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@ public class TableSchemaRequest {
1616

1717
private Long dataSourceId;
1818

19-
private List<List<BigDecimal>> contentVector;
19+
private String databaseName;
2020

21-
private List<String> sentenceList;
21+
private List<java.util.List<BigDecimal>> schemaVector;
22+
23+
private List<String> schemaList;
2224

2325
private Boolean deleteBeforeInsert;
2426
}

0 commit comments

Comments
 (0)