poimen-workflows/action/query_graph_rag.go

package action

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

// GraphRAGQueryInput for Memory System endpoint
type GraphRAGQueryInput struct {
	WorkflowID      string                    `json:"workflow_id"`
	Query           string                    `json:"query"`
	SearchType      string                    `json:"search_type"`
	RelationType    string                    `json:"relation_type"`
	ConfidenceFloor float64                   `json:"confidence_floor"`
	TopK            int                       `json:"top_k"`
	RankingProfile  string                    `json:"ranking_profile"`
	Canvas          CanvasWithRelationsData   `json:"canvas"`
}

// GraphRAGQueryOutput from Memory System
type GraphRAGQueryOutput struct {
	WorkflowID   string             `json:"workflow_id"`
	Query        string             `json:"query"`
	Edges        []EdgeWithWording  `json:"edges"`
	Paths        []QueryPathData    `json:"paths"`
	TotalCount   int                `json:"total_count"`
	HasMore      bool               `json:"has_more"`
	ExecutionMs  int64              `json:"execution_time_ms"`
}

type QueryPathData struct {
	SourceID   string   `json:"source_id"`
	TargetID   string   `json:"target_id"`
	Distance   int      `json:"distance"`
	PathCount  int      `json:"path_count"`
	NodeIDs    []string `json:"node_ids"`
	Confidence float64  `json:"total_confidence"`
}

// QueryGraphRAGActivity queries Memory System for semantic relations
func QueryGraphRAGActivity(ctx context.Context, input GraphRAGQueryInput) (GraphRAGQueryOutput, error) {
	logger := newActivityLogger(ctx)
	output := GraphRAGQueryOutput{
		WorkflowID: input.WorkflowID,
		Query:      input.Query,
		Edges:      []EdgeWithWording{},
		Paths:      []QueryPathData{},
	}

	logger.logf("info", "Querying GraphRAG: %s", input.Query)

	// Get Memory Service URL from env
	memoryURL := os.Getenv("MEMORY_SERVICE_URL")
	if memoryURL == "" {
		memoryURL = "http://localhost:8000"
	}

	// Build payload for Memory System
	payload := map[string]interface{}{
		"workflow_id":      input.WorkflowID,
		"query":            input.Query,
		"search_type":      input.SearchType,
		"relation_type":    input.RelationType,
		"confidence_floor": input.ConfidenceFloor,
		"top_k":            input.TopK,
		"ranking_profile":  input.RankingProfile,
		"canvas_nodes":     input.Canvas.Nodes,
		"canvas_edges":     input.Canvas.Edges,
		"relations":        input.Canvas.Relations,
	}

	reqBody, err := json.Marshal(payload)
	if err != nil {
		return output, fmt.Errorf("failed to marshal payload: %w", err)
	}

	// Call Memory System unified query endpoint
	req, err := http.NewRequestWithContext(
		ctx,
		"POST",
		memoryURL+"/workflows/query",
		bytes.NewReader(reqBody),
	)
	if err != nil {
		return output, fmt.Errorf("failed to create request: %w", err)
	}

	req.Header.Set("Content-Type", "application/json")
	if token := ctx.Value("jwt_token"); token != nil {
		req.Header.Set("Authorization", fmt.Sprintf("Bearer %v", token))
	}

	startTime := time.Now()
	client := &http.Client{Timeout: 30 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return output, fmt.Errorf("failed to call Memory Service: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != 200 {
		body, _ := io.ReadAll(resp.Body)
		return output, fmt.Errorf("Memory Service returned %d: %s", resp.StatusCode, string(body))
	}

	// Parse response
	var graphResp struct {
		Edges      []EdgeWithWording `json:"edges"`
		Paths      []QueryPathData   `json:"paths"`
		TotalCount int               `json:"total_count"`
		HasMore    bool              `json:"has_more"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&graphResp); err != nil {
		return output, fmt.Errorf("failed to decode response: %w", err)
	}

	output.Edges = graphResp.Edges
	output.Paths = graphResp.Paths
	output.TotalCount = graphResp.TotalCount
	output.HasMore = graphResp.HasMore
	output.ExecutionMs = time.Since(startTime).Milliseconds()

	logger.logf("info", "GraphRAG returned %d edges, %d paths in %dms", 
		len(output.Edges), len(output.Paths), output.ExecutionMs)
	return output, nil
}