rag

package
v0.1.5 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Aug 9, 2026 License: MIT Imports: 35 Imported by: 0

Documentation

Overview

Package rag implements a knowledge-base store for the coWork profile: import documents, chunk them, and search them. Phase 3 ships full-text search via SQLite FTS5 (reusing the same proven CJK-aware tokenizer as the memory subsystem) — no external embedding API or vector store required, so it works offline with zero new deps. An embedding/vector layer can be added behind the same rag_search interface later without changing the tool surface.

The store lives in the user config dir (persists across restarts), one DB per collection so collections are independent and deletable by file.

Index

Constants

View Source 
const (
	JobPending    = "pending"    // queued, not yet extracting
	JobExtracting = "extracting" // worker is processing chunks
	JobDone       = "done"       // all chunks processed (some may have errored)
	JobError      = "error"      // fatal: could not even start / all chunks failed
	JobCancelled  = "cancelled"  // user cancelled
)

JobStatus constants for rag_jobs.status.

View Source 
const (
	ChunkPending = "pending"
	ChunkRunning = "running"
	ChunkDone    = "done"
	ChunkError   = "error"
)

ChunkStatus constants for rag_chunks.status.

View Source 
const EdgeExtractionPrompt = `` /* 775-byte string literal not displayed */

EdgeExtractionPrompt is stage 2: given the entities already extracted from THIS chunk (injected as {known_nodes}), extract relations constrained to those endpoints — mirroring HE's graph.py:611 pattern. %s = the known-nodes list, %s = the chunk text.

View Source 
const ExtractionPrompt = `` /* 1126-byte string literal not displayed */

ExtractionPrompt is the system+user prompt sent to the LLM for each chunk. Adapted from Hyper-Extract's AutoGraph default (types/graph.py:41). Kept as a constant so the LLM impl and any future provider impl share it.

View Source 
const NodeExtractionPrompt = `` /* 670-byte string literal not displayed */

NodeExtractionPrompt is stage 1 of the two-stage extraction (borrowed from HE graph.py:510): extract ONLY entities first. The returned entity list then seeds stage 2's {known_nodes} so relations are forced to reference real entities, dramatically cutting hallucinated edges. %s = the chunk text.

Variables

View Source 
var ErrVectorScaleExceeded = errors.New("semantic search unavailable: entity count exceeds brute-force limit, use keyword search")

ErrVectorScaleExceeded is retained for backward compatibility but is no longer returned — the parallel in-memory cache handles any scale.

Functions

func DetectCommunities 

func DetectCommunities(g *louvainGraph, nameToIdx map[string]int, idxToName []string) map[string]int

DetectCommunities runs the full multi-level Louvain algorithm on the graph and returns a map from entity name → community ID. Community IDs are renumbered 0..C-1 in the final result.

The input graph is built from rag_relations (undirected, weighted by relation weight). Returns an empty map if there are no relations.

func ExportToObsidian 

func ExportToObsidian(store *Store, collection, outputDir string) error

ExportToObsidian exports a collection's knowledge graph as an Obsidian vault. Each entity becomes a markdown file with YAML front matter, wikilinks for relations, backlinks, and source citations. A MOC (_目录.md) is generated.

func FormatKnowledgeRef 

func FormatKnowledgeRef(store *Store, collection string, entityNames []string, relationKeys []string) string

FormatKnowledgeRef formats selected entities and relations into a markdown document suitable for passing to a skill as context. The output is structured as: entity list (name, type, description) then relation list (source → type → target).

func GetTemplatePrompt 

func GetTemplatePrompt(name string) (nodePrompt, edgePrompt string)

GetTemplatePrompt returns the domain-specific prompts for a template. Falls back to empty strings (caller should use default prompts) when the template is not found or has no custom prompts.

func IsTemplate 

func IsTemplate(name string) bool

IsTemplate reports whether name is a known extraction template.

func ReadDoc 

func ReadDoc(path string) (string, string, error)

readDoc reads a file and returns its text + an extension hint for chunking. Text-like formats (txt, md, code, csv, json, html) are parsed inline; binary Office formats (docx/xlsx/pptx/pdf) go through docconv/markitdown with a Go fallback. ReadDoc reads a document file and returns its text content and extension.

func ReadFileForPreview 

func ReadFileForPreview(path string) (string, string, error)

ReadFileForPreview reads a text file for document preview. Returns the body and extension. Delegates to readDoc which handles markitdown + Go fallback for all formats.

func RecommendTemplate 

func RecommendTemplate(content string) string

RecommendTemplate recommends a template based on document content analysis. Uses simple keyword matching for now.

func RelationKey 

func RelationKey(source, typ, target string) string

RelationKey formats a relation as a unique key string for selection tracking.

func SplitForPreview 

func SplitForPreview(body string, path string) []string

SplitForPreview splits a document body into chunks for preview highlighting. Uses the same chunking strategy as chunkDoc (paragraph split for md/txt, fixed windows for code).

Types

type BudgetAcquirer 

type BudgetAcquirer interface {
	Acquire(ctx context.Context, key string, priority bool) error
}

BudgetAcquirer gates a request through the global RPM limiter. It's the subset of *provider.RequestBudget this package needs, as an interface to avoid a rag→provider dependency.

type BudgetSetter 

type BudgetSetter interface {
	SetBudget(acquirer BudgetAcquirer, key string)
}

BudgetSetter is an optional capability an Extractor may implement so boot can install the global RPM limiter. Extractors that talk HTTP directly (instead of going through the provider layer) need this to share the per-minute quota with all other LLM calls. Type-assert to discover support.

type CollectionInfo 

type CollectionInfo struct {
	Name      string
	Documents int   // distinct paths
	Chunks    int   // total chunks
	Size      int64 // bytes of body text indexed (approx)
}

CollectionInfo describes one collection for rag_list.

type Embedder 

type Embedder interface {
	Embed(ctx context.Context, texts []string) ([][]float32, error)
}

Embedder produces a vector for a text string, enabling semantic reranking of FTS5 hits. Implementations call an embedding API. When nil, rag_search uses FTS5 alone. The interface is minimal so any provider can implement it; the store never assumes a dimension or model.

type Entity 

type Entity struct {
	ID          int64
	Collection  string
	Name        string // normalized key (lower+trim)
	NameRaw     string // display form
	Type        string
	Description string
	Sources     []Source
	RelationCnt int // number of relations touching this entity (0 if unknown)
	Community   int // Louvain community ID (-1 = unassigned)
}

Entity is one extracted entity (person/org/project/...). Name is the normalized key; NameRaw preserves the original surface form for display.

type EntityRelationView 

type EntityRelationView struct {
	Direction   string  `json:"direction"` // "out" | "in"
	Peer        string  `json:"peer"`
	Type        string  `json:"type"`
	Description string  `json:"description"`
	Weight      float64 `json:"weight"`
	Strength    float64 `json:"strength"`
}

EntityRelationView is a relation with direction info relative to an entity.

type EntityWithRelCount 

type EntityWithRelCount struct {
	Entity
	RelationCount int
}

EntityWithRelCount is an Entity augmented with its relation count.

type ExtractResult 

type ExtractResult struct {
	Entities  []Entity
	Relations []Relation
}

ExtractResult is the parsed LLM output for one chunk.

func ParseExtractJSON 

func ParseExtractJSON(b []byte) (ExtractResult, error)

ParseExtractJSON unmarshals an LLM JSON response into ExtractResult. The LLM may return either the canonical {entities:[...], relations:[...]} shape or wrap it; we tolerate both. Exposed so the LLM extractor impl can share parsing logic.

type Extractor 

type Extractor interface {
	Extract(ctx context.Context, chunk string, nodePrompt, edgePrompt string) (ExtractResult, error)
}

Extractor turns one chunk of text into entities + relations. The default implementation calls an LLM (OpenAI-compatible /chat/completions with a JSON schema); a no-op stub is used in tests.

nodePrompt and edgePrompt override the default extraction prompts when non-empty. edgePrompt must contain exactly two %s verbs (known-nodes list, chunk text). Pass "" for both to use the built-in general prompts.

func NewLLMExtractor 

func NewLLMExtractor(cfg LLMExtractorConfig) Extractor

NewLLMExtractor builds the default extractor. The cfg passed from boot.go carries the model from [cowork] extract_model (or the main model when empty).

type FieldMeta 

type FieldMeta struct {
	Name        string `json:"name"`
	Description string `json:"description"`
}

FieldMeta describes a field in a template's entity or relation schema.

type HEClient 

type HEClient struct {
	// contains filtered or unexported fields
}

HEClient is the HTTP client for the Hyper-Extract server.

func NewHEClient 

func NewHEClient(port int) *HEClient

NewHEClient creates a new Hyper-Extract client.

func (*HEClient) Embed 

func (c *HEClient) Embed(ctx context.Context, texts []string) ([][]float32, error)

Embed generates embedding vectors for a list of texts via the HE server.

func (*HEClient) ExportObsidian 

func (c *HEClient) ExportObsidian(ctx context.Context, kaPath string, outputDir string) error

ExportObsidian exports a KA to an Obsidian vault.

func (*HEClient) Extract 

func (c *HEClient) Extract(ctx context.Context, text string, template string, lang string) (*HEResult, error)

Extract extracts entities and relations from text.

func (*HEClient) Health 

func (c *HEClient) Health(ctx context.Context) (bool, bool, error)

Health checks if the server is running and Hyper-Extract is available.

func (*HEClient) ListTemplates 

func (c *HEClient) ListTemplates(ctx context.Context) ([]HETemplate, error)

ListTemplates returns available extraction templates.

func (*HEClient) Summarize 

func (c *HEClient) Summarize(ctx context.Context, entities []HEEntity, relations []HERelation, lang string) (*HESummary, error)

Summarize generates a knowledge summary from entities and relations.

type HEEntity 

type HEEntity struct {
	Name        string `json:"name"`
	Type        string `json:"type"`
	Description string `json:"description"`
}

HEEntity is an extracted entity from Hyper-Extract.

type HEFieldMeta 

type HEFieldMeta struct {
	Name        string `json:"name"`
	Description string `json:"description"`
}

HEFieldMeta is a field description from a template YAML.

type HERelation 

type HERelation struct {
	Source      string  `json:"source"`
	Target      string  `json:"target"`
	Type        string  `json:"type"`
	Description string  `json:"description"`
	Strength    float64 `json:"strength,omitempty"`
}

HERelation is an extracted relation from Hyper-Extract.

type HEResult 

type HEResult struct {
	Entities  []HEEntity   `json:"entities"`
	Relations []HERelation `json:"relations"`
}

HEResult is the extraction result from Hyper-Extract.

type HESummary 

type HESummary struct {
	Summary string   `json:"summary"`
	Themes  []string `json:"themes"`
	Error   string   `json:"error,omitempty"`
}

HESummary is the response from the /summarize endpoint.

type HETemplate 

type HETemplate struct {
	Name           string        `json:"name"`
	Category       string        `json:"category"`
	File           string        `json:"file"`
	Available      bool          `json:"available"`
	Description    string        `json:"description"`
	TemplateType   string        `json:"templateType"`
	EntityFields   []HEFieldMeta `json:"entityFields"`
	RelationFields []HEFieldMeta `json:"relationFields"`
}

HETemplate represents a Hyper-Extract template.

type JobRow 

type JobRow struct {
	ID          string
	Collection  string
	Path        string
	RelPath     string
	RootPath    string
	IsDir       bool
	Status      string
	TotalChunks int
	DoneChunks  int
	ErrorMsg    string
	ContentHash string // sha256 of chunked body; used for change-based dedup
	StatKey     string // "size:mtime" of the source file; cheap re-import dedup
	NodePrompt  string // persisted so Resume restores the original prompt
	EdgePrompt  string
	CreatedAt   time.Time
	UpdatedAt   time.Time
}

JobRow is one rag_jobs row, for the pipeline + UI.

type LLMExtractorConfig 

type LLMExtractorConfig struct {
	BaseURL  string // e.g. "https://api.example.com/v1" ("" = uses apihelper.BaseURL)
	APIKey   string // env var name to read the key from (e.g. "FAIRPEER_API_KEY"); if empty, reads FAIRPEER_API_KEY
	Model    string // chat model to use (e.g. the cowork main model)
	TwoStage bool   // extract entities then relations in two LLM calls (higher quality, 2× tokens); false = single combined call
}

LLMExtractorConfig configures the LLM-backed extractor.

type LLMSemantic 

type LLMSemantic struct {
	// contains filtered or unexported fields
}

LLMSemantic wraps a provider for query expansion + reranking. A nil provider (or nil LLMSemantic) means "off" — every method degrades to a no-op/pass-through.

func NewLLMSemantic 

func NewLLMSemantic(prov provider.Provider) *LLMSemantic

NewLLMSemantic returns an LLMSemantic over the given provider. prov may be nil (all methods then no-op), letting the caller wire unconditionally.

func (*LLMSemantic) ExpandQuery 

func (l *LLMSemantic) ExpandQuery(ctx context.Context, query string) []string

ExpandQuery returns the original query plus LLM-generated synonyms/variants, suitable for OR-ing into an FTS5 MATCH. On any failure (no provider, timeout, unparseable reply) it returns []string{query} — the caller always gets at least the original query, so search quality never drops below plain FTS5.

The LLM is prompted to produce bilingual variants, which simultaneously fixes CJK synonyms, English stemming, and cross-language alignment.

func (*LLMSemantic) Rerank 

func (l *LLMSemantic) Rerank(ctx context.Context, query string, results []Result) []Result

type MergeCandidate 

type MergeCandidate struct {
	KeepName  string  `json:"keepName"`  // higher-degree entity (the "canonical" one)
	MergeName string  `json:"mergeName"` // lower-degree entity to merge in
	KeepRaw   string  `json:"keepRaw"`
	MergeRaw  string  `json:"mergeRaw"`
	Score     float32 `json:"score"` // cosine similarity 0..1
}

MergeCandidate is a pair of entities that are semantically similar (high cosine similarity of their embeddings) and may be the same entity under different names. The UI presents these as suggestions; the user confirms.

type Pipeline 

type Pipeline struct {
	// contains filtered or unexported fields
}

Pipeline owns the extraction worker(s). One per app; call Start once at boot.

func NewPipeline 

func NewPipeline(store *Store, extractor Extractor, cfg PipelineConfig, emit ProgressEmitter) *Pipeline

NewPipeline constructs a pipeline. store + extractor are required; emit/logf may be nil. Call Start to begin processing.

func (*Pipeline) CancelJob 

func (p *Pipeline) CancelJob(jobID string) error

CancelJob marks a job cancelled and drops its pending tasks from the queue. In-flight chunks finish (we don't interrupt an LLM call mid-flight).

func (*Pipeline) EnqueuePaths 

func (p *Pipeline) EnqueuePaths(collection string, paths []string, nodePrompt, edgePrompt string, force bool) ([]string, error)

EnqueuePaths scans files under each path (recursing folders), imports them into FTS5 (synchronous, seconds), then enqueues their chunks for extraction. Returns the job IDs created. A path may be a file or a folder.

This is the "import" entrypoint from the UI: the user gets the file tree + FTS5 search immediately, and extraction runs in the background with progress.

func (*Pipeline) LatencyAvgMs 

func (p *Pipeline) LatencyAvgMs() int64

LatencyAvgMs exposes the in-memory sliding-average chunk latency, for the UI's ETA tooltip. Returns 0 when no chunks have completed yet (callers fall back to the persisted AvgChunkLatencyMs on the store).

func (*Pipeline) Resume 

func (p *Pipeline) Resume() int

Resume rehydrates the in-memory queue from durable state after a restart. It finds all jobs left in pending/extracting status (interrupted mid-run), re-reads each job's chunk text from FTS5, and re-enqueues the chunks that were still pending or errored. Call this once after Start() at boot.

This fulfills the restart-safety contract documented at the top of this file and in Store.PendingChunksForJob. Without Resume, an interrupted extraction leaves jobs stuck forever (workers gone, no tasks in memory). The prompt overrides ARE persisted on the job row (node_prompt/edge_prompt, v2 schema), so resumed tasks restore the original extraction prompts (e.g. a domain template like finance/graph survives a restart).

Returns the number of chunks re-enqueued.

func (*Pipeline) SetLogger 

func (p *Pipeline) SetLogger(logf func(format string, args ...any))

SetLogger installs a diagnostic logger (slog-style format string).

func (*Pipeline) Start 

func (p *Pipeline) Start()

Start launches the worker goroutine(s). Idempotent.

func (*Pipeline) Stop 

func (p *Pipeline) Stop()

Stop signals workers to drain and exit. Pending tasks remain in the queue (rehydrated on next Start via Resume).

type PipelineConfig 

type PipelineConfig struct {
	Concurrency int           // simultaneous chunk extractions (default 1)
	Interval    time.Duration // pause between chunks (default 3s)
	MaxRetries  int           // per-chunk retry count (default 3)
	RetryBase   time.Duration // exponential backoff base (default 2s: 2/4/8s)
	ChunkSize   int           // override store.chunkDoc default for extraction (0 = 1200)
}

PipelineConfig tunes the worker loop. Defaults favor stability over speed (concurrency=1, interval=3s) so we never trip rate limits — the user can raise these in [cowork] to speed extraction up on a beefy connection.

func DefaultPipelineConfig 

func DefaultPipelineConfig() PipelineConfig

DefaultPipelineConfig returns conservative defaults that prioritize "no errors" over throughput. Low concurrency (1) avoids API rate limits (429).

type ProgressEmitter 

type ProgressEmitter func(ProgressEvent)

ProgressEmitter pushes a ProgressEvent to the frontend. The desktop app supplies one backed by runtime.EventsEmit("rag:progress"). Nil = no events.

type ProgressEvent 

type ProgressEvent struct {
	JobID        string `json:"jobId"`
	Collection   string `json:"collection"`
	Path         string `json:"path"`
	Status       string `json:"status"` // job status
	DoneChunks   int    `json:"doneChunks"`
	TotalChunks  int    `json:"totalChunks"`
	AvgLatencyMs int64  `json:"avgLatencyMs"` // sliding-average ms/chunk
	Message      string `json:"message"`      // human-readable summary
}

ProgressEvent is emitted to the UI on each chunk completion. The frontend computes the visible ETA from AvgLatencyMs × (TotalChunks - DoneChunks) so the backend doesn't have to push a clock that ticks every second.

type Relation 

type Relation struct {
	ID          int64
	Collection  string
	Source      string // normalized entity name
	Target      string // normalized entity name
	Type        string
	Description string
	Sources     []Source
	Weight      float64 // co-occurrence frequency (incremented per source chunk)
	Strength    float64 // LLM-assigned semantic strength 1-10 (5=neutral)
}

Relation is one directed edge between two entity names.

type Result 

type Result struct {
	Collection string
	Path       string
	Chunk      int
	Snippet    string
	Score      float64
}

Result is one search hit.

type SessionRAGContext 

type SessionRAGContext struct {
	// contains filtered or unexported fields
}

SessionRAGContext tracks which collections are active for the current session. When non-empty, rag_search auto-scopes to these collections. Empty = search all collections (default behavior).

func NewSessionRAGContext 

func NewSessionRAGContext() *SessionRAGContext

NewSessionRAGContext creates a new session context.

func (*SessionRAGContext) ActiveCollectionsOrAll 

func (c *SessionRAGContext) ActiveCollectionsOrAll() []string

ActiveCollectionsOrAll returns the active collections, or a single empty string (meaning "all") when none are set.

func (*SessionRAGContext) GetActiveCollections 

func (c *SessionRAGContext) GetActiveCollections() []string

GetActiveCollections returns the currently active collections.

func (*SessionRAGContext) ResolveCollection 

func (c *SessionRAGContext) ResolveCollection(explicit string) string

ResolveCollection returns the effective collection scope for a search. If the caller specifies an explicit collection, it takes precedence. Otherwise, returns the session's active collections (may be empty = all).

func (*SessionRAGContext) SetActiveCollections 

func (c *SessionRAGContext) SetActiveCollections(collections []string)

SetActiveCollections sets the active collections for this session. Pass nil or empty to search all.

type Source 

type Source struct {
	Path  string `json:"path"`
	Chunk int    `json:"chunk"`
}

Source records where an entity/relation came from (for provenance/溯源).

type Store 

type Store struct {
	// contains filtered or unexported fields
}

Store is a FTS5-backed knowledge base. One store holds multiple collections; each collection is a set of imported documents. Search is scoped to a collection (or all collections when empty).

func Open 

func Open(dbPath string) (*Store, error)

Open creates/opens the RAG store at dbPath. The schema is a single FTS5 table keyed by collection + path + chunk index, plus four side tables for the structured-extraction layer (jobs/chunks/entities/relations). Existing FTS5 data and behavior are unchanged; the new tables simply add alongside.

func (*Store) AllJobs 

func (s *Store) AllJobs() ([]JobRow, error)

AllJobs returns every job (for the file-tree status view).

func (*Store) AvgChunkLatencyMs 

func (s *Store) AvgChunkLatencyMs(collection string) (int64, error)

AvgChunkLatencyMs returns the mean latency_ms across all done chunks in a collection (or all collections when empty). Used as a fallback ETA estimate when the in-memory sliding window is empty (e.g. right after restart).

func (*Store) ChunksByPath 

func (s *Store) ChunksByPath(collection, path string) ([]string, error)

ChunksByPath returns the ordered list of chunk bodies for a given collection+path by reading from FTS5. Used by Pipeline.Resume to rehydrate chunk text that is not persisted on rag_chunks (only status is).

func (*Store) ClearEmbeddings 

func (s *Store) ClearEmbeddings() error

ClearEmbeddings drops the whole embedding cache (e.g. on a collection clear, since stale entries for deleted chunks would otherwise linger forever).

func (*Store) Close 

func (s *Store) Close() error

func (*Store) CommunityCount 

func (s *Store) CommunityCount(collection string) (int, error)

CommunityCount returns the number of distinct communities in a collection.

func (*Store) CountEntities 

func (s *Store) CountEntities(collection string) (int, error)

CountEntities returns the number of entities in a collection.

func (*Store) CountRelations 

func (s *Store) CountRelations(collection string) (int, error)

CountRelations returns the number of relations in a collection.

func (*Store) CreateCollection 

func (s *Store) CreateCollection(name string) error

CreateCollection creates an empty collection by inserting a placeholder FTS5 row so it appears in List(). The placeholder is filtered from Search() results by the "path NOT LIKE 'placeholder://%'" clause. When the user imports real documents the placeholder is replaced.

func (*Store) CreateJob 

func (s *Store) CreateJob(j JobRow, chunkTexts []string) (string, error)

CreateJob inserts a new extraction job + its pending chunks. Returns the job id. Chunks is the list of chunk texts (the pipeline reads them back by id when processing, so we store the text on the chunk row to keep the pipeline stateless across restarts).

func (*Store) Delete 

func (s *Store) Delete(collection, path string) error

Delete removes all chunks for a path in a collection (or the whole collection when path is empty). It cascades across every RAG table so no structured knowledge is orphaned: rag_fts (text), rag_jobs + rag_chunks (extraction state), and rag_entities + rag_relations (the knowledge graph). For a single-path delete the entity/relation rows are pruned precisely by their sources JSON — rows whose only source was this path are dropped; rows that also came from other files keep those sources and survive.

func (*Store) DeleteCollectionEntities 

func (s *Store) DeleteCollectionEntities(collection string) error

DeleteCollectionEntities removes all entities and relations for a collection. Used before re-extracting with a different template. Wrapped in a transaction so a crash mid-delete doesn't leave dangling edges.

func (*Store) DeleteCollectionTree 

func (s *Store) DeleteCollectionTree(name string) error

DeleteCollectionTree removes a collection and all its path-prefix children (e.g. deleting "工作" also deletes "工作/领导材料"). Delegates to Delete with empty path for the exact collection, then deletes children.

func (*Store) DetectCommunitiesInStore 

func (s *Store) DetectCommunitiesInStore(collection string) (map[string]int, int, error)

DetectCommunitiesInStore builds the graph from rag_relations and runs Louvain. Returns the community map and the number of communities found.

func (*Store) EntityCount 

func (s *Store) EntityCount(collection string) (int, error)

EntityCount returns the number of extracted entities in a collection (or all when collection is empty). Used by the UI badge "✅ 已抽取 N 实体".

func (*Store) EntityEmbeddingStatus 

func (s *Store) EntityEmbeddingStatus(collection, model string) (map[int64]bool, error)

AllEntitiesWithEmbeddings returns entity IDs that already have embeddings for a given model.

func (*Store) EntityExists 

func (s *Store) EntityExists(collection, name string) (int64, bool)

EntityExists checks if an entity exists and returns its ID.

func (*Store) FindMergeCandidates 

func (s *Store) FindMergeCandidates(collection, model string, threshold float32) ([]MergeCandidate, error)

FindMergeCandidates finds entity pairs whose embeddings have cosine similarity ≥ threshold, suggesting they may be aliases of the same entity. Uses the in-memory vector cache for O(n²) pairwise comparison — practical up to ~50K entities (a few seconds). Returns candidates sorted by score desc.

The caller should ensure embeddings exist (RagEmbedEntities) before calling; entities without embeddings are skipped.

func (*Store) GraphBatch 

func (s *Store) GraphBatch(collection string, limit int) ([]Entity, map[string][]Relation, error)

GraphBatch returns entities + their relations + relation counts in a single query pass (one LEFT JOIN), eliminating the N+1 pattern where GetGraphData called RelationsOf once per entity. Designed for large knowledge bases (thousands to tens of thousands of entities).

The result is a flat row stream: each entity may appear multiple times (once per relation). The caller groups rows by entity name in memory. This is far cheaper than N separate RelationsOf queries, each of which acquires the store mutex and runs an independent SQL.

When limit > 0, entities are ordered by relation_cnt DESC (highest-connected "hub" entities first) so the caller gets the most visually important nodes upfront — ideal for paginated / progressive graph rendering.

func (*Store) HasEntities 

func (s *Store) HasEntities(collection string) (bool, error)

HasEntities reports whether a collection has any extracted entities (i.e. deep extraction has been run). Used to decide whether rag_search should query the structured layer at all.

func (*Store) Import 

func (s *Store) Import(collection, path string, tags []string) (int, error)

Import adds a document to a collection, splitting it into chunks first. Re- importing the same path replaces its chunks (delete-then-insert). Returns the number of chunks stored.

The tags parameter is currently ignored (deprecated): it was reserved for a metadata feature that was never wired through to storage or search. It is kept in the signature for source compatibility; pass nil. If per-document tags become wanted later, add a rag_doc_meta side table rather than reusing this param.

func (*Store) ImportContent 

func (s *Store) ImportContent(collection, path, body, ext string) (int, error)

ImportContent imports pre-read content into FTS5 (avoids re-reading the file). Use this when the caller already has the body and ext from a prior readDoc call.

func (*Store) ImportText 

func (s *Store) ImportText(collection, virtualPath, text string) (int, error)

ImportText adds raw text directly to a collection (for incremental updates). The text is chunked and indexed into FTS5 immediately.

func (*Store) JobByID 

func (s *Store) JobByID(jobID string) (JobRow, bool, error)

JobByID returns one job row. ok=false if not found.

func (*Store) JobContentHashForPath 

func (s *Store) JobContentHashForPath(collection, path string) (string, error)

JobContentHashForPath returns the stored content_hash for the job matching a collection+path (empty if no job or no hash). Used by the dedup check to detect content changes even when the chunk count is unchanged.

func (*Store) JobStatKeyForPath 

func (s *Store) JobStatKeyForPath(collection, path string) (string, error)

JobStatKeyForPath returns the stored stat_key ("size:mtime") for the job matching a collection+path (empty if no job or no key). Used by the cheap re-import dedup: if the file's on-disk size+mtime are unchanged, the body is guaranteed identical, so we skip the expensive readDoc (markitdown/OCR) and re-extraction entirely. This is the first-line dedup; content_hash is the second-line (catches edits that don't change the stat key).

func (*Store) JobStatusForPath 

func (s *Store) JobStatusForPath(collection, path string) (jobID, status string, totalChunks, doneChunks int, err error)

JobStatusForPath returns (jobID, status, totalChunks, doneChunks) for the job matching a collection+path, or ("", "", 0, 0) if none exists. Used by the pipeline's re-import dedup check: if a job is already done with the same chunk count, re-extraction is skipped to avoid burning LLM quota on unchanged files.

func (*Store) JobsByPath 

func (s *Store) JobsByPath(collection, path string) ([]JobRow, error)

JobsByPath returns all jobs for a given collection+path (usually 1, but a re-extract creates a new job row via ON CONFLICT replace).

func (*Store) List 

func (s *Store) List(name string) ([]CollectionInfo, error)

List returns a summary per collection (all collections when name is empty).

func (*Store) MarkChunkDone 

func (s *Store) MarkChunkDone(chunkID string, jobID string, latencyMs int64, err error) error

MarkChunkDone updates a chunk's status + the job's done_chunks counter. When all chunks are processed, the job status flips to done (or error if all failed). err != nil marks the chunk as errored with the message.

func (*Store) MergeEntities 

func (s *Store) MergeEntities(collection, keepName string, mergeNames []string) error

MergeEntities merges mergeNames into keepName within a collection. All relations referencing any merged entity are rewired to keepName, then the merged entity rows are deleted. Sources are concatenated (deduped).

func (*Store) PendingChunksForJob 

func (s *Store) PendingChunksForJob(jobID string) ([]struct {
	ChunkID string
	Idx     int
}, error)

PendingChunksForJob returns the (chunkID, idx) pairs still pending/errored for a job — used by Pipeline.Resume to rehydrate the queue after a restart. The chunk TEXT is re-read from FTS5 by (path, idx).

func (*Store) PruneDanglingRelations 

func (s *Store) PruneDanglingRelations(collection string) (int, error)

PruneDanglingRelations deletes relation rows whose source or target entity no longer exists in rag_entities. The relations table has no FK constraint (source/target are name strings, not ID references), so dangling edges can accumulate from cross-chunk extraction or HE-side mutations. Returns the number of deleted rows. When collection is empty, prunes all collections.

func (*Store) RecalcRelationCounts 

func (s *Store) RecalcRelationCounts(collection string) error

RecalcRelationCounts recomputes the relation_cnt column for all entities in a collection (or all collections when empty). Called by the v3 migration to backfill the column for pre-existing data, and after bulk deletes.

func (*Store) RelationsOf 

func (s *Store) RelationsOf(collection, entityName string, includeInverse bool) ([]Relation, error)

RelationsOf returns all relations touching entityName — outgoing (source), and when includeInverse, also incoming (target). entityName is normalized.

func (*Store) RelationsOfEntity 

func (s *Store) RelationsOfEntity(collection, entityName string) ([]EntityRelationView, error)

RelationsOfEntity returns all relations for an entity with direction info.

func (*Store) RenameCollection 

func (s *Store) RenameCollection(oldName, newName string) error

RenameCollection updates all tables from oldName to newName (a path prefix rename: "工作" → "工作资料" also updates "工作/领导材料" → "工作资料/领导材料"). Used by the collection tree's right-click rename. Transaction-wrapped so the rename either fully applies or not at all.

func (*Store) Rerank 

func (s *Store) Rerank(ctx context.Context, query string, results []Result, emb Embedder, ftsWeight float64) []Result

rerankWithEmbeddings upgrades FTS5 results using semantic similarity: it embeds the query + each hit's chunk, recomputes scores as a blend of FTS5 BM25 and cosine similarity, and re-sorts. If the embedder is nil or fails, it returns the original FTS5 ranking unchanged (graceful degradation — FTS5 is always a valid baseline).

The blend weight (ftsWeight) favors FTS5 for exact-term recall while letting semantics surface topically-related hits that don't share exact words. This matches the typical hybrid-RAG finding: BM25 for precision, embeddings for recall, blend for both.

func (*Store) ResumableJobs 

func (s *Store) ResumableJobs() ([]JobRow, error)

ResumableJobs returns all jobs whose status is pending or extracting (i.e. interrupted mid-extraction). Used by Pipeline.Resume at startup to rebuild the in-memory queue from durable state.

func (*Store) Search 

func (s *Store) Search(query, collection string, limit int) ([]Result, error)

Search runs an FTS5 MATCH query scoped to collection (empty = all). Returns ranked results (BM25, higher = better) up to limit.

func (*Store) SearchEntities 

func (s *Store) SearchEntities(query, collection string, limit int) ([]Entity, error)

SearchEntities returns entities whose normalized name contains the query (substring match, case-insensitive) OR whose description contains it. This is a lightweight lexical match over the (small) extracted set — for large entity counts an FTS5 mirror table would be better, but office collections rarely exceed thousands of entities. Limited to `limit` results.

func (*Store) SearchEntitiesByVector 

func (s *Store) SearchEntitiesByVector(collection, model string, queryVec []float32, topK int) ([]Entity, error)

SearchEntitiesByVector finds the topK entities most similar to the query vector using parallel brute-force cosine similarity over an in-memory vector cache. The cache is lazily loaded on first call and invalidated on any entity/embedding mutation. Supports 100K+ entities with sub-200ms queries on multi-core CPUs.

func (*Store) SetCommunity 

func (s *Store) SetCommunity(collection string, nameMap map[string]int) error

SetCommunity writes community IDs back to rag_entities. Called by DetectCommunities (the Store wrapper) after running the algorithm.

func (*Store) SetJobStatus 

func (s *Store) SetJobStatus(jobID, status string) error

SetJobStatus sets a job's status (used for extracting/cancelled transitions).

func (*Store) TopEntities 

func (s *Store) TopEntities(collection string, limit int) ([]EntityWithRelCount, error)

TopEntities returns the N entities with the most relations in a collection.

func (*Store) TopRelations 

func (s *Store) TopRelations(collection string, limit int) ([]Relation, error)

func (*Store) UpdateEntity 

func (s *Store) UpdateEntity(collection, name string, nameRaw, typ, desc string) error

UpdateEntity patches an entity's display fields (name_raw, type, description). The normalized name key is immutable; to change it, merge into a different entity.

func (*Store) UpsertEntity 

func (s *Store) UpsertEntity(collection string, e Entity, src Source) error

UpsertEntity inserts a new entity or merges into an existing one. Merging:

  • append the new Source to Sources (dedup by path+chunk)
  • keep the longer non-empty Description (later chunks may refine it)
  • keep the first non-empty Type (stable classification)

func (*Store) UpsertEntityEmbedding 

func (s *Store) UpsertEntityEmbedding(entityID int64, collection, model string, vec []float32) error

UpsertEntityEmbedding stores or updates the embedding vector for an entity.

func (*Store) UpsertRelation 

func (s *Store) UpsertRelation(collection string, r Relation, src Source) error

UpsertRelation inserts a new relation or merges into an existing one. The unique key is (collection, normalized source, normalized target, type).

func (*Store) Vacuum 

func (s *Store) Vacuum() error

Vacuum reclaims free space left by deletions. SQLite marks deleted rows as free internally but never shrinks the file without VACUUM, so a knowledge base that's been heavily imported-then-deleted can balloon. Safe but moderately expensive (rewrites the whole DB) — call after a collection clear or a large prune, not on every delete.

type TemplateInfo 

type TemplateInfo struct {
	Name           string      `json:"name"`
	Category       string      `json:"category"`
	DisplayName    string      `json:"displayName"`
	Description    string      `json:"description"`
	Available      bool        `json:"available"`
	TemplateType   string      `json:"templateType"`
	EntityFields   []FieldMeta `json:"entityFields"`
	RelationFields []FieldMeta `json:"relationFields"`
	// NodePrompt is the stage-1 entity extraction prompt for this template.
	// Empty = use the default NodeExtractionPrompt from extract.go.
	NodePrompt string `json:"-"`
	// EdgePrompt is the stage-2 relation extraction prompt. Must contain exactly
	// two %s verbs: first for the known-nodes list, second for the chunk text.
	// Empty = use the default EdgeExtractionPrompt from extract.go.
	EdgePrompt string `json:"-"`
}

TemplateInfo describes an extraction template.

func ListTemplates 

func ListTemplates(heClient *HEClient) []TemplateInfo

ListTemplates returns all available templates, ensuring built-in domain templates remain first-class citizens with rich schemas regardless of whether an external Hyper-Extract server is running.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.