github.com/weaviate/weaviate@v1.24.6/modules/generative-cohere/module.go (about)

     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  
    12  package modgenerativecohere
    13  
    14  import (
    15  	"context"
    16  	"net/http"
    17  	"os"
    18  	"time"
    19  
    20  	"github.com/pkg/errors"
    21  	"github.com/sirupsen/logrus"
    22  	"github.com/weaviate/weaviate/entities/modulecapabilities"
    23  	"github.com/weaviate/weaviate/entities/moduletools"
    24  	"github.com/weaviate/weaviate/modules/generative-cohere/clients"
    25  	additionalprovider "github.com/weaviate/weaviate/usecases/modulecomponents/additional"
    26  	generativemodels "github.com/weaviate/weaviate/usecases/modulecomponents/additional/models"
    27  )
    28  
    29  const Name = "generative-cohere"
    30  
    31  func New() *GenerativeCohereModule {
    32  	return &GenerativeCohereModule{}
    33  }
    34  
    35  type GenerativeCohereModule struct {
    36  	generative                   generativeClient
    37  	additionalPropertiesProvider modulecapabilities.AdditionalProperties
    38  }
    39  
    40  type generativeClient interface {
    41  	GenerateSingleResult(ctx context.Context, textProperties map[string]string, prompt string, cfg moduletools.ClassConfig) (*generativemodels.GenerateResponse, error)
    42  	GenerateAllResults(ctx context.Context, textProperties []map[string]string, task string, cfg moduletools.ClassConfig) (*generativemodels.GenerateResponse, error)
    43  	Generate(ctx context.Context, cfg moduletools.ClassConfig, prompt string) (*generativemodels.GenerateResponse, error)
    44  	MetaInfo() (map[string]interface{}, error)
    45  }
    46  
    47  func (m *GenerativeCohereModule) Name() string {
    48  	return Name
    49  }
    50  
    51  func (m *GenerativeCohereModule) Type() modulecapabilities.ModuleType {
    52  	return modulecapabilities.Text2TextGenerative
    53  }
    54  
    55  func (m *GenerativeCohereModule) Init(ctx context.Context,
    56  	params moduletools.ModuleInitParams,
    57  ) error {
    58  	if err := m.initAdditional(ctx, params.GetConfig().ModuleHttpClientTimeout, params.GetLogger()); err != nil {
    59  		return errors.Wrap(err, "init q/a")
    60  	}
    61  
    62  	return nil
    63  }
    64  
    65  func (m *GenerativeCohereModule) initAdditional(ctx context.Context, timeout time.Duration,
    66  	logger logrus.FieldLogger,
    67  ) error {
    68  	apiKey := os.Getenv("COHERE_APIKEY")
    69  
    70  	client := clients.New(apiKey, timeout, logger)
    71  
    72  	m.generative = client
    73  
    74  	m.additionalPropertiesProvider = additionalprovider.NewGenerativeProvider(m.generative, logger)
    75  
    76  	return nil
    77  }
    78  
    79  func (m *GenerativeCohereModule) MetaInfo() (map[string]interface{}, error) {
    80  	return m.generative.MetaInfo()
    81  }
    82  
    83  func (m *GenerativeCohereModule) RootHandler() http.Handler {
    84  	// TODO: remove once this is a capability interface
    85  	return nil
    86  }
    87  
    88  func (m *GenerativeCohereModule) AdditionalProperties() map[string]modulecapabilities.AdditionalProperty {
    89  	return m.additionalPropertiesProvider.AdditionalProperties()
    90  }
    91  
    92  // verify we implement the modules.Module interface
    93  var (
    94  	_ = modulecapabilities.Module(New())
    95  	_ = modulecapabilities.AdditionalProperties(New())
    96  	_ = modulecapabilities.MetaProvider(New())
    97  )