Skip to content

Commit

Permalink
Add entities.
Browse files Browse the repository at this point in the history
  • Loading branch information
AlekSi committed Mar 3, 2019
1 parent c335d63 commit e81cd03
Show file tree
Hide file tree
Showing 3 changed files with 178 additions and 5 deletions.
4 changes: 2 additions & 2 deletions request.go
Original file line number Diff line number Diff line change
Expand Up @@ -34,8 +34,8 @@ type RequestPayload struct {
}

type RequestNLU struct {
Tokens []string `json:"tokens"`
Entities []interface{} `json:"entities"` // TODO
Tokens []string `json:"tokens"`
Entities []*Entity `json:"entities"`
}

type RequestSession struct {
Expand Down
148 changes: 148 additions & 0 deletions request_entity.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,148 @@
package alice

import (
"encoding/json"
"strconv"
)

type Entity struct {
Tokens struct {
Start int `json:"start"`
End int `json:"end"`
}
Type string `json:"type"`
Value interface{} `json:"value"`
}

type YandexFio struct {
FirstName string
PatronymicName string
LastName string
}

func (e *Entity) YandexFio() *YandexFio {
if e.Type != "YANDEX.FIO" {
return nil
}
v, _ := e.Value.(map[string]interface{})
if v == nil {
return nil
}

f, _ := v["first_name"].(string)
p, _ := v["patronymic_name"].(string)
l, _ := v["last_name"].(string)

return &YandexFio{
FirstName: f,
PatronymicName: p,
LastName: l,
}
}

type YandexGeo struct {
Country string
City string
Street string
HouseNumber string
Airport string
}

func (e *Entity) YandexGeo() *YandexGeo {
if e.Type != "YANDEX.GEO" {
return nil
}
v, _ := e.Value.(map[string]interface{})
if v == nil {
return nil
}

co, _ := v["country"].(string)
ci, _ := v["city"].(string)
s, _ := v["street"].(string)
h, _ := v["house_number"].(string)
a, _ := v["airport"].(string)

return &YandexGeo{
Country: co,
City: ci,
Street: s,
HouseNumber: h,
Airport: a,
}
}

type YandexDateTime struct {
Year int
YearIsRelative bool
Month int
MonthIsRelative bool
Day int
DayIsRelative bool
Hour int
HourIsRelative bool
Minute int
MinuteIsRelative bool
}

func (e *Entity) YandexDateTime() *YandexDateTime {
if e.Type != "YANDEX.DATETIME" {
return nil
}
v, _ := e.Value.(map[string]interface{})
if v == nil {
return nil
}

keys := []string{"year", "month", "day", "hour", "minute"}

// extract absolute values
abs := make(map[string]int)
for _, k := range keys {
switch v := v[k].(type) {
case json.Number:
i64, _ := v.Int64()
abs[k] = int(i64)
case float64:
abs[k] = int(v)
}
}

// extract relative flags
rel := make(map[string]bool)
for _, k := range keys {
rel[k], _ = v[k+"_is_relative"].(bool)
}

return &YandexDateTime{
Year: abs["year"],
YearIsRelative: rel["year"],
Month: abs["month"],
MonthIsRelative: rel["month"],
Day: abs["day"],
DayIsRelative: rel["day"],
Hour: abs["hour"],
HourIsRelative: rel["hour"],
Minute: abs["minute"],
MinuteIsRelative: rel["minute"],
}
}

type YandexNumber struct {
json.Number
}

func (e *Entity) YandexNumber() *YandexNumber {
if e.Type != "YANDEX.NUMBER" {
return nil
}

switch v := e.Value.(type) {
case json.Number:
return &YandexNumber{v}
case float64:
return &YandexNumber{json.Number(strconv.FormatFloat(v, 'f', -1, 64))}
default:
return nil
}
}
31 changes: 28 additions & 3 deletions request_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func TestRequestDecode(t *testing.T) {
Expand Down Expand Up @@ -97,8 +98,32 @@ func TestRequestDecode(t *testing.T) {
d := json.NewDecoder(bytes.NewReader(b))
d.DisallowUnknownFields()
var req Request
if err := d.Decode(&req); err != nil {
t.Fatal(err)
}
require.NoError(t, d.Decode(&req))
assert.True(t, req.Meta.HasScreen())
assert.Equal(t, []string{"закажи", "пиццу", "на", "льва", "толстого", "16", "на", "завтра"}, req.Request.NLU.Tokens)
require.Len(t, req.Request.NLU.Entities, 4)

geo := req.Request.NLU.Entities[0]
assert.Equal(t, 2, geo.Tokens.Start)
assert.Equal(t, 6, geo.Tokens.End)
assert.Equal(t, "YANDEX.GEO", geo.Type)
assert.Equal(t, &YandexGeo{Street: "льва толстого", HouseNumber: "16"}, geo.YandexGeo())

fio := req.Request.NLU.Entities[1]
assert.Equal(t, 3, fio.Tokens.Start)
assert.Equal(t, 5, fio.Tokens.End)
assert.Equal(t, "YANDEX.FIO", fio.Type)
assert.Equal(t, &YandexFio{FirstName: "лев", LastName: "толстой"}, fio.YandexFio())

num := req.Request.NLU.Entities[2]
assert.Equal(t, 5, num.Tokens.Start)
assert.Equal(t, 6, num.Tokens.End)
assert.Equal(t, "YANDEX.NUMBER", num.Type)
assert.Equal(t, "16", num.YandexNumber().String())

dt := req.Request.NLU.Entities[3]
assert.Equal(t, 6, dt.Tokens.Start)
assert.Equal(t, 8, dt.Tokens.End)
assert.Equal(t, "YANDEX.DATETIME", dt.Type)
assert.Equal(t, &YandexDateTime{Day: 1, DayIsRelative: true}, dt.YandexDateTime())
}

0 comments on commit e81cd03

Please sign in to comment.