mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-08-15 05:04:27 +08:00
Implement OpenAI chat completions in GO (#16177)
### What problem does this PR solve? Implement OpenAI chat completions in GO POST /api/v1/openai/<chat_id>/chat/completions OpenAI chat cli: internal/development.md ### Type of change - [x] Refactoring
This commit is contained in:
qinling0210
@@ -49,7 +49,7 @@ var operatorMapping = map[string]string{
|
||||
">=": "≥",
|
||||
"<=": "≤",
|
||||
"!=": "≠",
|
||||
"==": "=",
|
||||
"==": "=",
|
||||
}
|
||||
|
||||
// ParseAndConvert converts raw API conditions into MetaFilterInput.
|
||||
@@ -76,10 +76,16 @@ func ParseAndConvert(metadataCondition map[string]interface{}) *MetaFilterInput
|
||||
continue
|
||||
}
|
||||
name, _ := cond["name"].(string)
|
||||
if name == "" {
|
||||
name, _ = cond["key"].(string) // OpenAI API metadata_condition uses "key"
|
||||
}
|
||||
if name == "" {
|
||||
continue
|
||||
}
|
||||
op, _ := cond["comparison_operator"].(string)
|
||||
if op == "" {
|
||||
op, _ = cond["operator"].(string) // OpenAI API uses "operator"
|
||||
}
|
||||
op = convertOperator(op)
|
||||
conditions = append(conditions, MetaCondition{
|
||||
Operator: op,
|
||||
|
||||
353
internal/common/multimodal.go
Normal file
353
internal/common/multimodal.go
Normal file
@@ -0,0 +1,353 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package common
|
||||
|
||||
import (
|
||||
"encoding/base64"
|
||||
"fmt"
|
||||
"regexp"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ContentPart is the internal representation of a multimodal content
|
||||
// fragment, decoupled from any provider's wire format. Drivers consume
|
||||
// the result of RenderContentPartsForFactory to produce their per-
|
||||
// provider JSON.
|
||||
type ContentPart struct {
|
||||
// Type is one of: "text", "image_url", "image", "inline_data".
|
||||
Type string
|
||||
// Text is set when Type == "text".
|
||||
Text string
|
||||
// ImageURL is set when Type == "image_url" (OpenAI shape).
|
||||
ImageURL *ImageURL
|
||||
// Source is set when Type == "image" (Anthropic) or
|
||||
// Type == "inline_data" (Gemini).
|
||||
Source *ContentSource
|
||||
}
|
||||
|
||||
// ImageURL is the OpenAI-shaped image reference.
|
||||
type ImageURL struct {
|
||||
URL string `json:"url"`
|
||||
}
|
||||
|
||||
// ContentSource is the Anthropic / Gemini source payload.
|
||||
type ContentSource struct {
|
||||
Type string `json:"type"` // "base64" or "url"
|
||||
MediaType string `json:"media_type"` // e.g. "image/png"
|
||||
Data string `json:"data,omitempty"` // base64 payload
|
||||
URL string `json:"url,omitempty"`
|
||||
}
|
||||
|
||||
// dataURIRE detects a "data:<mediatype>;base64,<data>" string.
|
||||
var dataURIRE = regexp.MustCompile(`^data:([^;,]+)(?:;base64)?,(.*)$`)
|
||||
|
||||
// parseDataURIOrB64 accepts a string and classifies it as a data URI,
|
||||
// a plain https URL, or a raw base64 payload
|
||||
func parseDataURIOrB64(s string) (ContentSource, error) {
|
||||
if s == "" {
|
||||
return ContentSource{}, fmt.Errorf("empty image source")
|
||||
}
|
||||
if m := dataURIRE.FindStringSubmatch(s); m != nil {
|
||||
mediaType := strings.TrimSpace(m[1])
|
||||
if mediaType == "" {
|
||||
mediaType = "image/png"
|
||||
}
|
||||
return ContentSource{
|
||||
Type: "base64",
|
||||
MediaType: mediaType,
|
||||
Data: m[2],
|
||||
}, nil
|
||||
}
|
||||
if strings.HasPrefix(s, "http://") || strings.HasPrefix(s, "https://") {
|
||||
return ContentSource{Type: "url", URL: s}, nil
|
||||
}
|
||||
// Assume raw base64 (no data URI, no http scheme). The provider
|
||||
// uses the file extension or a content-type hint from the call site
|
||||
// to pick the right media type; we default to image/png.
|
||||
if _, err := base64.StdEncoding.DecodeString(s); err != nil {
|
||||
return ContentSource{}, fmt.Errorf("not a valid data URI, URL, or base64: %w", err)
|
||||
}
|
||||
return ContentSource{
|
||||
Type: "base64",
|
||||
MediaType: "image/png",
|
||||
Data: s,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// normalizeTextFromContent extracts a single text string from a content
|
||||
// value that may be a string, []map[string]interface{}, or []interface{}.
|
||||
func normalizeTextFromContent(content interface{}) string {
|
||||
switch v := content.(type) {
|
||||
case string:
|
||||
return v
|
||||
case []map[string]interface{}:
|
||||
var parts []string
|
||||
for _, p := range v {
|
||||
if t, ok := p["type"].(string); ok && (t == "text" || t == "input_text") {
|
||||
if txt, ok := p["text"].(string); ok {
|
||||
parts = append(parts, txt)
|
||||
}
|
||||
} else if txt, ok := p["text"]; ok {
|
||||
// Fallback: "text" key present even though type didn't match.
|
||||
switch tv := txt.(type) {
|
||||
case string:
|
||||
parts = append(parts, tv)
|
||||
case float64:
|
||||
parts = append(parts, fmt.Sprintf("%v", tv))
|
||||
case int:
|
||||
parts = append(parts, fmt.Sprintf("%v", tv))
|
||||
}
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
case []interface{}:
|
||||
var parts []string
|
||||
for _, item := range v {
|
||||
switch p := item.(type) {
|
||||
case map[string]interface{}:
|
||||
if t, ok := p["type"].(string); ok && (t == "text" || t == "input_text") {
|
||||
if txt, ok := p["text"].(string); ok {
|
||||
parts = append(parts, txt)
|
||||
}
|
||||
} else if txt, ok := p["text"]; ok {
|
||||
// Fallback: "text" key present even though type didn't match.
|
||||
switch tv := txt.(type) {
|
||||
case string:
|
||||
parts = append(parts, tv)
|
||||
case float64:
|
||||
parts = append(parts, fmt.Sprintf("%v", tv))
|
||||
case int:
|
||||
parts = append(parts, fmt.Sprintf("%v", tv))
|
||||
}
|
||||
}
|
||||
case string:
|
||||
parts = append(parts, p)
|
||||
}
|
||||
}
|
||||
return strings.Join(parts, "\n")
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// extractImageURLs pulls image_url values out of a content value. Used
|
||||
// by ConvertLastUserMsgToMultimodal to assemble the ContentPart slice.
|
||||
func extractImageURLs(content interface{}) []string {
|
||||
var urls []string
|
||||
process := func(p map[string]interface{}) {
|
||||
t, _ := p["type"].(string)
|
||||
if t == "image_url" {
|
||||
if u, ok := p["image_url"].(string); ok && u != "" {
|
||||
urls = append(urls, u)
|
||||
} else if obj, ok := p["image_url"].(map[string]interface{}); ok {
|
||||
if u, ok := obj["url"].(string); ok && u != "" {
|
||||
urls = append(urls, u)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
switch v := content.(type) {
|
||||
case []map[string]interface{}:
|
||||
for _, p := range v {
|
||||
process(p)
|
||||
}
|
||||
case []interface{}:
|
||||
for _, item := range v {
|
||||
if p, ok := item.(map[string]interface{}); ok {
|
||||
process(p)
|
||||
}
|
||||
}
|
||||
}
|
||||
return urls
|
||||
}
|
||||
|
||||
// ConvertLastUserMsgToMultimodal converts a user message whose content
|
||||
// is a multimodal parts array into a message whose content is a
|
||||
// driver-ready content-parts value, dispatched by `factory` (provider
|
||||
// name).
|
||||
//
|
||||
// `imageAttachments` is an additional list of image URLs from the
|
||||
// `messages[-1]["files"]` array.
|
||||
// When non-empty, each URL is added to the content as an image
|
||||
// part regardless of the original message content.
|
||||
//
|
||||
// factory values supported:
|
||||
// - "gemini" → {"text": ...} / {"inline_data": {...}}
|
||||
// - "anthropic" → {"type": "text", ...} / {"type": "image", "source": {...}}
|
||||
// - default → {"type": "text", ...} / {"type": "image_url", "image_url": {...}}
|
||||
//
|
||||
// If the message is already a string, it is returned unchanged.
|
||||
// If the message has no image parts and `imageAttachments` is empty,
|
||||
// the text is returned as a string for compatibility with providers
|
||||
// that don't accept content arrays.
|
||||
func ConvertLastUserMsgToMultimodal(msg map[string]interface{}, imageAttachments []string, factory string) (map[string]interface{}, error) {
|
||||
if msg == nil {
|
||||
return nil, fmt.Errorf("nil message")
|
||||
}
|
||||
originalContent, ok := msg["content"]
|
||||
if !ok {
|
||||
return msg, nil
|
||||
}
|
||||
// If the content is already a plain string and there are no
|
||||
// imageAttachments to add, leave it alone.
|
||||
if _, isString := originalContent.(string); isString && len(imageAttachments) == 0 {
|
||||
return msg, nil
|
||||
}
|
||||
|
||||
// Combine images from the content array and from imageAttachments
|
||||
// (the `files` array on the last user message).
|
||||
// Order: content-array images first, then files-array images.
|
||||
textPart := normalizeTextFromContent(originalContent)
|
||||
imageURLs := extractImageURLs(originalContent)
|
||||
allImageURLs := append(imageURLs, imageAttachments...)
|
||||
if len(allImageURLs) == 0 {
|
||||
// No images — collapse to a string for compatibility.
|
||||
out := make(map[string]interface{}, len(msg))
|
||||
for k, v := range msg {
|
||||
out[k] = v
|
||||
}
|
||||
out["content"] = textPart
|
||||
return out, nil
|
||||
}
|
||||
|
||||
// Build ContentPart slice.
|
||||
parts := make([]ContentPart, 0, 1+len(allImageURLs))
|
||||
if textPart != "" {
|
||||
parts = append(parts, ContentPart{Type: "text", Text: textPart})
|
||||
}
|
||||
for _, u := range allImageURLs {
|
||||
src, err := parseDataURIOrB64(u)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("image_url %q: %w", u, err)
|
||||
}
|
||||
// OpenAI / default: pass the raw URL through (provider accepts
|
||||
// both data: and http(s):). Anthropic / Gemini need a Source.
|
||||
if factory == "anthropic" || factory == "gemini" {
|
||||
parts = append(parts, ContentPart{
|
||||
Type: pickImageType(factory),
|
||||
Source: &src,
|
||||
})
|
||||
} else {
|
||||
parts = append(parts, ContentPart{
|
||||
Type: "image_url",
|
||||
ImageURL: &ImageURL{URL: u},
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Render to the driver's wire format.
|
||||
rendered, err := RenderContentPartsForFactory(parts, factory)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out := make(map[string]interface{}, len(msg))
|
||||
for k, v := range msg {
|
||||
out[k] = v
|
||||
}
|
||||
out["content"] = rendered
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func pickImageType(factory string) string {
|
||||
if factory == "gemini" {
|
||||
return "inline_data"
|
||||
}
|
||||
return "image"
|
||||
}
|
||||
|
||||
// RenderContentPartsForFactory converts internal ContentPart values
|
||||
// into the per-provider JSON wire format:
|
||||
//
|
||||
// - gemini: [{"text": ...}, {"inline_data": {"mime_type": ..., "data": ...}}]
|
||||
// - anthropic: [{"type": "text", "text": ...}, {"type": "image", "source": {...}}]
|
||||
// - default: [{"type": "text", "text": ...}, {"type": "image_url", "image_url": {"url": ...}}]
|
||||
//
|
||||
// The return value is suitable for direct assignment to a Message's
|
||||
// `Content` field (`interface{}`).
|
||||
func RenderContentPartsForFactory(parts []ContentPart, factory string) (interface{}, error) {
|
||||
factory = strings.ToLower(factory)
|
||||
switch factory {
|
||||
case "gemini":
|
||||
out := make([]map[string]interface{}, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
switch p.Type {
|
||||
case "text":
|
||||
out = append(out, map[string]interface{}{"text": p.Text})
|
||||
case "image", "inline_data":
|
||||
if p.Source == nil {
|
||||
return nil, fmt.Errorf("gemini image part missing source")
|
||||
}
|
||||
if p.Source.Type == "url" {
|
||||
out = append(out, map[string]interface{}{
|
||||
"file_data": map[string]interface{}{
|
||||
"file_uri": p.Source.URL,
|
||||
"mime_type": p.Source.MediaType,
|
||||
},
|
||||
})
|
||||
} else {
|
||||
out = append(out, map[string]interface{}{
|
||||
"inline_data": map[string]interface{}{
|
||||
"mime_type": p.Source.MediaType,
|
||||
"data": p.Source.Data,
|
||||
},
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
case "anthropic":
|
||||
out := make([]map[string]interface{}, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
switch p.Type {
|
||||
case "text":
|
||||
out = append(out, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": p.Text,
|
||||
})
|
||||
case "image":
|
||||
if p.Source == nil {
|
||||
return nil, fmt.Errorf("anthropic image part missing source")
|
||||
}
|
||||
out = append(out, map[string]interface{}{
|
||||
"type": "image",
|
||||
"source": p.Source,
|
||||
})
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
default:
|
||||
// OpenAI-compatible.
|
||||
out := make([]map[string]interface{}, 0, len(parts))
|
||||
for _, p := range parts {
|
||||
switch p.Type {
|
||||
case "text":
|
||||
out = append(out, map[string]interface{}{
|
||||
"type": "text",
|
||||
"text": p.Text,
|
||||
})
|
||||
case "image_url":
|
||||
if p.ImageURL == nil {
|
||||
return nil, fmt.Errorf("openai image_url part missing URL")
|
||||
}
|
||||
out = append(out, map[string]interface{}{
|
||||
"type": "image_url",
|
||||
"image_url": p.ImageURL,
|
||||
})
|
||||
}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
}
|
||||
169
internal/common/timer.go
Normal file
169
internal/common/timer.go
Normal file
@@ -0,0 +1,169 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package common
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
)
|
||||
|
||||
// Phase is a named timing bucket in the RAG pipeline
|
||||
type Phase string
|
||||
|
||||
const (
|
||||
PhaseCheckLLM Phase = "check_llm"
|
||||
PhaseCheckLangfuse Phase = "check_langfuse"
|
||||
PhaseBindModels Phase = "bind_models"
|
||||
PhaseQueryRefinement Phase = "query_refinement"
|
||||
PhaseRetrieval Phase = "retrieval"
|
||||
PhaseGenerateAnswer Phase = "generate_answer"
|
||||
)
|
||||
|
||||
// allPhases ordered for Markdown() display.
|
||||
var allPhases = []Phase{
|
||||
PhaseCheckLLM,
|
||||
PhaseCheckLangfuse,
|
||||
PhaseBindModels,
|
||||
PhaseQueryRefinement,
|
||||
PhaseRetrieval,
|
||||
PhaseGenerateAnswer,
|
||||
}
|
||||
|
||||
// Timer tracks elapsed wall-clock time per named Phase.
|
||||
// Supports reentrant Enter/Exit on the same phase (inner span's duration
|
||||
// adds to the outer span's accumulated total).
|
||||
type Timer struct {
|
||||
mu sync.Mutex
|
||||
start time.Time
|
||||
phases map[Phase]time.Duration
|
||||
entries map[Phase][]time.Time
|
||||
}
|
||||
|
||||
// NewTimer constructs a Timer.
|
||||
func NewTimer() *Timer {
|
||||
return &Timer{
|
||||
phases: make(map[Phase]time.Duration, len(allPhases)),
|
||||
entries: make(map[Phase][]time.Time, len(allPhases)),
|
||||
}
|
||||
}
|
||||
|
||||
// Start anchors the timer. Calling Start() twice resets all state.
|
||||
func (t *Timer) Start() {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
t.start = time.Now()
|
||||
t.phases = make(map[Phase]time.Duration, len(allPhases))
|
||||
t.entries = make(map[Phase][]time.Time, len(allPhases))
|
||||
}
|
||||
|
||||
// Enter marks the start of phase p. Reentrant calls push a new anchor.
|
||||
func (t *Timer) Enter(p Phase) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
t.entries[p] = append(t.entries[p], time.Now())
|
||||
}
|
||||
|
||||
// Exit records the duration since the most recent Enter(p). No-op if no Enter.
|
||||
func (t *Timer) Exit(p Phase) {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
stack := t.entries[p]
|
||||
if len(stack) == 0 {
|
||||
return
|
||||
}
|
||||
open := stack[len(stack)-1]
|
||||
t.entries[p] = stack[:len(stack)-1]
|
||||
t.phases[p] += time.Since(open)
|
||||
}
|
||||
|
||||
// Phase returns the accumulated duration for phase p.
|
||||
func (t *Timer) Phase(p Phase) time.Duration {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
return t.phases[p]
|
||||
}
|
||||
|
||||
// Total returns the elapsed time since Start().
|
||||
func (t *Timer) Total() time.Duration {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
if t.start.IsZero() {
|
||||
return 0
|
||||
}
|
||||
return time.Since(t.start)
|
||||
}
|
||||
|
||||
// PhaseReport is the JSON-serializable view of a Timer's state.
|
||||
type PhaseReport struct {
|
||||
PhasesMs map[string]float64 `json:"phases_ms"`
|
||||
TotalMs float64 `json:"total_ms"`
|
||||
}
|
||||
|
||||
// Report returns a JSON-marshalable snapshot with microsecond precision.
|
||||
func (t *Timer) Report() *PhaseReport {
|
||||
t.mu.Lock()
|
||||
defer t.mu.Unlock()
|
||||
phases := make(map[string]float64, len(allPhases))
|
||||
for _, p := range allPhases {
|
||||
phases[string(p)] = float64(t.phases[p].Microseconds()) / 1000.0
|
||||
}
|
||||
var totalMs float64
|
||||
if !t.start.IsZero() {
|
||||
totalMs = float64(time.Since(t.start).Microseconds()) / 1000.0
|
||||
}
|
||||
return &PhaseReport{PhasesMs: phases, TotalMs: totalMs}
|
||||
}
|
||||
|
||||
func (t *Timer) MarshalJSON() ([]byte, error) {
|
||||
return json.Marshal(t.Report())
|
||||
}
|
||||
|
||||
// Markdown renders the Timer as a "## Time elapsed:" block matching
|
||||
func (t *Timer) Markdown() string {
|
||||
r := t.Report()
|
||||
var b strings.Builder
|
||||
b.WriteString("\n## Time elapsed:\n")
|
||||
b.WriteString(fmt.Sprintf(" - Total: %.1fms\n", r.TotalMs))
|
||||
for _, p := range allPhases {
|
||||
ms := r.PhasesMs[string(p)]
|
||||
b.WriteString(fmt.Sprintf(" - %s: %.1fms\n", displayName(p), ms))
|
||||
}
|
||||
b.WriteString("\n")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func displayName(p Phase) string {
|
||||
switch p {
|
||||
case PhaseCheckLLM:
|
||||
return "Check LLM"
|
||||
case PhaseCheckLangfuse:
|
||||
return "Check Langfuse tracer"
|
||||
case PhaseBindModels:
|
||||
return "Bind models"
|
||||
case PhaseQueryRefinement:
|
||||
return "Query refinement(LLM)"
|
||||
case PhaseRetrieval:
|
||||
return "Retrieval"
|
||||
case PhaseGenerateAnswer:
|
||||
return "Generate answer"
|
||||
default:
|
||||
return string(p)
|
||||
}
|
||||
}
|
||||
224
internal/common/timer_test.go
Normal file
224
internal/common/timer_test.go
Normal file
@@ -0,0 +1,224 @@
|
||||
//
|
||||
// Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
//
|
||||
|
||||
package common
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestTimer_BasicSequentialPhases(t *testing.T) {
|
||||
tm := NewTimer()
|
||||
tm.Start()
|
||||
|
||||
tm.Enter(PhaseCheckLLM)
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
tm.Exit(PhaseCheckLLM)
|
||||
|
||||
tm.Enter(PhaseBindModels)
|
||||
time.Sleep(3 * time.Millisecond)
|
||||
tm.Exit(PhaseBindModels)
|
||||
|
||||
got := tm.Phase(PhaseCheckLLM)
|
||||
if got < 4*time.Millisecond || got > 50*time.Millisecond {
|
||||
t.Errorf("PhaseCheckLLM = %v, want ~5ms", got)
|
||||
}
|
||||
got = tm.Phase(PhaseBindModels)
|
||||
if got < 2*time.Millisecond || got > 50*time.Millisecond {
|
||||
t.Errorf("PhaseBindModels = %v, want ~3ms", got)
|
||||
}
|
||||
|
||||
// Untouched phase should be 0.
|
||||
if d := tm.Phase(PhaseRetrieval); d != 0 {
|
||||
t.Errorf("PhaseRetrieval = %v, want 0", d)
|
||||
}
|
||||
|
||||
total := tm.Total()
|
||||
if total < 7*time.Millisecond {
|
||||
t.Errorf("Total = %v, want >= 7ms", total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimer_NestedPhasesAddUp(t *testing.T) {
|
||||
tm := NewTimer()
|
||||
tm.Start()
|
||||
|
||||
tm.Enter(PhaseQueryRefinement) // outer
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
tm.Enter(PhaseGenerateAnswer) // inner (LLM call inside pre-retrieval)
|
||||
time.Sleep(3 * time.Millisecond)
|
||||
tm.Exit(PhaseGenerateAnswer)
|
||||
time.Sleep(1 * time.Millisecond)
|
||||
tm.Exit(PhaseQueryRefinement)
|
||||
|
||||
// Generate answer records the inner 3ms.
|
||||
got := tm.Phase(PhaseGenerateAnswer)
|
||||
if got < 2*time.Millisecond || got > 50*time.Millisecond {
|
||||
t.Errorf("PhaseGenerateAnswer = %v, want ~3ms", got)
|
||||
}
|
||||
// Pre-retrieval processing records the WHOLE outer span (2 + 3 + 1 ≈ 6ms).
|
||||
got = tm.Phase(PhaseQueryRefinement)
|
||||
if got < 5*time.Millisecond || got > 50*time.Millisecond {
|
||||
t.Errorf("PhaseQueryRefinement = %v, want ~6ms (outer span)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimer_ExitWithoutEnterIsNoop(t *testing.T) {
|
||||
tm := NewTimer()
|
||||
tm.Start()
|
||||
// Should not panic, should not record anything.
|
||||
tm.Exit(PhaseRetrieval)
|
||||
if d := tm.Phase(PhaseRetrieval); d != 0 {
|
||||
t.Errorf("PhaseRetrieval = %v, want 0", d)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimer_StartResetsState(t *testing.T) {
|
||||
tm := NewTimer()
|
||||
tm.Start()
|
||||
tm.Enter(PhaseCheckLLM)
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
tm.Exit(PhaseCheckLLM)
|
||||
if tm.Phase(PhaseCheckLLM) == 0 {
|
||||
t.Fatal("precondition: phase must be non-zero before reset")
|
||||
}
|
||||
tm.Start()
|
||||
if d := tm.Phase(PhaseCheckLLM); d != 0 {
|
||||
t.Errorf("after Start, PhaseCheckLLM = %v, want 0", d)
|
||||
}
|
||||
if total := tm.Total(); total > 50*time.Millisecond {
|
||||
t.Errorf("after Start, Total = %v, want tiny", total)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimer_ConcurrentAccess(t *testing.T) {
|
||||
tm := NewTimer()
|
||||
tm.Start()
|
||||
var wg sync.WaitGroup
|
||||
for i := 0; i < 10; i++ {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
tm.Enter(PhaseRetrieval)
|
||||
time.Sleep(time.Millisecond)
|
||||
tm.Exit(PhaseRetrieval)
|
||||
}()
|
||||
}
|
||||
wg.Wait()
|
||||
got := tm.Phase(PhaseRetrieval)
|
||||
if got < 9*time.Millisecond {
|
||||
t.Errorf("PhaseRetrieval = %v, want ~10ms (10 parallel spans)", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimer_Report(t *testing.T) {
|
||||
tm := NewTimer()
|
||||
tm.Start()
|
||||
tm.Enter(PhaseCheckLLM)
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
tm.Exit(PhaseCheckLLM)
|
||||
tm.Enter(PhaseBindModels)
|
||||
time.Sleep(1 * time.Millisecond)
|
||||
tm.Exit(PhaseBindModels)
|
||||
|
||||
r := tm.Report()
|
||||
// Required fields
|
||||
if _, ok := r.PhasesMs[string(PhaseCheckLLM)]; !ok {
|
||||
t.Errorf("Report missing PhaseCheckLLM: %+v", r.PhasesMs)
|
||||
}
|
||||
if _, ok := r.PhasesMs[string(PhaseBindModels)]; !ok {
|
||||
t.Errorf("Report missing PhaseBindModels: %+v", r.PhasesMs)
|
||||
}
|
||||
if _, ok := r.PhasesMs[string(PhaseGenerateAnswer)]; !ok {
|
||||
t.Errorf("Report missing PhaseGenerateAnswer: %+v", r.PhasesMs)
|
||||
}
|
||||
if r.PhasesMs[string(PhaseCheckLLM)] < 1.0 {
|
||||
t.Errorf("Report PhaseCheckLLM_ms = %v, want >= 1.0", r.PhasesMs[string(PhaseCheckLLM)])
|
||||
}
|
||||
if r.TotalMs < 2.0 {
|
||||
t.Errorf("Report TotalMs = %v, want >= 2.0", r.TotalMs)
|
||||
}
|
||||
|
||||
// JSON round-trip
|
||||
b, err := json.Marshal(r)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal failed: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(b), `"phases_ms"`) || !strings.Contains(string(b), `"total_ms"`) {
|
||||
t.Errorf("JSON missing expected keys: %s", b)
|
||||
}
|
||||
|
||||
// Direct Marshal of the Timer
|
||||
b2, err := json.Marshal(tm)
|
||||
if err != nil {
|
||||
t.Fatalf("Marshal(Timer) failed: %v", err)
|
||||
}
|
||||
if !strings.Contains(string(b2), `"phases_ms"`) {
|
||||
t.Errorf("Timer JSON missing phases_ms: %s", b2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimer_Markdown(t *testing.T) {
|
||||
tm := NewTimer()
|
||||
tm.Start()
|
||||
tm.Enter(PhaseCheckLLM)
|
||||
time.Sleep(2 * time.Millisecond)
|
||||
tm.Exit(PhaseCheckLLM)
|
||||
tm.Enter(PhaseRetrieval)
|
||||
time.Sleep(5 * time.Millisecond)
|
||||
tm.Exit(PhaseRetrieval)
|
||||
tm.Enter(PhaseGenerateAnswer)
|
||||
time.Sleep(50 * time.Millisecond)
|
||||
tm.Exit(PhaseGenerateAnswer)
|
||||
|
||||
md := tm.Markdown()
|
||||
|
||||
// Should start with newline + "## Time elapsed:" header
|
||||
if !strings.HasPrefix(md, "\n## Time elapsed:") {
|
||||
t.Errorf("Markdown missing header: %q", md)
|
||||
}
|
||||
// Should contain all 6 phase labels
|
||||
for _, label := range []string{"Check LLM", "Check Langfuse tracer", "Bind models", "Query refinement(LLM)", "Retrieval", "Generate answer", "Total"} {
|
||||
if !strings.Contains(md, label+":") {
|
||||
t.Errorf("Markdown missing label %q: %q", label, md)
|
||||
}
|
||||
}
|
||||
// Phase durations should be numeric with "ms" suffix.
|
||||
mdRE := regexp.MustCompile(`(?m)^\s*-\s+([A-Za-z ()\.]+):\s+([0-9.]+)ms$`)
|
||||
matches := mdRE.FindAllStringSubmatch(md, -1)
|
||||
if len(matches) < 7 {
|
||||
t.Errorf("expected 7 phase lines, found %d in:\n%s", len(matches), md)
|
||||
}
|
||||
// Total should be the sum-ish of the three measured phases.
|
||||
totalRE := regexp.MustCompile(`Total:\s+([0-9.]+)ms`)
|
||||
totalMatch := totalRE.FindStringSubmatch(md)
|
||||
if len(totalMatch) < 2 {
|
||||
t.Fatalf("Markdown missing Total line: %q", md)
|
||||
}
|
||||
}
|
||||
|
||||
func TestTimer_TotalBeforeStart(t *testing.T) {
|
||||
tm := NewTimer()
|
||||
// No Start() called.
|
||||
if total := tm.Total(); total != 0 {
|
||||
t.Errorf("Total before Start = %v, want 0", total)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user