add xai support (#135)

4 files changed, 74 insertions, 10 deletions

diff --git a/internal/llm/models/models.go b/internal/llm/models/models.go
index fd0a2b41b..1e1cbde50 100644
--- a/internal/llm/models/models.go
+++ b/internal/llm/models/models.go
@@ -89,4 +89,5 @@ func init() {
 	maps.Copy(SupportedModels, GroqModels)
 	maps.Copy(SupportedModels, AzureModels)
 	maps.Copy(SupportedModels, OpenRouterModels)
+	maps.Copy(SupportedModels, XAIModels)
 }
diff --git a/internal/llm/models/xai.go b/internal/llm/models/xai.go
new file mode 100644
index 000000000..00caf3b89
--- /dev/null
+++ b/internal/llm/models/xai.go
@@ -0,0 +1,61 @@
+package models
+
+const (
+	ProviderXAI ModelProvider = "xai"
+
+	XAIGrok3Beta         ModelID = "grok-3-beta"
+	XAIGrok3MiniBeta     ModelID = "grok-3-mini-beta"
+	XAIGrok3FastBeta     ModelID = "grok-3-fast-beta"
+	XAiGrok3MiniFastBeta ModelID = "grok-3-mini-fast-beta"
+)
+
+var XAIModels = map[ModelID]Model{
+	XAIGrok3Beta: {
+		ID:                 XAIGrok3Beta,
+		Name:               "Grok3 Beta",
+		Provider:           ProviderXAI,
+		APIModel:           "grok-3-beta",
+		CostPer1MIn:        3.0,
+		CostPer1MInCached:  0,
+		CostPer1MOut:       15,
+		CostPer1MOutCached: 0,
+		ContextWindow:      131_072,
+		DefaultMaxTokens:   20_000,
+	},
+	XAIGrok3MiniBeta: {
+		ID:                 XAIGrok3MiniBeta,
+		Name:               "Grok3 Mini Beta",
+		Provider:           ProviderXAI,
+		APIModel:           "grok-3-mini-beta",
+		CostPer1MIn:        0.3,
+		CostPer1MInCached:  0,
+		CostPer1MOut:       0.5,
+		CostPer1MOutCached: 0,
+		ContextWindow:      131_072,
+		DefaultMaxTokens:   20_000,
+	},
+	XAIGrok3FastBeta: {
+		ID:                 XAIGrok3FastBeta,
+		Name:               "Grok3 Fast Beta",
+		Provider:           ProviderXAI,
+		APIModel:           "grok-3-fast-beta",
+		CostPer1MIn:        5,
+		CostPer1MInCached:  0,
+		CostPer1MOut:       25,
+		CostPer1MOutCached: 0,
+		ContextWindow:      131_072,
+		DefaultMaxTokens:   20_000,
+	},
+	XAiGrok3MiniFastBeta: {
+		ID:                 XAiGrok3MiniFastBeta,
+		Name:               "Grok3 Mini Fast Beta",
+		Provider:           ProviderXAI,
+		APIModel:           "grok-3-mini-fast-beta",
+		CostPer1MIn:        0.6,
+		CostPer1MInCached:  0,
+		CostPer1MOut:       4.0,
+		CostPer1MOutCached: 0,
+		ContextWindow:      131_072,
+		DefaultMaxTokens:   20_000,
+	},
+}
diff --git a/internal/llm/provider/openai.go b/internal/llm/provider/openai.go
index b557df535..d68cfbc2d 100644
--- a/internal/llm/provider/openai.go
+++ b/internal/llm/provider/openai.go
@@ -258,15 +258,6 @@ func (o *openaiClient) stream(ctx context.Context, messages []message.Message, t
 				chunk := openaiStream.Current()
 				acc.AddChunk(chunk)
 
-				if tool, ok := acc.JustFinishedToolCall(); ok {
-					toolCalls = append(toolCalls, message.ToolCall{
-						ID:    tool.Id,
-						Name:  tool.Name,
-						Input: tool.Arguments,
-						Type:  "function",
-					})
-				}
-
 				for _, choice := range chunk.Choices {
 					if choice.Delta.Content != "" {
 						eventChan <- ProviderEvent{
@@ -282,7 +273,9 @@ func (o *openaiClient) stream(ctx context.Context, messages []message.Message, t
 			if err == nil || errors.Is(err, io.EOF) {
 				// Stream completed successfully
 				finishReason := o.finishReason(string(acc.ChatCompletion.Choices[0].FinishReason))
-
+				if len(acc.ChatCompletion.Choices[0].Message.ToolCalls) > 0 {
+					toolCalls = append(toolCalls, o.toolCalls(acc.ChatCompletion)...)
+				}
 				if len(toolCalls) > 0 {
 					finishReason = message.FinishReasonToolUse
 				}
diff --git a/internal/llm/provider/provider.go b/internal/llm/provider/provider.go
index 1545bc27a..cad11eeb3 100644
--- a/internal/llm/provider/provider.go
+++ b/internal/llm/provider/provider.go
@@ -132,6 +132,15 @@ func NewProvider(providerName models.ModelProvider, opts ...ProviderClientOption
 			options: clientOptions,
 			client:  newOpenAIClient(clientOptions),
 		}, nil
+	case models.ProviderXAI:
+		clientOptions.openaiOptions = append(clientOptions.openaiOptions,
+			WithOpenAIBaseURL("https://api.x.ai/v1"),
+		)
+		return &baseProvider[OpenAIClient]{
+			options: clientOptions,
+			client:  newOpenAIClient(clientOptions),
+		}, nil
+
 	case models.ProviderMock:
 		// TODO: implement mock client for test
 		panic("not implemented")