content.go 6.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323
  1. package message
  2. import (
  3. "encoding/base64"
  4. "slices"
  5. "time"
  6. )
  7. type MessageRole string
  8. const (
  9. Assistant MessageRole = "assistant"
  10. User MessageRole = "user"
  11. System MessageRole = "system"
  12. Tool MessageRole = "tool"
  13. )
  14. type FinishReason string
  15. const (
  16. FinishReasonEndTurn FinishReason = "end_turn"
  17. FinishReasonMaxTokens FinishReason = "max_tokens"
  18. FinishReasonToolUse FinishReason = "tool_use"
  19. FinishReasonCanceled FinishReason = "canceled"
  20. FinishReasonError FinishReason = "error"
  21. FinishReasonPermissionDenied FinishReason = "permission_denied"
  22. // Should never happen
  23. FinishReasonUnknown FinishReason = "unknown"
  24. )
  25. type ContentPart interface {
  26. isPart()
  27. }
  28. type ReasoningContent struct {
  29. Thinking string `json:"thinking"`
  30. }
  31. func (tc ReasoningContent) String() string {
  32. return tc.Thinking
  33. }
  34. func (ReasoningContent) isPart() {}
  35. type TextContent struct {
  36. Text string `json:"text"`
  37. }
  38. func (tc *TextContent) String() string {
  39. if tc == nil {
  40. return ""
  41. }
  42. return tc.Text
  43. }
  44. func (TextContent) isPart() {}
  45. type ImageURLContent struct {
  46. URL string `json:"url"`
  47. Detail string `json:"detail,omitempty"`
  48. }
  49. func (iuc ImageURLContent) String() string {
  50. return iuc.URL
  51. }
  52. func (ImageURLContent) isPart() {}
  53. type BinaryContent struct {
  54. Path string
  55. MIMEType string
  56. Data []byte
  57. }
  58. func (bc BinaryContent) String(provider string) string {
  59. base64Encoded := base64.StdEncoding.EncodeToString(bc.Data)
  60. // if provider == models.ProviderOpenAI {
  61. // return "data:" + bc.MIMEType + ";base64," + base64Encoded
  62. // }
  63. return base64Encoded
  64. }
  65. func (BinaryContent) isPart() {}
  66. type ToolCall struct {
  67. ID string `json:"id"`
  68. Name string `json:"name"`
  69. Input string `json:"input"`
  70. Type string `json:"type"`
  71. Finished bool `json:"finished"`
  72. }
  73. func (ToolCall) isPart() {}
  74. type ToolResult struct {
  75. ToolCallID string `json:"tool_call_id"`
  76. Name string `json:"name"`
  77. Content string `json:"content"`
  78. Metadata string `json:"metadata"`
  79. IsError bool `json:"is_error"`
  80. }
  81. func (ToolResult) isPart() {}
  82. type Finish struct {
  83. Reason FinishReason `json:"reason"`
  84. Time time.Time `json:"time"`
  85. }
  86. type DBFinish struct {
  87. Reason FinishReason `json:"reason"`
  88. Time int64 `json:"time"`
  89. }
  90. func (Finish) isPart() {}
  91. func (m *Message) Content() *TextContent {
  92. for _, part := range m.Parts {
  93. if c, ok := part.(TextContent); ok {
  94. return &c
  95. }
  96. }
  97. return nil
  98. }
  99. func (m *Message) ReasoningContent() ReasoningContent {
  100. for _, part := range m.Parts {
  101. if c, ok := part.(ReasoningContent); ok {
  102. return c
  103. }
  104. }
  105. return ReasoningContent{}
  106. }
  107. func (m *Message) ImageURLContent() []ImageURLContent {
  108. imageURLContents := make([]ImageURLContent, 0)
  109. for _, part := range m.Parts {
  110. if c, ok := part.(ImageURLContent); ok {
  111. imageURLContents = append(imageURLContents, c)
  112. }
  113. }
  114. return imageURLContents
  115. }
  116. func (m *Message) BinaryContent() []BinaryContent {
  117. binaryContents := make([]BinaryContent, 0)
  118. for _, part := range m.Parts {
  119. if c, ok := part.(BinaryContent); ok {
  120. binaryContents = append(binaryContents, c)
  121. }
  122. }
  123. return binaryContents
  124. }
  125. func (m *Message) ToolCalls() []ToolCall {
  126. toolCalls := make([]ToolCall, 0)
  127. for _, part := range m.Parts {
  128. if c, ok := part.(ToolCall); ok {
  129. toolCalls = append(toolCalls, c)
  130. }
  131. }
  132. return toolCalls
  133. }
  134. func (m *Message) ToolResults() []ToolResult {
  135. toolResults := make([]ToolResult, 0)
  136. for _, part := range m.Parts {
  137. if c, ok := part.(ToolResult); ok {
  138. toolResults = append(toolResults, c)
  139. }
  140. }
  141. return toolResults
  142. }
  143. func (m *Message) IsFinished() bool {
  144. for _, part := range m.Parts {
  145. if _, ok := part.(Finish); ok {
  146. return true
  147. }
  148. }
  149. return false
  150. }
  151. func (m *Message) FinishPart() *Finish {
  152. for _, part := range m.Parts {
  153. if c, ok := part.(Finish); ok {
  154. return &c
  155. }
  156. }
  157. return nil
  158. }
  159. func (m *Message) FinishReason() FinishReason {
  160. for _, part := range m.Parts {
  161. if c, ok := part.(Finish); ok {
  162. return c.Reason
  163. }
  164. }
  165. return ""
  166. }
  167. func (m *Message) IsThinking() bool {
  168. if m.ReasoningContent().Thinking != "" && m.Content().Text == "" && !m.IsFinished() {
  169. return true
  170. }
  171. return false
  172. }
  173. func (m *Message) AppendContent(delta string) {
  174. found := false
  175. for i, part := range m.Parts {
  176. if c, ok := part.(TextContent); ok {
  177. m.Parts[i] = TextContent{Text: c.Text + delta}
  178. found = true
  179. }
  180. }
  181. if !found {
  182. m.Parts = append(m.Parts, TextContent{Text: delta})
  183. }
  184. }
  185. func (m *Message) AppendReasoningContent(delta string) {
  186. found := false
  187. for i, part := range m.Parts {
  188. if c, ok := part.(ReasoningContent); ok {
  189. m.Parts[i] = ReasoningContent{Thinking: c.Thinking + delta}
  190. found = true
  191. }
  192. }
  193. if !found {
  194. m.Parts = append(m.Parts, ReasoningContent{Thinking: delta})
  195. }
  196. }
  197. func (m *Message) FinishToolCall(toolCallID string) {
  198. for i, part := range m.Parts {
  199. if c, ok := part.(ToolCall); ok {
  200. if c.ID == toolCallID {
  201. m.Parts[i] = ToolCall{
  202. ID: c.ID,
  203. Name: c.Name,
  204. Input: c.Input,
  205. Type: c.Type,
  206. Finished: true,
  207. }
  208. return
  209. }
  210. }
  211. }
  212. }
  213. func (m *Message) AppendToolCallInput(toolCallID string, inputDelta string) {
  214. for i, part := range m.Parts {
  215. if c, ok := part.(ToolCall); ok {
  216. if c.ID == toolCallID {
  217. m.Parts[i] = ToolCall{
  218. ID: c.ID,
  219. Name: c.Name,
  220. Input: c.Input + inputDelta,
  221. Type: c.Type,
  222. Finished: c.Finished,
  223. }
  224. return
  225. }
  226. }
  227. }
  228. }
  229. func (m *Message) AddToolCall(tc ToolCall) {
  230. for i, part := range m.Parts {
  231. if c, ok := part.(ToolCall); ok {
  232. if c.ID == tc.ID {
  233. m.Parts[i] = tc
  234. return
  235. }
  236. }
  237. }
  238. m.Parts = append(m.Parts, tc)
  239. }
  240. func (m *Message) SetToolCalls(tc []ToolCall) {
  241. // remove any existing tool call part it could have multiple
  242. parts := make([]ContentPart, 0)
  243. for _, part := range m.Parts {
  244. if _, ok := part.(ToolCall); ok {
  245. continue
  246. }
  247. parts = append(parts, part)
  248. }
  249. m.Parts = parts
  250. for _, toolCall := range tc {
  251. m.Parts = append(m.Parts, toolCall)
  252. }
  253. }
  254. func (m *Message) AddToolResult(tr ToolResult) {
  255. m.Parts = append(m.Parts, tr)
  256. }
  257. func (m *Message) SetToolResults(tr []ToolResult) {
  258. for _, toolResult := range tr {
  259. m.Parts = append(m.Parts, toolResult)
  260. }
  261. }
  262. func (m *Message) AddFinish(reason FinishReason) {
  263. // remove any existing finish part
  264. for i, part := range m.Parts {
  265. if _, ok := part.(Finish); ok {
  266. m.Parts = slices.Delete(m.Parts, i, i+1)
  267. break
  268. }
  269. }
  270. m.Parts = append(m.Parts, Finish{Reason: reason, Time: time.Now()})
  271. }
  272. func (m *Message) AddImageURL(url, detail string) {
  273. m.Parts = append(m.Parts, ImageURLContent{URL: url, Detail: detail})
  274. }
  275. func (m *Message) AddBinary(mimeType string, data []byte) {
  276. m.Parts = append(m.Parts, BinaryContent{MIMEType: mimeType, Data: data})
  277. }