選択できるのは25トピックまでです。 トピックは、先頭が英数字で、英数字とダッシュ('-')を使用した35文字以内のものにしてください。
 
 
 

528 行
14 KiB

  1. package ollama
  2. import (
  3. "bufio"
  4. "encoding/json"
  5. "fmt"
  6. "io"
  7. "net/http"
  8. "strings"
  9. "time"
  10. "github.com/QuantumNous/new-api/common"
  11. "github.com/QuantumNous/new-api/dto"
  12. relaycommon "github.com/QuantumNous/new-api/relay/common"
  13. "github.com/QuantumNous/new-api/service"
  14. "github.com/QuantumNous/new-api/types"
  15. "github.com/gin-gonic/gin"
  16. )
  17. func openAIChatToOllamaChat(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaChatRequest, error) {
  18. chatReq := &OllamaChatRequest{
  19. Model: r.Model,
  20. Stream: r.Stream,
  21. Options: map[string]any{},
  22. Think: r.Think,
  23. }
  24. if r.ResponseFormat != nil {
  25. if r.ResponseFormat.Type == "json" {
  26. chatReq.Format = "json"
  27. } else if r.ResponseFormat.Type == "json_schema" {
  28. if len(r.ResponseFormat.JsonSchema) > 0 {
  29. var schema any
  30. _ = json.Unmarshal(r.ResponseFormat.JsonSchema, &schema)
  31. chatReq.Format = schema
  32. }
  33. }
  34. }
  35. // options mapping
  36. if r.Temperature != nil {
  37. chatReq.Options["temperature"] = r.Temperature
  38. }
  39. if r.TopP != 0 {
  40. chatReq.Options["top_p"] = r.TopP
  41. }
  42. if r.TopK != 0 {
  43. chatReq.Options["top_k"] = r.TopK
  44. }
  45. if r.FrequencyPenalty != 0 {
  46. chatReq.Options["frequency_penalty"] = r.FrequencyPenalty
  47. }
  48. if r.PresencePenalty != 0 {
  49. chatReq.Options["presence_penalty"] = r.PresencePenalty
  50. }
  51. if r.Seed != 0 {
  52. chatReq.Options["seed"] = int(r.Seed)
  53. }
  54. if mt := r.GetMaxTokens(); mt != 0 {
  55. chatReq.Options["num_predict"] = int(mt)
  56. }
  57. if r.Stop != nil {
  58. switch v := r.Stop.(type) {
  59. case string:
  60. chatReq.Options["stop"] = []string{v}
  61. case []string:
  62. chatReq.Options["stop"] = v
  63. case []any:
  64. arr := make([]string, 0, len(v))
  65. for _, i := range v {
  66. if s, ok := i.(string); ok {
  67. arr = append(arr, s)
  68. }
  69. }
  70. if len(arr) > 0 {
  71. chatReq.Options["stop"] = arr
  72. }
  73. }
  74. }
  75. if len(r.Tools) > 0 {
  76. tools := make([]OllamaTool, 0, len(r.Tools))
  77. for _, t := range r.Tools {
  78. tools = append(tools, OllamaTool{Type: "function", Function: OllamaToolFunction{Name: t.Function.Name, Description: t.Function.Description, Parameters: t.Function.Parameters}})
  79. }
  80. chatReq.Tools = tools
  81. }
  82. chatReq.Messages = make([]OllamaChatMessage, 0, len(r.Messages))
  83. for _, m := range r.Messages {
  84. var textBuilder strings.Builder
  85. var images []string
  86. if m.IsStringContent() {
  87. textBuilder.WriteString(m.StringContent())
  88. } else {
  89. parts := m.ParseContent()
  90. for _, part := range parts {
  91. if part.Type == dto.ContentTypeImageURL {
  92. img := part.GetImageMedia()
  93. if img != nil && img.Url != "" {
  94. // 使用统一的文件服务获取图片数据
  95. var source *types.FileSource
  96. if strings.HasPrefix(img.Url, "http") {
  97. source = types.NewURLFileSource(img.Url)
  98. } else {
  99. source = types.NewBase64FileSource(img.Url, "")
  100. }
  101. base64Data, _, err := service.GetBase64Data(c, source, "fetch image for ollama chat")
  102. if err != nil {
  103. return nil, err
  104. }
  105. if base64Data != "" {
  106. images = append(images, base64Data)
  107. }
  108. }
  109. } else if part.Type == dto.ContentTypeText {
  110. textBuilder.WriteString(part.Text)
  111. }
  112. }
  113. }
  114. cm := OllamaChatMessage{Role: m.Role, Content: textBuilder.String()}
  115. if len(images) > 0 {
  116. cm.Images = images
  117. }
  118. if m.Role == "tool" && m.Name != nil {
  119. cm.ToolName = *m.Name
  120. }
  121. if m.ToolCalls != nil && len(m.ToolCalls) > 0 {
  122. parsed := m.ParseToolCalls()
  123. if len(parsed) > 0 {
  124. calls := make([]OllamaToolCall, 0, len(parsed))
  125. for _, tc := range parsed {
  126. var args interface{}
  127. if tc.Function.Arguments != "" {
  128. _ = json.Unmarshal([]byte(tc.Function.Arguments), &args)
  129. }
  130. if args == nil {
  131. args = map[string]any{}
  132. }
  133. oc := OllamaToolCall{}
  134. oc.Function.Name = tc.Function.Name
  135. oc.Function.Arguments = args
  136. calls = append(calls, oc)
  137. }
  138. cm.ToolCalls = calls
  139. }
  140. }
  141. chatReq.Messages = append(chatReq.Messages, cm)
  142. }
  143. return chatReq, nil
  144. }
  145. // openAIToGenerate converts OpenAI completions request to Ollama generate
  146. func openAIToGenerate(c *gin.Context, r *dto.GeneralOpenAIRequest) (*OllamaGenerateRequest, error) {
  147. gen := &OllamaGenerateRequest{
  148. Model: r.Model,
  149. Stream: r.Stream,
  150. Options: map[string]any{},
  151. Think: r.Think,
  152. }
  153. // Prompt may be in r.Prompt (string or []any)
  154. if r.Prompt != nil {
  155. switch v := r.Prompt.(type) {
  156. case string:
  157. gen.Prompt = v
  158. case []any:
  159. var sb strings.Builder
  160. for _, it := range v {
  161. if s, ok := it.(string); ok {
  162. sb.WriteString(s)
  163. }
  164. }
  165. gen.Prompt = sb.String()
  166. default:
  167. gen.Prompt = fmt.Sprintf("%v", r.Prompt)
  168. }
  169. }
  170. if r.Suffix != nil {
  171. if s, ok := r.Suffix.(string); ok {
  172. gen.Suffix = s
  173. }
  174. }
  175. if r.ResponseFormat != nil {
  176. if r.ResponseFormat.Type == "json" {
  177. gen.Format = "json"
  178. } else if r.ResponseFormat.Type == "json_schema" {
  179. var schema any
  180. _ = json.Unmarshal(r.ResponseFormat.JsonSchema, &schema)
  181. gen.Format = schema
  182. }
  183. }
  184. if r.Temperature != nil {
  185. gen.Options["temperature"] = r.Temperature
  186. }
  187. if r.TopP != 0 {
  188. gen.Options["top_p"] = r.TopP
  189. }
  190. if r.TopK != 0 {
  191. gen.Options["top_k"] = r.TopK
  192. }
  193. if r.FrequencyPenalty != 0 {
  194. gen.Options["frequency_penalty"] = r.FrequencyPenalty
  195. }
  196. if r.PresencePenalty != 0 {
  197. gen.Options["presence_penalty"] = r.PresencePenalty
  198. }
  199. if r.Seed != 0 {
  200. gen.Options["seed"] = int(r.Seed)
  201. }
  202. if mt := r.GetMaxTokens(); mt != 0 {
  203. gen.Options["num_predict"] = int(mt)
  204. }
  205. if r.Stop != nil {
  206. switch v := r.Stop.(type) {
  207. case string:
  208. gen.Options["stop"] = []string{v}
  209. case []string:
  210. gen.Options["stop"] = v
  211. case []any:
  212. arr := make([]string, 0, len(v))
  213. for _, i := range v {
  214. if s, ok := i.(string); ok {
  215. arr = append(arr, s)
  216. }
  217. }
  218. if len(arr) > 0 {
  219. gen.Options["stop"] = arr
  220. }
  221. }
  222. }
  223. return gen, nil
  224. }
  225. func requestOpenAI2Embeddings(r dto.EmbeddingRequest) *OllamaEmbeddingRequest {
  226. opts := map[string]any{}
  227. if r.Temperature != nil {
  228. opts["temperature"] = r.Temperature
  229. }
  230. if r.TopP != 0 {
  231. opts["top_p"] = r.TopP
  232. }
  233. if r.FrequencyPenalty != 0 {
  234. opts["frequency_penalty"] = r.FrequencyPenalty
  235. }
  236. if r.PresencePenalty != 0 {
  237. opts["presence_penalty"] = r.PresencePenalty
  238. }
  239. if r.Seed != 0 {
  240. opts["seed"] = int(r.Seed)
  241. }
  242. if r.Dimensions != 0 {
  243. opts["dimensions"] = r.Dimensions
  244. }
  245. input := r.ParseInput()
  246. if len(input) == 1 {
  247. return &OllamaEmbeddingRequest{Model: r.Model, Input: input[0], Options: opts, Dimensions: r.Dimensions}
  248. }
  249. return &OllamaEmbeddingRequest{Model: r.Model, Input: input, Options: opts, Dimensions: r.Dimensions}
  250. }
  251. func ollamaEmbeddingHandler(c *gin.Context, info *relaycommon.RelayInfo, resp *http.Response) (*dto.Usage, *types.NewAPIError) {
  252. var oResp OllamaEmbeddingResponse
  253. body, err := io.ReadAll(resp.Body)
  254. if err != nil {
  255. return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
  256. }
  257. service.CloseResponseBodyGracefully(resp)
  258. if err = common.Unmarshal(body, &oResp); err != nil {
  259. return nil, types.NewOpenAIError(err, types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
  260. }
  261. if oResp.Error != "" {
  262. return nil, types.NewOpenAIError(fmt.Errorf("ollama error: %s", oResp.Error), types.ErrorCodeBadResponseBody, http.StatusInternalServerError)
  263. }
  264. data := make([]dto.OpenAIEmbeddingResponseItem, 0, len(oResp.Embeddings))
  265. for i, emb := range oResp.Embeddings {
  266. data = append(data, dto.OpenAIEmbeddingResponseItem{Index: i, Object: "embedding", Embedding: emb})
  267. }
  268. usage := &dto.Usage{PromptTokens: oResp.PromptEvalCount, CompletionTokens: 0, TotalTokens: oResp.PromptEvalCount}
  269. embResp := &dto.OpenAIEmbeddingResponse{Object: "list", Data: data, Model: info.UpstreamModelName, Usage: *usage}
  270. out, _ := common.Marshal(embResp)
  271. service.IOCopyBytesGracefully(c, resp, out)
  272. return usage, nil
  273. }
  274. func FetchOllamaModels(baseURL, apiKey string) ([]OllamaModel, error) {
  275. url := fmt.Sprintf("%s/api/tags", baseURL)
  276. client := &http.Client{}
  277. request, err := http.NewRequest("GET", url, nil)
  278. if err != nil {
  279. return nil, fmt.Errorf("创建请求失败: %v", err)
  280. }
  281. // Ollama 通常不需要 Bearer token,但为了兼容性保留
  282. if apiKey != "" {
  283. request.Header.Set("Authorization", "Bearer "+apiKey)
  284. }
  285. response, err := client.Do(request)
  286. if err != nil {
  287. return nil, fmt.Errorf("请求失败: %v", err)
  288. }
  289. defer response.Body.Close()
  290. if response.StatusCode != http.StatusOK {
  291. body, _ := io.ReadAll(response.Body)
  292. return nil, fmt.Errorf("服务器返回错误 %d: %s", response.StatusCode, string(body))
  293. }
  294. var tagsResponse OllamaTagsResponse
  295. body, err := io.ReadAll(response.Body)
  296. if err != nil {
  297. return nil, fmt.Errorf("读取响应失败: %v", err)
  298. }
  299. err = common.Unmarshal(body, &tagsResponse)
  300. if err != nil {
  301. return nil, fmt.Errorf("解析响应失败: %v", err)
  302. }
  303. return tagsResponse.Models, nil
  304. }
  305. // 拉取 Ollama 模型 (非流式)
  306. func PullOllamaModel(baseURL, apiKey, modelName string) error {
  307. url := fmt.Sprintf("%s/api/pull", baseURL)
  308. pullRequest := OllamaPullRequest{
  309. Name: modelName,
  310. Stream: false, // 非流式,简化处理
  311. }
  312. requestBody, err := common.Marshal(pullRequest)
  313. if err != nil {
  314. return fmt.Errorf("序列化请求失败: %v", err)
  315. }
  316. client := &http.Client{
  317. Timeout: 30 * 60 * 1000 * time.Millisecond, // 30分钟超时,支持大模型
  318. }
  319. request, err := http.NewRequest("POST", url, strings.NewReader(string(requestBody)))
  320. if err != nil {
  321. return fmt.Errorf("创建请求失败: %v", err)
  322. }
  323. request.Header.Set("Content-Type", "application/json")
  324. if apiKey != "" {
  325. request.Header.Set("Authorization", "Bearer "+apiKey)
  326. }
  327. response, err := client.Do(request)
  328. if err != nil {
  329. return fmt.Errorf("请求失败: %v", err)
  330. }
  331. defer response.Body.Close()
  332. if response.StatusCode != http.StatusOK {
  333. body, _ := io.ReadAll(response.Body)
  334. return fmt.Errorf("拉取模型失败 %d: %s", response.StatusCode, string(body))
  335. }
  336. return nil
  337. }
  338. // 流式拉取 Ollama 模型 (支持进度回调)
  339. func PullOllamaModelStream(baseURL, apiKey, modelName string, progressCallback func(OllamaPullResponse)) error {
  340. url := fmt.Sprintf("%s/api/pull", baseURL)
  341. pullRequest := OllamaPullRequest{
  342. Name: modelName,
  343. Stream: true, // 启用流式
  344. }
  345. requestBody, err := common.Marshal(pullRequest)
  346. if err != nil {
  347. return fmt.Errorf("序列化请求失败: %v", err)
  348. }
  349. client := &http.Client{
  350. Timeout: 60 * 60 * 1000 * time.Millisecond, // 1小时超时,支持超大模型
  351. }
  352. request, err := http.NewRequest("POST", url, strings.NewReader(string(requestBody)))
  353. if err != nil {
  354. return fmt.Errorf("创建请求失败: %v", err)
  355. }
  356. request.Header.Set("Content-Type", "application/json")
  357. if apiKey != "" {
  358. request.Header.Set("Authorization", "Bearer "+apiKey)
  359. }
  360. response, err := client.Do(request)
  361. if err != nil {
  362. return fmt.Errorf("请求失败: %v", err)
  363. }
  364. defer response.Body.Close()
  365. if response.StatusCode != http.StatusOK {
  366. body, _ := io.ReadAll(response.Body)
  367. return fmt.Errorf("拉取模型失败 %d: %s", response.StatusCode, string(body))
  368. }
  369. // 读取流式响应
  370. scanner := bufio.NewScanner(response.Body)
  371. successful := false
  372. for scanner.Scan() {
  373. line := scanner.Text()
  374. if strings.TrimSpace(line) == "" {
  375. continue
  376. }
  377. var pullResponse OllamaPullResponse
  378. if err := common.Unmarshal([]byte(line), &pullResponse); err != nil {
  379. continue // 忽略解析失败的行
  380. }
  381. if progressCallback != nil {
  382. progressCallback(pullResponse)
  383. }
  384. // 检查是否出现错误或完成
  385. if strings.EqualFold(pullResponse.Status, "error") {
  386. return fmt.Errorf("拉取模型失败: %s", strings.TrimSpace(line))
  387. }
  388. if strings.EqualFold(pullResponse.Status, "success") {
  389. successful = true
  390. break
  391. }
  392. }
  393. if err := scanner.Err(); err != nil {
  394. return fmt.Errorf("读取流式响应失败: %v", err)
  395. }
  396. if !successful {
  397. return fmt.Errorf("拉取模型未完成: 未收到成功状态")
  398. }
  399. return nil
  400. }
  401. // 删除 Ollama 模型
  402. func DeleteOllamaModel(baseURL, apiKey, modelName string) error {
  403. url := fmt.Sprintf("%s/api/delete", baseURL)
  404. deleteRequest := OllamaDeleteRequest{
  405. Name: modelName,
  406. }
  407. requestBody, err := common.Marshal(deleteRequest)
  408. if err != nil {
  409. return fmt.Errorf("序列化请求失败: %v", err)
  410. }
  411. client := &http.Client{}
  412. request, err := http.NewRequest("DELETE", url, strings.NewReader(string(requestBody)))
  413. if err != nil {
  414. return fmt.Errorf("创建请求失败: %v", err)
  415. }
  416. request.Header.Set("Content-Type", "application/json")
  417. if apiKey != "" {
  418. request.Header.Set("Authorization", "Bearer "+apiKey)
  419. }
  420. response, err := client.Do(request)
  421. if err != nil {
  422. return fmt.Errorf("请求失败: %v", err)
  423. }
  424. defer response.Body.Close()
  425. if response.StatusCode != http.StatusOK {
  426. body, _ := io.ReadAll(response.Body)
  427. return fmt.Errorf("删除模型失败 %d: %s", response.StatusCode, string(body))
  428. }
  429. return nil
  430. }
  431. func FetchOllamaVersion(baseURL, apiKey string) (string, error) {
  432. trimmedBase := strings.TrimRight(baseURL, "/")
  433. if trimmedBase == "" {
  434. return "", fmt.Errorf("baseURL 为空")
  435. }
  436. url := fmt.Sprintf("%s/api/version", trimmedBase)
  437. client := &http.Client{Timeout: 10 * time.Second}
  438. request, err := http.NewRequest("GET", url, nil)
  439. if err != nil {
  440. return "", fmt.Errorf("创建请求失败: %v", err)
  441. }
  442. if apiKey != "" {
  443. request.Header.Set("Authorization", "Bearer "+apiKey)
  444. }
  445. response, err := client.Do(request)
  446. if err != nil {
  447. return "", fmt.Errorf("请求失败: %v", err)
  448. }
  449. defer response.Body.Close()
  450. body, err := io.ReadAll(response.Body)
  451. if err != nil {
  452. return "", fmt.Errorf("读取响应失败: %v", err)
  453. }
  454. if response.StatusCode != http.StatusOK {
  455. return "", fmt.Errorf("查询版本失败 %d: %s", response.StatusCode, string(body))
  456. }
  457. var versionResp struct {
  458. Version string `json:"version"`
  459. }
  460. if err := json.Unmarshal(body, &versionResp); err != nil {
  461. return "", fmt.Errorf("解析响应失败: %v", err)
  462. }
  463. if versionResp.Version == "" {
  464. return "", fmt.Errorf("未返回版本信息")
  465. }
  466. return versionResp.Version, nil
  467. }