llms: Add ChatLLM interface

chains/chains.go

@@ -12,7 +12,7 @@ type Chain interface {
 	// Call runs the logic of the chain and returns the output. This method should
 	// not be called directly. Use rather the Call function that handles the memory
 	// of the chain.
-	Call(context.Context, map[string]any, ...ChainCallOption) (map[string]any, error)
+	Call(ctx context.Context, inputs map[string]any, options ...ChainCallOption) (map[string]any, error)
 	// GetMemory gets the memory of the chain.
 	GetMemory() schema.Memory
 	// InputKeys returns the input keys the chain expects.

llms/llms.go

@@ -1,7 +1,11 @@
 // Package llms defines the types for langchaingo LLMs.
 package llms
 
-import "context"
+import (
+	"context"
+
+	"github.com/tmc/langchaingo/schema"
+)
 
 // LLM is a langchaingo Large Language Model.
 type LLM interface {
@@ -9,6 +13,11 @@ type LLM interface {
 	Generate(ctx context.Context, prompts []string, options ...CallOption) ([]*Generation, error)
 }
 
+// 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)
+}
+
 // Generation is a single generation from a langchaingo LLM.
 type Generation struct {
 	// Text is the generated text.
@@ -16,3 +25,10 @@ type Generation struct {
 	// 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"`
+}

llms/openai/internal/openaiclient/chat.go

@@ -13,12 +13,12 @@ const (
 
 // ChatRequest is a request to create an embedding.
 type ChatRequest struct {
-	Model       string        `json:"model"`
-	Messages    []ChatMessage `json:"messages"`
-	Temperature float64       `json:"temperature,omitempty"`
-	TopP        int           `json:"top_p,omitempty"`
-	N           int           `json:"n,omitempty"`
-	StopWords   []string      `json:"stop,omitempty"`
+	Model       string         `json:"model"`
+	Messages    []*ChatMessage `json:"messages"`
+	Temperature float64        `json:"temperature,omitempty"`
+	TopP        int            `json:"top_p,omitempty"`
+	N           int            `json:"n,omitempty"`
+	StopWords   []string       `json:"stop,omitempty"`
 }
 
 // ChatMessage is a message in a chat request.

llms/openai/openaillm.go

@@ -6,6 +6,7 @@ import (
 
 	"github.com/tmc/langchaingo/llms"
 	"github.com/tmc/langchaingo/llms/openai/internal/openaiclient"
+	"github.com/tmc/langchaingo/schema"
 )
 
 var (
@@ -20,6 +21,7 @@ type LLM struct {
 }
 
 var _ llms.LLM = (*LLM)(nil)
+var _ llms.ChatLLM = (*LLM)(nil)
 
 // Call requests a completion for the given prompt.
 func (o *LLM) Call(ctx context.Context, prompt string, options ...llms.CallOption) (string, error) {
@@ -55,20 +57,48 @@ func (o *LLM) Generate(ctx context.Context, prompts []string, options ...llms.Ca
 type ChatMessage = openaiclient.ChatMessage
 
 // Chat requests a chat response for the given prompt.
-func (o *LLM) Chat(ctx context.Context, messages []ChatMessage, options ...llms.CallOption) (*ChatMessage, error) {
+func (o *LLM) Chat(ctx context.Context, messages []schema.ChatMessage, options ...llms.CallOption) (*llms.ChatGeneration, error) {
 	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
+	}
+
 	result, err := o.client.CreateChat(ctx, &openaiclient.ChatRequest{
 		Model:     opts.Model,
 		StopWords: opts.StopWords,
-		Messages:  messages,
+		Messages:  msgs,
 	})
 	if err != nil {
 		return nil, err
 	}
-	return &result.Choices[0].Message, nil
+	if len(result.Choices) == 0 {
+		return nil, ErrEmptyResponse
+	}
+	return &llms.ChatGeneration{
+		Message: &schema.AIChatMessage{
+			Text: result.Choices[0].Message.Content,
+		},
+		// TODO: fill in generation info
+	}, nil
 }
 
 // CreateEmbedding creates embeddings for the given input texts.