refactor: add compile method

- useful for streaming server impl work on #9
bsorrentino · Jul 15, 2024 · 7b795ff · 7b795ff
1 parent dc8ff48
commit 7b795ff
Showing 1 changed file with 14 additions and 8 deletions.
diff --git a/agent-executor/src/main/java/dev/langchain4j/agentexecutor/AgentExecutor.java b/agent-executor/src/main/java/dev/langchain4j/agentexecutor/AgentExecutor.java
@@ -5,6 +5,7 @@
 import dev.langchain4j.model.output.FinishReason;
 import lombok.var;
 import org.bsc.async.AsyncGenerator;
+import org.bsc.langgraph4j.CompiledGraph;
 import org.bsc.langgraph4j.StateGraph;
 import org.bsc.langgraph4j.NodeOutput;
 import org.bsc.langgraph4j.state.AgentState;
@@ -93,29 +94,28 @@ String shouldContinue(State state) {
         return "continue";
     }
 
-    public AsyncGenerator<NodeOutput<State>> execute(ChatLanguageModel chatLanguageModel, Map<String, Object> inputs, List<Object> objectsWithTools) throws Exception {
-
+    public CompiledGraph<State> compile(ChatLanguageModel chatLanguageModel, List<Object> objectsWithTools) throws Exception {
         var toolInfoList = ToolInfo.fromList( objectsWithTools );
 
         final List<ToolSpecification> toolSpecifications = toolInfoList.stream()
                 .map(ToolInfo::specification)
                 .collect(Collectors.toList());
 
         var agentRunnable = Agent.builder()
-                                .chatLanguageModel(chatLanguageModel)
-                                .tools( toolSpecifications )
-                                .build();
+                .chatLanguageModel(chatLanguageModel)
+                .tools( toolSpecifications )
+                .build();
 
         var workflow = new StateGraph<>(State::new);
 
         workflow.setEntryPoint("agent");
 
         workflow.addNode( "agent", node_async( state ->
-            runAgent(agentRunnable, state))
+                runAgent(agentRunnable, state))
         );
 
         workflow.addNode( "action", node_async( state ->
-            executeTools(toolInfoList, state))
+                executeTools(toolInfoList, state))
         );
 
         workflow.addConditionalEdges(
@@ -126,7 +126,13 @@ public AsyncGenerator<NodeOutput<State>> execute(ChatLanguageModel chatLanguageM
 
         workflow.addEdge("action", "agent");
 
-        var app = workflow.compile();
+        return workflow.compile();
+
+    }
+
+    public AsyncGenerator<NodeOutput<State>> execute(ChatLanguageModel chatLanguageModel, Map<String, Object> inputs, List<Object> objectsWithTools) throws Exception {
+
+        var app = compile(chatLanguageModel, objectsWithTools);
 
         return  app.stream( inputs );
     }