mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-20 22:51:06 +08:00
Implement OpenAI chat completions in GO (#16177)
### What problem does this PR solve? Implement OpenAI chat completions in GO POST /api/v1/openai/<chat_id>/chat/completions OpenAI chat cli: internal/development.md ### Type of change - [x] Refactoring
This commit is contained in:
qinling0210
@@ -1916,6 +1916,7 @@ func convertMatchingField(fieldWeightStr string) string {
|
||||
"authors_tks": "authors@ft_authors_rag_coarse",
|
||||
"authors_sm_tks": "authors@ft_authors_rag_fine",
|
||||
"tag_kwd": "tag_kwd@ft_tag_kwd_whitespace__",
|
||||
"toc_kwd": "toc_kwd@ft_toc_kwd_whitespace__",
|
||||
// Skill index fields
|
||||
"name": "name@ft_name_rag_coarse",
|
||||
"tags": "tags@ft_tags_rag_coarse",
|
||||
|
||||
@@ -30,10 +30,18 @@ import (
|
||||
infinity "github.com/infiniflow/infinity-go-sdk"
|
||||
)
|
||||
|
||||
// infinityClient Infinity SDK client wrapper
|
||||
type infinityClient struct {
|
||||
conn *infinity.InfinityConnection
|
||||
dbName string
|
||||
|
||||
// Original URI from config, used by RunSQL to extract the host.
|
||||
hostURI string
|
||||
|
||||
// Port for psql wire-protocol listener (default 5432).
|
||||
postgresPort int
|
||||
|
||||
// JSON file (under conf/) with the field-name alias map.
|
||||
mappingFileName string
|
||||
}
|
||||
|
||||
// NewInfinityClient creates a new Infinity client using the SDK
|
||||
@@ -70,8 +78,11 @@ func NewInfinityClient(cfg *server.InfinityConfig) (*infinityClient, error) {
|
||||
}
|
||||
|
||||
client := &infinityClient{
|
||||
conn: conn,
|
||||
dbName: cfg.DBName,
|
||||
conn: conn,
|
||||
dbName: cfg.DBName,
|
||||
hostURI: cfg.URI,
|
||||
postgresPort: cfg.PostgresPort,
|
||||
mappingFileName: cfg.MappingFileName,
|
||||
}
|
||||
|
||||
return client, nil
|
||||
|
||||
@@ -24,12 +24,13 @@ import (
|
||||
"path/filepath"
|
||||
"strings"
|
||||
|
||||
infinity "github.com/infiniflow/infinity-go-sdk"
|
||||
"ragflow/internal/common"
|
||||
"ragflow/internal/dao"
|
||||
"ragflow/internal/engine/types"
|
||||
"ragflow/internal/utility"
|
||||
|
||||
infinity "github.com/infiniflow/infinity-go-sdk"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
@@ -568,25 +569,113 @@ func (e *infinityEngine) SearchMetadata(ctx context.Context, req *types.SearchMe
|
||||
}, nil
|
||||
}
|
||||
|
||||
// Build search request for metadata - simpler than chunk search, no match expressions
|
||||
searchReq := &types.SearchRequest{
|
||||
IndexNames: []string{tableName},
|
||||
Offset: req.Offset,
|
||||
Limit: req.Limit,
|
||||
SelectFields: req.SelectFields,
|
||||
Filter: req.Filter,
|
||||
MatchExprs: nil, // No match expressions for metadata
|
||||
OrderBy: req.OrderBy,
|
||||
RankFeature: nil,
|
||||
// Build output columns: use caller-specified fields, or "*" for all columns
|
||||
var outputColumns []string
|
||||
if len(req.SelectFields) > 0 {
|
||||
outputColumns = req.SelectFields
|
||||
} else {
|
||||
outputColumns = []string{"*"}
|
||||
}
|
||||
|
||||
result, err := e.Search(ctx, searchReq)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
// Pagination defaults
|
||||
pageSize := req.Limit
|
||||
if pageSize <= 0 {
|
||||
pageSize = 30
|
||||
}
|
||||
offset := req.Offset
|
||||
if offset < 0 {
|
||||
offset = 0
|
||||
}
|
||||
|
||||
// Build filter from req.Filter
|
||||
var filterStr string
|
||||
if req.Filter != nil {
|
||||
filterStr = equivalentConditionToStr(req.Filter)
|
||||
}
|
||||
|
||||
// Get database and table
|
||||
db, err := e.client.conn.GetDatabase(e.client.dbName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get database: %w", err)
|
||||
}
|
||||
|
||||
tbl, err := db.GetTable(tableName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("failed to get metadata table %s: %w", tableName, err)
|
||||
}
|
||||
|
||||
// Build Infinity query (chainable API)
|
||||
table := tbl.Output(outputColumns)
|
||||
if filterStr != "" {
|
||||
table = table.Filter(filterStr)
|
||||
}
|
||||
|
||||
// Add order_by if provided
|
||||
if req.OrderBy != nil && len(req.OrderBy.Fields) > 0 {
|
||||
var sortFields [][2]interface{}
|
||||
for _, orderField := range req.OrderBy.Fields {
|
||||
sortType := infinity.SortTypeAsc
|
||||
if orderField.Type == types.SortDesc {
|
||||
sortType = infinity.SortTypeDesc
|
||||
}
|
||||
sortFields = append(sortFields, [2]interface{}{orderField.Field, sortType})
|
||||
}
|
||||
table = table.Sort(sortFields)
|
||||
}
|
||||
|
||||
table = table.Limit(pageSize)
|
||||
if offset > 0 {
|
||||
table = table.Offset(offset)
|
||||
}
|
||||
table = table.Option(map[string]interface{}{"total_hits_count": true})
|
||||
|
||||
// Execute query
|
||||
df, err := table.ToDataFrame()
|
||||
if err != nil {
|
||||
common.Warn("Infinity SearchMetadata query failed",
|
||||
zap.String("tableName", tableName),
|
||||
zap.Error(err))
|
||||
return nil, fmt.Errorf("metadata query failed: %w", err)
|
||||
}
|
||||
|
||||
// Convert column-oriented DataFrame to row-oriented records
|
||||
records := make([]map[string]interface{}, 0)
|
||||
for colName, colData := range df.ColumnData {
|
||||
for i, val := range colData {
|
||||
for len(records) <= i {
|
||||
records = append(records, make(map[string]interface{}))
|
||||
}
|
||||
records[i][colName] = val
|
||||
}
|
||||
}
|
||||
|
||||
// Handle ROW_ID -> row_id() mapping (Infinity internal column)
|
||||
for _, rec := range records {
|
||||
if val, ok := rec["ROW_ID"]; ok {
|
||||
rec["row_id()"] = val
|
||||
delete(rec, "ROW_ID")
|
||||
}
|
||||
}
|
||||
|
||||
// Realign meta_fields column for multi-row queries (Infinity may
|
||||
// concatenate values into one blob with 4-byte length prefix)
|
||||
realignMetaFieldsColumn(records)
|
||||
|
||||
// Parse total_hits_count from ExtraInfo
|
||||
var totalHits int64
|
||||
if df.ExtraInfo != "" {
|
||||
if t, ok := totalHitsFromInfinityExtraInfo(df.ExtraInfo); ok {
|
||||
totalHits = t
|
||||
}
|
||||
}
|
||||
|
||||
common.Debug("SearchMetadata in Infinity completed",
|
||||
zap.Int("rows", len(records)),
|
||||
zap.Int64("total", totalHits))
|
||||
|
||||
return &types.SearchMetadataResult{
|
||||
MetadataRecords: result.Chunks,
|
||||
Total: result.Total,
|
||||
MetadataRecords: records,
|
||||
Total: totalHits,
|
||||
}, nil
|
||||
}
|
||||
|
||||
|
||||
330
internal/engine/infinity/sql.go
Normal file
330
internal/engine/infinity/sql.go
Normal file
@@ -0,0 +1,330 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package infinity
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"os"
|
||||
"os/exec"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"ragflow/internal/common"
|
||||
|
||||
"go.uber.org/zap"
|
||||
)
|
||||
|
||||
const (
|
||||
psqlTimeout = 10 * time.Second
|
||||
defaultPsqlPath = "/usr/bin/psql"
|
||||
defaultPsqlHost = "infinity"
|
||||
defaultPsqlPort = "5432"
|
||||
)
|
||||
|
||||
var whitespaceRe = regexp.MustCompile("[ `]+")
|
||||
|
||||
var rowCountFooterRe = regexp.MustCompile(`^\(\d+ rows?`)
|
||||
|
||||
// fieldMappingEntry is one entry in infinity_mapping.json.
|
||||
type fieldMappingEntry struct {
|
||||
Type string `json:"type"`
|
||||
Comment string `json:"comment"`
|
||||
}
|
||||
|
||||
// loadFieldMapping reads infinity_mapping.json and returns alias→actual
|
||||
// and actual→firstAlias maps. Silently returns empty maps on missing file.
|
||||
func loadFieldMapping(mappingFileName string) (aliasToActual map[string]string, actualToFirstAlias map[string]string, err error) {
|
||||
if mappingFileName == "" {
|
||||
mappingFileName = "infinity_mapping.json"
|
||||
}
|
||||
confPath := filepath.Join(projectBaseDir(), "conf", mappingFileName)
|
||||
data, err := os.ReadFile(confPath)
|
||||
if err != nil {
|
||||
if errors.Is(err, os.ErrNotExist) {
|
||||
return map[string]string{}, map[string]string{}, nil
|
||||
}
|
||||
return nil, nil, fmt.Errorf("load field mapping %q: %w", confPath, err)
|
||||
}
|
||||
|
||||
fields := map[string]fieldMappingEntry{}
|
||||
if err := json.Unmarshal(data, &fields); err != nil {
|
||||
return nil, nil, fmt.Errorf("parse field mapping %q: %w", confPath, err)
|
||||
}
|
||||
|
||||
aliasToActual = make(map[string]string, len(fields)*2)
|
||||
actualToFirstAlias = make(map[string]string, len(fields))
|
||||
for actual, info := range fields {
|
||||
if info.Comment == "" {
|
||||
continue
|
||||
}
|
||||
var firstAlias string
|
||||
for _, raw := range strings.Split(info.Comment, ",") {
|
||||
alias := strings.TrimSpace(raw)
|
||||
if alias == "" {
|
||||
continue
|
||||
}
|
||||
aliasToActual[alias] = actual
|
||||
if firstAlias == "" {
|
||||
firstAlias = alias
|
||||
}
|
||||
}
|
||||
if firstAlias != "" {
|
||||
actualToFirstAlias[actual] = firstAlias
|
||||
}
|
||||
}
|
||||
return aliasToActual, actualToFirstAlias, nil
|
||||
}
|
||||
|
||||
// projectBaseDir returns the project root. Honors RAG_PROJECT_BASE and
|
||||
// RAG_DEPLOY_BASE env vars; falls back to working directory.
|
||||
func projectBaseDir() string {
|
||||
if v := os.Getenv("RAG_PROJECT_BASE"); v != "" {
|
||||
return v
|
||||
}
|
||||
if v := os.Getenv("RAG_DEPLOY_BASE"); v != "" {
|
||||
return v
|
||||
}
|
||||
// Fall back to the repository root. The Go engine package lives at
|
||||
// internal/engine/infinity/; the repo root is three levels up.
|
||||
wd, err := os.Getwd()
|
||||
if err != nil {
|
||||
return "."
|
||||
}
|
||||
return wd
|
||||
}
|
||||
|
||||
// preprocessSQL collapses spaces/backticks and strips '%'.
|
||||
func preprocessSQL(sql string) string {
|
||||
sql = whitespaceRe.ReplaceAllString(sql, " ")
|
||||
sql = strings.ReplaceAll(sql, "%", "")
|
||||
return sql
|
||||
}
|
||||
|
||||
// rewriteFieldAliases rewrites alias field names to actual stored names
|
||||
// in SELECT, WHERE, ORDER BY, GROUP BY, and HAVING clauses.
|
||||
func rewriteFieldAliases(sql string, aliasToActual map[string]string) string {
|
||||
if len(aliasToActual) == 0 {
|
||||
return sql
|
||||
}
|
||||
selectRe := regexp.MustCompile(`(?si)(select\s+)(.+?)(\s+from\b)`)
|
||||
sql = selectRe.ReplaceAllStringFunc(sql, func(m string) string {
|
||||
parts := selectRe.FindStringSubmatch(m)
|
||||
prefix, cols, suffix := parts[1], parts[2], parts[3]
|
||||
for alias, actual := range aliasToActual {
|
||||
pat := regexp.MustCompile(`(^|[,\s])` + regexp.QuoteMeta(alias) + `($|[,\s])`)
|
||||
cols = pat.ReplaceAllString(cols, "${1}"+actual+"${2}")
|
||||
}
|
||||
return prefix + cols + suffix
|
||||
})
|
||||
|
||||
clauseAliases := func(sql, keyword string) string {
|
||||
return rewriteFirstAliasAfterKeyword(sql, keyword, aliasToActual)
|
||||
}
|
||||
sql = clauseAliases(sql, "where")
|
||||
sql = clauseAliases(sql, "order by")
|
||||
sql = clauseAliases(sql, "group by")
|
||||
sql = clauseAliases(sql, "having")
|
||||
return sql
|
||||
}
|
||||
|
||||
func rewriteFirstAliasAfterKeyword(sql, keyword string, aliasToActual map[string]string) string {
|
||||
for alias, actual := range aliasToActual {
|
||||
aliasPat := regexp.MustCompile(`\b` + regexp.QuoteMeta(alias) + `\b`)
|
||||
kwIdx := regexp.MustCompile(`(?i)\b` + regexp.QuoteMeta(keyword) + `\b`).FindStringIndex(sql)
|
||||
if kwIdx == nil {
|
||||
continue
|
||||
}
|
||||
tail := sql[kwIdx[1]:]
|
||||
aliasIdx := aliasPat.FindStringIndex(tail)
|
||||
if aliasIdx == nil {
|
||||
continue
|
||||
}
|
||||
absStart := kwIdx[1] + aliasIdx[0]
|
||||
absEnd := kwIdx[1] + aliasIdx[1]
|
||||
sql = sql[:absStart] + actual + sql[absEnd:]
|
||||
}
|
||||
return sql
|
||||
}
|
||||
|
||||
// psqlResult is the structured parse of a psql table-format output.
|
||||
type psqlResult struct {
|
||||
Columns []string
|
||||
Rows [][]string
|
||||
}
|
||||
|
||||
// runPsql shells out to psql and parses the table-format output.
|
||||
func runPsql(ctx context.Context, host, port, sql string) (*psqlResult, error) {
|
||||
psqlPath, err := findPsqlBinary()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
ctx, cancel := context.WithTimeout(ctx, psqlTimeout)
|
||||
defer cancel()
|
||||
|
||||
cmd := exec.CommandContext(ctx, psqlPath, "-h", host, "-p", port, "-c", sql)
|
||||
common.Debug("executing psql",
|
||||
zap.String("path", psqlPath),
|
||||
zap.String("host", host),
|
||||
zap.String("port", port),
|
||||
)
|
||||
var stdout, stderr bytes.Buffer
|
||||
cmd.Stdout = &stdout
|
||||
cmd.Stderr = &stderr
|
||||
|
||||
if err := cmd.Run(); err != nil {
|
||||
if ctx.Err() != nil {
|
||||
return nil, fmt.Errorf("SQL timeout\n\nSQL: %s", sql)
|
||||
}
|
||||
return nil, fmt.Errorf("psql command failed: %s\nSQL: %s", strings.TrimSpace(stderr.String()), sql)
|
||||
}
|
||||
return parsePsqlTable(stdout.String()), nil
|
||||
}
|
||||
|
||||
// findPsqlBinary checks PATH first, then falls back to defaultPsqlPath.
|
||||
func findPsqlBinary() (string, error) {
|
||||
if path, err := exec.LookPath("psql"); err == nil {
|
||||
return path, nil
|
||||
}
|
||||
if _, err := os.Stat(defaultPsqlPath); err == nil {
|
||||
return defaultPsqlPath, nil
|
||||
}
|
||||
return "", fmt.Errorf("psql not found on PATH and not at %q", defaultPsqlPath)
|
||||
}
|
||||
|
||||
// parsePsqlTable parses psql's pipe-delimited output:
|
||||
//
|
||||
// col1 | col2
|
||||
// -----+-----
|
||||
// val1 | val2
|
||||
func parsePsqlTable(output string) *psqlResult {
|
||||
res := &psqlResult{}
|
||||
out := strings.TrimSpace(output)
|
||||
if out == "" {
|
||||
return res
|
||||
}
|
||||
lines := strings.Split(out, "\n")
|
||||
if len(lines) == 0 {
|
||||
return res
|
||||
}
|
||||
|
||||
for _, raw := range strings.Split(lines[0], "|") {
|
||||
if col := strings.TrimSpace(raw); col != "" {
|
||||
res.Columns = append(res.Columns, col)
|
||||
}
|
||||
}
|
||||
|
||||
dataStart := 1
|
||||
if len(lines) >= 2 && strings.Contains(lines[1], "-") {
|
||||
dataStart = 2
|
||||
}
|
||||
for i := dataStart; i < len(lines); i++ {
|
||||
line := strings.TrimSpace(lines[i])
|
||||
if line == "" || rowCountFooterRe.MatchString(line) {
|
||||
continue
|
||||
}
|
||||
cells := strings.Split(line, "|")
|
||||
for j := range cells {
|
||||
cells[j] = strings.TrimSpace(cells[j])
|
||||
}
|
||||
switch {
|
||||
case len(cells) == len(res.Columns):
|
||||
res.Rows = append(res.Rows, cells)
|
||||
case len(cells) > len(res.Columns):
|
||||
res.Rows = append(res.Rows, cells[:len(res.Columns)])
|
||||
default:
|
||||
padded := make([]string, len(res.Columns))
|
||||
copy(padded, cells)
|
||||
for k := len(cells); k < len(res.Columns); k++ {
|
||||
padded[k] = ""
|
||||
}
|
||||
res.Rows = append(res.Rows, padded)
|
||||
}
|
||||
}
|
||||
return res
|
||||
}
|
||||
|
||||
// toRowMaps converts psqlResult to a slice of column-keyed maps.
|
||||
func toRowMaps(res *psqlResult) []map[string]interface{} {
|
||||
if res == nil || len(res.Rows) == 0 {
|
||||
return nil
|
||||
}
|
||||
out := make([]map[string]interface{}, 0, len(res.Rows))
|
||||
for _, row := range res.Rows {
|
||||
m := make(map[string]interface{}, len(res.Columns))
|
||||
for j, col := range res.Columns {
|
||||
if j < len(row) {
|
||||
m[col] = row[j]
|
||||
}
|
||||
}
|
||||
out = append(out, m)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
func resolvePsqlHostPort(hostURI string, postgresPort int) (host, port string) {
|
||||
host = defaultPsqlHost
|
||||
port = defaultPsqlPort
|
||||
if postgresPort > 0 {
|
||||
port = strconv.Itoa(postgresPort)
|
||||
}
|
||||
if hostURI != "" {
|
||||
if h, _, ok := strings.Cut(hostURI, ":"); ok && h != "" {
|
||||
host = h
|
||||
}
|
||||
}
|
||||
return host, port
|
||||
}
|
||||
|
||||
// RunSQL implements the SQL retrieval path: preprocess, rewrite aliases,
|
||||
// run psql subprocess, parse output.
|
||||
func (e *infinityEngine) RunSQL(ctx context.Context, tableName string, sqlText string, kbIDs []string, _ string) ([]map[string]interface{}, error) {
|
||||
if e == nil || e.client == nil {
|
||||
return nil, fmt.Errorf("infinity RunSQL: client not initialized")
|
||||
}
|
||||
sqlText = strings.TrimSpace(sqlText)
|
||||
if sqlText == "" {
|
||||
return nil, fmt.Errorf("infinity RunSQL: empty SQL")
|
||||
}
|
||||
|
||||
common.Debug("InfinityConnection.sql get sql", zap.String("sql", sqlText))
|
||||
|
||||
sqlText = preprocessSQL(sqlText)
|
||||
|
||||
aliasMap, _, err := loadFieldMapping(e.client.mappingFileName)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("infinity RunSQL: %w", err)
|
||||
}
|
||||
sqlText = rewriteFieldAliases(sqlText, aliasMap)
|
||||
|
||||
common.Debug("InfinityConnection.sql to execute", zap.String("sql", sqlText))
|
||||
|
||||
host, port := resolvePsqlHostPort(e.client.hostURI, e.client.postgresPort)
|
||||
res, err := runPsql(ctx, host, port, sqlText)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
return toRowMaps(res), nil
|
||||
}
|
||||
394
internal/engine/infinity/sql_test.go
Normal file
394
internal/engine/infinity/sql_test.go
Normal file
@@ -0,0 +1,394 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package infinity
|
||||
|
||||
import (
|
||||
"os"
|
||||
"path/filepath"
|
||||
"reflect"
|
||||
"testing"
|
||||
)
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// preprocessSQL — mirrors infinity_conn_base.py:788-789.
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
func TestPreprocessSQL_WhitespaceAndBackticks(t *testing.T) {
|
||||
cases := []struct {
|
||||
in, want string
|
||||
}{
|
||||
{"a b", "a b"},
|
||||
{"a b c", "a b c"},
|
||||
{"a`b`c", "a b c"},
|
||||
{"a `` b", "a b"},
|
||||
// The regex collapses ALL runs of spaces/backticks — including
|
||||
// leading and trailing whitespace. Trimming is a separate step
|
||||
// in RunSQL (strings.TrimSpace before the preprocessing pass).
|
||||
{" leading and trailing ", " leading and trailing "},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := preprocessSQL(c.in); got != c.want {
|
||||
t.Errorf("preprocessSQL(%q) = %q, want %q", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreprocessSQL_StripsPercent(t *testing.T) {
|
||||
cases := []struct {
|
||||
in, want string
|
||||
}{
|
||||
{"count > 0 %", "count > 0 "},
|
||||
{"100% match", "100 match"},
|
||||
{"%%%", ""},
|
||||
}
|
||||
for _, c := range cases {
|
||||
if got := preprocessSQL(c.in); got != c.want {
|
||||
t.Errorf("preprocessSQL(%q) = %q, want %q", c.in, got, c.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPreprocessSQL_Combined(t *testing.T) {
|
||||
in := "SELECT docnm_kwd FROM `ragflow_t1` WHERE count > 0 %"
|
||||
got := preprocessSQL(in)
|
||||
want := "SELECT docnm_kwd FROM ragflow_t1 WHERE count > 0 "
|
||||
if got != want {
|
||||
t.Errorf("preprocessSQL(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// rewriteFieldAliases — mirrors infinity_conn_base.py:809-830.
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
func TestRewriteFieldAliases_SelectClause(t *testing.T) {
|
||||
aliases := map[string]string{
|
||||
"docnm_kwd": "docnm",
|
||||
"title_tks": "docnm",
|
||||
"title_sm_tks": "docnm",
|
||||
"content_ltks": "content",
|
||||
}
|
||||
in := "select docnm_kwd, title_tks, content_ltks from ragflow_t1"
|
||||
got := rewriteFieldAliases(in, aliases)
|
||||
want := "select docnm, docnm, content from ragflow_t1"
|
||||
if got != want {
|
||||
t.Errorf("rewriteFieldAliases(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRewriteFieldAliases_WhereClause(t *testing.T) {
|
||||
aliases := map[string]string{
|
||||
"docnm_kwd": "docnm",
|
||||
}
|
||||
in := "select doc_id from ragflow_t1 where docnm_kwd = 'foo'"
|
||||
got := rewriteFieldAliases(in, aliases)
|
||||
want := "select doc_id from ragflow_t1 where docnm = 'foo'"
|
||||
if got != want {
|
||||
t.Errorf("rewriteFieldAliases(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRewriteFieldAliases_OrderGroupHaving(t *testing.T) {
|
||||
aliases := map[string]string{
|
||||
"docnm_kwd": "docnm",
|
||||
"important_kwd": "important_keywords",
|
||||
}
|
||||
in := "select doc_id from ragflow_t1 order by docnm_kwd group by important_kwd having important_kwd > 0"
|
||||
got := rewriteFieldAliases(in, aliases)
|
||||
want := "select doc_id from ragflow_t1 order by docnm group by important_keywords having important_keywords > 0"
|
||||
if got != want {
|
||||
t.Errorf("rewriteFieldAliases(%q) = %q, want %q", in, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRewriteFieldAliases_EmptyMapIsNoop(t *testing.T) {
|
||||
in := "select docnm_kwd from ragflow_t1"
|
||||
if got := rewriteFieldAliases(in, map[string]string{}); got != in {
|
||||
t.Errorf("empty alias map should not modify SQL; got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRewriteFieldAliases_WordBoundaryProtected(t *testing.T) {
|
||||
// "title" is an alias; "title_sm_tks" should NOT match because
|
||||
// word boundary is enforced.
|
||||
aliases := map[string]string{
|
||||
"title": "docnm",
|
||||
}
|
||||
in := "select title_sm_tks from ragflow_t1"
|
||||
got := rewriteFieldAliases(in, aliases)
|
||||
// "title" inside "title_sm_tks" should NOT be rewritten.
|
||||
want := "select title_sm_tks from ragflow_t1"
|
||||
if got != want {
|
||||
t.Errorf("rewriteFieldAliases(%q) = %q, want %q (title_sm_tks must NOT be touched)", in, got, want)
|
||||
}
|
||||
}
|
||||
|
||||
func TestRewriteFieldAliases_NoAliasMatchLeavesSQLAlone(t *testing.T) {
|
||||
aliases := map[string]string{
|
||||
"docnm_kwd": "docnm",
|
||||
}
|
||||
in := "select content_with_weight from ragflow_t1"
|
||||
got := rewriteFieldAliases(in, aliases)
|
||||
if got != in {
|
||||
t.Errorf("unrelated SQL should be unchanged; got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// parsePsqlTable — mirrors infinity_conn_base.py:894-934.
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
func TestParsePsqlTable_StandardOutput(t *testing.T) {
|
||||
// Sample psql table output for `select 1 as a, 2 as b;`
|
||||
out := ` a | b
|
||||
---+---
|
||||
1 | 2
|
||||
(1 row)`
|
||||
|
||||
res := parsePsqlTable(out)
|
||||
wantCols := []string{"a", "b"}
|
||||
if !reflect.DeepEqual(res.Columns, wantCols) {
|
||||
t.Errorf("columns: got %v, want %v", res.Columns, wantCols)
|
||||
}
|
||||
wantRows := [][]string{{"1", "2"}}
|
||||
if !reflect.DeepEqual(res.Rows, wantRows) {
|
||||
t.Errorf("rows: got %v, want %v", res.Rows, wantRows)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParsePsqlTable_EmptyOutput(t *testing.T) {
|
||||
res := parsePsqlTable("")
|
||||
if len(res.Columns) != 0 || len(res.Rows) != 0 {
|
||||
t.Errorf("empty output should yield (0 cols, 0 rows); got %+v", res)
|
||||
}
|
||||
}
|
||||
|
||||
func TestParsePsqlTable_NoSeparatorLine(t *testing.T) {
|
||||
// Some psql configurations skip the separator line; the parser
|
||||
// should still recover (data starts at line 1 in that case).
|
||||
out := "a | b\n1 | 2"
|
||||
res := parsePsqlTable(out)
|
||||
if len(res.Rows) != 1 {
|
||||
t.Errorf("rows: got %d, want 1", len(res.Rows))
|
||||
}
|
||||
}
|
||||
|
||||
func TestParsePsqlTable_MultipleRowsAndRowCountFooter(t *testing.T) {
|
||||
out := ` id | name
|
||||
----+------
|
||||
1 | foo
|
||||
2 | bar
|
||||
(2 rows)`
|
||||
res := parsePsqlTable(out)
|
||||
wantCols := []string{"id", "name"}
|
||||
if !reflect.DeepEqual(res.Columns, wantCols) {
|
||||
t.Errorf("columns: got %v, want %v", res.Columns, wantCols)
|
||||
}
|
||||
if len(res.Rows) != 2 {
|
||||
t.Errorf("rows: got %d, want 2", len(res.Rows))
|
||||
}
|
||||
if res.Rows[0][0] != "1" || res.Rows[0][1] != "foo" {
|
||||
t.Errorf("row[0]: got %v, want [1 foo]", res.Rows[0])
|
||||
}
|
||||
if res.Rows[1][0] != "2" || res.Rows[1][1] != "bar" {
|
||||
t.Errorf("row[1]: got %v, want [2 bar]", res.Rows[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestParsePsqlTable_PadsAndTruncatesRows(t *testing.T) {
|
||||
// Row with fewer cells → pad with empty strings.
|
||||
// Row with more cells → truncate.
|
||||
out := ` a | b | c
|
||||
---+---+---
|
||||
1 | 2
|
||||
1 | 2 | 3 | 4
|
||||
(2 rows)`
|
||||
res := parsePsqlTable(out)
|
||||
if len(res.Rows) != 2 {
|
||||
t.Fatalf("rows: got %d, want 2", len(res.Rows))
|
||||
}
|
||||
// First row: ["1", "2", ""] (padded)
|
||||
if !reflect.DeepEqual(res.Rows[0], []string{"1", "2", ""}) {
|
||||
t.Errorf("padded row: got %v, want [1 2 ]", res.Rows[0])
|
||||
}
|
||||
// Second row: ["1", "2", "3"] (truncated)
|
||||
if !reflect.DeepEqual(res.Rows[1], []string{"1", "2", "3"}) {
|
||||
t.Errorf("truncated row: got %v, want [1 2 3]", res.Rows[1])
|
||||
}
|
||||
}
|
||||
|
||||
func TestParsePsqlTable_SkipsRowCountFooter(t *testing.T) {
|
||||
out := " a \n---\n 1 \n(1 row)"
|
||||
res := parsePsqlTable(out)
|
||||
if len(res.Rows) != 1 {
|
||||
t.Errorf("row count footer should be skipped; got %d rows", len(res.Rows))
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// toRowMaps — chunk-shape conversion.
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
func TestToRowMaps_EmptyResultsReturnsNil(t *testing.T) {
|
||||
if rows := toRowMaps(nil); rows != nil {
|
||||
t.Errorf("nil result: got %v, want nil", rows)
|
||||
}
|
||||
if rows := toRowMaps(&psqlResult{}); rows != nil {
|
||||
t.Errorf("empty result: got %v, want nil", rows)
|
||||
}
|
||||
}
|
||||
|
||||
func TestToRowMaps_ConvertsToRowMaps(t *testing.T) {
|
||||
res := &psqlResult{
|
||||
Columns: []string{"id", "name"},
|
||||
Rows: [][]string{
|
||||
{"1", "foo"},
|
||||
{"2", "bar"},
|
||||
},
|
||||
}
|
||||
got := toRowMaps(res)
|
||||
want := []map[string]interface{}{
|
||||
{"id": "1", "name": "foo"},
|
||||
{"id": "2", "name": "bar"},
|
||||
}
|
||||
if !reflect.DeepEqual(got, want) {
|
||||
t.Errorf("toRowMaps: got %v, want %v", got, want)
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// resolvePsqlHostPort — mirrors infinity_conn_base.py:838-858.
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
func TestResolvePsqlHostPort_DefaultsWhenConfigEmpty(t *testing.T) {
|
||||
host, port := resolvePsqlHostPort("", 0)
|
||||
if host != defaultPsqlHost {
|
||||
t.Errorf("host: got %q, want %q", host, defaultPsqlHost)
|
||||
}
|
||||
if port != defaultPsqlPort {
|
||||
t.Errorf("port: got %q, want %q", port, defaultPsqlPort)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolvePsqlHostPort_OverridesFromConfig(t *testing.T) {
|
||||
host, port := resolvePsqlHostPort("10.0.0.1:23817", 5433)
|
||||
if host != "10.0.0.1" {
|
||||
t.Errorf("host: got %q, want 10.0.0.1", host)
|
||||
}
|
||||
if port != "5433" {
|
||||
t.Errorf("port: got %q, want 5433", port)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolvePsqlHostPort_EmptyHostInURIFallsBackToDefault(t *testing.T) {
|
||||
// ":23817" parses via strings.Cut to ("", "23817") — the empty
|
||||
// host doesn't override the default, matching Python's
|
||||
// `re.search(r"host=(\S+)", ...)` which only matches a non-empty
|
||||
// value.
|
||||
host, port := resolvePsqlHostPort(":23817", 5432)
|
||||
if host != defaultPsqlHost {
|
||||
t.Errorf("host: got %q, want default %q (empty host in URI should not override)", host, defaultPsqlHost)
|
||||
}
|
||||
if port != "5432" {
|
||||
t.Errorf("port: got %q, want 5432", port)
|
||||
}
|
||||
}
|
||||
|
||||
// -----------------------------------------------------------------------------
|
||||
// loadFieldMapping — mirrors infinity_conn_base.py:793-807.
|
||||
// -----------------------------------------------------------------------------
|
||||
|
||||
func TestLoadFieldMapping_MissingFileReturnsEmpty(t *testing.T) {
|
||||
// Use a name that doesn't exist; the function should silently
|
||||
// return empty maps (matching Python's `os.path.exists` guard).
|
||||
a2a, r2a, err := loadFieldMapping("nonexistent_mapping_xyz.json")
|
||||
if err != nil {
|
||||
t.Fatalf("missing file should be a no-op, got error: %v", err)
|
||||
}
|
||||
if len(a2a) != 0 || len(r2a) != 0 {
|
||||
t.Errorf("missing file should yield empty maps; got a2a=%v r2a=%v", a2a, r2a)
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadFieldMapping_ParsesAliases(t *testing.T) {
|
||||
// Write a temporary mapping file.
|
||||
dir := t.TempDir()
|
||||
mappingPath := filepath.Join(dir, "test_mapping.json")
|
||||
contents := `{
|
||||
"docnm": {"type": "varchar", "comment": "docnm_kwd, title_tks, title_sm_tks"},
|
||||
"content": {"type": "varchar", "comment": "content_with_weight, content_ltks"},
|
||||
"plain": {"type": "varchar"}
|
||||
}`
|
||||
if err := os.WriteFile(mappingPath, []byte(contents), 0o644); err != nil {
|
||||
t.Fatalf("write mapping: %v", err)
|
||||
}
|
||||
|
||||
// Set RAG_PROJECT_BASE to the temp dir's parent so loadFieldMapping
|
||||
// finds the file at <base>/conf/<filename>.
|
||||
os.Setenv("RAG_PROJECT_BASE", dir)
|
||||
defer os.Unsetenv("RAG_PROJECT_BASE")
|
||||
|
||||
// Need to create conf/ subdir.
|
||||
if err := os.MkdirAll(filepath.Join(dir, "conf"), 0o755); err != nil {
|
||||
t.Fatalf("mkdir conf: %v", err)
|
||||
}
|
||||
if err := os.WriteFile(filepath.Join(dir, "conf", "test_mapping.json"), []byte(contents), 0o644); err != nil {
|
||||
t.Fatalf("write conf/mapping: %v", err)
|
||||
}
|
||||
|
||||
a2a, r2a, err := loadFieldMapping("test_mapping.json")
|
||||
if err != nil {
|
||||
t.Fatalf("loadFieldMapping: %v", err)
|
||||
}
|
||||
|
||||
// alias → actual
|
||||
expectedAliases := map[string]string{
|
||||
"docnm_kwd": "docnm",
|
||||
"title_tks": "docnm",
|
||||
"title_sm_tks": "docnm",
|
||||
"content_with_weight": "content",
|
||||
"content_ltks": "content",
|
||||
}
|
||||
if !reflect.DeepEqual(a2a, expectedAliases) {
|
||||
t.Errorf("aliasToActual: got %v, want %v", a2a, expectedAliases)
|
||||
}
|
||||
|
||||
// actual → first alias (mirrors Python at line 807)
|
||||
if r2a["docnm"] != "docnm_kwd" {
|
||||
t.Errorf("actualToFirstAlias[docnm]: got %q, want docnm_kwd", r2a["docnm"])
|
||||
}
|
||||
if r2a["content"] != "content_with_weight" {
|
||||
t.Errorf("actualToFirstAlias[content]: got %q, want content_with_weight", r2a["content"])
|
||||
}
|
||||
// "plain" has no comment, so it shouldn't appear in the reverse map.
|
||||
if _, ok := r2a["plain"]; ok {
|
||||
t.Errorf("actualToFirstAlias should not include fields without comments")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoadFieldMapping_EmptyNameDefaultsToInfinityMappingJSON(t *testing.T) {
|
||||
// Empty name → defaults to "infinity_mapping.json" (line 145).
|
||||
// We just verify the function doesn't panic and the file-not-found
|
||||
// path is taken silently.
|
||||
a2a, r2a, err := loadFieldMapping("")
|
||||
if err != nil {
|
||||
t.Fatalf("empty name: %v", err)
|
||||
}
|
||||
if len(a2a) != 0 || len(r2a) != 0 {
|
||||
t.Errorf("empty name + no file should yield empty maps; got a2a=%v r2a=%v", a2a, r2a)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user