From 39f122467d3d539f45ee580f40f8e5e8f49f3db5 Mon Sep 17 00:00:00 2001 From: Dmitry Ng <19asdek91@gmail.com> Date: Mon, 18 May 2026 18:26:52 +0300 Subject: [PATCH] feat(config): add new embedding and rename database connection pool settings - Introduced `EMBEDDING_MAX_TEXT_BYTES` to limit the maximum byte size of text sent to the embedding model. - Renamed database connection pool settings: `DATABASE_MAX_OPEN_CONNS`, `DATABASE_MAX_IDLE_CONNS`, and `DATABASE_VECTOR_MAX_CONNS` for improved PostgreSQL connection management. - Updated relevant documentation to reflect these new configuration options and their usage. - Adjusted various components to utilize the new settings for enhanced performance and resource management. --- .env.example | 6 + README.md | 19 +- backend/cmd/ftester/worker/executor.go | 11 + .../installer/wizard/controller/controller.go | 8 + backend/cmd/installer/wizard/locale/locale.go | 4 + .../installer/wizard/models/embedder_form.go | 37 ++++ backend/cmd/pentagi/main.go | 3 + backend/docs/config.md | 23 +- backend/docs/database.md | 18 +- backend/pkg/config/config.go | 11 +- backend/pkg/config/config_test.go | 2 +- backend/pkg/database/knowledge.sql.go | 45 +++- backend/pkg/database/knowledge/knowledge.go | 79 +++++-- backend/pkg/database/querier.go | 4 + backend/pkg/server/router.go | 2 +- backend/pkg/tools/code.go | 207 +++++++++++------- backend/pkg/tools/guide.go | 183 ++++++++++------ backend/pkg/tools/search.go | 191 ++++++++++------ backend/pkg/tools/tools.go | 23 ++ backend/pkg/tools/vecstore_helper.go | 78 +++++++ backend/sqlc/models/knowledge.sql | 15 ++ docker-compose.yml | 4 + 22 files changed, 702 insertions(+), 271 deletions(-) create mode 100644 backend/pkg/tools/vecstore_helper.go diff --git a/.env.example b/.env.example index 46a0a564..4b95dc8f 100644 --- a/.env.example +++ b/.env.example @@ -73,6 +73,7 @@ EMBEDDING_KEY= EMBEDDING_MODEL= EMBEDDING_PROVIDER= EMBEDDING_BATCH_SIZE= +EMBEDDING_MAX_TEXT_BYTES= EMBEDDING_STRIP_NEW_LINES= ## Summarizer @@ -226,6 +227,11 @@ PENTAGI_POSTGRES_USER=postgres PENTAGI_POSTGRES_PASSWORD=postgres # change this to improve security PENTAGI_POSTGRES_DB=pentagidb +## Postgres (pgvector) connection pool settings +DATABASE_MAX_OPEN_CONNS= +DATABASE_MAX_IDLE_CONNS= +DATABASE_VECTOR_MAX_CONNS= + ## Graphiti knowledge graph settings ## Set GRAPHITI_ENABLED=true and GRAPHITI_URL=http://graphiti:8000 to enable embedded Graphiti GRAPHITI_ENABLED=false diff --git a/README.md b/README.md index 5f7ef98d..485bc701 100644 --- a/README.md +++ b/README.md @@ -2619,23 +2619,23 @@ PentAGI opens two independent connection pools to the same Postgres instance: | Pool | Env var | Default | Used by | |---|---|---|---| -| Shared `sql.DB` | `DB_MAX_OPEN_CONNS` | `25` | All sqlc queries and GORM handlers share a single `*sql.DB` | -| Shared `pgxpool` | `DB_VECTOR_MAX_CONNS` | `10` | All pgvector stores (agent memory + knowledge API) share a single pool | +| Shared `sql.DB` | `DATABASE_MAX_OPEN_CONNS` | `25` | All sqlc queries and GORM handlers share a single `*sql.DB` | +| Shared `pgxpool` | `DATABASE_VECTOR_MAX_CONNS` | `10` | All pgvector stores (agent memory + knowledge API) share a single pool | Additional tuning knob: -- `DB_MAX_IDLE_CONNS` — maximum idle connections kept open in the `sql.DB` pool between requests (default: `5`). +- `DATABASE_MAX_IDLE_CONNS` — maximum idle connections kept open in the `sql.DB` pool between requests (default: `5`). **Budget for the stock `vxcontrol/pgvector` image** (`max_connections = 100`, `superuser_reserved_connections = 3`): ``` Available for client connections = 97 - pentagi sql.DB (DB_MAX_OPEN_CONNS) = 25 - pentagi pgxpool (DB_VECTOR_MAX_CONNS) = 10 - pgexporter = 3 - autovacuum workers = 3 + pentagi sql.DB (DATABASE_MAX_OPEN_CONNS) = 25 + pentagi pgxpool (DATABASE_VECTOR_MAX_CONNS) = 10 + pgexporter = 3 + autovacuum workers = 3 ───────────────────────────────────────── - Total consumed = 41 - Free buffer = 56 (≈ 58 %) + Total consumed = 41 + Free buffer = 56 (≈ 58 %) ``` The defaults are sized for **10 parallel flows** with concurrent API requests. If you run more flows or deploy multiple PentAGI instances against the same Postgres, raise `max_connections` via the `command` override in `docker-compose.yml` and increase the pool sizes proportionally: @@ -3084,6 +3084,7 @@ EMBEDDING_URL= # Optional custom API endpoint EMBEDDING_KEY= # API key for the provider (if required) EMBEDDING_BATCH_SIZE=100 # Number of documents to process in a batch EMBEDDING_STRIP_NEW_LINES=true # Whether to remove new lines from text before embedding +EMBEDDING_MAX_TEXT_BYTES=8192 # Max bytes of text sent to embedding model per document (byte proxy for token limit) # Advanced settings PROXY_URL= # Optional proxy for all API calls diff --git a/backend/cmd/ftester/worker/executor.go b/backend/cmd/ftester/worker/executor.go index 2cdb5dcf..2d2f6b6a 100644 --- a/backend/cmd/ftester/worker/executor.go +++ b/backend/cmd/ftester/worker/executor.go @@ -46,6 +46,7 @@ type toolExecutor struct { dockerClient docker.DockerClient handlers providers.FlowProviderHandlers store *pgvector.Store + embedder embeddings.Embedder graphitiClient *graphiti.Client proxies mocks.ProxyProviders userID int64 @@ -103,6 +104,7 @@ func newToolExecutor( dockerClient: dockerClient, handlers: handlers, store: store, + embedder: embedder, graphitiClient: graphitiClient, proxies: proxies, userID: userID, @@ -247,6 +249,9 @@ func (te *toolExecutor) GetTool(ctx context.Context, funcName string) (tools.Too te.subtaskID, te.replacer, te.store, + te.embedder, + te.db, + te.cfg.EmbeddingMaxTextBytes, te.proxies.GetVectorStoreLogProvider(), te.proxies.GetKnowledgeProvider(), ), nil @@ -259,6 +264,9 @@ func (te *toolExecutor) GetTool(ctx context.Context, funcName string) (tools.Too te.subtaskID, te.replacer, te.store, + te.embedder, + te.db, + te.cfg.EmbeddingMaxTextBytes, te.proxies.GetVectorStoreLogProvider(), te.proxies.GetKnowledgeProvider(), ), nil @@ -271,6 +279,9 @@ func (te *toolExecutor) GetTool(ctx context.Context, funcName string) (tools.Too te.subtaskID, te.replacer, te.store, + te.embedder, + te.db, + te.cfg.EmbeddingMaxTextBytes, te.proxies.GetVectorStoreLogProvider(), te.proxies.GetKnowledgeProvider(), ), nil diff --git a/backend/cmd/installer/wizard/controller/controller.go b/backend/cmd/installer/wizard/controller/controller.go index 3080a632..fe98d179 100644 --- a/backend/cmd/installer/wizard/controller/controller.go +++ b/backend/cmd/installer/wizard/controller/controller.go @@ -1197,6 +1197,7 @@ type EmbedderConfig struct { Model loader.EnvVar // EMBEDDING_MODEL BatchSize loader.EnvVar // EMBEDDING_BATCH_SIZE StripNewLines loader.EnvVar // EMBEDDING_STRIP_NEW_LINES + MaxTextBytes loader.EnvVar // EMBEDDING_MAX_TEXT_BYTES // computed fields (not directly mapped to env vars) Configured bool @@ -1212,6 +1213,7 @@ func (c *controller) GetEmbedderConfig() *EmbedderConfig { config.Model, _ = c.GetVar("EMBEDDING_MODEL") config.BatchSize, _ = c.GetVar("EMBEDDING_BATCH_SIZE") config.StripNewLines, _ = c.GetVar("EMBEDDING_STRIP_NEW_LINES") + config.MaxTextBytes, _ = c.GetVar("EMBEDDING_MAX_TEXT_BYTES") config.Installed = c.checker.PentagiInstalled // Determine if configured based on provider requirements @@ -1259,6 +1261,9 @@ func (c *controller) UpdateEmbedderConfig(config *EmbedderConfig) error { if err := c.SetVar("EMBEDDING_STRIP_NEW_LINES", config.StripNewLines.Value); err != nil { return fmt.Errorf("failed to set EMBEDDING_STRIP_NEW_LINES: %w", err) } + if err := c.SetVar("EMBEDDING_MAX_TEXT_BYTES", config.MaxTextBytes.Value); err != nil { + return fmt.Errorf("failed to set EMBEDDING_MAX_TEXT_BYTES: %w", err) + } return nil } @@ -1271,6 +1276,7 @@ func (c *controller) ResetEmbedderConfig() *EmbedderConfig { "EMBEDDING_MODEL", "EMBEDDING_BATCH_SIZE", "EMBEDDING_STRIP_NEW_LINES", + "EMBEDDING_MAX_TEXT_BYTES", } if err := c.ResetVars(vars); err != nil { @@ -2215,6 +2221,7 @@ func (c *controller) getVariableDescription(varName string) string { "EMBEDDING_MODEL": locale.EnvDesc_EMBEDDING_MODEL, "EMBEDDING_BATCH_SIZE": locale.EnvDesc_EMBEDDING_BATCH_SIZE, "EMBEDDING_STRIP_NEW_LINES": locale.EnvDesc_EMBEDDING_STRIP_NEW_LINES, + "EMBEDDING_MAX_TEXT_BYTES": locale.EnvDesc_EMBEDDING_MAX_TEXT_BYTES, "ASK_USER": locale.EnvDesc_ASK_USER, @@ -2455,6 +2462,7 @@ var criticalVariables = map[string]bool{ "EMBEDDING_MODEL": true, "EMBEDDING_BATCH_SIZE": true, "EMBEDDING_STRIP_NEW_LINES": true, + "EMBEDDING_MAX_TEXT_BYTES": true, // Docker configuration changes "DOCKER_INSIDE": true, diff --git a/backend/cmd/installer/wizard/locale/locale.go b/backend/cmd/installer/wizard/locale/locale.go index 372572a4..5ae20f77 100644 --- a/backend/cmd/installer/wizard/locale/locale.go +++ b/backend/cmd/installer/wizard/locale/locale.go @@ -1662,6 +1662,9 @@ Choose carefully as changing providers requires reindexing all stored data.` EmbedderFormStripNewLines = "Strip New Lines" EmbedderFormStripNewLinesDesc = "Remove line breaks from text before embedding (true/false)" + EmbedderFormMaxTextBytes = "Max Text Bytes" + EmbedderFormMaxTextBytesDesc = "Maximum number of bytes per text chunk sent to the embedding API (e.g. 8192)" + EmbedderFormHelpTitle = "Embedding Configuration" EmbedderFormHelpContent = `Configure text vectorization for semantic search and knowledge storage. @@ -2242,6 +2245,7 @@ const ( EnvDesc_EMBEDDING_MODEL = "Embedding Model" EnvDesc_EMBEDDING_BATCH_SIZE = "Embedding Batch Size" EnvDesc_EMBEDDING_STRIP_NEW_LINES = "Embedding Strip New Lines" + EnvDesc_EMBEDDING_MAX_TEXT_BYTES = "Embedding Max Text Bytes" EnvDesc_ASK_USER = "Human-in-the-loop" diff --git a/backend/cmd/installer/wizard/models/embedder_form.go b/backend/cmd/installer/wizard/models/embedder_form.go index ba3f8324..784a0aad 100644 --- a/backend/cmd/installer/wizard/models/embedder_form.go +++ b/backend/cmd/installer/wizard/models/embedder_form.go @@ -262,6 +262,11 @@ func (m *EmbedderFormModel) BuildForm() tea.Cmd { fields = append(fields, m.createStripNewLinesField(config)) } + // Max text bytes field (always show except for disabled) + if providerInfo.ID != locale.EmbedderProviderIDDisabled { + fields = append(fields, m.createMaxTextBytesField(config)) + } + m.SetFormFields(fields) return nil } @@ -352,6 +357,20 @@ func (m *EmbedderFormModel) createStripNewLinesField(config *controller.Embedder } } +func (m *EmbedderFormModel) createMaxTextBytesField(config *controller.EmbedderConfig) FormField { + input := NewTextInput(m.GetStyles(), m.GetWindow(), config.MaxTextBytes) + + return FormField{ + Key: "max_text_bytes", + Title: locale.EmbedderFormMaxTextBytes, + Description: locale.EmbedderFormMaxTextBytesDesc, + Required: false, + Masked: false, + Input: input, + Value: input.Value(), + } +} + func (m *EmbedderFormModel) GetFormTitle() string { return locale.EmbedderFormTitle } @@ -432,6 +451,15 @@ func (m *EmbedderFormModel) GetCurrentConfiguration() string { locale.EmbedderFormStripNewLines, m.GetStyles().Info.Render(stripNewLines))) } + maxTextBytes := config.MaxTextBytes.Value + if maxTextBytes == "" { + maxTextBytes = config.MaxTextBytes.Default + } + if maxTextBytes != "" { + sections = append(sections, fmt.Sprintf("• %s: %s", + locale.EmbedderFormMaxTextBytes, m.GetStyles().Info.Render(maxTextBytes))) + } + return strings.Join(sections, "\n") } @@ -477,6 +505,7 @@ func (m *EmbedderFormModel) HandleSave() error { Model: config.Model, BatchSize: config.BatchSize, StripNewLines: config.StripNewLines, + MaxTextBytes: config.MaxTextBytes, } // set provider @@ -514,6 +543,14 @@ func (m *EmbedderFormModel) HandleSave() error { return fmt.Errorf("invalid boolean value for strip newlines: %s (must be 'true' or 'false')", value) } newConfig.StripNewLines.Value = value + case "max_text_bytes": + // validate numeric input + if value != "" { + if intVal, err := strconv.Atoi(value); err != nil || intVal <= 0 { + return fmt.Errorf("invalid max text bytes: %s (must be a positive number)", value) + } + } + newConfig.MaxTextBytes.Value = value } } diff --git a/backend/cmd/pentagi/main.go b/backend/cmd/pentagi/main.go index 1f1fec9c..1cebf233 100644 --- a/backend/cmd/pentagi/main.go +++ b/backend/cmd/pentagi/main.go @@ -20,6 +20,7 @@ import ( "pentagi/pkg/docker" "pentagi/pkg/graph/subscriptions" obs "pentagi/pkg/observability" + "pentagi/pkg/observability/profiling" "pentagi/pkg/providers" router "pentagi/pkg/server" "pentagi/pkg/version" @@ -124,6 +125,8 @@ func main() { log.Println("Database schema updated successfully") + go profiling.Start() + client, err := docker.NewDockerClient(ctx, queries, cfg) if err != nil { log.Fatalf("Docker runtime client initialization failed: %v", err) diff --git a/backend/docs/config.md b/backend/docs/config.md index 42b90ab5..c53f41b0 100644 --- a/backend/docs/config.md +++ b/backend/docs/config.md @@ -127,17 +127,17 @@ The environment variables documented below remain the source of truth for config These settings control basic application behavior and are foundational for the system's operation. -| Option | Environment Variable | Default Value | Description | -| -------------- | -------------------- | ---------------------------------------------------------------------------- | ------------------------------------------------------------------------ | -| DatabaseURL | `DATABASE_URL` | `postgres://pentagiuser:pentagipass@pgvector:5432/pentagidb?sslmode=disable` | Connection string for the PostgreSQL database with pgvector extension | -| DBMaxOpenConns | `DB_MAX_OPEN_CONNS` | `25` | Maximum open connections in the shared `sql.DB` pool (sqlc + GORM combined). See [database.md §Connection Pooling](database.md#connection-pooling). | -| DBMaxIdleConns | `DB_MAX_IDLE_CONNS` | `5` | Maximum idle connections kept open between requests | -| DBVectorMaxConns | `DB_VECTOR_MAX_CONNS` | `10` | Maximum connections in the shared `pgxpool` for all pgvector stores | -| Debug | `DEBUG` | `false` | Enables debug mode with additional logging | -| DataDir | `DATA_DIR` | `./data` | Directory for storing persistent data | -| AskUser | `ASK_USER` | `false` | When enabled, requires explicit user confirmation for certain operations | -| InstallationID | `INSTALLATION_ID` | *(none)* | Unique installation identifier for PentAGI Cloud API communication | -| LicenseKey | `LICENSE_KEY` | *(none)* | License key for PentAGI Cloud API authentication and feature activation | +| Option | Environment Variable | Default Value | Description | +| ---------------- | --------------------------- | ---------------------------------------------------------------------------- | ------------------------------------------------------------------------ | +| DatabaseURL | `DATABASE_URL` | `postgres://pentagiuser:pentagipass@pgvector:5432/pentagidb?sslmode=disable` | Connection string for the PostgreSQL database with pgvector extension | +| DBMaxOpenConns | `DATABASE_MAX_OPEN_CONNS` | `25` | Maximum open connections in the shared `sql.DB` pool (sqlc + GORM combined). See [database.md §Connection Pooling](database.md#connection-pooling). | +| DBMaxIdleConns | `DATABASE_MAX_IDLE_CONNS` | `5` | Maximum idle connections kept open between requests | +| DBVectorMaxConns | `DATABASE_VECTOR_MAX_CONNS` | `10` | Maximum connections in the shared `pgxpool` for all pgvector stores | +| Debug | `DEBUG` | `false` | Enables debug mode with additional logging | +| DataDir | `DATA_DIR` | `./data` | Directory for storing persistent data | +| AskUser | `ASK_USER` | `false` | When enabled, requires explicit user confirmation for certain operations | +| InstallationID | `INSTALLATION_ID` | *(none)* | Unique installation identifier for PentAGI Cloud API communication | +| LicenseKey | `LICENSE_KEY` | *(none)* | License key for PentAGI Cloud API authentication and feature activation | ### Usage Details @@ -870,6 +870,7 @@ These settings control the vector embedding service used for semantic search and | EmbeddingStripNewLines | `EMBEDDING_STRIP_NEW_LINES` | `true` | Whether to strip newlines before embedding (improves quality) | | EmbeddingBatchSize | `EMBEDDING_BATCH_SIZE` | `512` | Batch size for embedding operations (affects memory usage and performance) | | EmbeddingProvider | `EMBEDDING_PROVIDER` | `openai` | Provider for embeddings (openai, ollama, mistral, jina, huggingface) | +| EmbeddingMaxTextBytes | `EMBEDDING_MAX_TEXT_BYTES` | `8192` | Maximum byte size of text sent to the embedding model per document. Acts as a byte-level proxy for token limits (e.g. 8192 tokens for OpenAI models). When a stored document exceeds this limit the heavy content field (Guide/Answer/Code) is truncated to fit before computing the vector; the full original text is always preserved in the database. Reduce if your model has a smaller context window. | ### Usage Details diff --git a/backend/docs/database.md b/backend/docs/database.md index 1074eede..c4e17fdc 100644 --- a/backend/docs/database.md +++ b/backend/docs/database.md @@ -1203,10 +1203,10 @@ PentAGI opens two independent connection pools to the same Postgres instance: | Pool | Env var | Default | Used by | |---|---|---|---| -| Shared `sql.DB` | `DB_MAX_OPEN_CONNS` | `25` | sqlc `Queries` and GORM — both clients are backed by the **same** `*sql.DB` created in `main.go` | -| Shared `pgxpool` | `DB_VECTOR_MAX_CONNS` | `10` | All `pgvector.Store` instances: every flow/assistant tool executor + knowledge API | +| Shared `sql.DB` | `DATABASE_MAX_OPEN_CONNS` | `25` | sqlc `Queries` and GORM — both clients are backed by the **same** `*sql.DB` created in `main.go` | +| Shared `pgxpool` | `DATABASE_VECTOR_MAX_CONNS` | `10` | All `pgvector.Store` instances: every flow/assistant tool executor + knowledge API | -Additional knob: `DB_MAX_IDLE_CONNS` (default `5`) — idle connections kept open between requests. +Additional knob: `DATABASE_MAX_IDLE_CONNS` (default `5`) — idle connections kept open between requests. `NewGorm` accepts the already-configured `*sql.DB` so GORM never opens its own pool: @@ -1234,13 +1234,13 @@ instead of opening individual `pgx.Connect` calls per executor. ``` Available for client connections = 97 - pentagi sql.DB (DB_MAX_OPEN_CONNS) = 25 - pentagi pgxpool (DB_VECTOR_MAX_CONNS) = 10 - pgexporter = 3 - autovacuum workers = 3 + pentagi sql.DB (DATABASE_MAX_OPEN_CONNS) = 25 + pentagi pgxpool (DATABASE_VECTOR_MAX_CONNS) = 10 + pgexporter = 3 + autovacuum workers = 3 ───────────────────────────────────────────── - Total consumed = 41 - Free buffer = 56 (≈ 58%) + Total consumed = 41 + Free buffer = 56 (≈ 58%) ``` Defaults are sized for 10 parallel flows with concurrent API requests. To inspect the live diff --git a/backend/pkg/config/config.go b/backend/pkg/config/config.go index 3f52212e..9ad47c97 100644 --- a/backend/pkg/config/config.go +++ b/backend/pkg/config/config.go @@ -72,6 +72,7 @@ type Config struct { EmbeddingStripNewLines bool `env:"EMBEDDING_STRIP_NEW_LINES" envDefault:"true"` EmbeddingBatchSize int `env:"EMBEDDING_BATCH_SIZE" envDefault:"512"` EmbeddingProvider string `env:"EMBEDDING_PROVIDER" envDefault:"openai"` + EmbeddingMaxTextBytes int `env:"EMBEDDING_MAX_TEXT_BYTES" envDefault:"8192"` // === Chain Summarization Engine === SummarizerPreserveLast bool `env:"SUMMARIZER_PRESERVE_LAST" envDefault:"true"` @@ -226,13 +227,9 @@ type Config struct { AgentPlanningStepEnabled bool `env:"AGENT_PLANNING_STEP_ENABLED" envDefault:"false"` // === Database Connection Pool Sizing === - // See backend/docs/database.md §Connection Pooling for budget calculation and - // operational commands. Both sqlc (Queries) and GORM share a single *sql.DB, - // so DB_MAX_OPEN_CONNS is the total sql.DB budget for the process. - // Ensure DB_MAX_OPEN_CONNS + DB_VECTOR_MAX_CONNS < Postgres max_connections. - DBMaxOpenConns int `env:"DB_MAX_OPEN_CONNS" envDefault:"25"` - DBMaxIdleConns int `env:"DB_MAX_IDLE_CONNS" envDefault:"5"` - DBVectorMaxConns int `env:"DB_VECTOR_MAX_CONNS" envDefault:"10"` + DBMaxOpenConns int `env:"DATABASE_MAX_OPEN_CONNS" envDefault:"25"` + DBMaxIdleConns int `env:"DATABASE_MAX_IDLE_CONNS" envDefault:"5"` + DBVectorMaxConns int `env:"DATABASE_VECTOR_MAX_CONNS" envDefault:"10"` // PgxPool is the shared pgxpool.Pool for all pgvector stores. Populated by // main after pool creation; NOT sourced from environment variables. diff --git a/backend/pkg/config/config_test.go b/backend/pkg/config/config_test.go index 68926522..85aa69f6 100644 --- a/backend/pkg/config/config_test.go +++ b/backend/pkg/config/config_test.go @@ -278,7 +278,7 @@ func clearConfigEnv(t *testing.T) { "OPEN_AI_KEY", "OPEN_AI_SERVER_URL", "ANTHROPIC_API_KEY", "ANTHROPIC_SERVER_URL", "EMBEDDING_URL", "EMBEDDING_KEY", "EMBEDDING_MODEL", - "EMBEDDING_STRIP_NEW_LINES", "EMBEDDING_BATCH_SIZE", "EMBEDDING_PROVIDER", + "EMBEDDING_STRIP_NEW_LINES", "EMBEDDING_BATCH_SIZE", "EMBEDDING_MAX_TEXT_BYTES", "EMBEDDING_PROVIDER", "SUMMARIZER_PRESERVE_LAST", "SUMMARIZER_USE_QA", "SUMMARIZER_SUM_MSG_HUMAN_IN_QA", "SUMMARIZER_LAST_SEC_BYTES", "SUMMARIZER_MAX_BP_BYTES", "SUMMARIZER_MAX_QA_SECTIONS", "SUMMARIZER_MAX_QA_BYTES", "SUMMARIZER_KEEP_QA_SECTIONS", diff --git a/backend/pkg/database/knowledge.sql.go b/backend/pkg/database/knowledge.sql.go index b6c6c8db..c11337d3 100644 --- a/backend/pkg/database/knowledge.sql.go +++ b/backend/pkg/database/knowledge.sql.go @@ -8,7 +8,9 @@ package database import ( "context" "database/sql" + "encoding/json" + "github.com/google/uuid" "github.com/sqlc-dev/pqtype" ) @@ -35,8 +37,8 @@ WHERE uuid::text = $1 ` // Delete a knowledge document by UUID (admin — no user_id check). -func (q *Queries) DeleteKnowledgeDocument(ctx context.Context, uuid sql.NullString) error { - _, err := q.db.ExecContext(ctx, deleteKnowledgeDocument, uuid) +func (q *Queries) DeleteKnowledgeDocument(ctx context.Context, argUuid sql.NullString) error { + _, err := q.db.ExecContext(ctx, deleteKnowledgeDocument, argUuid) return err } @@ -76,8 +78,8 @@ type GetKnowledgeDocumentRow struct { } // Fetch a single knowledge document by its UUID (admin view — no user_id check). -func (q *Queries) GetKnowledgeDocument(ctx context.Context, uuid string) (GetKnowledgeDocumentRow, error) { - row := q.db.QueryRowContext(ctx, getKnowledgeDocument, uuid) +func (q *Queries) GetKnowledgeDocument(ctx context.Context, argUuid string) (GetKnowledgeDocumentRow, error) { + row := q.db.QueryRowContext(ctx, getKnowledgeDocument, argUuid) var i GetKnowledgeDocumentRow err := row.Scan(&i.ID, &i.Document, &i.Cmetadata) return i, err @@ -114,6 +116,41 @@ func (q *Queries) GetUserKnowledgeDocument(ctx context.Context, arg GetUserKnowl return i, err } +const insertKnowledgeDocument = `-- name: InsertKnowledgeDocument :one +INSERT INTO langchain_pg_embedding (uuid, document, embedding, cmetadata, collection_id) +SELECT + $1::uuid, + $2, + $3::vector, + $4::json, + c.uuid +FROM langchain_pg_collection c +WHERE c.name = 'langchain' +RETURNING uuid::text AS id +` + +type InsertKnowledgeDocumentParams struct { + Uuid uuid.UUID `json:"uuid"` + Document sql.NullString `json:"document"` + Embedding interface{} `json:"embedding"` + Cmetadata json.RawMessage `json:"cmetadata"` +} + +// Insert a document with a pre-computed embedding vector and return its UUID. +// embedding must be formatted as a PostgreSQL vector literal: '[f1,f2,...]' +// cmetadata must be valid JSON text. +func (q *Queries) InsertKnowledgeDocument(ctx context.Context, arg InsertKnowledgeDocumentParams) (string, error) { + row := q.db.QueryRowContext(ctx, insertKnowledgeDocument, + arg.Uuid, + arg.Document, + arg.Embedding, + arg.Cmetadata, + ) + var id string + err := row.Scan(&id) + return id, err +} + const listAllKnowledgeDocuments = `-- name: ListAllKnowledgeDocuments :many SELECT e.uuid::text AS id, diff --git a/backend/pkg/database/knowledge/knowledge.go b/backend/pkg/database/knowledge/knowledge.go index 794210a8..b8e13845 100644 --- a/backend/pkg/database/knowledge/knowledge.go +++ b/backend/pkg/database/knowledge/knowledge.go @@ -14,8 +14,8 @@ import ( "pentagi/pkg/graph/subscriptions" "pentagi/pkg/providers/embeddings" + "github.com/google/uuid" "github.com/sqlc-dev/pqtype" - "github.com/vxcontrol/langchaingo/schema" "github.com/vxcontrol/langchaingo/vectorstores" ) @@ -54,10 +54,11 @@ type KnowledgeStore interface { } type knowledgeStore struct { - db database.Querier - store vectorstores.VectorStore // may be nil when no embedder is configured - embedder embeddings.Embedder // used for computing new embeddings on create/update - newKnp PublisherFactory + db database.Querier + store vectorstores.VectorStore // may be nil when no embedder is configured + embedder embeddings.Embedder // used for computing new embeddings on create/update + newKnp PublisherFactory + maxEmbeddingBytes int } // NewKnowledgeStore constructs a KnowledgeStore. @@ -66,17 +67,25 @@ type knowledgeStore struct { // list/get/delete still work. // - newKnp is called with the acting user's ID on each write to create a // correctly scoped event publisher. +// - maxEmbeddingBytes is the maximum byte size of text sent to the embedding +// model. Text is truncated to this limit before embedding to avoid token +// limit errors; the full original text is always stored in the database. func NewKnowledgeStore( db database.Querier, store vectorstores.VectorStore, embedder embeddings.Embedder, newKnp PublisherFactory, + maxEmbeddingBytes int, ) KnowledgeStore { + if maxEmbeddingBytes <= 0 { + maxEmbeddingBytes = 8192 + } return &knowledgeStore{ - db: db, - store: store, - embedder: embedder, - newKnp: newKnp, + db: db, + store: store, + embedder: embedder, + newKnp: newKnp, + maxEmbeddingBytes: maxEmbeddingBytes, } } @@ -347,6 +356,12 @@ func (ks *knowledgeStore) doSearch(ctx context.Context, userID int64, query stri limit = defaultSearchLimit } + // Truncate query to embedding size limit to avoid token limit errors. + // The search quality is preserved since the most relevant context is at the start. + if len(query) > ks.maxEmbeddingBytes { + query = query[:ks.maxEmbeddingBytes] + } + // Compute query embedding. vecs, err := ks.embedder.EmbedDocuments(ctx, []string{query}) if err != nil { @@ -443,7 +458,7 @@ func passesSearchFilter(doc *model.KnowledgeDocument, filter *model.KnowledgeFil // ---- CreateDocument --------------------------------------------------------- func (ks *knowledgeStore) CreateDocument(ctx context.Context, userID int64, input model.CreateKnowledgeDocumentInput) (*model.KnowledgeDocument, error) { - if err := ks.requireStore(); err != nil { + if err := ks.requireEmbedder(); err != nil { return nil, err } @@ -470,19 +485,36 @@ func (ks *knowledgeStore) CreateDocument(ctx context.Context, userID int64, inpu meta.PartSize = len(content) meta.TotalSize = len(content) - metaMap := metaToMap(meta) - ids, err := ks.store.AddDocuments(ctx, []schema.Document{{ - PageContent: content, - Metadata: metaMap, - }}) + // Truncate to embedding size limit for vector computation; full content goes to DB. + embeddingText := content + if len(embeddingText) > ks.maxEmbeddingBytes { + embeddingText = embeddingText[:ks.maxEmbeddingBytes] + } + vecs, err := ks.embedder.EmbedDocuments(ctx, []string{embeddingText}) + if err != nil { + return nil, fmt.Errorf("knowledge: compute embedding: %w", err) + } + if len(vecs) == 0 { + return nil, fmt.Errorf("knowledge: embedder returned no vectors") + } + + cmJSON, err := metaToJSON(meta) + if err != nil { + return nil, fmt.Errorf("knowledge: marshal cmetadata: %w", err) + } + + id := uuid.New() + docID, err := ks.db.InsertKnowledgeDocument(ctx, database.InsertKnowledgeDocumentParams{ + Uuid: id, + Document: nsOf(content), + Embedding: formatVector(vecs[0]), + Cmetadata: cmJSON.RawMessage, + }) if err != nil { return nil, fmt.Errorf("knowledge: create document: %w", err) } - if len(ids) == 0 { - return nil, fmt.Errorf("knowledge: create returned no IDs") - } - doc := metaToModelDoc(ids[0], content, meta) + doc := metaToModelDoc(docID, content, meta) ks.newKnp(userID).KnowledgeDocumentCreated(ctx, doc) return doc, nil } @@ -591,8 +623,13 @@ func (ks *knowledgeStore) doUpdate(ctx context.Context, userID int64, id string, meta.TotalSize = existing.TotalSize + deltaContentLen } - // Compute new embedding. - vecs, err := ks.embedder.EmbedDocuments(ctx, []string{content}) + // Compute new embedding. Truncate to maxEmbeddingBytes to avoid token limit + // errors; the full content is stored in the document column. + embeddingText := content + if len(embeddingText) > ks.maxEmbeddingBytes { + embeddingText = embeddingText[:ks.maxEmbeddingBytes] + } + vecs, err := ks.embedder.EmbedDocuments(ctx, []string{embeddingText}) if err != nil { return nil, fmt.Errorf("knowledge: compute embedding: %w", err) } diff --git a/backend/pkg/database/querier.go b/backend/pkg/database/querier.go index 5517efbb..8d80eb45 100644 --- a/backend/pkg/database/querier.go +++ b/backend/pkg/database/querier.go @@ -237,6 +237,10 @@ type Querier interface { GetUserTotalToolcallsStats(ctx context.Context, userID int64) (GetUserTotalToolcallsStatsRow, error) GetUserTotalUsageStats(ctx context.Context, userID int64) (GetUserTotalUsageStatsRow, error) GetUsers(ctx context.Context) ([]GetUsersRow, error) + // Insert a document with a pre-computed embedding vector and return its UUID. + // embedding must be formatted as a PostgreSQL vector literal: '[f1,f2,...]' + // cmetadata must be valid JSON text. + InsertKnowledgeDocument(ctx context.Context, arg InsertKnowledgeDocumentParams) (string, error) // List all knowledge documents excluding the noisy memory type (admin view). ListAllKnowledgeDocuments(ctx context.Context) ([]ListAllKnowledgeDocumentsRow, error) // List non-memory knowledge documents belonging to a specific flow (admin scoped). diff --git a/backend/pkg/server/router.go b/backend/pkg/server/router.go index 628f2c7c..5404e35d 100644 --- a/backend/pkg/server/router.go +++ b/backend/pkg/server/router.go @@ -148,7 +148,7 @@ func NewRouter( } } var knowledgeStore knowledge.KnowledgeStore - knowledgeStore = knowledge.NewKnowledgeStore(db, pgStore, embedder, subscriptions.NewKnowledgePublisher) + knowledgeStore = knowledge.NewKnowledgeStore(db, pgStore, embedder, subscriptions.NewKnowledgePublisher, cfg.EmbeddingMaxTextBytes) // ---- Anonymizer replacer ------------------------------------------------ // Shared singleton used by the GraphQL anonymizeText mutation. diff --git a/backend/pkg/tools/code.go b/backend/pkg/tools/code.go index c8eb06ee..238860f2 100644 --- a/backend/pkg/tools/code.go +++ b/backend/pkg/tools/code.go @@ -4,12 +4,14 @@ import ( "context" "encoding/json" "fmt" + "maps" "strings" "pentagi/pkg/database" "pentagi/pkg/graph/model" obs "pentagi/pkg/observability" "pentagi/pkg/observability/langfuse" + "pentagi/pkg/providers/embeddings" "github.com/sirupsen/logrus" "github.com/vxcontrol/cloud/anonymizer" @@ -27,14 +29,17 @@ const ( ) type code struct { - userID int64 - flowID int64 - taskID *int64 - subtaskID *int64 - replacer anonymizer.Replacer - store *pgvector.Store - vslp VectorStoreLogProvider - knp KnowledgeProvider + userID int64 + flowID int64 + taskID *int64 + subtaskID *int64 + replacer anonymizer.Replacer + store *pgvector.Store + embedder embeddings.Embedder + db database.Querier + maxEmbeddingBytes int + vslp VectorStoreLogProvider + knp KnowledgeProvider } func NewCodeTool( @@ -43,18 +48,24 @@ func NewCodeTool( taskID, subtaskID *int64, replacer anonymizer.Replacer, store *pgvector.Store, + embedder embeddings.Embedder, + db database.Querier, + maxEmbeddingBytes int, vslp VectorStoreLogProvider, knp KnowledgeProvider, ) Tool { return &code{ - userID: userID, - flowID: flowID, - taskID: taskID, - subtaskID: subtaskID, - replacer: replacer, - store: store, - vslp: vslp, - knp: knp, + userID: userID, + flowID: flowID, + taskID: taskID, + subtaskID: subtaskID, + replacer: replacer, + store: store, + embedder: embedder, + db: db, + maxEmbeddingBytes: maxEmbeddingBytes, + vslp: vslp, + knp: knp, } } @@ -215,22 +226,28 @@ func (c *code) Handle(ctx context.Context, name string, args json.RawMessage) (s return "", fmt.Errorf("failed to unmarshal %s store code action arguments: %w", name, err) } - buffer := strings.Builder{} - buffer.WriteString(action.Explanation) - buffer.WriteString(fmt.Sprintf("\n\n```%s\n\n", action.Lang)) - buffer.WriteString(action.Code) - buffer.WriteString("\n```") + renderedCode := c.renderCode(action.Explanation, action.Lang, action.Code) + // Anonymize before anything else so all downstream paths (including error + // branches that emit langfuse events) only ever expose the anonymized form. + var ( + anonymizedCode = c.replacer.ReplaceString(renderedCode) + anonymizedQuestion = c.replacer.ReplaceString(action.Question) + anonymizedDescription = c.replacer.ReplaceString(action.Description) + ) + + eventMetadata := map[string]any{ + "tool_name": name, + "code_lang": action.Lang, + "message": action.Message, + "description": anonymizedDescription, + "doc_type": codeVectorStoreDefaultType, + } opts := []langfuse.EventOption{ langfuse.WithEventName("store code samples to vector store"), langfuse.WithEventInput(action.Question), - langfuse.WithEventOutput(buffer.String()), - langfuse.WithEventMetadata(map[string]any{ - "tool_name": name, - "code_lang": action.Lang, - "message": action.Message, - "doc_type": codeVectorStoreDefaultType, - }), + langfuse.WithEventOutput(anonymizedCode), + langfuse.WithEventMetadata(eventMetadata), } logger = logger.WithFields(logrus.Fields{ @@ -239,49 +256,81 @@ func (c *code) Handle(ctx context.Context, name string, args json.RawMessage) (s "code": action.Code[:min(len(action.Code), 1000)], }) + // Build common metadata for the document. + metadata := map[string]any{ + "user_id": c.userID, + "flow_id": c.flowID, + "doc_type": codeVectorStoreDefaultType, + "code_lang": action.Lang, + "question": anonymizedQuestion, + "description": anonymizedDescription, + "part_size": len(anonymizedCode), + "total_size": len(anonymizedCode), + } + if c.taskID != nil { + metadata["task_id"] = *c.taskID + } + if c.subtaskID != nil { + metadata["subtask_id"] = *c.subtaskID + } + var ( - anonymizedCode = c.replacer.ReplaceString(buffer.String()) - anonymizedQuestion = c.replacer.ReplaceString(action.Question) + docs []schema.Document + err error + ids []string ) - docs, err := documentloaders.NewText(strings.NewReader(anonymizedCode)).Load(ctx) - if err != nil { - observation.Event(append(opts, - langfuse.WithEventStatus(err.Error()), - langfuse.WithEventLevel(langfuse.ObservationLevelError), - )...) - logger.WithError(err).Error("failed to load document") - return "", fmt.Errorf("failed to load document: %w", err) - } + if len(anonymizedCode) <= c.maxEmbeddingBytes || c.embedder == nil { + // Fast path: document fits within the embedding limit — use AddDocuments normally. + docs, err = documentloaders.NewText(strings.NewReader(anonymizedCode)).Load(ctx) + if err != nil { + observation.Event(append(opts, + langfuse.WithEventStatus(err.Error()), + langfuse.WithEventLevel(langfuse.ObservationLevelError), + )...) + logger.WithError(err).Error("failed to load document") + return "", fmt.Errorf("failed to load document: %w", err) + } + for i := range docs { + if docs[i].Metadata == nil { + docs[i].Metadata = map[string]any{} + } + maps.Copy(docs[i].Metadata, metadata) + docs[i].Metadata["part_size"] = len(docs[i].PageContent) + } + ids, err = c.store.AddDocuments(ctx, docs) + eventMetadata["ids"] = ids + if err != nil { + observation.Event(append(opts, + langfuse.WithEventStatus(err.Error()), + langfuse.WithEventLevel(langfuse.ObservationLevelError), + )...) + logger.WithError(err).Error("failed to store code sample") + return "", fmt.Errorf("failed to store code sample: %w", err) + } + } else { + // Slow path: document exceeds embedding limit — embed truncated text, + // store the full anonymized document in the database. + embeddingText := truncateForEmbedding(anonymizedCode, c.maxEmbeddingBytes) - for _, doc := range docs { - if doc.Metadata == nil { - doc.Metadata = map[string]any{} + id, err := storeDocumentWithEmbeddingLimit(ctx, c.db, c.embedder, + embeddingText, anonymizedCode, metadata) + if err != nil { + observation.Event(append(opts, + langfuse.WithEventStatus(err.Error()), + langfuse.WithEventLevel(langfuse.ObservationLevelError), + )...) + logger.WithError(err).Error("failed to store code sample with embedding limit") + return "", fmt.Errorf("failed to store code sample: %w", err) } - doc.Metadata["user_id"] = c.userID - doc.Metadata["flow_id"] = c.flowID - if c.taskID != nil { - doc.Metadata["task_id"] = *c.taskID + ids = []string{id} + docs = []schema.Document{ + { + PageContent: anonymizedCode, + Metadata: metadata, + }, } - if c.subtaskID != nil { - doc.Metadata["subtask_id"] = *c.subtaskID - } - doc.Metadata["doc_type"] = codeVectorStoreDefaultType - doc.Metadata["code_lang"] = action.Lang - doc.Metadata["question"] = anonymizedQuestion - doc.Metadata["description"] = action.Description - doc.Metadata["part_size"] = len(doc.PageContent) - doc.Metadata["total_size"] = len(anonymizedCode) - } - - ids, err := c.store.AddDocuments(ctx, docs) - if err != nil { - observation.Event(append(opts, - langfuse.WithEventStatus(err.Error()), - langfuse.WithEventLevel(langfuse.ObservationLevelError), - )...) - logger.WithError(err).Error("failed to store code sample") - return "", fmt.Errorf("failed to store code sample: %w", err) + eventMetadata["ids"] = ids } observation.Event(append(opts, @@ -292,20 +341,16 @@ func (c *code) Handle(ctx context.Context, name string, args json.RawMessage) (s if c.knp != nil { codeLang := action.Lang - for i, doc := range docs { - if i >= len(ids) { - break - } - desc := action.Description + for _, id := range ids { knDoc := &model.KnowledgeDocument{ - ID: ids[i], + ID: id, UserID: c.userID, DocType: model.KnowledgeDocTypeCode, - Content: doc.PageContent, + Content: anonymizedCode, Question: anonymizedQuestion, - Description: &desc, + Description: &anonymizedDescription, CodeLang: &codeLang, - PartSize: len(doc.PageContent), + PartSize: len(anonymizedCode), TotalSize: len(anonymizedCode), Manual: false, } @@ -341,7 +386,7 @@ func (c *code) Handle(ctx context.Context, name string, args json.RawMessage) (s filtersData, action.Question, database.VecstoreActionTypeStore, - buffer.String(), + renderedCode, c.taskID, c.subtaskID, ) @@ -358,3 +403,15 @@ func (c *code) Handle(ctx context.Context, name string, args json.RawMessage) (s func (c *code) IsAvailable() bool { return c.store != nil } + +func (c *code) renderCode(explanation, lang, code string) string { + buffer := strings.Builder{} + buffer.WriteString(explanation) + buffer.WriteString("\n\n") + buffer.WriteString("```") + buffer.WriteString(lang) + buffer.WriteString("\n") + buffer.WriteString(code) + buffer.WriteString("\n```") + return buffer.String() +} diff --git a/backend/pkg/tools/guide.go b/backend/pkg/tools/guide.go index 07e77030..a0d304c5 100644 --- a/backend/pkg/tools/guide.go +++ b/backend/pkg/tools/guide.go @@ -4,12 +4,14 @@ import ( "context" "encoding/json" "fmt" + "maps" "strings" "pentagi/pkg/database" "pentagi/pkg/graph/model" obs "pentagi/pkg/observability" "pentagi/pkg/observability/langfuse" + "pentagi/pkg/providers/embeddings" "github.com/sirupsen/logrus" "github.com/vxcontrol/cloud/anonymizer" @@ -27,32 +29,41 @@ const ( ) type guide struct { - userID int64 - flowID int64 - taskID *int64 - subtaskID *int64 - replacer anonymizer.Replacer - store *pgvector.Store - vslp VectorStoreLogProvider - knp KnowledgeProvider + userID int64 + flowID int64 + taskID *int64 + subtaskID *int64 + replacer anonymizer.Replacer + store *pgvector.Store + embedder embeddings.Embedder + db database.Querier + maxEmbeddingBytes int + vslp VectorStoreLogProvider + knp KnowledgeProvider } func NewGuideTool( userID int64, flowID int64, taskID, subtaskID *int64, replacer anonymizer.Replacer, store *pgvector.Store, + embedder embeddings.Embedder, + db database.Querier, + maxEmbeddingBytes int, vslp VectorStoreLogProvider, knp KnowledgeProvider, ) Tool { return &guide{ - userID: userID, - flowID: flowID, - taskID: taskID, - subtaskID: subtaskID, - replacer: replacer, - store: store, - vslp: vslp, - knp: knp, + userID: userID, + flowID: flowID, + taskID: taskID, + subtaskID: subtaskID, + replacer: replacer, + store: store, + embedder: embedder, + db: db, + maxEmbeddingBytes: maxEmbeddingBytes, + vslp: vslp, + knp: knp, } } @@ -214,16 +225,25 @@ func (g *guide) Handle(ctx context.Context, name string, args json.RawMessage) ( guide := fmt.Sprintf("Question:\n%s\n\nGuide:\n%s", action.Question, action.Guide) + // Anonymize before anything else so all downstream paths (including error + // branches that emit langfuse events) only ever expose the anonymized form. + var ( + anonymizedGuide = g.replacer.ReplaceString(guide) + anonymizedQuestion = g.replacer.ReplaceString(action.Question) + anonymizedGuideOnly = g.replacer.ReplaceString(action.Guide) // used in slow-path embedding text + ) + + eventMetadata := map[string]any{ + "tool_name": name, + "message": action.Message, + "doc_type": guideVectorStoreDefaultType, + "guide_type": action.Type, + } opts := []langfuse.EventOption{ langfuse.WithEventName("store guide to vector store"), langfuse.WithEventInput(action.Question), - langfuse.WithEventOutput(guide), - langfuse.WithEventMetadata(map[string]any{ - "tool_name": name, - "message": action.Message, - "doc_type": guideVectorStoreDefaultType, - "guide_type": action.Type, - }), + langfuse.WithEventOutput(anonymizedGuide), + langfuse.WithEventMetadata(eventMetadata), } logger = logger.WithFields(logrus.Fields{ @@ -232,48 +252,82 @@ func (g *guide) Handle(ctx context.Context, name string, args json.RawMessage) ( "guide": action.Guide[:min(len(action.Guide), 1000)], }) + // Build common metadata for the document. + metadata := map[string]any{ + "user_id": g.userID, + "flow_id": g.flowID, + "doc_type": guideVectorStoreDefaultType, + "guide_type": action.Type, + "question": anonymizedQuestion, + "part_size": len(anonymizedGuide), + "total_size": len(anonymizedGuide), + } + if g.taskID != nil { + metadata["task_id"] = *g.taskID + } + if g.subtaskID != nil { + metadata["subtask_id"] = *g.subtaskID + } + var ( - anonymizedGuide = g.replacer.ReplaceString(guide) - anonymizedQuestion = g.replacer.ReplaceString(action.Question) + docs []schema.Document + ids []string + err error ) - docs, err := documentloaders.NewText(strings.NewReader(anonymizedGuide)).Load(ctx) - if err != nil { - observation.Event(append(opts, - langfuse.WithEventStatus(err.Error()), - langfuse.WithEventLevel(langfuse.ObservationLevelError), - )...) - logger.WithError(err).Error("failed to load document") - return "", fmt.Errorf("failed to load document: %w", err) - } + if len(anonymizedGuide) <= g.maxEmbeddingBytes || g.embedder == nil { + // Fast path: document fits within the embedding limit. + docs, err = documentloaders.NewText(strings.NewReader(anonymizedGuide)).Load(ctx) + if err != nil { + observation.Event(append(opts, + langfuse.WithEventStatus(err.Error()), + langfuse.WithEventLevel(langfuse.ObservationLevelError), + )...) + logger.WithError(err).Error("failed to load document") + return "", fmt.Errorf("failed to load document: %w", err) + } + for i := range docs { + if docs[i].Metadata == nil { + docs[i].Metadata = map[string]any{} + } + maps.Copy(docs[i].Metadata, metadata) + docs[i].Metadata["part_size"] = len(docs[i].PageContent) + } + ids, err = g.store.AddDocuments(ctx, docs) + eventMetadata["ids"] = ids + if err != nil { + observation.Event(append(opts, + langfuse.WithEventStatus(err.Error()), + langfuse.WithEventLevel(langfuse.ObservationLevelError), + )...) + logger.WithError(err).Error("failed to store guide") + return "", fmt.Errorf("failed to store guide: %w", err) + } + } else { + // Slow path: Guide field exceeds embedding limit. + // Template: "Question:\n{question}\n\nGuide:\n{guide}" + prefix := "Question:\n" + anonymizedQuestion + "\n\nGuide:\n" + available := max(g.maxEmbeddingBytes-len(prefix), 0) + embeddingText := prefix + truncateForEmbedding(anonymizedGuideOnly, available) - for _, doc := range docs { - if doc.Metadata == nil { - doc.Metadata = map[string]any{} + id, err := storeDocumentWithEmbeddingLimit(ctx, g.db, g.embedder, + embeddingText, anonymizedGuide, metadata) + if err != nil { + observation.Event(append(opts, + langfuse.WithEventStatus(err.Error()), + langfuse.WithEventLevel(langfuse.ObservationLevelError), + )...) + logger.WithError(err).Error("failed to store guide with embedding limit") + return "", fmt.Errorf("failed to store guide: %w", err) } - doc.Metadata["user_id"] = g.userID - doc.Metadata["flow_id"] = g.flowID - if g.taskID != nil { - doc.Metadata["task_id"] = *g.taskID + ids = []string{id} + docs = []schema.Document{ + { + PageContent: anonymizedGuide, + Metadata: metadata, + }, } - if g.subtaskID != nil { - doc.Metadata["subtask_id"] = *g.subtaskID - } - doc.Metadata["doc_type"] = guideVectorStoreDefaultType - doc.Metadata["guide_type"] = action.Type - doc.Metadata["question"] = anonymizedQuestion - doc.Metadata["part_size"] = len(doc.PageContent) - doc.Metadata["total_size"] = len(anonymizedGuide) - } - - ids, err := g.store.AddDocuments(ctx, docs) - if err != nil { - observation.Event(append(opts, - langfuse.WithEventStatus(err.Error()), - langfuse.WithEventLevel(langfuse.ObservationLevelError), - )...) - logger.WithError(err).Error("failed to store guide") - return "", fmt.Errorf("failed to store guide: %w", err) + eventMetadata["ids"] = ids } observation.Event(append(opts, @@ -284,18 +338,15 @@ func (g *guide) Handle(ctx context.Context, name string, args json.RawMessage) ( if g.knp != nil { guideType := model.KnowledgeGuideType(action.Type) - for i, doc := range docs { - if i >= len(ids) { - break - } + for _, id := range ids { knDoc := &model.KnowledgeDocument{ - ID: ids[i], + ID: id, UserID: g.userID, DocType: model.KnowledgeDocTypeGuide, - Content: doc.PageContent, + Content: anonymizedGuide, Question: anonymizedQuestion, GuideType: &guideType, - PartSize: len(doc.PageContent), + PartSize: len(anonymizedGuide), TotalSize: len(anonymizedGuide), Manual: false, } diff --git a/backend/pkg/tools/search.go b/backend/pkg/tools/search.go index e1c9deca..ab852d89 100644 --- a/backend/pkg/tools/search.go +++ b/backend/pkg/tools/search.go @@ -4,12 +4,14 @@ import ( "context" "encoding/json" "fmt" + "maps" "strings" "pentagi/pkg/database" "pentagi/pkg/graph/model" obs "pentagi/pkg/observability" "pentagi/pkg/observability/langfuse" + "pentagi/pkg/providers/embeddings" "github.com/sirupsen/logrus" "github.com/vxcontrol/cloud/anonymizer" @@ -27,14 +29,17 @@ const ( ) type search struct { - userID int64 - flowID int64 - taskID *int64 - subtaskID *int64 - replacer anonymizer.Replacer - store *pgvector.Store - vslp VectorStoreLogProvider - knp KnowledgeProvider + userID int64 + flowID int64 + taskID *int64 + subtaskID *int64 + replacer anonymizer.Replacer + store *pgvector.Store + embedder embeddings.Embedder + db database.Querier + maxEmbeddingBytes int + vslp VectorStoreLogProvider + knp KnowledgeProvider } func NewSearchTool( @@ -43,18 +48,24 @@ func NewSearchTool( taskID, subtaskID *int64, replacer anonymizer.Replacer, store *pgvector.Store, + embedder embeddings.Embedder, + db database.Querier, + maxEmbeddingBytes int, vslp VectorStoreLogProvider, knp KnowledgeProvider, ) Tool { return &search{ - userID: userID, - flowID: flowID, - taskID: taskID, - subtaskID: subtaskID, - replacer: replacer, - store: store, - vslp: vslp, - knp: knp, + userID: userID, + flowID: flowID, + taskID: taskID, + subtaskID: subtaskID, + replacer: replacer, + store: store, + embedder: embedder, + db: db, + maxEmbeddingBytes: maxEmbeddingBytes, + vslp: vslp, + knp: knp, } } @@ -210,16 +221,24 @@ func (s *search) Handle(ctx context.Context, name string, args json.RawMessage) return "", fmt.Errorf("failed to unmarshal %s store answer action arguments: %w", name, err) } + // Anonymize before anything else so all downstream paths (including error + // branches that emit langfuse events) only ever expose the anonymized form. + var ( + anonymizedAnswer = s.replacer.ReplaceString(action.Answer) + anonymizedQuestion = s.replacer.ReplaceString(action.Question) + ) + + eventMetadata := map[string]any{ + "tool_name": name, + "message": action.Message, + "doc_type": searchVectorStoreDefaultType, + "answer_type": action.Type, + } opts := []langfuse.EventOption{ langfuse.WithEventName("store search answer to vector store"), langfuse.WithEventInput(action.Question), - langfuse.WithEventOutput(action.Answer), - langfuse.WithEventMetadata(map[string]any{ - "tool_name": name, - "message": action.Message, - "doc_type": searchVectorStoreDefaultType, - "answer_type": action.Type, - }), + langfuse.WithEventOutput(anonymizedAnswer), + langfuse.WithEventMetadata(eventMetadata), } logger = logger.WithFields(logrus.Fields{ @@ -228,44 +247,80 @@ func (s *search) Handle(ctx context.Context, name string, args json.RawMessage) "answer": action.Answer[:min(len(action.Answer), 1000)], }) + // Build common metadata for the document. + metadata := map[string]any{ + "user_id": s.userID, + "flow_id": s.flowID, + "doc_type": searchVectorStoreDefaultType, + "answer_type": action.Type, + "question": anonymizedQuestion, + "part_size": len(anonymizedAnswer), + "total_size": len(anonymizedAnswer), + } + if s.taskID != nil { + metadata["task_id"] = *s.taskID + } + if s.subtaskID != nil { + metadata["subtask_id"] = *s.subtaskID + } + var ( - anonymizedAnswer = s.replacer.ReplaceString(action.Answer) - anonymizedQuestion = s.replacer.ReplaceString(action.Question) + docs []schema.Document + ids []string + err error ) - docs, err := documentloaders.NewText(strings.NewReader(anonymizedAnswer)).Load(ctx) - if err != nil { - observation.Event(append(opts, - langfuse.WithEventStatus(err.Error()), - langfuse.WithEventLevel(langfuse.ObservationLevelError), - )...) - logger.WithError(err).Error("failed to load document") - return "", fmt.Errorf("failed to load document: %w", err) - } - - for _, doc := range docs { - if doc.Metadata == nil { - doc.Metadata = map[string]any{} + if len(anonymizedAnswer) <= s.maxEmbeddingBytes || s.embedder == nil { + // Fast path: answer fits within the embedding limit. + docs, err = documentloaders.NewText(strings.NewReader(anonymizedAnswer)).Load(ctx) + if err != nil { + observation.Event(append(opts, + langfuse.WithEventStatus(err.Error()), + langfuse.WithEventLevel(langfuse.ObservationLevelError), + )...) + logger.WithError(err).Error("failed to load document") + return "", fmt.Errorf("failed to load document: %w", err) } - doc.Metadata["user_id"] = s.userID - doc.Metadata["flow_id"] = s.flowID - doc.Metadata["task_id"] = s.taskID - doc.Metadata["subtask_id"] = s.subtaskID - doc.Metadata["doc_type"] = searchVectorStoreDefaultType - doc.Metadata["answer_type"] = action.Type - doc.Metadata["question"] = anonymizedQuestion - doc.Metadata["part_size"] = len(doc.PageContent) - doc.Metadata["total_size"] = len(anonymizedAnswer) - } + for i := range docs { + if docs[i].Metadata == nil { + docs[i].Metadata = map[string]any{} + } + maps.Copy(docs[i].Metadata, metadata) + docs[i].Metadata["part_size"] = len(docs[i].PageContent) + } + ids, err = s.store.AddDocuments(ctx, docs) + eventMetadata["ids"] = ids + if err != nil { + observation.Event(append(opts, + langfuse.WithEventStatus(err.Error()), + langfuse.WithEventLevel(langfuse.ObservationLevelError), + )...) + logger.WithError(err).Error("failed to store answer for question") + return "", fmt.Errorf("failed to store answer for question: %w", err) + } + } else { + // Slow path: Answer field exceeds embedding limit. + // PageContent is just the answer text, so overhead = 0. + embeddingText := truncateForEmbedding(anonymizedAnswer, s.maxEmbeddingBytes) - ids, err := s.store.AddDocuments(ctx, docs) - if err != nil { - observation.Event(append(opts, - langfuse.WithEventStatus(err.Error()), - langfuse.WithEventLevel(langfuse.ObservationLevelError), - )...) - logger.WithError(err).Error("failed to store answer for question") - return "", fmt.Errorf("failed to store answer for question: %w", err) + id, err := storeDocumentWithEmbeddingLimit(ctx, s.db, s.embedder, + embeddingText, anonymizedAnswer, metadata) + if err != nil { + observation.Event(append(opts, + langfuse.WithEventStatus(err.Error()), + langfuse.WithEventLevel(langfuse.ObservationLevelError), + )...) + logger.WithError(err).Error("failed to store answer with embedding limit") + return "", fmt.Errorf("failed to store answer for question: %w", err) + } + ids = []string{id} + docs = []schema.Document{ + { + PageContent: anonymizedAnswer, + Metadata: metadata, + }, + } + eventMetadata["ids"] = ids } observation.Event(append(opts, @@ -276,18 +331,15 @@ func (s *search) Handle(ctx context.Context, name string, args json.RawMessage) if s.knp != nil { answerType := model.KnowledgeAnswerType(action.Type) - for i, doc := range docs { - if i >= len(ids) { - break - } + for _, id := range ids { knDoc := &model.KnowledgeDocument{ - ID: ids[i], + ID: id, UserID: s.userID, DocType: model.KnowledgeDocTypeAnswer, - Content: doc.PageContent, + Content: anonymizedAnswer, Question: anonymizedQuestion, AnswerType: &answerType, - PartSize: len(doc.PageContent), + PartSize: len(anonymizedAnswer), TotalSize: len(anonymizedAnswer), Manual: false, } @@ -301,12 +353,17 @@ func (s *search) Handle(ctx context.Context, name string, args json.RawMessage) } if agentCtx, ok := GetAgentContext(ctx); ok { - filtersData, err := json.Marshal(map[string]any{ + data := map[string]any{ "doc_type": searchVectorStoreDefaultType, "answer_type": action.Type, - "task_id": s.taskID, - "subtask_id": s.subtaskID, - }) + } + if s.taskID != nil { + data["task_id"] = *s.taskID + } + if s.subtaskID != nil { + data["subtask_id"] = *s.subtaskID + } + filtersData, err := json.Marshal(data) if err != nil { logger.WithError(err).Error("failed to marshal filters") return "", fmt.Errorf("failed to marshal filters: %w", err) diff --git a/backend/pkg/tools/tools.go b/backend/pkg/tools/tools.go index 6e3bc80c..2c72884b 100644 --- a/backend/pkg/tools/tools.go +++ b/backend/pkg/tools/tools.go @@ -172,6 +172,7 @@ type flowToolsExecutor struct { db database.Querier cfg *config.Config + embedder embeddings.Embedder store *pgvector.Store graphitiClient *graphiti.Client image string @@ -392,6 +393,7 @@ func (fte *flowToolsExecutor) SetImage(image string) { } func (fte *flowToolsExecutor) SetEmbedder(embedder embeddings.Embedder) { + fte.embedder = embedder if !embedder.IsAvailable() { return } @@ -865,6 +867,9 @@ func (fte *flowToolsExecutor) GetAssistantExecutor(cfg AssistantExecutorConfig) fte.userID, fte.flowID, nil, nil, fte.replacer, fte.store, + fte.embedder, + fte.db, + fte.cfg.EmbeddingMaxTextBytes, fte.vslp, fte.knp, ) @@ -877,6 +882,9 @@ func (fte *flowToolsExecutor) GetAssistantExecutor(cfg AssistantExecutorConfig) fte.userID, fte.flowID, nil, nil, fte.replacer, fte.store, + fte.embedder, + fte.db, + fte.cfg.EmbeddingMaxTextBytes, fte.vslp, fte.knp, ) @@ -889,6 +897,9 @@ func (fte *flowToolsExecutor) GetAssistantExecutor(cfg AssistantExecutorConfig) fte.userID, fte.flowID, nil, nil, fte.replacer, fte.store, + fte.embedder, + fte.db, + fte.cfg.EmbeddingMaxTextBytes, fte.vslp, fte.knp, ) @@ -1174,6 +1185,9 @@ func (fte *flowToolsExecutor) GetInstallerExecutor(cfg InstallerExecutorConfig) cfg.SubtaskID, fte.replacer, fte.store, + fte.embedder, + fte.db, + fte.cfg.EmbeddingMaxTextBytes, fte.vslp, fte.knp, ) @@ -1279,6 +1293,9 @@ func (fte *flowToolsExecutor) GetCoderExecutor(cfg CoderExecutorConfig) (Context cfg.SubtaskID, fte.replacer, fte.store, + fte.embedder, + fte.db, + fte.cfg.EmbeddingMaxTextBytes, fte.vslp, fte.knp, ) @@ -1401,6 +1418,9 @@ func (fte *flowToolsExecutor) GetPentesterExecutor(cfg PentesterExecutorConfig) cfg.SubtaskID, fte.replacer, fte.store, + fte.embedder, + fte.db, + fte.cfg.EmbeddingMaxTextBytes, fte.vslp, fte.knp, ) @@ -1578,6 +1598,9 @@ func (fte *flowToolsExecutor) GetSearcherExecutor(cfg SearcherExecutorConfig) (C cfg.SubtaskID, fte.replacer, fte.store, + fte.embedder, + fte.db, + fte.cfg.EmbeddingMaxTextBytes, fte.vslp, fte.knp, ) diff --git a/backend/pkg/tools/vecstore_helper.go b/backend/pkg/tools/vecstore_helper.go new file mode 100644 index 00000000..a7c99616 --- /dev/null +++ b/backend/pkg/tools/vecstore_helper.go @@ -0,0 +1,78 @@ +package tools + +import ( + "context" + "database/sql" + "encoding/json" + "fmt" + "strconv" + "strings" + + "pentagi/pkg/database" + "pentagi/pkg/providers/embeddings" + + "github.com/google/uuid" +) + +// truncateForEmbedding truncates text to maxBytes bytes for sending to the +// embedding model API. Returns text unchanged if it is already within the limit. +// This avoids token-limit errors (e.g. OpenAI: max 8192 tokens) when a heavy +// field exceeds the model's context window. +func truncateForEmbedding(text string, maxBytes int) string { + if maxBytes <= 0 || len(text) <= maxBytes { + return text + } + return text[:maxBytes] +} + +// storeDocumentWithEmbeddingLimit computes the embedding vector from embeddingText +// (the truncated/prefix version of the content, ≤ maxBytes) and inserts the +// document into the vector store with the full original text as the stored content. +// +// Use this when len(fullText) > maxBytes to avoid embedding API token-limit errors +// while still persisting the complete document for retrieval. +// +// Returns the UUID string of the newly inserted document. +func storeDocumentWithEmbeddingLimit( + ctx context.Context, + db database.Querier, + embedder embeddings.Embedder, + embeddingText string, + fullText string, + metadata map[string]any, +) (string, error) { + vecs, err := embedder.EmbedDocuments(ctx, []string{embeddingText}) + if err != nil { + return "", fmt.Errorf("failed to compute embedding: %w", err) + } + if len(vecs) == 0 { + return "", fmt.Errorf("embedder returned no vectors") + } + + metaJSON, err := json.Marshal(metadata) + if err != nil { + return "", fmt.Errorf("failed to marshal metadata: %w", err) + } + + id := uuid.New() + docID, err := db.InsertKnowledgeDocument(ctx, database.InsertKnowledgeDocumentParams{ + Uuid: id, + Document: sql.NullString{String: fullText, Valid: true}, + Embedding: formatVectorFromFloat32s(vecs[0]), + Cmetadata: json.RawMessage(metaJSON), + }) + if err != nil { + return "", fmt.Errorf("failed to insert document with pre-computed embedding: %w", err) + } + + return docID, nil +} + +// formatVectorFromFloat32s converts a float32 slice into a PostgreSQL vector literal '[f1,f2,...]'. +func formatVectorFromFloat32s(v []float32) string { + strs := make([]string, len(v)) + for i, f := range v { + strs[i] = strconv.FormatFloat(float64(f), 'f', -1, 32) + } + return "[" + strings.Join(strs, ",") + "]" +} diff --git a/backend/sqlc/models/knowledge.sql b/backend/sqlc/models/knowledge.sql index a0812409..b06d18bc 100644 --- a/backend/sqlc/models/knowledge.sql +++ b/backend/sqlc/models/knowledge.sql @@ -138,3 +138,18 @@ DELETE FROM langchain_pg_embedding WHERE collection_id = (SELECT uuid FROM langchain_pg_collection WHERE name = 'langchain') AND COALESCE(cmetadata ->> 'doc_type', '') = 'memory' AND (cmetadata ->> 'flow_id') = sqlc.arg(flow_id); + +-- name: InsertKnowledgeDocument :one +-- Insert a document with a pre-computed embedding vector and return its UUID. +-- embedding must be formatted as a PostgreSQL vector literal: '[f1,f2,...]' +-- cmetadata must be valid JSON text. +INSERT INTO langchain_pg_embedding (uuid, document, embedding, cmetadata, collection_id) +SELECT + sqlc.arg(uuid)::uuid, + sqlc.arg(document), + sqlc.arg(embedding)::vector, + sqlc.arg(cmetadata)::json, + c.uuid +FROM langchain_pg_collection c +WHERE c.name = 'langchain' +RETURNING uuid::text AS id; diff --git a/docker-compose.yml b/docker-compose.yml index f706b166..20cc37d4 100644 --- a/docker-compose.yml +++ b/docker-compose.yml @@ -86,6 +86,7 @@ services: - EMBEDDING_MODEL=${EMBEDDING_MODEL:-} - EMBEDDING_PROVIDER=${EMBEDDING_PROVIDER:-} - EMBEDDING_BATCH_SIZE=${EMBEDDING_BATCH_SIZE:-} + - EMBEDDING_MAX_TEXT_BYTES=${EMBEDDING_MAX_TEXT_BYTES:-} - EMBEDDING_STRIP_NEW_LINES=${EMBEDDING_STRIP_NEW_LINES:-} - SUMMARIZER_PRESERVE_LAST=${SUMMARIZER_PRESERVE_LAST:-} - SUMMARIZER_USE_QA=${SUMMARIZER_USE_QA:-} @@ -131,6 +132,9 @@ services: - OAUTH_GITHUB_CLIENT_ID=${OAUTH_GITHUB_CLIENT_ID:-} - OAUTH_GITHUB_CLIENT_SECRET=${OAUTH_GITHUB_CLIENT_SECRET:-} - DATABASE_URL=postgres://${PENTAGI_POSTGRES_USER:-postgres}:${PENTAGI_POSTGRES_PASSWORD:-postgres}@pgvector:5432/${PENTAGI_POSTGRES_DB:-pentagidb}?sslmode=disable + - DATABASE_MAX_OPEN_CONNS=${DATABASE_MAX_OPEN_CONNS:-} + - DATABASE_MAX_IDLE_CONNS=${DATABASE_MAX_IDLE_CONNS:-} + - DATABASE_VECTOR_MAX_CONNS=${DATABASE_VECTOR_MAX_CONNS:-} - DUCKDUCKGO_ENABLED=${DUCKDUCKGO_ENABLED:-} - DUCKDUCKGO_REGION=${DUCKDUCKGO_REGION:-} - DUCKDUCKGO_SAFESEARCH=${DUCKDUCKGO_SAFESEARCH:-}