mirror of
https://github.com/Tencent/WeKnora.git
synced 2026-09-01 14:53:07 +08:00
feat: support minio storage
This commit is contained in:
+5
-11
@@ -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
|
||||
|
||||
|
||||
@@ -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
@@ -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:
|
||||
|
||||
@@ -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 \
|
||||
|
||||
@@ -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
@@ -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
@@ -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/AppID;PathPrefix 可选
|
||||
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 {
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
@@ -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}"
|
||||
|
||||
@@ -17,6 +17,7 @@ paddleocr==3.0.0
|
||||
markdown
|
||||
pypdf
|
||||
cos-python-sdk-v5
|
||||
minio
|
||||
textract
|
||||
antiword
|
||||
openai
|
||||
|
||||
Regular → Executable
@@ -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
|
||||
|
||||
@@ -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 = [
|
||||
|
||||
@@ -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
@@ -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""
|
||||
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""
|
||||
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:
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"` // 访问密钥 ID(MinIO/S3 使用)
|
||||
SecretAccessKey string `protobuf:"bytes,5,opt,name=secret_access_key,json=secretAccessKey,proto3" json:"secret_access_key,omitempty"` // 访问密钥 Secret(MinIO/S3 使用)
|
||||
AppId string `protobuf:"bytes,6,opt,name=app_id,json=appId,proto3" json:"app_id,omitempty"` // 应用 ID(COS 使用)
|
||||
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
|
||||
|
||||
@@ -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; // 访问密钥 ID(MinIO/S3 使用)
|
||||
string secret_access_key = 5; // 访问密钥 Secret(MinIO/S3 使用)
|
||||
string app_id = 6; // 应用 ID(COS 使用)
|
||||
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 配置
|
||||
}
|
||||
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user