mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-10-12 02:16:39 +00:00
refactor(go): default Go cache/queue backend to Kvrocks and fail fast when down (#19804)
This commit is contained in:
@@ -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
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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 }
|
||||
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+1
-1
@@ -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"
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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 {
|
||||
|
||||
@@ -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{
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
})
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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")
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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):
|
||||
|
||||
Reference in New Issue
Block a user