weaviate
109 строк · 3.4 Кб
1// _ _
2// __ _____ __ ___ ___ __ _| |_ ___
3// \ \ /\ / / _ \/ _` \ \ / / |/ _` | __/ _ \
4// \ V V / __/ (_| |\ V /| | (_| | || __/
5// \_/\_/ \___|\__,_| \_/ |_|\__,_|\__\___|
6//
7// Copyright © 2016 - 2024 Weaviate B.V. All rights reserved.
8//
9// CONTACT: hello@weaviate.io
10//
11
12package modspellcheck
13
14import (
15"context"
16"net/http"
17"os"
18"time"
19
20"github.com/pkg/errors"
21"github.com/weaviate/weaviate/entities/modulecapabilities"
22"github.com/weaviate/weaviate/entities/moduletools"
23spellcheckadditional "github.com/weaviate/weaviate/modules/text-spellcheck/additional"
24spellcheckadditionalspellcheck "github.com/weaviate/weaviate/modules/text-spellcheck/additional/spellcheck"
25"github.com/weaviate/weaviate/modules/text-spellcheck/clients"
26"github.com/weaviate/weaviate/modules/text-spellcheck/ent"
27spellchecktexttransformer "github.com/weaviate/weaviate/modules/text-spellcheck/transformer"
28spellchecktexttransformerautocorrect "github.com/weaviate/weaviate/modules/text-spellcheck/transformer/autocorrect"
29)
30
31func New() *SpellCheckModule {
32return &SpellCheckModule{}
33}
34
35type SpellCheckModule struct {
36spellCheck spellCheckClient
37additionalPropertiesProvider modulecapabilities.AdditionalProperties
38textTransformersProvider modulecapabilities.TextTransformers
39}
40
41type spellCheckClient interface {
42Check(ctx context.Context, text []string) (*ent.SpellCheckResult, error)
43MetaInfo() (map[string]interface{}, error)
44}
45
46func (m *SpellCheckModule) Name() string {
47return "text-spellcheck"
48}
49
50func (m *SpellCheckModule) Type() modulecapabilities.ModuleType {
51return modulecapabilities.Extension
52}
53
54func (m *SpellCheckModule) Init(ctx context.Context,
55params moduletools.ModuleInitParams,
56) error {
57uri := os.Getenv("SPELLCHECK_INFERENCE_API")
58if uri == "" {
59return errors.Errorf("required variable SPELLCHECK_INFERENCE_API is not set")
60}
61
62client := clients.New(uri, params.GetConfig().ModuleHttpClientTimeout, params.GetLogger())
63
64if err := client.WaitForStartup(ctx, 1*time.Second); err != nil {
65return errors.Wrap(err, "init remote spell check module")
66}
67
68m.spellCheck = client
69
70m.initTextTransformers()
71m.initAdditional()
72
73return nil
74}
75
76func (m *SpellCheckModule) initTextTransformers() {
77autocorrectProvider := spellchecktexttransformerautocorrect.New(m.spellCheck)
78m.textTransformersProvider = spellchecktexttransformer.New(autocorrectProvider)
79}
80
81func (m *SpellCheckModule) initAdditional() {
82spellCheckProvider := spellcheckadditionalspellcheck.New(m.spellCheck)
83m.additionalPropertiesProvider = spellcheckadditional.New(spellCheckProvider)
84}
85
86func (m *SpellCheckModule) RootHandler() http.Handler {
87// TODO: remove once this is a capability interface
88return nil
89}
90
91func (m *SpellCheckModule) MetaInfo() (map[string]interface{}, error) {
92return m.spellCheck.MetaInfo()
93}
94
95func (m *SpellCheckModule) AdditionalProperties() map[string]modulecapabilities.AdditionalProperty {
96return m.additionalPropertiesProvider.AdditionalProperties()
97}
98
99func (m *SpellCheckModule) TextTransformers() map[string]modulecapabilities.TextTransform {
100return m.textTransformersProvider.TextTransformers()
101}
102
103// verify we implement the modules.Module interface
104var (
105_ = modulecapabilities.Module(New())
106_ = modulecapabilities.AdditionalProperties(New())
107_ = modulecapabilities.MetaProvider(New())
108_ = modulecapabilities.TextTransformers(New())
109)
110