Merge branch 'master' into disable_grammar_by_default

Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

.github/workflows/comment-pr.yaml

@@ -8,8 +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:
+        ref: "${{ github.event.pull_request.merge_commit_sha }}"
         fetch-depth: 0 # needed to checkout all branches for this Action to work
     - uses: mudler/localai-github-action@v1
       with:
@@ -21,6 +23,7 @@ jobs:
             json_diff_file_output: diff.json
             raw_diff_file_output: diff.txt
             file_output_only: "true"
+            base_branch: ${{ github.event.pull_request.base.sha }}
     - name: Show diff
       env:
         DIFF: ${{ steps.git-diff-action.outputs.raw-diff-path }}

Makefile

@@ -8,7 +8,7 @@ DETECT_LIBS?=true
 # llama.cpp versions
 GOLLAMA_REPO?=https://github.com/go-skynet/go-llama.cpp
 GOLLAMA_VERSION?=2b57a8ae43e4699d3dc5d1496a1ccd42922993be
-CPPLLAMA_VERSION?=b3283448ce9a5098226afe1d8648ccc578511fe4
+CPPLLAMA_VERSION?=705b7ecf60e667ced57c15d67aa86865e3cc7aa7
 
 # gpt4all version
 GPT4ALL_REPO?=https://github.com/nomic-ai/gpt4all

core/cli/federated.go

@@ -2,129 +2,20 @@ package cli
 
 import (
 	"context"
-	"errors"
-	"fmt"
-	"io"
-	"net"
-	"time"
-
-	"math/rand/v2"
 
 	cliContext "github.com/mudler/LocalAI/core/cli/context"
 	"github.com/mudler/LocalAI/core/p2p"
-	"github.com/mudler/edgevpn/pkg/node"
-	"github.com/mudler/edgevpn/pkg/protocol"
-	"github.com/mudler/edgevpn/pkg/types"
-	"github.com/rs/zerolog/log"
 )
 
 type FederatedCLI struct {
 	Address        string `env:"LOCALAI_ADDRESS,ADDRESS" default:":8080" help:"Bind address for the API server" group:"api"`
 	Peer2PeerToken string `env:"LOCALAI_P2P_TOKEN,P2P_TOKEN,TOKEN" name:"p2ptoken" help:"Token for P2P mode (optional)" group:"p2p"`
+	LoadBalanced   bool   `env:"LOCALAI_LOAD_BALANCED,LOAD_BALANCED" default:"false" help:"Enable load balancing" group:"p2p"`
 }
 
 func (f *FederatedCLI) Run(ctx *cliContext.Context) error {
 
-	n, err := p2p.NewNode(f.Peer2PeerToken)
-	if err != nil {
-		return fmt.Errorf("creating a new node: %w", err)
-	}
-	err = n.Start(context.Background())
-	if err != nil {
-		return fmt.Errorf("creating a new node: %w", err)
-	}
+	fs := p2p.NewFederatedServer(f.Address, p2p.FederatedID, f.Peer2PeerToken, f.LoadBalanced)
 
-	if err := p2p.ServiceDiscoverer(context.Background(), n, f.Peer2PeerToken, p2p.FederatedID, nil); err != nil {
-		return err
-	}
-
-	return Proxy(context.Background(), n, f.Address, p2p.FederatedID)
-}
-
-func Proxy(ctx context.Context, node *node.Node, listenAddr, service string) error {
-
-	log.Info().Msgf("Allocating service '%s' on: %s", service, listenAddr)
-	// Open local port for listening
-	l, err := net.Listen("tcp", listenAddr)
-	if err != nil {
-		log.Error().Err(err).Msg("Error listening")
-		return err
-	}
-	//	ll.Info("Binding local port on", srcaddr)
-
-	ledger, _ := node.Ledger()
-
-	// Announce ourselves so nodes accepts our connection
-	ledger.Announce(
-		ctx,
-		10*time.Second,
-		func() {
-			// Retrieve current ID for ip in the blockchain
-			//_, found := ledger.GetKey(protocol.UsersLedgerKey, node.Host().ID().String())
-			// If mismatch, update the blockchain
-			//if !found {
-			updatedMap := map[string]interface{}{}
-			updatedMap[node.Host().ID().String()] = &types.User{
-				PeerID:    node.Host().ID().String(),
-				Timestamp: time.Now().String(),
-			}
-			ledger.Add(protocol.UsersLedgerKey, updatedMap)
-			//	}
-		},
-	)
-
-	defer l.Close()
-	for {
-		select {
-		case <-ctx.Done():
-			return errors.New("context canceled")
-		default:
-			log.Debug().Msg("New for connection")
-			// Listen for an incoming connection.
-			conn, err := l.Accept()
-			if err != nil {
-				fmt.Println("Error accepting: ", err.Error())
-				continue
-			}
-
-			// Handle connections in a new goroutine, forwarding to the p2p service
-			go func() {
-				var tunnelAddresses []string
-				for _, v := range p2p.GetAvailableNodes(p2p.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()
-				//	ll.Infof("(service %s) Done handling %s", serviceID, l.Addr().String())
-			}()
-		}
-	}
-
-}
-
-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)
+	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

@@ -201,7 +201,7 @@ func ChatEndpoint(cl *config.BackendConfigLoader, ml *model.ModelLoader, startup
 		}
 
 		switch {
-		case !config.FunctionsConfig.GrammarConfig.NoGrammar && shouldUseFn:
+		case config.FunctionsConfig.GrammarConfig.EnableGrammar && shouldUseFn:
 			noActionGrammar := functions.Function{
 				Name:        noActionName,
 				Description: noActionDescription,
@@ -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"
-			if config.FunctionsConfig.FunctionName {
-				jsStruct := funcs.ToJSONNameStructure()
-				config.Grammar = jsStruct.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
-			} else {
-				jsStruct := funcs.ToJSONFunctionStructure()
-				config.Grammar = jsStruct.Grammar(config.FunctionsConfig.GrammarConfig.Options()...)
-			}
+			jsStruct := funcs.ToJSONStructure(config.FunctionsConfig.FunctionNameKey, config.FunctionsConfig.FunctionNameKey)
+			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

@@ -0,0 +1,47 @@
+package p2p
+
+const FederatedID = "federated"
+
+type FederatedServer struct {
+	listenAddr, service, p2ptoken string
+	requestTable                  map[string]int
+	loadBalanced                  bool
+}
+
+func NewFederatedServer(listenAddr, service, p2pToken string, loadBalanced bool) *FederatedServer {
+	return &FederatedServer{
+		listenAddr:   listenAddr,
+		service:      service,
+		p2ptoken:     p2pToken,
+		requestTable: map[string]int{},
+		loadBalanced: loadBalanced,
+	}
+}
+
+func (fs *FederatedServer) SelectLeastUsedServer() string {
+	// cycle over requestTable and find the entry with the lower number
+	// if there are multiple entries with the same number, select one randomly
+	// if there are no entries, return an empty string
+	var min int
+	var minKey string
+	for k, v := range fs.requestTable {
+		if min == 0 || v < min {
+			min = v
+			minKey = k
+		}
+	}
+	return minKey
+}
+
+func (fs *FederatedServer) RecordRequest(nodeID string) {
+	// increment the counter for the nodeID in the requestTable
+	fs.requestTable[nodeID]++
+}
+
+func (fs *FederatedServer) EnsureRecordExist(nodeID string) {
+	// if the nodeID is not in the requestTable, add it with a counter of 0
+	_, ok := fs.requestTable[nodeID]
+	if !ok {
+		fs.requestTable[nodeID] = 0
+	}
+}

core/p2p/federated_server.go

@@ -0,0 +1,140 @@
+//go:build p2p
+// +build p2p
+
+package p2p
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net"
+	"time"
+
+	"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 (f *FederatedServer) Start(ctx context.Context) error {
+
+	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("creating a new node: %w", err)
+	}
+
+	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 f.proxy(ctx, n)
+}
+
+func (fs *FederatedServer) proxy(ctx context.Context, node *node.Node) error {
+
+	log.Info().Msgf("Allocating service '%s' on: %s", fs.service, fs.listenAddr)
+	// Open local port for listening
+	l, err := net.Listen("tcp", fs.listenAddr)
+	if err != nil {
+		log.Error().Err(err).Msg("Error listening")
+		return err
+	}
+	//	ll.Info("Binding local port on", srcaddr)
+
+	ledger, _ := node.Ledger()
+
+	// Announce ourselves so nodes accepts our connection
+	ledger.Announce(
+		ctx,
+		10*time.Second,
+		func() {
+			// Retrieve current ID for ip in the blockchain
+			//_, found := ledger.GetKey(protocol.UsersLedgerKey, node.Host().ID().String())
+			// If mismatch, update the blockchain
+			//if !found {
+			updatedMap := map[string]interface{}{}
+			updatedMap[node.Host().ID().String()] = &types.User{
+				PeerID:    node.Host().ID().String(),
+				Timestamp: time.Now().String(),
+			}
+			ledger.Add(protocol.UsersLedgerKey, updatedMap)
+			//	}
+		},
+	)
+
+	defer l.Close()
+	for {
+		select {
+		case <-ctx.Done():
+			return errors.New("context canceled")
+		default:
+			log.Debug().Msg("New for connection")
+			// Listen for an incoming connection.
+			conn, err := l.Accept()
+			if err != nil {
+				fmt.Println("Error accepting: ", err.Error())
+				continue
+			}
+
+			// Handle connections in a new goroutine, forwarding to the p2p service
+			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
+				}
+
+				tunnelAddr := ""
+
+				if fs.loadBalanced {
+					for _, t := range tunnelAddresses {
+						fs.EnsureRecordExist(t)
+					}
+
+					tunnelAddr = fs.SelectLeastUsedServer()
+					log.Debug().Msgf("Selected tunnel %s", tunnelAddr)
+					if tunnelAddr == "" {
+						tunnelAddr = tunnelAddresses[rand.IntN(len(tunnelAddresses))]
+					}
+
+					fs.RecordRequest(tunnelAddr)
+				} else {
+					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())
+			}()
+		}
+	}
+
+}

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

@@ -137,14 +137,9 @@ func allocateLocalService(ctx context.Context, node *node.Node, listenAddr, serv
 
 }
 
-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)
-}
-
 // 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
 	}
@@ -166,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)
 				}
 			}
 		}
@@ -396,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,7 +14,11 @@ func GenerateToken() string {
 	return "not implemented"
 }
 
-func ServiceDiscoverer(ctx context.Context, node *node.Node, token, servicesID string, fn func()) 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(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"`
-	JSONFunctionGrammarObjectName *functions.JSONFunctionStructureName     `json:"grammar_json_name" yaml:"grammar_json_name"`
+	JSONFunctionGrammarObject *functions.JSONFunctionStructure `json:"grammar_json_functions" yaml:"grammar_json_functions"`
 
 	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: {}

gallery/index.yaml

@@ -579,6 +579,21 @@
     - filename: EZO-Common-9B-gemma-2-it.Q4_K_M.gguf
       sha256: 57678b1828673dccb15f76e52b00672c74aa6169421bbb8620b8955955322cfd
       uri: huggingface://QuantFactory/EZO-Common-9B-gemma-2-it-GGUF/EZO-Common-9B-gemma-2-it.Q4_K_M.gguf
+- !!merge <<: *gemma
+  name: "big-tiger-gemma-27b-v1"
+  icon: https://cdn-uploads.huggingface.co/production/uploads/65f2fd1c25b848bd061b5c2e/A97OlLKeT4XOnv4IG1b6m.png
+  urls:
+    - https://huggingface.co/TheDrummer/Big-Tiger-Gemma-27B-v1
+    - https://huggingface.co/TheDrummer/Big-Tiger-Gemma-27B-v1-GGUF
+  description: |
+    Big Tiger Gemma 27B v1 is a Decensored Gemma 27B model with no refusals, except for some rare instances from the 9B model. It does not appear to have any brain damage. The model is available from various sources, including Hugging Face, and comes in different variations such as GGUF, iMatrix, and EXL2.
+  overrides:
+    parameters:
+      model: Big-Tiger-Gemma-27B-v1c-Q4_K_M.gguf
+  files:
+    - filename: Big-Tiger-Gemma-27B-v1c-Q4_K_M.gguf
+      sha256: c5fc5605d36ae280c1c908c9b4bcb12b28abbe2692f317edeb83ab1104657fe5
+      uri: huggingface://TheDrummer/Big-Tiger-Gemma-27B-v1-GGUF/Big-Tiger-Gemma-27B-v1c-Q4_K_M.gguf
 - &llama3
   url: "github:mudler/LocalAI/gallery/llama3-instruct.yaml@master"
   icon: https://cdn-uploads.huggingface.co/production/uploads/642cc1c253e76b4c2286c58e/aJJxKus1wP5N-euvHEUq7.png

gallery/mudler.yaml

@@ -10,7 +10,8 @@ config_file: |-
     - <|end_of_text|>
 
   function:
-    return_name_in_function_response: true
+    grammar:
+      enable: true
 
   template:
     chat: |

pkg/concurrency/concurrency_suite_test.go

@@ -0,0 +1,13 @@
+package concurrency
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestConcurrency(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Concurrency test suite")
+}

pkg/concurrency/jobresult.go

@@ -0,0 +1,69 @@
+package concurrency
+
+import (
+	"context"
+	"sync"
+)
+
+// This is a Read-ONLY structure that contains the result of an arbitrary asynchronous action
+type JobResult[RequestType any, ResultType any] struct {
+	request *RequestType
+	result  *ResultType
+	err     error
+	once    sync.Once
+	done    *chan struct{}
+}
+
+// This structure is returned in a pair with a JobResult and serves as the structure that has access to be updated.
+type WritableJobResult[RequestType any, ResultType any] struct {
+	*JobResult[RequestType, ResultType]
+}
+
+// Wait blocks until the result is ready and then returns the result, or the context expires.
+// Returns *ResultType instead of ResultType since its possible we have only an error and nil for ResultType.
+// Is this correct and idiomatic?
+func (jr *JobResult[RequestType, ResultType]) Wait(ctx context.Context) (*ResultType, error) {
+	if jr.done == nil { // If the channel is blanked out, result is ready.
+		return jr.result, jr.err
+	}
+	select {
+	case <-*jr.done: // Wait for the result to be ready
+		jr.done = nil
+		if jr.err != nil {
+			return nil, jr.err
+		}
+		return jr.result, nil
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
+}
+
+// Accessor function to allow holders of JobResults to access the associated request, without allowing the pointer to be updated.
+func (jr *JobResult[RequestType, ResultType]) Request() *RequestType {
+	return jr.request
+}
+
+// This is the function that actually updates the Result and Error on the JobResult... but it's normally not accessible
+func (jr *JobResult[RequestType, ResultType]) setResult(result ResultType, err error) {
+	jr.once.Do(func() {
+		jr.result = &result
+		jr.err = err
+		close(*jr.done) // Signal that the result is ready - since this is only ran once, jr.done cannot be set to nil yet.
+	})
+}
+
+// Only the WritableJobResult can actually call setResult - prevents accidental corruption
+func (wjr *WritableJobResult[RequestType, ResultType]) SetResult(result ResultType, err error) {
+	wjr.JobResult.setResult(result, err)
+}
+
+// NewJobResult binds a request to a matched pair of JobResult and WritableJobResult
+func NewJobResult[RequestType any, ResultType any](request RequestType) (*JobResult[RequestType, ResultType], *WritableJobResult[RequestType, ResultType]) {
+	done := make(chan struct{})
+	jr := &JobResult[RequestType, ResultType]{
+		once:    sync.Once{},
+		request: &request,
+		done:    &done,
+	}
+	return jr, &WritableJobResult[RequestType, ResultType]{JobResult: jr}
+}

pkg/concurrency/jobresult_test.go

@@ -0,0 +1,80 @@
+package concurrency_test
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	. "github.com/mudler/LocalAI/pkg/concurrency"
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+var _ = Describe("pkg/concurrency unit tests", func() {
+	It("can be used to recieve a result across goroutines", func() {
+		jr, wjr := NewJobResult[string, string]("foo")
+		Expect(jr).ToNot(BeNil())
+		Expect(wjr).ToNot(BeNil())
+
+		go func(wjr *WritableJobResult[string, string]) {
+			time.Sleep(time.Second * 5)
+			wjr.SetResult("bar", nil)
+		}(wjr)
+
+		resPtr, err := jr.Wait(context.Background())
+		Expect(err).To(BeNil())
+		Expect(jr.Request).ToNot(BeNil())
+		Expect(*jr.Request()).To(Equal("foo"))
+		Expect(resPtr).ToNot(BeNil())
+		Expect(*resPtr).To(Equal("bar"))
+
+	})
+
+	It("can be used to recieve an error across goroutines", func() {
+		jr, wjr := NewJobResult[string, string]("foo")
+		Expect(jr).ToNot(BeNil())
+		Expect(wjr).ToNot(BeNil())
+
+		go func(wjr *WritableJobResult[string, string]) {
+			time.Sleep(time.Second * 5)
+			wjr.SetResult("", fmt.Errorf("test"))
+		}(wjr)
+
+		_, err := jr.Wait(context.Background())
+		Expect(jr.Request).ToNot(BeNil())
+		Expect(*jr.Request()).To(Equal("foo"))
+		Expect(err).ToNot(BeNil())
+		Expect(err).To(MatchError("test"))
+	})
+
+	It("can properly handle timeouts", func() {
+		jr, wjr := NewJobResult[string, string]("foo")
+		Expect(jr).ToNot(BeNil())
+		Expect(wjr).ToNot(BeNil())
+
+		go func(wjr *WritableJobResult[string, string]) {
+			time.Sleep(time.Second * 5)
+			wjr.SetResult("bar", nil)
+		}(wjr)
+
+		timeout1s, c1 := context.WithTimeoutCause(context.Background(), time.Second, fmt.Errorf("timeout"))
+		timeout10s, c2 := context.WithTimeoutCause(context.Background(), time.Second*10, fmt.Errorf("timeout"))
+
+		_, err := jr.Wait(timeout1s)
+		Expect(jr.Request).ToNot(BeNil())
+		Expect(*jr.Request()).To(Equal("foo"))
+		Expect(err).ToNot(BeNil())
+		Expect(err).To(MatchError(context.DeadlineExceeded))
+
+		resPtr, err := jr.Wait(timeout10s)
+		Expect(jr.Request).ToNot(BeNil())
+		Expect(*jr.Request()).To(Equal("foo"))
+		Expect(err).To(BeNil())
+		Expect(resPtr).ToNot(BeNil())
+		Expect(*resPtr).To(Equal("bar"))
+
+		// Is this needed? Cleanup Either Way.
+		c1()
+		c2()
+	})
+})

pkg/downloader/downloader_suite_test.go

@@ -0,0 +1,13 @@
+package downloader
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestDownloader(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Downloader test suite")
+}

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 {
-	js := JSONFunctionStructureName{}
+func (f Functions) ToJSONStructure(name, args string) JSONFunctionStructure {
+	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",
-			Properties: NameProperties{
-				Function: FunctionName{Const: function.Name},
-				Arguments: Argument{
-					Type:       "object",
-					Properties: prop,
-				},
-			},
+
+		property := map[string]interface{}{}
+		property[nameKey] = FunctionName{Const: function.Name}
+		property[argsKey] = Argument{
+			Type:       "object",
+			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"))
-			Expect(js.OneOf[0].Properties.Arguments.Properties["event_name"].(map[string]interface{})["type"]).To(Equal("string"))
-			Expect(js.OneOf[0].Properties.Arguments.Properties["event_date"].(map[string]interface{})["type"]).To(Equal("string"))
-			Expect(js.OneOf[1].Properties.Function.Const).To(Equal("search"))
-			Expect(js.OneOf[1].Properties.Arguments.Properties["query"].(map[string]interface{})["type"]).To(Equal("string"))
+			fnName := js.OneOf[0].Properties["function"].(FunctionName)
+			fnArgs := js.OneOf[0].Properties["arguments"].(Argument)
+			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"))
 
-			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"))
-			Expect(jsN.OneOf[0].Properties.Arguments.Properties["event_date"].(map[string]interface{})["type"]).To(Equal("string"))
-			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"))
+
+			fnName = jsN.OneOf[0].Properties["name"].(FunctionName)
+			fnArgs = jsN.OneOf[0].Properties["arguments"].(Argument)
+
+			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       string         `json:"type"`
-	Properties NameProperties `json:"properties"`
+type Item struct {
+	Type       string                 `json:"type"`
+	Properties map[string]interface{} `json:"properties"`
 }
 
-type ItemFunction struct {
-	Type       string             `json:"type"`
-	Properties FunctionProperties `json:"properties"`
-}
-
-type JSONFunctionStructureName struct {
-	OneOf []ItemName             `json:"oneOf,omitempty"`
-	AnyOf []ItemName             `json:"anyOf,omitempty"`
+type JSONFunctionStructure struct {
+	OneOf []Item                 `json:"oneOf,omitempty"`
+	AnyOf []Item                 `json:"anyOf,omitempty"`
 	Defs  map[string]interface{} `json:"$defs,omitempty"`
 }
 
-func (j JSONFunctionStructureName) 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 {
+func (j JSONFunctionStructure) 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{
-			Function: FunctionName{
-				Const: "create_event",
+		Properties: createFunction(
+			"function",
+			"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{
-			Function: FunctionName{
-				Const: "search",
-			},
-			Arguments: Argument{
-				Type: "object",
-				Properties: map[string]interface{}{
-					"query": map[string]string{"type": "string"},
-				},
-			},
-		},
+		Properties: createFunction(
+			"function",
+			"arguments",
+			"search",
+			map[string]interface{}{
+				"query": map[string]string{"type": "string"},
+			}),
 	},
 }
 
-var testFunctionsName = []ItemName{
+var testFunctionsName = []Item{
 	{
 		Type: "object",
-		Properties: NameProperties{
-			Function: FunctionName{
-				Const: "create_event",
+		Properties: createFunction(
+			"name",
+			"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{
-			Function: FunctionName{
-				Const: "search",
-			},
-			Arguments: Argument{
-				Type: "object",
-				Properties: map[string]interface{}{
-					"query": map[string]string{"type": "string"},
-				},
-			},
-		},
+		Properties: createFunction(
+			"name",
+			"arguments",
+			"search",
+			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"
 
@@ -23,8 +25,8 @@ type GrammarConfig struct {
 	// In this way if the LLM selects a free string, it won't be mixed necessarly with JSON objects
 	NoMixedFreeString bool `yaml:"no_mixed_free_string"`
 
-	// NoGrammar disables the grammar parsing and parses the responses directly from the LLM
-	NoGrammar bool `yaml:"disable"`
+	// EnableGrammar disables the grammar parsing and parses the responses directly from the LLM
+	EnableGrammar bool `yaml:"enable"`
 
 	// Prefix is the suffix to append to the grammar when being generated
 	// This is useful when models prepend a tag before returning JSON
@@ -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"
-	if functionConfig.FunctionName {
-		functionNameKey = "name"
+	functionNameKey := defaultFunctionNameKey
+	functionArgumentsKey := defaultFunctionArgumentsKey
+	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
-				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().Err(err).Str("escapedLLMResult", s).Msg("unable to unmarshal llm result in a single object or an array of JSON objects")
 			}
 
 			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}}`
-			functionConfig.FunctionName = true
+			input := `{"function": "add", "arguments": {"x": 5, "y": 3}}`
+			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>{"function": "subtract", "arguments": {"x": 10, "y": 7}}</tool_call>
+<tool_call>{"name": "add", "arguments": {"x": 5, "y": 3}}</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",

swagger/docs.go

@@ -700,18 +700,6 @@ const docTemplate = `{
                 }
             }
         },
-        "functions.Argument": {
-            "type": "object",
-            "properties": {
-                "properties": {
-                    "type": "object",
-                    "additionalProperties": true
-                },
-                "type": {
-                    "type": "string"
-                }
-            }
-        },
         "functions.Function": {
             "type": "object",
             "properties": {
@@ -727,48 +715,19 @@ const docTemplate = `{
                 }
             }
         },
-        "functions.FunctionName": {
-            "type": "object",
-            "properties": {
-                "const": {
-                    "type": "string"
-                }
-            }
-        },
-        "functions.FunctionProperties": {
-            "type": "object",
-            "properties": {
-                "arguments": {
-                    "$ref": "#/definitions/functions.Argument"
-                },
-                "function": {
-                    "$ref": "#/definitions/functions.FunctionName"
-                }
-            }
-        },
-        "functions.ItemFunction": {
+        "functions.Item": {
             "type": "object",
             "properties": {
                 "properties": {
-                    "$ref": "#/definitions/functions.FunctionProperties"
+                    "type": "object",
+                    "additionalProperties": true
                 },
                 "type": {
                     "type": "string"
                 }
             }
         },
-        "functions.ItemName": {
-            "type": "object",
-            "properties": {
-                "properties": {
-                    "$ref": "#/definitions/functions.NameProperties"
-                },
-                "type": {
-                    "type": "string"
-                }
-            }
-        },
-        "functions.JSONFunctionStructureFunction": {
+        "functions.JSONFunctionStructure": {
             "type": "object",
             "properties": {
                 "$defs": {
@@ -778,49 +737,17 @@ const docTemplate = `{
                 "anyOf": {
                     "type": "array",
                     "items": {
-                        "$ref": "#/definitions/functions.ItemFunction"
+                        "$ref": "#/definitions/functions.Item"
                     }
                 },
                 "oneOf": {
                     "type": "array",
                     "items": {
-                        "$ref": "#/definitions/functions.ItemFunction"
+                        "$ref": "#/definitions/functions.Item"
                     }
                 }
             }
         },
-        "functions.JSONFunctionStructureName": {
-            "type": "object",
-            "properties": {
-                "$defs": {
-                    "type": "object",
-                    "additionalProperties": true
-                },
-                "anyOf": {
-                    "type": "array",
-                    "items": {
-                        "$ref": "#/definitions/functions.ItemName"
-                    }
-                },
-                "oneOf": {
-                    "type": "array",
-                    "items": {
-                        "$ref": "#/definitions/functions.ItemName"
-                    }
-                }
-            }
-        },
-        "functions.NameProperties": {
-            "type": "object",
-            "properties": {
-                "arguments": {
-                    "$ref": "#/definitions/functions.Argument"
-                },
-                "name": {
-                    "$ref": "#/definitions/functions.FunctionName"
-                }
-            }
-        },
         "functions.Tool": {
             "type": "object",
             "properties": {
@@ -1488,10 +1415,7 @@ const docTemplate = `{
                     "type": "string"
                 },
                 "grammar_json_functions": {
-                    "$ref": "#/definitions/functions.JSONFunctionStructureFunction"
-                },
-                "grammar_json_name": {
-                    "$ref": "#/definitions/functions.JSONFunctionStructureName"
+                    "$ref": "#/definitions/functions.JSONFunctionStructure"
                 },
                 "ignore_eos": {
                     "type": "boolean"

swagger/swagger.json

@@ -693,18 +693,6 @@
                 }
             }
         },
-        "functions.Argument": {
-            "type": "object",
-            "properties": {
-                "properties": {
-                    "type": "object",
-                    "additionalProperties": true
-                },
-                "type": {
-                    "type": "string"
-                }
-            }
-        },
         "functions.Function": {
             "type": "object",
             "properties": {
@@ -720,48 +708,19 @@
                 }
             }
         },
-        "functions.FunctionName": {
-            "type": "object",
-            "properties": {
-                "const": {
-                    "type": "string"
-                }
-            }
-        },
-        "functions.FunctionProperties": {
-            "type": "object",
-            "properties": {
-                "arguments": {
-                    "$ref": "#/definitions/functions.Argument"
-                },
-                "function": {
-                    "$ref": "#/definitions/functions.FunctionName"
-                }
-            }
-        },
-        "functions.ItemFunction": {
+        "functions.Item": {
             "type": "object",
             "properties": {
                 "properties": {
-                    "$ref": "#/definitions/functions.FunctionProperties"
+                    "type": "object",
+                    "additionalProperties": true
                 },
                 "type": {
                     "type": "string"
                 }
             }
         },
-        "functions.ItemName": {
-            "type": "object",
-            "properties": {
-                "properties": {
-                    "$ref": "#/definitions/functions.NameProperties"
-                },
-                "type": {
-                    "type": "string"
-                }
-            }
-        },
-        "functions.JSONFunctionStructureFunction": {
+        "functions.JSONFunctionStructure": {
             "type": "object",
             "properties": {
                 "$defs": {
@@ -771,49 +730,17 @@
                 "anyOf": {
                     "type": "array",
                     "items": {
-                        "$ref": "#/definitions/functions.ItemFunction"
+                        "$ref": "#/definitions/functions.Item"
                     }
                 },
                 "oneOf": {
                     "type": "array",
                     "items": {
-                        "$ref": "#/definitions/functions.ItemFunction"
+                        "$ref": "#/definitions/functions.Item"
                     }
                 }
             }
         },
-        "functions.JSONFunctionStructureName": {
-            "type": "object",
-            "properties": {
-                "$defs": {
-                    "type": "object",
-                    "additionalProperties": true
-                },
-                "anyOf": {
-                    "type": "array",
-                    "items": {
-                        "$ref": "#/definitions/functions.ItemName"
-                    }
-                },
-                "oneOf": {
-                    "type": "array",
-                    "items": {
-                        "$ref": "#/definitions/functions.ItemName"
-                    }
-                }
-            }
-        },
-        "functions.NameProperties": {
-            "type": "object",
-            "properties": {
-                "arguments": {
-                    "$ref": "#/definitions/functions.Argument"
-                },
-                "name": {
-                    "$ref": "#/definitions/functions.FunctionName"
-                }
-            }
-        },
         "functions.Tool": {
             "type": "object",
             "properties": {
@@ -1481,10 +1408,7 @@
                     "type": "string"
                 },
                 "grammar_json_functions": {
-                    "$ref": "#/definitions/functions.JSONFunctionStructureFunction"
-                },
-                "grammar_json_name": {
-                    "$ref": "#/definitions/functions.JSONFunctionStructureName"
+                    "$ref": "#/definitions/functions.JSONFunctionStructure"
                 },
                 "ignore_eos": {
                     "type": "boolean"

swagger/swagger.yaml

@@ -7,14 +7,6 @@ definitions:
       url:
         type: string
     type: object
-  functions.Argument:
-    properties:
-      properties:
-        additionalProperties: true
-        type: object
-      type:
-        type: string
-    type: object
   functions.Function:
     properties:
       description:
@@ -25,67 +17,28 @@ definitions:
         additionalProperties: true
         type: object
     type: object
-  functions.FunctionName:
-    properties:
-      const:
-        type: string
-    type: object
-  functions.FunctionProperties:
-    properties:
-      arguments:
-        $ref: '#/definitions/functions.Argument'
-      function:
-        $ref: '#/definitions/functions.FunctionName'
-    type: object
-  functions.ItemFunction:
+  functions.Item:
     properties:
       properties:
-        $ref: '#/definitions/functions.FunctionProperties'
+        additionalProperties: true
+        type: object
       type:
         type: string
     type: object
-  functions.ItemName:
-    properties:
-      properties:
-        $ref: '#/definitions/functions.NameProperties'
-      type:
-        type: string
-    type: object
-  functions.JSONFunctionStructureFunction:
+  functions.JSONFunctionStructure:
     properties:
       $defs:
         additionalProperties: true
         type: object
       anyOf:
         items:
-          $ref: '#/definitions/functions.ItemFunction'
+          $ref: '#/definitions/functions.Item'
         type: array
       oneOf:
         items:
-          $ref: '#/definitions/functions.ItemFunction'
+          $ref: '#/definitions/functions.Item'
         type: array
     type: object
-  functions.JSONFunctionStructureName:
-    properties:
-      $defs:
-        additionalProperties: true
-        type: object
-      anyOf:
-        items:
-          $ref: '#/definitions/functions.ItemName'
-        type: array
-      oneOf:
-        items:
-          $ref: '#/definitions/functions.ItemName'
-        type: array
-    type: object
-  functions.NameProperties:
-    properties:
-      arguments:
-        $ref: '#/definitions/functions.Argument'
-      name:
-        $ref: '#/definitions/functions.FunctionName'
-    type: object
   functions.Tool:
     properties:
       function:
@@ -538,9 +491,7 @@ definitions:
         description: A grammar to constrain the LLM output
         type: string
       grammar_json_functions:
-        $ref: '#/definitions/functions.JSONFunctionStructureFunction'
-      grammar_json_name:
-        $ref: '#/definitions/functions.JSONFunctionStructureName'
+        $ref: '#/definitions/functions.JSONFunctionStructure'
       ignore_eos:
         type: boolean
       input: {}