姚章浩
516 lines
16 KiB
Go
516 lines
16 KiB
Go
package weknora
|
|
|
|
import (
|
|
"bufio"
|
|
"bytes"
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"net/url"
|
|
"sort"
|
|
"strings"
|
|
"time"
|
|
|
|
"gitea.zheromiku.top/Weknora/dobotai-platform/internal/core"
|
|
)
|
|
|
|
var ErrIncompleteStream = errors.New("weknora stream ended without a terminal event")
|
|
var ErrEmptyAnswer = errors.New("weknora stream completed with an empty answer")
|
|
var ErrSessionNotFound = errors.New("weknora session not found")
|
|
|
|
// AnswerObserver is called with the accumulated full answer content after each
|
|
// answer chunk. The observer is purely a notification; its return is ignored.
|
|
type AnswerObserver func(fullContent string)
|
|
|
|
// StreamEvent is a safe, structured representation of a WeKnora streaming event.
|
|
// It deliberately does not expose raw Content, Data, or Reference content to
|
|
// prevent leaking internal thinking/tool details to external channels.
|
|
type StreamEvent struct {
|
|
ResponseType string
|
|
Done bool
|
|
FullAnswer string
|
|
ReferenceCount int
|
|
ToolName string
|
|
}
|
|
|
|
// StreamObserver is called with a safe StreamEvent for each SSE event.
|
|
type StreamObserver func(StreamEvent)
|
|
|
|
type Client struct {
|
|
baseURL string
|
|
apiKey string
|
|
http *http.Client
|
|
}
|
|
|
|
type AgentChatRequest struct {
|
|
Query string `json:"query"`
|
|
AgentEnabled bool `json:"agent_enabled"`
|
|
AgentID string `json:"agent_id"`
|
|
Channel string `json:"channel,omitempty"`
|
|
EnableMemory bool `json:"enable_memory,omitempty"`
|
|
DisableTitle bool `json:"disable_title,omitempty"`
|
|
KnowledgeBaseIDs []string `json:"knowledge_base_ids,omitempty"`
|
|
Images []Image `json:"images,omitempty"`
|
|
}
|
|
|
|
type Image struct {
|
|
Data string `json:"data"`
|
|
}
|
|
|
|
type ChatResult struct {
|
|
RequestID string
|
|
Answer string
|
|
References []core.Reference
|
|
}
|
|
|
|
func NewClient(baseURL, apiKey string, timeout time.Duration) (*Client, error) {
|
|
parsed, err := url.Parse(strings.TrimRight(strings.TrimSpace(baseURL), "/"))
|
|
if err != nil || parsed.Scheme == "" || parsed.Host == "" {
|
|
return nil, errors.New("invalid WEKNORA_BASE_URL")
|
|
}
|
|
if strings.TrimSpace(apiKey) == "" {
|
|
return nil, errors.New("WEKNORA_API_KEY is required")
|
|
}
|
|
if timeout <= 0 {
|
|
timeout = 120 * time.Second
|
|
}
|
|
return &Client{baseURL: parsed.String(), apiKey: apiKey, http: &http.Client{Timeout: timeout}}, nil
|
|
}
|
|
|
|
func (c *Client) CreateSession(ctx context.Context, title, description string) (string, error) {
|
|
var response struct {
|
|
Success bool `json:"success"`
|
|
Data struct {
|
|
ID string `json:"id"`
|
|
} `json:"data"`
|
|
}
|
|
if err := c.doJSON(ctx, http.MethodPost, "/api/v1/sessions", map[string]string{
|
|
"title": title, "description": description,
|
|
}, &response); err != nil {
|
|
return "", err
|
|
}
|
|
if !response.Success || strings.TrimSpace(response.Data.ID) == "" {
|
|
return "", errors.New("weknora returned an empty session id")
|
|
}
|
|
return response.Data.ID, nil
|
|
}
|
|
|
|
// UpdateSessionTitle updates only the readable title/description of an
|
|
// existing WeKnora session. The stable session ID is never changed.
|
|
func (c *Client) UpdateSessionTitle(ctx context.Context, sessionID, title, description string) error {
|
|
sessionID = strings.TrimSpace(sessionID)
|
|
title = strings.TrimSpace(title)
|
|
if sessionID == "" || title == "" {
|
|
return errors.New("session id and title are required")
|
|
}
|
|
var response struct {
|
|
Success bool `json:"success"`
|
|
}
|
|
if err := c.doJSON(ctx, http.MethodPut, "/api/v1/sessions/"+url.PathEscape(sessionID), map[string]string{
|
|
"title": title, "description": description,
|
|
}, &response); err != nil {
|
|
return err
|
|
}
|
|
if !response.Success {
|
|
return errors.New("weknora session update failed")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *Client) AgentChat(ctx context.Context, sessionID string, request AgentChatRequest) (ChatResult, error) {
|
|
return c.AgentChatStream(ctx, sessionID, request, nil)
|
|
}
|
|
|
|
func (c *Client) AgentChatStream(ctx context.Context, sessionID string, request AgentChatRequest, observer AnswerObserver) (ChatResult, error) {
|
|
return c.agentChat(ctx, sessionID, request, observer, nil)
|
|
}
|
|
|
|
func (c *Client) AgentChatEvents(ctx context.Context, sessionID string, request AgentChatRequest, observer StreamObserver) (ChatResult, error) {
|
|
return c.agentChat(ctx, sessionID, request, nil, observer)
|
|
}
|
|
|
|
func (c *Client) agentChat(ctx context.Context, sessionID string, request AgentChatRequest, answerObserver AnswerObserver, eventObserver StreamObserver) (ChatResult, error) {
|
|
if strings.TrimSpace(sessionID) == "" || strings.TrimSpace(request.AgentID) == "" || strings.TrimSpace(request.Query) == "" {
|
|
return ChatResult{}, errors.New("session id, agent id and query are required")
|
|
}
|
|
request.AgentEnabled = true
|
|
body, err := json.Marshal(request)
|
|
if err != nil {
|
|
return ChatResult{}, fmt.Errorf("encode agent chat request: %w", err)
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, http.MethodPost, c.baseURL+"/api/v1/agent-chat/"+url.PathEscape(sessionID), bytes.NewReader(body))
|
|
if err != nil {
|
|
return ChatResult{}, fmt.Errorf("create agent chat request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("Accept", "text/event-stream")
|
|
req.Header.Set("X-API-Key", c.apiKey)
|
|
resp, err := c.http.Do(req)
|
|
if err != nil {
|
|
return ChatResult{}, fmt.Errorf("call weknora agent chat: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
message, _ := io.ReadAll(io.LimitReader(resp.Body, 32<<10))
|
|
if isSessionNotFound(resp.StatusCode, message) {
|
|
return ChatResult{}, fmt.Errorf("%w: HTTP %d", ErrSessionNotFound, resp.StatusCode)
|
|
}
|
|
return ChatResult{}, fmt.Errorf("weknora agent chat returned HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(message)))
|
|
}
|
|
return parseSSEWithEventObserver(resp.Body, answerObserver, eventObserver)
|
|
}
|
|
|
|
func isSessionNotFound(status int, body []byte) bool {
|
|
if status != http.StatusNotFound {
|
|
return false
|
|
}
|
|
var response struct {
|
|
Error struct {
|
|
Code int `json:"code"`
|
|
Message string `json:"message"`
|
|
} `json:"error"`
|
|
}
|
|
if json.Unmarshal(body, &response) != nil {
|
|
return false
|
|
}
|
|
return response.Error.Code == 1003 && strings.EqualFold(strings.TrimSpace(response.Error.Message), "Session not found")
|
|
}
|
|
|
|
func (c *Client) doJSON(ctx context.Context, method, path string, input, output any) error {
|
|
body, err := json.Marshal(input)
|
|
if err != nil {
|
|
return fmt.Errorf("encode weknora request: %w", err)
|
|
}
|
|
req, err := http.NewRequestWithContext(ctx, method, c.baseURL+path, bytes.NewReader(body))
|
|
if err != nil {
|
|
return fmt.Errorf("create weknora request: %w", err)
|
|
}
|
|
req.Header.Set("Content-Type", "application/json")
|
|
req.Header.Set("X-API-Key", c.apiKey)
|
|
resp, err := c.http.Do(req)
|
|
if err != nil {
|
|
return fmt.Errorf("call weknora: %w", err)
|
|
}
|
|
defer resp.Body.Close()
|
|
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
|
return fmt.Errorf("weknora returned HTTP %d", resp.StatusCode)
|
|
}
|
|
if err := json.NewDecoder(io.LimitReader(resp.Body, 2<<20)).Decode(output); err != nil {
|
|
return fmt.Errorf("decode weknora response: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func parseSSE(reader io.Reader) (ChatResult, error) {
|
|
return parseSSEWithEventObserver(reader, nil, nil)
|
|
}
|
|
|
|
func parseSSEWithObserver(reader io.Reader, observer AnswerObserver) (ChatResult, error) {
|
|
return parseSSEWithEventObserver(reader, observer, nil)
|
|
}
|
|
|
|
func parseSSEWithEventObserver(reader io.Reader, answerObserver AnswerObserver, eventObserver StreamObserver) (ChatResult, error) {
|
|
scanner := bufio.NewScanner(reader)
|
|
scanner.Buffer(make([]byte, 64<<10), 2<<20)
|
|
var result ChatResult
|
|
var dataLines []string
|
|
terminal := false
|
|
|
|
notify := func(e StreamEvent) {
|
|
if eventObserver != nil {
|
|
eventObserver(e)
|
|
}
|
|
}
|
|
|
|
flush := func() error {
|
|
if len(dataLines) == 0 {
|
|
return nil
|
|
}
|
|
raw := strings.Join(dataLines, "\n")
|
|
dataLines = dataLines[:0]
|
|
var event struct {
|
|
ID string `json:"id"`
|
|
ResponseType string `json:"response_type"`
|
|
Content string `json:"content"`
|
|
Done bool `json:"done"`
|
|
KnowledgeReferences []core.Reference `json:"knowledge_references"`
|
|
Data json.RawMessage `json:"data"`
|
|
}
|
|
if err := json.Unmarshal([]byte(raw), &event); err != nil {
|
|
return fmt.Errorf("decode weknora SSE event: %w", err)
|
|
}
|
|
if event.ID != "" {
|
|
result.RequestID = event.ID
|
|
}
|
|
switch event.ResponseType {
|
|
case "references":
|
|
result.References = append([]core.Reference(nil), event.KnowledgeReferences...)
|
|
notify(StreamEvent{
|
|
ResponseType: "references",
|
|
ReferenceCount: len(event.KnowledgeReferences),
|
|
})
|
|
case "answer":
|
|
hasContent := strings.TrimSpace(event.Content) != ""
|
|
result.Answer += event.Content
|
|
if hasContent && answerObserver != nil {
|
|
answerObserver(result.Answer)
|
|
}
|
|
notify(StreamEvent{
|
|
ResponseType: "answer",
|
|
Done: event.Done,
|
|
FullAnswer: result.Answer,
|
|
})
|
|
if event.Done {
|
|
terminal = true
|
|
}
|
|
case "error":
|
|
if event.Content == "" {
|
|
event.Content = "weknora returned an unspecified error"
|
|
}
|
|
// WeKnora also uses response_type=error, done=false for a failed
|
|
// agent tool that the reasoning loop can recover from. Preserve the
|
|
// safe event signal, but only abort on an explicitly terminal error.
|
|
notify(StreamEvent{
|
|
ResponseType: "error",
|
|
Done: event.Done,
|
|
})
|
|
if event.Done {
|
|
return errors.New(event.Content)
|
|
}
|
|
case "thinking", "reflection":
|
|
notify(StreamEvent{
|
|
ResponseType: event.ResponseType,
|
|
Done: event.Done,
|
|
})
|
|
case "tool_call", "tool_result":
|
|
toolName := extractToolName(event.Data)
|
|
if event.ResponseType == "tool_result" {
|
|
result.References = appendUniqueReferences(result.References, extractToolResultReferences(event.Data))
|
|
}
|
|
notify(StreamEvent{
|
|
ResponseType: event.ResponseType,
|
|
Done: event.Done,
|
|
ToolName: toolName,
|
|
})
|
|
case "complete":
|
|
if event.Done {
|
|
terminal = true
|
|
}
|
|
notify(StreamEvent{
|
|
ResponseType: "complete",
|
|
Done: event.Done,
|
|
})
|
|
default:
|
|
notify(StreamEvent{
|
|
ResponseType: event.ResponseType,
|
|
Done: event.Done,
|
|
})
|
|
}
|
|
return nil
|
|
}
|
|
|
|
for scanner.Scan() {
|
|
line := scanner.Text()
|
|
if line == "" {
|
|
if err := flush(); err != nil {
|
|
return ChatResult{}, err
|
|
}
|
|
continue
|
|
}
|
|
if strings.HasPrefix(line, "data:") {
|
|
dataLines = append(dataLines, strings.TrimSpace(strings.TrimPrefix(line, "data:")))
|
|
}
|
|
}
|
|
if err := scanner.Err(); err != nil {
|
|
return ChatResult{}, fmt.Errorf("read weknora SSE stream: %w", err)
|
|
}
|
|
if err := flush(); err != nil {
|
|
return ChatResult{}, err
|
|
}
|
|
if !terminal {
|
|
return ChatResult{}, ErrIncompleteStream
|
|
}
|
|
if strings.TrimSpace(result.Answer) == "" {
|
|
return ChatResult{}, ErrEmptyAnswer
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
// extractToolName extracts the tool_name string from a JSON Data field.
|
|
// The Data may be any JSON shape; we look for the "tool_name" key and return
|
|
// its string value. Returns empty string if not found or not a string.
|
|
func extractToolName(data json.RawMessage) string {
|
|
if len(data) == 0 {
|
|
return ""
|
|
}
|
|
var m map[string]json.RawMessage
|
|
if err := json.Unmarshal(data, &m); err != nil {
|
|
return ""
|
|
}
|
|
v, ok := m["tool_name"]
|
|
if !ok {
|
|
v, ok = m["name"]
|
|
if !ok {
|
|
return ""
|
|
}
|
|
}
|
|
var s string
|
|
if err := json.Unmarshal(v, &s); err != nil {
|
|
return ""
|
|
}
|
|
return s
|
|
}
|
|
|
|
func extractToolResultReferences(data json.RawMessage) []core.Reference {
|
|
var envelope struct {
|
|
ToolName string `json:"tool_name"`
|
|
Name string `json:"name"`
|
|
Success bool `json:"success"`
|
|
Data json.RawMessage `json:"data"`
|
|
Result *struct {
|
|
Success bool `json:"success"`
|
|
Data json.RawMessage `json:"data"`
|
|
} `json:"result"`
|
|
}
|
|
if json.Unmarshal(data, &envelope) != nil {
|
|
return nil
|
|
}
|
|
success := envelope.Success
|
|
payloadData := envelope.Data
|
|
if envelope.Result != nil {
|
|
success = envelope.Result.Success
|
|
payloadData = envelope.Result.Data
|
|
} else if len(payloadData) == 0 || string(payloadData) == "null" {
|
|
// Live WeKnora SSE flattens sanitized tool result fields into the
|
|
// event data object alongside tool_name and success. Older fixtures
|
|
// and SDK variants wrap the result fields in data or result.data.
|
|
payloadData = data
|
|
}
|
|
if !success || len(payloadData) == 0 {
|
|
return nil
|
|
}
|
|
toolName := strings.TrimSpace(envelope.ToolName)
|
|
if toolName == "" {
|
|
toolName = strings.TrimSpace(envelope.Name)
|
|
}
|
|
|
|
var payload struct {
|
|
Count int `json:"count"`
|
|
DocumentCount int `json:"document_count"`
|
|
FetchedChunks int `json:"fetched_chunks"`
|
|
TotalDocs int `json:"total_docs"`
|
|
KnowledgeID string `json:"knowledge_id"`
|
|
KnowledgeTitle string `json:"knowledge_title"`
|
|
Results []struct {
|
|
KnowledgeID string `json:"knowledge_id"`
|
|
KnowledgeTitle string `json:"knowledge_title"`
|
|
} `json:"results"`
|
|
KnowledgeResults []struct {
|
|
KnowledgeID string `json:"knowledge_id"`
|
|
KnowledgeTitle string `json:"knowledge_title"`
|
|
ChunkHitCount int `json:"chunk_hit_count"`
|
|
} `json:"knowledge_results"`
|
|
Documents []struct {
|
|
KnowledgeID string `json:"knowledge_id"`
|
|
Title string `json:"title"`
|
|
ChunkCount int `json:"chunk_count"`
|
|
} `json:"documents"`
|
|
FoundKBs map[string][]string `json:"found_kbs"`
|
|
}
|
|
if json.Unmarshal(payloadData, &payload) != nil {
|
|
return nil
|
|
}
|
|
|
|
references := make([]core.Reference, 0)
|
|
switch toolName {
|
|
case "knowledge_search":
|
|
if payload.Count <= 0 && len(payload.Results) == 0 {
|
|
return nil
|
|
}
|
|
for _, item := range payload.Results {
|
|
if id := strings.TrimSpace(item.KnowledgeID); id != "" {
|
|
references = append(references, core.Reference{KnowledgeID: id, KnowledgeTitle: strings.TrimSpace(item.KnowledgeTitle)})
|
|
}
|
|
}
|
|
case "grep_chunks":
|
|
if payload.DocumentCount <= 0 && len(payload.KnowledgeResults) == 0 {
|
|
return nil
|
|
}
|
|
for _, item := range payload.KnowledgeResults {
|
|
if id := strings.TrimSpace(item.KnowledgeID); id != "" && item.ChunkHitCount > 0 {
|
|
references = append(references, core.Reference{KnowledgeID: id, KnowledgeTitle: strings.TrimSpace(item.KnowledgeTitle)})
|
|
}
|
|
}
|
|
case "list_knowledge_chunks":
|
|
if payload.FetchedChunks > 0 && strings.TrimSpace(payload.KnowledgeID) != "" {
|
|
references = append(references, core.Reference{
|
|
KnowledgeID: strings.TrimSpace(payload.KnowledgeID), KnowledgeTitle: strings.TrimSpace(payload.KnowledgeTitle),
|
|
})
|
|
}
|
|
case "get_document_info":
|
|
if payload.TotalDocs <= 0 && len(payload.Documents) == 0 {
|
|
return nil
|
|
}
|
|
for _, item := range payload.Documents {
|
|
if id := strings.TrimSpace(item.KnowledgeID); id != "" && item.ChunkCount > 0 {
|
|
references = append(references, core.Reference{KnowledgeID: id, KnowledgeTitle: strings.TrimSpace(item.Title)})
|
|
}
|
|
}
|
|
case "wiki_search", "wiki_read_page":
|
|
slugs := make([]string, 0, len(payload.FoundKBs))
|
|
for slug := range payload.FoundKBs {
|
|
slugs = append(slugs, slug)
|
|
}
|
|
sort.Strings(slugs)
|
|
for _, slug := range slugs {
|
|
kbIDs := append([]string(nil), payload.FoundKBs[slug]...)
|
|
sort.Strings(kbIDs)
|
|
for _, kbID := range kbIDs {
|
|
if strings.TrimSpace(slug) == "" || strings.TrimSpace(kbID) == "" {
|
|
continue
|
|
}
|
|
references = append(references, core.Reference{
|
|
ID: "wiki:" + strings.TrimSpace(kbID) + ":" + strings.TrimSpace(slug),
|
|
KnowledgeTitle: strings.TrimSpace(slug),
|
|
})
|
|
}
|
|
}
|
|
}
|
|
return references
|
|
}
|
|
|
|
func appendUniqueReferences(current, incoming []core.Reference) []core.Reference {
|
|
result := append([]core.Reference(nil), current...)
|
|
seen := make(map[string]struct{}, len(result)+len(incoming))
|
|
for _, reference := range result {
|
|
if key := referenceKey(reference); key != "" {
|
|
seen[key] = struct{}{}
|
|
}
|
|
}
|
|
for _, reference := range incoming {
|
|
key := referenceKey(reference)
|
|
if key == "" {
|
|
continue
|
|
}
|
|
if _, exists := seen[key]; exists {
|
|
continue
|
|
}
|
|
seen[key] = struct{}{}
|
|
result = append(result, reference)
|
|
}
|
|
return result
|
|
}
|
|
|
|
func referenceKey(reference core.Reference) string {
|
|
if value := strings.TrimSpace(reference.KnowledgeID); value != "" {
|
|
return "knowledge:" + value
|
|
}
|
|
if value := strings.TrimSpace(reference.ID); value != "" {
|
|
return "id:" + value
|
|
}
|
|
return ""
|
|
}
|