feat: GraphRAG query workflow and indexing

2026-09-05 05:52:56 -07:00
parent 84b4ca120f
commit e01dad4e8c
2 changed files with 389 additions and 0 deletions
@@ -0,0 +1,283 @@
+package action
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"time"
+)
+
+// IndexGraphRAGInput sends workflow relations to GraphRAG for indexing
+type IndexGraphRAGInput struct {
+	WorkflowID string           `json:"workflow_id"`
+	Version    int              `json:"version"`
+	Nodes      []db.WorkflowNode     `json:"nodes"`
+	Relations  []EdgeWithWording `json:"relations"`
+}
+
+// IndexGraphRAGOutput confirms indexing status
+type IndexGraphRAGOutput struct {
+	WorkflowID       string `json:"workflow_id"`
+	Version          int    `json:"version"`
+	IndexedEntities  int    `json:"indexed_entities"`
+	IndexedEdges     int    `json:"indexed_edges"`
+	Status           string `json:"status"`           // indexed, partial, failed
+	GraphRAGChecksum string `json:"graph_rag_checksum"`
+	IndexedAt        string `json:"indexed_at"`
+}
+
+// IndexGraphRAGActivity indexes workflow canvas to GraphRAG
+func IndexGraphRAGActivity(ctx context.Context, input IndexGraphRAGInput) (IndexGraphRAGOutput, error) {
+	logger := newActivityLogger(ctx)
+	output := IndexGraphRAGOutput{
+		WorkflowID: input.WorkflowID,
+		Version:    input.Version,
+		Status:     "pending",
+	}
+
+	logger.logf("info", "Indexing workflow to GraphRAG: %s v%d", input.WorkflowID, input.Version)
+
+	// Build GraphRAG payload
+	payload := buildGraphRAGPayload(input)
+
+	// Call GraphRAG indexing endpoint
+	graphRAGURL := getEnv("GRAPH_RAG_URL", "http://localhost:8090") // GraphRAG service
+	token := ctx.Value("jwt_token").(string)
+
+	reqBody, err := json.Marshal(payload)
+	if err != nil {
+		return output, fmt.Errorf("failed to marshal payload: %w", err)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, "POST", graphRAGURL+"/index/workflow", bytes.NewReader(reqBody))
+	if err != nil {
+		return output, fmt.Errorf("failed to create request: %w", err)
+	}
+
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
+	req.Header.Set("Content-Type", "application/json")
+
+	client := &http.Client{Timeout: 60 * time.Second}
+	resp, err := client.Do(req)
+	if err != nil {
+		return output, fmt.Errorf("failed to call GraphRAG: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != 200 && resp.StatusCode != 202 {
+		body, _ := io.ReadAll(resp.Body)
+		return output, fmt.Errorf("GraphRAG returned %d: %s", resp.StatusCode, string(body))
+	}
+
+	// Parse response
+	var graphResp struct {
+		Status           string `json:"status"`
+		IndexedEntities  int    `json:"indexed_entities"`
+		IndexedEdges     int    `json:"indexed_edges"`
+		GraphRAGChecksum string `json:"graph_rag_checksum"`
+	}
+
+	if err := json.NewDecoder(resp.Body).Decode(&graphResp); err != nil {
+		return output, fmt.Errorf("failed to decode GraphRAG response: %w", err)
+	}
+
+	output.Status = graphResp.Status
+	output.IndexedEntities = graphResp.IndexedEntities
+	output.IndexedEdges = graphResp.IndexedEdges
+	output.GraphRAGChecksum = graphResp.GraphRAGChecksum
+	output.IndexedAt = time.Now().UTC().Format(time.RFC3339)
+
+	logger.logf("info", "GraphRAG indexed %d entities, %d edges (status: %s)", output.IndexedEntities, output.IndexedEdges, output.Status)
+	return output, nil
+}
+
+// buildGraphRAGPayload converts workflow to GraphRAG format
+func buildGraphRAGPayload(input IndexGraphRAGInput) map[string]interface{} {
+	// Convert nodes to entities
+	entities := []map[string]interface{}{}
+	for _, node := range input.Nodes {
+		entity := map[string]interface{}{
+			"id":   node.ID,
+			"name": node.Label,
+			"type": node.Type,
+			"metadata": map[string]interface{}{
+				"workflow_id": input.WorkflowID,
+				"version":     input.Version,
+				"node_id":     node.ID,
+				"node_type":   node.Type,
+			},
+		}
+		if node.Data != nil {
+			entity["data"] = node.Data
+		}
+		entities = append(entities, entity)
+	}
+
+	// Convert relations to edges
+	edges := []map[string]interface{}{}
+	for _, rel := range input.Relations {
+		edge := map[string]interface{}{
+			"id":     rel.ID,
+			"source": rel.Source,
+			"target": rel.Target,
+			"type":   rel.RelationType,
+			"label":  rel.RelationLabel,
+			"properties": map[string]interface{}{
+				"verb":           rel.RelationWording.Verb,
+				"source_output":  rel.RelationWording.SourceOutput,
+				"target_input":   rel.RelationWording.TargetInput,
+				"connection_type": rel.RelationWording.ConnectionType,
+				"confidence":     rel.RelationWording.Confidence,
+				"semantic_match": rel.RelationWording.SemanticMatch,
+			},
+			"metadata": map[string]interface{}{
+				"workflow_id":     input.WorkflowID,
+				"version":         input.Version,
+				"relation_type":   rel.RelationType,
+				"created_at":      rel.CreatedAt,
+			},
+		}
+		edges = append(edges, edge)
+	}
+
+	return map[string]interface{}{
+		"workflow_id": input.WorkflowID,
+		"version":     input.Version,
+		"entities":    entities,
+		"edges":       edges,
+		"metadata": map[string]interface{}{
+			"project":        "poimen",
+			"type":           "workflow_canvas",
+			"indexed_at":     time.Now().UTC().Format(time.RFC3339),
+			"total_entities": len(entities),
+			"total_edges":    len(edges),
+		},
+	}
+}
+
+// QueryGraphRAGRelationsInput for direct relation discovery
+type QueryGraphRAGRelationsInput struct {
+	WorkflowID string   `json:"workflow_id"`
+	Version    int      `json:"version"`
+	Query      string   `json:"query"`
+	TopK       int      `json:"top_k"`
+	Filters    map[string]interface{} `json:"filters,omitempty"`
+}
+
+// QueryGraphRAGRelationsOutput returns discovered relations
+type QueryGraphRAGRelationsOutput struct {
+	Query       string           `json:"query"`
+	Results     []EdgeWithWording `json:"results"`
+	TotalCount  int              `json:"total_count"`
+	ExecutionMs int64            `json:"execution_time_ms"`
+}
+
+// QueryGraphRAGRelationsActivity queries GraphRAG for relation patterns
+func QueryGraphRAGRelationsActivity(ctx context.Context, input QueryGraphRAGRelationsInput) (QueryGraphRAGRelationsOutput, error) {
+	logger := newActivityLogger(ctx)
+	output := QueryGraphRAGRelationsOutput{
+		Query:   input.Query,
+		Results: []EdgeWithWording{},
+	}
+
+	logger.logf("info", "Querying GraphRAG relations: %s", input.Query)
+
+	// Build query payload
+	payload := map[string]interface{}{
+		"workflow_id": input.WorkflowID,
+		"version":     input.Version,
+		"query":       input.Query,
+		"top_k":       input.TopK,
+	}
+	if input.Filters != nil {
+		payload["filters"] = input.Filters
+	}
+
+	// Call GraphRAG query endpoint
+	graphRAGURL := getEnv("GRAPH_RAG_URL", "http://localhost:8090")
+	token := ctx.Value("jwt_token").(string)
+
+	reqBody, err := json.Marshal(payload)
+	if err != nil {
+		return output, fmt.Errorf("failed to marshal payload: %w", err)
+	}
+
+	req, err := http.NewRequestWithContext(ctx, "POST", graphRAGURL+"/query/relations", bytes.NewReader(reqBody))
+	if err != nil {
+		return output, fmt.Errorf("failed to create request: %w", err)
+	}
+
+	req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", token))
+	req.Header.Set("Content-Type", "application/json")
+
+	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 GraphRAG: %w", err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != 200 {
+		body, _ := io.ReadAll(resp.Body)
+		return output, fmt.Errorf("GraphRAG returned %d: %s", resp.StatusCode, string(body))
+	}
+
+	// Parse response
+	var graphResp struct {
+		Results []map[string]interface{} `json:"results"`
+		Count   int                      `json:"count"`
+	}
+
+	if err := json.NewDecoder(resp.Body).Decode(&graphResp); err != nil {
+		return output, fmt.Errorf("failed to decode GraphRAG response: %w", err)
+	}
+
+	// Map results to EdgeWithWording
+	for _, result := range graphResp.Results {
+		edge := EdgeWithWording{}
+
+		if source, ok := result["source"].(string); ok {
+			edge.Source = source
+		}
+		if target, ok := result["target"].(string); ok {
+			edge.Target = target
+		}
+		if relType, ok := result["type"].(string); ok {
+			edge.RelationType = relType
+		}
+		if label, ok := result["label"].(string); ok {
+			edge.RelationLabel = label
+		}
+
+		// Extract wording from properties
+		if props, ok := result["properties"].(map[string]interface{}); ok {
+			edge.RelationWording.Verb, _ = props["verb"].(string)
+			edge.RelationWording.SourceOutput, _ = props["source_output"].(string)
+			edge.RelationWording.TargetInput, _ = props["target_input"].(string)
+			edge.RelationWording.ConnectionType, _ = props["connection_type"].(string)
+			if conf, ok := props["confidence"].(float64); ok {
+				edge.RelationWording.Confidence = conf
+			}
+			edge.RelationWording.SemanticMatch, _ = props["semantic_match"].(string)
+		}
+
+		output.Results = append(output.Results, edge)
+	}
+
+	output.TotalCount = graphResp.Count
+	output.ExecutionMs = time.Since(startTime).Milliseconds()
+
+	logger.logf("info", "GraphRAG returned %d relations in %dms", output.TotalCount, output.ExecutionMs)
+	return output, nil
+}
+
+func getEnv(key, defaultVal string) string {
+	if val := os.Getenv(key); val != "" {
+		return val
+	}
+	return defaultVal
+}
@@ -0,0 +1,106 @@
+package statemachine
+
+import (
+	"context"
+	"time"
+
+	"go.temporal.io/sdk/workflow"
+	"github.com/rockliang/poimen/workflows/action"
+	"github.com/rockliang/poimen/workflows/pkg/db"
+)
+
+type WorkflowGraphQueryInput struct {
+	WorkflowID      string  `json:"workflow_id"`
+	Query           string  `json:"query"`
+	SearchType      string  `json:"search_type"`
+	RelationType    string  `json:"relation_type"`
+	Version         int     `json:"version"`
+	ConfidenceFloor float64 `json:"confidence_floor"`
+	TopK            int     `json:"top_k"`
+	FindPaths       bool    `json:"find_paths"`
+	TargetNodeID    string  `json:"target_node_id"`
+	MaxPathDepth    int     `json:"max_path_depth"`
+	RankingProfile  string  `json:"ranking_profile"`
+	IncludeReasoning bool   `json:"include_reasoning"`
+}
+
+type WorkflowGraphQueryOutput struct {
+	WorkflowID      string                `json:"workflow_id"`
+	Query           string                `json:"query"`
+	Version         int                   `json:"version"`
+	ExecutionTimeMs int64                 `json:"execution_time_ms"`
+	Results         []action.EdgeWithWording `json:"results"`
+	Paths           []QueryPath           `json:"paths"`
+	TotalCount      int                   `json:"total_count"`
+	HasMore         bool                  `json:"has_more"`
+	RankingProfile  string                `json:"ranking_profile"`
+}
+
+type QueryPath 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"`
+}
+
+func WorkflowGraphQuery(ctx workflow.Context, input WorkflowGraphQueryInput) (WorkflowGraphQueryOutput, error) {
+	startTime := time.Now()
+	output := WorkflowGraphQueryOutput{
+		WorkflowID:     input.WorkflowID,
+		Query:          input.Query,
+		Version:        input.Version,
+		RankingProfile: input.RankingProfile,
+		Results:        []action.EdgeWithWording{},
+		Paths:          []QueryPath{},
+	}
+
+	opts := workflow.ActivityOptions{
+		StartToCloseTimeout: 120 * time.Second,
+		RetryPolicy: &workflow.RetryPolicy{
+			InitialInterval:    2 * time.Second,
+			BackoffCoefficient: 2.0,
+			MaxInterval:        10 * time.Second,
+			MaxAttempts:        3,
+		},
+	}
+	ctx = workflow.WithActivityOptions(ctx, opts)
+
+	// Fetch canvas + relations
+	var canvasData action.CanvasWithRelationsData
+	err := workflow.ExecuteActivity(ctx, action.FetchCanvasRelationsActivity,
+		action.FetchCanvasRelationsInput{
+			WorkflowID: input.WorkflowID,
+			Version:    input.Version,
+		},
+	).Get(ctx, &canvasData)
+	if err != nil {
+		return output, err
+	}
+
+	// Query Memory System via unified endpoint
+	var graphResults action.GraphRAGQueryOutput
+	err = workflow.ExecuteActivity(ctx, action.QueryGraphRAGActivity,
+		action.GraphRAGQueryInput{
+			WorkflowID:      input.WorkflowID,
+			Query:           input.Query,
+			SearchType:      input.SearchType,
+			RelationType:    input.RelationType,
+			ConfidenceFloor: input.ConfidenceFloor,
+			TopK:            input.TopK,
+			RankingProfile:  input.RankingProfile,
+			Canvas:          canvasData,
+		},
+	).Get(ctx, &graphResults)
+	if err != nil {
+		return output, err
+	}
+
+	output.Results = graphResults.Edges
+	output.TotalCount = graphResults.TotalCount
+	output.HasMore = graphResults.HasMore
+
+	output.ExecutionTimeMs = time.Since(startTime).Milliseconds()
+	return output, nil
+}