Ollamaを用いたチャットバックエンドのストリーミング応答実装

環境設定

Ollamaコンテナの起動:

docker run -d -p 3000:8080 -p 11434:11434 -v ollama:/root/.ollama -v open-webui:/app/backend/data --name open-webui --restart always ollama/ollama

Gemmaモデルのダウンロード:

docker exec -it open-webui ollama run gemma:2b

ストリーミング応答の仕組み

クライアント側の実装

const [response, controller] = await createChatSession(authToken, {
  model: selectedModel,
  messages: conversationHistory,
  options: {...configuration}
});

if (response && response.ok) {
  const streamReader = response.body
    .pipeThrough(new TextDecoderStream())
    .pipeThrough(lineSplitter('\n'))
    .getReader();
}

この実装ではサーバーからのストリームデータを逐次処理し、チャットメッセージを段階的に表示します。

サーバー側の実装

Ginフレームワークによるストリーミング応答:

c.Stream(func(w io.Writer) bool {
  select {
  case message, active := <-messageChannel:
    if !active {
      return false
    }
    c.SSEvent("message_update", message)
    return true
  case <-c.Request.Context().Done():
    return false
  }
})

Ollama APIとの連携:

streamHandler := llms.WithStreamingFunc(func(ctx context.Context, dataChunk []byte) error {
  select {
  case messageChannel <- string(dataChunk):
    return nil
  case <-ctx.Done():
    return ctx.Err()
  }
})

_, err := ollamaClient.Invoke(ctx, userInput, streamHandler)

完全な実装例

package main

import (
  "context"
  "net/http"
  
  "github.com/gin-gonic/gin"
  "github.com/tmc/langchaingo/llms"
  "github.com/tmc/langchaingo/llms/ollama"
)

func main() {
  server := gin.Default()
  server.POST("/chat", handleChatRequest)
  server.Run(":8080")
}

type ChatInput struct {
  Content string `json:"content"`
}

func handleChatRequest(c *gin.Context) {
  var input ChatInput
  if err := c.BindJSON(&input); err != nil {
    c.AbortWithStatusJSON(http.StatusBadRequest, gin.H{"error": "invalid request"})
    return
  }

  responseChannel := make(chan string)
  go processQuery(input.Content, responseChannel)
  
  c.Stream(func(w io.Writer) bool {
    select {
    case responsePart, valid := <-responseChannel:
      if !valid {
        return false
      }
      c.SSEvent("response_chunk", responsePart)
      return true
    case <-c.Request.Context().Done():
      return false
    }
  })
}

var ollamaInstance *ollama.LLM

func initOllamaConnection() {
  conn, err := ollama.New(
    ollama.WithModel("gemma:2b"),
    ollama.WithServerURL("http://ollama-host:11434")
  )
  if err != nil {
    panic("connection failed: " + err.Error())
  }
  ollamaInstance = conn
}

func processQuery(query string, outputChan chan string) {
  streamCallback := llms.WithStreamingFunc(func(ctx context.Context, chunk []byte) error {
    select {
    case outputChan <- string(chunk):
      return nil
    case <-ctx.Done():
      return ctx.Err()
    }
  })

  _, err := ollamaInstance.Call(context.Background(), query, streamCallback)
  if err != nil {
    log.Println("execution error:", err)
  }
  close(outputChan)
}

タグ: Ollama Gin Docker ストリーミング バックエンド

7月25日 03:11 投稿