From ea8a4bfc2688b73f6766ec94edff5cbcd5f4a9e0 Mon Sep 17 00:00:00 2001 From: sunnyji-coder Date: Tue, 18 Aug 2026 08:48:08 +0800 Subject: [PATCH] Add Python parity features to veadk-java --- README.md | 74 +- core/pom.xml | 9 + .../com/volcengine/veadk/agent/Agent.java | 560 +++++++++++++ .../veadk/agent/AgentComponent.java | 42 + .../veadk/agent/AgentComponentProvider.java | 35 + .../volcengine/veadk/agent/AgentMetadata.java | 181 +++++ .../veadk/agent/SaveSessionPolicy.java | 64 ++ .../agent/SaveSessionToMemoryCallback.java | 196 ++++- .../volcengine/veadk/config/ConfigLoader.java | 159 ++++ .../volcengine/veadk/config/VeADKConfig.java | 207 +++++ .../InMemoryKnowledgebaseBackend.java | 84 ++ .../veadk/knowledgebase/KnowledgeBase.java | 157 ++++ .../knowledgebase/KnowledgebaseBackend.java | 35 + .../knowledgebase/KnowledgebaseEntry.java | 33 + .../KnowledgebaseServiceBackendAdapter.java | 63 ++ .../memory/InMemoryLongTermMemoryBackend.java | 109 +++ .../InMemoryShortTermMemoryBackend.java | 28 + .../veadk/memory/LongTermMemory.java | 244 ++++++ .../veadk/memory/LongTermMemoryBackend.java | 42 + .../memory/MemoryServiceBackendAdapter.java | 59 ++ .../memory/SessionServiceBackendAdapter.java | 32 + .../veadk/memory/ShortTermMemory.java | 153 ++++ .../veadk/memory/ShortTermMemoryBackend.java | 29 + .../veadk/memory/ShortTermMemoryMessage.java | 27 + .../memory/ShortTermMemoryProcessor.java | 82 ++ .../memory/viking/VikingMemoryService.java | 78 +- .../volcengine/veadk/model/ArkEmbedding.java | 230 ++++++ .../com/volcengine/veadk/model/ArkLlm.java | 766 ++++++++++++++---- .../veadk/processors/BaseRunProcessor.java | 27 + .../veadk/processors/NoOpRunProcessor.java | 34 + .../veadk/processors/RunContext.java | 57 ++ .../com/volcengine/veadk/runner/Runner.java | 120 ++- .../veadk/agent/AgentMetadataTest.java | 146 ++++ .../com/volcengine/veadk/agent/AgentTest.java | 140 ++++ .../SaveSessionToMemoryCallbackTest.java | 214 +++++ .../compat/PublicApiCompatibilityTest.java | 213 +++++ .../veadk/config/VeADKConfigTest.java | 112 +++ .../knowledgebase/KnowledgeBaseTest.java | 101 +++ .../veadk/memory/LongTermMemoryTest.java | 157 ++++ .../veadk/memory/ShortTermMemoryTest.java | 150 ++++ .../viking/VikingMemoryServiceTest.java | 47 ++ .../veadk/model/ArkEmbeddingTest.java | 145 ++++ .../volcengine/veadk/model/ArkLlmTest.java | 628 +++++++++++++- .../veadk/runner/RunnerProcessorTest.java | 119 +++ docs/parity/BASELINE.md | 46 ++ docs/parity/COMPATIBILITY.md | 43 + docs/parity/DECISIONS.md | 53 ++ docs/parity/PARITY_MATRIX.md | 49 ++ docs/parity/PROGRESS.md | 68 ++ example/pom.xml | 1 + memory-sqlite/pom.xml | 54 ++ .../memory/sqlite/SQLiteSessionService.java | 332 ++++++++ .../sqlite/SQLiteShortTermMemoryBackend.java | 45 + .../SQLiteShortTermMemoryBackendTest.java | 188 +++++ pom.xml | 20 + 55 files changed, 6934 insertions(+), 153 deletions(-) create mode 100644 core/src/main/java/com/volcengine/veadk/agent/Agent.java create mode 100644 core/src/main/java/com/volcengine/veadk/agent/AgentComponent.java create mode 100644 core/src/main/java/com/volcengine/veadk/agent/AgentComponentProvider.java create mode 100644 core/src/main/java/com/volcengine/veadk/agent/AgentMetadata.java create mode 100644 core/src/main/java/com/volcengine/veadk/agent/SaveSessionPolicy.java create mode 100644 core/src/main/java/com/volcengine/veadk/config/ConfigLoader.java create mode 100644 core/src/main/java/com/volcengine/veadk/config/VeADKConfig.java create mode 100644 core/src/main/java/com/volcengine/veadk/knowledgebase/InMemoryKnowledgebaseBackend.java create mode 100644 core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgeBase.java create mode 100644 core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseBackend.java create mode 100644 core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseEntry.java create mode 100644 core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseServiceBackendAdapter.java create mode 100644 core/src/main/java/com/volcengine/veadk/memory/InMemoryLongTermMemoryBackend.java create mode 100644 core/src/main/java/com/volcengine/veadk/memory/InMemoryShortTermMemoryBackend.java create mode 100644 core/src/main/java/com/volcengine/veadk/memory/LongTermMemory.java create mode 100644 core/src/main/java/com/volcengine/veadk/memory/LongTermMemoryBackend.java create mode 100644 core/src/main/java/com/volcengine/veadk/memory/MemoryServiceBackendAdapter.java create mode 100644 core/src/main/java/com/volcengine/veadk/memory/SessionServiceBackendAdapter.java create mode 100644 core/src/main/java/com/volcengine/veadk/memory/ShortTermMemory.java create mode 100644 core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryBackend.java create mode 100644 core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryMessage.java create mode 100644 core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryProcessor.java create mode 100644 core/src/main/java/com/volcengine/veadk/model/ArkEmbedding.java create mode 100644 core/src/main/java/com/volcengine/veadk/processors/BaseRunProcessor.java create mode 100644 core/src/main/java/com/volcengine/veadk/processors/NoOpRunProcessor.java create mode 100644 core/src/main/java/com/volcengine/veadk/processors/RunContext.java create mode 100644 core/src/test/java/com/volcengine/veadk/agent/AgentMetadataTest.java create mode 100644 core/src/test/java/com/volcengine/veadk/agent/AgentTest.java create mode 100644 core/src/test/java/com/volcengine/veadk/agent/SaveSessionToMemoryCallbackTest.java create mode 100644 core/src/test/java/com/volcengine/veadk/compat/PublicApiCompatibilityTest.java create mode 100644 core/src/test/java/com/volcengine/veadk/config/VeADKConfigTest.java create mode 100644 core/src/test/java/com/volcengine/veadk/knowledgebase/KnowledgeBaseTest.java create mode 100644 core/src/test/java/com/volcengine/veadk/memory/LongTermMemoryTest.java create mode 100644 core/src/test/java/com/volcengine/veadk/memory/ShortTermMemoryTest.java create mode 100644 core/src/test/java/com/volcengine/veadk/model/ArkEmbeddingTest.java create mode 100644 core/src/test/java/com/volcengine/veadk/runner/RunnerProcessorTest.java create mode 100644 docs/parity/BASELINE.md create mode 100644 docs/parity/COMPATIBILITY.md create mode 100644 docs/parity/DECISIONS.md create mode 100644 docs/parity/PARITY_MATRIX.md create mode 100644 docs/parity/PROGRESS.md create mode 100644 memory-sqlite/pom.xml create mode 100644 memory-sqlite/src/main/java/com/volcengine/veadk/memory/sqlite/SQLiteSessionService.java create mode 100644 memory-sqlite/src/main/java/com/volcengine/veadk/memory/sqlite/SQLiteShortTermMemoryBackend.java create mode 100644 memory-sqlite/src/test/java/com/volcengine/veadk/memory/sqlite/SQLiteShortTermMemoryBackendTest.java diff --git a/README.md b/README.md index 42354b7..dadd5ec 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ An open-source Agent development toolkit that integrates the powerful capabiliti com.volcengine.veadk veadk-java - 0.0.1 + 0.0.2 ``` ```java @@ -56,6 +56,78 @@ export MODEL_AGENT_API_KEY="" ## Run the Project Examples +### Local Knowledge and Long-term Memory + +The backend-neutral local implementations do not require cloud credentials: + +```java +KnowledgeBase knowledgeBase = new KnowledgeBase("docs"); +knowledgeBase.addFromText("VeADK supports deterministic local retrieval."); + +LongTermMemory longTermMemory = new LongTermMemory("my_app"); +ShortTermMemory shortTermMemory = new ShortTermMemory(); +Agent agent = Agent.builder() + .name("memory-agent") + .model(new ArkLlm("doubao-seed-1-8-preview-251115")) + .knowledgebase(knowledgeBase) + .shortTermMemory(shortTermMemory) + .longTermMemory(longTermMemory) + .build(); + +Runner runner = agent.newRunner(); +``` + +Existing `VikingMemoryService` and `VikingKnowledgebaseService` entry points remain supported. + +Ark LLM also accepts an ordered model list for Python-style fallbacks. The first model is tried +first; later models are only used when an attempt fails before any streaming output is emitted. + +```java +ArkLlm model = new ArkLlm(List.of("doubao-seed-1-8-251228", "deepseek-r1-250528")); +``` + +ADK generation config values such as temperature, top-p, max output tokens, stop sequences, +penalties, candidate count, log probabilities and JSON/JSON-schema response formats are forwarded +to Ark requests when set. Text, inline image bytes, image URLs and video URLs from ADK `Part` +values are also mapped into Ark chat content parts. Tool-call history is round-tripped through Ark +assistant `tool_calls` and tool-result messages, including ids and parallel streaming tool calls. + +To automatically copy completed sessions into long-term memory, enable `autoSaveSession` when a +long-term memory service is configured. The default policy matches Python's thresholds: first save +immediately, then save after 10 new events or 60 seconds, and flush the previous session when the +active session changes. + +Direct construction of `SaveSessionToMemoryCallback()` keeps the original Java behavior: every +invocation starts a background save and returns immediately. Use the policy constructor when the +save must participate in the reactive callback chain. + +```java +Agent agent = Agent.builder() + .name("memory-agent") + .model(new ArkLlm("doubao-seed-1-8-preview-251115")) + .longTermMemory(longTermMemory) + .autoSaveSession(true) + .build(); +``` + +For persistent local sessions, add the optional SQLite module: + +```xml + + com.volcengine.veadk + veadk-memory-sqlite + 0.0.2 + +``` + +```java +ShortTermMemory shortTermMemory = new ShortTermMemory( + new SQLiteShortTermMemoryBackend(Path.of("./data/sessions.db"))); +``` + +The core artifact does not transitively require the SQLite JDBC driver. SQLite appends reload the +canonical stored session before writing, so stale loaded views do not replace newer history. + ### Build the Project In the repository root, run: `./mvnw clean -DskipTests package` diff --git a/core/pom.xml b/core/pom.xml index 27b72f2..baa3697 100644 --- a/core/pom.xml +++ b/core/pom.xml @@ -36,6 +36,10 @@ limitations under the License. com.fasterxml.jackson.core jackson-core + + org.yaml + snakeyaml + org.slf4j slf4j-api @@ -72,6 +76,10 @@ limitations under the License. org.mockito mockito-junit-jupiter + + net.bytebuddy + byte-buddy-agent + @@ -82,6 +90,7 @@ limitations under the License. @{argLine} + -javaagent:${settings.localRepository}/net/bytebuddy/byte-buddy-agent/${byte-buddy.version}/byte-buddy-agent-${byte-buddy.version}.jar --add-opens java.base/java.util=ALL-UNNAMED --add-opens java.base/java.lang=ALL-UNNAMED diff --git a/core/src/main/java/com/volcengine/veadk/agent/Agent.java b/core/src/main/java/com/volcengine/veadk/agent/Agent.java new file mode 100644 index 0000000..db9bcab --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/agent/Agent.java @@ -0,0 +1,560 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.agent; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.Callbacks; +import com.google.adk.agents.Instruction; +import com.google.adk.agents.LlmAgent; +import com.google.adk.artifacts.InMemoryArtifactService; +import com.google.adk.codeexecutors.BaseCodeExecutor; +import com.google.adk.examples.BaseExampleProvider; +import com.google.adk.examples.Example; +import com.google.adk.memory.BaseMemoryService; +import com.google.adk.memory.InMemoryMemoryService; +import com.google.adk.models.BaseLlm; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.sessions.BaseSessionService; +import com.google.adk.sessions.InMemorySessionService; +import com.google.common.collect.ImmutableList; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.Schema; +import com.volcengine.veadk.config.VeADKConfig; +import com.volcengine.veadk.knowledgebase.BaseKnowledgebaseService; +import com.volcengine.veadk.model.ArkLlm; +import com.volcengine.veadk.processors.BaseRunProcessor; +import com.volcengine.veadk.processors.NoOpRunProcessor; +import com.volcengine.veadk.runner.Runner; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.Objects; +import java.util.UUID; +import java.util.concurrent.Executor; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** VeADK agent facade that adds configuration, memory, processors, and metadata to ADK. */ +public class Agent extends LlmAgent implements AgentComponentProvider { + + private static final Logger log = LoggerFactory.getLogger(Agent.class); + + public static final String DEFAULT_NAME = "veAgent"; + public static final String DEFAULT_DESCRIPTION = + "An AI agent developed by the VeADK team, specialized in data science, " + + "documentation, and software development."; + public static final String DEFAULT_INSTRUCTION = + "You are an AI agent created by the VeADK team. Use the available tools and " + + "resources to solve the user's task accurately."; + + private final String id; + private final BaseKnowledgebaseService knowledgebase; + private final BaseSessionService shortTermMemory; + private final BaseMemoryService longTermMemory; + private final BaseRunProcessor runProcessor; + private final Object promptManager; + private final List tracers; + private final List skills; + private final List plugins; + private final boolean enableAuthz; + private final boolean autoSaveSession; + private final SaveSessionPolicy saveSessionPolicy; + + protected Agent(Builder builder) { + super(builder); + id = builder.id; + knowledgebase = builder.knowledgebase; + shortTermMemory = builder.shortTermMemory; + longTermMemory = builder.longTermMemory; + runProcessor = builder.runProcessor; + promptManager = builder.promptManager; + tracers = List.copyOf(builder.tracers); + skills = List.copyOf(builder.skills); + plugins = List.copyOf(builder.plugins); + enableAuthz = builder.enableAuthz; + autoSaveSession = builder.autoSaveSession; + saveSessionPolicy = builder.saveSessionPolicy; + } + + public static Builder builder() { + return new Builder(); + } + + /** Creates a builder whose Ark model is resolved from VeADK configuration. */ + public static Builder builder(VeADKConfig config) { + Objects.requireNonNull(config, "config"); + VeADKConfig.ModelConfig model = config.model(); + String apiKey = config.require("MODEL_AGENT_API_KEY"); + return new Builder().model(new ArkLlm(model.name(), apiKey, model.apiBase())); + } + + public String id() { + return id; + } + + public BaseKnowledgebaseService knowledgebase() { + return knowledgebase; + } + + public BaseSessionService shortTermMemory() { + return shortTermMemory; + } + + public BaseMemoryService longTermMemory() { + return longTermMemory; + } + + public BaseRunProcessor runProcessor() { + return runProcessor; + } + + public Object promptManager() { + return promptManager; + } + + public List tracers() { + return tracers; + } + + public List skills() { + return skills; + } + + public List plugins() { + return plugins; + } + + public boolean enableAuthz() { + return enableAuthz; + } + + public boolean autoSaveSession() { + return autoSaveSession; + } + + public SaveSessionPolicy saveSessionPolicy() { + return saveSessionPolicy; + } + + public AgentMetadata metadata() { + return AgentMetadata.from(this); + } + + /** Creates a Runner using this agent's configured memory, plugins, and run processor. */ + public Runner newRunner() { + BaseSessionService sessions = + shortTermMemory == null ? new InMemorySessionService() : shortTermMemory; + BaseMemoryService memory = + longTermMemory == null ? new InMemoryMemoryService() : longTermMemory; + return new Runner( + this, + name(), + new InMemoryArtifactService(), + sessions, + memory, + plugins, + runProcessor); + } + + @Override + public List agentComponents() { + List components = new ArrayList<>(); + addComponent(components, "knowledgebase", knowledgebase, "knowledgebase"); + addComponent(components, "short_term_memory", shortTermMemory, "shortTermMemory"); + addComponent(components, "long_term_memory", longTermMemory, "longTermMemory"); + addComponent(components, "prompt_manager", promptManager, "promptManager"); + if (!(runProcessor instanceof NoOpRunProcessor)) { + addComponent(components, "run_processor", runProcessor, "runProcessor"); + } + tracers.forEach(tracer -> addComponent(components, "tracer", tracer, "tracers")); + plugins.forEach(plugin -> addComponent(components, "plugin", plugin, "plugins")); + return List.copyOf(components); + } + + @Override + public List agentSkills() { + return skills.stream() + .filter(name -> name.matches("[A-Za-z0-9][A-Za-z0-9_-]{0,127}")) + .map(name -> new AgentMetadata.SkillSummary(name, "")) + .toList(); + } + + private static void addComponent( + List components, String kind, Object component, String source) { + if (component != null) { + components.add( + new AgentComponent(kind, component.getClass().getSimpleName(), source, "", "")); + } + } + + public static class Builder extends LlmAgent.Builder { + + private String id = UUID.randomUUID().toString().substring(0, 8); + private BaseKnowledgebaseService knowledgebase; + private BaseSessionService shortTermMemory; + private BaseMemoryService longTermMemory; + private BaseRunProcessor runProcessor = NoOpRunProcessor.INSTANCE; + private Object promptManager; + private List tracers = List.of(); + private List skills = List.of(); + private List plugins = List.of(); + private boolean enableAuthz; + private boolean autoSaveSession; + private SaveSessionPolicy saveSessionPolicy = SaveSessionPolicy.defaults(); + private boolean saveSessionCallbackAttached; + + public Builder() { + name(DEFAULT_NAME); + description(DEFAULT_DESCRIPTION); + instruction(DEFAULT_INSTRUCTION); + } + + public Builder id(String id) { + this.id = requireText(id, "id"); + return this; + } + + public Builder knowledgebase(BaseKnowledgebaseService knowledgebase) { + this.knowledgebase = knowledgebase; + return this; + } + + public Builder shortTermMemory(BaseSessionService shortTermMemory) { + this.shortTermMemory = shortTermMemory; + return this; + } + + public Builder longTermMemory(BaseMemoryService longTermMemory) { + this.longTermMemory = longTermMemory; + return this; + } + + public Builder runProcessor(BaseRunProcessor runProcessor) { + this.runProcessor = Objects.requireNonNull(runProcessor, "runProcessor"); + return this; + } + + public Builder promptManager(Object promptManager) { + this.promptManager = promptManager; + return this; + } + + public Builder tracers(List tracers) { + this.tracers = List.copyOf(tracers); + return this; + } + + public Builder tracers(Object... tracers) { + return tracers(Arrays.asList(tracers)); + } + + public Builder skills(List skills) { + this.skills = List.copyOf(skills); + return this; + } + + public Builder skills(String... skills) { + return skills(Arrays.asList(skills)); + } + + public Builder plugins(List plugins) { + this.plugins = List.copyOf(plugins); + return this; + } + + public Builder plugins(BasePlugin... plugins) { + return plugins(Arrays.asList(plugins)); + } + + public Builder enableAuthz(boolean enableAuthz) { + this.enableAuthz = enableAuthz; + return this; + } + + public Builder autoSaveSession(boolean autoSaveSession) { + this.autoSaveSession = autoSaveSession; + return this; + } + + public Builder saveSessionPolicy(SaveSessionPolicy saveSessionPolicy) { + this.saveSessionPolicy = Objects.requireNonNull(saveSessionPolicy, "saveSessionPolicy"); + return this; + } + + @Override + public Builder name(String name) { + super.name(name); + return this; + } + + @Override + public Builder description(String description) { + super.description(description); + return this; + } + + @Override + public Builder subAgents(List subAgents) { + super.subAgents(subAgents); + return this; + } + + @Override + public Builder subAgents(BaseAgent... subAgents) { + super.subAgents(subAgents); + return this; + } + + @Override + public Builder model(String model) { + super.model(model); + return this; + } + + @Override + public Builder model(BaseLlm model) { + super.model(model); + return this; + } + + @Override + public Builder instruction(Instruction instruction) { + super.instruction(instruction); + return this; + } + + @Override + public Builder instruction(String instruction) { + super.instruction(instruction); + return this; + } + + @Override + public Builder globalInstruction(Instruction instruction) { + super.globalInstruction(instruction); + return this; + } + + @Override + public Builder globalInstruction(String instruction) { + super.globalInstruction(instruction); + return this; + } + + @Override + public Builder tools(List tools) { + super.tools(tools); + return this; + } + + @Override + public Builder tools(Object... tools) { + super.tools(tools); + return this; + } + + @Override + public Builder generateContentConfig(GenerateContentConfig config) { + super.generateContentConfig(config); + return this; + } + + @Override + public Builder exampleProvider(BaseExampleProvider provider) { + super.exampleProvider(provider); + return this; + } + + @Override + public Builder exampleProvider(List examples) { + super.exampleProvider(examples); + return this; + } + + @Override + public Builder exampleProvider(Example... examples) { + super.exampleProvider(examples); + return this; + } + + @Override + public Builder includeContents(IncludeContents includeContents) { + super.includeContents(includeContents); + return this; + } + + @Override + public Builder planning(boolean planning) { + super.planning(planning); + return this; + } + + @Override + public Builder maxSteps(int maxSteps) { + super.maxSteps(maxSteps); + return this; + } + + @Override + public Builder disallowTransferToParent(boolean disallow) { + super.disallowTransferToParent(disallow); + return this; + } + + @Override + public Builder disallowTransferToPeers(boolean disallow) { + super.disallowTransferToPeers(disallow); + return this; + } + + @Override + public Builder beforeAgentCallback(Callbacks.BeforeAgentCallback callback) { + super.beforeAgentCallback(callback); + return this; + } + + @Override + public Builder afterAgentCallback(Callbacks.AfterAgentCallback callback) { + super.afterAgentCallback(callback); + return this; + } + + @Override + public Builder beforeAgentCallbackSync(Callbacks.BeforeAgentCallbackSync callback) { + super.beforeAgentCallbackSync(callback); + return this; + } + + @Override + public Builder afterAgentCallbackSync(Callbacks.AfterAgentCallbackSync callback) { + super.afterAgentCallbackSync(callback); + return this; + } + + @Override + public Builder beforeModelCallback(Callbacks.BeforeModelCallback callback) { + super.beforeModelCallback(callback); + return this; + } + + @Override + public Builder beforeModelCallbackSync(Callbacks.BeforeModelCallbackSync callback) { + super.beforeModelCallbackSync(callback); + return this; + } + + @Override + public Builder afterModelCallback(Callbacks.AfterModelCallback callback) { + super.afterModelCallback(callback); + return this; + } + + @Override + public Builder afterModelCallbackSync(Callbacks.AfterModelCallbackSync callback) { + super.afterModelCallbackSync(callback); + return this; + } + + @Override + public Builder beforeToolCallback(Callbacks.BeforeToolCallback callback) { + super.beforeToolCallback(callback); + return this; + } + + @Override + public Builder beforeToolCallbackSync(Callbacks.BeforeToolCallbackSync callback) { + super.beforeToolCallbackSync(callback); + return this; + } + + @Override + public Builder afterToolCallback(Callbacks.AfterToolCallback callback) { + super.afterToolCallback(callback); + return this; + } + + @Override + public Builder afterToolCallbackSync(Callbacks.AfterToolCallbackSync callback) { + super.afterToolCallbackSync(callback); + return this; + } + + @Override + public Builder inputSchema(Schema schema) { + super.inputSchema(schema); + return this; + } + + @Override + public Builder outputSchema(Schema schema) { + super.outputSchema(schema); + return this; + } + + @Override + public Builder executor(Executor executor) { + super.executor(executor); + return this; + } + + @Override + public Builder outputKey(String outputKey) { + super.outputKey(outputKey); + return this; + } + + @Override + public Builder codeExecutor(BaseCodeExecutor codeExecutor) { + super.codeExecutor(codeExecutor); + return this; + } + + @Override + public Agent build() { + configureAutoSaveSession(); + validate(); + return new Agent(this); + } + + private void configureAutoSaveSession() { + if (!autoSaveSession || saveSessionCallbackAttached) { + return; + } + if (longTermMemory == null) { + log.warn( + "autoSaveSession is enabled, but longTermMemory is not configured; " + + "the save callback was not installed"); + return; + } + ImmutableList.Builder callbacks = ImmutableList.builder(); + if (this.afterAgentCallback != null) { + callbacks.addAll(this.afterAgentCallback); + } + this.afterAgentCallback = + callbacks + .add(new SaveSessionToMemoryCallback(saveSessionPolicy, longTermMemory)) + .build(); + saveSessionCallbackAttached = true; + } + + private static String requireText(String value, String field) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(field + " must not be blank"); + } + return value; + } + } +} diff --git a/core/src/main/java/com/volcengine/veadk/agent/AgentComponent.java b/core/src/main/java/com/volcengine/veadk/agent/AgentComponent.java new file mode 100644 index 0000000..f0df169 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/agent/AgentComponent.java @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.agent; + +import java.util.Objects; + +/** Stable, serialization-friendly description of a component mounted on an agent. */ +public record AgentComponent( + String kind, String name, String source, String backend, String description) { + + public AgentComponent { + kind = requireText(kind, "kind"); + name = requireText(name, "name"); + source = Objects.requireNonNullElse(source, ""); + backend = Objects.requireNonNullElse(backend, ""); + description = Objects.requireNonNullElse(description, ""); + } + + public AgentComponent(String kind, String name) { + this(kind, name, "", "", ""); + } + + private static String requireText(String value, String field) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(field + " must not be blank"); + } + return value.trim(); + } +} diff --git a/core/src/main/java/com/volcengine/veadk/agent/AgentComponentProvider.java b/core/src/main/java/com/volcengine/veadk/agent/AgentComponentProvider.java new file mode 100644 index 0000000..ce5fcc5 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/agent/AgentComponentProvider.java @@ -0,0 +1,35 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.agent; + +import java.util.List; +import java.util.Set; + +/** Optional contract for agents that expose VeADK components, skills, and search sources. */ +public interface AgentComponentProvider { + + default List agentComponents() { + return List.of(); + } + + default List agentSkills() { + return List.of(); + } + + default Set additionalSearchSources() { + return Set.of(); + } +} diff --git a/core/src/main/java/com/volcengine/veadk/agent/AgentMetadata.java b/core/src/main/java/com/volcengine/veadk/agent/AgentMetadata.java new file mode 100644 index 0000000..c8cf97f --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/agent/AgentMetadata.java @@ -0,0 +1,181 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.agent; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.LlmAgent; +import com.google.adk.tools.BaseTool; +import com.google.adk.tools.BaseToolset; +import java.util.ArrayList; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Objects; +import java.util.Set; + +/** Stable, JSON-safe metadata and recursive topology extracted from an ADK agent. */ +public record AgentMetadata( + String name, + String description, + String agentClass, + String model, + List searchSources, + List tools, + List components, + List skills, + List subAgents) { + + private static final Set WEB_SEARCH_TOOL_NAMES = + Set.of("parallel_web_search", "vesearch", "web_search"); + + public AgentMetadata { + name = Objects.requireNonNullElse(name, ""); + description = Objects.requireNonNullElse(description, ""); + agentClass = Objects.requireNonNullElse(agentClass, ""); + model = Objects.requireNonNullElse(model, ""); + searchSources = List.copyOf(searchSources); + tools = List.copyOf(tools); + components = List.copyOf(components); + skills = List.copyOf(skills); + subAgents = List.copyOf(subAgents); + } + + public static AgentMetadata from(BaseAgent agent) { + Objects.requireNonNull(agent, "agent"); + return extract(agent, Collections.newSetFromMap(new IdentityHashMap<>())); + } + + private static AgentMetadata extract(BaseAgent agent, Set path) { + if (!path.add(agent)) { + return new AgentMetadata( + agent.name(), + agent.description(), + agent.getClass().getSimpleName(), + "", + List.of(), + List.of(), + List.of(), + List.of(), + List.of()); + } + + List tools = tools(agent); + List components = new ArrayList<>(); + List skills = new ArrayList<>(); + Set searchSources = new LinkedHashSet<>(); + if (tools.stream().map(ToolSummary::name).anyMatch(WEB_SEARCH_TOOL_NAMES::contains)) { + searchSources.add("web"); + } + if (agent instanceof AgentComponentProvider provider) { + components.addAll(deduplicateComponents(provider.agentComponents())); + skills.addAll(deduplicateSkills(provider.agentSkills())); + searchSources.addAll(provider.additionalSearchSources()); + components.stream() + .map(AgentMetadata::componentSearchSource) + .filter(source -> !source.isEmpty()) + .forEach(searchSources::add); + } + + List children = + agent.subAgents().stream().map(child -> extract(child, path)).toList(); + path.remove(agent); + return new AgentMetadata( + agent.name(), + agent.description(), + agent.getClass().getSimpleName(), + modelName(agent), + List.copyOf(searchSources), + tools, + components, + skills, + children); + } + + private static List tools(BaseAgent agent) { + if (!(agent instanceof LlmAgent llmAgent)) { + return List.of(); + } + Map summaries = new LinkedHashMap<>(); + for (BaseTool tool : llmAgent.tools()) { + summaries.putIfAbsent( + tool.name(), + new ToolSummary( + tool.name(), tool.description(), tool.getClass().getSimpleName())); + } + for (BaseToolset toolset : llmAgent.toolsets()) { + String name = toolset.getClass().getSimpleName(); + summaries.putIfAbsent(name, new ToolSummary(name, "", "toolset")); + } + return List.copyOf(summaries.values()); + } + + private static String modelName(BaseAgent agent) { + if (!(agent instanceof LlmAgent llmAgent) || llmAgent.model().isEmpty()) { + return ""; + } + com.google.adk.models.Model model = llmAgent.model().get(); + return model.modelName() + .orElseGet(() -> model.model().map(modelLlm -> modelLlm.model()).orElse("")); + } + + private static List deduplicateComponents(List components) { + Map unique = new LinkedHashMap<>(); + for (AgentComponent component : components) { + unique.putIfAbsent(component.kind() + "\u0000" + component.name(), component); + } + return List.copyOf(unique.values()); + } + + private static List deduplicateSkills(List skills) { + Map unique = new LinkedHashMap<>(); + for (SkillSummary skill : skills) { + unique.putIfAbsent(skill.name(), skill); + } + return List.copyOf(unique.values()); + } + + private static String componentSearchSource(AgentComponent component) { + String kind = component.kind().toLowerCase(Locale.ROOT); + if (kind.equals("knowledgebase") || kind.equals("knowledge")) { + return "knowledge"; + } + if (kind.equals("long_term_memory") || kind.equals("memory")) { + return "memory"; + } + return ""; + } + + public record ToolSummary(String name, String description, String type) { + public ToolSummary { + name = Objects.requireNonNullElse(name, ""); + description = Objects.requireNonNullElse(description, ""); + type = Objects.requireNonNullElse(type, ""); + } + } + + public record SkillSummary(String name, String description) { + public SkillSummary { + if (name == null || !name.matches("[A-Za-z0-9][A-Za-z0-9_-]{0,127}")) { + throw new IllegalArgumentException("Invalid skill name: " + name); + } + description = Objects.requireNonNullElse(description, ""); + } + } +} diff --git a/core/src/main/java/com/volcengine/veadk/agent/SaveSessionPolicy.java b/core/src/main/java/com/volcengine/veadk/agent/SaveSessionPolicy.java new file mode 100644 index 0000000..40838a3 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/agent/SaveSessionPolicy.java @@ -0,0 +1,64 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.agent; + +import java.time.Duration; + +/** Controls when completed sessions are copied into long-term memory. */ +public record SaveSessionPolicy( + int minNewEvents, + Duration minInterval, + boolean saveOnSessionSwitch, + boolean suppressErrors) { + + public static final int DEFAULT_MIN_NEW_EVENTS = 10; + public static final Duration DEFAULT_MIN_INTERVAL = Duration.ofSeconds(60); + + public SaveSessionPolicy { + if (minNewEvents < 0) { + throw new IllegalArgumentException("minNewEvents must not be negative"); + } + if (minInterval == null || minInterval.isNegative()) { + throw new IllegalArgumentException("minInterval must not be negative"); + } + } + + public static SaveSessionPolicy defaults() { + return new SaveSessionPolicy( + environmentInt("MIN_MESSAGES_THRESHOLD", DEFAULT_MIN_NEW_EVENTS), + Duration.ofSeconds( + environmentInt( + "MIN_TIME_THRESHOLD", + Math.toIntExact(DEFAULT_MIN_INTERVAL.toSeconds()))), + true, + true); + } + + /** Policy matching the original Java callback: save every invocation without switch flushing. */ + public static SaveSessionPolicy legacy() { + return new SaveSessionPolicy(0, Duration.ZERO, false, true); + } + + private static int environmentInt(String name, int fallback) { + String value = System.getenv(name); + if (value == null || value.isBlank()) { + return fallback; + } + try { + return Integer.parseInt(value.trim()); + } catch (NumberFormatException exception) { + return fallback; + } + } +} diff --git a/core/src/main/java/com/volcengine/veadk/agent/SaveSessionToMemoryCallback.java b/core/src/main/java/com/volcengine/veadk/agent/SaveSessionToMemoryCallback.java index dd3e7ef..c02f459 100644 --- a/core/src/main/java/com/volcengine/veadk/agent/SaveSessionToMemoryCallback.java +++ b/core/src/main/java/com/volcengine/veadk/agent/SaveSessionToMemoryCallback.java @@ -18,9 +18,18 @@ import com.google.adk.agents.CallbackContext; import com.google.adk.agents.Callbacks; import com.google.adk.agents.InvocationContext; +import com.google.adk.memory.BaseMemoryService; +import com.google.adk.sessions.BaseSessionService; +import com.google.adk.sessions.Session; import com.google.genai.types.Content; import com.volcengine.veadk.utils.ReadonlyContextAccessorUtil; +import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.subjects.CompletableSubject; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.function.LongSupplier; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -28,16 +37,191 @@ public class SaveSessionToMemoryCallback implements Callbacks.AfterAgentCallback private static final Logger log = LoggerFactory.getLogger(SaveSessionToMemoryCallback.class); + private final SaveSessionPolicy policy; + private final LongSupplier currentTimeMillis; + private final BaseMemoryService memoryServiceOverride; + private final boolean legacyFireAndForget; + private final ConcurrentMap saveStates = new ConcurrentHashMap<>(); + private final ConcurrentMap activeSessions = new ConcurrentHashMap<>(); + private final ConcurrentMap saveQueues = new ConcurrentHashMap<>(); + + public SaveSessionToMemoryCallback() { + this(SaveSessionPolicy.legacy(), System::currentTimeMillis, null, true); + } + + public SaveSessionToMemoryCallback(SaveSessionPolicy policy) { + this(policy, System::currentTimeMillis, null, false); + } + + SaveSessionToMemoryCallback(SaveSessionPolicy policy, LongSupplier currentTimeMillis) { + this(policy, currentTimeMillis, null, false); + } + + SaveSessionToMemoryCallback(SaveSessionPolicy policy, BaseMemoryService memoryServiceOverride) { + this(policy, System::currentTimeMillis, memoryServiceOverride, false); + } + + private SaveSessionToMemoryCallback( + SaveSessionPolicy policy, + LongSupplier currentTimeMillis, + BaseMemoryService memoryServiceOverride, + boolean legacyFireAndForget) { + this.policy = Objects.requireNonNull(policy, "policy"); + this.currentTimeMillis = Objects.requireNonNull(currentTimeMillis, "currentTimeMillis"); + this.memoryServiceOverride = memoryServiceOverride; + this.legacyFireAndForget = legacyFireAndForget; + } + + public SaveSessionPolicy policy() { + return policy; + } + @Override public Maybe call(CallbackContext callbackContext) { - InvocationContext invocationContext = + if (legacyFireAndForget) { + return legacyCall(callbackContext); + } + Completable save = + Completable.defer( + () -> + save( + ReadonlyContextAccessorUtil.getInvocationContext( + callbackContext))); + if (policy.suppressErrors()) { + save = + save.onErrorComplete( + error -> { + log.error("Failed to save session", error); + return true; + }); + } + return save.andThen(Maybe.empty()); + } + + private Maybe legacyCall(CallbackContext callbackContext) { + InvocationContext context = ReadonlyContextAccessorUtil.getInvocationContext(callbackContext); - invocationContext - .memoryService() - .addSessionToMemory(invocationContext.session()) + memoryService(context) + .addSessionToMemory(context.session()) .subscribe( - () -> log.info("Saved session {}", invocationContext.session().id()), - err -> log.error("Failed to save session", err)); + () -> log.info("Saved session {}", context.session().id()), + error -> log.error("Failed to save session", error)); return Maybe.empty(); } + + private Completable save(InvocationContext context) { + UserKey userKey = new UserKey(context.appName(), context.userId()); + return Completable.defer( + () -> { + CompletableSubject completion = CompletableSubject.create(); + CompletableSubject previous = saveQueues.put(userKey, completion); + Completable turn = previous == null ? Completable.complete() : previous; + return turn.andThen(saveOne(context, userKey)) + .doFinally( + () -> { + completion.onComplete(); + saveQueues.remove(userKey, completion); + }); + }); + } + + private Completable saveOne(InvocationContext context, UserKey userKey) { + return Completable.defer( + () -> { + Session currentSession = context.session(); + long now = currentTimeMillis.getAsLong(); + String previousSessionId = activeSessions.get(userKey); + Completable previousSave = + savePreviousSessionIfSwitched( + context, previousSessionId, currentSession.id(), now); + return previousSave.andThen( + Completable.defer( + () -> { + activeSessions.put(userKey, currentSession.id()); + return saveCurrentSession(context, now); + })); + }); + } + + private Completable savePreviousSessionIfSwitched( + InvocationContext context, + String previousSessionId, + String currentSessionId, + long now) { + if (!policy.saveOnSessionSwitch() + || previousSessionId == null + || previousSessionId.equals(currentSessionId)) { + return Completable.complete(); + } + BaseSessionService sessions = context.sessionService(); + BaseMemoryService memory = memoryService(context); + return sessions.getSession( + context.appName(), + context.userId(), + previousSessionId, + java.util.Optional.empty()) + .flatMapCompletable( + session -> + memory.addSessionToMemory(session) + .doOnComplete( + () -> { + saveStates.put( + new SessionKey( + context.appName(), + context.userId(), + previousSessionId), + new SaveState( + now, session.events().size())); + log.info( + "Saved previous session {} after" + + " session switch", + previousSessionId); + })); + } + + private Completable saveCurrentSession(InvocationContext context, long now) { + String sessionId = context.session().id(); + SessionKey key = new SessionKey(context.appName(), context.userId(), sessionId); + return context.sessionService() + .getSession( + context.appName(), context.userId(), sessionId, java.util.Optional.empty()) + .switchIfEmpty( + Maybe.error( + new IllegalStateException( + "Session not found in session service: " + sessionId))) + .flatMapCompletable( + session -> { + int eventCount = session.events().size(); + SaveState previous = saveStates.get(key); + if (shouldSkip(previous, eventCount, now)) { + return Completable.complete(); + } + return memoryService(context) + .addSessionToMemory(session) + .doOnComplete( + () -> { + saveStates.put(key, new SaveState(now, eventCount)); + log.info("Saved session {}", sessionId); + }); + }); + } + + private BaseMemoryService memoryService(InvocationContext context) { + return memoryServiceOverride == null ? context.memoryService() : memoryServiceOverride; + } + + private boolean shouldSkip(SaveState previous, int currentEventCount, long now) { + if (previous == null) { + return false; + } + long elapsed = now - previous.savedAtMillis(); + int newEvents = currentEventCount - previous.eventCount(); + return elapsed < policy.minInterval().toMillis() && newEvents < policy.minNewEvents(); + } + + private record UserKey(String appName, String userId) {} + + private record SessionKey(String appName, String userId, String sessionId) {} + + private record SaveState(long savedAtMillis, int eventCount) {} } diff --git a/core/src/main/java/com/volcengine/veadk/config/ConfigLoader.java b/core/src/main/java/com/volcengine/veadk/config/ConfigLoader.java new file mode 100644 index 0000000..5a599fc --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/config/ConfigLoader.java @@ -0,0 +1,159 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.config; + +import java.io.IOException; +import java.io.Reader; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import org.yaml.snakeyaml.LoaderOptions; +import org.yaml.snakeyaml.Yaml; +import org.yaml.snakeyaml.constructor.SafeConstructor; + +final class ConfigLoader { + + private ConfigLoader() {} + + static Map load(Path workingDirectory, Map environment) { + Path directory = workingDirectory.toAbsolutePath().normalize(); + Map values = new LinkedHashMap<>(); + + findUpwards(directory, "config.yaml").ifPresent(path -> values.putAll(loadYaml(path))); + Path dotenv = directory.resolve(".env"); + if (Files.isRegularFile(dotenv)) { + values.putAll(loadDotenv(dotenv)); + } + environment.forEach((key, value) -> values.put(normalize(key), value)); + applyProviderAliases(values); + return values; + } + + private static java.util.Optional findUpwards(Path directory, String filename) { + Path current = directory; + while (current != null) { + Path candidate = current.resolve(filename); + if (Files.isRegularFile(candidate)) { + return java.util.Optional.of(candidate); + } + current = current.getParent(); + } + return java.util.Optional.empty(); + } + + private static Map loadYaml(Path path) { + LoaderOptions options = new LoaderOptions(); + options.setAllowDuplicateKeys(false); + Yaml yaml = new Yaml(new SafeConstructor(options)); + try (Reader reader = Files.newBufferedReader(path, StandardCharsets.UTF_8)) { + Object root = yaml.load(reader); + Map flattened = new LinkedHashMap<>(); + if (root instanceof Map map) { + flatten(map, "", flattened); + } + return flattened; + } catch (IOException | RuntimeException exception) { + throw new IllegalArgumentException("Unable to load VeADK config: " + path, exception); + } + } + + private static void flatten(Map source, String prefix, Map target) { + source.forEach( + (rawKey, value) -> { + String key = String.valueOf(rawKey); + String path = prefix.isEmpty() ? key : prefix + "_" + key; + if (value instanceof Map nested) { + flatten(nested, path, target); + } else { + target.put(normalize(path), stringify(value)); + } + }); + } + + private static String stringify(Object value) { + if (value == null) { + return ""; + } + if (value instanceof List list) { + return list.stream() + .map(String::valueOf) + .reduce((left, right) -> left + "," + right) + .orElse(""); + } + return String.valueOf(value); + } + + private static Map loadDotenv(Path path) { + Map values = new LinkedHashMap<>(); + try { + for (String rawLine : Files.readAllLines(path, StandardCharsets.UTF_8)) { + String line = rawLine.trim(); + if (line.isEmpty() || line.startsWith("#")) { + continue; + } + if (line.startsWith("export ")) { + line = line.substring("export ".length()).trim(); + } + int separator = line.indexOf('='); + if (separator <= 0) { + continue; + } + String key = normalize(line.substring(0, separator).trim()); + String value = unquote(line.substring(separator + 1).trim()); + values.put(key, value); + } + return values; + } catch (IOException exception) { + throw new IllegalArgumentException("Unable to load VeADK dotenv: " + path, exception); + } + } + + private static String unquote(String value) { + if (value.length() >= 2) { + char first = value.charAt(0); + char last = value.charAt(value.length() - 1); + if ((first == '\'' && last == '\'') || (first == '"' && last == '"')) { + return value.substring(1, value.length() - 1); + } + } + return value; + } + + private static void applyProviderAliases(Map values) { + if (!"byteplus".equalsIgnoreCase(values.getOrDefault("CLOUD_PROVIDER", ""))) { + return; + } + copyIfAbsent(values, "BYTEPLUS_ACCESS_KEY", "VOLCENGINE_ACCESS_KEY"); + copyIfAbsent(values, "BYTEPLUS_SECRET_KEY", "VOLCENGINE_SECRET_KEY"); + } + + private static void copyIfAbsent( + Map values, String source, String destination) { + String existing = values.get(destination); + String alias = values.get(source); + if ((existing == null || existing.isBlank()) && alias != null && !alias.isBlank()) { + values.put(destination, alias); + } + } + + static String normalize(String key) { + return key.trim().toUpperCase(Locale.ROOT).replace('.', '_').replace('-', '_'); + } +} diff --git a/core/src/main/java/com/volcengine/veadk/config/VeADKConfig.java b/core/src/main/java/com/volcengine/veadk/config/VeADKConfig.java new file mode 100644 index 0000000..6ec4ca6 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/config/VeADKConfig.java @@ -0,0 +1,207 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.config; + +import java.nio.file.Path; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Locale; +import java.util.Map; +import java.util.Optional; +import java.util.Set; + +/** Loads VeADK settings from {@code config.yaml}, {@code .env}, and process environment. */ +public final class VeADKConfig { + + public static final String DEFAULT_MODEL_NAME = "doubao-seed-2-1-pro-260628"; + public static final String DEFAULT_MODEL_PROVIDER = "openai"; + public static final String DEFAULT_MODEL_API_BASE = "https://ark.cn-beijing.volces.com/api/v3/"; + public static final String DEFAULT_EMBEDDING_MODEL = "doubao-embedding-vision-250615"; + public static final int DEFAULT_EMBEDDING_DIMENSION = 2048; + + private static final Set SECRET_MARKERS = + Set.of("API_KEY", "SECRET", "PASSWORD", "TOKEN", "CREDENTIAL"); + + private final Map values; + + private VeADKConfig(Map values) { + this.values = Collections.unmodifiableMap(new LinkedHashMap<>(values)); + } + + public static VeADKConfig load() { + return load(Path.of("")); + } + + public static VeADKConfig load(Path workingDirectory) { + return load(workingDirectory, System.getenv()); + } + + /** Loads with an explicit environment map, useful for embedded runtimes and deterministic tests. */ + public static VeADKConfig load(Path workingDirectory, Map environment) { + return new VeADKConfig(ConfigLoader.load(workingDirectory, environment)); + } + + public static VeADKConfig from(Map values) { + Map normalized = new LinkedHashMap<>(); + values.forEach((key, value) -> normalized.put(ConfigLoader.normalize(key), value)); + return new VeADKConfig(normalized); + } + + public Optional get(String key) { + return Optional.ofNullable(values.get(ConfigLoader.normalize(key))); + } + + public String get(String key, String defaultValue) { + String value = values.get(ConfigLoader.normalize(key)); + return value == null || value.isBlank() ? defaultValue : value; + } + + public String require(String key) { + String normalized = ConfigLoader.normalize(key); + String value = values.get(normalized); + if (value == null || value.isBlank()) { + throw new IllegalStateException( + "Missing required configuration: " + + normalized + + ". Configure it in the environment, .env, or config.yaml."); + } + return value; + } + + public boolean getBoolean(String key, boolean defaultValue) { + String value = values.get(ConfigLoader.normalize(key)); + if (value == null || value.isBlank()) { + return defaultValue; + } + return switch (value.trim().toLowerCase(Locale.ROOT)) { + case "true", "1", "yes", "on", "enabled" -> true; + case "false", "0", "no", "off", "disabled" -> false; + default -> + throw new IllegalArgumentException( + "Configuration " + ConfigLoader.normalize(key) + " must be a boolean"); + }; + } + + public int getInt(String key, int defaultValue) { + String value = values.get(ConfigLoader.normalize(key)); + if (value == null || value.isBlank()) { + return defaultValue; + } + try { + return Integer.parseInt(value.trim()); + } catch (NumberFormatException exception) { + throw new IllegalArgumentException( + "Configuration " + ConfigLoader.normalize(key) + " must be an integer", + exception); + } + } + + public ModelConfig model() { + boolean byteplus = "byteplus".equalsIgnoreCase(get("CLOUD_PROVIDER", "")); + String defaultName = byteplus ? "seed-2-0-lite-260228" : DEFAULT_MODEL_NAME; + String defaultBase = + byteplus + ? "https://ark.ap-southeast.bytepluses.com/api/v3" + : DEFAULT_MODEL_API_BASE; + return new ModelConfig( + get("MODEL_AGENT_NAME", defaultName), + get("MODEL_AGENT_PROVIDER", DEFAULT_MODEL_PROVIDER), + get("MODEL_AGENT_API_BASE", defaultBase), + get("MODEL_AGENT_API_KEY", ""), + get("MODEL_AGENT_API_KEY_NAME", "")); + } + + public EmbeddingConfig embedding() { + boolean byteplus = "byteplus".equalsIgnoreCase(get("CLOUD_PROVIDER", "")); + return new EmbeddingConfig( + get( + "MODEL_EMBEDDING_NAME", + byteplus ? "skylark-embedding-vision-250615" : DEFAULT_EMBEDDING_MODEL), + getInt("MODEL_EMBEDDING_DIM", DEFAULT_EMBEDDING_DIMENSION), + get("MODEL_EMBEDDING_API_BASE", model().apiBase()), + get("MODEL_EMBEDDING_API_KEY", model().apiKey())); + } + + public OpenTelemetryConfig openTelemetry() { + return new OpenTelemetryConfig( + getBoolean("OBSERVABILITY_OPENTELEMETRY_TRACE_CONTENT", true), + get( + "OBSERVABILITY_OPENTELEMETRY_TLS_ENDPOINT", + "https://tls-cn-beijing.volces.com:4317"), + get("OBSERVABILITY_OPENTELEMETRY_TLS_REGION", "cn-beijing"), + get("OBSERVABILITY_OPENTELEMETRY_TLS_SERVICE_NAME", "")); + } + + /** Returns all resolved values. Prefer {@link #redactedValues()} for diagnostics. */ + public Map values() { + return values; + } + + public Map redactedValues() { + Map redacted = new LinkedHashMap<>(); + values.forEach( + (key, value) -> + redacted.put( + key, isSecret(key) && !value.isBlank() ? "" : value)); + return Collections.unmodifiableMap(redacted); + } + + private static boolean isSecret(String key) { + String upper = key.toUpperCase(Locale.ROOT); + return SECRET_MARKERS.stream().anyMatch(upper::contains); + } + + @Override + public String toString() { + return "VeADKConfig" + redactedValues(); + } + + public record ModelConfig( + String name, String provider, String apiBase, String apiKey, String apiKeyName) { + @Override + public String toString() { + return "ModelConfig[name=" + + name + + ", provider=" + + provider + + ", apiBase=" + + apiBase + + ", apiKey=" + + (apiKey.isBlank() ? "" : "") + + ", apiKeyName=" + + apiKeyName + + "]"; + } + } + + public record EmbeddingConfig(String name, int dimension, String apiBase, String apiKey) { + @Override + public String toString() { + return "EmbeddingConfig[name=" + + name + + ", dimension=" + + dimension + + ", apiBase=" + + apiBase + + ", apiKey=" + + (apiKey.isBlank() ? "" : "") + + "]"; + } + } + + public record OpenTelemetryConfig( + boolean traceContent, String tlsEndpoint, String tlsRegion, String tlsServiceName) {} +} diff --git a/core/src/main/java/com/volcengine/veadk/knowledgebase/InMemoryKnowledgebaseBackend.java b/core/src/main/java/com/volcengine/veadk/knowledgebase/InMemoryKnowledgebaseBackend.java new file mode 100644 index 0000000..bdd4780 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/knowledgebase/InMemoryKnowledgebaseBackend.java @@ -0,0 +1,84 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.knowledgebase; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Set; +import java.util.concurrent.CopyOnWriteArrayList; + +/** Deterministic in-memory backend suitable for local development and tests. */ +public final class InMemoryKnowledgebaseBackend implements KnowledgebaseBackend { + + private final String index; + private final CopyOnWriteArrayList entries = new CopyOnWriteArrayList<>(); + + public InMemoryKnowledgebaseBackend(String index) { + if (index == null || index.isBlank()) { + throw new IllegalArgumentException("index must not be blank"); + } + this.index = index; + } + + @Override + public String index() { + return index; + } + + @Override + public boolean add(List entries) { + this.entries.addAll(List.copyOf(entries)); + return true; + } + + @Override + public List search(String query, int topK) { + if (topK <= 0) { + throw new IllegalArgumentException("topK must be positive"); + } + Set queryTerms = terms(query); + List ranked = new ArrayList<>(entries.size()); + for (int index = 0; index < entries.size(); index++) { + KnowledgebaseEntry entry = entries.get(index); + Set contentTerms = terms(entry.content()); + long overlap = queryTerms.stream().filter(contentTerms::contains).count(); + ranked.add(new RankedEntry(entry, overlap, index)); + } + ranked.sort( + Comparator.comparingLong(RankedEntry::score) + .reversed() + .thenComparingInt(RankedEntry::insertionOrder)); + return ranked.stream().limit(topK).map(RankedEntry::entry).toList(); + } + + private static Set terms(String text) { + Set terms = new HashSet<>(); + if (text == null) { + return terms; + } + for (String term : text.toLowerCase(Locale.ROOT).split("[^\\p{L}\\p{N}_]+")) { + if (!term.isBlank()) { + terms.add(term); + } + } + return terms; + } + + private record RankedEntry(KnowledgebaseEntry entry, long score, int insertionOrder) {} +} diff --git a/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgeBase.java b/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgeBase.java new file mode 100644 index 0000000..97a0379 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgeBase.java @@ -0,0 +1,157 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.knowledgebase; + +import io.reactivex.rxjava3.core.Single; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** Backend-neutral knowledge-base facade aligned with the Python VeADK contract. */ +public final class KnowledgeBase implements BaseKnowledgebaseService, AutoCloseable { + + public static final String DEFAULT_NAME = "user_knowledgebase"; + public static final String DEFAULT_DESCRIPTION = + "This knowledgebase stores some user-related information."; + + private final String name; + private final String description; + private final KnowledgebaseBackend backend; + private final int topK; + + public KnowledgeBase(String index) { + this(new InMemoryKnowledgebaseBackend(index), DEFAULT_NAME, DEFAULT_DESCRIPTION, 10); + } + + public KnowledgeBase(KnowledgebaseBackend backend) { + this(backend, DEFAULT_NAME, DEFAULT_DESCRIPTION, 10); + } + + public KnowledgeBase(KnowledgebaseBackend backend, String name, String description, int topK) { + this.backend = Objects.requireNonNull(backend, "backend"); + this.name = requireText(name, "name"); + this.description = Objects.requireNonNullElse(description, ""); + if (topK <= 0) { + throw new IllegalArgumentException("topK must be positive"); + } + this.topK = topK; + } + + public String name() { + return name; + } + + public String description() { + return description; + } + + public String index() { + return backend.index(); + } + + public String backend() { + return backend.backendName(); + } + + public int topK() { + return topK; + } + + public boolean addFromText(String text) { + return addFromText(List.of(text)); + } + + public boolean addFromText(List texts) { + return backend.add(texts.stream().map(KnowledgebaseEntry::new).toList()); + } + + public boolean addFromFiles(List files) { + List entries = new ArrayList<>(); + for (Path file : files) { + try { + entries.add( + new KnowledgebaseEntry( + Files.readString(file, StandardCharsets.UTF_8), + Map.of("file_path", file.toAbsolutePath().normalize().toString()))); + } catch (IOException exception) { + throw new IllegalArgumentException( + "Unable to read knowledge file: " + file, exception); + } + } + return backend.add(entries); + } + + public boolean addFromDirectory(Path directory) { + if (!Files.isDirectory(directory)) { + throw new IllegalArgumentException("Knowledge directory does not exist: " + directory); + } + try (java.util.stream.Stream paths = Files.walk(directory)) { + return addFromFiles(paths.filter(Files::isRegularFile).sorted().toList()); + } catch (IOException exception) { + throw new IllegalArgumentException( + "Unable to read knowledge directory: " + directory, exception); + } + } + + public List search(String query) { + return search(query, topK); + } + + public List search(String query, int topK) { + return List.copyOf(backend.search(Objects.requireNonNull(query, "query"), topK)); + } + + @Override + public Single searchKnowledgebase(String query) { + return Single.fromCallable( + () -> { + List + legacyEntries = + search(query).stream() + .map(KnowledgeBase::toLegacyEntry) + .toList(); + SearchKnowledgebaseResponse response = new SearchKnowledgebaseResponse(); + response.setKnowledgebaseEntries(legacyEntries); + return response; + }); + } + + @Override + public void close() { + backend.close(); + } + + private static com.volcengine.veadk.integration.vikingknowledgebase.KnowledgebaseEntry + toLegacyEntry(KnowledgebaseEntry entry) { + Map metadata = new LinkedHashMap<>(); + entry.metadata().forEach((key, value) -> metadata.put(key, String.valueOf(value))); + return new com.volcengine.veadk.integration.vikingknowledgebase.KnowledgebaseEntry( + entry.content(), metadata); + } + + private static String requireText(String value, String field) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(field + " must not be blank"); + } + return value; + } +} diff --git a/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseBackend.java b/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseBackend.java new file mode 100644 index 0000000..44e896b --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseBackend.java @@ -0,0 +1,35 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.knowledgebase; + +import java.util.List; + +/** Storage and retrieval contract implemented by all knowledge-base backends. */ +public interface KnowledgebaseBackend extends AutoCloseable { + + String index(); + + boolean add(List entries); + + List search(String query, int topK); + + default String backendName() { + return getClass().getSimpleName(); + } + + @Override + default void close() {} +} diff --git a/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseEntry.java b/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseEntry.java new file mode 100644 index 0000000..14e3e5b --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseEntry.java @@ -0,0 +1,33 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.knowledgebase; + +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Objects; + +/** Backend-neutral knowledge-base entry. */ +public record KnowledgebaseEntry(String content, Map metadata) { + + public KnowledgebaseEntry { + content = Objects.requireNonNull(content, "content"); + metadata = Map.copyOf(new LinkedHashMap<>(Objects.requireNonNullElse(metadata, Map.of()))); + } + + public KnowledgebaseEntry(String content) { + this(content, Map.of()); + } +} diff --git a/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseServiceBackendAdapter.java b/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseServiceBackendAdapter.java new file mode 100644 index 0000000..d4d98d5 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/knowledgebase/KnowledgebaseServiceBackendAdapter.java @@ -0,0 +1,63 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.knowledgebase; + +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** Read-only adapter that exposes an existing Java knowledge-base service as a new backend. */ +public final class KnowledgebaseServiceBackendAdapter implements KnowledgebaseBackend { + + private final String index; + private final BaseKnowledgebaseService service; + + public KnowledgebaseServiceBackendAdapter(String index, BaseKnowledgebaseService service) { + if (index == null || index.isBlank()) { + throw new IllegalArgumentException("index must not be blank"); + } + this.index = index; + this.service = Objects.requireNonNull(service, "service"); + } + + @Override + public String index() { + return index; + } + + @Override + public boolean add(List entries) { + throw new UnsupportedOperationException("The legacy knowledge-base service is read-only"); + } + + @Override + public List search(String query, int topK) { + SearchKnowledgebaseResponse response = service.searchKnowledgebase(query).blockingGet(); + if (response == null || response.getKnowledgebaseEntries() == null) { + return List.of(); + } + return response.getKnowledgebaseEntries().stream() + .limit(topK) + .map( + entry -> + new KnowledgebaseEntry( + entry.getContent(), + entry.getMetadata() == null + ? Map.of() + : Map.copyOf(entry.getMetadata()))) + .toList(); + } +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/InMemoryLongTermMemoryBackend.java b/core/src/main/java/com/volcengine/veadk/memory/InMemoryLongTermMemoryBackend.java new file mode 100644 index 0000000..f796199 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/memory/InMemoryLongTermMemoryBackend.java @@ -0,0 +1,109 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import java.util.ArrayList; +import java.util.Comparator; +import java.util.HashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; + +/** Deterministic local backend suitable for development and tests. */ +public final class InMemoryLongTermMemoryBackend implements LongTermMemoryBackend { + + private final String index; + private final Map> memoriesByUser = + new ConcurrentHashMap<>(); + + public InMemoryLongTermMemoryBackend(String index) { + if (index == null || index.isBlank()) { + throw new IllegalArgumentException("index must not be blank"); + } + this.index = index; + } + + @Override + public String index() { + return index; + } + + @Override + public boolean saveMemory( + String userId, List eventStrings, Map options) { + String normalizedUserId = requireText(userId, "userId"); + List events = List.copyOf(eventStrings); + if (events.stream().anyMatch(event -> event == null || event.isBlank())) { + throw new IllegalArgumentException("eventStrings must not contain blank values"); + } + memoriesByUser + .computeIfAbsent(normalizedUserId, ignored -> new CopyOnWriteArrayList<>()) + .addAll(events); + return true; + } + + @Override + public List searchMemory( + String userId, String query, int topK, Map options) { + if (topK <= 0) { + throw new IllegalArgumentException("topK must be positive"); + } + requireText(query, "query"); + List memories = + memoriesByUser.getOrDefault( + requireText(userId, "userId"), new CopyOnWriteArrayList<>()); + Set queryTerms = terms(query); + List ranked = new ArrayList<>(memories.size()); + for (int position = 0; position < memories.size(); position++) { + String memory = memories.get(position); + Set memoryTerms = terms(memory); + long overlap = queryTerms.stream().filter(memoryTerms::contains).count(); + if (memory.toLowerCase(Locale.ROOT).contains(query.toLowerCase(Locale.ROOT))) { + overlap++; + } + ranked.add(new RankedMemory(memory, overlap, position)); + } + ranked.sort( + Comparator.comparingLong(RankedMemory::score) + .reversed() + .thenComparingInt(RankedMemory::insertionOrder)); + return ranked.stream().limit(topK).map(RankedMemory::value).toList(); + } + + public int size(String userId) { + return memoriesByUser.getOrDefault(userId, new CopyOnWriteArrayList<>()).size(); + } + + private static Set terms(String text) { + Set terms = new HashSet<>(); + for (String term : text.toLowerCase(Locale.ROOT).split("[^\\p{L}\\p{N}_]+")) { + if (!term.isBlank()) { + terms.add(term); + } + } + return terms; + } + + private static String requireText(String value, String field) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(field + " must not be blank"); + } + return value; + } + + private record RankedMemory(String value, long score, int insertionOrder) {} +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/InMemoryShortTermMemoryBackend.java b/core/src/main/java/com/volcengine/veadk/memory/InMemoryShortTermMemoryBackend.java new file mode 100644 index 0000000..6189c6b --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/memory/InMemoryShortTermMemoryBackend.java @@ -0,0 +1,28 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import com.google.adk.sessions.BaseSessionService; +import com.google.adk.sessions.InMemorySessionService; + +/** Local backend that keeps sessions in process memory. */ +public final class InMemoryShortTermMemoryBackend implements ShortTermMemoryBackend { + + private final InMemorySessionService sessionService = new InMemorySessionService(); + + @Override + public BaseSessionService sessionService() { + return sessionService; + } +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/LongTermMemory.java b/core/src/main/java/com/volcengine/veadk/memory/LongTermMemory.java new file mode 100644 index 0000000..d4ecc74 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/memory/LongTermMemory.java @@ -0,0 +1,244 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import com.fasterxml.jackson.databind.JsonNode; +import com.google.adk.events.Event; +import com.google.adk.memory.BaseMemoryService; +import com.google.adk.memory.MemoryEntry; +import com.google.adk.memory.SearchMemoryResponse; +import com.google.adk.sessions.Session; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import com.volcengine.veadk.utils.JSONUtil; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Single; +import java.io.IOException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** Backend-neutral long-term-memory facade aligned with the Python VeADK contract. */ +public final class LongTermMemory implements BaseMemoryService, AutoCloseable { + + public static final String DEFAULT_INDEX = "default_app"; + public static final int DEFAULT_TOP_K = 5; + + private static final Logger log = LoggerFactory.getLogger(LongTermMemory.class); + + private final LongTermMemoryBackend backend; + private final String appName; + private final int topK; + + public LongTermMemory() { + this(DEFAULT_INDEX); + } + + public LongTermMemory(String index) { + this(new InMemoryLongTermMemoryBackend(index), index, DEFAULT_TOP_K); + } + + public LongTermMemory(LongTermMemoryBackend backend) { + this(backend, backend.index(), DEFAULT_TOP_K); + } + + public LongTermMemory(LongTermMemoryBackend backend, int topK) { + this(backend, backend.index(), topK); + } + + public LongTermMemory(LongTermMemoryBackend backend, String appName, int topK) { + this.backend = Objects.requireNonNull(backend, "backend"); + this.appName = requireText(appName, "appName"); + if (topK <= 0) { + throw new IllegalArgumentException("topK must be positive"); + } + this.topK = topK; + } + + public String index() { + return backend.index(); + } + + public String appName() { + return appName; + } + + public String backend() { + return backend.backendName(); + } + + public int topK() { + return topK; + } + + @Override + public Completable addSessionToMemory(Session session) { + Objects.requireNonNull(session, "session"); + return Completable.fromAction( + () -> { + List eventStrings = serializeUserEvents(session.events()); + if (eventStrings.isEmpty()) { + return; + } + Map options = new LinkedHashMap<>(); + options.put( + "appName", + session.appName() == null || session.appName().isBlank() + ? appName + : session.appName()); + options.put("sessionId", session.id()); + backend.saveMemory(session.userId(), eventStrings, Map.copyOf(options)); + }); + } + + @Override + public Single searchMemory( + String requestedAppName, String userId, String query) { + return Single.fromCallable( + () -> { + List chunks; + try { + chunks = + backend.searchMemory( + userId, + query, + topK, + Map.of( + "appName", + requestedAppName == null + || requestedAppName.isBlank() + ? appName + : requestedAppName)); + } catch (Exception exception) { + log.warn( + "Long-term memory search failed; returning no memories", exception); + chunks = List.of(); + } + List entries = new ArrayList<>(); + chunks.forEach(chunk -> entries.addAll(toMemoryEntries(chunk))); + return SearchMemoryResponse.builder().setMemories(entries).build(); + }); + } + + @Override + public void close() { + backend.close(); + } + + private static List serializeUserEvents(List events) { + List serialized = new ArrayList<>(); + for (Event event : events) { + if (!"user".equals(event.author()) || event.content().isEmpty()) { + continue; + } + List parts = event.content().get().parts().orElse(List.of()); + List> textParts = + parts.stream() + .filter(part -> part.text().isPresent()) + .map(part -> Map.of("text", part.text().get())) + .toList(); + if (textParts.isEmpty()) { + continue; + } + String value = JSONUtil.toJson(Map.of("role", "user", "parts", textParts)); + if (!value.isBlank()) { + serialized.add(value); + } + } + return serialized; + } + + private static List toMemoryEntries(String chunk) { + try { + return toMemoryEntries(JSONUtil.parseJson(chunk)); + } catch (IOException exception) { + return List.of(memoryEntry("user", chunk)); + } + } + + private static List toMemoryEntries(JsonNode node) { + if (node == null || node.isNull()) { + return List.of(); + } + if (node.isTextual() || node.isNumber() || node.isBoolean()) { + return List.of(memoryEntry("user", node.asText())); + } + if (node.isArray()) { + List entries = new ArrayList<>(); + node.forEach(item -> entries.addAll(toMemoryEntries(item))); + return entries; + } + + JsonNode memories = node.get("memories"); + if (memories != null && memories.isArray()) { + return toMemoryEntries(memories); + } + + JsonNode content = node.path("content"); + String role = firstText(content.path("role"), node.path("role"), "user"); + JsonNode parts = content.has("parts") ? content.path("parts") : node.path("parts"); + List textParts = new ArrayList<>(); + if (parts.isArray()) { + for (JsonNode part : parts) { + String text = part.isTextual() ? part.asText() : part.path("text").asText(""); + if (!text.isBlank()) { + textParts.add(text); + } + } + } + if (textParts.isEmpty()) { + for (String field : List.of("text", "abstract", "summary")) { + String text = node.path(field).asText(""); + if (!text.isBlank()) { + textParts.add(text); + break; + } + } + } + if (textParts.isEmpty() && content.isTextual()) { + textParts.add(content.asText()); + } + return textParts.isEmpty() + ? List.of() + : List.of(memoryEntry(role, String.join("\n", textParts))); + } + + private static String firstText(JsonNode first, JsonNode second, String fallback) { + if (first != null && first.isTextual() && !first.asText().isBlank()) { + return first.asText(); + } + if (second != null && second.isTextual() && !second.asText().isBlank()) { + return second.asText(); + } + return fallback; + } + + private static MemoryEntry memoryEntry(String role, String text) { + return MemoryEntry.builder() + .author(role) + .content(Content.builder().role(role).parts(List.of(Part.fromText(text))).build()) + .build(); + } + + private static String requireText(String value, String field) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(field + " must not be blank"); + } + return value; + } +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/LongTermMemoryBackend.java b/core/src/main/java/com/volcengine/veadk/memory/LongTermMemoryBackend.java new file mode 100644 index 0000000..c014344 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/memory/LongTermMemoryBackend.java @@ -0,0 +1,42 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import java.util.List; +import java.util.Map; + +/** Storage and retrieval contract implemented by long-term-memory backends. */ +public interface LongTermMemoryBackend extends AutoCloseable { + + String index(); + + boolean saveMemory(String userId, List eventStrings, Map options); + + List searchMemory(String userId, String query, int topK, Map options); + + default boolean saveMemory(String userId, List eventStrings) { + return saveMemory(userId, eventStrings, Map.of()); + } + + default List searchMemory(String userId, String query, int topK) { + return searchMemory(userId, query, topK, Map.of()); + } + + default String backendName() { + return getClass().getSimpleName(); + } + + @Override + default void close() {} +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/MemoryServiceBackendAdapter.java b/core/src/main/java/com/volcengine/veadk/memory/MemoryServiceBackendAdapter.java new file mode 100644 index 0000000..78e2663 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/memory/MemoryServiceBackendAdapter.java @@ -0,0 +1,59 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import com.google.adk.memory.BaseMemoryService; +import com.google.adk.memory.SearchMemoryResponse; +import com.volcengine.veadk.utils.JSONUtil; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** Read-only adapter that exposes an existing ADK memory service as a new backend. */ +public final class MemoryServiceBackendAdapter implements LongTermMemoryBackend { + + private final String index; + private final BaseMemoryService service; + + public MemoryServiceBackendAdapter(String index, BaseMemoryService service) { + if (index == null || index.isBlank()) { + throw new IllegalArgumentException("index must not be blank"); + } + this.index = index; + this.service = Objects.requireNonNull(service, "service"); + } + + @Override + public String index() { + return index; + } + + @Override + public boolean saveMemory( + String userId, List eventStrings, Map options) { + throw new UnsupportedOperationException("The legacy memory service is read-only"); + } + + @Override + public List searchMemory( + String userId, String query, int topK, Map options) { + Object configuredAppName = options.get("appName"); + String appName = configuredAppName == null ? index : configuredAppName.toString(); + SearchMemoryResponse response = service.searchMemory(appName, userId, query).blockingGet(); + if (response == null || response.memories() == null) { + return List.of(); + } + return response.memories().stream().limit(topK).map(JSONUtil::toJson).toList(); + } +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/SessionServiceBackendAdapter.java b/core/src/main/java/com/volcengine/veadk/memory/SessionServiceBackendAdapter.java new file mode 100644 index 0000000..56dd5d0 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/memory/SessionServiceBackendAdapter.java @@ -0,0 +1,32 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import com.google.adk.sessions.BaseSessionService; +import java.util.Objects; + +/** Adapts an existing ADK session service to the VeADK short-term-memory backend SPI. */ +public final class SessionServiceBackendAdapter implements ShortTermMemoryBackend { + + private final BaseSessionService sessionService; + + public SessionServiceBackendAdapter(BaseSessionService sessionService) { + this.sessionService = Objects.requireNonNull(sessionService, "sessionService"); + } + + @Override + public BaseSessionService sessionService() { + return sessionService; + } +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemory.java b/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemory.java new file mode 100644 index 0000000..74fa20b --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemory.java @@ -0,0 +1,153 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import com.google.adk.events.Event; +import com.google.adk.sessions.BaseSessionService; +import com.google.adk.sessions.GetSessionConfig; +import com.google.adk.sessions.ListEventsResponse; +import com.google.adk.sessions.ListSessionsResponse; +import com.google.adk.sessions.Session; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.function.Consumer; + +/** ADK-compatible short-term-memory facade aligned with the Python VeADK contract. */ +public final class ShortTermMemory implements BaseSessionService, AutoCloseable { + + private final ShortTermMemoryBackend backend; + private final BaseSessionService delegate; + private final ShortTermMemoryProcessor processor; + private final Consumer afterLoadMemoryCallback; + + public ShortTermMemory() { + this(new InMemoryShortTermMemoryBackend()); + } + + public ShortTermMemory(BaseSessionService sessionService) { + this(new SessionServiceBackendAdapter(sessionService)); + } + + public ShortTermMemory(ShortTermMemoryBackend backend) { + this(backend, null, session -> {}); + } + + public ShortTermMemory( + ShortTermMemoryBackend backend, Consumer afterLoadMemoryCallback) { + this(backend, null, afterLoadMemoryCallback); + } + + public ShortTermMemory(ShortTermMemoryBackend backend, ShortTermMemoryProcessor processor) { + this(backend, processor, session -> {}); + } + + public ShortTermMemory( + ShortTermMemoryBackend backend, + ShortTermMemoryProcessor processor, + Consumer afterLoadMemoryCallback) { + this.backend = Objects.requireNonNull(backend, "backend"); + this.delegate = Objects.requireNonNull(backend.sessionService(), "sessionService"); + this.processor = processor; + this.afterLoadMemoryCallback = + Objects.requireNonNull(afterLoadMemoryCallback, "afterLoadMemoryCallback"); + } + + public BaseSessionService sessionService() { + return this; + } + + public String backend() { + return backend.backendName(); + } + + /** Creates a session, or returns the existing session when the identifier is already present. */ + public Single createSession(String appName, String userId, String sessionId) { + return getSession(appName, userId, sessionId, Optional.empty()) + .switchIfEmpty( + createSession(appName, userId, new ConcurrentHashMap<>(), sessionId)); + } + + @Override + public Single createSession( + String appName, String userId, ConcurrentMap state, String sessionId) { + return delegate.createSession(appName, userId, state, sessionId); + } + + @Override + public Maybe getSession( + String appName, String userId, String sessionId, Optional config) { + return delegate.getSession(appName, userId, sessionId, config) + .map(session -> processor == null ? session : processor.afterLoadSession(session)) + .doOnSuccess(afterLoadMemoryCallback::accept); + } + + @Override + public Single listSessions(String appName, String userId) { + return delegate.listSessions(appName, userId); + } + + @Override + public Completable deleteSession(String appName, String userId, String sessionId) { + return delegate.deleteSession(appName, userId, sessionId); + } + + @Override + public Single listEvents(String appName, String userId, String sessionId) { + return delegate.listEvents(appName, userId, sessionId); + } + + @Override + public Completable closeSession(Session session) { + return delegate.closeSession(session); + } + + @Override + public Single appendEvent(Session session, Event event) { + if (processor != null) { + // The processor returns a presentation-only session. Keep that view current for the + // active invocation, but append to a freshly loaded canonical session so optimized or + // filtered history never replaces the persisted history. + return BaseSessionService.super + .appendEvent(session, event) + .flatMap( + appendedEvent -> + delegate.getSession( + session.appName(), + session.userId(), + session.id(), + Optional.empty()) + .switchIfEmpty( + Single.error( + new IllegalStateException( + "Session not found: " + + session.id()))) + .flatMap( + persistedSession -> + delegate.appendEvent( + persistedSession, + appendedEvent))); + } + return delegate.appendEvent(session, event); + } + + @Override + public void close() { + backend.close(); + } +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryBackend.java b/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryBackend.java new file mode 100644 index 0000000..59938db --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryBackend.java @@ -0,0 +1,29 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import com.google.adk.sessions.BaseSessionService; + +/** Contract implemented by short-term-memory storage backends. */ +public interface ShortTermMemoryBackend extends AutoCloseable { + + BaseSessionService sessionService(); + + default String backendName() { + return getClass().getSimpleName(); + } + + @Override + default void close() {} +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryMessage.java b/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryMessage.java new file mode 100644 index 0000000..e09c155 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryMessage.java @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +/** A provider-neutral chat message used by short-term-memory optimizers. */ +public record ShortTermMemoryMessage(String role, String content) { + + public ShortTermMemoryMessage { + if (role == null || role.isBlank()) { + throw new IllegalArgumentException("role must not be blank"); + } + if (content == null || content.isBlank()) { + throw new IllegalArgumentException("content must not be blank"); + } + } +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryProcessor.java b/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryProcessor.java new file mode 100644 index 0000000..4a81dca --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/memory/ShortTermMemoryProcessor.java @@ -0,0 +1,82 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import com.google.adk.events.Event; +import com.google.adk.sessions.Session; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Function; + +/** Rewrites loaded session history using a caller-provided message optimizer. */ +public final class ShortTermMemoryProcessor { + + private final Function, List> optimizer; + + public ShortTermMemoryProcessor( + Function, List> optimizer) { + this.optimizer = Objects.requireNonNull(optimizer, "optimizer"); + } + + public Session afterLoadSession(Session session) { + Objects.requireNonNull(session, "session"); + List sourceMessages = extractMessages(session.events()); + List optimizedMessages = + List.copyOf(Objects.requireNonNull(optimizer.apply(sourceMessages), "messages")); + List optimizedEvents = + optimizedMessages.stream().map(ShortTermMemoryProcessor::toEvent).toList(); + return Session.builder(session.id()) + .appName(session.appName()) + .userId(session.userId()) + .state(new ConcurrentHashMap<>(session.state())) + // ADK's default appendEvent implementation appends in place. Stream.toList() + // returns an unmodifiable list, so keep processor-created session views mutable. + .events(new ArrayList<>(optimizedEvents)) + .lastUpdateTime(session.lastUpdateTime()) + .build(); + } + + private static List extractMessages(List events) { + List messages = new ArrayList<>(); + for (Event event : events) { + if (event.content().isEmpty()) { + continue; + } + Content content = event.content().get(); + List parts = content.parts().orElse(List.of()); + if (parts.isEmpty() || parts.get(0).text().isEmpty()) { + continue; + } + messages.add( + new ShortTermMemoryMessage( + content.role().orElse(event.author()), parts.get(0).text().get())); + } + return List.copyOf(messages); + } + + private static Event toEvent(ShortTermMemoryMessage message) { + return Event.builder() + .author("memory_optimizer") + .content( + Content.builder() + .role(message.role()) + .parts(List.of(Part.fromText(message.content()))) + .build()) + .build(); + } +} diff --git a/core/src/main/java/com/volcengine/veadk/memory/viking/VikingMemoryService.java b/core/src/main/java/com/volcengine/veadk/memory/viking/VikingMemoryService.java index 858925c..e429a14 100644 --- a/core/src/main/java/com/volcengine/veadk/memory/viking/VikingMemoryService.java +++ b/core/src/main/java/com/volcengine/veadk/memory/viking/VikingMemoryService.java @@ -15,6 +15,7 @@ */ package com.volcengine.veadk.memory.viking; +import com.fasterxml.jackson.databind.JsonNode; import com.google.adk.memory.BaseMemoryService; import com.google.adk.memory.MemoryEntry; import com.google.adk.memory.SearchMemoryResponse; @@ -22,15 +23,20 @@ import com.volcengine.veadk.integration.vikingmemory.Message; import com.volcengine.veadk.integration.vikingmemory.Metadata; import com.volcengine.veadk.integration.vikingmemory.VikingMemoryWrapper; +import com.volcengine.veadk.memory.LongTermMemoryBackend; import com.volcengine.veadk.utils.EnvUtil; +import com.volcengine.veadk.utils.JSONUtil; import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Single; +import java.io.IOException; +import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.stream.Collectors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; -public class VikingMemoryService implements BaseMemoryService { +public class VikingMemoryService implements BaseMemoryService, LongTermMemoryBackend { private static final Logger log = LoggerFactory.getLogger(VikingMemoryService.class); @@ -115,4 +121,74 @@ public Single searchMemory(String appName, String userId, return SearchMemoryResponse.builder().setMemories(memoryEntries).build(); }); } + + @Override + public String index() { + return appName; + } + + @Override + public boolean saveMemory( + String userId, List eventStrings, Map options) { + List messages = + eventStrings.stream() + .map(VikingMemoryService::toMessage) + .collect(Collectors.toList()); + if (messages.isEmpty()) { + return true; + } + try { + Metadata metadata = new Metadata(userId, "assistant", System.currentTimeMillis()); + return vikingMemoryWrapper.addSession(appName, messages, metadata); + } catch (Exception exception) { + throw new IllegalStateException("Unable to save Viking long-term memory", exception); + } + } + + @Override + public List searchMemory( + String userId, String query, int requestedTopK, Map options) { + try { + return vikingMemoryWrapper + .searchMemory(appName, userId, query, requestedTopK, builtinEventTypes) + .stream() + .map(JSONUtil::toJson) + .toList(); + } catch (Exception exception) { + throw new IllegalStateException("Unable to search Viking long-term memory", exception); + } + } + + private static Message toMessage(String eventString) { + try { + JsonNode event = JSONUtil.parseJson(eventString); + JsonNode content = event.path("content"); + String role = content.path("role").asText(event.path("role").asText("user")); + JsonNode parts = content.has("parts") ? content.path("parts") : event.path("parts"); + List textParts = new ArrayList<>(); + if (parts.isArray()) { + for (JsonNode part : parts) { + String text = part.isTextual() ? part.asText() : part.path("text").asText(""); + if (!text.isBlank()) { + textParts.add(text); + } + } + } + String text = String.join("\n", textParts); + if (!text.isBlank()) { + return new Message(normalizeRole(role), text); + } + } catch (IOException exception) { + // Plain text is a valid backend input and is handled below. + } + return new Message("user", eventString); + } + + private static String normalizeRole(String role) { + return switch (role) { + case "assistant", "model" -> "assistant"; + case "system" -> "system"; + default -> "user"; + }; + } } diff --git a/core/src/main/java/com/volcengine/veadk/model/ArkEmbedding.java b/core/src/main/java/com/volcengine/veadk/model/ArkEmbedding.java new file mode 100644 index 0000000..0897c3f --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/model/ArkEmbedding.java @@ -0,0 +1,230 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.model; + +import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbeddingInput; +import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbeddingRequest; +import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbeddingResult; +import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbeddingUsage; +import com.volcengine.ark.runtime.service.ArkService; +import com.volcengine.veadk.config.VeADKConfig; +import io.reactivex.rxjava3.core.Single; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; + +/** Ark multimodal embedding adapter with synchronous, reactive, and batched text APIs. */ +public final class ArkEmbedding { + + public enum Model { + DOUBAO_EMBEDDING_VISION_251215("doubao-embedding-vision-251215"), + DOUBAO_EMBEDDING_VISION_250615("doubao-embedding-vision-250615"); + + private final String id; + + Model(String id) { + this.id = id; + } + + public String id() { + return id; + } + } + + private final String modelName; + private final Integer dimensions; + private final ArkService arkService; + + public ArkEmbedding(String modelName, String apiKey) { + this(builder().modelName(modelName).apiKey(apiKey)); + } + + public ArkEmbedding(String modelName, String apiKey, String apiBase, Integer dimensions) { + this(builder().modelName(modelName).apiKey(apiKey).apiBase(apiBase).dimensions(dimensions)); + } + + private ArkEmbedding(Builder builder) { + modelName = requireText(builder.modelName, "modelName"); + dimensions = builder.dimensions; + if (dimensions != null && dimensions <= 0) { + throw new IllegalArgumentException("dimensions must be positive"); + } + arkService = builder.arkService != null ? builder.arkService : createService(builder); + } + + public static Builder builder() { + return new Builder(); + } + + public static Builder builder(VeADKConfig config) { + Objects.requireNonNull(config, "config"); + VeADKConfig.EmbeddingConfig embedding = config.embedding(); + return new Builder() + .modelName(embedding.name()) + .apiKey(embedding.apiKey()) + .apiBase(embedding.apiBase()) + .dimensions(embedding.dimension()); + } + + public String modelName() { + return modelName; + } + + public Integer dimensions() { + return dimensions; + } + + public List embed(String text) { + return embedWithUsage(text).embedding(); + } + + public EmbeddingResponse embedWithUsage(String text) { + Objects.requireNonNull(text, "text"); + MultimodalEmbeddingInput input = + MultimodalEmbeddingInput.builder().type("text").text(text).build(); + MultimodalEmbeddingRequest.Builder request = + MultimodalEmbeddingRequest.builder().model(modelName).input(List.of(input)); + if (dimensions != null) { + request.dimensions(dimensions); + } + MultimodalEmbeddingResult result = arkService.createMultiModalEmbeddings(request.build()); + if (result == null || result.getData() == null || result.getData().getEmbedding() == null) { + throw new IllegalStateException("Ark embedding response did not contain an embedding"); + } + MultimodalEmbeddingUsage usage = result.getUsage(); + return new EmbeddingResponse( + result.getModel() == null ? modelName : result.getModel(), + result.getData().getEmbedding(), + usage == null ? 0 : usage.getPromptTokens(), + usage == null ? 0 : usage.getTotalTokens()); + } + + public List> embedAll(List texts) { + Objects.requireNonNull(texts, "texts"); + List> embeddings = new ArrayList<>(texts.size()); + for (String text : texts) { + embeddings.add(embed(text)); + } + return List.copyOf(embeddings); + } + + public Single> embedAsync(String text) { + return Single.fromCallable(() -> embed(text)); + } + + public Single>> embedAllAsync(List texts) { + return Single.fromCallable(() -> embedAll(texts)); + } + + public List getTextEmbedding(String text) { + return embed(text); + } + + public List> getTextEmbeddings(List texts) { + return embedAll(texts); + } + + public List getQueryEmbedding(String query) { + return embed(query); + } + + private static ArkService createService(Builder builder) { + ArkService.Builder service = + ArkService.builder() + .apiKey(requireText(builder.apiKey, "apiKey")) + .timeout(builder.timeout) + .retryTimes(builder.maxRetries); + if (builder.apiBase != null && !builder.apiBase.isBlank()) { + service.baseUrl(builder.apiBase); + } + return service.build(); + } + + private static String requireText(String value, String field) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(field + " must not be blank"); + } + return value; + } + + public record EmbeddingResponse( + String model, List embedding, long promptTokens, long totalTokens) { + public EmbeddingResponse { + model = Objects.requireNonNullElse(model, ""); + embedding = List.copyOf(embedding); + } + } + + public static final class Builder { + private String modelName = VeADKConfig.DEFAULT_EMBEDDING_MODEL; + private String apiKey; + private String apiBase; + private Integer dimensions; + private int maxRetries = 10; + private Duration timeout = Duration.ofSeconds(60); + private ArkService arkService; + + private Builder() {} + + public Builder modelName(String modelName) { + this.modelName = modelName; + return this; + } + + public Builder model(Model model) { + return modelName(Objects.requireNonNull(model, "model").id()); + } + + public Builder apiKey(String apiKey) { + this.apiKey = apiKey; + return this; + } + + public Builder apiBase(String apiBase) { + this.apiBase = apiBase; + return this; + } + + public Builder dimensions(Integer dimensions) { + this.dimensions = dimensions; + return this; + } + + public Builder maxRetries(int maxRetries) { + if (maxRetries < 0) { + throw new IllegalArgumentException("maxRetries must not be negative"); + } + this.maxRetries = maxRetries; + return this; + } + + public Builder timeout(Duration timeout) { + this.timeout = Objects.requireNonNull(timeout, "timeout"); + return this; + } + + /** Supplies a preconfigured client, primarily for custom transports and deterministic tests. */ + public Builder arkService(ArkService arkService) { + this.arkService = Objects.requireNonNull(arkService, "arkService"); + return this; + } + + public ArkEmbedding build() { + return new ArkEmbedding(this); + } + } +} diff --git a/core/src/main/java/com/volcengine/veadk/model/ArkLlm.java b/core/src/main/java/com/volcengine/veadk/model/ArkLlm.java index c9d8196..90b0204 100644 --- a/core/src/main/java/com/volcengine/veadk/model/ArkLlm.java +++ b/core/src/main/java/com/volcengine/veadk/model/ArkLlm.java @@ -24,27 +24,40 @@ import com.google.adk.tools.BaseTool; import com.google.common.collect.ImmutableMap; import com.google.genai.types.Content; +import com.google.genai.types.FinishReason; +import com.google.genai.types.FunctionCall; import com.google.genai.types.FunctionDeclaration; import com.google.genai.types.FunctionResponse; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.Part; import com.google.genai.types.Schema; +import com.google.genai.types.VideoMetadata; +import com.volcengine.ark.runtime.model.Usage; import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChunk; +import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionContentPart; import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest; import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionResult; import com.volcengine.ark.runtime.model.completion.chat.ChatFunction; +import com.volcengine.ark.runtime.model.completion.chat.ChatFunctionCall; import com.volcengine.ark.runtime.model.completion.chat.ChatMessage; import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole; import com.volcengine.ark.runtime.model.completion.chat.ChatTool; import com.volcengine.ark.runtime.model.completion.chat.ChatToolCall; +import com.volcengine.ark.runtime.model.completion.chat.ResponseFormatJSONSchemaJSONSchemaParam; import com.volcengine.ark.runtime.service.ArkService; import com.volcengine.veadk.utils.EnvUtil; import com.volcengine.veadk.utils.JSONUtil; import io.reactivex.rxjava3.core.Flowable; import java.util.ArrayList; +import java.util.Base64; +import java.util.LinkedHashMap; import java.util.List; import java.util.Map; import java.util.Objects; import java.util.Optional; +import java.util.TreeMap; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Collectors; import java.util.stream.Stream; import org.apache.commons.lang3.StringUtils; @@ -70,6 +83,7 @@ public final class ArkLlm extends BaseLlm { .build(); private final ArkService arkService; + private final List fallbacks; private ChatCompletionRequest.ChatCompletionRequestThinking thinking = null; public ArkLlm(String modelName) { @@ -77,14 +91,88 @@ public ArkLlm(String modelName) { } public ArkLlm(String modelName, String thinking) { + this(modelName, thinking, EnvUtil.getAgentApiKey(), null); + } + + public ArkLlm(String modelName, String apiKey, String apiBase) { + this(modelName, null, apiKey, apiBase); + } + + public ArkLlm(String modelName, String thinking, String apiKey, String apiBase) { + this(modelName, thinking, apiKey, apiBase, List.of()); + } + + public ArkLlm(List modelNames) { + this(modelNames, null); + } + + public ArkLlm(List modelNames, String thinking) { + this(primaryModelName(modelNames), thinking, EnvUtil.getAgentApiKey(), null, modelNames); + } + + public ArkLlm(List modelNames, String apiKey, String apiBase) { + this(primaryModelName(modelNames), null, apiKey, apiBase, modelNames); + } + + public ArkLlm(List modelNames, String thinking, String apiKey, String apiBase) { + this(primaryModelName(modelNames), thinking, apiKey, apiBase, modelNames); + } + + private ArkLlm( + String modelName, String thinking, String apiKey, String apiBase, List models) { super(modelName); - Objects.requireNonNull(modelName, "modelName must be set."); - this.arkService = ArkService.builder().apiKey(EnvUtil.getAgentApiKey()).build(); + requireModelName(modelName); + this.fallbacks = fallbackModelNames(models); + ArkService.Builder serviceBuilder = + ArkService.builder().apiKey(Objects.requireNonNull(apiKey, "apiKey must be set.")); + if (StringUtils.isNotBlank(apiBase)) { + serviceBuilder.baseUrl(apiBase); + } + this.arkService = serviceBuilder.build(); if (StringUtils.isNotBlank(thinking)) { this.thinking = new ChatCompletionRequest.ChatCompletionRequestThinking(thinking); } } + public List fallbacks() { + return fallbacks; + } + + private static String primaryModelName(List modelNames) { + Objects.requireNonNull(modelNames, "modelNames must be set."); + if (modelNames.isEmpty()) { + throw new IllegalArgumentException("modelNames must not be empty"); + } + return requireModelName(modelNames.get(0)); + } + + private static List fallbackModelNames(List modelNames) { + if (modelNames == null || modelNames.size() <= 1) { + return List.of(); + } + return modelNames.subList(1, modelNames.size()).stream() + .map(ArkLlm::requireModelName) + .toList(); + } + + private static String requireModelName(String modelName) { + Objects.requireNonNull(modelName, "modelName must be set."); + if (modelName.isBlank()) { + throw new IllegalArgumentException("modelName must not be blank"); + } + return modelName; + } + + private List modelCandidates(LlmRequest llmRequest) { + Optional requestedModel = llmRequest.model().filter(StringUtils::isNotBlank); + // ADK's Basic request processor always copies BaseLlm.model() into the request. Treat that + // value as the configured primary model so fallbacks remain available in normal Agent runs. + if (requestedModel.isPresent() && !requestedModel.get().equals(model())) { + return List.of(requestedModel.get()); + } + return Stream.concat(Stream.of(model()), fallbacks.stream()).toList(); + } + /** * Generate content based on LLM request * @param llmRequest The request containing prompts and parameters @@ -93,19 +181,49 @@ public ArkLlm(String modelName, String thinking) { */ @Override public Flowable generateContent(LlmRequest llmRequest, boolean stream) { - // Convert ADK request to Ark request format - ChatCompletionRequest arkRequest = toArkRequest(llmRequest); + return Flowable.defer( + () -> + generateContentWithFallbacks( + llmRequest, stream, modelCandidates(llmRequest), 0)); + } + + private Flowable generateContentWithFallbacks( + LlmRequest llmRequest, boolean stream, List modelCandidates, int index) { + String modelName = modelCandidates.get(index); + ChatCompletionRequest arkRequest = toArkRequest(llmRequest, modelName); + Flowable attempt = + Flowable.defer(() -> generateContentWithModel(arkRequest, stream)); + AtomicBoolean emittedResponse = new AtomicBoolean(false); + return attempt.doOnNext(response -> emittedResponse.set(true)) + .onErrorResumeNext( + error -> { + if (emittedResponse.get() || index + 1 >= modelCandidates.size()) { + return Flowable.error(error); + } + String nextModel = modelCandidates.get(index + 1); + log.warn( + "Ark request with model {} failed before emitting a response;" + + " falling back to {}", + modelName, + nextModel, + error); + return generateContentWithFallbacks( + llmRequest, stream, modelCandidates, index + 1); + }); + } + + private Flowable generateContentWithModel( + ChatCompletionRequest arkRequest, boolean stream) { if (stream) { log.debug( "Sending streaming generateContent request to model {}", arkRequest.getModel()); - // Handle streaming response + arkRequest.setStreamOptions( + ChatCompletionRequest.ChatCompletionRequestStreamOptions.of(true)); return generateContentStreaming(arkRequest); - } else { - log.debug("Sending generateContent request to model {}", arkRequest.getModel()); - // Handle non-streaming response - return Flowable.fromCallable(() -> arkService.createChatCompletion(arkRequest)) - .map(this::toLlmResponse); } + log.debug("Sending generateContent request to model {}", arkRequest.getModel()); + return Flowable.fromCallable(() -> arkService.createChatCompletion(arkRequest)) + .map(this::toLlmResponse); } /** @@ -114,70 +232,64 @@ public Flowable generateContent(LlmRequest llmRequest, boolean stre * @return Flowable of LlmResponse objects */ private Flowable generateContentStreaming(ChatCompletionRequest arkRequest) { - // Get streaming response from Ark service io.reactivex.Flowable streamResponse = arkService.streamChatCompletion(arkRequest); return Flowable.defer( () -> { - // Accumulate complete text response final StringBuilder accumulatedText = new StringBuilder(); - // Buffer partial text for incremental responses final StringBuilder partialText = new StringBuilder(); - // Hold the last chunk for final processing - final ChatCompletionChunk[] lastChunkHolder = {null}; - // Accumulate tool calls if any - final List accumulatedToolCalls = new ArrayList<>(); + final Map accumulatedToolCalls = new TreeMap<>(); + final String[] finishReason = {null}; + final Usage[] usage = {null}; return Flowable.fromPublisher(streamResponse) .concatMap( chunk -> { - lastChunkHolder[0] = chunk; log.debug("Raw Ark streaming chunk: {}", chunk); + if (chunk.getUsage() != null) { + usage[0] = chunk.getUsage(); + } + String chunkFinishReason = finishReason(chunk); + if (StringUtils.isNotBlank(chunkFinishReason)) { + finishReason[0] = chunkFinishReason; + } - // Prepare list of responses to emit List responsesToEmit = new ArrayList<>(); - // Process text content from chunk processTextContent( chunk, accumulatedText, partialText, responsesToEmit); - // Process tool calls from chunk processToolCalls(chunk, accumulatedToolCalls); - // Handle stop chunk (final chunk) - if (isStopChunk(chunk)) { + if (StringUtils.isNotBlank(chunkFinishReason)) { processStopChunk(partialText, responsesToEmit); } - // Emit responses if any if (responsesToEmit.isEmpty()) { return Flowable.empty(); - } else { - log.debug("Responses to emit: {}", responsesToEmit); - return Flowable.fromIterable(responsesToEmit); } + log.debug("Responses to emit: {}", responsesToEmit); + return Flowable.fromIterable(responsesToEmit); }) .concatWith( Flowable.defer( () -> { - // Process final response after stream ends - ChatCompletionChunk lastChunk = lastChunkHolder[0]; - if (lastChunk == null - || accumulatedText.length() == 0) { + if (StringUtils.isBlank(finishReason[0]) + || (accumulatedText.isEmpty() + && accumulatedToolCalls + .isEmpty())) { return Flowable.empty(); } - - if (isStopChunk(lastChunk)) { - // Build and emit final aggregated response - return Flowable.just( - buildFinalResponse( - accumulatedText.toString(), - accumulatedToolCalls)); - } - - return Flowable.empty(); + return Flowable.just( + buildFinalResponse( + accumulatedText.toString(), + new ArrayList<>( + accumulatedToolCalls + .values()), + finishReason[0], + usage[0])); })); }); } @@ -194,18 +306,18 @@ private void processTextContent( StringBuilder accumulatedText, StringBuilder partialText, List responsesToEmit) { - // Extract text content from chunk - String content = (String) chunk.getChoices().get(0).getMessage().getContent(); + if (!hasMessageChoice(chunk)) { + return; + } + Object rawContent = chunk.getChoices().get(0).getMessage().getContent(); + String content = rawContent instanceof String ? (String) rawContent : null; if (StringUtils.isNotEmpty(content)) { - // Add to accumulated text accumulatedText.append(content); - // Add to partial text buffer partialText.append(content); - // Emit partial response when buffer exceeds threshold if (partialText.length() > 30) { responsesToEmit.add(buildPartialResponse(partialText.toString())); - partialText.setLength(0); // Clear buffer + partialText.setLength(0); } } } @@ -228,22 +340,73 @@ private LlmResponse buildPartialResponse(String text) { * @param accumulatedToolCalls List of accumulated tool calls */ private void processToolCalls( - ChatCompletionChunk chunk, List accumulatedToolCalls) { - // Extract tool calls from chunk + ChatCompletionChunk chunk, Map accumulatedToolCalls) { + if (!hasMessageChoice(chunk)) { + return; + } List toolCalls = chunk.getChoices().get(0).getMessage().getToolCalls(); - if (null != toolCalls && !toolCalls.isEmpty()) { - ChatToolCall toolCall = toolCalls.get(0); - // If tool call has ID, it's a new tool call - if (StringUtils.isNotBlank(toolCall.getId())) { - accumulatedToolCalls.add(toolCall); - } else { - // Otherwise, it's continuation of existing tool call - int index = toolCall.getIndex(); - String arguments = - accumulatedToolCalls.get(index).getFunction().getArguments() - + toolCall.getFunction().getArguments(); - accumulatedToolCalls.get(index).getFunction().setArguments(arguments); + if (toolCalls == null) { + return; + } + for (ChatToolCall toolCall : toolCalls) { + if (toolCall == null) { + continue; } + int index = resolveToolCallIndex(toolCall, accumulatedToolCalls); + ChatToolCall accumulated = + accumulatedToolCalls.computeIfAbsent(index, unused -> new ChatToolCall()); + mergeToolCall(accumulated, toolCall); + } + } + + private int resolveToolCallIndex( + ChatToolCall toolCall, Map accumulatedToolCalls) { + if (toolCall.getIndex() != null && toolCall.getIndex() >= 0) { + return toolCall.getIndex(); + } + if (StringUtils.isNotBlank(toolCall.getId())) { + return accumulatedToolCalls.entrySet().stream() + .filter(entry -> toolCall.getId().equals(entry.getValue().getId())) + .map(Map.Entry::getKey) + .findFirst() + .orElseGet(() -> nextToolCallIndex(accumulatedToolCalls)); + } + return accumulatedToolCalls.size() == 1 + ? accumulatedToolCalls.keySet().iterator().next() + : nextToolCallIndex(accumulatedToolCalls); + } + + private int nextToolCallIndex(Map accumulatedToolCalls) { + int index = 0; + while (accumulatedToolCalls.containsKey(index)) { + index++; + } + return index; + } + + private void mergeToolCall(ChatToolCall target, ChatToolCall fragment) { + if (StringUtils.isNotBlank(fragment.getId())) { + target.setId(fragment.getId()); + } + if (StringUtils.isNotBlank(fragment.getType())) { + target.setType(fragment.getType()); + } + ChatFunctionCall fragmentFunction = fragment.getFunction(); + if (fragmentFunction == null) { + return; + } + ChatFunctionCall targetFunction = target.getFunction(); + if (targetFunction == null) { + targetFunction = new ChatFunctionCall(); + target.setFunction(targetFunction); + } + if (StringUtils.isNotBlank(fragmentFunction.getName())) { + targetFunction.setName(fragmentFunction.getName()); + } + if (fragmentFunction.getArguments() != null) { + targetFunction.setArguments( + Objects.requireNonNullElse(targetFunction.getArguments(), "") + + fragmentFunction.getArguments()); } } @@ -256,7 +419,7 @@ private void processStopChunk(StringBuilder partialText, List respo // Emit any remaining partial text if (!partialText.isEmpty()) { responsesToEmit.add(buildPartialResponse(partialText.toString())); - partialText.setLength(0); // Clear buffer + partialText.setLength(0); } } @@ -267,22 +430,26 @@ private void processStopChunk(StringBuilder partialText, List respo * @return Final LlmResponse object */ private LlmResponse buildFinalResponse( - String accumulatedText, List accumulatedToolCalls) { + String accumulatedText, + List accumulatedToolCalls, + String finishReason, + Usage usage) { List parts = new ArrayList<>(); - // Add text part - parts.add(Part.fromText(accumulatedText)); + if (StringUtils.isNotEmpty(accumulatedText)) { + parts.add(Part.fromText(accumulatedText)); + } - // Add tool call parts if any if (!accumulatedToolCalls.isEmpty()) { parts.addAll(parseToolCalls(accumulatedToolCalls)); } - // Build final response - LlmResponse finalAggregatedResponse = + LlmResponse.Builder builder = LlmResponse.builder() .content(Content.builder().role("model").parts(parts).build()) - .partial(false) // Mark as complete response - .build(); + .partial(false); + toFinishReason(finishReason).ifPresent(builder::finishReason); + toUsageMetadata(usage).ifPresent(builder::usageMetadata); + LlmResponse finalAggregatedResponse = builder.build(); log.debug("finalAggregatedResponse to emit: {}", finalAggregatedResponse); return finalAggregatedResponse; } @@ -292,9 +459,16 @@ private LlmResponse buildFinalResponse( * @param chunk The streaming chunk * @return True if chunk is stop chunk, false otherwise */ - private boolean isStopChunk(ChatCompletionChunk chunk) { - String finishReason = chunk.getChoices().get(0).getFinishReason(); - return StringUtils.isNotBlank(finishReason); + private String finishReason(ChatCompletionChunk chunk) { + return chunk.getChoices() == null || chunk.getChoices().isEmpty() + ? null + : chunk.getChoices().get(0).getFinishReason(); + } + + private boolean hasMessageChoice(ChatCompletionChunk chunk) { + return chunk.getChoices() != null + && !chunk.getChoices().isEmpty() + && chunk.getChoices().get(0).getMessage() != null; } /** @@ -304,54 +478,89 @@ private boolean isStopChunk(ChatCompletionChunk chunk) { */ private LlmResponse toLlmResponse(ChatCompletionResult arkResponse) { log.debug("Raw Ark response:{}", arkResponse); - LlmResponse response = null; // Check finish reason to determine response type String finishReason = arkResponse.getChoices().get(0).getFinishReason(); + List parts = new ArrayList<>(); + String text = (String) arkResponse.getChoices().get(0).getMessage().getContent(); + if (StringUtils.isNotEmpty(text)) { + parts.add(Part.fromText(text)); + } if ("tool_calls".equalsIgnoreCase(finishReason)) { - // Handle tool call response - List parts = new ArrayList<>(); - - // Add text content if any - String text = (String) arkResponse.getChoices().get(0).getMessage().getContent(); - if (StringUtils.isNotEmpty(text)) { - parts.add(Part.fromText(text)); - } - // Add tool call parts parts.addAll( parseToolCalls(arkResponse.getChoices().get(0).getMessage().getToolCalls())); - - response = - LlmResponse.builder() - .content(Content.builder().role("model").parts(parts).build()) - .build(); - } else { - // Handle regular text response - String text = (String) arkResponse.getChoices().get(0).getMessage().getContent(); - response = - LlmResponse.builder() - .content( - Content.builder() - .role("model") - .parts(Part.fromText(text)) - .build()) - .build(); } + LlmResponse response = buildLlmResponse(parts, finishReason, arkResponse.getUsage()); log.debug("LlmResponse:{}", response); return response; } + private LlmResponse buildLlmResponse(String text, String finishReason, Usage usage) { + List parts = new ArrayList<>(); + if (StringUtils.isNotEmpty(text)) { + parts.add(Part.fromText(text)); + } + return buildLlmResponse(parts, finishReason, usage); + } + + private LlmResponse buildLlmResponse(List parts, String finishReason, Usage usage) { + LlmResponse.Builder builder = + LlmResponse.builder().content(Content.builder().role("model").parts(parts).build()); + toFinishReason(finishReason).ifPresent(builder::finishReason); + toUsageMetadata(usage).ifPresent(builder::usageMetadata); + return builder.build(); + } + + private Optional toFinishReason(String finishReason) { + if (StringUtils.isBlank(finishReason)) { + return Optional.empty(); + } + FinishReason.Known known = + switch (finishReason.toLowerCase()) { + case "stop", "tool_calls", "function_call" -> FinishReason.Known.STOP; + case "length" -> FinishReason.Known.MAX_TOKENS; + case "content_filter" -> FinishReason.Known.SAFETY; + default -> FinishReason.Known.OTHER; + }; + return Optional.of(new FinishReason(known)); + } + + private Optional toUsageMetadata(Usage usage) { + if (usage == null) { + return Optional.empty(); + } + GenerateContentResponseUsageMetadata.Builder builder = + GenerateContentResponseUsageMetadata.builder() + .promptTokenCount(toIntTokenCount(usage.getPromptTokens())) + .candidatesTokenCount(toIntTokenCount(usage.getCompletionTokens())) + .totalTokenCount(toIntTokenCount(usage.getTotalTokens())); + if (usage.getPromptTokensDetails() != null + && usage.getPromptTokensDetails().getCachedTokens() != null) { + builder.cachedContentTokenCount(usage.getPromptTokensDetails().getCachedTokens()); + } + if (usage.getCompletionTokensDetails() != null + && usage.getCompletionTokensDetails().getReasoningTokens() != null) { + builder.thoughtsTokenCount(usage.getCompletionTokensDetails().getReasoningTokens()); + } + return Optional.of(builder.build()); + } + + private Integer toIntTokenCount(long value) { + return Math.toIntExact(Math.min(value, Integer.MAX_VALUE)); + } + /** * Convert ADK LlmRequest to Ark ChatCompletionRequest * @param llmRequest The ADK request * @return ChatCompletionRequest object for Ark API */ private ChatCompletionRequest toArkRequest(LlmRequest llmRequest) { - // Determine model name to use - String effectiveModelName = llmRequest.model().orElse(model()); + return toArkRequest(llmRequest, llmRequest.model().orElse(model())); + } + private ChatCompletionRequest toArkRequest(LlmRequest llmRequest, String effectiveModelName) { // Build chat messages from request List messages = buildChatMessages(llmRequest); @@ -367,6 +576,8 @@ private ChatCompletionRequest toArkRequest(LlmRequest llmRequest) { request.setThinking(thinking); } + llmRequest.config().ifPresent(config -> applyGenerateContentConfig(request, config)); + // Add tools if any if (llmRequest.tools() != null && !llmRequest.tools().isEmpty()) { List chatTools = buildChatTools(llmRequest); @@ -378,6 +589,73 @@ private ChatCompletionRequest toArkRequest(LlmRequest llmRequest) { return request; } + private void applyGenerateContentConfig( + ChatCompletionRequest request, GenerateContentConfig config) { + config.temperature().map(Float::doubleValue).ifPresent(request::setTemperature); + config.topP().map(Float::doubleValue).ifPresent(request::setTopP); + config.maxOutputTokens().ifPresent(request::setMaxTokens); + config.stopSequences().filter(stop -> !stop.isEmpty()).ifPresent(request::setStop); + config.presencePenalty().map(Float::doubleValue).ifPresent(request::setPresencePenalty); + config.frequencyPenalty().map(Float::doubleValue).ifPresent(request::setFrequencyPenalty); + config.candidateCount().ifPresent(request::setN); + config.responseLogprobs().ifPresent(request::setLogprobs); + config.logprobs() + .ifPresent( + topLogprobs -> { + request.setLogprobs(true); + request.setTopLogprobs(topLogprobs); + }); + config.responseSchema() + .ifPresentOrElse( + schema -> request.setResponseFormat(buildJsonSchemaResponseFormat(schema)), + () -> + config.responseJsonSchema() + .ifPresentOrElse( + schema -> + request.setResponseFormat( + buildJsonSchemaResponseFormat( + schema)), + () -> + applyJsonObjectResponseFormat( + request, config))); + } + + private ChatCompletionRequest.ChatCompletionRequestResponseFormat buildJsonSchemaResponseFormat( + Schema schema) { + Map schemaMap = + JSONUtil.convertValue(schema, new TypeReference>() {}); + updateTypeString(schemaMap); + ResponseFormatJSONSchemaJSONSchemaParam jsonSchema = + new ResponseFormatJSONSchemaJSONSchemaParam( + schema.title().filter(StringUtils::isNotBlank).orElse("response_schema"), + schema.description().orElse(null), + JSONUtil.valueToTree(schemaMap), + true); + return new ChatCompletionRequest.ChatCompletionRequestResponseFormat( + "json_schema", jsonSchema); + } + + private ChatCompletionRequest.ChatCompletionRequestResponseFormat buildJsonSchemaResponseFormat( + Object schema) { + ResponseFormatJSONSchemaJSONSchemaParam jsonSchema = + new ResponseFormatJSONSchemaJSONSchemaParam( + "response_schema", null, JSONUtil.valueToTree(schema), true); + return new ChatCompletionRequest.ChatCompletionRequestResponseFormat( + "json_schema", jsonSchema); + } + + private void applyJsonObjectResponseFormat( + ChatCompletionRequest request, GenerateContentConfig config) { + config.responseMimeType() + .filter(mimeType -> "application/json".equalsIgnoreCase(mimeType)) + .ifPresent( + unused -> + request.setResponseFormat( + new ChatCompletionRequest + .ChatCompletionRequestResponseFormat( + "json_object"))); + } + /** * Build chat messages from LlmRequest * @param llmRequest The ADK request @@ -415,12 +693,95 @@ private Stream buildSystemMessages(LlmRequest llmRequest) { */ private Stream buildContentMessages(LlmRequest llmRequest) { return llmRequest.contents().stream() - .map( - content -> - ChatMessage.builder() - .role(toArkRole(content.role().orElse("user"))) - .content(extractText(content)) - .build()); + .flatMap(content -> buildContentMessages(content).stream()); + } + + private List buildContentMessages(Content content) { + if (hasFunctionResponsePart(content)) { + return buildFunctionResponseMessages(content); + } + if (hasFunctionCallPart(content)) { + return List.of(buildFunctionCallMessage(content)); + } + return List.of(buildContentMessage(content)); + } + + private ChatMessage buildContentMessage(Content content) { + ChatMessage.Builder builder = + ChatMessage.builder().role(toArkRole(content.role().orElse("user"))); + if (hasSupportedMediaPart(content)) { + return builder.multiContent(extractContentParts(content)).build(); + } + return builder.content(extractText(content)).build(); + } + + private boolean hasFunctionResponsePart(Content content) { + return content.parts().stream() + .flatMap(List::stream) + .filter(Objects::nonNull) + .anyMatch(part -> part.functionResponse().isPresent()); + } + + private boolean hasFunctionCallPart(Content content) { + return content.parts().stream() + .flatMap(List::stream) + .filter(Objects::nonNull) + .anyMatch(part -> part.functionCall().isPresent()); + } + + private List buildFunctionResponseMessages(Content content) { + return content.parts().stream() + .flatMap(List::stream) + .filter(Objects::nonNull) + .flatMap(part -> part.functionResponse().stream()) + .map(this::buildFunctionResponseMessage) + .toList(); + } + + private ChatMessage buildFunctionResponseMessage(FunctionResponse functionResponse) { + ChatMessage.Builder builder = + ChatMessage.builder() + .role(ChatMessageRole.TOOL) + .content(functionResponse.response().map(JSONUtil::toJson).orElse("{}")); + functionResponse.id().filter(StringUtils::isNotBlank).ifPresent(builder::toolCallId); + functionResponse.name().filter(StringUtils::isNotBlank).ifPresent(builder::name); + return builder.build(); + } + + private ChatMessage buildFunctionCallMessage(Content content) { + ChatMessage.Builder builder = ChatMessage.builder().role(ChatMessageRole.ASSISTANT); + String text = extractText(content); + if (StringUtils.isNotEmpty(text)) { + builder.content(text); + } + List toolCalls = + content.parts().stream() + .flatMap(List::stream) + .filter(Objects::nonNull) + .flatMap(part -> part.functionCall().stream()) + .map(this::toChatToolCall) + .filter(Optional::isPresent) + .map(Optional::get) + .toList(); + if (!toolCalls.isEmpty()) { + builder.toolCalls(toolCalls); + } + return builder.build(); + } + + private Optional toChatToolCall(FunctionCall functionCall) { + Optional name = functionCall.name().filter(StringUtils::isNotBlank); + if (name.isEmpty()) { + return Optional.empty(); + } + ChatFunctionCall chatFunctionCall = new ChatFunctionCall(); + chatFunctionCall.setName(name.get()); + chatFunctionCall.setArguments(functionCall.args().map(JSONUtil::toJson).orElse("{}")); + ChatToolCall chatToolCall = new ChatToolCall(); + functionCall.id().filter(StringUtils::isNotBlank).ifPresent(chatToolCall::setId); + chatToolCall.setType("function"); + chatToolCall.setFunction(chatFunctionCall); + return Optional.of(chatToolCall); } /** @@ -442,27 +803,49 @@ private List buildChatTools(LlmRequest llmRequest) { * @return Optional ChatTool object */ private Optional convertToChatTool(BaseTool tool) { - // Get tool parameters schema - Optional parameters = tool.declaration().flatMap(FunctionDeclaration::parameters); - return parameters.map( - schema -> { - // Convert schema to map - Map schemaMap = - JSONUtil.convertValue( - schema, new TypeReference>() {}); - - // Normalize type strings in schema - updateTypeString(schemaMap); - - // Create chat function - ChatFunction chatFunction = new ChatFunction(); - chatFunction.setName(tool.name()); - chatFunction.setDescription(tool.description()); - chatFunction.setParameters(JSONUtil.valueToTree(schemaMap)); - - // Return chat tool - return new ChatTool("function", chatFunction); - }); + Optional declaration = tool.declaration(); + if (declaration.isEmpty()) { + return Optional.empty(); + } + Map schemaMap = + declaration + .get() + .parameters() + .map( + schema -> + JSONUtil.convertValue( + schema, + new TypeReference>() {})) + .orElseGet( + () -> + declaration + .get() + .parametersJsonSchema() + .map( + schema -> + JSONUtil.convertValue( + schema, + new TypeReference< + Map< + String, + Object>>() {})) + .orElseGet( + () -> { + Map emptySchema = + new LinkedHashMap<>(); + emptySchema.put("type", "object"); + emptySchema.put( + "properties", + new LinkedHashMap<>()); + return emptySchema; + })); + updateTypeString(schemaMap); + + ChatFunction chatFunction = new ChatFunction(); + chatFunction.setName(tool.name()); + chatFunction.setDescription(tool.description()); + chatFunction.setParameters(JSONUtil.valueToTree(schemaMap)); + return Optional.of(new ChatTool("function", chatFunction)); } /** @@ -499,11 +882,39 @@ private void updateTypeString(Map valueDict) { } } - /** - * Extract text content from Content object - * @param content The Content object - * @return Extracted text - */ + private boolean hasSupportedMediaPart(Content content) { + return content.parts().stream() + .flatMap(List::stream) + .filter(Objects::nonNull) + .anyMatch( + part -> + part.inlineData() + .flatMap(blob -> blob.mimeType()) + .filter(this::isSupportedMediaMimeType) + .isPresent() + || part.fileData() + .flatMap(fileData -> fileData.mimeType()) + .filter(this::isSupportedMediaMimeType) + .isPresent()); + } + + private List extractContentParts(Content content) { + List contentParts = new ArrayList<>(); + content.parts() + .ifPresent( + parts -> { + for (Part part : parts) { + if (part == null) { + continue; + } + appendTextContentPart(part, contentParts); + appendInlineDataContentPart(part, contentParts); + appendFileDataContentPart(part, contentParts); + } + }); + return contentParts; + } + private String extractText(Content content) { StringBuilder textBuilder = new StringBuilder(); // Use ifPresent with a lambda for a more functional and readable style @@ -523,6 +934,78 @@ private String extractText(Content content) { return textBuilder.toString(); } + private void appendTextContentPart(Part part, List contentParts) { + part.text() + .filter(StringUtils::isNotEmpty) + .ifPresent( + text -> + contentParts.add( + ChatCompletionContentPart.builder() + .type("text") + .text(text) + .build())); + } + + private void appendInlineDataContentPart( + Part part, List contentParts) { + part.inlineData() + .filter(blob -> blob.data().isPresent()) + .flatMap( + blob -> + toMediaContentPart( + dataUri( + blob.mimeType().orElse("application/octet-stream"), + blob.data().get()), + blob.mimeType().orElse("application/octet-stream"), + part.videoMetadata())) + .ifPresent(contentParts::add); + } + + private void appendFileDataContentPart( + Part part, List contentParts) { + part.fileData() + .filter(fileData -> fileData.fileUri().isPresent()) + .flatMap( + fileData -> + toMediaContentPart( + fileData.fileUri().get(), + fileData.mimeType().orElse("application/octet-stream"), + part.videoMetadata())) + .ifPresent(contentParts::add); + } + + private Optional toMediaContentPart( + String url, String mimeType, Optional videoMetadata) { + if (mimeType.startsWith("image/")) { + return Optional.of( + ChatCompletionContentPart.builder() + .type("image_url") + .imageUrl( + new ChatCompletionContentPart.ChatCompletionContentPartImageURL( + url, "auto")) + .build()); + } + if (mimeType.startsWith("video/")) { + return Optional.of( + ChatCompletionContentPart.builder() + .type("video_url") + .videoUrl( + new ChatCompletionContentPart.ChatCompletionContentPartVideoURL( + url, + videoMetadata.flatMap(VideoMetadata::fps).orElse(1.0))) + .build()); + } + return Optional.empty(); + } + + private boolean isSupportedMediaMimeType(String mimeType) { + return mimeType.startsWith("image/") || mimeType.startsWith("video/"); + } + + private String dataUri(String mimeType, byte[] data) { + return "data:" + mimeType + ";base64," + Base64.getEncoder().encodeToString(data); + } + /** * Append text part to StringBuilder * @param part The Part object @@ -571,11 +1054,18 @@ private ChatMessageRole toArkRole(String adkRole) { * @throws JsonProcessingException If JSON parsing fails */ private Part parseToolCallPart(ChatToolCall toolCall) throws JsonProcessingException { - return Part.fromFunctionCall( - toolCall.getFunction().getName(), - JSONUtil.fromJson( - toolCall.getFunction().getArguments(), - new TypeReference>() {})); + ChatFunctionCall function = toolCall.getFunction(); + Map args = + StringUtils.isBlank(function.getArguments()) + ? Map.of() + : JSONUtil.fromJson( + function.getArguments(), + new TypeReference>() {}); + FunctionCall.Builder builder = FunctionCall.builder().name(function.getName()).args(args); + if (StringUtils.isNotBlank(toolCall.getId())) { + builder.id(toolCall.getId()); + } + return Part.builder().functionCall(builder.build()).build(); } /** diff --git a/core/src/main/java/com/volcengine/veadk/processors/BaseRunProcessor.java b/core/src/main/java/com/volcengine/veadk/processors/BaseRunProcessor.java new file mode 100644 index 0000000..ed8c941 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/processors/BaseRunProcessor.java @@ -0,0 +1,27 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.processors; + +import com.google.adk.events.Event; +import io.reactivex.rxjava3.core.Flowable; +import java.util.function.Supplier; + +/** Intercepts an agent run and may observe, transform, retry, or replace its event stream. */ +@FunctionalInterface +public interface BaseRunProcessor { + + Flowable processRun(RunContext context, Supplier> eventGenerator); +} diff --git a/core/src/main/java/com/volcengine/veadk/processors/NoOpRunProcessor.java b/core/src/main/java/com/volcengine/veadk/processors/NoOpRunProcessor.java new file mode 100644 index 0000000..d398e5c --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/processors/NoOpRunProcessor.java @@ -0,0 +1,34 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.processors; + +import com.google.adk.events.Event; +import io.reactivex.rxjava3.core.Flowable; +import java.util.function.Supplier; + +/** Default run processor that preserves the original event stream. */ +public final class NoOpRunProcessor implements BaseRunProcessor { + + public static final NoOpRunProcessor INSTANCE = new NoOpRunProcessor(); + + private NoOpRunProcessor() {} + + @Override + public Flowable processRun( + RunContext context, Supplier> eventGenerator) { + return Flowable.defer(eventGenerator::get); + } +} diff --git a/core/src/main/java/com/volcengine/veadk/processors/RunContext.java b/core/src/main/java/com/volcengine/veadk/processors/RunContext.java new file mode 100644 index 0000000..fd26060 --- /dev/null +++ b/core/src/main/java/com/volcengine/veadk/processors/RunContext.java @@ -0,0 +1,57 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.processors; + +import com.google.adk.agents.RunConfig; +import com.google.adk.sessions.Session; +import com.google.genai.types.Content; +import com.volcengine.veadk.runner.Runner; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Objects; + +/** Immutable inputs made available to a {@link BaseRunProcessor}. */ +public record RunContext( + Runner runner, + Session session, + Content message, + RunConfig runConfig, + Map invocationState) { + + public RunContext { + Objects.requireNonNull(runner, "runner"); + Objects.requireNonNull(session, "session"); + Objects.requireNonNull(message, "message"); + Objects.requireNonNull(runConfig, "runConfig"); + invocationState = + invocationState == null + ? Map.of() + : Collections.unmodifiableMap(new LinkedHashMap<>(invocationState)); + } + + public String appName() { + return session.appName(); + } + + public String userId() { + return session.userId(); + } + + public String sessionId() { + return session.id(); + } +} diff --git a/core/src/main/java/com/volcengine/veadk/runner/Runner.java b/core/src/main/java/com/volcengine/veadk/runner/Runner.java index ca554ff..cec1de5 100644 --- a/core/src/main/java/com/volcengine/veadk/runner/Runner.java +++ b/core/src/main/java/com/volcengine/veadk/runner/Runner.java @@ -16,20 +16,36 @@ package com.volcengine.veadk.runner; import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.RunConfig; +import com.google.adk.artifacts.BaseArtifactService; import com.google.adk.artifacts.InMemoryArtifactService; +import com.google.adk.events.Event; import com.google.adk.memory.BaseMemoryService; import com.google.adk.memory.InMemoryMemoryService; +import com.google.adk.plugins.BasePlugin; +import com.google.adk.sessions.BaseSessionService; import com.google.adk.sessions.InMemorySessionService; -import com.google.common.collect.ImmutableList; +import com.google.adk.sessions.Session; +import com.google.genai.types.Content; +import com.volcengine.veadk.agent.Agent; +import com.volcengine.veadk.processors.BaseRunProcessor; +import com.volcengine.veadk.processors.NoOpRunProcessor; +import com.volcengine.veadk.processors.RunContext; +import io.reactivex.rxjava3.core.Flowable; +import java.util.List; +import java.util.Map; +import java.util.Objects; public class Runner extends com.google.adk.runner.Runner { + private final BaseRunProcessor runProcessor; + public Runner(BaseAgent agent) { this(agent, agent.name()); } public Runner(BaseAgent agent, String appName) { - this(agent, appName, null); + this(agent, appName, (BaseMemoryService) null); } public Runner(BaseAgent agent, BaseMemoryService baseMemoryService) { @@ -37,12 +53,104 @@ public Runner(BaseAgent agent, BaseMemoryService baseMemoryService) { } public Runner(BaseAgent agent, String appName, BaseMemoryService baseMemoryService) { - super( + this( agent, appName, new InMemoryArtifactService(), - new InMemorySessionService(), - null != baseMemoryService ? baseMemoryService : new InMemoryMemoryService(), - ImmutableList.of()); + defaultSessionService(agent), + null != baseMemoryService ? baseMemoryService : defaultMemoryService(agent), + defaultPlugins(agent), + defaultRunProcessor(agent)); + } + + public static Runner withProcessor(BaseAgent agent, BaseRunProcessor runProcessor) { + return withProcessor(agent, agent.name(), runProcessor); + } + + public static Runner withProcessor( + BaseAgent agent, String appName, BaseRunProcessor runProcessor) { + return new Runner( + agent, + appName, + new InMemoryArtifactService(), + defaultSessionService(agent), + defaultMemoryService(agent), + defaultPlugins(agent), + runProcessor); + } + + public Runner( + BaseAgent agent, + String appName, + BaseArtifactService artifactService, + BaseSessionService sessionService, + BaseMemoryService memoryService, + List plugins) { + this( + agent, + appName, + artifactService, + sessionService, + memoryService, + plugins, + NoOpRunProcessor.INSTANCE); + } + + public Runner( + BaseAgent agent, + String appName, + BaseArtifactService artifactService, + BaseSessionService sessionService, + BaseMemoryService memoryService, + List plugins, + BaseRunProcessor runProcessor) { + super( + agent, + appName, + Objects.requireNonNull(artifactService, "artifactService"), + Objects.requireNonNull(sessionService, "sessionService"), + Objects.requireNonNull(memoryService, "memoryService"), + List.copyOf(Objects.requireNonNull(plugins, "plugins"))); + this.runProcessor = Objects.requireNonNull(runProcessor, "runProcessor"); + } + + public BaseRunProcessor runProcessor() { + return runProcessor; + } + + private static BaseSessionService defaultSessionService(BaseAgent agent) { + return agent instanceof Agent veadkAgent && veadkAgent.shortTermMemory() != null + ? veadkAgent.shortTermMemory() + : new InMemorySessionService(); + } + + private static BaseMemoryService defaultMemoryService(BaseAgent agent) { + return agent instanceof Agent veadkAgent && veadkAgent.longTermMemory() != null + ? veadkAgent.longTermMemory() + : new InMemoryMemoryService(); + } + + private static List defaultPlugins(BaseAgent agent) { + return agent instanceof Agent veadkAgent ? veadkAgent.plugins() : List.of(); + } + + private static BaseRunProcessor defaultRunProcessor(BaseAgent agent) { + return agent instanceof Agent veadkAgent + ? veadkAgent.runProcessor() + : NoOpRunProcessor.INSTANCE; + } + + @Override + public Flowable runAsync( + Session session, + Content newMessage, + RunConfig runConfig, + Map stateDelta) { + RunContext context = new RunContext(this, session, newMessage, runConfig, stateDelta); + return Flowable.defer( + () -> + runProcessor.processRun( + context, + () -> super.runAsync(session, newMessage, runConfig, stateDelta))); } } diff --git a/core/src/test/java/com/volcengine/veadk/agent/AgentMetadataTest.java b/core/src/test/java/com/volcengine/veadk/agent/AgentMetadataTest.java new file mode 100644 index 0000000..276da17 --- /dev/null +++ b/core/src/test/java/com/volcengine/veadk/agent/AgentMetadataTest.java @@ -0,0 +1,146 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.agent; + +import static org.assertj.core.api.Assertions.assertThat; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.LlmAgent; +import com.google.adk.models.BaseLlm; +import com.google.adk.models.BaseLlmConnection; +import com.google.adk.models.LlmRequest; +import com.google.adk.models.LlmResponse; +import com.google.adk.tools.BaseTool; +import com.volcengine.veadk.agent.AgentMetadata.SkillSummary; +import io.reactivex.rxjava3.core.Flowable; +import java.util.List; +import java.util.Set; +import org.junit.jupiter.api.Test; + +class AgentMetadataTest { + + @Test + void from_extractsToolsSearchSourcesAndTopology() { + BaseAgent child = new EmptyAgent("child"); + LlmAgent root = + LlmAgent.builder() + .name("root") + .description("root agent") + .model(new EmptyLlm()) + .tools(new NamedTool("web_search"), new NamedTool("calculator")) + .subAgents(child) + .build(); + + AgentMetadata metadata = AgentMetadata.from(root); + + assertThat(metadata.name()).isEqualTo("root"); + assertThat(metadata.model()).isEqualTo("test-model"); + assertThat(metadata.tools()) + .extracting(AgentMetadata.ToolSummary::name) + .containsExactly("web_search", "calculator"); + assertThat(metadata.searchSources()).containsExactly("web"); + assertThat(metadata.subAgents()).extracting(AgentMetadata::name).containsExactly("child"); + } + + @Test + void from_extractsAndDeduplicatesProviderComponentsAndSkills() { + AgentMetadata metadata = AgentMetadata.from(new ComponentAgent()); + + assertThat(metadata.searchSources()).containsExactly("custom", "knowledge", "memory"); + assertThat(metadata.components()).hasSize(2); + assertThat(metadata.skills()).containsExactly(new SkillSummary("review", "Review code")); + } + + private static final class ComponentAgent extends BaseAgent implements AgentComponentProvider { + + private ComponentAgent() { + super("component-agent", "", List.of(), List.of(), List.of()); + } + + @Override + public List agentComponents() { + return List.of( + new AgentComponent("knowledgebase", "docs", "knowledgebase", "memory", ""), + new AgentComponent("knowledgebase", "docs", "knowledgebase", "memory", ""), + new AgentComponent("long_term_memory", "user-memory")); + } + + @Override + public List agentSkills() { + return List.of( + new SkillSummary("review", "Review code"), + new SkillSummary("review", "Duplicate")); + } + + @Override + public Set additionalSearchSources() { + return Set.of("custom"); + } + + @Override + protected Flowable runAsyncImpl( + InvocationContext invocationContext) { + return Flowable.empty(); + } + + @Override + protected Flowable runLiveImpl( + InvocationContext invocationContext) { + return Flowable.empty(); + } + } + + private static final class NamedTool extends BaseTool { + private NamedTool(String name) { + super(name, name); + } + } + + private static final class EmptyLlm extends BaseLlm { + private EmptyLlm() { + super("test-model"); + } + + @Override + public Flowable generateContent(LlmRequest llmRequest, boolean stream) { + return Flowable.empty(); + } + + @Override + public BaseLlmConnection connect(LlmRequest llmRequest) { + return null; + } + } + + private static final class EmptyAgent extends BaseAgent { + private EmptyAgent(String name) { + super(name, "", List.of(), List.of(), List.of()); + } + + @Override + protected Flowable runAsyncImpl( + InvocationContext invocationContext) { + return Flowable.empty(); + } + + @Override + protected Flowable runLiveImpl( + InvocationContext invocationContext) { + return Flowable.empty(); + } + } +} diff --git a/core/src/test/java/com/volcengine/veadk/agent/AgentTest.java b/core/src/test/java/com/volcengine/veadk/agent/AgentTest.java new file mode 100644 index 0000000..aa08c3d --- /dev/null +++ b/core/src/test/java/com/volcengine/veadk/agent/AgentTest.java @@ -0,0 +1,140 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.agent; + +import static org.assertj.core.api.Assertions.assertThat; + +import com.google.adk.agents.Callbacks; +import com.google.adk.agents.LlmAgent; +import com.google.adk.memory.InMemoryMemoryService; +import com.google.adk.models.BaseLlm; +import com.google.adk.models.BaseLlmConnection; +import com.google.adk.models.LlmRequest; +import com.google.adk.models.LlmResponse; +import com.google.adk.sessions.InMemorySessionService; +import com.volcengine.veadk.config.VeADKConfig; +import com.volcengine.veadk.processors.BaseRunProcessor; +import com.volcengine.veadk.runner.Runner; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Maybe; +import java.time.Duration; +import java.util.Map; +import org.junit.jupiter.api.Test; + +class AgentTest { + + @Test + void builderCreatesLlmAgentAndPreservesVeadkComponents() { + InMemorySessionService sessions = new InMemorySessionService(); + InMemoryMemoryService memory = new InMemoryMemoryService(); + BaseRunProcessor processor = (context, eventGenerator) -> eventGenerator.get(); + + Agent agent = + Agent.builder() + .id("agent-id") + .name("assistant") + .description("test assistant") + .instruction("help the user") + .model(new EmptyLlm("test-model")) + .shortTermMemory(sessions) + .longTermMemory(memory) + .runProcessor(processor) + .skills("review", "testing") + .enableAuthz(true) + .autoSaveSession(true) + .build(); + + assertThat(agent).isInstanceOf(LlmAgent.class); + assertThat(agent.id()).isEqualTo("agent-id"); + assertThat(agent.name()).isEqualTo("assistant"); + assertThat(agent.shortTermMemory()).isSameAs(sessions); + assertThat(agent.longTermMemory()).isSameAs(memory); + assertThat(agent.runProcessor()).isSameAs(processor); + assertThat(agent.skills()).containsExactly("review", "testing"); + assertThat(agent.enableAuthz()).isTrue(); + assertThat(agent.autoSaveSession()).isTrue(); + assertThat(agent.metadata().model()).isEqualTo("test-model"); + assertThat(agent.metadata().searchSources()).containsExactly("memory"); + + Runner runner = agent.newRunner(); + assertThat(runner.agent()).isSameAs(agent); + assertThat(runner.sessionService()).isSameAs(sessions); + assertThat(runner.memoryService()).isSameAs(memory); + assertThat(runner.runProcessor()).isSameAs(processor); + + Runner directRunner = new Runner(agent); + assertThat(directRunner.sessionService()).isSameAs(sessions); + assertThat(directRunner.memoryService()).isSameAs(memory); + assertThat(directRunner.runProcessor()).isSameAs(processor); + } + + @Test + void configuredBuilderCreatesArkModelWithoutReadingGlobalEnvironment() { + VeADKConfig config = + VeADKConfig.from( + Map.of( + "MODEL_AGENT_NAME", "configured-model", + "MODEL_AGENT_API_BASE", "https://ark.example/api/v3", + "MODEL_AGENT_API_KEY", "configured-key")); + + Agent agent = Agent.builder(config).name("configured-agent").build(); + + assertThat(agent.name()).isEqualTo("configured-agent"); + assertThat(agent.metadata().model()).isEqualTo("configured-model"); + } + + @Test + void autoSaveSessionAddsCallbackWithoutReplacingExistingCallbacks() { + Callbacks.AfterAgentCallback existing = context -> Maybe.empty(); + SaveSessionPolicy policy = new SaveSessionPolicy(3, Duration.ofSeconds(5), true, false); + + Agent agent = + Agent.builder() + .name("assistant") + .model(new EmptyLlm("test-model")) + .longTermMemory(new InMemoryMemoryService()) + .afterAgentCallback(existing) + .autoSaveSession(true) + .saveSessionPolicy(policy) + .build(); + + assertThat(agent.saveSessionPolicy()).isEqualTo(policy); + assertThat(agent.afterAgentCallback()) + .hasValueSatisfying( + callbacks -> { + assertThat(callbacks).hasSize(2); + assertThat(callbacks.get(0)).isSameAs(existing); + assertThat(callbacks.get(1)) + .isInstanceOf(SaveSessionToMemoryCallback.class); + }); + } + + private static final class EmptyLlm extends BaseLlm { + private EmptyLlm(String model) { + super(model); + } + + @Override + public Flowable generateContent(LlmRequest llmRequest, boolean stream) { + return Flowable.empty(); + } + + @Override + public BaseLlmConnection connect(LlmRequest llmRequest) { + return null; + } + } +} diff --git a/core/src/test/java/com/volcengine/veadk/agent/SaveSessionToMemoryCallbackTest.java b/core/src/test/java/com/volcengine/veadk/agent/SaveSessionToMemoryCallbackTest.java new file mode 100644 index 0000000..7f353b7 --- /dev/null +++ b/core/src/test/java/com/volcengine/veadk/agent/SaveSessionToMemoryCallbackTest.java @@ -0,0 +1,214 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.agent; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import com.google.adk.agents.CallbackContext; +import com.google.adk.agents.InvocationContext; +import com.google.adk.events.Event; +import com.google.adk.memory.BaseMemoryService; +import com.google.adk.sessions.BaseSessionService; +import com.google.adk.sessions.Session; +import com.volcengine.veadk.utils.ReadonlyContextAccessorUtil; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.subjects.CompletableSubject; +import java.time.Duration; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicLong; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; +import org.mockito.Mockito; + +class SaveSessionToMemoryCallbackTest { + + @Test + void legacyConstructorReturnsImmediatelyAndSavesEveryInvocation() { + Fixture fixture = new Fixture("session"); + CompletableSubject pendingSave = CompletableSubject.create(); + when(fixture.memory.addSessionToMemory(fixture.session)).thenReturn(pendingSave); + + try (MockedStatic accessor = + Mockito.mockStatic(ReadonlyContextAccessorUtil.class)) { + accessor.when(() -> ReadonlyContextAccessorUtil.getInvocationContext(fixture.callback)) + .thenReturn(fixture.invocation); + + SaveSessionToMemoryCallback callback = new SaveSessionToMemoryCallback(); + callback.call(fixture.callback).test().assertComplete(); + callback.call(fixture.callback).test().assertComplete(); + + verify(fixture.memory, times(2)).addSessionToMemory(fixture.session); + } + } + + @Test + void policySkipsFrequentSmallUpdatesAndSavesAtEitherThreshold() { + AtomicLong now = new AtomicLong(1_000L); + SaveSessionToMemoryCallback callback = + new SaveSessionToMemoryCallback( + new SaveSessionPolicy(10, Duration.ofSeconds(60), true, true), now::get); + Fixture fixture = new Fixture("session"); + List events = new ArrayList<>(); + events.add(mock(Event.class)); + when(fixture.session.events()).thenReturn(events); + + try (MockedStatic accessor = + Mockito.mockStatic(ReadonlyContextAccessorUtil.class)) { + accessor.when(() -> ReadonlyContextAccessorUtil.getInvocationContext(fixture.callback)) + .thenReturn(fixture.invocation); + + callback.call(fixture.callback).test().assertComplete(); + callback.call(fixture.callback).test().assertComplete(); + verify(fixture.memory, times(1)).addSessionToMemory(fixture.session); + + for (int index = 0; index < 10; index++) { + events.add(mock(Event.class)); + } + callback.call(fixture.callback).test().assertComplete(); + verify(fixture.memory, times(2)).addSessionToMemory(fixture.session); + + now.addAndGet(Duration.ofSeconds(60).toMillis()); + callback.call(fixture.callback).test().assertComplete(); + verify(fixture.memory, times(3)).addSessionToMemory(fixture.session); + } + } + + @Test + void sessionSwitchForcesPreviousSessionSaveBeforeCurrentSession() { + SaveSessionToMemoryCallback callback = + new SaveSessionToMemoryCallback( + new SaveSessionPolicy(10, Duration.ofHours(1), true, true), () -> 1_000L); + Fixture first = new Fixture("first"); + Fixture second = new Fixture("second", first.sessions, first.memory); + when(first.sessions.getSession("app", "user", "second", Optional.empty())) + .thenReturn(Maybe.just(second.session)); + + try (MockedStatic accessor = + Mockito.mockStatic(ReadonlyContextAccessorUtil.class)) { + accessor.when(() -> ReadonlyContextAccessorUtil.getInvocationContext(first.callback)) + .thenReturn(first.invocation); + accessor.when(() -> ReadonlyContextAccessorUtil.getInvocationContext(second.callback)) + .thenReturn(second.invocation); + + callback.call(first.callback).test().assertComplete(); + callback.call(second.callback).test().assertComplete(); + + verify(first.memory, times(2)).addSessionToMemory(first.session); + verify(first.memory).addSessionToMemory(second.session); + } + } + + @Test + void errorHandlingCanPreserveLegacySuppressionOrPropagate() { + Fixture fixture = new Fixture("session"); + when(fixture.memory.addSessionToMemory(fixture.session)) + .thenReturn(Completable.error(new IllegalStateException("backend unavailable"))); + + try (MockedStatic accessor = + Mockito.mockStatic(ReadonlyContextAccessorUtil.class)) { + accessor.when(() -> ReadonlyContextAccessorUtil.getInvocationContext(fixture.callback)) + .thenReturn(fixture.invocation); + + new SaveSessionToMemoryCallback(new SaveSessionPolicy(0, Duration.ZERO, true, true)) + .call(fixture.callback) + .test() + .assertComplete(); + new SaveSessionToMemoryCallback(new SaveSessionPolicy(0, Duration.ZERO, true, false)) + .call(fixture.callback) + .test() + .assertError(IllegalStateException.class); + } + } + + @Test + void concurrentSavesForSameUserAreSerialized() { + Fixture fixture = new Fixture("session"); + CompletableSubject firstSave = CompletableSubject.create(); + when(fixture.memory.addSessionToMemory(fixture.session)).thenReturn(firstSave); + SaveSessionToMemoryCallback callback = + new SaveSessionToMemoryCallback( + new SaveSessionPolicy(10, Duration.ofMinutes(1), true, true)); + + try (MockedStatic accessor = + Mockito.mockStatic(ReadonlyContextAccessorUtil.class)) { + accessor.when(() -> ReadonlyContextAccessorUtil.getInvocationContext(fixture.callback)) + .thenReturn(fixture.invocation); + + var first = callback.call(fixture.callback).test(); + var second = callback.call(fixture.callback).test(); + verify(fixture.memory, times(1)).addSessionToMemory(fixture.session); + + firstSave.onComplete(); + + first.assertComplete(); + second.assertComplete(); + verify(fixture.memory, times(1)).addSessionToMemory(fixture.session); + } + } + + @Test + void configuredMemoryOverridesRunnerContextMemory() { + Fixture fixture = new Fixture("session"); + BaseMemoryService configuredMemory = mock(BaseMemoryService.class); + when(configuredMemory.addSessionToMemory(fixture.session)) + .thenReturn(Completable.complete()); + SaveSessionToMemoryCallback callback = + new SaveSessionToMemoryCallback( + new SaveSessionPolicy(0, Duration.ZERO, true, true), configuredMemory); + + try (MockedStatic accessor = + Mockito.mockStatic(ReadonlyContextAccessorUtil.class)) { + accessor.when(() -> ReadonlyContextAccessorUtil.getInvocationContext(fixture.callback)) + .thenReturn(fixture.invocation); + + callback.call(fixture.callback).test().assertComplete(); + + verify(configuredMemory).addSessionToMemory(fixture.session); + verify(fixture.memory, times(0)).addSessionToMemory(fixture.session); + } + } + + private static final class Fixture { + private final CallbackContext callback = mock(CallbackContext.class); + private final InvocationContext invocation = mock(InvocationContext.class); + private final BaseSessionService sessions; + private final BaseMemoryService memory; + private final Session session = mock(Session.class); + + private Fixture(String sessionId) { + this(sessionId, mock(BaseSessionService.class), mock(BaseMemoryService.class)); + } + + private Fixture(String sessionId, BaseSessionService sessions, BaseMemoryService memory) { + this.sessions = sessions; + this.memory = memory; + when(session.id()).thenReturn(sessionId); + when(session.events()).thenReturn(new ArrayList<>()); + when(invocation.appName()).thenReturn("app"); + when(invocation.userId()).thenReturn("user"); + when(invocation.session()).thenReturn(session); + when(invocation.sessionService()).thenReturn(sessions); + when(invocation.memoryService()).thenReturn(memory); + when(sessions.getSession("app", "user", sessionId, Optional.empty())) + .thenReturn(Maybe.just(session)); + when(memory.addSessionToMemory(session)).thenReturn(Completable.complete()); + } + } +} diff --git a/core/src/test/java/com/volcengine/veadk/compat/PublicApiCompatibilityTest.java b/core/src/test/java/com/volcengine/veadk/compat/PublicApiCompatibilityTest.java new file mode 100644 index 0000000..dfb8c41 --- /dev/null +++ b/core/src/test/java/com/volcengine/veadk/compat/PublicApiCompatibilityTest.java @@ -0,0 +1,213 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.compat; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.mockito.Mockito.mock; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.memory.BaseMemoryService; +import com.google.adk.memory.InMemoryMemoryService; +import com.google.adk.models.BaseLlmConnection; +import com.google.adk.models.LlmRequest; +import com.google.adk.models.LlmResponse; +import com.google.adk.sessions.BaseSessionService; +import com.google.adk.tools.ToolContext; +import com.volcengine.veadk.agent.SaveSessionPolicy; +import com.volcengine.veadk.agent.SaveSessionToMemoryCallback; +import com.volcengine.veadk.knowledgebase.BaseKnowledgebaseService; +import com.volcengine.veadk.memory.LongTermMemory; +import com.volcengine.veadk.memory.LongTermMemoryBackend; +import com.volcengine.veadk.memory.ShortTermMemory; +import com.volcengine.veadk.memory.ShortTermMemoryBackend; +import com.volcengine.veadk.memory.ShortTermMemoryProcessor; +import com.volcengine.veadk.model.ArkLlm; +import com.volcengine.veadk.runner.Runner; +import com.volcengine.veadk.tools.knowledgebase.LoadKnowledgebaseTool; +import com.volcengine.veadk.tools.sandbox.CodeSandboxToolset; +import com.volcengine.veadk.trace.OpenTelemetry; +import com.volcengine.veadk.utils.EnvUtil; +import io.reactivex.rxjava3.core.Flowable; +import io.reactivex.rxjava3.core.Single; +import java.lang.reflect.Modifier; +import java.util.List; +import org.junit.jupiter.api.Test; + +class PublicApiCompatibilityTest { + + private static final List BASELINE_PUBLIC_TYPES = + List.of( + "com.volcengine.veadk.Version", + "com.volcengine.veadk.agent.SaveSessionToMemoryCallback", + "com.volcengine.veadk.integration.agentkit.AgentKitWrapper", + "com.volcengine.veadk.integration.vikingknowledgebase.KnowledgebaseEntry", + "com.volcengine.veadk.integration.vikingknowledgebase.VikingKnowledgebaseWrapper", + "com.volcengine.veadk.integration.vikingmemory.Message", + "com.volcengine.veadk.integration.vikingmemory.Metadata", + "com.volcengine.veadk.integration.vikingmemory.VikingMemoryWrapper", + "com.volcengine.veadk.integration.websearch.WebSearchWrapper", + "com.volcengine.veadk.knowledgebase.BaseKnowledgebaseService", + "com.volcengine.veadk.knowledgebase.SearchKnowledgebaseResponse", + "com.volcengine.veadk.knowledgebase.viking.VikingKnowledgebaseService", + "com.volcengine.veadk.memory.viking.VikingMemoryService", + "com.volcengine.veadk.model.ArkLlm", + "com.volcengine.veadk.runner.Runner", + "com.volcengine.veadk.tools.knowledgebase.LoadKnowledgebaseResponse", + "com.volcengine.veadk.tools.knowledgebase.LoadKnowledgebaseTool", + "com.volcengine.veadk.tools.sandbox.CodeSandboxToolset", + "com.volcengine.veadk.tools.sandbox.RunCodeTool", + "com.volcengine.veadk.tools.websearch.WebSearchTool", + "com.volcengine.veadk.trace.OpenTelemetry", + "com.volcengine.veadk.trace.exporter.AttributeRewritingSpanExporter", + "com.volcengine.veadk.trace.exporter.ExporterFactory", + "com.volcengine.veadk.trace.exporter.TLSExporter", + "com.volcengine.veadk.utils.EnvUtil", + "com.volcengine.veadk.utils.JSONUtil", + "com.volcengine.veadk.utils.ReadonlyContextAccessorUtil"); + + @Test + void baselinePublicTypesRemainLoadableAndPublic() throws ClassNotFoundException { + for (String typeName : BASELINE_PUBLIC_TYPES) { + Class type = Class.forName(typeName); + assertThat(Modifier.isPublic(type.getModifiers())).as(typeName).isTrue(); + } + } + + @Test + void runnerInheritanceAndConstructorsRemainCompatible() throws NoSuchMethodException { + assertThat(com.google.adk.runner.Runner.class).isAssignableFrom(Runner.class); + assertThat(Runner.class.getConstructor(BaseAgent.class)).isNotNull(); + assertThat(Runner.class.getConstructor(BaseAgent.class, String.class)).isNotNull(); + assertThat(Runner.class.getConstructor(BaseAgent.class, BaseMemoryService.class)) + .isNotNull(); + assertThat( + Runner.class.getConstructor( + BaseAgent.class, String.class, BaseMemoryService.class)) + .isNotNull(); + } + + @Test + void runnerLegacyThreeArgumentNullCallRemainsSourceCompatible() { + Runner runner = new Runner(mock(BaseAgent.class), "app", null); + + assertThat(runner.memoryService()).isInstanceOf(InMemoryMemoryService.class); + } + + @Test + void arkLlmConstructorsAndGenerationMethodsRemainCompatible() throws NoSuchMethodException { + assertThat(ArkLlm.class.getConstructor(String.class)).isNotNull(); + assertThat(ArkLlm.class.getConstructor(String.class, String.class)).isNotNull(); + assertThat(ArkLlm.class.getConstructor(List.class)).isNotNull(); + assertThat(ArkLlm.class.getConstructor(List.class, String.class)).isNotNull(); + assertThat(ArkLlm.class.getConstructor(List.class, String.class, String.class)).isNotNull(); + assertThat( + ArkLlm.class.getConstructor( + List.class, String.class, String.class, String.class)) + .isNotNull(); + assertThat(ArkLlm.class.getMethod("fallbacks").getReturnType()).isEqualTo(List.class); + assertThat( + ArkLlm.class + .getMethod("generateContent", LlmRequest.class, boolean.class) + .getReturnType()) + .isEqualTo(Flowable.class); + assertThat(ArkLlm.class.getMethod("connect", LlmRequest.class).getReturnType()) + .isEqualTo(BaseLlmConnection.class); + assertThat(ArkLlm.class.getSuperclass().getName()) + .isEqualTo("com.google.adk.models.BaseLlm"); + assertThat(LlmResponse.class).isNotNull(); + } + + @Test + void knowledgebaseAndToolContractsRemainCompatible() throws NoSuchMethodException { + assertThat( + BaseKnowledgebaseService.class + .getMethod("searchKnowledgebase", String.class) + .getReturnType()) + .isEqualTo(Single.class); + assertThat(LoadKnowledgebaseTool.class.getConstructor(BaseKnowledgebaseService.class)) + .isNotNull(); + assertThat( + LoadKnowledgebaseTool.class + .getMethod("loadKnowledgebase", String.class, ToolContext.class) + .getReturnType()) + .isEqualTo(Single.class); + assertThat(CodeSandboxToolset.class.getMethod("create")).isNotNull(); + assertThat(CodeSandboxToolset.class.getMethod("create", String.class)).isNotNull(); + } + + @Test + void newLongTermMemoryFacadeKeepsAdkMemoryContract() throws NoSuchMethodException { + assertThat(BaseMemoryService.class).isAssignableFrom(LongTermMemory.class); + assertThat(LongTermMemory.class.getConstructor()).isNotNull(); + assertThat(LongTermMemory.class.getConstructor(String.class)).isNotNull(); + assertThat(LongTermMemory.class.getConstructor(LongTermMemoryBackend.class)).isNotNull(); + assertThat( + LongTermMemory.class + .getMethod("searchMemory", String.class, String.class, String.class) + .getReturnType()) + .isEqualTo(Single.class); + } + + @Test + void newShortTermMemoryFacadeKeepsAdkSessionContract() throws NoSuchMethodException { + assertThat(com.google.adk.sessions.BaseSessionService.class) + .isAssignableFrom(ShortTermMemory.class); + assertThat(ShortTermMemory.class.getConstructor()).isNotNull(); + assertThat(ShortTermMemory.class.getConstructor(BaseSessionService.class)).isNotNull(); + assertThat(ShortTermMemory.class.getConstructor(ShortTermMemoryBackend.class)).isNotNull(); + assertThat( + ShortTermMemory.class.getConstructor( + ShortTermMemoryBackend.class, ShortTermMemoryProcessor.class)) + .isNotNull(); + assertThat( + ShortTermMemory.class + .getMethod( + "createSession", String.class, String.class, String.class) + .getReturnType()) + .isEqualTo(io.reactivex.rxjava3.core.Single.class); + } + + @Test + void saveSessionCallbackKeepsLegacyConstructorAndAddsPolicyConfiguration() + throws NoSuchMethodException { + assertThat(SaveSessionToMemoryCallback.class.getConstructor()).isNotNull(); + assertThat(SaveSessionToMemoryCallback.class.getConstructor(SaveSessionPolicy.class)) + .isNotNull(); + assertThat(SaveSessionToMemoryCallback.class.getMethod("policy")).isNotNull(); + } + + @Test + void environmentAndTelemetryEntryPointsRemainCompatible() throws NoSuchMethodException { + List environmentMethods = + List.of( + "getAgentKitToolId", + "getAgentKitService", + "getAgentKitRegion", + "getAgentKitHost", + "getAgentApiKey", + "getAccessKey", + "getSecretKey", + "getTLSEndpoint", + "getTLSServiceName", + "getTLSRegion", + "getVikingMmemoryType", + "getCodeSandboxUrl"); + for (String methodName : environmentMethods) { + assertThat(EnvUtil.class.getMethod(methodName).getReturnType()) + .as(methodName) + .isEqualTo(String.class); + } + assertThat(OpenTelemetry.class.getMethod("initOpenTelemetry", List.class)).isNotNull(); + } +} diff --git a/core/src/test/java/com/volcengine/veadk/config/VeADKConfigTest.java b/core/src/test/java/com/volcengine/veadk/config/VeADKConfigTest.java new file mode 100644 index 0000000..a8b5f22 --- /dev/null +++ b/core/src/test/java/com/volcengine/veadk/config/VeADKConfigTest.java @@ -0,0 +1,112 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.config; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class VeADKConfigTest { + + @TempDir Path directory; + + @Test + void load_resolvesYamlDotenvAndEnvironmentWithExpectedPrecedence() throws Exception { + Files.writeString( + directory.resolve("config.yaml"), + """ + model: + agent: + name: yaml-model + api_base: https://yaml.example/v3 + api_key: yaml-secret + embedding: + dim: 1024 + observability: + opentelemetry: + trace_content: false + """); + Files.writeString( + directory.resolve(".env"), + "MODEL_AGENT_NAME=dotenv-model\nMODEL_AGENT_API_KEY=dotenv-secret\n"); + + VeADKConfig config = + VeADKConfig.load( + directory, + Map.of( + "MODEL_AGENT_NAME", "environment-model", + "MODEL_AGENT_API_KEY", "environment-secret")); + + assertThat(config.model().name()).isEqualTo("environment-model"); + assertThat(config.model().apiBase()).isEqualTo("https://yaml.example/v3"); + assertThat(config.model().apiKey()).isEqualTo("environment-secret"); + assertThat(config.embedding().dimension()).isEqualTo(1024); + assertThat(config.openTelemetry().traceContent()).isFalse(); + } + + @Test + void load_findsParentYamlAndSupportsByteplusCredentialAliases() throws Exception { + Files.writeString( + directory.resolve("config.yaml"), + "cloud_provider: byteplus\n" + + "byteplus_access_key: bp-ak\n" + + "byteplus_secret_key: bp-sk\n"); + Path child = Files.createDirectories(directory.resolve("nested/agent")); + + VeADKConfig config = VeADKConfig.load(child, Map.of()); + + assertThat(config.require("VOLCENGINE_ACCESS_KEY")).isEqualTo("bp-ak"); + assertThat(config.require("VOLCENGINE_SECRET_KEY")).isEqualTo("bp-sk"); + assertThat(config.model().name()).isEqualTo("seed-2-0-lite-260228"); + assertThat(config.embedding().name()).isEqualTo("skylark-embedding-vision-250615"); + } + + @Test + void diagnosticsRedactSecretsWithoutHidingNormalConfiguration() { + VeADKConfig config = + VeADKConfig.from( + Map.of( + "MODEL_AGENT_NAME", "test-model", + "MODEL_AGENT_API_KEY", "super-secret", + "DATABASE_REDIS_PASSWORD", "redis-secret")); + + assertThat(config.toString()).contains("test-model", ""); + assertThat(config.toString()).doesNotContain("super-secret", "redis-secret"); + assertThat(config.model().toString()).doesNotContain("super-secret"); + } + + @Test + void typedAccessorsRejectInvalidValuesAndMissingRequiredValues() { + VeADKConfig config = + VeADKConfig.from( + Map.of( + "FEATURE_ENABLED", "sometimes", + "RETRY_COUNT", "many")); + + assertThatThrownBy(() -> config.getBoolean("FEATURE_ENABLED", false)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> config.getInt("RETRY_COUNT", 1)) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> config.require("MISSING_VALUE")) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("MISSING_VALUE"); + } +} diff --git a/core/src/test/java/com/volcengine/veadk/knowledgebase/KnowledgeBaseTest.java b/core/src/test/java/com/volcengine/veadk/knowledgebase/KnowledgeBaseTest.java new file mode 100644 index 0000000..e4c868c --- /dev/null +++ b/core/src/test/java/com/volcengine/veadk/knowledgebase/KnowledgeBaseTest.java @@ -0,0 +1,101 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.knowledgebase; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import com.volcengine.veadk.integration.vikingknowledgebase.KnowledgebaseEntry; +import io.reactivex.rxjava3.core.Single; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class KnowledgeBaseTest { + + @TempDir Path directory; + + @Test + void inMemoryFacadeAddsAndRanksTextDeterministically() { + KnowledgeBase knowledgeBase = new KnowledgeBase("docs"); + knowledgeBase.addFromText( + List.of( + "Java reactive streams and Flowable", + "Python data science", + "Java build compatibility")); + + List results = + knowledgeBase.search("Java reactive", 2); + + assertThat(results) + .extracting(com.volcengine.veadk.knowledgebase.KnowledgebaseEntry::content) + .containsExactly("Java reactive streams and Flowable", "Java build compatibility"); + assertThat(knowledgeBase.index()).isEqualTo("docs"); + } + + @Test + void facadeLoadsFilesAndKeepsSourceMetadata() throws Exception { + Path nested = Files.createDirectories(directory.resolve("nested")); + Files.writeString(directory.resolve("one.txt"), "first document"); + Files.writeString(nested.resolve("two.txt"), "second document"); + KnowledgeBase knowledgeBase = new KnowledgeBase("files"); + + assertThat(knowledgeBase.addFromDirectory(directory)).isTrue(); + + List results = + knowledgeBase.search("second", 1); + assertThat(results).hasSize(1); + assertThat(results.get(0).metadata().get("file_path").toString()).endsWith("two.txt"); + } + + @Test + void facadePreservesLegacySearchServiceContract() { + KnowledgeBase knowledgeBase = new KnowledgeBase("legacy"); + knowledgeBase.addFromText("legacy compatible content"); + + SearchKnowledgebaseResponse response = + knowledgeBase.searchKnowledgebase("compatible").blockingGet(); + + assertThat(response.getKnowledgebaseEntries()) + .extracting(KnowledgebaseEntry::getContent) + .containsExactly("legacy compatible content"); + } + + @Test + void legacyServiceAdapterSupportsReadOnlyMigration() { + SearchKnowledgebaseResponse response = new SearchKnowledgebaseResponse(); + response.setKnowledgebaseEntries( + List.of(new KnowledgebaseEntry("from viking", Map.of("source", "legacy")))); + BaseKnowledgebaseService service = query -> Single.just(response); + KnowledgebaseServiceBackendAdapter adapter = + new KnowledgebaseServiceBackendAdapter("legacy", service); + + assertThat(adapter.search("query", 5)) + .containsExactly( + new com.volcengine.veadk.knowledgebase.KnowledgebaseEntry( + "from viking", Map.of("source", "legacy"))); + assertThatThrownBy( + () -> + adapter.add( + List.of( + new com.volcengine.veadk.knowledgebase + .KnowledgebaseEntry("new")))) + .isInstanceOf(UnsupportedOperationException.class); + } +} diff --git a/core/src/test/java/com/volcengine/veadk/memory/LongTermMemoryTest.java b/core/src/test/java/com/volcengine/veadk/memory/LongTermMemoryTest.java new file mode 100644 index 0000000..8281628 --- /dev/null +++ b/core/src/test/java/com/volcengine/veadk/memory/LongTermMemoryTest.java @@ -0,0 +1,157 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.adk.events.Event; +import com.google.adk.memory.BaseMemoryService; +import com.google.adk.memory.MemoryEntry; +import com.google.adk.memory.SearchMemoryResponse; +import com.google.adk.sessions.Session; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import io.reactivex.rxjava3.core.Single; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; + +class LongTermMemoryTest { + + @Test + void facadeSavesOnlyUserTextAndKeepsUsersIsolated() { + InMemoryLongTermMemoryBackend backend = new InMemoryLongTermMemoryBackend("memory"); + LongTermMemory memory = new LongTermMemory(backend); + Event userEvent = + Event.builder() + .author("user") + .content( + Content.builder() + .role("user") + .parts(List.of(Part.fromText("I prefer Java"))) + .build()) + .build(); + Event assistantEvent = + Event.builder() + .author("assistant") + .content( + Content.builder() + .role("model") + .parts(List.of(Part.fromText("I will remember"))) + .build()) + .build(); + Session session = + Session.builder("session-1") + .appName("sample-app") + .userId("user-1") + .events(List.of(userEvent, assistantEvent)) + .build(); + + memory.addSessionToMemory(session).blockingAwait(); + + SearchMemoryResponse response = + memory.searchMemory("sample-app", "user-1", "Java").blockingGet(); + assertThat(backend.size("user-1")).isEqualTo(1); + assertThat(backend.size("user-2")).isZero(); + assertThat(response.memories()) + .extracting(entry -> entry.content().parts().orElseThrow().get(0).text().orElse("")) + .containsExactly("I prefer Java"); + } + + @Test + void inMemoryBackendRanksResultsAndValidatesTopK() { + InMemoryLongTermMemoryBackend backend = new InMemoryLongTermMemoryBackend("memory"); + backend.saveMemory( + "user-1", + List.of( + "{\"role\":\"user\",\"parts\":[{\"text\":\"Java reactive streams\"}]}", + "{\"role\":\"user\",\"parts\":[{\"text\":\"Python data science\"}]}", + "{\"role\":\"user\",\"parts\":[{\"text\":\"Java build tools\"}]}")); + + assertThat(backend.searchMemory("user-1", "Java reactive", 2)) + .containsExactly( + "{\"role\":\"user\",\"parts\":[{\"text\":\"Java reactive streams\"}]}", + "{\"role\":\"user\",\"parts\":[{\"text\":\"Java build tools\"}]}"); + assertThat(backend.searchMemory("other-user", "Java", 5)).isEmpty(); + assertThatThrownBy(() -> backend.searchMemory("user-1", "Java", 0)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void facadeConvertsNestedAndPlainBackendResults() { + LongTermMemoryBackend backend = + new LongTermMemoryBackend() { + @Override + public String index() { + return "memory"; + } + + @Override + public boolean saveMemory( + String userId, List eventStrings, Map options) { + return true; + } + + @Override + public List searchMemory( + String userId, String query, int topK, Map options) { + return List.of( + "{\"memories\":[{\"summary\":\"nested" + + " memory\",\"role\":\"assistant\"}]}", + "plain memory"); + } + }; + + List entries = + new LongTermMemory(backend) + .searchMemory("app", "user", "query") + .blockingGet() + .memories(); + + assertThat(entries).extracting(MemoryEntry::author).containsExactly("assistant", "user"); + assertThat(entries) + .extracting(entry -> entry.content().parts().orElseThrow().get(0).text().orElse("")) + .containsExactly("nested memory", "plain memory"); + } + + @Test + void legacyMemoryServiceAdapterSupportsReadOnlyMigration() { + BaseMemoryService service = mock(BaseMemoryService.class); + MemoryEntry entry = + MemoryEntry.builder() + .author("user") + .content( + Content.builder() + .role("user") + .parts(List.of(Part.fromText("legacy memory"))) + .build()) + .build(); + when(service.searchMemory("app", "user", "query")) + .thenReturn( + Single.just( + SearchMemoryResponse.builder() + .setMemories(List.of(entry)) + .build())); + MemoryServiceBackendAdapter adapter = new MemoryServiceBackendAdapter("default", service); + + List results = adapter.searchMemory("user", "query", 5, Map.of("appName", "app")); + assertThat(results).hasSize(1); + assertThat(results.get(0)).contains("legacy memory"); + assertThatThrownBy(() -> adapter.saveMemory("user", List.of("new memory"))) + .isInstanceOf(UnsupportedOperationException.class); + } +} diff --git a/core/src/test/java/com/volcengine/veadk/memory/ShortTermMemoryTest.java b/core/src/test/java/com/volcengine/veadk/memory/ShortTermMemoryTest.java new file mode 100644 index 0000000..7e9721e --- /dev/null +++ b/core/src/test/java/com/volcengine/veadk/memory/ShortTermMemoryTest.java @@ -0,0 +1,150 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory; + +import static org.assertj.core.api.Assertions.assertThat; + +import com.google.adk.events.Event; +import com.google.adk.sessions.InMemorySessionService; +import com.google.adk.sessions.Session; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import java.util.List; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; + +class ShortTermMemoryTest { + + @Test + void createSessionIsIdempotentAndListsSingleSession() { + ShortTermMemory memory = new ShortTermMemory(); + + Session first = memory.createSession("app", "user", "session").blockingGet(); + Session second = memory.createSession("app", "user", "session").blockingGet(); + + assertThat(second.id()).isEqualTo(first.id()); + assertThat(memory.listSessions("app", "user").blockingGet().sessions()) + .extracting(Session::id) + .containsExactly("session"); + assertThat(memory.backend()).isEqualTo("InMemoryShortTermMemoryBackend"); + } + + @Test + void callbackRunsAfterExistingSessionIsLoaded() { + AtomicInteger callbackCount = new AtomicInteger(); + ShortTermMemory memory = + new ShortTermMemory( + new InMemoryShortTermMemoryBackend(), + session -> callbackCount.incrementAndGet()); + memory.createSession("app", "user", "session").blockingGet(); + + Session loaded = + memory.getSession("app", "user", "session", Optional.empty()).blockingGet(); + + assertThat(loaded).isNotNull(); + assertThat(callbackCount).hasValue(1); + } + + @Test + void facadeDelegatesAppendListAndDeleteOperations() { + ShortTermMemory memory = new ShortTermMemory(); + Session session = memory.createSession("app", "user", "session").blockingGet(); + Event event = + Event.builder() + .author("user") + .content( + Content.builder() + .role("user") + .parts(List.of(Part.fromText("hello"))) + .build()) + .build(); + + memory.appendEvent(session, event).blockingGet(); + + assertThat(memory.listEvents("app", "user", "session").blockingGet().events()) + .extracting(Event::author) + .containsExactly("user"); + memory.deleteSession("app", "user", "session").blockingAwait(); + assertThat( + memory.getSession("app", "user", "session", Optional.empty()) + .isEmpty() + .blockingGet()) + .isTrue(); + } + + @Test + void existingSessionServiceCanBeAdaptedWithoutCopyingStorage() { + InMemorySessionService existing = new InMemorySessionService(); + ShortTermMemory memory = new ShortTermMemory(existing); + + Session created = memory.createSession("app", "user", "session").blockingGet(); + Session loaded = + existing.getSession("app", "user", "session", Optional.empty()).blockingGet(); + + assertThat(loaded.id()).isEqualTo(created.id()); + assertThat(existing.listSessions("app", "user").blockingGet().sessionIds()) + .containsExactly("session"); + assertThat(memory.sessionService()).isSameAs(memory); + } + + @Test + void processorRewritesLoadedHistoryWithoutMutatingStoredEvents() { + InMemoryShortTermMemoryBackend backend = new InMemoryShortTermMemoryBackend(); + ShortTermMemoryProcessor processor = + new ShortTermMemoryProcessor( + messages -> { + assertThat(messages) + .extracting(ShortTermMemoryMessage::content) + .containsExactly("first", "second"); + return List.of(new ShortTermMemoryMessage("user", "compacted history")); + }); + ShortTermMemory memory = new ShortTermMemory(backend, processor); + Session session = memory.createSession("app", "user", "session").blockingGet(); + memory.appendEvent(session, event("user", "first")).blockingGet(); + memory.appendEvent(session, event("assistant", "second")).blockingGet(); + + Session optimized = + memory.getSession("app", "user", "session", Optional.empty()).blockingGet(); + Session stored = + backend.sessionService() + .getSession("app", "user", "session", Optional.empty()) + .blockingGet(); + + assertThat(optimized.events()).hasSize(1); + assertThat(optimized.events().get(0).author()).isEqualTo("memory_optimizer"); + assertThat(optimized.events().get(0).stringifyContent()).contains("compacted history"); + assertThat(stored.events()).hasSize(2); + + memory.appendEvent(optimized, event("user", "third")).blockingGet(); + Session storedAfterAppend = + backend.sessionService() + .getSession("app", "user", "session", Optional.empty()) + .blockingGet(); + + assertThat(optimized.events()) + .extracting(Event::stringifyContent) + .containsExactly("compacted history", "third"); + assertThat(storedAfterAppend.events()) + .extracting(Event::stringifyContent) + .containsExactly("first", "second", "third"); + } + + private static Event event(String role, String text) { + return Event.builder() + .author(role) + .content(Content.builder().role(role).parts(List.of(Part.fromText(text))).build()) + .build(); + } +} diff --git a/core/src/test/java/com/volcengine/veadk/memory/viking/VikingMemoryServiceTest.java b/core/src/test/java/com/volcengine/veadk/memory/viking/VikingMemoryServiceTest.java index 91a369c..2234365 100644 --- a/core/src/test/java/com/volcengine/veadk/memory/viking/VikingMemoryServiceTest.java +++ b/core/src/test/java/com/volcengine/veadk/memory/viking/VikingMemoryServiceTest.java @@ -2,6 +2,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -10,12 +11,14 @@ import com.google.adk.sessions.Session; import com.google.genai.types.Content; import com.google.genai.types.Part; +import com.volcengine.veadk.integration.vikingmemory.Message; import com.volcengine.veadk.integration.vikingmemory.Metadata; import com.volcengine.veadk.integration.vikingmemory.VikingMemoryWrapper; import com.volcengine.veadk.utils.EnvUtil; import io.reactivex.rxjava3.observers.TestObserver; import java.util.Collections; import java.util.List; +import java.util.Map; import java.util.Optional; import org.junit.jupiter.api.Test; import org.mockito.ArgumentCaptor; @@ -280,4 +283,48 @@ void searchMemory_returnsResponseWithEntries_and_callsWrapperWithExpectedArgs() assertEquals(List.of("sys_event_v1", "user_event_v1"), eventTypes); } } + + @Test + void backendContractKeepsLegacyVikingServiceUsable() throws Exception { + String appName = "AppMem"; + try (MockedStatic mockedEnv = Mockito.mockStatic(EnvUtil.class); + MockedConstruction mockedCtor = + Mockito.mockConstruction( + VikingMemoryWrapper.class, + (mock, context) -> { + Mockito.when(mock.isCollectionExists(appName)).thenReturn(true); + Mockito.when( + mock.addSession( + Mockito.eq(appName), + Mockito.anyList(), + Mockito.any(Metadata.class))) + .thenReturn(true); + })) { + mockedEnv.when(EnvUtil::getAccessKey).thenReturn("ak"); + mockedEnv.when(EnvUtil::getSecretKey).thenReturn("sk"); + mockedEnv.when(EnvUtil::getVikingMmemoryType).thenReturn("sys_event_v1"); + VikingMemoryService service = new VikingMemoryService(appName); + + boolean saved = + service.saveMemory( + "user-1", + List.of( + "{\"role\":\"model\",\"parts\":[{\"text\":\"first\"},{\"text\":\"second\"}]}"), + Map.of()); + + assertTrue(saved); + assertEquals(appName, service.index()); + VikingMemoryWrapper wrapperMock = mockedCtor.constructed().get(0); + ArgumentCaptor messagesCaptor = ArgumentCaptor.forClass(List.class); + verify(wrapperMock) + .addSession( + Mockito.eq(appName), + messagesCaptor.capture(), + Mockito.any(Metadata.class)); + @SuppressWarnings("unchecked") + List messages = messagesCaptor.getValue(); + assertEquals("assistant", messages.get(0).getRole()); + assertEquals("first\nsecond", messages.get(0).getContent()); + } + } } diff --git a/core/src/test/java/com/volcengine/veadk/model/ArkEmbeddingTest.java b/core/src/test/java/com/volcengine/veadk/model/ArkEmbeddingTest.java new file mode 100644 index 0000000..fe9c025 --- /dev/null +++ b/core/src/test/java/com/volcengine/veadk/model/ArkEmbeddingTest.java @@ -0,0 +1,145 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.model; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbedding; +import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbeddingRequest; +import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbeddingResult; +import com.volcengine.ark.runtime.model.multimodalembeddings.MultimodalEmbeddingUsage; +import com.volcengine.ark.runtime.service.ArkService; +import com.volcengine.veadk.config.VeADKConfig; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; + +class ArkEmbeddingTest { + + @Test + void embedBuildsMultimodalRequestAndReturnsUsage() { + FakeArkService service = new FakeArkService(); + service.responses.add(response("test-model", List.of(0.1, 0.2), 3, 4)); + ArkEmbedding embedding = + ArkEmbedding.builder() + .modelName("test-model") + .dimensions(2) + .arkService(service) + .build(); + + ArkEmbedding.EmbeddingResponse response = embedding.embedWithUsage("hello"); + + assertThat(response.embedding()).containsExactly(0.1, 0.2); + assertThat(response.promptTokens()).isEqualTo(3); + assertThat(response.totalTokens()).isEqualTo(4); + MultimodalEmbeddingRequest request = service.requests.get(0); + assertThat(request.getModel()).isEqualTo("test-model"); + assertThat(request.getDimensions()).isEqualTo(2); + assertThat(request.getInput()).hasSize(1); + assertThat(request.getInput().get(0).getType()).isEqualTo("text"); + assertThat(request.getInput().get(0).getText()).isEqualTo("hello"); + } + + @Test + void batchAndReactiveApisPreserveInputOrder() { + FakeArkService service = new FakeArkService(); + service.responses.add(response("test-model", List.of(1.0), 0, 0)); + service.responses.add(response("test-model", List.of(2.0), 0, 0)); + service.responses.add(response("test-model", List.of(3.0), 0, 0)); + ArkEmbedding embedding = ArkEmbedding.builder().arkService(service).build(); + + assertThat(embedding.getTextEmbeddings(List.of("one", "two"))) + .containsExactly(List.of(1.0), List.of(2.0)); + assertThat(embedding.embedAsync("three").blockingGet()).containsExactly(3.0); + assertThat(service.requests) + .extracting(request -> request.getInput().get(0).getText()) + .containsExactly("one", "two", "three"); + } + + @Test + void malformedResponseFailsWithActionableError() { + FakeArkService service = new FakeArkService(); + service.responses.add(new MultimodalEmbeddingResult()); + ArkEmbedding embedding = ArkEmbedding.builder().arkService(service).build(); + + assertThatThrownBy(() -> embedding.embed("text")) + .isInstanceOf(IllegalStateException.class) + .hasMessageContaining("did not contain an embedding"); + } + + @Test + void configuredBuilderUsesEmbeddingSettings() { + VeADKConfig config = + VeADKConfig.from( + Map.of( + "MODEL_EMBEDDING_NAME", "configured-embedding", + "MODEL_EMBEDDING_DIM", "512", + "MODEL_EMBEDDING_API_KEY", "test-key", + "MODEL_EMBEDDING_API_BASE", "https://ark.example/api/v3")); + + ArkEmbedding embedding = ArkEmbedding.builder(config).build(); + + assertThat(embedding.modelName()).isEqualTo("configured-embedding"); + assertThat(embedding.dimensions()).isEqualTo(512); + } + + @Test + void configuredBuilderFallsBackToAgentApiKey() { + VeADKConfig config = + VeADKConfig.from( + Map.of( + "MODEL_EMBEDDING_NAME", "configured-embedding", + "MODEL_AGENT_API_KEY", "agent-key")); + + ArkEmbedding embedding = ArkEmbedding.builder(config).build(); + + assertThat(embedding.modelName()).isEqualTo("configured-embedding"); + } + + private static MultimodalEmbeddingResult response( + String model, List vector, long promptTokens, long totalTokens) { + MultimodalEmbedding data = new MultimodalEmbedding(); + data.setEmbedding(vector); + MultimodalEmbeddingUsage usage = new MultimodalEmbeddingUsage(); + usage.setPromptTokens(promptTokens); + usage.setTotalTokens(totalTokens); + MultimodalEmbeddingResult result = new MultimodalEmbeddingResult(); + result.setModel(model); + result.setData(data); + result.setUsage(usage); + return result; + } + + private static final class FakeArkService extends ArkService { + private final List requests = new ArrayList<>(); + private final Deque responses = new ArrayDeque<>(); + + private FakeArkService() { + super("test-key"); + } + + @Override + public MultimodalEmbeddingResult createMultiModalEmbeddings( + MultimodalEmbeddingRequest request) { + requests.add(request); + return responses.removeFirst(); + } + } +} diff --git a/core/src/test/java/com/volcengine/veadk/model/ArkLlmTest.java b/core/src/test/java/com/volcengine/veadk/model/ArkLlmTest.java index f40af8c..0adb22f 100644 --- a/core/src/test/java/com/volcengine/veadk/model/ArkLlmTest.java +++ b/core/src/test/java/com/volcengine/veadk/model/ArkLlmTest.java @@ -18,20 +18,39 @@ import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertTrue; import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; import static org.mockito.Mockito.mockStatic; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.google.adk.models.LlmRequest; import com.google.adk.models.LlmResponse; +import com.google.adk.tools.BaseTool; +import com.google.genai.types.Blob; import com.google.genai.types.Content; +import com.google.genai.types.FileData; +import com.google.genai.types.FinishReason; import com.google.genai.types.FunctionCall; +import com.google.genai.types.FunctionDeclaration; +import com.google.genai.types.FunctionResponse; +import com.google.genai.types.GenerateContentConfig; +import com.google.genai.types.GenerateContentResponseUsageMetadata; import com.google.genai.types.Part; +import com.google.genai.types.Schema; +import com.google.genai.types.Type; +import com.google.genai.types.VideoMetadata; +import com.volcengine.ark.runtime.model.CompletionTokensDetails; +import com.volcengine.ark.runtime.model.PromptTokensDetails; +import com.volcengine.ark.runtime.model.Usage; import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChoice; import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChunk; +import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionContentPart; import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest; import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionResult; import com.volcengine.ark.runtime.model.completion.chat.ChatFunctionCall; import com.volcengine.ark.runtime.model.completion.chat.ChatMessage; +import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole; import com.volcengine.ark.runtime.model.completion.chat.ChatToolCall; import com.volcengine.ark.runtime.service.ArkService; import com.volcengine.veadk.utils.EnvUtil; @@ -41,10 +60,12 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.concurrent.TimeUnit; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; import org.mockito.Mock; import org.mockito.MockedStatic; import org.mockito.junit.jupiter.MockitoExtension; @@ -81,6 +102,14 @@ void generateContent_nonStreaming_textResponse() throws InterruptedException { .build(); ChatCompletionResult mockResult = createMockTextResult("Hi there!"); + Usage usage = new Usage(5, 2, 7); + PromptTokensDetails promptTokensDetails = new PromptTokensDetails(); + promptTokensDetails.setCachedTokens(3); + usage.setPromptTokensDetails(promptTokensDetails); + CompletionTokensDetails completionTokensDetails = new CompletionTokensDetails(); + completionTokensDetails.setReasoningTokens(1); + usage.setCompletionTokensDetails(completionTokensDetails); + mockResult.setUsage(usage); when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) .thenReturn(mockResult); @@ -94,6 +123,15 @@ void generateContent_nonStreaming_textResponse() throws InterruptedException { assertTrue(response.content().isPresent()); assertTrue(response.content().get().parts().isPresent()); assertEquals("Hi there!", response.content().get().parts().get().get(0).text().get()); + assertTrue(response.finishReason().isPresent()); + assertEquals(FinishReason.Known.STOP, response.finishReason().get().knownEnum()); + assertTrue(response.usageMetadata().isPresent()); + GenerateContentResponseUsageMetadata usageMetadata = response.usageMetadata().get(); + assertEquals(Integer.valueOf(5), usageMetadata.promptTokenCount().get()); + assertEquals(Integer.valueOf(2), usageMetadata.candidatesTokenCount().get()); + assertEquals(Integer.valueOf(7), usageMetadata.totalTokenCount().get()); + assertEquals(Integer.valueOf(3), usageMetadata.cachedContentTokenCount().get()); + assertEquals(Integer.valueOf(1), usageMetadata.thoughtsTokenCount().get()); } @Test @@ -135,10 +173,395 @@ void generateContent_nonStreaming_toolCallResponse() throws InterruptedException FunctionCall fc = response.content().get().parts().get().get(0).functionCall().get(); assertTrue(fc.name().isPresent()); assertEquals("search", fc.name().get()); + assertTrue(fc.id().isPresent()); + assertEquals("tool-123", fc.id().get()); assertTrue(fc.args().isPresent()); assertEquals(Map.of("query", "cats"), fc.args().get()); } + @Test + void generateContent_mapsFunctionCallAndResponseHistoryToArkToolMessages() + throws InterruptedException { + Content assistantCall = + Content.builder() + .role("model") + .parts( + List.of( + Part.fromText("Searching"), + Part.builder() + .functionCall( + FunctionCall.builder() + .id("call-1") + .name("search") + .args(Map.of("query", "cats")) + .build()) + .build())) + .build(); + Content toolResult = + Content.builder() + .role("user") + .parts( + Part.builder() + .functionResponse( + FunctionResponse.builder() + .id("call-1") + .name("search") + .response(Map.of("result", "cats")) + .build()) + .build()) + .build(); + LlmRequest llmRequest = + LlmRequest.builder() + .model("test-model") + .contents(List.of(assistantCall, toolResult)) + .build(); + when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) + .thenReturn(createMockTextResult("done")); + + TestSubscriber testSubscriber = + arkLlm.generateContent(llmRequest, false).test(); + + testSubscriber.awaitDone(5, TimeUnit.SECONDS); + testSubscriber.assertNoErrors(); + + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService).createChatCompletion(captor.capture()); + List messages = captor.getValue().getMessages(); + assertEquals(2, messages.size()); + assertEquals(ChatMessageRole.ASSISTANT, messages.get(0).getRole()); + assertEquals("Searching", messages.get(0).getContent()); + assertEquals("call-1", messages.get(0).getToolCalls().get(0).getId()); + assertEquals("function", messages.get(0).getToolCalls().get(0).getType()); + assertEquals("search", messages.get(0).getToolCalls().get(0).getFunction().getName()); + assertEquals( + "{\"query\":\"cats\"}", + messages.get(0).getToolCalls().get(0).getFunction().getArguments()); + assertEquals(ChatMessageRole.TOOL, messages.get(1).getRole()); + assertEquals("call-1", messages.get(1).getToolCallId()); + assertEquals("search", messages.get(1).getName()); + assertEquals("{\"result\":\"cats\"}", messages.get(1).getContent()); + } + + @Test + void generateContent_mapsGenerationConfigToArkRequest() throws InterruptedException { + GenerateContentConfig config = + GenerateContentConfig.builder() + .temperature(0.2f) + .topP(0.9f) + .maxOutputTokens(128) + .stopSequences(List.of("END")) + .presencePenalty(0.1f) + .frequencyPenalty(0.3f) + .candidateCount(2) + .responseLogprobs(true) + .logprobs(4) + .build(); + LlmRequest llmRequest = + LlmRequest.builder() + .model("test-model") + .config(config) + .contents( + Collections.singletonList( + Content.builder() + .role("user") + .parts(Part.fromText("Hello")) + .build())) + .build(); + when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) + .thenReturn(createMockTextResult("configured response")); + + TestSubscriber testSubscriber = + arkLlm.generateContent(llmRequest, false).test(); + + testSubscriber.awaitDone(5, TimeUnit.SECONDS); + testSubscriber.assertNoErrors(); + + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService).createChatCompletion(captor.capture()); + ChatCompletionRequest request = captor.getValue(); + assertEquals(0.2, request.getTemperature(), 0.0001); + assertEquals(0.9, request.getTopP(), 0.0001); + assertEquals(Integer.valueOf(128), request.getMaxTokens()); + assertEquals(List.of("END"), request.getStop()); + assertEquals(0.1, request.getPresencePenalty(), 0.0001); + assertEquals(0.3, request.getFrequencyPenalty(), 0.0001); + assertEquals(Integer.valueOf(2), request.getN()); + assertEquals(Boolean.TRUE, request.getLogprobs()); + assertEquals(Integer.valueOf(4), request.getTopLogprobs()); + } + + @Test + void generateContent_mapsJsonMimeTypeToArkResponseFormat() throws InterruptedException { + GenerateContentConfig config = + GenerateContentConfig.builder().responseMimeType("application/json").build(); + LlmRequest llmRequest = + LlmRequest.builder() + .model("test-model") + .config(config) + .contents( + Collections.singletonList( + Content.builder() + .role("user") + .parts(Part.fromText("Hello")) + .build())) + .build(); + when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) + .thenReturn(createMockTextResult("{\"answer\":\"ok\"}")); + + TestSubscriber testSubscriber = + arkLlm.generateContent(llmRequest, false).test(); + + testSubscriber.awaitDone(5, TimeUnit.SECONDS); + testSubscriber.assertNoErrors(); + + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService).createChatCompletion(captor.capture()); + assertEquals("json_object", captor.getValue().getResponseFormat().getType()); + } + + @Test + void generateContent_mapsResponseSchemaToArkJsonSchemaFormat() throws InterruptedException { + Schema schema = + Schema.builder() + .title("Answer") + .description("Structured answer") + .type(Type.Known.OBJECT) + .properties( + Map.of( + "answer", + Schema.builder() + .type(Type.Known.STRING) + .description("Answer text") + .build())) + .required("answer") + .build(); + GenerateContentConfig config = + GenerateContentConfig.builder().responseSchema(schema).build(); + LlmRequest llmRequest = + LlmRequest.builder() + .model("test-model") + .config(config) + .contents( + Collections.singletonList( + Content.builder() + .role("user") + .parts(Part.fromText("Hello")) + .build())) + .build(); + when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) + .thenReturn(createMockTextResult("{\"answer\":\"ok\"}")); + + TestSubscriber testSubscriber = + arkLlm.generateContent(llmRequest, false).test(); + + testSubscriber.awaitDone(5, TimeUnit.SECONDS); + testSubscriber.assertNoErrors(); + + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService).createChatCompletion(captor.capture()); + ChatCompletionRequest request = captor.getValue(); + assertEquals("json_schema", request.getResponseFormat().getType()); + assertEquals("Answer", request.getResponseFormat().getJsonSchema().getName()); + assertTrue(request.getResponseFormat().getJsonSchema().isStrict()); + assertEquals( + "object", + request.getResponseFormat().getJsonSchema().getSchema().get("type").asText()); + assertEquals( + "string", + request.getResponseFormat() + .getJsonSchema() + .getSchema() + .get("properties") + .get("answer") + .get("type") + .asText()); + } + + @Test + void generateContent_mapsResponseJsonSchemaToArkJsonSchemaFormat() throws InterruptedException { + GenerateContentConfig config = + GenerateContentConfig.builder() + .responseJsonSchema( + Map.of( + "type", + "object", + "properties", + Map.of("answer", Map.of("type", "string")), + "required", + List.of("answer"))) + .build(); + LlmRequest llmRequest = requestWithConfig(config); + when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) + .thenReturn(createMockTextResult("{\"answer\":\"ok\"}")); + + arkLlm.generateContent(llmRequest, false).test().awaitDone(5, TimeUnit.SECONDS); + + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService).createChatCompletion(captor.capture()); + assertEquals("json_schema", captor.getValue().getResponseFormat().getType()); + assertEquals( + "string", + captor.getValue() + .getResponseFormat() + .getJsonSchema() + .getSchema() + .get("properties") + .get("answer") + .get("type") + .asText()); + } + + @Test + void generateContent_mapsToolParametersJsonSchema() throws InterruptedException { + BaseTool tool = mock(BaseTool.class); + when(tool.name()).thenReturn("search"); + when(tool.description()).thenReturn("Search documents"); + when(tool.declaration()) + .thenReturn( + Optional.of( + FunctionDeclaration.builder() + .name("search") + .parametersJsonSchema( + Map.of( + "type", + "OBJECT", + "properties", + Map.of("query", Map.of("type", "STRING")))) + .build())); + LlmRequest llmRequest = + LlmRequest.builder() + .model("test-model") + .contents( + List.of( + Content.builder() + .role("user") + .parts(Part.fromText("Hello")) + .build())) + .appendTools(List.of(tool)) + .build(); + when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) + .thenReturn(createMockTextResult("ok")); + + arkLlm.generateContent(llmRequest, false).test().awaitDone(5, TimeUnit.SECONDS); + + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService).createChatCompletion(captor.capture()); + assertEquals(1, captor.getValue().getTools().size()); + assertEquals( + "object", + captor.getValue() + .getTools() + .get(0) + .getFunction() + .getParameters() + .get("type") + .asText()); + assertEquals( + "string", + captor.getValue() + .getTools() + .get(0) + .getFunction() + .getParameters() + .get("properties") + .get("query") + .get("type") + .asText()); + } + + @Test + void generateContent_keepsTextOnlyMessageContentAsString() throws InterruptedException { + LlmRequest llmRequest = + LlmRequest.builder() + .model("test-model") + .contents( + Collections.singletonList( + Content.builder() + .role("user") + .parts(Part.fromText("plain text")) + .build())) + .build(); + when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) + .thenReturn(createMockTextResult("ok")); + + TestSubscriber testSubscriber = + arkLlm.generateContent(llmRequest, false).test(); + + testSubscriber.awaitDone(5, TimeUnit.SECONDS); + testSubscriber.assertNoErrors(); + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService).createChatCompletion(captor.capture()); + assertEquals("plain text", captor.getValue().getMessages().get(0).getContent()); + } + + @Test + void generateContent_mapsInlineImageAndFileVideoToArkContentParts() + throws InterruptedException { + Part imagePart = + Part.builder() + .inlineData( + Blob.builder() + .mimeType("image/png") + .data(new byte[] {1, 2, 3}) + .build()) + .build(); + Part videoPart = + Part.builder() + .fileData( + FileData.builder() + .mimeType("video/mp4") + .fileUri("https://example.com/video.mp4") + .build()) + .videoMetadata(VideoMetadata.builder().fps(2.5).build()) + .build(); + LlmRequest llmRequest = + LlmRequest.builder() + .model("test-model") + .contents( + Collections.singletonList( + Content.builder() + .role("user") + .parts( + List.of( + Part.fromText("describe"), + imagePart, + videoPart)) + .build())) + .build(); + when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) + .thenReturn(createMockTextResult("ok")); + + TestSubscriber testSubscriber = + arkLlm.generateContent(llmRequest, false).test(); + + testSubscriber.awaitDone(5, TimeUnit.SECONDS); + testSubscriber.assertNoErrors(); + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService).createChatCompletion(captor.capture()); + + Object content = captor.getValue().getMessages().get(0).getContent(); + assertTrue(content instanceof List); + @SuppressWarnings("unchecked") + List parts = (List) content; + assertEquals(3, parts.size()); + assertEquals("text", parts.get(0).getType()); + assertEquals("describe", parts.get(0).getText()); + assertEquals("image_url", parts.get(1).getType()); + assertEquals("auto", parts.get(1).getImageUrl().getDetail()); + assertEquals("data:image/png;base64,AQID", parts.get(1).getImageUrl().getUrl()); + assertEquals("video_url", parts.get(2).getType()); + assertEquals("https://example.com/video.mp4", parts.get(2).getVideoUrl().getUrl()); + assertEquals(2.5, parts.get(2).getVideoUrl().getFps(), 0.0001); + } + @Test void generateContent_streaming_textResponse() throws InterruptedException { LlmRequest llmRequest = @@ -184,6 +607,168 @@ void generateContent_streaming_textResponse() throws InterruptedException { "Hello World!", finalResponse.content().get().parts().get().get(0).text().get()); } + @Test + void generateContent_streaming_parallelToolCallsWithoutText() throws InterruptedException { + LlmRequest llmRequest = + LlmRequest.builder() + .model("test-model") + .contents( + List.of( + Content.builder() + .role("user") + .parts(Part.fromText("Use tools")) + .build())) + .build(); + Usage usage = new Usage(3, 4, 7); + io.reactivex.Flowable chunks = + io.reactivex.Flowable.just( + createToolCallChunk( + List.of( + toolCall(0, "call-1", "first", "{\"value\":"), + toolCall(1, "call-2", "second", "{\"value\":"))), + createToolCallChunk( + List.of( + toolCall(0, null, null, "1}"), + toolCall(1, null, null, "2}"))), + createStopChunk("tool_calls"), + createUsageChunk(usage)); + when(arkService.streamChatCompletion(any(ChatCompletionRequest.class))).thenReturn(chunks); + + TestSubscriber subscriber = arkLlm.generateContent(llmRequest, true).test(); + + subscriber.awaitDone(5, TimeUnit.SECONDS); + subscriber.assertNoErrors(); + subscriber.assertValueCount(1); + LlmResponse response = subscriber.values().get(0); + assertTrue(response.partial().isPresent() && !response.partial().get()); + List parts = response.content().get().parts().get(); + assertEquals(2, parts.size()); + assertEquals("call-1", parts.get(0).functionCall().get().id().get()); + assertEquals(Map.of("value", 1), parts.get(0).functionCall().get().args().get()); + assertEquals("call-2", parts.get(1).functionCall().get().id().get()); + assertEquals(Map.of("value", 2), parts.get(1).functionCall().get().args().get()); + assertEquals(FinishReason.Known.STOP, response.finishReason().get().knownEnum()); + assertEquals(Integer.valueOf(7), response.usageMetadata().get().totalTokenCount().get()); + + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService).streamChatCompletion(captor.capture()); + assertEquals(Boolean.TRUE, captor.getValue().getStreamOptions().getIncludeUsage()); + } + + @Test + void generateContent_nonStreaming_usesFallbackWhenPrimaryFails() throws Exception { + ArkLlm fallbackLlm = + new ArkLlm(List.of("primary-model", "fallback-model"), "test-api-key", null); + injectArkService(fallbackLlm); + LlmRequest llmRequest = + LlmRequest.builder() + .model("primary-model") + .contents( + Collections.singletonList( + Content.builder() + .role("user") + .parts(Part.fromText("Hello")) + .build())) + .build(); + + when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) + .thenThrow(new RuntimeException("primary failed")) + .thenReturn(createMockTextResult("fallback response")); + + TestSubscriber testSubscriber = + fallbackLlm.generateContent(llmRequest, false).test(); + + testSubscriber.awaitDone(5, TimeUnit.SECONDS); + testSubscriber.assertNoErrors(); + testSubscriber.assertValueCount(1); + assertEquals( + "fallback response", + testSubscriber.values().get(0).content().get().parts().get().get(0).text().get()); + + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService, times(2)).createChatCompletion(captor.capture()); + assertEquals( + List.of("primary-model", "fallback-model"), + captor.getAllValues().stream().map(ChatCompletionRequest::getModel).toList()); + } + + @Test + void generateContent_streaming_doesNotFallbackAfterEmittingOutput() throws Exception { + ArkLlm fallbackLlm = + new ArkLlm(List.of("primary-model", "fallback-model"), "test-api-key", null); + injectArkService(fallbackLlm); + LlmRequest llmRequest = + LlmRequest.builder() + .contents( + Collections.singletonList( + Content.builder() + .role("user") + .parts(Part.fromText("Hello")) + .build())) + .build(); + RuntimeException streamFailure = new RuntimeException("stream interrupted"); + io.reactivex.Flowable failingStream = + io.reactivex.Flowable.concat( + io.reactivex.Flowable.just( + createMockTextChunk("abcdefghijklmnopqrstuvwxyz12345")), + io.reactivex.Flowable.error(streamFailure)); + when(arkService.streamChatCompletion(any(ChatCompletionRequest.class))) + .thenReturn(failingStream); + + TestSubscriber testSubscriber = + fallbackLlm.generateContent(llmRequest, true).test(); + + testSubscriber.awaitDone(5, TimeUnit.SECONDS); + testSubscriber.assertValueCount(1); + testSubscriber.assertError(streamFailure); + assertTrue(testSubscriber.values().get(0).partial().orElse(false)); + verify(arkService, times(1)).streamChatCompletion(any(ChatCompletionRequest.class)); + } + + @Test + void generateContent_respectsExplicitRequestModelWithoutFallbacks() throws Exception { + ArkLlm fallbackLlm = + new ArkLlm(List.of("primary-model", "fallback-model"), "test-api-key", null); + injectArkService(fallbackLlm); + LlmRequest llmRequest = + LlmRequest.builder() + .model("request-model") + .contents( + Collections.singletonList( + Content.builder() + .role("user") + .parts(Part.fromText("Hello")) + .build())) + .build(); + when(arkService.createChatCompletion(any(ChatCompletionRequest.class))) + .thenReturn(createMockTextResult("request response")); + + TestSubscriber testSubscriber = + fallbackLlm.generateContent(llmRequest, false).test(); + + testSubscriber.awaitDone(5, TimeUnit.SECONDS); + testSubscriber.assertNoErrors(); + ArgumentCaptor captor = + ArgumentCaptor.forClass(ChatCompletionRequest.class); + verify(arkService).createChatCompletion(captor.capture()); + assertEquals("request-model", captor.getValue().getModel()); + } + + private LlmRequest requestWithConfig(GenerateContentConfig config) { + return LlmRequest.builder() + .model("test-model") + .config(config) + .contents( + List.of( + Content.builder() + .role("user") + .parts(Part.fromText("Hello")) + .build())) + .build(); + } + private ChatCompletionResult createMockTextResult(String content) { ChatCompletionResult mockResult = new ChatCompletionResult(); ChatCompletionChoice mockChoice = new ChatCompletionChoice(); @@ -226,12 +811,53 @@ private ChatCompletionChunk createMockTextChunk(String content) { return chunk; } + private ChatCompletionChunk createToolCallChunk(List toolCalls) { + ChatCompletionChunk chunk = new ChatCompletionChunk(); + ChatCompletionChoice choice = new ChatCompletionChoice(); + ChatMessage message = new ChatMessage(); + message.setToolCalls(toolCalls); + choice.setMessage(message); + chunk.setChoices(List.of(choice)); + return chunk; + } + + private ChatToolCall toolCall( + int index, String id, String functionName, String argumentsFragment) { + ChatFunctionCall function = new ChatFunctionCall(); + function.setName(functionName); + function.setArguments(argumentsFragment); + ChatToolCall toolCall = new ChatToolCall(); + toolCall.setIndex(index); + toolCall.setId(id); + toolCall.setType("function"); + toolCall.setFunction(function); + return toolCall; + } + private ChatCompletionChunk createStopChunk() { + return createStopChunk("stop"); + } + + private ChatCompletionChunk createStopChunk(String finishReason) { ChatCompletionChunk chunk = new ChatCompletionChunk(); ChatCompletionChoice choice = new ChatCompletionChoice(); - choice.setFinishReason("stop"); + choice.setFinishReason(finishReason); choice.setMessage(new ChatMessage()); chunk.setChoices(Collections.singletonList(choice)); return chunk; } + + private ChatCompletionChunk createUsageChunk(Usage usage) { + ChatCompletionChunk chunk = new ChatCompletionChunk(); + chunk.setChoices(List.of()); + chunk.setUsage(usage); + return chunk; + } + + private void injectArkService(ArkLlm target) + throws NoSuchFieldException, IllegalAccessException { + Field field = ArkLlm.class.getDeclaredField("arkService"); + field.setAccessible(true); + field.set(target, arkService); + } } diff --git a/core/src/test/java/com/volcengine/veadk/runner/RunnerProcessorTest.java b/core/src/test/java/com/volcengine/veadk/runner/RunnerProcessorTest.java new file mode 100644 index 0000000..c23eebd --- /dev/null +++ b/core/src/test/java/com/volcengine/veadk/runner/RunnerProcessorTest.java @@ -0,0 +1,119 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.runner; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import com.google.adk.agents.BaseAgent; +import com.google.adk.agents.InvocationContext; +import com.google.adk.agents.RunConfig; +import com.google.adk.artifacts.InMemoryArtifactService; +import com.google.adk.events.Event; +import com.google.adk.memory.InMemoryMemoryService; +import com.google.adk.sessions.InMemorySessionService; +import com.google.adk.sessions.Session; +import com.google.genai.types.Content; +import com.volcengine.veadk.processors.BaseRunProcessor; +import com.volcengine.veadk.processors.RunContext; +import io.reactivex.rxjava3.core.Flowable; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Test; + +class RunnerProcessorTest { + + @Test + void runAsync_passesImmutableContextThroughConfiguredProcessor() { + AtomicReference observed = new AtomicReference<>(); + Event processedEvent = Event.builder().author("processor").build(); + BaseRunProcessor processor = + (context, eventGenerator) -> { + observed.set(context); + return Flowable.just(processedEvent); + }; + Runner runner = createRunner(processor); + Session session = + Session.builder("test-session").appName("test-app").userId("test-user").build(); + Content message = Content.builder().role("user").parts(List.of()).build(); + RunConfig runConfig = RunConfig.builder().build(); + + List events = + runner.runAsync(session, message, runConfig, Map.of("tenant", "one")) + .toList() + .blockingGet(); + + assertThat(events).containsExactly(processedEvent); + assertThat(observed.get().runner()).isSameAs(runner); + assertThat(observed.get().message()).isSameAs(message); + assertThat(observed.get().appName()).isEqualTo("test-app"); + assertThat(observed.get().userId()).isEqualTo("test-user"); + assertThat(observed.get().sessionId()).isEqualTo("test-session"); + assertThat(observed.get().invocationState()).containsEntry("tenant", "one"); + assertThatThrownBy(() -> observed.get().invocationState().put("other", "two")) + .isInstanceOf(UnsupportedOperationException.class); + } + + @Test + void processorExecutionIsDeferredUntilTheEventStreamIsSubscribed() { + AtomicReference observed = new AtomicReference<>(); + BaseRunProcessor processor = + (context, eventGenerator) -> { + observed.set(context); + return Flowable.empty(); + }; + Runner runner = createRunner(processor); + Session session = + Session.builder("test-session").appName("test-app").userId("test-user").build(); + Content message = Content.builder().role("user").parts(List.of()).build(); + + Flowable events = + runner.runAsync(session, message, RunConfig.builder().build(), Map.of()); + + assertThat(observed.get()).isNull(); + events.test().assertComplete(); + assertThat(observed.get()).isNotNull(); + } + + private static Runner createRunner(BaseRunProcessor processor) { + return new Runner( + new EmptyAgent(), + "test-app", + new InMemoryArtifactService(), + new InMemorySessionService(), + new InMemoryMemoryService(), + List.of(), + processor); + } + + private static final class EmptyAgent extends BaseAgent { + + private EmptyAgent() { + super("empty-agent", "", List.of(), List.of(), List.of()); + } + + @Override + protected Flowable runAsyncImpl(InvocationContext invocationContext) { + return Flowable.empty(); + } + + @Override + protected Flowable runLiveImpl(InvocationContext invocationContext) { + return Flowable.empty(); + } + } +} diff --git a/docs/parity/BASELINE.md b/docs/parity/BASELINE.md new file mode 100644 index 0000000..37b849a --- /dev/null +++ b/docs/parity/BASELINE.md @@ -0,0 +1,46 @@ +# VeADK parity baseline + +Baseline captured on 2026-08-14 before parity work starts. + +## Source revisions + +| Repository | Revision | Role | +| --- | --- | --- | +| `volcengine/veadk-java` | `dc77c38dd2ab35d43535af13543f7c2b9b0a7189` | Compatibility baseline and implementation target | +| `volcengine/veadk-python` | `36f55ee24a18057300f564e20a63fba686badb47` | Read-only behavior and feature reference | + +The Python revision is intentionally pinned. New upstream Python commits are not part of the +current parity target until this baseline is completed. + +## Java baseline + +- Maven coordinates: `com.volcengine.veadk:veadk-java:0.0.2`. +- Minimum Java release: 17. +- Google ADK dependency: 0.4.0. +- Main source files: 27. +- Test source files: 12. +- Example source/resource files: 4. +- Public packages already exposed under `com.volcengine.veadk` must remain available. + +At the pinned revision the README showed version `0.0.1`. The current README corrects that +historical documentation mismatch to the Maven project version, `0.0.2`, without changing the +compatibility baseline. + +## Python reference size + +- Package files under `veadk/`: 492. +- Test files under `tests/`: 154. +- Google ADK requirement: `>=1.34.0`. + +The large ADK version difference means behavior must be ported through explicit Java adapters. +Upgrading Google ADK Java is a separate compatibility decision, not an automatic prerequisite. + +## Baseline test status + +The user ran the unchanged baseline with `./mvnw -B -ntp clean verify` after configuring the +already-installed OpenJDK 17. All three reactor modules passed with 73 tests and no failures. + +After the first parity batch, the same command also passes inside the restricted environment: +all three modules build, 88 tests pass, examples compile, and binary, source, and Javadoc jars are +created. Mockito now loads Byte Buddy at test-JVM startup so the test suite does not depend on +runtime agent attachment. diff --git a/docs/parity/COMPATIBILITY.md b/docs/parity/COMPATIBILITY.md new file mode 100644 index 0000000..ca300ec --- /dev/null +++ b/docs/parity/COMPATIBILITY.md @@ -0,0 +1,43 @@ +# Backward compatibility policy + +Parity work is additive. Existing Java users must be able to upgrade without source changes. + +## Protected compatibility surface + +The following are protected from incompatible changes: + +- Maven coordinates and Java 17 bytecode baseline. +- Existing `com.volcengine.veadk` packages, public classes, interfaces, records, constructors, + methods, fields, and generic return types. +- `Runner` constructors and its inheritance from Google ADK's runner. +- Existing environment-variable names and their current defaults. +- JSON property names and response shapes used by tools and Volcengine wrappers. +- Tool names, descriptions, argument names, and result keys. +- Existing examples and documented invocation patterns. + +## Rules for implementation + +1. Do not remove, rename, narrow, or relocate an existing public API. +2. Add overloads, builders, adapters, and new types instead of changing existing signatures. +3. Add only `default` methods to an existing public interface. +4. Keep old environment variables as aliases when introducing Python-compatible configuration. +5. Preserve old defaults unless an explicit opt-in selects new behavior. +6. Keep heavy or backend-specific dependencies out of the existing core artifact where possible. +7. Deprecations remain functional for the entire 0.x parity line; removal requires a future major + compatibility decision. +8. Exceptions may gain more context, but existing exception categories and success result shapes + must not change unexpectedly. + +## Automated gates + +The build will gain the following gates before broad feature work: + +- `japicmp` or Revapi comparison against the captured 0.0.2 API. +- Characterization tests for existing public constructors, environment lookup, JSON shapes, + tools, model streaming, memory, and knowledge-base behavior. +- Compilation of unchanged legacy examples. +- JDK 17 and 21 build matrix. +- Offline unit tests by default; external-service tests in an explicit integration profile. + +An implementation batch is complete only when unit, compatibility, and relevant integration +tests pass together. diff --git a/docs/parity/DECISIONS.md b/docs/parity/DECISIONS.md new file mode 100644 index 0000000..c32b144 --- /dev/null +++ b/docs/parity/DECISIONS.md @@ -0,0 +1,53 @@ +# Parity architecture decisions + +## D1: Port contracts, not Python syntax + +Python source and tests define observable behavior. Java implementations use Java conventions, +Google ADK Java types, reactive APIs already exposed by the project, and stable service-provider +interfaces. A matching directory tree is not considered parity by itself. + +## D2: Preserve the current core artifact + +`com.volcengine.veadk:veadk-java` remains the compatibility artifact. Core contracts and +lightweight implementations stay there. Integrations that require large database, browser, +realtime, or cloud SDK dependency graphs should be optional Maven modules. + +## D3: Introduce SPIs before backend implementations + +Knowledge-base, long-term-memory, short-term-memory, evaluation, tracing, authentication, and +runtime functionality receive stable Java interfaces first. Existing Viking services are adapted +to those interfaces without removing their current APIs. + +## D4: Separate deterministic and live tests + +The normal build must not require cloud credentials or network access. WireMock, fake clients, +in-memory implementations, and Testcontainers cover deterministic behavior. Credentialed smoke +tests run only through an opt-in integration profile. + +## D5: Treat Google ADK differences as an adapter boundary + +Python currently depends on Google ADK `>=1.34.0`, while Java pins Google ADK `0.4.0`. New VeADK +APIs must not expose avoidable version-specific internals. Any ADK upgrade must first pass the +captured public API and behavior tests. + +## D6: Deliver in reversible batches + +Each feature batch includes its contracts, implementation, tests, documentation, and parity +matrix update. No batch may depend on an untested later batch to restore a green build. + +## D7: Configuration precedence is explicit + +Java resolves process environment over `.env`, then flattened `config.yaml`, matching the Python +observable behavior. Typed access is additive; existing `EnvUtil` methods keep their signatures +and defaults. Diagnostic output always redacts credentials. + +## D8: VeADK Agent remains an ADK Agent + +The Java `Agent` extends Google ADK's `LlmAgent`, so it works wherever existing ADK APIs expect a +`BaseAgent`. VeADK capabilities are builder options and metadata contracts rather than a parallel +runtime hierarchy. + +## D9: Test instrumentation starts with the JVM + +Mockito inline mocking uses the Byte Buddy Java agent supplied at Surefire JVM startup. This +avoids runtime self-attachment, which is unavailable in some containers and managed CI sandboxes. diff --git a/docs/parity/PARITY_MATRIX.md b/docs/parity/PARITY_MATRIX.md new file mode 100644 index 0000000..b6cad3a --- /dev/null +++ b/docs/parity/PARITY_MATRIX.md @@ -0,0 +1,49 @@ +# Python to Java parity matrix + +Reference revisions are recorded in [BASELINE.md](BASELINE.md). Status and implementation notes +describe the current Java working tree relative to that pinned baseline. + +| Area | Python reference | Current Java implementation | Status | Priority | Required Java outcome | +| --- | --- | --- | --- | --- | --- | +| Compatibility/build | `pyproject.toml`, Python CI/tests | Reflection API guard, legacy-null overload compile guard, portable Mockito agent, Spotless, JaCoCo, 130 green tests across four modules and legacy example compilation | Partial | P0 | Add binary API comparison and JDK CI matrix | +| Configuration | `config.py`, `configs/*`, `.env`, `config.yaml` | Typed YAML/dotenv/environment loading, precedence, BytePlus aliases, defaults and secret redaction | Complete | P0 | Maintain aliases as new domains are added | +| Agent facade | `agent.py` | VeADK `Agent` extends `LlmAgent`; builder supports model, tools, sub-agents, memories, plugins, processors, skills and metadata | Partial | P0 | Add tracing/prompt interfaces and remaining Python opt-in features | +| Runner | `runner.py`, runner contract tests | Old constructors preserved; injectable services/plugins and run processors added; direct `new Runner(agent)` inherits a VeADK Agent's configured memory, plugins and processor | Partial | P0 | Add convenience streaming/session helpers and multimodal messages | +| Multi-agent primitives | `agents/{loop,parallel,sequential,supervise}_agent.py`, `flows/*` | Google ADK primitives only | Missing | P0 | VeADK wrappers and supervisor flows with contract tests | +| Agent metadata/search | `agent_metadata.py`, `agent_search.py` | Stable recursive metadata for agents, models, tools, components, skills, sources and topology; backend search execution pending | Partial | P0 | Add knowledge/memory search adapters after neutral SPIs | +| Ark LLM | `models/ark_llm.py` and context/fallback tests | Chat Completions adapter with effective primary-model fallback, streaming and parallel tool calls, tool-call history round-trip, generation parameters/logprobs, typed and raw JSON schemas, image/video input parts, finishReason and usage metadata | Partial | P0 | Preserve old constructors; align Responses API caching and remaining Chat Completions edge cases | +| Ark embedding | `models/ark_embedding.py` | Ark multimodal text embeddings with sync/reactive/batch APIs, dimensions, usage, model enum and credential fallback | Complete | P0 | Extend only when Python adds new input modalities/contracts | +| Knowledge-base facade | `knowledgebase/{knowledgebase,entry,types}.py` | Backend-neutral facade/SPI, text/file/directory ingestion, deterministic search and legacy service adapter | Complete | P0 | Maintain facade contracts while adding external backends | +| Knowledge-base backends | `knowledgebase/backends/*` | In-memory and existing Viking service adapter | Partial | P1 | Add Redis, Milvus, OpenSearch, OpenViking, TOS and context-search adapters | +| Long-term memory | `memory/long_term_memory.py`, backend SPI | ADK-compatible facade, event filtering/conversion, backend SPI, graceful search errors and legacy service adapter | Complete | P0 | Maintain facade contracts while adding external backends | +| Long-term-memory backends | `memory/long_term_memory_backends/*` | Deterministic per-user in-memory backend; Viking service also implements the new SPI | Partial | P1 | Add Mem0, OpenSearch, OpenViking, Redis and TOS adapters | +| Short-term memory | `memory/short_term_memory.py`, processor and backends | ADK-compatible facade/SPI, idempotent creation, load callback, history processor, canonical-history append, in-memory backend, service adapter and optional persistent SQLite module with stale-view protection | Partial | P0 | Add profile-based compaction; MySQL/PostgreSQL remain optional adapters | +| Save-session behavior | `memory/save_session_callback.py` | Python-compatible policy with first-save, thresholds and session-switch flush; per-user saves are serialized and awaited, Agent-bound memory is honored, while the zero-argument constructor preserves legacy fire-and-forget behavior | Complete | P0 | Maintain callback compatibility as memory contracts evolve | +| Built-in tools | `tools/builtin_tools/*` | Knowledge-base, web-search and run-code tools | Partial | P1 | Registry plus compatible implementations for applicable Python built-ins | +| MCP tools | `tools/mcp_tool/*` | Raw Google ADK MCP toolset used by sandbox | Partial | P1 | Trusted session manager/toolset, auth propagation, retry and lifecycle contracts | +| Sandbox tools | `tools/sandbox/*` | Code Sandbox MCP and AgentKit run-code | Partial | P1 | Code, browser and computer sandbox contracts with optional implementations | +| Skills | `skills/*`, `tools/skills_tools/*` | Missing | Missing | P1 | Parser, registry, materializer, file safety, toolset and checklist callback | +| Vanna/data tools | `tools/vanna_tools/*` | Missing | Missing | P2 | Optional data-analysis module and backend-neutral SQL/tool contracts | +| Authentication | `auth/*`, `integrations/ve_identity/*` | AK/SK passed directly to wrappers | Partial | P1 | Credential service, VeAuth providers, OAuth middleware and identity-aware tools | +| A2A | `a2a/*` | Missing | Missing | P1 | Agent card, executor, server, task store, registry client and middleware | +| A2UI | `a2ui/*` | Missing | Missing | P2 | Catalog and send-to-client toolset with optional dependency boundary | +| AgentKit application | `integrations/agentkit/app.py` | AgentKit tool invocation wrapper only | Partial | P1 | Application factory/server, health, metadata, topology and component summaries | +| AgentKit evaluation/feedback | `integrations/agentkit/evaluation/*` | Missing | Missing | P2 | Evaluation client, idempotent feedback and session capabilities | +| Evaluation | `evaluation/*` | Missing | Missing | P1 | Evaluator SPI, ADK evaluator, dataset loading/recording and Java-native metric adapters | +| Prompt management | `prompts/*`, PromptPilot integration | Missing | Missing | P1 | Prompt manager, evaluation, memory processor and optional PromptPilot adapter | +| Run processors | `processors/*` | Composable deferred event-stream processor, immutable run context and no-op default | Complete | P0 | Add domain-specific processors in their feature batches | +| Tracing/telemetry | `tracing/*` | OTel initialization, TLS and attribute rewriting | Partial | P1 | Preserve current API; add content policy, attributes, metrics and exporter SPI | +| Extensions/harness | `extensions/harness/*` | Missing | Missing | P2 | Extension lifecycle, invocation context, compaction and response verification | +| Feishu channel | `extensions/feishu_channel.py` | Missing | Missing | P2 | Runner/channel adapter with stable conversation/session mapping | +| Multimodal | `multimodal/*` | Missing | Missing | P2 | Attachment models, storage, transport, service and AgentKit routes | +| Realtime voice | `realtime/*` | Missing | Missing | P2 | Realtime protocol/client and Doubao voice model connection | +| Alternative runtimes | `runtime/{codex,piagent}/*` | Google ADK runtime only | Missing | P2 | Runtime SPI plus optional runtime bridges; default remains ADK | +| Tunnel | `tunnel/*` | Missing | Missing | P2 | Connector, registry, server protocol and MCP toolset | +| CLI/Studio | `cli/*`, `frontend/*`, `webui/*` | CLI example and Google ADK dev web dependency | Missing | P2 | Java CLI, generated project flow, runtime management APIs and reusable web assets | +| Cloud integrations | `cloud/*`, `integrations/ve_*` | Direct Volcengine wrappers for three services | Partial | P2 | Optional deployment/container/FaaS/pipeline/TOS/TLS integrations | +| Audio/dataset toolkits | `toolkits/*` | Missing | Missing | P2 | Optional ASR/TTS clients and dataset-generation callback | + +## Completion rule + +Every row must eventually be `Complete` or carry a documented, reviewed Java-specific exclusion. +Priority controls implementation order only; it does not remove P2 features from the parity goal. diff --git a/docs/parity/PROGRESS.md b/docs/parity/PROGRESS.md new file mode 100644 index 0000000..35732f7 --- /dev/null +++ b/docs/parity/PROGRESS.md @@ -0,0 +1,68 @@ +# Parity progress + +## Completed + +- [x] Clone and pin the Java and Python repositories. +- [x] Verify both worktrees started clean. +- [x] Create the `codex/python-parity` implementation branch. +- [x] Inventory Java source, tests, public APIs, Maven configuration and examples. +- [x] Inventory Python feature domains, dependencies and test domains. +- [x] Record baseline, compatibility policy, architecture decisions and initial parity matrix. +- [x] Establish a green baseline: 73 tests, three Maven modules, examples and packaging. +- [x] Add public API compatibility characterization tests. +- [x] Make Mockito tests portable to restricted CI environments. +- [x] Add typed `config.yaml`/`.env`/environment loading with BytePlus aliases and secret redaction. +- [x] Add a VeADK `Agent` facade over `LlmAgent` with memory, plugins and processor assembly. +- [x] Add injectable Runner services and a composable run-processor lifecycle. +- [x] Add recursive Agent metadata, tools, components, skills, search sources and topology. +- [x] Add Ark multimodal text embeddings with sync/reactive/batch APIs, dimensions and usage. +- [x] Add a backend-neutral KnowledgeBase facade, deterministic in-memory backend and legacy + service adapter. +- [x] Add a backend-neutral LongTermMemory facade, per-user in-memory backend and legacy service + adapter. +- [x] Make the existing Viking memory service implement the new backend SPI without changing its + original ADK API. +- [x] Add an ADK-compatible ShortTermMemory facade, backend SPI, idempotent session creation, + after-load callback, deterministic in-memory backend and existing-service adapter. +- [x] Add a provider-neutral short-term-memory history processor that rewrites loaded sessions + without mutating persisted events, while appending new events to canonical stored history. +- [x] Add an optional `veadk-memory-sqlite` module with complete ADK session CRUD, event append, + recent-event filtering, cross-instance persistence and stale-session overwrite protection. +- [x] Align save-session behavior with Python: first save, event/time thresholds, session-switch + flush, configurable error handling, per-user async serialization and automatic + `Agent.autoSaveSession(true)` callback wiring, while retaining the legacy zero-argument callback. +- [x] Add Ark LLM model fallback constructors, pre-output fallback retry, streaming no-switch + safety after partial output, parallel streaming tool-call assembly, and finishReason/usage + metadata propagation. +- [x] Map Ark LLM ADK generation config into Ark requests: temperature, topP, max output tokens, + stop sequences, penalties, candidate/logprob settings and typed/raw structured JSON formats. +- [x] Map Ark LLM ADK multimodal input parts into Ark chat content parts for inline/file image + and video inputs while preserving plain text request compatibility. +- [x] Map Ark LLM tool-call history into Ark assistant/tool messages and preserve returned tool-call + ids in ADK function-call responses. + +## In progress + +- [ ] Complete Ark LLM fallback, response-cache and multimodal behavior parity. +- [ ] Add an automated binary API comparison profile in addition to reflection contracts. + +## Latest validation + +- [x] `./mvnw -o -B -ntp clean verify` + +Result: 130 tests, zero failures/errors/skips; all four reactor modules, Spotless, JaCoCo, example +compilation, main JARs, sources JARs and Javadoc JARs passed with OpenJDK 17.0.17. + +## Planned batches + +1. Build and compatibility gates. +2. Typed configuration and environment compatibility. +3. Agent, Runner, processors and metadata. +4. Ark model and embedding parity. +5. Knowledge-base and memory SPIs plus deterministic backends. +6. External storage integrations. +7. Tools, MCP, sandbox and skills. +8. A2A, AgentKit application and evaluation. +9. Tracing, prompts and authentication. +10. A2UI, realtime, multimodal, tunnel, runtimes, CLI and cloud integrations. +11. Full regression, examples, documentation and release candidate. diff --git a/example/pom.xml b/example/pom.xml index a5fe5e7..c59fcd3 100644 --- a/example/pom.xml +++ b/example/pom.xml @@ -45,6 +45,7 @@ limitations under the License. org.sonatype.central central-publishing-maven-plugin + ${central.publishing.maven.version} true diff --git a/memory-sqlite/pom.xml b/memory-sqlite/pom.xml new file mode 100644 index 0000000..32c3104 --- /dev/null +++ b/memory-sqlite/pom.xml @@ -0,0 +1,54 @@ + + + + 4.0.0 + + + com.volcengine.veadk + veadk-parent + 0.0.2 + + + veadk-memory-sqlite + Volcengine Agent Development Kit SQLite Memory + Optional SQLite short-term-memory backend for VeADK Java + + + + com.volcengine.veadk + veadk-java + ${project.parent.version} + + + org.xerial + sqlite-jdbc + + + org.junit.jupiter + junit-jupiter-api + + + org.junit.jupiter + junit-jupiter-engine + + + org.assertj + assertj-core + + + diff --git a/memory-sqlite/src/main/java/com/volcengine/veadk/memory/sqlite/SQLiteSessionService.java b/memory-sqlite/src/main/java/com/volcengine/veadk/memory/sqlite/SQLiteSessionService.java new file mode 100644 index 0000000..3553beb --- /dev/null +++ b/memory-sqlite/src/main/java/com/volcengine/veadk/memory/sqlite/SQLiteSessionService.java @@ -0,0 +1,332 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory.sqlite; + +import com.google.adk.events.Event; +import com.google.adk.sessions.BaseSessionService; +import com.google.adk.sessions.GetSessionConfig; +import com.google.adk.sessions.ListEventsResponse; +import com.google.adk.sessions.ListSessionsResponse; +import com.google.adk.sessions.Session; +import io.reactivex.rxjava3.core.Completable; +import io.reactivex.rxjava3.core.Maybe; +import io.reactivex.rxjava3.core.Single; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; + +/** SQLite implementation of the Google ADK session-service contract. */ +public final class SQLiteSessionService implements BaseSessionService { + + private static final ConcurrentMap DATABASE_LOCKS = new ConcurrentHashMap<>(); + private static final String BEGIN_IMMEDIATE_SQL = "BEGIN IMMEDIATE"; + private static final String COMMIT_SQL = "COMMIT"; + private static final String ROLLBACK_SQL = "ROLLBACK"; + + private static final String CREATE_TABLE_SQL = + """ + CREATE TABLE IF NOT EXISTS veadk_sessions ( + app_name TEXT NOT NULL, + user_id TEXT NOT NULL, + session_id TEXT NOT NULL, + session_json TEXT NOT NULL, + updated_at INTEGER NOT NULL, + PRIMARY KEY (app_name, user_id, session_id) + ) + """; + + private final Path databasePath; + private final String jdbcUrl; + private final Object databaseLock; + + public SQLiteSessionService(Path databasePath) { + this.databasePath = + Objects.requireNonNull(databasePath, "databasePath").toAbsolutePath().normalize(); + this.databaseLock = + DATABASE_LOCKS.computeIfAbsent(this.databasePath, unused -> new Object()); + if (this.databasePath.getParent() == null) { + throw new IllegalArgumentException("databasePath must have a parent directory"); + } + try { + Files.createDirectories(this.databasePath.getParent()); + Class.forName("org.sqlite.JDBC"); + } catch (IOException | ClassNotFoundException exception) { + throw new IllegalStateException( + "Unable to initialize SQLite session storage", exception); + } + this.jdbcUrl = "jdbc:sqlite:" + this.databasePath; + initializeSchema(); + } + + public Path databasePath() { + return databasePath; + } + + @Override + public Single createSession( + String appName, String userId, ConcurrentMap state, String sessionId) { + return Single.fromCallable( + () -> { + String resolvedSessionId = + sessionId == null || sessionId.isBlank() + ? java.util.UUID.randomUUID().toString() + : sessionId; + Session session = + Session.builder(resolvedSessionId) + .appName(requireText(appName, "appName")) + .userId(requireText(userId, "userId")) + .state( + state == null + ? new ConcurrentHashMap<>() + : new ConcurrentHashMap<>(state)) + .lastUpdateTime(Instant.now()) + .build(); + save(session); + return session; + }); + } + + @Override + public Maybe getSession( + String appName, String userId, String sessionId, Optional config) { + return Maybe.fromCallable( + () -> { + Session session = load(appName, userId, sessionId); + return session == null || config.isEmpty() + ? session + : applyConfig(session, config.get()); + }); + } + + @Override + public Single listSessions(String appName, String userId) { + return Single.fromCallable( + () -> { + List sessions = new ArrayList<>(); + synchronized (databaseLock) { + try (Connection connection = openConnection(); + PreparedStatement statement = + connection.prepareStatement( + "SELECT session_json FROM veadk_sessions " + + "WHERE app_name = ? AND user_id = ? " + + "ORDER BY updated_at, session_id")) { + statement.setString(1, requireText(appName, "appName")); + statement.setString(2, requireText(userId, "userId")); + try (ResultSet results = statement.executeQuery()) { + while (results.next()) { + sessions.add(Session.fromJson(results.getString(1))); + } + } + } + } + return ListSessionsResponse.builder().sessions(sessions).build(); + }); + } + + @Override + public Completable deleteSession(String appName, String userId, String sessionId) { + return Completable.fromAction( + () -> { + synchronized (databaseLock) { + try (Connection connection = openConnection(); + PreparedStatement statement = + connection.prepareStatement( + "DELETE FROM veadk_sessions WHERE app_name = ? " + + "AND user_id = ? AND session_id = ?")) { + statement.setString(1, requireText(appName, "appName")); + statement.setString(2, requireText(userId, "userId")); + statement.setString(3, requireText(sessionId, "sessionId")); + statement.executeUpdate(); + } + } + }); + } + + @Override + public Single listEvents(String appName, String userId, String sessionId) { + return Single.fromCallable( + () -> { + Session session = load(appName, userId, sessionId); + List events = session == null ? List.of() : session.events(); + return ListEventsResponse.builder().events(events).build(); + }); + } + + @Override + public Single appendEvent(Session session, Event event) { + Objects.requireNonNull(session, "session"); + Objects.requireNonNull(event, "event"); + return Single.fromCallable( + () -> { + synchronized (databaseLock) { + try (Connection connection = openConnection()) { + execute(connection, BEGIN_IMMEDIATE_SQL); + try { + Session persisted = + load( + connection, + session.appName(), + session.userId(), + session.id()); + if (persisted == null) { + throw new IllegalStateException( + "Session not found: " + session.id()); + } + Event appendedEvent = + BaseSessionService.super + .appendEvent(persisted, event) + .blockingGet(); + BaseSessionService.super.appendEvent(session, event).blockingGet(); + save(connection, persisted); + execute(connection, COMMIT_SQL); + return appendedEvent; + } catch (Exception exception) { + rollback(connection, exception); + throw exception; + } + } + } + }); + } + + private void initializeSchema() { + synchronized (databaseLock) { + try (Connection connection = openConnection(); + Statement statement = connection.createStatement()) { + statement.execute(CREATE_TABLE_SQL); + } catch (SQLException exception) { + throw new IllegalStateException( + "Unable to create SQLite session schema", exception); + } + } + } + + private Connection openConnection() throws SQLException { + Connection connection = DriverManager.getConnection(jdbcUrl); + try { + execute(connection, "PRAGMA busy_timeout = 5000"); + return connection; + } catch (SQLException exception) { + connection.close(); + throw exception; + } + } + + private Session load(String appName, String userId, String sessionId) throws SQLException { + synchronized (databaseLock) { + try (Connection connection = openConnection()) { + return load(connection, appName, userId, sessionId); + } + } + } + + private Session load(Connection connection, String appName, String userId, String sessionId) + throws SQLException { + try (PreparedStatement statement = + connection.prepareStatement( + "SELECT session_json FROM veadk_sessions WHERE app_name = ? " + + "AND user_id = ? AND session_id = ?")) { + statement.setString(1, requireText(appName, "appName")); + statement.setString(2, requireText(userId, "userId")); + statement.setString(3, requireText(sessionId, "sessionId")); + try (ResultSet result = statement.executeQuery()) { + return result.next() ? Session.fromJson(result.getString(1)) : null; + } + } + } + + private void save(Session session) throws SQLException { + synchronized (databaseLock) { + try (Connection connection = openConnection()) { + save(connection, session); + } + } + } + + private void save(Connection connection, Session session) throws SQLException { + session.lastUpdateTime(Instant.now()); + try (PreparedStatement statement = + connection.prepareStatement( + "INSERT OR REPLACE INTO veadk_sessions (app_name, user_id, session_id," + + " session_json, updated_at) VALUES (?, ?, ?, ?, ?)")) { + statement.setString(1, session.appName()); + statement.setString(2, session.userId()); + statement.setString(3, session.id()); + statement.setString(4, session.toJson()); + statement.setLong(5, session.lastUpdateTime().toEpochMilli()); + statement.executeUpdate(); + } + } + + private static void execute(Connection connection, String sql) throws SQLException { + try (Statement statement = connection.createStatement()) { + statement.execute(sql); + } + } + + private static void rollback(Connection connection, Exception originalException) { + try { + execute(connection, ROLLBACK_SQL); + } catch (SQLException rollbackException) { + originalException.addSuppressed(rollbackException); + } + } + + private static Session applyConfig(Session session, GetSessionConfig config) { + List events = new ArrayList<>(session.events()); + config.afterTimestamp() + .ifPresent( + timestamp -> + events.removeIf( + event -> event.timestamp() <= timestamp.toEpochMilli())); + config.numRecentEvents() + .ifPresent( + limit -> { + if (limit < 0) { + throw new IllegalArgumentException( + "numRecentEvents must not be negative"); + } + if (events.size() > limit) { + events.subList(0, events.size() - limit).clear(); + } + }); + return Session.builder(session.id()) + .appName(session.appName()) + .userId(session.userId()) + .state(new ConcurrentHashMap<>(session.state())) + .events(events) + .lastUpdateTime(session.lastUpdateTime()) + .build(); + } + + private static String requireText(String value, String field) { + if (value == null || value.isBlank()) { + throw new IllegalArgumentException(field + " must not be blank"); + } + return value; + } +} diff --git a/memory-sqlite/src/main/java/com/volcengine/veadk/memory/sqlite/SQLiteShortTermMemoryBackend.java b/memory-sqlite/src/main/java/com/volcengine/veadk/memory/sqlite/SQLiteShortTermMemoryBackend.java new file mode 100644 index 0000000..9bbc35f --- /dev/null +++ b/memory-sqlite/src/main/java/com/volcengine/veadk/memory/sqlite/SQLiteShortTermMemoryBackend.java @@ -0,0 +1,45 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory.sqlite; + +import com.google.adk.sessions.BaseSessionService; +import com.volcengine.veadk.memory.ShortTermMemoryBackend; +import java.nio.file.Path; +import java.util.Objects; + +/** Optional SQLite backend for {@code ShortTermMemory}. */ +public final class SQLiteShortTermMemoryBackend implements ShortTermMemoryBackend { + + private final Path databasePath; + private final SQLiteSessionService sessionService; + + public SQLiteShortTermMemoryBackend(Path databasePath) { + this.databasePath = + Objects.requireNonNull(databasePath, "databasePath").toAbsolutePath().normalize(); + this.sessionService = new SQLiteSessionService(this.databasePath); + } + + public SQLiteShortTermMemoryBackend(String databasePath) { + this(Path.of(databasePath)); + } + + public Path databasePath() { + return databasePath; + } + + @Override + public BaseSessionService sessionService() { + return sessionService; + } +} diff --git a/memory-sqlite/src/test/java/com/volcengine/veadk/memory/sqlite/SQLiteShortTermMemoryBackendTest.java b/memory-sqlite/src/test/java/com/volcengine/veadk/memory/sqlite/SQLiteShortTermMemoryBackendTest.java new file mode 100644 index 0000000..8b359de --- /dev/null +++ b/memory-sqlite/src/test/java/com/volcengine/veadk/memory/sqlite/SQLiteShortTermMemoryBackendTest.java @@ -0,0 +1,188 @@ +/** + * Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates. + * + *

Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file + * except in compliance with the License. You may obtain a copy of the License at + * + *

http://www.apache.org/licenses/LICENSE-2.0 + * + *

Unless required by applicable law or agreed to in writing, software distributed under the + * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, + * either express or implied. See the License for the specific language governing permissions and + * limitations under the License. + */ +package com.volcengine.veadk.memory.sqlite; + +import static org.assertj.core.api.Assertions.assertThat; + +import com.google.adk.events.Event; +import com.google.adk.sessions.GetSessionConfig; +import com.google.adk.sessions.Session; +import com.google.genai.types.Content; +import com.google.genai.types.Part; +import com.volcengine.veadk.memory.ShortTermMemory; +import com.volcengine.veadk.memory.ShortTermMemoryMessage; +import com.volcengine.veadk.memory.ShortTermMemoryProcessor; +import java.nio.file.Path; +import java.time.Instant; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +class SQLiteShortTermMemoryBackendTest { + + @TempDir Path directory; + + @Test + void sessionsAndEventsPersistAcrossBackendInstances() { + Path database = directory.resolve("nested").resolve("sessions.db"); + SQLiteShortTermMemoryBackend firstBackend = new SQLiteShortTermMemoryBackend(database); + ShortTermMemory first = new ShortTermMemory(firstBackend); + Session session = + first.createSession( + "app", + "user", + new ConcurrentHashMap<>(Map.of("theme", "dark")), + "session") + .blockingGet(); + first.appendEvent(session, event("hello")).blockingGet(); + + SQLiteShortTermMemoryBackend secondBackend = new SQLiteShortTermMemoryBackend(database); + ShortTermMemory second = new ShortTermMemory(secondBackend); + Session restored = + second.getSession("app", "user", "session", Optional.empty()).blockingGet(); + + assertThat(database).exists(); + assertThat(restored.state()).containsEntry("theme", "dark"); + assertThat(restored.events()).extracting(Event::stringifyContent).containsExactly("hello"); + assertThat(second.listSessions("app", "user").blockingGet().sessionIds()) + .containsExactly("session"); + } + + @Test + void recentEventFilteringAndDeletionDoNotRewriteStoredHistory() { + SQLiteShortTermMemoryBackend backend = + new SQLiteShortTermMemoryBackend(directory.resolve("sessions.db")); + ShortTermMemory memory = new ShortTermMemory(backend); + Session session = memory.createSession("app", "user", "session").blockingGet(); + Event first = event("first"); + first.setTimestamp(Instant.now().minusSeconds(5).toEpochMilli()); + Event second = event("second"); + second.setTimestamp(Instant.now().toEpochMilli()); + memory.appendEvent(session, first).blockingGet(); + memory.appendEvent(session, second).blockingGet(); + GetSessionConfig recentOnly = GetSessionConfig.builder().numRecentEvents(1).build(); + + Session filtered = + memory.getSession("app", "user", "session", Optional.of(recentOnly)).blockingGet(); + + assertThat(filtered.events()).extracting(Event::stringifyContent).containsExactly("second"); + assertThat(memory.listEvents("app", "user", "session").blockingGet().events()).hasSize(2); + memory.deleteSession("app", "user", "session").blockingAwait(); + assertThat( + memory.getSession("app", "user", "session", Optional.empty()) + .isEmpty() + .blockingGet()) + .isTrue(); + } + + @Test + void staleSessionsFromDifferentBackendInstancesDoNotLoseEvents() { + Path database = directory.resolve("shared.db"); + ShortTermMemory first = new ShortTermMemory(new SQLiteShortTermMemoryBackend(database)); + ShortTermMemory second = new ShortTermMemory(new SQLiteShortTermMemoryBackend(database)); + Session firstView = first.createSession("app", "user", "session").blockingGet(); + Session staleSecondView = + second.getSession("app", "user", "session", Optional.empty()).blockingGet(); + + first.appendEvent(firstView, event("first")).blockingGet(); + second.appendEvent(staleSecondView, event("second")).blockingGet(); + + assertThat(first.listEvents("app", "user", "session").blockingGet().events()) + .extracting(Event::stringifyContent) + .containsExactly("first", "second"); + } + + @Test + void concurrentAppendsFromDifferentBackendInstancesAreSerialized() throws Exception { + Path database = directory.resolve("concurrent.db"); + ShortTermMemory first = new ShortTermMemory(new SQLiteShortTermMemoryBackend(database)); + ShortTermMemory second = new ShortTermMemory(new SQLiteShortTermMemoryBackend(database)); + Session firstView = first.createSession("app", "user", "session").blockingGet(); + Session secondView = + second.getSession("app", "user", "session", Optional.empty()).blockingGet(); + CountDownLatch ready = new CountDownLatch(2); + CountDownLatch start = new CountDownLatch(1); + ExecutorService executor = Executors.newFixedThreadPool(2); + + try { + Future firstAppend = + executor.submit( + () -> { + ready.countDown(); + start.await(); + return first.appendEvent(firstView, event("first")).blockingGet(); + }); + Future secondAppend = + executor.submit( + () -> { + ready.countDown(); + start.await(); + return second.appendEvent(secondView, event("second")) + .blockingGet(); + }); + + assertThat(ready.await(5, TimeUnit.SECONDS)).isTrue(); + start.countDown(); + firstAppend.get(5, TimeUnit.SECONDS); + secondAppend.get(5, TimeUnit.SECONDS); + } finally { + start.countDown(); + executor.shutdownNow(); + } + + assertThat(first.listEvents("app", "user", "session").blockingGet().events()) + .extracting(Event::stringifyContent) + .containsExactlyInAnyOrder("first", "second"); + } + + @Test + void processorViewAppendPreservesFullPersistedHistory() { + SQLiteShortTermMemoryBackend backend = + new SQLiteShortTermMemoryBackend(directory.resolve("processed.db")); + ShortTermMemoryProcessor processor = + new ShortTermMemoryProcessor( + messages -> List.of(new ShortTermMemoryMessage("user", "summary"))); + ShortTermMemory memory = new ShortTermMemory(backend, processor); + Session created = memory.createSession("app", "user", "session").blockingGet(); + memory.appendEvent(created, event("first")).blockingGet(); + memory.appendEvent(created, event("second")).blockingGet(); + Session processed = + memory.getSession("app", "user", "session", Optional.empty()).blockingGet(); + + memory.appendEvent(processed, event("third")).blockingGet(); + + assertThat( + backend.sessionService() + .listEvents("app", "user", "session") + .blockingGet() + .events()) + .extracting(Event::stringifyContent) + .containsExactly("first", "second", "third"); + } + + private static Event event(String text) { + return Event.builder() + .author("user") + .content(Content.builder().role("user").parts(List.of(Part.fromText(text))).build()) + .build(); + } +} diff --git a/pom.xml b/pom.xml index 8dd6df7..bb1dc94 100644 --- a/pom.xml +++ b/pom.xml @@ -50,6 +50,7 @@ limitations under the License. core + memory-sqlite example @@ -63,11 +64,14 @@ limitations under the License. 0.2.48 1.0.250 2.18.2 + 2.4 2.0.16 5.8.2 3.23.1 1.7.1 4.5.1 + 1.12.9 + 3.50.3.0 3.13.0 3.5.2 @@ -124,6 +128,16 @@ limitations under the License. jackson-core ${jackson.version} + + org.yaml + snakeyaml + ${snakeyaml.version} + + + org.xerial + sqlite-jdbc + ${sqlite-jdbc.version} + org.slf4j @@ -173,6 +187,12 @@ limitations under the License. ${mockito.version} test + + net.bytebuddy + byte-buddy-agent + ${byte-buddy.version} + test +