/
expensebot.go
71 lines (63 loc) · 1.82 KB
/
expensebot.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
package bot
import (
"context"
"encoding/json"
"fmt"
"log"
documentai "cloud.google.com/go/documentai/apiv1"
documentaipb "cloud.google.com/go/documentai/apiv1/documentaipb"
"github.com/koshkaj/expensebot/config"
"github.com/koshkaj/expensebot/types"
"google.golang.org/api/option"
)
type GoogleProcessor struct {
*config.GoogleProcessorConfig
client *documentai.DocumentProcessorClient
}
func NewGoogleProcessor(cfg *config.GoogleProcessorConfig) Processor {
ctx := context.Background()
client, err := documentai.NewDocumentProcessorClient(ctx, option.WithEndpoint(cfg.Endpoint))
if err != nil {
log.Fatal(err)
}
return &GoogleProcessor{
GoogleProcessorConfig: cfg,
client: client,
}
}
func (gp *GoogleProcessor) getProcessorName() string {
return fmt.Sprintf("projects/%s/locations/%s/processors/%s", gp.ProjectID, gp.Location, gp.ProcessorID)
}
func (gp *GoogleProcessor) createProcessRequest(ctx context.Context, file *types.File) (*documentaipb.ProcessRequest, error) {
r := &documentaipb.ProcessRequest{
SkipHumanReview: true,
Name: gp.getProcessorName(),
Source: &documentaipb.ProcessRequest_RawDocument{
RawDocument: &documentaipb.RawDocument{
MimeType: file.MimeType,
Content: file.Data,
},
},
}
return r, nil
}
func (gp *GoogleProcessor) Process(fileDocument *types.File) ([]byte, error) {
ctx := context.Background()
defer gp.client.Close()
request, err := gp.createProcessRequest(ctx, fileDocument)
if err != nil {
return nil, err
}
response, err := gp.client.ProcessDocument(ctx, request)
if err != nil {
log.Print("Error processing document: ", err)
return nil, err
}
document := response.GetDocument()
jsoned, err := json.Marshal(document)
if err != nil {
log.Print("Error marshaling document: ", err)
return nil, err
}
return jsoned, nil
}