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,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
+}