ai-operator/internal/kb/jsonl_parser.go

package kb

import (
	"bufio"
	"bytes"
	"context"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"strings"

	"ai-operator/internal/dialogue/region"
)

var requiredFields = []string{"external_id", "title", "language", "region_code", "status", "content", "source_hash"}

func ParseJSONLFile(path string, allowDisabled bool) ([]Document, []string, error) {
	f, err := os.Open(path)
	if err != nil {
		return nil, nil, err
	}
	defer f.Close()
	return ParseJSONL(f, path, allowDisabled)
}

func ParseJSONL(r io.Reader, sourcePath string, allowDisabled bool) ([]Document, []string, error) {
	resolver := region.NewDefaultResolver()
	sc := bufio.NewScanner(r)
	sc.Buffer(make([]byte, 1024), 10*1024*1024)
	var docs []Document
	var errs []string
	line := 0
	for sc.Scan() {
		line++
		raw := bytes.TrimSpace(sc.Bytes())
		if len(raw) == 0 {
			continue
		}
		var m map[string]any
		if err := json.Unmarshal(raw, &m); err != nil {
			errs = append(errs, fmt.Sprintf("%s:%d invalid json", sourcePath, line))
			continue
		}
		missing := ""
		for _, k := range requiredFields {
			if strings.TrimSpace(asString(m[k])) == "" {
				missing = k
				break
			}
		}
		if missing != "" {
			errs = append(errs, fmt.Sprintf("%s:%d missing %s", sourcePath, line, missing))
			continue
		}
		lang := asString(m["language"])
		if lang != "ru" && lang != "kk" {
			errs = append(errs, fmt.Sprintf("%s:%d invalid language", sourcePath, line))
			continue
		}
		regCode := asString(m["region_code"])
		if regCode != "global" {
			reg, ok := resolver.GetByCode(regCode)
			if !ok || (!reg.Enabled && !allowDisabled) {
				errs = append(errs, fmt.Sprintf("%s:%d invalid region_code", sourcePath, line))
				continue
			}
		}
		status := asString(m["status"])
		if status != "draft" && status != "published" && status != "archived" {
			errs = append(errs, fmt.Sprintf("%s:%d invalid status", sourcePath, line))
			continue
		}
		metadata := map[string]any{}
		if mm, ok := m["metadata"].(map[string]any); ok {
			for k, v := range mm {
				metadata[k] = v
			}
		}
		known := map[string]bool{"external_id": true, "title": true, "language": true, "region_code": true, "category": true, "status": true, "content": true, "question": true, "short_answer": true, "full_answer": true, "keywords": true, "source": true, "source_type": true, "source_uri": true, "source_hash": true, "metadata": true}
		for k, v := range m {
			if !known[k] {
				metadata[k] = v
			}
		}
		docs = append(docs, Document{ExternalID: asString(m["external_id"]), Title: asString(m["title"]), Language: lang, RegionCode: regCode, Category: defaultString(asString(m["category"]), "general"), SourceType: defaultString(asString(m["source_type"]), "jsonl"), SourceURI: defaultString(asString(m["source_uri"]), filepath.Base(sourcePath)), SourceHash: asString(m["source_hash"]), Status: status, Content: asString(m["content"]), Question: asString(m["question"]), ShortAnswer: asString(m["short_answer"]), FullAnswer: asString(m["full_answer"]), Source: asString(m["source"]), Keywords: nonNilStrings(asStringSlice(m["keywords"])), Metadata: metadata})
	}
	return docs, errs, sc.Err()
}

func LoadJSONLDocuments(ctx context.Context, root string, allowDisabled bool) ([]Document, []string, error) {
	var docs []Document
	var errs []string
	err := filepath.WalkDir(root, func(path string, d os.DirEntry, err error) error {
		if err != nil {
			return err
		}
		if d.IsDir() || !strings.HasSuffix(path, ".jsonl") {
			return nil
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		default:
		}
		dd, ee, err := ParseJSONLFile(path, allowDisabled)
		if err != nil {
			return err
		}
		docs = append(docs, dd...)
		errs = append(errs, ee...)
		return nil
	})
	return docs, errs, err
}

func BuildChunkText(doc Document) string {
	parts := []string{doc.Title}
	if doc.Question != "" {
		parts = append(parts, "Вопрос: "+doc.Question)
	}
	if doc.ShortAnswer != "" {
		parts = append(parts, "Краткий ответ: "+doc.ShortAnswer)
	}
	if doc.FullAnswer != "" {
		parts = append(parts, "Расширенное пояснение: "+doc.FullAnswer)
	}
	if len(doc.Keywords) > 0 {
		parts = append(parts, "Ключевые слова: "+strings.Join(doc.Keywords, ", "))
	}
	if doc.Source != "" {
		parts = append(parts, "Источник: "+doc.Source)
	}
	if doc.Content != "" {
		parts = append(parts, doc.Content)
	}
	return strings.Join(parts, "\n\n")
}
func ContentHash(s string) string { h := sha256.Sum256([]byte(s)); return hex.EncodeToString(h[:]) }
func asString(v any) string {
	if s, ok := v.(string); ok {
		return strings.TrimSpace(s)
	}
	return ""
}
func defaultString(v, d string) string {
	if v == "" {
		return d
	}
	return v
}
func asStringSlice(v any) []string {
	arr, ok := v.([]any)
	if !ok {
		return nil
	}
	out := []string{}
	for _, x := range arr {
		if s := asString(x); s != "" {
			out = append(out, s)
		}
	}
	return out
}

func nonNilStrings(v []string) []string {
	if v == nil {
		return []string{}
	}
	return v
}