From 005eda8b9b1283b9baea2dc8ca40a4fb948ee677 Mon Sep 17 00:00:00 2001
From: Simon Vergauwen <nomisRev@users.noreply.github.com>
Date: Thu, 29 Jun 2023 08:36:40 +0200
Subject: [PATCH] Fix contextScope, and remove undefined (#211)

---
 .../xebia/functional/xef/java/auto/AIScope.java  | 16 ++++++++++++----
 1 file changed, 12 insertions(+), 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
index 0819f41fa..674cbce27 100644
--- 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) {