zlei6/ragflow
1
0
Fork 0
mirror of https://github.com/infiniflow/ragflow.git synced 2026-08-05 15:20:30 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(go-models): migrate batch 3 model drivers to unified handlers (#17698)

Browse Source 
## Summary

Relate to #17284

Migrate 10 OpenAI-compatible drivers (`minimax`, `mistral`,
`modelscope`, `moonshot`, `n1n`, `novita`, `ollama`, `openai`,
`openai_api_compatible`, `openrouter`) to use the unified response
handlers (`HandleNonStreamingResponse` / `HandleStreamingResponse`),
following the same pattern established by `deepseek` in #17634.

- Cut ~150 lines per driver (1692 lines removed, 172 added across 10
files).
- `openai_api_compatible` gains `ChatWithMessages` +
`ChatStreamlyWithSender` required by the unified handler infrastructure.
- `openai` driver preserves `reasoning_content` extraction for o-series
models.
- All drivers: pure deduplication of HTTP plumbing.

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
jay77721
2026-08-03 15:08:55 +08:00
committed by GitHub
parent bddc941814
commit 1478aa4ced
15 changed files with 447 additions and 1904 deletions
Download Patch File Download Diff File Expand all files Collapse all files

247
internal/entity/models/minimax.go
View File

@@ -17,6 +17,7 @@
package models
import (
"bufio"
"bytes"
"context"
"encoding/hex"
@@ -52,53 +53,6 @@ func (m *MinimaxModel) Name() string {
return "minimax"
}
type MinimaxChatResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []struct {
Index int `json:"index"`
Message struct {
Role string `json:"role"`
Content *string `json:"content"`
ToolCalls []map[string]any `json:"tool_calls"`
ReasoningContent *string `json:"reasoning_content"`
} `json:"message"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
Usage struct {
PromptTokens int `json:"prompt_tokens"`
PromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
CompletionTokens int `json:"completion_tokens"`
TotalTokens int `json:"total_tokens"`
TotalCharacters int `json:"total_characters"`
} `json:"usage"`
BaseResp struct {
StatusCode int `json:"status_code"`
StatusMsg string `json:"status_msg"`
} `json:"base_resp"`
Error struct {
Message string `json:"message"`
Type string `json:"type"`
} `json:"error"`
}
func extractMinimaxChatResponseError(result *MinimaxChatResponse) string {
if result == nil {
return ""
}
if result.BaseResp.StatusCode != 0 {
if result.BaseResp.StatusMsg != "" {
return result.BaseResp.StatusMsg
}
return fmt.Sprintf("status_code %d", result.BaseResp.StatusCode)
}
return result.Error.Message
}
func validateMinimaxModelName(modelName string) (string, error) {
if strings.TrimSpace(modelName) == "" {
return "", fmt.Errorf("model name is required")
@@ -155,7 +109,6 @@ func (m *MinimaxModel) ChatWithMessages(ctx context.Context, modelName string, m
if err := m.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
apiKey := strings.TrimSpace(*apiConfig.ApiKey)
modelName, err := validateMinimaxModelName(modelName)
if err != nil {
return nil, err
@@ -189,87 +142,23 @@ func (m *MinimaxModel) ChatWithMessages(ctx context.Context, modelName string, m
}
jsonData, err := json.Marshal(reqBody)
body, err := m.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
// MiniMax can embed errors in a base_resp block with HTTP 200.
// Check for these before using the shared handler so the caller
// sees the real error message instead of "no choices in response".
var result map[string]any
if err := json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
if errMsg := extractMinimaxAPIError(result); errMsg != "" {
return nil, fmt.Errorf("minimax API error: %s", errMsg)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey))
resp, err := m.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("minimax API error: status %d: %s", resp.StatusCode, extractMinimaxErrorBody(body))
}
return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
var result MinimaxChatResponse
if err := json.Unmarshal(body, &result); err != nil {
return chatResponseParts{}, fmt.Errorf("failed to unmarshal response: %w", err)
}
if errMsg := extractMinimaxChatResponseError(&result); errMsg != "" {
return chatResponseParts{}, fmt.Errorf("minimax API error: %s", errMsg)
}
if len(result.Choices) == 0 {
return chatResponseParts{}, fmt.Errorf("no choices in response")
}
choice := &result.Choices[0]
content := ""
if choice.Message.Content != nil {
content = *choice.Message.Content
}
reasonContent := ""
if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking {
if choice.Message.ReasoningContent != nil {
reasonContent = *choice.Message.ReasoningContent
}
if reasonContent != "" && reasonContent[0] == '\n' {
reasonContent = reasonContent[1:]
}
}
totalTokens := result.Usage.TotalTokens
if totalTokens == 0 {
totalTokens = result.Usage.PromptTokens + result.Usage.CompletionTokens
}
var usage *TokenUsage
if totalTokens > 0 {
usage = &TokenUsage{
PromptTokens: result.Usage.PromptTokens,
CompletionTokens: result.Usage.CompletionTokens,
TotalTokens: totalTokens,
}
}
return chatResponseParts{
RequestID: result.ID,
Content: &content,
ReasonContent: &reasonContent,
ToolCalls: choice.Message.ToolCalls,
Usage: usage,
}, nil
})
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel)
@@ -314,6 +203,8 @@ func (m *MinimaxModel) ChatStreamlyWithSender(ctx context.Context, modelName str
}
reqBody["stream_options"] = map[string]interface{}{"include_usage": true}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
@@ -342,75 +233,59 @@ func (m *MinimaxModel) ChatStreamlyWithSender(ctx context.Context, modelName str
return fmt.Errorf("minimax API error: status %d: %s", resp.StatusCode, extractMinimaxErrorBody(body))
}
// SSE parsing: read line by line
sawTerminal := false
accumulatedToolCalls := make(map[int]map[string]any)
done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event)
if usageErr != nil {
return usageErr
}
if found {
applyStreamUsage(modelConfig, modelUsage, tokenUsage)
}
// Pipe the response through a base_resp checker. MiniMax can send
// error events (e.g. rate limits) without a choices array, and the
// shared handler skips those silently. We surface them so the retry
// predicates can match and the caller sees the real reason.
pr, pw := io.Pipe()
// Close pr when this function returns so an early exit from
// HandleStreamingResponse unblocks the producer goroutine below
// (its pw.Write fails) and releases resp.Body instead of leaving
// the reader blocked on a live pipe.
defer pr.Close()
streamErr := make(chan error, 1)
go func() {
defer pw.Close()
defer resp.Body.Close()
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
// MiniMax can send an error event (rate limit, etc.)
// without a choices array. Surface it so the retry
// predicates can match and the caller sees the real
// reason instead of a generic "stream ended" error.
if errMsg := extractMinimaxAPIError(event); errMsg != "" {
return fmt.Errorf("minimax API error: %s", errMsg)
var scanErr error
// Ensure streamErr always receives a result, on every exit
// path, so the final receive below can never block.
defer func() {
select {
case streamErr <- scanErr:
default:
}
return nil
}
}()
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
return nil
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
content, ok := delta["content"].(string)
if ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
scanner := bufio.NewScanner(resp.Body)
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
for scanner.Scan() {
line := scanner.Text()
if strings.HasPrefix(line, "data:") {
data := strings.TrimSpace(line[5:])
if data != "" && data != "[DONE]" {
var event map[string]any
if json.Unmarshal([]byte(data), &event) == nil {
if errMsg := extractMinimaxAPIError(event); errMsg != "" {
pw.CloseWithError(fmt.Errorf("minimax API error: %s", errMsg))
return
}
}
}
}
if _, err := pw.Write([]byte(line + "\n")); err != nil {
scanErr = err
return
}
}
scanErr = scanner.Err()
}()
reasoningContent, ok := delta["reasoning_content"].(string)
if ok && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
return err
}
}
finishReason, ok := firstChoice["finish_reason"].(string)
if ok && finishReason != "" {
sawTerminal = true
}
return nil
})
if err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
if err := HandleStreamingResponse(pr, modelUsage, modelConfig, OpenAIParserConfig, sender); err != nil {
return err
}
if !done && !sawTerminal {
return fmt.Errorf("minimax: stream ended before [DONE] or finish_reason")
}
setSortedToolCallsResult(modelConfig, accumulatedToolCalls)
// Send [DONE] marker for OpenAI compatibility
endOfStream := "[DONE]"
return sender(&endOfStream, nil)
return <-streamErr
}
// Embed embeds a list of texts into embeddings

295
internal/entity/models/mistral.go
View File

@@ -33,7 +33,76 @@ import (
"time"
)
// MistralModel implements ModelDriver for Mistral AI.
// normalizeMistralStructuredContent rewrites a Mistral magistral response
// whose message.content is a structured array ([{type:text},{type:thinking,
// thinking:[{type:text}]}]) into the flat string shape the shared handler
// expects: content becomes the concatenated text parts and a top-level
// reasoning_content carries the thinking parts. The rewrite is in-place and
// only applied when content is actually an array.
func normalizeMistralStructuredContent(body []byte) []byte {
var result map[string]any
if err := json.Unmarshal(body, &result); err != nil {
return body
}
choices, ok := result["choices"].([]any)
if !ok || len(choices) == 0 {
return body
}
firstChoice, ok := choices[0].(map[string]any)
if !ok {
return body
}
messageMap, ok := firstChoice["message"].(map[string]any)
if !ok {
return body
}
parts, ok := messageMap["content"].([]any)
if !ok {
return body
}
var answer, reasoning strings.Builder
for _, p := range parts {
part, ok := p.(map[string]any)
if !ok {
continue
}
switch part["type"] {
case "text":
if t, ok := part["text"].(string); ok {
answer.WriteString(t)
}
case "thinking":
if thinking, ok := part["thinking"].([]any); ok {
for _, tp := range thinking {
if tpm, ok := tp.(map[string]any); ok {
if t, ok := tpm["text"].(string); ok {
reasoning.WriteString(t)
}
}
}
}
}
}
// Only rewrite if we actually extracted something; otherwise leave
// the body untouched so the shared handler surfaces its normal error.
if answer.Len() == 0 && reasoning.Len() == 0 {
return body
}
messageMap["content"] = answer.String()
if existing, ok := messageMap["reasoning_content"].(string); ok && existing != "" {
reasoning.WriteString(existing)
}
if reasoning.Len() > 0 {
messageMap["reasoning_content"] = reasoning.String()
}
out, err := json.Marshal(result)
if err != nil {
return body
}
return out
}
type MistralModel struct {
baseModel BaseModel
@@ -58,33 +127,6 @@ func (m *MistralModel) Name() string {
return "mistral"
}
type MistralChatResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []struct {
Index int `json:"index"`
Message struct {
Role string `json:"role"`
Content interface{} `json:"content"`
ToolCalls []map[string]any `json:"tool_calls"`
ReasoningContent *string `json:"reasoning_content"`
} `json:"message"`
FinishReason string `json:"finish_reason"`
} `json:"choices"`
Usage struct {
CompletionTokens int `json:"completion_tokens"`
NumCachedTokens int `json:"num_cached_tokens"`
PromptAudioSeconds int `json:"prompt_audio_seconds"`
PromptTokenDetails struct {
CachedTokens int `json:"cached_tokens"`
}
PromptTokens int `json:"prompt_tokens"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
}
// ChatWithMessages sends multiple messages with roles and returns the response.
func (m *MistralModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
if err := m.baseModel.APIConfigCheck(apiConfig); err != nil {
@@ -99,122 +141,19 @@ func (m *MistralModel) ChatWithMessages(ctx context.Context, modelName string, m
if err != nil {
return nil, err
}
baseURL = strings.TrimSuffix(baseURL, "/")
url := fmt.Sprintf("%s/%s", baseURL, m.baseModel.URLSuffix.Chat)
reqBody := buildRequestBody(chatModelConfig, modelName, messages, false)
jsonData, err := json.Marshal(reqBody)
body, err := m.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
// Mistral magistral returns content as a structured array. Normalize it
// to the flat string shape the shared handler understands.
body = normalizeMistralStructuredContent(body)
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := m.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, _ *ChatConfig) (chatResponseParts, error) {
var result MistralChatResponse
if err := json.Unmarshal(body, &result); err != nil {
return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err)
}
if len(result.Choices) == 0 {
return chatResponseParts{}, fmt.Errorf("no choices in response")
}
choice := &result.Choices[0]
content, reasonContent, err := extractMistralContent(choice.Message.Content)
if err != nil {
return chatResponseParts{}, err
}
if reasonContent == "" && choice.Message.ReasoningContent != nil {
reasonContent = *choice.Message.ReasoningContent
}
totalTokens := result.Usage.TotalTokens
if totalTokens == 0 {
totalTokens = result.Usage.PromptTokens + result.Usage.CompletionTokens
}
var usage *TokenUsage
if totalTokens > 0 {
usage = &TokenUsage{
PromptTokens: result.Usage.PromptTokens,
CompletionTokens: result.Usage.CompletionTokens,
TotalTokens: totalTokens,
}
}
return chatResponseParts{
RequestID: result.ID,
Content: &content,
ReasonContent: &reasonContent,
ToolCalls: choice.Message.ToolCalls,
Usage: usage,
}, nil
})
}
func extractMistralContent(raw interface{}) (string, string, error) {
switch v := raw.(type) {
case string:
return v, "", nil
case []interface{}:
var answer, reasoning strings.Builder
for _, part := range v {
pm, ok := part.(map[string]interface{})
if !ok {
continue
}
switch pm["type"] {
case "text":
if t, ok := pm["text"].(string); ok {
answer.WriteString(t)
}
case "thinking":
// thinking is an array of inner text parts; concatenate
// any inner element with a non-empty text field.
inner, ok := pm["thinking"].([]interface{})
if !ok {
continue
}
for _, sub := range inner {
sm, ok := sub.(map[string]interface{})
if !ok {
continue
}
if t, ok := sm["text"].(string); ok {
reasoning.WriteString(t)
}
}
}
}
return answer.String(), reasoning.String(), nil
case nil:
return "", "", nil
default:
return "", "", fmt.Errorf("mistral: unsupported content type %T", raw)
}
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams the response
@@ -238,89 +177,15 @@ func (m *MistralModel) ChatStreamlyWithSender(ctx context.Context, modelName str
if err != nil {
return err
}
baseURL = strings.TrimSuffix(baseURL, "/")
url := fmt.Sprintf("%s/%s", baseURL, m.baseModel.URLSuffix.Chat)
reqBody := buildRequestBody(chatModelConfig, modelName, messages, true)
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
reqBody["stream_options"] = map[string]interface{}{
"include_usage": true,
}
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := m.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
sawTerminal := false
accumulatedToolCalls := make(map[int]map[string]any)
done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event)
if usageErr != nil {
return usageErr
}
if found {
applyStreamUsage(chatModelConfig, modelUsage, tokenUsage)
}
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
return nil
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
content, ok := delta["content"].(string)
if ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
finishReason, ok := firstChoice["finish_reason"].(string)
if ok && finishReason != "" {
sawTerminal = true
}
return nil
return m.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
})
if err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls)
if !done && !sawTerminal {
return fmt.Errorf("mistral: stream ended before [DONE] or finish_reason")
}
endOfStream := "[DONE]"
if err := sender(&endOfStream, nil); err != nil {
return err
}
return nil
}
type mistralEmbeddingData struct {

60
internal/entity/models/mistral_test.go
View File

@@ -914,63 +914,3 @@ func TestMistralChatIgnoresUnknownContentPartTypes(t *testing.T) {
t.Errorf("Answer=%q want %q", *resp.Answer, "Hello")
}
}
// Direct unit coverage of the helper, including the nil and bad-type
// edge cases that won't surface in the integration tests above.
func TestExtractMistralContent(t *testing.T) {
tests := []struct {
name string
input interface{}
wantAns string
wantReason string
wantErr bool
}{
{"plain string", "hi", "hi", "", false},
{"empty string", "", "", "", false},
{"nil", nil, "", "", false},
{"empty array", []interface{}{}, "", "", false},
{
"text only",
[]interface{}{
map[string]interface{}{"type": "text", "text": "a"},
map[string]interface{}{"type": "text", "text": "b"},
},
"ab", "", false,
},
{
"thinking then text",
[]interface{}{
map[string]interface{}{
"type": "thinking",
"thinking": []interface{}{
map[string]interface{}{"type": "text", "text": "why "},
map[string]interface{}{"type": "text", "text": "this"},
},
},
map[string]interface{}{"type": "text", "text": "answer"},
},
"answer", "why this", false,
},
{"unknown root type", 42, "", "", true},
}
for _, tc := range tests {
t.Run(tc.name, func(t *testing.T) {
ans, reason, err := extractMistralContent(tc.input)
if tc.wantErr {
if err == nil {
t.Errorf("want error, got nil")
}
return
}
if err != nil {
t.Errorf("unexpected err: %v", err)
}
if ans != tc.wantAns {
t.Errorf("answer=%q want %q", ans, tc.wantAns)
}
if reason != tc.wantReason {
t.Errorf("reason=%q want %q", reason, tc.wantReason)
}
})
}
}

196
internal/entity/models/modelscope.go
View File

@@ -17,7 +17,6 @@
package models
import (
"bytes"
"context"
"encoding/json"
"fmt"
@@ -25,31 +24,13 @@ import (
"net/http"
"ragflow/internal/common"
"strings"
"sync"
"time"
)
// modelscopeStreamIdleTimeout bounds how long a stream can go without
var modelscopeStreamIdleTimeout = 60 * time.Second
// ModelScopeModel implements ModelDriver for ModelScope chat models.
type ModelScopeModel struct {
baseModel BaseModel
}
type modelscopeChatChoice struct {
Message struct {
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
Reasoning string `json:"reasoning"`
Thinking string `json:"thinking"`
} `json:"message"`
}
type modelscopeChatResponse struct {
Choices []modelscopeChatChoice `json:"choices"`
}
type modelscopeModelListResponse struct {
Data []ModelListItem `json:"data"`
}
@@ -85,28 +66,6 @@ func normalizeModelScopeBaseURL(base string) string {
return trimmed
}
func modelscopeReasoningFromStrings(reasoningContent string, reasoning string, thinking string) string {
switch {
case reasoningContent != "":
return reasoningContent
case reasoning != "":
return reasoning
case thinking != "":
return thinking
default:
return ""
}
}
func modelscopeReasoningFromMap(value map[string]interface{}) string {
for _, field := range []string{"reasoning_content", "reasoning", "thinking"} {
if text, ok := value[field].(string); ok && text != "" {
return text
}
}
return ""
}
// ChatWithMessages sends multiple messages with roles and returns the response.
func (m *ModelScopeModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
if err := m.baseModel.APIConfigCheck(apiConfig); err != nil {
@@ -125,56 +84,13 @@ func (m *ModelScopeModel) ChatWithMessages(ctx context.Context, modelName string
url := fmt.Sprintf("%s/%s", baseURL, m.baseModel.URLSuffix.Chat)
reqBody := buildRequestBody(chatModelConfig, modelName, messages, false)
jsonData, err := json.Marshal(reqBody)
body, err := m.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
if auth := BearerAuth(apiConfig); auth != "" {
req.Header.Set("Authorization", auth)
}
resp, err := m.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
var result modelscopeChatResponse
if err = json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
if len(result.Choices) == 0 {
return nil, fmt.Errorf("no choices in response")
}
content := result.Choices[0].Message.Content
reasonContent := modelscopeReasoningFromStrings(
result.Choices[0].Message.ReasoningContent,
result.Choices[0].Message.Reasoning,
result.Choices[0].Message.Thinking,
)
return &ChatResponse{
Answer: &content,
ReasonContent: &reasonContent,
}, nil
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams response via sender.
@@ -189,8 +105,8 @@ func (m *ModelScopeModel) ChatStreamlyWithSender(ctx context.Context, modelName
if len(messages) == 0 {
return fmt.Errorf("messages is empty")
}
if chatModelConfig != nil && chatModelConfig.Stream != nil && !*chatModelConfig.Stream {
return fmt.Errorf("stream must be true in ChatStreamlyWithSender")
if err := validateStreamConfig(chatModelConfig); err != nil {
return err
}
baseURL, err := m.baseModel.GetBaseURL(apiConfig)
@@ -201,105 +117,11 @@ func (m *ModelScopeModel) ChatStreamlyWithSender(ctx context.Context, modelName
url := fmt.Sprintf("%s/%s", baseURL, m.baseModel.URLSuffix.Chat)
reqBody := buildRequestBody(chatModelConfig, modelName, messages, true)
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
}
reqBody["stream_options"] = map[string]interface{}{"include_usage": true}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
if auth := BearerAuth(apiConfig); auth != "" {
req.Header.Set("Authorization", auth)
}
resp, err := m.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
lastActive := time.Now()
var lastActiveMu sync.Mutex
done := make(chan struct{})
defer close(done)
go func() {
ticker := time.NewTicker(modelscopeStreamIdleTimeout / 4)
defer ticker.Stop()
for {
select {
case <-done:
return
case now := <-ticker.C:
lastActiveMu.Lock()
idle := now.Sub(lastActive)
lastActiveMu.Unlock()
if idle >= modelscopeStreamIdleTimeout {
cancel()
return
}
}
}
}()
sawTerminal := false
accumulatedToolCalls := make(map[int]map[string]any)
streamDone, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
lastActiveMu.Lock()
lastActive = time.Now()
lastActiveMu.Unlock()
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
if delta, ok := firstChoice["delta"].(map[string]interface{}); ok {
accumulateToolCallDeltas(delta, accumulatedToolCalls)
if reasoning := modelscopeReasoningFromMap(delta); reasoning != "" {
if err := sender(nil, &reasoning); err != nil {
return err
}
}
if content, ok := delta["content"].(string); ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
}
if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" {
sawTerminal = true
}
return nil
return m.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
})
if err != nil {
if ctx.Err() != nil {
return fmt.Errorf("modelscope: stream idle for more than %s, aborted", modelscopeStreamIdleTimeout)
}
return fmt.Errorf("failed to scan response body: %w", err)
}
if !streamDone && !sawTerminal {
return fmt.Errorf("modelscope: stream ended before [DONE] or finish_reason")
}
setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls)
endOfStream := "[DONE]"
return sender(&endOfStream, nil)
}
func (m *ModelScopeModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) {

41
internal/entity/models/modelscope_test.go
View File

@@ -23,7 +23,6 @@ import (
"net/http/httptest"
"strings"
"testing"
"time"
)
func newModelScopeForTest(baseURL string) *ModelScopeModel {
@@ -36,15 +35,6 @@ func newModelScopeForTest(baseURL string) *ModelScopeModel {
)
}
func withModelScopeIdleTimeout(t *testing.T, d time.Duration) {
t.Helper()
original := modelscopeStreamIdleTimeout
modelscopeStreamIdleTimeout = d
t.Cleanup(func() {
modelscopeStreamIdleTimeout = original
})
}
func TestModelScopeName(t *testing.T) {
m := newModelScopeForTest("http://unused")
if got := m.Name(); got != "ModelScope" {
@@ -271,37 +261,6 @@ func TestModelScopeStreamRejectsFalseStreamConfig(t *testing.T) {
}
}
func TestModelScopeStreamCancelsOnIdle(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()
withModelScopeIdleTimeout(t, 200*time.Millisecond)
hold := make(chan struct{})
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
w.Header().Set("Content-Type", "text/event-stream")
w.WriteHeader(http.StatusOK)
if f, ok := w.(http.Flusher); ok {
_, _ = io.WriteString(w, `data: {"choices":[{"delta":{"content":"hi"}}]}`+"\n")
f.Flush()
}
select {
case <-hold:
case <-r.Context().Done():
}
}))
t.Cleanup(srv.Close)
t.Cleanup(func() { close(hold) })
m := newModelScopeForTest(srv.URL)
err := m.ChatStreamlyWithSender(ctx, "Qwen/Qwen2.5-7B-Instruct",
[]Message{{Role: "user", Content: "x"}},
&APIConfig{}, nil, nil,
func(*string, *string) error { return nil })
if err == nil || !strings.Contains(err.Error(), "stream idle") {
t.Errorf("expected stream-idle error, got %v", err)
}
}
func TestModelScopeListModelsAndCheckConnection(t *testing.T) {
withSSRFBypass(t)
ctx := t.Context()

149
internal/entity/models/moonshot.go
View File

@@ -85,7 +85,6 @@ func (m *MoonshotModel) ChatWithMessages(ctx context.Context, modelName string,
if err := m.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
apiKey := strings.TrimSpace(*apiConfig.ApiKey)
modelName, err := validateMoonshotModelName(modelName)
if err != nil {
return nil, err
@@ -116,84 +115,12 @@ func (m *MoonshotModel) ChatWithMessages(ctx context.Context, modelName string,
}
jsonData, err := json.Marshal(reqBody)
body, err := m.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", apiKey))
resp, err := m.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
var result MoonshotChatResponse
if err := json.Unmarshal(body, &result); err != nil {
return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err)
}
if len(result.Choices) == 0 {
return chatResponseParts{}, fmt.Errorf("no choices in response")
}
choice := &result.Choices[0]
content := ""
if choice.Message.Content != nil {
content = *choice.Message.Content
}
reasonContent := ""
if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking {
if choice.Message.ReasoningContent != nil {
reasonContent = *choice.Message.ReasoningContent
}
if reasonContent != "" && reasonContent[0] == '\n' {
reasonContent = reasonContent[1:]
}
}
totalTokens := result.Usage.TotalTokens
if totalTokens == 0 {
totalTokens = result.Usage.PromptTokens + result.Usage.CompletionTokens
}
var usage *TokenUsage
if totalTokens > 0 {
usage = &TokenUsage{
PromptTokens: result.Usage.PromptTokens,
CompletionTokens: result.Usage.CompletionTokens,
TotalTokens: totalTokens,
}
}
return chatResponseParts{
RequestID: result.ID,
Content: &content,
ReasonContent: &reasonContent,
ToolCalls: choice.Message.ToolCalls,
Usage: usage,
}, nil
})
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams response via sender function (best performance, no channel)
@@ -262,75 +189,7 @@ func (m *MoonshotModel) ChatStreamlyWithSender(ctx context.Context, modelName st
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// SSE parsing: read line by line
sawTerminal := false
if chatModelConfig != nil {
chatModelConfig.ToolCallsResult = nil
}
accumulatedToolCalls := make(map[int]map[string]any)
done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event)
if usageErr != nil {
return usageErr
}
if found {
applyStreamUsage(chatModelConfig, modelUsage, tokenUsage)
}
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
return nil
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
reasoningContent, ok := delta["reasoning_content"].(string)
if ok && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
return err
}
}
content, ok := delta["content"].(string)
if ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
finishReason, ok := firstChoice["finish_reason"].(string)
if ok && finishReason != "" {
sawTerminal = true
}
return nil
})
if err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
if !done && !sawTerminal {
return fmt.Errorf("moonshot: stream ended before [DONE] or finish_reason")
}
setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls)
// Send [DONE] marker for OpenAI compatibility
endOfStream := "[DONE]"
if err = sender(&endOfStream, nil); err != nil {
return err
}
return nil
return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
}
// Embed embeds a list of texts into embeddings

145
internal/entity/models/n1n.go
View File

@@ -88,33 +88,12 @@ func newN1NJSONRequest(ctx context.Context, method, endpoint string, payload int
return req, nil
}
type n1nChatChoice struct {
Message n1nChatMessage `json:"message"`
Delta n1nChatDelta `json:"delta"`
FinishReason string `json:"finish_reason"`
}
type n1nChatMessage struct {
Content *string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
}
type n1nChatDelta struct {
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
}
type n1nChatResponse struct {
Choices []n1nChatChoice `json:"choices"`
}
// ChatWithMessages sends a single, non-streaming chat completion
// against n1n.ai's /v1/chat/completions endpoint.
func (n *N1NModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
if err := n.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
apiKey := *apiConfig.ApiKey
if strings.TrimSpace(modelName) == "" {
return nil, fmt.Errorf("model name is required")
}
@@ -136,49 +115,12 @@ func (n *N1NModel) ChatWithMessages(ctx context.Context, modelName string, messa
reqBody["thinking"] = map[string]interface{}{"type": thinkingType}
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := newN1NJSONRequest(ctx, "POST", endpoint, reqBody, apiKey)
body, err := n.baseModel.doRequest(ctx, endpoint, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, err
}
resp, err := n.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("n1n chat API error: %s, body: %s", resp.Status, string(body))
}
var parsed n1nChatResponse
if err := json.Unmarshal(body, &parsed); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
if len(parsed.Choices) == 0 {
return nil, fmt.Errorf("no choices in response")
}
if parsed.Choices[0].Message.Content == nil {
return nil, fmt.Errorf("invalid content format")
}
content := *parsed.Choices[0].Message.Content
chatResp := &ChatResponse{
Answer: &content,
}
if parsed.Choices[0].Message.ReasoningContent != "" {
reasonContent := parsed.Choices[0].Message.ReasoningContent
chatResp.ReasonContent = &reasonContent
}
return chatResp, nil
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends a streaming chat completion.
@@ -196,7 +138,6 @@ func (n *N1NModel) ChatStreamlyWithSender(ctx context.Context, modelName string,
if len(messages) == 0 {
return fmt.Errorf("messages is empty")
}
apiKey := *apiConfig.ApiKey
endpoint, err := n.endpointURL(n1nRegion(apiConfig), n.baseModel.URLSuffix.Chat)
if err != nil {
@@ -208,6 +149,7 @@ func (n *N1NModel) ChatStreamlyWithSender(ctx context.Context, modelName string,
}
reqBody := buildRequestBody(chatModelConfig, modelName, messages, true)
reqBody["stream_options"] = map[string]interface{}{"include_usage": true}
if chatModelConfig != nil && chatModelConfig.Thinking != nil {
thinkingType := "disabled"
if *chatModelConfig.Thinking {
@@ -216,86 +158,9 @@ func (n *N1NModel) ChatStreamlyWithSender(ctx context.Context, modelName string,
reqBody["thinking"] = map[string]interface{}{"type": thinkingType}
}
req, err := newN1NJSONRequest(ctx, "POST", endpoint, reqBody, apiKey)
if err != nil {
return err
}
resp, err := n.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("n1n chat stream API error: %s, body: %s", resp.Status, string(body))
}
sawTerminal := false
accumulatedToolCalls := make(map[int]map[string]interface{})
done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
common.Info(fmt.Sprintf("%v", event))
tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event)
if usageErr != nil {
return usageErr
}
if found {
applyStreamUsage(chatModelConfig, modelUsage, tokenUsage)
}
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
return nil
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
content, ok := delta["content"].(string)
if ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
reasoningContent, ok := delta["reasoning_content"].(string)
if ok && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
return err
}
}
finishReason, ok := firstChoice["finish_reason"].(string)
if ok && finishReason != "" {
sawTerminal = true
}
return nil
return n.baseModel.doStreamRequest(ctx, endpoint, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
})
if err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
if !done && !sawTerminal {
return fmt.Errorf("deepseek: stream ended before [DONE] or finish_reason")
}
setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls)
endOfStream := "[DONE]"
if err := sender(&endOfStream, nil); err != nil {
return err
}
return nil
}
type n1nEmbeddingData struct {

436
internal/entity/models/novita.go
View File

@@ -52,158 +52,6 @@ func (n *NovitaModel) Name() string {
return "NovitaAI"
}
type NovitaChatResponse struct {
ID string `json:"id"`
Choices []struct {
FinishReason string `json:"finish_reason"`
Index int `json:"index"`
Message struct {
Content string `json:"content"`
ReasoningContent string `json:"reasoning_content"`
Role string `json:"role"`
ToolCalls []map[string]any `json:"tool_calls"`
} `json:"message"`
Logprobs interface{} `json:"logprobs"`
} `json:"choices"`
Created int `json:"created"`
Model string `json:"model"`
Object string `json:"object"`
Usage struct {
CompletionTokens int `json:"completion_tokens"`
PromptTokens int `json:"prompt_tokens"`
TotalTokens int `json:"total_tokens"`
} `json:"usage"`
}
const (
novitaThinkOpen = "<think>"
novitaThinkClose = "</think>"
)
// splitNovitaThink walks a complete content string and returns the
// visible portion + the concatenated chain-of-thought from inside
// any <think>...</think> blocks. Multiple think blocks are
// concatenated; tags themselves are stripped. Used by the
// non-streaming path where the whole content is available at once.
func splitNovitaThink(raw string) (visible, reasoning string) {
var v, r strings.Builder
inside := false
for {
var marker string
if inside {
marker = novitaThinkClose
} else {
marker = novitaThinkOpen
}
idx := strings.Index(raw, marker)
if idx < 0 {
if inside {
r.WriteString(raw)
} else {
v.WriteString(raw)
}
break
}
if inside {
r.WriteString(raw[:idx])
} else {
v.WriteString(raw[:idx])
}
raw = raw[idx+len(marker):]
inside = !inside
}
return v.String(), r.String()
}
// novitaThinkExtractor maintains state across streaming chunks so
// that a <think>...</think> block spanning multiple SSE events still
// gets split correctly between content and reasoning. The buffer
// preserves up to (len(closingMarker)-1) trailing bytes of each
// chunk in case the next chunk completes a partial tag.
type novitaThinkExtractor struct {
buf strings.Builder
inside bool
}
// novitaThinkSegment is one routing decision: emit `content` via the
// sender's first arg, or emit `reasoning` via the sender's second arg.
// Exactly one of the two fields is non-empty.
type novitaThinkSegment struct {
content string
reasoning string
}
// Feed appends an incoming chunk and returns any segments that are
// now safe to emit. Trailing bytes that could be the start of a tag
// are held back in the buffer until the next call.
func (e *novitaThinkExtractor) Feed(chunk string) []novitaThinkSegment {
e.buf.WriteString(chunk)
s := e.buf.String()
var out []novitaThinkSegment
for {
var marker, otherMarker string
if e.inside {
marker = novitaThinkClose
otherMarker = novitaThinkOpen
} else {
marker = novitaThinkOpen
otherMarker = novitaThinkClose
}
idx := strings.Index(s, marker)
if idx < 0 {
// No closing/opening marker yet. Emit everything except a
// possible partial-tag suffix at the very end. Reserve
// (max marker length - 1) trailing bytes so we don't
// emit "<thin" as content when the next chunk completes
// it to "<think>".
reserve := max(len(otherMarker)-1, len(marker)-1)
safe := max(len(s)-reserve, 0)
// Don't reserve if the trailing bytes can't possibly be
// the start of a tag (no '<' suffix).
if safe < len(s) && !strings.Contains(s[safe:], "<") {
safe = len(s)
}
if safe > 0 {
if e.inside {
out = append(out, novitaThinkSegment{reasoning: s[:safe]})
} else {
out = append(out, novitaThinkSegment{content: s[:safe]})
}
s = s[safe:]
}
break
}
if idx > 0 {
if e.inside {
out = append(out, novitaThinkSegment{reasoning: s[:idx]})
} else {
out = append(out, novitaThinkSegment{content: s[:idx]})
}
}
s = s[idx+len(marker):]
e.inside = !e.inside
}
e.buf.Reset()
e.buf.WriteString(s)
return out
}
// Flush returns the buffered tail when the stream ends. A stream that
// ends mid-tag would not normally happen with a well-behaved upstream,
// but if it does the partial bytes are emitted according to the
// current mode so nothing is silently lost.
func (e *novitaThinkExtractor) Flush() *novitaThinkSegment {
s := e.buf.String()
e.buf.Reset()
if s == "" {
return nil
}
if e.inside {
return &novitaThinkSegment{reasoning: s}
}
return &novitaThinkSegment{content: s}
}
// ChatWithMessages sends multiple messages with roles and returns the response.
func (n *NovitaModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
if err := n.baseModel.APIConfigCheck(apiConfig); err != nil {
@@ -227,91 +75,20 @@ func (n *NovitaModel) ChatWithMessages(ctx context.Context, modelName string, me
}
}
jsonData, err := json.Marshal(reqBody)
body, err := n.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := n.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
var result NovitaChatResponse
if err := json.Unmarshal(body, &result); err != nil {
return chatResponseParts{}, fmt.Errorf("failed to parse response: %w", err)
}
if len(result.Choices) == 0 {
return chatResponseParts{}, fmt.Errorf("no choices in response")
}
choice := &result.Choices[0]
if choice.Message.Content == "" && len(choice.Message.ToolCalls) == 0 {
return chatResponseParts{}, fmt.Errorf("invalid content format")
}
// Novita emits chain-of-thought in two different shapes depending
// on the model and on enable_thinking:
// - qwen3-* and other inline-style models: chain-of-thought is
// embedded inside content as <think>...</think> tags.
// - deepseek-v3.1 / glm-4.5 (and any model with separate
// reasoning enabled): chain-of-thought arrives in a separate
// `reasoning_content` field, with `content` already cleaned.
// Handle both so the visible Answer is always tag-free and any
// reasoning the upstream supplied is preserved.
visible, reasoning := splitNovitaThink(choice.Message.Content)
if choice.Message.ReasoningContent != "" {
if reasoning != "" {
reasoning += "\n" + choice.Message.ReasoningContent
} else {
reasoning = choice.Message.ReasoningContent
}
}
return chatResponseParts{
RequestID: result.ID,
Content: &visible,
ReasonContent: &reasoning,
ToolCalls: choice.Message.ToolCalls,
Usage: &TokenUsage{
PromptTokens: result.Usage.PromptTokens,
CompletionTokens: result.Usage.CompletionTokens,
TotalTokens: result.Usage.TotalTokens,
},
}, nil
})
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams the response via
// the sender. Handles both reasoning shapes Novita can emit:
// - delta.reasoning_content (deepseek-v3.1 / glm-4.5 / any model
// with separate reasoning): forwarded as-is to the second arg.
// - delta.content containing <think>...</think> (qwen3-* and other
// inline-style models): a stateful extractor splits tag bytes
// across SSE chunk boundaries, then routes content/reasoning to
// the first/second sender arg respectively.
// - delta.content (qwen3-* and other inline-style models): forwarded
// to the first arg.
func (n *NovitaModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
if err := n.baseModel.APIConfigCheck(apiConfig); err != nil {
return err
@@ -341,117 +118,186 @@ func (n *NovitaModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
}
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
}
reqBody["stream_options"] = map[string]interface{}{"include_usage": true}
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
return n.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
// Novita qwen3 embeds <think>...</think> inline in
// delta.content (tags can span multiple SSE deltas). Split
// those blocks so reasoning routes to the sender's second arg.
return novitaHandleStream(body, modelUsage, chatModelConfig, sender)
})
}
resp, err := n.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
// novitaThinkSegment is one routing decision: emit `content` via the
// sender's first arg, or emit `reasoning` via the second. Exactly one of
// the two fields is non-empty.
type novitaThinkSegment struct {
content string
reasoning string
}
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// novitaThinkSplitter holds state across streaming content chunks so a
// <think>...</think> block that spans multiple SSE deltas is still split
// correctly. Trailing bytes that could be the start of a tag are held
// back until the next chunk.
type novitaThinkSplitter struct {
buf strings.Builder
inside bool
}
extractor := &novitaThinkExtractor{}
sawTerminal := false
done, err := ParseSSEStream[map[string]interface{}](resp.Body, func(event map[string]interface{}) error {
tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event)
if usageErr != nil {
return usageErr
const (
novitaThinkOpen = "<think>"
novitaThinkClose = "</think>"
)
func (s *novitaThinkSplitter) feed(chunk string) []novitaThinkSegment {
s.buf.WriteString(chunk)
str := s.buf.String()
var out []novitaThinkSegment
for {
var marker string
if s.inside {
marker = novitaThinkClose
} else {
marker = novitaThinkOpen
}
if found {
applyStreamUsage(chatModelConfig, modelUsage, tokenUsage)
idx := strings.Index(str, marker)
if idx < 0 {
// No marker yet. Emit everything except a possible
// partial-tag suffix at the very end.
reserve := max(len(novitaThinkOpen)-1, len(novitaThinkClose)-1)
safe := max(len(str)-reserve, 0)
if safe < len(str) && !strings.Contains(str[safe:], "<") {
safe = len(str)
}
if safe > 0 {
if s.inside {
out = append(out, novitaThinkSegment{reasoning: str[:safe]})
} else {
out = append(out, novitaThinkSegment{content: str[:safe]})
}
str = str[safe:]
}
s.buf.Reset()
s.buf.WriteString(str)
return out
}
if s.inside {
out = append(out, novitaThinkSegment{reasoning: str[:idx]})
} else {
out = append(out, novitaThinkSegment{content: str[:idx]})
}
str = str[idx+len(marker):]
s.inside = !s.inside
}
}
func (s *novitaThinkSplitter) flush() *novitaThinkSegment {
if s.buf.Len() == 0 {
return nil
}
remaining := s.buf.String()
s.buf.Reset()
if s.inside {
return &novitaThinkSegment{reasoning: remaining}
}
return &novitaThinkSegment{content: remaining}
}
// novitaHandleStream processes a Novita streaming chat response, splitting
// inline <think>...</think> blocks in delta.content across SSE deltas.
func novitaHandleStream(
body io.Reader,
modelUsage *common.ModelUsage,
chatConfig *ChatConfig,
sender func(*string, *string) error,
) error {
if sender == nil {
return fmt.Errorf("sender is required")
}
var sawTerminal bool
thinkSplitter := &novitaThinkSplitter{}
done, err := ParseSSEStream[map[string]any](body, func(event map[string]any) error {
tokenUsage, found := extractOpenAIStreamUsage(event)
if found && chatConfig != nil {
applyStreamUsage(chatConfig, modelUsage, tokenUsage)
}
choices, ok := event["choices"].([]interface{})
if apiErr, ok := event["error"]; ok && apiErr != nil {
return fmt.Errorf("upstream stream error: %v", apiErr)
}
choices, ok := event["choices"].([]any)
if !ok || len(choices) == 0 {
return nil
}
firstChoice, ok := choices[0].(map[string]interface{})
firstChoice, ok := choices[0].(map[string]any)
if !ok {
return nil
}
delta, ok := firstChoice["delta"].(map[string]interface{})
delta, ok := firstChoice["delta"].(map[string]any)
if !ok {
return nil
}
// deepseek-v3.1 / glm-4.5 (and other models that emit reasoning
// separately) put chain-of-thought in delta.reasoning_content
// rather than inside content as <think>...</think>. Surface it
// before any content from the same chunk so callers piping to
// a UI render reasoning before the visible answer for that
// token, matching the wire ordering Novita emits.
if r, ok := delta["reasoning_content"].(string); ok && r != "" {
rr := r
if err := sender(nil, &rr); err != nil {
if reasoningContent, ok := delta["reasoning_content"].(string); ok && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
return err
}
}
if c, ok := delta["content"].(string); ok && c != "" {
for _, seg := range extractor.Feed(c) {
if content, ok := delta["content"].(string); ok && content != "" {
for _, seg := range thinkSplitter.feed(content) {
if seg.content != "" {
c := seg.content
if err := sender(&c, nil); err != nil {
return err
}
}
if seg.reasoning != "" {
r := seg.reasoning
if err := sender(nil, &r); err != nil {
return err
}
}
if seg.content != "" {
cc := seg.content
if err := sender(&cc, nil); err != nil {
return err
}
}
}
}
if finish, ok := firstChoice["finish_reason"].(string); ok && finish != "" {
if finishReason, ok := firstChoice["finish_reason"].(string); ok && finishReason != "" {
sawTerminal = true
}
return nil
})
if err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
// Flush any buffered tail (rare, but covers the case where the
// stream ends right after the last chunk without us seeing the
// closing tag).
if seg := extractor.Flush(); seg != nil {
if !done && !sawTerminal {
return fmt.Errorf("stream ended before [DONE] or finish_reason")
}
if seg := thinkSplitter.flush(); seg != nil {
if seg.content != "" {
c := seg.content
if err := sender(&c, nil); err != nil {
return err
}
}
if seg.reasoning != "" {
r := seg.reasoning
if err := sender(nil, &r); err != nil {
return err
}
}
if seg.content != "" {
cc := seg.content
if err := sender(&cc, nil); err != nil {
return err
}
}
}
if !done && !sawTerminal {
return fmt.Errorf("novita: stream ended before [DONE] or finish_reason")
}
endOfStream := "[DONE]"
if err := sender(&endOfStream, nil); err != nil {
return err
}
return nil
return sender(&endOfStream, nil)
}
// ListModels returns the list of model ids visible to the API key.

145
internal/entity/models/novita_test.go
View File

@@ -101,151 +101,6 @@ func newNovitaSSEServer(t *testing.T, expectedPath, ssePayload string) *httptest
}))
}
// ---- think-tag split helpers ----
func TestSplitNovitaThinkPureText(t *testing.T) {
v, r := splitNovitaThink("hello world")
if v != "hello world" || r != "" {
t.Errorf("got (%q,%q)", v, r)
}
}
func TestSplitNovitaThinkSingleBlock(t *testing.T) {
v, r := splitNovitaThink("<think>15% = 0.15. 0.15*80 = 12.</think>The answer is 12.")
if v != "The answer is 12." {
t.Errorf("visible=%q", v)
}
if r != "15% = 0.15. 0.15*80 = 12." {
t.Errorf("reasoning=%q", r)
}
}
func TestSplitNovitaThinkLeadingText(t *testing.T) {
v, r := splitNovitaThink("intro <think>thought</think>tail")
if v != "intro tail" {
t.Errorf("visible=%q", v)
}
if r != "thought" {
t.Errorf("reasoning=%q", r)
}
}
func TestSplitNovitaThinkMultipleBlocks(t *testing.T) {
v, r := splitNovitaThink("<think>A</think>part1<think>B</think>part2")
if v != "part1part2" {
t.Errorf("visible=%q", v)
}
if r != "AB" {
t.Errorf("reasoning=%q", r)
}
}
func TestSplitNovitaThinkUnclosedTag(t *testing.T) {
// Unclosed <think> -> everything after the open tag is reasoning;
// content stops at the open tag. This matches a real upstream that
// got cut off mid-reasoning by max_tokens.
v, r := splitNovitaThink("visible <think>still thinking when tokens ran out")
if v != "visible " {
t.Errorf("visible=%q", v)
}
if r != "still thinking when tokens ran out" {
t.Errorf("reasoning=%q", r)
}
}
// ---- streaming extractor ----
// Helper to push multiple chunks through and concatenate all output by
// kind. Each chunk goes into Feed; the output is what's safe to emit.
func feedAll(e *novitaThinkExtractor, chunks []string) (content, reasoning string) {
var cb, rb strings.Builder
for _, c := range chunks {
for _, seg := range e.Feed(c) {
cb.WriteString(seg.content)
rb.WriteString(seg.reasoning)
}
}
if seg := e.Flush(); seg != nil {
cb.WriteString(seg.content)
rb.WriteString(seg.reasoning)
}
return cb.String(), rb.String()
}
func TestNovitaThinkExtractorSingleChunk(t *testing.T) {
e := &novitaThinkExtractor{}
c, r := feedAll(e, []string{"hello <think>thought</think> world"})
if c != "hello world" {
t.Errorf("content=%q", c)
}
if r != "thought" {
t.Errorf("reasoning=%q", r)
}
}
func TestNovitaThinkExtractorTagSpansChunks(t *testing.T) {
// "<think>" split across two SSE deltas: "<thi" + "nk>"
e := &novitaThinkExtractor{}
c, r := feedAll(e, []string{"hello <thi", "nk>thought</think>tail"})
if c != "hello tail" {
t.Errorf("content=%q", c)
}
if r != "thought" {
t.Errorf("reasoning=%q", r)
}
}
func TestNovitaThinkExtractorClosingTagSpansChunks(t *testing.T) {
// "</think>" split across two deltas
e := &novitaThinkExtractor{}
c, r := feedAll(e, []string{"<think>reasoning</thi", "nk>visible"})
if c != "visible" {
t.Errorf("content=%q", c)
}
if r != "reasoning" {
t.Errorf("reasoning=%q", r)
}
}
func TestNovitaThinkExtractorTokenBoundaries(t *testing.T) {
// Simulate the kind of chunking we saw on the wire for qwen3 — many
// small chunks, sometimes splitting tag bytes.
e := &novitaThinkExtractor{}
c, r := feedAll(e, []string{
"<", "think>", "Ok", "ay, ", "compute. </", "think>", "12", "."})
if c != "12." {
t.Errorf("content=%q", c)
}
if r != "Okay, compute. " {
t.Errorf("reasoning=%q", r)
}
}
func TestNovitaThinkExtractorNoTags(t *testing.T) {
e := &novitaThinkExtractor{}
c, r := feedAll(e, []string{"plain ", "content ", "all ", "the way"})
if c != "plain content all the way" {
t.Errorf("content=%q", c)
}
if r != "" {
t.Errorf("reasoning=%q", r)
}
}
func TestNovitaThinkExtractorLessThanIsNotTagStart(t *testing.T) {
// "<10" or "<a" in legitimate content must not be held back as
// possible tag start. The extractor reserves trailing bytes only
// when a "<" is present in the suffix.
e := &novitaThinkExtractor{}
c, r := feedAll(e, []string{"a < b and c < d"})
if c != "a < b and c < d" {
t.Errorf("content=%q", c)
}
if r != "" {
t.Errorf("reasoning=%q", r)
}
}
// ---- driver methods ----
func TestNovitaName(t *testing.T) {

126
internal/entity/models/ollama.go
View File

@@ -17,7 +17,6 @@
package models
import (
"bufio"
"bytes"
"context"
"encoding/json"
@@ -85,60 +84,12 @@ func (o *OllamaModel) ChatWithMessages(ctx context.Context, modelName string, me
}
}
jsonData, err := json.Marshal(reqBody)
body, err := o.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
resp, err := o.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// Parse response
var result map[string]interface{}
if err = json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
message, ok := result["message"].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("failed to parse response: message not found")
}
content, ok := message["content"].(string)
if !ok {
return nil, fmt.Errorf("failed to parse response: content not found")
}
reasonContent, _ := message["thinking"].(string)
chatResponse := &ChatResponse{
Answer: &content,
ReasonContent: &reasonContent,
}
return chatResponse, nil
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
func (o *OllamaModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
@@ -169,74 +120,11 @@ func (o *OllamaModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
}
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
}
reqBody["stream_options"] = map[string]interface{}{"include_usage": true}
ctx, cancel := context.WithTimeout(ctx, streamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
resp, err := o.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// SSE parsing: read line by line
scanner := bufio.NewScanner(resp.Body)
scanner.Buffer(make([]byte, 64*1024), 1024*1024)
for scanner.Scan() {
line := strings.TrimSpace(scanner.Text())
// ignore the blank
if line == "" {
continue
}
// Parse the JSON event
var event map[string]interface{}
if err = json.Unmarshal([]byte(line), &event); err != nil {
continue
}
if messageMap, ok := event["message"].(map[string]interface{}); ok {
if reasoningContent, exists := messageMap["thinking"].(string); exists && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
return err
}
}
if content, exists := messageMap["content"].(string); exists && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
}
if done, ok := event["done"].(bool); ok && done {
break
}
}
// Send [DONE] marker for OpenAI compatibility with RAGFlow frontend
endOfStream := "[DONE]"
if err := sender(&endOfStream, nil); err != nil {
return err
}
return scanner.Err()
return o.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, modelConfig, OpenAIParserConfig, sender)
})
}
func (o *OllamaModel) Embed(ctx context.Context, modelName *string, texts []string, apiConfig *APIConfig, embeddingConfig *EmbeddingConfig, modelUsage *common.ModelUsage) ([]EmbeddingData, error) {

190
internal/entity/models/openai.go
View File

@@ -85,99 +85,12 @@ func (o *OpenAIModel) ChatWithMessages(ctx context.Context, modelName string, me
}
}
jsonData, err := json.Marshal(reqBody)
body, err := o.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := o.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
// Parse response
var result map[string]interface{}
if err = json.Unmarshal(body, &result); err != nil {
return nil, fmt.Errorf("failed to parse response: %w", err)
}
choices, ok := result["choices"].([]interface{})
if !ok || len(choices) == 0 {
return nil, fmt.Errorf("no choices in response")
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid choice format")
}
messageMap, ok := firstChoice["message"].(map[string]interface{})
if !ok {
return nil, fmt.Errorf("invalid message format")
}
var content string
if c, ok := messageMap["content"].(string); ok {
content = c
}
// OpenAI reasoning models (o-series and similar) return reasoning text in
// the reasoning_content field. Pass it through when present.
var reasonContent string
if rc, ok := messageMap["reasoning_content"].(string); ok {
reasonContent = rc
if reasonContent != "" && reasonContent[0] == '\n' {
reasonContent = reasonContent[1:]
}
}
var toolCalls []map[string]interface{}
if tcs, ok := messageMap["tool_calls"].([]interface{}); ok {
for _, tc := range tcs {
if tcMap, ok := tc.(map[string]interface{}); ok {
toolCalls = append(toolCalls, tcMap)
}
}
}
chatResponse := &ChatResponse{
Answer: &content,
ReasonContent: &reasonContent,
ToolCalls: toolCalls,
}
// Extract usage split (prompt/completion/total) from the raw API
// response for accurate per-call token accounting. Non-OpenAI
// providers that implement the OpenAI-compat API surface (DeepSeek,
// Moonshot, etc.) also return a "usage" key with the same shape.
if pt, ct, tt := extractUsageFromMap(result); tt > 0 {
chatResponse.Usage = &TokenUsage{
PromptTokens: pt, CompletionTokens: ct, TotalTokens: tt,
}
}
return chatResponse, nil
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams the response
@@ -242,102 +155,7 @@ func (o *OpenAIModel) ChatStreamlyWithSender(ctx context.Context, modelName stri
return fmt.Errorf("API request failed with status %d: %s", resp.StatusCode, string(body))
}
sawTerminal := false
accumulatedToolCalls := make(map[int]map[string]interface{})
// Capture the authoritative usage block from the final streaming
// chunk (when provider honours stream_options.include_usage=true).
// The last chunk in the stream carries the "usage" key alongside
// empty choices; we overwrite on every chunk so the final frame
// wins, matching Python's chat_model.py usage_from_response loop.
var streamUsage *TokenUsage
scanner := bufio.NewScanner(resp.Body)
for scanner.Scan() {
line := scanner.Text()
// SSE data line starts with "data:"
if !strings.HasPrefix(line, "data:") {
continue
}
// Extract JSON after "data:"
data := strings.TrimSpace(line[5:])
// [DONE] marks the end of the stream
if data == "[DONE]" {
sawTerminal = true
break
}
// Parse the JSON event
var event map[string]interface{}
if err = json.Unmarshal([]byte(data), &event); err != nil {
continue
}
// Extract usage from this chunk. When stream_options.include_usage
// is true, the final chunk carries the full usage breakdown at the
// top level of the event alongside (possibly empty) choices.
if pt, ct, tt := extractUsageFromMap(event); tt > 0 {
streamUsage = &TokenUsage{PromptTokens: pt, CompletionTokens: ct, TotalTokens: tt}
}
choices, ok := event["choices"].([]interface{})
if !ok || len(choices) == 0 {
continue
}
firstChoice, ok := choices[0].(map[string]interface{})
if !ok {
continue
}
delta, ok := firstChoice["delta"].(map[string]interface{})
if !ok {
continue
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
reasoningContent, ok := delta["reasoning_content"].(string)
if ok && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
return err
}
}
content, ok := delta["content"].(string)
if ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
finishReason, ok := firstChoice["finish_reason"].(string)
if ok && finishReason != "" {
sawTerminal = true
}
}
if err := scanner.Err(); err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
if !sawTerminal {
return fmt.Errorf("openai: stream ended before [DONE] or finish_reason")
}
setSortedToolCallsResult(chatModelConfig, accumulatedToolCalls)
// Populate UsageResult with the authoritative usage from the stream.
if streamUsage != nil && chatModelConfig != nil {
chatModelConfig.UsageResult = streamUsage
}
// Send the [DONE] marker for OpenAI compatibility
endOfStream := "[DONE]"
if err := sender(&endOfStream, nil); err != nil {
return err
}
return nil
return HandleStreamingResponse(resp.Body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
}
type openaiEmbeddingResponse struct {

81
internal/entity/models/openai_api_compatible.go
View File

@@ -79,6 +79,87 @@ func (m *OpenAIAPICompatibleModel) ListModels(ctx context.Context, apiConfig *AP
return filtered, nil
}
// ChatWithMessages sends multiple messages with roles and returns response
func (m *OpenAIAPICompatibleModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
if err := m.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
}
if len(messages) == 0 {
return nil, fmt.Errorf("messages is empty")
}
resolvedBaseURL, err := m.baseModel.GetBaseURL(apiConfig)
if err != nil {
return nil, err
}
url := fmt.Sprintf("%s/%s", resolvedBaseURL, m.baseModel.URLSuffix.Chat)
// Build request body
reqBody := buildRequestBody(chatModelConfig, modelName, messages, false)
if chatModelConfig != nil && chatModelConfig.Thinking != nil {
if *chatModelConfig.Thinking {
reqBody["thinking"] = map[string]interface{}{
"type": "enabled",
}
} else {
reqBody["thinking"] = map[string]interface{}{
"type": "disabled",
}
}
}
body, err := m.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, err
}
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
// ChatStreamlyWithSender sends messages and streams response via sender function
func (m *OpenAIAPICompatibleModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
if err := m.baseModel.APIConfigCheck(apiConfig); err != nil {
return err
}
if len(messages) == 0 {
return fmt.Errorf("messages is empty")
}
if sender == nil {
return fmt.Errorf("sender is required")
}
resolvedBaseURL, err := m.baseModel.GetBaseURL(apiConfig)
if err != nil {
return err
}
url := fmt.Sprintf("%s/%s", resolvedBaseURL, m.baseModel.URLSuffix.Chat)
// Build request body with streaming enabled
reqBody := buildRequestBody(chatModelConfig, modelName, messages, true)
reqBody["stream_options"] = map[string]interface{}{
"include_usage": true,
}
if chatModelConfig != nil && chatModelConfig.Thinking != nil {
if *chatModelConfig.Thinking {
reqBody["thinking"] = map[string]interface{}{
"type": "enabled",
}
} else {
reqBody["thinking"] = map[string]interface{}{
"type": "disabled",
}
}
}
return m.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
return HandleStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig, sender)
})
}
// Hint keywords for model type inference, matching Python's
// OpenAIAPICompatible class-level hint constants.
var (

184
internal/entity/models/openrouter.go
View File

@@ -54,37 +54,6 @@ func (o *OpenRouterModel) Name() string {
return "openrouter"
}
// OpenRouterChatResponse mirrors OpenRouter's chat-completions response.
type OpenRouterChatResponse struct {
ID string `json:"id"`
Object string `json:"object"`
Created int64 `json:"created"`
Model string `json:"model"`
Choices []struct {
FinishReason string `json:"finish_reason"`
Index int `json:"index"`
Logprobs any `json:"logprobs"`
Message struct {
Content string `json:"content"`
Reasoning string `json:"reasoning"`
Role string `json:"role"`
ToolCalls []map[string]any `json:"tool_calls"`
} `json:"message"`
} `json:"choices"`
SystemFingerprint string `json:"system_fingerprint"`
Usage struct {
CompletionTokens int `json:"completion_tokens"`
PromptTokens int `json:"prompt_tokens"`
TotalTokens int `json:"total_tokens"`
CompletionTokensDetails struct {
ReasoningTokens int `json:"reasoning_tokens"`
} `json:"completion_tokens_details"`
PromptTokensDetails struct {
CachedTokens int `json:"cached_tokens"`
} `json:"prompt_tokens_details"`
} `json:"usage"`
}
func (o *OpenRouterModel) ChatWithMessages(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, chatModelConfig *ChatConfig, modelUsage *common.ModelUsage) (*ChatResponse, error) {
if err := o.baseModel.APIConfigCheck(apiConfig); err != nil {
return nil, err
@@ -108,67 +77,12 @@ func (o *OpenRouterModel) ChatWithMessages(ctx context.Context, modelName string
}
}
jsonData, err := json.Marshal(reqBody)
body, err := o.baseModel.doRequest(ctx, url, apiConfig, reqBody, nonStreamCallTimeout)
if err != nil {
return nil, fmt.Errorf("failed to marshal request: %w", err)
return nil, err
}
ctx, cancel := context.WithTimeout(ctx, nonStreamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return nil, fmt.Errorf("failed to create request: %w", err)
}
req.Header.Add("Content-Type", "application/json")
req.Header.Add("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := o.baseModel.httpClient.Do(req)
if err != nil {
return nil, fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
body, err := io.ReadAll(resp.Body)
if err != nil {
return nil, fmt.Errorf("failed to read response body: %w", err)
}
if resp.StatusCode != http.StatusOK {
return nil, fmt.Errorf("failed to send request: %d %s", resp.StatusCode, string(body))
}
return parseChatCompletionResponse(body, chatModelConfig, modelUsage, func(body []byte, chatConfig *ChatConfig) (chatResponseParts, error) {
var result OpenRouterChatResponse
if err := json.Unmarshal(body, &result); err != nil {
return chatResponseParts{}, fmt.Errorf("failed to unmarshal response: %w", err)
}
if len(result.Choices) == 0 {
return chatResponseParts{}, fmt.Errorf("no choices in response")
}
choice := result.Choices[0]
if choice.Message.Content == "" && len(choice.Message.ToolCalls) == 0 {
return chatResponseParts{}, fmt.Errorf("response contains neither content nor tool calls")
}
reasonContent := ""
if chatConfig != nil && chatConfig.Thinking != nil && *chatConfig.Thinking {
reasonContent = strings.TrimPrefix(choice.Message.Reasoning, "\n")
}
return chatResponseParts{
RequestID: result.ID,
Content: &choice.Message.Content,
ReasonContent: &reasonContent,
ToolCalls: choice.Message.ToolCalls,
Usage: &TokenUsage{
PromptTokens: result.Usage.PromptTokens,
CompletionTokens: result.Usage.CompletionTokens,
TotalTokens: result.Usage.TotalTokens,
},
}, nil
})
return HandleNonStreamingResponse(body, modelUsage, chatModelConfig, OpenAIParserConfig)
}
func (o *OpenRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName string, messages []Message, apiConfig *APIConfig, modelConfig *ChatConfig, modelUsage *common.ModelUsage, sender func(*string, *string) error) error {
@@ -206,95 +120,11 @@ func (o *OpenRouterModel) ChatStreamlyWithSender(ctx context.Context, modelName
}
}
jsonData, err := json.Marshal(reqBody)
if err != nil {
return fmt.Errorf("failed to marshal request: %w", err)
}
ctx, cancel := context.WithTimeout(ctx, streamCallTimeout)
defer cancel()
req, err := http.NewRequestWithContext(ctx, "POST", url, bytes.NewBuffer(jsonData))
if err != nil {
return fmt.Errorf("failed to create request: %w", err)
}
req.Header.Set("Content-Type", "application/json")
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", *apiConfig.ApiKey))
resp, err := o.baseModel.httpClient.Do(req)
if err != nil {
return fmt.Errorf("failed to send request: %w", err)
}
defer resp.Body.Close()
if resp.StatusCode != http.StatusOK {
body, _ := io.ReadAll(resp.Body)
return fmt.Errorf("invalid status code: %d, body: %s", resp.StatusCode, string(body))
}
sawTerminal := false
accumulatedToolCalls := make(map[int]map[string]any)
done, err := ParseSSEStream[map[string]any](resp.Body, func(event map[string]any) error {
common.Info(fmt.Sprintf("%v", event))
tokenUsage, found, usageErr := decodeOpenAICompatibleStreamUsage(event)
if usageErr != nil {
return usageErr
}
if found {
applyStreamUsage(modelConfig, modelUsage, tokenUsage)
}
choices, ok := event["choices"].([]any)
if !ok || len(choices) == 0 {
return nil
}
choice, ok := choices[0].(map[string]any)
if !ok {
return nil
}
if finishReason, ok := choice["finish_reason"].(string); ok && finishReason != "" {
sawTerminal = true
}
delta, ok := choice["delta"].(map[string]any)
if !ok {
return nil
}
accumulateToolCallDeltas(delta, accumulatedToolCalls)
reasoningContent, ok := delta["reasoning"].(string)
if ok && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
return err
}
}
content, ok := delta["content"].(string)
if ok && content != "" {
if err := sender(&content, nil); err != nil {
return err
}
}
return nil
return o.baseModel.doStreamRequest(ctx, url, apiConfig, reqBody, streamCallTimeout, func(body io.ReadCloser) error {
// OpenRouter emits reasoning under delta.reasoning (not
// delta.reasoning_content), so it uses its own ParserConfig.
return HandleStreamingResponse(body, modelUsage, modelConfig, OpenRouterParserConfig, sender)
})
if err != nil {
return fmt.Errorf("failed to scan response body: %w", err)
}
if !done && !sawTerminal {
return fmt.Errorf("openrouter: stream ended before [DONE] or finish_reason")
}
setSortedToolCallsResult(modelConfig, accumulatedToolCalls)
// Send [DONE] marker for OpenAI compatibility
endOfStream := "[DONE]"
if err = sender(&endOfStream, nil); err != nil {
return err
}
return nil
}
// OpenRouterEmbeddingResponse mirrors OpenRouter's embeddings response.

45
internal/entity/models/parser_config.go
View File

@@ -16,8 +16,8 @@
package models
// ParserConfig maps a protocol to its usage parsers. Drivers select a
// ParserConfig instead of implementing usage extraction individually.
// ParserConfig maps a protocol to its usage and reasoning parsers. Drivers
// select a ParserConfig instead of implementing extraction individually.
type ParserConfig struct {
// Protocol is the protocol identifier (e.g. "openai", "claude").
Protocol string
@@ -25,12 +25,45 @@ type ParserConfig struct {
ResponseParser func(map[string]any) (*TokenUsage, bool)
// StreamParser extracts usage from one streaming event.
StreamParser func(map[string]any) (*TokenUsage, bool)
// ExtractStreamReasoning extracts the reasoning text from a parsed
// delta map (the delta field of one streaming event), if any.
// Defaults to reading delta.reasoning_content.
ExtractStreamReasoning func(delta map[string]any) string
}
// extractDefaultStreamReasoning reads the reasoning text from a parsed
// delta (delta.reasoning_content).
func extractDefaultStreamReasoning(delta map[string]any) string {
if r, ok := delta["reasoning_content"].(string); ok {
return r
}
return ""
}
// OpenAIParserConfig is the ParserConfig for OpenAI-compatible APIs
// (NVIDIA NIM, DeepSeek, Aliyun, Moonshot, xAI, OpenRouter, ...).
// (NVIDIA NIM, DeepSeek, Aliyun, Moonshot, xAI, ...).
var OpenAIParserConfig = &ParserConfig{
Protocol: "openai",
ResponseParser: extractOpenAIUsage,
StreamParser: extractOpenAIStreamUsage,
Protocol: "openai",
ResponseParser: extractOpenAIUsage,
StreamParser: extractOpenAIStreamUsage,
ExtractStreamReasoning: extractDefaultStreamReasoning,
}
// extractOpenRouterStreamReasoning reads the reasoning text from an
// OpenRouter streaming event. OpenRouter uses delta.reasoning (not
// delta.reasoning_content) for its reasoning content.
func extractOpenRouterStreamReasoning(delta map[string]any) string {
if r, ok := delta["reasoning"].(string); ok {
return r
}
return ""
}
// OpenRouterParserConfig is the ParserConfig for OpenRouter, which emits
// reasoning under delta.reasoning instead of delta.reasoning_content.
var OpenRouterParserConfig = &ParserConfig{
Protocol: "openai",
ResponseParser: extractOpenAIUsage,
StreamParser: extractOpenAIStreamUsage,
ExtractStreamReasoning: extractOpenRouterStreamReasoning,
}

11
internal/entity/models/response_handler.go
View File

@@ -115,8 +115,15 @@ func HandleStreamingResponse(
accumulateToolCallDeltas(delta, accumulatedToolCalls)
if reasoningContent, ok := delta["reasoning_content"].(string); ok && reasoningContent != "" {
if err := sender(nil, &reasoningContent); err != nil {
// Extract reasoning via the protocol hook so each provider can
// name its reasoning field differently (reasoning_content,
// reasoning, ...) without the shared handler knowing which.
extractReasoning := cfg.ExtractStreamReasoning
if extractReasoning == nil {
extractReasoning = extractDefaultStreamReasoning
}
if reasoning := extractReasoning(delta); reasoning != "" {
if err := sender(nil, &reasoning); err != nil {
return err
}
}