summaryrefslogtreecommitdiffhomepage
path: root/internal/llm/provider
diff options
context:
space:
mode:
Diffstat (limited to 'internal/llm/provider')
-rw-r--r--internal/llm/provider/anthropic.go21
-rw-r--r--internal/llm/provider/gemini.go256
-rw-r--r--internal/llm/provider/openai.go15
-rw-r--r--internal/llm/provider/provider.go7
4 files changed, 133 insertions, 166 deletions
diff --git a/internal/llm/provider/anthropic.go b/internal/llm/provider/anthropic.go
index 63a68b92b..2b960ebca 100644
--- a/internal/llm/provider/anthropic.go
+++ b/internal/llm/provider/anthropic.go
@@ -111,7 +111,7 @@ func (a *anthropicProvider) StreamResponse(ctx context.Context, messages []messa
var thinkingParam anthropic.ThinkingConfigParamUnion
lastMessage := messages[len(messages)-1]
temperature := anthropic.Float(0)
- if lastMessage.Role == message.User && strings.Contains(strings.ToLower(lastMessage.Content), "think") {
+ if lastMessage.Role == message.User && strings.Contains(strings.ToLower(lastMessage.Content().String()), "think") {
thinkingParam = anthropic.ThinkingConfigParamUnion{
OfThinkingConfigEnabled: &anthropic.ThinkingConfigEnabledParam{
BudgetTokens: int64(float64(a.maxTokens) * 0.8),
@@ -187,9 +187,10 @@ func (a *anthropicProvider) StreamResponse(ctx context.Context, messages []messa
eventChan <- ProviderEvent{
Type: EventComplete,
Response: &ProviderResponse{
- Content: content,
- ToolCalls: toolCalls,
- Usage: tokenUsage,
+ Content: content,
+ ToolCalls: toolCalls,
+ Usage: tokenUsage,
+ FinishReason: string(accumulatedMessage.StopReason),
},
}
}
@@ -263,7 +264,7 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag
for i, msg := range messages {
switch msg.Role {
case message.User:
- content := anthropic.NewTextBlock(msg.Content)
+ content := anthropic.NewTextBlock(msg.Content().String())
if cachedBlocks < 2 {
content.OfRequestTextBlock.CacheControl = anthropic.CacheControlEphemeralParam{
Type: "ephemeral",
@@ -274,8 +275,8 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag
case message.Assistant:
blocks := []anthropic.ContentBlockParamUnion{}
- if msg.Content != "" {
- content := anthropic.NewTextBlock(msg.Content)
+ if msg.Content().String() != "" {
+ content := anthropic.NewTextBlock(msg.Content().String())
if cachedBlocks < 2 {
content.OfRequestTextBlock.CacheControl = anthropic.CacheControlEphemeralParam{
Type: "ephemeral",
@@ -285,7 +286,7 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag
blocks = append(blocks, content)
}
- for _, toolCall := range msg.ToolCalls {
+ for _, toolCall := range msg.ToolCalls() {
var inputMap map[string]any
err := json.Unmarshal([]byte(toolCall.Input), &inputMap)
if err != nil {
@@ -297,8 +298,8 @@ func (a *anthropicProvider) convertToAnthropicMessages(messages []message.Messag
anthropicMessages[i] = anthropic.NewAssistantMessage(blocks...)
case message.Tool:
- results := make([]anthropic.ContentBlockParamUnion, len(msg.ToolResults))
- for i, toolResult := range msg.ToolResults {
+ results := make([]anthropic.ContentBlockParamUnion, len(msg.ToolResults()))
+ for i, toolResult := range msg.ToolResults() {
results[i] = anthropic.NewToolResultBlock(toolResult.ToolCallID, toolResult.Content, toolResult.IsError)
}
anthropicMessages[i] = anthropic.NewUserMessage(results...)
diff --git a/internal/llm/provider/gemini.go b/internal/llm/provider/gemini.go
index 6b252b581..53ffa154e 100644
--- a/internal/llm/provider/gemini.go
+++ b/internal/llm/provider/gemini.go
@@ -78,7 +78,6 @@ func (p *geminiProvider) Close() {
}
}
-// convertToGeminiHistory converts the message history to Gemini's format
func (p *geminiProvider) convertToGeminiHistory(messages []message.Message) []*genai.Content {
var history []*genai.Content
@@ -86,7 +85,7 @@ func (p *geminiProvider) convertToGeminiHistory(messages []message.Message) []*g
switch msg.Role {
case message.User:
history = append(history, &genai.Content{
- Parts: []genai.Part{genai.Text(msg.Content)},
+ Parts: []genai.Part{genai.Text(msg.Content().String())},
Role: "user",
})
case message.Assistant:
@@ -95,14 +94,12 @@ func (p *geminiProvider) convertToGeminiHistory(messages []message.Message) []*g
Parts: []genai.Part{},
}
- // Handle regular content
- if msg.Content != "" {
- content.Parts = append(content.Parts, genai.Text(msg.Content))
+ if msg.Content().String() != "" {
+ content.Parts = append(content.Parts, genai.Text(msg.Content().String()))
}
- // Handle tool calls if any
- if len(msg.ToolCalls) > 0 {
- for _, call := range msg.ToolCalls {
+ if len(msg.ToolCalls()) > 0 {
+ for _, call := range msg.ToolCalls() {
args, _ := parseJsonToMap(call.Input)
content.Parts = append(content.Parts, genai.FunctionCall{
Name: call.Name,
@@ -113,8 +110,7 @@ func (p *geminiProvider) convertToGeminiHistory(messages []message.Message) []*g
history = append(history, content)
case message.Tool:
- for _, result := range msg.ToolResults {
- // Parse response content to map if possible
+ for _, result := range msg.ToolResults() {
response := map[string]interface{}{"result": result.Content}
parsed, err := parseJsonToMap(result.Content)
if err == nil {
@@ -123,7 +119,7 @@ func (p *geminiProvider) convertToGeminiHistory(messages []message.Message) []*g
var toolCall message.ToolCall
for _, msg := range messages {
if msg.Role == message.Assistant {
- for _, call := range msg.ToolCalls {
+ for _, call := range msg.ToolCalls() {
if call.ID == result.ToolCallID {
toolCall = call
break
@@ -146,108 +142,6 @@ func (p *geminiProvider) convertToGeminiHistory(messages []message.Message) []*g
return history
}
-// convertToolsToGeminiFunctionDeclarations converts tool definitions to Gemini's function declarations
-func (p *geminiProvider) convertToolsToGeminiFunctionDeclarations(tools []tools.BaseTool) []*genai.FunctionDeclaration {
- declarations := make([]*genai.FunctionDeclaration, len(tools))
-
- for i, tool := range tools {
- info := tool.Info()
-
- // Convert parameters to genai.Schema format
- properties := make(map[string]*genai.Schema)
- for name, param := range info.Parameters {
- // Try to extract type and description from the parameter
- paramMap, ok := param.(map[string]interface{})
- if !ok {
- // Default to string if unable to determine type
- properties[name] = &genai.Schema{Type: genai.TypeString}
- continue
- }
-
- schemaType := genai.TypeString // Default
- var description string
- var itemsTypeSchema *genai.Schema
- if typeVal, found := paramMap["type"]; found {
- if typeStr, ok := typeVal.(string); ok {
- switch typeStr {
- case "string":
- schemaType = genai.TypeString
- case "number":
- schemaType = genai.TypeNumber
- case "integer":
- schemaType = genai.TypeInteger
- case "boolean":
- schemaType = genai.TypeBoolean
- case "array":
- schemaType = genai.TypeArray
- items, found := paramMap["items"]
- if found {
- itemsMap, ok := items.(map[string]interface{})
- if ok {
- itemsType, found := itemsMap["type"]
- if found {
- itemsTypeStr, ok := itemsType.(string)
- if ok {
- switch itemsTypeStr {
- case "string":
- itemsTypeSchema = &genai.Schema{
- Type: genai.TypeString,
- }
- case "number":
- itemsTypeSchema = &genai.Schema{
- Type: genai.TypeNumber,
- }
- case "integer":
- itemsTypeSchema = &genai.Schema{
- Type: genai.TypeInteger,
- }
- case "boolean":
- itemsTypeSchema = &genai.Schema{
- Type: genai.TypeBoolean,
- }
- }
- }
- }
- }
- }
- case "object":
- schemaType = genai.TypeObject
- if _, found := paramMap["properties"]; !found {
- continue
- }
- // TODO: Add support for other types
- }
- }
- }
-
- if desc, found := paramMap["description"]; found {
- if descStr, ok := desc.(string); ok {
- description = descStr
- }
- }
-
- properties[name] = &genai.Schema{
- Type: schemaType,
- Description: description,
- Items: itemsTypeSchema,
- }
- }
-
- declarations[i] = &genai.FunctionDeclaration{
- Name: info.Name,
- Description: info.Description,
- Parameters: &genai.Schema{
- Type: genai.TypeObject,
- Properties: properties,
- Required: info.Required,
- },
- }
- }
-
- return declarations
-}
-
-// extractTokenUsage extracts token usage information from Gemini's response
func (p *geminiProvider) extractTokenUsage(resp *genai.GenerateContentResponse) TokenUsage {
if resp == nil || resp.UsageMetadata == nil {
return TokenUsage{}
@@ -261,41 +155,28 @@ func (p *geminiProvider) extractTokenUsage(resp *genai.GenerateContentResponse)
}
}
-// SendMessages sends a batch of messages to Gemini and returns the response
func (p *geminiProvider) SendMessages(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (*ProviderResponse, error) {
- // Create a generative model
model := p.client.GenerativeModel(p.model.APIModel)
model.SetMaxOutputTokens(p.maxTokens)
- // Set system instruction
model.SystemInstruction = genai.NewUserContent(genai.Text(p.systemMessage))
- // Set up tools if provided
if len(tools) > 0 {
declarations := p.convertToolsToGeminiFunctionDeclarations(tools)
- model.Tools = []*genai.Tool{{FunctionDeclarations: declarations}}
+ for _, declaration := range declarations {
+ model.Tools = append(model.Tools, &genai.Tool{FunctionDeclarations: []*genai.FunctionDeclaration{declaration}})
+ }
}
- // Create chat session and set history
chat := model.StartChat()
chat.History = p.convertToGeminiHistory(messages[:len(messages)-1]) // Exclude last message
- // Get the most recent user message
- var lastUserMsg message.Message
- for i := len(messages) - 1; i >= 0; i-- {
- if messages[i].Role == message.User {
- lastUserMsg = messages[i]
- break
- }
- }
-
- // Send the message
- resp, err := chat.SendMessage(ctx, genai.Text(lastUserMsg.Content))
+ lastUserMsg := messages[len(messages)-1]
+ resp, err := chat.SendMessage(ctx, genai.Text(lastUserMsg.Content().String()))
if err != nil {
return nil, err
}
- // Process the response
var content string
var toolCalls []message.ToolCall
@@ -317,7 +198,6 @@ func (p *geminiProvider) SendMessages(ctx context.Context, messages []message.Me
}
}
- // Extract token usage
tokenUsage := p.extractTokenUsage(resp)
return &ProviderResponse{
@@ -327,16 +207,12 @@ func (p *geminiProvider) SendMessages(ctx context.Context, messages []message.Me
}, nil
}
-// StreamResponse streams the response from Gemini
func (p *geminiProvider) StreamResponse(ctx context.Context, messages []message.Message, tools []tools.BaseTool) (<-chan ProviderEvent, error) {
- // Create a generative model
model := p.client.GenerativeModel(p.model.APIModel)
model.SetMaxOutputTokens(p.maxTokens)
- // Set system instruction
model.SystemInstruction = genai.NewUserContent(genai.Text(p.systemMessage))
- // Set up tools if provided
if len(tools) > 0 {
declarations := p.convertToolsToGeminiFunctionDeclarations(tools)
for _, declaration := range declarations {
@@ -344,14 +220,12 @@ func (p *geminiProvider) StreamResponse(ctx context.Context, messages []message.
}
}
- // Create chat session and set history
chat := model.StartChat()
chat.History = p.convertToGeminiHistory(messages[:len(messages)-1]) // Exclude last message
lastUserMsg := messages[len(messages)-1]
- // Start streaming
- iter := chat.SendMessageStream(ctx, genai.Text(lastUserMsg.Content))
+ iter := chat.SendMessageStream(ctx, genai.Text(lastUserMsg.Content().String()))
eventChan := make(chan ProviderEvent)
@@ -392,7 +266,6 @@ func (p *geminiProvider) StreamResponse(ctx context.Context, messages []message.
}
currentContent += newText
case genai.FunctionCall:
- // For function calls, we assume they come complete, not streamed in parts
id := "call_" + uuid.New().String()
args, _ := json.Marshal(p.Args)
newCall := message.ToolCall{
@@ -402,7 +275,6 @@ func (p *geminiProvider) StreamResponse(ctx context.Context, messages []message.
Type: "function",
}
- // Check if this is a new tool call
isNew := true
for _, existing := range toolCalls {
if existing.Name == newCall.Name && existing.Input == newCall.Input {
@@ -419,15 +291,15 @@ func (p *geminiProvider) StreamResponse(ctx context.Context, messages []message.
}
}
- // Extract token usage from the final response
tokenUsage := p.extractTokenUsage(finalResp)
eventChan <- ProviderEvent{
Type: EventComplete,
Response: &ProviderResponse{
- Content: currentContent,
- ToolCalls: toolCalls,
- Usage: tokenUsage,
+ Content: currentContent,
+ ToolCalls: toolCalls,
+ Usage: tokenUsage,
+ FinishReason: string(finalResp.Candidates[0].FinishReason.String()),
},
}
}()
@@ -435,7 +307,99 @@ func (p *geminiProvider) StreamResponse(ctx context.Context, messages []message.
return eventChan, nil
}
-// Helper function to parse JSON string into map
+func (p *geminiProvider) convertToolsToGeminiFunctionDeclarations(tools []tools.BaseTool) []*genai.FunctionDeclaration {
+ declarations := make([]*genai.FunctionDeclaration, len(tools))
+
+ for i, tool := range tools {
+ info := tool.Info()
+ declarations[i] = &genai.FunctionDeclaration{
+ Name: info.Name,
+ Description: info.Description,
+ Parameters: &genai.Schema{
+ Type: genai.TypeObject,
+ Properties: convertSchemaProperties(info.Parameters),
+ Required: info.Required,
+ },
+ }
+ }
+
+ return declarations
+}
+
+func convertSchemaProperties(parameters map[string]interface{}) map[string]*genai.Schema {
+ properties := make(map[string]*genai.Schema)
+
+ for name, param := range parameters {
+ properties[name] = convertToSchema(param)
+ }
+
+ return properties
+}
+
+func convertToSchema(param interface{}) *genai.Schema {
+ schema := &genai.Schema{Type: genai.TypeString}
+
+ paramMap, ok := param.(map[string]interface{})
+ if !ok {
+ return schema
+ }
+
+ if desc, ok := paramMap["description"].(string); ok {
+ schema.Description = desc
+ }
+
+ typeVal, hasType := paramMap["type"]
+ if !hasType {
+ return schema
+ }
+
+ typeStr, ok := typeVal.(string)
+ if !ok {
+ return schema
+ }
+
+ schema.Type = mapJSONTypeToGenAI(typeStr)
+
+ switch typeStr {
+ case "array":
+ schema.Items = processArrayItems(paramMap)
+ case "object":
+ if props, ok := paramMap["properties"].(map[string]interface{}); ok {
+ schema.Properties = convertSchemaProperties(props)
+ }
+ }
+
+ return schema
+}
+
+func processArrayItems(paramMap map[string]interface{}) *genai.Schema {
+ items, ok := paramMap["items"].(map[string]interface{})
+ if !ok {
+ return nil
+ }
+
+ return convertToSchema(items)
+}
+
+func mapJSONTypeToGenAI(jsonType string) genai.Type {
+ switch jsonType {
+ case "string":
+ return genai.TypeString
+ case "number":
+ return genai.TypeNumber
+ case "integer":
+ return genai.TypeInteger
+ case "boolean":
+ return genai.TypeBoolean
+ case "array":
+ return genai.TypeArray
+ case "object":
+ return genai.TypeObject
+ default:
+ return genai.TypeString // Default to string for unknown types
+ }
+}
+
func parseJsonToMap(jsonStr string) (map[string]interface{}, error) {
var result map[string]interface{}
err := json.Unmarshal([]byte(jsonStr), &result)
diff --git a/internal/llm/provider/openai.go b/internal/llm/provider/openai.go
index d86a58690..c8e04d5ee 100644
--- a/internal/llm/provider/openai.go
+++ b/internal/llm/provider/openai.go
@@ -84,22 +84,22 @@ func (p *openaiProvider) convertToOpenAIMessages(messages []message.Message) []o
for _, msg := range messages {
switch msg.Role {
case message.User:
- chatMessages = append(chatMessages, openai.UserMessage(msg.Content))
+ chatMessages = append(chatMessages, openai.UserMessage(msg.Content().String()))
case message.Assistant:
assistantMsg := openai.ChatCompletionAssistantMessageParam{
Role: "assistant",
}
- if msg.Content != "" {
+ if msg.Content().String() != "" {
assistantMsg.Content = openai.ChatCompletionAssistantMessageParamContentUnion{
- OfString: openai.String(msg.Content),
+ OfString: openai.String(msg.Content().String()),
}
}
- if len(msg.ToolCalls) > 0 {
- assistantMsg.ToolCalls = make([]openai.ChatCompletionMessageToolCallParam, len(msg.ToolCalls))
- for i, call := range msg.ToolCalls {
+ if len(msg.ToolCalls()) > 0 {
+ assistantMsg.ToolCalls = make([]openai.ChatCompletionMessageToolCallParam, len(msg.ToolCalls()))
+ for i, call := range msg.ToolCalls() {
assistantMsg.ToolCalls[i] = openai.ChatCompletionMessageToolCallParam{
ID: call.ID,
Type: "function",
@@ -116,7 +116,7 @@ func (p *openaiProvider) convertToOpenAIMessages(messages []message.Message) []o
})
case message.Tool:
- for _, result := range msg.ToolResults {
+ for _, result := range msg.ToolResults() {
chatMessages = append(chatMessages,
openai.ToolMessage(result.Content, result.ToolCallID),
)
@@ -276,3 +276,4 @@ func (p *openaiProvider) StreamResponse(ctx context.Context, messages []message.
return eventChan, nil
}
+
diff --git a/internal/llm/provider/provider.go b/internal/llm/provider/provider.go
index 9ac1def37..f40429738 100644
--- a/internal/llm/provider/provider.go
+++ b/internal/llm/provider/provider.go
@@ -27,9 +27,10 @@ type TokenUsage struct {
}
type ProviderResponse struct {
- Content string
- ToolCalls []message.ToolCall
- Usage TokenUsage
+ Content string
+ ToolCalls []message.ToolCall
+ Usage TokenUsage
+ FinishReason string
}
type ProviderEvent struct {