Skip to content

Commit 5865099

Browse files
paul-paliychukpaul-paliychuk
committed
feat: Update context string utility to include entity attributes and episodes
1 parent 0008e55 commit 5865099

3 files changed

Lines changed: 239 additions & 6 deletions

File tree

‎.fernignore‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,4 +10,5 @@ graph/client/schema.go
1010
graph/client/ontology.go
1111
entity_types.go
1212
context_string.go
13+
context_string_test.go
1314
LICENSE

‎context_string.go‎

Lines changed: 87 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ const dateFormat = "2006-01-02 15:04:05"
1111

1212
// templateString defines the template for context information
1313
const templateString = `
14-
FACTS and ENTITIES represent relevant context to the current conversation.
14+
FACTS and ENTITIES%s represent relevant context to the current conversation.
1515
1616
# These are the most relevant facts and their valid date ranges
1717
# format: FACT (Date range: from - to)
@@ -20,10 +20,15 @@ FACTS and ENTITIES represent relevant context to the current conversation.
2020
</FACTS>
2121
2222
# These are the most relevant entities
23-
# ENTITY_NAME: entity summary
23+
# Name: ENTITY_NAME
24+
# Label: entity_label (if present)
25+
# Attributes: (if present)
26+
# attr_name: attr_value
27+
# Summary: entity summary
2428
<ENTITIES>
2529
%s
2630
</ENTITIES>
31+
%s
2732
`
2833

2934
// formatEdgeDateRange formats the date range of an entity edge.
@@ -46,8 +51,8 @@ func formatEdgeDateRange(edge *EntityEdge) string {
4651
return fmt.Sprintf("%s - %s", validAt, invalidAt)
4752
}
4853

49-
// ComposeContextString composes a search context from entity edges and nodes.
50-
func ComposeContextString(edges []*EntityEdge, nodes []*EntityNode) string {
54+
// ComposeContextString composes a search context from entity edges, nodes, and episodes.
55+
func ComposeContextString(edges []*EntityEdge, nodes []*EntityNode, episodes []*Episode) string {
5156
var facts []string
5257
for _, edge := range edges {
5358
fact := fmt.Sprintf(" - %s (%s)", edge.Fact, formatEdgeDateRange(edge))
@@ -56,12 +61,88 @@ func ComposeContextString(edges []*EntityEdge, nodes []*EntityNode) string {
5661

5762
var entities []string
5863
for _, node := range nodes {
59-
entity := fmt.Sprintf(" - %s: %s", node.Name, node.Summary)
64+
var entityParts []string
65+
entityParts = append(entityParts, fmt.Sprintf("Name: %s", node.Name))
66+
67+
// Add label if present (excluding 'Entity' from labels)
68+
if node.Labels != nil && len(node.Labels) > 0 {
69+
labels := make([]string, 0, len(node.Labels))
70+
for _, label := range node.Labels {
71+
if label != "Entity" {
72+
labels = append(labels, label)
73+
}
74+
}
75+
if len(labels) > 0 {
76+
entityParts = append(entityParts, fmt.Sprintf("Label: %s", labels[0]))
77+
}
78+
}
79+
80+
// Add attributes if present (excluding 'labels' attribute)
81+
if node.Attributes != nil && len(node.Attributes) > 0 {
82+
hasNonLabelAttributes := false
83+
for key := range node.Attributes {
84+
if key != "labels" {
85+
hasNonLabelAttributes = true
86+
break
87+
}
88+
}
89+
if hasNonLabelAttributes {
90+
entityParts = append(entityParts, "Attributes:")
91+
for key, value := range node.Attributes {
92+
if key != "labels" {
93+
entityParts = append(entityParts, fmt.Sprintf(" %s: %v", key, value))
94+
}
95+
}
96+
}
97+
}
98+
99+
// Add summary if present
100+
if node.Summary != "" {
101+
entityParts = append(entityParts, fmt.Sprintf("Summary: %s", node.Summary))
102+
}
103+
104+
entity := strings.Join(entityParts, "\n")
60105
entities = append(entities, entity)
61106
}
62107

108+
// Format episodes
109+
var episodesList []string
110+
if episodes != nil {
111+
for _, episode := range episodes {
112+
var rolePrefix string
113+
if episode.Role != nil && *episode.Role != "" {
114+
if episode.RoleType != nil && *episode.RoleType != "" {
115+
rolePrefix = fmt.Sprintf("%s (%s): ", *episode.Role, *episode.RoleType)
116+
} else {
117+
rolePrefix = fmt.Sprintf("%s: ", *episode.Role)
118+
}
119+
} else if episode.RoleType != nil && *episode.RoleType != "" {
120+
rolePrefix = fmt.Sprintf("(%s): ", *episode.RoleType)
121+
}
122+
123+
timestamp := "date unknown"
124+
if episode.CreatedAt != "" {
125+
if t, err := time.Parse(time.RFC3339, episode.CreatedAt); err == nil {
126+
timestamp = t.Format(dateFormat)
127+
}
128+
}
129+
130+
episodeStr := fmt.Sprintf(" - %s%s (%s)", rolePrefix, episode.Content, timestamp)
131+
episodesList = append(episodesList, episodeStr)
132+
}
133+
}
134+
63135
factsStr := strings.Join(facts, "\n")
64136
entitiesStr := strings.Join(entities, "\n")
137+
episodesStr := strings.Join(episodesList, "\n")
138+
139+
// Determine if episodes section should be included
140+
episodesHeader := ""
141+
episodesSection := ""
142+
if len(episodesList) > 0 {
143+
episodesHeader = ", and EPISODES"
144+
episodesSection = fmt.Sprintf("\n# These are the most relevant episodes\n<EPISODES>\n%s\n</EPISODES>", episodesStr)
145+
}
65146

66-
return fmt.Sprintf(templateString, factsStr, entitiesStr)
147+
return fmt.Sprintf(templateString, episodesHeader, factsStr, entitiesStr, episodesSection)
67148
}

‎context_string_test.go‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
package zep
2+
3+
import (
4+
"testing"
5+
6+
"github.com/stretchr/testify/assert"
7+
)
8+
9+
func TestFormatEdgeDateRange(t *testing.T) {
10+
t.Run("both dates present", func(t *testing.T) {
11+
validAt := "2023-01-01T10:00:00Z"
12+
invalidAt := "2023-01-02T15:00:00Z"
13+
edge := &EntityEdge{
14+
ValidAt: &validAt,
15+
InvalidAt: &invalidAt,
16+
}
17+
result := formatEdgeDateRange(edge)
18+
assert.Equal(t, "2023-01-01 10:00:00 - 2023-01-02 15:00:00", result)
19+
})
20+
21+
t.Run("only valid_at present", func(t *testing.T) {
22+
validAt := "2023-01-01T10:00:00Z"
23+
edge := &EntityEdge{
24+
ValidAt: &validAt,
25+
}
26+
result := formatEdgeDateRange(edge)
27+
assert.Equal(t, "2023-01-01 10:00:00 - present", result)
28+
})
29+
30+
t.Run("no dates present", func(t *testing.T) {
31+
edge := &EntityEdge{}
32+
result := formatEdgeDateRange(edge)
33+
assert.Equal(t, "date unknown - present", result)
34+
})
35+
}
36+
37+
func TestComposeContextString(t *testing.T) {
38+
t.Run("with facts, entities, and episodes", func(t *testing.T) {
39+
validAt := "2023-01-01T10:00:00Z"
40+
edges := []*EntityEdge{
41+
{
42+
Fact: "John likes coffee",
43+
ValidAt: &validAt,
44+
},
45+
}
46+
47+
labels := []string{"Person", "Entity"}
48+
attributes := map[string]interface{}{
49+
"age": 30,
50+
"labels": []string{"Person"},
51+
}
52+
nodes := []*EntityNode{
53+
{
54+
Name: "John",
55+
Labels: labels,
56+
Attributes: attributes,
57+
Summary: "A coffee enthusiast",
58+
},
59+
}
60+
61+
role := "user"
62+
roleType := RoleTypeUserRole
63+
episodes := []*Episode{
64+
{
65+
Role: &role,
66+
RoleType: &roleType,
67+
Content: "I love coffee",
68+
CreatedAt: "2023-01-01T12:00:00Z",
69+
},
70+
}
71+
72+
result := ComposeContextString(edges, nodes, episodes)
73+
74+
assert.Contains(t, result, "FACTS and ENTITIES, and EPISODES represent")
75+
assert.Contains(t, result, "John likes coffee (2023-01-01 10:00:00 - present)")
76+
assert.Contains(t, result, "Name: John")
77+
assert.Contains(t, result, "Label: Person")
78+
assert.Contains(t, result, "Attributes:")
79+
assert.Contains(t, result, " age: 30")
80+
assert.Contains(t, result, "Summary: A coffee enthusiast")
81+
assert.Contains(t, result, "user (user): I love coffee (2023-01-01 12:00:00)")
82+
assert.Contains(t, result, "<EPISODES>")
83+
})
84+
85+
t.Run("without episodes", func(t *testing.T) {
86+
validAt := "2023-01-01T10:00:00Z"
87+
edges := []*EntityEdge{
88+
{
89+
Fact: "John likes coffee",
90+
ValidAt: &validAt,
91+
},
92+
}
93+
94+
nodes := []*EntityNode{
95+
{
96+
Name: "John",
97+
Summary: "A person",
98+
},
99+
}
100+
101+
result := ComposeContextString(edges, nodes, nil)
102+
103+
assert.Contains(t, result, "FACTS and ENTITIES represent")
104+
assert.NotContains(t, result, ", and EPISODES")
105+
assert.NotContains(t, result, "<EPISODES>")
106+
assert.Contains(t, result, "John likes coffee")
107+
assert.Contains(t, result, "Name: John")
108+
assert.Contains(t, result, "Summary: A person")
109+
})
110+
111+
t.Run("entity with only Entity label filtered out", func(t *testing.T) {
112+
labels := []string{"Entity"}
113+
nodes := []*EntityNode{
114+
{
115+
Name: "Test",
116+
Labels: labels,
117+
Summary: "Test entity",
118+
},
119+
}
120+
121+
result := ComposeContextString(nil, nodes, nil)
122+
123+
assert.Contains(t, result, "Name: Test")
124+
assert.NotContains(t, result, "Label: Entity")
125+
assert.Contains(t, result, "Summary: Test entity")
126+
})
127+
128+
t.Run("episode with only role_type", func(t *testing.T) {
129+
roleType := RoleTypeAssistantRole
130+
episodes := []*Episode{
131+
{
132+
RoleType: &roleType,
133+
Content: "Hello there",
134+
CreatedAt: "2023-01-01T12:00:00Z",
135+
},
136+
}
137+
138+
result := ComposeContextString(nil, nil, episodes)
139+
140+
assert.Contains(t, result, "(assistant): Hello there")
141+
})
142+
143+
t.Run("empty inputs", func(t *testing.T) {
144+
result := ComposeContextString(nil, nil, nil)
145+
146+
assert.Contains(t, result, "FACTS and ENTITIES represent")
147+
assert.NotContains(t, result, "EPISODES")
148+
assert.Contains(t, result, "<FACTS>\n\n</FACTS>")
149+
assert.Contains(t, result, "<ENTITIES>\n\n</ENTITIES>")
150+
})
151+
}

0 commit comments

Comments
 (0)