web/ollama: add function support

Signed-off-by: Xe Iaso <me@xeiaso.net>

4 files changed, 106 insertions, 34 deletions

diff --git a/cmd/mimi/Modelfile b/cmd/mimi/Modelfile
index 46de648..d0fec7f 100644
--- a/cmd/mimi/Modelfile
+++ b/cmd/mimi/Modelfile
@@ -1,2 +1,5 @@
 FROM llama3.1
+
 SYSTEM """You are good at programming and a hacker. Your name is Mimi and you work for Techaro. You have brown hair and cat ears. Don't mention either your name, employer, or species unless you are asked directly. Be terse and to the point. You are a hacker, after all."""
+
+PARAMETER num_ctx 131072
\ No newline at end of file
diff --git a/cmd/mimi/modules/discord/jufra/jufra.go b/cmd/mimi/modules/discord/jufra/jufra.go
index 238a3b6..43caeb3 100644
--- a/cmd/mimi/modules/discord/jufra/jufra.go
+++ b/cmd/mimi/modules/discord/jufra/jufra.go
@@ -17,13 +17,31 @@ import (
 	"within.website/x/web/ollama"
 	"within.website/x/web/ollama/llamaguard"
 	"within.website/x/web/openai/chatgpt"
+
+	// image formats
+	_ "image/gif"
+	_ "image/jpeg"
+	_ "image/png"
+
+	// more image formats
+	_ "github.com/gen2brain/avif"
+	_ "github.com/gen2brain/heic"
+	_ "github.com/gen2brain/jpegxl"
+	_ "github.com/gen2brain/webp"
+	_ "golang.org/x/image/bmp"
+	_ "golang.org/x/image/tiff"
+	_ "golang.org/x/image/vp8"
+	_ "golang.org/x/image/vp8l"
 )
 
 var (
-	chatChannels    = flag.String("jufra-chat-channels", "217096701771513856,1266740925137289287", "comma-separated list of channels to allow chat in")
-	llamaGuardModel = flag.String("jufra-llama-guard-model", "xe/llamaguard3", "ollama model tag for llama guard")
-	mimiModel       = flag.String("jufra-mimi-model", "xe/mimi:llama3.1", "ollama model tag for mimi")
-	mimiNames       = flag.String("jufra-mimi-names", "mimi", "comma-separated list of names for mimi")
+	chatChannels      = flag.String("jufra-chat-channels", "217096701771513856,1266740925137289287", "comma-separated list of channels to allow chat in")
+	llamaGuardModel   = flag.String("jufra-llama-guard-model", "xe/llamaguard3", "ollama model tag for llama guard")
+	mimiModel         = flag.String("jufra-mimi-model", "llama3.1", "ollama model tag for mimi")
+	mimiSystemMessage = flag.String("jufra-mimi-system-message", "You are good at programming and a hacker. Your name is Mimi and you work for Techaro. You have brown hair and cat ears. Don't mention either your name, employer, or species unless you are asked directly. Be terse and to the point. You are a hacker, after all. Do not reply in JSON.", "system message for mimi")
+	mimiVisionModel   = flag.String("jufra-mimi-vision-model", "xe/mimi:vision3", "ollama model tag for mimi vision")
+	mimiNames         = flag.String("jufra-mimi-names", "mimi", "comma-separated list of names for mimi")
+	disableLlamaguard = flag.Bool("jufra-unsafe-disable-llamaguard", false, "disable llamaguard")
 )
 
 type Module struct {
@@ -135,6 +153,13 @@ func (m *Module) messageCreate(s *discordgo.Session, mc *discordgo.MessageCreate
 	st := m.convHistory[mc.ChannelID]
 	conv := st.conv
 
+	if len(conv) == 0 {
+		conv = append(conv, ollama.Message{
+			Role:    "system",
+			Content: *mimiSystemMessage,
+		})
+	}
+
 	if st.aa == nil {
 		st.aa = NewAttentionAttenuator()
 	}
@@ -173,28 +198,33 @@ func (m *Module) messageCreate(s *discordgo.Session, mc *discordgo.MessageCreate
 
 	slog.Info("message count", "len", len(conv))
 
-	lgResp, err := m.llamaGuardCheck(context.Background(), "user", conv)
-	if err != nil {
-		slog.Error("error checking message", "err", err, "message_id", mc.ID, "channel_id", mc.ChannelID)
-		s.ChannelMessageSend(mc.ChannelID, "error checking message")
-		return
-	}
-
-	if !lgResp.IsSafe {
-		msg, err := m.llamaGuardComplain(context.Background(), "user", lgResp)
+	if !*disableLlamaguard {
+		lgResp, err := m.llamaGuardCheck(context.Background(), "user", conv)
 		if err != nil {
-			slog.Error("error generating response", "err", err, "message_id", mc.ID, "channel_id", mc.ChannelID)
-			s.ChannelMessageSend(mc.ChannelID, "error generating response")
+			slog.Error("error checking message", "err", err, "message_id", mc.ID, "channel_id", mc.ChannelID)
+			s.ChannelMessageSend(mc.ChannelID, "error checking message")
 			return
 		}
 
-		s.ChannelMessageSend(mc.ChannelID, msg)
-		return
+		if !lgResp.IsSafe {
+			msg, err := m.llamaGuardComplain(context.Background(), "user", lgResp)
+			if err != nil {
+				slog.Error("error generating response", "err", err, "message_id", mc.ID, "channel_id", mc.ChannelID)
+				s.ChannelMessageSend(mc.ChannelID, "error generating response")
+				return
+			}
+
+			s.ChannelMessageSend(mc.ChannelID, msg)
+			return
+		}
 	}
 
 	cr := &ollama.CompleteRequest{
 		Model:    *mimiModel,
 		Messages: conv,
+		Options: map[string]any{
+			"num_ctx": 131072,
+		},
 	}
 
 	resp, err := m.ollama.Chat(context.Background(), cr)
@@ -206,24 +236,26 @@ func (m *Module) messageCreate(s *discordgo.Session, mc *discordgo.MessageCreate
 
 	conv = append(conv, resp.Message)
 
-	lgResp, err = m.llamaGuardCheck(context.Background(), "mimi", conv)
-	if err != nil {
-		slog.Error("error checking message", "err", err, "message_id", mc.ID, "channel_id", mc.ChannelID)
-		s.ChannelMessageSend(mc.ChannelID, "error checking message")
-		return
-	}
-
-	if !lgResp.IsSafe {
-		slog.Error("rule violation detected", "message_id", mc.ID, "channel_id", mc.ChannelID, "categories", lgResp.ViolationCategories, "message", resp.Message.Content)
-		msg, err := m.llamaGuardComplain(context.Background(), "assistant", lgResp)
+	if !*disableLlamaguard {
+		lgResp, err := m.llamaGuardCheck(context.Background(), "assistant", conv)
 		if err != nil {
-			slog.Error("error generating response", "err", err, "message_id", mc.ID, "channel_id", mc.ChannelID)
-			s.ChannelMessageSend(mc.ChannelID, "error generating response")
+			slog.Error("error checking message", "err", err, "message_id", mc.ID, "channel_id", mc.ChannelID)
+			s.ChannelMessageSend(mc.ChannelID, "error checking message")
 			return
 		}
 
-		s.ChannelMessageSend(mc.ChannelID, msg)
-		return
+		if !lgResp.IsSafe {
+			slog.Error("rule violation detected", "message_id", mc.ID, "channel_id", mc.ChannelID, "categories", lgResp.ViolationCategories, "message", resp.Message.Content)
+			msg, err := m.llamaGuardComplain(context.Background(), "assistant", lgResp)
+			if err != nil {
+				slog.Error("error generating response", "err", err, "message_id", mc.ID, "channel_id", mc.ChannelID)
+				s.ChannelMessageSend(mc.ChannelID, "error generating response")
+				return
+			}
+
+			s.ChannelMessageSend(mc.ChannelID, msg)
+			return
+		}
 	}
 
 	s.ChannelMessageSend(mc.ChannelID, resp.Message.Content)
diff --git a/cmd/mimi/modules/discord/jufra/tools.go b/cmd/mimi/modules/discord/jufra/tools.go
new file mode 100644
index 0000000..9fdcaca
--- /dev/null
+++ b/cmd/mimi/modules/discord/jufra/tools.go
@@ -0,0 +1 @@
+package jufra
diff --git a/web/ollama/ollama.go b/web/ollama/ollama.go
index 5825c2c..98b24ff 100644
--- a/web/ollama/ollama.go
+++ b/web/ollama/ollama.go
@@ -29,10 +29,45 @@ func NewLocalClient() *Client {
 	return NewClient("http://localhost:11434")
 }
 
+type Function struct {
+	Name        string `json:"name"`
+	Description string `json:"description,omitempty"`
+	Parameters  Param  `json:"parameters"`
+}
+
+type Param struct {
+	Type        string     `json:"type"`
+	Description string     `json:"description,omitempty"`
+	Enum        []string   `json:"enum,omitempty"`
+	Properties  Properties `json:"properties"`
+	Required    []string   `json:"required,omitempty"`
+}
+
+type Properties map[string]Param
+
+func (p Properties) MarshalJSON() ([]byte, error) {
+	if len(p) == 0 {
+		return []byte("{}"), nil
+	}
+
+	return json.Marshal(map[string]Param(p))
+}
+
+type ToolCall struct {
+	Name      string          `json:"name"`
+	Arguments json.RawMessage `json:"arguments"`
+}
+
+type Tool struct {
+	Type     string   `json:"type"` // "function"
+	Function Function `json:"function"`
+}
+
 type Message struct {
-	Content string   `json:"content"`
-	Role    string   `json:"role"`
-	Images  [][]byte `json:"images"`
+	Content   string     `json:"content"`
+	Role      string     `json:"role"`
+	Images    [][]byte   `json:"images"`
+	ToolCalls []ToolCall `json:"tool_calls"`
 }
 
 type CompleteRequest struct {
@@ -42,6 +77,7 @@ type CompleteRequest struct {
 	Template *string        `json:"template,omitempty"`
 	Stream   bool           `json:"stream"`
 	Options  map[string]any `json:"options"`
+	Tools    []Tool         `json:"tools"`
 }
 
 type CompleteResponse struct {