Merge branch 'master' into ci/static-check

Signed-off-by: Ettore Di Giacinto <mudler@users.noreply.github.com>

.github/workflows/comment-pr.yaml

@@ -8,9 +8,10 @@ jobs:
         MODEL_NAME: hermes-2-theta-llama-3-8b
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v4
+    - name: Checkout code
+      uses: actions/checkout@v3
       with:
-        fetch-depth: 0 # needed to checkout all branches for this Action to work
+        ref: "${{ github.event.pull_request.merge_commit_sha }}"
     - uses: mudler/localai-github-action@v1
       with:
         model: 'hermes-2-theta-llama-3-8b' # Any from models.localai.io, or from huggingface.com with: "huggingface://<repository>/file"

.github/workflows/static-check.yml

@@ -0,0 +1,70 @@
+name: static check
+on: pull_request
+
+jobs:
+  imports:
+    name: Imports
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - name: check
+      uses: danhunsaker/golang-github-actions@v1.3.0
+      with:
+        run: imports
+        token: ${{ secrets.GITHUB_TOKEN }}
+
+  errcheck:
+    name: Errcheck
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - name: check
+      uses: danhunsaker/golang-github-actions@v1.3.0
+      with:
+        run: errcheck
+        token: ${{ secrets.GITHUB_TOKEN }}
+
+  lint:
+    name: Lint
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - name: check
+      uses: danhunsaker/golang-github-actions@v1.3.0
+      with:
+        run: lint
+        token: ${{ secrets.GITHUB_TOKEN }}
+
+  shadow:
+    name: Shadow
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - name: check
+      uses: danhunsaker/golang-github-actions@v1.3.0
+      with:
+        run: shadow
+        token: ${{ secrets.GITHUB_TOKEN }}
+
+  staticcheck:
+    name: StaticCheck
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - name: check
+      uses: danhunsaker/golang-github-actions@v1.3.0
+      with:
+        run: staticcheck
+        token: ${{ secrets.GITHUB_TOKEN }}
+
+  sec:
+    name: Sec
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@master
+    - name: check
+      uses: danhunsaker/golang-github-actions@v1.3.0
+      with:
+        run: sec
+        token: ${{ secrets.GITHUB_TOKEN }}
+        flags: "-exclude=G104"

core/cli/federated.go

@@ -13,6 +13,7 @@ type FederatedCLI struct {
 }
 
 func (f *FederatedCLI) Run(ctx *cliContext.Context) error {
+
 	fs := p2p.NewFederatedServer(f.Address, p2p.FederatedID, f.Peer2PeerToken)
 
 	return fs.Start(context.Background())

core/cli/run.go

@@ -119,7 +119,7 @@ func (r *RunCMD) Run(ctx *cliContext.Context) error {
 		}
 
 		log.Info().Msg("Starting P2P server discovery...")
-		if err := p2p.ServiceDiscoverer(context.Background(), node, token, "", func() {
+		if err := p2p.ServiceDiscoverer(context.Background(), node, token, "", func(serviceID string, node p2p.NodeData) {
 			var tunnelAddresses []string
 			for _, v := range p2p.GetAvailableNodes("") {
 				if v.IsOnline() {

core/http/endpoints/openai/chat.go

@@ -225,18 +225,10 @@ func ChatEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, startup
 			}
 
 			// Update input grammar
-			// Handle if we should return "name" instead of "functions"
+			jsStruct := funcs.ToJSONStructure(config.FunctionsConfig.FunctionNameKey, config.FunctionsConfig.FunctionNameKey)
-			if config.FunctionsConfig.FunctionName {
+			config.Grammar = jsStruct.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
-				jsStruct := funcs.ToJSONNameStructure()
-				config.Grammar = jsStruct.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
-			} else {
-				jsStruct := funcs.ToJSONFunctionStructure()
-				config.Grammar = jsStruct.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
-			}
 		case input.JSONFunctionGrammarObject != nil:
 			config.Grammar = input.JSONFunctionGrammarObject.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
-		case input.JSONFunctionGrammarObjectName != nil:
-			config.Grammar = input.JSONFunctionGrammarObjectName.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
 		default:
 			// Force picking one of the functions by the request
 			if config.FunctionToCall() != "" {

core/p2p/federated.go

@@ -1,5 +1,7 @@
 package p2p
 
+const FederatedID = "federated"
+
 type FederatedServer struct {
 	listenAddr, service, p2ptoken string
 }

core/p2p/federated_server.go

@@ -7,34 +7,35 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	"io"
 	"net"
 	"time"
 
-	"github.com/rs/zerolog/log"
-
 	"math/rand/v2"
 
 	"github.com/mudler/edgevpn/pkg/node"
 	"github.com/mudler/edgevpn/pkg/protocol"
 	"github.com/mudler/edgevpn/pkg/types"
+	"github.com/rs/zerolog/log"
 )
 
-func (fs *FederatedServer) Start(ctx context.Context) error {
+func (f *FederatedServer) Start(ctx context.Context) error {
-	n, err := NewNode(fs.p2ptoken)
+
+	n, err := NewNode(f.p2ptoken)
 	if err != nil {
 		return fmt.Errorf("creating a new node: %w", err)
 	}
 	err = n.Start(ctx)
 	if err != nil {
-		return fmt.Errorf("starting a new node: %w", err)
+		return fmt.Errorf("creating a new node: %w", err)
 	}
 
-	if err := ServiceDiscoverer(ctx, n, fs.p2ptoken, FederatedID, nil); err != nil {
+	if err := ServiceDiscoverer(ctx, n, f.p2ptoken, f.service, func(servicesID string, tunnel NodeData) {
+		log.Debug().Msgf("Discovered node: %s", tunnel.ID)
+	}); err != nil {
 		return err
 	}
 
-	return fs.proxy(ctx, n)
+	return f.proxy(ctx, n)
 }
 
 func (fs *FederatedServer) proxy(ctx context.Context, node *node.Node) error {
@@ -84,44 +85,43 @@ func (fs *FederatedServer) proxy(ctx context.Context, node *node.Node) error {
 			}
 
 			// Handle connections in a new goroutine, forwarding to the p2p service
-			go handleConn(conn)
+			go func() {
+				var tunnelAddresses []string
+				for _, v := range GetAvailableNodes(fs.service) {
+					if v.IsOnline() {
+						tunnelAddresses = append(tunnelAddresses, v.TunnelAddress)
+					} else {
+						log.Info().Msgf("Node %s is offline", v.ID)
+					}
+				}
+
+				if len(tunnelAddresses) == 0 {
+					log.Error().Msg("No available nodes yet")
+					return
+				}
+
+				// open a TCP stream to one of the tunnels
+				// chosen randomly
+				// TODO: optimize this and track usage
+				tunnelAddr := tunnelAddresses[rand.IntN(len(tunnelAddresses))]
+
+				tunnelConn, err := net.Dial("tcp", tunnelAddr)
+				if err != nil {
+					log.Error().Err(err).Msg("Error connecting to tunnel")
+					return
+				}
+
+				log.Info().Msgf("Redirecting %s to %s", conn.LocalAddr().String(), tunnelConn.RemoteAddr().String())
+				closer := make(chan struct{}, 2)
+				go copyStream(closer, tunnelConn, conn)
+				go copyStream(closer, conn, tunnelConn)
+				<-closer
+
+				tunnelConn.Close()
+				conn.Close()
+				//	ll.Infof("(service %s) Done handling %s", serviceID, l.Addr().String())
+			}()
 		}
 	}
 
 }
-
-func handleConn(conn net.Conn) {
-	var tunnelAddresses []string
-	for _, v := range GetAvailableNodes(FederatedID) {
-		if v.IsOnline() {
-			tunnelAddresses = append(tunnelAddresses, v.TunnelAddress)
-		} else {
-			log.Info().Msgf("Node %s is offline", v.ID)
-		}
-	}
-
-	// open a TCP stream to one of the tunnels
-	// chosen randomly
-	// TODO: optimize this and track usage
-	tunnelAddr := tunnelAddresses[rand.IntN(len(tunnelAddresses))]
-
-	tunnelConn, err := net.Dial("tcp", tunnelAddr)
-	if err != nil {
-		log.Error().Err(err).Msg("Error connecting to tunnel")
-		return
-	}
-
-	log.Info().Msgf("Redirecting %s to %s", conn.LocalAddr().String(), tunnelConn.RemoteAddr().String())
-	closer := make(chan struct{}, 2)
-	go copyStream(closer, tunnelConn, conn)
-	go copyStream(closer, conn, tunnelConn)
-	<-closer
-
-	tunnelConn.Close()
-	conn.Close()
-}
-
-func copyStream(closer chan struct{}, dst io.Writer, src io.Reader) {
-	defer func() { closer <- struct{}{} }() // connection is closed, send signal to stop proxy
-	io.Copy(dst, src)
-}

core/p2p/node.go

@@ -6,7 +6,6 @@ import (
 )
 
 const defaultServicesID = "services_localai"
-const FederatedID = "federated"
 
 type NodeData struct {
 	Name          string

core/p2p/p2p.go

@@ -7,6 +7,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"io"
 	"net"
 	"os"
 	"sync"
@@ -138,7 +139,7 @@ func allocateLocalService(ctx context.Context, node *node.Node, listenAddr, serv
 
 // This is the main of the server (which keeps the env variable updated)
 // This starts a goroutine that keeps LLAMACPP_GRPC_SERVERS updated with the discovered services
-func ServiceDiscoverer(ctx context.Context, n *node.Node, token, servicesID string, discoveryFunc func()) error {
+func ServiceDiscoverer(ctx context.Context, n *node.Node, token, servicesID string, discoveryFunc func(serviceID string, node NodeData)) error {
 	if servicesID == "" {
 		servicesID = defaultServicesID
 	}
@@ -160,7 +161,7 @@ func ServiceDiscoverer(ctx context.Context, n *node.Node, token, servicesID stri
 			case tunnel := <-tunnels:
 				AddNode(servicesID, tunnel)
 				if discoveryFunc != nil {
-					discoveryFunc()
+					discoveryFunc(servicesID, tunnel)
 				}
 			}
 		}
@@ -390,3 +391,8 @@ func newNodeOpts(token string) ([]node.Option, error) {
 
 	return nodeOpts, nil
 }
+
+func copyStream(closer chan struct{}, dst io.Writer, src io.Reader) {
+	defer func() { closer <- struct{}{} }() // connection is closed, send signal to stop proxy
+	io.Copy(dst, src)
+}

core/p2p/p2p_disabled.go

@@ -14,11 +14,11 @@ func GenerateToken() string {
 	return "not implemented"
 }
 
-func (fs *FederatedServer) Start(ctx context.Context) error {
+func (f *FederatedServer) Start(ctx context.Context) error {
 	return fmt.Errorf("not implemented")
 }
 
-func ServiceDiscoverer(ctx context.Context, node *node.Node, token, servicesID string, fn func()) error {
+func ServiceDiscoverer(ctx context.Context, node *node.Node, token, servicesID string, fn func(string, NodeData)) error {
 	return fmt.Errorf("not implemented")
 }
 

core/schema/openai.go

@@ -179,8 +179,7 @@ type OpenAIRequest struct {
 	// A grammar to constrain the LLM output
 	Grammar string `json:"grammar" yaml:"grammar"`
 
-	JSONFunctionGrammarObject     *functions.JSONFunctionStructureFunction `json:"grammar_json_functions" yaml:"grammar_json_functions"`
+	JSONFunctionGrammarObject *functions.JSONFunctionStructure `json:"grammar_json_functions" yaml:"grammar_json_functions"`
-	JSONFunctionGrammarObjectName *functions.JSONFunctionStructureName     `json:"grammar_json_name" yaml:"grammar_json_name"`
 
 	Backend string `json:"backend" yaml:"backend"`
 

docs/content/docs/advanced/advanced-usage.md

@@ -152,7 +152,8 @@ function:
     replace_function_results: [] # Placeholder to replace function call results with arbitrary strings or patterns.
     replace_llm_results: [] # Replace language model results with arbitrary strings or patterns.
     capture_llm_results: [] # Capture language model results as text result, among JSON, in function calls. For instance, if a model returns a block for "thinking" and a block for "response", this will allow you to capture the thinking block.
-    return_name_in_function_response: false # Some models might prefer to use "name" rather then "function" when returning JSON data. This will allow to use "name" as a key in the JSON response.
+    function_name_key: "name"
+    function_arguments_key: "arguments"
 
 # Feature gating flags to enable experimental or optional features.
 feature_flags: {}

pkg/functions/functions.go

@@ -6,6 +6,11 @@ import (
 	"github.com/rs/zerolog/log"
 )
 
+const (
+	defaultFunctionNameKey      = "name"
+	defaultFunctionArgumentsKey = "arguments"
+)
+
 type Function struct {
 	Name        string                 `json:"name"`
 	Description string                 `json:"description"`
@@ -19,50 +24,18 @@ type Tool struct {
 }
 type Tools []Tool
 
-// ToJSONFunctionStructure converts a list of functions to a JSON structure that can be parsed to a grammar
-// This allows the LLM to return a response of the type: { "function": "function_name", "arguments": { "arg1": "value1", "arg2": "value2" } }
-func (f Functions) ToJSONFunctionStructure() JSONFunctionStructureFunction {
-	js := JSONFunctionStructureFunction{}
-	for _, function := range f {
-		//	t := function.Parameters["type"]
-		//tt := t.(string)
-
-		properties := function.Parameters["properties"]
-		defs := function.Parameters["$defs"]
-		dat, _ := json.Marshal(properties)
-		dat2, _ := json.Marshal(defs)
-		prop := map[string]interface{}{}
-		defsD := map[string]interface{}{}
-
-		err := json.Unmarshal(dat, &prop)
-		if err != nil {
-			log.Error().Err(err).Msg("error unmarshalling dat")
-		}
-		err = json.Unmarshal(dat2, &defsD)
-		if err != nil {
-			log.Error().Err(err).Msg("error unmarshalling dat2")
-		}
-		if js.Defs == nil {
-			js.Defs = defsD
-		}
-		js.OneOf = append(js.OneOf, ItemFunction{
-			Type: "object",
-			Properties: FunctionProperties{
-				Function: FunctionName{Const: function.Name},
-				Arguments: Argument{
-					Type:       "object",
-					Properties: prop,
-				},
-			},
-		})
-	}
-	return js
-}
-
 // ToJSONNameStructure converts a list of functions to a JSON structure that can be parsed to a grammar
 // This allows the LLM to return a response of the type: { "name": "function_name", "arguments": { "arg1": "value1", "arg2": "value2" } }
-func (f Functions) ToJSONNameStructure() JSONFunctionStructureName {
+func (f Functions) ToJSONStructure(name, args string) JSONFunctionStructure {
-	js := JSONFunctionStructureName{}
+	nameKey := defaultFunctionNameKey
+	argsKey := defaultFunctionArgumentsKey
+	if name != "" {
+		nameKey = name
+	}
+	if args != "" {
+		argsKey = args
+	}
+	js := JSONFunctionStructure{}
 	for _, function := range f {
 		//	t := function.Parameters["type"]
 		//tt := t.(string)
@@ -85,15 +58,16 @@ func (f Functions) ToJSONNameStructure() JSONFunctionStructureName {
 		if js.Defs == nil {
 			js.Defs = defsD
 		}
-		js.OneOf = append(js.OneOf, ItemName{
+
-			Type: "object",
+		property := map[string]interface{}{}
-			Properties: NameProperties{
+		property[nameKey] = FunctionName{Const: function.Name}
-				Function: FunctionName{Const: function.Name},
+		property[argsKey] = Argument{
-				Arguments: Argument{
+			Type:       "object",
-					Type:       "object",
+			Properties: prop,
-					Properties: prop,
+		}
-				},
+		js.OneOf = append(js.OneOf, Item{
-			},
+			Type:       "object",
+			Properties: property,
 		})
 	}
 	return js

pkg/functions/functions_test.go

@@ -35,21 +35,35 @@ var _ = Describe("LocalAI grammar functions", func() {
 				},
 			}
 
-			js := functions.ToJSONFunctionStructure()
+			js := functions.ToJSONStructure("function", "arguments")
 			Expect(len(js.OneOf)).To(Equal(2))
-			Expect(js.OneOf[0].Properties.Function.Const).To(Equal("create_event"))
+			fnName := js.OneOf[0].Properties["function"].(FunctionName)
-			Expect(js.OneOf[0].Properties.Arguments.Properties["event_name"].(map[string]interface{})["type"]).To(Equal("string"))
+			fnArgs := js.OneOf[0].Properties["arguments"].(Argument)
-			Expect(js.OneOf[0].Properties.Arguments.Properties["event_date"].(map[string]interface{})["type"]).To(Equal("string"))
+			Expect(fnName.Const).To(Equal("create_event"))
-			Expect(js.OneOf[1].Properties.Function.Const).To(Equal("search"))
+			Expect(fnArgs.Properties["event_name"].(map[string]interface{})["type"]).To(Equal("string"))
-			Expect(js.OneOf[1].Properties.Arguments.Properties["query"].(map[string]interface{})["type"]).To(Equal("string"))
+			Expect(fnArgs.Properties["event_date"].(map[string]interface{})["type"]).To(Equal("string"))
 
-			jsN := functions.ToJSONNameStructure()
+			fnName = js.OneOf[1].Properties["function"].(FunctionName)
+			fnArgs = js.OneOf[1].Properties["arguments"].(Argument)
+			Expect(fnName.Const).To(Equal("search"))
+			Expect(fnArgs.Properties["query"].(map[string]interface{})["type"]).To(Equal("string"))
+
+			// Test with custom keys
+			jsN := functions.ToJSONStructure("name", "arguments")
 			Expect(len(jsN.OneOf)).To(Equal(2))
-			Expect(jsN.OneOf[0].Properties.Function.Const).To(Equal("create_event"))
+
-			Expect(jsN.OneOf[0].Properties.Arguments.Properties["event_name"].(map[string]interface{})["type"]).To(Equal("string"))
+			fnName = jsN.OneOf[0].Properties["name"].(FunctionName)
-			Expect(jsN.OneOf[0].Properties.Arguments.Properties["event_date"].(map[string]interface{})["type"]).To(Equal("string"))
+			fnArgs = jsN.OneOf[0].Properties["arguments"].(Argument)
-			Expect(jsN.OneOf[1].Properties.Function.Const).To(Equal("search"))
+
-			Expect(jsN.OneOf[1].Properties.Arguments.Properties["query"].(map[string]interface{})["type"]).To(Equal("string"))
+			Expect(fnName.Const).To(Equal("create_event"))
+			Expect(fnArgs.Properties["event_name"].(map[string]interface{})["type"]).To(Equal("string"))
+			Expect(fnArgs.Properties["event_date"].(map[string]interface{})["type"]).To(Equal("string"))
+
+			fnName = jsN.OneOf[1].Properties["name"].(FunctionName)
+			fnArgs = jsN.OneOf[1].Properties["arguments"].(Argument)
+
+			Expect(fnName.Const).To(Equal("search"))
+			Expect(fnArgs.Properties["query"].(map[string]interface{})["type"]).To(Equal("string"))
 		})
 	})
 	Context("Select()", func() {

pkg/functions/grammar_json_schema.go

@@ -331,6 +331,7 @@ func (sc *JSONSchemaConverter) resolveReference(ref string, rootSchema map[strin
 
 	return def
 }
+
 func (sc *JSONSchemaConverter) Grammar(schema map[string]interface{}, options ...func(*GrammarOption)) string {
 	sc.addRule("freestring", PRIMITIVE_RULES["freestring"])
 	sc.visit(schema, "", schema)
@@ -352,52 +353,23 @@ type FunctionName struct {
 	Const string `json:"const"`
 }
 
-type FunctionProperties struct {
-	Function  FunctionName `json:"function"`
-	Arguments Argument     `json:"arguments"`
-}
-
-type NameProperties struct {
-	Function  FunctionName `json:"name"`
-	Arguments Argument     `json:"arguments"`
-}
-
 type Argument struct {
 	Type       string                 `json:"type"`
 	Properties map[string]interface{} `json:"properties"`
 }
 
-type ItemName struct {
+type Item struct {
-	Type       string         `json:"type"`
+	Type       string                 `json:"type"`
-	Properties NameProperties `json:"properties"`
+	Properties map[string]interface{} `json:"properties"`
 }
 
-type ItemFunction struct {
+type JSONFunctionStructure struct {
-	Type       string             `json:"type"`
+	OneOf []Item                 `json:"oneOf,omitempty"`
-	Properties FunctionProperties `json:"properties"`
+	AnyOf []Item                 `json:"anyOf,omitempty"`
-}
-
-type JSONFunctionStructureName struct {
-	OneOf []ItemName             `json:"oneOf,omitempty"`
-	AnyOf []ItemName             `json:"anyOf,omitempty"`
 	Defs  map[string]interface{} `json:"$defs,omitempty"`
 }
 
-func (j JSONFunctionStructureName) Grammar(options ...func(*GrammarOption)) string {
+func (j JSONFunctionStructure) Grammar(options ...func(*GrammarOption)) string {
-	grammarOpts := &GrammarOption{}
-	grammarOpts.Apply(options...)
-
-	dat, _ := json.Marshal(j)
-	return NewJSONSchemaConverter(grammarOpts.PropOrder).GrammarFromBytes(dat, options...)
-}
-
-type JSONFunctionStructureFunction struct {
-	OneOf []ItemFunction         `json:"oneOf,omitempty"`
-	AnyOf []ItemFunction         `json:"anyOf,omitempty"`
-	Defs  map[string]interface{} `json:"$defs,omitempty"`
-}
-
-func (j JSONFunctionStructureFunction) Grammar(options ...func(*GrammarOption)) string {
 	grammarOpts := &GrammarOption{}
 	grammarOpts.Apply(options...)
 

pkg/functions/grammar_json_schema_test.go

@@ -9,69 +9,65 @@ import (
 	. "github.com/onsi/gomega"
 )
 
-var testFunctions = []ItemFunction{
+func createFunction(field1 string, field2 string, name string, properties map[string]interface{}) map[string]interface{} {
+	property := map[string]interface{}{}
+	property[field1] = FunctionName{Const: name}
+	property[field2] = Argument{
+		Type:       "object",
+		Properties: properties,
+	}
+	return property
+}
+
+var testFunctions = []Item{
 	{
 		Type: "object",
-		Properties: FunctionProperties{
+		Properties: createFunction(
-			Function: FunctionName{
+			"function",
-				Const: "create_event",
+			"arguments",
+			"create_event",
+			map[string]interface{}{
+				"title": map[string]string{"type": "string"},
+				"date":  map[string]string{"type": "string"},
+				"time":  map[string]string{"type": "string"},
 			},
-			Arguments: Argument{ // this is OpenAI's parameter
+		),
-				Type: "object",
-				Properties: map[string]interface{}{
-					"title": map[string]string{"type": "string"},
-					"date":  map[string]string{"type": "string"},
-					"time":  map[string]string{"type": "string"},
-				},
-			},
-		},
 	},
 	{
 		Type: "object",
-		Properties: FunctionProperties{
+		Properties: createFunction(
-			Function: FunctionName{
+			"function",
-				Const: "search",
+			"arguments",
-			},
+			"search",
-			Arguments: Argument{
+			map[string]interface{}{
-				Type: "object",
+				"query": map[string]string{"type": "string"},
-				Properties: map[string]interface{}{
+			}),
-					"query": map[string]string{"type": "string"},
-				},
-			},
-		},
 	},
 }
 
-var testFunctionsName = []ItemName{
+var testFunctionsName = []Item{
 	{
 		Type: "object",
-		Properties: NameProperties{
+		Properties: createFunction(
-			Function: FunctionName{
+			"name",
-				Const: "create_event",
+			"arguments",
+			"create_event",
+			map[string]interface{}{
+				"title": map[string]string{"type": "string"},
+				"date":  map[string]string{"type": "string"},
+				"time":  map[string]string{"type": "string"},
 			},
-			Arguments: Argument{ // this is OpenAI's parameter
+		),
-				Type: "object",
-				Properties: map[string]interface{}{
-					"title": map[string]string{"type": "string"},
-					"date":  map[string]string{"type": "string"},
-					"time":  map[string]string{"type": "string"},
-				},
-			},
-		},
 	},
 	{
 		Type: "object",
-		Properties: NameProperties{
+		Properties: createFunction(
-			Function: FunctionName{
+			"name",
-				Const: "search",
+			"arguments",
-			},
+			"search",
-			Arguments: Argument{
+			map[string]interface{}{
-				Type: "object",
+				"query": map[string]string{"type": "string"},
-				Properties: map[string]interface{}{
+			}),
-					"query": map[string]string{"type": "string"},
-				},
-			},
-		},
 	},
 }
 
@@ -270,7 +266,7 @@ var _ = Describe("JSON schema grammar tests", func() {
 		})
 		It("generates a valid grammar from JSON Objects", func() {
 
-			structuredGrammar := JSONFunctionStructureFunction{
+			structuredGrammar := JSONFunctionStructure{
 				OneOf: testFunctions}
 
 			grammar := structuredGrammar.Grammar()
@@ -284,7 +280,7 @@ var _ = Describe("JSON schema grammar tests", func() {
 		})
 
 		It("generates a valid grammar from JSON Objects for multiple function return", func() {
-			structuredGrammar := JSONFunctionStructureFunction{
+			structuredGrammar := JSONFunctionStructure{
 				OneOf: testFunctions}
 
 			grammar := structuredGrammar.Grammar(functions.EnableMaybeArray)
@@ -302,7 +298,7 @@ var _ = Describe("JSON schema grammar tests", func() {
 		})
 
 		It("generates a valid grammar from JSON Objects for multiple function return", func() {
-			structuredGrammar := JSONFunctionStructureName{
+			structuredGrammar := JSONFunctionStructure{
 				OneOf: testFunctionsName}
 
 			grammar := structuredGrammar.Grammar(functions.EnableMaybeArray)
@@ -320,7 +316,7 @@ var _ = Describe("JSON schema grammar tests", func() {
 		})
 
 		It("generates a valid grammar from JSON Objects for multiple function return with a suffix and array", func() {
-			structuredGrammar := JSONFunctionStructureName{
+			structuredGrammar := JSONFunctionStructure{
 				OneOf: testFunctionsName}
 
 			grammar := structuredGrammar.Grammar(
@@ -340,7 +336,7 @@ var _ = Describe("JSON schema grammar tests", func() {
 			Expect(len(results)).To(Equal(len(strings.Split(grammar, "\n"))), grammar)
 		})
 		It("generates a valid grammar from JSON Objects with a suffix", func() {
-			structuredGrammar := JSONFunctionStructureName{
+			structuredGrammar := JSONFunctionStructure{
 				OneOf: testFunctionsName}
 
 			grammar := structuredGrammar.Grammar(functions.SetPrefix("suffix"))
@@ -357,7 +353,7 @@ var _ = Describe("JSON schema grammar tests", func() {
 			Expect(len(results)).To(Equal(len(strings.Split(grammar, "\n"))), grammar)
 		})
 		It("generates a valid grammar from JSON Objects with a suffix and could return string", func() {
-			structuredGrammar := JSONFunctionStructureName{
+			structuredGrammar := JSONFunctionStructure{
 				OneOf: testFunctionsName}
 
 			grammar := structuredGrammar.Grammar(functions.SetPrefix("suffix"), functions.EnableMaybeString)
@@ -374,7 +370,7 @@ var _ = Describe("JSON schema grammar tests", func() {
 			Expect(len(results)).To(Equal(len(strings.Split(grammar, "\n"))), grammar)
 		})
 		It("generates a valid grammar from JSON Objects with a suffix that could return text or an array of tools", func() {
-			structuredGrammar := JSONFunctionStructureName{
+			structuredGrammar := JSONFunctionStructure{
 				OneOf: testFunctionsName}
 
 			grammar := structuredGrammar.Grammar(functions.SetPrefix("suffix"), functions.EnableMaybeString, functions.EnableMaybeArray)
@@ -393,7 +389,7 @@ var _ = Describe("JSON schema grammar tests", func() {
 		})
 
 		It("generates a valid grammar from JSON Objects without a suffix that could return text or an array of tools or just string", func() {
-			structuredGrammar := JSONFunctionStructureName{
+			structuredGrammar := JSONFunctionStructure{
 				OneOf: testFunctionsName}
 
 			grammar := structuredGrammar.Grammar(functions.EnableMaybeString, functions.EnableMaybeArray)
@@ -411,7 +407,7 @@ var _ = Describe("JSON schema grammar tests", func() {
 		})
 
 		It("generates a valid grammar from JSON Objects without a suffix that could return text or an array of tools or just string. Disables mixedstring", func() {
-			structuredGrammar := JSONFunctionStructureName{
+			structuredGrammar := JSONFunctionStructure{
 				OneOf: testFunctionsName}
 
 			grammar := structuredGrammar.Grammar(functions.EnableMaybeString, functions.EnableMaybeArray, functions.NoMixedFreeString)
@@ -429,7 +425,7 @@ var _ = Describe("JSON schema grammar tests", func() {
 		})
 
 		It("generates parallel tools without newlines in JSON", func() {
-			structuredGrammar := JSONFunctionStructureName{
+			structuredGrammar := JSONFunctionStructure{
 				OneOf: testFunctionsName}
 			content := `arr  ::=
 "["  (

pkg/functions/parse.go

@@ -2,6 +2,8 @@ package functions
 
 import (
 	"encoding/json"
+	"errors"
+	"io"
 	"regexp"
 	"strings"
 
@@ -76,7 +78,8 @@ type FunctionsConfig struct {
 	// FunctionName enable the LLM to return { "name": "function_name", "arguments": { "arg1": "value1", "arg2": "value2" } }
 	// instead of { "function": "function_name", "arguments": { "arg1": "value1", "arg2": "value2" } }.
 	// This might be useful for certain models trained with the function name as the first token.
-	FunctionName bool `yaml:"return_name_in_function_response"`
+	FunctionNameKey      string `yaml:"function_name_key"`
+	FunctionArgumentsKey string `yaml:"function_arguments_key"`
 }
 
 type ReplaceResult struct {
@@ -145,6 +148,47 @@ func ParseTextContent(llmresult string, functionConfig FunctionsConfig) string {
 	return ""
 }
 
+// ParseJSON is a function that parses a JSON string that might contain multiple JSON objects
+// and syntax errors in between by shifting the offset
+// This for e.g. allow to parse
+// { "foo": "bar" } invalid { "baz": "qux" }
+// into
+// [ { "foo": "bar" }, { "baz": "qux" } ]
+// Credits to Michael Yang (https://github.com/mxyng) for the original implementation
+// This is a slighly reworked version, improved for readability and error handling
+func ParseJSON(s string) ([]map[string]any, error) {
+	var objs []map[string]any
+	offset := 0
+
+	for offset < len(s) {
+		var obj map[string]any
+		decoder := json.NewDecoder(strings.NewReader(s[offset:]))
+
+		err := decoder.Decode(&obj)
+		switch {
+		case errors.Is(err, io.EOF):
+			return objs, nil
+		case err == nil:
+			offset += int(decoder.InputOffset())
+			objs = append(objs, obj)
+		default: // handle the error type
+			var syntaxErr *json.SyntaxError
+			var unmarshalTypeErr *json.UnmarshalTypeError
+
+			switch {
+			case errors.As(err, &syntaxErr):
+				offset += int(syntaxErr.Offset)
+			case errors.As(err, &unmarshalTypeErr):
+				offset += int(unmarshalTypeErr.Offset)
+			default:
+				return objs, err
+			}
+		}
+	}
+
+	return objs, nil
+}
+
 func ParseFunctionCall(llmresult string, functionConfig FunctionsConfig) []FuncCallResults {
 
 	log.Debug().Msgf("LLM result: %s", llmresult)
@@ -157,9 +201,13 @@ func ParseFunctionCall(llmresult string, functionConfig FunctionsConfig) []FuncC
 	}
 	log.Debug().Msgf("LLM result(function cleanup): %s", llmresult)
 
-	functionNameKey := "function"
+	functionNameKey := defaultFunctionNameKey
-	if functionConfig.FunctionName {
+	functionArgumentsKey := defaultFunctionArgumentsKey
-		functionNameKey = "name"
+	if functionConfig.FunctionNameKey != "" {
+		functionNameKey = functionConfig.FunctionNameKey
+	}
+	if functionConfig.FunctionArgumentsKey != "" {
+		functionArgumentsKey = functionConfig.FunctionArgumentsKey
 	}
 
 	results := []FuncCallResults{}
@@ -170,19 +218,13 @@ func ParseFunctionCall(llmresult string, functionConfig FunctionsConfig) []FuncC
 		result = make([]FuncCallResults, 0)
 
 		for _, s := range results {
-			var ss []map[string]interface{}
+			var ss []map[string]any
 
 			s = utils.EscapeNewLines(s)
-			err := json.Unmarshal([]byte(s), &ss)
+			ss, err := ParseJSON(s)
+			//err := json.Unmarshal([]byte(s), &ss)
 			if err != nil {
-				// If the LLM result is a single object, try unmarshaling it into a single map
+				log.Debug().Err(err).Str("escapedLLMResult", s).Msg("unable to unmarshal llm result in a single object or an array of JSON objects")
-				var singleObj map[string]interface{}
-				err = json.Unmarshal([]byte(s), &singleObj)
-				if err != nil {
-					log.Debug().Err(err).Str("escapedLLMResult", s).Msg("unable to unmarshal llm result in a single object or an array of JSON objects")
-				} else {
-					ss = []map[string]interface{}{singleObj}
-				}
 			}
 
 			log.Debug().Msgf("Function return: %s %+v", s, ss)
@@ -195,7 +237,7 @@ func ParseFunctionCall(llmresult string, functionConfig FunctionsConfig) []FuncC
 					//return result, fmt.Errorf("unable to find function name in result")
 				}
 				// Similarly, while here arguments is a map[string]interface{}, OpenAI actually want a stringified object
-				args, ok := s["arguments"] // arguments needs to be a string, but we return an object from the grammar result (TODO: fix)
+				args, ok := s[functionArgumentsKey] // arguments needs to be a string, but we return an object from the grammar result (TODO: fix)
 				if !ok {
 					continue
 					//return result, fmt.Errorf("unable to find arguments in result")
@@ -253,7 +295,7 @@ func ParseFunctionCall(llmresult string, functionConfig FunctionsConfig) []FuncC
 				if functionName == "" {
 					return results
 				}
-				results = append(results, FuncCallResults{Name: result[functionNameKey], Arguments: result["arguments"]})
+				results = append(results, FuncCallResults{Name: result[functionNameKey], Arguments: result[functionArgumentsKey]})
 			}
 		}
 	} else {

pkg/functions/parse_test.go

@@ -16,7 +16,7 @@ var _ = Describe("LocalAI function parse tests", func() {
 
 	Context("when using grammars and single result expected", func() {
 		It("should parse the function name and arguments correctly", func() {
-			input := `{"function": "add", "arguments": {"x": 5, "y": 3}}`
+			input := `{"name": "add", "arguments": {"x": 5, "y": 3}}`
 
 			results := ParseFunctionCall(input, functionConfig)
 			Expect(results).To(HaveLen(1))
@@ -28,13 +28,22 @@ var _ = Describe("LocalAI function parse tests", func() {
 	Context("when not using grammars and regex is needed", func() {
 		It("should extract function name and arguments from the regex", func() {
 			input := `add({"x":5,"y":3})`
-			functionConfig.ResponseRegex = []string{`(?P<function>\w+)\s*\((?P<arguments>.*)\)`}
+			functionConfig.ResponseRegex = []string{`(?P<name>\w+)\s*\((?P<arguments>.*)\)`}
 
 			results := ParseFunctionCall(input, functionConfig)
 			Expect(results).To(HaveLen(1))
 			Expect(results[0].Name).To(Equal("add"))
 			Expect(results[0].Arguments).To(Equal(`{"x":5,"y":3}`))
 		})
+		It("should extract function name and arguments from the regex", func() {
+			input := `add({"x":5,"y":3})`
+			functionConfig.ResponseRegex = []string{`(?P<function>\w+)\s*\((?P<arguments>.*)\)`}
+			functionConfig.FunctionNameKey = "function"
+			results := ParseFunctionCall(input, functionConfig)
+			Expect(results).To(HaveLen(1))
+			Expect(results[0].Name).To(Equal("add"))
+			Expect(results[0].Arguments).To(Equal(`{"x":5,"y":3}`))
+		})
 	})
 
 	Context("when having invalid input", func() {
@@ -53,7 +62,7 @@ var _ = Describe("LocalAI function parse tests", func() {
 
 	Context("when parallel calls are enabled", func() {
 		It("should handle multiple function calls", func() {
-			input := `[{"function": "add", "arguments": {"x": 5, "y": 3}}, {"function": "subtract", "arguments": {"x": 10, "y": 7}}]`
+			input := `[{"name": "add", "arguments": {"x": 5, "y": 3}}, {"name": "subtract", "arguments": {"x": 10, "y": 7}}]`
 
 			results := ParseFunctionCall(input, functionConfig)
 			Expect(results).To(HaveLen(2))
@@ -66,8 +75,8 @@ var _ = Describe("LocalAI function parse tests", func() {
 
 	Context("without grammars and without regex", func() {
 		It("should parse the function name and arguments correctly with the name key", func() {
-			input := `{"name": "add", "arguments": {"x": 5, "y": 3}}`
+			input := `{"function": "add", "arguments": {"x": 5, "y": 3}}`
-			functionConfig.FunctionName = true
+			functionConfig.FunctionNameKey = "function"
 
 			results := ParseFunctionCall(input, functionConfig)
 			Expect(results).To(HaveLen(1))
@@ -76,7 +85,7 @@ var _ = Describe("LocalAI function parse tests", func() {
 		})
 
 		It("should parse the function name and arguments correctly with the function key", func() {
-			input := `{"function": "add", "arguments": {"x": 5, "y": 3}}`
+			input := `{"name": "add", "arguments": {"x": 5, "y": 3}}`
 
 			results := ParseFunctionCall(input, functionConfig)
 			Expect(results).To(HaveLen(1))
@@ -87,7 +96,7 @@ var _ = Describe("LocalAI function parse tests", func() {
 		It("should parse the result by matching the JSONRegexMatch", func() {
 			input := `
 <tool_call>
-{"function": "add", "arguments": {"x": 5, "y": 3}}
+{"name": "add", "arguments": {"x": 5, "y": 3}}
 </tool_call>`
 
 			functionConfig.JSONRegexMatch = []string{`(?s)<tool_call>(.*?)</tool_call>`}
@@ -100,7 +109,7 @@ var _ = Describe("LocalAI function parse tests", func() {
 
 		It("should parse the result by matching the JSONRegexMatch", func() {
 			input := `
-{"function": "add", "arguments": {"x": 5, "y": 3}}
+{"name": "add", "arguments": {"x": 5, "y": 3}}
 </tool_call>`
 
 			functionConfig.JSONRegexMatch = []string{`(?s)(.*?)</tool_call>`}
@@ -110,13 +119,21 @@ var _ = Describe("LocalAI function parse tests", func() {
 			Expect(results[0].Name).To(Equal("add"))
 			Expect(results[0].Arguments).To(Equal(`{"x":5,"y":3}`))
 		})
+
+		It("should parse the result even with invalid JSON", func() {
+			input := `{"name": "add", "arguments": {"x": 5, "y": 3}} invalid {"name": "add", "arguments": {"x": 5, "y": 3}}`
+			results := ParseFunctionCall(input, functionConfig)
+			Expect(results).To(HaveLen(2))
+			Expect(results[0].Name).To(Equal("add"))
+			Expect(results[0].Arguments).To(Equal(`{"x":5,"y":3}`))
+		})
 	})
 
 	Context("when using ReplaceResults to clean up input", func() {
 		It("should replace text before and after JSON blob", func() {
 			input := `
 Some text before the JSON
-{"function": "add", "arguments": {"x": 5, "y": 3}}
+{"name": "add", "arguments": {"x": 5, "y": 3}}
 Some text after the JSON
 `
 
@@ -134,7 +151,7 @@ Some text after the JSON
 		It("should replace text before and after array JSON blob", func() {
 			input := `
 Some text before the JSON
-[{"function": "add", "arguments": {"x": 5, "y": 3}}, {"function": "subtract", "arguments": {"x": 10, "y": 7}}]
+[{"name": "add", "arguments": {"x": 5, "y": 3}}, {"name": "subtract", "arguments": {"x": 10, "y": 7}}]
 Some text after the JSON
 `
 			functionConfig.ReplaceFunctionResults = []ReplaceResult{
@@ -153,7 +170,7 @@ Some text after the JSON
 		It("should convert single-quoted key-value pairs to double-quoted and escape double quotes within values", func() {
 			input := `
 Some text before the JSON
-{'function': '"add"', 'arguments': {'x': 5, 'z': '"v"', 'y': 'v"value"'}}
+{'name': '"add"', 'arguments': {'x': 5, 'z': '"v"', 'y': 'v"value"'}}
 Some text after the JSON
 `
 			functionConfig.JSONRegexMatch = []string{`(?s)<tool_call>(.*?)</tool_call>`}
@@ -186,7 +203,7 @@ Some text after the JSON
 		It("should convert single-quoted key-value pairs to double-quoted and escape double quotes within values", func() {
 			input := `
 Some text before the JSON
-<tool_call>{'function': '"add"', 'arguments': {'x': 5, 'z': '"v"', 'y': 'v"value"'}}</tool_call>
+<tool_call>{'name': '"add"', 'arguments': {'x': 5, 'z': '"v"', 'y': 'v"value"'}}</tool_call>
 Some text after the JSON
 `
 			functionConfig.JSONRegexMatch = []string{`(?s)<tool_call>(.*?)</tool_call>`}
@@ -219,8 +236,8 @@ Some text after the JSON
 		It("should detect multiple functions call where the JSONRegexMatch is repeated", func() {
 			input := `
 Some text before the JSON
-<tool_call>{"function": "add", "arguments": {"x": 5, "y": 3}}</tool_call>
+<tool_call>{"name": "add", "arguments": {"x": 5, "y": 3}}</tool_call>
-<tool_call>{"function": "subtract", "arguments": {"x": 10, "y": 7}}</tool_call>
+<tool_call>{"name": "subtract", "arguments": {"x": 10, "y": 7}}</tool_call>
 Some text after the JSON
 `
 			functionConfig.JSONRegexMatch = []string{`(?s)<tool_call>(.*?)</tool_call>`}
@@ -240,7 +257,7 @@ Some text after the JSON
 <sketchpad>
 roses are red
 </sketchpad>
-		<tool_call>{"function": "subtract", "arguments": {"x": 10, "y": 7}}</tool_call>
+		<tool_call>{"name": "subtract", "arguments": {"x": 10, "y": 7}}</tool_call>
 		Some text after the JSON
 		`
 			functionConfig.CaptureLLMResult = []string{`(?s)<sketchpad>(.*?)</sketchpad>`}
@@ -251,7 +268,7 @@ roses are red
 		It("Defaults to empty if doesn't catch any", func() {
 			input := `
 		Some text before the JSON
-		<tool_call>{"function": "subtract", "arguments": {"x": 10, "y": 7}}</tool_call>
+		<tool_call>{"name": "subtract", "arguments": {"x": 10, "y": 7}}</tool_call>
 		Some text after the JSON
 		`
 			functionConfig.CaptureLLMResult = []string{`(?s)<sketchpad>(.*?)</sketchpad>`}
@@ -259,4 +276,74 @@ roses are red
 			Expect(results).To(Equal(""))
 		})
 	})
+	Context("ParseJSON - when given valid JSON strings", func() {
+		It("should parse multiple JSON objects", func() {
+			input := `{"key1": "value1"} {"key2": "value2"}`
+			expected := []map[string]any{
+				{"key1": "value1"},
+				{"key2": "value2"},
+			}
+			result, err := ParseJSON(input)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(result).To(Equal(expected))
+		})
+
+		It("should parse a single JSON object with various types", func() {
+			input := `{"key1": "value1", "key2": 2}`
+			expected := []map[string]any{
+				{"key1": "value1", "key2": float64(2)},
+			}
+			result, err := ParseJSON(input)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(result).To(Equal(expected))
+		})
+		It("should handle JSON without syntax errors gracefully", func() {
+			input := `{"key1": "value1"}`
+			expected := []map[string]any{
+				{"key1": "value1"},
+			}
+			result, err := ParseJSON(input)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(result).To(Equal(expected))
+		})
+		It("should handle JSON without syntax errors gracefully", func() {
+			input := `[{"key1": "value1"}]`
+			expected := []map[string]any{
+				{"key1": "value1"},
+			}
+			result, err := ParseJSON(input)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(result).To(Equal(expected))
+		})
+	})
+
+	Context("ParseJSON - when given invalid JSON strings", func() {
+		It("should return an error for completely invalid JSON", func() {
+			input := `invalid json`
+			result, err := ParseJSON(input)
+			Expect(err).To(HaveOccurred())
+			Expect(result).To(BeNil())
+		})
+
+		It("should skip invalid JSON parts and parse valid parts", func() {
+			input := `{"key1": "value1"} invalid {"key2": "value2"}`
+			expected := []map[string]any{
+				{"key1": "value1"},
+				{"key2": "value2"},
+			}
+			result, err := ParseJSON(input)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(result).To(Equal(expected))
+		})
+
+		PIt("should handle JSON with syntax errors gracefully", func() {
+			input := `{"key1": "value1", "key2": }`
+			expected := []map[string]any{
+				{"key1": "value1"},
+			}
+			result, err := ParseJSON(input)
+			Expect(err).NotTo(HaveOccurred())
+			Expect(result).To(Equal(expected))
+		})
+	})
 })

pkg/model/loader.go

@@ -98,6 +98,9 @@ var knownModelsNameSuffixToSkip []string = []string{
 	".yaml",
 	".yml",
 	".json",
+	".txt",
+	".md",
+	".MD",
 	".DS_Store",
 	".",
 	".partial",