Adding llm_chat function to starlark stored procedure. (#22300)

Adding `llm_chat` to starlark stored procedure.

Approved by: @zhangxu19830126, @qingxinhome, @daviszhen, @heni02

Author: fengttt

go.mod

@@ -87,14 +87,15 @@ require (
 	github.com/ti-mo/netfilter v0.5.2
 	github.com/tidwall/btree v1.7.0
 	github.com/tidwall/pretty v1.2.1
+	github.com/tmc/langchaingo v0.1.13
 	github.com/unum-cloud/usearch/golang v0.0.0-20250207215718-306d6646b8f5
 	go.starlark.net v0.0.0-20250701195324-d457b4515e0e
 	go.uber.org/automaxprocs v1.5.3
 	go.uber.org/ratelimit v0.2.0
 	go.uber.org/zap v1.24.0
 	golang.org/x/exp v0.0.0-20241009180824-f66d83c29e7c
-	golang.org/x/sync v0.8.0
-	golang.org/x/sys v0.26.0
+	golang.org/x/sync v0.9.0
+	golang.org/x/sys v0.27.0
 	gonum.org/v1/gonum v0.14.0
 	google.golang.org/grpc v1.65.0
 	google.golang.org/protobuf v1.36.0
@@ -109,6 +110,7 @@ require (
 	github.com/cilium/ebpf v0.9.1 // indirect
 	github.com/clbanning/mxj v1.8.4 // indirect
 	github.com/coreos/go-systemd/v22 v22.3.2 // indirect
+	github.com/dlclark/regexp2 v1.10.0 // indirect
 	github.com/dustin/go-humanize v1.0.1 // indirect
 	github.com/dylibso/observe-sdk/go v0.0.0-20240819160327-2d926c5d788a // indirect
 	github.com/go-ini/ini v1.67.0 // indirect
@@ -142,17 +144,18 @@ require (
 	github.com/opencontainers/runtime-spec v1.0.2 // indirect
 	github.com/opentracing/opentracing-go v1.2.1-0.20220228012449-10b1cf09e00b // indirect
 	github.com/pbnjay/memory v0.0.0-20210728143218-7b4eea64cf58 // indirect
+	github.com/pkoukk/tiktoken-go v0.1.6 // indirect
 	github.com/rivo/uniseg v0.4.7 // indirect
 	github.com/rs/xid v1.6.0 // indirect
 	github.com/segmentio/asm v1.1.3 // indirect
 	github.com/shoenig/go-m1cpu v0.1.6 // indirect
 	github.com/tetratelabs/wabin v0.0.0-20230304001439-f6f874872834 // indirect
 	github.com/tetratelabs/wazero v1.8.1-0.20240916092830-1353ca24fef0 // indirect
 	go.opentelemetry.io/proto/otlp v1.3.1 // indirect
-	golang.org/x/crypto v0.28.0 // indirect
+	golang.org/x/crypto v0.29.0 // indirect
 	golang.org/x/net v0.30.0 // indirect
-	golang.org/x/text v0.19.0 // indirect
-	golang.org/x/time v0.3.0 // indirect
+	golang.org/x/text v0.20.0 // indirect
+	golang.org/x/time v0.5.0 // indirect
 )
 
 require (
@@ -209,7 +212,7 @@ require (
 	github.com/power-devops/perfstat v0.0.0-20210106213030-5aafc221ea8c // indirect
 	github.com/prometheus/common v0.44.0 // indirect
 	github.com/prometheus/procfs v0.11.1 // indirect
-	github.com/rogpeppe/go-internal v1.10.0 // indirect
+	github.com/rogpeppe/go-internal v1.11.0 // indirect
 	github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529 // indirect
 	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/smartystreets/assertions v1.13.1 // indirect

go.sum

@@ -20,6 +20,7 @@ import (
 
 	"github.com/matrixorigin/matrixone/pkg/common/moerr"
 	"github.com/matrixorigin/matrixone/pkg/logutil"
+	"github.com/matrixorigin/matrixone/pkg/monlp/llm"
 	"github.com/matrixorigin/matrixone/pkg/sql/parsers/tree"
 	"github.com/matrixorigin/matrixone/pkg/sql/plan"
 	ujson "github.com/matrixorigin/matrixone/pkg/util/json"
@@ -33,6 +34,8 @@ type starlarkInterpreter struct {
 	interp      *Interpreter
 	thread      *starlark.Thread
 	predeclared starlark.StringDict
+
+	llm llm.LLMClient
 }
 
 func convertToStarlarkValue(ctx context.Context, v any) (starlark.Value, error) {
@@ -202,11 +205,13 @@ func (si *starlarkInterpreter) buildModule() starlark.Value {
 	return &starlarkstruct.Module{
 		Name: "mo",
 		Members: starlark.StringDict{
-			"sql":    starlark.NewBuiltin("mo.sql", si.moSql),
-			"jq":     starlark.NewBuiltin("mo.jq", si.moJq),
-			"quote":  starlark.NewBuiltin("mo.quote", si.moQuote),
-			"getvar": starlark.NewBuiltin("mo.getvar", si.moGetVar),
-			"setvar": starlark.NewBuiltin("mo.getvar", si.moSetVar),
+			"sql":         starlark.NewBuiltin("mo.sql", si.moSql),
+			"jq":          starlark.NewBuiltin("mo.jq", si.moJq),
+			"quote":       starlark.NewBuiltin("mo.quote", si.moQuote),
+			"getvar":      starlark.NewBuiltin("mo.getvar", si.moGetVar),
+			"setvar":      starlark.NewBuiltin("mo.getvar", si.moSetVar),
+			"llm_connect": starlark.NewBuiltin("mo.llm_connect", si.moLlmConnect),
+			"llm_chat":    starlark.NewBuiltin("mo.llm_chat", si.moLlmChat),
 		},
 	}
 }
@@ -339,3 +344,66 @@ func (si *starlarkInterpreter) moSetVar(thread *starlark.Thread, b *starlark.Bui
 	}
 	return starlark.NewList(ret), nil
 }
+
+func (si *starlarkInterpreter) llmGetVar(name, orig string) string {
+	if orig == "" {
+		value, err := si.interp.ses.GetUserDefinedVar(name)
+		if err != nil {
+			return ""
+		}
+		return value.Value.(string)
+	}
+	return orig
+}
+
+func (si *starlarkInterpreter) llmConnect(server, addr, model, options string) (llm.LLMClient, error) {
+	server = si.llmGetVar("llm_server", server)
+	addr = si.llmGetVar("llm_addr", addr)
+	model = si.llmGetVar("llm_model", model)
+	options = si.llmGetVar("llm_options", options)
+	return llm.NewLLMClient(server, addr, model, options)
+}
+
+func (si *starlarkInterpreter) moLlmConnect(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var err error
+	var server, addr, model, options string
+	var ret = []starlark.Value{starlark.None, starlark.None}
+	if err := starlark.UnpackPositionalArgs("llm_connect", args, kwargs, 3, &server, &addr, &model, &options); err != nil {
+		ret[1] = starlark.String(err.Error())
+		return starlark.NewList(ret), nil
+	}
+
+	si.llm, err = si.llmConnect(server, addr, model, options)
+	if err != nil {
+		ret[1] = starlark.String(err.Error())
+		return starlark.NewList(ret), nil
+	}
+	return starlark.NewList(ret), nil
+}
+
+func (si *starlarkInterpreter) moLlmChat(thread *starlark.Thread, b *starlark.Builtin, args starlark.Tuple, kwargs []starlark.Tuple) (starlark.Value, error) {
+	var err error
+	var prompt string
+	var ret = []starlark.Value{starlark.None, starlark.None}
+	if err := starlark.UnpackPositionalArgs("llm_chat", args, kwargs, 1, &prompt); err != nil {
+		ret[1] = starlark.String(err.Error())
+		return starlark.NewList(ret), nil
+	}
+
+	if si.llm == nil {
+		si.llm, err = si.llmConnect("", "", "", "")
+		if err != nil {
+			ret[1] = starlark.String(err.Error())
+			return starlark.NewList(ret), nil
+		}
+	}
+
+	reply, err := si.llm.Chat(si.interp.ctx, prompt)
+	if err != nil {
+		ret[1] = starlark.String(err.Error())
+		return starlark.NewList(ret), nil
+	}
+
+	ret[0] = starlark.String(reply)
+	return starlark.NewList(ret), nil
+}

pkg/frontend/starlark_interpreter.go

@@ -0,0 +1,63 @@
+// Copyright 2024 Matrix Origin
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package llm
+
+import (
+	"context"
+	"encoding/json"
+
+	"github.com/matrixorigin/matrixone/pkg/common/moerr"
+)
+
+const (
+	MockServer   = ""
+	OllamaServer = "ollama"
+
+	MockEchoModel = "echo"
+
+	LLMRoleSystem = "system"
+	LLMRoleUser   = "user"
+	LLMRoleAI     = "ai"
+)
+
+type Message struct {
+	Role    string `json:"role"`
+	Content string `json:"content"`
+}
+
+type LLMClient interface {
+	ChatMsg(ctx context.Context, messages []Message) (string, error)
+	Chat(ctx context.Context, prompt string) (string, error)
+}
+
+func NewLLMClient(server string, addr string, model string, options string) (LLMClient, error) {
+	switch server {
+	case MockServer:
+		return NewMockClient(model, options)
+	case OllamaServer:
+		return NewOllamaClient(addr, model, options)
+	default:
+		return nil, moerr.NewInvalidInputf(context.TODO(), "invalid server: %s", server)
+	}
+}
+
+func stringToMessage(prompt string) ([]Message, error) {
+	var messages []Message
+	err := json.Unmarshal([]byte(prompt), &messages)
+	if err != nil {
+		return nil, err
+	}
+	return messages, nil
+}

pkg/monlp/llm/llm.go

@@ -0,0 +1,54 @@
+// Copyright 2024 Matrix Origin
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package llm
+
+import (
+	"context"
+
+	"github.com/matrixorigin/matrixone/pkg/common/moerr"
+)
+
+type MockClient struct {
+	model   string
+	options string
+}
+
+func NewMockClient(model string, options string) (*MockClient, error) {
+	switch model {
+	case MockEchoModel:
+		return &MockClient{model: model, options: options}, nil
+	default:
+		return nil, moerr.NewInvalidInputf(context.TODO(), "invalid model: %s", model)
+	}
+}
+
+func (c *MockClient) ChatMsg(ctx context.Context, messages []Message) (string, error) {
+	switch c.model {
+	case MockEchoModel:
+		if len(messages) == 0 {
+			return "", moerr.NewInvalidInputf(ctx, "no messages")
+		}
+		return messages[len(messages)-1].Content, nil
+	}
+	return "", moerr.NewInvalidInputf(ctx, "invalid model: %s", c.model)
+}
+
+func (c *MockClient) Chat(ctx context.Context, prompt string) (string, error) {
+	messages, err := stringToMessage(prompt)
+	if err != nil {
+		return "", err
+	}
+	return c.ChatMsg(ctx, messages)
+}

pkg/monlp/llm/mockllm.go

@@ -0,0 +1,74 @@
+// Copyright 2024 Matrix Origin
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package llm
+
+import (
+	"context"
+	"encoding/json"
+	"testing"
+)
+
+func TestMockEchoModel(t *testing.T) {
+	client, err := NewMockClient(MockEchoModel, "")
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	msgs := []Message{
+		{
+			Role:    "system",
+			Content: "You are a helpful assistant.",
+		},
+		{
+			Role:    "user",
+			Content: "Hello, world!",
+		},
+	}
+
+	reply, err := client.ChatMsg(context.Background(), msgs)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if reply != "Hello, world!" {
+		t.Fatal("reply is not correct")
+	}
+
+	prompt, err := json.Marshal(msgs)
+	if err != nil {
+		t.Fatal(err)
+	}
+	reply, err = client.Chat(context.Background(), string(prompt))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if reply != "Hello, world!" {
+		t.Fatal("reply is not correct")
+	}
+
+	msgs = append(msgs, Message{
+		Role:    "assistant",
+		Content: reply + " again",
+	})
+
+	reply, err = client.ChatMsg(context.Background(), msgs)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if reply != "Hello, world! again" {
+		t.Fatal("reply is not correct")
+	}
+}