CHORUS/pkg/mcp/lightrag_client.go

package mcp

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"time"
)

// LightRAGClient provides access to LightRAG MCP server
type LightRAGClient struct {
	baseURL    string
	httpClient *http.Client
	apiKey     string // Optional API key for authentication
}

// LightRAGConfig holds configuration for LightRAG client
type LightRAGConfig struct {
	BaseURL string        // e.g., "http://127.0.0.1:9621"
	Timeout time.Duration // HTTP timeout
	APIKey  string        // Optional API key
}

// QueryMode represents LightRAG query modes
type QueryMode string

const (
	QueryModeNaive  QueryMode = "naive"    // Simple semantic search
	QueryModeLocal  QueryMode = "local"    // Local graph traversal
	QueryModeGlobal QueryMode = "global"   // Global graph analysis
	QueryModeHybrid QueryMode = "hybrid"   // Combined approach
)

// QueryRequest represents a LightRAG query request
type QueryRequest struct {
	Query      string    `json:"query"`
	Mode       QueryMode `json:"mode"`
	OnlyNeedContext bool  `json:"only_need_context,omitempty"`
}

// QueryResponse represents a LightRAG query response
type QueryResponse struct {
	Response string `json:"response"`
	Context  string `json:"context,omitempty"`
}

// InsertRequest represents a LightRAG document insertion request
type InsertRequest struct {
	Text        string `json:"text"`
	Description string `json:"description,omitempty"`
}

// InsertResponse represents a LightRAG insertion response
type InsertResponse struct {
	Success bool   `json:"success"`
	Message string `json:"message"`
	Status  string `json:"status"`
}

// HealthResponse represents LightRAG health check response
type HealthResponse struct {
	Status            string                 `json:"status"`
	WorkingDirectory  string                 `json:"working_directory"`
	InputDirectory    string                 `json:"input_directory"`
	Configuration     map[string]interface{} `json:"configuration"`
	AuthMode          string                 `json:"auth_mode"`
	PipelineBusy      bool                   `json:"pipeline_busy"`
	KeyedLocks        map[string]interface{} `json:"keyed_locks"`
	CoreVersion       string                 `json:"core_version"`
	APIVersion        string                 `json:"api_version"`
	WebUITitle        string                 `json:"webui_title"`
	WebUIDescription  string                 `json:"webui_description"`
}

// NewLightRAGClient creates a new LightRAG MCP client
func NewLightRAGClient(config LightRAGConfig) *LightRAGClient {
	if config.Timeout == 0 {
		config.Timeout = 30 * time.Second
	}

	return &LightRAGClient{
		baseURL: config.BaseURL,
		httpClient: &http.Client{
			Timeout: config.Timeout,
		},
		apiKey: config.APIKey,
	}
}

// Query performs a RAG query against LightRAG
func (c *LightRAGClient) Query(ctx context.Context, query string, mode QueryMode) (*QueryResponse, error) {
	req := QueryRequest{
		Query: query,
		Mode:  mode,
	}

	respData, err := c.post(ctx, "/query", req)
	if err != nil {
		return nil, fmt.Errorf("query failed: %w", err)
	}

	var response QueryResponse
	if err := json.Unmarshal(respData, &response); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	return &response, nil
}

// QueryWithContext performs a RAG query and returns both response and context
func (c *LightRAGClient) QueryWithContext(ctx context.Context, query string, mode QueryMode) (*QueryResponse, error) {
	req := QueryRequest{
		Query:           query,
		Mode:            mode,
		OnlyNeedContext: false, // Get both response and context
	}

	respData, err := c.post(ctx, "/query", req)
	if err != nil {
		return nil, fmt.Errorf("query with context failed: %w", err)
	}

	var response QueryResponse
	if err := json.Unmarshal(respData, &response); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	return &response, nil
}

// GetContext retrieves context without generating a response
func (c *LightRAGClient) GetContext(ctx context.Context, query string, mode QueryMode) (string, error) {
	req := QueryRequest{
		Query:           query,
		Mode:            mode,
		OnlyNeedContext: true,
	}

	respData, err := c.post(ctx, "/query", req)
	if err != nil {
		return "", fmt.Errorf("get context failed: %w", err)
	}

	var response QueryResponse
	if err := json.Unmarshal(respData, &response); err != nil {
		return "", fmt.Errorf("failed to parse response: %w", err)
	}

	return response.Context, nil
}

// Insert adds a document to the LightRAG knowledge base
func (c *LightRAGClient) Insert(ctx context.Context, text, description string) error {
	req := InsertRequest{
		Text:        text,
		Description: description,
	}

	respData, err := c.post(ctx, "/insert", req)
	if err != nil {
		return fmt.Errorf("insert failed: %w", err)
	}

	var response InsertResponse
	if err := json.Unmarshal(respData, &response); err != nil {
		return fmt.Errorf("failed to parse insert response: %w", err)
	}

	if !response.Success {
		return fmt.Errorf("insert failed: %s", response.Message)
	}

	return nil
}

// Health checks the health of the LightRAG server
func (c *LightRAGClient) Health(ctx context.Context) (*HealthResponse, error) {
	respData, err := c.get(ctx, "/health")
	if err != nil {
		return nil, fmt.Errorf("health check failed: %w", err)
	}

	var response HealthResponse
	if err := json.Unmarshal(respData, &response); err != nil {
		return nil, fmt.Errorf("failed to parse health response: %w", err)
	}

	return &response, nil
}

// IsHealthy checks if LightRAG server is healthy
func (c *LightRAGClient) IsHealthy(ctx context.Context) bool {
	health, err := c.Health(ctx)
	if err != nil {
		return false
	}
	return health.Status == "healthy"
}

// post performs an HTTP POST request
func (c *LightRAGClient) post(ctx context.Context, endpoint string, body interface{}) ([]byte, error) {
	jsonData, err := json.Marshal(body)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	req, err := http.NewRequestWithContext(ctx, "POST", c.baseURL+endpoint, bytes.NewBuffer(jsonData))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/json")
	if c.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.apiKey)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	respData, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("request failed with status %d: %s", resp.StatusCode, string(respData))
	}

	return respData, nil
}

// get performs an HTTP GET request
func (c *LightRAGClient) get(ctx context.Context, endpoint string) ([]byte, error) {
	req, err := http.NewRequestWithContext(ctx, "GET", c.baseURL+endpoint, nil)
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	if c.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.apiKey)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	respData, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("request failed with status %d: %s", resp.StatusCode, string(respData))
	}

	return respData, nil
}