-
Notifications
You must be signed in to change notification settings - Fork 98
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Make @RegisterAiService beans request scoped by default
This is done because otherwise the chat memory does not get cleared properly. Furthermore, add a way to remove memory entries when the service goes out of scope Fixes: #95
- Loading branch information
Showing
16 changed files
with
416 additions
and
32 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
13 changes: 13 additions & 0 deletions
13
core/runtime/src/main/java/io/quarkiverse/langchain4j/RemovableChatMemoryProvider.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,13 @@ | ||
package io.quarkiverse.langchain4j; | ||
|
||
import dev.langchain4j.memory.ChatMemory; | ||
import dev.langchain4j.memory.chat.ChatMemoryProvider; | ||
|
||
/** | ||
* Extends {@link ChatMemoryProvider} to allow for removing {@link ChatMemory} | ||
* when it is no longer needed. | ||
*/ | ||
public interface RemovableChatMemoryProvider extends ChatMemoryProvider { | ||
|
||
void remove(Object id); | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
24 changes: 24 additions & 0 deletions
24
.../java/io/quarkiverse/langchain4j/runtime/aiservice/DeclarativeAiServiceBeanDestroyer.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,24 @@ | ||
package io.quarkiverse.langchain4j.runtime.aiservice; | ||
|
||
import java.util.Map; | ||
|
||
import jakarta.enterprise.context.spi.CreationalContext; | ||
|
||
import org.jboss.logging.Logger; | ||
|
||
import io.quarkus.arc.BeanDestroyer; | ||
|
||
public class DeclarativeAiServiceBeanDestroyer implements BeanDestroyer<AutoCloseable> { | ||
|
||
private static final Logger log = Logger.getLogger(DeclarativeAiServiceBeanDestroyer.class); | ||
|
||
@Override | ||
public void destroy(AutoCloseable instance, CreationalContext<AutoCloseable> creationalContext, | ||
Map<String, Object> params) { | ||
try { | ||
instance.close(); | ||
} catch (Exception e) { | ||
log.error("Unable to close " + instance); | ||
} | ||
} | ||
} |
31 changes: 31 additions & 0 deletions
31
...e/src/main/java/io/quarkiverse/langchain4j/runtime/aiservice/QuarkusAiServiceContext.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,13 +1,44 @@ | ||
package io.quarkiverse.langchain4j.runtime.aiservice; | ||
|
||
import java.util.Set; | ||
import java.util.concurrent.ConcurrentHashMap; | ||
|
||
import dev.langchain4j.service.AiServiceContext; | ||
import io.quarkiverse.langchain4j.RegisterAiService; | ||
import io.quarkiverse.langchain4j.RemovableChatMemoryProvider; | ||
import io.quarkiverse.langchain4j.audit.AuditService; | ||
|
||
public class QuarkusAiServiceContext extends AiServiceContext { | ||
|
||
public AuditService auditService; | ||
|
||
public Set<Object> usedMemoryIds = ConcurrentHashMap.newKeySet(); | ||
|
||
public QuarkusAiServiceContext(Class<?> aiServiceClass) { | ||
super(aiServiceClass); | ||
} | ||
|
||
/** | ||
* This is called by the {@code close} method of AiServices registered with {@link RegisterAiService} | ||
* when the bean's scope is closed | ||
*/ | ||
public void close() { | ||
removeChatMemories(); | ||
} | ||
|
||
private void removeChatMemories() { | ||
if (usedMemoryIds.isEmpty()) { | ||
return; | ||
} | ||
RemovableChatMemoryProvider removableChatMemoryProvider = null; | ||
if (chatMemoryProvider instanceof RemovableChatMemoryProvider) { | ||
removableChatMemoryProvider = (RemovableChatMemoryProvider) chatMemoryProvider; | ||
} | ||
for (Object memoryId : usedMemoryIds) { | ||
if (removableChatMemoryProvider != null) { | ||
removableChatMemoryProvider.remove(memoryId); | ||
} | ||
chatMemories.remove(memoryId); | ||
} | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.