107 lines
3.3 KiB
Go
107 lines
3.3 KiB
Go
|
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
|
||
|
|
}
|