DESKTOP-72TV0V4\caoxiaozhu
0a9f6e278e
- 更新文档解析客户端 - 优化知识库服务逻辑 - 更新 protobuf 定义 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
115 lines
2.6 KiB
Go
115 lines
2.6 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"time"
|
|
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
|
|
docparser "x-agents/server"
|
|
)
|
|
|
|
// AICoreClient AI-Core 文档解析服务客户端
|
|
type AICoreClient struct {
|
|
conn *grpc.ClientConn
|
|
address string
|
|
}
|
|
|
|
// ParseResult 解析结果
|
|
type ParseResult struct {
|
|
Success bool
|
|
Content string
|
|
Message string
|
|
ContentLength int32
|
|
FileType string
|
|
ParserEngine string
|
|
}
|
|
|
|
// VLMConfig VLM 模型配置
|
|
type VLMConfig struct {
|
|
Enabled bool
|
|
Provider string // openai, anthropic, local 等
|
|
Model string
|
|
APIKey string
|
|
BaseURL string
|
|
Prompt string
|
|
}
|
|
|
|
// NewAICoreClient 创建 AI-Core 客户端
|
|
func NewAICoreClient(address string) (*AICoreClient, error) {
|
|
return &AICoreClient{address: address}, nil
|
|
}
|
|
|
|
// Connect 连接到 gRPC 服务
|
|
func (c *AICoreClient) Connect() error {
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
defer cancel()
|
|
|
|
conn, err := grpc.DialContext(ctx, c.address,
|
|
grpc.WithTransportCredentials(insecure.NewCredentials()),
|
|
grpc.WithBlock(),
|
|
)
|
|
if err != nil {
|
|
return fmt.Errorf("failed to connect to AI-Core service: %w", err)
|
|
}
|
|
c.conn = conn
|
|
return nil
|
|
}
|
|
|
|
// Close 关闭连接
|
|
func (c *AICoreClient) Close() {
|
|
if c.conn != nil {
|
|
c.conn.Close()
|
|
}
|
|
}
|
|
|
|
// ParseDocument 解析文档 - 使用生成的 protobuf 代码
|
|
// vlmConfig 可选，如果不使用 VLM 传 nil
|
|
func (c *AICoreClient) ParseDocument(fileURL, fileName, fileType string, vlmConfig *VLMConfig) (*ParseResult, error) {
|
|
if c.conn == nil {
|
|
if err := c.Connect(); err != nil {
|
|
return nil, err
|
|
}
|
|
}
|
|
|
|
// 使用生成的 protobuf 客户端
|
|
client := docparser.NewDocumentParserClient(c.conn)
|
|
|
|
req := &docparser.ParseRequest{
|
|
FileUrl: fileURL,
|
|
FileName: fileName,
|
|
FileType: fileType,
|
|
}
|
|
|
|
// 如果提供了 VLM 配置，添加到请求中
|
|
if vlmConfig != nil {
|
|
req.VlmConfig = &docparser.VLMConfig{
|
|
Enabled: vlmConfig.Enabled,
|
|
Provider: vlmConfig.Provider,
|
|
Model: vlmConfig.Model,
|
|
ApiKey: vlmConfig.APIKey,
|
|
BaseUrl: vlmConfig.BaseURL,
|
|
Prompt: vlmConfig.Prompt,
|
|
}
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
|
|
defer cancel()
|
|
|
|
resp, err := client.ParseDocument(ctx, req)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("failed to parse document: %w", err)
|
|
}
|
|
|
|
return &ParseResult{
|
|
Success: resp.GetSuccess(),
|
|
Content: resp.GetContent(),
|
|
Message: resp.GetMessage(),
|
|
ContentLength: resp.GetContentLength(),
|
|
FileType: resp.GetFileType(),
|
|
ParserEngine: resp.GetParserEngine(),
|
|
}, nil
|
|
}
|