routes_create_test.go 22 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712
  1. package server
  2. import (
  3. "bytes"
  4. "cmp"
  5. "encoding/json"
  6. "fmt"
  7. "io"
  8. "net/http"
  9. "net/http/httptest"
  10. "os"
  11. "path/filepath"
  12. "slices"
  13. "strings"
  14. "testing"
  15. "github.com/gin-gonic/gin"
  16. "github.com/ollama/ollama/api"
  17. "github.com/ollama/ollama/envconfig"
  18. "github.com/ollama/ollama/llm"
  19. )
  20. var stream bool = false
  21. func createBinFile(t *testing.T, kv map[string]any, ti []llm.Tensor) (string, string) {
  22. t.Helper()
  23. t.Setenv("OLLAMA_MODELS", cmp.Or(os.Getenv("OLLAMA_MODELS"), t.TempDir()))
  24. modelDir := envconfig.Models()
  25. f, err := os.CreateTemp(t.TempDir(), "")
  26. if err != nil {
  27. t.Fatal(err)
  28. }
  29. defer f.Close()
  30. if err := llm.WriteGGUF(f, kv, ti); err != nil {
  31. t.Fatal(err)
  32. }
  33. // Calculate sha256 of file
  34. if _, err := f.Seek(0, 0); err != nil {
  35. t.Fatal(err)
  36. }
  37. digest, _ := GetSHA256Digest(f)
  38. if err := f.Close(); err != nil {
  39. t.Fatal(err)
  40. }
  41. if err := createLink(f.Name(), filepath.Join(modelDir, "blobs", fmt.Sprintf("sha256-%s", strings.TrimPrefix(digest, "sha256:")))); err != nil {
  42. t.Fatal(err)
  43. }
  44. return f.Name(), digest
  45. }
  46. type responseRecorder struct {
  47. *httptest.ResponseRecorder
  48. http.CloseNotifier
  49. }
  50. func NewRecorder() *responseRecorder {
  51. return &responseRecorder{
  52. ResponseRecorder: httptest.NewRecorder(),
  53. }
  54. }
  55. func (t *responseRecorder) CloseNotify() <-chan bool {
  56. return make(chan bool)
  57. }
  58. func createRequest(t *testing.T, fn func(*gin.Context), body any) *httptest.ResponseRecorder {
  59. t.Helper()
  60. // if OLLAMA_MODELS is not set, set it to the temp directory
  61. t.Setenv("OLLAMA_MODELS", cmp.Or(os.Getenv("OLLAMA_MODELS"), t.TempDir()))
  62. w := NewRecorder()
  63. c, _ := gin.CreateTestContext(w)
  64. var b bytes.Buffer
  65. if err := json.NewEncoder(&b).Encode(body); err != nil {
  66. t.Fatal(err)
  67. }
  68. c.Request = &http.Request{
  69. Body: io.NopCloser(&b),
  70. }
  71. fn(c)
  72. return w.ResponseRecorder
  73. }
  74. func checkFileExists(t *testing.T, p string, expect []string) {
  75. t.Helper()
  76. actual, err := filepath.Glob(p)
  77. if err != nil {
  78. t.Fatal(err)
  79. }
  80. if !slices.Equal(actual, expect) {
  81. t.Fatalf("expected slices to be equal %v", actual)
  82. }
  83. }
  84. func TestCreateFromBin(t *testing.T) {
  85. gin.SetMode(gin.TestMode)
  86. p := t.TempDir()
  87. t.Setenv("OLLAMA_MODELS", p)
  88. var s Server
  89. _, digest := createBinFile(t, nil, nil)
  90. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  91. Name: "test",
  92. Files: map[string]string{"test.gguf": digest},
  93. Stream: &stream,
  94. })
  95. if w.Code != http.StatusOK {
  96. fmt.Println(w)
  97. t.Fatalf("expected status code 200, actual %d", w.Code)
  98. }
  99. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  100. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  101. })
  102. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  103. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  104. filepath.Join(p, "blobs", "sha256-ca239d7bd8ea90e4a5d2e6bf88f8d74a47b14336e73eb4e18bed4dd325018116"),
  105. })
  106. }
  107. func TestCreateFromModel(t *testing.T) {
  108. gin.SetMode(gin.TestMode)
  109. p := t.TempDir()
  110. t.Setenv("OLLAMA_MODELS", p)
  111. var s Server
  112. _, digest := createBinFile(t, nil, nil)
  113. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  114. Name: "test",
  115. Files: map[string]string{"test.gguf": digest},
  116. Stream: &stream,
  117. })
  118. if w.Code != http.StatusOK {
  119. t.Fatalf("expected status code 200, actual %d", w.Code)
  120. }
  121. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  122. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  123. })
  124. w = createRequest(t, s.CreateHandler, api.CreateRequest{
  125. Name: "test2",
  126. From: "test",
  127. Stream: &stream,
  128. })
  129. if w.Code != http.StatusOK {
  130. t.Fatalf("expected status code 200, actual %d", w.Code)
  131. }
  132. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  133. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  134. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
  135. })
  136. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  137. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  138. filepath.Join(p, "blobs", "sha256-ca239d7bd8ea90e4a5d2e6bf88f8d74a47b14336e73eb4e18bed4dd325018116"),
  139. })
  140. }
  141. func TestCreateRemovesLayers(t *testing.T) {
  142. gin.SetMode(gin.TestMode)
  143. p := t.TempDir()
  144. t.Setenv("OLLAMA_MODELS", p)
  145. var s Server
  146. _, digest := createBinFile(t, nil, nil)
  147. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  148. Name: "test",
  149. Files: map[string]string{"test.gguf": digest},
  150. Template: "{{ .Prompt }}",
  151. Stream: &stream,
  152. })
  153. if w.Code != http.StatusOK {
  154. t.Fatalf("expected status code 200, actual %d", w.Code)
  155. }
  156. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  157. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  158. })
  159. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  160. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  161. filepath.Join(p, "blobs", "sha256-b507b9c2f6ca642bffcd06665ea7c91f235fd32daeefdf875a0f938db05fb315"),
  162. filepath.Join(p, "blobs", "sha256-bc80b03733773e0728011b2f4adf34c458b400e1aad48cb28d61170f3a2ad2d6"),
  163. })
  164. w = createRequest(t, s.CreateHandler, api.CreateRequest{
  165. Name: "test",
  166. Files: map[string]string{"test.gguf": digest},
  167. Template: "{{ .System }} {{ .Prompt }}",
  168. Stream: &stream,
  169. })
  170. if w.Code != http.StatusOK {
  171. t.Fatalf("expected status code 200, actual %d", w.Code)
  172. }
  173. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  174. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  175. })
  176. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  177. filepath.Join(p, "blobs", "sha256-8f2c2167d789c6b2302dff965160fa5029f6a24096d262c1cbb469f21a045382"),
  178. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  179. filepath.Join(p, "blobs", "sha256-fe7ac77b725cda2ccad03f88a880ecdfd7a33192d6cae08fce2c0ee1455991ed"),
  180. })
  181. }
  182. func TestCreateUnsetsSystem(t *testing.T) {
  183. gin.SetMode(gin.TestMode)
  184. p := t.TempDir()
  185. t.Setenv("OLLAMA_MODELS", p)
  186. var s Server
  187. _, digest := createBinFile(t, nil, nil)
  188. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  189. Name: "test",
  190. Files: map[string]string{"test.gguf": digest},
  191. System: "Say hi!",
  192. Stream: &stream,
  193. })
  194. if w.Code != http.StatusOK {
  195. t.Fatalf("expected status code 200, actual %d", w.Code)
  196. }
  197. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  198. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  199. })
  200. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  201. filepath.Join(p, "blobs", "sha256-8585df945d1069bc78b79bd10bb73ba07fbc29b0f5479a31a601c0d12731416e"),
  202. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  203. filepath.Join(p, "blobs", "sha256-f29e82a8284dbdf5910b1555580ff60b04238b8da9d5e51159ada67a4d0d5851"),
  204. })
  205. w = createRequest(t, s.CreateHandler, api.CreateRequest{
  206. Name: "test",
  207. Files: map[string]string{"test.gguf": digest},
  208. System: "",
  209. Stream: &stream,
  210. })
  211. if w.Code != http.StatusOK {
  212. t.Fatalf("expected status code 200, actual %d", w.Code)
  213. }
  214. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  215. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  216. })
  217. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  218. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  219. filepath.Join(p, "blobs", "sha256-ca239d7bd8ea90e4a5d2e6bf88f8d74a47b14336e73eb4e18bed4dd325018116"),
  220. })
  221. }
  222. func TestCreateMergeParameters(t *testing.T) {
  223. gin.SetMode(gin.TestMode)
  224. p := t.TempDir()
  225. t.Setenv("OLLAMA_MODELS", p)
  226. var s Server
  227. _, digest := createBinFile(t, nil, nil)
  228. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  229. Name: "test",
  230. Files: map[string]string{"test.gguf": digest},
  231. Parameters: map[string]any{
  232. "temperature": 1,
  233. "top_k": 10,
  234. "stop": []string{"USER:", "ASSISTANT:"},
  235. },
  236. Stream: &stream,
  237. })
  238. if w.Code != http.StatusOK {
  239. t.Fatalf("expected status code 200, actual %d", w.Code)
  240. }
  241. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  242. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  243. })
  244. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  245. filepath.Join(p, "blobs", "sha256-1d0ad71299d48c2fb7ae2b98e683643e771f8a5b72be34942af90d97a91c1e37"),
  246. filepath.Join(p, "blobs", "sha256-4a384beaf47a9cbe452dfa5ab70eea691790f3b35a832d12933a1996685bf2b6"),
  247. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  248. })
  249. // in order to merge parameters, the second model must be created FROM the first
  250. w = createRequest(t, s.CreateHandler, api.CreateRequest{
  251. Name: "test2",
  252. From: "test",
  253. Parameters: map[string]any{
  254. "temperature": 0.6,
  255. "top_p": 0.7,
  256. },
  257. Stream: &stream,
  258. })
  259. if w.Code != http.StatusOK {
  260. t.Fatalf("expected status code 200, actual %d", w.Code)
  261. }
  262. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  263. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  264. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
  265. })
  266. // Display contents of each blob in the directory
  267. blobDir := filepath.Join(p, "blobs")
  268. entries, err := os.ReadDir(blobDir)
  269. if err != nil {
  270. t.Fatalf("failed to read blobs directory: %v", err)
  271. }
  272. for _, entry := range entries {
  273. blobPath := filepath.Join(blobDir, entry.Name())
  274. content, err := os.ReadFile(blobPath)
  275. if err != nil {
  276. t.Fatalf("failed to read blob %s: %v", entry.Name(), err)
  277. }
  278. t.Logf("Contents of %s:\n%s", entry.Name(), string(content))
  279. }
  280. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  281. filepath.Join(p, "blobs", "sha256-1d0ad71299d48c2fb7ae2b98e683643e771f8a5b72be34942af90d97a91c1e37"),
  282. filepath.Join(p, "blobs", "sha256-4a384beaf47a9cbe452dfa5ab70eea691790f3b35a832d12933a1996685bf2b6"),
  283. filepath.Join(p, "blobs", "sha256-4cd9d4ba6b734d9b4cbd1e5caa60374c00722e993fce5e1e2d15a33698f71187"),
  284. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  285. filepath.Join(p, "blobs", "sha256-e29a7b3c47287a2489c895d21fe413c20f859a85d20e749492f52a838e36e1ba"),
  286. })
  287. actual, err := os.ReadFile(filepath.Join(p, "blobs", "sha256-e29a7b3c47287a2489c895d21fe413c20f859a85d20e749492f52a838e36e1ba"))
  288. if err != nil {
  289. t.Fatal(err)
  290. }
  291. expect, err := json.Marshal(map[string]any{"temperature": 0.6, "top_k": 10, "top_p": 0.7, "stop": []string{"USER:", "ASSISTANT:"}})
  292. if err != nil {
  293. t.Fatal(err)
  294. }
  295. if !bytes.Equal(bytes.TrimSpace(expect), bytes.TrimSpace(actual)) {
  296. t.Errorf("expected %s, actual %s", string(expect), string(actual))
  297. }
  298. // slices are replaced
  299. w = createRequest(t, s.CreateHandler, api.CreateRequest{
  300. Name: "test2",
  301. From: "test",
  302. Parameters: map[string]any{
  303. "temperature": 0.6,
  304. "top_p": 0.7,
  305. "stop": []string{"<|endoftext|>"},
  306. },
  307. Stream: &stream,
  308. })
  309. if w.Code != http.StatusOK {
  310. t.Fatalf("expected status code 200, actual %d", w.Code)
  311. }
  312. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  313. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  314. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
  315. })
  316. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  317. filepath.Join(p, "blobs", "sha256-12f58bb75cb3042d69a7e013ab87fb3c3c7088f50ddc62f0c77bd332f0d44d35"),
  318. filepath.Join(p, "blobs", "sha256-1d0ad71299d48c2fb7ae2b98e683643e771f8a5b72be34942af90d97a91c1e37"),
  319. filepath.Join(p, "blobs", "sha256-257aa726584f24970a4f240765e75a7169bfbe7f4966c1f04513d6b6c860583a"),
  320. filepath.Join(p, "blobs", "sha256-4a384beaf47a9cbe452dfa5ab70eea691790f3b35a832d12933a1996685bf2b6"),
  321. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  322. })
  323. actual, err = os.ReadFile(filepath.Join(p, "blobs", "sha256-12f58bb75cb3042d69a7e013ab87fb3c3c7088f50ddc62f0c77bd332f0d44d35"))
  324. if err != nil {
  325. t.Fatal(err)
  326. }
  327. expect, err = json.Marshal(map[string]any{"temperature": 0.6, "top_k": 10, "top_p": 0.7, "stop": []string{"<|endoftext|>"}})
  328. if err != nil {
  329. t.Fatal(err)
  330. }
  331. if !bytes.Equal(bytes.TrimSpace(expect), bytes.TrimSpace(actual)) {
  332. t.Errorf("expected %s, actual %s", string(expect), string(actual))
  333. }
  334. }
  335. func TestCreateReplacesMessages(t *testing.T) {
  336. gin.SetMode(gin.TestMode)
  337. p := t.TempDir()
  338. t.Setenv("OLLAMA_MODELS", p)
  339. var s Server
  340. _, digest := createBinFile(t, nil, nil)
  341. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  342. Name: "test",
  343. Files: map[string]string{"test.gguf": digest},
  344. Messages: []api.Message{
  345. {
  346. Role: "assistant",
  347. Content: "What is my purpose?",
  348. },
  349. {
  350. Role: "user",
  351. Content: "You run tests.",
  352. },
  353. {
  354. Role: "assistant",
  355. Content: "Oh, my god.",
  356. },
  357. },
  358. Stream: &stream,
  359. })
  360. if w.Code != http.StatusOK {
  361. t.Fatalf("expected status code 200, actual %d", w.Code)
  362. }
  363. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  364. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  365. })
  366. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  367. filepath.Join(p, "blobs", "sha256-298baeaf6928a60cf666d88d64a1ba606feb43a2865687c39e40652e407bffc4"),
  368. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  369. filepath.Join(p, "blobs", "sha256-e0e27d47045063ccb167ae852c51d49a98eab33fabaee4633fdddf97213e40b5"),
  370. })
  371. w = createRequest(t, s.CreateHandler, api.CreateRequest{
  372. Name: "test2",
  373. From: "test",
  374. Messages: []api.Message{
  375. {
  376. Role: "assistant",
  377. Content: "You're a test, Harry.",
  378. },
  379. {
  380. Role: "user",
  381. Content: "I-I'm a what?",
  382. },
  383. {
  384. Role: "assistant",
  385. Content: "A test. And a thumping good one at that, I'd wager.",
  386. },
  387. },
  388. Stream: &stream,
  389. })
  390. if w.Code != http.StatusOK {
  391. t.Fatalf("expected status code 200, actual %d", w.Code)
  392. }
  393. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  394. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  395. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test2", "latest"),
  396. })
  397. // Old layers will not have been pruned
  398. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  399. filepath.Join(p, "blobs", "sha256-298baeaf6928a60cf666d88d64a1ba606feb43a2865687c39e40652e407bffc4"),
  400. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  401. filepath.Join(p, "blobs", "sha256-a60ecc9da299ec7ede453f99236e5577fd125e143689b646d9f0ddc9971bf4db"),
  402. filepath.Join(p, "blobs", "sha256-e0e27d47045063ccb167ae852c51d49a98eab33fabaee4633fdddf97213e40b5"),
  403. filepath.Join(p, "blobs", "sha256-f4e2c3690efef1b4b63ba1e1b2744ffeb6a7438a0110b86596069f6d9999c80b"),
  404. })
  405. type message struct {
  406. Role string `json:"role"`
  407. Content string `json:"content"`
  408. }
  409. f, err := os.Open(filepath.Join(p, "blobs", "sha256-a60ecc9da299ec7ede453f99236e5577fd125e143689b646d9f0ddc9971bf4db"))
  410. if err != nil {
  411. t.Fatal(err)
  412. }
  413. defer f.Close()
  414. var actual []message
  415. if err := json.NewDecoder(f).Decode(&actual); err != nil {
  416. t.Fatal(err)
  417. }
  418. expect := []message{
  419. {Role: "assistant", Content: "You're a test, Harry."},
  420. {Role: "user", Content: "I-I'm a what?"},
  421. {Role: "assistant", Content: "A test. And a thumping good one at that, I'd wager."},
  422. }
  423. if !slices.Equal(actual, expect) {
  424. t.Errorf("expected %s, actual %s", expect, actual)
  425. }
  426. }
  427. func TestCreateTemplateSystem(t *testing.T) {
  428. gin.SetMode(gin.TestMode)
  429. p := t.TempDir()
  430. t.Setenv("OLLAMA_MODELS", p)
  431. var s Server
  432. _, digest := createBinFile(t, nil, nil)
  433. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  434. Name: "test",
  435. Files: map[string]string{"test.gguf": digest},
  436. Template: "{{ .System }} {{ .Prompt }}",
  437. System: "Say bye!",
  438. Stream: &stream,
  439. })
  440. if w.Code != http.StatusOK {
  441. t.Fatalf("expected status code 200, actual %d", w.Code)
  442. }
  443. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  444. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  445. })
  446. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  447. filepath.Join(p, "blobs", "sha256-2b5e330885117c82f3fd75169ea323e141070a2947c11ddb9f79ee0b01c589c1"),
  448. filepath.Join(p, "blobs", "sha256-4c5f51faac758fecaff8db42f0b7382891a4d0c0bb885f7b86be88c814a7cc86"),
  449. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  450. filepath.Join(p, "blobs", "sha256-fe7ac77b725cda2ccad03f88a880ecdfd7a33192d6cae08fce2c0ee1455991ed"),
  451. })
  452. template, err := os.ReadFile(filepath.Join(p, "blobs", "sha256-fe7ac77b725cda2ccad03f88a880ecdfd7a33192d6cae08fce2c0ee1455991ed"))
  453. if err != nil {
  454. t.Fatal(err)
  455. }
  456. if string(template) != "{{ .System }} {{ .Prompt }}" {
  457. t.Errorf("expected \"{{ .System }} {{ .Prompt }}\", actual %s", template)
  458. }
  459. system, err := os.ReadFile(filepath.Join(p, "blobs", "sha256-4c5f51faac758fecaff8db42f0b7382891a4d0c0bb885f7b86be88c814a7cc86"))
  460. if err != nil {
  461. t.Fatal(err)
  462. }
  463. if string(system) != "Say bye!" {
  464. t.Errorf("expected \"Say bye!\", actual %s", system)
  465. }
  466. t.Run("incomplete template", func(t *testing.T) {
  467. _, digest := createBinFile(t, nil, nil)
  468. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  469. Name: "test",
  470. Files: map[string]string{"test.gguf": digest},
  471. Template: "{{ .Prompt",
  472. Stream: &stream,
  473. })
  474. if w.Code != http.StatusBadRequest {
  475. t.Fatalf("expected status code 400, actual %d", w.Code)
  476. }
  477. })
  478. t.Run("template with unclosed if", func(t *testing.T) {
  479. _, digest := createBinFile(t, nil, nil)
  480. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  481. Name: "test",
  482. Files: map[string]string{"test.gguf": digest},
  483. Template: "{{ if .Prompt }}",
  484. Stream: &stream,
  485. })
  486. if w.Code != http.StatusBadRequest {
  487. t.Fatalf("expected status code 400, actual %d", w.Code)
  488. }
  489. })
  490. t.Run("template with undefined function", func(t *testing.T) {
  491. _, digest := createBinFile(t, nil, nil)
  492. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  493. Name: "test",
  494. Files: map[string]string{"test.gguf": digest},
  495. Template: "{{ Prompt }}",
  496. Stream: &stream,
  497. })
  498. if w.Code != http.StatusBadRequest {
  499. t.Fatalf("expected status code 400, actual %d", w.Code)
  500. }
  501. })
  502. }
  503. func TestCreateLicenses(t *testing.T) {
  504. gin.SetMode(gin.TestMode)
  505. p := t.TempDir()
  506. t.Setenv("OLLAMA_MODELS", p)
  507. var s Server
  508. _, digest := createBinFile(t, nil, nil)
  509. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  510. Name: "test",
  511. Files: map[string]string{"test.gguf": digest},
  512. License: []string{"MIT", "Apache-2.0"},
  513. Stream: &stream,
  514. })
  515. if w.Code != http.StatusOK {
  516. t.Fatalf("expected status code 200, actual %d", w.Code)
  517. }
  518. checkFileExists(t, filepath.Join(p, "manifests", "*", "*", "*", "*"), []string{
  519. filepath.Join(p, "manifests", "registry.ollama.ai", "library", "test", "latest"),
  520. })
  521. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  522. filepath.Join(p, "blobs", "sha256-2af71558e438db0b73a20beab92dc278a94e1bbe974c00c1a33e3ab62d53a608"),
  523. filepath.Join(p, "blobs", "sha256-79a39c37536ddee29cbadd5d5e2dcba8ed7f03e431f626ff38432c1c866bb7e2"),
  524. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  525. filepath.Join(p, "blobs", "sha256-e5dcffe836b6ec8a58e492419b550e65fb8cbdc308503979e5dacb33ac7ea3b7"),
  526. })
  527. mit, err := os.ReadFile(filepath.Join(p, "blobs", "sha256-e5dcffe836b6ec8a58e492419b550e65fb8cbdc308503979e5dacb33ac7ea3b7"))
  528. if err != nil {
  529. t.Fatal(err)
  530. }
  531. if string(mit) != "MIT" {
  532. t.Errorf("expected MIT, actual %s", mit)
  533. }
  534. apache, err := os.ReadFile(filepath.Join(p, "blobs", "sha256-2af71558e438db0b73a20beab92dc278a94e1bbe974c00c1a33e3ab62d53a608"))
  535. if err != nil {
  536. t.Fatal(err)
  537. }
  538. if string(apache) != "Apache-2.0" {
  539. t.Errorf("expected Apache-2.0, actual %s", apache)
  540. }
  541. }
  542. func TestCreateDetectTemplate(t *testing.T) {
  543. gin.SetMode(gin.TestMode)
  544. p := t.TempDir()
  545. t.Setenv("OLLAMA_MODELS", p)
  546. var s Server
  547. t.Run("matched", func(t *testing.T) {
  548. _, digest := createBinFile(t, llm.KV{
  549. "tokenizer.chat_template": "{{ bos_token }}{% for message in messages %}{{'<|' + message['role'] + '|>' + '\n' + message['content'] + '<|end|>\n' }}{% endfor %}{% if add_generation_prompt %}{{ '<|assistant|>\n' }}{% else %}{{ eos_token }}{% endif %}",
  550. }, nil)
  551. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  552. Name: "test",
  553. Files: map[string]string{"test.gguf": digest},
  554. Stream: &stream,
  555. })
  556. if w.Code != http.StatusOK {
  557. t.Fatalf("expected status code 200, actual %d", w.Code)
  558. }
  559. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  560. filepath.Join(p, "blobs", "sha256-0d79f567714c62c048378f2107fb332dabee0135d080c302d884317da9433cc5"),
  561. filepath.Join(p, "blobs", "sha256-35360843d0c84fb1506952a131bbef13cd2bb4a541251f22535170c05b56e672"),
  562. filepath.Join(p, "blobs", "sha256-553c4a3f747b3d22a4946875f1cc8ed011c2930d83f864a0c7265f9ec0a20413"),
  563. filepath.Join(p, "blobs", "sha256-de3959f841e9ef6b4b6255fa41cb9e0a45da89c3066aa72bdd07a4747f848990"),
  564. })
  565. })
  566. t.Run("unmatched", func(t *testing.T) {
  567. _, digest := createBinFile(t, nil, nil)
  568. w := createRequest(t, s.CreateHandler, api.CreateRequest{
  569. Name: "test",
  570. Files: map[string]string{"test.gguf": digest},
  571. Stream: &stream,
  572. })
  573. if w.Code != http.StatusOK {
  574. t.Fatalf("expected status code 200, actual %d", w.Code)
  575. }
  576. checkFileExists(t, filepath.Join(p, "blobs", "*"), []string{
  577. filepath.Join(p, "blobs", "sha256-a4e5e156ddec27e286f75328784d7106b60a4eb1d246e950a001a3f944fbda99"),
  578. filepath.Join(p, "blobs", "sha256-ca239d7bd8ea90e4a5d2e6bf88f8d74a47b14336e73eb4e18bed4dd325018116"),
  579. })
  580. })
  581. }