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}) }