MuyueWorkspace/internal/api/handlers_rag.go

package api

import (
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/muyue/muyue/internal/config"
	"github.com/muyue/muyue/internal/rag"
)

func (s *Server) handleRAGIndex(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, `{"error":"method not allowed"}`, http.StatusMethodNotAllowed)
		return
	}

	s.ensureRAGStore()

	if r.Header.Get("Content-Type") == "application/json" {
		var req struct {
			Text string `json:"text"`
			Name string `json:"name"`
			Type string `json:"type"`
		}
		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
			jsonError(w, "invalid request: "+err.Error())
			return
		}
		if req.Text == "" {
			jsonError(w, "text is required")
			return
		}
		if req.Name == "" {
			req.Name = "document-" + time.Now().Format("20060102-150405")
		}
		if req.Type == "" {
			req.Type = "text"
		}

		s.indexText(w, req.Text, req.Name, req.Type)
		return
	}

	if err := r.ParseMultipartForm(32 << 20); err != nil {
		jsonError(w, "invalid multipart: "+err.Error())
		return
	}

	file, header, err := r.FormFile("file")
	if err != nil {
		jsonError(w, "file is required")
		return
	}
	defer file.Close()

	data, err := io.ReadAll(file)
	if err != nil {
		jsonError(w, "reading file: "+err.Error())
		return
	}

	name := header.Filename
	ext := strings.ToLower(filepath.Ext(name))
	docType := "text"
	switch ext {
	case ".md", ".markdown":
		docType = "markdown"
	case ".go", ".js", ".ts", ".py", ".java", ".rs", ".jsx", ".tsx":
		docType = "code"
	}

	s.indexText(w, string(data), name, docType)
}

func (s *Server) indexText(w http.ResponseWriter, text, name, docType string) {
	var chunks []rag.Chunk
	switch docType {
	case "markdown":
		chunks = rag.ChunkMarkdown(text, 500)
	case "code":
		lang := strings.TrimPrefix(filepath.Ext(name), ".")
		chunks = rag.ChunkCode(text, lang, 300)
	default:
		chunks = rag.ChunkText(text, 500)
	}

	if len(chunks) == 0 {
		jsonError(w, "no content to index")
		return
	}

	docID := uuid.New().String()[:8]
	doc := rag.Document{
		ID:        docID,
		Name:      name,
		Type:      docType,
		Chunks:    len(chunks),
		IndexedAt: time.Now(),
		Size:      int64(len(text)),
	}

	var chunkRecords []rag.ChunkRecord
	var texts []string
	for _, c := range chunks {
		texts = append(texts, c.Content)
		chunkRecords = append(chunkRecords, rag.ChunkRecord{
			DocumentID: docID,
			Content:    c.Content,
			StartPos:   c.StartPos,
			EndPos:     c.EndPos,
			Metadata:   c.Metadata,
		})
	}

	embClient := s.getEmbeddingClient()
	if embClient != nil {
		embeddings, err := embClient.Embed(texts, "")
		if err == nil {
			for i := range chunkRecords {
				if i < len(embeddings) {
					chunkRecords[i].Embedding = embeddings[i]
				}
			}
		}
	}

	if err := s.ragStore.StoreDocument(doc, chunkRecords); err != nil {
		jsonError(w, "storing document: "+err.Error())
		return
	}

	jsonResp(w, map[string]interface{}{
		"id":     docID,
		"name":   name,
		"chunks": len(chunks),
		"type":   docType,
	})
}

func (s *Server) handleRAGSearch(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		http.Error(w, `{"error":"method not allowed"}`, http.StatusMethodNotAllowed)
		return
	}

	s.ensureRAGStore()

	var req struct {
		Query string `json:"query"`
		Limit int    `json:"limit"`
	}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		jsonError(w, "invalid request: "+err.Error())
		return
	}
	if req.Query == "" {
		jsonError(w, "query is required")
		return
	}
	if req.Limit <= 0 {
		req.Limit = 5
	}

	embClient := s.getEmbeddingClient()
	var results []rag.SearchResult
	var err error

	if embClient != nil {
		queryEmb, embErr := embClient.EmbedSingle(req.Query, "")
		if embErr == nil {
			results, err = s.ragStore.Search(queryEmb, req.Limit)
		}
	}

	if err != nil || len(results) == 0 {
		results, err = s.ragStore.SearchKeyword(req.Query, req.Limit)
		if err != nil {
			jsonError(w, "search error: "+err.Error())
			return
		}
	}

	jsonResp(w, map[string]interface{}{
		"results": results,
		"query":   req.Query,
		"count":   len(results),
	})
}

func (s *Server) handleRAGStatus(w http.ResponseWriter, r *http.Request) {
	s.ensureRAGStore()
	status, err := s.ragStore.Status()
	if err != nil {
		jsonError(w, "status error: "+err.Error())
		return
	}
	jsonResp(w, status)
}

func (s *Server) handleRAGDelete(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodDelete {
		http.Error(w, `{"error":"method not allowed"}`, http.StatusMethodNotAllowed)
		return
	}

	s.ensureRAGStore()

	id := strings.TrimPrefix(r.URL.Path, "/api/rag/index/")
	if id == "" {
		jsonError(w, "document id is required")
		return
	}

	if err := s.ragStore.DeleteDocument(id); err != nil {
		jsonError(w, "delete error: "+err.Error())
		return
	}
	jsonResp(w, map[string]interface{}{"deleted": id})
}

func (s *Server) ensureRAGStore() {
	if s.ragStore != nil {
		return
	}
	configDir, err := config.ConfigDir()
	if err != nil {
		return
	}
	store, err := rag.NewStore(configDir)
	if err != nil {
		fmt.Fprintf(os.Stderr, "RAG store init error: %v\n", err)
		return
	}
	s.ragStore = store
}

func (s *Server) getEmbeddingClient() *rag.EmbeddingClient {
	for _, p := range s.config.AI.Providers {
		if p.Active && p.APIKey != "" {
			baseURL := p.BaseURL
			if baseURL == "" {
				baseURL = "https://api.openai.com/v1"
			}
			return rag.NewEmbeddingClient(p.APIKey, baseURL)
		}
	}
	return nil
}

func (s *Server) handleRAGDocuments(w http.ResponseWriter, r *http.Request) {
	s.ensureRAGStore()
	docs, err := s.ragStore.ListDocuments()
	if err != nil {
		jsonError(w, "list error: "+err.Error())
		return
	}
	if docs == nil {
		docs = []rag.Document{}
	}
	jsonResp(w, map[string]interface{}{"documents": docs})
}