133 lines
3.1 KiB
Go
133 lines
3.1 KiB
Go
|
package service
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"context"
|
|||
|
|
"fmt"
|
|||
|
|
"time"
|
|||
|
|
|
|||
|
|
"google.golang.org/grpc"
|
|||
|
|
"google.golang.org/grpc/credentials/insecure"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
// 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
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 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 解析文档
|
|||
|
|
func (c *AICoreClient) ParseDocument(fileURL, fileName, fileType string) (*ParseResult, error) {
|
|||
|
|
if c.conn == nil {
|
|||
|
|
if err := c.Connect(); err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 使用 gRPC raw bytes 调用
|
|||
|
|
// 由于没有生成 protobuf 代码，使用 raw bytes 方式调用
|
|||
|
|
client := NewDocumentParserClient(c.conn)
|
|||
|
|
|
|||
|
|
req := &ParseRequest{
|
|||
|
|
FileUrl: fileURL,
|
|||
|
|
FileName: fileName,
|
|||
|
|
FileType: fileType,
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
ctx, cancel := context.WithTimeout(context.Background(), 60*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.Success,
|
|||
|
|
Content: resp.Content,
|
|||
|
|
Message: resp.Message,
|
|||
|
|
ContentLength: resp.ContentLength,
|
|||
|
|
FileType: resp.FileType,
|
|||
|
|
ParserEngine: resp.ParserEngine,
|
|||
|
|
}, nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 以下是手动定义的 protobuf messages（与 proto 文件一致）
|
|||
|
|
// 不需要生成 .pb.go 文件，直接手动定义
|
|||
|
|
|
|||
|
|
type ParseRequest struct {
|
|||
|
|
FileUrl string
|
|||
|
|
FileName string
|
|||
|
|
FileType string
|
|||
|
|
ParserEngine string
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type ParseResponse struct {
|
|||
|
|
Success bool
|
|||
|
|
Content string
|
|||
|
|
Message string
|
|||
|
|
ContentLength int32
|
|||
|
|
FileType string
|
|||
|
|
ParserEngine string
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// DocumentParserClient gRPC 客户端接口（手动实现）
|
|||
|
|
type DocumentParserClient interface {
|
|||
|
|
ParseDocument(ctx context.Context, in *ParseRequest, opts ...grpc.CallOption) (*ParseResponse, error)
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
type documentParserClient struct {
|
|||
|
|
cc grpc.ClientConnInterface
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// NewDocumentParserClient 创建 DocumentParser 客户端
|
|||
|
|
func NewDocumentParserClient(cc grpc.ClientConnInterface) DocumentParserClient {
|
|||
|
|
return &documentParserClient{cc: cc}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
func (c *documentParserClient) ParseDocument(ctx context.Context, in *ParseRequest, opts ...grpc.CallOption) (*ParseResponse, error) {
|
|||
|
|
out := new(ParseResponse)
|
|||
|
|
err := c.cc.Invoke(ctx, "/docparser.DocumentParser/ParseDocument", in, out, opts...)
|
|||
|
|
if err != nil {
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
return out, nil
|
|||
|
|
}
|