feat: support minio storage

This commit is contained in:
wizardchen
2025-08-13 21:56:20 +08:00
committed by lyingbug
parent 621d9aad37
commit 8b43931886
29 changed files with 2735 additions and 1494 deletions
+5 -11
View File
@@ -7,8 +7,7 @@
# 可选值: debug(开发模式,有详细日志), release(生产模式)
GIN_MODE=debug
# ollama
# Ollama 服务的基准 URL,用于连接本地运行的 Ollama LLM 服务
# Ollama 服务的基准 URL,用于连接本地/其他服务器上运行的 Ollama 服务
OLLAMA_BASE_URL=http://host.docker.internal:11435
# 存储配置
@@ -25,9 +24,6 @@ STORAGE_TYPE=local
STREAM_MANAGER_TYPE=redis
# 主数据库配置
# 数据库地址,可以是主机名或IP地址
DB_HOST=postgres
# 数据库端口,默认为5432
DB_PORT=5432
@@ -41,9 +37,6 @@ DB_PASSWORD=postgres123!@#
DB_NAME=WeKnora
# 如果使用 redis 作为流处理后端,需要配置以下参数
# Redis地址,可以是主机名或IP地址
REDIS_HOST=redis
# Redis端口,默认为6379
REDIS_PORT=6379
@@ -64,6 +57,10 @@ TENANT_AES_KEY=weknorarag-api-key-secret-secret
# 是否开启知识图谱构建和检索(构建阶段需调用大模型,耗时较长)
ENABLE_GRAPH_RAG=false
MINIO_PORT=9000
MINIO_CONSOLE_PORT=9001
# 如果使用ElasticSearch作为向量存储,需要配置以下参数
# ElasticSearch地址,例如 http://localhost:9200
# ELASTICSEARCH_ADDR=your_elasticsearch_addr
@@ -78,9 +75,6 @@ ENABLE_GRAPH_RAG=false
# ELASTICSEARCH_INDEX=WeKnora
# 如果使用MinIO作为文件存储,需要配置以下参数
# MinIO服务器地址,例如 minio:9000
# MINIO_ENDPOINT=minio:9000
# MinIO访问密钥
# MINIO_ACCESS_KEY_ID=your_minio_access_key
+6
View File
@@ -90,5 +90,11 @@ clean-db:
@if [ $$(docker volume ls -q -f name=weknora_postgres-data) ]; then \
docker volume rm weknora_postgres-data; \
fi
@if [ $$(docker volume ls -q -f name=weknora_minio_data) ]; then \
docker volume rm weknora_minio_data; \
fi
@if [ $$(docker volume ls -q -f name=weknora_redis_data) ]; then \
docker volume rm weknora_redis_data; \
fi
+38 -12
View File
@@ -1,5 +1,3 @@
version: "3.8"
services:
app:
build:
@@ -13,9 +11,9 @@ services:
- ./config:/app/config
environment:
- GIN_MODE=${GIN_MODE}
- DB_DRIVER=${DB_DRIVER}
- DB_HOST=${DB_HOST}
- DB_PORT=${DB_PORT}
- DB_DRIVER=postgres
- DB_HOST=postgres
- DB_PORT=5432
- DB_USER=${DB_USER}
- DB_PASSWORD=${DB_PASSWORD}
- DB_NAME=${DB_NAME}
@@ -34,13 +32,13 @@ services:
- DOCREADER_ADDR=docreader:50051
- STORAGE_TYPE=${STORAGE_TYPE}
- LOCAL_STORAGE_BASE_DIR=${LOCAL_STORAGE_BASE_DIR}
- MINIO_ENDPOINT=${MINIO_ENDPOINT}
- MINIO_ACCESS_KEY_ID=${MINIO_ACCESS_KEY_ID}
- MINIO_SECRET_ACCESS_KEY=${MINIO_SECRET_ACCESS_KEY}
- MINIO_ENDPOINT=minio:9000
- MINIO_ACCESS_KEY_ID=${MINIO_ACCESS_KEY_ID:-minioadmin}
- MINIO_SECRET_ACCESS_KEY=${MINIO_SECRET_ACCESS_KEY:-minioadmin}
- MINIO_BUCKET_NAME=${MINIO_BUCKET_NAME}
- OLLAMA_BASE_URL=${OLLAMA_BASE_URL}
- OLLAMA_BASE_URL=${OLLAMA_BASE_URL:-http://host.docker.internal:11434}
- STREAM_MANAGER_TYPE=${STREAM_MANAGER_TYPE}
- REDIS_ADDR=${REDIS_HOST}:${REDIS_PORT}
- REDIS_ADDR=redis:6379
- REDIS_PASSWORD=${REDIS_PASSWORD}
- REDIS_DB=${REDIS_DB}
- REDIS_PREFIX=${REDIS_PREFIX}
@@ -57,19 +55,39 @@ services:
- INIT_RERANK_MODEL_NAME=${INIT_RERANK_MODEL_NAME}
- INIT_RERANK_MODEL_BASE_URL=${INIT_RERANK_MODEL_BASE_URL}
- INIT_RERANK_MODEL_API_KEY=${INIT_RERANK_MODEL_API_KEY}
- INIT_TEST_TENANT_ID=${INIT_TEST_TENANT_ID}
- INIT_TEST_KNOWLEDGE_BASE_ID=${INIT_TEST_KNOWLEDGE_BASE_ID}
depends_on:
redis:
condition: service_started
postgres:
condition: service_healthy
minio:
condition: service_started
networks:
- WeKnora-network
restart: unless-stopped
extra_hosts:
- "host.docker.internal:host-gateway"
minio:
image: minio/minio:latest
container_name: WeKnora-minio
ports:
- "${MINIO_PORT:-9000}:9000"
- "${MINIO_CONSOLE_PORT:-9001}:9001"
environment:
- MINIO_ROOT_USER=${MINIO_ACCESS_KEY_ID:-minioadmin}
- MINIO_ROOT_PASSWORD=${MINIO_SECRET_ACCESS_KEY:-minioadmin}
command: server --console-address ":9001" /data
volumes:
- minio_data:/data
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"]
interval: 30s
timeout: 20s
retries: 3
networks:
- WeKnora-network
frontend:
build:
context: ./frontend
@@ -103,6 +121,13 @@ services:
- VLM_MODEL_BASE_URL=${VLM_MODEL_BASE_URL}
- VLM_MODEL_NAME=${VLM_MODEL_NAME}
- VLM_MODEL_API_KEY=${VLM_MODEL_API_KEY}
- STORAGE_TYPE=${STORAGE_TYPE}
- MINIO_PUBLIC_ENDPOINT=http://localhost:${MINIO_PORT:-9000}
- MINIO_ENDPOINT=minio:9000
- MINIO_ACCESS_KEY_ID=${MINIO_ACCESS_KEY_ID:-minioadmin}
- MINIO_SECRET_ACCESS_KEY=${MINIO_SECRET_ACCESS_KEY:-minioadmin}
- MINIO_BUCKET_NAME=${MINIO_BUCKET_NAME}
- MINIO_USE_SSL=${MINIO_USE_SSL}
networks:
- WeKnora-network
restart: unless-stopped
@@ -175,3 +200,4 @@ volumes:
data-files:
jaeger_data:
redis_data:
minio_data:
+2 -2
View File
@@ -99,8 +99,8 @@ RUN mkdir -p /tmp/libreoffice && \
cd LibreOffice_25.2.4*_Linux_x86-64_deb/DEBS && \
dpkg -i *.deb; \
elif [ "$(uname -m)" = "aarch64" ] || [ "$(uname -m)" = "arm64" ]; then \
wget -q https://mirrors.tuna.tsinghua.edu.cn/libreoffice/libreoffice/testing/25.8.0/deb/aarch64/LibreOffice_25.8.0.2_Linux_aarch64_deb.tar.gz && \
tar -xzf LibreOffice_25.8.0.2_Linux_aarch64_deb.tar.gz && \
wget -q https://mirrors.aliyun.com/libreoffice/testing/25.8.0/deb/aarch64/LibreOffice_25.8.0.3_Linux_aarch64_deb.tar.gz && \
tar -xzf LibreOffice_25.8.0.3_Linux_aarch64_deb.tar.gz && \
cd LibreOffice_25.8.0*_Linux_aarch64_deb/DEBS && \
dpkg -i *.deb; \
else \
+46 -13
View File
@@ -22,6 +22,7 @@ export interface InitializationConfig {
};
multimodal: {
enabled: boolean;
storageType: 'cos' | 'minio';
vlm?: {
modelName: string;
baseUrl: string;
@@ -36,12 +37,18 @@ export interface InitializationConfig {
appId: string;
pathPrefix?: string;
};
minio?: {
bucketName: string;
pathPrefix?: string;
};
};
documentSplitting: {
chunkSize: number;
chunkOverlap: number;
separators: string[];
};
// Frontend-only hint for storage selection UI
storageType?: 'cos' | 'minio';
}
// 下载任务状态类型
@@ -88,7 +95,7 @@ export function initializeSystem(config: InitializationConfig): Promise<any> {
}
// 检查Ollama服务状态
export function checkOllamaStatus(): Promise<{ available: boolean; version?: string; error?: string }> {
export function checkOllamaStatus(): Promise<{ available: boolean; version?: string; error?: string; baseUrl?: string }> {
return new Promise((resolve, reject) => {
get('/api/v1/initialization/ollama/status')
.then((response: any) => {
@@ -101,6 +108,20 @@ export function checkOllamaStatus(): Promise<{ available: boolean; version?: str
});
}
// 列出已安装的 Ollama 模型
export function listOllamaModels(): Promise<string[]> {
return new Promise((resolve, reject) => {
get('/api/v1/initialization/ollama/models')
.then((response: any) => {
resolve((response.data && response.data.models) || []);
})
.catch((error: any) => {
console.error('获取 Ollama 模型列表失败:', error);
resolve([]);
});
});
}
// 检查Ollama模型状态
export function checkOllamaModels(models: string[]): Promise<{ models: Record<string, boolean> }> {
return new Promise((resolve, reject) => {
@@ -218,12 +239,17 @@ export function testMultimodalFunction(testData: {
vlm_base_url: string;
vlm_api_key?: string;
vlm_interface_type?: string;
cos_secret_id: string;
cos_secret_key: string;
cos_region: string;
cos_bucket_name: string;
cos_app_id: string;
storage_type?: 'cos'|'minio';
// COS optional fields (required only when storage_type === 'cos')
cos_secret_id?: string;
cos_secret_key?: string;
cos_region?: string;
cos_bucket_name?: string;
cos_app_id?: string;
cos_path_prefix?: string;
// MinIO optional fields
minio_bucket_name?: string;
minio_path_prefix?: string;
chunk_size: number;
chunk_overlap: number;
separators: string[];
@@ -245,14 +271,21 @@ export function testMultimodalFunction(testData: {
if (testData.vlm_interface_type) {
formData.append('vlm_interface_type', testData.vlm_interface_type);
}
formData.append('cos_secret_id', testData.cos_secret_id);
formData.append('cos_secret_key', testData.cos_secret_key);
formData.append('cos_region', testData.cos_region);
formData.append('cos_bucket_name', testData.cos_bucket_name);
formData.append('cos_app_id', testData.cos_app_id);
if (testData.cos_path_prefix) {
formData.append('cos_path_prefix', testData.cos_path_prefix);
if (testData.storage_type) {
formData.append('storage_type', testData.storage_type);
}
// Append COS fields only when storage_type is COS
if (testData.storage_type === 'cos') {
if (testData.cos_secret_id) formData.append('cos_secret_id', testData.cos_secret_id);
if (testData.cos_secret_key) formData.append('cos_secret_key', testData.cos_secret_key);
if (testData.cos_region) formData.append('cos_region', testData.cos_region);
if (testData.cos_bucket_name) formData.append('cos_bucket_name', testData.cos_bucket_name);
if (testData.cos_app_id) formData.append('cos_app_id', testData.cos_app_id);
if (testData.cos_path_prefix) formData.append('cos_path_prefix', testData.cos_path_prefix);
}
// MinIO fields
if (testData.minio_bucket_name) formData.append('minio_bucket_name', testData.minio_bucket_name);
if (testData.minio_path_prefix) formData.append('minio_path_prefix', testData.minio_path_prefix);
formData.append('chunk_size', testData.chunk_size.toString());
formData.append('chunk_overlap', testData.chunk_overlap.toString());
formData.append('separators', JSON.stringify(testData.separators));
@@ -158,7 +158,7 @@ function splitMarkdownByImages(markdown) {
function isLink(str) {
const trimmedStr = str.trim();
// 正则表达式匹配常见链接格式
const urlPattern = /^(https?:\/\/|ftp:\/\/|www\.)[\w-]+(\.[\w-]+)+([\w.,@?^=%&:/~+#-]*[\w@?^=%&/~+#-])?/i;
const urlPattern = /^(https?:\/\/|ftp:\/\/|www\.)(?:(?:[\w-]+(?:\.[\w-]+)*)|(?:\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})|(?:\[[a-fA-F0-9:]+\]))(?::\d{1,5})?(?:[\/\w.,@?^=%&:~+#-]*[\w@?^=%&\/~+#-])?/i;
return urlPattern.test(trimmedStr);
}
File diff suppressed because it is too large Load Diff
+31 -20
View File
@@ -112,11 +112,19 @@ func (s *knowledgeService) CreateKnowledgeFromFile(ctx context.Context,
logger.Info(ctx, "Non-image file with multimodal enabled, skipping COS/VLM validation")
} else {
// 检查COS配置
if kb.COSConfig.SecretID == "" || kb.COSConfig.SecretKey == "" ||
kb.COSConfig.Region == "" || kb.COSConfig.BucketName == "" ||
kb.COSConfig.AppID == "" {
logger.Error(ctx, "COS configuration incomplete for image multimodal processing")
return nil, werrors.NewBadRequestError("上传图片文件需要完整的COS配置信息, 请前往系统设置页面进行补全")
switch kb.StorageConfig.Provider {
case "cos":
if kb.StorageConfig.SecretID == "" || kb.StorageConfig.SecretKey == "" ||
kb.StorageConfig.Region == "" || kb.StorageConfig.BucketName == "" ||
kb.StorageConfig.AppID == "" {
logger.Error(ctx, "COS configuration incomplete for image multimodal processing")
return nil, werrors.NewBadRequestError("上传图片文件需要完整的对象存储配置信息, 请前往系统设置页面进行补全")
}
case "minio":
if kb.StorageConfig.BucketName == "" {
logger.Error(ctx, "MinIO configuration incomplete for image multimodal processing")
return nil, werrors.NewBadRequestError("上传图片文件需要完整的对象存储配置信息, 请前往系统设置页面进行补全")
}
}
// 检查VLM配置
@@ -317,7 +325,8 @@ func (s *knowledgeService) CreateKnowledgeFromURL(ctx context.Context,
if enableMultimodel == nil {
enableMultimodel = &kb.ChunkingConfig.EnableMultimodal
}
go s.processDocumentFromURL(ctx, kb, knowledge, url, *enableMultimodel)
newCtx := logger.CloneContext(ctx)
go s.processDocumentFromURL(newCtx, kb, knowledge, url, *enableMultimodel)
logger.Infof(ctx, "Knowledge from URL created successfully, ID: %s", knowledge.ID)
return knowledge, nil
@@ -670,13 +679,14 @@ func (s *knowledgeService) processDocument(ctx context.Context,
ChunkOverlap: int32(kb.ChunkingConfig.ChunkOverlap),
Separators: kb.ChunkingConfig.Separators,
EnableMultimodal: enableMultimodel,
CosConfig: &proto.COSConfig{
SecretId: kb.COSConfig.SecretID,
SecretKey: kb.COSConfig.SecretKey,
Region: kb.COSConfig.Region,
BucketName: kb.COSConfig.BucketName,
AppId: kb.COSConfig.AppID,
PathPrefix: kb.COSConfig.PathPrefix,
StorageConfig: &proto.StorageConfig{
Provider: proto.StorageProvider(proto.StorageProvider_value[strings.ToUpper(kb.StorageConfig.Provider)]),
Region: kb.StorageConfig.Region,
BucketName: kb.StorageConfig.BucketName,
AccessKeyId: kb.StorageConfig.SecretID,
SecretAccessKey: kb.StorageConfig.SecretKey,
AppId: kb.StorageConfig.AppID,
PathPrefix: kb.StorageConfig.PathPrefix,
},
VlmConfig: &proto.VLMConfig{
ModelName: kb.VLMConfig.ModelName,
@@ -724,13 +734,14 @@ func (s *knowledgeService) processDocumentFromURL(ctx context.Context,
ChunkOverlap: int32(kb.ChunkingConfig.ChunkOverlap),
Separators: kb.ChunkingConfig.Separators,
EnableMultimodal: enableMultimodel,
CosConfig: &proto.COSConfig{
SecretId: kb.COSConfig.SecretID,
SecretKey: kb.COSConfig.SecretKey,
Region: kb.COSConfig.Region,
BucketName: kb.COSConfig.BucketName,
AppId: kb.COSConfig.AppID,
PathPrefix: kb.COSConfig.PathPrefix,
StorageConfig: &proto.StorageConfig{
Provider: proto.StorageProvider(proto.StorageProvider_value[strings.ToUpper(kb.StorageConfig.Provider)]),
Region: kb.StorageConfig.Region,
BucketName: kb.StorageConfig.BucketName,
AccessKeyId: kb.StorageConfig.SecretID,
SecretAccessKey: kb.StorageConfig.SecretKey,
AppId: kb.StorageConfig.AppID,
PathPrefix: kb.StorageConfig.PathPrefix,
},
VlmConfig: &proto.VLMConfig{
ModelName: kb.VLMConfig.ModelName,
@@ -249,7 +249,7 @@ func (s *knowledgeBaseService) CopyKnowledgeBase(ctx context.Context,
SummaryModelID: sourceKB.SummaryModelID,
RerankModelID: sourceKB.RerankModelID,
VLMModelID: sourceKB.VLMModelID,
COSConfig: sourceKB.COSConfig,
StorageConfig: sourceKB.StorageConfig,
}
if err := s.repo.CreateKnowledgeBase(ctx, targetKB); err != nil {
return nil, nil, err
+241 -167
View File
@@ -80,7 +80,9 @@ func NewInitializationHandler(
// InitializationRequest 初始化请求结构
type InitializationRequest struct {
LLM struct {
// 前端传入的存储类型:cos 或 minio
StorageType string `json:"storageType"`
LLM struct {
Source string `json:"source" binding:"required"`
ModelName string `json:"modelName" binding:"required"`
BaseURL string `json:"baseUrl"`
@@ -118,6 +120,10 @@ type InitializationRequest struct {
AppID string `json:"appId"`
PathPrefix string `json:"pathPrefix"`
} `json:"cos,omitempty"`
Minio *struct {
BucketName string `json:"bucketName"`
PathPrefix string `json:"pathPrefix"`
} `json:"minio,omitempty"`
} `json:"multimodal"`
DocumentSplitting struct {
@@ -195,9 +201,10 @@ func (h *InitializationHandler) Initialize(c *gin.Context) {
// 验证多模态配置
if req.Multimodal.Enabled {
if req.Multimodal.VLM == nil || req.Multimodal.COS == nil {
logger.Error(ctx, "Multimodal enabled but missing VLM or COS configuration")
c.Error(errors.NewBadRequestError("启用多模态时需要配置VLM和COS信息"))
storageType := strings.ToLower(req.StorageType)
if req.Multimodal.VLM == nil {
logger.Error(ctx, "Multimodal enabled but missing VLM configuration")
c.Error(errors.NewBadRequestError("启用多模态时需要配置VLM信息"))
return
}
if req.Multimodal.VLM.InterfaceType == "ollama" {
@@ -208,12 +215,22 @@ func (h *InitializationHandler) Initialize(c *gin.Context) {
c.Error(errors.NewBadRequestError("VLM配置不完整"))
return
}
if req.Multimodal.COS.SecretID == "" || req.Multimodal.COS.SecretKey == "" ||
req.Multimodal.COS.Region == "" || req.Multimodal.COS.BucketName == "" ||
req.Multimodal.COS.AppID == "" {
logger.Error(ctx, "COS configuration incomplete")
c.Error(errors.NewBadRequestError("COS配置不完整"))
return
switch storageType {
case "cos":
if req.Multimodal.COS == nil || req.Multimodal.COS.SecretID == "" || req.Multimodal.COS.SecretKey == "" ||
req.Multimodal.COS.Region == "" || req.Multimodal.COS.BucketName == "" ||
req.Multimodal.COS.AppID == "" {
logger.Error(ctx, "COS configuration incomplete")
c.Error(errors.NewBadRequestError("COS配置不完整"))
return
}
case "minio":
if req.Multimodal.Minio == nil || req.Multimodal.Minio.BucketName == "" ||
os.Getenv("MINIO_ACCESS_KEY_ID") == "" || os.Getenv("MINIO_SECRET_ACCESS_KEY") == "" {
logger.Error(ctx, "MinIO configuration incomplete")
c.Error(errors.NewBadRequestError("MinIO配置不完整"))
return
}
}
}
@@ -520,14 +537,30 @@ func (h *InitializationHandler) Initialize(c *gin.Context) {
APIKey: req.Multimodal.VLM.APIKey,
InterfaceType: req.Multimodal.VLM.InterfaceType,
},
COSConfig: types.COSConfig{
SecretID: req.Multimodal.COS.SecretID,
SecretKey: req.Multimodal.COS.SecretKey,
Region: req.Multimodal.COS.Region,
BucketName: req.Multimodal.COS.BucketName,
AppID: req.Multimodal.COS.AppID,
PathPrefix: req.Multimodal.COS.PathPrefix,
},
}
switch req.StorageType {
case "cos":
if req.Multimodal.COS != nil {
kb.StorageConfig = types.StorageConfig{
Provider: req.StorageType,
BucketName: req.Multimodal.COS.BucketName,
AppID: req.Multimodal.COS.AppID,
PathPrefix: req.Multimodal.COS.PathPrefix,
SecretID: req.Multimodal.COS.SecretID,
SecretKey: req.Multimodal.COS.SecretKey,
Region: req.Multimodal.COS.Region,
}
}
case "minio":
if req.Multimodal.Minio != nil {
kb.StorageConfig = types.StorageConfig{
Provider: req.StorageType,
BucketName: req.Multimodal.Minio.BucketName,
PathPrefix: req.Multimodal.Minio.PathPrefix,
SecretID: os.Getenv("MINIO_ACCESS_KEY_ID"),
SecretKey: os.Getenv("MINIO_SECRET_ACCESS_KEY"),
}
}
}
_, err = h.kbService.CreateKnowledgeBase(newCtx, kb)
@@ -562,19 +595,35 @@ func (h *InitializationHandler) Initialize(c *gin.Context) {
APIKey: req.Multimodal.VLM.APIKey,
InterfaceType: req.Multimodal.VLM.InterfaceType,
}
kb.COSConfig = types.COSConfig{
SecretID: req.Multimodal.COS.SecretID,
SecretKey: req.Multimodal.COS.SecretKey,
Region: req.Multimodal.COS.Region,
BucketName: req.Multimodal.COS.BucketName,
AppID: req.Multimodal.COS.AppID,
PathPrefix: req.Multimodal.COS.PathPrefix,
switch req.StorageType {
case "cos":
if req.Multimodal.COS != nil {
kb.StorageConfig = types.StorageConfig{
Provider: req.StorageType,
SecretID: req.Multimodal.COS.SecretID,
SecretKey: req.Multimodal.COS.SecretKey,
Region: req.Multimodal.COS.Region,
BucketName: req.Multimodal.COS.BucketName,
AppID: req.Multimodal.COS.AppID,
PathPrefix: req.Multimodal.COS.PathPrefix,
}
}
case "minio":
if req.Multimodal.Minio != nil {
kb.StorageConfig = types.StorageConfig{
Provider: req.StorageType,
BucketName: req.Multimodal.Minio.BucketName,
PathPrefix: req.Multimodal.Minio.PathPrefix,
SecretID: os.Getenv("MINIO_ACCESS_KEY_ID"),
SecretKey: os.Getenv("MINIO_SECRET_ACCESS_KEY"),
}
}
}
} else {
kb.VLMModelID = "" // 清空VLM模型ID
// 清空VLM配置
kb.VLMConfig = types.VLMConfig{}
kb.COSConfig = types.COSConfig{}
kb.StorageConfig = types.StorageConfig{}
}
if !hasFiles {
kb.EmbeddingModelID = embeddingModelID
@@ -642,6 +691,12 @@ func (h *InitializationHandler) CheckOllamaStatus(c *gin.Context) {
logger.Info(ctx, "Checking Ollama service status")
// Determine Ollama base URL for display
baseURL := os.Getenv("OLLAMA_BASE_URL")
if baseURL == "" {
baseURL = "http://host.docker.internal:11434"
}
// 检查Ollama服务是否可用
err := h.ollamaService.StartService(ctx)
if err != nil {
@@ -651,6 +706,7 @@ func (h *InitializationHandler) CheckOllamaStatus(c *gin.Context) {
"data": gin.H{
"available": false,
"error": err.Error(),
"baseUrl": baseURL,
},
})
return
@@ -668,6 +724,7 @@ func (h *InitializationHandler) CheckOllamaStatus(c *gin.Context) {
"data": gin.H{
"available": h.ollamaService.IsAvailable(),
"version": version,
"baseUrl": baseURL,
},
})
}
@@ -702,7 +759,11 @@ func (h *InitializationHandler) CheckOllamaModels(c *gin.Context) {
// 检查每个模型是否存在
for _, modelName := range req.Models {
available, err := h.ollamaService.IsModelAvailable(ctx, modelName)
checkModelName := modelName
if !strings.Contains(modelName, ":") {
checkModelName = modelName + ":latest"
}
available, err := h.ollamaService.IsModelAvailable(ctx, checkModelName)
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"model_name": modelName,
@@ -867,6 +928,36 @@ func (h *InitializationHandler) ListDownloadTasks(c *gin.Context) {
})
}
// ListOllamaModels 列出已安装的 Ollama 模型
func (h *InitializationHandler) ListOllamaModels(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Listing installed Ollama models")
// 确保服务可用
if !h.ollamaService.IsAvailable() {
if err := h.ollamaService.StartService(ctx); err != nil {
logger.ErrorWithFields(ctx, err, nil)
c.Error(errors.NewInternalServerError("Ollama服务不可用: " + err.Error()))
return
}
}
models, err := h.ollamaService.ListModels(ctx)
if err != nil {
logger.ErrorWithFields(ctx, err, nil)
c.Error(errors.NewInternalServerError("获取模型列表失败: " + err.Error()))
return
}
c.JSON(http.StatusOK, gin.H{
"success": true,
"data": gin.H{
"models": models,
},
})
}
// downloadModelAsync 异步下载模型
func (h *InitializationHandler) downloadModelAsync(ctx context.Context,
taskID, modelName string,
@@ -974,20 +1065,6 @@ func (h *InitializationHandler) updateTaskStatus(
}
}
// 清理过期任务 (可以在后台定期执行)
func (h *InitializationHandler) cleanupExpiredTasks() {
tasksMutex.Lock()
defer tasksMutex.Unlock()
cutoff := time.Now().Add(-24 * time.Hour) // 保留24小时内的任务
for id, task := range downloadTasks {
if task.EndTime != nil && task.EndTime.Before(cutoff) {
delete(downloadTasks, id)
}
}
}
// GetCurrentConfig 获取当前系统配置信息
func (h *InitializationHandler) GetCurrentConfig(c *gin.Context) {
ctx := c.Request.Context()
@@ -1109,20 +1186,29 @@ func buildConfigResponse(models []*types.Model,
}
// 添加多模态的COS配置信息
if kb.COSConfig.SecretID != "" {
if kb.StorageConfig.SecretID != "" {
if config["multimodal"] == nil {
config["multimodal"] = map[string]interface{}{
"enabled": true,
}
}
multimodal := config["multimodal"].(map[string]interface{})
multimodal["cos"] = map[string]interface{}{
"secretId": kb.COSConfig.SecretID,
"secretKey": kb.COSConfig.SecretKey,
"region": kb.COSConfig.Region,
"bucketName": kb.COSConfig.BucketName,
"appId": kb.COSConfig.AppID,
"pathPrefix": kb.COSConfig.PathPrefix,
multimodal["storageType"] = kb.StorageConfig.Provider
switch kb.StorageConfig.Provider {
case "cos":
multimodal["cos"] = map[string]interface{}{
"secretId": kb.StorageConfig.SecretID,
"secretKey": kb.StorageConfig.SecretKey,
"region": kb.StorageConfig.Region,
"bucketName": kb.StorageConfig.BucketName,
"appId": kb.StorageConfig.AppID,
"pathPrefix": kb.StorageConfig.PathPrefix,
}
case "minio":
multimodal["minio"] = map[string]interface{}{
"bucketName": kb.StorageConfig.BucketName,
"pathPrefix": kb.StorageConfig.PathPrefix,
}
}
}
}
@@ -1434,110 +1520,84 @@ func (h *InitializationHandler) CheckRerankModel(c *gin.Context) {
})
}
// 使用结构体解析表单数据
type testMultimodalForm struct {
VLMModel string `form:"vlm_model"`
VLMBaseURL string `form:"vlm_base_url"`
VLMAPIKey string `form:"vlm_api_key"`
VLMInterfaceType string `form:"vlm_interface_type"`
StorageType string `form:"storage_type"`
// COS 配置
COSSecretID string `form:"cos_secret_id"`
COSSecretKey string `form:"cos_secret_key"`
COSRegion string `form:"cos_region"`
COSBucketName string `form:"cos_bucket_name"`
COSAppID string `form:"cos_app_id"`
COSPathPrefix string `form:"cos_path_prefix"`
// MinIO 配置(当存储为 minio 时)
MinioBucketName string `form:"minio_bucket_name"`
MinioPathPrefix string `form:"minio_path_prefix"`
// 文档切分配置(字符串后续自行解析,以避免类型绑定失败)
ChunkSize string `form:"chunk_size"`
ChunkOverlap string `form:"chunk_overlap"`
SeparatorsRaw string `form:"separators"`
}
// TestMultimodalFunction 测试多模态功能
func (h *InitializationHandler) TestMultimodalFunction(c *gin.Context) {
ctx := c.Request.Context()
logger.Info(ctx, "Testing multimodal functionality")
// 解析表单数据
vlmModel := c.PostForm("vlm_model")
vlmBaseURL := c.PostForm("vlm_base_url")
vlmAPIKey := c.PostForm("vlm_api_key")
vlmInterfaceType := c.PostForm("vlm_interface_type")
if vlmInterfaceType == "ollama" {
vlmBaseURL = os.Getenv("OLLAMA_BASE_URL") + "/v1"
var req testMultimodalForm
if err := c.ShouldBind(&req); err != nil {
logger.Error(ctx, "Failed to parse form data", err)
c.Error(errors.NewBadRequestError("表单参数解析失败"))
return
}
// ollama 场景自动拼接 base url
if req.VLMInterfaceType == "ollama" {
req.VLMBaseURL = os.Getenv("OLLAMA_BASE_URL") + "/v1"
}
// 如果没有提供VLM配置,尝试从KnowledgeBase获取
if vlmModel == "" || vlmBaseURL == "" {
logger.Info(ctx, "VLM configuration not provided, trying to get from KnowledgeBase")
req.StorageType = strings.ToLower(req.StorageType)
// 获取默认知识库
kb, err := h.kbService.GetKnowledgeBaseByID(ctx, types.InitDefaultKnowledgeBaseID)
if err != nil {
logger.Error(ctx, "Failed to get KnowledgeBase", err)
c.Error(errors.NewBadRequestError("获取知识库配置失败"))
return
}
// 使用知识库中的VLM配置
if kb.VLMConfig.ModelName != "" && kb.VLMConfig.BaseURL != "" {
vlmModel = kb.VLMConfig.ModelName
vlmBaseURL = kb.VLMConfig.BaseURL
vlmAPIKey = kb.VLMConfig.APIKey
vlmInterfaceType = kb.VLMConfig.InterfaceType
logger.Infof(ctx, "Using VLM config from KnowledgeBase: Model=%s, URL=%s, Type=%s",
vlmModel, vlmBaseURL, vlmInterfaceType)
} else {
logger.Error(ctx, "VLM configuration not found in KnowledgeBase")
c.Error(errors.NewBadRequestError("知识库中未找到VLM配置信息"))
return
}
}
// COS配置
cosSecretID := c.PostForm("cos_secret_id")
cosSecretKey := c.PostForm("cos_secret_key")
cosRegion := c.PostForm("cos_region")
cosBucketName := c.PostForm("cos_bucket_name")
cosAppID := c.PostForm("cos_app_id")
cosPathPrefix := c.PostForm("cos_path_prefix")
// 如果没有提供COS配置,尝试从KnowledgeBase获取
if cosSecretID == "" || cosSecretKey == "" ||
cosRegion == "" || cosBucketName == "" || cosAppID == "" {
logger.Info(ctx, "COS configuration not provided, trying to get from KnowledgeBase")
// 获取默认知识库
kb, err := h.kbService.GetKnowledgeBaseByID(ctx, types.InitDefaultKnowledgeBaseID)
if err != nil {
logger.Error(ctx, "Failed to get KnowledgeBase", err)
c.Error(errors.NewBadRequestError("获取知识库配置失败"))
return
}
// 使用知识库中的COS配置
if kb.COSConfig.SecretID != "" && kb.COSConfig.SecretKey != "" {
cosSecretID = kb.COSConfig.SecretID
cosSecretKey = kb.COSConfig.SecretKey
cosRegion = kb.COSConfig.Region
cosBucketName = kb.COSConfig.BucketName
cosAppID = kb.COSConfig.AppID
cosPathPrefix = kb.COSConfig.PathPrefix
logger.Infof(ctx, "Using COS config from KnowledgeBase: Region=%s, Bucket=%s, App=%s",
cosRegion, cosBucketName, cosAppID)
} else {
logger.Error(ctx, "COS configuration not found in KnowledgeBase")
c.Error(errors.NewBadRequestError("知识库中未找到COS配置信息"))
return
}
}
// 文档分割配置
chunkSizeStr := c.PostForm("chunk_size")
chunkOverlapStr := c.PostForm("chunk_overlap")
separatorsStr := c.PostForm("separators")
if vlmModel == "" || vlmBaseURL == "" {
if req.VLMModel == "" || req.VLMBaseURL == "" {
logger.Error(ctx, "VLM model name and base URL are required")
c.Error(errors.NewBadRequestError("VLM模型名称和Base URL不能为空"))
return
}
if cosSecretID == "" || cosSecretKey == "" ||
cosRegion == "" || cosBucketName == "" || cosAppID == "" {
logger.Error(ctx, "COS configuration is required")
c.Error(errors.NewBadRequestError("COS配置信息不能为空"))
switch req.StorageType {
case "cos":
logger.Infof(ctx, "COS config: ID=%s, Region=%s, Bucket=%s, App=%s, Prefix=%s",
req.COSSecretID, req.COSRegion, req.COSBucketName, req.COSAppID, req.COSPathPrefix)
// 必填:SecretID/SecretKey/Region/BucketName/AppIDPathPrefix 可选
if req.COSSecretID == "" || req.COSSecretKey == "" ||
req.COSRegion == "" || req.COSBucketName == "" ||
req.COSAppID == "" {
logger.Error(ctx, "COS configuration is required")
c.Error(errors.NewBadRequestError("COS配置信息不能为空"))
return
}
case "minio":
logger.Infof(ctx, "MinIO config: Bucket=%s, PathPrefix=%s", req.MinioBucketName, req.MinioPathPrefix)
if req.MinioBucketName == "" {
logger.Error(ctx, "MinIO configuration is required")
c.Error(errors.NewBadRequestError("MinIO配置信息不能为空"))
return
}
default:
logger.Error(ctx, "Invalid storage type")
c.Error(errors.NewBadRequestError("无效的存储类型"))
return
}
// 记录COS配置信息用于日志
logger.Infof(ctx, "COS config: ID=%s, Region=%s, Bucket=%s, App=%s, Prefix=%s",
cosSecretID, cosRegion, cosBucketName, cosAppID, cosPathPrefix)
logger.Infof(ctx, "VLM config: Model=%s, URL=%s, HasKey=%v, Type=%s",
vlmModel, vlmBaseURL, vlmAPIKey != "", vlmInterfaceType)
req.VLMModel, req.VLMBaseURL, req.VLMAPIKey != "", req.VLMInterfaceType)
// 获取上传的图片文件
file, header, err := c.Request.FormFile("image")
@@ -1561,27 +1621,26 @@ func (h *InitializationHandler) TestMultimodalFunction(c *gin.Context) {
c.Error(errors.NewBadRequestError("图片文件大小不能超过10MB"))
return
}
logger.Infof(ctx, "Processing image: %s, size: %d bytes", header.Filename, header.Size)
// 解析文档分割配置
chunkSize, err := strconv.Atoi(chunkSizeStr)
chunkSize, err := strconv.Atoi(req.ChunkSize)
if err != nil || chunkSize < 100 || chunkSize > 10000 {
chunkSize = 1000 // 默认值
chunkSize = 1000
}
chunkOverlap, err := strconv.Atoi(chunkOverlapStr)
chunkOverlap, err := strconv.Atoi(req.ChunkOverlap)
if err != nil || chunkOverlap < 0 || chunkOverlap >= chunkSize {
chunkOverlap = 200 // 默认值
chunkOverlap = 200
}
var separators []string
if separatorsStr != "" {
if err := json.Unmarshal([]byte(separatorsStr), &separators); err != nil {
separators = []string{"\n\n", "\n", "。", "", "", ";", ""} // 默认值
if req.SeparatorsRaw != "" {
if err := json.Unmarshal([]byte(req.SeparatorsRaw), &separators); err != nil {
separators = []string{"\n\n", "\n", "。", "", "", ";", ""}
}
} else {
separators = []string{"\n\n", "\n", "。", "", "", ";", ""} // 默认值
separators = []string{"\n\n", "\n", "。", "", "", ";", ""}
}
// 读取图片文件内容
@@ -1594,16 +1653,17 @@ func (h *InitializationHandler) TestMultimodalFunction(c *gin.Context) {
// 调用多模态测试
startTime := time.Now()
result, err := h.testMultimodalWithDocReader(ctx, imageContent, header.Filename,
chunkSize, chunkOverlap, separators,
vlmModel, vlmBaseURL, vlmAPIKey, vlmInterfaceType,
cosSecretID, cosSecretKey, cosRegion, cosBucketName, cosAppID, cosPathPrefix)
result, err := h.testMultimodalWithDocReader(
ctx,
imageContent, header.Filename,
chunkSize, chunkOverlap, separators, &req,
)
processingTime := time.Since(startTime).Milliseconds()
if err != nil {
logger.ErrorWithFields(ctx, err, map[string]interface{}{
"vlm_model": vlmModel,
"vlm_base_url": vlmBaseURL,
"vlm_model": req.VLMModel,
"vlm_base_url": req.VLMBaseURL,
"filename": header.Filename,
})
c.JSON(http.StatusOK, gin.H{
@@ -1631,12 +1691,12 @@ func (h *InitializationHandler) TestMultimodalFunction(c *gin.Context) {
}
// testMultimodalWithDocReader 调用docreader服务进行多模态处理
func (h *InitializationHandler) testMultimodalWithDocReader(ctx context.Context,
func (h *InitializationHandler) testMultimodalWithDocReader(
ctx context.Context,
imageContent []byte, filename string,
chunkSize, chunkOverlap int, separators []string,
vlmModel, vlmBaseURL, vlmAPIKey, vlmInterfaceType,
cosSecretID, cosSecretKey, cosRegion, cosBucketName, cosAppID, cosPathPrefix string) (
map[string]string, error) {
req *testMultimodalForm,
) (map[string]string, error) {
// 获取文件扩展名
fileExt := ""
if idx := strings.LastIndex(filename, "."); idx != -1 {
@@ -1659,23 +1719,37 @@ func (h *InitializationHandler) testMultimodalWithDocReader(ctx context.Context,
Separators: separators,
EnableMultimodal: true, // 启用多模态处理
VlmConfig: &proto.VLMConfig{
ModelName: vlmModel,
BaseUrl: vlmBaseURL,
ApiKey: vlmAPIKey,
InterfaceType: vlmInterfaceType,
},
CosConfig: &proto.COSConfig{
SecretId: cosSecretID,
SecretKey: cosSecretKey,
Region: cosRegion,
BucketName: cosBucketName,
AppId: cosAppID,
PathPrefix: cosPathPrefix,
ModelName: req.VLMModel,
BaseUrl: req.VLMBaseURL,
ApiKey: req.VLMAPIKey,
InterfaceType: req.VLMInterfaceType,
},
},
RequestId: ctx.Value(types.RequestIDContextKey).(string),
}
// 设置对象存储配置(通用)
switch strings.ToLower(req.StorageType) {
case "cos":
request.ReadConfig.StorageConfig = &proto.StorageConfig{
Provider: proto.StorageProvider_COS,
Region: req.COSRegion,
BucketName: req.COSBucketName,
AccessKeyId: req.COSSecretID,
SecretAccessKey: req.COSSecretKey,
AppId: req.COSAppID,
PathPrefix: req.COSPathPrefix,
}
case "minio":
request.ReadConfig.StorageConfig = &proto.StorageConfig{
Provider: proto.StorageProvider_MINIO,
BucketName: req.MinioBucketName,
PathPrefix: req.MinioPathPrefix,
AccessKeyId: os.Getenv("MINIO_ACCESS_KEY_ID"),
SecretAccessKey: os.Getenv("MINIO_SECRET_ACCESS_KEY"),
}
}
// 调用docreader服务
response, err := h.docReaderClient.ReadFromFile(ctx, request)
if err != nil {
+1
View File
@@ -70,6 +70,7 @@ func NewRouter(params RouterParams) *gin.Engine {
// Ollama相关接口(不需要认证)
r.GET("/api/v1/initialization/ollama/status", params.InitializationHandler.CheckOllamaStatus)
r.GET("/api/v1/initialization/ollama/models", params.InitializationHandler.ListOllamaModels)
r.POST("/api/v1/initialization/ollama/models/check", params.InitializationHandler.CheckOllamaModels)
r.POST("/api/v1/initialization/ollama/models/download", params.InitializationHandler.DownloadOllamaModel)
r.GET("/api/v1/initialization/ollama/download/progress/:taskId", params.InitializationHandler.GetDownloadProgress)
+7 -5
View File
@@ -36,8 +36,8 @@ type KnowledgeBase struct {
VLMModelID string `yaml:"vlm_model_id" json:"vlm_model_id"`
// VLM config
VLMConfig VLMConfig `yaml:"vlm_config" json:"vlm_config" gorm:"type:json"`
// COS config
COSConfig COSConfig `yaml:"cos_config" json:"cos_config" gorm:"type:json"`
// Storage config
StorageConfig StorageConfig `yaml:"cos_config" json:"cos_config" gorm:"column:cos_config;type:json"`
// Creation time of the knowledge base
CreatedAt time.Time `yaml:"created_at" json:"created_at"`
// Last updated time of the knowledge base
@@ -67,7 +67,7 @@ type ChunkingConfig struct {
}
// COSConfig represents the COS configuration
type COSConfig struct {
type StorageConfig struct {
// Secret ID
SecretID string `yaml:"secret_id" json:"secret_id"`
// Secret Key
@@ -80,13 +80,15 @@ type COSConfig struct {
AppID string `yaml:"app_id" json:"app_id"`
// Path Prefix
PathPrefix string `yaml:"path_prefix" json:"path_prefix"`
// Provider
Provider string `yaml:"provider" json:"provider"`
}
func (c *COSConfig) Value() (driver.Value, error) {
func (c *StorageConfig) Value() (driver.Value, error) {
return json.Marshal(c)
}
func (c *COSConfig) Scan(value interface{}) error {
func (c *StorageConfig) Scan(value interface{}) error {
if value == nil {
return nil
}
+106 -63
View File
@@ -80,18 +80,20 @@ check_env_file() {
if [ -z "$DB_DRIVER" ]; then missing_vars+=("DB_DRIVER"); fi
if [ -z "$STORAGE_TYPE" ]; then missing_vars+=("STORAGE_TYPE"); fi
if [ ${#missing_vars[@]} -gt 0 ]; then
log_warning "以下环境变量未设置,将使用默认值: ${missing_vars[*]}"
else
log_success "所有必要的环境变量已设置"
fi
return 0
}
# 安装Ollama(根据平台不同采用不同方法)
install_ollama() {
log_info "Ollama未安装,正在安装..."
# 检查是否为远程服务
get_ollama_base_url
if [ $IS_REMOTE -eq 1 ]; then
log_info "检测到远程Ollama服务配置,无需在本地安装Ollama"
return 0
fi
log_info "本地Ollama未安装,正在安装..."
OS=$(uname)
if [ "$OS" = "Darwin" ]; then
@@ -114,18 +116,54 @@ install_ollama() {
fi
if [ $? -eq 0 ]; then
log_success "Ollama安装完成"
log_success "本地Ollama安装完成"
return 0
else
log_error "Ollama安装失败"
log_error "本地Ollama安装失败"
return 1
fi
}
# 获取Ollama基础URL,检查是否为远程服务
get_ollama_base_url() {
check_env_file
# 从环境变量获取Ollama基础URL
OLLAMA_URL=${OLLAMA_BASE_URL:-"http://host.docker.internal:11434"}
# 提取主机部分
OLLAMA_HOST=$(echo "$OLLAMA_URL" | sed -E 's|^https?://||' | sed -E 's|:[0-9]+$||' | sed -E 's|/.*$||')
# 提取端口部分
OLLAMA_PORT=$(echo "$OLLAMA_URL" | grep -oE ':[0-9]+' | grep -oE '[0-9]+' || echo "11434")
# 检查是否为localhost或127.0.0.1
IS_REMOTE=0
if [ "$OLLAMA_HOST" = "localhost" ] || [ "$OLLAMA_HOST" = "127.0.0.1" ] || [ "$OLLAMA_HOST" = "host.docker.internal" ]; then
IS_REMOTE=0 # 本地服务
else
IS_REMOTE=1 # 远程服务
fi
}
# 启动Ollama服务
start_ollama() {
log_info "正在检查Ollama服务..."
# 提取主机和端口
get_ollama_base_url
log_info "Ollama服务地址: $OLLAMA_URL"
if [ $IS_REMOTE -eq 1 ]; then
log_info "检测到远程Ollama服务,将直接使用远程服务,不进行本地安装和启动"
# 检查远程服务是否可用
if curl -s "$OLLAMA_URL/api/tags" &> /dev/null; then
log_success "远程Ollama服务可访问"
return 0
else
log_warning "远程Ollama服务不可访问,请确认服务地址正确且已启动"
return 1
fi
fi
# 以下为本地服务的处理
# 检查Ollama是否已安装
if ! command -v ollama &> /dev/null; then
install_ollama
@@ -135,19 +173,19 @@ start_ollama() {
fi
# 检查Ollama服务是否已运行
if curl -s http://localhost:11434/api/tags &> /dev/null; then
log_success "Ollama服务已经在运行"
if curl -s "http://localhost:$OLLAMA_PORT/api/tags" &> /dev/null; then
log_success "本地Ollama服务已经在运行,端口:$OLLAMA_PORT"
else
log_info "启动Ollama服务..."
log_info "启动本地Ollama服务..."
# 注意:官方推荐使用 systemctl 或 launchctl 管理服务,直接后台运行仅用于临时场景
systemctl restart ollama || (ollama serve > /dev/null 2>&1 &)
systemctl restart ollama || (ollama serve > /dev/null 2>&1 < /dev/null &)
# 等待服务启动
MAX_RETRIES=30
COUNT=0
while [ $COUNT -lt $MAX_RETRIES ]; do
if curl -s http://localhost:11434/api/tags &> /dev/null; then
log_success "Ollama服务已成功启动"
if curl -s "http://localhost:$OLLAMA_PORT/api/tags" &> /dev/null; then
log_success "本地Ollama服务已成功启动,端口:$OLLAMA_PORT"
break
fi
echo -ne "等待Ollama服务启动... ($COUNT/$MAX_RETRIES)\r"
@@ -157,12 +195,12 @@ start_ollama() {
echo "" # 换行
if [ $COUNT -eq $MAX_RETRIES ]; then
log_error "Ollama服务启动失败"
log_error "本地Ollama服务启动失败"
return 1
fi
fi
log_success "Ollama服务地址: http://localhost:11434"
log_success "本地Ollama服务地址: http://localhost:$OLLAMA_PORT"
return 0
}
@@ -170,9 +208,17 @@ start_ollama() {
stop_ollama() {
log_info "正在停止Ollama服务..."
# 检查是否为远程服务
get_ollama_base_url
if [ $IS_REMOTE -eq 1 ]; then
log_info "检测到远程Ollama服务,无需在本地停止"
return 0
fi
# 检查Ollama是否已安装
if ! command -v ollama &> /dev/null; then
log_info "Ollama未安装,无需停止"
log_info "本地Ollama未安装,无需停止"
return 0
fi
@@ -184,9 +230,9 @@ stop_ollama() {
else
pkill -f "ollama serve"
fi
log_success "Ollama服务已停止"
log_success "本地Ollama服务已停止"
else
log_info "Ollama服务未运行"
log_info "本地Ollama服务未运行"
fi
return 0
@@ -217,6 +263,20 @@ check_docker() {
return 0
}
check_platform() {
# 检测当前系统平台
log_info "检测系统平台信息..."
if [ "$(uname -m)" = "x86_64" ]; then
export PLATFORM="linux/amd64"
elif [ "$(uname -m)" = "aarch64" ] || [ "$(uname -m)" = "arm64" ]; then
export PLATFORM="linux/arm64"
else
log_warning "未识别的平台类型:$(uname -m),将使用默认平台 linux/amd64"
export PLATFORM="linux/amd64"
fi
log_info "当前平台:$PLATFORM"
}
# 启动Docker容器
start_docker() {
log_info "正在启动Docker容器..."
@@ -233,18 +293,8 @@ start_docker() {
# 读取.env文件
source "$PROJECT_ROOT/.env"
storage_type=${STORAGE_TYPE:-local}
# 检测当前系统平台
log_info "检测系统平台信息..."
if [ "$(uname -m)" = "x86_64" ]; then
export PLATFORM="linux/amd64"
elif [ "$(uname -m)" = "aarch64" ] || [ "$(uname -m)" = "arm64" ]; then
export PLATFORM="linux/arm64"
else
log_warning "未识别的平台类型:$(uname -m),将使用默认平台 linux/amd64"
export PLATFORM="linux/amd64"
fi
log_info "当前平台:$PLATFORM"
check_platform
# 进入项目根目录再执行docker-compose命令
cd "$PROJECT_ROOT"
@@ -258,20 +308,6 @@ start_docker() {
return 1
fi
# 如果存储类型是minio,则启动MinIO服务
if [ "$storage_type" == "minio" ]; then
log_info "检测到MinIO存储配置,启动MinIO服务..."
# *** 修改点: 使用新的 docker compose 命令 ***
docker compose -f ./docker/docker-compose.minio.yml up --build -d
if [ $? -ne 0 ]; then
log_error "MinIO服务启动失败"
return 1
fi
log_success "MinIO服务已启动"
else
log_info "使用本地存储,不启动MinIO服务"
fi
log_success "所有Docker容器已成功启动"
# 显示容器状态
@@ -304,12 +340,6 @@ stop_docker() {
return 1
fi
# 如果存在minio配置,也停止minio
if [ -f "$PROJECT_ROOT/docker/docker-compose.minio.yml" ]; then
# *** 修改点: 使用新的 docker compose 命令 ***
docker compose -f "./docker/docker-compose.minio.yml" down
fi
log_success "所有Docker容器已停止"
return 0
}
@@ -354,6 +384,8 @@ restart_container() {
return 1
fi
check_platform
# 进入项目根目录再执行docker-compose命令
cd "$PROJECT_ROOT"
@@ -369,7 +401,7 @@ restart_container() {
# 构建并重启容器
log_info "正在重新构建容器 '$container_name'..."
# *** 修改点: 使用新的 docker compose 命令 ***
docker compose build "$container_name"
PLATFORM=$PLATFORM docker compose build "$container_name"
if [ $? -ne 0 ]; then
log_error "容器 '$container_name' 构建失败"
return 1
@@ -377,7 +409,7 @@ restart_container() {
log_info "正在重启容器 '$container_name'..."
# *** 修改点: 使用新的 docker compose 命令 ***
docker compose up -d --no-deps "$container_name"
PLATFORM=$PLATFORM docker compose up -d --no-deps "$container_name"
if [ $? -ne 0 ]; then
log_error "容器 '$container_name' 重启失败"
return 1
@@ -401,17 +433,28 @@ check_environment() {
# 检查.env文件
check_env_file
# 检查Ollama
if command -v ollama &> /dev/null; then
log_success "Ollama已安装"
if curl -s http://localhost:11434/api/tags &> /dev/null; then
version=$(curl -s http://localhost:11434/api/tags | grep -o '"version":"[^"]*"' | cut -d'"' -f4)
log_success "Ollama服务正在运行,版本: $version"
get_ollama_base_url
if [ $IS_REMOTE -eq 1 ]; then
log_info "检测到远程Ollama服务配置"
if curl -s "$OLLAMA_URL/api/tags" &> /dev/null; then
version=$(curl -s "$OLLAMA_URL/api/tags" | grep -o '"version":"[^"]*"' | cut -d'"' -f4)
log_success "远程Ollama服务可访问,版本: $version"
else
log_warning "Ollama已安装但服务未运行"
log_warning "远程Ollama服务不可访问,请确认服务地址正确且已启动"
fi
else
log_warning "Ollama未安装"
if command -v ollama &> /dev/null; then
log_success "本地Ollama已安装"
if curl -s "http://localhost:$OLLAMA_PORT/api/tags" &> /dev/null; then
version=$(curl -s "http://localhost:$OLLAMA_PORT/api/tags" | grep -o '"version":"[^"]*"' | cut -d'"' -f4)
log_success "本地Ollama服务正在运行,版本: $version"
else
log_warning "本地Ollama已安装但服务未运行"
fi
else
log_warning "本地Ollama未安装"
fi
fi
# 检查磁盘空间
@@ -580,7 +623,7 @@ else
fi
elif [ "$START_OLLAMA" = true ] && [ $OLLAMA_RESULT -eq 0 ]; then
log_success "Ollama服务启动完成,可通过以下地址访问:"
echo -e "${GREEN} - Ollama API: http://localhost:11434${NC}"
echo -e "${GREEN} - Ollama API: http://localhost:$OLLAMA_PORT${NC}"
elif [ "$START_DOCKER" = true ] && [ $DOCKER_RESULT -eq 0 ]; then
log_success "Docker容器启动完成,可通过以下地址访问:"
echo -e "${GREEN} - 前端界面: http://localhost${NC}"
+1
View File
@@ -17,6 +17,7 @@ paddleocr==3.0.0
markdown
pypdf
cos-python-sdk-v5
minio
textract
antiword
openai
View File
+2 -1
View File
@@ -21,7 +21,8 @@ from .markdown_parser import MarkdownParser
from .text_parser import TextParser
from .image_parser import ImageParser
from .web_parser import WebParser
from .parser import Parser, ChunkingConfig
from .parser import Parser
from .config import ChunkingConfig
from .ocr_engine import OCREngine
# Export public classes and modules
+65 -148
View File
@@ -1,19 +1,22 @@
# -*- coding: utf-8 -*-
import re
import os
import uuid
import asyncio
from typing import List, Dict, Any, Optional, Tuple, Union
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
import logging
from qcloud_cos import CosConfig, CosS3Client
import sys
import traceback
import numpy as np
import time
import io
import json
from .ocr_engine import OCREngine
from .image_utils import image_to_base64
from .config import ChunkingConfig
from .storage import create_storage
from PIL import Image
# Add parent directory to Python path for src imports
current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -101,7 +104,7 @@ class BaseParser(ABC):
max_image_size: int = 1920, # Maximum image size
max_concurrent_tasks: int = 5, # Max concurrent tasks
max_chunks: int = 1000, # Max number of returned chunks
chunking_config: object = None, # Chunking configuration object
chunking_config: ChunkingConfig = None, # Chunking configuration object
):
"""Initialize parser
@@ -118,6 +121,8 @@ class BaseParser(ABC):
max_concurrent_tasks: Max concurrent tasks
max_chunks: Max number of returned chunks
"""
# Storage client instance
self._storage = None
self.file_name = file_name
self.file_type = file_type or os.path.splitext(file_name)[1]
self.enable_multimodal = enable_multimodal
@@ -144,12 +149,7 @@ class BaseParser(ABC):
# Only initialize Caption service if multimodal is enabled
if self.enable_multimodal:
try:
# Get VLM config from chunking config if available
vlm_config = None
if self.chunking_config and hasattr(self.chunking_config, 'vlm_config'):
vlm_config = self.chunking_config.vlm_config
self.caption_parser = Caption(vlm_config)
self.caption_parser = Caption(self.chunking_config.vlm_config)
except Exception as e:
logger.warning(f"Failed to initialize Caption service: {str(e)}")
self.caption_parser = None
@@ -526,70 +526,16 @@ class BaseParser(ABC):
caption = self.get_image_caption(image_data)
return image_data, caption
def _init_cos_client(self, cos_config=None):
"""Initialize Tencent Cloud COS client"""
try:
# Use provided COS config if available, otherwise fall back to environment variables
if cos_config:
secret_id = cos_config.get("secret_id")
secret_key = cos_config.get("secret_key")
region = cos_config.get("region")
bucket_name = cos_config.get("bucket_name")
appid = cos_config.get("app_id")
prefix = cos_config.get("path_prefix", "")
enable_old_domain = cos_config.get("enable_old_domain", "true").lower() == "true"
else:
# Get COS configuration from environment variables
secret_id = os.getenv("COS_SECRET_ID")
secret_key = os.getenv("COS_SECRET_KEY")
region = os.getenv("COS_REGION")
bucket_name = os.getenv("COS_BUCKET_NAME")
appid = os.getenv("COS_APP_ID")
prefix = os.getenv("COS_PATH_PREFIX")
enable_old_domain = (
os.getenv("COS_ENABLE_OLD_DOMAIN", "true").lower() == "true"
)
if not all([secret_id, secret_key, region, bucket_name, appid]):
logger.error(
"Incomplete COS configuration, missing required environment variables"
)
return None, None, None, None
# Initialize COS configuration
logger.info(
f"Initializing COS client with region: {region}, bucket: {bucket_name}"
)
config = CosConfig(
Appid=appid,
Region=region,
SecretId=secret_id,
SecretKey=secret_key,
EnableOldDomain=enable_old_domain,
)
# Create client
client = CosS3Client(config)
return client, bucket_name, region, prefix
except Exception as e:
logger.error(f"Failed to initialize COS client: {str(e)}")
return None, None, None, None
def _get_file_url(self, bucket_name, region, object_key):
"""Generate COS object URL
Args:
bucket_name: Bucket name
region: Region
object_key: Object key
Returns:
File URL
"""
return f"https://{bucket_name}.cos.{region}.myqcloud.com/{object_key}"
def __init_storage(self):
"""Initialize storage client based on configuration"""
if self._storage is None:
storage_config = self.chunking_config.storage_config if self.chunking_config else None
self._storage = create_storage(storage_config)
logger.info(f"Initialized storage client: {self._storage.__class__.__name__}")
return self._storage
def upload_file(self, file_path: str) -> str:
"""Upload file to Tencent Cloud COS
"""Upload file to object storage
Args:
file_path: File path
@@ -597,84 +543,43 @@ class BaseParser(ABC):
Returns:
File URL
"""
logger.info(f"Uploading file to COS: {file_path}")
logger.info(f"Uploading file: {file_path}")
try:
client, bucket_name, region, prefix = self._init_cos_client()
if not client:
return ""
# Generate object key, use UUID to avoid conflicts
file_name = os.path.basename(file_path)
object_key = (
f"{prefix}/images/{uuid.uuid4().hex}{os.path.splitext(file_name)[1]}"
)
logger.info(f"Generated object key: {object_key}")
# Upload file
logger.info("Attempting to upload file to COS")
response = client.upload_file(
Bucket=bucket_name, LocalFilePath=file_path, Key=object_key
)
# Get file URL
file_url = self._get_file_url(bucket_name, region, object_key)
logger.info(f"Successfully uploaded file to COS: {file_url}")
return file_url
storage = self.__init_storage()
return storage.upload_file(file_path)
except Exception as e:
logger.error(f"Failed to upload file to COS: {str(e)}")
logger.error(f"Failed to upload file: {str(e)}")
return ""
def upload_bytes(self, content: bytes, file_ext: str = ".png", cos_config=None) -> str:
"""Directly upload file content to Tencent Cloud COS
def upload_bytes(self, content: bytes, file_ext: str = ".png") -> str:
"""Upload bytes to object storage
Args:
content: File byte content
file_ext: File extension, default is .png
cos_config: COS configuration dictionary
content: Byte content to upload
file_ext: File extension
Returns:
File URL
"""
logger.info(
f"Uploading bytes content to COS, content size: {len(content)} bytes"
)
logger.info(f"Uploading bytes content, size: {len(content)} bytes")
try:
client, bucket_name, region, prefix = self._init_cos_client(cos_config)
if not client:
return ""
# Generate object key, use UUID to avoid conflicts
object_key = f"{prefix}/images/{uuid.uuid4().hex}{file_ext}"
logger.info(f"Generated object key: {object_key}")
# Directly upload file content
logger.info("Attempting to upload bytes content to COS")
response = client.put_object(
Bucket=bucket_name, Body=content, Key=object_key
)
# Get file URL
file_url = self._get_file_url(bucket_name, region, object_key)
logger.info(f"Successfully uploaded bytes to COS: {file_url}")
return file_url
storage = self.__init_storage()
return storage.upload_bytes(content, file_ext)
except Exception as e:
logger.error(f"Failed to upload bytes to COS: {str(e)}")
logger.error(f"Failed to upload bytes to storage: {str(e)}")
traceback.print_exc()
return ""
@abstractmethod
def parse_into_text(self, content: bytes) -> str:
def parse_into_text(self, content: bytes) -> Union[str, Tuple[str, Dict[str, Any]]]:
"""Parse document content
Args:
content: Document content
Returns:
Parse result
Either a string containing the parsed text, or a tuple of (text, image_map)
where image_map is a dict mapping image URLs to Image objects
"""
pass
@@ -690,7 +595,12 @@ class BaseParser(ABC):
logger.info(
f"Parsing document with {self.__class__.__name__}, content size: {len(content)} bytes"
)
text = self.parse_into_text(content)
parse_result = self.parse_into_text(content)
if isinstance(parse_result, tuple):
text, image_map = parse_result
else:
text = parse_result
image_map = {}
logger.info(f"Extracted {len(text)} characters of text from {self.file_name}")
logger.info(f"Beginning chunking process for text")
chunks = self.chunk_text(text)
@@ -725,7 +635,7 @@ class BaseParser(ABC):
if file_ext in allowed_types:
logger.info(f"Processing images in each chunk for file type: {file_ext}")
chunks = self.process_chunks_images(chunks)
chunks = self.process_chunks_images(chunks, image_map)
else:
logger.info(f"Skipping image processing for unsupported file type: {file_ext}")
@@ -1053,15 +963,16 @@ class BaseParser(ABC):
return images_info
async def download_and_upload_image(self, img_url: str, current_request_id=None):
"""Download image and upload to COS, if it's already a COS path or local path, use directly
async def download_and_upload_image(self, img_url: str, current_request_id=None, image_map=None):
"""Download image and upload to object storage, if it's already an object storage path or local path, use directly
Args:
img_url: Image URL or local path
current_request_id: Current request ID
image_map: Optional dictionary mapping image URLs to Image objects
Returns:
tuple: (original URL, COS URL, image object), if failed returns (original URL, None, None)
tuple: (original URL, storage URL, image object), if failed returns (original URL, None, None)
"""
# Set request ID context in the asynchronous task
try:
@@ -1077,8 +988,14 @@ class BaseParser(ABC):
from PIL import Image
import io
# Check if it's already a COS path
if "cos" in img_url and "myqcloud.com" in img_url:
# Check if image is already in the image_map
if image_map and img_url in image_map:
logger.info(f"Image already in image_map: {img_url}, using cached object")
return img_url, img_url, image_map[img_url]
# Check if it's already a storage URL (COS or MinIO)
is_storage_url = any(pattern in img_url for pattern in ["cos", "myqcloud.com", "minio", ".s3."])
if is_storage_url:
logger.info(f"Image already on COS: {img_url}, no need to re-upload")
try:
# Still need to get image object for OCR processing
@@ -1101,10 +1018,10 @@ class BaseParser(ABC):
# Image will be closed by the caller
pass
else:
logger.warning(f"Failed to get COS image: {response.status_code}")
logger.warning(f"Failed to get storage image: {response.status_code}")
return img_url, img_url, None
except Exception as e:
logger.error(f"Error getting COS image: {str(e)}")
logger.error(f"Error getting storage image: {str(e)}")
return img_url, img_url, None
# Check if it's a local file path
@@ -1114,12 +1031,12 @@ class BaseParser(ABC):
try:
# Read local image
image = Image.open(img_url)
# Upload to COS
# Upload to storage
with open(img_url, 'rb') as f:
content = f.read()
cos_url = self.upload_bytes(content)
logger.info(f"Successfully uploaded local image to COS: {cos_url}")
return img_url, cos_url, image
storage_url = self.upload_bytes(content)
logger.info(f"Successfully uploaded local image to storage: {storage_url}")
return img_url, storage_url, image
except Exception as e:
logger.error(f"Error processing local image: {str(e)}")
if image and hasattr(image, 'close'):
@@ -1144,10 +1061,10 @@ class BaseParser(ABC):
# Download successful, create image object
image = Image.open(io.BytesIO(response.content))
try:
# Upload to COS using the method in BaseParser
cos_url = self.upload_bytes(response.content)
logger.info(f"Successfully uploaded image to COS: {cos_url}")
return img_url, cos_url, image
# Upload to storage using the method in BaseParser
storage_url = self.upload_bytes(response.content)
logger.info(f"Successfully uploaded image to storage: {storage_url}")
return img_url, storage_url, image
finally:
# Image will be closed by the caller
pass
@@ -1159,7 +1076,7 @@ class BaseParser(ABC):
logger.error(f"Error downloading or processing image: {str(e)}")
return img_url, None, None
async def process_chunk_images_async(self, chunk, chunk_idx, total_chunks, current_request_id=None):
async def process_chunk_images_async(self, chunk, chunk_idx, total_chunks, current_request_id=None, image_map=None):
"""Asynchronously process images in a single Chunk
Args:
@@ -1201,7 +1118,7 @@ class BaseParser(ABC):
loop = asyncio.get_event_loop()
# Concurrent download and upload of images
tasks = [self.download_and_upload_image(url, current_request_id) for url in url_to_info_map.keys()]
tasks = [self.download_and_upload_image(url, current_request_id, image_map) for url in url_to_info_map.keys()]
results = await asyncio.gather(*tasks)
# Process download results, prepare for OCR processing
@@ -1248,7 +1165,7 @@ class BaseParser(ABC):
logger.info(f"Completed image processing in Chunk #{chunk_idx+1}")
return chunk
def process_chunks_images(self, chunks: List[Chunk]) -> List[Chunk]:
def process_chunks_images(self, chunks: List[Chunk], image_map=None) -> List[Chunk]:
"""Concurrent processing of images in all Chunks
Args:
@@ -1282,7 +1199,7 @@ class BaseParser(ABC):
async def process_with_limit(chunk, idx, total):
"""Use semaphore to control concurrent processing of Chunks"""
async with semaphore:
return await self.process_chunk_images_async(chunk, idx, total, current_request_id)
return await self.process_chunk_images_async(chunk, idx, total, current_request_id, image_map)
# Create tasks for all Chunks
tasks = [
+21
View File
@@ -0,0 +1,21 @@
from dataclasses import dataclass, field
@dataclass
class ChunkingConfig:
"""
Configuration for text chunking process.
Controls how documents are split into smaller pieces for processing.
"""
chunk_size: int = 512 # Maximum size of each chunk in tokens/chars
chunk_overlap: int = 50 # Number of tokens/chars to overlap between chunks
separators: list = field(
default_factory=lambda: ["\n\n", "\n", ""]
) # Text separators in order of priority
enable_multimodal: bool = (
False # Whether to enable multimodal processing (text + images)
)
storage_config: dict = None # Preferred field name going forward
vlm_config: dict = None # VLM configuration for image captioning
File diff suppressed because it is too large Load Diff
+21 -17
View File
@@ -3,12 +3,13 @@ import os
import asyncio
from PIL import Image
import io
from typing import Dict, Any, Tuple, Union
from .base_parser import BaseParser, ParseResult
import numpy as np
# Set up logger for this module
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
class ImageParser(BaseParser):
"""
@@ -22,43 +23,46 @@ class ImageParser(BaseParser):
4. Returning a combined result with both text and image reference
"""
def parse_into_text(self, content: bytes) -> str:
def parse_into_text(self, content: bytes) -> Union[str, Tuple[str, Dict[str, Any]]]:
"""
Parse image content, only upload the image and return Markdown reference, no OCR or caption processing.
Parse image content, upload the image and return Markdown reference along with image map.
Args:
content: Raw image data (bytes)
Returns:
String containing Markdown image reference
Tuple of (markdown_text, image_map) where image_map maps image URLs to PIL Image objects
"""
logger.info(f"Parsing image content, size: {len(content)} bytes")
image_map = {}
try:
# Upload image to storage service
logger.info("Uploading image to storage")
_, ext = os.path.splitext(self.file_name)
# Get COS config from chunking config if available
cos_config = None
if hasattr(self, 'chunking_config') and self.chunking_config and hasattr(self.chunking_config, 'cos_config'):
cos_config = self.chunking_config.cos_config
image_url = self.upload_bytes(content, file_ext=ext, cos_config=cos_config)
image_url = self.upload_bytes(content, file_ext=ext)
if not image_url:
logger.error("Failed to upload image to storage")
return ""
return "", {}
logger.info(
f"Successfully uploaded image, URL: {image_url[:50]}..."
if len(image_url) > 50
else f"Successfully uploaded image, URL: {image_url}"
)
# Directly generate Markdown image reference, no OCR or caption processing
markdown_text = f"![{self.file_name}]({image_url})"
logger.info("Generated Markdown image reference without OCR or caption processing")
# Create image object and add to map
try:
from PIL import Image
import io
image = Image.open(io.BytesIO(content))
image_map[image_url] = image
logger.info(f"Added image to image_map for URL: {image_url}")
except Exception as img_err:
logger.error(f"Error creating image object: {str(img_err)}")
return markdown_text
markdown_text = f"![{self.file_name}]({image_url})"
return markdown_text, image_map
except Exception as e:
logger.error(f"Error parsing image: {str(e)}")
return ""
return "", {}
@@ -3,7 +3,7 @@ import re
import logging
import numpy as np
import os # Import os module to get environment variables
from typing import Dict, List, Optional, Tuple
from typing import Dict, List, Optional, Tuple, Union, Any
from .base_parser import BaseParser
# Get logger object
@@ -13,7 +13,7 @@ logger = logging.getLogger(__name__)
class MarkdownParser(BaseParser):
"""Markdown document parser"""
def parse_into_text(self, content: bytes) -> str:
def parse_into_text(self, content: bytes) -> Union[str, Tuple[str, Dict[str, Any]]]:
"""Parse Markdown document, only extract text content, do not process images
Args:
+1 -21
View File
@@ -10,30 +10,11 @@ from .markdown_parser import MarkdownParser
from .text_parser import TextParser
from .image_parser import ImageParser
from .web_parser import WebParser
from .config import ChunkingConfig
import traceback
logger = logging.getLogger(__name__)
@dataclass
class ChunkingConfig:
"""
Configuration for text chunking process.
Controls how documents are split into smaller pieces for processing.
"""
chunk_size: int = 512 # Maximum size of each chunk in tokens/chars
chunk_overlap: int = 50 # Number of tokens/chars to overlap between chunks
separators: list = field(
default_factory=lambda: ["\n\n", "\n", ""]
) # Text separators in order of priority
enable_multimodal: bool = (
False # Whether to enable multimodal processing (text + images)
)
cos_config: dict = None # COS configuration for file storage
vlm_config: dict = None # VLM configuration for image captioning
@dataclass
class Chunk:
"""
@@ -100,7 +81,6 @@ class Parser:
file_type: str,
content: bytes,
config: ChunkingConfig,
enable_multimodal: bool = False,
) -> Optional[ParseResult]:
"""
Parse file content using appropriate parser based on file type.
+2 -2
View File
@@ -1,7 +1,7 @@
import logging
import os
import io
from typing import Any, List, Iterator, Optional, Mapping, Tuple, Dict
from typing import Any, List, Iterator, Optional, Mapping, Tuple, Dict, Union
from pypdf import PdfReader
from .base_parser import BaseParser
@@ -16,7 +16,7 @@ class PDFParser(BaseParser):
It uses the pypdf library for simple text extraction.
"""
def parse_into_text(self, content: bytes) -> str:
def parse_into_text(self, content: bytes) -> Union[str, Tuple[str, Dict[str, Any]]]:
"""
Parse PDF document content into text
+363
View File
@@ -0,0 +1,363 @@
# -*- coding: utf-8 -*-
import os
import uuid
import logging
import io
import traceback
from abc import ABC, abstractmethod
from typing import Tuple, Optional
from qcloud_cos import CosConfig, CosS3Client
from minio import Minio
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
class Storage(ABC):
"""Abstract base class for object storage operations"""
@abstractmethod
def upload_file(self, file_path: str) -> str:
"""Upload file to object storage
Args:
file_path: File path
Returns:
File URL
"""
pass
@abstractmethod
def upload_bytes(self, content: bytes, file_ext: str = ".png") -> str:
"""Upload bytes to object storage
Args:
content: Byte content to upload
file_ext: File extension
Returns:
File URL
"""
pass
class CosStorage(Storage):
"""Tencent Cloud COS storage implementation"""
def __init__(self, storage_config=None):
"""Initialize COS storage
Args:
storage_config: Storage configuration
"""
self.storage_config = storage_config
self.client, self.bucket_name, self.region, self.prefix = self._init_cos_client()
def _init_cos_client(self):
"""Initialize Tencent Cloud COS client"""
try:
# Use provided COS config if available, otherwise fall back to environment variables
if self.storage_config and self.storage_config.get("access_key_id") != "":
cos_config = self.storage_config
secret_id = cos_config.get("access_key_id")
secret_key = cos_config.get("secret_access_key")
region = cos_config.get("region")
bucket_name = cos_config.get("bucket_name")
appid = cos_config.get("app_id")
prefix = cos_config.get("path_prefix", "")
else:
# Get COS configuration from environment variables
secret_id = os.getenv("COS_SECRET_ID")
secret_key = os.getenv("COS_SECRET_KEY")
region = os.getenv("COS_REGION")
bucket_name = os.getenv("COS_BUCKET_NAME")
appid = os.getenv("COS_APP_ID")
prefix = os.getenv("COS_PATH_PREFIX")
enable_old_domain = (
os.getenv("COS_ENABLE_OLD_DOMAIN", "true").lower() == "true"
)
if not all([secret_id, secret_key, region, bucket_name, appid]):
logger.error(
"Incomplete COS configuration, missing required environment variables"
f"secret_id: {secret_id}, secret_key: {secret_key}, region: {region}, bucket_name: {bucket_name}, appid: {appid}"
)
return None, None, None, None
# Initialize COS configuration
logger.info(
f"Initializing COS client with region: {region}, bucket: {bucket_name}"
)
config = CosConfig(
Appid=appid,
Region=region,
SecretId=secret_id,
SecretKey=secret_key,
EnableOldDomain=enable_old_domain,
)
# Create client
client = CosS3Client(config)
return client, bucket_name, region, prefix
except Exception as e:
logger.error(f"Failed to initialize COS client: {str(e)}")
return None, None, None, None
def _get_download_url(self, bucket_name, region, object_key):
"""Generate COS object URL
Args:
bucket_name: Bucket name
region: Region
object_key: Object key
Returns:
File URL
"""
return f"https://{bucket_name}.cos.{region}.myqcloud.com/{object_key}"
def upload_file(self, file_path: str) -> str:
"""Upload file to Tencent Cloud COS
Args:
file_path: File path
Returns:
File URL
"""
logger.info(f"Uploading file to COS: {file_path}")
try:
if not self.client:
return ""
# Generate object key, use UUID to avoid conflicts
file_name = os.path.basename(file_path)
object_key = (
f"{self.prefix}/images/{uuid.uuid4().hex}{os.path.splitext(file_name)[1]}"
)
logger.info(f"Generated object key: {object_key}")
# Upload file
logger.info("Attempting to upload file to COS")
response = self.client.upload_file(
Bucket=self.bucket_name, LocalFilePath=file_path, Key=object_key
)
# Get file URL
file_url = self._get_download_url(self.bucket_name, self.region, object_key)
logger.info(f"Successfully uploaded file to COS: {file_url}")
return file_url
except Exception as e:
logger.error(f"Failed to upload file to COS: {str(e)}")
return ""
def upload_bytes(self, content: bytes, file_ext: str = ".png") -> str:
"""Upload bytes to Tencent Cloud COS
Args:
content: Byte content to upload
file_ext: File extension
Returns:
File URL
"""
try:
logger.info(f"Uploading bytes content to COS, size: {len(content)} bytes")
if not self.client:
return ""
object_key = f"{self.prefix}/images/{uuid.uuid4().hex}{file_ext}" if self.prefix else f"images/{uuid.uuid4().hex}{file_ext}"
logger.info(f"Generated object key: {object_key}")
self.client.put_object(Bucket=self.bucket_name, Body=content, Key=object_key)
file_url = self._get_download_url(self.bucket_name, self.region, object_key)
logger.info(f"Successfully uploaded bytes to COS: {file_url}")
return file_url
except Exception as e:
logger.error(f"Failed to upload bytes to COS: {str(e)}")
traceback.print_exc()
return ""
class MinioStorage(Storage):
"""MinIO storage implementation"""
def __init__(self, storage_config=None):
"""Initialize MinIO storage
Args:
storage_config: Storage configuration
"""
self.storage_config = storage_config
self.client, self.bucket_name, self.use_ssl, self.endpoint, self.path_prefix = self._init_minio_client()
def _init_minio_client(self):
"""Initialize MinIO client from environment variables or injected config.
If storage_config.path_prefix contains JSON from server (for minio case),
prefer those values to override envs.
"""
try:
endpoint = os.getenv("MINIO_ENDPOINT")
access_key = os.getenv("MINIO_ACCESS_KEY_ID")
secret_key = os.getenv("MINIO_SECRET_ACCESS_KEY")
bucket_name = os.getenv("MINIO_BUCKET_NAME")
use_ssl = os.getenv("MINIO_USE_SSL", "false").lower() == "true"
path_prefix = os.getenv("MINIO_PATH_PREFIX")
# Attempt to override with new generic storage fields
if self.storage_config:
storage_config = self.storage_config
# Direct fields
endpoint = storage_config.get("endpoint", endpoint)
bucket_name = storage_config.get("bucket_name", bucket_name)
path_prefix = storage_config.get("path_prefix", path_prefix)
access_key = storage_config.get("access_key_id", access_key)
secret_key = storage_config.get("secret_access_key", secret_key)
if not all([endpoint, access_key, secret_key, bucket_name]):
logger.error("Incomplete MinIO configuration, missing required environment variables")
return None, None, None, None, None
# Initialize client
client = Minio(endpoint, access_key=access_key, secret_key=secret_key, secure=use_ssl)
# Ensure bucket exists
found = client.bucket_exists(bucket_name)
if not found:
client.make_bucket(bucket_name)
policy = '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"AWS":["*"]},"Action":["s3:GetBucketLocation","s3:ListBucket"],"Resource":["arn:aws:s3:::%s"]},{"Effect":"Allow","Principal":{"AWS":["*"]},"Action":["s3:GetObject"],"Resource":["arn:aws:s3:::%s/*"]}]}' % (bucket_name, bucket_name)
client.set_bucket_policy(bucket_name, policy)
return client, bucket_name, use_ssl, endpoint, path_prefix
except Exception as e:
logger.error(f"Failed to initialize MinIO client: {str(e)}")
return None, None, None, None, None
def _get_download_url(self, bucket_name: str, object_key: str, use_ssl: bool, endpoint: str, public_endpoint: str = None):
"""Construct a public URL for MinIO object.
If MINIO_PUBLIC_ENDPOINT is provided, use it; otherwise fallback to endpoint.
"""
if public_endpoint:
base = public_endpoint
else:
scheme = "https" if use_ssl else "http"
base = f"{scheme}://{endpoint}"
# Path-style URL for MinIO
return f"{base}/{bucket_name}/{object_key}"
def upload_file(self, file_path: str) -> str:
"""Upload file to MinIO
Args:
file_path: File path
Returns:
File URL
"""
logger.info(f"Uploading file to MinIO: {file_path}")
try:
if not self.client:
return ""
# Generate object key, use UUID to avoid conflicts
file_name = os.path.basename(file_path)
object_key = f"{self.path_prefix}/images/{uuid.uuid4().hex}{os.path.splitext(file_name)[1]}" if self.path_prefix else f"images/{uuid.uuid4().hex}{os.path.splitext(file_name)[1]}"
logger.info(f"Generated MinIO object key: {object_key}")
# Upload file
logger.info("Attempting to upload file to MinIO")
with open(file_path, 'rb') as file_data:
file_size = os.path.getsize(file_path)
self.client.put_object(
bucket_name=self.bucket_name,
object_name=object_key,
data=file_data,
length=file_size,
content_type='application/octet-stream'
)
# Get file URL
file_url = self._get_download_url(
self.bucket_name,
object_key,
self.use_ssl,
self.endpoint,
os.getenv("MINIO_PUBLIC_ENDPOINT", None)
)
logger.info(f"Successfully uploaded file to MinIO: {file_url}")
return file_url
except Exception as e:
logger.error(f"Failed to upload file to MinIO: {str(e)}")
return ""
def upload_bytes(self, content: bytes, file_ext: str = ".png") -> str:
"""Upload bytes to MinIO
Args:
content: Byte content to upload
file_ext: File extension
Returns:
File URL
"""
try:
logger.info(f"Uploading bytes content to MinIO, size: {len(content)} bytes")
if not self.client:
return ""
object_key = f"{self.path_prefix}/images/{uuid.uuid4().hex}{file_ext}" if self.path_prefix else f"images/{uuid.uuid4().hex}{file_ext}"
logger.info(f"Generated MinIO object key: {object_key}")
self.client.put_object(
self.bucket_name,
object_key,
data=io.BytesIO(content),
length=len(content),
content_type="application/octet-stream"
)
file_url = self._get_download_url(
self.bucket_name,
object_key,
self.use_ssl,
self.endpoint,
os.getenv("MINIO_PUBLIC_ENDPOINT", None)
)
logger.info(f"Successfully uploaded bytes to MinIO: {file_url}")
return file_url
except Exception as e:
logger.error(f"Failed to upload bytes to MinIO: {str(e)}")
traceback.print_exc()
return ""
def create_storage(storage_config=None) -> Storage:
"""Create a storage instance based on configuration or environment variables
Args:
storage_config: Storage configuration dictionary
Returns:
Storage instance
"""
storage_type = os.getenv("STORAGE_TYPE", "cos").lower()
if storage_config:
storage_type = str(storage_config.get("provider", storage_type)).lower()
logger.info(f"Creating {storage_type} storage instance")
if storage_type == "minio":
return MinioStorage(storage_config)
elif storage_type == "cos":
# Default to COS
return CosStorage(storage_config)
else:
return None
+2 -1
View File
@@ -1,5 +1,6 @@
import logging
from .base_parser import BaseParser
from typing import Dict, Any, Tuple, Union
logger = logging.getLogger(__name__)
@@ -10,7 +11,7 @@ class TextParser(BaseParser):
This parser handles text extraction and chunking from plain text documents.
"""
def parse_into_text(self, content: bytes) -> str:
def parse_into_text(self, content: bytes) -> Union[str, Tuple[str, Dict[str, Any]]]:
"""
Parse text document content by decoding bytes to string.
+2 -2
View File
@@ -1,4 +1,4 @@
from typing import Any, Optional
from typing import Any, Optional, Tuple, Dict, Union
from playwright.async_api import async_playwright
from bs4 import BeautifulSoup
@@ -66,7 +66,7 @@ class WebParser(BaseParser):
# Return empty BeautifulSoup object on error
return BeautifulSoup("", "html.parser")
def parse_into_text(self, content: bytes) -> str:
def parse_into_text(self, content: bytes) -> Union[str, Tuple[str, Dict[str, Any]]]:
"""Parse web page
Args:
+140 -75
View File
@@ -1,6 +1,6 @@
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.36.7
// protoc-gen-go v1.36.6
// protoc v5.29.3
// source: docreader.proto
@@ -21,33 +21,84 @@ const (
_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
)
// COS 配置
type COSConfig struct {
state protoimpl.MessageState `protogen:"open.v1"`
SecretId string `protobuf:"bytes,1,opt,name=secret_id,json=secretId,proto3" json:"secret_id,omitempty"` // COS Secret ID
SecretKey string `protobuf:"bytes,2,opt,name=secret_key,json=secretKey,proto3" json:"secret_key,omitempty"` // COS Secret Key
Region string `protobuf:"bytes,3,opt,name=region,proto3" json:"region,omitempty"` // COS Region
BucketName string `protobuf:"bytes,4,opt,name=bucket_name,json=bucketName,proto3" json:"bucket_name,omitempty"` // COS Bucket Name
AppId string `protobuf:"bytes,5,opt,name=app_id,json=appId,proto3" json:"app_id,omitempty"` // COS App ID
PathPrefix string `protobuf:"bytes,6,opt,name=path_prefix,json=pathPrefix,proto3" json:"path_prefix,omitempty"` // COS Path Prefix
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
// 对象存储提供方
type StorageProvider int32
const (
StorageProvider_STORAGE_PROVIDER_UNSPECIFIED StorageProvider = 0
StorageProvider_COS StorageProvider = 1 // 腾讯云 COS
StorageProvider_MINIO StorageProvider = 2 // MinIO/S3 兼容
)
// Enum value maps for StorageProvider.
var (
StorageProvider_name = map[int32]string{
0: "STORAGE_PROVIDER_UNSPECIFIED",
1: "COS",
2: "MINIO",
}
StorageProvider_value = map[string]int32{
"STORAGE_PROVIDER_UNSPECIFIED": 0,
"COS": 1,
"MINIO": 2,
}
)
func (x StorageProvider) Enum() *StorageProvider {
p := new(StorageProvider)
*p = x
return p
}
func (x *COSConfig) Reset() {
*x = COSConfig{}
func (x StorageProvider) String() string {
return protoimpl.X.EnumStringOf(x.Descriptor(), protoreflect.EnumNumber(x))
}
func (StorageProvider) Descriptor() protoreflect.EnumDescriptor {
return file_docreader_proto_enumTypes[0].Descriptor()
}
func (StorageProvider) Type() protoreflect.EnumType {
return &file_docreader_proto_enumTypes[0]
}
func (x StorageProvider) Number() protoreflect.EnumNumber {
return protoreflect.EnumNumber(x)
}
// Deprecated: Use StorageProvider.Descriptor instead.
func (StorageProvider) EnumDescriptor() ([]byte, []int) {
return file_docreader_proto_rawDescGZIP(), []int{0}
}
// 通用对象存储配置,兼容 COS 与 MinIO
type StorageConfig struct {
state protoimpl.MessageState `protogen:"open.v1"`
Provider StorageProvider `protobuf:"varint,1,opt,name=provider,proto3,enum=docreader.StorageProvider" json:"provider,omitempty"` // 存储提供方
Region string `protobuf:"bytes,2,opt,name=region,proto3" json:"region,omitempty"` // 区域(COS 使用)
BucketName string `protobuf:"bytes,3,opt,name=bucket_name,json=bucketName,proto3" json:"bucket_name,omitempty"` // 桶名
AccessKeyId string `protobuf:"bytes,4,opt,name=access_key_id,json=accessKeyId,proto3" json:"access_key_id,omitempty"` // 访问密钥 IDMinIO/S3 使用)
SecretAccessKey string `protobuf:"bytes,5,opt,name=secret_access_key,json=secretAccessKey,proto3" json:"secret_access_key,omitempty"` // 访问密钥 SecretMinIO/S3 使用)
AppId string `protobuf:"bytes,6,opt,name=app_id,json=appId,proto3" json:"app_id,omitempty"` // 应用 IDCOS 使用)
PathPrefix string `protobuf:"bytes,7,opt,name=path_prefix,json=pathPrefix,proto3" json:"path_prefix,omitempty"` // 路径前缀
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
}
func (x *StorageConfig) Reset() {
*x = StorageConfig{}
mi := &file_docreader_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
func (x *COSConfig) String() string {
func (x *StorageConfig) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*COSConfig) ProtoMessage() {}
func (*StorageConfig) ProtoMessage() {}
func (x *COSConfig) ProtoReflect() protoreflect.Message {
func (x *StorageConfig) ProtoReflect() protoreflect.Message {
mi := &file_docreader_proto_msgTypes[0]
if x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
@@ -59,47 +110,54 @@ func (x *COSConfig) ProtoReflect() protoreflect.Message {
return mi.MessageOf(x)
}
// Deprecated: Use COSConfig.ProtoReflect.Descriptor instead.
func (*COSConfig) Descriptor() ([]byte, []int) {
// Deprecated: Use StorageConfig.ProtoReflect.Descriptor instead.
func (*StorageConfig) Descriptor() ([]byte, []int) {
return file_docreader_proto_rawDescGZIP(), []int{0}
}
func (x *COSConfig) GetSecretId() string {
func (x *StorageConfig) GetProvider() StorageProvider {
if x != nil {
return x.SecretId
return x.Provider
}
return ""
return StorageProvider_STORAGE_PROVIDER_UNSPECIFIED
}
func (x *COSConfig) GetSecretKey() string {
if x != nil {
return x.SecretKey
}
return ""
}
func (x *COSConfig) GetRegion() string {
func (x *StorageConfig) GetRegion() string {
if x != nil {
return x.Region
}
return ""
}
func (x *COSConfig) GetBucketName() string {
func (x *StorageConfig) GetBucketName() string {
if x != nil {
return x.BucketName
}
return ""
}
func (x *COSConfig) GetAppId() string {
func (x *StorageConfig) GetAccessKeyId() string {
if x != nil {
return x.AccessKeyId
}
return ""
}
func (x *StorageConfig) GetSecretAccessKey() string {
if x != nil {
return x.SecretAccessKey
}
return ""
}
func (x *StorageConfig) GetAppId() string {
if x != nil {
return x.AppId
}
return ""
}
func (x *COSConfig) GetPathPrefix() string {
func (x *StorageConfig) GetPathPrefix() string {
if x != nil {
return x.PathPrefix
}
@@ -181,7 +239,7 @@ type ReadConfig struct {
ChunkOverlap int32 `protobuf:"varint,2,opt,name=chunk_overlap,json=chunkOverlap,proto3" json:"chunk_overlap,omitempty"` // 分块重叠
Separators []string `protobuf:"bytes,3,rep,name=separators,proto3" json:"separators,omitempty"` // 分隔符
EnableMultimodal bool `protobuf:"varint,4,opt,name=enable_multimodal,json=enableMultimodal,proto3" json:"enable_multimodal,omitempty"` // 多模态处理
CosConfig *COSConfig `protobuf:"bytes,5,opt,name=cos_config,json=cosConfig,proto3" json:"cos_config,omitempty"` // COS 配置
StorageConfig *StorageConfig `protobuf:"bytes,5,opt,name=storage_config,json=storageConfig,proto3" json:"storage_config,omitempty"` // 对象存储配置(通用)
VlmConfig *VLMConfig `protobuf:"bytes,6,opt,name=vlm_config,json=vlmConfig,proto3" json:"vlm_config,omitempty"` // VLM 配置
unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache
@@ -245,9 +303,9 @@ func (x *ReadConfig) GetEnableMultimodal() bool {
return false
}
func (x *ReadConfig) GetCosConfig() *COSConfig {
func (x *ReadConfig) GetStorageConfig() *StorageConfig {
if x != nil {
return x.CosConfig
return x.StorageConfig
}
return nil
}
@@ -623,23 +681,23 @@ var File_docreader_proto protoreflect.FileDescriptor
const file_docreader_proto_rawDesc = "" +
"\n" +
"\x0fdocreader.proto\x12\tdocreader\"\xb8\x01\n" +
"\tCOSConfig\x12\x1b\n" +
"\tsecret_id\x18\x01 \x01(\tR\bsecretId\x12\x1d\n" +
"\n" +
"secret_key\x18\x02 \x01(\tR\tsecretKey\x12\x16\n" +
"\x06region\x18\x03 \x01(\tR\x06region\x12\x1f\n" +
"\vbucket_name\x18\x04 \x01(\tR\n" +
"bucketName\x12\x15\n" +
"\x06app_id\x18\x05 \x01(\tR\x05appId\x12\x1f\n" +
"\vpath_prefix\x18\x06 \x01(\tR\n" +
"\x0fdocreader.proto\x12\tdocreader\"\x88\x02\n" +
"\rStorageConfig\x126\n" +
"\bprovider\x18\x01 \x01(\x0e2\x1a.docreader.StorageProviderR\bprovider\x12\x16\n" +
"\x06region\x18\x02 \x01(\tR\x06region\x12\x1f\n" +
"\vbucket_name\x18\x03 \x01(\tR\n" +
"bucketName\x12\"\n" +
"\raccess_key_id\x18\x04 \x01(\tR\vaccessKeyId\x12*\n" +
"\x11secret_access_key\x18\x05 \x01(\tR\x0fsecretAccessKey\x12\x15\n" +
"\x06app_id\x18\x06 \x01(\tR\x05appId\x12\x1f\n" +
"\vpath_prefix\x18\a \x01(\tR\n" +
"pathPrefix\"\x85\x01\n" +
"\tVLMConfig\x12\x1d\n" +
"\n" +
"model_name\x18\x01 \x01(\tR\tmodelName\x12\x19\n" +
"\bbase_url\x18\x02 \x01(\tR\abaseUrl\x12\x17\n" +
"\aapi_key\x18\x03 \x01(\tR\x06apiKey\x12%\n" +
"\x0einterface_type\x18\x04 \x01(\tR\rinterfaceType\"\x87\x02\n" +
"\x0einterface_type\x18\x04 \x01(\tR\rinterfaceType\"\x93\x02\n" +
"\n" +
"ReadConfig\x12\x1d\n" +
"\n" +
@@ -648,9 +706,8 @@ const file_docreader_proto_rawDesc = "" +
"\n" +
"separators\x18\x03 \x03(\tR\n" +
"separators\x12+\n" +
"\x11enable_multimodal\x18\x04 \x01(\bR\x10enableMultimodal\x123\n" +
"\n" +
"cos_config\x18\x05 \x01(\v2\x14.docreader.COSConfigR\tcosConfig\x123\n" +
"\x11enable_multimodal\x18\x04 \x01(\bR\x10enableMultimodal\x12?\n" +
"\x0estorage_config\x18\x05 \x01(\v2\x18.docreader.StorageConfigR\rstorageConfig\x123\n" +
"\n" +
"vlm_config\x18\x06 \x01(\v2\x14.docreader.VLMConfigR\tvlmConfig\"\xc9\x01\n" +
"\x13ReadFromFileRequest\x12!\n" +
@@ -683,7 +740,11 @@ const file_docreader_proto_rawDesc = "" +
"\x06images\x18\x05 \x03(\v2\x10.docreader.ImageR\x06images\"N\n" +
"\fReadResponse\x12(\n" +
"\x06chunks\x18\x01 \x03(\v2\x10.docreader.ChunkR\x06chunks\x12\x14\n" +
"\x05error\x18\x02 \x01(\tR\x05error2\x9f\x01\n" +
"\x05error\x18\x02 \x01(\tR\x05error*G\n" +
"\x0fStorageProvider\x12 \n" +
"\x1cSTORAGE_PROVIDER_UNSPECIFIED\x10\x00\x12\a\n" +
"\x03COS\x10\x01\x12\t\n" +
"\x05MINIO\x10\x022\x9f\x01\n" +
"\tDocReader\x12I\n" +
"\fReadFromFile\x12\x1e.docreader.ReadFromFileRequest\x1a\x17.docreader.ReadResponse\"\x00\x12G\n" +
"\vReadFromURL\x12\x1d.docreader.ReadFromURLRequest\x1a\x17.docreader.ReadResponse\"\x00B5Z3github.com/Tencent/WeKnora/internal/docreader/protob\x06proto3"
@@ -700,33 +761,36 @@ func file_docreader_proto_rawDescGZIP() []byte {
return file_docreader_proto_rawDescData
}
var file_docreader_proto_enumTypes = make([]protoimpl.EnumInfo, 1)
var file_docreader_proto_msgTypes = make([]protoimpl.MessageInfo, 8)
var file_docreader_proto_goTypes = []any{
(*COSConfig)(nil), // 0: docreader.COSConfig
(*VLMConfig)(nil), // 1: docreader.VLMConfig
(*ReadConfig)(nil), // 2: docreader.ReadConfig
(*ReadFromFileRequest)(nil), // 3: docreader.ReadFromFileRequest
(*ReadFromURLRequest)(nil), // 4: docreader.ReadFromURLRequest
(*Image)(nil), // 5: docreader.Image
(*Chunk)(nil), // 6: docreader.Chunk
(*ReadResponse)(nil), // 7: docreader.ReadResponse
(StorageProvider)(0), // 0: docreader.StorageProvider
(*StorageConfig)(nil), // 1: docreader.StorageConfig
(*VLMConfig)(nil), // 2: docreader.VLMConfig
(*ReadConfig)(nil), // 3: docreader.ReadConfig
(*ReadFromFileRequest)(nil), // 4: docreader.ReadFromFileRequest
(*ReadFromURLRequest)(nil), // 5: docreader.ReadFromURLRequest
(*Image)(nil), // 6: docreader.Image
(*Chunk)(nil), // 7: docreader.Chunk
(*ReadResponse)(nil), // 8: docreader.ReadResponse
}
var file_docreader_proto_depIdxs = []int32{
0, // 0: docreader.ReadConfig.cos_config:type_name -> docreader.COSConfig
1, // 1: docreader.ReadConfig.vlm_config:type_name -> docreader.VLMConfig
2, // 2: docreader.ReadFromFileRequest.read_config:type_name -> docreader.ReadConfig
2, // 3: docreader.ReadFromURLRequest.read_config:type_name -> docreader.ReadConfig
5, // 4: docreader.Chunk.images:type_name -> docreader.Image
6, // 5: docreader.ReadResponse.chunks:type_name -> docreader.Chunk
3, // 6: docreader.DocReader.ReadFromFile:input_type -> docreader.ReadFromFileRequest
4, // 7: docreader.DocReader.ReadFromURL:input_type -> docreader.ReadFromURLRequest
7, // 8: docreader.DocReader.ReadFromFile:output_type -> docreader.ReadResponse
7, // 9: docreader.DocReader.ReadFromURL:output_type -> docreader.ReadResponse
8, // [8:10] is the sub-list for method output_type
6, // [6:8] is the sub-list for method input_type
6, // [6:6] is the sub-list for extension type_name
6, // [6:6] is the sub-list for extension extendee
0, // [0:6] is the sub-list for field type_name
0, // 0: docreader.StorageConfig.provider:type_name -> docreader.StorageProvider
1, // 1: docreader.ReadConfig.storage_config:type_name -> docreader.StorageConfig
2, // 2: docreader.ReadConfig.vlm_config:type_name -> docreader.VLMConfig
3, // 3: docreader.ReadFromFileRequest.read_config:type_name -> docreader.ReadConfig
3, // 4: docreader.ReadFromURLRequest.read_config:type_name -> docreader.ReadConfig
6, // 5: docreader.Chunk.images:type_name -> docreader.Image
7, // 6: docreader.ReadResponse.chunks:type_name -> docreader.Chunk
4, // 7: docreader.DocReader.ReadFromFile:input_type -> docreader.ReadFromFileRequest
5, // 8: docreader.DocReader.ReadFromURL:input_type -> docreader.ReadFromURLRequest
8, // 9: docreader.DocReader.ReadFromFile:output_type -> docreader.ReadResponse
8, // 10: docreader.DocReader.ReadFromURL:output_type -> docreader.ReadResponse
9, // [9:11] is the sub-list for method output_type
7, // [7:9] is the sub-list for method input_type
7, // [7:7] is the sub-list for extension type_name
7, // [7:7] is the sub-list for extension extendee
0, // [0:7] is the sub-list for field type_name
}
func init() { file_docreader_proto_init() }
@@ -739,13 +803,14 @@ func file_docreader_proto_init() {
File: protoimpl.DescBuilder{
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
RawDescriptor: unsafe.Slice(unsafe.StringData(file_docreader_proto_rawDesc), len(file_docreader_proto_rawDesc)),
NumEnums: 0,
NumEnums: 1,
NumMessages: 8,
NumExtensions: 0,
NumServices: 1,
},
GoTypes: file_docreader_proto_goTypes,
DependencyIndexes: file_docreader_proto_depIdxs,
EnumInfos: file_docreader_proto_enumTypes,
MessageInfos: file_docreader_proto_msgTypes,
}.Build()
File_docreader_proto = out.File
+17 -9
View File
@@ -12,14 +12,22 @@ service DocReader {
rpc ReadFromURL(ReadFromURLRequest) returns (ReadResponse) {}
}
// COS 配置
message COSConfig {
string secret_id = 1; // COS Secret ID
string secret_key = 2; // COS Secret Key
string region = 3; // COS Region
string bucket_name = 4; // COS Bucket Name
string app_id = 5; // COS App ID
string path_prefix = 6; // COS Path Prefix
// 对象存储提供方
enum StorageProvider {
STORAGE_PROVIDER_UNSPECIFIED = 0;
COS = 1; // 腾讯云 COS
MINIO = 2; // MinIO/S3 兼容
}
// 通用对象存储配置,兼容 COS 与 MinIO
message StorageConfig {
StorageProvider provider = 1; // 存储提供方
string region = 2; // 区域(COS 使用)
string bucket_name = 3; // 桶名
string access_key_id = 4; // 访问密钥 IDMinIO/S3 使用)
string secret_access_key = 5; // 访问密钥 SecretMinIO/S3 使用)
string app_id = 6; // 应用 IDCOS 使用)
string path_prefix = 7; // 路径前缀
}
// VLM 配置
@@ -35,7 +43,7 @@ message ReadConfig {
int32 chunk_overlap = 2; // 分块重叠
repeated string separators = 3; // 分隔符
bool enable_multimodal = 4; // 多模态处理
COSConfig cos_config = 5; // COS 配置
StorageConfig storage_config = 5; // 对象存储配置(通用)
VLMConfig vlm_config = 6; // VLM 配置
}
+49 -47
View File
@@ -14,7 +14,8 @@ if parent_dir not in sys.path:
from proto.docreader_pb2 import ReadResponse, Chunk, Image
from proto import docreader_pb2_grpc
from parser import Parser, ChunkingConfig, OCREngine
from parser import Parser, OCREngine
from parser.config import ChunkingConfig
from utils.request import request_id_context, init_logging_request_id
# Ensure no existing handlers
@@ -74,38 +75,39 @@ class DocReaderServicer(docreader_pb2_grpc.DocReaderServicer):
f"multimodal={enable_multimodal}"
)
# Get COS and VLM config from request
cos_config = None
# Get Storage and VLM config from request
storage_config = None
vlm_config = None
if hasattr(request.read_config, 'cos_config') and request.read_config.cos_config:
cos_config = {
'secret_id': request.read_config.cos_config.secret_id,
'secret_key': request.read_config.cos_config.secret_key,
'region': request.read_config.cos_config.region,
'bucket_name': request.read_config.cos_config.bucket_name,
'app_id': request.read_config.cos_config.app_id,
'path_prefix': request.read_config.cos_config.path_prefix or '',
}
logger.info(f"Using COS config: region={cos_config['region']}, bucket={cos_config['bucket_name']}")
sc = request.read_config.storage_config
# Keep parser-side key name as cos_config for backward compatibility
storage_config = {
'provider': 'minio' if sc.provider == 2 else 'cos',
'region': sc.region,
'bucket_name': sc.bucket_name,
'access_key_id': sc.access_key_id,
'secret_access_key': sc.secret_access_key,
'app_id': sc.app_id,
'path_prefix': sc.path_prefix,
}
logger.info(f"Using Storage config: provider={storage_config.get('provider')}, bucket={storage_config['bucket_name']}")
if hasattr(request.read_config, 'vlm_config') and request.read_config.vlm_config:
vlm_config = {
'model_name': request.read_config.vlm_config.model_name,
'base_url': request.read_config.vlm_config.base_url,
'api_key': request.read_config.vlm_config.api_key or '',
'interface_type': request.read_config.vlm_config.interface_type or 'openai',
}
logger.info(f"Using VLM config: model={vlm_config['model_name']}, "
f"base_url={vlm_config['base_url']}, "
f"interface_type={vlm_config['interface_type']}")
vlm_config = {
'model_name': request.read_config.vlm_config.model_name,
'base_url': request.read_config.vlm_config.base_url,
'api_key': request.read_config.vlm_config.api_key or '',
'interface_type': request.read_config.vlm_config.interface_type or 'openai',
}
logger.info(f"Using VLM config: model={vlm_config['model_name']}, "
f"base_url={vlm_config['base_url']}, "
f"interface_type={vlm_config['interface_type']}")
chunking_config = ChunkingConfig(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
separators=separators,
enable_multimodal=enable_multimodal,
cos_config=cos_config,
storage_config=storage_config,
vlm_config=vlm_config,
)
@@ -166,38 +168,38 @@ class DocReaderServicer(docreader_pb2_grpc.DocReaderServicer):
f"multimodal={enable_multimodal}"
)
# Get COS and VLM config from request
cos_config = None
# Get Storage and VLM config from request
storage_config = None
vlm_config = None
if hasattr(request.read_config, 'cos_config') and request.read_config.cos_config:
cos_config = {
'secret_id': request.read_config.cos_config.secret_id,
'secret_key': request.read_config.cos_config.secret_key,
'region': request.read_config.cos_config.region,
'bucket_name': request.read_config.cos_config.bucket_name,
'app_id': request.read_config.cos_config.app_id,
'path_prefix': request.read_config.cos_config.path_prefix or '',
}
logger.info(f"Using COS config: region={cos_config['region']}, bucket={cos_config['bucket_name']}")
sc = request.read_config.storage_config
storage_config = {
'provider': 'minio' if sc.provider == 2 else 'cos',
'region': sc.region,
'bucket_name': sc.bucket_name,
'access_key_id': sc.access_key_id,
'secret_access_key': sc.secret_access_key,
'app_id': sc.app_id,
'path_prefix': sc.path_prefix,
}
logger.info(f"Using Storage config: provider={storage_config.get('provider')}, bucket={storage_config['bucket_name']}")
if hasattr(request.read_config, 'vlm_config') and request.read_config.vlm_config:
vlm_config = {
'model_name': request.read_config.vlm_config.model_name,
'base_url': request.read_config.vlm_config.base_url,
'api_key': request.read_config.vlm_config.api_key or '',
'interface_type': request.read_config.vlm_config.interface_type or 'openai',
}
logger.info(f"Using VLM config: model={vlm_config['model_name']}, "
f"base_url={vlm_config['base_url']}, "
f"interface_type={vlm_config['interface_type']}")
vlm_config = {
'model_name': request.read_config.vlm_config.model_name,
'base_url': request.read_config.vlm_config.base_url,
'api_key': request.read_config.vlm_config.api_key or '',
'interface_type': request.read_config.vlm_config.interface_type or 'openai',
}
logger.info(f"Using VLM config: model={vlm_config['model_name']}, "
f"base_url={vlm_config['base_url']}, "
f"interface_type={vlm_config['interface_type']}")
chunking_config = ChunkingConfig(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
separators=separators,
enable_multimodal=enable_multimodal,
cos_config=cos_config,
storage_config=storage_config,
vlm_config=vlm_config,
)