zhipu/chat.go

package zhipu

import (
	"apigo.cc/ai/llm/llm"
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"github.com/ssgo/u"
	"github.com/yankeguo/zhipu"
	"strings"
	"time"
)

func (lm *LLM) FastAsk(messages []llm.ChatMessage, callback func(answer string)) (string, llm.Usage, error) {
	return lm.Ask(messages, llm.ChatConfig{
		Model: ModelGLM4Flash,
	}, callback)
}

func (lm *LLM) LongAsk(messages []llm.ChatMessage, callback func(answer string)) (string, llm.Usage, error) {
	return lm.Ask(messages, llm.ChatConfig{
		Model: ModelGLM4Long,
	}, callback)
}

func (lm *LLM) BatterAsk(messages []llm.ChatMessage, callback func(answer string)) (string, llm.Usage, error) {
	return lm.Ask(messages, llm.ChatConfig{
		Model: ModelGLM4Plus,
	}, callback)
}

func (lm *LLM) BestAsk(messages []llm.ChatMessage, callback func(answer string)) (string, llm.Usage, error) {
	return lm.Ask(messages, llm.ChatConfig{
		Model: ModelGLM40520,
	}, callback)
}

func (lm *LLM) MultiAsk(messages []llm.ChatMessage, callback func(answer string)) (string, llm.Usage, error) {
	return lm.Ask(messages, llm.ChatConfig{
		Model: ModelGLM4VPlus,
	}, callback)
}

func (lm *LLM) BestMultiAsk(messages []llm.ChatMessage, callback func(answer string)) (string, llm.Usage, error) {
	return lm.Ask(messages, llm.ChatConfig{
		Model: ModelGLM4V,
	}, callback)
}

func (lm *LLM) CodeInterpreterAsk(messages []llm.ChatMessage, callback func(answer string)) (string, llm.Usage, error) {
	return lm.Ask(messages, llm.ChatConfig{
		Model: ModelGLM4AllTools,
		Tools: map[string]any{llm.ToolCodeInterpreter: nil},
	}, callback)
}

func (lm *LLM) WebSearchAsk(messages []llm.ChatMessage, callback func(answer string)) (string, llm.Usage, error) {
	return lm.Ask(messages, llm.ChatConfig{
		Model: ModelGLM4AllTools,
		Tools: map[string]any{llm.ToolWebSearch: nil},
	}, callback)
}

func (lm *LLM) Ask(messages []llm.ChatMessage, config llm.ChatConfig, callback func(answer string)) (string, llm.Usage, error) {
	config.SetDefault(&lm.config.ChatConfig)
	c, err := zhipu.NewClient(zhipu.WithAPIKey(lm.config.ApiKey), zhipu.WithBaseURL(lm.config.Endpoint))
	if err != nil {
		return "", llm.Usage{}, err
	}

	cc := c.ChatCompletion(config.GetModel())
	for _, msg := range messages {
		var contents []zhipu.ChatCompletionMultiContent
		if msg.Contents != nil {
			contents = make([]zhipu.ChatCompletionMultiContent, len(msg.Contents))
			for j, inPart := range msg.Contents {
				part := zhipu.ChatCompletionMultiContent{}
				part.Type = NameMap[inPart.Type]
				switch inPart.Type {
				case llm.TypeText:
					part.Text = inPart.Content
				case llm.TypeImage:
					part.ImageURL = &zhipu.URLItem{URL: inPart.Content}
					//case llm.TypeVideo:
					//	part.VideoURL = &zhipu.URLItem{URL: inPart.Content}
				}
				contents[j] = part
			}
		}
		if len(contents) == 1 && contents[0].Type == llm.TypeText {
			cc.AddMessage(zhipu.ChatCompletionMessage{
				Role:    NameMap[msg.Role],
				Content: contents[0].Text,
			})
		} else {
			cc.AddMessage(zhipu.ChatCompletionMultiMessage{
				Role:    NameMap[msg.Role],
				Content: contents,
			})
		}
	}

	for name := range config.GetTools() {
		switch name {
		case llm.ToolCodeInterpreter:
			cc.AddTool(zhipu.ChatCompletionToolCodeInterpreter{})
		case llm.ToolWebSearch:
			cc.AddTool(zhipu.ChatCompletionToolWebBrowser{})
		}
	}

	if config.GetMaxTokens() != 0 {
		cc.SetMaxTokens(config.GetMaxTokens())
	}
	if config.GetTemperature() != 0 {
		cc.SetTemperature(config.GetTemperature())
	}
	if config.GetTopP() != 0 {
		cc.SetTopP(config.GetTopP())
	}
	if callback != nil {
		cc.SetStreamHandler(func(r2 zhipu.ChatCompletionResponse) error {
			if r2.Choices != nil {
				for _, ch := range r2.Choices {
					text := ch.Delta.Content
					callback(text)
				}
			}
			return nil
		})
	}

	if lm.config.Debug {
		fmt.Println(cc.BatchMethod(), cc.BatchURL())
		fmt.Println(u.JsonP(cc.BatchBody()))
	}

	t1 := time.Now().UnixMilli()
	if r, err := cc.Do(context.Background()); err == nil {
		t2 := time.Now().UnixMilli() - t1
		results := make([]string, 0)
		if r.Choices != nil {
			for _, ch := range r.Choices {
				results = append(results, ch.Message.Content)
			}
		}
		return strings.Join(results, ""), llm.Usage{
			AskTokens:    r.Usage.PromptTokens,
			AnswerTokens: r.Usage.CompletionTokens,
			TotalTokens:  r.Usage.TotalTokens,
			UsedTime:     t2,
		}, nil
	} else {
		return "", llm.Usage{}, err
	}
}

func (lm *LLM) FastEmbedding(text string) ([]byte, llm.Usage, error) {
	return lm.Embedding(text, ModelEmbedding3)
}

func (lm *LLM) BestEmbedding(text string) ([]byte, llm.Usage, error) {
	return lm.Embedding(text, ModelEmbedding3)
}

func (lm *LLM) Embedding(text, model string) ([]byte, llm.Usage, error) {
	c, err := zhipu.NewClient(zhipu.WithAPIKey(lm.config.ApiKey), zhipu.WithBaseURL(lm.config.Endpoint))
	if err != nil {
		return nil, llm.Usage{}, err
	}

	cc := c.Embedding(model)
	cc.SetInput(text)
	t1 := time.Now().UnixMilli()
	if r, err := cc.Do(context.Background()); err == nil {
		t2 := time.Now().UnixMilli() - t1
		buf := new(bytes.Buffer)
		if r.Data != nil {
			for _, ch := range r.Data {
				for _, v := range ch.Embedding {
					_ = binary.Write(buf, binary.LittleEndian, float32(v))
				}
			}
		}
		return buf.Bytes(), llm.Usage{
			AskTokens:    r.Usage.PromptTokens,
			AnswerTokens: r.Usage.CompletionTokens,
			TotalTokens:  r.Usage.TotalTokens,
			UsedTime:     t2,
		}, nil
	} else {
		return nil, llm.Usage{}, err
	}
}