feat: add openai function call

Author: sunls24

internal/cron/v2ex.go

@@ -2,7 +2,6 @@ package cron
 
 import (
 	"fmt"
-	"github.com/tidwall/gjson"
 	"log/slog"
 	"net/http"
 	"strings"
@@ -13,10 +12,12 @@ import (
 	"webot/internal/types"
 	"webot/pkg/client"
 	"webot/pkg/openai"
+
+	"github.com/tidwall/gjson"
 )
 
 const (
-	v2exSpec       = "30 8-18/3 * * *"
+	v2exSpec       = "30 10-18/3 * * *"
 	v2exHotAPI     = "https://www.v2ex.com/api/topics/hot.json"
 	v2exRepliesAPI = "https://www.v2ex.com/api/replies/show.json?topic_id="
 )
@@ -62,7 +63,7 @@ func (ctx v2ex) run() {
 			Title: title,
 		})
 	}
-	if err = PushToAll("V2EX 热帖推送"+v2exSplit+strings.Join(summarizeList, v2exSplit), types.PushV2ex); err != nil {
+	if err = PushToAll(v2exHeader+strings.Join(summarizeList, v2exSplit)+v2exSplit+v2exFooter, types.PushV2ex); err != nil {
 		slog.Error("push v2ex post failed", slog.Any("err", err))
 	}
 }
@@ -106,5 +107,10 @@ AI 点评：%s
 
 const (
 	v2exSystem = "你擅长点评 V2EX 论坛热门帖子，从主题内容和回复中提取关键信息，并使用 60 个汉字以内的简洁描述。"
-	v2exSplit  = "\n---------------\n"
+	v2exSplit  = "\n-------------------\n"
+
+	v2exHeader = `V2EX 热帖推送
+---------------
+`
+	v2exFooter = "输入 v2ex 关闭推送"
 )

internal/handler/handler.go

@@ -105,7 +105,7 @@ func onThink(msg *openwechat.Message, done <-chan struct{}, atUser string) {
 	select {
 	case <-done:
 		return
-	case <-time.After(time.Second * 2):
+	case <-time.After(time.Second * 5):
 		_, err := msg.ReplyText(atUser + thinkText())
 		if err != nil {
 			slog.Error("reply think failed", slog.Any("err", err))
@@ -132,7 +132,7 @@ func (h *Handler) onText(msg *openwechat.Message, ctx params) {
 	done := make(chan struct{}, 1)
 	go onThink(msg, done, ctx.atUser)
 	start := time.Now()
-	result, err := h.AI.Chat(h.Cfg.GetModel(false), messages)
+	result, err := h.AI.Chat(h.Cfg.GetModel(false), messages, openai.Google)
 	if err != nil {
 		done <- struct{}{}
 		slog.Error("AI chat failed", attr, slog.Any("err", err))

pkg/openai/google.go

@@ -0,0 +1,55 @@
+package openai
+
+import (
+	"fmt"
+	"github.com/tidwall/gjson"
+	"log/slog"
+	"net/http"
+	"net/url"
+	"os"
+	"webot/pkg/client"
+)
+
+const (
+	nameGoogle = "GoogleSearch"
+)
+
+var Google = Tool{
+	Name: nameGoogle,
+	Func: fmt.Sprintf(`{
+  "type": "function",
+  "function": {
+    "name": "%s",
+    "strict": true,
+    "description": "Using Google to search the internet",
+    "parameters": {
+      "type": "object",
+      "properties": {
+        "keyword": { "type": "string", "description": "Search keyword" }
+      },
+      "required": ["keyword"]
+    }
+  }
+}`, nameGoogle),
+	Call: callGoogle,
+}
+
+func callGoogle(keyword string) string {
+	keyword = gjson.Get(keyword, "keyword").String()
+	slog.Info("call google search", slog.String("keyword", keyword))
+
+	var apiKey, _ = os.LookupEnv("GOOGLE_API_KEY")
+	var engineId, _ = os.LookupEnv("GOOGLE_ENGINE_ID")
+	if apiKey == "" || engineId == "" {
+		slog.Warn("not found env GOOGLE_API_KEY or GOOGLE_ENGINE_ID")
+		return "nothing"
+	}
+	req, _ := http.NewRequest(http.MethodGet, fmt.Sprintf("https://www.googleapis.com/customsearch/v1?&fields=items(title,link,snippet,pagemap/metatags(og:description))&key=%s&cx=%s&q=%s", apiKey, engineId, url.QueryEscape(keyword)), nil)
+	resp, err := client.Do(req)
+	if err != nil {
+		slog.Error("google search failed", slog.Any("err", err))
+		return "nothing"
+	}
+
+	return gjson.GetBytes(resp, "items").String()
+}

pkg/openai/openai.go

@@ -1,10 +1,12 @@
 package openai
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
 	"github.com/tidwall/gjson"
 	"net/http"
+	"strings"
 	"webot/pkg/client"
 )
 
@@ -23,13 +25,23 @@ func New(opts Options) *OpenAI {
 	}
 }
 
-func (oai *OpenAI) Chat(model string, messages []Message) (*Message, error) {
+func (oai *OpenAI) Chat(model string, messages []Message, tools ...Tool) (*Message, error) {
 	const api = "/v1/chat/completions"
 
+	var toolsBody json.RawMessage
+	if len(tools) != 0 {
+		var funcList = make([]string, 0, len(tools))
+		for _, t := range tools {
+			funcList = append(funcList, t.Func)
+		}
+		toolsBody = json.RawMessage(fmt.Sprintf("[%s]", strings.Join(funcList, ",")))
+	}
+
 	reqBody := RequestBody{
 		Model:       model,
 		Messages:    messages,
 		Temperature: 0.5,
+		Tools:       toolsBody,
 	}
 
 	req, _ := http.NewRequest(http.MethodPost, oai.opts.BaseURL+api, client.MarshalBody(reqBody))
@@ -39,7 +51,37 @@ func (oai *OpenAI) Chat(model string, messages []Message) (*Message, error) {
 	if err != nil {
 		return nil, err
 	}
-	message := NewMessage(gjson.GetBytes(resp, "choices.0.message"))
+
+	result := gjson.GetBytes(resp, "choices.0.message")
+	var toolCalls = result.Get("tool_calls")
+	var toolCallsResult []Message
+	for _, call := range toolCalls.Array() {
+		callId := call.Get("id").String()
+		call = call.Get("function")
+		name := call.Get("name").String()
+		args := call.Get("arguments").String()
+
+		for _, t := range tools {
+			if t.Name != name {
+				continue
+			}
+			toolCallsResult = append(toolCallsResult, Message{
+				Role:       RTool,
+				Content:    t.Call(args),
+				ToolCallId: callId,
+			})
+		}
+	}
+	if len(toolCallsResult) != 0 {
+		messages = append(messages, Message{
+			Role:      RAssistant,
+			ToolCalls: json.RawMessage(toolCalls.Raw),
+		})
+		messages = append(messages, toolCallsResult...)
+		return oai.Chat(model, messages, tools...)
+	}
+
+	message := NewMessage(result)
 	if message.Role == "" {
 		return nil, errors.New(string(resp))
 	}

pkg/openai/openai_test.go

@@ -1,7 +1,6 @@
 package openai
 
 import (
-	"fmt"
 	"os"
 	"testing"
 )
@@ -24,11 +23,11 @@ func TestOpenAI_Chat(t *testing.T) {
 	msg, err := openai.Chat(model, []Message{
 		{
 			Role:    RUser,
-			Content: "help me calculate 1+2*3",
+			Content: "help me search linux.do",
 		},
-	})
+	}, Google)
 	if err != nil {
 		t.Fatal(err)
 	}
-	fmt.Println(msg.Role, msg.Content)
+	t.Logf("%s: %s\n", msg.Role, msg.Content)
 }

pkg/openai/types.go

@@ -17,6 +17,9 @@ const (
 type Message struct {
 	Role    Role   `json:"role"`
 	Content string `json:"content"`
+
+	ToolCalls  json.RawMessage `json:"tool_calls,omitempty"`
+	ToolCallId string          `json:"tool_call_id,omitempty"`
 }
 
 func NewMessage(msg gjson.Result) *Message {
@@ -32,3 +35,11 @@ type RequestBody struct {
 	Temperature float64         `json:"temperature"`
 	Tools       json.RawMessage `json:"tools,omitempty"`
 }
+
+type ToolCall func(string) string
+
+type Tool struct {
+	Name string
+	Func string
+	Call ToolCall
+}