Skip to content

Commit

Permalink
feat(kb): support knowledge base file related api (#23)
Browse files Browse the repository at this point in the history
Because

user should be able to add files into knowledge base.

This commit

add file related APIs
  • Loading branch information
Yougigun authored Jun 12, 2024
1 parent c57330c commit 3912028
Show file tree
Hide file tree
Showing 21 changed files with 2,724 additions and 248 deletions.
2 changes: 1 addition & 1 deletion Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ run-local:
echo "A container named ${SERVICE_NAME} is already running. \nRestarting..."; \
make rm; \
fi
@docker run -d --rm \
@docker run -d --rm \
-p ${SERVICE_PORT}:${SERVICE_PORT} \
--network instill-network \
--name ${SERVICE_NAME} \
Expand Down
7 changes: 6 additions & 1 deletion cmd/main/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@ import (
httpclient "github.com/instill-ai/artifact-backend/pkg/client/http"
database "github.com/instill-ai/artifact-backend/pkg/db"
custom_otel "github.com/instill-ai/artifact-backend/pkg/logger/otel"
minio "github.com/instill-ai/artifact-backend/pkg/minio"
artifactPB "github.com/instill-ai/protogen-go/artifact/artifact/v1alpha"
)

Expand Down Expand Up @@ -178,7 +179,11 @@ func main() {

repository := repository.NewRepository(db)

service := service.NewService(repository, httpclient.NewRegistryClient(ctx))
// Initialize Minio client
minioClient, err := minio.NewMinioClientAndInitBucket()


service := service.NewService(repository, minioClient, mgmtPrivateServiceClient, httpclient.NewRegistryClient(ctx))

publicGrpcS := grpc.NewServer(grpcServerOpts...)
reflection.Register(publicGrpcS)
Expand Down
10 changes: 10 additions & 0 deletions config/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@ type AppConfig struct {
MgmtBackend MgmtBackendConfig `koanf:"mgmtbackend"`
Registry RegistryConfig `koanf:"registry"`
OpenFGA OpenFGAConfig `koanf:"openfga"`
Minio MinioConfig `koanf:"minio"`
}

// OpenFGA config
Expand Down Expand Up @@ -121,6 +122,15 @@ type RegistryConfig struct {
Port int `koanf:"port"`
}

// MinioConfig is the minio configuration.
type MinioConfig struct {
Host string `koanf:"host"`
Port string `koanf:"port"`
RootUser string `koanf:"rootuser"`
RootPwd string `koanf:"rootpwd"`
BucketName string `koanf:"bucketname"`
}

// Init - Assign global config to decoded config struct
func Init() error {
k := koanf.New(".")
Expand Down
8 changes: 7 additions & 1 deletion config/config.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@ database:
host: pg-sql
port: 5432
name: artifact
version: 2
version: 4
timezone: Etc/UTC
pool:
idleconnections: 5
Expand Down Expand Up @@ -59,3 +59,9 @@ registry:
openfga:
host: openfga
port: 8080
minio:
host: minio
port: 9000
rootuser: minioadmin
rootpwd: minioadmin
bucketname: instill-ai-knowledge-bases
9 changes: 7 additions & 2 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,15 @@ require (
github.com/gojuno/minimock/v3 v3.3.6
github.com/golang-migrate/migrate/v4 v4.17.0
github.com/google/go-cmp v0.6.0
github.com/google/uuid v1.6.0
github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
github.com/grpc-ecosystem/grpc-gateway/v2 v2.19.1
github.com/influxdata/influxdb-client-go/v2 v2.12.3
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20240530063823-abcb1c583452
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20240612051236-4bd57e13ff0b
github.com/instill-ai/usage-client v0.3.0-alpha.0.20240319060111-4a3a39f2fd61
github.com/instill-ai/x v0.3.0-alpha.0.20231219052200-6230a89e386c
github.com/knadh/koanf v1.5.0
github.com/minio/minio-go v6.0.14+incompatible
github.com/openfga/go-sdk v0.2.3
github.com/redis/go-redis/v9 v9.5.1
go.opentelemetry.io/contrib/propagators/b3 v1.17.0
Expand All @@ -38,7 +40,10 @@ require (
gorm.io/gorm v1.24.7-0.20230306060331-85eaf9eeda11
)

require github.com/google/uuid v1.6.0 // indirect
require (
github.com/go-ini/ini v1.67.0 // indirect
github.com/mitchellh/go-homedir v1.1.0 // indirect
)

require (
github.com/catalinc/hashcash v0.0.0-20220723060415-5e3ec3e24f67 // indirect
Expand Down
9 changes: 7 additions & 2 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -138,6 +138,8 @@ github.com/go-chi/chi/v5 v5.0.0/go.mod h1:BBug9lr0cqtdAhsu6R4AAdvufI0/XBzAQSsUqJ
github.com/go-gl/glfw v0.0.0-20190409004039-e6da0acd62b1/go.mod h1:vR7hzQXu2zJy9AVAgeJqvqgH9Q5CA+iKCZ2gyEVpxRU=
github.com/go-gl/glfw/v3.3/glfw v0.0.0-20191125211704-12ad95a8df72/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=
github.com/go-gl/glfw/v3.3/glfw v0.0.0-20200222043503-6f7a984d4dc4/go.mod h1:tQ2UAYgL5IevRw8kRxooKSPJfGvJ9fJQFa0TUsXzTg8=
github.com/go-ini/ini v1.67.0 h1:z6ZrTEZqSWOTyH2FlglNbNgARyHG8oLW9gMELqKr06A=
github.com/go-ini/ini v1.67.0/go.mod h1:ByCAeIL28uOIIG0E3PJtZPDL8WnHpFKFOtgjp+3Ies8=
github.com/go-kit/kit v0.8.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
github.com/go-kit/kit v0.9.0/go.mod h1:xBxKIO96dXMWWy0MnWVtmwkA9/13aqxPnvrjFYMA2as=
github.com/go-kit/log v0.1.0/go.mod h1:zbhenjAZHb184qTLMA9ZjW7ThYL0H2mk7Q6pNt4vbaY=
Expand Down Expand Up @@ -284,8 +286,8 @@ github.com/influxdata/influxdb-client-go/v2 v2.12.3 h1:28nRlNMRIV4QbtIUvxhWqaxn0
github.com/influxdata/influxdb-client-go/v2 v2.12.3/go.mod h1:IrrLUbCjjfkmRuaCiGQg4m2GbkaeJDcuWoxiWdQEbA0=
github.com/influxdata/line-protocol v0.0.0-20200327222509-2487e7298839 h1:W9WBk7wlPfJLvMCdtV4zPulc4uCPrlywQOmbFOhgQNU=
github.com/influxdata/line-protocol v0.0.0-20200327222509-2487e7298839/go.mod h1:xaLFMmpvUxqXtVkUJfg9QmT88cDaCJ3ZKgdZ78oO8Qo=
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20240530063823-abcb1c583452 h1:uAkZouMby0pjaH0spCa5nUHMIKGjAxaBXRSVhe35M44=
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20240530063823-abcb1c583452/go.mod h1:2blmpUwiTwxIDnrjIqT6FhR5ewshZZF554wzjXFvKpQ=
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20240612051236-4bd57e13ff0b h1:YsH/M29+hdGLU5GS6Ty/JcLY1RMnEPT2lhfE43JUR5E=
github.com/instill-ai/protogen-go v0.3.3-alpha.0.20240612051236-4bd57e13ff0b/go.mod h1:2blmpUwiTwxIDnrjIqT6FhR5ewshZZF554wzjXFvKpQ=
github.com/instill-ai/usage-client v0.3.0-alpha.0.20240319060111-4a3a39f2fd61 h1:smPTvmXDhn/QC7y/TPXyMTqbbRd0gvzmFgWBChwTfhE=
github.com/instill-ai/usage-client v0.3.0-alpha.0.20240319060111-4a3a39f2fd61/go.mod h1:/TAHs4ybuylk5icuy+MQtHRc4XUnIyXzeNKxX9qDFhw=
github.com/instill-ai/x v0.3.0-alpha.0.20231219052200-6230a89e386c h1:a2RVkpIV2QcrGnSHAou+t/L+vBsaIfFvk5inVg5Uh4s=
Expand Down Expand Up @@ -407,11 +409,14 @@ github.com/mattn/go-isatty v0.0.12/go.mod h1:cbi8OIDigv2wuxKPP5vlRcQ1OAZbq2CE4Ky
github.com/matttproud/golang_protobuf_extensions v1.0.1/go.mod h1:D8He9yQNgCq6Z5Ld7szi9bcBfOoFv/3dc6xSMkL2PC0=
github.com/miekg/dns v1.1.26/go.mod h1:bPDLeHnStXmXAq1m/Ch/hvfNHr14JKNPMBo3VZKjuso=
github.com/miekg/dns v1.1.41/go.mod h1:p6aan82bvRIyn+zDIv9xYNUpwa73JcSh9BKwknJysuI=
github.com/minio/minio-go v6.0.14+incompatible h1:fnV+GD28LeqdN6vT2XdGKW8Qe/IfjJDswNVuni6km9o=
github.com/minio/minio-go v6.0.14+incompatible/go.mod h1:7guKYtitv8dktvNUGrhzmNlA5wrAABTQXCoesZdFQO8=
github.com/mitchellh/cli v1.0.0/go.mod h1:hNIlj7HEI86fIcpObd7a0FcrxTWetlwJDGcceTlRvqc=
github.com/mitchellh/cli v1.1.0/go.mod h1:xcISNoH86gajksDmfB23e/pu+B+GeFRMYmoHXxx3xhI=
github.com/mitchellh/copystructure v1.0.0/go.mod h1:SNtv71yrdKgLRyLFxmLdkAbkKEFWgYaq1OVrnRcwhnw=
github.com/mitchellh/copystructure v1.2.0 h1:vpKXTN4ewci03Vljg/q9QvCGUDttBOGBIa15WveJJGw=
github.com/mitchellh/copystructure v1.2.0/go.mod h1:qLl+cE2AmVv+CoeAwDPye/v+N2HKCj9FbZEVFJRxO9s=
github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
github.com/mitchellh/go-testing-interface v0.0.0-20171004221916-a61a99592b77/go.mod h1:kRemZodwjscx+RGhAo8eIhFbs2+BFgRtFPeD/KE+zxI=
github.com/mitchellh/go-testing-interface v1.0.0/go.mod h1:kRemZodwjscx+RGhAo8eIhFbs2+BFgRtFPeD/KE+zxI=
Expand Down
4 changes: 3 additions & 1 deletion pkg/db/migration/000002_create_knowledge_base_table.down.sql
Original file line number Diff line number Diff line change
@@ -1 +1,3 @@
DROP TABLE knowledge_base;
BEGIN;
DROP TABLE knowledge_base;
COMMIT;
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
BEGIN;
-- Drop the table knowledge_base_files
DROP TABLE IF EXISTS knowledge_base;

COMMIT;
33 changes: 33 additions & 0 deletions pkg/db/migration/000003_re_create_knowledge_base_tables.up.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
BEGIN;
-- Drop the existing table if exists to avoid conflicts
DROP TABLE IF EXISTS knowledge_base;
-- Create the new knowledge_base table
CREATE TABLE knowledge_base (
uid UUID PRIMARY KEY NOT NULL DEFAULT gen_random_uuid(),
id VARCHAR(255) NOT NULL,
name VARCHAR(255) NOT NULL,
description VARCHAR(1023),
tags VARCHAR(255) [],
owner UUID NOT NULL,
create_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
update_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
delete_time TIMESTAMP
);
-- Add the index
CREATE UNIQUE INDEX idx_unique_owner_name_delete_time ON knowledge_base (owner, name)
WHERE delete_time IS NULL;
CREATE UNIQUE INDEX unique_owner_id_delete_time ON knowledge_base (owner, id)
WHERE delete_time IS NULL;

-- Comments for the table and columns
COMMENT ON TABLE knowledge_base IS 'Table to store knowledge base information';
COMMENT ON COLUMN knowledge_base.uid IS 'Primary key, auto-generated UUID';
COMMENT ON COLUMN knowledge_base.id IS 'Unique identifier from name for the knowledge base, up to 255 characters created from name';
COMMENT ON COLUMN knowledge_base.name IS 'Name of the knowledge base, up to 255 characters';
COMMENT ON COLUMN knowledge_base.description IS 'Description of the knowledge base, up to 1023 characters';
COMMENT ON COLUMN knowledge_base.tags IS 'Array of tags associated with the knowledge base';
COMMENT ON COLUMN knowledge_base.owner IS 'Owner of the knowledge base. It is a UUID referencing the owner table(uid field).';
COMMENT ON COLUMN knowledge_base.create_time IS 'Timestamp when the entry was created, stored in UTC';
COMMENT ON COLUMN knowledge_base.update_time IS 'Timestamp of the last update, stored in UTC';
COMMENT ON COLUMN knowledge_base.delete_time IS 'Timestamp when the entry was marked as deleted, stored in UTC';
COMMIT;
Original file line number Diff line number Diff line change
@@ -0,0 +1,5 @@
BEGIN;
-- Drop the table knowledge_base_files
DROP TABLE IF EXISTS knowledge_base_files;

COMMIT;
37 changes: 37 additions & 0 deletions pkg/db/migration/000004_create_knowledge_base_file_tables.up.sql
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
BEGIN;
-- Create the table knowledge_base_files
CREATE TABLE knowledge_base_file (
uid UUID PRIMARY KEY DEFAULT gen_random_uuid(),
owner UUID NOT NULL,
kb_uid UUID NOT NULL,
creator_uid UUID NOT NULL,
name VARCHAR(255) NOT NULL,
type VARCHAR(100) NOT NULL,
destination VARCHAR(255) NOT NULL,
process_status VARCHAR(100) NOT NULL,
extra_meta_data JSONB,
content BYTEA,
create_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
update_time TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP,
delete_time TIMESTAMP
);

-- Add the unique index on (kb_uid, name)
CREATE UNIQUE INDEX idx_unique_kb_uid_name_delete_time ON knowledge_base_file (kb_uid, name)
WHERE delete_time IS NULL;

-- Comments for the table and columns
COMMENT ON TABLE knowledge_base_file IS 'Table to store knowledge base files with metadata';
COMMENT ON COLUMN knowledge_base_file.uid IS 'Unique identifier for the file';
COMMENT ON COLUMN knowledge_base_file.create_time IS 'Timestamp when the record was created';
COMMENT ON COLUMN knowledge_base_file.update_time IS 'Timestamp when the record was last updated';
COMMENT ON COLUMN knowledge_base_file.delete_time IS 'Timestamp when the record was deleted (soft delete)';
COMMENT ON COLUMN knowledge_base_file.owner IS 'Owner of the file, references owner table''s uid field';
COMMENT ON COLUMN knowledge_base_file.kb_uid IS 'Knowledge base unique identifier, references knowledge base table''s uid field';
COMMENT ON COLUMN knowledge_base_file.creator_uid IS 'Creator unique identifier, references owner table''s uid field';
COMMENT ON COLUMN knowledge_base_file.name IS 'Name of the file';
COMMENT ON COLUMN knowledge_base_file.process_status IS 'Processing status of the file';
COMMENT ON COLUMN knowledge_base_file.extra_meta_data IS 'Extra metadata in JSON format, e.g., word count, image count';
COMMENT ON COLUMN knowledge_base_file.content IS 'File content stored as byte array';

COMMIT;
68 changes: 41 additions & 27 deletions pkg/handler/knowledgebase.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,7 @@ const ErrorDeleteKnowledgeBaseMsg = "failed to delete knowledge base: %w"

func (ph *PublicHandler) CreateKnowledgeBase(ctx context.Context, req *artifactpb.CreateKnowledgeBaseRequest) (*artifactpb.CreateKnowledgeBaseResponse, error) {

uid, err := getUserIDFromContext(ctx)
uid, err := getUserUIDFromContext(ctx)
if err != nil {
err := fmt.Errorf("failed to get user id from header: %v. err: %w", err, customerror.ErrUnauthenticated)
return nil, err
Expand All @@ -38,10 +38,11 @@ func (ph *PublicHandler) CreateKnowledgeBase(ctx context.Context, req *artifactp
msg := "name is invalid: %v. err: %w"
return nil, fmt.Errorf(msg, req.Name, customerror.ErrInvalidArgument)
}
res, err := ph.service.Repository.CreateKnowledgeBase(ctx,

dbData, err := ph.service.Repository.CreateKnowledgeBase(ctx,
repository.KnowledgeBase{
Name: req.Name,
KbID: toIDStyle(req.Name),
ID: toIDStyle(req.Name),
Description: req.Description,
Tags: req.Tags,
Owner: uid,
Expand All @@ -50,36 +51,43 @@ func (ph *PublicHandler) CreateKnowledgeBase(ctx context.Context, req *artifactp
if err != nil {
return nil, err
}

// TODO: ACL - set the owner of the knowledge base
// ....

return &artifactpb.CreateKnowledgeBaseResponse{
Body: &artifactpb.KnowledgeBase{
Name: res.Name,
Id: res.KbID,
Description: res.Description,
Tags: res.Tags,
OwnerName: res.Owner,
CreateTime: res.CreateTime.String(),
UpdateTime: res.UpdateTime.String(),
Name: dbData.Name,
Id: dbData.ID,
Description: dbData.Description,
Tags: dbData.Tags,
OwnerName: dbData.Owner,
CreateTime: dbData.CreateTime.String(),
UpdateTime: dbData.UpdateTime.String(),
}, ErrorMsg: "", StatusCode: 0,
}, nil
}
func (ph *PublicHandler) ListKnowledgeBases(ctx context.Context, _ *artifactpb.ListKnowledgeBasesRequest) (*artifactpb.ListKnowledgeBasesResponse, error) {

// get user id from context
uid, err := getUserIDFromContext(ctx)
uid, err := getUserUIDFromContext(ctx)
if err != nil {

return nil, fmt.Errorf(ErrorListKnowledgeBasesMsg, err)
}
res, err := ph.service.Repository.ListKnowledgeBases(ctx, uid)

// TODO: ACL - check user's permission to list knowledge bases

dbData, err := ph.service.Repository.ListKnowledgeBases(ctx, uid)
if err != nil {
return nil, fmt.Errorf(ErrorListKnowledgeBasesMsg, err)
}

kbs := make([]*artifactpb.KnowledgeBase, len(res))
for i, kb := range res {
kbs := make([]*artifactpb.KnowledgeBase, len(dbData))
for i, kb := range dbData {
kbs[i] = &artifactpb.KnowledgeBase{
Name: kb.Name,
Id: kb.KbID,
Id: kb.ID,
Description: kb.Description,
Tags: kb.Tags,
CreateTime: kb.CreateTime.String(),
Expand All @@ -95,7 +103,7 @@ func (ph *PublicHandler) ListKnowledgeBases(ctx context.Context, _ *artifactpb.L
}, nil
}
func (ph *PublicHandler) UpdateKnowledgeBase(ctx context.Context, req *artifactpb.UpdateKnowledgeBaseRequest) (*artifactpb.UpdateKnowledgeBaseResponse, error) {
uid, err := getUserIDFromContext(ctx)
uid, err := getUserUIDFromContext(ctx)
if err != nil {
return nil, err
}
Expand All @@ -107,13 +115,16 @@ func (ph *PublicHandler) UpdateKnowledgeBase(ctx context.Context, req *artifactp
if !nameOk {
return nil, fmt.Errorf("name: %s is invalid. err: %w", req.Name, customerror.ErrInvalidArgument)
}

// TODO: ACL - check user's permission to update knowledge base

// check if knowledge base exists
res, err := ph.service.Repository.UpdateKnowledgeBase(
dbData, err := ph.service.Repository.UpdateKnowledgeBase(
ctx,
uid,
repository.KnowledgeBase{
Name: req.Name,
KbID: req.Id,
ID: req.Id,
Description: req.Description,
Tags: req.Tags,
Owner: uid,
Expand All @@ -125,23 +136,26 @@ func (ph *PublicHandler) UpdateKnowledgeBase(ctx context.Context, req *artifactp
// populate response
return &artifactpb.UpdateKnowledgeBaseResponse{
Body: &artifactpb.KnowledgeBase{
Name: res.Name,
Id: res.KbID,
Description: res.Description,
Tags: res.Tags,
CreateTime: res.CreateTime.String(),
UpdateTime: res.UpdateTime.String(),
OwnerName: res.Owner,
Name: dbData.Name,
Id: dbData.ID,
Description: dbData.Description,
Tags: dbData.Tags,
CreateTime: dbData.CreateTime.String(),
UpdateTime: dbData.UpdateTime.String(),
OwnerName: dbData.Owner,
}, ErrorMsg: "", StatusCode: 0,
}, nil
}
func (ph *PublicHandler) DeleteKnowledgeBase(ctx context.Context, req *artifactpb.DeleteKnowledgeBaseRequest) (*artifactpb.DeleteKnowledgeBaseResponse, error) {

uid, err := getUserIDFromContext(ctx)
uid, err := getUserUIDFromContext(ctx)
if err != nil {

return nil, err
}

// TODO: ACL - check user's permission to delete knowledge base

err = ph.service.Repository.DeleteKnowledgeBase(ctx, uid, req.Id)
if err != nil {

Expand All @@ -151,7 +165,7 @@ func (ph *PublicHandler) DeleteKnowledgeBase(ctx context.Context, req *artifactp
ErrorMsg: "", StatusCode: 0,
}, nil
}
func getUserIDFromContext(ctx context.Context) (string, error) {
func getUserUIDFromContext(ctx context.Context) (string, error) {
md, _ := metadata.FromIncomingContext(ctx)
if v, ok := md[strings.ToLower(constant.HeaderUserUIDKey)]; ok {
return v[0], nil
Expand Down
Loading

0 comments on commit 3912028

Please sign in to comment.