diff --git a/chains/chains.go b/chains/chains.go index 8a121bd5..5ec1c448 100644 --- a/chains/chains.go +++ b/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. diff --git a/llms/llms.go b/llms/llms.go index 0d8a76bd..e9321a3d 100644 --- a/llms/llms.go +++ b/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"` +} diff --git a/llms/openai/internal/openaiclient/chat.go b/llms/openai/internal/openaiclient/chat.go index 1ddc4b86..841cbd91 100644 --- a/llms/openai/internal/openaiclient/chat.go +++ b/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. diff --git a/llms/openai/openaillm.go b/llms/openai/openaillm.go index 143fede0..e842fa40 100644 --- a/llms/openai/openaillm.go +++ b/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.