streaming.WriteCompaction packages a 'compaction' event with the kept/older turn counts, summary tokens and provider-reported window/used tokens so the client can hint 'context optimized' to the user without parsing the stream body. streamChat runs Compact before BuildMessages and writes the event right after start, ensuring the client sees it before any chunk is emitted. Add a runner test that exercises limitRAGContext to keep the system prompt + RAG block under the configured window.
75 lines
No EOL
1.8 KiB
Go
75 lines
No EOL
1.8 KiB
Go
package streaming
|
|
|
|
import (
|
|
"encoding/json"
|
|
"fmt"
|
|
"net/http"
|
|
)
|
|
|
|
type Event struct {
|
|
Type string `json:"type"`
|
|
Content string `json:"content,omitempty"`
|
|
}
|
|
|
|
func WriteEvent(w http.ResponseWriter, eventType string, payload any) error {
|
|
data, err := json.Marshal(payload)
|
|
if err != nil {
|
|
return fmt.Errorf("marshal event: %w", err)
|
|
}
|
|
if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", eventType, data); err != nil {
|
|
return err
|
|
}
|
|
if f, ok := w.(http.Flusher); ok {
|
|
f.Flush()
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func WriteChunk(w http.ResponseWriter, content string) error {
|
|
return WriteEvent(w, "chunk", Event{Type: "chunk", Content: content})
|
|
}
|
|
|
|
func WriteStart(w http.ResponseWriter, conversationID string) error {
|
|
return WriteEvent(w, "start", map[string]any{
|
|
"type": "start",
|
|
"conversation_id": conversationID,
|
|
})
|
|
}
|
|
|
|
func WriteSources(w http.ResponseWriter, sources []string) error {
|
|
return WriteEvent(w, "sources", map[string]any{
|
|
"type": "sources",
|
|
"documents": sources,
|
|
})
|
|
}
|
|
|
|
// WriteCompaction surfaces a context-compaction event so the client can
|
|
// render a "summary injected" hint. Always emitted right after the
|
|
// `sources` event (before the streamed chunks), and only when compaction
|
|
// actually fired for this request.
|
|
func WriteCompaction(w http.ResponseWriter, payload map[string]any) error {
|
|
out := map[string]any{"type": "compaction"}
|
|
for k, v := range payload {
|
|
out[k] = v
|
|
}
|
|
return WriteEvent(w, "compaction", out)
|
|
}
|
|
|
|
type Usage struct {
|
|
InputTokens int `json:"input_tokens"`
|
|
OutputTokens int `json:"output_tokens"`
|
|
}
|
|
|
|
func WriteDone(w http.ResponseWriter, usage Usage) error {
|
|
return WriteEvent(w, "done", map[string]any{
|
|
"type": "done",
|
|
"usage": usage,
|
|
})
|
|
}
|
|
|
|
func WriteError(w http.ResponseWriter, errMsg string) error {
|
|
return WriteEvent(w, "error", map[string]string{
|
|
"type": "error",
|
|
"error": errMsg,
|
|
})
|
|
} |