routes_create_test.go 21 KB

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