Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,9 @@
import org.jetbrains.annotations.Nullable;

import java.time.Instant;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
* A memory record in the system.
Expand Down Expand Up @@ -72,6 +74,17 @@ public class MemoryRecord {
@JsonProperty("event_date")
private Instant eventDate;

@NotNull
@JsonProperty("extraction_strategy")
private String extractionStrategy;

@NotNull
@JsonProperty("extraction_strategy_config")
private Map<String, Object> extractionStrategyConfig;

@NotNull
private Map<String, Object> metadata;

public MemoryRecord() {
this.id = UlidCreator.getUlid().toString();
Instant now = Instant.now();
Expand All @@ -80,6 +93,9 @@ public MemoryRecord() {
this.updatedAt = now;
this.discreteMemoryExtracted = "f";
this.memoryType = MemoryType.MESSAGE;
this.extractionStrategy = "discrete";
this.extractionStrategyConfig = new HashMap<>();
this.metadata = new HashMap<>();
}

public MemoryRecord(@NotNull String text) {
Expand Down Expand Up @@ -233,6 +249,33 @@ public void setEventDate(@Nullable Instant eventDate) {
this.eventDate = eventDate;
}

@NotNull
public String getExtractionStrategy() {
return extractionStrategy;
}

public void setExtractionStrategy(@NotNull String extractionStrategy) {
this.extractionStrategy = extractionStrategy;
}

@NotNull
public Map<String, Object> getExtractionStrategyConfig() {
return extractionStrategyConfig;
}

public void setExtractionStrategyConfig(@NotNull Map<String, Object> extractionStrategyConfig) {
this.extractionStrategyConfig = extractionStrategyConfig;
}

@NotNull
public Map<String, Object> getMetadata() {
return metadata;
}

public void setMetadata(@NotNull Map<String, Object> metadata) {
this.metadata = metadata;
}

@Override
public String toString() {
return "MemoryRecord{" +
Expand All @@ -251,6 +294,9 @@ public String toString() {
", memoryType=" + memoryType +
", extractedFrom=" + extractedFrom +
", eventDate=" + eventDate +
", extractionStrategy='" + extractionStrategy + '\'' +
", extractionStrategyConfig=" + extractionStrategyConfig +
", metadata=" + metadata +
'}';
}

Expand Down Expand Up @@ -282,6 +328,9 @@ public static class Builder {
private Instant persistedAt;
private List<String> extractedFrom;
private Instant eventDate;
private String extractionStrategy;
private Map<String, Object> extractionStrategyConfig;
private Map<String, Object> metadata;

private Builder() {
// Initialize with defaults for extracted memories (client-created long-term memories)
Expand All @@ -292,6 +341,9 @@ private Builder() {
this.updatedAt = now;
this.discreteMemoryExtracted = "t"; // "t" for extracted memories
this.memoryType = MemoryType.SEMANTIC; // SEMANTIC for long-term memories
this.extractionStrategy = "manual";
this.extractionStrategyConfig = new HashMap<>();
this.metadata = new HashMap<>();
}

/**
Expand All @@ -316,6 +368,9 @@ public Builder from(MemoryRecord record) {
this.persistedAt = record.persistedAt;
this.extractedFrom = record.extractedFrom;
this.eventDate = record.eventDate;
this.extractionStrategy = record.extractionStrategy;
this.extractionStrategyConfig = record.extractionStrategyConfig;
this.metadata = record.metadata;
return this;
}

Expand Down Expand Up @@ -479,6 +534,36 @@ public Builder eventDate(@Nullable Instant eventDate) {
return this;
}

/**
* Sets the extraction strategy.
* @param extractionStrategy the extraction strategy
* @return this builder
*/
public Builder extractionStrategy(@NotNull String extractionStrategy) {
this.extractionStrategy = extractionStrategy;
return this;
}

/**
* Sets the extraction strategy configuration.
* @param extractionStrategyConfig the extraction strategy configuration
* @return this builder
*/
public Builder extractionStrategyConfig(@NotNull Map<String, Object> extractionStrategyConfig) {
this.extractionStrategyConfig = extractionStrategyConfig;
return this;
}

/**
* Sets additional metadata.
* @param metadata the metadata map
* @return this builder
*/
public Builder metadata(@NotNull Map<String, Object> metadata) {
this.metadata = metadata;
return this;
}

/**
* Builds the MemoryRecord instance.
* @return a new MemoryRecord
Expand Down Expand Up @@ -506,6 +591,9 @@ public MemoryRecord build() {
record.persistedAt = this.persistedAt;
record.extractedFrom = this.extractedFrom;
record.eventDate = this.eventDate;
record.extractionStrategy = this.extractionStrategy;
record.extractionStrategyConfig = this.extractionStrategyConfig;
record.metadata = this.metadata;
return record;
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,10 @@ public class SearchRequest {
@JsonProperty("distance_threshold")
private Double distanceThreshold;

@Nullable
@JsonProperty("extraction_strategy")
private String extractionStrategy;

private int limit = 10;

private int offset = 0;
Expand Down Expand Up @@ -178,6 +182,15 @@ public void setDistanceThreshold(@Nullable Double distanceThreshold) {
this.distanceThreshold = distanceThreshold;
}

@Nullable
public String getExtractionStrategy() {
return extractionStrategy;
}

public void setExtractionStrategy(@Nullable String extractionStrategy) {
this.extractionStrategy = extractionStrategy;
}

public int getLimit() {
return limit;
}
Expand Down Expand Up @@ -279,6 +292,7 @@ public String toString() {
", entities=" + entities +
", userId='" + userId + '\'' +
", distanceThreshold=" + distanceThreshold +
", extractionStrategy='" + extractionStrategy + '\'' +
", limit=" + limit +
", offset=" + offset +
", recencyBoost=" + recencyBoost +
Expand Down Expand Up @@ -356,6 +370,11 @@ public Builder distanceThreshold(@Nullable Double distanceThreshold) {
return this;
}

public Builder extractionStrategy(@Nullable String extractionStrategy) {
request.extractionStrategy = extractionStrategy;
return this;
}

public Builder limit(int limit) {
request.limit = limit;
return this;
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,12 @@ public class SummaryViewPartitionResult {
@JsonProperty("memory_count")
private int memoryCount;

private boolean empty;

@Nullable
@JsonProperty("empty_reason")
private String emptyReason;

@Nullable
@JsonProperty("computed_at")
private String computedAt;
Expand Down Expand Up @@ -77,6 +83,23 @@ public void setMemoryCount(int memoryCount) {
this.memoryCount = memoryCount;
}

public boolean isEmpty() {
return empty;
}

public void setEmpty(boolean empty) {
this.empty = empty;
}

@Nullable
public String getEmptyReason() {
return emptyReason;
}

public void setEmptyReason(@Nullable String emptyReason) {
this.emptyReason = emptyReason;
}

@Nullable
public String getComputedAt() {
return computedAt;
Expand All @@ -91,6 +114,8 @@ public String toString() {
return "SummaryViewPartitionResult{" +
"viewId='" + viewId + '\'' +
", group=" + group +
", empty=" + empty +
", emptyReason='" + emptyReason + '\'' +
", memoryCount=" + memoryCount +
'}';
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -106,6 +106,9 @@ public MemoryRecordResults searchLongTermMemories(@NotNull SearchRequest request
if (request.getDistanceThreshold() != null) {
payload.put("distance_threshold", request.getDistanceThreshold());
}
if (request.getExtractionStrategy() != null) {
payload.put("extraction_strategy", Map.of("eq", request.getExtractionStrategy()));
}

// Add recency boost parameters if present
if (request.getRecencyBoost() != null) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -54,17 +54,25 @@ void testMemoryRecordSerialization() throws Exception {
record.setMemoryType(MemoryType.SEMANTIC);
record.setTopics(Arrays.asList("topic1", "topic2"));
record.setEntities(Arrays.asList("entity1", "entity2"));
record.setExtractionStrategy("summary");
record.setExtractionStrategyConfig(Map.of("summary_version", "v1"));
record.setMetadata(Map.of("message_count", 2));

String json = objectMapper.writeValueAsString(record);
assertNotNull(json);
assertTrue(json.contains("\"text\":\"Test memory\""));
assertTrue(json.contains("\"user_id\":\"user-123\""));
assertTrue(json.contains("\"memory_type\":\"semantic\""));
assertTrue(json.contains("\"extraction_strategy\":\"summary\""));
assertTrue(json.contains("\"metadata\""));

MemoryRecord deserialized = objectMapper.readValue(json, MemoryRecord.class);
assertEquals("Test memory", deserialized.getText());
assertEquals("user-123", deserialized.getUserId());
assertEquals(MemoryType.SEMANTIC, deserialized.getMemoryType());
assertEquals("summary", deserialized.getExtractionStrategy());
assertEquals("v1", deserialized.getExtractionStrategyConfig().get("summary_version"));
assertEquals(2, deserialized.getMetadata().get("message_count"));
assertNotNull(deserialized.getTopics());
assertEquals(2, deserialized.getTopics().size());
}
Expand Down Expand Up @@ -105,7 +113,6 @@ void testWorkingMemorySerialization() throws Exception {
assertNotNull(json);
assertTrue(json.contains("\"session_id\":\"session-123\""));
assertTrue(json.contains("\"user_id\":\"user-456\""));

WorkingMemory deserialized = objectMapper.readValue(json, WorkingMemory.class);
assertEquals("session-123", deserialized.getSessionId());
assertEquals("user-456", deserialized.getUserId());
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@
import org.junit.jupiter.api.Test;

import java.util.Arrays;
import java.util.Map;
import java.util.List;

import static org.junit.jupiter.api.Assertions.*;
Expand All @@ -19,6 +20,9 @@ void testDefaultConstructor() {
assertNotNull(record.getUpdatedAt());
assertEquals("f", record.getDiscreteMemoryExtracted());
assertEquals(MemoryType.MESSAGE, record.getMemoryType());
assertEquals("discrete", record.getExtractionStrategy());
assertNotNull(record.getExtractionStrategyConfig());
assertNotNull(record.getMetadata());
}

@Test
Expand All @@ -45,6 +49,9 @@ void testSettersAndGetters() {
record.setEntities(entities);

record.setMemoryType(MemoryType.SEMANTIC);
record.setExtractionStrategy("summary");
record.setExtractionStrategyConfig(Map.of("summary_version", "v1"));
record.setMetadata(Map.of("message_count", 2));

assertEquals("Test memory", record.getText());
assertEquals("session-123", record.getSessionId());
Expand All @@ -53,5 +60,8 @@ void testSettersAndGetters() {
assertEquals(topics, record.getTopics());
assertEquals(entities, record.getEntities());
assertEquals(MemoryType.SEMANTIC, record.getMemoryType());
assertEquals("summary", record.getExtractionStrategy());
assertEquals("v1", record.getExtractionStrategyConfig().get("summary_version"));
assertEquals(2, record.getMetadata().get("message_count"));
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -461,4 +461,23 @@ void testSearchRequestBuilder_AllRecencyFields() {
assertFalse(request.getServerSideRecency());
}

@Test
void testSearchRequestBuilder_ExtractionStrategy() throws Exception {
mockServer.enqueue(new MockResponse()
.setResponseCode(200)
.setHeader("Content-Type", "application/json")
.setBody("{\"memories\":[],\"total\":0}"));

SearchRequest request = SearchRequest.builder()
.text("query")
.extractionStrategy("summary")
.build();

client.longTermMemory().searchLongTermMemories(request);

RecordedRequest recorded = mockServer.takeRequest();
String requestBody = recorded.getBody().readUtf8();
assertTrue(requestBody.contains("\"extraction_strategy\":{\"eq\":\"summary\"}"));
}

}
Loading
Loading