Fix contextScope, and remove undefined (#211)

xebia-functional · Jun 29, 2023 · 005eda8 · 005eda8
1 parent 76420f6
commit 005eda8
Showing 1 changed file with 12 additions and 4 deletions.
diff --git a/java/src/main/java/com/xebia/functional/xef/java/auto/AIScope.java b/java/src/main/java/com/xebia/functional/xef/java/auto/AIScope.java
@@ -68,8 +68,13 @@ public AIScope() {
         this(new ObjectMapper(), new OpenAIConfig(), Executors.newCachedThreadPool(new AIScopeThreadFactory()));
     }
 
-    private <T> T undefined() {
-        throw new RuntimeException("Method is undefined");
+    private AIScope(CoreAIScope nested, AIScope outer) {
+        this.om = outer.om;
+        this.executorService = outer.executorService;
+        this.coroutineScope = outer.coroutineScope;
+        this.schemaGen = outer.schemaGen;
+        this.client = outer.client;
+        this.scope = nested;
     }
 
     public <A> CompletableFuture<A> prompt(String prompt, Class<A> cls) {
@@ -107,8 +112,11 @@ public CompletableFuture<List<String>> promptMessage(String prompt, LLMModel llm
         return future(continuation -> scope.promptMessage(prompt, llmModel, functions, user, echo, n, temperature, bringFromContext, minResponseTokens, continuation));
     }
 
-    public <T> CompletableFuture<T> contextScope(List<String> docs) {
-        return future(continuation -> scope.contextScopeWithDocs(docs, undefined(), continuation));
+    public <A> CompletableFuture<A> contextScope(List<String> docs, Function1<AIScope, CompletableFuture<A>> f) {
+        return future(continuation -> scope.contextScopeWithDocs(docs, (coreAIScope, continuation1) -> {
+            AIScope nestedScope = new AIScope(coreAIScope, AIScope.this);
+            return FutureKt.await(f.invoke(nestedScope), continuation);
+        }, continuation));
     }
 
     public CompletableFuture<List<String>> pdf(String url, TextSplitter splitter) {