server.go 24 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899
  1. package llm
  2. import (
  3. "bufio"
  4. "bytes"
  5. "context"
  6. "encoding/json"
  7. "errors"
  8. "fmt"
  9. "io"
  10. "log"
  11. "log/slog"
  12. "math/rand"
  13. "net"
  14. "net/http"
  15. "os"
  16. "os/exec"
  17. "path/filepath"
  18. "runtime"
  19. "strconv"
  20. "strings"
  21. "time"
  22. "github.com/ollama/ollama/api"
  23. "github.com/ollama/ollama/format"
  24. "github.com/ollama/ollama/gpu"
  25. )
  26. // LlamaServer is an instance of the llama.cpp server
  27. type LlamaServer struct {
  28. port int
  29. cmd *exec.Cmd
  30. done chan error // Channel to signal when the process exits
  31. status *StatusWriter
  32. options api.Options
  33. }
  34. func NewLlamaServer(model string, adapters, projectors []string, opts api.Options) (*LlamaServer, error) {
  35. f, err := os.Open(model)
  36. if err != nil {
  37. return nil, err
  38. }
  39. defer f.Close()
  40. ggml, _, err := DecodeGGML(f)
  41. if err != nil {
  42. return nil, err
  43. }
  44. if opts.NumCtx > int(ggml.KV().ContextLength()) {
  45. slog.Warn("requested context length is greater than model max context length", "requested", opts.NumCtx, "model", ggml.KV().ContextLength())
  46. opts.NumCtx = int(ggml.KV().ContextLength())
  47. }
  48. if opts.NumCtx < 4 {
  49. opts.NumCtx = 4
  50. }
  51. memoryAvailable, _ := gpu.CheckVRAM()
  52. info := gpu.GetGPUInfo()
  53. memoryMinimum := info.MinimumMemory
  54. for _, projector := range projectors {
  55. memoryMinimum += projectorMemoryRequirements(projector)
  56. // multimodal models require at least 2048 context
  57. opts.NumCtx = max(opts.NumCtx, 2048)
  58. }
  59. // fp16 k,v = (1 (k) + 1 (v)) * sizeof(float16) * n_ctx * n_layer * n_embd / n_head * n_head_kv
  60. var kv uint64 = 2 * 2 * uint64(opts.NumCtx) * ggml.KV().BlockCount() * ggml.KV().EmbeddingLength() / ggml.KV().HeadCount() * ggml.KV().HeadCountKV()
  61. graphPartialOffload, graphFullOffload := ggml.GraphSize(uint64(opts.NumCtx), uint64(min(opts.NumCtx, opts.NumBatch)))
  62. if graphPartialOffload == 0 {
  63. graphPartialOffload = ggml.KV().GQA() * kv / 6
  64. }
  65. if graphFullOffload == 0 {
  66. graphFullOffload = graphPartialOffload
  67. }
  68. graphFullOffload *= uint64(info.DeviceCount)
  69. graphPartialOffload *= uint64(info.DeviceCount)
  70. // memoryRequiredTotal represents the memory required for full GPU offloading (all layers)
  71. memoryRequiredTotal := memoryMinimum + graphFullOffload
  72. // memoryRequiredPartial represents the memory required for partial GPU offloading (n > 0, n < layers)
  73. memoryRequiredPartial := memoryMinimum + graphPartialOffload
  74. if info.Library != "metal" {
  75. if memoryRequiredPartial > memoryAvailable {
  76. info.Library = "cpu"
  77. }
  78. }
  79. var layerCount int
  80. layers := ggml.Tensors().Layers()
  81. for i := 0; i < int(ggml.KV().BlockCount()); i++ {
  82. memoryLayer := layers[fmt.Sprintf("blk.%d", i)].size()
  83. // KV is proportional to the number of layers
  84. memoryLayer += kv / ggml.KV().BlockCount()
  85. memoryRequiredTotal += memoryLayer
  86. if memoryAvailable > memoryRequiredPartial+memoryLayer {
  87. memoryRequiredPartial += memoryLayer
  88. layerCount++
  89. }
  90. }
  91. var memoryLayerOutput uint64
  92. for k, v := range layers {
  93. if !strings.HasPrefix(k, "blk.") {
  94. memoryLayerOutput += v.size()
  95. }
  96. }
  97. memoryRequiredTotal += memoryLayerOutput
  98. if info.Library == "metal" && memoryRequiredTotal > info.TotalMemory {
  99. // disable partial offloading when model is greater than total system memory
  100. opts.NumGPU = 0
  101. } else if memoryAvailable > memoryRequiredTotal {
  102. layerCount = int(ggml.KV().BlockCount()) + 1
  103. memoryRequiredPartial = memoryRequiredTotal
  104. }
  105. if opts.NumGPU < 0 {
  106. opts.NumGPU = layerCount
  107. }
  108. memoryWeights := memoryRequiredTotal - memoryMinimum - graphFullOffload - kv
  109. slog.Info(
  110. "offload to gpu",
  111. slog.Group(
  112. "layers",
  113. // actual number of layers offloaded
  114. "real", opts.NumGPU,
  115. // estimated number of layers that can be offloaded
  116. "estimate", layerCount,
  117. ),
  118. slog.Group(
  119. "memory",
  120. // memory available for offloading
  121. "available", format.HumanBytes2(memoryAvailable),
  122. slog.Group(
  123. "required",
  124. // memory required for full offloading
  125. "full", format.HumanBytes2(memoryRequiredTotal),
  126. // memory required to offload layers.estimate layers
  127. "partial", format.HumanBytes2(memoryRequiredPartial),
  128. // memory of KV cache
  129. "kv", format.HumanBytes2(kv),
  130. ),
  131. slog.Group(
  132. "weights",
  133. // memory of the weights
  134. "total", format.HumanBytes2(memoryWeights),
  135. // memory of repeating layers
  136. "repeating", format.HumanBytes2(memoryWeights-memoryLayerOutput),
  137. // memory of non-repeating layers
  138. "nonrepeating", format.HumanBytes2(memoryLayerOutput),
  139. ),
  140. slog.Group(
  141. "graph",
  142. // memory of graph when fully offloaded
  143. "full", format.HumanBytes2(graphFullOffload),
  144. // memory of graph when not fully offloaded
  145. "partial", format.HumanBytes2(graphPartialOffload),
  146. ),
  147. ),
  148. )
  149. if len(adapters) > 1 {
  150. return nil, errors.New("ollama supports only one lora adapter, but multiple were provided")
  151. }
  152. availableServers := availableServers()
  153. servers := serversForGpu(info)
  154. demandLib := os.Getenv("OLLAMA_LLM_LIBRARY")
  155. if demandLib != "" {
  156. serverPath := availableServers[demandLib]
  157. if serverPath == "" {
  158. slog.Info(fmt.Sprintf("Invalid OLLAMA_LLM_LIBRARY %s - not found", demandLib))
  159. } else {
  160. slog.Info("user override", "OLLAMA_LLM_LIBRARY", demandLib, "path", serverPath)
  161. servers = []string{demandLib}
  162. }
  163. }
  164. if len(servers) == 0 {
  165. return nil, fmt.Errorf("no servers found for %v", info)
  166. }
  167. params := []string{
  168. "--model", model,
  169. "--ctx-size", fmt.Sprintf("%d", opts.NumCtx),
  170. "--batch-size", fmt.Sprintf("%d", opts.NumBatch),
  171. "--embedding",
  172. }
  173. if debug := os.Getenv("OLLAMA_DEBUG"); debug != "" {
  174. params = append(params, "--log-format", "json")
  175. } else {
  176. params = append(params, "--log-disable")
  177. }
  178. if opts.NumGPU >= 0 {
  179. params = append(params, "--n-gpu-layers", fmt.Sprintf("%d", opts.NumGPU))
  180. }
  181. if debug := os.Getenv("OLLAMA_DEBUG"); debug != "" {
  182. params = append(params, "--verbose")
  183. }
  184. if opts.MainGPU > 0 {
  185. params = append(params, "--main-gpu", fmt.Sprintf("%d", opts.MainGPU))
  186. }
  187. if len(adapters) > 0 {
  188. // TODO: applying multiple adapters is not supported by the llama.cpp server yet
  189. params = append(params, "--lora", adapters[0])
  190. }
  191. if len(projectors) > 0 {
  192. // TODO: applying multiple projectors is not supported by the llama.cpp server yet
  193. params = append(params, "--mmproj", projectors[0])
  194. }
  195. if opts.NumThread > 0 {
  196. params = append(params, "--threads", fmt.Sprintf("%d", opts.NumThread))
  197. }
  198. if !opts.F16KV {
  199. params = append(params, "--memory-f32")
  200. }
  201. if opts.UseMLock {
  202. params = append(params, "--mlock")
  203. }
  204. if !opts.UseMMap {
  205. params = append(params, "--no-mmap")
  206. }
  207. if opts.UseNUMA {
  208. params = append(params, "--numa")
  209. }
  210. // Loop through potential servers
  211. var finalErr error
  212. for i := 0; i < len(servers); i++ {
  213. dir := availableServers[servers[i]]
  214. // Find an availableServers port, retry on each iterration in case the failure was a port conflict race
  215. port := 0
  216. if a, err := net.ResolveTCPAddr("tcp", "localhost:0"); err == nil {
  217. var l *net.TCPListener
  218. if l, err = net.ListenTCP("tcp", a); err == nil {
  219. port = l.Addr().(*net.TCPAddr).Port
  220. l.Close()
  221. }
  222. }
  223. if port == 0 {
  224. slog.Debug("ResolveTCPAddr failed ", "error", err)
  225. port = rand.Intn(65535-49152) + 49152 // get a random port in the ephemeral range
  226. }
  227. finalParams := append(params, "--port", strconv.Itoa(port))
  228. pathEnv := "LD_LIBRARY_PATH"
  229. if runtime.GOOS == "windows" {
  230. pathEnv = "PATH"
  231. }
  232. // append the server directory to LD_LIBRARY_PATH/PATH
  233. libraryPaths := []string{dir}
  234. if libraryPath, ok := os.LookupEnv(pathEnv); ok {
  235. // Append our runner directory to the path
  236. // This will favor system libraries over our bundled library dependencies
  237. libraryPaths = append(filepath.SplitList(libraryPath), libraryPaths...)
  238. }
  239. server := filepath.Join(dir, "ollama_llama_server")
  240. if runtime.GOOS == "windows" {
  241. server = server + ".exe"
  242. }
  243. s := &LlamaServer{
  244. port: port,
  245. cmd: exec.Command(server, finalParams...),
  246. status: NewStatusWriter(os.Stderr),
  247. options: opts,
  248. }
  249. libEnv := fmt.Sprintf("%s=%s", pathEnv, strings.Join(libraryPaths, string(filepath.ListSeparator)))
  250. slog.Debug(libEnv)
  251. s.cmd.Env = append(os.Environ(), libEnv)
  252. s.cmd.Stdout = os.Stdout
  253. s.cmd.Stderr = s.status
  254. slog.Info("starting llama server", "cmd", s.cmd.String())
  255. if err = s.cmd.Start(); err != nil {
  256. msg := ""
  257. if s.status != nil && s.status.LastErrMsg != "" {
  258. msg = s.status.LastErrMsg
  259. }
  260. err = fmt.Errorf("error starting the external llama server: %v %s", err, msg)
  261. finalErr = err
  262. continue
  263. }
  264. // reap subprocess when it exits
  265. go func() {
  266. // Exit status managed via getServerStatus
  267. _ = s.cmd.Wait()
  268. }()
  269. return s, nil
  270. }
  271. slog.Error("unable to load any llama server", "error", finalErr)
  272. return nil, finalErr
  273. }
  274. func projectorMemoryRequirements(filename string) uint64 {
  275. file, err := os.Open(filename)
  276. if err != nil {
  277. return 0
  278. }
  279. defer file.Close()
  280. ggml, _, err := DecodeGGML(file)
  281. if err != nil {
  282. return 0
  283. }
  284. var mem uint64
  285. for _, layer := range ggml.Tensors().Layers() {
  286. mem += layer.size()
  287. }
  288. return mem
  289. }
  290. type ServerStatus int
  291. const ( // iota is reset to 0
  292. ServerStatusReady ServerStatus = iota
  293. ServerStatusNoSlotsAvaialble
  294. ServerStatusLoadingModel
  295. ServerStatusNotResponding
  296. ServerStatusError
  297. )
  298. type ServerStatusResp struct {
  299. Status string `json:"status"`
  300. SlotsIdle int `json:"slots_idle"`
  301. SlotsProcessing int `json:"slots_processing"`
  302. Error string `json:"error"`
  303. }
  304. func (s *LlamaServer) getServerStatus(ctx context.Context) (ServerStatus, error) {
  305. // Fail fast if its exited
  306. if s.cmd.ProcessState != nil {
  307. msg := ""
  308. if s.status != nil && s.status.LastErrMsg != "" {
  309. msg = s.status.LastErrMsg
  310. }
  311. return ServerStatusError, fmt.Errorf("llama runner process no longer running: %d %s", s.cmd.ProcessState.ExitCode(), msg)
  312. }
  313. req, err := http.NewRequestWithContext(ctx, http.MethodGet, fmt.Sprintf("http://127.0.0.1:%d/health", s.port), nil)
  314. if err != nil {
  315. return ServerStatusError, fmt.Errorf("error creating GET request: %v", err)
  316. }
  317. req.Header.Set("Content-Type", "application/json")
  318. resp, err := http.DefaultClient.Do(req)
  319. if err != nil {
  320. if errors.Is(err, context.DeadlineExceeded) {
  321. return ServerStatusNotResponding, fmt.Errorf("server not responding")
  322. }
  323. return ServerStatusError, fmt.Errorf("health resp: %w", err)
  324. }
  325. defer resp.Body.Close()
  326. body, err := io.ReadAll(resp.Body)
  327. if err != nil {
  328. return ServerStatusError, fmt.Errorf("read health request: %w", err)
  329. }
  330. var status ServerStatusResp
  331. if err := json.Unmarshal(body, &status); err != nil {
  332. return ServerStatusError, fmt.Errorf("health unmarshal encode response: %w", err)
  333. }
  334. switch status.Status {
  335. case "ok":
  336. return ServerStatusReady, nil
  337. case "no slot available":
  338. return ServerStatusNoSlotsAvaialble, nil
  339. case "loading model":
  340. return ServerStatusLoadingModel, nil
  341. default:
  342. return ServerStatusError, fmt.Errorf("server error: %+v", status)
  343. }
  344. }
  345. func (s *LlamaServer) Ping(ctx context.Context) error {
  346. _, err := s.getServerStatus(ctx)
  347. if err != nil {
  348. slog.Debug("server unhealthy", "error", err)
  349. return err
  350. }
  351. return nil
  352. }
  353. func (s *LlamaServer) WaitUntilRunning() error {
  354. start := time.Now()
  355. // TODO we need to wire up a better way to detect hangs during model load and startup of the server
  356. expiresAt := time.Now().Add(10 * time.Minute) // be generous with timeout, large models can take a while to load
  357. ticker := time.NewTicker(50 * time.Millisecond)
  358. defer ticker.Stop()
  359. slog.Info("waiting for llama runner to start responding")
  360. var lastStatus ServerStatus = -1
  361. for {
  362. select {
  363. case err := <-s.done:
  364. msg := ""
  365. if s.status != nil && s.status.LastErrMsg != "" {
  366. msg = s.status.LastErrMsg
  367. }
  368. return fmt.Errorf("llama runner process has terminated: %v %s", err, msg)
  369. case <-ticker.C:
  370. if time.Now().After(expiresAt) {
  371. // timeout
  372. msg := ""
  373. if s.status != nil && s.status.LastErrMsg != "" {
  374. msg = s.status.LastErrMsg
  375. }
  376. return fmt.Errorf("timed out waiting for llama runner to start: %s", msg)
  377. }
  378. if s.cmd.ProcessState != nil {
  379. msg := ""
  380. if s.status != nil && s.status.LastErrMsg != "" {
  381. msg = s.status.LastErrMsg
  382. }
  383. return fmt.Errorf("llama runner process no longer running: %d %s", s.cmd.ProcessState.ExitCode(), msg)
  384. }
  385. ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond)
  386. defer cancel()
  387. status, err := s.getServerStatus(ctx)
  388. if err != nil && lastStatus != status {
  389. slog.Debug("server not yet available", "error", err)
  390. lastStatus = status
  391. continue
  392. }
  393. switch status {
  394. case ServerStatusLoadingModel:
  395. // TODO - this state never seems to happen with the current server.cpp code (bug?)
  396. // it doesn't respond to the health endpoint until after the model is loaded
  397. slog.Debug("loading model")
  398. case ServerStatusReady:
  399. slog.Debug(fmt.Sprintf("llama runner started in %f seconds", time.Since(start).Seconds()))
  400. return nil
  401. }
  402. }
  403. }
  404. }
  405. const jsonGrammar = `
  406. root ::= object
  407. value ::= object | array | string | number | ("true" | "false" | "null") ws
  408. object ::=
  409. "{" ws (
  410. string ":" ws value
  411. ("," ws string ":" ws value)*
  412. )? "}" ws
  413. array ::=
  414. "[" ws (
  415. value
  416. ("," ws value)*
  417. )? "]" ws
  418. string ::=
  419. "\"" (
  420. [^"\\] |
  421. "\\" (["\\/bfnrt] | "u" [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F] [0-9a-fA-F]) # escapes
  422. )* "\"" ws
  423. number ::= ("-"? ([0-9] | [1-9] [0-9]*)) ("." [0-9]+)? ([eE] [-+]? [0-9]+)? ws
  424. # Optional space: by convention, applied in this grammar after literal chars when allowed
  425. ws ::= ([ \t\n] ws)?
  426. `
  427. const maxBufferSize = 512 * format.KiloByte
  428. const maxRetries = 3
  429. type ImageData struct {
  430. Data []byte `json:"data"`
  431. ID int `json:"id"`
  432. }
  433. type completion struct {
  434. Content string `json:"content"`
  435. Model string `json:"model"`
  436. Prompt string `json:"prompt"`
  437. Stop bool `json:"stop"`
  438. StoppedEos bool `json:"stopped_eos"`
  439. StoppedWord bool `json:"stopped_word"`
  440. StoppedLimit bool `json:"stopped_limit"`
  441. Timings struct {
  442. PredictedN int `json:"predicted_n"`
  443. PredictedMS float64 `json:"predicted_ms"`
  444. PromptN int `json:"prompt_n"`
  445. PromptMS float64 `json:"prompt_ms"`
  446. }
  447. }
  448. type CompletionRequest struct {
  449. Prompt string
  450. Format string
  451. Images []ImageData
  452. Options api.Options
  453. }
  454. type CompletionResponse struct {
  455. Content string
  456. Done bool
  457. DoneReason string
  458. PromptEvalCount int
  459. PromptEvalDuration time.Duration
  460. EvalCount int
  461. EvalDuration time.Duration
  462. }
  463. func (s *LlamaServer) Completion(ctx context.Context, req CompletionRequest, fn func(CompletionResponse)) error {
  464. request := map[string]any{
  465. "prompt": req.Prompt,
  466. "stream": true,
  467. "n_predict": req.Options.NumPredict,
  468. "n_keep": req.Options.NumKeep,
  469. "main_gpu": req.Options.MainGPU,
  470. "temperature": req.Options.Temperature,
  471. "top_k": req.Options.TopK,
  472. "top_p": req.Options.TopP,
  473. "tfs_z": req.Options.TFSZ,
  474. "typical_p": req.Options.TypicalP,
  475. "repeat_last_n": req.Options.RepeatLastN,
  476. "repeat_penalty": req.Options.RepeatPenalty,
  477. "presence_penalty": req.Options.PresencePenalty,
  478. "frequency_penalty": req.Options.FrequencyPenalty,
  479. "mirostat": req.Options.Mirostat,
  480. "mirostat_tau": req.Options.MirostatTau,
  481. "mirostat_eta": req.Options.MirostatEta,
  482. "penalize_nl": req.Options.PenalizeNewline,
  483. "seed": req.Options.Seed,
  484. "stop": req.Options.Stop,
  485. "image_data": req.Images,
  486. "cache_prompt": true,
  487. }
  488. // Make sure the server is ready
  489. status, err := s.getServerStatus(ctx)
  490. if err != nil {
  491. return err
  492. } else if status != ServerStatusReady {
  493. return fmt.Errorf("unexpected server status: %d", status)
  494. }
  495. if req.Format == "json" {
  496. request["grammar"] = jsonGrammar
  497. if !strings.Contains(strings.ToLower(req.Prompt), "json") {
  498. slog.Warn("Prompt does not specify that the LLM should response in JSON, but JSON format is expected. For best results specify that JSON is expected in the system prompt.")
  499. }
  500. }
  501. retryDelay := 100 * time.Microsecond
  502. for retries := 0; retries < maxRetries; retries++ {
  503. if retries > 0 {
  504. time.Sleep(retryDelay) // wait before retrying
  505. retryDelay *= 2 // exponential backoff
  506. }
  507. // Handling JSON marshaling with special characters unescaped.
  508. buffer := &bytes.Buffer{}
  509. enc := json.NewEncoder(buffer)
  510. enc.SetEscapeHTML(false)
  511. if err := enc.Encode(request); err != nil {
  512. return fmt.Errorf("failed to marshal data: %v", err)
  513. }
  514. endpoint := fmt.Sprintf("http://127.0.0.1:%d/completion", s.port)
  515. req, err := http.NewRequestWithContext(ctx, http.MethodPost, endpoint, buffer)
  516. if err != nil {
  517. return fmt.Errorf("error creating POST request: %v", err)
  518. }
  519. req.Header.Set("Content-Type", "application/json")
  520. resp, err := http.DefaultClient.Do(req)
  521. if err != nil {
  522. return fmt.Errorf("POST predict: %v", err)
  523. }
  524. defer resp.Body.Close()
  525. if resp.StatusCode >= 400 {
  526. bodyBytes, err := io.ReadAll(resp.Body)
  527. if err != nil {
  528. return fmt.Errorf("failed reading llm error response: %w", err)
  529. }
  530. log.Printf("llm predict error: %s", bodyBytes)
  531. return fmt.Errorf("%s", bodyBytes)
  532. }
  533. scanner := bufio.NewScanner(resp.Body)
  534. buf := make([]byte, 0, maxBufferSize)
  535. scanner.Buffer(buf, maxBufferSize)
  536. retryNeeded := false
  537. // keep track of the last token generated, this is used to abort if the model starts looping
  538. var lastToken string
  539. var tokenRepeat int
  540. for scanner.Scan() {
  541. select {
  542. case <-ctx.Done():
  543. // This handles the request cancellation
  544. return ctx.Err()
  545. default:
  546. line := scanner.Bytes()
  547. if len(line) == 0 {
  548. continue
  549. }
  550. // try again on slot unavailable
  551. if bytes.Contains(line, []byte("slot unavailable")) {
  552. retryNeeded = true
  553. break
  554. }
  555. evt, ok := bytes.CutPrefix(line, []byte("data: "))
  556. if !ok {
  557. return fmt.Errorf("error parsing llm response stream: %s", line)
  558. }
  559. fmt.Println("c", string(evt))
  560. var c completion
  561. if err := json.Unmarshal(evt, &c); err != nil {
  562. return fmt.Errorf("error unmarshaling llm prediction response: %v", err)
  563. }
  564. switch {
  565. case strings.TrimSpace(c.Content) == lastToken:
  566. tokenRepeat++
  567. default:
  568. lastToken = strings.TrimSpace(c.Content)
  569. tokenRepeat = 0
  570. }
  571. // 30 picked as an arbitrary max token repeat limit, modify as needed
  572. if tokenRepeat > 30 {
  573. slog.Debug("prediction aborted, token repeat limit reached")
  574. return ctx.Err()
  575. }
  576. if c.Content != "" {
  577. fn(CompletionResponse{
  578. Content: c.Content,
  579. })
  580. }
  581. if c.Stop {
  582. var doneReason string
  583. switch {
  584. case c.StoppedEos:
  585. doneReason = "stop"
  586. case c.StoppedWord:
  587. doneReason = "stop"
  588. case c.StoppedLimit:
  589. doneReason = "limit"
  590. }
  591. fn(CompletionResponse{
  592. Done: true,
  593. DoneReason: doneReason,
  594. PromptEvalCount: c.Timings.PromptN,
  595. PromptEvalDuration: parseDurationMs(c.Timings.PromptMS),
  596. EvalCount: c.Timings.PredictedN,
  597. EvalDuration: parseDurationMs(c.Timings.PredictedMS),
  598. })
  599. return nil
  600. }
  601. }
  602. }
  603. if err := scanner.Err(); err != nil {
  604. if strings.Contains(err.Error(), "unexpected EOF") {
  605. s.Close()
  606. msg := ""
  607. if s.status != nil && s.status.LastErrMsg != "" {
  608. msg = s.status.LastErrMsg
  609. }
  610. return fmt.Errorf("an unknown error was encountered while running the model %s", msg)
  611. }
  612. return fmt.Errorf("error reading llm response: %v", err)
  613. }
  614. if !retryNeeded {
  615. return nil // success
  616. }
  617. }
  618. // should never reach here ideally
  619. return fmt.Errorf("max retries exceeded")
  620. }
  621. type EmbeddingRequest struct {
  622. Content string `json:"content"`
  623. }
  624. type EmbeddingResponse struct {
  625. Embedding []float64 `json:"embedding"`
  626. }
  627. func (s *LlamaServer) Embedding(ctx context.Context, prompt string) ([]float64, error) {
  628. // Make sure the server is ready
  629. status, err := s.getServerStatus(ctx)
  630. if err != nil {
  631. return nil, err
  632. } else if status != ServerStatusReady {
  633. return nil, fmt.Errorf("unexpected server status: %d", status)
  634. }
  635. data, err := json.Marshal(TokenizeRequest{Content: prompt})
  636. if err != nil {
  637. return nil, fmt.Errorf("error marshaling embed data: %w", err)
  638. }
  639. req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/embedding", s.port), bytes.NewBuffer(data))
  640. if err != nil {
  641. return nil, fmt.Errorf("error creating embed request: %w", err)
  642. }
  643. req.Header.Set("Content-Type", "application/json")
  644. resp, err := http.DefaultClient.Do(req)
  645. if err != nil {
  646. return nil, fmt.Errorf("do embedding request: %w", err)
  647. }
  648. defer resp.Body.Close()
  649. body, err := io.ReadAll(resp.Body)
  650. if err != nil {
  651. return nil, fmt.Errorf("error reading embed response: %w", err)
  652. }
  653. if resp.StatusCode >= 400 {
  654. log.Printf("llm encode error: %s", body)
  655. return nil, fmt.Errorf("%s", body)
  656. }
  657. var embedding EmbeddingResponse
  658. if err := json.Unmarshal(body, &embedding); err != nil {
  659. return nil, fmt.Errorf("unmarshal tokenize response: %w", err)
  660. }
  661. return embedding.Embedding, nil
  662. }
  663. type TokenizeRequest struct {
  664. Content string `json:"content"`
  665. }
  666. type TokenizeResponse struct {
  667. Tokens []int `json:"tokens"`
  668. }
  669. func (s *LlamaServer) Tokenize(ctx context.Context, content string) ([]int, error) {
  670. // Make sure the server is ready
  671. status, err := s.getServerStatus(ctx)
  672. if err != nil {
  673. return nil, err
  674. } else if status != ServerStatusReady {
  675. return nil, fmt.Errorf("unexpected server status: %d", status)
  676. }
  677. data, err := json.Marshal(TokenizeRequest{Content: content})
  678. if err != nil {
  679. return nil, fmt.Errorf("marshaling encode data: %w", err)
  680. }
  681. req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/tokenize", s.port), bytes.NewBuffer(data))
  682. if err != nil {
  683. return nil, fmt.Errorf("encode request: %w", err)
  684. }
  685. req.Header.Set("Content-Type", "application/json")
  686. resp, err := http.DefaultClient.Do(req)
  687. if err != nil {
  688. return nil, fmt.Errorf("do encode request: %w", err)
  689. }
  690. defer resp.Body.Close()
  691. body, err := io.ReadAll(resp.Body)
  692. if err != nil {
  693. return nil, fmt.Errorf("read encode request: %w", err)
  694. }
  695. if resp.StatusCode >= 400 {
  696. log.Printf("llm encode error: %s", body)
  697. return nil, fmt.Errorf("%s", body)
  698. }
  699. var encoded TokenizeResponse
  700. if err := json.Unmarshal(body, &encoded); err != nil {
  701. return nil, fmt.Errorf("unmarshal encode response: %w", err)
  702. }
  703. return encoded.Tokens, nil
  704. }
  705. type DetokenizeRequest struct {
  706. Tokens []int `json:"tokens"`
  707. }
  708. type DetokenizeResponse struct {
  709. Content string `json:"content"`
  710. }
  711. func (s *LlamaServer) Detokenize(ctx context.Context, tokens []int) (string, error) {
  712. // Make sure the server is ready
  713. status, err := s.getServerStatus(ctx)
  714. if err != nil {
  715. return "", err
  716. } else if status != ServerStatusReady {
  717. return "", fmt.Errorf("unexpected server status: %d", status)
  718. }
  719. data, err := json.Marshal(DetokenizeRequest{Tokens: tokens})
  720. if err != nil {
  721. return "", fmt.Errorf("marshaling decode data: %w", err)
  722. }
  723. req, err := http.NewRequestWithContext(ctx, http.MethodPost, fmt.Sprintf("http://127.0.0.1:%d/detokenize", s.port), bytes.NewBuffer(data))
  724. if err != nil {
  725. return "", fmt.Errorf("decode request: %w", err)
  726. }
  727. req.Header.Set("Content-Type", "application/json")
  728. resp, err := http.DefaultClient.Do(req)
  729. if err != nil {
  730. return "", fmt.Errorf("do decode request: %w", err)
  731. }
  732. defer resp.Body.Close()
  733. body, err := io.ReadAll(resp.Body)
  734. if err != nil {
  735. return "", fmt.Errorf("read decode request: %w", err)
  736. }
  737. if resp.StatusCode >= 400 {
  738. log.Printf("llm decode error: %s", body)
  739. return "", fmt.Errorf("%s", body)
  740. }
  741. var decoded DetokenizeResponse
  742. if err := json.Unmarshal(body, &decoded); err != nil {
  743. return "", fmt.Errorf("unmarshal encode response: %w", err)
  744. }
  745. return decoded.Content, nil
  746. }
  747. func (s *LlamaServer) Close() error {
  748. if s.cmd != nil {
  749. slog.Debug("stopping llama server")
  750. return s.cmd.Process.Kill()
  751. }
  752. return nil
  753. }
  754. func parseDurationMs(ms float64) time.Duration {
  755. dur, err := time.ParseDuration(fmt.Sprintf("%fms", ms))
  756. if err != nil {
  757. panic(err)
  758. }
  759. return dur
  760. }