WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

internal/models/utils/ollama/ollama.go

@@ -0,0 +1,365 @@
+package ollama
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net"
+	"net/http"
+	"net/url"
+	"os"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/Tencent/WeKnora/internal/logger"
+	"github.com/ollama/ollama/api"
+)
+
+// OllamaService manages Ollama service
+type OllamaService struct {
+	client      *api.Client
+	baseURL     string
+	mu          sync.Mutex
+	isAvailable bool
+	isOptional  bool // Added: marks if Ollama service is optional
+}
+
+// GetOllamaService gets Ollama service instance (singleton pattern)
+func GetOllamaService() (*OllamaService, error) {
+	// Get Ollama base URL from environment variable, if not set use provided baseURL or default value
+	logger.GetLogger(context.Background()).Infof("Ollama base URL: %s", os.Getenv("OLLAMA_BASE_URL"))
+	baseURL := "http://localhost:11434"
+	envURL := os.Getenv("OLLAMA_BASE_URL")
+	if envURL != "" {
+		baseURL = envURL
+	}
+
+	// Create URL object
+	parsedURL, err := url.Parse(baseURL)
+	if err != nil {
+		return nil, fmt.Errorf("invalid Ollama service URL: %w", err)
+	}
+
+	// Dedicated HTTP client for Ollama instead of http.DefaultClient.
+	// - Dial timeout prevents hanging when Ollama process is down or port unreachable
+	// - No overall Timeout so long-running streaming calls are controlled by context cancellation
+	ollamaHTTPClient := &http.Client{
+		Transport: &http.Transport{
+			DialContext:         (&net.Dialer{Timeout: 10 * time.Second}).DialContext,
+			TLSHandshakeTimeout: 10 * time.Second,
+			IdleConnTimeout:     90 * time.Second,
+			MaxIdleConns:        10,
+		},
+	}
+	client := api.NewClient(parsedURL, ollamaHTTPClient)
+
+	// Check if Ollama is set as optional
+	isOptional := false
+	if os.Getenv("OLLAMA_OPTIONAL") == "true" {
+		isOptional = true
+		logger.GetLogger(context.Background()).Info("Ollama service set to optional mode")
+	}
+
+	service := &OllamaService{
+		client:     client,
+		baseURL:    baseURL,
+		isOptional: isOptional,
+	}
+
+	return service, nil
+}
+
+// StartService checks if Ollama service is available
+func (s *OllamaService) StartService(ctx context.Context) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	// Check if service is available
+	err := s.client.Heartbeat(ctx)
+	if err != nil {
+		logger.GetLogger(ctx).Warnf("ollama service unavailable: %v", err)
+		s.isAvailable = false
+
+		// If configured as optional, don't return an error
+		if s.isOptional {
+			logger.GetLogger(ctx).Info("ollama service set as optional, will continue running the application")
+			return nil
+		}
+
+		return fmt.Errorf("ollama service unavailable: %w", err)
+	}
+
+	s.isAvailable = true
+	return nil
+}
+
+// IsAvailable returns whether the service is available
+func (s *OllamaService) IsAvailable() bool {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	return s.isAvailable
+}
+
+// IsModelAvailable checks if a model is available
+func (s *OllamaService) IsModelAvailable(ctx context.Context, modelName string) (bool, error) {
+	// First check if the service is available
+	if err := s.StartService(ctx); err != nil {
+		return false, err
+	}
+
+	// If service is not available but set as optional, return false but no error
+	if !s.isAvailable && s.isOptional {
+		return false, nil
+	}
+
+	// Get model list
+	listResp, err := s.client.List(ctx)
+	if err != nil {
+		return false, fmt.Errorf("failed to get model list: %w", err)
+	}
+
+	// If no version is specified for the model, add ":latest" by default
+	checkModelName := modelName
+	if !strings.Contains(modelName, ":") {
+		checkModelName = modelName + ":latest"
+	}
+	// Check if model is in the list
+	for _, model := range listResp.Models {
+		if model.Name == checkModelName {
+			return true, nil
+		}
+	}
+
+	return false, nil
+}
+
+// PullModel pulls a model
+func (s *OllamaService) PullModel(ctx context.Context, modelName string) error {
+	// First check if the service is available
+	if err := s.StartService(ctx); err != nil {
+		return err
+	}
+
+	// If service is not available but set as optional, return nil without further operations
+	if !s.isAvailable && s.isOptional {
+		logger.GetLogger(ctx).Warnf("Ollama service unavailable, unable to pull model %s", modelName)
+		return nil
+	}
+
+	// Check if model already exists
+	available, err := s.IsModelAvailable(ctx, modelName)
+	if err != nil {
+		return err
+	}
+	if available {
+		logger.GetLogger(ctx).Infof("Model %s already exists", modelName)
+		return nil
+	}
+
+	// Use official client to pull model
+	pullReq := &api.PullRequest{
+		Name: modelName,
+	}
+
+	err = s.client.Pull(ctx, pullReq, func(progress api.ProgressResponse) error {
+		if progress.Status != "" {
+			if progress.Total > 0 && progress.Completed > 0 {
+				percentage := float64(progress.Completed) / float64(progress.Total) * 100
+				logger.GetLogger(ctx).Infof("Pull progress: %s (%.2f%%)",
+					progress.Status, percentage)
+			} else {
+				logger.GetLogger(ctx).Infof("Pull status: %s", progress.Status)
+			}
+		}
+
+		if progress.Total > 0 && progress.Completed == progress.Total {
+			logger.GetLogger(ctx).Infof("Model %s pull completed", modelName)
+		}
+
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("failed to pull model: %w", err)
+	}
+
+	return nil
+}
+
+// EnsureModelAvailable ensures the model is available, pulls it if not available
+func (s *OllamaService) EnsureModelAvailable(ctx context.Context, modelName string) error {
+	// If service is not available but set as optional, return nil directly
+	if !s.IsAvailable() && s.isOptional {
+		logger.GetLogger(ctx).Warnf("Ollama service unavailable, skipping ensuring model %s availability", modelName)
+		return nil
+	}
+
+	available, err := s.IsModelAvailable(ctx, modelName)
+	if err != nil {
+		if s.isOptional {
+			logger.GetLogger(ctx).
+				Warnf("Failed to check model %s availability, but Ollama is set as optional", modelName)
+			return nil
+		}
+		return err
+	}
+
+	if !available {
+		return s.PullModel(ctx, modelName)
+	}
+
+	return nil
+}
+
+// GetVersion gets Ollama version
+func (s *OllamaService) GetVersion(ctx context.Context) (string, error) {
+	// If service is not available but set as optional, return empty version info
+	if !s.IsAvailable() && s.isOptional {
+		return "unavailable", nil
+	}
+
+	version, err := s.client.Version(ctx)
+	if err != nil {
+		return "", fmt.Errorf("failed to get Ollama version: %w", err)
+	}
+	return version, nil
+}
+
+// CreateModel creates a custom model
+func (s *OllamaService) CreateModel(ctx context.Context, name, modelfile string) error {
+	req := &api.CreateRequest{
+		Model:    name,
+		Template: modelfile, // Use Template field instead of Modelfile
+	}
+
+	err := s.client.Create(ctx, req, func(progress api.ProgressResponse) error {
+		if progress.Status != "" {
+			logger.GetLogger(ctx).Infof("Model creation status: %s", progress.Status)
+		}
+		return nil
+	})
+	if err != nil {
+		return fmt.Errorf("failed to create model: %w", err)
+	}
+
+	return nil
+}
+
+// GetModelInfo gets model information
+func (s *OllamaService) GetModelInfo(ctx context.Context, modelName string) (*api.ShowResponse, error) {
+	req := &api.ShowRequest{
+		Name: modelName,
+	}
+
+	resp, err := s.client.Show(ctx, req)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get model information: %w", err)
+	}
+
+	return resp, nil
+}
+
+// OllamaModelInfo represents detailed information about an Ollama model
+type OllamaModelInfo struct {
+	Name       string    `json:"name"`
+	Size       int64     `json:"size"`
+	Digest     string    `json:"digest"`
+	ModifiedAt time.Time `json:"modified_at"`
+}
+
+// ListModels lists all available models with basic info (names only)
+func (s *OllamaService) ListModels(ctx context.Context) ([]string, error) {
+	listResp, err := s.client.List(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get model list: %w", err)
+	}
+
+	modelNames := make([]string, len(listResp.Models))
+	for i, model := range listResp.Models {
+		modelNames[i] = model.Name
+	}
+
+	return modelNames, nil
+}
+
+// ListModelsDetailed lists all available models with detailed information
+func (s *OllamaService) ListModelsDetailed(ctx context.Context) ([]OllamaModelInfo, error) {
+	listResp, err := s.client.List(ctx)
+	if err != nil {
+		return nil, fmt.Errorf("failed to get model list: %w", err)
+	}
+	jsonData, err := json.Marshal(listResp.Models)
+	if err != nil {
+		return nil, fmt.Errorf("failed to marshal model list: %w", err)
+	}
+	logger.GetLogger(ctx).Infof("List models detailed: %s", string(jsonData))
+
+	models := make([]OllamaModelInfo, len(listResp.Models))
+	for i, model := range listResp.Models {
+		models[i] = OllamaModelInfo{
+			Name:       model.Name,
+			Size:       model.Size,
+			Digest:     model.Digest,
+			ModifiedAt: model.ModifiedAt,
+		}
+	}
+
+	return models, nil
+}
+
+// DeleteModel deletes a model
+func (s *OllamaService) DeleteModel(ctx context.Context, modelName string) error {
+	req := &api.DeleteRequest{
+		Name: modelName,
+	}
+
+	err := s.client.Delete(ctx, req)
+	if err != nil {
+		return fmt.Errorf("failed to delete model: %w", err)
+	}
+
+	return nil
+}
+
+// IsValidModelName checks if model name is valid
+func IsValidModelName(name string) bool {
+	// Simple check for model name format
+	return name != "" && !strings.Contains(name, " ")
+}
+
+// Chat uses Ollama chat
+func (s *OllamaService) Chat(ctx context.Context, req *api.ChatRequest, fn api.ChatResponseFunc) error {
+	// First check if service is available
+	if err := s.StartService(ctx); err != nil {
+		return err
+	}
+
+	// Use official client Chat method
+	return s.client.Chat(ctx, req, fn)
+}
+
+// Embeddings gets text embedding vectors
+func (s *OllamaService) Embeddings(ctx context.Context, req *api.EmbedRequest) (*api.EmbedResponse, error) {
+	// First check if service is available
+	if err := s.StartService(ctx); err != nil {
+		return nil, err
+	}
+	// Use official client Embed method
+	return s.client.Embed(ctx, req)
+}
+
+// Generate generates text (used for Rerank)
+func (s *OllamaService) Generate(ctx context.Context, req *api.GenerateRequest, fn api.GenerateResponseFunc) error {
+	// First check if service is available
+	if err := s.StartService(ctx); err != nil {
+		return err
+	}
+
+	// Use official client Generate method
+	return s.client.Generate(ctx, req, fn)
+}
+
+// GetClient returns the underlying ollama client for advanced operations
+func (s *OllamaService) GetClient() *api.Client {
+	return s.client
+}

internal/models/utils/signer.go

@@ -0,0 +1,93 @@
+package utils
+
+import (
+	"bytes"
+	"crypto/md5"
+	"fmt"
+	"math/rand"
+	"sort"
+	"strconv"
+	"strings"
+	"time"
+)
+
+const (
+	nonceChars  = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
+	nonceLength = 16
+)
+
+// Sign 按 WeKnoraCloud 参考实现生成请求头。
+// appID: 上游 APPID
+// apiKey: 上游 API Key（当前沿用 AppSecret 字段承载）
+// requestID: 每次请求唯一的 UUID 字符串
+// bodyJSON: 请求体 JSON 字符串，空请求体传 "{}"
+func Sign(appID, apiKey, requestID, bodyJSON string) map[string]string {
+	timestamp := strconv.FormatInt(time.Now().Unix(), 10)
+	nonce := generateNonce(nonceLength)
+
+	bodyForHash := bodyJSON
+	if bodyForHash == "" {
+		bodyForHash = "{}"
+	}
+	bodyMD5 := md5Hex(bodyForHash)
+
+	params := map[string]string{
+		"x-appid":      appID,
+		"x-api-key":    apiKey,
+		"x-request-id": requestID,
+		"x-timestamp":  timestamp,
+		"x-nonce":      nonce,
+		"body":         bodyMD5,
+	}
+
+	keys := make([]string, 0, len(params))
+	for k := range params {
+		keys = append(keys, k)
+	}
+	sort.Strings(keys)
+
+	parts := make([]string, 0, len(keys))
+	for _, k := range keys {
+		parts = append(parts, rfc3986Encode(k)+"="+rfc3986Encode(params[k]))
+	}
+	signature := md5Hex(strings.Join(parts, "&"))
+
+	return map[string]string{
+		"X-APPID":      appID,
+		"X-API-Key":    apiKey,
+		"X-Request-ID": requestID,
+		"X-Timestamp":  timestamp,
+		"X-Nonce":      nonce,
+		"X-Signature":  signature,
+	}
+}
+
+func md5Hex(s string) string {
+	h := md5.New()
+	h.Write([]byte(s))
+	return fmt.Sprintf("%x", h.Sum(nil))
+}
+
+func generateNonce(length int) string {
+	b := make([]byte, length)
+	for i := range b {
+		b[i] = nonceChars[rand.Intn(len(nonceChars))]
+	}
+	return string(b)
+}
+
+// rfc3986Encode 对字符串做 RFC3986 编码
+// 保留字符：A-Z a-z 0-9 - _ . ~
+func rfc3986Encode(s string) string {
+	var buf bytes.Buffer
+	for _, r := range s {
+		if (r >= 'A' && r <= 'Z') || (r >= 'a' && r <= 'z') ||
+			(r >= '0' && r <= '9') || r == '-' || r == '_' ||
+			r == '.' || r == '~' {
+			buf.WriteRune(r)
+		} else {
+			fmt.Fprintf(&buf, "%%%02X", r)
+		}
+	}
+	return buf.String()
+}

internal/models/utils/slices.go

@@ -0,0 +1,36 @@
+package utils
+
+// ChunkSlice splits a slice into multiple sub-slices of the specified size
+func ChunkSlice[T any](slice []T, chunkSize int) [][]T {
+	// Handle edge cases
+	if len(slice) == 0 {
+		return [][]T{}
+	}
+
+	if chunkSize <= 0 {
+		panic("chunkSize must be greater than 0")
+	}
+
+	// Calculate how many sub-slices are needed
+	chunks := make([][]T, 0, (len(slice)+chunkSize-1)/chunkSize)
+
+	// Split the slice
+	for i := 0; i < len(slice); i += chunkSize {
+		end := i + chunkSize
+		if end > len(slice) {
+			end = len(slice)
+		}
+		chunks = append(chunks, slice[i:end])
+	}
+
+	return chunks
+}
+
+// MapSlice applies a function to each element of a slice and returns a new slice with the results
+func MapSlice[A any, B any](in []A, f func(A) B) []B {
+	out := make([]B, 0, len(in))
+	for _, item := range in {
+		out = append(out, f(item))
+	}
+	return out
+}