feat: GraphRAG query API handlers and activities

2026-09-05 05:58:18 -07:00
parent da44923c5c
commit 447951daca
3 changed files with 393 additions and 0 deletions
@@ -0,0 +1,96 @@
+package action
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"github.com/rockliang/poimen/workflows/pkg/db"
+)
+
+// CanvasWithRelationsData combines canvas nodes/edges with relation wording
+type CanvasWithRelationsData struct {
+	WorkflowID  string              `json:"workflow_id"`
+	Version     int                 `json:"version"`
+	Nodes       []db.WorkflowNode   `json:"nodes"`
+	Edges       []db.WorkflowEdge   `json:"edges"`
+	Relations   []EdgeWithWording   `json:"relations"`
+	UpdatedAt   string              `json:"updated_at"`
+}
+
+// FetchCanvasRelationsInput parameters
+type FetchCanvasRelationsInput struct {
+	WorkflowID string `json:"workflow_id"`
+	Version    int    `json:"version"`
+}
+
+// FetchCanvasRelationsActivity fetches canvas + relations from DB
+func FetchCanvasRelationsActivity(ctx context.Context, input FetchCanvasRelationsInput) (CanvasWithRelationsData, error) {
+	logger := newActivityLogger(ctx)
+	output := CanvasWithRelationsData{
+		WorkflowID: input.WorkflowID,
+		Version:    input.Version,
+		Nodes:      []db.WorkflowNode{},
+		Edges:      []db.WorkflowEdge{},
+		Relations:  []EdgeWithWording{},
+	}
+
+	logger.logf("info", "Fetching canvas relations: %s v%d", input.WorkflowID, input.Version)
+
+	// Get database client from context or activity manager
+	dbClient, ok := ctx.Value("db_client").(*db.Client)
+	if !ok {
+		return output, fmt.Errorf("database client not in context")
+	}
+
+	// Fetch workflow
+	workflow, err := dbClient.GetWorkflow(ctx, input.WorkflowID)
+	if err != nil {
+		return output, fmt.Errorf("failed to get workflow: %w", err)
+	}
+
+	// Parse canvas nodes and edges
+	var nodes []db.WorkflowNode
+	if err := json.Unmarshal([]byte(workflow.Nodes), &nodes); err != nil {
+		return output, fmt.Errorf("failed to parse nodes: %w", err)
+	}
+
+	var edges []db.WorkflowEdge
+	if err := json.Unmarshal([]byte(workflow.Edges), &edges); err != nil {
+		return output, fmt.Errorf("failed to parse edges: %w", err)
+	}
+
+	output.Nodes = nodes
+	output.Edges = edges
+	output.UpdatedAt = workflow.UpdatedAt.String()
+
+	// Fetch workflow relations
+	relations, err := dbClient.GetWorkflowRelations(ctx, input.WorkflowID, input.Version)
+	if err != nil {
+		// Relations may not exist for old canvases - this is OK
+		logger.logf("warn", "Failed to fetch relations: %v", err)
+		return output, nil
+	}
+
+	// Map to EdgeWithWording
+	for _, rel := range relations {
+		edge := EdgeWithWording{
+			ID:             rel.ID,
+			Source:         rel.SourceNodeID,
+			Target:         rel.TargetNodeID,
+			RelationType:   rel.RelationType,
+			RelationLabel:  rel.Label,
+			CreatedAt:      rel.CreatedAt.String(),
+		}
+
+		// Parse relation wording JSON
+		if err := json.Unmarshal(rel.RelationWording, &edge.RelationWording); err != nil {
+			logger.logf("warn", "Failed to parse relation wording: %v", err)
+		}
+
+		output.Relations = append(output.Relations, edge)
+	}
+
+	logger.logf("info", "Fetched %d nodes, %d edges, %d relations", len(output.Nodes), len(output.Edges), len(output.Relations))
+	return output, nil
+}
@@ -0,0 +1,133 @@
+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
+}
@@ -0,0 +1,164 @@
+package api
+
+import (
+	"encoding/json"
+	"net/http"
+	"strconv"
+
+	"github.com/gorilla/mux"
+	"go.temporal.io/sdk/client"
+	"github.com/rockliang/poimen/workflows/pkg/db"
+	"github.com/rockliang/poimen/workflows/statemachine"
+)
+
+// QueryWorkflowGraphRequest matches frontend payload
+type QueryWorkflowGraphRequest struct {
+	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"`
+}
+
+// QueryWorkflowGraphHandler handles POST /workflows/{id}/query
+func QueryWorkflowGraphHandler(temporalClient client.Client, dbClient *db.Client) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodPost {
+			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+			return
+		}
+
+		workflowID := mux.Vars(r)["id"]
+		if workflowID == "" {
+			http.Error(w, "Missing workflow ID", http.StatusBadRequest)
+			return
+		}
+
+		var req QueryWorkflowGraphRequest
+		if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
+			http.Error(w, "Invalid request body", http.StatusBadRequest)
+			return
+		}
+
+		// Defaults
+		if req.SearchType == "" {
+			req.SearchType = "edges"
+		}
+		if req.ConfidenceFloor == 0 {
+			req.ConfidenceFloor = 0.5
+		}
+		if req.TopK == 0 {
+			req.TopK = 10
+		}
+		if req.MaxPathDepth == 0 {
+			req.MaxPathDepth = 3
+		}
+		if req.RankingProfile == "" {
+			req.RankingProfile = "default"
+		}
+
+		// Get latest version if not specified
+		if req.Version == 0 {
+			workflow, err := dbClient.GetWorkflow(r.Context(), workflowID)
+			if err != nil {
+				http.Error(w, "Failed to get workflow", http.StatusInternalServerError)
+				return
+			}
+			req.Version = workflow.Version
+		}
+
+		// Start Temporal workflow
+		workflowInput := statemachine.WorkflowGraphQueryInput{
+			WorkflowID:       workflowID,
+			Query:            req.Query,
+			SearchType:       req.SearchType,
+			RelationType:     req.RelationType,
+			Version:          req.Version,
+			ConfidenceFloor:  req.ConfidenceFloor,
+			TopK:             req.TopK,
+			FindPaths:        req.FindPaths,
+			TargetNodeID:     req.TargetNodeID,
+			MaxPathDepth:     req.MaxPathDepth,
+			RankingProfile:   req.RankingProfile,
+			IncludeReasoning: req.IncludeReasoning,
+		}
+
+		options := client.StartWorkflowOptions{
+			ID:       "graph-query-" + workflowID + "-" + strconv.FormatInt(int64(req.Version), 10),
+			TaskQueue: "poimen-default",
+		}
+
+		// Execute workflow (blocking)
+		run, err := temporalClient.ExecuteWorkflow(
+			r.Context(),
+			options,
+			statemachine.WorkflowGraphQuery,
+			workflowInput,
+		)
+		if err != nil {
+			http.Error(w, "Failed to start workflow", http.StatusInternalServerError)
+			return
+		}
+
+		var output statemachine.WorkflowGraphQueryOutput
+		if err := run.Get(r.Context(), &output); err != nil {
+			http.Error(w, "Workflow execution failed", http.StatusInternalServerError)
+			return
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(output)
+	}
+}
+
+// GetWorkflowRelationVersionsHandler handles GET /workflows/{id}/relations/{edge_id}/versions
+func GetWorkflowRelationVersionsHandler(dbClient *db.Client) http.HandlerFunc {
+	return func(w http.ResponseWriter, r *http.Request) {
+		if r.Method != http.MethodGet {
+			http.Error(w, "Method not allowed", http.StatusMethodNotAllowed)
+			return
+		}
+
+		workflowID := mux.Vars(r)["id"]
+		edgeID := mux.Vars(r)["edge_id"]
+
+		if workflowID == "" || edgeID == "" {
+			http.Error(w, "Missing parameters", http.StatusBadRequest)
+			return
+		}
+
+		// Query relation versions from DB
+		versions, err := dbClient.GetRelationVersions(r.Context(), workflowID, edgeID)
+		if err != nil {
+			http.Error(w, "Failed to fetch versions", http.StatusInternalServerError)
+			return
+		}
+
+		w.Header().Set("Content-Type", "application/json")
+		json.NewEncoder(w).Encode(map[string]interface{}{
+			"workflow_id": workflowID,
+			"edge_id":     edgeID,
+			"versions":    versions,
+			"total_count": len(versions),
+		})
+	}
+}
+
+// RegisterGraphQueryHandlers registers all graph query endpoints
+func RegisterGraphQueryHandlers(
+	router *mux.Router,
+	temporalClient client.Client,
+	dbClient *db.Client,
+) {
+	// Query workflow relations via GraphRAG
+	router.HandleFunc("/workflows/{id}/query", QueryWorkflowGraphHandler(temporalClient, dbClient)).Methods("POST")
+
+	// Get relation version history
+	router.HandleFunc("/workflows/{id}/relations/{edge_id}/versions", GetWorkflowRelationVersionsHandler(dbClient)).Methods("GET")
+}