Skip to content

Commit 4b28b57

Browse files
heiruwuheiruwu
authored
feat(conversion): support instill model in doc conversion process (#154)
Because - host docling model in Instill Model with CUDA acceleration This commit - support instill model in doc conversion process
1 parent 12eb600 commit 4b28b57

File tree

14 files changed

+201
-17
lines changed

14 files changed

+201
-17
lines changed

‎cmd/main/main.go‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@ import (
5555
servicepkg "github.com/instill-ai/artifact-backend/pkg/service"
5656
artifactpb "github.com/instill-ai/protogen-go/artifact/artifact/v1alpha"
5757
mgmtpb "github.com/instill-ai/protogen-go/core/mgmt/v1beta"
58+
modelpb "github.com/instill-ai/protogen-go/model/model/v1alpha"
5859
pipelinepb "github.com/instill-ai/protogen-go/pipeline/pipeline/v1beta"
5960
miniox "github.com/instill-ai/x/minio"
6061
)
@@ -124,11 +125,14 @@ func main() {
124125
grpczap.ReplaceGrpcLoggerV2WithVerbosity(logger, 3)
125126

126127
// Initialize clients needed for service
127-
pipelinePublicServiceClient, pipelinePublicGrpcConn, _, mgmtPublicServiceClientConn, mgmtPrivateServiceClient, mgmtPrivateServiceGrpcConn,
128+
pipelinePublicServiceClient, pipelinePublicGrpcConn, modelPublicServiceClient, modelPublicGrpcConn, _, mgmtPublicServiceClientConn, mgmtPrivateServiceClient, mgmtPrivateServiceGrpcConn,
128129
redisClient, influxDBClient, db, minioClient, milvusClient, aclClient, fgaClientConn, fgaReplicaClientConn := newClients(ctx, logger)
129130
if pipelinePublicGrpcConn != nil {
130131
defer pipelinePublicGrpcConn.Close()
131132
}
133+
if modelPublicGrpcConn != nil {
134+
defer modelPublicGrpcConn.Close()
135+
}
132136
if mgmtPublicServiceClientConn != nil {
133137
defer mgmtPublicServiceClientConn.Close()
134138
}
@@ -151,6 +155,7 @@ func main() {
151155
minioClient,
152156
mgmtPrivateServiceClient,
153157
pipelinePublicServiceClient,
158+
modelPublicServiceClient,
154159
httpclient.NewRegistryClient(ctx),
155160
redisClient,
156161
milvusClient,
@@ -325,6 +330,8 @@ func main() {
325330
func newClients(ctx context.Context, logger *zap.Logger) (
326331
pipelinepb.PipelinePublicServiceClient,
327332
*grpc.ClientConn,
333+
modelpb.ModelPublicServiceClient,
334+
*grpc.ClientConn,
328335
mgmtpb.MgmtPublicServiceClient,
329336
*grpc.ClientConn,
330337
mgmtpb.MgmtPrivateServiceClient,
@@ -350,6 +357,8 @@ func newClients(ctx context.Context, logger *zap.Logger) (
350357
}
351358
pipelinePublicServiceClient := pipelinepb.NewPipelinePublicServiceClient(pipelinePublicGrpcConn)
352359

360+
modelPublicServiceClient, modelPublicServiceClientConn := grpcclient.NewModelPublicClient(ctx)
361+
353362
// initialize mgmt clients
354363
mgmtPrivateServiceClient, mgmtPrivateServiceClientConn := grpcclient.NewMGMTPrivateClient(ctx)
355364
mgmtPublicServiceClient, mgmtPublicServiceClientConn := grpcclient.NewMGMTPublicClient(ctx)
@@ -400,7 +409,7 @@ func newClients(ctx context.Context, logger *zap.Logger) (
400409

401410
}
402411
aclClient := acl.NewACLClient(fgaClient, fgaReplicaClient, redisClient)
403-
return pipelinePublicServiceClient, pipelinePublicGrpcConn, mgmtPublicServiceClient, mgmtPrivateServiceClientConn, mgmtPrivateServiceClient, mgmtPublicServiceClientConn, redisClient, influxDBClient, db, minioClient, milvusClient, aclClient, fgaClientConn, fgaReplicaClientConn
412+
return pipelinePublicServiceClient, pipelinePublicGrpcConn, modelPublicServiceClient, modelPublicServiceClientConn, mgmtPublicServiceClient, mgmtPrivateServiceClientConn, mgmtPrivateServiceClient, mgmtPublicServiceClientConn, redisClient, influxDBClient, db, minioClient, milvusClient, aclClient, fgaClientConn, fgaReplicaClientConn
404413
}
405414

406415
func newGrpcOptionAndCreds(logger *zap.Logger) ([]grpc.ServerOption, credentials.TransportCredentials) {

‎config/config.go‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@ type AppConfig struct {
2828
Log LogConfig `koanf:"log"`
2929
MgmtBackend MgmtBackendConfig `koanf:"mgmtbackend"`
3030
PipelineBackend PipelineBackendConfig `koanf:"pipelinebackend"`
31+
ModelBackend ModelBackendConfig `koanf:"modelbackend"`
3132
Registry RegistryConfig `koanf:"registry"`
3233
OpenFGA OpenFGAConfig `koanf:"openfga"`
3334
Minio miniox.Config `koanf:"minio"`
@@ -138,6 +139,18 @@ type PipelineBackendConfig struct {
138139
}
139140
}
140141

142+
// ModelBackendConfig related to model-backend
143+
type ModelBackendConfig struct {
144+
Host string `koanf:"host"`
145+
PublicPort int `koanf:"publicport"`
146+
PrivatePort int `koanf:"privateport"`
147+
Namespace string `koanf:"namespace"`
148+
HTTPS struct {
149+
Cert string `koanf:"cert"`
150+
Key string `koanf:"key"`
151+
}
152+
}
153+
141154
// RegistryConfig is the registry configuration.
142155
type RegistryConfig struct {
143156
Host string `koanf:"host"`

‎config/config.yaml‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,14 @@ pipelinebackend:
6060
https:
6161
cert:
6262
key:
63+
modelbackend:
64+
host: model-backend
65+
publicport: 8083
66+
privateport: 3083
67+
namespace: admin
68+
https:
69+
cert:
70+
key:
6371
registry:
6472
host: registry
6573
port: 5000

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ require (
1313
github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
1414
github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.1
1515
github.com/influxdata/influxdb-client-go/v2 v2.12.3
16-
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20250218192547-887cb37e3b6e
16+
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20250219010801-09f3eb30f063
1717
github.com/instill-ai/usage-client v0.3.0-alpha.0.20240319060111-4a3a39f2fd61
1818
github.com/instill-ai/x v0.6.0-alpha.0.20250220113648-be48bc78368d
1919
github.com/knadh/koanf v1.5.0

‎go.sum‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -350,8 +350,8 @@ github.com/influxdata/influxdb-client-go/v2 v2.12.3 h1:28nRlNMRIV4QbtIUvxhWqaxn0
350350
github.com/influxdata/influxdb-client-go/v2 v2.12.3/go.mod h1:IrrLUbCjjfkmRuaCiGQg4m2GbkaeJDcuWoxiWdQEbA0=
351351
github.com/influxdata/line-protocol v0.0.0-20200327222509-2487e7298839 h1:W9WBk7wlPfJLvMCdtV4zPulc4uCPrlywQOmbFOhgQNU=
352352
github.com/influxdata/line-protocol v0.0.0-20200327222509-2487e7298839/go.mod h1:xaLFMmpvUxqXtVkUJfg9QmT88cDaCJ3ZKgdZ78oO8Qo=
353-
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20250218192547-887cb37e3b6e h1:DJUavDxLUzayfMYeysdDhLh9Ajf4UPcduShB7uPdKjs=
354-
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20250218192547-887cb37e3b6e/go.mod h1:fusT92ceR5+GVn1LT5mT4XcOq1DlemBjpb6JpodLLdc=
353+
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20250219010801-09f3eb30f063 h1:zf6H3XmrHiVRhsYw44oeDVh4ATS0T6HchWWXBFKF0a0=
354+
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20250219010801-09f3eb30f063/go.mod h1:fusT92ceR5+GVn1LT5mT4XcOq1DlemBjpb6JpodLLdc=
355355
github.com/instill-ai/usage-client v0.3.0-alpha.0.20240319060111-4a3a39f2fd61 h1:smPTvmXDhn/QC7y/TPXyMTqbbRd0gvzmFgWBChwTfhE=
356356
github.com/instill-ai/usage-client v0.3.0-alpha.0.20240319060111-4a3a39f2fd61/go.mod h1:/TAHs4ybuylk5icuy+MQtHRc4XUnIyXzeNKxX9qDFhw=
357357
github.com/instill-ai/x v0.6.0-alpha.0.20250220113648-be48bc78368d h1:HsnJ1inBFWXurGVLsH7jJFkqRYU1O2Nwg3ixtOCk98E=

‎pkg/client/grpc/model.go‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package grpcclient
2+
3+
import (
4+
"context"
5+
"fmt"
6+
7+
"google.golang.org/grpc"
8+
"google.golang.org/grpc/credentials"
9+
"google.golang.org/grpc/credentials/insecure"
10+
11+
"github.com/instill-ai/artifact-backend/config"
12+
"github.com/instill-ai/artifact-backend/pkg/constant"
13+
"github.com/instill-ai/artifact-backend/pkg/logger"
14+
15+
pb "github.com/instill-ai/protogen-go/model/model/v1alpha"
16+
)
17+
18+
// NewModelPublicClient returns an initialized gRPC client for the Model public
19+
// API.
20+
func NewModelPublicClient(ctx context.Context) (pb.ModelPublicServiceClient, *grpc.ClientConn) {
21+
logger, _ := logger.GetZapLogger(ctx)
22+
23+
var clientDialOpts grpc.DialOption
24+
if config.Config.ModelBackend.HTTPS.Cert != "" && config.Config.ModelBackend.HTTPS.Key != "" {
25+
creds, err := credentials.NewServerTLSFromFile(config.Config.ModelBackend.HTTPS.Cert, config.Config.ModelBackend.HTTPS.Key)
26+
if err != nil {
27+
logger.Fatal(err.Error())
28+
}
29+
clientDialOpts = grpc.WithTransportCredentials(creds)
30+
} else {
31+
clientDialOpts = grpc.WithTransportCredentials(insecure.NewCredentials())
32+
}
33+
34+
clientConn, err := grpc.NewClient(fmt.Sprintf("%v:%v", config.Config.ModelBackend.Host, config.Config.ModelBackend.PublicPort), clientDialOpts, grpc.WithDefaultCallOptions(grpc.MaxCallRecvMsgSize(constant.MaxPayloadSize), grpc.MaxCallSendMsgSize(constant.MaxPayloadSize)))
35+
if err != nil {
36+
logger.Error(err.Error())
37+
return nil, nil
38+
}
39+
40+
return pb.NewModelPublicServiceClient(clientConn), clientConn
41+
}

‎pkg/constant/constant.go‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,4 +31,3 @@ type ContentType string
3131
var ChunkContentType ContentType = "chunk"
3232
var SummaryContentType ContentType = "summary"
3333
var AugmentedContentType ContentType = "augmented"
34-

‎pkg/handler/knowledgebase.go‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import (
1010
"go.uber.org/zap"
1111
"google.golang.org/grpc/metadata"
1212

13+
"github.com/instill-ai/artifact-backend/config"
1314
"github.com/instill-ai/artifact-backend/pkg/constant"
1415
"github.com/instill-ai/artifact-backend/pkg/customerror"
1516
"github.com/instill-ai/artifact-backend/pkg/logger"
@@ -151,7 +152,7 @@ func (ph *PublicHandler) CreateCatalog(ctx context.Context, req *artifactpb.Crea
151152
CreateTime: dbData.CreateTime.String(),
152153
UpdateTime: dbData.UpdateTime.String(),
153154
ConvertingPipelines: []string{
154-
service.NamespaceID + "/" + service.ConvertDocToMDPipelineID + "@" + service.DocToMDVersion,
155+
config.Config.ModelBackend.Namespace + "/" + service.ChunkMdPipelineID + "@" + service.ConvertDocToMDModelVersion,
155156
},
156157
SummarizingPipelines: []string{
157158
service.NamespaceID + "/" + service.GenerateSummaryPipelineID + "@" + service.GenerateSummaryVersion,
@@ -234,7 +235,7 @@ func (ph *PublicHandler) ListCatalogs(ctx context.Context, req *artifactpb.ListC
234235
UpdateTime: kb.UpdateTime.String(),
235236
OwnerName: kb.Owner,
236237
ConvertingPipelines: []string{
237-
service.NamespaceID + "/" + service.ConvertDocToMDPipelineID + "@" + service.DocToMDVersion,
238+
config.Config.ModelBackend.Namespace + "/" + service.ChunkMdPipelineID + "@" + service.ConvertDocToMDModelVersion,
238239
},
239240
SummarizingPipelines: []string{
240241
service.NamespaceID + "/" + service.GenerateSummaryPipelineID + "@" + service.GenerateSummaryVersion,
@@ -331,7 +332,7 @@ func (ph *PublicHandler) UpdateCatalog(ctx context.Context, req *artifactpb.Upda
331332
UpdateTime: kb.UpdateTime.String(),
332333
OwnerName: kb.Owner,
333334
ConvertingPipelines: []string{
334-
service.NamespaceID + "/" + service.ConvertDocToMDPipelineID,
335+
config.Config.ModelBackend.Namespace + "/" + service.ChunkMdPipelineID + "@" + service.ConvertDocToMDModelVersion,
335336
},
336337
SummarizingPipelines: []string{
337338
service.NamespaceID + "/" + service.GenerateSummaryPipelineID + "@" + service.GenerateSummaryVersion,

‎pkg/repository/knowledgebasefile.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -693,7 +693,7 @@ func (r *Repository) UpdateKbFileExtraMetaData(
693693
summarizingPipe,
694694
chunkingPipe,
695695
embeddingPipe string,
696-
processingTime, summarizingTime, convertingTime, chunkingTime, embeddingTime *int64) error {
696+
processingTime, convertingTime, summarizingTime, chunkingTime, embeddingTime *int64) error {
697697
var kb KnowledgeBaseFile
698698

699699
// Use GORM's Transaction function

‎pkg/service/model.go‎

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
package service
2+
3+
import (
4+
"context"
5+
"fmt"
6+
"strings"
7+
8+
"github.com/gofrs/uuid"
9+
"go.uber.org/zap"
10+
"google.golang.org/grpc/metadata"
11+
"google.golang.org/protobuf/types/known/structpb"
12+
13+
"github.com/instill-ai/artifact-backend/config"
14+
"github.com/instill-ai/artifact-backend/pkg/constant"
15+
"github.com/instill-ai/artifact-backend/pkg/logger"
16+
17+
artifactpb "github.com/instill-ai/protogen-go/artifact/artifact/v1alpha"
18+
modelpb "github.com/instill-ai/protogen-go/model/model/v1alpha"
19+
)
20+
21+
const ConvertDocToMDModelID = "docling"
22+
const ConvertDocToMDModelVersion = "v0.1.0"
23+
24+
// ConvertToMDModel using docling model to convert some file type to MD and consume caller's credits
25+
func (s *Service) ConvertToMDModel(ctx context.Context, fileUID uuid.UUID, caller uuid.UUID, requester uuid.UUID, fileBase64 string, fileType artifactpb.FileType) (string, error) {
26+
logger, _ := logger.GetZapLogger(ctx)
27+
var md metadata.MD
28+
if requester != uuid.Nil {
29+
md = metadata.New(map[string]string{
30+
constant.HeaderUserUIDKey: caller.String(),
31+
constant.HeaderAuthTypeKey: "user",
32+
constant.HeaderRequesterUIDKey: requester.String(),
33+
})
34+
} else {
35+
md = metadata.New(map[string]string{
36+
constant.HeaderUserUIDKey: caller.String(),
37+
constant.HeaderAuthTypeKey: "user",
38+
})
39+
}
40+
ctx = metadata.NewOutgoingContext(ctx, md)
41+
42+
// Get the appropriate prefix for the file type
43+
prefix := getFileTypePrefix(fileType)
44+
45+
req := &modelpb.TriggerNamespaceModelRequest{
46+
NamespaceId: config.Config.ModelBackend.Namespace,
47+
ModelId: ConvertDocToMDModelID,
48+
Version: ConvertDocToMDModelVersion,
49+
TaskInputs: []*structpb.Struct{
50+
{
51+
Fields: map[string]*structpb.Value{
52+
"data": {Kind: &structpb.Value_StructValue{
53+
StructValue: &structpb.Struct{
54+
Fields: map[string]*structpb.Value{
55+
"doc_content": {Kind: &structpb.Value_StringValue{StringValue: prefix + fileBase64}},
56+
},
57+
},
58+
}},
59+
},
60+
},
61+
},
62+
}
63+
64+
resp, err := s.ModelPub.TriggerNamespaceModel(ctx, req)
65+
if err != nil {
66+
logger.Error("failed to trigger model", zap.Error(err))
67+
return "", fmt.Errorf("failed to trigger %s model: %w", ConvertDocToMDModelID, err)
68+
}
69+
70+
result, err := getModelConvertResult(resp)
71+
if err != nil {
72+
logger.Error("failed to get convert result", zap.Error(err))
73+
return "", fmt.Errorf("failed to get convert result: %w", err)
74+
}
75+
return result, nil
76+
}
77+
78+
func getModelConvertResult(resp *modelpb.TriggerNamespaceModelResponse) (string, error) {
79+
if resp == nil || len(resp.TaskOutputs) == 0 {
80+
return "", fmt.Errorf("response is nil or has no outputs. resp: %v", resp)
81+
}
82+
fields := resp.TaskOutputs[0].GetFields()["data"].GetStructValue().GetFields()
83+
if fields == nil {
84+
return "", fmt.Errorf("fields in the output are nil. resp: %v", resp)
85+
}
86+
markdownPages, ok := fields["markdown_pages"]
87+
if !ok {
88+
return "", fmt.Errorf("markdown_pages not found in the output fields. resp: %v", resp)
89+
}
90+
91+
// Not used at the moment
92+
// extractedImages, _ := fields["extracted_images"]
93+
// pagesWithImages, _ := fields["pages_with_images"]
94+
95+
markdownPagesSlice := []string{}
96+
for _, v := range markdownPages.GetListValue().GetValues() {
97+
markdownPagesSlice = append(markdownPagesSlice, v.GetStringValue())
98+
}
99+
100+
return strings.Join(markdownPagesSlice, "\n"), nil
101+
}

0 commit comments

Comments
 (0)