refactor(go): default Go cache/queue backend to Kvrocks and fail fast when down (#19804)

This commit is contained in:
xugangqiang
2026-09-18 19:16:07 +08:00
committed by GitHub
parent a1da66010b
commit 7aa012fc5c
51 changed files with 605 additions and 332 deletions
+8 -4
View File
@@ -603,7 +603,7 @@ jobs:
# Engine-specific offset partitions keep concurrent engine jobs from
# choosing the same host ports when they land on the same self-hosted runner.
# A lock plus reservation file closes the check/start race between parallel jobs.
PORT_BASES=(1200 1201 23817 23820 5432 5455 9000 9001 6379 6380 6601 9380 9381 9382 9384 9383 9385 80 443 4222 8222 9900 8123)
PORT_BASES=(1200 1201 23817 23820 5432 5455 9000 9001 6379 6380 6381 6601 9380 9381 9382 9384 9383 9385 80 443 4222 8222 9900 8123)
PARTITION_SIZE=6000
case "${DOC_ENGINE}" in
elasticsearch) PARTITION_BASE=1000 ;;
@@ -754,6 +754,7 @@ jobs:
EXPOSE_NATS_PORT=$((4222 + PORT_OFFSET))
NATS_MONITORING_PORT=$((8222 + PORT_OFFSET))
TEI_PORT=$((6380 + PORT_OFFSET))
KVROCKS_PORT=$((6381 + PORT_OFFSET))
KIBANA_PORT=$((6601 + PORT_OFFSET))
SVR_HTTP_PORT=$((9380 + PORT_OFFSET))
ADMIN_SVR_HTTP_PORT=$((9381 + PORT_OFFSET))
@@ -768,7 +769,7 @@ jobs:
# Persist computed ports into .env so docker-compose uses the correct host bindings.
# Remove previous CI overrides first; docker compose uses the last duplicate key.
sed -i '/^ES_PORT=/d;/^OS_PORT=/d;/^INFINITY_THRIFT_PORT=/d;/^INFINITY_HTTP_PORT=/d;/^INFINITY_PSQL_PORT=/d;/^EXPOSE_MYSQL_PORT=/d;/^MINIO_PORT=/d;/^MINIO_CONSOLE_PORT=/d;/^REDIS_PORT=/d;/^EXPOSE_NATS_PORT=/d;/^NATS_MONITORING_PORT=/d;/^TEI_PORT=/d;/^KIBANA_PORT=/d;/^SVR_HTTP_PORT=/d;/^ADMIN_SVR_HTTP_PORT=/d;/^SVR_MCP_PORT=/d;/^GO_HTTP_PORT=/d;/^GO_ADMIN_PORT=/d;/^SANDBOX_EXECUTOR_MANAGER_PORT=/d;/^SVR_WEB_HTTP_PORT=/d;/^SVR_WEB_HTTPS_PORT=/d;/^EXPOSE_CLICKHOUSE_TCP_PORT=/d;/^CLICKHOUSE_HTTP_PORT=/d;/^COMPOSE_PROFILES=/d;/^TEI_MODEL=/d;/^RAGFLOW_IMAGE=/d;/^DOC_ENGINE=/d;/^API_PROXY_SCHEME=/d' .env
sed -i '/^ES_PORT=/d;/^OS_PORT=/d;/^INFINITY_THRIFT_PORT=/d;/^INFINITY_HTTP_PORT=/d;/^INFINITY_PSQL_PORT=/d;/^EXPOSE_MYSQL_PORT=/d;/^MINIO_PORT=/d;/^MINIO_CONSOLE_PORT=/d;/^REDIS_PORT=/d;/^KVROCKS_PORT=/d;/^EXPOSE_NATS_PORT=/d;/^NATS_MONITORING_PORT=/d;/^TEI_PORT=/d;/^KIBANA_PORT=/d;/^SVR_HTTP_PORT=/d;/^ADMIN_SVR_HTTP_PORT=/d;/^SVR_MCP_PORT=/d;/^GO_HTTP_PORT=/d;/^GO_ADMIN_PORT=/d;/^SANDBOX_EXECUTOR_MANAGER_PORT=/d;/^SVR_WEB_HTTP_PORT=/d;/^SVR_WEB_HTTPS_PORT=/d;/^EXPOSE_CLICKHOUSE_TCP_PORT=/d;/^CLICKHOUSE_HTTP_PORT=/d;/^COMPOSE_PROFILES=/d;/^TEI_MODEL=/d;/^RAGFLOW_IMAGE=/d;/^DOC_ENGINE=/d;/^API_PROXY_SCHEME=/d' .env
# docker-compose-go.yml uses env_file: .env-go for the ragflow-cpu
# container, so the Go server process reads COMPOSE_PROFILES and
# TEI_MODEL from .env-go (not .env). Without syncing these two,
@@ -788,6 +789,7 @@ jobs:
echo "MINIO_PORT=${MINIO_PORT}"
echo "MINIO_CONSOLE_PORT=${MINIO_CONSOLE_PORT}"
echo "REDIS_PORT=${REDIS_PORT}"
echo "KVROCKS_PORT=${KVROCKS_PORT}"
echo "EXPOSE_NATS_PORT=${EXPOSE_NATS_PORT}"
echo "NATS_MONITORING_PORT=${NATS_MONITORING_PORT}"
echo "TEI_PORT=${TEI_PORT}"
@@ -1337,7 +1339,7 @@ jobs:
# Engine-specific offset partitions keep concurrent engine jobs from
# choosing the same host ports when they land on the same self-hosted runner.
# A lock plus reservation file closes the check/start race between parallel jobs.
PORT_BASES=(1200 1201 23817 23820 5432 5455 9000 9001 6379 6380 6601 9380 9381 9382 9384 9383 9385 80 443 4222 8222 9900 8123)
PORT_BASES=(1200 1201 23817 23820 5432 5455 9000 9001 6379 6380 6381 6601 9380 9381 9382 9384 9383 9385 80 443 4222 8222 9900 8123)
PARTITION_SIZE=6000
case "${DOC_ENGINE}" in
elasticsearch) PARTITION_BASE=1000 ;;
@@ -1488,6 +1490,7 @@ jobs:
EXPOSE_NATS_PORT=$((4222 + PORT_OFFSET))
NATS_MONITORING_PORT=$((8222 + PORT_OFFSET))
TEI_PORT=$((6380 + PORT_OFFSET))
KVROCKS_PORT=$((6381 + PORT_OFFSET))
KIBANA_PORT=$((6601 + PORT_OFFSET))
SVR_HTTP_PORT=$((9380 + PORT_OFFSET))
ADMIN_SVR_HTTP_PORT=$((9381 + PORT_OFFSET))
@@ -1502,7 +1505,7 @@ jobs:
# Persist computed ports into .env so docker-compose uses the correct host bindings.
# Remove previous CI overrides first; docker compose uses the last duplicate key.
sed -i '/^ES_PORT=/d;/^OS_PORT=/d;/^INFINITY_THRIFT_PORT=/d;/^INFINITY_HTTP_PORT=/d;/^INFINITY_PSQL_PORT=/d;/^EXPOSE_MYSQL_PORT=/d;/^MINIO_PORT=/d;/^MINIO_CONSOLE_PORT=/d;/^REDIS_PORT=/d;/^EXPOSE_NATS_PORT=/d;/^NATS_MONITORING_PORT=/d;/^TEI_PORT=/d;/^KIBANA_PORT=/d;/^SVR_HTTP_PORT=/d;/^ADMIN_SVR_HTTP_PORT=/d;/^SVR_MCP_PORT=/d;/^GO_HTTP_PORT=/d;/^GO_ADMIN_PORT=/d;/^SANDBOX_EXECUTOR_MANAGER_PORT=/d;/^SVR_WEB_HTTP_PORT=/d;/^SVR_WEB_HTTPS_PORT=/d;/^EXPOSE_CLICKHOUSE_TCP_PORT=/d;/^CLICKHOUSE_HTTP_PORT=/d;/^COMPOSE_PROFILES=/d;/^TEI_MODEL=/d;/^RAGFLOW_IMAGE=/d;/^DOC_ENGINE=/d;/^API_PROXY_SCHEME=/d' .env
sed -i '/^ES_PORT=/d;/^OS_PORT=/d;/^INFINITY_THRIFT_PORT=/d;/^INFINITY_HTTP_PORT=/d;/^INFINITY_PSQL_PORT=/d;/^EXPOSE_MYSQL_PORT=/d;/^MINIO_PORT=/d;/^MINIO_CONSOLE_PORT=/d;/^REDIS_PORT=/d;/^KVROCKS_PORT=/d;/^EXPOSE_NATS_PORT=/d;/^NATS_MONITORING_PORT=/d;/^TEI_PORT=/d;/^KIBANA_PORT=/d;/^SVR_HTTP_PORT=/d;/^ADMIN_SVR_HTTP_PORT=/d;/^SVR_MCP_PORT=/d;/^GO_HTTP_PORT=/d;/^GO_ADMIN_PORT=/d;/^SANDBOX_EXECUTOR_MANAGER_PORT=/d;/^SVR_WEB_HTTP_PORT=/d;/^SVR_WEB_HTTPS_PORT=/d;/^EXPOSE_CLICKHOUSE_TCP_PORT=/d;/^CLICKHOUSE_HTTP_PORT=/d;/^COMPOSE_PROFILES=/d;/^TEI_MODEL=/d;/^RAGFLOW_IMAGE=/d;/^DOC_ENGINE=/d;/^API_PROXY_SCHEME=/d' .env
# docker-compose-go.yml uses env_file: .env-go for the ragflow-cpu
# container, so the Go server process reads COMPOSE_PROFILES and
# TEI_MODEL from .env-go (not .env). Without syncing these two,
@@ -1522,6 +1525,7 @@ jobs:
echo "MINIO_PORT=${MINIO_PORT}"
echo "MINIO_CONSOLE_PORT=${MINIO_CONSOLE_PORT}"
echo "REDIS_PORT=${REDIS_PORT}"
echo "KVROCKS_PORT=${KVROCKS_PORT}"
echo "EXPOSE_NATS_PORT=${EXPOSE_NATS_PORT}"
echo "NATS_MONITORING_PORT=${NATS_MONITORING_PORT}"
echo "TEI_PORT=${TEI_PORT}"
+10 -10
View File
@@ -68,7 +68,7 @@ import (
"ragflow/internal/dao"
"ragflow/internal/deepdoc/parser/pdf/inference/native_analyzer"
"ragflow/internal/engine"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
et "ragflow/internal/engine/types"
"ragflow/internal/entity"
_ "ragflow/internal/ingestion/wire"
@@ -627,11 +627,11 @@ func main() {
}
defer engine.Close()
// Initialize Redis cache
if err = redis.Init(ctx); err != nil {
common.Fatal("Failed to initialize Redis", zap.Error(err))
// Initialize Kvrocks cache
if err = kvrocks.Init(ctx); err != nil {
common.Fatal("Failed to initialize Kvrocks", zap.Error(err))
}
defer redis.Close()
defer kvrocks.Close()
if err = storage.Init(ctx); err != nil {
common.Error("Failed to initialize storage factory", err)
@@ -644,7 +644,7 @@ func main() {
// Initialize server variables (runtime variables that can change during operation)
// This must be done after Cache is initialized
if err = server.InitVariables(redis.Get()); err != nil {
if err = server.InitVariables(kvrocks.Get()); err != nil {
common.Warn("Failed to initialize server variables from Redis, using defaults", zap.String("error", err.Error()))
}
@@ -1402,7 +1402,7 @@ func startServer(ctx context.Context, args *serverArgs) error {
)
skillSearchHandler := handler.NewSkillSearchHandler(docEngine, documentService)
providerHandler := handler.NewProviderHandler(userService, modelProviderService)
// Install the agent service's Redis-backed run infrastructure
// Install the agent service's Kvrocks-backed run infrastructure
// (CheckPointStore / StateSerializer / RunTracker). When Redis
// is unreachable (degraded boot, stand-alone mode, no-redis CI)
// the constructors return errors, and we fall through to the
@@ -1718,7 +1718,7 @@ type agentRunOptions struct {
runTracker *canvas.RunTracker
}
// buildAgentRunOptions installs the Redis-backed run infrastructure
// buildAgentRunOptions installs the Kvrocks-backed run infrastructure
// when Redis is available. The Redis client is the one already
// initialized at the top of main; the TTL is a conservative 24h for
// both the checkpoint store and the run tracker. On any error
@@ -1727,11 +1727,11 @@ type agentRunOptions struct {
// to the in-memory path transparently.
func buildAgentRunOptions() agentRunOptions {
var out agentRunOptions
if !redis.IsEnabled() || redis.Get() == nil {
if !kvrocks.IsEnabled() || kvrocks.Get() == nil {
common.Info("agent: redis client not initialised; agent run infra in in-memory mode (no checkpoints, no run tracker)")
return out
}
cp := canvas.NewRedisCheckPointStore(24 * time.Hour)
cp := canvas.NewKvrocksCheckPointStore(24 * time.Hour)
out.checkpointStore = cp
// stateSerializer is intentionally left nil. eino's default
// InternalSerializer (used when no compose.WithSerializer is
+9 -7
View File
@@ -14,14 +14,13 @@
# limitations under the License.
#
from docx import Document
import re
import pandas as pd
from collections import Counter
from rag.nlp import rag_tokenizer
from io import BytesIO
import logging
from common.constants import MAXIMUM_PAGE_NUMBER
import re
from collections import Counter
from io import BytesIO
import pandas as pd
from docx import Document
from docx.image.exceptions import (
InvalidImageStreamError,
UnexpectedEndOfFileError,
@@ -29,6 +28,9 @@ from docx.image.exceptions import (
)
from docx.oxml.ns import qn
from docx.text.run import Run
from common.constants import MAXIMUM_PAGE_NUMBER
from rag.nlp import rag_tokenizer
from rag.utils.lazy_image import LazyImage
# Markup Compatibility namespace. Word stores every text box twice inside an
+3
View File
@@ -197,6 +197,9 @@ REDIS_HOST=redis
REDIS_PORT=6379
# The password for Redis.
REDIS_PASSWORD=infini_rag_flow
# Host port that maps to the Kvrocks container port 6379. The Go deployment
# disables the Valkey/Redis service, so Kvrocks reuses the conventional 6379.
KVROCKS_PORT=6379
NATS_HOST=nats
# Port used by RAGFlow's Go services to connect to NATS inside the container network.
+8 -6
View File
@@ -195,12 +195,14 @@ MINIO_USER=rag_flow
# When updated, you must revise the `minio.password` entry in service_conf.yaml accordingly.
MINIO_PASSWORD=infini_rag_flow
# The hostname where the Redis service is exposed
REDIS_HOST=redis
# The port used to expose the Redis service to the host machine,
# allowing EXTERNAL access to the Redis service running inside the Docker container.
REDIS_PORT=6379
# The password for Redis.
# The hostname where the Kvrocks service is exposed. Kvrocks is the only
# cache/queue backend the Go services use (Redis-protocol, RocksDB-backed).
# For a host-run Go binary, point this at the published Kvrocks port on localhost.
KVROCKS_HOST=127.0.0.1
# The host port that maps to the Kvrocks container port 6379. The Go deployment
# disables the Valkey/Redis service, so Kvrocks reuses the conventional 6379.
KVROCKS_PORT=6379
# The password for Kvrocks (shared with the Valkey service; see kvrocks-entrypoint.sh).
REDIS_PASSWORD=infini_rag_flow
NATS_HOST=nats
+11 -3
View File
@@ -332,9 +332,17 @@ services:
entrypoint: ["/bin/sh", "/etc/kvrocks/entrypoint.sh"]
user: "0:0"
env_file: .env
# No host port publish: the Go services reach Kvrocks over the internal
# docker network at `kvrocks:6379`. Gated to the Go stack profile so the
# Python-only stack does not spin up an unused container.
# Publish the Kvrocks port so a host-run Go binary (e.g. ./bin/ragflow_server
# --ingestor) can reach it. The host-facing port uses KVROCKS_PORT (default 6379
# locally, so a local host binary connects to localhost:6379). In CI the port is
# offset per job (KVROCKS_PORT = 6381 + PORT_OFFSET, reserved alongside the other
# CI ports) so concurrent test runs on a shared runner do not clash, and it stays
# distinct from REDIS_PORT (used by the Valkey `redis` service that the Python
# deployment still runs next to Kvrocks). The container port (right side) stays
# 6379, and the Go services reach `kvrocks:6379` inside the docker network
# regardless of the host mapping (KVROCKS_PORT stays 6379, set by the Go overlay).
ports:
- "127.0.0.1:${KVROCKS_PORT:-6379}:6379"
profiles:
- ragflow-go
volumes:
+26 -6
View File
@@ -2,6 +2,16 @@ include:
- ./docker-compose-base.yml
# To ensure that the container processes the locally modified `service_conf.yaml.template` instead of the one included in its image, you need to mount the local `service_conf.yaml.template` to the container.
services:
# The Go stack uses Kvrocks exclusively (see the `kvrocks` service) and does
# not need Valkey. Disable the Valkey `redis` service inherited from the base
# compose by giving it a profile the Go deployment never activates, so it does
# not occupy the Redis port slot and Kvrocks can use a dedicated, CI-offset
# host-published port (see docker-compose-base.yml; the variable is KVROCKS_PORT,
# offset per CI job and distinct from REDIS_PORT).
redis:
profiles:
- ragflow-py
deepdoc:
image: ${DEEPDOC_IMAGE:-deepdoc_oss:latest}
profiles:
@@ -55,10 +65,15 @@ services:
- ./service_conf.yaml.template:/ragflow/conf/service_conf.yaml.template
- ./entrypoint-go.sh:/ragflow/entrypoint.sh
env_file: .env-go
# Route the Go services' cache/queue backend to the separate Kvrocks service
# (overriding REDIS_HOST=redis from .env-go). The Python services keep Valkey.
# Route the Go services' cache/queue backend to the separate Kvrocks service.
# The Python services use the Valkey `redis` service, which the Go deployment
# disables (see the `redis` profile override above). KVROCKS_HOST/KVROCKS_PORT
# are the IN-NETWORK coordinates the Go process uses to reach `kvrocks:6379`; the
# host-published port is KVROCKS_PORT (see docker-compose-base.yml), which CI
# offsets per job and keeps distinct from REDIS_PORT.
environment:
- REDIS_HOST=kvrocks
- KVROCKS_HOST=kvrocks
- KVROCKS_PORT=6379
networks:
- ragflow
restart: unless-stopped
@@ -100,10 +115,15 @@ services:
- ./service_conf.yaml.template:/ragflow/conf/service_conf.yaml.template
- ./entrypoint-go.sh:/ragflow/entrypoint.sh
env_file: .env-go
# Route the Go services' cache/queue backend to the separate Kvrocks service
# (overriding REDIS_HOST=redis from .env-go). The Python services keep Valkey.
# Route the Go services' cache/queue backend to the separate Kvrocks service.
# The Python services use the Valkey `redis` service, which the Go deployment
# disables (see the `redis` profile override above). KVROCKS_HOST/KVROCKS_PORT
# are the IN-NETWORK coordinates the Go process uses to reach `kvrocks:6379`; the
# host-published port is KVROCKS_PORT (see docker-compose-base.yml), which CI
# offsets per job and keeps distinct from REDIS_PORT.
environment:
- REDIS_HOST=kvrocks
- KVROCKS_HOST=kvrocks
- KVROCKS_PORT=6379
networks:
- ragflow
restart: unless-stopped
+9
View File
@@ -81,6 +81,15 @@ redis:
username: '${REDIS_USERNAME:-}'
password: '${REDIS_PASSWORD:-infini_rag_flow}'
host: '${REDIS_HOST:-redis}:6379'
# Kvrocks is the RocksDB-backed, Redis-protocol store used ONLY by the Go
# services. It has no in-memory maxmemory cap, so it does not OOM the way
# Valkey/Redis does. The Go services read this section; the Python services
# keep using the `redis:` section above.
kvrocks:
db: 1
username: '${KVROCKS_USERNAME:-}'
password: '${REDIS_PASSWORD:-infini_rag_flow}'
host: '${KVROCKS_HOST:-kvrocks}:${KVROCKS_PORT:-6379}'
nats:
host: ${NATS_HOST:-0.0.0.0}
port: ${NATS_PORT:-4222}
+3 -3
View File
@@ -24,7 +24,7 @@ import (
"net/http"
"ragflow/internal/common"
"ragflow/internal/engine"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"ragflow/internal/handler"
"ragflow/internal/server"
"ragflow/internal/service"
@@ -113,7 +113,7 @@ func (h *Handler) Login(c *gin.Context) {
return
}
secretKey, err := server.GetSecretKey(ctx, redis.Get())
secretKey, err := server.GetSecretKey(ctx, kvrocks.Get())
if err != nil {
common.ErrorWithCode(c, common.CodeServerError, fmt.Sprintf("Failed to get secret key: %s", err.Error()))
return
@@ -1250,7 +1250,7 @@ func (h *Handler) PingStore(c *gin.Context) {
func (h *Handler) PingCache(c *gin.Context) {
ctx := c.Request.Context()
redisClient := redis.Get()
redisClient := kvrocks.Get()
if redisClient.Health(ctx) {
common.SuccessNoMessage(c, "SUCCESS")
} else {
+6 -3
View File
@@ -32,7 +32,7 @@ import (
"ragflow/internal/engine"
"ragflow/internal/engine/clickhouse"
"ragflow/internal/engine/elasticsearch"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"ragflow/internal/entity"
modelModule "ragflow/internal/entity/models"
"ragflow/internal/server"
@@ -1080,7 +1080,10 @@ func (s *Service) ListServices(ctx context.Context) ([]ServiceStatus, error) {
// cache engine
cacheType := globalConfig.CacheEngineType()
switch cacheType {
case "redis":
// The Go stack talks to Kvrocks. "redis" is accepted for backwards
// compatibility with the shared service_conf.yaml.template; both map to
// the same Kvrocks backend, so probe the same connection.
case "redis", "kvrocks":
mysqlStatus := s.getRedisInfo(ctx)
results = append(results, mysqlStatus)
default:
@@ -1190,7 +1193,7 @@ func (s *Service) getRedisInfo(ctx context.Context) ServiceStatus {
startTime := time.Now()
redisConfig := server.GetConfig().GetRedisConfig()
redisClient := redis.Get()
redisClient := kvrocks.Get()
if redisClient.Health(ctx) {
return newServiceStatus(serviceType, name, redisConfig.Host, redisConfig.Port, "alive", startTime, "")
}
@@ -45,7 +45,7 @@ import (
"time"
"ragflow/internal/agent/runtime"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
)
// ModelProviderFunc is the contract the audio package uses to
@@ -80,14 +80,14 @@ type ModelProviderRequest struct {
func SetModelProviderSynthesizer(fn ModelProviderFunc) {
var s Synthesizer = stubSynthesizer{}
if fn != nil {
s = &modelProviderSynthesizer{fn: fn, redis: redis.Get()}
s = &modelProviderSynthesizer{fn: fn, kvrocks: kvrocks.Get()}
}
SetSynthesizer(s)
}
type modelProviderSynthesizer struct {
fn ModelProviderFunc
redis *redis.Client
fn ModelProviderFunc
kvrocks *kvrocks.Client
}
func (m *modelProviderSynthesizer) Synthesize(ctx context.Context, req SynthesizeRequest) (*SynthesizeResponse, error) {
@@ -107,8 +107,8 @@ func (m *modelProviderSynthesizer) Synthesize(ctx context.Context, req Synthesiz
// rag/utils/tts_cache.py). Cache failures are non-fatal —
// log and fall through to the model provider.
cacheKey := buildTTSCacheKey(tenantID, req)
if cacheKey != "" && m.redis != nil {
if cached, _ := m.redis.Get(ctx, cacheKey); cached != "" {
if cacheKey != "" && m.kvrocks != nil {
if cached, _ := m.kvrocks.Get(ctx, cacheKey); cached != "" {
if b, err := hex.DecodeString(cached); err == nil && len(b) > 0 {
return &SynthesizeResponse{Audio: b, MediaType: "audio/mpeg"}, nil
}
@@ -136,10 +136,10 @@ func (m *modelProviderSynthesizer) Synthesize(ctx context.Context, req Synthesiz
// Store in cache. TTL defaults to 7 days; env override via
// RAGFLOW_TTS_CACHE_TTL_SECONDS (positive integer seconds;
// 0 or invalid → default).
if cacheKey != "" && m.redis != nil {
if cacheKey != "" && m.kvrocks != nil {
ttl := ttsCacheTTL()
if ttl > 0 {
m.redis.Set(ctx, cacheKey, hex.EncodeToString(resp.Audio), ttl)
m.kvrocks.Set(ctx, cacheKey, hex.EncodeToString(resp.Audio), ttl)
}
}
return resp, nil
+3 -3
View File
@@ -25,7 +25,7 @@ import (
"go.uber.org/zap"
"ragflow/internal/common"
redis2 "ragflow/internal/engine/redis"
kvrocks "ragflow/internal/engine/kvrocks"
)
// cancelKeySuffix is appended to the session id to form the Redis key.
@@ -47,13 +47,13 @@ const RequestCancelTTL = 24 * time.Hour
// a package-level variable so tests can override it with a miniredis
// client (the production path goes through cache.Get()).
var cancelClientFn = func() (*redis.Client, error) {
rc := redis2.Get()
rc := kvrocks.Get()
if rc == nil {
return nil, errors.New("cancel: redis cache not initialized")
}
c := rc.GetClient()
if c == nil {
return nil, errors.New("cancel: redis client not initialized")
return nil, errors.New("cancel: kvrocks client not initialized")
}
return c, nil
}
+21 -21
View File
@@ -15,7 +15,7 @@
//
// checkpoint_store.go implements the eino CheckPointStore / CheckPointDeleter
// interfaces backed by Redis. See plan §2.6 (Redis-backed CheckPointStore).
// interfaces backed by Redis. See plan §2.6 (Kvrocks-backed CheckPointStore).
//
// The store holds raw eino-serialized checkpoint bytes keyed by
// "agent:cp:{id}". Business metadata (canvas_id, run_id, status, ...) lives
@@ -25,7 +25,7 @@ package canvas
import (
"context"
"errors"
redis2 "ragflow/internal/engine/redis"
kvrocks "ragflow/internal/engine/kvrocks"
"time"
"github.com/redis/go-redis/v9"
@@ -35,44 +35,44 @@ import (
// The full key is "agent:cp:{id}".
const checkpointKeyPrefix = "agent:cp:"
// RedisCheckPointStore is a Redis-backed eino CheckPointStore /
// KvrocksCheckPointStore is a Kvrocks-backed eino CheckPointStore /
// CheckPointDeleter. Values are stored as raw bytes — the eino Serializer
// has already marshaled the structured payload, so we do not re-encode.
type RedisCheckPointStore struct {
type KvrocksCheckPointStore struct {
client *redis.Client
ttl time.Duration
}
// NewRedisCheckPointStore returns a store wired to the global Redis client
// NewKvrocksCheckPointStore returns a store wired to the global Kvrocks client
// from internal/cache. Returns a non-nil store even when the cache is
// uninitialized (client is nil); Get/Set/Delete will return an error in that
// case rather than nil-deref, but the type stays usable for tests that
// inject their own client via struct-literal construction.
func NewRedisCheckPointStore(ttl time.Duration) *RedisCheckPointStore {
func NewKvrocksCheckPointStore(ttl time.Duration) *KvrocksCheckPointStore {
var client *redis.Client
if rc := redis2.Get(); rc != nil {
if rc := kvrocks.Get(); rc != nil {
client = rc.GetClient()
}
return &RedisCheckPointStore{client: client, ttl: ttl}
return &KvrocksCheckPointStore{client: client, ttl: ttl}
}
// NewRedisCheckPointStoreWithClient returns a store wired to a
// NewKvrocksCheckPointStoreWithClient returns a store wired to a
// caller-supplied redis.Client. Same rationale as
// NewRunTrackerWithClient: enables test code (or any code that
// needs a dedicated Redis pool) to inject a client without going
// through the global cache singleton.
func NewRedisCheckPointStoreWithClient(client *redis.Client, ttl time.Duration) *RedisCheckPointStore {
return &RedisCheckPointStore{client: client, ttl: ttl}
func NewKvrocksCheckPointStoreWithClient(client *redis.Client, ttl time.Duration) *KvrocksCheckPointStore {
return &KvrocksCheckPointStore{client: client, ttl: ttl}
}
// RedisCheckpointExists reports whether a pipeline checkpoint is present for
// KvrocksCheckpointExists reports whether a pipeline checkpoint is present for
// id. It is used by ingestion progress handling to distinguish a fresh run
// from a resume: resumed nodes may not emit lifecycle events again, so their
// previous completed progress rows must be retained.
func RedisCheckpointExists(ctx context.Context, id string) (bool, error) {
rc := redis2.Get()
func KvrocksCheckpointExists(ctx context.Context, id string) (bool, error) {
rc := kvrocks.Get()
if rc == nil || rc.GetClient() == nil {
return false, errors.New("checkpoint store: redis client not initialized")
return false, errors.New("checkpoint store: kvrocks client not initialized")
}
found, err := rc.GetClient().Exists(ctx, checkpointKeyPrefix+id).Result()
return found > 0, err
@@ -81,9 +81,9 @@ func RedisCheckpointExists(ctx context.Context, id string) (bool, error) {
// Get implements eino's CheckPointStore.Get. Returns (nil, false, nil) when
// the key does not exist (redis.Nil) so callers can distinguish "missing"
// from "present-but-error".
func (s *RedisCheckPointStore) Get(ctx context.Context, id string) ([]byte, bool, error) {
func (s *KvrocksCheckPointStore) Get(ctx context.Context, id string) ([]byte, bool, error) {
if s == nil || s.client == nil {
return nil, false, errors.New("checkpoint store: redis client not initialized")
return nil, false, errors.New("checkpoint store: kvrocks client not initialized")
}
data, err := s.client.Get(ctx, checkpointKeyPrefix+id).Bytes()
if errors.Is(err, redis.Nil) {
@@ -97,18 +97,18 @@ func (s *RedisCheckPointStore) Get(ctx context.Context, id string) ([]byte, bool
// Set implements eino's CheckPointStore.Set. The TTL is applied on every
// call so a frequently-updated checkpoint does not expire mid-run.
func (s *RedisCheckPointStore) Set(ctx context.Context, id string, payload []byte) error {
func (s *KvrocksCheckPointStore) Set(ctx context.Context, id string, payload []byte) error {
if s == nil || s.client == nil {
return errors.New("checkpoint store: redis client not initialized")
return errors.New("checkpoint store: kvrocks client not initialized")
}
return s.client.Set(ctx, checkpointKeyPrefix+id, payload, s.ttl).Err()
}
// Delete implements eino's optional CheckPointDeleter. It is safe to call
// on a non-existent key (Del returns 0, no error).
func (s *RedisCheckPointStore) Delete(ctx context.Context, id string) error {
func (s *KvrocksCheckPointStore) Delete(ctx context.Context, id string) error {
if s == nil || s.client == nil {
return errors.New("checkpoint store: redis client not initialized")
return errors.New("checkpoint store: kvrocks client not initialized")
}
return s.client.Del(ctx, checkpointKeyPrefix+id).Err()
}
@@ -27,9 +27,9 @@ import (
// newTestStore spins up a miniredis-backed store for table-driven tests.
// Returns the store, the miniredis handle (caller must Close()), and a
// cleanup function. We construct the struct directly so we can inject the
// *redis.Client — NewRedisCheckPointStore reads from the global cache
// *redis.Client — NewKvrocksCheckPointStore reads from the global cache
// which is nil in unit tests.
func newTestStore(t *testing.T, ttl time.Duration) (*RedisCheckPointStore, *miniredis.Miniredis) {
func newTestStore(t *testing.T, ttl time.Duration) (*KvrocksCheckPointStore, *miniredis.Miniredis) {
t.Helper()
mr, err := miniredis.Run()
if err != nil {
@@ -40,7 +40,7 @@ func newTestStore(t *testing.T, ttl time.Duration) (*RedisCheckPointStore, *mini
client := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = client.Close() })
return &RedisCheckPointStore{client: client, ttl: ttl}, mr
return &KvrocksCheckPointStore{client: client, ttl: ttl}, mr
}
func TestRedisCheckPointStore_RoundTrip(t *testing.T) {
@@ -123,9 +123,9 @@ func TestRedisCheckPointStore_Delete(t *testing.T) {
}
func TestRedisCheckPointStore_NilClient(t *testing.T) {
// Cache uninitialized → NewRedisCheckPointStore returns a store with
// Cache uninitialized → NewKvrocksCheckPointStore returns a store with
// nil client. Operations must error rather than panic.
store := &RedisCheckPointStore{client: nil, ttl: time.Minute}
store := &KvrocksCheckPointStore{client: nil, ttl: time.Minute}
ctx := t.Context()
if _, _, err := store.Get(ctx, "x"); err == nil {
+2 -2
View File
@@ -20,7 +20,7 @@ import (
)
// CheckPointStore is the minimal interface Compile needs at compile time.
// RedisCheckPointStore satisfies this; tests can pass any in-memory
// KvrocksCheckPointStore satisfies this; tests can pass any in-memory
// implementation. Matches eino's compose.CheckPointStore (an alias for
// core.CheckPointStore) and adds a Delete method.
type CheckPointStore interface {
@@ -314,7 +314,7 @@ func dedupeStrings(in []string) []string {
}
// checkPointAdapter drops the Delete method that compose.CheckPointStore
// does not declare. The RedisCheckPointStore in this package has deleted;
// does not declare. The KvrocksCheckPointStore in this package has deleted;
// the adapter is a thin passthrough.
type checkPointAdapter struct{ inner CheckPointStore }
+19 -19
View File
@@ -31,7 +31,7 @@ import (
"github.com/redis/go-redis/v9"
redis2 "ragflow/internal/engine/redis"
kvrocks "ragflow/internal/engine/kvrocks"
)
// runKeyPrefix is the Redis Hash key namespace for run metadata.
@@ -78,13 +78,13 @@ type RunTracker struct {
ttl time.Duration
}
// NewRunTracker returns a tracker wired to the global Redis client. When
// NewRunTracker returns a tracker wired to the global Kvrocks client. When
// the cache is uninitialized, client is nil; methods error in that case
// rather than panicking, and tests can inject a client via struct-literal
// construction.
func NewRunTracker(ttl time.Duration) *RunTracker {
var client *redis.Client
if rc := redis2.Get(); rc != nil {
if rc := kvrocks.Get(); rc != nil {
client = rc.GetClient()
}
return &RunTracker{client: client, ttl: ttl}
@@ -93,7 +93,7 @@ func NewRunTracker(ttl time.Duration) *RunTracker {
// NewRunTrackerWithClient returns a tracker wired to a caller-supplied
// redis.Client. The intended use is tests that want to drive the
// RunTracker against an in-memory miniredis without touching the
// global Redis cache, but the helper is exported so non-test callers
// global Kvrocks cache, but the helper is exported so non-test callers
// (multi-tenant deployments, custom Redis pools) can inject a
// dedicated client without going through the global cache singleton.
func NewRunTrackerWithClient(client *redis.Client, ttl time.Duration) *RunTracker {
@@ -118,7 +118,7 @@ func (t *RunTracker) setRunFields(ctx context.Context, runID string, values ...i
// run would leak the hash.
func (t *RunTracker) Start(ctx context.Context, runID, canvasID, tenantID, parentRunID string) error {
if t == nil || t.client == nil {
return errors.New("run tracker: redis client not initialized")
return errors.New("run tracker: kvrocks client not initialized")
}
now := time.Now().UnixMilli()
key := runKey(runID)
@@ -141,7 +141,7 @@ func (t *RunTracker) Start(ctx context.Context, runID, canvasID, tenantID, paren
// §2.6) must call this once before the run goroutine returns.
func (t *RunTracker) AttachCheckpoint(ctx context.Context, runID, checkpointID string) error {
if t == nil || t.client == nil {
return errors.New("run tracker: redis client not initialized")
return errors.New("run tracker: kvrocks client not initialized")
}
return t.setRunFields(ctx, runID, "checkpoint_id", checkpointID)
}
@@ -154,7 +154,7 @@ func (t *RunTracker) AttachCheckpoint(ctx context.Context, runID, checkpointID s
// calls compose.ResumeWithData with it. The ONLY writer of "interrupt_id".
func (t *RunTracker) AttachInterrupt(ctx context.Context, runID, interruptID string) error {
if t == nil || t.client == nil {
return errors.New("run tracker: redis client not initialized")
return errors.New("run tracker: kvrocks client not initialized")
}
return t.setRunFields(ctx, runID, runFieldInterruptID, interruptID)
}
@@ -165,7 +165,7 @@ func (t *RunTracker) AttachInterrupt(ctx context.Context, runID, interruptID str
// "fresh run, do a normal Invoke".
func (t *RunTracker) GetInterruptID(ctx context.Context, runID string) (string, bool, error) {
if t == nil || t.client == nil {
return "", false, errors.New("run tracker: redis client not initialized")
return "", false, errors.New("run tracker: kvrocks client not initialized")
}
v, err := t.client.HGet(ctx, runKey(runID), runFieldInterruptID).Result()
if errors.Is(err, redis.Nil) {
@@ -185,7 +185,7 @@ func (t *RunTracker) GetInterruptID(ctx context.Context, runID string) (string,
// UserFillUp checkpoint.
func (t *RunTracker) ClearInterruptID(ctx context.Context, runID string) error {
if t == nil || t.client == nil {
return errors.New("run tracker: redis client not initialized")
return errors.New("run tracker: kvrocks client not initialized")
}
return t.client.HDel(ctx, runKey(runID), runFieldInterruptID).Err()
}
@@ -193,7 +193,7 @@ func (t *RunTracker) ClearInterruptID(ctx context.Context, runID string) error {
// MarkSucceeded transitions the run to status=1 and stamps finished_at.
func (t *RunTracker) MarkSucceeded(ctx context.Context, runID string) error {
if t == nil || t.client == nil {
return errors.New("run tracker: redis client not initialized")
return errors.New("run tracker: kvrocks client not initialized")
}
return t.setRunFields(ctx, runID,
"status", runStatusSucceeded,
@@ -204,7 +204,7 @@ func (t *RunTracker) MarkSucceeded(ctx context.Context, runID string) error {
// MarkFailed transitions the run to status=2 and records the reason.
func (t *RunTracker) MarkFailed(ctx context.Context, runID, reason string) error {
if t == nil || t.client == nil {
return errors.New("run tracker: redis client not initialized")
return errors.New("run tracker: kvrocks client not initialized")
}
return t.setRunFields(ctx, runID,
"status", runStatusFailed,
@@ -216,7 +216,7 @@ func (t *RunTracker) MarkFailed(ctx context.Context, runID, reason string) error
// MarkCancelled transitions the run to status=3 and sets the cancel flag.
func (t *RunTracker) MarkCancelled(ctx context.Context, runID string) error {
if t == nil || t.client == nil {
return errors.New("run tracker: redis client not initialized")
return errors.New("run tracker: kvrocks client not initialized")
}
key := runKey(runID)
pipe := t.client.Pipeline()
@@ -236,7 +236,7 @@ func (t *RunTracker) MarkCancelled(ctx context.Context, runID string) error {
// means another instance already owns the session.
func (t *RunTracker) RegisterActiveSession(ctx context.Context, active ActiveSession) (bool, error) {
if t == nil || t.client == nil {
return false, errors.New("run tracker: redis client not initialized")
return false, errors.New("run tracker: kvrocks client not initialized")
}
const script = `
if redis.call("EXISTS", KEYS[1]) == 1 then return 0 end
@@ -258,7 +258,7 @@ return 1`
// result means the session is not running.
func (t *RunTracker) GetActiveSession(ctx context.Context, sessionID string) (*ActiveSession, error) {
if t == nil || t.client == nil {
return nil, errors.New("run tracker: redis client not initialized")
return nil, errors.New("run tracker: kvrocks client not initialized")
}
fields, err := t.client.HGetAll(ctx, activeSessionKey(sessionID)).Result()
if err != nil || len(fields) == 0 {
@@ -276,7 +276,7 @@ func (t *RunTracker) GetActiveSession(ctx context.Context, sessionID string) (*A
// RefreshActiveSession extends the lease only while token still owns it.
func (t *RunTracker) RefreshActiveSession(ctx context.Context, sessionID, token string) (bool, error) {
if t == nil || t.client == nil {
return false, errors.New("run tracker: redis client not initialized")
return false, errors.New("run tracker: kvrocks client not initialized")
}
const script = `if redis.call("HGET", KEYS[1], "token") == ARGV[1] then return redis.call("PEXPIRE", KEYS[1], ARGV[2]) end return 0`
result, err := t.client.Eval(ctx, script, []string{activeSessionKey(sessionID)}, token,
@@ -289,7 +289,7 @@ func (t *RunTracker) RefreshActiveSession(ctx context.Context, sessionID, token
// run.
func (t *RunTracker) ReleaseActiveSession(ctx context.Context, sessionID, token string) (bool, error) {
if t == nil || t.client == nil {
return false, errors.New("run tracker: redis client not initialized")
return false, errors.New("run tracker: kvrocks client not initialized")
}
const script = `
if redis.call("HGET", KEYS[1], "token") ~= ARGV[1] then return 0 end
@@ -306,7 +306,7 @@ return 1`
// a newer run after its lease has been replaced or released.
func (t *RunTracker) RequestCancelActiveSession(ctx context.Context, sessionID, token string) (bool, error) {
if t == nil || t.client == nil {
return false, errors.New("run tracker: redis client not initialized")
return false, errors.New("run tracker: kvrocks client not initialized")
}
const script = `
if redis.call("HGET", KEYS[1], "token") ~= ARGV[1] then return 0 end
@@ -365,7 +365,7 @@ func (t *RunTracker) leaseTTL() time.Duration {
// MarkWaiting records a normal UserFillUp pause. It is not a failure.
func (t *RunTracker) MarkWaiting(ctx context.Context, runID string) error {
if t == nil || t.client == nil {
return errors.New("run tracker: redis client not initialized")
return errors.New("run tracker: kvrocks client not initialized")
}
key := runKey(runID)
pipe := t.client.Pipeline()
@@ -380,7 +380,7 @@ func (t *RunTracker) MarkWaiting(ctx context.Context, runID string) error {
// if they need to distinguish from a key that exists with no fields.
func (t *RunTracker) Get(ctx context.Context, runID string) (map[string]string, error) {
if t == nil || t.client == nil {
return nil, errors.New("run tracker: redis client not initialized")
return nil, errors.New("run tracker: kvrocks client not initialized")
}
return t.client.HGetAll(ctx, runKey(runID)).Result()
}
+1 -1
View File
@@ -147,7 +147,7 @@ type canvasStateJSON struct {
// MarshalJSON serialises the CanvasState for eino's StatePre/Post
// handler chain (which JSON-encodes the state on every node boundary
// when a StateSerializer is wired) and for Redis-backed CheckPointStore
// when a StateSerializer is wired) and for Kvrocks-backed CheckPointStore
// payloads.
//
// Eino's interrupt path hit "failed to marshal state: unknown
+13
View File
@@ -201,6 +201,19 @@ missing from the model directory.
docker compose -f docker/docker-compose-base.yml --profile ragflow-go --profile infinity up -d
```
- Point the host-run Go binaries at Kvrocks. Kvrocks is published on
`127.0.0.1:6379` (loopback, reusing the conventional Redis port; the Go
deployment disables the Valkey/Redis service so there is no clash). Export these
before running `./bin/ragflow_server ...` on the host, or load them from
`docker/.env-go`:
```bash
export KVROCKS_HOST=127.0.0.1
export KVROCKS_PORT=6379
```
In docker the Go services override these to `kvrocks:6379` automatically
(see `docker-compose-go.yml`). If Kvrocks is unreachable at startup the
process fails fast with `failed to connect to Kvrocks`.
- Start RAGFlow
Note: Database migrations must complete before starting any server mode.
@@ -0,0 +1,134 @@
//
// 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.
//
//go:build integration
// Package kvrocks integration tests exercise the real connection path used by
// the server (cmd/ragflow_server.go calls kvrocks.Init). By default they run
// against an in-process miniredis (Redis-protocol) instance so they need no
// external service. Set KVROCKS_TEST_ADDR (host:port) to point them at a real
// Kvrocks deployment and prove the Go stack actually talks to Kvrocks.
package kvrocks
import (
"context"
"fmt"
"os"
"path/filepath"
"strings"
"sync"
"testing"
"github.com/alicebob/miniredis/v2"
"ragflow/internal/server"
)
// loadConfigForTest writes a minimal service_conf.yaml whose `kvrocks` section
// points at kvrocksAddr (host:port) and loads it via server.Init. server.Init
// only parses config (no connections), so a minimal file is enough to drive
// kvrocks.Init's connection path.
func loadConfigForTest(t *testing.T, kvrocksAddr string) {
t.Helper()
dir := t.TempDir()
cfgPath := filepath.Join(dir, "service_conf.yaml")
content := fmt.Sprintf(`general:
cache_engine: kvrocks
kvrocks:
db: 1
password: 'infini_rag_flow'
host: '%s'
`, kvrocksAddr)
if err := os.WriteFile(cfgPath, []byte(content), 0o644); err != nil {
t.Fatalf("write config: %v", err)
}
if err := server.Init(cfgPath); err != nil {
t.Fatalf("server.Init: %v", err)
}
}
// resetKvrocksState clears the package-level init once and client so each test
// can re-run kvrocks.Init against a freshly parsed config. kvrocks.Init is
// guarded by a sync.Once, while server.Init re-parses the global config on every
// call, so resetting here lets both tests share one process safely.
func resetKvrocksState() {
once = sync.Once{}
globalClient = nil
}
// TestKvrocksInitConnects proves the Go stack initializes a client against a
// reachable Kvrocks and can round-trip a key. This is the exact Init call the
// server makes at startup.
func TestKvrocksInitConnects(t *testing.T) {
resetKvrocksState()
addr := os.Getenv("KVROCKS_TEST_ADDR")
var mr *miniredis.Miniredis
if addr == "" {
var err error
mr, err = miniredis.Run()
if err != nil {
t.Fatalf("miniredis: %v", err)
}
t.Cleanup(mr.Close)
addr = mr.Addr()
}
loadConfigForTest(t, addr)
if err := Init(context.Background()); err != nil {
t.Fatalf("kvrocks.Init to %s should succeed, got: %v", addr, err)
}
defer Close()
if Get() == nil {
t.Fatal("kvrocks client should be initialized")
}
if !IsEnabled() {
t.Fatal("kvrocks should report enabled")
}
// Round-trip a key to prove the connection is real, not just a Ping.
ctx := context.Background()
const probeKey = "kvrocks_integration_probe"
if !Get().Set(ctx, probeKey, "ok", 0) {
t.Fatal("Set failed against backend")
}
val, err := Get().Get(ctx, probeKey)
if err != nil {
t.Fatalf("Get failed: %v", err)
}
if val != "ok" {
t.Fatalf("round-trip mismatch: got %q", val)
}
Get().Delete(ctx, probeKey)
}
// TestKvrocksInitFailFast proves that when Kvrocks is unreachable the process
// fails fast with a clear error instead of silently falling back to another
// backend (the Valkey OOM regression this migration fixes).
func TestKvrocksInitFailFast(t *testing.T) {
resetKvrocksState()
// 127.0.0.1:1 has nothing listening, so Ping fails immediately.
const deadAddr = "127.0.0.1:1"
loadConfigForTest(t, deadAddr)
err := Init(context.Background())
if err == nil {
t.Fatal("kvrocks.Init should fail when Kvrocks is unreachable")
}
if !strings.Contains(err.Error(), "failed to connect to Kvrocks") {
t.Fatalf("expected fail-fast error, got: %v", err)
}
}
@@ -18,7 +18,7 @@
// Valkey->Kvrocks migration. These are plain unit tests (no external service):
// they assert that the deployment artifacts (CI workflow, compose file) encode
// the required fixes so a future edit cannot silently regress the migration.
package redis
package kvrocks
import (
"os"
@@ -14,7 +14,7 @@
// limitations under the License.
//
package redis
package kvrocks
import (
"context"
@@ -47,7 +47,7 @@ type Client struct {
luaDeleteIfEqual *redis.Script
luaTokenBucket *redis.Script
luaAutoIncrement *redis.Script
config config.RedisConfig
config config.KvrocksConfig
}
// Message represents a message from Redis Stream
@@ -106,22 +106,23 @@ const (
`
)
// Init InitRedis initializes Redis client
// Init initializes the Kvrocks client. Kvrocks is the only supported cache/queue
// backend for the Go services; if it is unreachable the process fails fast.
func Init(ctx context.Context) error {
var initErr error
once.Do(func() {
globalConfig := server.GetConfig()
redisConfig := globalConfig.GetRedisConfig()
kvrocksConfig := globalConfig.GetKvrocksConfig()
if redisConfig.Host == "" {
common.Info("Redis host not configured, skipping Redis initialization")
if kvrocksConfig.Host == "" {
initErr = fmt.Errorf("kvrocks host not configured")
return
}
client := redis.NewClient(&redis.Options{
Addr: fmt.Sprintf("%s:%d", redisConfig.Host, redisConfig.Port),
Password: redisConfig.Password,
DB: redisConfig.DB,
Addr: fmt.Sprintf("%s:%d", kvrocksConfig.Host, kvrocksConfig.Port),
Password: kvrocksConfig.Password,
DB: kvrocksConfig.DB,
})
// Test connection
@@ -129,21 +130,21 @@ func Init(ctx context.Context) error {
defer cancel()
if err := client.Ping(redisCtx).Err(); err != nil {
initErr = fmt.Errorf("failed to connect to Redis: %w", err)
initErr = fmt.Errorf("failed to connect to Kvrocks: %w", err)
return
}
globalClient = &Client{
client: client,
config: redisConfig,
config: kvrocksConfig,
luaDeleteIfEqual: redis.NewScript(luaDeleteIfEqualScript),
luaTokenBucket: redis.NewScript(luaTokenBucketScript),
}
common.Info("Redis client initialized",
zap.String("host", redisConfig.Host),
zap.Int("port", redisConfig.Port),
zap.Int("db", redisConfig.DB),
common.Info("Kvrocks client initialized",
zap.String("host", kvrocksConfig.Host),
zap.Int("port", kvrocksConfig.Port),
zap.Int("db", kvrocksConfig.DB),
)
})
return initErr
@@ -14,7 +14,7 @@
// limitations under the License.
//
package redis
package kvrocks
import (
"context"
+2 -2
View File
@@ -28,7 +28,7 @@ import (
"sync"
"time"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
)
const (
@@ -324,7 +324,7 @@ func (e *Engine) withDDLLock(ctx context.Context, lockName string, check func()
}
}
timeout := time.Duration(timeoutSeconds) * time.Second
distributed := redis.NewDistributedLock(lockName, "", timeout, timeout)
distributed := kvrocks.NewDistributedLock(lockName, "", timeout, timeout)
if distributed != nil && !distributed.Acquire(ctx) {
deadline := time.NewTimer(timeout)
defer deadline.Stop()
+6 -6
View File
@@ -23,7 +23,7 @@ import (
"fmt"
"io"
"mime/multipart"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"strconv"
"strings"
"time"
@@ -78,12 +78,12 @@ type AgentHandler struct {
fileService agentFileService
loader canvasLoader
// redisGet fetches a raw string from Redis. Defaults to the global
// client (redis.Get) so production behaviour is unchanged; tests inject
// client (kvrocks.Get) so production behaviour is unchanged; tests inject
// a miniredis-backed getter to exercise Agent log endpoints without a
// live Redis (mirrors the newExecutor injection pattern).
redisGet func(key string) (string, error)
// redisStore writes the debug-run log array. Defaults to the global
// client (redis.Get, which satisfies task.DebugLogStore); tests inject a
// client (kvrocks.Get, which satisfies task.DebugLogStore); tests inject a
// miniredis-backed writer so runCanvasPipelineDebug can be exercised without a
// live Redis.
redisStore task.DebugLogStore
@@ -112,8 +112,8 @@ func NewAgentHandler(ctx context.Context, agentService *service.AgentService, fi
chatRunner: agentService,
fileService: fileService,
loader: agentService,
redisGet: func(key string) (string, error) { return redis.Get().Get(ctx, key) },
redisStore: redis.Get(),
redisGet: func(key string) (string, error) { return kvrocks.Get().Get(ctx, key) },
redisStore: kvrocks.Get(),
newExecutor: func(taskCtx *task.TaskContext, canvasID string, docBulkSize int) (debugExecutor, error) {
return task.NewPipelineExecutor(taskCtx, canvasID, docBulkSize)
},
@@ -1635,7 +1635,7 @@ func (h *AgentHandler) checkCanvasAccessForHandler(c *gin.Context, userID, canva
// api/apps/restful_apis/agent_api.py:992 — but unlike the Python
// implementation this handler does not sync a Canvas replica.
// `api.apps.services.canvas_replica_service.CanvasReplicaService` is
// the Python Redis-backed runtime replica (distributed lock + 3h TTL);
// the Python Kvrocks-backed runtime replica (distributed lock + 3h TTL);
// it is intentionally NOT ported to Go. The Go agent port runs every
// agent through eino's compose.Workflow.Invoke, which is reconstructed
// from the DSL on each run, so the replica's read-side acceleration
+2 -2
View File
@@ -68,7 +68,7 @@ import (
"net/http"
"ragflow/internal/agent/canvas"
"ragflow/internal/common"
rediscli "ragflow/internal/engine/redis"
kvrocks "ragflow/internal/engine/kvrocks"
"ragflow/internal/service"
"ragflow/internal/utility"
"strconv"
@@ -820,7 +820,7 @@ const (
// appendWebhookTrace appends a single RunEvent to the per-canvas trace
// key in Redis. Each event is recorded as {"ts": <float>, "event": <type>, ...}.
func appendWebhookTrace(ctx context.Context, agentID string, startTs time.Time, ev canvas.RunEvent) {
rdb := rediscli.Get()
rdb := kvrocks.Get()
if rdb == nil || rdb.GetClient() == nil {
return
}
+4 -4
View File
@@ -24,8 +24,8 @@ package handler
//
// - max body size (with the 10 MB cap at agent_api.py:1652)
// - IP whitelist (CIDR + exact match)
// - rate limit (token bucket via redis.EvalTokenBucketStrict — strict
// fail-closed; see redis.go)
// - rate limit (token bucket via kvrocks.EvalTokenBucketStrict — strict
// fail-closed; see kvrocks.go)
// - token auth (header check)
// - basic auth (HTTP Basic)
// - JWT (HS/RS256, audience/issuer/required-claims)
@@ -48,7 +48,7 @@ import (
"github.com/golang-jwt/jwt/v5"
"go.uber.org/zap"
rediscli "ragflow/internal/engine/redis"
kvrocks "ragflow/internal/engine/kvrocks"
)
const (
@@ -291,7 +291,7 @@ func validateRateLimit(ctx context.Context, canvasID string, cfg map[string]any)
newCtx, cancel := context.WithTimeout(ctx, webhookRateLimitTimeout)
defer cancel()
rdb := rediscli.Get()
rdb := kvrocks.Get()
if rdb == nil {
return fmt.Errorf("rate limit error: redis not initialised")
}
+5 -5
View File
@@ -21,7 +21,7 @@ import (
"net/http"
"ragflow/internal/common"
"ragflow/internal/engine/clickhouse"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"ragflow/internal/server"
"ragflow/internal/server/local"
"ragflow/internal/utility"
@@ -99,7 +99,7 @@ func (h *UserHandler) Register(c *gin.Context) {
return
}
secretKey, err := server.GetSecretKey(ctx, redis.Get())
secretKey, err := server.GetSecretKey(ctx, kvrocks.Get())
if err != nil {
common.ResponseWithCodeData(c, common.CodeServerError, false, err.Error())
return
@@ -169,7 +169,7 @@ func (h *UserHandler) Login(c *gin.Context) {
operationLog.UserID = user.ID
// Sign the access_token using itsdangerous (compatible with Python)
secretKey, err := server.GetSecretKey(ctx, redis.Get())
secretKey, err := server.GetSecretKey(ctx, kvrocks.Get())
if err != nil {
errMessage := fmt.Sprintf("Failed to get secret key: %s", err.Error())
common.ResponseWithCodeData(c, common.CodeServerError, false, errMessage)
@@ -256,7 +256,7 @@ func (h *UserHandler) LoginByEmail(c *gin.Context) {
}
operationLog.UserID = user.ID
secretKey, err := server.GetSecretKey(ctx, redis.Get())
secretKey, err := server.GetSecretKey(ctx, kvrocks.Get())
if err != nil {
errorMessage := fmt.Sprintf("Failed to get secret key: %s", err.Error())
common.ResponseWithCodeData(c, common.CodeServerError, false, errorMessage)
@@ -713,7 +713,7 @@ func (h *UserHandler) ForgotResetPassword(c *gin.Context) {
return
}
secretKey, err := server.GetSecretKey(ctx, redis.Get())
secretKey, err := server.GetSecretKey(ctx, kvrocks.Get())
if err != nil {
common.ResponseWithCodeData(c, common.CodeServerError, false, fmt.Sprintf("Failed to get secret key: %s", err.Error()))
return
+4 -4
View File
@@ -38,7 +38,7 @@ import (
"github.com/cespare/xxhash/v2"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"ragflow/internal/ingestion/component/globals"
)
@@ -52,7 +52,7 @@ const TTL = 7 * 24 * time.Hour
const keyNamespace = "kc"
// Store is the slice of the Redis client this package needs.
// *redis.Client satisfies it; tests supply an in-memory double.
// *kvrocks.Client satisfies it; tests supply an in-memory double.
type Store interface {
Get(ctx context.Context, key string) (string, error)
Set(ctx context.Context, key, value string, exp time.Duration) bool
@@ -63,11 +63,11 @@ type Store interface {
}
// Client resolves the process-wide Redis client as a Store, or nil when Redis
// is not configured. Always resolve through this function: passing redis.Get()
// is not configured. Always resolve through this function: passing kvrocks.Get()
// directly into a Store parameter would wrap a typed nil pointer in a non-nil
// interface, defeating the nil guards below.
func Client() Store {
if c := redis.Get(); c != nil {
if c := kvrocks.Get(); c != nil {
return c
}
return nil
+5 -5
View File
@@ -19,21 +19,21 @@ package chunkcache
import (
"testing"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
)
// Compile-time proof that the real client still satisfies Store. Without it a
// signature change in the redis package would only surface at the call sites.
var _ Store = (*redis.Client)(nil)
var _ Store = (*kvrocks.Client)(nil)
// TestClient_NilInterfaceWhenRedisAbsent is the reason Client() exists. Handing
// redis.Get() straight to a Store parameter would wrap a typed nil pointer in a
// kvrocks.Get() straight to a Store parameter would wrap a typed nil pointer in a
// non-nil interface, defeating every `s == nil` guard and panicking on the
// first call in a Redis-less deployment. Client() must return an interface that
// compares equal to nil.
func TestClient_NilInterfaceWhenRedisAbsent(t *testing.T) {
if redis.Get() != nil {
t.Skip("a global Redis client is configured in this process")
if kvrocks.Get() != nil {
t.Skip("a global Kvrocks client is configured in this process")
}
if got := Client(); got != nil {
t.Fatalf("Client() = %#v, want a nil interface", got)
+10 -10
View File
@@ -29,7 +29,7 @@ import (
_ "ragflow/internal/agent/component"
"ragflow/internal/agent/runtime"
"ragflow/internal/common"
redis2 "ragflow/internal/engine/redis"
kvrocks "ragflow/internal/engine/kvrocks"
"ragflow/internal/ingestion/component"
"ragflow/internal/ingestion/component/globals"
"ragflow/internal/utility"
@@ -53,7 +53,7 @@ type Pipeline struct {
store canvas.CheckPointStore // optional injected; nil -> resolve at Run
tracker *canvas.RunTracker // optional injected; nil -> resolve at Run
// requireResume, when true, makes Run refuse to start if no checkpoint
// store can be resolved (no injected store AND no global Redis client).
// store can be resolved (no injected store AND no global Kvrocks client).
// Plan §6.a M4: a deployment that cannot persist checkpoints must
// not silently degrade to a non-resumable run — it must surface a clear,
// distinguishable error so the caller knows resume is unavailable.
@@ -74,20 +74,20 @@ var ErrResumeUnavailable = errors.New("resume unavailable: no checkpoint store (
type PipelineOption func(*Pipeline)
// WithCheckPointStore injects a checkpoint store. When unset, Run resolves
// one from the global Redis client (and degrades to a non-resumable run when
// one from the global Kvrocks client (and degrades to a non-resumable run when
// Redis is unavailable — plan §6.a).
func WithCheckPointStore(s canvas.CheckPointStore) PipelineOption {
return func(p *Pipeline) { p.store = s }
}
// WithRunTracker injects a RunTracker for interrupt-id persistence / crash
// recovery. When unset, Run resolves one from the global Redis client.
// recovery. When unset, Run resolves one from the global Kvrocks client.
func WithRunTracker(t *canvas.RunTracker) PipelineOption {
return func(p *Pipeline) { p.tracker = t }
}
// WithRequireResume makes Run refuse to start when no checkpoint store can be
// resolved (no injected store AND no global Redis client). This is plan A: a
// resolved (no injected store AND no global Kvrocks client). This is plan A: a
// deployment that cannot persist checkpoints must not silently
// degrade to a non-resumable run — it must surface a clear, distinguishable
// error (ErrResumeUnavailable) so the caller knows resume is unavailable.
@@ -513,15 +513,15 @@ func (p *Pipeline) Run(ctx context.Context, inputs map[string]any, overrideParam
return p.runResumable(ctx, runCtx, current, compiled, store, tracker, runState)
}
// resolveStore returns the injected store, or a Redis-backed one when the
// global Redis client is available. Returns nil (degraded, non-resumable)
// resolveStore returns the injected store, or a Kvrocks-backed one when the
// global Kvrocks client is available. Returns nil (degraded, non-resumable)
// when neither is present.
func (p *Pipeline) resolveStore() canvas.CheckPointStore {
if p.store != nil {
return p.store
}
if redis2.Get() != nil {
return canvas.NewRedisCheckPointStore(defaultCheckpointTTL)
if kvrocks.Get() != nil {
return canvas.NewKvrocksCheckPointStore(defaultCheckpointTTL)
}
return nil
}
@@ -531,7 +531,7 @@ func (p *Pipeline) resolveTracker() *canvas.RunTracker {
if p.tracker != nil {
return p.tracker
}
if redis2.Get() != nil {
if kvrocks.Get() != nil {
return canvas.NewRunTracker(defaultCheckpointTTL)
}
return nil
@@ -29,7 +29,7 @@ import (
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/engine"
redis2 "ragflow/internal/engine/redis"
kvrocks "ragflow/internal/engine/kvrocks"
"ragflow/internal/entity"
"ragflow/internal/ingestion/knowledge_compile"
pipelinepkg "ragflow/internal/ingestion/pipeline"
@@ -768,7 +768,7 @@ func (e *Ingestor) markStopped(ctx context.Context, taskID string) bool {
common.Error(fmt.Sprintf("markStopped: MarkStopped task %s: %v", taskID, err), err)
return false
}
if rc := redis2.Get(); rc != nil {
if rc := kvrocks.Get(); rc != nil {
utility.BestEffort(fmt.Sprintf("clear cancel flag for %s", taskID), func() error {
rc.Delete(ctx, fmt.Sprintf("%s-cancel", taskID))
return nil // Delete returns bool; the bool does not distinguish "not found" from "error"
@@ -818,7 +818,7 @@ func (e *Ingestor) runTask(ctx context.Context, task *entity.IngestionTask) bool
// by the DB status (STOPPING), which defaultCancelCheck falls back to
// when the Redis flag is absent. Clearing a stale flag here is safe:
// a genuine concurrent cancel sets the task to STOPPING in DB.
if rc := redis2.Get(); rc != nil {
if rc := kvrocks.Get(); rc != nil {
key := fmt.Sprintf("%s-cancel", task.ID)
utility.BestEffort(fmt.Sprintf("clear stale cancel flag for %s", task.ID), func() error {
rc.Delete(ctx, key)
@@ -1036,7 +1036,7 @@ func (e *Ingestor) ackOrNack(taskCtx *taskpkg.TaskContext, terminal bool) {
// task status in DB when Redis is unavailable — a STOPPING status
// (set by RequestStop) is treated as a cancel signal.
func (e *Ingestor) defaultCancelCheck(ctx context.Context, taskID string) bool {
rc := redis2.Get()
rc := kvrocks.Get()
if rc != nil {
if ok, _ := rc.Exist(ctx, fmt.Sprintf("%s-cancel", taskID)); ok {
return true
+5 -2
View File
@@ -230,8 +230,11 @@ func GetAllConfigs() ([]map[string]interface{}, error) {
// cache engine
cacheType := globalConfig.CacheEngineType()
switch cacheType {
case "redis":
redisConfig := globalConfig.GetRedisConfig()
// The Go stack talks to Kvrocks. "redis" is accepted for backwards
// compatibility with the shared service_conf.yaml.template; both map to
// the same Kvrocks backend, so export the same connection settings.
case "redis", "kvrocks":
redisConfig := globalConfig.GetKvrocksConfig()
exportedRedisConfigs := redisConfig.ExportConfigs()
allConfigs = append(allConfigs, exportedRedisConfigs)
default:
+37 -32
View File
@@ -24,12 +24,15 @@ import (
"github.com/spf13/viper"
)
// CacheEngineConfig holds the Go services' cache/queue backend settings. The Go
// stack always talks to Kvrocks (a RocksDB-backed, Redis-protocol store) so it
// is not subject to the in-memory maxmemory cap that Valkey/Redis enforces.
type CacheEngineConfig struct {
Redis RedisConfig `mapstructure:"redis"`
Kvrocks KvrocksConfig `mapstructure:"kvrocks"`
}
// RedisConfig Redis configuration
type RedisConfig struct {
// KvrocksConfig connection settings for the Kvrocks backend.
type KvrocksConfig struct {
Host string `mapstructure:"host"`
Port int `mapstructure:"port"`
Username string `mapstructure:"username"`
@@ -41,8 +44,11 @@ func (c *Config) ParseCacheEngineConfig(v *viper.Viper) error {
cacheEngineType := c.general.CacheEngine
var err error
switch cacheEngineType {
case "redis":
err = c.parseRedisConfig(v)
// The Go stack connects to Kvrocks. "redis" is accepted for backwards
// compatibility with the shared service_conf.yaml.template (its `redis:`
// section still drives the Python/Valkey path); both map to Kvrocks here.
case "redis", "kvrocks":
err = c.parseKvrocksConfig(v)
default:
return fmt.Errorf("cache engine type %s is not supported", cacheEngineType)
}
@@ -50,18 +56,18 @@ func (c *Config) ParseCacheEngineConfig(v *viper.Viper) error {
return err
}
func (c *Config) parseRedisConfig(v *viper.Viper) error {
// Default Redis config
c.cacheEngine.Redis.Host = "localhost"
c.cacheEngine.Redis.Port = 6379
c.cacheEngine.Redis.DB = 1
c.cacheEngine.Redis.Username = ""
c.cacheEngine.Redis.Password = "infini_rag_flow"
func (c *Config) parseKvrocksConfig(v *viper.Viper) error {
// Sensible defaults; deployed values come from the `kvrocks` section.
c.cacheEngine.Kvrocks.Host = "localhost"
c.cacheEngine.Kvrocks.Port = 6379
c.cacheEngine.Kvrocks.DB = 1
c.cacheEngine.Kvrocks.Username = ""
c.cacheEngine.Kvrocks.Password = "infini_rag_flow"
if !v.IsSet("redis") {
if !v.IsSet("kvrocks") {
return nil
}
sub := v.Sub("redis")
sub := v.Sub("kvrocks")
if sub == nil {
return nil
}
@@ -71,48 +77,47 @@ func (c *Config) parseRedisConfig(v *viper.Viper) error {
// Handle host:port format (e.g., "localhost:6379")
host, portStr, err := net.SplitHostPort(hostStr)
if err != nil {
return fmt.Errorf("error address format of Redis: %s", hostStr)
return fmt.Errorf("error address format of Kvrocks: %s", hostStr)
}
if host == "" {
return fmt.Errorf("empty host of Redis configuration")
return fmt.Errorf("empty host of Kvrocks configuration")
}
c.cacheEngine.Redis.Host = host
c.cacheEngine.Kvrocks.Host = host
if portStr != "" {
var port int
if port, err = strconv.Atoi(portStr); err == nil {
c.cacheEngine.Redis.Port = port
c.cacheEngine.Kvrocks.Port = port
}
}
}
if sub.IsSet("db") {
c.cacheEngine.Redis.DB = sub.GetInt("db")
c.cacheEngine.Kvrocks.DB = sub.GetInt("db")
}
if sub.IsSet("username") {
c.cacheEngine.Redis.Username = sub.GetString("username")
c.cacheEngine.Kvrocks.Username = sub.GetString("username")
}
if sub.IsSet("password") {
c.cacheEngine.Redis.Password = sub.GetString("password")
c.cacheEngine.Kvrocks.Password = sub.GetString("password")
}
return nil
}
func (c *Config) GetRedisConfig() RedisConfig {
return c.cacheEngine.Redis
func (c *Config) GetKvrocksConfig() KvrocksConfig {
return c.cacheEngine.Kvrocks
}
func (r RedisConfig) ExportConfigs() map[string]interface{} {
var redisConfigs map[string]interface{}
redisConfigs = make(map[string]interface{})
redisConfigs["host"] = r.Host
redisConfigs["port"] = r.Port
redisConfigs["username"] = r.Username
redisConfigs["password"] = r.Password
redisConfigs["db"] = r.DB
return redisConfigs
func (r KvrocksConfig) ExportConfigs() map[string]interface{} {
kvrocksConfigs := make(map[string]interface{})
kvrocksConfigs["host"] = r.Host
kvrocksConfigs["port"] = r.Port
kvrocksConfigs["username"] = r.Username
kvrocksConfigs["password"] = r.Password
kvrocksConfigs["db"] = r.DB
return kvrocksConfigs
}
@@ -0,0 +1,64 @@
//
// 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 server
import (
"os"
"path/filepath"
"testing"
)
// TestGetAllConfigsKvrocksCacheEngine reproduces Finding 1: when the Go stack is
// configured with `cache_engine: kvrocks`, GetAllConfigs must export the Kvrocks
// connection settings instead of failing with "not supported cache engine".
func TestGetAllConfigsKvrocksCacheEngine(t *testing.T) {
dir := t.TempDir()
cfgPath := filepath.Join(dir, "service_conf.yaml")
// Only cache_engine is overridden; the remaining engine types keep the
// ParseGeneralConfig defaults (mysql/elasticsearch/minio/nats/clickhouse),
// so GetAllConfigs only fails on the cache-engine dispatch.
content := `
general:
cache_engine: kvrocks
kvrocks:
host: 'localhost:6379'
`
if err := os.WriteFile(cfgPath, []byte(content), 0o600); err != nil {
t.Fatalf("write temp config: %v", err)
}
if err := Init(cfgPath); err != nil {
t.Fatalf("Init: %v", err)
}
allConfigs, err := GetAllConfigs()
if err != nil {
t.Fatalf("GetAllConfigs with cache_engine=kvrocks: %v", err)
}
var foundKvrocks bool
for _, cfg := range allConfigs {
if host, ok := cfg["host"]; ok && host == "localhost" {
if port, ok := cfg["port"]; ok && port == 6379 {
foundKvrocks = true
}
}
}
if !foundKvrocks {
t.Fatalf("Kvrocks config not exported in GetAllConfigs: %#v", allConfigs)
}
}
+4 -4
View File
@@ -349,7 +349,7 @@ type AgentService struct {
// in-memory; a follow-up phase moves to Redis per plan §4.9.
runner *canvas.Runner
// Phase 4.4 V2 — Redis-backed run infrastructure. nil = in-memory
// Phase 4.4 V2 — Kvrocks-backed run infrastructure. nil = in-memory
// / no-tracking (test path, current production boot path until
// cmd/server_main.go wires them in v3.6.0).
//
@@ -391,7 +391,7 @@ func NewAgentService() *AgentService {
}
// NewAgentServiceWithOptions is the production constructor that
// injects the Redis-backed run infrastructure. The zero-arg
// injects the Kvrocks-backed run infrastructure. The zero-arg
// NewAgentService() remains as a thin wrapper that calls this with
// all-nil options so existing call sites (cmd/server_main.go,
// handler tests, agent_test.go) keep compiling.
@@ -2847,7 +2847,7 @@ func shouldTreatAsCompletedLoopRun(err error, answer string) bool {
}
// markRunSucceeded records the run as completed successfully via
// the Redis-backed RunTracker. No-op when tracker is nil (test path)
// the Kvrocks-backed RunTracker. No-op when tracker is nil (test path)
// or when the underlying Redis call fails (degraded boot).
func (s *AgentService) markRunSucceeded(ctx context.Context, runID string) {
if s.runTracker == nil {
@@ -2862,7 +2862,7 @@ func (s *AgentService) markRunSucceeded(ctx context.Context, runID string) {
}
// markRunFailed records the run as failed (with reason) via the
// Redis-backed RunTracker. No-op when tracker is nil or the
// Kvrocks-backed RunTracker. No-op when tracker is nil or the
// underlying Redis call fails.
func (s *AgentService) markRunFailed(ctx context.Context, runID, reason string) {
if s.runTracker == nil {
+12 -12
View File
@@ -20,7 +20,7 @@
// These tests pin the production chain end-to-end: loadCanvasForUser
// → versionDAO.GetLatest → decodeCanvasFromDSL → canvas.Compile →
// cc.Workflow.Invoke → orchestrator answer extraction. They also
// cover the boot-wiring surface (Redis-backed CheckPointStore +
// cover the boot-wiring surface (Kvrocks-backed CheckPointStore +
// RunTracker) and the failure paths (compile error, invoke
// error, wait-for-user resume cycle). If any of these tests
// fails, the RunAgent path has regressed.
@@ -197,7 +197,7 @@ func TestRunAgent_RealCanvas_BeginMessage(t *testing.T) {
tracker, mr := newRunTrackerForTest(t, 30*24*time.Hour)
cpClient := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = cpClient.Close() })
cp := canvas.NewRedisCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
cp := canvas.NewKvrocksCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
svc := NewAgentServiceWithOptions(cp, nil, tracker)
events, err := svc.RunAgent(
t.Context(),
@@ -589,7 +589,7 @@ func TestRunAgent_RealCanvas_WaitForUserResume(t *testing.T) {
tracker, mr := newRunTrackerForTest(t, 30*24*time.Hour)
cpClient := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = cpClient.Close() })
cp := canvas.NewRedisCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
cp := canvas.NewKvrocksCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
// stateSerializer is intentionally nil so eino's default
// InternalSerializer is used (which knows about CanvasState
// via runtime.RegisterSerializableType[CanvasState]). The
@@ -794,7 +794,7 @@ func TestRunAgent_RealCanvas_WaitForUserResume_EventSemantics(t *testing.T) {
tracker, mr := newRunTrackerForTest(t, 30*24*time.Hour)
cpClient := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = cpClient.Close() })
cp := canvas.NewRedisCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
cp := canvas.NewKvrocksCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
svc := NewAgentServiceWithOptions(cp, nil, tracker)
events1, err := svc.RunAgent(
@@ -1005,7 +1005,7 @@ func TestRunAgent_AllFixture_LoopInterruptResume(t *testing.T) {
tracker, mr := newRunTrackerForTest(t, 30*24*time.Hour)
cpClient := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = cpClient.Close() })
cp := canvas.NewRedisCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
cp := canvas.NewKvrocksCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
svc := NewAgentServiceWithOptions(cp, nil, tracker)
events1, err := svc.RunAgent(
@@ -1126,7 +1126,7 @@ func TestRunAgent_AllFixture_LoopInterruptResume_MultiTurn(t *testing.T) {
tracker, mr := newRunTrackerForTest(t, 30*24*time.Hour)
cpClient := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = cpClient.Close() })
cp := canvas.NewRedisCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
cp := canvas.NewKvrocksCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
svc := NewAgentServiceWithOptions(cp, nil, tracker)
sessionID := "session-all-loop-multi"
@@ -1222,7 +1222,7 @@ func TestRunAgent_AllFixture_IterationFormatsItems(t *testing.T) {
tracker, mr := newRunTrackerForTest(t, 30*24*time.Hour)
cpClient := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = cpClient.Close() })
cp := canvas.NewRedisCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
cp := canvas.NewKvrocksCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
svc := NewAgentServiceWithOptions(cp, nil, tracker)
sessionID := "session-all-iteration"
@@ -1335,7 +1335,7 @@ func TestRunAgent_AllFixture_VarAssigner(t *testing.T) {
tracker, mr := newRunTrackerForTest(t, 30*24*time.Hour)
cpClient := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = cpClient.Close() })
cp := canvas.NewRedisCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
cp := canvas.NewKvrocksCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
svc := NewAgentServiceWithOptions(cp, nil, tracker)
events, err := svc.RunAgent(
t.Context(),
@@ -1419,7 +1419,7 @@ func TestRunAgent_AllFixture_DataOps(t *testing.T) {
tracker, mr := newRunTrackerForTest(t, 30*24*time.Hour)
cpClient := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = cpClient.Close() })
cp := canvas.NewRedisCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
cp := canvas.NewKvrocksCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
svc := NewAgentServiceWithOptions(cp, nil, tracker)
events, err := svc.RunAgent(
t.Context(),
@@ -1593,7 +1593,7 @@ func TestRunAgent_AllFixture_CategorizeResume(t *testing.T) {
tracker, mr := newRunTrackerForTest(t, 30*24*time.Hour)
cpClient := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = cpClient.Close() })
cp := canvas.NewRedisCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
cp := canvas.NewKvrocksCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
svc := NewAgentServiceWithOptions(cp, nil, tracker)
events1, err := svc.RunAgent(
t.Context(),
@@ -1794,7 +1794,7 @@ func newRunTrackerForTest(t *testing.T, ttl time.Duration) (*canvas.RunTracker,
// TestRunAgent_RunTracker_AttachCheckpoint_CallSequence pins the
// production boot path that v3.6.0 enables: an AgentService
// constructed via NewAgentServiceWithOptions (with a real
// RedisCheckPointStore + CanvasStateSerializer + RunTracker) must
// KvrocksCheckPointStore + CanvasStateSerializer + RunTracker) must
// record the full Start → AttachCheckpoint → MarkSucceeded sequence
// against the run hash during a single successful run.
//
@@ -1836,7 +1836,7 @@ func TestRunAgent_RunTracker_AttachCheckpoint_CallSequence(t *testing.T) {
tracker, mr := newRunTrackerForTest(t, 30*24*time.Hour)
cpClient := redis.NewClient(&redis.Options{Addr: mr.Addr()})
t.Cleanup(func() { _ = cpClient.Close() })
cp := canvas.NewRedisCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
cp := canvas.NewKvrocksCheckPointStoreWithClient(cpClient, 30*24*time.Hour)
dsl := map[string]any{
"components": map[string]any{
+2 -2
View File
@@ -32,7 +32,7 @@ import (
"ragflow/internal/agent/dsl"
"ragflow/internal/common"
"ragflow/internal/dao"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"ragflow/internal/entity"
)
@@ -177,7 +177,7 @@ func (s *BotService) AgentbotLogs(ctx context.Context, tenantID, agentID, messag
if _, err := s.loadCanvas(ctx, tenantID, agentID); err != nil {
return nil, common.CodeDataError, err
}
payload, err := redis.Get().Get(ctx, fmt.Sprintf("%s-%s-logs", agentID, messageID))
payload, err := kvrocks.Get().Get(ctx, fmt.Sprintf("%s-%s-logs", agentID, messageID))
if err != nil {
return nil, common.CodeServerError, errors.New("failed to read agent logs")
}
+5 -5
View File
@@ -28,7 +28,7 @@ import (
"io"
"net/http"
"net/url"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
syncerconnector "ragflow/internal/syncer/connector"
"ragflow/internal/utility"
"strings"
@@ -68,7 +68,7 @@ var (
"https://www.googleapis.com/auth/admin.directory.user.readonly",
"https://www.googleapis.com/auth/admin.directory.group.readonly",
}
connectorRedisGet = redis.Get
connectorRedisGet = kvrocks.Get
)
// Sentinel errors so handlers can map to the proper response codes.
@@ -544,7 +544,7 @@ func (s *ConnectorService) StartGoogleWebOAuth(ctx context.Context, userID, sour
return nil, common.CodeServerError, fmt.Errorf("failed to initialize Google OAuth flow. Please verify the uploaded client configuration")
}
redisClient := redis.Get()
redisClient := kvrocks.Get()
if redisClient == nil {
return nil, common.CodeServerError, fmt.Errorf("no configure Redis on the server")
}
@@ -578,7 +578,7 @@ func (s *ConnectorService) GoogleWebOAuthCallback(ctx context.Context, source, s
return renderWebOAuthPopup("", false, "Missing OAuth state parameter.", source)
}
redisClient := redis.Get()
redisClient := kvrocks.Get()
if redisClient == nil {
return renderWebOAuthPopup(stateID, false, "Authorization session expired. Please restart from the main window.", source)
}
@@ -639,7 +639,7 @@ func (s *ConnectorService) PollGoogleWebOAuthResult(ctx context.Context, userID,
return nil, common.CodeArgumentError, fmt.Errorf("required argument is missing: flow_id")
}
redisClient := redis.Get()
redisClient := kvrocks.Get()
if redisClient == nil {
return nil, common.CodeRunning, fmt.Errorf("authorization is still pending")
}
+2 -2
View File
@@ -8,7 +8,7 @@ import (
"testing"
"ragflow/internal/common"
redisengine "ragflow/internal/engine/redis"
kvrocks "ragflow/internal/engine/kvrocks"
)
type roundTripFunc func(*http.Request) (*http.Response, error)
@@ -20,7 +20,7 @@ func (f roundTripFunc) RoundTrip(req *http.Request) (*http.Response, error) {
func forceNilConnectorRedis(t *testing.T) {
t.Helper()
previous := connectorRedisGet
connectorRedisGet = func() *redisengine.Client { return nil }
connectorRedisGet = func() *kvrocks.Client { return nil }
t.Cleanup(func() {
connectorRedisGet = previous
})
+3 -3
View File
@@ -9,7 +9,7 @@ import (
"ragflow/internal/common"
"ragflow/internal/dao"
redis2 "ragflow/internal/engine/redis"
kvrocks "ragflow/internal/engine/kvrocks"
"ragflow/internal/entity"
"go.uber.org/zap"
@@ -320,7 +320,7 @@ func (s *IngestionTaskService) RequestStop(ctx context.Context, taskID string) (
// Mirror Python's cancel_all_task_of: set Redis cancel flag so the
// running worker's pollCancel detects the stop immediately rather
// than waiting for the next DB poll (up to 3s).
if rc := redis2.Get(); rc != nil {
if rc := kvrocks.Get(); rc != nil {
rc.Set(ctx, fmt.Sprintf("%s-cancel", taskID), "x", 1*time.Hour)
}
return task, nil
@@ -732,7 +732,7 @@ func (s *IngestionTaskService) settlePublishFailure(ctx context.Context, task *e
// RequestStop sets for a RUNNING task. No-op when Redis is unavailable —
// the DB STOPPING status remains the fallback cancel signal.
func clearCancelFlag(ctx context.Context, taskID string) {
if rc := redis2.Get(); rc != nil {
if rc := kvrocks.Get(); rc != nil {
rc.Delete(ctx, fmt.Sprintf("%s-cancel", taskID))
}
}
+2 -2
View File
@@ -61,7 +61,7 @@ import (
"ragflow/internal/common"
"ragflow/internal/dao"
redisengine "ragflow/internal/engine/redis"
kvrocks "ragflow/internal/engine/kvrocks"
"ragflow/internal/entity"
models "ragflow/internal/entity/models"
"ragflow/internal/utility"
@@ -230,7 +230,7 @@ func (s *MemoryMessageService) ReconcileMemoryTasks(ctx context.Context, limit i
// generateRawMessageID returns the Redis auto-increment id used by the Python
// side (`REDIS_CONN.generate_auto_increment_id(namespace="memory")`).
func generateRawMessageID(ctx context.Context) int64 {
if redisClient := redisengine.Get(); redisClient != nil {
if redisClient := kvrocks.Get(); redisClient != nil {
if id := redisClient.GenerateAutoIncrementID(ctx, "id_generator", "memory", 1, nil); id > 0 {
return id
}
@@ -260,7 +260,7 @@ func TestBuildMemoryTaskRecordsShape(t *testing.T) {
}
// TestGenerateRawMessageID_Unique: two calls produce different
// values. (Wall-clock based today; the Redis-backed counter will
// values. (Wall-clock based today; the Kvrocks-backed counter will
// be added when the project's Redis client lands.)
func TestGenerateRawMessageID_Unique(t *testing.T) {
ctx := t.Context()
+3 -3
View File
@@ -20,7 +20,7 @@ import (
"context"
"fmt"
"ragflow/internal/common"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"ragflow/internal/entity"
"time"
)
@@ -43,10 +43,10 @@ type OAuthCallbackResult struct {
IsNewUser bool
}
func (s *UserService) OAuthLoginInitiate(channel string, redis *redis.Client) (*OAuthLoginInit, common.ErrorCode, error) {
func (s *UserService) OAuthLoginInitiate(channel string, redis *kvrocks.Client) (*OAuthLoginInit, common.ErrorCode, error) {
return nil, common.CodeServerError, fmt.Errorf("oauth login initiate not implemented")
}
func (s *UserService) OAuthCallback(ctx context.Context, channel, code, callbackState, expectedState string, redis *redis.Client) (*OAuthCallbackResult, common.ErrorCode, error) {
func (s *UserService) OAuthCallback(ctx context.Context, channel, code, callbackState, expectedState string, redis *kvrocks.Client) (*OAuthCallbackResult, common.ErrorCode, error) {
return nil, common.CodeDataError, fmt.Errorf("oauth callback not implemented")
}
+4 -4
View File
@@ -21,7 +21,7 @@ import (
"encoding/json"
"fmt"
"ragflow/internal/common"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"ragflow/internal/entity"
"time"
@@ -230,7 +230,7 @@ func (s *SystemService) getDatabaseStatus(ctx context.Context) ComponentStatus {
func (s *SystemService) getRedisStatus(ctx context.Context) ComponentStatus {
startedAt := time.Now()
redisClient := redis.Get()
redisClient := kvrocks.Get()
if redisClient == nil {
return ComponentStatus{
"status": "red",
@@ -254,7 +254,7 @@ func (s *SystemService) getRedisStatus(ctx context.Context) ComponentStatus {
func (s *SystemService) getTaskExecutorHeartbeats(ctx context.Context) map[string][]interface{} {
heartbeatsByExecutor := map[string][]interface{}{}
redisClient := redis.Get()
redisClient := kvrocks.Get()
if redisClient == nil {
return heartbeatsByExecutor
}
@@ -328,7 +328,7 @@ func GetComponentsHealthz(ctx context.Context) (*HealthzResponse, bool) {
}
redisOK, redisMeta := timedHealthCheck(func() error {
redisClient := redis.Get()
redisClient := kvrocks.Get()
if redisClient == nil || !redisClient.Health(ctx) {
return fmt.Errorf("redis is not healthy")
}
+3 -3
View File
@@ -21,7 +21,7 @@ import (
"encoding/json"
"fmt"
"ragflow/internal/common"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"sort"
"strings"
"time"
@@ -60,7 +60,7 @@ func GetTagsFromCache(ctx context.Context, kbIDs []string) (map[string]float64,
return nil, nil
}
redisClient := redis.Get()
redisClient := kvrocks.Get()
if redisClient == nil {
common.Warn("Redis client not available, skipping cache lookup")
return nil, nil
@@ -88,7 +88,7 @@ func SetTagsToCache(ctx context.Context, kbIDs []string, tags map[string]float64
return nil
}
redisClient := redis.Get()
redisClient := kvrocks.Get()
if redisClient == nil {
common.Warn("Redis client not available, skipping cache store")
return nil
+6 -6
View File
@@ -26,7 +26,7 @@ import (
"fmt"
"hash"
"ragflow/internal/common"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"ragflow/internal/entity"
"ragflow/internal/server"
"ragflow/internal/server/config"
@@ -595,7 +595,7 @@ func defaultUserLanguage() string {
// using itsdangerous URLSafeTimedSerializer to get the actual access_token
func (s *UserService) GetUserByToken(ctx context.Context, authorization string) (*entity.User, common.ErrorCode, error) {
// Get secret key from config
secretKey, err := server.GetSecretKey(ctx, redis.Get())
secretKey, err := server.GetSecretKey(ctx, kvrocks.Get())
if err != nil {
return nil, common.CodeUnauthorized, err
}
@@ -1085,7 +1085,7 @@ func (s *UserService) ForgotIssueCaptcha(ctx context.Context, email string) (cap
return "", "", common.CodeServerError, err
}
captchaID = utility.GenerateToken()
if ok := redis.Get().Set(ctx, utility.CaptchaIDRedisKey(captchaID), text, 60*time.Second); !ok {
if ok := kvrocks.Get().Set(ctx, utility.CaptchaIDRedisKey(captchaID), text, 60*time.Second); !ok {
return "", "", common.CodeServerError, fmt.Errorf("failed to store captcha")
}
imageDataURL = utility.RenderCaptchaPNGDataURL(text)
@@ -1105,7 +1105,7 @@ func (s *UserService) ForgotSendOTP(ctx context.Context, email, captchaID, captc
return common.CodeDataError, fmt.Errorf("invalid email")
}
rc := redis.Get()
rc := kvrocks.Get()
captchaKey := utility.CaptchaIDRedisKey(captchaID)
stored, _ := rc.Get(ctx, captchaKey)
if stored == "" {
@@ -1204,7 +1204,7 @@ func (s *UserService) ForgotVerifyOTP(ctx context.Context, email, otp string) (c
return common.CodeDataError, fmt.Errorf("invalid email")
}
rc := redis.Get()
rc := kvrocks.Get()
codeKey, attemptsKey, lastSentKey, lockKey := utility.OTPRedisKeys(email)
if locked, _ := rc.Get(ctx, lockKey); locked != "" {
@@ -1274,7 +1274,7 @@ func (s *UserService) ForgotResetPassword(ctx context.Context, req *ForgotResetP
return nil, common.CodeArgumentError, fmt.Errorf("email and passwords are required")
}
rc := redis.Get()
rc := kvrocks.Get()
verifiedKey := utility.OTPVerifiedRedisKey(req.Email)
if v, _ := rc.Get(ctx, verifiedKey); v != "1" {
return nil, common.CodeAuthenticationError, fmt.Errorf("email not verified")
+5 -5
View File
@@ -19,7 +19,7 @@ package syncer
import (
"context"
"fmt"
"ragflow/internal/engine/redis"
"ragflow/internal/engine/kvrocks"
"ragflow/internal/utility"
"sync"
"time"
@@ -43,12 +43,12 @@ type ConnectorLock struct {
holder string
mu sync.Mutex
local map[string]struct{}
redis map[string]*redis.DistributedLock
redis map[string]*kvrocks.DistributedLock
}
// NewConnectorLock creates an empty connector/KB lock.
func NewConnectorLock() *ConnectorLock {
return &ConnectorLock{holder: utility.GenerateUUID(), local: map[string]struct{}{}, redis: map[string]*redis.DistributedLock{}}
return &ConnectorLock{holder: utility.GenerateUUID(), local: map[string]struct{}{}, redis: map[string]*kvrocks.DistributedLock{}}
}
// TryLock attempts to acquire the connector/KB lock without blocking.
@@ -65,8 +65,8 @@ func (l *ConnectorLock) TryLock(connectorID, kbID string) (ConnectorLockLease, b
l.local[key] = struct{}{}
l.mu.Unlock()
if client := redis.Get(); client != nil {
lock := redis.NewDistributedLock(key, l.holder, connectorLockTTL, 0)
if client := kvrocks.Get(); client != nil {
lock := kvrocks.NewDistributedLock(key, l.holder, connectorLockTTL, 0)
if lock == nil || !lock.Acquire(context.Background()) {
l.mu.Lock()
delete(l.local, key)
+10 -10
View File
@@ -18,18 +18,18 @@ import logging
import re
from html import escape as html_escape
from io import BytesIO
from docx import Document
from docx.table import Table as DocxTable
from docx.text.paragraph import Paragraph
from common.constants import ParserType, MAXIMUM_PAGE_NUMBER
from deepdoc.parser.utils import get_text
from rag.nlp import bullets_category, remove_contents_table, make_colon_as_title, tokenize_chunks, docx_question_level, tree_merge, DEFAULT_DELIMITER
from rag.nlp import rag_tokenizer, Node
from deepdoc.parser import PdfParser, DocxParser, HtmlParser
from api.db.joint_services.tenant_model_service import get_composite_model_name_by_id
from rag.app.naive import by_plaintext, PARSERS
from common.constants import MAXIMUM_PAGE_NUMBER, ParserType
from common.parser_config_utils import normalize_layout_recognizer
from deepdoc.parser import DocxParser, HtmlParser, PdfParser
from deepdoc.parser.utils import get_text
from rag.app.naive import PARSERS, by_plaintext
from rag.nlp import DEFAULT_DELIMITER, Node, bullets_category, docx_question_level, make_colon_as_title, rag_tokenizer, remove_contents_table, tokenize_chunks, tree_merge
class Docx(DocxParser):
@@ -136,15 +136,15 @@ class Pdf(PdfParser):
start = timer()
callback(msg="OCR started")
self.__images__(filename if binary is None else binary, zoomin, from_page, to_page, callback)
callback(msg="OCR finished ({:.2f}s)".format(timer() - start))
callback(msg=f"OCR finished ({timer() - start:.2f}s)")
start = timer()
self._layouts_rec(zoomin)
callback(0.67, "Layout analysis ({:.2f}s)".format(timer() - start))
logging.debug("layouts: {}".format((timer() - start)))
callback(0.67, f"Layout analysis ({timer() - start:.2f}s)")
logging.debug(f"layouts: {timer() - start}")
self._naive_vertical_merge()
callback(0.8, "Text extraction ({:.2f}s)".format(timer() - start))
callback(0.8, f"Text extraction ({timer() - start:.2f}s)")
return [(b["text"], self._line_tag(b, zoomin)) for b in self.boxes], None
+17 -16
View File
@@ -14,21 +14,22 @@
# limitations under the License.
#
import logging
import copy
import logging
import re
from common.constants import ParserType, MAXIMUM_PAGE_NUMBER
from io import BytesIO
from deepdoc.parser.utils import extract_pdf_outlines
from rag.nlp import rag_tokenizer, tokenize, tokenize_table, bullets_category, title_frequency, tokenize_chunks, docx_question_level, attach_media_context, concat_img, DEFAULT_DELIMITER
from common.token_utils import num_tokens_from_string
from deepdoc.parser import PdfParser, DocxParser
from deepdoc.parser.figure_parser import vision_figure_parser_pdf_wrapper, vision_figure_parser_docx_wrapper
from docx import Document
from api.db.joint_services.tenant_model_service import get_composite_model_name_by_id
from rag.app.naive import by_plaintext, PARSERS
from common.constants import MAXIMUM_PAGE_NUMBER, ParserType
from common.parser_config_utils import normalize_layout_recognizer
from common.token_utils import num_tokens_from_string
from deepdoc.parser import DocxParser, PdfParser
from deepdoc.parser.figure_parser import vision_figure_parser_docx_wrapper, vision_figure_parser_pdf_wrapper
from deepdoc.parser.utils import extract_pdf_outlines
from rag.app.naive import PARSERS, by_plaintext
from rag.nlp import DEFAULT_DELIMITER, attach_media_context, bullets_category, concat_img, docx_question_level, rag_tokenizer, title_frequency, tokenize, tokenize_chunks, tokenize_table
class Pdf(PdfParser):
@@ -42,24 +43,24 @@ class Pdf(PdfParser):
start = timer()
callback(msg="OCR started")
self.__images__(filename if binary is None else binary, zoomin, from_page, to_page, callback)
callback(msg="OCR finished ({:.2f}s)".format(timer() - start))
logging.debug("OCR: {}".format(timer() - start))
callback(msg=f"OCR finished ({timer() - start:.2f}s)")
logging.debug(f"OCR: {timer() - start}")
start = timer()
self._layouts_rec(zoomin)
callback(0.65, "Layout analysis ({:.2f}s)".format(timer() - start))
logging.debug("layouts: {}".format(timer() - start))
callback(0.65, f"Layout analysis ({timer() - start:.2f}s)")
logging.debug(f"layouts: {timer() - start}")
start = timer()
self._table_transformer_job(zoomin)
callback(0.67, "Table analysis ({:.2f}s)".format(timer() - start))
callback(0.67, f"Table analysis ({timer() - start:.2f}s)")
start = timer()
self._text_merge()
tbls = self._extract_table_figure(True, zoomin, True, True)
self._concat_downward()
self._filter_forpages()
callback(0.68, "Text merged ({:.2f}s)".format(timer() - start))
callback(0.68, f"Text merged ({timer() - start:.2f}s)")
# clean mess
for b in self.boxes:
@@ -258,7 +259,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
# MinerU tags stay local and one-based for crop(); crop() adds
# the task's page offset when it emits final positions.
pn += 1
return "@@{}\t{:.1f}\t{:.1f}\t{:.1f}\t{:.1f}##".format(pn, left, right, top, bottom)
return f"@@{pn}\t{left:.1f}\t{right:.1f}\t{top:.1f}\t{bottom:.1f}##"
chunks = []
last_sid = -2
+41 -39
View File
@@ -15,57 +15,59 @@
#
import logging
import re
import os
import re
from collections.abc import Callable
from functools import reduce
from io import BytesIO
from timeit import default_timer as timer
from typing import Any, Callable
from typing import Any
from docx import Document
from docx.opc.pkgreader import _SerializedRelationships, _SerializedRelationship
from docx.opc.oxml import parse_xml
from docx.opc.pkgreader import _SerializedRelationship, _SerializedRelationships
from docx.table import Table as DocxTable
from docx.text.paragraph import Paragraph
from docx.opc.oxml import parse_xml
from markdown import markdown
from PIL import Image
from common.token_utils import num_tokens_from_string
from common.constants import LLMType, MAXIMUM_PAGE_NUMBER
from api.db.services.llm_service import LLMBundle
from api.db.joint_services.tenant_model_service import (
ensure_mineru_from_env,
ensure_opendataloader_from_env,
ensure_paddleocr_from_env,
get_composite_model_name_by_id,
get_first_provider_model_name,
resolve_model_config,
get_tenant_default_model_by_type,
resolve_model_config,
)
from rag.utils.file_utils import extract_embed_file, extract_links_from_pdf, extract_links_from_docx, extract_html
from deepdoc.parser import DocxParser, EpubParser, ExcelParser, HtmlParser, JsonParser, MarkdownElementExtractor, MarkdownParser, PdfParser, TxtParser
from deepdoc.parser.figure_parser import VisionFigureParser, vision_figure_parser_docx_wrapper_naive, vision_figure_parser_pdf_wrapper
from deepdoc.parser.pdf_parser import PlainParser, VisionParser
from deepdoc.parser.docling_parser import DoclingParser
from deepdoc.parser.monkeyocrv2_parser import MonkeyOCRv2Parser
from deepdoc.parser.tcadp_parser import TCADPParser
from api.db.services.llm_service import LLMBundle
from common.constants import MAXIMUM_PAGE_NUMBER, LLMType
from common.float_utils import normalize_overlapped_percent
from common.parser_config_utils import has_mineru_options, is_tenant_model_id, normalize_layout_recognizer
from common.text_utils import normalize_arabic_presentation_forms
from common.token_utils import num_tokens_from_string
from deepdoc.parser import DocxParser, EpubParser, ExcelParser, HtmlParser, JsonParser, MarkdownElementExtractor, MarkdownParser, PdfParser, TxtParser
from deepdoc.parser.docling_parser import DoclingParser
from deepdoc.parser.figure_parser import VisionFigureParser, vision_figure_parser_docx_wrapper_naive, vision_figure_parser_pdf_wrapper
from deepdoc.parser.monkeyocrv2_parser import MonkeyOCRv2Parser
from deepdoc.parser.pdf_parser import PlainParser, VisionParser
from deepdoc.parser.tcadp_parser import TCADPParser
from rag.nlp import (
concat_img,
DEFAULT_DELIMITER,
append_context2table_image4pdf,
concat_img,
decode_text,
doc_tokenize_chunks_with_images,
naive_merge,
naive_merge_with_images,
naive_merge_docx,
naive_merge_with_images,
rag_tokenizer,
tokenize_chunks,
tokenize_chunks_with_positions,
doc_tokenize_chunks_with_images,
tokenize_table,
append_context2table_image4pdf,
tokenize_chunks_with_images,
) # noqa: F401
tokenize_chunks_with_positions,
tokenize_table,
)
from rag.utils.file_utils import extract_embed_file, extract_html, extract_links_from_docx, extract_links_from_pdf
def _is_short_header(text, max_tokens=50):
@@ -596,7 +598,6 @@ def _is_toc_paragraph(text: str, style_name: str, paragraph=None) -> bool:
class Docx(DocxParser):
def __init__(self):
"""Initialize the naive DOCX parser."""
pass
def __clean(self, line):
line = re.sub(r"\u3000", " ", line).strip()
@@ -605,6 +606,7 @@ class Docx(DocxParser):
def __get_nearest_title(self, table_index, filename):
"""Get the hierarchical title structure before the table"""
import re
from docx.text.paragraph import Paragraph
titles = []
@@ -651,7 +653,7 @@ class Docx(DocxParser):
if block_type != "p":
continue
if block.style and block.style.name and re.search(r"Heading\s*(\d+)", block.style.name, re.I):
if block.style and block.style.name and re.search(r"Heading\s*(\d+)", block.style.name, re.IGNORECASE):
try:
level_match = re.search(r"(\d+)", block.style.name)
if level_match:
@@ -680,7 +682,7 @@ class Docx(DocxParser):
if block_type != "p":
continue
if block.style and re.search(r"Heading\s*(\d+)", block.style.name, re.I):
if block.style and re.search(r"Heading\s*(\d+)", block.style.name, re.IGNORECASE):
try:
level_match = re.search(r"(\d+)", block.style.name)
if level_match:
@@ -892,25 +894,25 @@ class Pdf(PdfParser):
first_start = start
callback(msg="OCR started")
self.__images__(filename if binary is None else binary, zoomin, from_page, to_page, callback)
callback(msg="OCR finished ({:.2f}s)".format(timer() - start))
logging.info("OCR({}~{}): {:.2f}s".format(from_page, to_page, timer() - start))
callback(msg=f"OCR finished ({timer() - start:.2f}s)")
logging.info(f"OCR({from_page}~{to_page}): {timer() - start:.2f}s")
start = timer()
self._layouts_rec(zoomin)
callback(0.63, "Layout analysis ({:.2f}s)".format(timer() - start))
callback(0.63, f"Layout analysis ({timer() - start:.2f}s)")
start = timer()
self._table_transformer_job(zoomin)
callback(0.65, "Table analysis ({:.2f}s)".format(timer() - start))
callback(0.65, f"Table analysis ({timer() - start:.2f}s)")
start = timer()
self._text_merge(zoomin=zoomin)
callback(0.67, "Text merged ({:.2f}s)".format(timer() - start))
callback(0.67, f"Text merged ({timer() - start:.2f}s)")
if separate_tables_figures:
tbls, figures = self._extract_table_figure(True, zoomin, True, True, True)
self._concat_downward()
logging.info("layouts cost: {}s".format(timer() - first_start))
logging.info(f"layouts cost: {timer() - first_start}s")
return [(b["text"], self._line_tag(b, zoomin)) for b in self.boxes], tbls, figures
else:
tbls = self._extract_table_figure(True, zoomin, True, True)
@@ -918,7 +920,7 @@ class Pdf(PdfParser):
self._concat_downward()
# self._final_reading_order_merge()
# self._filter_forpages()
logging.info("layouts cost: {}s".format(timer() - first_start))
logging.info(f"layouts cost: {timer() - first_start}s")
return [(b["text"], self._line_tag(b, zoomin)) for b in self.boxes], tbls
@@ -931,9 +933,9 @@ class Markdown(MarkdownParser):
def md_to_html(self, sections):
if not sections:
return []
if isinstance(sections, type("")):
if isinstance(sections, str):
text = sections
elif isinstance(sections[0], type("")):
elif isinstance(sections[0], str):
text = sections[0]
else:
return []
@@ -990,16 +992,16 @@ class Markdown(MarkdownParser):
urls.append({"url": src, "line": line_no})
seen.add((src, line_no))
except Exception as e:
logging.error("Failed to extract image urls: {}".format(e))
pass
logging.error(f"Failed to extract image urls: {e}")
return urls
def load_images_from_urls(self, urls, cache=None):
import requests
from pathlib import Path
from urllib.parse import urljoin
import requests
from common.ssrf_guard import assert_url_is_safe, pin_dns
cache = cache or {}
@@ -1207,7 +1209,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
st = timer()
res.extend(doc_tokenize_chunks_with_images(chunks, doc, is_english, child_delimiters_pattern=child_deli, language=lang))
logging.info("naive_merge({}): {}".format(filename, timer() - st))
logging.info(f"naive_merge({filename}): {timer() - st}")
res.extend(embed_res)
res.extend(url_res)
return res
@@ -1507,7 +1509,7 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
sub_url_res = chunk(f"{index}.html", html_bytes, callback=callback, lang=lang, is_root=False, **kwargs)
url_res.extend(sub_url_res)
logging.info("naive_merge({}): {}".format(filename, timer() - st))
logging.info(f"naive_merge({filename}): {timer() - st}")
if embed_res:
res.extend(embed_res)
+20 -21
View File
@@ -14,23 +14,22 @@
# limitations under the License.
#
import csv
import logging
import re
import csv
from copy import deepcopy
from io import BytesIO
from timeit import default_timer as timer
from docx import Document
from markdown import markdown
from openpyxl import load_workbook
from common.constants import MAXIMUM_PAGE_NUMBER
from deepdoc.parser.utils import get_text
from rag.nlp import is_english, random_choices, qbullets_category, add_positions, has_qbullet, docx_question_level
from rag.nlp import rag_tokenizer, tokenize_table, concat_img
from deepdoc.parser import PdfParser, ExcelParser, DocxParser
from docx import Document
from markdown import markdown
from common.float_utils import get_float
from deepdoc.parser import DocxParser, ExcelParser, PdfParser
from deepdoc.parser.utils import get_text
from rag.nlp import add_positions, concat_img, docx_question_level, has_qbullet, is_english, qbullets_category, rag_tokenizer, random_choices, tokenize_table
class Excel(ExcelParser):
@@ -64,9 +63,9 @@ class Excel(ExcelParser):
else:
fails.append(str(i + 1))
if len(res) % 999 == 0:
callback(len(res) * 0.6 / total, ("Extract pairs: {}".format(len(res)) + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
callback(len(res) * 0.6 / total, (f"Extract pairs: {len(res)}" + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
callback(0.6, ("Extract pairs: {}. ".format(len(res)) + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
callback(0.6, (f"Extract pairs: {len(res)}. " + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
self.is_english = is_english([rmPrefix(q) for q, _ in random_choices(res, k=30) if len(q) > 1])
return res
@@ -76,24 +75,24 @@ class Pdf(PdfParser):
start = timer()
callback(msg="OCR started")
self.__images__(filename if binary is None else binary, zoomin, from_page, to_page, callback)
callback(msg="OCR finished ({:.2f}s)".format(timer() - start))
logging.debug("OCR({}~{}): {:.2f}s".format(from_page, to_page, timer() - start))
callback(msg=f"OCR finished ({timer() - start:.2f}s)")
logging.debug(f"OCR({from_page}~{to_page}): {timer() - start:.2f}s")
start = timer()
self._layouts_rec(zoomin, drop=False)
callback(0.63, "Layout analysis ({:.2f}s)".format(timer() - start))
callback(0.63, f"Layout analysis ({timer() - start:.2f}s)")
start = timer()
self._table_transformer_job(zoomin)
callback(0.65, "Table analysis ({:.2f}s)".format(timer() - start))
callback(0.65, f"Table analysis ({timer() - start:.2f}s)")
start = timer()
self._text_merge()
callback(0.67, "Text merged ({:.2f}s)".format(timer() - start))
callback(0.67, f"Text merged ({timer() - start:.2f}s)")
tbls = self._extract_table_figure(True, zoomin, True, True)
# self._naive_vertical_merge()
# self._concat_downward()
# self._filter_forpages()
logging.debug("layouts: {}".format(timer() - start))
logging.debug(f"layouts: {timer() - start}")
sections = [b["text"] for b in self.boxes]
bull_x0_list = []
q_bull, reg = qbullets_category(sections)
@@ -168,7 +167,7 @@ class Pdf(PdfParser):
tbl_right = tbls[tbl_index][1][0][2]
tbl_top = tbls[tbl_index][1][0][3]
tbl_bottom = tbls[tbl_index][1][0][4]
tbl_tag = "@@{}\t{:.1f}\t{:.1f}\t{:.1f}\t{:.1f}##".format(tbl_pn, tbl_left, tbl_right, tbl_top, tbl_bottom)
tbl_tag = f"@@{tbl_pn}\t{tbl_left:.1f}\t{tbl_right:.1f}\t{tbl_top:.1f}\t{tbl_bottom:.1f}##"
_tbl_text = "".join(tbls[tbl_index][0][1])
return tbl_pn, tbl_left, tbl_right, tbl_top, tbl_bottom, tbl_tag, _tbl_text
@@ -343,12 +342,12 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
question, answer = arr
i += 1
if len(res) % 999 == 0:
callback(len(res) * 0.6 / len(lines), ("Extract Q&A: {}".format(len(res)) + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
callback(len(res) * 0.6 / len(lines), (f"Extract Q&A: {len(res)}" + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
if question:
res.append(beAdoc(deepcopy(doc), question, answer, eng, len(lines)))
callback(0.6, ("Extract Q&A: {}".format(len(res)) + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
callback(0.6, (f"Extract Q&A: {len(res)}" + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
return res
@@ -378,12 +377,12 @@ def chunk(filename, binary=None, from_page=0, to_page=MAXIMUM_PAGE_NUMBER, lang=
res.append(beAdoc(deepcopy(doc), question, answer, eng, i))
question, answer = row
if len(res) % 999 == 0:
callback(len(res) * 0.6 / len(lines), ("Extract Q&A: {}".format(len(res)) + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
callback(len(res) * 0.6 / len(lines), (f"Extract Q&A: {len(res)}" + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
if question:
res.append(beAdoc(deepcopy(doc), question, answer, eng, len(lines)))
callback(0.6, ("Extract Q&A: {}".format(len(res)) + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
callback(0.6, (f"Extract Q&A: {len(res)}" + (f"{len(fails)} failure, line: %s..." % (",".join(fails[:3])) if fails else "")))
return res
elif re.search(r"\.pdf$", filename, re.IGNORECASE):