Add Azure OpenAI support and improve message parsing

- Implement Azure OpenAI authentication with `azure` flag in OpenAI struct
- Update HTTP request headers to support Azure API key authentication
- Enhance message parsing to handle additional edge cases in OpenAI responses
- Add debug logging for OpenAI messages when debug environment is set
- Improve error and usage detection in stream message processing

neo/message/message.go

@@ -2,6 +2,7 @@ package message
 
 import (
 	"fmt"
+	"os"
 	"strings"
 
 	"github.com/fatih/color"
@@ -184,7 +185,16 @@ func NewAny(content interface{}) (*Message, error) {
 }
 
 // NewOpenAI create a new message from OpenAI response
+// @todo:
+//
+//	this function need to be refactored
 func NewOpenAI(data []byte, isThinking bool) *Message {
+
+	// For Debug
+	if os.Getenv("YAO_AGENT_DEBUG") == "true" {
+		fmt.Printf("%s\n", string(data))
+	}
+
 	if data == nil || len(data) == 0 {
 		return nil
 	}
@@ -194,7 +204,7 @@ func NewOpenAI(data []byte, isThinking bool) *Message {
 	data = []byte(strings.TrimPrefix(text, "data: "))
 	switch {
 
-	case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"tool_calls"`):
+	case strings.Contains(text, `"delta":{`) && strings.Contains(text, `"tool_calls"`) && !strings.Contains(text, `"tool_calls":null`):
 		var toolCalls openai.ToolCalls
 		if err := jsoniter.Unmarshal(data, &toolCalls); err != nil {
 			color.Red("JSON parse error: %s", err.Error())
@@ -247,7 +257,7 @@ func NewOpenAI(data []byte, isThinking bool) *Message {
 			return msg
 		}
 
-	case strings.Index(text, `{"code":`) == 0:
+	case strings.Index(text, `{"code":`) == 0 || strings.Index(text, `"statusCode":`) > 0:
 		var errorMessage openai.Error
 		if err := jsoniter.UnmarshalFromString(text, &errorMessage); err != nil {
 			color.Red("JSON parse error: %s", err.Error())
@@ -277,6 +287,10 @@ func NewOpenAI(data []byte, isThinking bool) *Message {
 		msg.IsDone = true
 		break
 
+	case strings.Contains(text, `"usage":`) && !strings.Contains(text, `"chat.completion.chunk`):
+		msg.IsDone = true
+		break
+
 	case strings.Contains(text, `[DONE]`):
 		msg.IsDone = true
 

openai/openai.go

@@ -31,6 +31,7 @@ type OpenAI struct {
 	baseURL      string
 	organization string
 	maxToken     int
+	azure        bool // Azure Credentials, "true" or "false" or ""
 }
 
 // New create a new OpenAI instance by connector id
@@ -94,6 +95,11 @@ func NewOpenAI(setting map[string]interface{}) (*OpenAI, error) {
 		maxToken = v
 	}
 
+	azure := false
+	if v, ok := setting["azure"].(string); ok {
+		azure = v == "true" || v == "1"
+	}
+
 	return &OpenAI{
 		key:          key,
 		model:        model,
@@ -101,6 +107,7 @@ func NewOpenAI(setting map[string]interface{}) (*OpenAI, error) {
 		baseURL:      baseURL,
 		organization: organization,
 		maxToken:     maxToken,
+		azure:        azure,
 	}, nil
 }
 
@@ -366,11 +373,14 @@ func (openai OpenAI) Stream(ctx context.Context, path string, payload map[string
 func (openai OpenAI) post(path string, payload map[string]interface{}) (interface{}, *exception.Exception) {
 
 	url := fmt.Sprintf("%s%s", openai.host, path)
-	key := fmt.Sprintf("Bearer %s", openai.key)
 	payload["model"] = openai.model
 
-	req := http.New(url).
-		WithHeader(map[string][]string{"Authorization": {key}})
+	req := http.New(url)
+	if openai.azure {
+		req.WithHeader(map[string][]string{"api-key": {openai.key}})
+	} else {
+		req.WithHeader(map[string][]string{"Authorization": {fmt.Sprintf("Bearer %s", openai.key)}})
+	}
 
 	res := req.Post(payload)
 	if err := openai.isError(res); err != nil {
@@ -383,10 +393,12 @@ func (openai OpenAI) post(path string, payload map[string]interface{}) (interfac
 func (openai OpenAI) postWithoutModel(path string, payload map[string]interface{}) (interface{}, *exception.Exception) {
 
 	url := fmt.Sprintf("%s%s", openai.host, path)
-	key := fmt.Sprintf("Bearer %s", openai.key)
-
-	req := http.New(url).
-		WithHeader(map[string][]string{"Authorization": {key}})
+	req := http.New(url)
+	if openai.azure {
+		req.WithHeader(map[string][]string{"api-key": {openai.key}})
+	} else {
+		req.WithHeader(map[string][]string{"Authorization": {fmt.Sprintf("Bearer %s", openai.key)}})
+	}
 
 	res := req.Post(payload)
 	if err := openai.isError(res); err != nil {
@@ -399,14 +411,15 @@ func (openai OpenAI) postWithoutModel(path string, payload map[string]interface{
 func (openai OpenAI) postFile(path string, files map[string][]byte, option map[string]interface{}) (interface{}, *exception.Exception) {
 
 	url := fmt.Sprintf("%s%s", openai.host, path)
-	key := fmt.Sprintf("Bearer %s", openai.key)
 	option["model"] = openai.model
 
-	req := http.New(url).
-		WithHeader(map[string][]string{
-			"Authorization": {key},
-			"Content-Type":  {"multipart/form-data"},
-		})
+	req := http.New(url).WithHeader(map[string][]string{"Content-Type": {"multipart/form-data"}})
+
+	if openai.azure {
+		req.WithHeader(map[string][]string{"api-key": {openai.key}})
+	} else {
+		req.WithHeader(map[string][]string{"Authorization": {fmt.Sprintf("Bearer %s", openai.key)}})
+	}
 
 	for name, data := range files {
 		req.AddFileBytes(name, fmt.Sprintf("%s.mp3", name), data)
@@ -425,11 +438,12 @@ func (openai OpenAI) postFileWithoutModel(path string, files map[string][]byte,
 	url := fmt.Sprintf("%s%s", openai.host, path)
 	key := fmt.Sprintf("Bearer %s", openai.key)
 
-	req := http.New(url).
-		WithHeader(map[string][]string{
-			"Authorization": {key},
-			"Content-Type":  {"multipart/form-data"},
-		})
+	req := http.New(url).WithHeader(map[string][]string{"Authorization": {key}})
+	if openai.azure {
+		req.WithHeader(map[string][]string{"api-key": {openai.key}})
+	} else {
+		req.WithHeader(map[string][]string{"Authorization": {fmt.Sprintf("Bearer %s", openai.key)}})
+	}
 
 	for name, data := range files {
 		req.AddFileBytes(name, fmt.Sprintf("%s.mp3", name), data)
@@ -445,16 +459,17 @@ func (openai OpenAI) postFileWithoutModel(path string, files map[string][]byte,
 // stream post request
 func (openai OpenAI) stream(ctx context.Context, path string, payload map[string]interface{}, cb func(data []byte) int) *exception.Exception {
 	url := fmt.Sprintf("%s%s", openai.host, path)
-	key := fmt.Sprintf("Bearer %s", openai.key)
 	payload["model"] = openai.model
 	req := http.New(url)
-	err := req.
-		WithHeader(map[string][]string{
-			"Content-Type":  {"application/json; charset=utf-8"},
-			"Authorization": {key},
-		}).
-		Stream(ctx, "POST", payload, cb)
+	req.WithHeader(map[string][]string{"Content-Type": {"application/json; charset=utf-8"}})
 
+	if openai.azure {
+		req.WithHeader(map[string][]string{"api-key": {openai.key}})
+	} else {
+		req.WithHeader(map[string][]string{"Authorization": {fmt.Sprintf("Bearer %s", openai.key)}})
+	}
+
+	err := req.Stream(ctx, "POST", payload, cb)
 	if err != nil {
 		return exception.New(err.Error(), 500)
 	}