Repository: tmc/langchaingo
Pull request #128 (2023-06-14 22:42:07)

llms: split LLMChat and LLM into different structs
I saw in the upstream versions that the there are different classes for the chat models and standard llms. Doing something similar in the go version would bring multiple benefits. Better prompt selection being one example. This therefore changes the chat interface to have a  call and a generate function, like the llm interface, and splits the llm structs into a chat and llm struct. Both the openai llm and chat struct implement the LanguageModel interface. To do this the chat generation and generation types is united into one type, and functionality for token counting is added. 

I also saw that some types are defined in both the schema package and the llm package. For example the generation type. What versions should be removed? I think that it would be best to have the types used for llms in the llms package.  

### References
https://github.com/hwchase17/langchain/blob/master/langchain/base_language.py
https://github.com/hwchase17/langchainjs/blob/main/langchain/src/base_language/index.ts 


### PR Checklist
- [x] Read the [Contributing documentation](https://github.com/tmc/langchaingo/blob/main/CONTRIBUTING.md).
- [x] Read the [Code of conduct documentation](https://github.com/tmc/langchaingo/blob/main/CODE_OF_CONDUCT.md).
- [x] Name your Pull Request title clearly, concisely, and prefixed with the name of primarily affected package you changed according to [Good commit messages](https://go.dev/doc/contribute#commit_messages) (such as `memory: add interfaces for X, Y` or `util: add whizzbang helpers`).
- [x] Check that there isn't already a PR that solves the problem the same way to avoid creating a duplicate.
- [x] Provide a description in this PR that addresses **what** the PR is solving, or reference the issue that it solves (e.g. `fixes #123`).
- [x] Describes source of new concepts.
- [x] References existing implementations as appropriate.
- [x] Contains test coverage for new functions.
- [x] Passes all [`golangci-lint`](https://golangci-lint.run/) checks.
- [x] I am an existing contributor.



Patch:
diff --git a/go.mod b/go.mod
index 2c6643e3..0511329b 100644
--- a/go.mod
+++ b/go.mod
@@ -19,6 +19,7 @@ require (
 	github.com/PuerkitoBio/urlesc v0.0.0-20170810143723-de5bf2ad4578 // indirect
 	github.com/asaskevich/govalidator v0.0.0-20210307081110-f21760c49a8d // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/dlclark/regexp2 v1.8.1 // indirect
 	github.com/go-openapi/analysis v0.21.2 // indirect
 	github.com/go-openapi/errors v0.20.3 // indirect
 	github.com/go-openapi/jsonpointer v0.19.5 // indirect
@@ -40,6 +41,7 @@ require (
 	github.com/mitchellh/copystructure v1.0.0 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/mitchellh/reflectwalk v1.0.0 // indirect
+	github.com/pkoukk/tiktoken-go v0.1.2 // indirect
 	github.com/oklog/ulid v1.3.1 // indirect
 	github.com/pmezard/go-difflib v1.0.0 // indirect
 	github.com/shopspring/decimal v1.2.0 // indirect
diff --git a/go.sum b/go.sum
index acbf9422..1317308f 100644
--- a/go.sum
+++ b/go.sum
@@ -78,6 +78,8 @@ github.com/creack/pty v1.1.9/go.mod h1:oKZEueFk5CKHvIhNR5MUki03XCEU+Q6VDXinZuGJ3
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/dlclark/regexp2 v1.8.1 h1:6Lcdwya6GjPUNsBct8Lg/yRPwMhABj269AAzdGSiR+0=
+github.com/dlclark/regexp2 v1.8.1/go.mod h1:DHkYz0B9wPfa6wondMfaivmHpzrQ3v9q8cnmRbL6yW8=
 github.com/envoyproxy/go-control-plane v0.9.0/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/go-control-plane v0.9.1-0.20191026205805-5f8ba28d4473/go.mod h1:YTl/9mNaCwkRvm6d1a2C3ymFceY/DCBVvsKhRF0iEA4=
 github.com/envoyproxy/go-control-plane v0.9.4/go.mod h1:6rpuAdCZL397s3pYoYcLgu1mIlRU8Am5FuJP05cCM98=
@@ -261,6 +263,8 @@ github.com/oklog/ulid v1.3.1/go.mod h1:CirwcVhetQ6Lv90oh/F+FBtV6XMibvdAFo93nm5qn
 github.com/pelletier/go-toml v1.7.0/go.mod h1:vwGMzjaWMwyfHwgIBhI2YUM4fB6nL6lVAvS1LBMMhTE=
 github.com/pinecone-io/go-pinecone v0.3.0 h1:+t0CiYaaA+JN6YM9QRNlvfLEr2kkGzcVEj/xNmSAON4=
 github.com/pinecone-io/go-pinecone v0.3.0/go.mod h1:VdSieE1r4jT3XydjFi+iL5w9qsGRz/x8LxWach2Hnv8=
+github.com/pkoukk/tiktoken-go v0.1.2 h1:u7PCSBiWJ3nJYoTGShyM9iHXz4dNyYkurwwp+GHtyHY=
+github.com/pkoukk/tiktoken-go v0.1.2/go.mod h1:boMWvk9pQCOTx11pgu0DrIdrAKgQzzJKUP6vLXaz7Rw=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
 github.com/pkg/errors v0.9.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=
diff --git a/llms/count_tokens.go b/llms/count_tokens.go
new file mode 100644
index 00000000..37969ba7
--- /dev/null
+++ b/llms/count_tokens.go
@@ -0,0 +1,67 @@
+package llms
+
+import (
+	"log"
+
+	"github.com/pkoukk/tiktoken-go"
+)
+
+const (
+	_tokenApproximation = 4
+)
+
+const (
+	_gpt35TurboContextSize   = 4096
+	_gpt432KContextSize      = 32768
+	_gpt4ContextSize         = 8192
+	_textDavinci3ContextSize = 4097
+	_textBabbage1ContextSize = 2048
+	_textAda1ContextSize     = 2048
+	_textCurie1ContextSize   = 2048
+	_codeDavinci2ContextSize = 8000
+	_codeCushman1ContextSize = 2048
+	_textBisonContextSize    = 2048
+	_chatBisonContextSize    = 2048
+	_defaultContextSize      = 2048
+)
+
+// nolint:gochecknoglobals
+var modelToContextSize = map[string]int{
+	"gpt-3.5-turbo":    _gpt35TurboContextSize,
+	"gpt-4-32k":        _gpt432KContextSize,
+	"gpt-4":            _gpt4ContextSize,
+	"text-davinci-003": _textDavinci3ContextSize,
+	"text-curie-001":   _textCurie1ContextSize,
+	"text-babbage-001": _textBabbage1ContextSize,
+	"text-ada-001":     _textBabbage1ContextSize,
+	"code-davinci-002": _codeDavinci2ContextSize,
+	"code-cushman-001": _codeCushman1ContextSize,
+}
+
+// ModelContextSize gets the max number of tokens for a language model. If the model
+// name isn't recognized the default value 2048 is returned.
+func GetModelContextSize(model string) int {
+	contextSize, ok := modelToContextSize[model]
+	if !ok {
+		return _defaultContextSize
+	}
+	return contextSize
+}
+
+// CountTokens gets the number of tokens the text contains.
+func CountTokens(model, text string) int {
+	e, err := tiktoken.EncodingForModel(model)
+	if err != nil {
+		e, err = tiktoken.GetEncoding("gpt2")
+		if err != nil {
+			log.Printf("[WARN] Failed to calculate number of tokens for model, falling back to approximate count")
+			return len([]rune(text)) / _tokenApproximation
+		}
+	}
+	return len(e.Encode(text, nil, nil))
+}
+
+// CalculateMaxTokens calculates the max number of tokens that could be added to a text.
+func CalculateMaxTokens(model, text string) int {
+	return GetModelContextSize(model) - CountTokens(model, text)
+}
diff --git a/llms/huggingface/huggingfacellm.go b/llms/huggingface/huggingfacellm.go
index e77fb8c4..53a68b42 100644
--- a/llms/huggingface/huggingfacellm.go
+++ b/llms/huggingface/huggingfacellm.go
@@ -7,6 +7,7 @@ import (
 
 	"github.com/tmc/langchaingo/llms"
 	"github.com/tmc/langchaingo/llms/huggingface/internal/huggingfaceclient"
+	"github.com/tmc/langchaingo/schema"
 )
 
 var (
@@ -18,7 +19,10 @@ type LLM struct {
 	client *huggingfaceclient.Client
 }
 
-var _ llms.LLM = (*LLM)(nil)
+var (
+	_ llms.LLM           = (*LLM)(nil)
+	_ llms.LanguageModel = (*LLM)(nil)
+)
 
 func (o *LLM) Call(ctx context.Context, prompt string, options ...llms.CallOption) (string, error) {
 	r, err := o.Generate(ctx, []string{prompt}, options...)
@@ -56,6 +60,14 @@ func (o *LLM) Generate(ctx context.Context, prompts []string, options ...llms.Ca
 	}, nil
 }
 
+func (o *LLM) GeneratePrompt(ctx context.Context, prompts []schema.PromptValue, options ...llms.CallOption) (llms.LLMResult, error) { //nolint:lll
+	return llms.GeneratePrompt(ctx, o, prompts, options...)
+}
+
+func (o *LLM) GetNumTokens(text string) int {
+	return llms.CountTokens("gpt2", text)
+}
+
 func New(opts ...Option) (*LLM, error) {
 	options := &options{
 		token: os.Getenv(tokenEnvVarName),
diff --git a/llms/llms.go b/llms/llms.go
index 3700d501..67b4e6c7 100644
--- a/llms/llms.go
+++ b/llms/llms.go
@@ -6,6 +6,14 @@ import (
 	"github.com/tmc/langchaingo/schema"
 )
 
+// LanguageModel is the interface all language models must implement.
+type LanguageModel interface {
+	// Take in a list of prompt values and return an LLMResult.
+	GeneratePrompt(ctx context.Context, prompts []schema.PromptValue, options ...CallOption) (LLMResult, error)
+	// Get the number of tokens present in the text.
+	GetNumTokens(text string) int
+}
+
 // LLM is a langchaingo Large Language Model.
 type LLM interface {
 	Call(ctx context.Context, prompt string, options ...CallOption) (string, error)
@@ -14,20 +22,44 @@ type LLM interface {
 
 // ChatLLM is a langchaingo LLM that can be used for chatting.
 type ChatLLM interface {
-	Chat(ctx context.Context, messages []schema.ChatMessage, options ...CallOption) (*ChatGeneration, error)
+	Call(ctx context.Context, messages []schema.ChatMessage, options ...CallOption) (string, error)
+	Generate(ctx context.Context, messages [][]schema.ChatMessage, options ...CallOption) ([]*Generation, error)
 }
 
 // Generation is a single generation from a langchaingo LLM.
 type Generation struct {
 	// Text is the generated text.
 	Text string `json:"text"`
+	// Message stores the potentially generated message.
+	Message *schema.AIChatMessage `json:"message"`
 	// GenerationInfo is the generation info. This can contain vendor-specific information.
 	GenerationInfo map[string]any `json:"generation_info"`
 }
 
-// ChatGeneration is a single generation from a langchaingo ChatLLM.
-type ChatGeneration struct {
-	Message *schema.AIChatMessage `json:"message"`
-	// GenerationInfo is the generation info. This can contain vendor-specific information.
-	GenerationInfo map[string]any `json:"generation_info"`
+// LLMResult is the class that contains all relevant information for an LLM Result.
+type LLMResult struct {
+	Generations [][]*Generation
+	LLMOutput   map[string]any
+}
+
+func GeneratePrompt(ctx context.Context, l LLM, promptValues []schema.PromptValue, options ...CallOption) (LLMResult, error) { //nolint:lll
+	prompts := make([]string, 0, len(promptValues))
+	for _, promptValue := range promptValues {
+		prompts = append(prompts, promptValue.String())
+	}
+	generations, err := l.Generate(ctx, prompts, options...)
+	return LLMResult{
+		Generations: [][]*Generation{generations},
+	}, err
+}
+
+func GenerateChatPrompt(ctx context.Context, l ChatLLM, promptValues []schema.PromptValue, options ...CallOption) (LLMResult, error) { //nolint:lll
+	messages := make([][]schema.ChatMessage, 0, len(promptValues))
+	for _, promptValue := range promptValues {
+		messages = append(messages, promptValue.Messages())
+	}
+	generations, err := l.Generate(ctx, messages, options...)
+	return LLMResult{
+		Generations: [][]*Generation{generations},
+	}, err
 }
diff --git a/llms/local/localllm.go b/llms/local/localllm.go
index 564f914f..1f3fdff1 100644
--- a/llms/local/localllm.go
+++ b/llms/local/localllm.go
@@ -9,6 +9,7 @@ import (
 
 	"github.com/tmc/langchaingo/llms"
 	"github.com/tmc/langchaingo/llms/local/internal/localclient"
+	"github.com/tmc/langchaingo/schema"
 )
 
 var (
@@ -23,8 +24,11 @@ type LLM struct {
 	client *localclient.Client
 }
 
-// _ ensures that LLM implements the llms.LLM interface.
-var _ llms.LLM = (*LLM)(nil)
+// _ ensures that LLM implements the llms.LLM and language model interface.
+var (
+	_ llms.LLM           = (*LLM)(nil)
+	_ llms.LanguageModel = (*LLM)(nil)
+)
 
 // Call calls the local LLM binary with the given prompt.
 func (o *LLM) Call(ctx context.Context, prompt string, options ...llms.CallOption) (string, error) {
@@ -88,6 +92,14 @@ func (o *LLM) Generate(ctx context.Context, prompts []string, options ...llms.Ca
 	}, nil
 }
 
+func (o *LLM) GeneratePrompt(ctx context.Context, prompts []schema.PromptValue, options ...llms.CallOption) (llms.LLMResult, error) { //nolint:lll
+	return llms.GeneratePrompt(ctx, o, prompts, options...)
+}
+
+func (o *LLM) GetNumTokens(text string) int {
+	return llms.CountTokens("gpt2", text)
+}
+
 // New creates a new local LLM implementation.
 func New(opts ...Option) (*LLM, error) {
 	options := &options{
diff --git a/llms/openai/internal/openaiclient/completions.go b/llms/openai/internal/openaiclient/completions.go
index 2dc23ac0..a9dfa298 100644
--- a/llms/openai/internal/openaiclient/completions.go
+++ b/llms/openai/internal/openaiclient/completions.go
@@ -61,8 +61,8 @@ func (c *Client) setCompletionDefaults(payload *completionPayload) {
 	case payload.Model != "":
 
 	// If no model is set in the payload, take the one specified in the client.
-	case c.model != "":
-		payload.Model = c.model
+	case c.Model != "":
+		payload.Model = c.Model
 	// Fallback: use the default model
 	default:
 		payload.Model = defaultCompletionModel
diff --git a/llms/openai/internal/openaiclient/openaiclient.go b/llms/openai/internal/openaiclient/openaiclient.go
index e281da9e..8d5bbb02 100644
--- a/llms/openai/internal/openaiclient/openaiclient.go
+++ b/llms/openai/internal/openaiclient/openaiclient.go
@@ -12,7 +12,7 @@ var ErrEmptyResponse = errors.New("empty response")
 // Client is a client for the OpenAI API.
 type Client struct {
 	token      string
-	model      string
+	Model      string
 	baseURL    string
 	httpClient Doer
 }
@@ -38,7 +38,7 @@ func WithHTTPClient(client Doer) Option {
 func New(token string, model string, baseURL string, opts ...Option) (*Client, error) {
 	c := &Client{
 		token:      token,
-		model:      model,
+		Model:      model,
 		baseURL:    baseURL,
 		httpClient: http.DefaultClient,
 	}
@@ -92,7 +92,7 @@ type EmbeddingRequest struct {
 
 // CreateEmbedding creates embeddings.
 func (c *Client) CreateEmbedding(ctx context.Context, r *EmbeddingRequest) ([][]float64, error) {
-	r.Model = c.model
+	r.Model = c.Model
 	if r.Model == "" {
 		r.Model = defaultEmbeddingModel
 	}
@@ -119,7 +119,7 @@ func (c *Client) CreateEmbedding(ctx context.Context, r *EmbeddingRequest) ([][]
 
 // CreateChat creates chat request.
 func (c *Client) CreateChat(ctx context.Context, r *ChatRequest) (*ChatResponse, error) {
-	r.Model = c.model
+	r.Model = c.Model
 	if r.Model == "" {
 		r.Model = defaultChatModel
 	}
diff --git a/llms/openai/openaillm.go b/llms/openai/openaillm.go
index 9b69644f..337cd99f 100644
--- a/llms/openai/openaillm.go
+++ b/llms/openai/openaillm.go
@@ -17,13 +17,16 @@ var (
 	ErrUnexpectedResponseLength = errors.New("unexpected length of response")
 )
 
+type ChatMessage = openaiclient.ChatMessage
+
 type LLM struct {
 	client *openaiclient.Client
 }
 
-var _ llms.LLM = (*LLM)(nil)
-
-var _ llms.ChatLLM = (*LLM)(nil)
+var (
+	_ llms.LLM           = (*LLM)(nil)
+	_ llms.LanguageModel = (*LLM)(nil)
+)
 
 // Call requests a completion for the given prompt.
 func (o *LLM) Call(ctx context.Context, prompt string, options ...llms.CallOption) (string, error) {
@@ -42,91 +45,150 @@ func (o *LLM) Generate(ctx context.Context, prompts []string, options ...llms.Ca
 	for _, opt := range options {
 		opt(&opts)
 	}
-	result, err := o.client.CreateCompletion(ctx, &openaiclient.CompletionRequest{
-		Model:     opts.Model,
-		Prompt:    prompts[0],
-		MaxTokens: opts.MaxTokens,
-		StopWords: opts.StopWords,
-	})
-	if err != nil {
-		return nil, err
+
+	generations := make([]*llms.Generation, 0, len(prompts))
+	for _, prompt := range prompts {
+		result, err := o.client.CreateCompletion(ctx, &openaiclient.CompletionRequest{
+			Model:     opts.Model,
+			Prompt:    prompt,
+			MaxTokens: opts.MaxTokens,
+			StopWords: opts.StopWords,
+		})
+		if err != nil {
+			return nil, err
+		}
+		generations = append(generations, &llms.Generation{
+			Text: result.Text,
+		})
 	}
-	return []*llms.Generation{
-		{Text: result.Text},
-	}, nil
+
+	return generations, nil
 }
 
-type ChatMessage = openaiclient.ChatMessage
+func (o *LLM) GeneratePrompt(ctx context.Context, promptValues []schema.PromptValue, options ...llms.CallOption) (llms.LLMResult, error) { //nolint:lll
+	return llms.GeneratePrompt(ctx, o, promptValues, options...)
+}
 
-// Chat requests a chat response for the given prompt.
-func (o *LLM) Chat(ctx context.Context, messages []schema.ChatMessage, options ...llms.CallOption) (*llms.ChatGeneration, error) { // nolint: lll
-	opts := llms.CallOptions{}
-	for _, opt := range options {
-		opt(&opts)
-	}
-	msgs := make([]*openaiclient.ChatMessage, len(messages))
-	for i, m := range messages {
-		msg := &openaiclient.ChatMessage{
-			Content: m.GetText(),
-		}
-		typ := m.GetType()
-		switch typ {
-		case schema.ChatMessageTypeSystem:
-			msg.Role = "system"
-		case schema.ChatMessageTypeAI:
-			msg.Role = "assistant"
-		case schema.ChatMessageTypeHuman:
-			msg.Role = "user"
-		case schema.ChatMessageTypeGeneric:
-			msg.Role = "user"
-			// TODO: support name
-		}
-		msgs[i] = msg
-	}
+func (o *LLM) GetNumTokens(text string) int {
+	return llms.CountTokens(o.client.Model, text)
+}
 
-	result, err := o.client.CreateChat(ctx, &openaiclient.ChatRequest{
-		Model:         opts.Model,
-		StopWords:     opts.StopWords,
-		Messages:      msgs,
-		StreamingFunc: opts.StreamingFunc,
+// CreateEmbedding creates embeddings for the given input texts.
+func (o *LLM) CreateEmbedding(ctx context.Context, inputTexts []string) ([][]float64, error) {
+	embeddings, err := o.client.CreateEmbedding(ctx, &openaiclient.EmbeddingRequest{
+		Input: inputTexts,
 	})
 	if err != nil {
 		return nil, err
 	}
-	if len(result.Choices) == 0 {
+	if len(embeddings) == 0 {
 		return nil, ErrEmptyResponse
 	}
-	return &llms.ChatGeneration{
-		Message: &schema.AIChatMessage{
-			Text: result.Choices[0].Message.Content,
-		},
-		// TODO: fill in generation info
-	}, nil
+	if len(inputTexts) != len(embeddings) {
+		return embeddings, ErrUnexpectedResponseLength
+	}
+	return embeddings, nil
 }
 
-// CreateEmbedding creates embeddings for the given input texts.
-func (o *LLM) CreateEmbedding(ctx context.Context, inputTexts []string) ([][]float64, error) {
-	embeddings, err := o.client.CreateEmbedding(ctx, &openaiclient.EmbeddingRequest{
-		Input: inputTexts,
-	})
+type Chat struct {
+	client *openaiclient.Client
+}
 
-	if len(embeddings) == 0 {
-		return [][]float64{}, ErrEmptyResponse
-	}
+var (
+	_ llms.ChatLLM       = (*Chat)(nil)
+	_ llms.LanguageModel = (*Chat)(nil)
+)
 
+// Chat requests a chat response for the given messages.
+func (o *Chat) Call(ctx context.Context, messages []schema.ChatMessage, options ...llms.CallOption) (string, error) { // nolint: lll
+	r, err := o.Generate(ctx, [][]schema.ChatMessage{messages}, options...)
 	if err != nil {
-		return [][]float64{}, err
+		return "", err
+	}
+	if len(r) == 0 {
+		return "", ErrEmptyResponse
 	}
+	return r[0].Message.Text, nil
+}
 
-	if len(inputTexts) != len(embeddings) {
-		return embeddings, ErrUnexpectedResponseLength
+func (o *Chat) Generate(ctx context.Context, messageSets [][]schema.ChatMessage, options ...llms.CallOption) ([]*llms.Generation, error) { // nolint:lll
+	opts := llms.CallOptions{}
+	for _, opt := range options {
+		opt(&opts)
 	}
 
-	return embeddings, nil
+	generations := make([]*llms.Generation, 0, len(messageSets))
+	for _, messageSet := range messageSets {
+		msgs := make([]*openaiclient.ChatMessage, len(messageSet))
+		for i, m := range messageSet {
+			msg := &openaiclient.ChatMessage{
+				Content: m.GetText(),
+			}
+			typ := m.GetType()
+			switch typ {
+			case schema.ChatMessageTypeSystem:
+				msg.Role = "system"
+			case schema.ChatMessageTypeAI:
+				msg.Role = "assistant"
+			case schema.ChatMessageTypeHuman:
+				msg.Role = "user"
+			case schema.ChatMessageTypeGeneric:
+				msg.Role = "user"
+				// TODO: support name
+			}
+			msgs[i] = msg
+		}
+
+		result, err := o.client.CreateChat(ctx, &openaiclient.ChatRequest{
+			Model:         opts.Model,
+			StopWords:     opts.StopWords,
+			Messages:      msgs,
+			StreamingFunc: opts.StreamingFunc,
+		})
+		if err != nil {
+			return nil, err
+		}
+		if len(result.Choices) == 0 {
+			return nil, ErrEmptyResponse
+		}
+		text := result.Choices[0].Message.Content
+		generations = append(generations, &llms.Generation{
+			Message: &schema.AIChatMessage{
+				Text: text,
+			},
+			Text: text,
+			// TODO: fill in generation info
+		})
+	}
+
+	return generations, nil
+}
+
+func (o *Chat) GetNumTokens(text string) int {
+	return llms.CountTokens(o.client.Model, text)
+}
+
+func (o *Chat) GeneratePrompt(ctx context.Context, promptValues []schema.PromptValue, options ...llms.CallOption) (llms.LLMResult, error) { //nolint:lll
+	return llms.GenerateChatPrompt(ctx, o, promptValues, options...)
 }
 
 // New returns a new OpenAI LLM.
 func New(opts ...Option) (*LLM, error) {
+	c, err := newClient(opts...)
+	return &LLM{
+		client: c,
+	}, err
+}
+
+// NewChat returns a new OpenAI chat LLM.
+func NewChat(opts ...Option) (*Chat, error) {
+	c, err := newClient(opts...)
+	return &Chat{
+		client: c,
+	}, err
+}
+
+func newClient(opts ...Option) (*openaiclient.Client, error) {
 	options := &options{
 		token:   os.Getenv(tokenEnvVarName),
 		model:   os.Getenv(modelEnvVarName),
@@ -141,12 +203,5 @@ func New(opts ...Option) (*LLM, error) {
 		return nil, ErrMissingToken
 	}
 
-	client, err := openaiclient.New(options.token, options.model, options.baseURL)
-	if err != nil {
-		return nil, err
-	}
-
-	return &LLM{
-		client: client,
-	}, nil
+	return openaiclient.New(options.token, options.model, options.baseURL)
 }
diff --git a/llms/vertexai/internal/vertexaiclient/palm_client.go b/llms/vertexai/internal/vertexaiclient/palm_client.go
index 6f1b04a7..84aa0820 100644
--- a/llms/vertexai/internal/vertexaiclient/palm_client.go
+++ b/llms/vertexai/internal/vertexaiclient/palm_client.go
@@ -35,8 +35,8 @@ var defaultParameters = map[string]interface{}{ //nolint:gochecknoglobals
 
 const (
 	embeddingModelName = "textembedding-gecko"
-	textModelName      = "text-bison"
-	chatModelName      = "chat-bison"
+	TextModelName      = "text-bison"
+	ChatModelName      = "chat-bison"
 
 	defaultMaxConns = 4
 )
@@ -95,7 +95,7 @@ func (c *PaLMClient) CreateCompletion(ctx context.Context, r *CompletionRequest)
 		"top_p":           r.TopP,
 		"top_k":           r.TopK,
 	}
-	predictions, err := c.batchPredict(ctx, textModelName, r.Prompts, params)
+	predictions, err := c.batchPredict(ctx, TextModelName, r.Prompts, params)
 	if err != nil {
 		return nil, err
 	}
@@ -169,9 +169,7 @@ type ChatMessage struct {
 }
 
 // Statically assert that the types implement the interface.
-var (
-	_ schema.ChatMessage = ChatMessage{}
-)
+var _ schema.ChatMessage = ChatMessage{}
 
 // GetType returns the type of the message.
 func (m ChatMessage) GetType() schema.ChatMessageType {
@@ -303,7 +301,7 @@ func (c *PaLMClient) chat(ctx context.Context, r *ChatRequest) ([]*structpb.Valu
 		structpb.NewStructValue(instance),
 	}
 	resp, err := c.client.Predict(ctx, &aiplatformpb.PredictRequest{
-		Endpoint:   c.projectLocationPublisherModelPath(c.projectID, "us-central1", "google", chatModelName),
+		Endpoint:   c.projectLocationPublisherModelPath(c.projectID, "us-central1", "google", ChatModelName),
 		Instances:  instances,
 		Parameters: structpb.NewStructValue(mergedParams),
 	})
diff --git a/llms/vertexai/vertexai_palm_llm.go b/llms/vertexai/vertexai_palm_llm.go
index c169e79a..638ccefc 100644
--- a/llms/vertexai/vertexai_palm_llm.go
+++ b/llms/vertexai/vertexai_palm_llm.go
@@ -20,35 +20,10 @@ type LLM struct {
 	client *vertexaiclient.PaLMClient
 }
 
-var _ llms.LLM = (*LLM)(nil)
-
-var _ llms.ChatLLM = (*LLM)(nil)
-
-// New returns a new VertexAI PaLM LLM.
-func New(opts ...Option) (*LLM, error) {
-	// Ensure options are initialized only once.
-	initOptions.Do(initOpts)
-
-	options := &options{}
-	*options = *defaultOptions // Copy default options.
-
-	for _, opt := range opts {
-		opt(options)
-	}
-
-	if len(options.projectID) == 0 {
-		return nil, ErrMissingProjectID
-	}
-
-	client, err := vertexaiclient.New(options.projectID, options.clientOptions...)
-	if err != nil {
-		return nil, err
-	}
-
-	return &LLM{
-		client: client,
-	}, nil
-}
+var (
+	_ llms.LLM           = (*LLM)(nil)
+	_ llms.LanguageModel = (*LLM)(nil)
+)
 
 // Call requests a completion for the given prompt.
 func (o *LLM) Call(ctx context.Context, prompt string, options ...llms.CallOption) (string, error) {
@@ -85,10 +60,58 @@ func (o *LLM) Generate(ctx context.Context, prompts []string, options ...llms.Ca
 	return generations, nil
 }
 
+// CreateEmbedding creates embeddings for the given input texts.
+func (o *LLM) CreateEmbedding(ctx context.Context, inputTexts []string) ([][]float64, error) {
+	embeddings, err := o.client.CreateEmbedding(ctx, &vertexaiclient.EmbeddingRequest{
+		Input: inputTexts,
+	})
+	if err != nil {
+		return [][]float64{}, err
+	}
+
+	if len(embeddings) == 0 {
+		return [][]float64{}, ErrEmptyResponse
+	}
+	if len(inputTexts) != len(embeddings) {
+		return embeddings, ErrUnexpectedResponseLength
+	}
+
+	return embeddings, nil
+}
+
+func (o *LLM) GeneratePrompt(ctx context.Context, promptValues []schema.PromptValue, options ...llms.CallOption) (llms.LLMResult, error) { //nolint:lll
+	return llms.GeneratePrompt(ctx, o, promptValues, options...)
+}
+
+func (o *LLM) GetNumTokens(text string) int {
+	return llms.CountTokens(vertexaiclient.TextModelName, text)
+}
+
 type ChatMessage = vertexaiclient.ChatMessage
 
-// Chat requests a chat response for the given prompt.
-func (o *LLM) Chat(ctx context.Context, messages []schema.ChatMessage, options ...llms.CallOption) (*llms.ChatGeneration, error) { // nolint: lll
+type Chat struct {
+	client *vertexaiclient.PaLMClient
+}
+
+var (
+	_ llms.ChatLLM       = (*Chat)(nil)
+	_ llms.LanguageModel = (*Chat)(nil)
+)
+
+// Chat requests a chat response for the given messages.
+func (o *Chat) Call(ctx context.Context, messages []schema.ChatMessage, options ...llms.CallOption) (string, error) { // nolint: lll
+	r, err := o.Generate(ctx, [][]schema.ChatMessage{messages}, options...)
+	if err != nil {
+		return "", err
+	}
+	if len(r) == 0 {
+		return "", ErrEmptyResponse
+	}
+	return r[0].Message.Text, nil
+}
+
+// Generate requests a chat response for each of the sets of messages.
+func (o *Chat) Generate(ctx context.Context, messageSets [][]schema.ChatMessage, options ...llms.CallOption) ([]*llms.Generation, error) { // nolint: lll
 	opts := llms.CallOptions{}
 	for _, opt := range options {
 		opt(&opts)
@@ -96,6 +119,40 @@ func (o *LLM) Chat(ctx context.Context, messages []schema.ChatMessage, options .
 	if opts.StreamingFunc != nil {
 		return nil, ErrNotImplemented
 	}
+
+	generations := make([]*llms.Generation, 0, len(messageSets))
+	for _, messages := range messageSets {
+		msgs := toClientChatMessage(messages)
+		result, err := o.client.CreateChat(ctx, &vertexaiclient.ChatRequest{
+			Temperature: opts.Temperature,
+			Messages:    msgs,
+		})
+		if err != nil {
+			return nil, err
+		}
+		if len(result.Candidates) == 0 {
+			return nil, ErrEmptyResponse
+		}
+		generations = append(generations, &llms.Generation{
+			Message: &schema.AIChatMessage{
+				Text: result.Candidates[0].Content,
+			},
+			Text: result.Candidates[0].Content,
+		})
+	}
+
+	return generations, nil
+}
+
+func (o *Chat) GeneratePrompt(ctx context.Context, promptValues []schema.PromptValue, options ...llms.CallOption) (llms.LLMResult, error) { //nolint:lll
+	return llms.GenerateChatPrompt(ctx, o, promptValues, options...)
+}
+
+func (o *Chat) GetNumTokens(text string) int {
+	return llms.CountTokens(vertexaiclient.TextModelName, text)
+}
+
+func toClientChatMessage(messages []schema.ChatMessage) []*vertexaiclient.ChatMessage {
 	msgs := make([]*vertexaiclient.ChatMessage, len(messages))
 	for i, m := range messages {
 		msg := &vertexaiclient.ChatMessage{
@@ -114,40 +171,33 @@ func (o *LLM) Chat(ctx context.Context, messages []schema.ChatMessage, options .
 		}
 		msgs[i] = msg
 	}
+	return msgs
+}
 
-	result, err := o.client.CreateChat(ctx, &vertexaiclient.ChatRequest{
-		Temperature: opts.Temperature,
-		Messages:    msgs,
-	})
-	if err != nil {
-		return nil, err
-	}
-	if len(result.Candidates) == 0 {
-		return nil, ErrEmptyResponse
-	}
-	return &llms.ChatGeneration{
-		Message: &schema.AIChatMessage{
-			Text: result.Candidates[0].Content,
-		},
-	}, nil
+// New returns a new VertexAI PaLM LLM.
+func New(opts ...Option) (*LLM, error) {
+	client, err := newClient(opts...)
+	return &LLM{client: client}, err
 }
 
-// CreateEmbedding creates embeddings for the given input texts.
-func (o *LLM) CreateEmbedding(ctx context.Context, inputTexts []string) ([][]float64, error) {
-	embeddings, err := o.client.CreateEmbedding(ctx, &vertexaiclient.EmbeddingRequest{
-		Input: inputTexts,
-	})
-	if err != nil {
-		return [][]float64{}, err
-	}
+// New returns a new VertexAI PaLM Chat LLM.
+func NewChat(opts ...Option) (*Chat, error) {
+	client, err := newClient(opts...)
+	return &Chat{client: client}, err
+}
 
-	if len(embeddings) == 0 {
-		return [][]float64{}, ErrEmptyResponse
-	}
+func newClient(opts ...Option) (*vertexaiclient.PaLMClient, error) {
+	// Ensure options are initialized only once.
+	initOptions.Do(initOpts)
+	options := &options{}
+	*options = *defaultOptions // Copy default options.
 
-	if len(inputTexts) != len(embeddings) {
-		return embeddings, ErrUnexpectedResponseLength
+	for _, opt := range opts {
+		opt(options)
+	}
+	if len(options.projectID) == 0 {
+		return nil, ErrMissingProjectID
 	}
 
-	return embeddings, nil
+	return vertexaiclient.New(options.projectID, options.clientOptions...)
 }
diff --git a/schema/language_model.go b/schema/language_model.go
deleted file mode 100644
index f7f28837..00000000
--- a/schema/language_model.go
+++ /dev/null
@@ -1,11 +0,0 @@
-package schema
-
-// LanguageModel is the base class for all language models.
-type LanguageModel interface {
-	// Take in a list of prompt values and return an LLMResult.
-	GeneratePrompt(prompts []PromptValue, stop []string) LLMResult
-	// Get the number of tokens present in the text.
-	GetNumTokens(text string) int
-	// Get the number of tokens in the message.
-	GetNumTokensFromMessages(messages []ChatMessage) int
-}


Test patch:
diff --git a/llms/count_tokens_test.go b/llms/count_tokens_test.go
new file mode 100644
index 00000000..b7089340
--- /dev/null
+++ b/llms/count_tokens_test.go
@@ -0,0 +1,14 @@
+package llms
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCountTokens(t *testing.T) {
+	t.Parallel()
+	numTokens := CountTokens("gpt-3.5-turbo", "test for counting tokens")
+	expectedNumTokens := 4
+	assert.Equal(t, expectedNumTokens, numTokens)
+}