Eino: Flow Integration

// instantiate the flow: multiquery.NewRetriever

vk, err := newVikingDBRetriever(ctx, vikingDBHost, vikingDBRegion, vikingDBAK, vikingDBSK)
if err != nil {
    logs.Errorf("newVikingDBRetriever failed, err=%v", err)
    return
}

llm, err := newChatModel(ctx, openAIBaseURL, openAIAPIKey, openAIModelName)
if err != nil {
    logs.Errorf("newChatModel failed, err=%v", err)
    return
}

// rewrite query by llm
mqr, err := multiquery.NewRetriever(ctx, &multiquery.Config{
    RewriteLLM:      llm,
    RewriteTemplate: nil, // use default
    QueryVar:        "",  // use default
    LLMOutputParser: nil, // use default
    MaxQueriesNum:   3,
    OrigRetriever:   vk,
    FusionFunc:      nil, // use default fusion, just deduplicate by doc id
})
if err != nil {
    logs.Errorf("NewMultiQueryRetriever failed, err=%v", err)
    return
}

// add the flow to graph
graph := compose.NewGraph[string, *schema.Message]()
_ = graph.AddRetrieverNode("multi_query_retriever", mqr, compose.WithOutputKey("context"))
_ = graph.AddEdge(compose._START_, "multi_query_retriever")
_ = graph.AddChatTemplateNode("template", prompt.FromMessages(schema._FString_, schema.UserMessage("{context}")))

// ...

// instantiate the host multi-agent
hostMA, err := NewMultiAgent(ctx, &MultiAgentConfig{
    Host: Host{
       ChatModel: mockHostLLM,
    },
    Specialists: []*Specialist{
       specialist1,
       specialist2,
    },
})
assert.Nil(t, err)

// export graph and []GraphAddNodeOption from host multi-agent
maGraph, opts := hostMA.ExportGraph()

// add to another graph 
fullGraph, err := compose.NewChain[map[string]any, *schema.Message]().
    AppendChatTemplate(prompt.FromMessages(schema._FString_, schema.UserMessage("what's the capital city of {country_name}"))).
    AppendGraph(maGraph, append(opts, compose.WithNodeKey("host_ma_node"))...).
    Compile(ctx)
assert.Nil(t, err)

// invoke the other graph
// convert the flow's own option to compose.Option if needed
// assign options to flow's nodes if needed
out, err := fullGraph.Invoke(ctx, map[string]any{"country_name": "China"}, 
    compose.WithCallbacks(ConvertCallbackHandlers(mockCallback)).
        DesignateNodeWithPath(compose.NewNodePath("host_ma_node", hostMA.HostNodeKey())))

// instantiate the flow
a, err := NewAgent(ctx, &AgentConfig{
    Model: cm,
    ToolsConfig: compose.ToolsNodeConfig{
       Tools: []tool.BaseTool{fakeTool, &fakeStreamToolGreetForTest{}},
    },

    MaxStep: 40,
})
assert.Nil(t, err)

chain := compose.NewChain[[]*schema.Message, string]()

// convert the flow to Lambda
agentLambda, err := compose.AnyLambda(a.Generate, a.Stream, nil, nil)
assert.Nil(t, err)

// add lambda to another graph
chain.
    AppendLambda(agentLambda).
    AppendLambda(compose.InvokableLambda(func(ctx context.Context, input *schema.Message) (string, error) {
       t.Log("got agent response: ", input.Content)
       return input.Content, nil
    }))
r, err := chain.Compile(ctx)
assert.Nil(t, err)

// invoke the graph
res, err := r.Invoke(ctx, []*schema.Message{{Role: schema._User_, Content: "hello"}},
    compose.WithCallbacks(callbackForTest))